sqlspec 0.17.1__py3-none-any.whl → 0.19.0__py3-none-any.whl

sqlspec/migrations/runner.py

@@ -5,13 +5,13 @@ This module handles migration file loading and execution using SQLFileLoader.
 
 import time
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Optional, cast
 
 from sqlspec.core.statement import SQL
 from sqlspec.migrations.base import BaseMigrationRunner
 from sqlspec.migrations.loaders import get_migration_loader
 from sqlspec.utils.logging import get_logger
-from sqlspec.utils.sync_tools import run_
+from sqlspec.utils.sync_tools import await_
 
 if TYPE_CHECKING:
     from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
@@ -22,7 +22,7 @@ logger = get_logger("migrations.runner")
 
 
 class SyncMigrationRunner(BaseMigrationRunner["SyncDriverAdapterBase"]):
-    """Executes migrations using SQLFileLoader."""
+    """Synchronous migration executor."""
 
     def get_migration_files(self) -> "list[tuple[str, Path]]":
         """Get all migration files sorted by version.
@@ -55,12 +55,15 @@ class SyncMigrationRunner(BaseMigrationRunner["SyncDriverAdapterBase"]):
         Returns:
             Tuple of (sql_content, execution_time_ms).
         """
-        upgrade_sql = self._get_migration_sql(migration, "up")
-        if upgrade_sql is None:
+        upgrade_sql_list = self._get_migration_sql(migration, "up")
+        if upgrade_sql_list is None:
             return None, 0
 
         start_time = time.time()
-        driver.execute(upgrade_sql)
+
+        for sql_statement in upgrade_sql_list:
+            if sql_statement.strip():
+                driver.execute_script(sql_statement)
         execution_time = int((time.time() - start_time) * 1000)
         return None, execution_time
 
@@ -76,12 +79,15 @@ class SyncMigrationRunner(BaseMigrationRunner["SyncDriverAdapterBase"]):
         Returns:
             Tuple of (sql_content, execution_time_ms).
         """
-        downgrade_sql = self._get_migration_sql(migration, "down")
-        if downgrade_sql is None:
+        downgrade_sql_list = self._get_migration_sql(migration, "down")
+        if downgrade_sql_list is None:
             return None, 0
 
         start_time = time.time()
-        driver.execute(downgrade_sql)
+
+        for sql_statement in downgrade_sql_list:
+            if sql_statement.strip():
+                driver.execute_script(sql_statement)
         execution_time = int((time.time() - start_time) * 1000)
         return None, execution_time
 
@@ -103,8 +109,8 @@ class SyncMigrationRunner(BaseMigrationRunner["SyncDriverAdapterBase"]):
                 loader = get_migration_loader(file_path, self.migrations_path, self.project_root)
 
                 try:
-                    up_sql = run_(loader.get_up_sql)(file_path)
-                    down_sql = run_(loader.get_down_sql)(file_path)
+                    up_sql = await_(loader.get_up_sql, raise_sync_error=False)(file_path)
+                    down_sql = await_(loader.get_down_sql, raise_sync_error=False)(file_path)
 
                     if up_sql:
                         all_queries[f"migrate-{version}-up"] = SQL(up_sql[0])
@@ -118,7 +124,7 @@ class SyncMigrationRunner(BaseMigrationRunner["SyncDriverAdapterBase"]):
 
 
 class AsyncMigrationRunner(BaseMigrationRunner["AsyncDriverAdapterBase"]):
-    """Executes migrations using SQLFileLoader."""
+    """Asynchronous migration executor."""
 
     async def get_migration_files(self) -> "list[tuple[str, Path]]":
         """Get all migration files sorted by version.
