knowledge-governance/internal/releaseadapter/v071/postgres_export.go

package v071

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"regexp"
	"sort"

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

var resourceReferencePattern = regexp.MustCompile(`resource://([A-Za-z0-9_-]{22})`)

type PostgresIdentity struct {
	TenantID        uint64
	KnowledgeBaseID string
	KnowledgeID     string
}

type TableSnapshot struct {
	Name     string
	Rows     []json.RawMessage
	RowCount int
	SHA256   string
}

type PostgresSnapshot struct {
	Tables []TableSnapshot
}

func (snapshot PostgresSnapshot) Table(name string) TableSnapshot {
	for _, table := range snapshot.Tables {
		if table.Name == name {
			return table
		}
	}
	return TableSnapshot{Name: name}
}

type PostgresExporter struct {
	db postgresQueryDB
}

type postgresQueryDB interface {
	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
	Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error)
}

func NewPostgresExporter(db postgresQueryDB) (*PostgresExporter, error) {
	if db == nil {
		return nil, errors.New("PostgreSQL export database is required")
	}
	return &PostgresExporter{db: db}, nil
}

const exportKnowledgeSQL = `
SELECT to_jsonb(knowledge)
FROM knowledges AS knowledge
WHERE knowledge.id = $1
  AND knowledge.tenant_id = $2
  AND knowledge.knowledge_base_id = $3
  AND knowledge.deleted_at IS NULL
  AND knowledge.parse_status = 'completed'
`

const exportChunksSQL = `
SELECT to_jsonb(chunk)
FROM chunks AS chunk
WHERE chunk.knowledge_id = $1
  AND chunk.tenant_id = $2
  AND chunk.deleted_at IS NULL
ORDER BY chunk.id
`

const exportTagRelationsSQL = `
SELECT to_jsonb(relation)
FROM knowledge_tag_relations AS relation
WHERE relation.knowledge_id = $1
  AND EXISTS (
    SELECT 1
    FROM knowledges AS knowledge
    WHERE knowledge.id = relation.knowledge_id
      AND knowledge.tenant_id = $2
      AND knowledge.knowledge_base_id = $3
      AND knowledge.deleted_at IS NULL
  )
ORDER BY relation.tag_id
`

const exportKnowledgeTagsSQL = `
SELECT to_jsonb(tag)
FROM knowledge_tags AS tag
WHERE tag.tenant_id = $2
  AND tag.knowledge_base_id = $3
  AND tag.deleted_at IS NULL
  AND EXISTS (
    SELECT 1
    FROM knowledge_tag_relations AS relation
    WHERE relation.knowledge_id = $1
      AND relation.tag_id = tag.id
  )
ORDER BY tag.id
`

const exportResourcesSQL = `
SELECT to_jsonb(resource)
FROM resources AS resource
WHERE resource.tenant_id = $1
  AND resource.deleted_at IS NULL
  AND (
    resource.handle = ANY($4::text[])
    OR EXISTS (
      SELECT 1
      FROM resource_bindings AS binding
      WHERE binding.resource_id = resource.id
        AND binding.tenant_id = $1
        AND (
          (binding.owner_type = 'knowledge' AND binding.owner_id = $2)
          OR (binding.owner_type = 'chunk' AND binding.owner_id = ANY($3::text[]))
        )
    )
  )
ORDER BY resource.id
`

const exportResourceBindingsSQL = `
SELECT to_jsonb(binding)
FROM resource_bindings AS binding
WHERE binding.tenant_id = $1
  AND binding.resource_id = ANY($2::text[])
  AND (
    (binding.owner_type = 'knowledge' AND binding.owner_id = $3)
    OR (binding.owner_type = 'chunk' AND binding.owner_id = ANY($4::text[]))
  )
ORDER BY binding.id
`

