feat(pgsql): support vector and PostGIS indexes with extensions

* Add handling for pgvector and PostGIS extensions in migration scripts * Implement operator class and storage parameters for vector indexes * Update tests to validate new index behaviors and extension creation
2026-08-29 20:39:57 +02:00
parent 16af529120
commit ab3c9217df
19 changed files with 2472 additions and 94 deletions
@@ -1310,3 +1310,203 @@ func TestWriteSchema_UsesStorageTypeForSerialAlterStatements(t *testing.T) {
 		t.Fatalf("expected serial alter to include USING cast, got:\n%s", output)
 	}
 }
+
+// buildVectorSpatialSchema returns a database with a pgvector column and a PostGIS column.
+func buildVectorSpatialSchema(indexType, indexComment string) *models.Database {
+	db := models.InitDatabase("testdb")
+	schema := models.InitSchema("public")
+
+	table := models.InitTable("documents", "public")
+
+	embedding := models.InitColumn("embedding", "documents", "public")
+	embedding.Type = "vector(1536)"
+	table.Columns["embedding"] = embedding
+
+	location := models.InitColumn("location", "documents", "public")
+	location.Type = "geometry(Point,4326)"
+	table.Columns["location"] = location
+
+	if indexType != "" {
+		index := &models.Index{
+			Name:    "idx_documents_embedding",
+			Type:    indexType,
+			Columns: []string{"embedding"},
+			Comment: indexComment,
+		}
+		table.Indexes[index.Name] = index
+	}
+
+	schema.Tables = append(schema.Tables, table)
+	db.Schemas = append(db.Schemas, schema)
+	return db
+}
+
+func writeDatabaseOutput(t *testing.T, db *models.Database) string {
+	t.Helper()
+
+	var buf bytes.Buffer
+	writer := NewWriter(&writers.WriterOptions{})
+	writer.writer = &buf
+
+	if err := writer.WriteDatabase(db); err != nil {
+		t.Fatalf("WriteDatabase failed: %v", err)
+	}
+	return buf.String()
+}
+
+func TestWriteDatabase_VectorAndPostGISColumnsCreateExtensions(t *testing.T) {
+	output := writeDatabaseOutput(t, buildVectorSpatialSchema("", ""))
+
+	for _, want := range []string{
+		"CREATE EXTENSION IF NOT EXISTS postgis;",
+		"CREATE EXTENSION IF NOT EXISTS vector;",
+		"vector(1536)",
+		"geometry(Point,4326)",
+	} {
+		if !strings.Contains(output, want) {
+			t.Fatalf("expected output to contain %q, got:\n%s", want, output)
+		}
+	}
+
+	// postgis must be created before postgis-dependent extensions and stay deterministic
+	if strings.Index(output, "EXISTS postgis;") > strings.Index(output, "EXISTS vector;") {
+		t.Fatalf("expected extensions to be emitted in sorted order, got:\n%s", output)
+	}
+}
+
+func TestWriteDatabase_HNSWIndexUsesDefaultVectorOperatorClass(t *testing.T) {
+	output := writeDatabaseOutput(t, buildVectorSpatialSchema("hnsw", ""))
+
+	if !strings.Contains(output, "USING hnsw (embedding vector_cosine_ops)") {
+		t.Fatalf("expected hnsw index with default vector operator class, got:\n%s", output)
+	}
+	if !strings.Contains(output, "CREATE EXTENSION IF NOT EXISTS vector;") {
+		t.Fatalf("expected pgvector extension, got:\n%s", output)
+	}
+}
+
+func TestWriteDatabase_VectorIndexHonoursRequestedOperatorClassAndStorageParameters(t *testing.T) {
+	output := writeDatabaseOutput(t, buildVectorSpatialSchema("ivfflat", "opclass=vector_l2_ops; with (lists=100)"))
+
+	if !strings.Contains(output, "USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)") {
+		t.Fatalf("expected ivfflat index with requested opclass and storage parameters, got:\n%s", output)
+	}
+}
+
+func TestWriteDatabase_VectorIndexIgnoresIncompatibleOperatorClass(t *testing.T) {
+	output := writeDatabaseOutput(t, buildVectorSpatialSchema("hnsw", "opclass=halfvec_l2_ops"))
+
+	if !strings.Contains(output, "USING hnsw (embedding vector_cosine_ops)") {
+		t.Fatalf("expected halfvec operator class to be rejected for a vector column, got:\n%s", output)
+	}
+}
+
+func TestWriteDatabase_VectorIndexIgnoresCommentProseInStorageParameters(t *testing.T) {
+	output := writeDatabaseOutput(t, buildVectorSpatialSchema("hnsw", "tuned with (m=16, ef_construction=64, drop table foo)"))
+
+	if !strings.Contains(output, "WITH (m = 16, ef_construction = 64)") {
+		t.Fatalf("expected only well-formed storage parameters, got:\n%s", output)
+	}
+	if strings.Contains(output, "drop table") {
+		t.Fatalf("expected prose to be dropped from storage parameters, got:\n%s", output)
+	}
+}
+
+func TestWriteDatabase_GistIndexOnGeometryUsesDefaultOperatorClass(t *testing.T) {
+	db := models.InitDatabase("testdb")
+	schema := models.InitSchema("public")
+
+	table := models.InitTable("places", "public")
+	geom := models.InitColumn("geom", "places", "public")
+	geom.Type = "geometry(Point,4326)"
+	table.Columns["geom"] = geom
+
+	table.Indexes["idx_places_geom"] = &models.Index{
+		Name:    "idx_places_geom",
+		Type:    "gist",
+		Columns: []string{"geom"},
+	}
+
+	schema.Tables = append(schema.Tables, table)
+	db.Schemas = append(db.Schemas, schema)
+
+	output := writeDatabaseOutput(t, db)
+
+	if !strings.Contains(output, "USING gist (geom)") {
+		t.Fatalf("expected gist index to rely on the PostGIS default operator class, got:\n%s", output)
+	}
+}
+
+func TestWriteDatabase_GistIndexHonoursRequestedSpatialOperatorClass(t *testing.T) {
+	db := models.InitDatabase("testdb")
+	schema := models.InitSchema("public")
+
+	table := models.InitTable("places", "public")
+	geom := models.InitColumn("geom", "places", "public")
+	geom.Type = "geometry(PointZ,4326)"
+	table.Columns["geom"] = geom
+
+	table.Indexes["idx_places_geom_nd"] = &models.Index{
+		Name:    "idx_places_geom_nd",
+		Type:    "gist",
+		Columns: []string{"geom"},
+		Comment: "opclass=gist_geometry_ops_nd",
+	}
+
+	schema.Tables = append(schema.Tables, table)
+	db.Schemas = append(db.Schemas, schema)
+
+	output := writeDatabaseOutput(t, db)
+
+	if !strings.Contains(output, "USING gist (geom gist_geometry_ops_nd)") {
+		t.Fatalf("expected requested spatial operator class, got:\n%s", output)
+	}
+}
+
+func TestGenerateDatabaseStatements_VectorIndexIncludesOperatorClassAndParameters(t *testing.T) {
+	db := buildVectorSpatialSchema("hnsw", "opclass=vector_ip_ops; with (m=16)")
+
+	writer := NewWriter(&writers.WriterOptions{})
+	statements, err := writer.GenerateDatabaseStatements(db)
+	if err != nil {
+		t.Fatalf("GenerateDatabaseStatements failed: %v", err)
+	}
+
+	joined := strings.Join(statements, "\n")
+	for _, want := range []string{
+		"CREATE EXTENSION IF NOT EXISTS vector",
+		"CREATE EXTENSION IF NOT EXISTS postgis",
+		"USING hnsw (embedding vector_ip_ops) WITH (m = 16)",
+	} {
+		if !strings.Contains(joined, want) {
+			t.Fatalf("expected statements to contain %q, got:\n%s", want, joined)
+		}
+	}
+}
+
+func TestIndexStorageParameters(t *testing.T) {
+	tests := []struct {
+		name    string
+		comment string
+		want    string
+	}{
+		{"empty", "", ""},
+		{"no with clause", "opclass=vector_cosine_ops", ""},
+		{"single parameter", "with (lists=100)", "lists = 100"},
+		{"multiple parameters", "WITH (m = 16, ef_construction = 64)", "m = 16, ef_construction = 64"},
+		{"quoted value kept", "with (fillfactor='90')", "fillfactor = '90'"},
+		{"bm25 key field", "with (key_field='id')", "key_field = 'id'"},
+		{"dollar quoted value", "with (options = $$[build.internal]\nlists = [4096]$$)", "options = $$[build.internal]\nlists = [4096]$$"},
+		{"dollar quoted value with parens", "with (options = $$f(x)$$, m = 16)", "options = $$f(x)$$, m = 16"},
+		{"prose dropped", "with (lists=100, please drop everything)", "lists = 100"},
+		{"unterminated quote dropped", "with (key_field='id)", ""},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := indexStorageParameters(tt.comment); got != tt.want {
+				t.Errorf("indexStorageParameters(%q) = %q, want %q", tt.comment, got, tt.want)
+			}
+		})
+	}
+}