@@ -137,7 +143,80 @@ class AsyncMigrationRunner(BaseMigrationRunner["AsyncDriverAdapterBase"]):
         Returns:
             Dictionary containing migration metadata.
         """
-        return self._load_migration_metadata(file_path)
+        return await self._load_migration_metadata_async(file_path)
+
+    async def _load_migration_metadata_async(self, file_path: Path) -> "dict[str, Any]":
+        """Load migration metadata from file (async version).
+
+        Args:
+            file_path: Path to the migration file.
+
+        Returns:
+            Migration metadata dictionary.
+        """
+        loader = get_migration_loader(file_path, self.migrations_path, self.project_root)
+        loader.validate_migration_file(file_path)
+        content = file_path.read_text(encoding="utf-8")
+        checksum = self._calculate_checksum(content)
+        version = self._extract_version(file_path.name)
+        description = file_path.stem.split("_", 1)[1] if "_" in file_path.stem else ""
+
+        has_upgrade, has_downgrade = True, False
+
+        if file_path.suffix == ".sql":
+            up_query, down_query = f"migrate-{version}-up", f"migrate-{version}-down"
+            self.loader.clear_cache()
+            self.loader.load_sql(file_path)
+            has_upgrade, has_downgrade = self.loader.has_query(up_query), self.loader.has_query(down_query)
+        else:
+            try:
+                has_downgrade = bool(await loader.get_down_sql(file_path))
+            except Exception:
+                has_downgrade = False
+
+        return {
+            "version": version,
+            "description": description,
+            "file_path": file_path,
+            "checksum": checksum,
+            "has_upgrade": has_upgrade,
+            "has_downgrade": has_downgrade,
+            "loader": loader,
+        }
+
+    async def _get_migration_sql_async(self, migration: "dict[str, Any]", direction: str) -> "Optional[list[str]]":
+        """Get migration SQL for given direction (async version).
+
+        Args:
+            migration: Migration metadata.
+            direction: Either 'up' or 'down'.
+
+        Returns:
+            SQL statements for the migration.
+        """
+        if not migration.get(f"has_{direction}grade"):
+            if direction == "down":
+                logger.warning("Migration %s has no downgrade query", migration["version"])
+                return None
+            msg = f"Migration {migration['version']} has no upgrade query"
+            raise ValueError(msg)
+
+        file_path, loader = migration["file_path"], migration["loader"]
+
+        try:
+            method = loader.get_up_sql if direction == "up" else loader.get_down_sql
+            sql_statements = await method(file_path)
+
+        except Exception as e:
+            if direction == "down":
+                logger.warning("Failed to load downgrade for migration %s: %s", migration["version"], e)
+                return None
+            msg = f"Failed to load upgrade for migration {migration['version']}: {e}"
+            raise ValueError(msg) from e
+        else:
+            if sql_statements:
+                return cast("list[str]", sql_statements)
+            return None
 
     async def execute_upgrade(
         self, driver: "AsyncDriverAdapterBase", migration: "dict[str, Any]"
@@ -151,12 +230,15 @@ class AsyncMigrationRunner(BaseMigrationRunner["AsyncDriverAdapterBase"]):
         Returns:
             Tuple of (sql_content, execution_time_ms).
         """
-        upgrade_sql = self._get_migration_sql(migration, "up")
-        if upgrade_sql is None:
+        upgrade_sql_list = await self._get_migration_sql_async(migration, "up")
+        if upgrade_sql_list is None:
             return None, 0
 
         start_time = time.time()
-        await driver.execute(upgrade_sql)
+
+        for sql_statement in upgrade_sql_list:
+            if sql_statement.strip():
+                await driver.execute_script(sql_statement)
         execution_time = int((time.time() - start_time) * 1000)
         return None, execution_time
 
