sqlspec 0.36.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

sqlspec/utils/correlation.py

@@ -0,0 +1,134 @@
+"""Correlation ID tracking for distributed tracing.
+
+This module provides utilities for tracking correlation IDs across
+database operations, enabling distributed tracing and debugging.
+"""
+
+import uuid
+from collections.abc import Generator, MutableMapping
+from contextlib import contextmanager
+from contextvars import ContextVar
+from logging import Logger, LoggerAdapter
+from typing import Any, ClassVar
+
+__all__ = ("CorrelationContext", "correlation_context", "get_correlation_adapter")
+
+correlation_id_var: "ContextVar[str | None]" = ContextVar("sqlspec_correlation_id", default=None)
+
+
+class CorrelationContext:
+    """Context manager for correlation ID tracking.
+
+    This class provides a context-aware way to track correlation IDs
+    across async and sync operations.
+    """
+
+    _correlation_id: ClassVar["ContextVar[str | None]"] = correlation_id_var
+
+    @classmethod
+    def get(cls) -> str | None:
+        """Get the current correlation ID.
+
+        Returns:
+            The current correlation ID or None if not set
+        """
+        return cls._correlation_id.get()
+
+    @classmethod
+    def set(cls, correlation_id: str | None) -> None:
+        """Set the correlation ID.
+
+        Args:
+            correlation_id: The correlation ID to set
+        """
+        cls._correlation_id.set(correlation_id)
+
+    @classmethod
+    def generate(cls) -> str:
+        """Generate a new correlation ID.
+
+        Returns:
+            A new UUID-based correlation ID
+        """
+        return str(uuid.uuid4())
+
+    @classmethod
+    @contextmanager
+    def context(cls, correlation_id: str | None = None) -> Generator[str, None, None]:
+        """Context manager for correlation ID scope.
+
+        Args:
+            correlation_id: The correlation ID to use. If None, generates a new one.
+
+        Yields:
+            The correlation ID being used
+        """
+        if correlation_id is None:
+            correlation_id = cls.generate()
+
+        previous_id = cls.get()
+
+        try:
+            cls.set(correlation_id)
+            yield correlation_id
+        finally:
+            cls.set(previous_id)
+
+    @classmethod
+    def clear(cls) -> None:
+        """Clear the current correlation ID."""
+        cls.set(None)
+
+    @classmethod
+    def to_dict(cls) -> "dict[str, Any]":
+        """Get correlation context as a dictionary.
+
+        Returns:
+            Dictionary with correlation_id key if set
+        """
+        correlation_id = cls.get()
+        return {"correlation_id": correlation_id} if correlation_id else {}
+
+
+@contextmanager
+def correlation_context(correlation_id: str | None = None) -> Generator[str, None, None]:
+    """Convenience context manager for correlation ID tracking.
+
+    Args:
+        correlation_id: Optional correlation ID. If None, generates a new one.
+
+    Yields:
+        The active correlation ID
+
+    Example:
+        ```python
+        with correlation_context() as correlation_id:
+            logger.info(
+                "Processing request",
+                extra={"correlation_id": correlation_id},
+            )
+        ```
+    """
+    with CorrelationContext.context(correlation_id) as cid:
+        yield cid
+
+
+class _CorrelationAdapter(LoggerAdapter):  # pyright: ignore
+    """Logger adapter that adds correlation ID to all logs."""
+
+    def process(self, msg: str, kwargs: MutableMapping[str, Any]) -> "tuple[str, dict[str, Any]]":
+        """Add correlation ID to the log record."""
+
+        extra = kwargs.get("extra", {})
+
+        if correlation_id := CorrelationContext.get():
+            extra["correlation_id"] = correlation_id
+
+        kwargs["extra"] = extra
+        return msg, dict(kwargs)
+
+
+def get_correlation_adapter(logger: Logger | LoggerAdapter) -> LoggerAdapter:  # pyright: ignore
+    """Get a logger adapter that automatically includes correlation ID."""
+
+    return _CorrelationAdapter(logger, {})

sqlspec/utils/deprecation.py

