diff --git a/internal/handlers/dto/province.go b/internal/handlers/dto/province.go index 8187640..b9af877 100644 --- a/internal/handlers/dto/province.go +++ b/internal/handlers/dto/province.go @@ -36,3 +36,7 @@ type ImportProvinceResponse struct { IsAsync bool `json:"is_async"` Message string `json:"message,omitempty"` } + +type StopImportRequest struct { + JobID string `json:"job_id" binding:"required"` +} diff --git a/internal/handlers/router/city/city.go b/internal/handlers/router/city/city.go index 8779ac4..51ff6eb 100644 --- a/internal/handlers/router/city/city.go +++ b/internal/handlers/router/city/city.go @@ -36,6 +36,7 @@ func (h *CityHandler) Router() { r.DELETE("/:id", h.Delete) r.GET("/template", h.Template) r.POST("/import", h.Import) + r.POST("/import/stop", h.ImportStop) r.GET("/import/status/:jobId", h.ImportStatus) } } @@ -254,3 +255,27 @@ func (h *CityHandler) ImportStatus(c *gin.Context) { "logs": logs, }) } + +func (h *CityHandler) ImportStop(c *gin.Context) { + userid, err := h.Helper.GetAuthInfo(c.GetString("auth")) + if err != nil { + response.Error(c, http.StatusUnauthorized, err.Error()) + return + } + + var req dto.StopImportRequest + if err := c.BindJSON(&req); err != nil { + response.Error(c, http.StatusBadRequest, "job_id is required") + return + } + + handler := cityimporter.NewCityImportHandler(*userid, h.Helper) + imp := importer.NewImporter(h.Helper, handler) + + if err := imp.StopJob(req.JobID); err != nil { + response.Error(c, http.StatusBadRequest, err.Error()) + return + } + + response.Success(c, map[string]interface{}{"message": "Job stopped"}) +} diff --git a/internal/handlers/router/province/province.go b/internal/handlers/router/province/province.go index 3926280..dc1936b 100644 --- a/internal/handlers/router/province/province.go +++ b/internal/handlers/router/province/province.go @@ -36,6 +36,7 @@ func (h *ProvinceHandler) Router() { r.DELETE("/:id", h.Delete) r.GET("/template", h.Template) r.POST("/import", h.Import) + r.POST("/import/stop", h.ImportStop) r.GET("/import/status/:jobId", h.ImportStatus) } } @@ -254,3 +255,27 @@ func (h *ProvinceHandler) ImportStatus(c *gin.Context) { "logs": logs, }) } + +func (h *ProvinceHandler) ImportStop(c *gin.Context) { + userid, err := h.Helper.GetAuthInfo(c.GetString("auth")) + if err != nil { + response.Error(c, http.StatusUnauthorized, err.Error()) + return + } + + var req dto.StopImportRequest + if err := c.BindJSON(&req); err != nil { + response.Error(c, http.StatusBadRequest, "job_id is required") + return + } + + handler := provinceimporter.NewProvinceImportHandler(*userid) + imp := importer.NewImporter(h.Helper, handler) + + if err := imp.StopJob(req.JobID); err != nil { + response.Error(c, http.StatusBadRequest, err.Error()) + return + } + + response.Success(c, map[string]interface{}{"message": "Job stopped"}) +} diff --git a/pkg/csv/csv.go b/pkg/csv/csv.go index 9ced353..16c78e9 100644 --- a/pkg/csv/csv.go +++ b/pkg/csv/csv.go @@ -1,6 +1,7 @@ package csv import ( + "context" "encoding/csv" "io" "strings" @@ -9,6 +10,10 @@ import ( type RowHandler func(row []string, rowNumber int) error func ReadCSV(file io.Reader, handler RowHandler) (int, []string, error) { + return ReadCSVWithContext(context.Background(), file, handler) +} + +func ReadCSVWithContext(ctx context.Context, file io.Reader, handler RowHandler) (int, []string, error) { reader := csv.NewReader(file) reader.LazyQuotes = true reader.TrimLeadingSpace = true @@ -18,6 +23,12 @@ func ReadCSV(file io.Reader, handler RowHandler) (int, []string, error) { totalRows := 0 for { + select { + case <-ctx.Done(): + return totalRows, errors, ctx.Err() + default: + } + row, err := reader.Read() if err == io.EOF { break @@ -29,17 +40,18 @@ func ReadCSV(file io.Reader, handler RowHandler) (int, []string, error) { continue } - // Skip empty rows if len(row) == 0 || (len(row) == 1 && strings.TrimSpace(row[0]) == "") { continue } - // Skip header row (first row) if rowNumber == 1 { continue } if err := handler(row, rowNumber); err != nil { + if ctx.Err() != nil { + return totalRows, errors, ctx.Err() + } errors = append(errors, "Row "+itoa(rowNumber)+": "+err.Error()) } totalRows++ diff --git a/pkg/importer/importer.go b/pkg/importer/importer.go index 2bc6aec..5d54933 100644 --- a/pkg/importer/importer.go +++ b/pkg/importer/importer.go @@ -12,12 +12,18 @@ import ( "io" "net/http" "strings" + "sync" "time" "github.com/google/uuid" "gorm.io/gorm" ) +var ( + cancelFuncs = make(map[string]context.CancelFunc) + cancelMu sync.RWMutex +) + type RowData interface { TableName() string } @@ -231,7 +237,11 @@ func (i *Importer) processSync(file io.Reader, userID string, totalRows int) (*I func (i *Importer) processAsync(file io.Reader, userID string, totalRows int) (*ImportResult, error) { jobID := uuid.New().String() - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + + cancelMu.Lock() + cancelFuncs[jobID] = cancel + cancelMu.Unlock() metadata := map[string]string{ "user_id": userID, @@ -239,10 +249,14 @@ func (i *Importer) processAsync(file io.Reader, userID string, totalRows int) (* } if err := i.JobQ.CreateJob(ctx, jobID, totalRows, metadata); err != nil { + cancel() + cancelMu.Lock() + delete(cancelFuncs, jobID) + cancelMu.Unlock() return nil, fmt.Errorf("failed to create job: %w", err) } - go i.processAsyncWorker(file, jobID, userID, totalRows) + go i.processAsyncWorker(ctx, file, jobID, userID, totalRows) return &ImportResult{ JobID: jobID, @@ -252,9 +266,8 @@ func (i *Importer) processAsync(file io.Reader, userID string, totalRows int) (* }, nil } -func (i *Importer) processAsyncWorker(file io.Reader, jobID string, userID string, totalRows int) { +func (i *Importer) processAsyncWorker(ctx context.Context, file io.Reader, jobID string, userID string, totalRows int) { db := i.Helper.GetDB("master") - ctx := context.Background() var imported, skipped int var logs []ImportLog @@ -263,9 +276,15 @@ func (i *Importer) processAsyncWorker(file io.Reader, jobID string, userID strin var batchRows []int batchSize := 500 + defer func() { + cancelMu.Lock() + delete(cancelFuncs, jobID) + cancelMu.Unlock() + }() + i.JobQ.UpdateProgress(ctx, jobID, 0, 0, nil) - csvlib.ReadCSV(file, func(row []string, rowNumber int) error { + csvlib.ReadCSVWithContext(ctx, file, func(row []string, rowNumber int) error { data, rowErrors := i.Handler.ValidateRow(row, rowNumber) if len(rowErrors) > 0 { for _, e := range rowErrors { @@ -350,10 +369,27 @@ func (i *Importer) processAsyncWorker(file io.Reader, jobID string, userID strin } } - i.JobQ.CompleteJob(ctx, jobID, imported, skipped, errors) + if ctx.Err() != nil { + i.JobQ.StopJob(ctx, jobID) + } else { + i.JobQ.CompleteJob(ctx, jobID, imported, skipped, errors) + } i.saveJobLogs(ctx, jobID, logs, imported, skipped, errors) } +func (i *Importer) StopJob(jobID string) error { + cancelMu.RLock() + cancel, exists := cancelFuncs[jobID] + cancelMu.RUnlock() + + if !exists { + return fmt.Errorf("job not found or already completed") + } + + cancel() + return nil +} + func (i *Importer) saveJobLogs(ctx context.Context, jobID string, logs []ImportLog, imported int, skipped int, errors []string) { logsJSON, _ := json.Marshal(logs) key := fmt.Sprintf("import_job:%s:logs", jobID) diff --git a/pkg/jobqueue/jobqueue.go b/pkg/jobqueue/jobqueue.go index 161d1fd..c80edf0 100644 --- a/pkg/jobqueue/jobqueue.go +++ b/pkg/jobqueue/jobqueue.go @@ -15,6 +15,7 @@ const ( StatusProcessing JobStatus = "processing" StatusCompleted JobStatus = "completed" StatusFailed JobStatus = "failed" + StatusStopped JobStatus = "stopped" ) type Job struct { @@ -124,6 +125,16 @@ func (jq *JobQueue) FailJob(ctx context.Context, jobID string, errMsg string) er }).Err() } +func (jq *JobQueue) StopJob(ctx context.Context, jobID string) error { + key := fmt.Sprintf("import_job:%s", jobID) + now := time.Now() + + return jq.Redis.HSet(ctx, key, map[string]interface{}{ + "status": string(StatusStopped), + "finished_at": now.Format(time.RFC3339), + }).Err() +} + func (jq *JobQueue) GetJob(ctx context.Context, jobID string) (*Job, error) { key := fmt.Sprintf("import_job:%s", jobID) data, err := jq.Redis.HGetAll(ctx, key).Result()