Agent-System/internal/projectworkflow/store_pg_lineage.go

package projectworkflow

import (
	"context"
	"time"
)

func (s *pgStore) RecordLineage(ctx context.Context, lineage *DocumentLineage) error {
	query := `
		INSERT INTO document_lineage (
			project_id, downstream_doc_id, downstream_version, upstream_doc_id, upstream_version,
			upstream_sha256, is_stale, stale_reason, created_at
		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, now())
		ON CONFLICT (downstream_doc_id, downstream_version, upstream_doc_id, upstream_version)
		DO UPDATE SET upstream_sha256 = EXCLUDED.upstream_sha256,
		              is_stale = EXCLUDED.is_stale,
		              stale_reason = EXCLUDED.stale_reason
	`
	_, err := s.pool.Exec(
		ctx, query,
		lineage.ProjectID, lineage.DownstreamDocID, lineage.DownstreamVersion,
		lineage.UpstreamDocID, lineage.UpstreamVersion, lineage.UpstreamSHA256,
		lineage.IsStale, lineage.StaleReason,
	)
	return err
}

func (s *pgStore) GetLineageForDownstream(ctx context.Context, docID string, version int) ([]*DocumentLineage, error) {
	query := `
		SELECT id, project_id, downstream_doc_id, downstream_version, upstream_doc_id,
		       upstream_version, upstream_sha256, is_stale, stale_reason, stale_detected_at, created_at
		FROM document_lineage
		WHERE downstream_doc_id = $1 AND ($2 <= 0 OR downstream_version = $2)
		ORDER BY id DESC
	`
	rows, err := s.pool.Query(ctx, query, docID, version)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var list []*DocumentLineage
	for rows.Next() {
		l := &DocumentLineage{}
		if err := rows.Scan(
			&l.ID, &l.ProjectID, &l.DownstreamDocID, &l.DownstreamVersion, &l.UpstreamDocID,
			&l.UpstreamVersion, &l.UpstreamSHA256, &l.IsStale, &l.StaleReason, &l.StaleDetectedAt, &l.CreatedAt,
		); err != nil {
			return nil, err
		}
		list = append(list, l)
	}
	return list, rows.Err()
}

func (s *pgStore) GetLineageForProjectDownstream(ctx context.Context, projectID int64, docID string, version int) ([]*DocumentLineage, error) {
	query := `
		SELECT id, project_id, downstream_doc_id, downstream_version, upstream_doc_id,
		       upstream_version, upstream_sha256, is_stale, stale_reason, stale_detected_at, created_at
		FROM document_lineage
		WHERE project_id = $1 AND downstream_doc_id = $2 AND ($3 <= 0 OR downstream_version = $3)
		ORDER BY id DESC
	`
	rows, err := s.pool.Query(ctx, query, projectID, docID, version)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var list []*DocumentLineage
	for rows.Next() {
		l := &DocumentLineage{}
		if err := rows.Scan(
			&l.ID, &l.ProjectID, &l.DownstreamDocID, &l.DownstreamVersion, &l.UpstreamDocID,
			&l.UpstreamVersion, &l.UpstreamSHA256, &l.IsStale, &l.StaleReason, &l.StaleDetectedAt, &l.CreatedAt,
		); err != nil {
			return nil, err
		}
		list = append(list, l)
	}
	return list, rows.Err()
}

func (s *pgStore) GetLineagesForUpstream(ctx context.Context, upstreamDocID string) ([]*DocumentLineage, error) {
	query := `
		SELECT id, project_id, downstream_doc_id, downstream_version, upstream_doc_id,
		       upstream_version, upstream_sha256, is_stale, stale_reason, stale_detected_at, created_at
		FROM document_lineage
		WHERE upstream_doc_id = $1
		ORDER BY id DESC
	`
	rows, err := s.pool.Query(ctx, query, upstreamDocID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var list []*DocumentLineage
	for rows.Next() {
		l := &DocumentLineage{}
		if err := rows.Scan(
			&l.ID, &l.ProjectID, &l.DownstreamDocID, &l.DownstreamVersion, &l.UpstreamDocID,
			&l.UpstreamVersion, &l.UpstreamSHA256, &l.IsStale, &l.StaleReason, &l.StaleDetectedAt, &l.CreatedAt,
		); err != nil {
			return nil, err
		}
		list = append(list, l)
	}
	return list, rows.Err()
}

func (s *pgStore) MarkDownstreamStale(ctx context.Context, upstreamDocID string, staleReason string) ([]*DocumentLineage, error) {
	now := time.Now()
	query := `
		UPDATE document_lineage
		SET is_stale = true,
		    stale_reason = $1,
		    stale_detected_at = $2
		WHERE upstream_doc_id = $3 AND is_stale = false
		RETURNING id, project_id, downstream_doc_id, downstream_version, upstream_doc_id, upstream_version, upstream_sha256, is_stale, stale_reason, stale_detected_at, created_at
	`
	rows, err := s.pool.Query(ctx, query, staleReason, now, upstreamDocID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var updated []*DocumentLineage
	for rows.Next() {
		l := &DocumentLineage{}
		if err := rows.Scan(
			&l.ID, &l.ProjectID, &l.DownstreamDocID, &l.DownstreamVersion, &l.UpstreamDocID,
			&l.UpstreamVersion, &l.UpstreamSHA256, &l.IsStale, &l.StaleReason, &l.StaleDetectedAt, &l.CreatedAt,
		); err != nil {
			return nil, err
		}
		updated = append(updated, l)
	}
	return updated, rows.Err()
}