sqlspec 0.11.0__py3-none-any.whl → 0.12.0__py3-none-any.whl

sqlspec/utils/correlation.py

@@ -0,0 +1,155 @@
+"""Correlation ID tracking for distributed tracing.
+
+This module provides utilities for tracking correlation IDs across
+database operations, enabling distributed tracing and debugging.
+"""
+
+from __future__ import annotations
+
+import uuid
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Generator, MutableMapping
+    from logging import LoggerAdapter
+
+__all__ = ("CorrelationContext", "correlation_context", "get_correlation_adapter")
+
+
+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: ContextVar[str | None] = ContextVar("sqlspec_correlation_id", default=None)
+
+    @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()
+
+        # Save the current correlation ID
+        previous_id = cls.get()
+
+        try:
+            # Set the new correlation ID
+            cls.set(correlation_id)
+            yield correlation_id
+        finally:
+            # Restore the previous correlation ID
+            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},
+            )
+            # All operations within this context will have the same correlation ID
+        ```
+    """
+    with CorrelationContext.context(correlation_id) as cid:
+        yield cid
+
+
+def get_correlation_adapter(logger: Any) -> LoggerAdapter:
+    """Get a logger adapter that automatically includes correlation ID.
+
+    Args:
+        logger: The base logger to wrap
+
+    Returns:
+        LoggerAdapter that includes correlation ID in all logs
+    """
+    from logging import LoggerAdapter
+
+    class CorrelationAdapter(LoggerAdapter):
+        """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.
+
+            Args:
+                msg: The log message
+                kwargs: Keyword arguments for the log record
+
+            Returns:
+                The message and updated kwargs
+            """
+            extra = kwargs.get("extra", {})
+
+            # Add correlation ID if available
+            if correlation_id := CorrelationContext.get():
+                extra["correlation_id"] = correlation_id
+
+            kwargs["extra"] = extra
+            return msg, dict(kwargs)
+
+    return CorrelationAdapter(logger, {})

sqlspec/utils/deprecation.py

@@ -44,15 +44,12 @@ def warn_deprecation(
         access_type = "Use of"
 
     if pending:
-        parts.append(f"{access_type} {kind} awaiting deprecation {deprecated_name!r}")  # pyright: ignore[reportUnknownMemberType]
+        parts.append(f"{access_type} {kind} awaiting deprecation '{deprecated_name}'")  # pyright: ignore[reportUnknownMemberType]
     else:
-        parts.append(f"{access_type} deprecated {kind} {deprecated_name!r}")  # pyright: ignore[reportUnknownMemberType]
+        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'}",
-        ),
+        (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]

sqlspec/utils/fixtures.py

@@ -1,21 +1,17 @@
-from typing import TYPE_CHECKING, Any, Union
+from pathlib import Path
+from typing import Any
 
 from sqlspec._serialization import decode_json
 from sqlspec.exceptions import MissingDependencyError
 
-if TYPE_CHECKING:
-    from pathlib import Path
-
-    from anyio import Path as AsyncPath
-
 __all__ = ("open_fixture", "open_fixture_async")
 
 
-def open_fixture(fixtures_path: "Union[Path, AsyncPath]", fixture_name: str) -> "Any":
+def open_fixture(fixtures_path: Any, fixture_name: str) -> Any:
     """Loads JSON file with the specified fixture name
 
     Args:
-        fixtures_path: :class:`pathlib.Path` | :class:`anyio.Path` The path to look for fixtures
+        fixtures_path: The path to look for fixtures (pathlib.Path or anyio.Path)
         fixture_name (str): The fixture name to load.
 
     Raises:
@@ -24,7 +20,6 @@ def open_fixture(fixtures_path: "Union[Path, AsyncPath]", fixture_name: str) ->
     Returns:
         Any: The parsed JSON data
     """
-    from pathlib import Path
 
     fixture = Path(fixtures_path / f"{fixture_name}.json")
     if fixture.exists():
@@ -35,11 +30,11 @@ def open_fixture(fixtures_path: "Union[Path, AsyncPath]", fixture_name: str) ->
     raise FileNotFoundError(msg)
 
 
-async def open_fixture_async(fixtures_path: "Union[Path, AsyncPath]", fixture_name: str) -> "Any":
+async def open_fixture_async(fixtures_path: Any, fixture_name: str) -> Any:
     """Loads JSON file with the specified fixture name
 
     Args:
-        fixtures_path: :class:`pathlib.Path` | :class:`anyio.Path` The path to look for fixtures
+        fixtures_path: The path to look for fixtures (pathlib.Path or anyio.Path)
         fixture_name (str): The fixture name to load.
 
     Raises:

sqlspec/utils/logging.py

@@ -0,0 +1,135 @@
+"""Logging utilities for SQLSpec.
+
+This module provides utilities for structured logging with correlation IDs.
+Users should configure their own logging handlers and levels as needed.
+SQLSpec provides StructuredFormatter for JSON-formatted logs if desired.
+"""
+
+from __future__ import annotations
+
+import logging
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Any
+
+from sqlspec._serialization import encode_json
+
+if TYPE_CHECKING:
+    from logging import LogRecord
+
+__all__ = ("StructuredFormatter", "correlation_id_var", "get_correlation_id", "get_logger", "set_correlation_id")
+
+# Context variable for correlation ID tracking
+correlation_id_var: ContextVar[str | None] = ContextVar("correlation_id", default=None)
+
+
+def set_correlation_id(correlation_id: str | None) -> None:
+    """Set the correlation ID for the current context.
+
+    Args:
+        correlation_id: The correlation ID to set, or None to clear
+    """
+    correlation_id_var.set(correlation_id)
+
+
+def get_correlation_id() -> str | None:
+    """Get the current correlation ID.
+
+    Returns:
+        The current correlation ID or None if not set
+    """
+    return correlation_id_var.get()
+
+
+class StructuredFormatter(logging.Formatter):
+    """Structured JSON formatter with correlation ID support."""
+
+    def format(self, record: LogRecord) -> str:
+        """Format log record as structured JSON.
+
+        Args:
+            record: The log record to format
+
+        Returns:
+            JSON formatted log entry
+        """
+        # Base log entry
+        log_entry = {
+            "timestamp": self.formatTime(record, self.datefmt),
+            "level": record.levelname,
+            "logger": record.name,
+            "message": record.getMessage(),
+            "module": record.module,
+            "function": record.funcName,
+            "line": record.lineno,
+        }
+
+        # Add correlation ID if available
+        if correlation_id := get_correlation_id():
+            log_entry["correlation_id"] = correlation_id
+
+        # Add any extra fields from the record
+        if hasattr(record, "extra_fields"):
+            log_entry.update(record.extra_fields)  # pyright: ignore
+
+        # Add exception info if present
+        if record.exc_info:
+            log_entry["exception"] = self.formatException(record.exc_info)
+
+        return encode_json(log_entry)
+
+
+class CorrelationIDFilter(logging.Filter):
+    """Filter that adds correlation ID to log records."""
+
+    def filter(self, record: LogRecord) -> bool:
+        """Add correlation ID to record if available.
+
+        Args:
+            record: The log record to filter
+
+        Returns:
+            Always True to pass the record through
+        """
+        if correlation_id := get_correlation_id():
+            record.correlation_id = correlation_id
+        return True
+
+
+def get_logger(name: str | None = None) -> logging.Logger:
+    """Get a logger instance with standardized configuration.
+
+    Args:
+        name: Logger name. If not provided, returns the root sqlspec logger.
+
+    Returns:
+        Configured logger instance
+    """
+    if name is None:
+        return logging.getLogger("sqlspec")
+
+    # Ensure all loggers are under the sqlspec namespace
+    if not name.startswith("sqlspec"):
+        name = f"sqlspec.{name}"
+
+    logger = logging.getLogger(name)
+
+    # Add correlation ID filter if not already present
+    if not any(isinstance(f, CorrelationIDFilter) for f in logger.filters):
+        logger.addFilter(CorrelationIDFilter())
+
+    return logger
+
+
+def log_with_context(logger: logging.Logger, level: int, message: str, **extra_fields: Any) -> None:
+    """Log a message with structured extra fields.
+
+    Args:
+        logger: The logger to use
+        level: Log level
+        message: Log message
+        **extra_fields: Additional fields to include in structured logs
+    """
+    # Create a LogRecord with extra fields
+    record = logger.makeRecord(logger.name, level, "(unknown file)", 0, message, (), None)
+    record.extra_fields = extra_fields
+    logger.handle(record)

sqlspec/utils/module_loader.py

@@ -1,18 +1,11 @@
 """General utility functions."""
 
-import sys
-from importlib import import_module
+import importlib
 from importlib.util import find_spec
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
+from typing import Any, Optional
 
-if TYPE_CHECKING:
-    from types import ModuleType
-
-__all__ = (
-    "import_string",
-    "module_to_os_path",
-)
+__all__ = ("import_string", "module_to_os_path")
 
 
 def module_to_os_path(dotted_path: str = "app") -> "Path":
@@ -51,42 +44,51 @@ def import_string(dotted_path: str) -> "Any":
     Args:
         dotted_path: The path of the module to import.
 
-    Raises:
-        ImportError: Could not import the module.
-
     Returns:
         object: The imported object.
     """
 
-    def _is_loaded(module: "Optional[ModuleType]") -> bool:
-        spec = getattr(module, "__spec__", None)
-        initializing = getattr(spec, "_initializing", False)
-        return bool(module and spec and not initializing)
-
-    def _cached_import(module_path: str, class_name: str) -> Any:
-        """Import and cache a class from a module.
-
-        Args:
-            module_path: dotted path to module.
-            class_name: Class or function name.
-
-        Returns:
-            object: The imported class or function
-        """
-        # Check whether module is loaded and fully initialized.
-        module = sys.modules.get(module_path)
-        if not _is_loaded(module):
-            module = import_module(module_path)
-        return getattr(module, class_name)
-
-    try:
-        module_path, class_name = dotted_path.rsplit(".", 1)
-    except ValueError as e:
-        msg = "%s doesn't look like a module path"
-        raise ImportError(msg, dotted_path) from e
+    def _raise_import_error(msg: str, exc: "Optional[Exception]" = None) -> None:
+        if exc is not None:
+            raise ImportError(msg) from exc
+        raise ImportError(msg)
 
+    obj: Any = None
     try:
-        return _cached_import(module_path, class_name)
-    except AttributeError as e:
-        msg = "Module '%s' does not define a '%s' attribute/class"
-        raise ImportError(msg, module_path, class_name) from e
+        parts = dotted_path.split(".")
+        module = None
+        i = len(parts)  # Initialize to full length
+
+        for i in range(len(parts), 0, -1):
+            module_path = ".".join(parts[:i])
+            try:
+                module = importlib.import_module(module_path)
+                break
+            except ModuleNotFoundError:
+                continue
+        else:
+            _raise_import_error(f"{dotted_path} doesn't look like a module path")
+
+        if module is None:
+            _raise_import_error(f"Failed to import any module from {dotted_path}")
+
+        obj = module
+        attrs = parts[i:]
+        if not attrs and i == len(parts) and len(parts) > 1:
+            parent_module_path = ".".join(parts[:-1])
+            attr = parts[-1]
+            try:
+                parent_module = importlib.import_module(parent_module_path)
+            except Exception:
+                return obj
+            if not hasattr(parent_module, attr):
+                _raise_import_error(f"Module '{parent_module_path}' has no attribute '{attr}' in '{dotted_path}'")
+        for attr in attrs:
+            if not hasattr(obj, attr):
+                _raise_import_error(
+                    f"Module '{module.__name__ if module is not None else 'unknown'}' has no attribute '{attr}' in '{dotted_path}'"
+                )
+            obj = getattr(obj, attr)
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        _raise_import_error(f"Could not import '{dotted_path}': {e}", e)
+    return obj

sqlspec/utils/serializers.py

@@ -0,0 +1,4 @@
+from sqlspec._serialization import decode_json as from_json
+from sqlspec._serialization import encode_json as to_json
+
+__all__ = ("from_json", "to_json")

sqlspec/utils/singleton.py

@@ -1,3 +1,4 @@
+import threading
 from typing import Any, TypeVar
 
 __all__ = ("SingletonMeta",)
@@ -11,6 +12,7 @@ class SingletonMeta(type):
 
     # We store instances keyed by the class type
     _instances: dict[type, object] = {}
+    _lock = threading.Lock()
 
     def __call__(cls: type[_T], *args: Any, **kwargs: Any) -> _T:
         """Call method for the singleton metaclass.
@@ -23,13 +25,9 @@ class SingletonMeta(type):
         Returns:
             The singleton instance of the class.
         """
-        # Use SingletonMeta._instances to access the class attribute
         if cls not in SingletonMeta._instances:  # pyright: ignore[reportUnnecessaryContains]
-            # Create the instance using super().__call__ which calls the class's __new__ and __init__
-            instance = super().__call__(*args, **kwargs)  # type: ignore[misc]
-            SingletonMeta._instances[cls] = instance
-
-        # Return the cached instance. We cast here because the dictionary stores `object`,
-        # but we know it's of type _T for the given cls key.
-        # Mypy might need an ignore here depending on configuration, but pyright should handle it.
+            with SingletonMeta._lock:
+                if cls not in SingletonMeta._instances:
+                    instance = super().__call__(*args, **kwargs)  # type: ignore[misc]
+                    SingletonMeta._instances[cls] = instance
         return SingletonMeta._instances[cls]  # type: ignore[return-value]

sqlspec/utils/sync_tools.py

@@ -3,15 +3,7 @@ import functools
 import inspect
 import sys
 from contextlib import AbstractAsyncContextManager, AbstractContextManager
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Generic,
-    Optional,
-    TypeVar,
-    Union,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar, Union, cast
 
 from typing_extensions import ParamSpec
 
@@ -44,7 +36,7 @@ class CapacityLimiter:
 
     @property
     def total_tokens(self) -> int:
-        return self._semaphore._value  # noqa: SLF001
+        return self._semaphore._value
 
     @total_tokens.setter
     def total_tokens(self, value: int) -> None:
@@ -55,9 +47,9 @@ class CapacityLimiter:
 
     async def __aexit__(
         self,
-        exc_type: "Optional[type[BaseException]]",  # noqa: PYI036
-        exc_val: "Optional[BaseException]",  # noqa: PYI036
-        exc_tb: "Optional[TracebackType]",  # noqa: PYI036
+        exc_type: "Optional[type[BaseException]]",
+        exc_val: "Optional[BaseException]",
+        exc_tb: "Optional[TracebackType]",
     ) -> None:
         self.release()
 
@@ -96,8 +88,7 @@ def run_(async_function: "Callable[ParamSpecT, Coroutine[Any, Any, ReturnT]]") -
 
 
 def await_(
-    async_function: "Callable[ParamSpecT, Coroutine[Any, Any, ReturnT]]",
-    raise_sync_error: bool = True,
+    async_function: "Callable[ParamSpecT, Coroutine[Any, Any, ReturnT]]", raise_sync_error: bool = True
 ) -> "Callable[ParamSpecT, ReturnT]":
     """Convert an async function to a blocking one, running in the main async loop.
 
@@ -118,7 +109,7 @@ def await_(
         except RuntimeError:
             # No running event loop
             if raise_sync_error:
-                msg = "await_ called without a running event loop and raise_sync_error=True"
+                msg = "Cannot run async function"
                 raise RuntimeError(msg) from None
             return asyncio.run(partial_f())
         else:
@@ -145,7 +136,7 @@ def await_(
             # but the loop isn't running, but handle defensively.
             # loop is not running
             if raise_sync_error:
-                msg = "await_ found a non-running loop via get_running_loop()"
+                msg = "Cannot run async function"
                 raise RuntimeError(msg)
             # Fallback to running in a new loop
             return asyncio.run(partial_f())
@@ -154,9 +145,7 @@ def await_(
 
 
 def async_(
-    function: "Callable[ParamSpecT, ReturnT]",
-    *,
-    limiter: "Optional[CapacityLimiter]" = None,
+    function: "Callable[ParamSpecT, ReturnT]", *, limiter: "Optional[CapacityLimiter]" = None
 ) -> "Callable[ParamSpecT, Awaitable[ReturnT]]":
     """Convert a blocking function to an async one using asyncio.to_thread().
 
@@ -169,10 +158,8 @@ def async_(
         Callable: An async function that runs the original function in a thread.
     """
 
-    async def wrapper(
-        *args: "ParamSpecT.args",
-        **kwargs: "ParamSpecT.kwargs",
-    ) -> "ReturnT":
+    @functools.wraps(function)
+    async def wrapper(*args: "ParamSpecT.args", **kwargs: "ParamSpecT.kwargs") -> "ReturnT":
         partial_f = functools.partial(function, *args, **kwargs)
         used_limiter = limiter or _default_limiter
         async with used_limiter:
@@ -195,6 +182,7 @@ def ensure_async_(
     if inspect.iscoroutinefunction(function):
         return function
 
+    @functools.wraps(function)
     async def wrapper(*args: "ParamSpecT.args", **kwargs: "ParamSpecT.kwargs") -> "ReturnT":
         result = function(*args, **kwargs)
         if inspect.isawaitable(result):
@@ -213,9 +201,9 @@ class _ContextManagerWrapper(Generic[T]):
 
     async def __aexit__(
         self,
-        exc_type: "Optional[type[BaseException]]",  # noqa: PYI036
-        exc_val: "Optional[BaseException]",  # noqa: PYI036
-        exc_tb: "Optional[TracebackType]",  # noqa: PYI036
+        exc_type: "Optional[type[BaseException]]",
+        exc_val: "Optional[BaseException]",
+        exc_tb: "Optional[TracebackType]",
     ) -> "Optional[bool]":
         return self._cm.__exit__(exc_type, exc_val, exc_tb)