91 lines
2.2 KiB
Go
91 lines
2.2 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"fmt"
|
|
|
|
"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/backend/internal/models"
|
|
"gitea.d3m0k1d.ru/d3m0k1d/HellreigN/backend/internal/storage"
|
|
)
|
|
|
|
type JobRepository struct {
|
|
DB *sql.DB
|
|
}
|
|
|
|
func NewJobRepository(db *sql.DB) *JobRepository {
|
|
return &JobRepository{DB: db}
|
|
}
|
|
|
|
func (r *JobRepository) Init(ctx context.Context) error {
|
|
_, err := r.DB.ExecContext(ctx, storage.CreateJobsTable)
|
|
return err
|
|
}
|
|
|
|
func (r *JobRepository) InitJob(ctx context.Context, agentID string, job models.JobForInsert) (int64, error) {
|
|
commandJSON, err := json.Marshal(job.Command)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("marshal command: %w", err)
|
|
}
|
|
|
|
var stdinVal *string
|
|
if job.Stdin != nil {
|
|
stdinVal = job.Stdin
|
|
}
|
|
|
|
result, err := r.DB.ExecContext(ctx,
|
|
`INSERT INTO jobs (agent_id, command, stdin, stdout, stderr, status) VALUES (?, ?, ?, '', '', 0)`,
|
|
agentID, string(commandJSON), stdinVal,
|
|
)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return result.LastInsertId()
|
|
}
|
|
|
|
func (r *JobRepository) UpdateJobInDB(ctx context.Context, jid int64, msg models.JobForUpdate) (models.Job, error) {
|
|
result, err := r.DB.ExecContext(ctx,
|
|
`UPDATE jobs SET stdout = ?, stderr = ?, status = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?`,
|
|
msg.Stdout, msg.Stderr, msg.Status, jid,
|
|
)
|
|
if err != nil {
|
|
return models.Job{}, err
|
|
}
|
|
|
|
affected, err := result.RowsAffected()
|
|
if err != nil {
|
|
return models.Job{}, err
|
|
}
|
|
if affected == 0 {
|
|
return models.Job{}, ErrNotFound
|
|
}
|
|
|
|
return r.GetJobByID(ctx, jid)
|
|
}
|
|
|
|
func (r *JobRepository) GetJobByID(ctx context.Context, jid int64) (models.Job, error) {
|
|
var job models.Job
|
|
var commandJSON string
|
|
var stdinVal *string
|
|
|
|
err := r.DB.QueryRowContext(ctx,
|
|
`SELECT id, command, stdin, stdout, stderr, status FROM jobs WHERE id = ?`,
|
|
jid,
|
|
).Scan(&job.ID, &commandJSON, &stdinVal, &job.Stdout, &job.Stderr, &job.Status)
|
|
if err != nil {
|
|
if err == sql.ErrNoRows {
|
|
return models.Job{}, ErrNotFound
|
|
}
|
|
return models.Job{}, err
|
|
}
|
|
|
|
if err := json.Unmarshal([]byte(commandJSON), &job.JobForInsert.Command); err != nil {
|
|
return models.Job{}, fmt.Errorf("unmarshal command: %w", err)
|
|
}
|
|
|
|
job.JobForInsert.Stdin = stdinVal
|
|
return job, nil
|
|
}
|