sqlspec 0.18.0__py3-none-any.whl → 0.20.0__py3-none-any.whl

sqlspec/driver/mixins/_sql_translator.py

@@ -1,3 +1,5 @@
+"""SQL translation mixin for cross-database compatibility."""
+
 from typing import Final, NoReturn, Optional
 
 from mypy_extensions import trait
@@ -33,8 +35,7 @@ class SQLTranslatorMixin:
         Returns:
             SQL string in target dialect
 
-        Raises:
-            SQLConversionError: If parsing or conversion fails
+
         """
 
         parsed_expression: Optional[exp.Expression] = None
@@ -53,7 +54,15 @@ class SQLTranslatorMixin:
         return self._generate_sql_safely(parsed_expression, target_dialect, pretty)
 
     def _parse_statement_safely(self, statement: "Statement") -> "exp.Expression":
-        """Parse statement with copy=False optimization and proper error handling."""
+        """Parse statement with error handling.
+
+        Args:
+            statement: SQL statement to parse
+
+        Returns:
+            Parsed expression
+
+        """
         try:
             sql_string = str(statement)
 
@@ -62,23 +71,52 @@ class SQLTranslatorMixin:
             self._raise_parse_error(e)
 
     def _generate_sql_safely(self, expression: "exp.Expression", dialect: DialectType, pretty: bool) -> str:
-        """Generate SQL with proper error handling."""
+        """Generate SQL with error handling.
+
+        Args:
+            expression: Parsed expression to convert
+            dialect: Target SQL dialect
+            pretty: Whether to format the output SQL
+
+        Returns:
+            Generated SQL string
+
+        """
         try:
             return expression.sql(dialect=dialect, pretty=pretty)
         except Exception as e:
             self._raise_conversion_error(dialect, e)
 
     def _raise_statement_parse_error(self) -> NoReturn:
-        """Raise error for unparsable statements."""
+        """Raise error for unparsable statements.
+
+        Raises:
+            SQLConversionError: Always raised
+        """
         msg = "Statement could not be parsed"
         raise SQLConversionError(msg)
 
     def _raise_parse_error(self, e: Exception) -> NoReturn:
-        """Raise error for parsing failures."""
+        """Raise error for parsing failures.
+
+        Args:
+            e: Original exception that caused the failure
+
+        Raises:
+            SQLConversionError: Always raised
+        """
         error_msg = f"Failed to parse SQL statement: {e!s}"
         raise SQLConversionError(error_msg) from e
 
     def _raise_conversion_error(self, dialect: DialectType, e: Exception) -> NoReturn:
-        """Raise error for conversion failures."""
+        """Raise error for conversion failures.
+
+        Args:
+            dialect: Target dialect that caused the failure
+            e: Original exception that caused the failure
+
+        Raises:
+            SQLConversionError: Always raised
+        """
         error_msg = f"Failed to convert SQL expression to {dialect}: {e!s}"
         raise SQLConversionError(error_msg) from e

sqlspec/extensions/aiosql/adapter.py

@@ -2,7 +2,7 @@
 
 This module provides adapter classes that implement the aiosql adapter protocols
 while using SQLSpec drivers under the hood. This enables users to load SQL queries
-from files using aiosql while leveraging all of SQLSpec's advanced features.
+from files using aiosql while using SQLSpec's features for execution and type mapping.
 """
 
 import logging

sqlspec/extensions/litestar/_utils.py

