chore: move configuration to alembic package

timescale · Jan 7, 2025 · 828347f · 828347f
1 parent 1c43792
commit 828347f
Show file tree

Hide file tree

Showing 6 changed files with 9 additions and 18 deletions.
diff --git a/docs/python-integration.md b/docs/python-integration.md
@@ -197,7 +197,7 @@ Then you can use the `create_vectorizer` operation to create a vectorizer for yo
 
 ```python
 from alembic import op
-from pgai.vectorizer.configuration import (
+from pgai.alembic.configuration import (
     OpenAIConfig,
     CharacterTextSplitterConfig,
     PythonTemplateConfig

diff --git a/projects/extension/Dockerfile b/projects/extension/Dockerfile
@@ -52,6 +52,10 @@ WORKDIR /pgai
 COPY . .
 RUN just build install
 
+RUN mkdir -p /docker-entrypoint-initdb.d && \
+    echo "#!/bin/bash" > /docker-entrypoint-initdb.d/configure-timescaledb.sh && \
+    echo "echo \"shared_preload_libraries = 'timescaledb'\" >> \${PGDATA}/postgresql.conf" >> /docker-entrypoint-initdb.d/configure-timescaledb.sh && \
+    chmod +x /docker-entrypoint-initdb.d/configure-timescaledb.sh
 
 ###############################################################################
 # image for use in extension development

diff --git a/...cts/pgai/pgai/vectorizer/configuration.py → projects/pgai/pgai/alembic/configuration.py b/...cts/pgai/pgai/vectorizer/configuration.py → projects/pgai/pgai/alembic/configuration.py
@@ -65,7 +65,6 @@ def to_sql_argument(self) -> str:
             base_name = re.sub("([a-z0-9])([A-Z])", r"\1_\2", base_name).lower()
             # Remove 'config' and clean up any double underscores
             base_name = base_name.replace("config", "").strip("_")
-            # Remove any duplicate underscores that might have been created
             fn_name = f"{self.arg_type}_{base_name}"
 
         return f", {self.arg_type} => ai.{fn_name}({format_sql_params(params)})"  # type: ignore

diff --git a/projects/pgai/pgai/alembic/operations.py b/projects/pgai/pgai/alembic/operations.py
@@ -4,7 +4,7 @@
 from sqlalchemy import text
 from typing_extensions import override
 
-from pgai.vectorizer.configuration import (
+from pgai.alembic.configuration import (
     CharacterTextSplitterConfig,
     ChunkValueConfig,
     CreateVectorizerParams,
@@ -23,8 +23,6 @@
 
 
 class CreateVectorizerOp(MigrateOperation):
-    """Create a vectorizer for automatic embedding generation."""
-
     def __init__(
         self,
         source_table: str | None,
@@ -68,19 +66,15 @@ def __init__(
 
     @classmethod
     def create_vectorizer(cls, operations: Operations, source_table: str, **kw: Any):
-        """Issue a CREATE VECTORIZER command."""
         op = CreateVectorizerOp(source_table, **kw)
         return operations.invoke(op)
 
     @override
     def reverse(self) -> MigrateOperation:
-        """Creates the downgrade operation"""
         return DropVectorizerOp(None, True)
 
 
 class DropVectorizerOp(MigrateOperation):
-    """Drop a vectorizer and its associated objects."""
-
     def __init__(self, vectorizer_id: int | None, drop_all: bool):
         self.vectorizer_id = vectorizer_id
         self.drop_all = drop_all
@@ -92,28 +86,22 @@ def drop_vectorizer(
         vectorizer_id: int | None = None,
         drop_all: bool = True,
     ):
-        """Issue a DROP VECTORIZER command."""
         op = DropVectorizerOp(vectorizer_id, drop_all)
         return operations.invoke(op)
 
     @override
     def reverse(self) -> MigrateOperation:
-        """Creates the upgrade operation"""
         return CreateVectorizerOp(None)
 
 
 def create_vectorizer(operations: Operations, operation: CreateVectorizerOp):
-    """Implement CREATE VECTORIZER."""
     params = operation.params
     operations.execute(params.to_sql())
 
 
 def drop_vectorizer(operations: Operations, operation: DropVectorizerOp):
-    """Implement DROP VECTORIZER with cleanup of dependent objects."""
     connection = operations.get_bind()
     vectorizer_id = operation.vectorizer_id
-
-    # Finally drop the vectorizer itself
     connection.execute(
         text("SELECT ai.drop_vectorizer(:id, drop_all=>:drop_all)"),
         {"id": vectorizer_id, "drop_all": operation.drop_all},

diff --git a/projects/pgai/tests/vectorizer/extensions/fixtures/migrations/generic_vectorizer.py.template b/projects/pgai/tests/vectorizer/extensions/fixtures/migrations/generic_vectorizer.py.template
@@ -7,7 +7,7 @@ Create Date: {create_date}
 from alembic import op
 from sqlalchemy import Column, Integer, String, Text
 from sqlalchemy import text
-from pgai.vectorizer.configuration import (
+from pgai.alembic.configuration import (
     OpenAIConfig,
     OllamaConfig,
     VoyageAIConfig,

diff --git a/projects/pgai/tests/vectorizer/extensions/test_alembic.py b/projects/pgai/tests/vectorizer/extensions/test_alembic.py
@@ -7,7 +7,7 @@
 
 If you add new configuration options, add a test for them here.
 Either in an existing test case or a new one.
-The config classes for migrations lie in pgai/vectorizer/configuration.py
+The config classes for migrations lie in pgai/alembic/configuration.py
 If the migration and the validation definitions overlap,
 it is possible to define a parent class in base.py
 see e.g. CharacterTextSplitterConfig as an example
@@ -22,7 +22,7 @@
 from alembic.config import Config
 from sqlalchemy import Engine, text
 
-from pgai.vectorizer.configuration import (
+from pgai.alembic.configuration import (
     CharacterTextSplitterConfig,
     HNSWIndexingConfig,
     OllamaConfig,