@@ -172,12 +254,15 @@ class AsyncMigrationRunner(BaseMigrationRunner["AsyncDriverAdapterBase"]):
         Returns:
             Tuple of (sql_content, execution_time_ms).
         """
-        downgrade_sql = self._get_migration_sql(migration, "down")
-        if downgrade_sql is None:
+        downgrade_sql_list = await self._get_migration_sql_async(migration, "down")
+        if downgrade_sql_list is None:
             return None, 0
 
         start_time = time.time()
-        await driver.execute(downgrade_sql)
+
+        for sql_statement in downgrade_sql_list:
+            if sql_statement.strip():
+                await driver.execute_script(sql_statement)
         execution_time = int((time.time() - start_time) * 1000)
         return None, execution_time
 

sqlspec/migrations/tracker.py

@@ -7,15 +7,18 @@ import os
 from typing import TYPE_CHECKING, Any, Optional
 
 from sqlspec.migrations.base import BaseMigrationTracker
+from sqlspec.utils.logging import get_logger
 
 if TYPE_CHECKING:
     from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
 
 __all__ = ("AsyncMigrationTracker", "SyncMigrationTracker")
 
+logger = get_logger("migrations.tracker")
+
 
 class SyncMigrationTracker(BaseMigrationTracker["SyncDriverAdapterBase"]):
-    """Tracks applied migrations in the database."""
+    """Synchronous migration version tracker."""
 
     def ensure_tracking_table(self, driver: "SyncDriverAdapterBase") -> None:
         """Create the migration tracking table if it doesn't exist.
@@ -24,6 +27,7 @@ class SyncMigrationTracker(BaseMigrationTracker["SyncDriverAdapterBase"]):
             driver: The database driver to use.
         """
         driver.execute(self._get_create_table_sql())
+        self._safe_commit(driver)
 
     def get_current_version(self, driver: "SyncDriverAdapterBase") -> Optional[str]:
         """Get the latest applied migration version.
@@ -66,6 +70,7 @@ class SyncMigrationTracker(BaseMigrationTracker["SyncDriverAdapterBase"]):
                 version, description, execution_time_ms, checksum, os.environ.get("USER", "unknown")
             )
         )
+        self._safe_commit(driver)
 
     def remove_migration(self, driver: "SyncDriverAdapterBase", version: str) -> None:
         """Remove a migration record (used during downgrade).
@@ -75,10 +80,30 @@ class SyncMigrationTracker(BaseMigrationTracker["SyncDriverAdapterBase"]):
             version: Version number to remove.
         """
         driver.execute(self._get_remove_migration_sql(version))
+        self._safe_commit(driver)
+
+    def _safe_commit(self, driver: "SyncDriverAdapterBase") -> None:
+        """Safely commit a transaction only if autocommit is disabled.
+
+        Args:
+            driver: The database driver to use.
+        """
+        try:
+            connection = getattr(driver, "connection", None)
+            if connection and hasattr(connection, "autocommit") and getattr(connection, "autocommit", False):
+                return
+
+            driver_features = getattr(driver, "driver_features", {})
+            if driver_features and driver_features.get("autocommit", False):
+                return
+
+            driver.commit()
+        except Exception:
+            logger.debug("Failed to commit transaction, likely due to autocommit being enabled")
 
 
 class AsyncMigrationTracker(BaseMigrationTracker["AsyncDriverAdapterBase"]):
-    """Tracks applied migrations in the database."""
+    """Asynchronous migration version tracker."""
 
     async def ensure_tracking_table(self, driver: "AsyncDriverAdapterBase") -> None:
         """Create the migration tracking table if it doesn't exist.
@@ -87,6 +112,7 @@ class AsyncMigrationTracker(BaseMigrationTracker["AsyncDriverAdapterBase"]):
             driver: The database driver to use.
         """
         await driver.execute(self._get_create_table_sql())
+        await self._safe_commit_async(driver)
 
     async def get_current_version(self, driver: "AsyncDriverAdapterBase") -> Optional[str]:
         """Get the latest applied migration version.
@@ -129,6 +155,7 @@ class AsyncMigrationTracker(BaseMigrationTracker["AsyncDriverAdapterBase"]):
                 version, description, execution_time_ms, checksum, os.environ.get("USER", "unknown")
             )
         )
+        await self._safe_commit_async(driver)
 
     async def remove_migration(self, driver: "AsyncDriverAdapterBase", version: str) -> None:
         """Remove a migration record (used during downgrade).