func (exporter *PostgresExporter) Export(ctx context.Context, identity PostgresIdentity) (PostgresSnapshot, error) {
	if identity.TenantID == 0 || !validSearchIdentifier(identity.KnowledgeBaseID) || !validSearchIdentifier(identity.KnowledgeID) {
		return PostgresSnapshot{}, errors.New("invalid PostgreSQL export identity")
	}
	var knowledgeRaw []byte
	if err := exporter.db.QueryRow(ctx, exportKnowledgeSQL, identity.KnowledgeID, identity.TenantID, identity.KnowledgeBaseID).Scan(&knowledgeRaw); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return PostgresSnapshot{}, errors.New("eligible PostgreSQL knowledge not found")
		}
		return PostgresSnapshot{}, errors.New("read PostgreSQL knowledge")
	}
	knowledge, err := canonicalJSON(knowledgeRaw)
	if err != nil || !matchesKnowledgeRow(knowledge, identity) {
		return PostgresSnapshot{}, errors.New("PostgreSQL returned a cross-scope knowledge row")
	}
	chunks, err := exporter.queryRows(ctx, exportChunksSQL, identity.KnowledgeID, identity.TenantID)
	if err != nil {
		return PostgresSnapshot{}, err
	}
	if len(chunks) == 0 {
		return PostgresSnapshot{}, errors.New("PostgreSQL knowledge has no chunks")
	}
	chunkIDs := make([]string, 0, len(chunks))
	for _, chunk := range chunks {
		chunkID, ok := matchesChunkRow(chunk, identity)
		if !ok {
			return PostgresSnapshot{}, errors.New("PostgreSQL returned a cross-scope chunk row")
		}
		chunkIDs = append(chunkIDs, chunkID)
	}
	knowledgeTags, err := exporter.queryRows(ctx, exportKnowledgeTagsSQL, identity.KnowledgeID, identity.TenantID, identity.KnowledgeBaseID)
	if err != nil {
		return PostgresSnapshot{}, err
	}
	for _, tag := range knowledgeTags {
		if !matchesKnowledgeTagRow(tag, identity) {
			return PostgresSnapshot{}, errors.New("PostgreSQL returned a cross-scope knowledge tag")
		}
	}
	tagRelations, err := exporter.queryRows(ctx, exportTagRelationsSQL, identity.KnowledgeID, identity.TenantID, identity.KnowledgeBaseID)
	if err != nil {
		return PostgresSnapshot{}, err
	}
	for _, relation := range tagRelations {
		if !matchesTagRelationRow(relation, identity.KnowledgeID) {
			return PostgresSnapshot{}, errors.New("PostgreSQL returned a cross-scope tag relation")
		}
	}
	handles := referencedResourceHandles(append([]json.RawMessage{knowledge}, chunks...))
	resources, err := exporter.queryRows(ctx, exportResourcesSQL, identity.TenantID, identity.KnowledgeID, chunkIDs, handles)
	if err != nil {
		return PostgresSnapshot{}, err
	}
	resourceIDs := make([]string, 0, len(resources))
	for _, resource := range resources {
		resourceID, ok := matchesResourceRow(resource, identity.TenantID)
		if !ok {
			return PostgresSnapshot{}, errors.New("PostgreSQL returned a cross-scope resource row")
		}
		resourceIDs = append(resourceIDs, resourceID)
	}
	bindings, err := exporter.queryRows(ctx, exportResourceBindingsSQL, identity.TenantID, resourceIDs, identity.KnowledgeID, chunkIDs)
	if err != nil {
		return PostgresSnapshot{}, err
	}
	for _, binding := range bindings {
		if !matchesBindingRow(binding, identity.TenantID, identity.KnowledgeID, chunkIDs, resourceIDs) {
			return PostgresSnapshot{}, errors.New("PostgreSQL returned a cross-scope resource binding")
		}
	}
	tables := []TableSnapshot{
		newTableSnapshot("knowledges", []json.RawMessage{knowledge}),
		newTableSnapshot("chunks", chunks),
		newTableSnapshot("knowledge_tags", knowledgeTags),
		newTableSnapshot("knowledge_tag_relations", tagRelations),
		newTableSnapshot("resources", resources),
		newTableSnapshot("resource_bindings", bindings),
	}
	return PostgresSnapshot{Tables: tables}, nil
}

func (exporter *PostgresExporter) queryRows(ctx context.Context, query string, arguments ...any) ([]json.RawMessage, error) {
	rows, err := exporter.db.Query(ctx, query, arguments...)
	if err != nil {
		return nil, errors.New("read PostgreSQL release rows")
	}
	defer rows.Close()
	result := make([]json.RawMessage, 0)
	for rows.Next() {
		var raw []byte
		if err := rows.Scan(&raw); err != nil {
			return nil, errors.New("scan PostgreSQL release row")
		}
		canonical, err := canonicalJSON(raw)
		if err != nil {
			return nil, errors.New("invalid PostgreSQL release row")
		}
		result = append(result, canonical)
	}
	if rows.Err() != nil {
		return nil, errors.New("iterate PostgreSQL release rows")
	}
	return result, nil
}

