sqlspec 0.16.2__cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

sqlspec/cli.py

@@ -0,0 +1,247 @@
+import sys
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
+
+if TYPE_CHECKING:
+    from click import Group
+
+    from sqlspec.config import AsyncDatabaseConfig, SyncDatabaseConfig
+
+__all__ = ("add_migration_commands", "get_sqlspec_group")
+
+
+def get_sqlspec_group() -> "Group":
+    """Get the SQLSpec CLI group.
+
+    Raises:
+        MissingDependencyError: If the `click` package is not installed.
+
+    Returns:
+        The SQLSpec CLI group.
+    """
+    from sqlspec.exceptions import MissingDependencyError
+
+    try:
+        import rich_click as click
+    except ImportError:
+        try:
+            import click  # type: ignore[no-redef]
+        except ImportError as e:
+            raise MissingDependencyError(package="click", install_package="cli") from e
+
+    @click.group(name="sqlspec")
+    @click.option(
+        "--config",
+        help="Dotted path to SQLSpec config(s) (e.g. 'myapp.config.sqlspec_configs')",
+        required=True,
+        type=str,
+    )
+    @click.pass_context
+    def sqlspec_group(ctx: "click.Context", config: str) -> None:
+        """SQLSpec CLI commands."""
+        from rich import get_console
+
+        from sqlspec.utils import module_loader
+
+        console = get_console()
+        ctx.ensure_object(dict)
+        try:
+            config_instance = module_loader.import_string(config)
+            if isinstance(config_instance, Sequence):
+                ctx.obj["configs"] = config_instance
+            else:
+                ctx.obj["configs"] = [config_instance]
+        except ImportError as e:
+            console.print(f"[red]Error loading config: {e}[/]")
+            ctx.exit(1)
+
+    return sqlspec_group
+
+
+def add_migration_commands(database_group: Optional["Group"] = None) -> "Group":
+    """Add migration commands to the database group.
+
+    Args:
+        database_group: The database group to add the commands to.
+
+    Raises:
+        MissingDependencyError: If the `click` package is not installed.
+
+    Returns:
+        The database group with the migration commands added.
+    """
+    from sqlspec.exceptions import MissingDependencyError
+
+    try:
+        import rich_click as click
+    except ImportError:
+        try:
+            import click  # type: ignore[no-redef]
+        except ImportError as e:
+            raise MissingDependencyError(package="click", install_package="cli") from e
+    from rich import get_console
+
+    console = get_console()
+
+    if database_group is None:
+        database_group = get_sqlspec_group()
+
+    bind_key_option = click.option(
+        "--bind-key", help="Specify which SQLSpec config to use by bind key", type=str, default=None
+    )
+    verbose_option = click.option("--verbose", help="Enable verbose output.", type=bool, default=False, is_flag=True)
+    no_prompt_option = click.option(
+        "--no-prompt",
+        help="Do not prompt for confirmation before executing the command.",
+        type=bool,
+        default=False,
+        required=False,
+        show_default=True,
+        is_flag=True,
+    )
+
+    def get_config_by_bind_key(
+        ctx: "click.Context", bind_key: Optional[str]
+    ) -> "Union[AsyncDatabaseConfig[Any, Any, Any], SyncDatabaseConfig[Any, Any, Any]]":
+        """Get the SQLSpec config for the specified bind key.
+
+        Args:
+            ctx: The click context.
+            bind_key: The bind key to get the config for.
+
+        Returns:
+            The SQLSpec config for the specified bind key.
+        """
+        configs = ctx.obj["configs"]
+        if bind_key is None:
+            return cast("Union[AsyncDatabaseConfig[Any, Any, Any], SyncDatabaseConfig[Any, Any, Any]]", configs[0])
+
+        for config in configs:
+            config_name = getattr(config, "name", None) or getattr(config, "bind_key", None)
+            if config_name == bind_key:
+                return cast("Union[AsyncDatabaseConfig[Any, Any, Any], SyncDatabaseConfig[Any, Any, Any]]", config)
+
+        console.print(f"[red]No config found for bind key: {bind_key}[/]")
+        sys.exit(1)
+
+    @database_group.command(name="show-current-revision", help="Shows the current revision for the database.")
+    @bind_key_option
+    @verbose_option
+    def show_database_revision(bind_key: Optional[str], verbose: bool) -> None:  # pyright: ignore[reportUnusedFunction]
+        """Show current database revision."""
+        from sqlspec.migrations.commands import MigrationCommands
+
+        ctx = click.get_current_context()
+        console.rule("[yellow]Listing current revision[/]", align="left")
+        sqlspec_config = get_config_by_bind_key(ctx, bind_key)
+        migration_commands = MigrationCommands(config=sqlspec_config)
+        migration_commands.current(verbose=verbose)
+
+    @database_group.command(name="downgrade", help="Downgrade database to a specific revision.")
+    @bind_key_option
+    @no_prompt_option
+    @click.argument("revision", type=str, default="-1")
+    def downgrade_database(  # pyright: ignore[reportUnusedFunction]
+        bind_key: Optional[str], revision: str, no_prompt: bool
+    ) -> None:
+        """Downgrade the database to the latest revision."""
+        from rich.prompt import Confirm
+
+        from sqlspec.migrations.commands import MigrationCommands
+
+        ctx = click.get_current_context()
+        console.rule("[yellow]Starting database downgrade process[/]", align="left")
+        input_confirmed = (
+            True
+            if no_prompt
+            else Confirm.ask(f"Are you sure you want to downgrade the database to the `{revision}` revision?")
+        )
+        if input_confirmed:
+            sqlspec_config = get_config_by_bind_key(ctx, bind_key)
+            migration_commands = MigrationCommands(config=sqlspec_config)
+            migration_commands.downgrade(revision=revision)
+
+    @database_group.command(name="upgrade", help="Upgrade database to a specific revision.")
+    @bind_key_option
+    @no_prompt_option
+    @click.argument("revision", type=str, default="head")
+    def upgrade_database(  # pyright: ignore[reportUnusedFunction]
+        bind_key: Optional[str], revision: str, no_prompt: bool
+    ) -> None:
+        """Upgrade the database to the latest revision."""
+        from rich.prompt import Confirm
+
+        from sqlspec.migrations.commands import MigrationCommands
+
+        ctx = click.get_current_context()
+        console.rule("[yellow]Starting database upgrade process[/]", align="left")
+        input_confirmed = (
+            True
+            if no_prompt
+            else Confirm.ask(f"[bold]Are you sure you want migrate the database to the `{revision}` revision?[/]")
+        )
+        if input_confirmed:
+            sqlspec_config = get_config_by_bind_key(ctx, bind_key)
+            migration_commands = MigrationCommands(config=sqlspec_config)
+            migration_commands.upgrade(revision=revision)
+
+    @database_group.command(help="Stamp the revision table with the given revision")
+    @click.argument("revision", type=str)
+    @bind_key_option
+    def stamp(bind_key: Optional[str], revision: str) -> None:  # pyright: ignore[reportUnusedFunction]
+        """Stamp the revision table with the given revision."""
+        from sqlspec.migrations.commands import MigrationCommands
+
+        ctx = click.get_current_context()
+        sqlspec_config = get_config_by_bind_key(ctx, bind_key)
+        migration_commands = MigrationCommands(config=sqlspec_config)
+        migration_commands.stamp(revision=revision)
+
+    @database_group.command(name="init", help="Initialize migrations for the project.")
+    @bind_key_option
+    @click.argument("directory", default=None, required=False)
+    @click.option("--package", is_flag=True, default=True, help="Create `__init__.py` for created folder")
+    @no_prompt_option
+    def init_sqlspec(  # pyright: ignore[reportUnusedFunction]
+        bind_key: Optional[str], directory: Optional[str], package: bool, no_prompt: bool
+    ) -> None:
+        """Initialize the database migrations."""
+        from rich.prompt import Confirm
+
+        from sqlspec.migrations.commands import MigrationCommands
+
+        ctx = click.get_current_context()
+        console.rule("[yellow]Initializing database migrations.", align="left")
+        input_confirmed = (
+            True if no_prompt else Confirm.ask("[bold]Are you sure you want initialize migrations for the project?[/]")
+        )
+        if input_confirmed:
+            configs = [get_config_by_bind_key(ctx, bind_key)] if bind_key is not None else ctx.obj["configs"]
+            for config in configs:
+                migration_config = getattr(config, "migration_config", {})
+                directory = migration_config.get("script_location", "migrations") if directory is None else directory
+                migration_commands = MigrationCommands(config=config)
+                migration_commands.init(directory=cast("str", directory), package=package)
+
+    @database_group.command(name="make-migrations", help="Create a new migration revision.")
+    @bind_key_option
+    @click.option("-m", "--message", default=None, help="Revision message")
+    @no_prompt_option
+    def create_revision(  # pyright: ignore[reportUnusedFunction]
+        bind_key: Optional[str], message: Optional[str], no_prompt: bool
+    ) -> None:
+        """Create a new database revision."""
+        from rich.prompt import Prompt
+
+        from sqlspec.migrations.commands import MigrationCommands
+
+        ctx = click.get_current_context()
+        console.rule("[yellow]Creating new migration revision[/]", align="left")
+        if message is None:
+            message = "new migration" if no_prompt else Prompt.ask("Please enter a message describing this revision")
+
+        sqlspec_config = get_config_by_bind_key(ctx, bind_key)
+        migration_commands = MigrationCommands(config=sqlspec_config)
+        migration_commands.revision(message=message)
+
+    return database_group