@@ -0,0 +1,190 @@
+"""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 collections.abc import Callable
+from functools import WRAPPER_ASSIGNMENTS, WRAPPER_UPDATES
+from typing import Generic, Literal, cast
+from warnings import warn
+
+from typing_extensions import ParamSpec, TypeVar
+
+__all__ = ("deprecated", "warn_deprecation")
+
+
+T = TypeVar("T")
+P = ParamSpec("P")
+DeprecatedKind = Literal["function", "method", "classmethod", "attribute", "property", "class", "parameter", "import"]
+
+
+def warn_deprecation(
+    version: str,
+    deprecated_name: str,
+    kind: DeprecatedKind,
+    *,
+    removal_in: str | None = None,
+    alternative: str | None = None,
+    info: str | None = None,
+    pending: bool = False,
+    stacklevel: int = 2,
+) -> None:
+    """Warn about a call to a deprecated function.
+
+    Args:
+        version: SQLSpec version where the deprecation will occur
+        deprecated_name: Name of the deprecated function
+        removal_in: SQLSpec version where the deprecated function will be removed
+        alternative: Name of a function that should be used instead
+        info: Additional information
+        pending: Use :class:`warnings.PendingDeprecationWarning` instead of :class:`warnings.DeprecationWarning`
+        kind: Type of the deprecated thing
+        stacklevel: Warning stacklevel to report the correct caller site.
+    """
+    parts = []
+
+    if kind == "import":
+        access_type = "Import of"
+    elif kind in {"function", "method"}:
+        access_type = "Call to"
+    else:
+        access_type = "Use of"
+
+    if pending:
+        parts.append(f"{access_type} {kind} awaiting deprecation '{deprecated_name}'")  # pyright: ignore[reportUnknownMemberType]
+    else:
+        parts.append(f"{access_type} deprecated {kind} '{deprecated_name}'")  # pyright: ignore[reportUnknownMemberType]
+
+    parts.extend(  # pyright: ignore[reportUnknownMemberType]
+        (f"Deprecated in SQLSpec {version}", f"This {kind} will be removed in {removal_in or 'the next major version'}")
+    )
+    if alternative:
+        parts.append(f"Use {alternative!r} instead")  # pyright: ignore[reportUnknownMemberType]
+
+    if info:
+        parts.append(info)  # pyright: ignore[reportUnknownMemberType]
+
+    text = ". ".join(parts)  # pyright: ignore[reportUnknownArgumentType]
+    warning_class = PendingDeprecationWarning if pending else DeprecationWarning
+
+    warn(text, warning_class, stacklevel=stacklevel)
+
+
+def deprecated(
+    version: str,
+    *,
+    removal_in: str | None = None,
+    alternative: str | None = None,
+    info: str | None = None,
+    pending: bool = False,
+    kind: Literal["function", "method", "classmethod", "property"] | None = None,
+) -> Callable[[Callable[P, T]], Callable[P, T]]:
+    """Create a decorator wrapping a function, method or property with a deprecation warning.
+
+    Args:
+        version: SQLSpec version where the deprecation will occur
+        removal_in: SQLSpec version where the deprecated function will be removed
+        alternative: Name of a function that should be used instead
+        info: Additional information
+        pending: Use :class:`warnings.PendingDeprecationWarning` instead of :class:`warnings.DeprecationWarning`
+        kind: Type of the deprecated callable. If ``None``, will use ``inspect`` to figure
+            out if it's a function or method
+
+    Returns:
+        A decorator wrapping the function call with a warning
+    """
+
+    return cast(
+        "Callable[[Callable[P, T]], Callable[P, T]]",
+        _DeprecatedFactory(
+            version=version, removal_in=removal_in, alternative=alternative, info=info, pending=pending, kind=kind
+        ),
+    )
+
+
+class _DeprecatedWrapper(Generic[P, T]):
+    __slots__ = ("__dict__", "_alternative", "_func", "_info", "_kind", "_pending", "_removal_in", "_version")
+
+    def __init__(
+        self,
+        func: Callable[P, T],
+        *,
+        version: str,
+        removal_in: str | None,
+        alternative: str | None,
+        info: str | None,
+        pending: bool,
+        kind: Literal["function", "method", "classmethod", "property"] | None,
+    ) -> None:
+        self._func = func
+        self._version = version
+        self._removal_in = removal_in
+        self._alternative = alternative
+        self._info = info
+        self._pending = pending
+        self._kind = kind
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
+        kind = cast("DeprecatedKind", self._kind or ("method" if inspect.ismethod(self._func) else "function"))
+        warn_deprecation(
+            version=self._version,
+            deprecated_name=self._func.__name__,
+            info=self._info,
+            alternative=self._alternative,
+            pending=self._pending,
+            removal_in=self._removal_in,
+            kind=kind,
+        )
+        return self._func(*args, **kwargs)
+
+
+class _DeprecatedFactory(Generic[P, T]):
+    __slots__ = ("_alternative", "_info", "_kind", "_pending", "_removal_in", "_version")
+
+    def __init__(
+        self,
+        *,
+        version: str,
+        removal_in: str | None,
+        alternative: str | None,
+        info: str | None,
+        pending: bool,
+        kind: Literal["function", "method", "classmethod", "property"] | None,
+    ) -> None:
+        self._version = version
+        self._removal_in = removal_in
+        self._alternative = alternative
+        self._info = info
+        self._pending = pending
+        self._kind: Literal["function", "method", "classmethod", "property"] | None = kind
+
+    def __call__(self, func: Callable[P, T]) -> Callable[P, T]:
+        wrapper = _DeprecatedWrapper(
+            func,
+            version=self._version,
+            removal_in=self._removal_in,
+            alternative=self._alternative,
+            info=self._info,
+            pending=self._pending,
+            kind=self._kind,
+        )
+        return cast("Callable[P, T]", _copy_wrapper_metadata(wrapper, func))
+
+
+def _copy_wrapper_metadata(wrapper: "_DeprecatedWrapper[P, T]", func: "Callable[P, T]") -> "_DeprecatedWrapper[P, T]":
+    assignments: tuple[str, ...] = tuple(WRAPPER_ASSIGNMENTS)
+    updates: tuple[str, ...] = tuple(WRAPPER_UPDATES)
+    wrapper_dict = wrapper.__dict__
+    wrapper_dict.update(getattr(func, "__dict__", {}))
+    for attr in assignments:
+        if hasattr(func, attr):
+            wrapper_dict[attr] = getattr(func, attr)
+    wrapper_dict["__wrapped__"] = func
+    for attr in updates:
+        if attr == "__dict__":
+            continue
+        if hasattr(func, attr):
+            wrapper_dict[attr] = getattr(func, attr)
+    return wrapper

sqlspec/utils/fixtures.py

@@ -0,0 +1,258 @@
+"""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 TYPE_CHECKING, Any
+
+from sqlspec.storage import storage_registry
+from sqlspec.utils.serializers import from_json as decode_json
+from sqlspec.utils.serializers import schema_dump
+from sqlspec.utils.serializers import to_json as encode_json
+from sqlspec.utils.sync_tools import async_
+
+if TYPE_CHECKING:
+    from sqlspec.typing import SupportedSchemaModel
+
+__all__ = ("open_fixture", "open_fixture_async", "write_fixture", "write_fixture_async")
+
+
+def _read_text_sync(path: "Path") -> str:
+    return path.read_text(encoding="utf-8")
+
+
+def _compress_text(content: str) -> bytes:
+    return gzip.compress(content.encode("utf-8"))
+
+
+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:
+            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")
+            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)
+
+    for extension in [".json", ".json.gz", ".json.zip"]:
+        fixture_path = base_path / f"{fixture_name}{extension}"
+        if fixture_path.exists():
+            return fixture_path
+
+    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 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)
+        fixture_name: The fixture name to load.
+
+
+    Returns:
+        The parsed JSON data
+    """
+    fixture_path = _find_fixture_file(fixtures_path, fixture_name)
+
+    if fixture_path.suffix in {".gz", ".zip"}:
+        f_data = _read_compressed_file(fixture_path)
+    else:
+        with fixture_path.open(mode="r", encoding="utf-8") as f:
+            f_data = f.read()
+
+    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 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)
+        fixture_name: The fixture name to load.
+
+
+    Returns:
+        The parsed JSON data
+    """
+    fixture_path = _find_fixture_file(fixtures_path, fixture_name)
+
+    if fixture_path.suffix in {".gz", ".zip"}:
+        read_func = async_(_read_compressed_file)
+        f_data = await read_func(fixture_path)
+    else:
+        async_read = async_(_read_text_sync)
+        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)):
+        serialized_items: list[Any] = []
+
+        for item in data:
+            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)
+    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: "list[SupportedSchemaModel] | 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
+    """
+    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
+
+    json_content = _serialize_data(data)
+
+    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: "list[SupportedSchemaModel] | 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
+    """
+    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_func = async_(_serialize_data)
+    json_content = await serialize_func(data)
+
+    if compress:
+        file_path = f"{table_name}.json.gz"
+        compress_func = async_(_compress_text)
+        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)