fix(migrations): various migration bugfixes and increased test coverage (#57)

Improve migration tests and fix issues related to tracking and execution. Enhance connection pool management by adding a method to close the pool and ensure proper execution in running loops.

Author: cofin

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
       - id: mixed-line-ending
       - id: trailing-whitespace
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: "v0.12.8"
+    rev: "v0.12.9"
     hooks:
       - id: ruff
         args: ["--fix"]

pyproject.toml

@@ -13,7 +13,7 @@ maintainers = [{ name = "Litestar Developers", email = "[email protected]" }]
 name = "sqlspec"
 readme = "README.md"
 requires-python = ">=3.9, <4.0"
-version = "0.17.1"
+version = "0.18.0"
 
 [project.urls]
 Discord = "https://discord.gg/litestar"

sqlspec/_sql.py

@@ -42,11 +42,11 @@
 )
 from sqlspec.builder.mixins._join_operations import JoinBuilder
 from sqlspec.builder.mixins._select_operations import Case, SubqueryBuilder, WindowFunctionBuilder
+from sqlspec.core.statement import SQL
 from sqlspec.exceptions import SQLBuilderError
 
 if TYPE_CHECKING:
     from sqlspec.builder._expression_wrappers import ExpressionWrapper
-    from sqlspec.core.statement import SQL
 
 
 __all__ = (
@@ -285,9 +285,7 @@ def create_table(self, table_name: str, dialect: DialectType = None) -> "CreateT
         Returns:
             CreateTable builder instance
         """
-        builder = CreateTable(table_name)
-        builder.dialect = dialect or self.dialect
-        return builder
+        return CreateTable(table_name, dialect=dialect or self.dialect)
 
     def create_table_as_select(self, dialect: DialectType = None) -> "CreateTableAsSelect":
         """Create a CREATE TABLE AS SELECT builder.
@@ -298,35 +296,31 @@ def create_table_as_select(self, dialect: DialectType = None) -> "CreateTableAsS
         Returns:
             CreateTableAsSelect builder instance
         """
-        builder = CreateTableAsSelect()
-        builder.dialect = dialect or self.dialect
-        return builder
+        return CreateTableAsSelect(dialect=dialect or self.dialect)
 
-    def create_view(self, dialect: DialectType = None) -> "CreateView":
+    def create_view(self, view_name: str, dialect: DialectType = None) -> "CreateView":
         """Create a CREATE VIEW builder.
 
         Args:
+            view_name: Name of the view to create
             dialect: Optional SQL dialect
 
         Returns:
             CreateView builder instance
         """
-        builder = CreateView()
-        builder.dialect = dialect or self.dialect
-        return builder
+        return CreateView(view_name, dialect=dialect or self.dialect)
 
-    def create_materialized_view(self, dialect: DialectType = None) -> "CreateMaterializedView":
+    def create_materialized_view(self, view_name: str, dialect: DialectType = None) -> "CreateMaterializedView":
         """Create a CREATE MATERIALIZED VIEW builder.
 
         Args:
+            view_name: Name of the materialized view to create
             dialect: Optional SQL dialect
 
         Returns:
             CreateMaterializedView builder instance
         """
-        builder = CreateMaterializedView()
-        builder.dialect = dialect or self.dialect
-        return builder
+        return CreateMaterializedView(view_name, dialect=dialect or self.dialect)
 
     def create_index(self, index_name: str, dialect: DialectType = None) -> "CreateIndex":
         """Create a CREATE INDEX builder.
@@ -340,18 +334,17 @@ def create_index(self, index_name: str, dialect: DialectType = None) -> "CreateI
         """
         return CreateIndex(index_name, dialect=dialect or self.dialect)
 
-    def create_schema(self, dialect: DialectType = None) -> "CreateSchema":
+    def create_schema(self, schema_name: str, dialect: DialectType = None) -> "CreateSchema":
         """Create a CREATE SCHEMA builder.
 
         Args:
+            schema_name: Name of the schema to create
             dialect: Optional SQL dialect
 
         Returns:
             CreateSchema builder instance
         """
-        builder = CreateSchema()
-        builder.dialect = dialect or self.dialect
-        return builder
+        return CreateSchema(schema_name, dialect=dialect or self.dialect)
 
     def drop_table(self, table_name: str, dialect: DialectType = None) -> "DropTable":
         """Create a DROP TABLE builder.
@@ -365,16 +358,17 @@ def drop_table(self, table_name: str, dialect: DialectType = None) -> "DropTable
         """
         return DropTable(table_name, dialect=dialect or self.dialect)
 
-    def drop_view(self, dialect: DialectType = None) -> "DropView":
+    def drop_view(self, view_name: str, dialect: DialectType = None) -> "DropView":
         """Create a DROP VIEW builder.
 
         Args:
+            view_name: Name of the view to drop
             dialect: Optional SQL dialect
 
         Returns:
             DropView builder instance
         """
-        return DropView(dialect=dialect or self.dialect)
+        return DropView(view_name, dialect=dialect or self.dialect)
 
     def drop_index(self, index_name: str, dialect: DialectType = None) -> "DropIndex":
         """Create a DROP INDEX builder.
@@ -388,16 +382,17 @@ def drop_index(self, index_name: str, dialect: DialectType = None) -> "DropIndex
         """
         return DropIndex(index_name, dialect=dialect or self.dialect)
 
-    def drop_schema(self, dialect: DialectType = None) -> "DropSchema":
+    def drop_schema(self, schema_name: str, dialect: DialectType = None) -> "DropSchema":
         """Create a DROP SCHEMA builder.
 
         Args:
+            schema_name: Name of the schema to drop
             dialect: Optional SQL dialect
 
         Returns:
             DropSchema builder instance
         """
-        return DropSchema(dialect=dialect or self.dialect)
+        return DropSchema(schema_name, dialect=dialect or self.dialect)
 
     def alter_table(self, table_name: str, dialect: DialectType = None) -> "AlterTable":
         """Create an ALTER TABLE builder.
@@ -409,22 +404,19 @@ def alter_table(self, table_name: str, dialect: DialectType = None) -> "AlterTab
         Returns:
             AlterTable builder instance
         """
-        builder = AlterTable(table_name)
-        builder.dialect = dialect or self.dialect
-        return builder
+        return AlterTable(table_name, dialect=dialect or self.dialect)
 
-    def rename_table(self, dialect: DialectType = None) -> "RenameTable":
+    def rename_table(self, old_name: str, dialect: DialectType = None) -> "RenameTable":
         """Create a RENAME TABLE builder.
 
         Args:
+            old_name: Current name of the table
             dialect: Optional SQL dialect
 
         Returns:
             RenameTable builder instance
         """
-        builder = RenameTable()
-        builder.dialect = dialect or self.dialect
-        return builder
+        return RenameTable(old_name, dialect=dialect or self.dialect)
 
     def comment_on(self, dialect: DialectType = None) -> "CommentOn":
         """Create a COMMENT ON builder.
@@ -435,9 +427,7 @@ def comment_on(self, dialect: DialectType = None) -> "CommentOn":
         Returns:
             CommentOn builder instance
         """
-        builder = CommentOn()
-        builder.dialect = dialect or self.dialect
-        return builder
+        return CommentOn(dialect=dialect or self.dialect)
 
     # ===================
     # SQL Analysis Helpers
@@ -746,7 +736,6 @@ def raw(sql_fragment: str, **parameters: Any) -> "Union[exp.Expression, SQL]":
                 raise SQLBuilderError(msg) from e
 
         # New behavior - return SQL statement with parameters
-        from sqlspec.core.statement import SQL
 
         return SQL(sql_fragment, parameters)
 
@@ -1331,9 +1320,7 @@ def truncate(self, table_name: str) -> "Truncate":
             )
             ```
         """
-        builder = Truncate(dialect=self.dialect)
-        builder._table_name = table_name
-        return builder
+        return Truncate(table_name, dialect=self.dialect)
 
     # ===================
     # Case Expressions

sqlspec/adapters/adbc/driver.py

@@ -48,7 +48,7 @@
     "postgres": (ParameterStyle.NUMERIC, [ParameterStyle.NUMERIC]),
     "postgresql": (ParameterStyle.NUMERIC, [ParameterStyle.NUMERIC]),
     "bigquery": (ParameterStyle.NAMED_AT, [ParameterStyle.NAMED_AT]),
-    "sqlite": (ParameterStyle.QMARK, [ParameterStyle.QMARK, ParameterStyle.NAMED_COLON]),
+    "sqlite": (ParameterStyle.QMARK, [ParameterStyle.QMARK]),
     "duckdb": (ParameterStyle.QMARK, [ParameterStyle.QMARK, ParameterStyle.NUMERIC, ParameterStyle.NAMED_DOLLAR]),
     "mysql": (ParameterStyle.POSITIONAL_PYFORMAT, [ParameterStyle.POSITIONAL_PYFORMAT, ParameterStyle.NAMED_PYFORMAT]),
     "snowflake": (ParameterStyle.QMARK, [ParameterStyle.QMARK, ParameterStyle.NUMERIC]),

sqlspec/adapters/aiosqlite/pool.py

@@ -123,15 +123,15 @@ class AiosqliteConnectionPool:
     """Multi-connection pool for aiosqlite with proper shutdown handling."""
 
     __slots__ = (
-        "_closed_event",
+        "_closed_event_instance",
         "_connect_timeout",
         "_connection_parameters",
         "_connection_registry",
         "_idle_timeout",
-        "_lock",
+        "_lock_instance",
         "_operation_timeout",
         "_pool_size",
-        "_queue",
+        "_queue_instance",
         "_tracked_threads",
         "_wal_initialized",
     )
@@ -159,21 +159,44 @@ def __init__(
         self._idle_timeout = idle_timeout
         self._operation_timeout = operation_timeout
 
-        self._queue: asyncio.Queue[AiosqlitePoolConnection] = asyncio.Queue(maxsize=pool_size)
         self._connection_registry: dict[str, AiosqlitePoolConnection] = {}
-        self._lock = asyncio.Lock()
-        self._closed_event = asyncio.Event()
         self._tracked_threads: set[Union[threading.Thread, AiosqliteConnection]] = set()
         self._wal_initialized = False
 
+        # Lazy initialization for Python 3.9 compatibility (asyncio objects can't be created without event loop)
+        self._queue_instance: Optional[asyncio.Queue[AiosqlitePoolConnection]] = None
+        self._lock_instance: Optional[asyncio.Lock] = None
+        self._closed_event_instance: Optional[asyncio.Event] = None
+
+    @property
+    def _queue(self) -> "asyncio.Queue[AiosqlitePoolConnection]":
+        """Lazy initialization of asyncio.Queue for Python 3.9 compatibility."""
+        if self._queue_instance is None:
+            self._queue_instance = asyncio.Queue(maxsize=self._pool_size)
+        return self._queue_instance
+
+    @property
+    def _lock(self) -> asyncio.Lock:
+        """Lazy initialization of asyncio.Lock for Python 3.9 compatibility."""
+        if self._lock_instance is None:
+            self._lock_instance = asyncio.Lock()
+        return self._lock_instance
+
+    @property
+    def _closed_event(self) -> asyncio.Event:
+        """Lazy initialization of asyncio.Event for Python 3.9 compatibility."""
+        if self._closed_event_instance is None:
+            self._closed_event_instance = asyncio.Event()
+        return self._closed_event_instance
+
     @property
     def is_closed(self) -> bool:
         """Check if pool is closed.
 
         Returns:
             True if pool is closed
         """
-        return self._closed_event.is_set()
+        return self._closed_event_instance is not None and self._closed_event.is_set()
 
     def size(self) -> int:
         """Get total number of connections in pool.
@@ -189,6 +212,8 @@ def checked_out(self) -> int:
         Returns:
             Number of connections currently in use
         """
+        if self._queue_instance is None:
+            return len(self._connection_registry)
         return len(self._connection_registry) - self._queue.qsize()
 
     def _track_aiosqlite_thread(self, connection: "AiosqliteConnection") -> None:

sqlspec/adapters/asyncmy/config.py

@@ -107,7 +107,11 @@ async def _create_pool(self) -> "Pool":  # pyright: ignore
     async def _close_pool(self) -> None:
         """Close the actual async connection pool."""
         if self.pool_instance:
-            await self.pool_instance.close()
+            self.pool_instance.close()
+
+    async def close_pool(self) -> None:
+        """Close the connection pool."""
+        await self._close_pool()
 
     async def create_connection(self) -> AsyncmyConnection:  # pyright: ignore
         """Create a single async connection (not from pool).

sqlspec/adapters/asyncpg/config.py

@@ -144,6 +144,10 @@ async def _close_pool(self) -> None:
         if self.pool_instance:
             await self.pool_instance.close()
 
+    async def close_pool(self) -> None:
+        """Close the connection pool."""
+        await self._close_pool()
+
     async def create_connection(self) -> "AsyncpgConnection":
         """Create a single async connection from the pool.
 

sqlspec/adapters/duckdb/driver.py

@@ -45,11 +45,7 @@
         default_parameter_style=ParameterStyle.QMARK,
         supported_parameter_styles={ParameterStyle.QMARK, ParameterStyle.NUMERIC, ParameterStyle.NAMED_DOLLAR},
         default_execution_parameter_style=ParameterStyle.QMARK,
-        supported_execution_parameter_styles={
-            ParameterStyle.QMARK,
-            ParameterStyle.NUMERIC,
-            ParameterStyle.NAMED_DOLLAR,
-        },
+        supported_execution_parameter_styles={ParameterStyle.QMARK, ParameterStyle.NUMERIC},
         type_coercion_map={},
         has_native_list_expansion=True,
         needs_static_script_compilation=False,

sqlspec/adapters/oracledb/config.py

@@ -239,6 +239,10 @@ async def _close_pool(self) -> None:
         if self.pool_instance:
             await self.pool_instance.close()
 
+    async def close_pool(self) -> None:
+        """Close the connection pool."""
+        await self._close_pool()
+
     async def create_connection(self) -> OracleAsyncConnection:
         """Create a single async connection (not from pool).