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

package v071

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"testing"

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

func TestPostgresExporterReadsOnlyDocumentAllowlist(t *testing.T) {
	handle1, handle2 := "ABCDEFGHIJKLMNOPQRSTUV", "abcdefghijklmnopqrstuv"
	db := &exportDB{
		knowledge: jsonRow{raw: []byte(`{"id":"knowledge-safe","tenant_id":10000,"knowledge_base_id":"kb-safe","parse_status":"completed","file_path":"resource://` + handle1 + `"}`)},
		rowsByMarker: map[string][]json.RawMessage{
			"select to_jsonb(chunk)":    {json.RawMessage(`{"id":"chunk-1","tenant_id":10000,"knowledge_id":"knowledge-safe","image_info":{"url":"resource://` + handle2 + `"}}`)},
			"select to_jsonb(tag)":      {json.RawMessage(`{"id":"tag-safe","tenant_id":10000,"knowledge_base_id":"kb-safe","name":"synthetic"}`)},
			"select to_jsonb(relation)": {json.RawMessage(`{"knowledge_id":"knowledge-safe","tag_id":"tag-safe"}`)},
			"select to_jsonb(resource)": {json.RawMessage(`{"id":"res-1","handle":"` + handle1 + `","tenant_id":10000,"physical_path":"tenant/knowledge/source.pdf"}`), json.RawMessage(`{"id":"res-2","handle":"` + handle2 + `","tenant_id":10000,"physical_path":"tenant/knowledge/image.png"}`)},
			"select to_jsonb(binding)":  {json.RawMessage(`{"id":"bind-1","resource_id":"res-1","tenant_id":10000,"owner_type":"knowledge","owner_id":"knowledge-safe"}`)},
		},
	}
	exporter, err := NewPostgresExporter(db)
	if err != nil {
		t.Fatal(err)
	}
	snapshot, err := exporter.Export(context.Background(), PostgresIdentity{TenantID: 10000, KnowledgeBaseID: "kb-safe", KnowledgeID: "knowledge-safe"})
	if err != nil {
		t.Fatal(err)
	}
	if snapshot.Table("knowledges").RowCount != 1 || snapshot.Table("chunks").RowCount != 1 || snapshot.Table("knowledge_tags").RowCount != 1 || snapshot.Table("knowledge_tag_relations").RowCount != 1 || snapshot.Table("resources").RowCount != 2 {
		t.Fatalf("snapshot=%#v", snapshot)
	}
	if got := db.resourceHandles; len(got) != 2 || got[0] != handle1 || got[1] != handle2 {
		t.Fatalf("resource handles=%v", got)
	}
	if len(db.chunkIDs) != 1 || db.chunkIDs[0] != "chunk-1" {
		t.Fatalf("chunk IDs=%v", db.chunkIDs)
	}
	joined := strings.ToLower(strings.Join(db.queries, "\n"))
	for _, required := range []string{"from knowledges", "from chunks", "from knowledge_tags", "from knowledge_tag_relations", "from resources", "from resource_bindings", "tenant_id", "knowledge_base_id", "knowledge_id"} {
		if !strings.Contains(joined, required) {
			t.Fatalf("SQL missing %q", required)
		}
	}
	for _, forbidden := range []string{"users", "messages", "sessions", "api_keys", "model_providers", "jobs", "feedback"} {
		if strings.Contains(joined, forbidden) {
			t.Fatalf("SQL included forbidden table %q", forbidden)
		}
	}
}

func TestPostgresExporterRejectsCrossScopeRows(t *testing.T) {
	db := &exportDB{knowledge: jsonRow{raw: []byte(`{"id":"knowledge-safe","tenant_id":10000,"knowledge_base_id":"wrong","parse_status":"completed"}`)}, rowsByMarker: map[string][]json.RawMessage{}}
	exporter, _ := NewPostgresExporter(db)
	_, err := exporter.Export(context.Background(), PostgresIdentity{TenantID: 10000, KnowledgeBaseID: "kb-safe", KnowledgeID: "knowledge-safe"})
	if err == nil || strings.Contains(err.Error(), "wrong") || strings.Contains(err.Error(), "knowledge-safe") {
		t.Fatalf("error=%v", err)
	}
}

type jsonRow struct {
	raw []byte
	err error
}

func (r jsonRow) Scan(dest ...any) error {
	if r.err != nil {
		return r.err
	}
	*dest[0].(*[]byte) = append([]byte(nil), r.raw...)
	return nil
}

type jsonRows struct {
	values []json.RawMessage
	index  int
}

func (r *jsonRows) Close()                                       {}
func (r *jsonRows) Err() error                                   { return nil }
func (r *jsonRows) CommandTag() pgconn.CommandTag                { return pgconn.CommandTag{} }
func (r *jsonRows) FieldDescriptions() []pgconn.FieldDescription { return nil }
func (r *jsonRows) RawValues() [][]byte                          { return nil }
func (r *jsonRows) Conn() *pgx.Conn                              { return nil }
func (r *jsonRows) Next() bool {
	if r.index >= len(r.values) {
		return false
	}
	r.index++
	return true
}
func (r *jsonRows) Scan(dest ...any) error {
	*dest[0].(*[]byte) = append([]byte(nil), r.values[r.index-1]...)
	return nil
}
func (r *jsonRows) Values() ([]any, error) { return nil, nil }

type exportDB struct {
	knowledge       pgx.Row
	rowsByMarker    map[string][]json.RawMessage
	queries         []string
	resourceHandles []string
	chunkIDs        []string
}

func (d *exportDB) QueryRow(_ context.Context, sql string, args ...any) pgx.Row {
	d.queries = append(d.queries, sql)
	return d.knowledge
}
func (d *exportDB) Query(_ context.Context, sql string, args ...any) (pgx.Rows, error) {
	d.queries = append(d.queries, sql)
	lower := strings.ToLower(sql)
	for marker, values := range d.rowsByMarker {
		if strings.Contains(lower, marker) {
			if marker == "select to_jsonb(resource)" && len(args) >= 4 {
				d.chunkIDs = append([]string(nil), args[2].([]string)...)
				d.resourceHandles = append([]string(nil), args[3].([]string)...)
			}
			return &jsonRows{values: values}, nil
		}
	}
	return nil, errors.New("unexpected query")
}
func (d *exportDB) Exec(context.Context, string, ...any) (pgconn.CommandTag, error) {
	return pgconn.CommandTag{}, nil
}