Files

536 lines
13 KiB
Go
Raw Permalink Normal View History

package db
import (
"context"
"database/sql"
"encoding/json"
"fmt"
"os"
"time"
_ "github.com/lib/pq"
)
// DB wraps the database connection
type DB struct {
conn *sql.DB
}
// New creates a new database connection to memory-db (K8s CNPG)
// Expected DSN format: postgresql://app:[email protected]:5432/memory?sslmode=disable
func New(dsn string) (*DB, error) {
if dsn == "" {
// Fallback: try to construct from K8s env vars
host := os.Getenv("DATABASE_HOST")
port := os.Getenv("DATABASE_PORT")
name := os.Getenv("DATABASE_NAME")
user := os.Getenv("DATABASE_USER")
password := os.Getenv("DATABASE_PASSWORD")
if host != "" && port != "" && name != "" && user != "" && password != "" {
dsn = fmt.Sprintf("postgresql://%s:%s@%s:%s/%s?sslmode=disable",
user, password, host, port, name)
} else {
return nil, fmt.Errorf("DATABASE_URL or K8s env vars (DATABASE_HOST, DATABASE_PORT, DATABASE_NAME, DATABASE_USER, DATABASE_PASSWORD) required")
}
}
conn, err := sql.Open("postgres", dsn)
if err != nil {
return nil, fmt.Errorf("failed to open database: %w", err)
}
// Test connection
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := conn.PingContext(ctx); err != nil {
return nil, fmt.Errorf("failed to ping database: %w", err)
}
// Set connection pool settings
conn.SetMaxOpenConns(25)
conn.SetMaxIdleConns(5)
conn.SetConnMaxLifetime(5 * time.Minute)
return &DB{conn: conn}, nil
}
// Close closes the database connection
func (db *DB) Close() error {
return db.conn.Close()
}
// SaveWorkflow saves or updates a workflow with canvas
func (db *DB) SaveWorkflow(ctx context.Context, wf *Workflow) error {
query := `
INSERT INTO workflows (id, customer_id, name, description, status, version, nodes, edges, created_by, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
ON CONFLICT (id) DO UPDATE SET
name = $3,
description = $4,
status = $5,
version = $6,
nodes = $7,
edges = $8,
updated_at = $11
`
_, err := db.conn.ExecContext(ctx, query,
wf.ID,
wf.CustomerID,
wf.Name,
wf.Description,
wf.Status,
wf.Version,
wf.Nodes,
wf.Edges,
wf.CreatedBy,
wf.CreatedAt,
wf.UpdatedAt,
)
return err
}
// SaveCanvasUpdate saves canvas (nodes + edges) for a workflow
func (db *DB) SaveCanvasUpdate(ctx context.Context, workflowID, customerID string, canvas *Canvas) error {
nodesJSON, err := json.Marshal(canvas.Nodes)
if err != nil {
return fmt.Errorf("failed to marshal nodes: %w", err)
}
edgesJSON, err := json.Marshal(canvas.Edges)
if err != nil {
return fmt.Errorf("failed to marshal edges: %w", err)
}
query := `
UPDATE workflows
SET nodes = $1, edges = $2, updated_at = now()
WHERE id = $3 AND customer_id = $4
`
result, err := db.conn.ExecContext(ctx, query, nodesJSON, edgesJSON, workflowID, customerID)
if err != nil {
return fmt.Errorf("failed to update canvas: %w", err)
}
rows, err := result.RowsAffected()
if err != nil {
return err
}
if rows == 0 {
return fmt.Errorf("workflow not found: %s", workflowID)
}
return nil
}
// FetchWorkflow retrieves a workflow by ID
func (db *DB) FetchWorkflow(ctx context.Context, workflowID, customerID string) (*Workflow, error) {
query := `
SELECT id, customer_id, name, description, status, version, nodes, edges, created_by, created_at, updated_at, last_executed_at
FROM workflows
WHERE id = $1 AND customer_id = $2
`
wf := &Workflow{}
err := db.conn.QueryRowContext(ctx, query, workflowID, customerID).Scan(
&wf.ID,
&wf.CustomerID,
&wf.Name,
&wf.Description,
&wf.Status,
&wf.Version,
&wf.Nodes,
&wf.Edges,
&wf.CreatedBy,
&wf.CreatedAt,
&wf.UpdatedAt,
&wf.LastExecutedAt,
)
if err != nil {
if err == sql.ErrNoRows {
return nil, fmt.Errorf("workflow not found: %s", workflowID)
}
return nil, fmt.Errorf("failed to fetch workflow: %w", err)
}
return wf, nil
}
// FetchCanvas retrieves canvas (nodes + edges) for a workflow
func (db *DB) FetchCanvas(ctx context.Context, workflowID, customerID string) (*Canvas, error) {
query := `
SELECT nodes, edges
FROM workflows
WHERE id = $1 AND customer_id = $2
`
var nodesJSON, edgesJSON []byte
err := db.conn.QueryRowContext(ctx, query, workflowID, customerID).Scan(&nodesJSON, &edgesJSON)
if err != nil {
if err == sql.ErrNoRows {
return nil, fmt.Errorf("workflow not found: %s", workflowID)
}
return nil, fmt.Errorf("failed to fetch canvas: %w", err)
}
var nodes []WorkflowNode
var edges []WorkflowEdge
if err := json.Unmarshal(nodesJSON, &nodes); err != nil {
return nil, fmt.Errorf("failed to unmarshal nodes: %w", err)
}
if err := json.Unmarshal(edgesJSON, &edges); err != nil {
return nil, fmt.Errorf("failed to unmarshal edges: %w", err)
}
return &Canvas{Nodes: nodes, Edges: edges}, nil
}
// ListWorkflows retrieves all workflows for a customer
func (db *DB) ListWorkflows(ctx context.Context, customerID string, limit, offset int) ([]Workflow, error) {
query := `
SELECT id, customer_id, name, description, status, version, nodes, edges, created_by, created_at, updated_at, last_executed_at
FROM workflows
WHERE customer_id = $1
ORDER BY created_at DESC
LIMIT $2 OFFSET $3
`
rows, err := db.conn.QueryContext(ctx, query, customerID, limit, offset)
if err != nil {
return nil, fmt.Errorf("failed to list workflows: %w", err)
}
defer rows.Close()
var workflows []Workflow
for rows.Next() {
wf := Workflow{}
err := rows.Scan(
&wf.ID,
&wf.CustomerID,
&wf.Name,
&wf.Description,
&wf.Status,
&wf.Version,
&wf.Nodes,
&wf.Edges,
&wf.CreatedBy,
&wf.CreatedAt,
&wf.UpdatedAt,
&wf.LastExecutedAt,
)
if err != nil {
return nil, fmt.Errorf("failed to scan workflow: %w", err)
}
workflows = append(workflows, wf)
}
return workflows, rows.Err()
}
// DeleteWorkflow deletes a workflow
func (db *DB) DeleteWorkflow(ctx context.Context, workflowID, customerID string) error {
query := `
DELETE FROM workflows
WHERE id = $1 AND customer_id = $2
`
result, err := db.conn.ExecContext(ctx, query, workflowID, customerID)
if err != nil {
return fmt.Errorf("failed to delete workflow: %w", err)
}
rows, err := result.RowsAffected()
if err != nil {
return err
}
if rows == 0 {
return fmt.Errorf("workflow not found: %s", workflowID)
}
return nil
}
// SaveExecution saves a workflow execution record
func (db *DB) SaveExecution(ctx context.Context, exec *WorkflowExecution) error {
query := `
INSERT INTO workflow_executions (id, workflow_id, customer_id, temporal_id, status, inputs, outputs, started_at, completed_at, duration_ms, error_message, error_count)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
ON CONFLICT (id) DO UPDATE SET
status = $5,
outputs = $7,
completed_at = $9,
duration_ms = $10,
error_message = $11,
error_count = $12
`
_, err := db.conn.ExecContext(ctx, query,
exec.ID,
exec.WorkflowID,
exec.CustomerID,
exec.TemporalID,
exec.Status,
exec.Inputs,
exec.Outputs,
exec.StartedAt,
exec.CompletedAt,
exec.DurationMs,
exec.ErrorMessage,
exec.ErrorCount,
)
return err
}
// FetchExecution retrieves a workflow execution
func (db *DB) FetchExecution(ctx context.Context, executionID string) (*WorkflowExecution, error) {
query := `
SELECT id, workflow_id, customer_id, temporal_id, status, inputs, outputs, started_at, completed_at, duration_ms, error_message, error_count
FROM workflow_executions
WHERE id = $1
`
exec := &WorkflowExecution{}
err := db.conn.QueryRowContext(ctx, query, executionID).Scan(
&exec.ID,
&exec.WorkflowID,
&exec.CustomerID,
&exec.TemporalID,
&exec.Status,
&exec.Inputs,
&exec.Outputs,
&exec.StartedAt,
&exec.CompletedAt,
&exec.DurationMs,
&exec.ErrorMessage,
&exec.ErrorCount,
)
if err != nil {
if err == sql.ErrNoRows {
return nil, fmt.Errorf("execution not found: %s", executionID)
}
return nil, fmt.Errorf("failed to fetch execution: %w", err)
}
return exec, nil
}
// SaveExecutionLog saves an activity log entry
func (db *DB) SaveExecutionLog(ctx context.Context, log *ExecutionLog) error {
query := `
INSERT INTO execution_logs (execution_id, node_id, activity_name, level, message, metadata, logged_at)
VALUES ($1, $2, $3, $4, $5, $6, $7)
`
_, err := db.conn.ExecContext(ctx, query,
log.ExecutionID,
log.NodeID,
log.ActivityName,
log.Level,
log.Message,
log.Metadata,
log.LoggedAt,
)
return err
}
// FetchExecutionLogs retrieves all logs for an execution
func (db *DB) FetchExecutionLogs(ctx context.Context, executionID string) ([]ExecutionLog, error) {
query := `
SELECT id, execution_id, node_id, activity_name, level, message, metadata, logged_at
FROM execution_logs
WHERE execution_id = $1
ORDER BY logged_at ASC
`
rows, err := db.conn.QueryContext(ctx, query, executionID)
if err != nil {
return nil, fmt.Errorf("failed to fetch execution logs: %w", err)
}
defer rows.Close()
var logs []ExecutionLog
for rows.Next() {
log := ExecutionLog{}
err := rows.Scan(
&log.ID,
&log.ExecutionID,
&log.NodeID,
&log.ActivityName,
&log.Level,
&log.Message,
&log.Metadata,
&log.LoggedAt,
)
if err != nil {
return nil, fmt.Errorf("failed to scan log: %w", err)
}
logs = append(logs, log)
}
return logs, rows.Err()
}
// SaveActivityTrace saves per-activity execution trace
func (db *DB) SaveActivityTrace(ctx context.Context, trace *ActivityTrace) error {
query := `
INSERT INTO activity_traces (execution_id, node_id, activity_name, parameters, result, started_at, completed_at, duration_ms, attempt, retry_reason, status, error_message)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
ON CONFLICT (id) DO UPDATE SET
status = $11,
result = $5,
completed_at = $7,
duration_ms = $8,
error_message = $12
`
_, err := db.conn.ExecContext(ctx, query,
trace.ExecutionID,
trace.NodeID,
trace.ActivityName,
trace.Parameters,
trace.Result,
trace.StartedAt,
trace.CompletedAt,
trace.DurationMs,
trace.Attempt,
trace.RetryReason,
trace.Status,
trace.ErrorMessage,
)
return err
}
// FetchActivityTraces retrieves all activity traces for an execution
func (db *DB) FetchActivityTraces(ctx context.Context, executionID string) ([]ActivityTrace, error) {
query := `
SELECT id, execution_id, node_id, activity_name, parameters, result, started_at, completed_at, duration_ms, attempt, retry_reason, status, error_message
FROM activity_traces
WHERE execution_id = $1
ORDER BY started_at ASC
`
rows, err := db.conn.QueryContext(ctx, query, executionID)
if err != nil {
return nil, fmt.Errorf("failed to fetch activity traces: %w", err)
}
defer rows.Close()
var traces []ActivityTrace
for rows.Next() {
trace := ActivityTrace{}
err := rows.Scan(
&trace.ID,
&trace.ExecutionID,
&trace.NodeID,
&trace.ActivityName,
&trace.Parameters,
&trace.Result,
&trace.StartedAt,
&trace.CompletedAt,
&trace.DurationMs,
&trace.Attempt,
&trace.RetryReason,
&trace.Status,
&trace.ErrorMessage,
)
if err != nil {
return nil, fmt.Errorf("failed to scan trace: %w", err)
}
traces = append(traces, trace)
}
return traces, rows.Err()
}
// GetWorkflowRelations retrieves all relations for a workflow version
func (db *DB) GetWorkflowRelations(ctx context.Context, workflowID string, version int) ([]WorkflowRelation, error) {
var relations []WorkflowRelation
query := `
SELECT id, workflow_id, version, source_node_id, target_node_id,
relation_type, label, relation_wording, metadata, created_at
FROM workflow_relations
WHERE workflow_id = $1 AND version = $2
ORDER BY created_at DESC
`
rows, err := db.conn.QueryContext(ctx, query, workflowID, version)
if err != nil {
return nil, fmt.Errorf("failed to query relations: %w", err)
}
defer rows.Close()
for rows.Next() {
var rel WorkflowRelation
if err := rows.Scan(
&rel.ID,
&rel.WorkflowID,
&rel.Version,
&rel.SourceNodeID,
&rel.TargetNodeID,
&rel.RelationType,
&rel.Label,
&rel.RelationWording,
&rel.Metadata,
&rel.CreatedAt,
); err != nil {
return nil, fmt.Errorf("failed to scan relation: %w", err)
}
relations = append(relations, rel)
}
return relations, rows.Err()
}
// GetRelationVersions retrieves version history for a specific relation
func (db *DB) GetRelationVersions(ctx context.Context, workflowID string, edgeID string) ([]WorkflowRelationVersion, error) {
var versions []WorkflowRelationVersion
query := `
SELECT id, workflow_id, edge_id, version_num, operation, snapshot,
changed_at, changed_by, fields_changed
FROM workflow_relation_versions
WHERE workflow_id = $1 AND edge_id = $2
ORDER BY version_num ASC
`
rows, err := db.conn.QueryContext(ctx, query, workflowID, edgeID)
if err != nil {
return nil, fmt.Errorf("failed to query relation versions: %w", err)
}
defer rows.Close()
for rows.Next() {
var v WorkflowRelationVersion
if err := rows.Scan(
&v.ID,
&v.WorkflowID,
&v.EdgeID,
&v.VersionNum,
&v.Operation,
&v.Snapshot,
&v.ChangedAt,
&v.ChangedBy,
&v.FieldsChanged,
); err != nil {
return nil, fmt.Errorf("failed to scan version: %w", err)
}
versions = append(versions, v)
}
return versions, rows.Err()
}