sqlspec 0.13.1__py3-none-any.whl → 0.16.2__py3-none-any.whl

sqlspec/migrations/tracker.py

@@ -0,0 +1,140 @@
+"""Migration version tracking for SQLSpec.
+
+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 AsyncDriverAdapterBase, SyncDriverAdapterBase
+
+__all__ = ("AsyncMigrationTracker", "SyncMigrationTracker")
+
+
+class SyncMigrationTracker(BaseMigrationTracker["SyncDriverAdapterBase"]):
+    """Tracks applied migrations in the database."""
+
+    def ensure_tracking_table(self, driver: "SyncDriverAdapterBase") -> None:
+        """Create the migration tracking table if it doesn't exist.
+
+        Args:
+            driver: The database driver to use.
+        """
+        driver.execute(self._get_create_table_sql())
+
+    def get_current_version(self, driver: "SyncDriverAdapterBase") -> Optional[str]:
+        """Get the latest applied migration version.
+
+        Args:
+            driver: The database driver to use.
+
+        Returns:
+            The current version number or None if no migrations applied.
+        """
+        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: "SyncDriverAdapterBase") -> "list[dict[str, Any]]":
+        """Get all applied migrations in order.
+
+        Args:
+            driver: The database driver to use.
+
+        Returns:
+            List of migration records.
+        """
+        result = driver.execute(self._get_applied_migrations_sql())
+        return result.data or []
+
+    def record_migration(
+        self, driver: "SyncDriverAdapterBase", version: str, description: str, execution_time_ms: int, checksum: str
+    ) -> None:
+        """Record a successfully applied migration.
+
+        Args:
+            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.
+        """
+        driver.execute(
+            self._get_record_migration_sql(
+                version, description, execution_time_ms, checksum, os.environ.get("USER", "unknown")
+            )
+        )
+
+    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.
+        """
+        driver.execute(self._get_remove_migration_sql(version))
+
+
+class AsyncMigrationTracker(BaseMigrationTracker["AsyncDriverAdapterBase"]):
+    """Tracks applied migrations in the database."""
+
+    async def ensure_tracking_table(self, driver: "AsyncDriverAdapterBase") -> None:
+        """Create the migration tracking table if it doesn't exist.
+
+        Args:
+            driver: The database driver to use.
+        """
+        await driver.execute(self._get_create_table_sql())
+
+    async def get_current_version(self, driver: "AsyncDriverAdapterBase") -> Optional[str]:
+        """Get the latest applied migration version.
+
+        Args:
+            driver: The database driver to use.
+
+        Returns:
+            The current version number or None if no migrations applied.
+        """
+        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: "AsyncDriverAdapterBase") -> "list[dict[str, Any]]":
+        """Get all applied migrations in order.
+
+        Args:
+            driver: The database driver to use.
+
+        Returns:
+            List of migration records.
+        """
+        result = await driver.execute(self._get_applied_migrations_sql())
+        return result.data or []
+
+    async def record_migration(
+        self, driver: "AsyncDriverAdapterBase", version: str, description: str, execution_time_ms: int, checksum: str
+    ) -> None:
+        """Record a successfully applied migration.
+
+        Args:
+            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.
+        """
+        await driver.execute(
+            self._get_record_migration_sql(
+                version, description, execution_time_ms, checksum, os.environ.get("USER", "unknown")
+            )
+        )
+
+    async def remove_migration(self, driver: "AsyncDriverAdapterBase", version: str) -> None:
+        """Remove a migration record (used during downgrade).
+
+        Args:
+            driver: The database driver to use.
+            version: Version number to remove.
+        """
+        await driver.execute(self._get_remove_migration_sql(version))

sqlspec/migrations/utils.py

