claw-code/.guardrails/mcp-server/internal/database/production_code_tracking.go

package database

import (
	"context"
	"database/sql"
	"errors"
	"fmt"

	"github.com/thearchitectit/guardrail-mcp/internal/models"
)

// ProductionCodeStore handles production code tracking database operations
type ProductionCodeStore struct {
	db *DB
}

// NewProductionCodeStore creates a new production code store
func NewProductionCodeStore(db *DB) *ProductionCodeStore {
	return &ProductionCodeStore{db: db}
}

// CreateOrUpdate inserts or updates a production code record
func (s *ProductionCodeStore) CreateOrUpdate(ctx context.Context, pc *models.ProductionCode) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	if err := pc.Validate(); err != nil {
		return fmt.Errorf("invalid production code record: %w", err)
	}

	query := `
		INSERT INTO production_code_tracking (session_id, file_path, code_type)
		VALUES ($1, $2, $3)
		ON CONFLICT (session_id, file_path) DO UPDATE
		SET code_type = EXCLUDED.code_type
		RETURNING id, created_at
	`
	err = tx.QueryRowContext(ctx, query, pc.SessionID, pc.FilePath, pc.CodeType).Scan(&pc.ID, &pc.CreatedAt)
	if err != nil {
		return fmt.Errorf("failed to create production code record: %w", err)
	}

	if err := tx.Commit(); err != nil {
		return fmt.Errorf("failed to commit transaction: %w", err)
	}

	return nil
}

// MarkAsVerified marks production code as verified
func (s *ProductionCodeStore) MarkAsVerified(ctx context.Context, sessionID, filePath string) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	result, err := tx.ExecContext(ctx, `
		UPDATE production_code_tracking
		SET verified_at = NOW()
		WHERE session_id = $1 AND file_path = $2 AND code_type = 'production'
	`, sessionID, filePath)
	if err != nil {
		return fmt.Errorf("failed to mark as verified: %w", err)
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("failed to get rows affected: %w", err)
	}

	if rows == 0 {
		return sql.ErrNoRows
	}

	if err := tx.Commit(); err != nil {
		return fmt.Errorf("failed to commit transaction: %w", err)
	}

	return nil
}

// GetBySessionAndPath retrieves a production code record by session ID and file path
func (s *ProductionCodeStore) GetBySessionAndPath(ctx context.Context, sessionID, filePath string) (*models.ProductionCode, error) {
	var pc models.ProductionCode
	query := `
		SELECT id, session_id, file_path, code_type, created_at, verified_at
		FROM production_code_tracking
		WHERE session_id = $1 AND file_path = $2
	`
	err := s.db.QueryRowContext(ctx, query, sessionID, filePath).Scan(
		&pc.ID, &pc.SessionID, &pc.FilePath, &pc.CodeType, &pc.CreatedAt, &pc.VerifiedAt,
	)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, sql.ErrNoRows
		}
		return nil, fmt.Errorf("failed to get production code record: %w", err)
	}
	return &pc, nil
}

// HasProductionCode checks if production code exists in the session
func (s *ProductionCodeStore) HasProductionCode(ctx context.Context, sessionID string) (bool, error) {
	var exists bool
	query := `
		SELECT EXISTS(
			SELECT 1 FROM production_code_tracking
			WHERE session_id = $1 AND code_type = 'production'
		)
	`
	err := s.db.QueryRowContext(ctx, query, sessionID).Scan(&exists)
	if err != nil {
		return false, fmt.Errorf("failed to check production code existence: %w", err)
	}
	return exists, nil
}

// ListBySession retrieves all production code records for a session
func (s *ProductionCodeStore) ListBySession(ctx context.Context, sessionID string, limit, offset int) ([]models.ProductionCode, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT id, session_id, file_path, code_type, created_at, verified_at
		FROM production_code_tracking
		WHERE session_id = $1
		ORDER BY created_at DESC
		LIMIT $2 OFFSET $3
	`, sessionID, limit, offset)
	if err != nil {
		return nil, fmt.Errorf("failed to list production code records: %w", err)
	}
	defer rows.Close()

	return scanProductionCodeRecords(rows)
}

// DeleteBySession removes all production code records for a session
func (s *ProductionCodeStore) DeleteBySession(ctx context.Context, sessionID string) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	result, err := tx.ExecContext(ctx, `DELETE FROM production_code_tracking WHERE session_id = $1`, sessionID)
	if err != nil {
		return fmt.Errorf("failed to delete production code records: %w", err)
	}

	_, err = result.RowsAffected()
	if err != nil {
		return fmt.Errorf("failed to get rows affected: %w", err)
	}

	if err := tx.Commit(); err != nil {
		return fmt.Errorf("failed to commit transaction: %w", err)
	}

	return nil
}

// CountBySession returns the number of production code records for a session
func (s *ProductionCodeStore) CountBySession(ctx context.Context, sessionID string) (int, error) {
	var count int
	err := s.db.QueryRowContext(ctx, `
		SELECT COUNT(*) FROM production_code_tracking WHERE session_id = $1
	`, sessionID).Scan(&count)
	if err != nil {
		return 0, fmt.Errorf("failed to count production code records: %w", err)
	}
	return count, nil
}

// CountByType returns the count of production code records by type for a session
func (s *ProductionCodeStore) CountByType(ctx context.Context, sessionID string) (map[string]int, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT code_type, COUNT(*) as count
		FROM production_code_tracking
		WHERE session_id = $1
		GROUP BY code_type
	`, sessionID)
	if err != nil {
		return nil, fmt.Errorf("failed to count by type: %w", err)
	}
	defer rows.Close()

	counts := make(map[string]int)
	for rows.Next() {
		var codeType string
		var count int
		if err := rows.Scan(&codeType, &count); err != nil {
			return nil, fmt.Errorf("failed to scan count: %w", err)
		}
		counts[codeType] = count
	}

	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("error iterating rows: %w", err)
	}

	return counts, nil
}

// scanProductionCodeRecords is a helper function to scan production code rows
func scanProductionCodeRecords(rows *sql.Rows) ([]models.ProductionCode, error) {
	var records []models.ProductionCode
	for rows.Next() {
		var pc models.ProductionCode
		err := rows.Scan(
			&pc.ID, &pc.SessionID, &pc.FilePath, &pc.CodeType, &pc.CreatedAt, &pc.VerifiedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to scan production code record: %w", err)
		}
		records = append(records, pc)
	}

	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("error iterating production code records: %w", err)
	}

	return records, nil
}