DocumentAIKit/FastExtract/go/dao/task_repository.go

package dao

import (
	"Zythum/common"
	"database/sql"
	_ "github.com/go-sql-driver/mysql"
	"log"
)

func GetNewTask() (common.Task, error) {
	tx, err := db.Begin()
	committed := false
	if err != nil {
		log.Println("Failed to begin transaction:", err)
		return common.Task{ID: -1, Path: ""}, err
	}
	defer func(tx *sql.Tx) {
		if !committed {
			err := tx.Rollback()
			if err != nil {
				log.Println("Failed to rollback transaction:", err)
			}
		}
	}(tx)

	rows, err := tx.Query("SELECT id, path FROM financial_reports WHERE status = 0 ORDER BY id LIMIT 1 FOR UPDATE;")
	if err != nil {
		return common.Task{ID: -1, Path: ""}, err
	}

	var task common.Task
	rows.Next()
	err = rows.Scan(&task.ID, &task.Path)
	if err != nil {
		return common.Task{ID: -2, Path: ""}, err
	}

	err = rows.Close()
	if err != nil {
		log.Println("Failed to close rows:", err)
	}

	_, err = tx.Exec("UPDATE financial_reports SET status = 1 WHERE id = ?", task.ID)
	if err != nil {
		return common.Task{ID: -3, Path: ""}, err
	}

	if err := tx.Commit(); err != nil {
		return common.Task{ID: -4, Path: ""}, err
	} else {
		committed = true
	}

	return task, nil
}

func UpdateTaskStatus(id int, status int) int {
	_, err := db.Exec("UPDATE financial_reports SET status = ? WHERE id = ?;", status, id)
	if err != nil {
		log.Println("Database update error:", err)
		return -1
	}
	return 0
}

func GetTaskStatus() common.TaskStatus {
	var taskStatus common.TaskStatus
	err := db.QueryRow(
		"SELECT (SELECT COUNT(id) FROM financial_reports WHERE status = 0) AS ready_task_count, "+
			"(SELECT COUNT(id) FROM financial_reports WHERE status = 1) AS pending_task_count, "+
			"(SELECT COUNT(id) FROM financial_reports WHERE status = 2) AS completed_task_count, "+
			"(SELECT COUNT(id) FROM financial_reports WHERE status = 3) AS failed_task_count;").Scan(&taskStatus.ReadyTaskCount, &taskStatus.PendingTaskCount, &taskStatus.CompletedTaskCount, &taskStatus.FailedTaskCount)

	if err != nil {
		log.Println("Database query error:", err)
		return common.TaskStatus{}
	}

	return taskStatus
}

func getTasksByStatus(status int) {
	rows, err := db.Query("SELECT id, path FROM financial_reports WHERE status = ?;", status)
	if err != nil {
		log.Println("Database query error:", err)
		return
	}

	for rows.Next() {
		var task common.Task
		err := rows.Scan(&task.ID, &task.Path)
		if err != nil {
			log.Println("Database scan error:", err)
			return
		}
		log.Println("Task:", task)
	}
	defer rows.Close()
}