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

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

sqlspec/utils/serializers.py

@@ -1,3 +1,9 @@
+"""JSON serialization utilities for SQLSpec.
+
+Re-exports common JSON encoding and decoding functions from the core
+serialization module for convenient access.
+"""
+
 from sqlspec._serialization import decode_json as from_json
 from sqlspec._serialization import encode_json as to_json
 

sqlspec/utils/singleton.py

@@ -1,3 +1,9 @@
+"""Singleton pattern implementation for SQLSpec.
+
+Provides a thread-safe metaclass for implementing the singleton pattern.
+Used for creating classes that should have only one instance per class type.
+"""
+
 import threading
 from typing import Any, TypeVar
 

sqlspec/utils/sync_tools.py

@@ -33,6 +33,11 @@ class CapacityLimiter:
     """Limits the number of concurrent operations using a semaphore."""
 
     def __init__(self, total_tokens: int) -> None:
+        """Initialize the capacity limiter.
+
+        Args:
+            total_tokens: Maximum number of concurrent operations allowed
+        """
         self._total_tokens = total_tokens
         self._semaphore_instance: Optional[asyncio.Semaphore] = None
 
@@ -44,13 +49,16 @@ class CapacityLimiter:
         return self._semaphore_instance
 
     async def acquire(self) -> None:
+        """Acquire a token from the semaphore."""
         await self._semaphore.acquire()
 
     def release(self) -> None:
+        """Release a token back to the semaphore."""
         self._semaphore.release()
 
     @property
     def total_tokens(self) -> int:
+        """Get the total number of tokens available."""
         return self._total_tokens
 
     @total_tokens.setter
@@ -59,6 +67,7 @@ class CapacityLimiter:
         self._semaphore_instance = None
 
     async def __aenter__(self) -> None:
+        """Async context manager entry."""
         await self.acquire()
 
     async def __aexit__(
@@ -67,6 +76,7 @@ class CapacityLimiter:
         exc_val: "Optional[BaseException]",
         exc_tb: "Optional[TracebackType]",
     ) -> None:
+        """Async context manager exit."""
         self.release()
 
 
@@ -225,7 +235,6 @@ def with_ensure_async_(
     Returns:
         An async context manager that runs the original context manager.
     """
-
     if isinstance(obj, AbstractContextManager):
         return cast("AbstractAsyncContextManager[T]", _ContextManagerWrapper(obj))
     return obj