@@ -138,3 +165,23 @@ class AsyncMigrationTracker(BaseMigrationTracker["AsyncDriverAdapterBase"]):
             version: Version number to remove.
         """
         await driver.execute(self._get_remove_migration_sql(version))
+        await self._safe_commit_async(driver)
+
+    async def _safe_commit_async(self, driver: "AsyncDriverAdapterBase") -> None:
+        """Safely commit a transaction only if autocommit is disabled.
+
+        Args:
+            driver: The database driver to use.
+        """
+        try:
+            connection = getattr(driver, "connection", None)
+            if connection and hasattr(connection, "autocommit") and getattr(connection, "autocommit", False):
+                return
+
+            driver_features = getattr(driver, "driver_features", {})
+            if driver_features and driver_features.get("autocommit", False):
+                return
+
+            await driver.commit()
+        except Exception:
+            logger.debug("Failed to commit transaction, likely due to autocommit being enabled")

sqlspec/protocols.py

@@ -1,7 +1,7 @@
-"""Runtime-checkable protocols for SQLSpec to replace duck typing.
+"""Runtime-checkable protocols for type safety and runtime checks.
 
 This module provides protocols that can be used for static type checking
-and runtime isinstance() checks, replacing defensive hasattr() patterns.
+and runtime isinstance() checks.
 """
 
 from typing import TYPE_CHECKING, Any, ClassVar, Optional, Protocol, Union, runtime_checkable
@@ -192,10 +192,7 @@ class ObjectStoreItemProtocol(Protocol):
 
 @runtime_checkable
 class ObjectStoreProtocol(Protocol):
-    """Protocol for object storage operations.
-
-    Defines the interface for storage backends with both sync and async operations.
-    """
+    """Protocol for object storage operations."""
 
     capabilities: ClassVar["StorageCapabilities"]
 

sqlspec/storage/__init__.py

@@ -1,10 +1,10 @@
 """Storage abstraction layer for SQLSpec.
 
-Provides a flexible storage system with:
+Provides a storage system with:
 - Multiple backend support (local, fsspec, obstore)
-- Lazy loading and configuration-based registration
-- URI scheme-based automatic backend resolution
-- Key-based named storage configurations
+- Configuration-based registration
+- URI scheme-based backend resolution
+- Named storage configurations
 - Capability-based backend selection
 """
 

sqlspec/storage/backends/fsspec.py

@@ -32,7 +32,7 @@ class _ArrowStreamer:
         return self
 
     async def _initialize(self) -> None:
-        """Initialize the paths iterator."""
+        """Initialize paths iterator."""
         if self.paths_iterator is None:
             paths = await async_(self.backend.glob)(self.pattern, **self.kwargs)
             self.paths_iterator = iter(paths)
