breakpilot-pwa/admin-v2/ai-compliance-sdk/internal/db/postgres.go

package db

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
)

// Pool wraps a pgxpool.Pool with SDK-specific methods
type Pool struct {
	*pgxpool.Pool
}

// SDKState represents the state stored in the database
type SDKState struct {
	ID        string          `json:"id"`
	TenantID  string          `json:"tenant_id"`
	UserID    string          `json:"user_id,omitempty"`
	State     json.RawMessage `json:"state"`
	Version   int             `json:"version"`
	CreatedAt time.Time       `json:"created_at"`
	UpdatedAt time.Time       `json:"updated_at"`
}

// NewPostgresPool creates a new database connection pool
func NewPostgresPool(connectionString string) (*Pool, error) {
	config, err := pgxpool.ParseConfig(connectionString)
	if err != nil {
		return nil, fmt.Errorf("failed to parse connection string: %w", err)
	}

	config.MaxConns = 10
	config.MinConns = 2
	config.MaxConnLifetime = 1 * time.Hour
	config.MaxConnIdleTime = 30 * time.Minute

	pool, err := pgxpool.NewWithConfig(context.Background(), config)
	if err != nil {
		return nil, fmt.Errorf("failed to create connection pool: %w", err)
	}

	// Test connection
	if err := pool.Ping(context.Background()); err != nil {
		return nil, fmt.Errorf("failed to ping database: %w", err)
	}

	return &Pool{Pool: pool}, nil
}

// GetState retrieves state for a tenant
func (p *Pool) GetState(ctx context.Context, tenantID string) (*SDKState, error) {
	query := `
		SELECT id, tenant_id, user_id, state, version, created_at, updated_at
		FROM sdk_states
		WHERE tenant_id = $1
	`

	var state SDKState
	err := p.QueryRow(ctx, query, tenantID).Scan(
		&state.ID,
		&state.TenantID,
		&state.UserID,
		&state.State,
		&state.Version,
		&state.CreatedAt,
		&state.UpdatedAt,
	)
	if err != nil {
		return nil, err
	}

	return &state, nil
}

// SaveState saves or updates state for a tenant with optimistic locking
func (p *Pool) SaveState(ctx context.Context, tenantID string, userID string, state json.RawMessage, expectedVersion *int) (*SDKState, error) {
	query := `
		INSERT INTO sdk_states (tenant_id, user_id, state, version)
		VALUES ($1, $2, $3, 1)
		ON CONFLICT (tenant_id) DO UPDATE SET
			state = $3,
			user_id = COALESCE($2, sdk_states.user_id),
			version = sdk_states.version + 1,
			updated_at = NOW()
		WHERE ($4::int IS NULL OR sdk_states.version = $4)
		RETURNING id, tenant_id, user_id, state, version, created_at, updated_at
	`

	var result SDKState
	err := p.QueryRow(ctx, query, tenantID, userID, state, expectedVersion).Scan(
		&result.ID,
		&result.TenantID,
		&result.UserID,
		&result.State,
		&result.Version,
		&result.CreatedAt,
		&result.UpdatedAt,
	)
	if err != nil {
		return nil, err
	}

	return &result, nil
}

// DeleteState deletes state for a tenant
func (p *Pool) DeleteState(ctx context.Context, tenantID string) error {
	query := `DELETE FROM sdk_states WHERE tenant_id = $1`
	_, err := p.Exec(ctx, query, tenantID)
	return err
}

// InMemoryStore provides an in-memory fallback when database is not available
type InMemoryStore struct {
	states map[string]*SDKState
}

// NewInMemoryStore creates a new in-memory store
func NewInMemoryStore() *InMemoryStore {
	return &InMemoryStore{
		states: make(map[string]*SDKState),
	}
}

// GetState retrieves state from memory
func (s *InMemoryStore) GetState(tenantID string) (*SDKState, error) {
	state, ok := s.states[tenantID]
	if !ok {
		return nil, fmt.Errorf("state not found")
	}
	return state, nil
}

// SaveState saves state to memory
func (s *InMemoryStore) SaveState(tenantID string, userID string, state json.RawMessage, expectedVersion *int) (*SDKState, error) {
	existing, exists := s.states[tenantID]

	// Optimistic locking check
	if expectedVersion != nil && exists && existing.Version != *expectedVersion {
		return nil, fmt.Errorf("version conflict")
	}

	now := time.Now()
	version := 1
	createdAt := now

	if exists {
		version = existing.Version + 1
		createdAt = existing.CreatedAt
	}

	newState := &SDKState{
		ID:        fmt.Sprintf("%s-%d", tenantID, time.Now().UnixNano()),
		TenantID:  tenantID,
		UserID:    userID,
		State:     state,
		Version:   version,
		CreatedAt: createdAt,
		UpdatedAt: now,
	}

	s.states[tenantID] = newState
	return newState, nil
}

// DeleteState deletes state from memory
func (s *InMemoryStore) DeleteState(tenantID string) error {
	delete(s.states, tenantID)
	return nil
}