@@ -12,7 +12,7 @@ def get_sqlspec_scope_state(scope: "Scope", key: str, default: Any = None, pop:
     """Get an internal value from connection scope state.
 
     Note:
-        If called with a default value, this method behaves like to `dict.set_default()`, both setting the key in the
+        If called with a default value, this method behaves like `dict.setdefault()`, both setting the key in the
         namespace to the default value, and returning it.
 
         If called without a default value, the method behaves like `dict.get()`, returning ``None`` if the key does not

sqlspec/extensions/litestar/handlers.py

@@ -199,6 +199,17 @@ def pool_provider_maker(
 def connection_provider_maker(
     config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]", pool_key: str, connection_key: str
 ) -> "Callable[[State, Scope], AsyncGenerator[ConnectionT, None]]":
+    """Create provider for database connections with proper lifecycle management.
+
+    Args:
+        config: The database configuration object.
+        pool_key: The key used to retrieve the connection pool from `app.state`.
+        connection_key: The key used to store the connection in the ASGI scope.
+
+    Returns:
+        The connection provider function.
+    """
+
     async def provide_connection(state: "State", scope: "Scope") -> "AsyncGenerator[ConnectionT, None]":
         if (db_pool := state.get(pool_key)) is None:
             msg = f"Database pool with key '{pool_key}' not found. Cannot create a connection."
@@ -230,6 +241,16 @@ def connection_provider_maker(
 def session_provider_maker(
     config: "DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]", connection_dependency_key: str
 ) -> "Callable[[Any], AsyncGenerator[DriverT, None]]":
+    """Create provider for database driver sessions.
+
+    Args:
+        config: The database configuration object.
+        connection_dependency_key: The key used for connection dependency injection.
+
+    Returns:
+        The session provider function.
+    """
+
     async def provide_session(*args: Any, **kwargs: Any) -> "AsyncGenerator[DriverT, None]":
         yield cast("DriverT", config.driver_type(connection=args[0] if args else kwargs.get(connection_dependency_key)))  # pyright: ignore
 

sqlspec/extensions/litestar/plugin.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 from litestar.di import Provide
 from litestar.plugins import CLIPlugin, InitPluginProtocol
@@ -14,23 +14,31 @@ if TYPE_CHECKING:
     from click import Group
     from litestar.config.app import AppConfig
 
+    from sqlspec.loader import SQLFileLoader
+
 logger = get_logger("extensions.litestar")
 
 
-class SQLSpec(InitPluginProtocol, CLIPlugin, SQLSpecBase):
+class SQLSpec(SQLSpecBase, InitPluginProtocol, CLIPlugin):
     """Litestar plugin for SQLSpec database integration."""
 
-    __slots__ = ("_config", "_plugin_configs")
+    __slots__ = ("_plugin_configs",)
 
-    def __init__(self, config: Union["SyncConfigT", "AsyncConfigT", "DatabaseConfig", list["DatabaseConfig"]]) -> None:
+    def __init__(
+        self,
+        config: Union["SyncConfigT", "AsyncConfigT", "DatabaseConfig", list["DatabaseConfig"]],
+        *,
+        loader: "Optional[SQLFileLoader]" = None,
+    ) -> None:
         """Initialize SQLSpec plugin.
 
         Args:
             config: Database configuration for SQLSpec plugin.
+            loader: Optional SQL file loader instance.
         """
-        self._configs: dict[Any, DatabaseConfigProtocol[Any, Any, Any]] = {}
+        super().__init__(loader=loader)
         if isinstance(config, DatabaseConfigProtocol):
-            self._plugin_configs: list[DatabaseConfig] = [DatabaseConfig(config=config)]
+            self._plugin_configs: list[DatabaseConfig] = [DatabaseConfig(config=config)]  # pyright: ignore
         elif isinstance(config, DatabaseConfig):
             self._plugin_configs = [config]
         else:
@@ -83,8 +91,7 @@ class SQLSpec(InitPluginProtocol, CLIPlugin, SQLSpecBase):
             app_config.signature_types.append(c.config.connection_type)  # type: ignore[union-attr]
             app_config.signature_types.append(c.config.driver_type)  # type: ignore[union-attr]
 
-            if hasattr(c.config, "get_signature_namespace"):
-                signature_namespace.update(c.config.get_signature_namespace())  # type: ignore[attr-defined]
+            signature_namespace.update(c.config.get_signature_namespace())  # type: ignore[union-attr]
 
             app_config.before_send.append(c.before_send_handler)
             app_config.lifespan.append(c.lifespan_handler)  # pyright: ignore[reportUnknownMemberType]

sqlspec/loader.py

@@ -55,10 +55,10 @@ def _normalize_query_name(name: str) -> str:
     """Normalize query name to be a valid Python identifier.
 
     Args:
-        name: Raw query name from SQL file
+        name: Raw query name from SQL file.
 
     Returns:
-        Normalized query name suitable as Python identifier
+        Normalized query name suitable as Python identifier.
     """
     return TRIM_SPECIAL_CHARS.sub("", name).replace("-", "_")
 
@@ -67,10 +67,10 @@ def _normalize_dialect(dialect: str) -> str:
     """Normalize dialect name with aliases.
 
     Args:
-        dialect: Raw dialect name from SQL file
+        dialect: Raw dialect name from SQL file.
 
     Returns:
-        Normalized dialect name
+        Normalized dialect name.
     """
     normalized = dialect.lower().strip()
     return DIALECT_ALIASES.get(normalized, normalized)
@@ -80,10 +80,10 @@ def _normalize_dialect_for_sqlglot(dialect: str) -> str:
     """Normalize dialect name for SQLGlot compatibility.
 
     Args:
-        dialect: Dialect name from SQL file or parameter
+        dialect: Dialect name from SQL file or parameter.
 
     Returns:
-        SQLGlot-compatible dialect name
+        SQLGlot-compatible dialect name.
     """
     normalized = dialect.lower().strip()
     return DIALECT_ALIASES.get(normalized, normalized)
@@ -125,7 +125,7 @@ class SQLFile:
         """Initialize SQLFile.
 
         Args:
-            content: The raw SQL content from the file.
+            content: Raw SQL content from the file.
             path: Path where the SQL file was loaded from.
             metadata: Optional metadata associated with the SQL file.
             loaded_at: Timestamp when the file was loaded.
@@ -150,7 +150,7 @@ class CachedSQLFile:
         """Initialize cached SQL file.
 
         Args:
-            sql_file: The original SQLFile with content and metadata.
+            sql_file: Original SQLFile with content and metadata.
             parsed_statements: Named statements from the file.
         """
         self.sql_file = sql_file
@@ -291,15 +291,15 @@ class SQLFileLoader:
         """Parse SQL content and extract named statements with dialect specifications.
 
         Args:
-            content: Raw SQL file content to parse
-            file_path: File path for error reporting
+            content: Raw SQL file content to parse.
+            file_path: File path for error reporting.
 
         Returns:
-            Dictionary mapping normalized statement names to NamedStatement objects
+            Dictionary mapping normalized statement names to NamedStatement objects.
 
         Raises:
             SQLFileParseError: If no named statements found, duplicate names exist,
-                              or invalid dialect names are specified
+                              or invalid dialect names are specified.
         """
         statements: dict[str, NamedStatement] = {}
 

sqlspec/migrations/loaders.py

@@ -1,4 +1,7 @@
-"""Migration loader abstractions for SQLSpec."""
+"""Migration loader implementations for SQLSpec.
+
+This module provides loader classes for different migration file formats.
+"""
 
 import abc
 import inspect
@@ -341,7 +344,7 @@ class PythonFileLoader(BaseMigrationLoader):
         return module
 
     def _normalize_and_validate_sql(self, sql: Any, migration_path: Path) -> list[str]:
-        """Validate return type and normalize to list of strings.
+        """Validate and normalize SQL return value to list of strings.
 
         Args:
             sql: Return value from migration function.

sqlspec/migrations/utils.py

@@ -42,7 +42,7 @@ Author: {get_author()}
 
 Migration functions can use either naming convention:
 - Preferred: up()/down()
-- Legacy: migrate_up()/migrate_down()
+- Alternate: migrate_up()/migrate_down()
 
 Both can be synchronous or asynchronous:
 - def up(): ...
@@ -71,7 +71,7 @@ def up() -> Union[str, List[str]]:
 
 
 def down() -> Union[str, List[str]]:
-    """Reverse the migration (downgrade).
+    """Reverse the migration.
 
     Returns:
         SQL statement(s) to execute for downgrade.

sqlspec/storage/backends/obstore.py

@@ -1,7 +1,7 @@
 """Object storage backend using obstore.
 
 Implements the ObjectStoreProtocol using obstore for S3, GCS, Azure,
-and local file storage with Arrow support.
+and local file storage.
 """
 
 from __future__ import annotations
@@ -58,8 +58,6 @@ class ObStoreBackend(ObjectStoreBase, HasStorageCapabilities):
     Uses obstore's Rust-based implementation for storage operations.
     Supports AWS S3, Google Cloud Storage, Azure Blob Storage,
     local filesystem, and HTTP endpoints.
-
-    Includes Arrow support.
     """
 
     capabilities: ClassVar[StorageCapabilities] = StorageCapabilities(

sqlspec/storage/registry.py

@@ -1,4 +1,4 @@
-"""Storage Registry for ObjectStore backends.
+"""Storage registry for ObjectStore backends.
 
 Provides a storage registry that supports URI-first access
 pattern with automatic backend detection, ObStore preferred with FSSpec fallback,

sqlspec/utils/__init__.py

@@ -1,3 +1,10 @@
+"""Utility functions and classes for SQLSpec.
+
+This package provides various utility modules for deprecation handling,
+fixture loading, logging, module loading, singleton patterns, sync/async
+conversion, text processing, and type guards.
+"""
+
 from sqlspec.utils import deprecation, fixtures, logging, module_loader, singleton, sync_tools, text, type_guards
 
 __all__ = ("deprecation", "fixtures", "logging", "module_loader", "singleton", "sync_tools", "text", "type_guards")

sqlspec/utils/deprecation.py

@@ -1,3 +1,9 @@
+"""Deprecation utilities for SQLSpec.
+
+Provides decorators and warning functions for marking deprecated functionality.
+Used to communicate API changes and migration paths to users.
+"""
+
 import inspect
 from functools import wraps
 from typing import Callable, Literal, Optional

sqlspec/utils/fixtures.py

@@ -1,58 +1,267 @@
+"""Fixture loading utilities for SQLSpec.
+
+Provides functions for writing, loading and parsing JSON fixture files
+used in testing and development. Supports both sync and async operations.
+"""
+
+import gzip
+import zipfile
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any, Union
+
+from sqlspec.storage import storage_registry
+from sqlspec.utils.serializers import from_json as decode_json
+from sqlspec.utils.serializers import to_json as encode_json
+from sqlspec.utils.sync_tools import async_
+from sqlspec.utils.type_guards import schema_dump
+
+if TYPE_CHECKING:
+    from sqlspec.typing import ModelDictList, SupportedSchemaModel
+
+__all__ = ("open_fixture", "open_fixture_async", "write_fixture", "write_fixture_async")
 
-from sqlspec._serialization import decode_json
-from sqlspec.exceptions import MissingDependencyError
 
-__all__ = ("open_fixture", "open_fixture_async")
+def _read_compressed_file(file_path: Path) -> str:
+    """Read and decompress a file based on its extension.
+
+    Args:
+        file_path: Path to the file to read
+
+    Returns:
+        The decompressed file content as a string
+
+    Raises:
+        ValueError: If the file format is not supported
+    """
+    if file_path.suffix == ".gz":
+        with gzip.open(file_path, mode="rt", encoding="utf-8") as f:
+            return f.read()
+    elif file_path.suffix == ".zip":
+        with zipfile.ZipFile(file_path, "r") as zf:
+            # Assume the JSON file inside has the same name without .zip
+            json_name = file_path.stem + ".json"
+            if json_name in zf.namelist():
+                with zf.open(json_name) as f:
+                    return f.read().decode("utf-8")
+            # If not found, try the first JSON file in the archive
+            json_files = [name for name in zf.namelist() if name.endswith(".json")]
+            if json_files:
+                with zf.open(json_files[0]) as f:
+                    return f.read().decode("utf-8")
+            msg = f"No JSON file found in ZIP archive: {file_path}"
+            raise ValueError(msg)
+    else:
+        msg = f"Unsupported compression format: {file_path.suffix}"
+        raise ValueError(msg)
+
+
+def _find_fixture_file(fixtures_path: Any, fixture_name: str) -> Path:
+    """Find a fixture file with various extensions.
+
+    Args:
+        fixtures_path: The path to look for fixtures
+        fixture_name: The fixture name to load
+
+    Returns:
+        Path to the found fixture file
+
+    Raises:
+        FileNotFoundError: If no fixture file is found
+    """
+    base_path = Path(fixtures_path)
+
+    # Try different file extensions in order of preference
+    for extension in [".json", ".json.gz", ".json.zip"]:
+        fixture_path = base_path / f"{fixture_name}{extension}"
+        if fixture_path.exists():
+            return fixture_path
+
+    # If no file found, raise error
+    msg = f"Could not find the {fixture_name} fixture"
+    raise FileNotFoundError(msg)
 
 
 def open_fixture(fixtures_path: Any, fixture_name: str) -> Any:
-    """Load and parse a JSON fixture file.
+    """Load and parse a JSON fixture file with compression support.
+
+    Supports reading from:
+    - Regular JSON files (.json)
+    - Gzipped JSON files (.json.gz)
+    - Zipped JSON files (.json.zip)
 
     Args:
-        fixtures_path: The path to look for fixtures (pathlib.Path or anyio.Path)
+        fixtures_path: The path to look for fixtures (pathlib.Path)
         fixture_name: The fixture name to load.
 
-    Raises:
-        FileNotFoundError: Fixtures not found.
 
     Returns:
         The parsed JSON data
     """
+    fixture_path = _find_fixture_file(fixtures_path, fixture_name)
 
-    fixture = Path(fixtures_path / f"{fixture_name}.json")
-    if fixture.exists():
-        with fixture.open(mode="r", encoding="utf-8") as f:
+    if fixture_path.suffix in {".gz", ".zip"}:
+        f_data = _read_compressed_file(fixture_path)
+    else:
+        # Regular JSON file
+        with fixture_path.open(mode="r", encoding="utf-8") as f:
             f_data = f.read()
-        return decode_json(f_data)
-    msg = f"Could not find the {fixture_name} fixture"
-    raise FileNotFoundError(msg)
+
+    return decode_json(f_data)
 
 
 async def open_fixture_async(fixtures_path: Any, fixture_name: str) -> Any:
-    """Load and parse a JSON fixture file asynchronously.
+    """Load and parse a JSON fixture file asynchronously with compression support.
+
+    Supports reading from:
+    - Regular JSON files (.json)
+    - Gzipped JSON files (.json.gz)
+    - Zipped JSON files (.json.zip)
+
+    For compressed files, uses sync reading in a thread pool since gzip and zipfile
+    don't have native async equivalents.
 
     Args:
-        fixtures_path: The path to look for fixtures (pathlib.Path or anyio.Path)
+        fixtures_path: The path to look for fixtures (pathlib.Path)
         fixture_name: The fixture name to load.
 
-    Raises:
-        FileNotFoundError: Fixtures not found.
-        MissingDependencyError: The `anyio` library is required to use this function.
 
     Returns:
         The parsed JSON data
     """
+    # Use sync path finding since it's fast
+    fixture_path = _find_fixture_file(fixtures_path, fixture_name)
+
+    if fixture_path.suffix in {".gz", ".zip"}:
+        # For compressed files, run in thread pool since they don't have async equivalents
+        read_func = async_(_read_compressed_file)
+        f_data = await read_func(fixture_path)
+    else:
+        # For regular JSON files, use async file reading
+        async_read = async_(lambda p: p.read_text(encoding="utf-8"))
+        f_data = await async_read(fixture_path)
+
+    return decode_json(f_data)
+
+
+def _serialize_data(data: Any) -> str:
+    """Serialize data to JSON string, handling different input types.
+
+    Args:
+        data: Data to serialize. Can be dict, list, or SQLSpec model types
+
+    Returns:
+        JSON string representation of the data
+    """
+    if isinstance(data, (list, tuple)):
+        # List of models or dicts - convert each item, handling primitives
+        serialized_items: list[Any] = []
+        for item in data:
+            # Use schema_dump for structured data, pass primitives through
+            if isinstance(item, (str, int, float, bool, type(None))):
+                serialized_items.append(item)
+            else:
+                serialized_items.append(schema_dump(item))
+        return encode_json(serialized_items)
+    # Single model, dict, or other type - try schema_dump first, fallback for primitives
+    if isinstance(data, (str, int, float, bool, type(None))):
+        return encode_json(data)
+    return encode_json(schema_dump(data))
+
+
+def write_fixture(
+    fixtures_path: str,
+    table_name: str,
+    data: "Union[ModelDictList, list[dict[str, Any]], SupportedSchemaModel]",
+    storage_backend: str = "local",
+    compress: bool = False,
+    **storage_kwargs: Any,
+) -> None:
+    """Write fixture data to storage using SQLSpec storage backend.
+
+    Args:
+        fixtures_path: Base path where fixtures should be stored
+        table_name: Name of the table/fixture (used as filename)
+        data: Data to write - can be list of dicts, models, or single model
+        storage_backend: Storage backend to use (default: "local")
+        compress: Whether to gzip compress the output
+        **storage_kwargs: Additional arguments for the storage backend
+
+    Raises:
+        ValueError: If storage backend is not found
+    """
+    # Get the storage backend using URI-based registration
+    # For "local" backend, use file:// URI with base_path parameter
+    if storage_backend == "local":
+        uri = "file://"
+        storage_kwargs["base_path"] = str(Path(fixtures_path).resolve())
+    else:
+        uri = storage_backend
+
     try:
-        from anyio import Path as AsyncPath
-    except ImportError as exc:
-        raise MissingDependencyError(package="anyio") from exc
-
-    fixture = AsyncPath(fixtures_path / f"{fixture_name}.json")
-    if await fixture.exists():
-        async with await fixture.open(mode="r", encoding="utf-8") as f:
-            f_data = await f.read()
-        return decode_json(f_data)
-    msg = f"Could not find the {fixture_name} fixture"
-    raise FileNotFoundError(msg)
+        storage = storage_registry.get(uri, **storage_kwargs)
+    except Exception as exc:
+        msg = f"Failed to get storage backend for '{storage_backend}': {exc}"
+        raise ValueError(msg) from exc
+
+    # Serialize the data
+    json_content = _serialize_data(data)
+
+    # Determine file path and content - use relative path from the base path
+    if compress:
+        file_path = f"{table_name}.json.gz"
+        content = gzip.compress(json_content.encode("utf-8"))
+        storage.write_bytes(file_path, content)
+    else:
+        file_path = f"{table_name}.json"
+        storage.write_text(file_path, json_content)
+
+
+async def write_fixture_async(
+    fixtures_path: str,
+    table_name: str,
+    data: "Union[ModelDictList, list[dict[str, Any]], SupportedSchemaModel]",
+    storage_backend: str = "local",
+    compress: bool = False,
+    **storage_kwargs: Any,
+) -> None:
+    """Write fixture data to storage using SQLSpec storage backend asynchronously.
+
+    Args:
+        fixtures_path: Base path where fixtures should be stored
+        table_name: Name of the table/fixture (used as filename)
+        data: Data to write - can be list of dicts, models, or single model
+        storage_backend: Storage backend to use (default: "local")
+        compress: Whether to gzip compress the output
+        **storage_kwargs: Additional arguments for the storage backend
+
+    Raises:
+        ValueError: If storage backend is not found
+    """
+    # Get the storage backend using URI-based registration
+    # For "local" backend, use file:// URI with base_path parameter
+    if storage_backend == "local":
+        uri = "file://"
+        storage_kwargs["base_path"] = str(Path(fixtures_path).resolve())
+    else:
+        uri = storage_backend
+
+    try:
+        storage = storage_registry.get(uri, **storage_kwargs)
+    except Exception as exc:
+        msg = f"Failed to get storage backend for '{storage_backend}': {exc}"
+        raise ValueError(msg) from exc
+
+    # Serialize the data in a thread pool since it might be CPU intensive
+    serialize_func = async_(_serialize_data)
+    json_content = await serialize_func(data)
+
+    # Determine file path and content
+    if compress:
+        file_path = f"{table_name}.json.gz"
+        # Compress in thread pool since gzip is CPU intensive
+        compress_func = async_(lambda content: gzip.compress(content.encode("utf-8")))
+        content = await compress_func(json_content)
+        await storage.write_bytes_async(file_path, content)
+    else:
+        file_path = f"{table_name}.json"
+        await storage.write_text_async(file_path, json_content)

sqlspec/utils/module_loader.py

@@ -1,4 +1,8 @@
-"""General utility functions."""
+"""Module loading utilities for SQLSpec.
+
+Provides functions for dynamic module imports and path resolution.
+Used for loading modules from dotted paths and converting module paths to filesystem paths.
+"""
 
 import importlib
 from importlib.util import find_spec