@@ -59,9 +59,8 @@ class _ArrowStreamer:
 class FSSpecBackend(ObjectStoreBase):
     """Storage backend using fsspec.
 
-    Implements the ObjectStoreProtocol using fsspec,
-    providing support for various protocols including HTTP, HTTPS, FTP,
-    and cloud storage services.
+    Implements the ObjectStoreProtocol using fsspec for various protocols
+    including HTTP, HTTPS, FTP, and cloud storage services.
     """
 
     _default_capabilities: ClassVar[StorageCapabilities] = StorageCapabilities(
@@ -116,7 +115,7 @@ class FSSpecBackend(ObjectStoreBase):
         return path_str
 
     def _detect_capabilities(self) -> StorageCapabilities:
-        """Detect capabilities based on underlying filesystem protocol."""
+        """Detect capabilities based on filesystem protocol."""
         protocol = self.protocol.lower()
 
         if protocol in {"s3", "s3a", "s3n"}:
@@ -138,7 +137,7 @@ class FSSpecBackend(ObjectStoreBase):
 
     @property
     def capabilities(self) -> StorageCapabilities:
-        """Return instance-specific capabilities based on detected protocol."""
+        """Return capabilities based on detected protocol."""
         return getattr(self, "_instance_capabilities", self.__class__._default_capabilities)
 
     @classmethod

sqlspec/storage/backends/obstore.py

@@ -1,8 +1,7 @@
 """Object storage backend using obstore.
 
-Implements the ObjectStoreProtocol using obstore,
-providing native support for S3, GCS, Azure, and local file storage
-with Arrow support.
+Implements the ObjectStoreProtocol using obstore for S3, GCS, Azure,
+and local file storage with Arrow support.
 """
 
 from __future__ import annotations
@@ -56,11 +55,11 @@ DEFAULT_OPTIONS: Final[dict[str, Any]] = {"connect_timeout": "30s", "request_tim
 class ObStoreBackend(ObjectStoreBase, HasStorageCapabilities):
     """Object storage backend using obstore.
 
-    Uses obstore's Rust-based implementation for storage operations,
-    providing native support for AWS S3, Google Cloud Storage, Azure Blob Storage,
+    Uses obstore's Rust-based implementation for storage operations.
+    Supports AWS S3, Google Cloud Storage, Azure Blob Storage,
     local filesystem, and HTTP endpoints.
 
-    Includes native Arrow support.
+    Includes Arrow support.
     """
 
     capabilities: ClassVar[StorageCapabilities] = StorageCapabilities(
@@ -203,7 +202,7 @@ class ObStoreBackend(ObjectStoreBase, HasStorageCapabilities):
             raise StorageOperationFailedError(msg) from exc
 
     def glob(self, pattern: str, **kwargs: Any) -> list[str]:
-        """Find objects matching pattern using obstore.
+        """Find objects matching pattern.
 
         Lists all objects and filters them client-side using the pattern.
         """
@@ -330,7 +329,7 @@ class ObStoreBackend(ObjectStoreBase, HasStorageCapabilities):
             raise StorageOperationFailedError(msg) from exc
 
     def stream_arrow(self, pattern: str, **kwargs: Any) -> Iterator[ArrowRecordBatch]:
-        """Stream Arrow record batches using obstore.
+        """Stream Arrow record batches.
 
         Yields:
             Iterator of Arrow record batches from matching objects.

sqlspec/storage/registry.py

@@ -1,8 +1,8 @@
-"""Unified Storage Registry for ObjectStore backends.
+"""Storage Registry for ObjectStore backends.
 
-Provides a flexible, lazy-loading storage registry that supports URI-first access
+Provides a storage registry that supports URI-first access
 pattern with automatic backend detection, ObStore preferred with FSSpec fallback,
-intelligent scheme-based routing, and named aliases for common configurations.
+scheme-based routing, and named aliases for common configurations.
 """
 
 import logging

sqlspec/utils/__init__.py

@@ -1,3 +1,3 @@
-from sqlspec.utils import deprecation, fixtures, module_loader, singleton, sync_tools, text, type_guards
+from sqlspec.utils import deprecation, fixtures, logging, module_loader, singleton, sync_tools, text, type_guards
 
-__all__ = ("deprecation", "fixtures", "module_loader", "singleton", "sync_tools", "text", "type_guards")
+__all__ = ("deprecation", "fixtures", "logging", "module_loader", "singleton", "sync_tools", "text", "type_guards")

sqlspec/utils/logging.py

@@ -5,23 +5,19 @@ Users should configure their own logging handlers and levels as needed.
 SQLSpec provides StructuredFormatter for JSON-formatted logs if desired.
 """
 
-from __future__ import annotations
-
 import logging
 from contextvars import ContextVar
-from typing import TYPE_CHECKING, Any
+from logging import LogRecord
+from typing import Any, Optional
 
 from sqlspec._serialization import encode_json
 
-if TYPE_CHECKING:
-    from logging import LogRecord
-
 __all__ = ("StructuredFormatter", "correlation_id_var", "get_correlation_id", "get_logger", "set_correlation_id")
 
-correlation_id_var: ContextVar[str | None] = ContextVar("correlation_id", default=None)
+correlation_id_var: "ContextVar[Optional[str]]" = ContextVar("correlation_id", default=None)
 
 
-def set_correlation_id(correlation_id: str | None) -> None:
+def set_correlation_id(correlation_id: "Optional[str]") -> None:
     """Set the correlation ID for the current context.
 
     Args:
@@ -30,7 +26,7 @@ def set_correlation_id(correlation_id: str | None) -> None:
     correlation_id_var.set(correlation_id)
 
 
-def get_correlation_id() -> str | None:
+def get_correlation_id() -> "Optional[str]":
     """Get the current correlation ID.
 
     Returns:
@@ -90,7 +86,7 @@ class CorrelationIDFilter(logging.Filter):
         return True
 
 
-def get_logger(name: str | None = None) -> logging.Logger:
+def get_logger(name: "Optional[str]" = None) -> logging.Logger:
     """Get a logger instance with standardized configuration.
 
     Args:

sqlspec/utils/sync_tools.py

@@ -1,3 +1,10 @@
+"""Utilities for async/sync interoperability in SQLSpec.
+
+This module provides utilities for converting between async and sync functions,
+managing concurrency limits, and handling context managers. Used primarily
+for adapter implementations that need to support both sync and async patterns.
+"""
+
 import asyncio
 import functools
 import inspect
@@ -26,7 +33,15 @@ class CapacityLimiter:
     """Limits the number of concurrent operations using a semaphore."""
 
     def __init__(self, total_tokens: int) -> None:
-        self._semaphore = asyncio.Semaphore(total_tokens)
+        self._total_tokens = total_tokens
+        self._semaphore_instance: Optional[asyncio.Semaphore] = None
+
+    @property
+    def _semaphore(self) -> asyncio.Semaphore:
+        """Lazy initialization of asyncio.Semaphore for Python 3.9 compatibility."""
+        if self._semaphore_instance is None:
+            self._semaphore_instance = asyncio.Semaphore(self._total_tokens)
+        return self._semaphore_instance
 
     async def acquire(self) -> None:
         await self._semaphore.acquire()
@@ -36,11 +51,12 @@ class CapacityLimiter:
 
     @property
     def total_tokens(self) -> int:
-        return self._semaphore._value
+        return self._total_tokens
 
     @total_tokens.setter
     def total_tokens(self, value: int) -> None:
-        self._semaphore = asyncio.Semaphore(value)
+        self._total_tokens = value
+        self._semaphore_instance = None
 
     async def __aenter__(self) -> None:
         await self.acquire()
@@ -76,7 +92,14 @@ def run_(async_function: "Callable[ParamSpecT, Coroutine[Any, Any, ReturnT]]") -
             loop = None
 
         if loop is not None:
-            return asyncio.run(partial_f())
+            if loop.is_running():
+                import concurrent.futures
+
+                with concurrent.futures.ThreadPoolExecutor() as executor:
+                    future = executor.submit(asyncio.run, partial_f())
+                    return future.result()
+            else:
+                return asyncio.run(partial_f())
         if uvloop and sys.platform != "win32":
             uvloop.install()  # pyright: ignore[reportUnknownMemberType]
         return asyncio.run(partial_f())

sqlspec/utils/text.py

@@ -1,4 +1,9 @@
-"""General utility functions."""
+"""Text processing utilities for SQLSpec.
+
+Provides functions for string manipulation including case conversion,
+slugification, and email validation. Used primarily for identifier
+generation and data validation.
+"""
 
 import re
 import unicodedata

{sqlspec-0.17.1.dist-info → sqlspec-0.19.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlspec
-Version: 0.17.1
+Version: 0.19.0
 Summary: SQL Experiments in Python
 Project-URL: Discord, https://discord.gg/litestar
 Project-URL: Issue, https://github.com/litestar-org/sqlspec/issues/