@@ -0,0 +1,129 @@
+"""Utility functions for SQLSpec migrations.
+
+This module provides helper functions for migration operations.
+"""
+
+import os
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+if TYPE_CHECKING:
+    from sqlspec.driver import AsyncDriverAdapterBase
+
+__all__ = ("create_migration_file", "drop_all", "get_author")
+
+
+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.
+    """
+    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]
+
+    if file_type == "py":
+        filename = f"{version}_{safe_message}.py"
+        file_path = migrations_dir / filename
+        template = f'''"""SQLSpec Migration - {message}
+
+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
+CREATE TABLE placeholder (
+    id INTEGER PRIMARY KEY
+);
+
+-- name: migrate-{version}-down
+DROP TABLE placeholder;
+"""
+
+    file_path.write_text(template)
+    return file_path
+
+
+def get_author() -> str:
+    """Get current user for migration metadata.
+
+    Returns:
+        Username from environment or 'unknown'.
+    """
+    return os.environ.get("USER", "unknown")
+
+
+async def drop_all(engine: "AsyncDriverAdapterBase", version_table_name: str, metadata: Optional[Any] = None) -> None:
+    """Drop all tables from the database.
+
+    Args:
+        engine: The database engine/driver.
+        version_table_name: Name of the version tracking table.
+        metadata: Optional metadata object.
+
+    Raises:
+        NotImplementedError: Always raised.
+    """
+    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,41 +14,31 @@ 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",
     "FilterParameterProtocol",
+    "HasExpressionProtocol",
     "HasExpressionsProtocol",
     "HasLimitProtocol",
     "HasOffsetProtocol",
     "HasOrderByProtocol",
-    "HasRiskLevelProtocol",
+    "HasParameterBuilderProtocol",
+    "HasSQLGlotExpressionProtocol",
     "HasSQLMethodProtocol",
+    "HasToStatementProtocol",
     "HasWhereProtocol",
     "IndexableRow",
     "IterableParameters",
     "ObjectStoreItemProtocol",
     "ObjectStoreProtocol",
     "ParameterValueProtocol",
-    "ProcessorProtocol",
     "SQLBuilderProtocol",
     "SelectBuilderProtocol",
-    "SyncCloseableConnectionProtocol",
-    "SyncCopyCapableConnectionProtocol",
-    "SyncPipelineCapableDriverProtocol",
-    "SyncTransactionCapableConnectionProtocol",
-    "SyncTransactionStateConnectionProtocol",
     "WithMethodProtocol",
 )
 
@@ -167,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."""
@@ -184,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."""
@@ -220,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.
+    """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.
-
-    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""
@@ -368,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
@@ -385,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 []
@@ -394,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
@@ -407,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"
@@ -501,9 +329,39 @@ class ObjectStoreProtocol(Protocol):
         raise NotImplementedError(msg)
 
 
-# =============================================================================
-# SQL Builder Protocols
-# =============================================================================
+@runtime_checkable
+class HasSQLGlotExpressionProtocol(Protocol):
+    """Protocol for objects with a sqlglot_expression property."""
+
+    @property
+    def sqlglot_expression(self) -> "Optional[exp.Expression]":
+        """Return the SQLGlot expression for this object."""
+        ...
+
+
+@runtime_checkable
+class HasParameterBuilderProtocol(Protocol):
+    """Protocol for objects that can add parameters."""
+
+    def add_parameter(self, value: Any, name: "Optional[str]" = None) -> tuple[Any, str]:
+        """Add a parameter to the builder."""
+        ...
+
+
+@runtime_checkable
+class HasExpressionProtocol(Protocol):
+    """Protocol for objects with an _expression attribute."""
+
+    _expression: "Optional[exp.Expression]"
+
+
+@runtime_checkable
+class HasToStatementProtocol(Protocol):
+    """Protocol for objects with a to_statement method."""
+
+    def to_statement(self) -> Any:
+        """Convert to SQL statement."""
+        ...
 
 
 @runtime_checkable
@@ -513,6 +371,9 @@ class SQLBuilderProtocol(Protocol):
     _expression: "Optional[exp.Expression]"
     _parameters: dict[str, Any]
     _parameter_counter: int
+    _columns: Any  # Optional attribute for some builders
+    _table: Any  # Optional attribute for some builders
+    _with_ctes: Any  # Optional attribute for some builders
     dialect: Any
     dialect_name: "Optional[str]"
 
@@ -525,6 +386,10 @@ class SQLBuilderProtocol(Protocol):
         """Add a parameter to the builder."""
         ...
 
+    def _generate_unique_parameter_name(self, base_name: str) -> str:
+        """Generate a unique parameter name."""
+        ...
+
     def _parameterize_expression(self, expression: "exp.Expression") -> "exp.Expression":
         """Replace literal values in an expression with bound parameters."""
         ...

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