relspecgo/pkg/writers/pgsql/extensions_test.go

package pgsql

import (
	"reflect"
	"strings"
	"testing"

	"git.warky.dev/wdevs/relspecgo/pkg/models"
	"git.warky.dev/wdevs/relspecgo/pkg/writers"
)

// buildExtensionSchema returns a single-table schema the extension detection tests mutate.
func buildExtensionSchema(t *testing.T) (*models.Schema, *models.Table) {
	t.Helper()

	schema := models.InitSchema("public")
	table := models.InitTable("documents", "public")
	schema.Tables = append(schema.Tables, table)
	return schema, table
}

func addColumn(table *models.Table, name, sqlType string) *models.Column {
	col := models.InitColumn(name, table.Name, table.Schema)
	col.Type = sqlType
	table.Columns[name] = col
	return col
}

func TestRequiredExtensions_Detection(t *testing.T) {
	tests := []struct {
		name  string
		build func(schema *models.Schema, table *models.Table)
		want  []string
	}{
		{
			name:  "no extensions",
			build: func(_ *models.Schema, table *models.Table) { addColumn(table, "id", "integer") },
			want:  nil,
		},
		{
			name: "column type",
			build: func(_ *models.Schema, table *models.Table) {
				addColumn(table, "embedding", "vector(1536)")
				addColumn(table, "name", "citext")
			},
			want: []string{"citext", "vector"},
		},
		{
			name: "column default function",
			build: func(_ *models.Schema, table *models.Table) {
				addColumn(table, "id", "uuid").Default = "uuid_generate_v4()"
			},
			want: []string{"uuid-ossp"},
		},
		{
			name: "check constraint expression",
			build: func(_ *models.Schema, table *models.Table) {
				addColumn(table, "geom", "geometry")
				table.Constraints["chk_geom"] = &models.Constraint{
					Name:       "chk_geom",
					Type:       models.CheckConstraint,
					Expression: "ST_IsValid(geom)",
				}
			},
			want: []string{"postgis"},
		},
		{
			name: "partial index predicate",
			build: func(_ *models.Schema, table *models.Table) {
				addColumn(table, "title", "text")
				table.Indexes["idx_title"] = &models.Index{
					Name:    "idx_title",
					Type:    "btree",
					Columns: []string{"title"},
					Where:   "similarity(title, 'x') > 0.3",
				}
			},
			want: []string{"pg_trgm"},
		},
		{
			name: "view definition",
			build: func(schema *models.Schema, table *models.Table) {
				addColumn(table, "title", "text")
				schema.Views = append(schema.Views, &models.View{
					Name:       "v_documents",
					Schema:     "public",
					Definition: "SELECT unaccent(title) FROM documents",
				})
			},
			want: []string{"unaccent"},
		},
		{
			name: "index access method",
			build: func(_ *models.Schema, table *models.Table) {
				addColumn(table, "body", "text")
				table.Indexes["idx_body"] = &models.Index{
					Name:    "idx_body",
					Type:    "bm25",
					Columns: []string{"body"},
					Comment: "with (key_field='id')",
				}
			},
			want: []string{"pg_search"},
		},
		{
			name: "vchord depends on vector",
			build: func(_ *models.Schema, table *models.Table) {
				addColumn(table, "embedding", "vector(3)")
				table.Indexes["idx_embedding"] = &models.Index{
					Name:    "idx_embedding",
					Type:    "vchordrq",
					Columns: []string{"embedding"},
				}
			},
			want: []string{"vector", "vchord"},
		},
		{
			name: "gin on scalar needs btree_gin",
			build: func(_ *models.Schema, table *models.Table) {
				addColumn(table, "views", "integer")
				table.Indexes["idx_views"] = &models.Index{
					Name:    "idx_views",
					Type:    "gin",
					Columns: []string{"views"},
				}
			},
			want: []string{"btree_gin"},
		},
		{
			name: "gist on scalar needs btree_gist",
			build: func(_ *models.Schema, table *models.Table) {
				addColumn(table, "views", "integer")
				table.Indexes["idx_views"] = &models.Index{
					Name:    "idx_views",
					Type:    "gist",
					Columns: []string{"views"},
				}
			},
			want: []string{"btree_gist"},
		},
		{
			name: "gist on geometry uses postgis operator classes",
			build: func(_ *models.Schema, table *models.Table) {
				addColumn(table, "location", "geometry(Point,4326)")
				table.Indexes["idx_location"] = &models.Index{
					Name:    "idx_location",
					Type:    "gist",
					Columns: []string{"location"},
				}
			},
			want: []string{"postgis"},
		},
		{
			name: "gin on jsonb needs no companion",
			build: func(_ *models.Schema, table *models.Table) {
				addColumn(table, "payload", "jsonb")
				table.Indexes["idx_payload"] = &models.Index{
					Name:    "idx_payload",
					Type:    "gin",
					Columns: []string{"payload"},
				}
			},
			want: nil,
		},
		{
			name: "gin on text uses pg_trgm",
			build: func(_ *models.Schema, table *models.Table) {
				addColumn(table, "title", "text")
				table.Indexes["idx_title"] = &models.Index{
					Name:    "idx_title",
					Type:    "gin",
					Columns: []string{"title"},
				}
			},
			want: []string{"pg_trgm"},
		},
		{
			name: "gin on array needs no companion",
			build: func(_ *models.Schema, table *models.Table) {
				addColumn(table, "tags", "text[]")
				table.Indexes["idx_tags"] = &models.Index{
					Name:    "idx_tags",
					Type:    "gin",
					Columns: []string{"tags"},
				}
			},
			want: nil,
		},
		{
			name: "declared in metadata as string",
			build: func(schema *models.Schema, _ *models.Table) {
				schema.Metadata = map[string]any{"extensions": "pg_cron, timescaledb"}
			},
			want: []string{"pg_cron", "timescaledb"},
		},
		{
			name: "declared in metadata as list",
			build: func(schema *models.Schema, _ *models.Table) {
				schema.Metadata = map[string]any{"extensions": []any{"postgis_topology", "pg_stat_statements"}}
			},
			// postgis is pulled in as a dependency of postgis_topology and emitted first.
			want: []string{"pg_stat_statements", "postgis", "postgis_topology"},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			schema, table := buildExtensionSchema(t)
			tt.build(schema, table)

			if got := requiredExtensions(schema); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("requiredExtensions() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestRequiredExtensions_NilSchema(t *testing.T) {
	if got := requiredExtensions(nil); got != nil {
		t.Errorf("requiredExtensions(nil) = %v, want nil", got)
	}
}

func TestWriteDatabase_QuotesExtensionNames(t *testing.T) {
	db := models.InitDatabase("testdb")
	schema, table := buildExtensionSchema(t)
	addColumn(table, "id", "uuid").Default = "uuid_generate_v4()"
	db.Schemas = append(db.Schemas, schema)

	output := writeDatabaseOutput(t, db)
	if !strings.Contains(output, `CREATE EXTENSION IF NOT EXISTS "uuid-ossp";`) {
		t.Fatalf("expected quoted extension name, got:\n%s", output)
	}
}

func TestGenerateSchemaStatements_ExtensionDependencyOrder(t *testing.T) {
	schema, table := buildExtensionSchema(t)
	addColumn(table, "embedding", "vector(3)")
	table.Indexes["idx_embedding"] = &models.Index{
		Name:    "idx_embedding",
		Type:    "vchordrq",
		Columns: []string{"embedding"},
	}

	writer := NewWriter(&writers.WriterOptions{})
	statements, err := writer.GenerateSchemaStatements(schema)
	if err != nil {
		t.Fatalf("GenerateSchemaStatements failed: %v", err)
	}

	joined := strings.Join(statements, "\n")
	vector := strings.Index(joined, "CREATE EXTENSION IF NOT EXISTS vector")
	vchord := strings.Index(joined, "CREATE EXTENSION IF NOT EXISTS vchord")
	if vector < 0 || vchord < 0 {
		t.Fatalf("expected vector and vchord extensions, got:\n%s", joined)
	}
	if vector > vchord {
		t.Fatalf("expected vector to be created before vchord, got:\n%s", joined)
	}
}