sqlspec 0.14.1__py3-none-any.whl → 0.16.0__py3-none-any.whl

sqlspec/migrations/tracker.py

@@ -3,20 +3,21 @@
 This module provides functionality to track applied migrations in the database.
 """
 
+import os
 from typing import TYPE_CHECKING, Any, Optional
 
 from sqlspec.migrations.base import BaseMigrationTracker
 
 if TYPE_CHECKING:
-    from sqlspec.driver import AsyncDriverAdapterProtocol, SyncDriverAdapterProtocol
+    from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
 
 __all__ = ("AsyncMigrationTracker", "SyncMigrationTracker")
 
 
-class SyncMigrationTracker(BaseMigrationTracker["SyncDriverAdapterProtocol[Any]"]):
-    """Sync version - tracks applied migrations in the database."""
+class SyncMigrationTracker(BaseMigrationTracker["SyncDriverAdapterBase"]):
+    """Tracks applied migrations in the database."""
 
-    def ensure_tracking_table(self, driver: "SyncDriverAdapterProtocol[Any]") -> None:
+    def ensure_tracking_table(self, driver: "SyncDriverAdapterBase") -> None:
         """Create the migration tracking table if it doesn't exist.
 
         Args:
@@ -24,7 +25,7 @@ class SyncMigrationTracker(BaseMigrationTracker["SyncDriverAdapterProtocol[Any]"
         """
         driver.execute(self._get_create_table_sql())
 
-    def get_current_version(self, driver: "SyncDriverAdapterProtocol[Any]") -> Optional[str]:
+    def get_current_version(self, driver: "SyncDriverAdapterBase") -> Optional[str]:
         """Get the latest applied migration version.
 
         Args:
@@ -36,7 +37,7 @@ class SyncMigrationTracker(BaseMigrationTracker["SyncDriverAdapterProtocol[Any]"
         result = driver.execute(self._get_current_version_sql())
         return result.data[0]["version_num"] if result.data else None
 
-    def get_applied_migrations(self, driver: "SyncDriverAdapterProtocol[Any]") -> "list[dict[str, Any]]":
+    def get_applied_migrations(self, driver: "SyncDriverAdapterBase") -> "list[dict[str, Any]]":
         """Get all applied migrations in order.
 
         Args:
@@ -46,15 +47,10 @@ class SyncMigrationTracker(BaseMigrationTracker["SyncDriverAdapterProtocol[Any]"
             List of migration records.
         """
         result = driver.execute(self._get_applied_migrations_sql())
-        return result.data
+        return result.data or []
 
     def record_migration(
-        self,
-        driver: "SyncDriverAdapterProtocol[Any]",
-        version: str,
-        description: str,
-        execution_time_ms: int,
-        checksum: str,
+        self, driver: "SyncDriverAdapterBase", version: str, description: str, execution_time_ms: int, checksum: str
     ) -> None:
         """Record a successfully applied migration.
 
@@ -64,41 +60,39 @@ class SyncMigrationTracker(BaseMigrationTracker["SyncDriverAdapterProtocol[Any]"
             description: Description of the migration.
             execution_time_ms: Execution time in milliseconds.
             checksum: MD5 checksum of the migration content.
-            connection: Optional connection to use for the operation.
         """
-        import os
-
-        applied_by = os.environ.get("USER", "unknown")
-
-        driver.execute(self._get_record_migration_sql(version, description, execution_time_ms, checksum, applied_by))
+        driver.execute(
+            self._get_record_migration_sql(
+                version, description, execution_time_ms, checksum, os.environ.get("USER", "unknown")
+            )
+        )
 
-    def remove_migration(self, driver: "SyncDriverAdapterProtocol[Any]", version: str) -> None:
+    def remove_migration(self, driver: "SyncDriverAdapterBase", version: str) -> None:
         """Remove a migration record (used during downgrade).
 
         Args:
             driver: The database driver to use.
             version: Version number to remove.
-            connection: Optional connection to use for the operation.
         """
         driver.execute(self._get_remove_migration_sql(version))
 
 
-class AsyncMigrationTracker(BaseMigrationTracker["AsyncDriverAdapterProtocol[Any]"]):
-    """Async version - tracks applied migrations in the database."""
+class AsyncMigrationTracker(BaseMigrationTracker["AsyncDriverAdapterBase"]):
+    """Tracks applied migrations in the database."""
 
-    async def ensure_tracking_table(self, driver: "AsyncDriverAdapterProtocol[Any]") -> None:
+    async def ensure_tracking_table(self, driver: "AsyncDriverAdapterBase") -> None:
         """Create the migration tracking table if it doesn't exist.
 
         Args:
-            driver: The async database driver to use.
+            driver: The database driver to use.
         """
         await driver.execute(self._get_create_table_sql())
 
-    async def get_current_version(self, driver: "AsyncDriverAdapterProtocol[Any]") -> Optional[str]:
+    async def get_current_version(self, driver: "AsyncDriverAdapterBase") -> Optional[str]:
         """Get the latest applied migration version.
 
         Args:
-            driver: The async database driver to use.
+            driver: The database driver to use.
 
         Returns:
             The current version number or None if no migrations applied.
@@ -106,48 +100,41 @@ class AsyncMigrationTracker(BaseMigrationTracker["AsyncDriverAdapterProtocol[Any
         result = await driver.execute(self._get_current_version_sql())
         return result.data[0]["version_num"] if result.data else None
 
-    async def get_applied_migrations(self, driver: "AsyncDriverAdapterProtocol[Any]") -> "list[dict[str, Any]]":
+    async def get_applied_migrations(self, driver: "AsyncDriverAdapterBase") -> "list[dict[str, Any]]":
         """Get all applied migrations in order.
 
         Args:
-            driver: The async database driver to use.
+            driver: The database driver to use.
 
         Returns:
             List of migration records.
         """
         result = await driver.execute(self._get_applied_migrations_sql())
-        return result.data
+        return result.data or []
 
     async def record_migration(
-        self,
-        driver: "AsyncDriverAdapterProtocol[Any]",
-        version: str,
-        description: str,
-        execution_time_ms: int,
-        checksum: str,
+        self, driver: "AsyncDriverAdapterBase", version: str, description: str, execution_time_ms: int, checksum: str
     ) -> None:
         """Record a successfully applied migration.
 
         Args:
-            driver: The async database driver to use.
+            driver: The database driver to use.
             version: Version number of the migration.
             description: Description of the migration.
             execution_time_ms: Execution time in milliseconds.
             checksum: MD5 checksum of the migration content.
         """
-        import os
-
-        applied_by = os.environ.get("USER", "unknown")
-
         await driver.execute(
-            self._get_record_migration_sql(version, description, execution_time_ms, checksum, applied_by)
+            self._get_record_migration_sql(
+                version, description, execution_time_ms, checksum, os.environ.get("USER", "unknown")
+            )
         )
 
-    async def remove_migration(self, driver: "AsyncDriverAdapterProtocol[Any]", version: str) -> None:
+    async def remove_migration(self, driver: "AsyncDriverAdapterBase", version: str) -> None:
         """Remove a migration record (used during downgrade).
 
         Args:
-            driver: The async database driver to use.
+            driver: The database driver to use.
             version: Version number to remove.
         """
         await driver.execute(self._get_remove_migration_sql(version))

sqlspec/migrations/utils.py

@@ -9,50 +9,96 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional
 
 if TYPE_CHECKING:
-    from sqlspec.driver import AsyncDriverAdapterProtocol
+    from sqlspec.driver import AsyncDriverAdapterBase
 
 __all__ = ("create_migration_file", "drop_all", "get_author")
 
 
-def create_migration_file(migrations_dir: Path, version: str, message: str) -> Path:
+def create_migration_file(migrations_dir: Path, version: str, message: str, file_type: str = "sql") -> Path:
     """Create a new migration file from template.
 
     Args:
         migrations_dir: Directory to create the migration in.
         version: Version number for the migration.
         message: Description message for the migration.
+        file_type: Type of migration file to create ('sql' or 'py').
 
     Returns:
         Path to the created migration file.
     """
-    # Sanitize message for filename
     safe_message = message.lower()
     safe_message = "".join(c if c.isalnum() or c in " -" else "" for c in safe_message)
     safe_message = safe_message.replace(" ", "_").replace("-", "_")
     safe_message = "_".join(filter(None, safe_message.split("_")))[:50]
 
-    filename = f"{version}_{safe_message}.sql"
-    file_path = migrations_dir / filename
+    if file_type == "py":
+        filename = f"{version}_{safe_message}.py"
+        file_path = migrations_dir / filename
+        template = f'''"""SQLSpec Migration - {message}
 
-    # Generate template content
-    template = f"""-- SQLSpec Migration
+Version: {version}
+Created: {datetime.now(timezone.utc).isoformat()}
+Author: {get_author()}
+
+Migration functions can use either naming convention:
+- Preferred: up()/down()
+- Legacy: migrate_up()/migrate_down()
+
+Both can be synchronous or asynchronous:
+- def up(): ...
+- async def up(): ...
+"""
+
+from typing import List, Union
+
+
+def up() -> Union[str, List[str]]:
+    """Apply the migration (upgrade).
+
+    Returns:
+        SQL statement(s) to execute for upgrade.
+        Can return a single string or list of strings.
+
+    Note: You can use either 'up()' or 'migrate_up()' for function names.
+    Both support async versions: 'async def up()' or 'async def migrate_up()'
+    """
+    return """
+    CREATE TABLE example (
+        id INTEGER PRIMARY KEY,
+        name TEXT NOT NULL
+    );
+    """
+
+
+def down() -> Union[str, List[str]]:
+    """Reverse the migration (downgrade).
+
+    Returns:
+        SQL statement(s) to execute for downgrade.
+        Can return a single string or list of strings.
+        Return empty string or empty list if downgrade is not supported.
+
+    Note: You can use either 'down()' or 'migrate_down()' for function names.
+    Both support async versions: 'async def down()' or 'async def migrate_down()'
+    """
+    return "DROP TABLE example;"
+'''
+    else:
+        filename = f"{version}_{safe_message}.sql"
+        file_path = migrations_dir / filename
+        template = f"""-- SQLSpec Migration
 -- Version: {version}
 -- Description: {message}
 -- Created: {datetime.now(timezone.utc).isoformat()}
 -- Author: {get_author()}
 
 -- name: migrate-{version}-up
--- TODO: Add your upgrade SQL statements here
--- Example:
--- CREATE TABLE example (
---     id INTEGER PRIMARY KEY,
---     name TEXT NOT NULL
--- );
+CREATE TABLE placeholder (
+    id INTEGER PRIMARY KEY
+);
 
 -- name: migrate-{version}-down
--- TODO: Add your downgrade SQL statements here (optional)
--- Example:
--- DROP TABLE example;
+DROP TABLE placeholder;
 """
 
     file_path.write_text(template)
@@ -68,22 +114,16 @@ def get_author() -> str:
     return os.environ.get("USER", "unknown")
 
 
-async def drop_all(
-    engine: "AsyncDriverAdapterProtocol[Any]", version_table_name: str, metadata: Optional[Any] = None
-) -> None:
+async def drop_all(engine: "AsyncDriverAdapterBase", version_table_name: str, metadata: Optional[Any] = None) -> None:
     """Drop all tables from the database.
 
-    This is a placeholder for database-specific implementations.
-
     Args:
         engine: The database engine/driver.
         version_table_name: Name of the version tracking table.
         metadata: Optional metadata object.
 
     Raises:
-        NotImplementedError: Always, as this requires database-specific logic.
+        NotImplementedError: Always raised.
     """
-    # This would need database-specific implementation
-    # For now, it's a placeholder
     msg = "drop_all functionality requires database-specific implementation"
     raise NotImplementedError(msg)

sqlspec/protocols.py

@@ -4,7 +4,6 @@ This module provides protocols that can be used for static type checking
 and runtime isinstance() checks, replacing defensive hasattr() patterns.
 """
 
-from abc import abstractmethod
 from typing import TYPE_CHECKING, Any, ClassVar, Optional, Protocol, Union, runtime_checkable
 
 from typing_extensions import Self
@@ -15,17 +14,10 @@ if TYPE_CHECKING:
 
     from sqlglot import exp
 
-    from sqlspec.statement.pipelines.context import SQLProcessingContext
     from sqlspec.storage.capabilities import StorageCapabilities
     from sqlspec.typing import ArrowRecordBatch, ArrowTable
 
 __all__ = (
-    # Database Connection Protocols
-    "AsyncCloseableConnectionProtocol",
-    "AsyncCopyCapableConnectionProtocol",
-    "AsyncPipelineCapableDriverProtocol",
-    "AsyncTransactionCapableConnectionProtocol",
-    "AsyncTransactionStateConnectionProtocol",
     "BytesConvertibleProtocol",
     "DictProtocol",
     "FilterAppenderProtocol",
@@ -36,7 +28,6 @@ __all__ = (
     "HasOffsetProtocol",
     "HasOrderByProtocol",
     "HasParameterBuilderProtocol",
-    "HasRiskLevelProtocol",
     "HasSQLGlotExpressionProtocol",
     "HasSQLMethodProtocol",
     "HasToStatementProtocol",
@@ -46,14 +37,8 @@ __all__ = (
     "ObjectStoreItemProtocol",
     "ObjectStoreProtocol",
     "ParameterValueProtocol",
-    "ProcessorProtocol",
     "SQLBuilderProtocol",
     "SelectBuilderProtocol",
-    "SyncCloseableConnectionProtocol",
-    "SyncCopyCapableConnectionProtocol",
-    "SyncPipelineCapableDriverProtocol",
-    "SyncTransactionCapableConnectionProtocol",
-    "SyncTransactionStateConnectionProtocol",
     "WithMethodProtocol",
 )
 
@@ -171,16 +156,6 @@ class ParameterValueProtocol(Protocol):
     type_hint: str
 
 
-@runtime_checkable
-class HasRiskLevelProtocol(Protocol):
-    """Protocol for objects with a risk_level attribute."""
-
-    @property
-    def risk_level(self) -> Any:
-        """Get the risk level of this object."""
-        ...
-
-
 @runtime_checkable
 class DictProtocol(Protocol):
     """Protocol for objects with a __dict__ attribute."""
@@ -188,25 +163,6 @@ class DictProtocol(Protocol):
     __dict__: dict[str, Any]
 
 
-class ProcessorProtocol(Protocol):
-    """Defines the interface for a single processing step in the SQL pipeline."""
-
-    @abstractmethod
-    def process(
-        self, expression: "Optional[exp.Expression]", context: "SQLProcessingContext"
-    ) -> "Optional[exp.Expression]":
-        """Processes an SQL expression.
-
-        Args:
-            expression: The SQL expression to process.
-            context: The SQLProcessingContext holding the current state and config.
-
-        Returns:
-            The (possibly modified) SQL expression for transformers, or None for validators/analyzers.
-        """
-        ...
-
-
 @runtime_checkable
 class BytesConvertibleProtocol(Protocol):
     """Protocol for objects that can be converted to bytes."""
@@ -224,138 +180,20 @@ class ObjectStoreItemProtocol(Protocol):
     key: "Optional[str]"
 
 
-@runtime_checkable
-class SyncTransactionCapableConnectionProtocol(Protocol):
-    """Protocol for sync connections that support transactions."""
-
-    def commit(self) -> None:
-        """Commit the current transaction."""
-        ...
-
-    def rollback(self) -> None:
-        """Rollback the current transaction."""
-        ...
-
-
-@runtime_checkable
-class AsyncTransactionCapableConnectionProtocol(Protocol):
-    """Protocol for async connections that support transactions."""
-
-    async def commit(self) -> None:
-        """Commit the current transaction."""
-        ...
-
-    async def rollback(self) -> None:
-        """Rollback the current transaction."""
-        ...
-
-
-@runtime_checkable
-class SyncTransactionStateConnectionProtocol(SyncTransactionCapableConnectionProtocol, Protocol):
-    """Protocol for sync connections that can report transaction state."""
-
-    def in_transaction(self) -> bool:
-        """Check if connection is currently in a transaction."""
-        ...
-
-    def begin(self) -> None:
-        """Begin a new transaction."""
-        ...
-
-
-@runtime_checkable
-class AsyncTransactionStateConnectionProtocol(AsyncTransactionCapableConnectionProtocol, Protocol):
-    """Protocol for async connections that can report transaction state."""
-
-    def in_transaction(self) -> bool:
-        """Check if connection is currently in a transaction."""
-        ...
-
-    async def begin(self) -> None:
-        """Begin a new transaction."""
-        ...
-
-
-@runtime_checkable
-class SyncCloseableConnectionProtocol(Protocol):
-    """Protocol for sync connections that can be closed."""
-
-    def close(self) -> None:
-        """Close the connection."""
-        ...
-
-
-@runtime_checkable
-class AsyncCloseableConnectionProtocol(Protocol):
-    """Protocol for async connections that can be closed."""
-
-    async def close(self) -> None:
-        """Close the connection."""
-        ...
-
-
-@runtime_checkable
-class SyncCopyCapableConnectionProtocol(Protocol):
-    """Protocol for sync connections that support COPY operations."""
-
-    def copy_from(self, table: str, file: Any, **kwargs: Any) -> None:
-        """Copy data from file to table."""
-        ...
-
-    def copy_to(self, table: str, file: Any, **kwargs: Any) -> None:
-        """Copy data from table to file."""
-        ...
-
-
-@runtime_checkable
-class AsyncCopyCapableConnectionProtocol(Protocol):
-    """Protocol for async connections that support COPY operations."""
-
-    async def copy_from(self, table: str, file: Any, **kwargs: Any) -> None:
-        """Copy data from file to table."""
-        ...
-
-    async def copy_to(self, table: str, file: Any, **kwargs: Any) -> None:
-        """Copy data from table to file."""
-        ...
-
-
-@runtime_checkable
-class SyncPipelineCapableDriverProtocol(Protocol):
-    """Protocol for sync drivers that support native pipeline execution."""
-
-    def _execute_pipeline_native(self, operations: list[Any], **options: Any) -> list[Any]:
-        """Execute pipeline operations natively."""
-        ...
-
-
-@runtime_checkable
-class AsyncPipelineCapableDriverProtocol(Protocol):
-    """Protocol for async drivers that support native pipeline execution."""
-
-    async def _execute_pipeline_native(self, operations: list[Any], **options: Any) -> list[Any]:
-        """Execute pipeline operations natively."""
-        ...
-
-
 @runtime_checkable
 class ObjectStoreProtocol(Protocol):
-    """Unified protocol for object storage operations.
-
-    This protocol defines the interface for all storage backends with built-in
-    instrumentation support. Backends must implement both sync and async operations
-    where possible, with async operations suffixed with _async.
+    """Protocol for object storage operations.
 
-    All methods use 'path' terminology for consistency with object store patterns.
+    Defines the interface for storage backends with both sync and async operations.
     """
 
-    # Class-level capability descriptor
     capabilities: ClassVar["StorageCapabilities"]
 
+    protocol: str
+
     def __init__(self, uri: str, **kwargs: Any) -> None:
         return
 
-    # Core Operations (sync)
     def read_bytes(self, path: "Union[str, Path]", **kwargs: Any) -> bytes:
         """Read bytes from an object."""
         return b""
@@ -372,7 +210,6 @@ class ObjectStoreProtocol(Protocol):
         """Write text to an object."""
         return
 
-    # Object Operations
     def exists(self, path: "Union[str, Path]", **kwargs: Any) -> bool:
         """Check if an object exists."""
         return False
@@ -389,7 +226,6 @@ class ObjectStoreProtocol(Protocol):
         """Move an object."""
         return
 
-    # Listing Operations
     def list_objects(self, prefix: str = "", recursive: bool = True, **kwargs: Any) -> list[str]:
         """List objects with optional prefix."""
         return []
@@ -398,7 +234,6 @@ class ObjectStoreProtocol(Protocol):
         """Find objects matching a glob pattern."""
         return []
 
-    # Path Operations
     def is_object(self, path: "Union[str, Path]") -> bool:
         """Check if path points to an object."""
         return False
@@ -411,32 +246,21 @@ class ObjectStoreProtocol(Protocol):
         """Get object metadata."""
         return {}
 
-    # Arrow Operations
     def read_arrow(self, path: "Union[str, Path]", **kwargs: Any) -> "ArrowTable":
-        """Read an Arrow table from storage.
-
-        For obstore backend, this should use native arrow operations when available.
-        """
+        """Read an Arrow table from storage."""
         msg = "Arrow reading not implemented"
         raise NotImplementedError(msg)
 
     def write_arrow(self, path: "Union[str, Path]", table: "ArrowTable", **kwargs: Any) -> None:
-        """Write an Arrow table to storage.
-
-        For obstore backend, this should use native arrow operations when available.
-        """
+        """Write an Arrow table to storage."""
         msg = "Arrow writing not implemented"
         raise NotImplementedError(msg)
 
     def stream_arrow(self, pattern: str, **kwargs: Any) -> "Iterator[ArrowRecordBatch]":
-        """Stream Arrow record batches from matching objects.
-
-        For obstore backend, this should use native streaming when available.
-        """
+        """Stream Arrow record batches from matching objects."""
         msg = "Arrow streaming not implemented"
         raise NotImplementedError(msg)
 
-    # Async versions
     async def read_bytes_async(self, path: "Union[str, Path]", **kwargs: Any) -> bytes:
         """Async read bytes from an object."""
         msg = "Async operations not implemented"

sqlspec/storage/__init__.py

@@ -1,6 +1,6 @@
 """Storage abstraction layer for SQLSpec.
 
-This module provides a flexible storage system with:
+Provides a flexible storage system with:
 - Multiple backend support (local, fsspec, obstore)
 - Lazy loading and configuration-based registration
 - URI scheme-based automatic backend resolution