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 }