From 9daa1ee90f8291dbf05f7b7e621694019f4b8a2b Mon Sep 17 00:00:00 2001
From: Gourab Singha <gaurabsingha16@gmail.com>
Date: Sat, 27 Jun 2026 21:20:15 +0530
Subject: [PATCH] feat(collection): add export_migration method to export SQL
 DDL for database migrations

---
 src/tests/test_collection.py | 29 ++++++++++++++++++
 src/vecs/collection.py       | 59 ++++++++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+)

diff --git a/src/tests/test_collection.py b/src/tests/test_collection.py
index 0653329..8b1367d 100644
--- a/src/tests/test_collection.py
+++ b/src/tests/test_collection.py
@@ -895,3 +895,32 @@ def test_hnsw_unavailable_error(client: vecs.Client) -> None:
     bar = client.get_or_create_collection(name="bar", dimension=dim)
     with pytest.raises(ArgError):
         bar.create_index(method=IndexMethod.hnsw)
+
+
+def test_export_migration(client: vecs.Client) -> None:
+    bar = client.get_or_create_collection(name="bar_migration", dimension=3)
+    
+    # 1. Test defaults (include_extension=True, include_schema=True, if_not_exists=True)
+    sql_default = bar.export_migration()
+    assert "CREATE EXTENSION IF NOT EXISTS vector;" in sql_default
+    assert "CREATE SCHEMA IF NOT EXISTS vecs;" in sql_default
+    assert "CREATE TABLE IF NOT EXISTS vecs.bar_migration (" in sql_default
+    assert "CREATE INDEX IF NOT EXISTS ix_meta_bar_migration ON vecs.\"bar_migration\" USING gin (metadata jsonb_path_ops);" in sql_default
+    
+    # 2. Test include_extension=False
+    sql_no_ext = bar.export_migration(include_extension=False)
+    assert "CREATE EXTENSION" not in sql_no_ext
+    assert "CREATE SCHEMA IF NOT EXISTS vecs;" in sql_no_ext
+    
+    # 3. Test include_schema=False
+    sql_no_schema = bar.export_migration(include_schema=False)
+    assert "CREATE EXTENSION IF NOT EXISTS vector;" in sql_no_schema
+    assert "CREATE SCHEMA" not in sql_no_schema
+    
+    # 4. Test if_not_exists=False
+    sql_no_exists = bar.export_migration(if_not_exists=False)
+    assert "CREATE EXTENSION vector;" in sql_no_exists
+    assert "CREATE SCHEMA vecs;" in sql_no_exists
+    assert "CREATE TABLE vecs.bar_migration (" in sql_no_exists
+    assert "CREATE INDEX ix_meta_bar_migration ON vecs.\"bar_migration\" USING gin (metadata jsonb_path_ops);" in sql_no_exists
+
diff --git a/src/vecs/collection.py b/src/vecs/collection.py
index eab8f93..07ea340 100644
--- a/src/vecs/collection.py
+++ b/src/vecs/collection.py
@@ -300,6 +300,65 @@ def _create(self):
             )
         return self
 
+    def export_migration(
+        self,
+        *,
+        include_extension: bool = True,
+        include_schema: bool = True,
+        if_not_exists: bool = True,
+    ) -> str:
+        """
+        Exports the SQL migration script for this collection.
+
+        Args:
+            include_extension (bool): Whether to include the `CREATE EXTENSION` statement for pgvector.
+            include_schema (bool): Whether to include the `CREATE SCHEMA` statement.
+            if_not_exists (bool): Whether to include `IF NOT EXISTS` clauses.
+
+        Returns:
+            str: The SQL DDL statements as a string.
+        """
+        from sqlalchemy.schema import CreateTable
+        from sqlalchemy.dialects import postgresql
+
+        schema = self.table.schema or "vecs"
+        statements = []
+
+        if include_extension:
+            ext_sql = (
+                "CREATE EXTENSION IF NOT EXISTS vector;"
+                if if_not_exists
+                else "CREATE EXTENSION vector;"
+            )
+            statements.append(ext_sql)
+
+        if include_schema:
+            schema_sql = (
+                f"CREATE SCHEMA IF NOT EXISTS {schema};"
+                if if_not_exists
+                else f"CREATE SCHEMA {schema};"
+            )
+            statements.append(schema_sql)
+
+        create_table_sql = str(
+            CreateTable(self.table).compile(dialect=postgresql.dialect())
+        ).strip()
+        if if_not_exists and create_table_sql.startswith("CREATE TABLE "):
+            create_table_sql = (
+                "CREATE TABLE IF NOT EXISTS " + create_table_sql[len("CREATE TABLE ") :]
+            )
+        create_table_sql += ";"
+        statements.append(create_table_sql)
+
+        index_name = f"ix_meta_{self.table.name}"
+        if if_not_exists:
+            index_sql = f'CREATE INDEX IF NOT EXISTS {index_name} ON {schema}."{self.table.name}" USING gin (metadata jsonb_path_ops);'
+        else:
+            index_sql = f'CREATE INDEX {index_name} ON {schema}."{self.table.name}" USING gin (metadata jsonb_path_ops);'
+        statements.append(index_sql)
+
+        return "\n\n".join(statements)
+
     def _drop(self):
         """
         PRIVATE