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

package v071

import (
	"context"
	"encoding/json"
	"os"
	"testing"

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

func TestPostgresImporterV071ReplacesAndRestoresDocument(t *testing.T) {
	sourceURL, targetURL := os.Getenv("TEST_WEKNORA_V071_DATABASE_URL"), os.Getenv("TEST_WEKNORA_V071_TARGET_DATABASE_URL")
	if sourceURL == "" || targetURL == "" {
		t.Skip("v0.7.1 source and target URLs are not set")
	}
	ctx := context.Background()
	source, err := pgxpool.New(ctx, sourceURL)
	if err != nil {
		t.Fatal(err)
	}
	defer source.Close()
	target, err := pgxpool.New(ctx, targetURL)
	if err != nil {
		t.Fatal(err)
	}
	defer target.Close()
	identity := PostgresIdentity{TenantID: 10000, KnowledgeBaseID: "kb-release-safe", KnowledgeID: "knowledge-release-safe"}
	sourceExporter, _ := NewPostgresExporter(source)
	newSnapshot, err := sourceExporter.Export(ctx, identity)
	if err != nil {
		t.Fatal(err)
	}
	oldSnapshot := clonePostgresSnapshot(newSnapshot)
	oldChunks := oldSnapshot.Table("chunks")
	var chunk map[string]any
	if err := json.Unmarshal(oldChunks.Rows[0], &chunk); err != nil {
		t.Fatal(err)
	}
	chunk["content"] = "Prior production content"
	oldRaw, err := json.Marshal(chunk)
	if err != nil {
		t.Fatal(err)
	}
	for index, table := range oldSnapshot.Tables {
		if table.Name == "chunks" {
			oldSnapshot.Tables[index] = newTableSnapshot("chunks", []json.RawMessage{oldRaw})
		}
	}

	importer, err := NewPostgresImporter(target)
	if err != nil {
		t.Fatal(err)
	}
	if err := importer.Replace(ctx, identity, oldSnapshot); err != nil {
		t.Fatal(err)
	}
	targetExporter, _ := NewPostgresExporter(target)
	oldActual, err := targetExporter.Export(ctx, identity)
	if err != nil {
		t.Fatal(err)
	}
	assertPostgresSnapshotDigests(t, oldActual, oldSnapshot)
	if err := importer.Replace(ctx, identity, newSnapshot); err != nil {
		t.Fatal(err)
	}
	newActual, err := targetExporter.Export(ctx, identity)
	if err != nil {
		t.Fatal(err)
	}
	assertPostgresSnapshotDigests(t, newActual, newSnapshot)
	if err := importer.Replace(ctx, identity, oldSnapshot); err != nil {
		t.Fatal(err)
	}
	restored, err := targetExporter.Export(ctx, identity)
	if err != nil {
		t.Fatal(err)
	}
	assertPostgresSnapshotDigests(t, restored, oldSnapshot)
	if err := importer.Disable(ctx, identity); err != nil {
		t.Fatal(err)
	}
	var knowledgeEnabled string
	var enabledChunks, knowledgeRows, chunkRows int
	if err := target.QueryRow(ctx, `SELECT enable_status FROM knowledges WHERE id=$1`, identity.KnowledgeID).Scan(&knowledgeEnabled); err != nil {
		t.Fatal(err)
	}
	if err := target.QueryRow(ctx, `SELECT COUNT(*) FILTER (WHERE is_enabled), COUNT(*) FROM chunks WHERE knowledge_id=$1`, identity.KnowledgeID).Scan(&enabledChunks, &chunkRows); err != nil {
		t.Fatal(err)
	}
	if err := target.QueryRow(ctx, `SELECT COUNT(*) FROM knowledges WHERE id=$1`, identity.KnowledgeID).Scan(&knowledgeRows); err != nil {
		t.Fatal(err)
	}
	if knowledgeEnabled != "disabled" || enabledChunks != 0 || knowledgeRows != 1 || chunkRows != 1 {
		t.Fatalf("status=%s enabled_chunks=%d rows=%d/%d", knowledgeEnabled, enabledChunks, knowledgeRows, chunkRows)
	}
}

func clonePostgresSnapshot(snapshot PostgresSnapshot) PostgresSnapshot {
	result := PostgresSnapshot{Tables: make([]TableSnapshot, len(snapshot.Tables))}
	for i, table := range snapshot.Tables {
		result.Tables[i] = table
		result.Tables[i].Rows = append([]json.RawMessage(nil), table.Rows...)
	}
	return result
}
func assertPostgresSnapshotDigests(t *testing.T, actual, want PostgresSnapshot) {
	t.Helper()
	for _, table := range want.Tables {
		got := actual.Table(table.Name)
		if got.RowCount != table.RowCount || got.SHA256 != table.SHA256 {
			t.Fatalf("%s got=%#v want=%#v", table.Name, got, table)
		}
	}
}