Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions src/tests/test_collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

59 changes: 59 additions & 0 deletions src/vecs/collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down