func newTableSnapshot(name string, rows []json.RawMessage) TableSnapshot {
	rows = append([]json.RawMessage(nil), rows...)
	sort.Slice(rows, func(left, right int) bool { return string(rows[left]) < string(rows[right]) })
	encoded, _ := json.Marshal(rows)
	digest := sha256.Sum256(encoded)
	return TableSnapshot{Name: name, Rows: rows, RowCount: len(rows), SHA256: hex.EncodeToString(digest[:])}
}

func referencedResourceHandles(rows []json.RawMessage) []string {
	seen := map[string]bool{}
	for _, row := range rows {
		for _, match := range resourceReferencePattern.FindAllSubmatch(row, -1) {
			seen[string(match[1])] = true
		}
	}
	handles := make([]string, 0, len(seen))
	for handle := range seen {
		handles = append(handles, handle)
	}
	sort.Strings(handles)
	return handles
}

func matchesKnowledgeRow(raw json.RawMessage, identity PostgresIdentity) bool {
	var row struct {
		ID              string `json:"id"`
		TenantID        uint64 `json:"tenant_id"`
		KnowledgeBaseID string `json:"knowledge_base_id"`
		ParseStatus     string `json:"parse_status"`
	}
	return json.Unmarshal(raw, &row) == nil && row.ID == identity.KnowledgeID && row.TenantID == identity.TenantID && row.KnowledgeBaseID == identity.KnowledgeBaseID && row.ParseStatus == "completed"
}

func matchesChunkRow(raw json.RawMessage, identity PostgresIdentity) (string, bool) {
	var row struct {
		ID          string `json:"id"`
		TenantID    uint64 `json:"tenant_id"`
		KnowledgeID string `json:"knowledge_id"`
	}
	if json.Unmarshal(raw, &row) != nil {
		return "", false
	}
	return row.ID, validSearchIdentifier(row.ID) && row.TenantID == identity.TenantID && row.KnowledgeID == identity.KnowledgeID
}

func matchesKnowledgeTagRow(raw json.RawMessage, identity PostgresIdentity) bool {
	var row struct {
		ID              string `json:"id"`
		TenantID        uint64 `json:"tenant_id"`
		KnowledgeBaseID string `json:"knowledge_base_id"`
	}
	return json.Unmarshal(raw, &row) == nil && validSearchIdentifier(row.ID) && row.TenantID == identity.TenantID && row.KnowledgeBaseID == identity.KnowledgeBaseID
}

func matchesTagRelationRow(raw json.RawMessage, knowledgeID string) bool {
	var row struct {
		KnowledgeID string `json:"knowledge_id"`
		TagID       string `json:"tag_id"`
	}
	return json.Unmarshal(raw, &row) == nil && row.KnowledgeID == knowledgeID && validSearchIdentifier(row.TagID)
}

func matchesResourceRow(raw json.RawMessage, tenantID uint64) (string, bool) {
	var row struct {
		ID       string `json:"id"`
		TenantID uint64 `json:"tenant_id"`
		Handle   string `json:"handle"`
	}
	if json.Unmarshal(raw, &row) != nil {
		return "", false
	}
	return row.ID, validSearchIdentifier(row.ID) && row.TenantID == tenantID && len(row.Handle) == 22
}

func matchesBindingRow(raw json.RawMessage, tenantID uint64, knowledgeID string, chunkIDs, resourceIDs []string) bool {
	var row struct {
		ResourceID string `json:"resource_id"`
		TenantID   uint64 `json:"tenant_id"`
		OwnerType  string `json:"owner_type"`
		OwnerID    string `json:"owner_id"`
	}
	if json.Unmarshal(raw, &row) != nil || row.TenantID != tenantID || !containsString(resourceIDs, row.ResourceID) {
		return false
	}
	return row.OwnerType == "knowledge" && row.OwnerID == knowledgeID || row.OwnerType == "chunk" && containsString(chunkIDs, row.OwnerID)
}

func containsString(values []string, target string) bool {
	for _, value := range values {
		if value == target {
			return true
		}
	}
	return false
}