sqlspec/config.py

@@ -0,0 +1,395 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Generic, Optional, TypeVar, Union
+
+from typing_extensions import NotRequired, TypedDict
+
+from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
+from sqlspec.core.statement import StatementConfig
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable
+    from contextlib import AbstractAsyncContextManager, AbstractContextManager
+
+    from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
+
+
+__all__ = (
+    "AsyncConfigT",
+    "AsyncDatabaseConfig",
+    "ConfigT",
+    "DatabaseConfigProtocol",
+    "DriverT",
+    "LifecycleConfig",
+    "NoPoolAsyncConfig",
+    "NoPoolSyncConfig",
+    "SyncConfigT",
+    "SyncDatabaseConfig",
+)
+
+AsyncConfigT = TypeVar("AsyncConfigT", bound="Union[AsyncDatabaseConfig[Any, Any, Any], NoPoolAsyncConfig[Any, Any]]")
+SyncConfigT = TypeVar("SyncConfigT", bound="Union[SyncDatabaseConfig[Any, Any, Any], NoPoolSyncConfig[Any, Any]]")
+ConfigT = TypeVar(
+    "ConfigT",
+    bound="Union[Union[AsyncDatabaseConfig[Any, Any, Any], NoPoolAsyncConfig[Any, Any]], SyncDatabaseConfig[Any, Any, Any], NoPoolSyncConfig[Any, Any]]",
+)
+
+# Define TypeVars for Generic classes
+ConnectionT = TypeVar("ConnectionT")
+PoolT = TypeVar("PoolT")
+DriverT = TypeVar("DriverT", bound="Union[SyncDriverAdapterBase, AsyncDriverAdapterBase]")
+
+logger = get_logger("config")
+
+
+class LifecycleConfig(TypedDict, total=False):
+    """Universal lifecycle hooks for all adapters.
+
+    Each hook accepts a list of callables to support multiple handlers.
+    """
+
+    on_connection_create: NotRequired[list[Callable[[Any], None]]]
+    on_connection_destroy: NotRequired[list[Callable[[Any], None]]]
+    on_pool_create: NotRequired[list[Callable[[Any], None]]]
+    on_pool_destroy: NotRequired[list[Callable[[Any], None]]]
+    on_session_start: NotRequired[list[Callable[[Any], None]]]
+    on_session_end: NotRequired[list[Callable[[Any], None]]]
+    on_query_start: NotRequired[list[Callable[[str, dict], None]]]
+    on_query_complete: NotRequired[list[Callable[[str, dict, Any], None]]]
+    on_error: NotRequired[list[Callable[[Exception, str, dict], None]]]
+
+
+class DatabaseConfigProtocol(ABC, Generic[ConnectionT, PoolT, DriverT]):
+    """Protocol defining the interface for database configurations."""
+
+    __slots__ = ("driver_features", "migration_config", "pool_instance", "statement_config")
+    driver_type: "ClassVar[type[Any]]"
+    connection_type: "ClassVar[type[Any]]"
+    is_async: "ClassVar[bool]" = False
+    supports_connection_pooling: "ClassVar[bool]" = False
+    supports_native_arrow_import: "ClassVar[bool]" = False
+    supports_native_arrow_export: "ClassVar[bool]" = False
+    supports_native_parquet_import: "ClassVar[bool]" = False
+    supports_native_parquet_export: "ClassVar[bool]" = False
+    statement_config: "StatementConfig"
+    pool_instance: "Optional[PoolT]"
+    migration_config: "dict[str, Any]"
+
+    def __hash__(self) -> int:
+        return id(self)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, type(self)):
+            return False
+        return bool(self.pool_instance == other.pool_instance and self.migration_config == other.migration_config)
+
+    def __repr__(self) -> str:
+        parts = ", ".join([f"pool_instance={self.pool_instance!r}", f"migration_config={self.migration_config!r}"])
+        return f"{type(self).__name__}({parts})"
+
+    @abstractmethod
+    def create_connection(self) -> "Union[ConnectionT, Awaitable[ConnectionT]]":
+        """Create and return a new database connection."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def provide_connection(
+        self, *args: Any, **kwargs: Any
+    ) -> "Union[AbstractContextManager[ConnectionT], AbstractAsyncContextManager[ConnectionT]]":
+        """Provide a database connection context manager."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def provide_session(
+        self, *args: Any, **kwargs: Any
+    ) -> "Union[AbstractContextManager[DriverT], AbstractAsyncContextManager[DriverT]]":
+        """Provide a database session context manager."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def create_pool(self) -> "Union[PoolT, Awaitable[PoolT]]":
+        """Create and return connection pool."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def close_pool(self) -> "Optional[Awaitable[None]]":
+        """Terminate the connection pool."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def provide_pool(
+        self, *args: Any, **kwargs: Any
+    ) -> "Union[PoolT, Awaitable[PoolT], AbstractContextManager[PoolT], AbstractAsyncContextManager[PoolT]]":
+        """Provide pool instance."""
+        raise NotImplementedError
+
+    def get_signature_namespace(self) -> "dict[str, type[Any]]":
+        """Get the signature namespace for this database configuration.
+
+        This method returns a dictionary of type names to types that should be
+        registered with Litestar's signature namespace to prevent serialization
+        attempts on database-specific types.
+
+        Returns:
+            Dictionary mapping type names to types.
+        """
+        return {}
+
+
+class NoPoolSyncConfig(DatabaseConfigProtocol[ConnectionT, None, DriverT]):
+    """Base class for a sync database configurations that do not implement a pool."""
+
+    __slots__ = ("connection_config",)
+    is_async: "ClassVar[bool]" = False
+    supports_connection_pooling: "ClassVar[bool]" = False
+
+    def __init__(
+        self,
+        *,
+        connection_config: Optional[dict[str, Any]] = None,
+        migration_config: "Optional[dict[str, Any]]" = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        driver_features: "Optional[dict[str, Any]]" = None,
+    ) -> None:
+        self.pool_instance = None
+        self.connection_config = connection_config or {}
+        self.migration_config: dict[str, Any] = migration_config if migration_config is not None else {}
+
+        if statement_config is None:
+            default_parameter_config = ParameterStyleConfig(
+                default_parameter_style=ParameterStyle.QMARK, supported_parameter_styles={ParameterStyle.QMARK}
+            )
+            self.statement_config = StatementConfig(dialect="sqlite", parameter_config=default_parameter_config)
+        else:
+            self.statement_config = statement_config
+        self.driver_features = driver_features or {}
+
+    def create_connection(self) -> ConnectionT:
+        """Create a database connection."""
+        raise NotImplementedError
+
+    def provide_connection(self, *args: Any, **kwargs: Any) -> "AbstractContextManager[ConnectionT]":
+        """Provide a database connection context manager."""
+        raise NotImplementedError
+
+    def provide_session(
+        self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
+    ) -> "AbstractContextManager[DriverT]":
+        """Provide a database session context manager."""
+        raise NotImplementedError
+
+    def create_pool(self) -> None:
+        return None
+
+    def close_pool(self) -> None:
+        return None
+
+    def provide_pool(self, *args: Any, **kwargs: Any) -> None:
+        return None
+
+
+class NoPoolAsyncConfig(DatabaseConfigProtocol[ConnectionT, None, DriverT]):
+    """Base class for an async database configurations that do not implement a pool."""
+
+    __slots__ = ("connection_config",)
+
+    is_async: "ClassVar[bool]" = True
+    supports_connection_pooling: "ClassVar[bool]" = False
+
+    def __init__(
+        self,
+        *,
+        connection_config: "Optional[dict[str, Any]]" = None,
+        migration_config: "Optional[dict[str, Any]]" = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        driver_features: "Optional[dict[str, Any]]" = None,
+    ) -> None:
+        self.pool_instance = None
+        self.connection_config = connection_config or {}
+        self.migration_config: dict[str, Any] = migration_config if migration_config is not None else {}
+
+        if statement_config is None:
+            default_parameter_config = ParameterStyleConfig(
+                default_parameter_style=ParameterStyle.QMARK, supported_parameter_styles={ParameterStyle.QMARK}
+            )
+            self.statement_config = StatementConfig(dialect="sqlite", parameter_config=default_parameter_config)
+        else:
+            self.statement_config = statement_config
+        self.driver_features = driver_features or {}
+
+    async def create_connection(self) -> ConnectionT:
+        """Create a database connection."""
+        raise NotImplementedError
+
+    def provide_connection(self, *args: Any, **kwargs: Any) -> "AbstractAsyncContextManager[ConnectionT]":
+        """Provide a database connection context manager."""
+        raise NotImplementedError
+
+    def provide_session(
+        self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
+    ) -> "AbstractAsyncContextManager[DriverT]":
+        """Provide a database session context manager."""
+        raise NotImplementedError
+
+    async def create_pool(self) -> None:
+        return None
+
+    async def close_pool(self) -> None:
+        return None
+
+    def provide_pool(self, *args: Any, **kwargs: Any) -> None:
+        return None
+
+
+class SyncDatabaseConfig(DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]):
+    """Generic Sync Database Configuration."""
+
+    __slots__ = ("pool_config",)
+
+    is_async: "ClassVar[bool]" = False
+    supports_connection_pooling: "ClassVar[bool]" = True
+
+    def __init__(
+        self,
+        *,
+        pool_config: "Optional[dict[str, Any]]" = None,
+        pool_instance: "Optional[PoolT]" = None,
+        migration_config: "Optional[dict[str, Any]]" = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        driver_features: "Optional[dict[str, Any]]" = None,
+    ) -> None:
+        self.pool_instance = pool_instance
+        self.pool_config = pool_config or {}
+        self.migration_config: dict[str, Any] = migration_config if migration_config is not None else {}
+
+        if statement_config is None:
+            default_parameter_config = ParameterStyleConfig(
+                default_parameter_style=ParameterStyle.QMARK, supported_parameter_styles={ParameterStyle.QMARK}
+            )
+            self.statement_config = StatementConfig(dialect="postgres", parameter_config=default_parameter_config)
+        else:
+            self.statement_config = statement_config
+        self.driver_features = driver_features or {}
+
+    def create_pool(self) -> PoolT:
+        """Create and return the connection pool.
+
+        Returns:
+            The created pool.
+        """
+        if self.pool_instance is not None:
+            return self.pool_instance
+        self.pool_instance = self._create_pool()
+        return self.pool_instance
+
+    def close_pool(self) -> None:
+        """Close the connection pool."""
+        self._close_pool()
+
+    def provide_pool(self, *args: Any, **kwargs: Any) -> PoolT:
+        """Provide pool instance."""
+        if self.pool_instance is None:
+            self.pool_instance = self.create_pool()
+        return self.pool_instance
+
+    def create_connection(self) -> ConnectionT:
+        """Create a database connection."""
+        raise NotImplementedError
+
+    def provide_connection(self, *args: Any, **kwargs: Any) -> "AbstractContextManager[ConnectionT]":
+        """Provide a database connection context manager."""
+        raise NotImplementedError
+
+    def provide_session(
+        self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
+    ) -> "AbstractContextManager[DriverT]":
+        """Provide a database session context manager."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _create_pool(self) -> PoolT:
+        """Actual pool creation implementation."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _close_pool(self) -> None:
+        """Actual pool destruction implementation."""
+        raise NotImplementedError
+
+
+class AsyncDatabaseConfig(DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]):
+    """Generic Async Database Configuration."""
+
+    __slots__ = ("pool_config",)
+
+    is_async: "ClassVar[bool]" = True
+    supports_connection_pooling: "ClassVar[bool]" = True
+
+    def __init__(
+        self,
+        *,
+        pool_config: "Optional[dict[str, Any]]" = None,
+        pool_instance: "Optional[PoolT]" = None,
+        migration_config: "Optional[dict[str, Any]]" = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        driver_features: "Optional[dict[str, Any]]" = None,
+    ) -> None:
+        self.pool_instance = pool_instance
+        self.pool_config = pool_config or {}
+        self.migration_config: dict[str, Any] = migration_config if migration_config is not None else {}
+
+        if statement_config is None:
+            self.statement_config = StatementConfig(
+                parameter_config=ParameterStyleConfig(
+                    default_parameter_style=ParameterStyle.QMARK, supported_parameter_styles={ParameterStyle.QMARK}
+                ),
+                dialect="postgres",
+            )
+        else:
+            self.statement_config = statement_config
+        self.driver_features = driver_features or {}
+
+    async def create_pool(self) -> PoolT:
+        """Create and return the connection pool.
+
+        Returns:
+            The created pool.
+        """
+        if self.pool_instance is not None:
+            return self.pool_instance
+        self.pool_instance = await self._create_pool()
+        return self.pool_instance
+
+    async def close_pool(self) -> None:
+        """Close the connection pool."""
+        await self._close_pool()
+
+    async def provide_pool(self, *args: Any, **kwargs: Any) -> PoolT:
+        """Provide pool instance."""
+        if self.pool_instance is None:
+            self.pool_instance = await self.create_pool()
+        return self.pool_instance
+
+    async def create_connection(self) -> ConnectionT:
+        """Create a database connection."""
+        raise NotImplementedError
+
+    def provide_connection(self, *args: Any, **kwargs: Any) -> "AbstractAsyncContextManager[ConnectionT]":
+        """Provide a database connection context manager."""
+        raise NotImplementedError
+
+    def provide_session(
+        self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
+    ) -> "AbstractAsyncContextManager[DriverT]":
+        """Provide a database session context manager."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def _create_pool(self) -> PoolT:
+        """Actual async pool creation implementation."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def _close_pool(self) -> None:
+        """Actual async pool destruction implementation."""
+        raise NotImplementedError