From a73409299b81aed118fcfd0e8285d1c3193161a8 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Wed, 21 Oct 2026 18:38:49 +0000 Subject: [PATCH] feat(observability): add wire-level sqlcommenter, zero-RTT Oracle trace piggybacking, and DuckDB statement profiling --- sqlspec/adapters/adbc/config.py | 36 ++- sqlspec/adapters/adbc/core.py | 23 ++ sqlspec/adapters/adbc/driver.py | 83 ++++++- sqlspec/adapters/duckdb/config.py | 4 + sqlspec/adapters/duckdb/core.py | 63 +++++- sqlspec/adapters/duckdb/driver.py | 76 ++++++- sqlspec/adapters/duckdb/pool.py | 36 ++- sqlspec/adapters/oracledb/_typing.py | 22 +- sqlspec/adapters/oracledb/config.py | 1 + sqlspec/adapters/oracledb/core.py | 90 ++++++++ sqlspec/adapters/oracledb/driver.py | 53 +++++ sqlspec/adapters/psqlpy/driver.py | 9 +- sqlspec/builder/_explain.py | 4 + sqlspec/core/sqlcommenter.py | 192 +++++++++++++++- sqlspec/core/statement.py | 6 +- sqlspec/driver/_async.py | 39 +++- sqlspec/driver/_common.py | 106 ++++++++- sqlspec/driver/_sync.py | 41 +++- sqlspec/observability/_observer.py | 42 +++- sqlspec/observability/_runtime.py | 99 +++++++-- .../adapters/_shared/_driver_type_system.py | 3 + .../unit/adapters/test_adbc/test_profiling.py | 57 +++++ .../adapters/test_duckdb/test_profiling.py | 88 ++++++++ .../adapters/test_oracledb/test_driver.py | 205 ++++++++++++++++++ tests/unit/adapters/test_psqlpy/test_core.py | 52 +++++ .../unit/adapters/test_psqlpy/test_driver.py | 52 +++++ tests/unit/builder/test_explain.py | 11 + tests/unit/core/test_sqlcommenter.py | 142 ++++++++++++ .../observability/test_observability_core.py | 100 +++++++++ tests/unit/observability/test_observer.py | 68 ++++++ 30 files changed, 1717 insertions(+), 86 deletions(-) create mode 100644 tests/unit/adapters/test_adbc/test_profiling.py create mode 100644 tests/unit/adapters/test_duckdb/test_profiling.py create mode 100644 tests/unit/adapters/test_oracledb/test_driver.py create mode 100644 tests/unit/adapters/test_psqlpy/test_driver.py create mode 100644 tests/unit/observability/test_observer.py diff --git a/sqlspec/adapters/adbc/config.py b/sqlspec/adapters/adbc/config.py index d96d7d91d..81ab70f61 100644 --- a/sqlspec/adapters/adbc/config.py +++ b/sqlspec/adapters/adbc/config.py @@ -1,11 +1,13 @@ """ADBC database configuration.""" +from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar, Literal, NotRequired, TypedDict, cast from sqlspec.adapters.adbc._typing import AdbcConnection, AdbcCursor, AdbcSessionContext from sqlspec.adapters.adbc.core import ( apply_driver_features, build_connection_config, + configure_adbc_duckdb_profiling, detect_postgres_extensions, get_statement_config, is_postgres_dialect, @@ -15,6 +17,7 @@ resolve_driver_connect_func, ) from sqlspec.adapters.adbc.driver import AdbcDriver, AdbcExceptionHandler +from sqlspec.adapters.duckdb.core import cleanup_duckdb_profile_path, create_duckdb_profile_path from sqlspec.config import ExtensionConfigs, NoPoolSyncConfig from sqlspec.core import StatementConfig from sqlspec.core.capabilities import TypeCoercionCapabilities @@ -135,6 +138,8 @@ class AdbcDriverFeatures(TypedDict): Only option: "poll_queue" (durable table-backed queue with lease-based retries and acknowledgements). ADBC does not have native pub/sub, so poll_queue is the only backend. Defaults to "poll_queue". + enable_statement_profiling: Enable DuckDB JSON statement profiling when using a DuckDB ADBC backend. + Defaults to False. """ json_serializer: "NotRequired[Callable[[Any], str]]" @@ -150,6 +155,7 @@ class AdbcDriverFeatures(TypedDict): enable_events: NotRequired[bool] on_connection_create: "NotRequired[Callable[[AdbcConnection], None]]" events_backend: NotRequired[Literal["poll_queue"]] + enable_statement_profiling: NotRequired[bool] class AdbcConnectionContext(SyncPoolConnectionContext): @@ -176,22 +182,42 @@ def __exit__( class _AdbcSessionConnectionHandler(SyncPoolSessionFactory): - __slots__ = ("_connection",) + __slots__ = ("_connection", "_driver_features", "_profile_path") def __init__(self, config: "AdbcConfig") -> None: super().__init__(config) self._connection: AdbcConnection | None = None + self._driver_features: dict[str, Any] = dict(config.driver_features) + self._profile_path: Path | None = None + + @property + def driver_features(self) -> "dict[str, Any]": + return self._driver_features def acquire_connection(self) -> "AdbcConnection": self._connection = self._config.create_connection() + if self._config._resolved_dialect == "duckdb" and bool( + self._driver_features.get("enable_statement_profiling", False) + ): + self._profile_path = create_duckdb_profile_path() + configure_adbc_duckdb_profiling(self._connection, self._profile_path) + self._driver_features["_duckdb_profile_path"] = str(self._profile_path) return cast("AdbcConnection", self._connection) def release_connection(self, _conn: "AdbcConnection", **kwargs: Any) -> None: if self._connection is None: return - if self._connection is not self._config.connection_instance: - self._connection.close() - self._connection = None + try: + if self._connection is not self._config.connection_instance: + self._connection.close() + finally: + self._connection = None + if self._profile_path is not None: + cleanup_duckdb_profile_path(self._profile_path) + self._profile_path = None + raw_path = self._driver_features.pop("_duckdb_profile_path", None) + if isinstance(raw_path, str): + cleanup_duckdb_profile_path(Path(raw_path)) class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]): @@ -369,7 +395,7 @@ def provide_session( acquire_connection=handler.acquire_connection, release_connection=handler.release_connection, statement_config=statement_config, - driver_features=self.driver_features, + driver_features=handler.driver_features, prepare_driver=self._prepare_driver, ) diff --git a/sqlspec/adapters/adbc/core.py b/sqlspec/adapters/adbc/core.py index a467ac91d..6ec3e0e35 100644 --- a/sqlspec/adapters/adbc/core.py +++ b/sqlspec/adapters/adbc/core.py @@ -1,5 +1,6 @@ """ADBC adapter compiled helpers.""" +import contextlib import datetime as dt import decimal import re @@ -50,6 +51,7 @@ if TYPE_CHECKING: from collections.abc import Callable, Mapping + from pathlib import Path from sqlspec.core import SQL from sqlspec.utils.dispatch import TypeDispatcher @@ -59,6 +61,7 @@ "build_connection_config", "build_profile", "collect_rows", + "configure_adbc_duckdb_profiling", "create_mapped_exception", "detect_dialect", "detect_postgres_extensions", @@ -983,6 +986,7 @@ def apply_driver_features( processed_features.setdefault("enable_paradedb", True) processed_features.setdefault("enable_pg_textsearch", True) processed_features.setdefault("enable_lowercase_column_names", resolved_dialect in {"oracle", "db2"}) + processed_features.setdefault("enable_statement_profiling", False) if json_serializer is not None: statement_config = _apply_adbc_json_serializer(statement_config, json_serializer) @@ -990,6 +994,25 @@ def apply_driver_features( return statement_config, processed_features +def configure_adbc_duckdb_profiling(connection: Any, profile_path: "Path") -> None: + """Enable DuckDB JSON statement profiling on an ADBC DuckDB connection. + + Args: + connection: Active ADBC DuckDB connection. + profile_path: Temporary file path where DuckDB writes statement profiles. + """ + escaped_path = str(profile_path).replace("'", "''") + cursor = connection.cursor() + try: + cursor.execute("PRAGMA enable_profiling='json'") + with contextlib.suppress(Exception): + cursor.execute("SET profiling_coverage='ALL'") + cursor.execute(f"SET profiling_output='{escaped_path}'") + finally: + with contextlib.suppress(Exception): + cursor.close() + + def normalize_column_names( column_names: "list[str]", driver_features: "Mapping[str, Any] | None" = None ) -> "list[str]": diff --git a/sqlspec/adapters/adbc/driver.py b/sqlspec/adapters/adbc/driver.py index 6cd715c69..1344b0550 100644 --- a/sqlspec/adapters/adbc/driver.py +++ b/sqlspec/adapters/adbc/driver.py @@ -5,12 +5,15 @@ """ import contextlib +from pathlib import Path +from time import perf_counter from typing import TYPE_CHECKING, Any, Literal, cast, final from sqlspec.adapters.adbc._typing import AdbcCursor, AdbcNativeError, AdbcSessionContext from sqlspec.adapters.adbc.core import ( apply_driver_features, collect_rows, + configure_adbc_duckdb_profiling, create_mapped_exception, detect_dialect, driver_profile, @@ -31,6 +34,11 @@ resolve_rowcount, ) from sqlspec.adapters.adbc.data_dictionary import AdbcDataDictionary +from sqlspec.adapters.duckdb.core import ( + cleanup_duckdb_profile_path, + create_duckdb_profile_path, + read_duckdb_profile_json, +) from sqlspec.core import ( SQL, StatementConfig, @@ -239,6 +247,7 @@ def dispatch_execute(self, cursor: "AdbcRawCursor", statement: SQL) -> "Executio Execution result with data or row count """ sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + self._ensure_statement_profiling() try: execute_parameters = normalize_postgres_empty_parameters(self._dialect_name, prepared_parameters) @@ -252,6 +261,7 @@ def dispatch_execute(self, cursor: "AdbcRawCursor", statement: SQL) -> "Executio if is_select_like: arrow_table = cursor.fetch_arrow_table() + self._capture_statement_profile(sql) raw_column_names = list(arrow_table.column_names) column_names = normalize_column_names(raw_column_names, self.driver_features) if column_names != raw_column_names and hasattr(arrow_table, "rename_columns"): @@ -270,6 +280,7 @@ def dispatch_execute(self, cursor: "AdbcRawCursor", statement: SQL) -> "Executio ) row_count = self._resolve_count_result_rowcount(cursor, fallback=resolve_rowcount(cursor)) + self._capture_statement_profile(sql) return self.create_execution_result(cursor, rowcount_override=row_count) def dispatch_execute_many(self, cursor: "AdbcRawCursor", statement: SQL) -> "ExecutionResult": @@ -338,28 +349,57 @@ def dispatch_execute_script(self, cursor: "AdbcRawCursor", statement: "SQL") -> successful_count = 0 last_rowcount = 0 + total_statements = len(statements) + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers + profiling_enabled = self._dialect_name == "duckdb" and bool( + self.driver_features.get("enable_statement_profiling", False) + ) + if profiling_enabled: + self._ensure_statement_profiling() try: - for stmt in statements: + for idx, stmt in enumerate(statements): + stmt_started = perf_counter() if has_observers else 0.0 + stmt_params: Any = () if prepared_parameters and ( - len(statements) == 1 or any(marker in stmt for marker in ("?", "$", "%", "@", ":")) + total_statements == 1 or any(marker in stmt for marker in ("?", "$", "%", "@", ":")) ): - postgres_compatible_params = normalize_postgres_empty_parameters( - self._dialect_name, prepared_parameters - ) - cursor.execute(stmt, parameters=postgres_compatible_params) + stmt_params = normalize_postgres_empty_parameters(self._dialect_name, prepared_parameters) + cursor.execute(stmt, parameters=stmt_params) else: cursor.execute(stmt) successful_count += 1 - last_rowcount = self._resolve_count_result_rowcount( + stmt_rowcount = self._resolve_count_result_rowcount( cursor, fallback=normalize_script_rowcount(last_rowcount, cursor) ) + last_rowcount = stmt_rowcount + stmt_profile = self._capture_statement_profile(stmt) if profiling_enabled else None + if has_observers: + stmt_duration = perf_counter() - stmt_started + single_stmt = self._sub_statement(stmt, stmt_params) + try: + stmt_operation = single_stmt.operation_type + finally: + self._release_pooled_statement(single_stmt) + self._emit_script_statement_event( + sql=stmt, + parameters=stmt_params, + operation=stmt_operation, + rows_affected=stmt_rowcount, + duration_s=stmt_duration, + started_at=stmt_started, + profile=stmt_profile, + script_ordinal=idx, + script_statement_count=total_statements, + ) except Exception: handle_postgres_rollback(self._dialect_name, cursor, logger) raise + self._pending_statement_profile = None return self.create_execution_result( cursor, - statement_count=len(statements), + statement_count=total_statements, successful_statements=successful_count, rowcount_override=last_rowcount, is_script_result=True, @@ -826,6 +866,33 @@ def _detect_flightsql_connection(connection: "AdbcConnection") -> bool: def _resolve_column_names(self, description: Any) -> list[str]: return resolve_column_names(description, self._column_name_cache, self.driver_features) + def _ensure_statement_profiling(self) -> "Path | None": + if self._dialect_name != "duckdb" or not bool(self.driver_features.get("enable_statement_profiling", False)): + return None + raw_path = self.driver_features.get("_duckdb_profile_path") + if isinstance(raw_path, Path): + return raw_path + if isinstance(raw_path, str) and raw_path: + return Path(raw_path) + profile_path = create_duckdb_profile_path() + try: + configure_adbc_duckdb_profiling(self.connection, profile_path) + except Exception: + cleanup_duckdb_profile_path(profile_path) + return None + self.driver_features["_duckdb_profile_path"] = str(profile_path) + return profile_path + + def _capture_statement_profile(self, executed_sql: str) -> "dict[str, Any] | None": + if self._dialect_name != "duckdb" or not bool(self.driver_features.get("enable_statement_profiling", False)): + return None + profile_path = self._ensure_statement_profiling() + if profile_path is None: + return None + profile = read_duckdb_profile_json(profile_path, executed_sql) + self._pending_statement_profile = profile + return profile + @staticmethod def _resolve_count_result_rowcount(cursor: "AdbcRawCursor", *, fallback: int) -> int: """Consume ADBC count result streams and return their row count.""" diff --git a/sqlspec/adapters/duckdb/config.py b/sqlspec/adapters/duckdb/config.py index 6cf4ddff7..b2db8778e 100644 --- a/sqlspec/adapters/duckdb/config.py +++ b/sqlspec/adapters/duckdb/config.py @@ -176,6 +176,8 @@ class DuckDBDriverFeatures(TypedDict): extension_config["events"] is set; requires it for migrations. events_backend: Only "poll_queue" is supported, with lease-based retries and acknowledgements. DuckDB has no native pub/sub. + enable_statement_profiling: Enable per-statement JSON operator tree + profiling attached to ``StatementEvent.profile`` (default False). """ extensions: NotRequired[Sequence[DuckDBExtensionConfig]] @@ -186,6 +188,7 @@ class DuckDBDriverFeatures(TypedDict): extension_flags: NotRequired["dict[str, Any]"] enable_events: NotRequired[bool] events_backend: NotRequired[Literal["poll_queue"]] + enable_statement_profiling: NotRequired[bool] class DuckDBConnectionContext(SyncPoolConnectionContext): @@ -369,6 +372,7 @@ def _create_pool(self) -> DuckDBConnectionPool: extensions=extensions_dicts, secrets=secrets_dicts, on_connection_create=self._user_connection_hook, + enable_statement_profiling=bool(self.driver_features.get("enable_statement_profiling", False)), **pool_kwargs, ) diff --git a/sqlspec/adapters/duckdb/core.py b/sqlspec/adapters/duckdb/core.py index 69ffcdfd5..3401dccef 100644 --- a/sqlspec/adapters/duckdb/core.py +++ b/sqlspec/adapters/duckdb/core.py @@ -2,8 +2,10 @@ import contextlib import re +import tempfile from datetime import date, datetime from decimal import Decimal +from pathlib import Path from typing import TYPE_CHECKING, Any, Final, cast from urllib.parse import urlsplit @@ -40,11 +42,11 @@ from sqlspec.storage import ResolvedStorageTarget from sqlspec.storage.backends.fsspec import FSSpecBackend from sqlspec.storage.backends.obstore import ObStoreBackend -from sqlspec.utils.serializers import to_json +from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import quote_identifier, split_qualified_identifier from sqlspec.utils.type_converters import build_decimal_converter, build_uuid_coercions, time_iso_convert from sqlspec.utils.type_guards import has_rowcount -from sqlspec.utils.uuids import uuid_from_string +from sqlspec.utils.uuids import uuid4, uuid_from_string if TYPE_CHECKING: from collections.abc import Callable, Mapping, Sequence @@ -58,13 +60,17 @@ "build_connection_config", "build_profile", "build_statement_config", + "cleanup_duckdb_profile_path", "collect_rows", + "configure_duckdb_statement_profiling", + "create_duckdb_profile_path", "create_mapped_exception", "default_statement_config", "description_has_arrow_divergent_types", "driver_profile", "format_identifier", "normalize_execute_parameters", + "read_duckdb_profile_json", "resolve_rowcount", ) @@ -271,6 +277,7 @@ def apply_driver_features( features: dict[str, Any] = dict(driver_features) if driver_features else {} features.setdefault("enable_events", True) features.setdefault("events_backend", "poll_queue") + features.setdefault("enable_statement_profiling", False) param_config = statement_config.parameter_config json_serializer = features.get("json_serializer") @@ -290,6 +297,58 @@ def apply_driver_features( return statement_config.replace(parameter_config=param_config), features +def create_duckdb_profile_path() -> Path: + """Allocate a unique temporary file path for DuckDB JSON statement profiling.""" + return Path(tempfile.gettempdir()) / f"sqlspec_duckdb_profile_{uuid4().hex}.json" + + +def configure_duckdb_statement_profiling(connection: Any, profile_path: Path) -> None: + """Enable DuckDB JSON statement profiling on a connection without observer emission.""" + escaped_path = str(profile_path).replace("'", "''") + connection.execute("PRAGMA enable_profiling='json'") + with contextlib.suppress(Exception): + connection.execute("SET profiling_coverage='ALL'") + connection.execute(f"SET profiling_output='{escaped_path}'") + + +def cleanup_duckdb_profile_path(profile_path: "Path | str | None") -> None: + """Remove a temporary DuckDB JSON statement profile file if it exists.""" + if not profile_path: + return + with contextlib.suppress(OSError): + Path(profile_path).unlink(missing_ok=True) + + +def read_duckdb_profile_json( + profile_path: "Path | str | None", executed_sql: str | None = None +) -> "dict[str, Any] | None": + """Read and validate the DuckDB JSON profile file written for executed_sql.""" + if not profile_path: + return None + path_obj = Path(profile_path) + if not path_obj.exists(): + return None + try: + raw_bytes = path_obj.read_bytes() + except OSError: + return None + if not raw_bytes: + return None + try: + parsed: Any = from_json(raw_bytes) + except Exception: + return None + if not isinstance(parsed, dict): + return None + query_name = parsed.get("query_name") + if isinstance(query_name, str) and executed_sql: + normalized_query = query_name.strip().rstrip(";").strip() + normalized_sql = executed_sql.strip().rstrip(";").strip() + if normalized_sql and normalized_query != normalized_sql: + return None + return cast("dict[str, Any]", parsed) + + _EXCEPTION_MAPPING: Final[dict[type[BaseException], tuple[type[SQLSpecError], str]]] = { DuckDBCatalogException: (NotFoundError, "catalog error"), DuckDBParserException: (SQLParsingError, "SQL parsing error"), diff --git a/sqlspec/adapters/duckdb/driver.py b/sqlspec/adapters/duckdb/driver.py index 7707a28dc..d5be135bc 100644 --- a/sqlspec/adapters/duckdb/driver.py +++ b/sqlspec/adapters/duckdb/driver.py @@ -16,12 +16,15 @@ _resolve_native_storage_target, _restore_uuid_columns, collect_rows, + configure_duckdb_statement_profiling, + create_duckdb_profile_path, create_mapped_exception, default_statement_config, description_has_arrow_divergent_types, driver_profile, format_identifier, normalize_execute_parameters, + read_duckdb_profile_json, resolve_rowcount, ) from sqlspec.adapters.duckdb.data_dictionary import DuckDBDataDictionary @@ -142,12 +145,17 @@ def dispatch_execute(self, cursor: "DuckDBConnection", statement: SQL) -> "Execu ExecutionResult with execution metadata """ sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + profiling_enabled = bool(self.driver_features.get("enable_statement_profiling", False)) + if profiling_enabled: + self._ensure_statement_profiling(cursor) cursor.execute(sql, normalize_execute_parameters(prepared_parameters)) is_select_like = statement.returns_rows() or self._should_force_select(statement, cursor) if is_select_like: data, column_names, row_count, row_format = self._materialize_select_rows(cursor, sql) + if profiling_enabled: + self._capture_statement_profile(sql) return self.create_execution_result( cursor, selected_data=data, @@ -158,6 +166,10 @@ def dispatch_execute(self, cursor: "DuckDBConnection", statement: SQL) -> "Execu ) row_count = resolve_rowcount(cursor) + if profiling_enabled: + with contextlib.suppress(Exception): + cursor.fetchall() + self._capture_statement_profile(sql) return self.create_execution_result(cursor, rowcount_override=row_count) @@ -222,15 +234,45 @@ def dispatch_execute_script(self, cursor: "DuckDBConnection", statement: SQL) -> sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) + statement_count = len(statements) successful_count = 0 last_result = None - - for stmt in statements: + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers + profiling_enabled = bool(self.driver_features.get("enable_statement_profiling", False)) + if profiling_enabled: + self._ensure_statement_profiling(cursor) + + for idx, stmt in enumerate(statements): + stmt_started = perf_counter() if has_observers else 0.0 last_result = cursor.execute(stmt, normalize_execute_parameters(prepared_parameters)) + if profiling_enabled and last_result is not None: + with contextlib.suppress(Exception): + last_result.fetchall() + stmt_profile = self._capture_statement_profile(stmt) if profiling_enabled else None + if has_observers: + stmt_duration = perf_counter() - stmt_started + single_stmt = self._sub_statement(stmt, prepared_parameters) + try: + stmt_operation = single_stmt.operation_type + finally: + self._release_pooled_statement(single_stmt) + self._emit_script_statement_event( + sql=stmt, + parameters=prepared_parameters, + operation=stmt_operation, + rows_affected=None, + duration_s=stmt_duration, + started_at=stmt_started, + profile=stmt_profile, + script_ordinal=idx, + script_statement_count=statement_count, + ) successful_count += 1 + self._pending_statement_profile = None return self.create_execution_result( - last_result, statement_count=len(statements), successful_statements=successful_count, is_script_result=True + last_result, statement_count=statement_count, successful_statements=successful_count, is_script_result=True ) def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": @@ -564,10 +606,16 @@ def _execute_cache_hit_with_cursor( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": """Execute a cached query using DuckDB's unified row materializer.""" + profiling_enabled = bool(self.driver_features.get("enable_statement_profiling", False)) with self.with_cursor(self.connection) as cursor: - cursor.execute(cached.compiled_sql, normalize_execute_parameters(params)) + if profiling_enabled: + self._ensure_statement_profiling(cursor) + wire_sql = self._format_wire_sql(cached.compiled_sql) + cursor.execute(wire_sql, normalize_execute_parameters(params)) if cached.operation_profile.returns_rows: data, column_names, row_count, row_format = self._materialize_select_rows(cursor, cached.compiled_sql) + if profiling_enabled: + self._capture_statement_profile(wire_sql) return self.build_cached_select_result( sql, params, @@ -576,10 +624,30 @@ def _execute_cache_hit_with_cursor( column_names=column_names, data_row_count=row_count, row_format=row_format, + compiled_sql=wire_sql, ) affected_rows = resolve_rowcount(cursor) + if profiling_enabled: + with contextlib.suppress(Exception): + cursor.fetchall() + self._capture_statement_profile(wire_sql) return DMLResult(cached.operation_type, affected_rows) + def _ensure_statement_profiling(self, cursor: Any) -> None: + """Ensure per-connection DuckDB JSON statement profiling is configured.""" + if self.driver_features.get("_duckdb_profile_path"): + return + profile_path = create_duckdb_profile_path() + configure_duckdb_statement_profiling(cursor, profile_path) + self.driver_features["_duckdb_profile_path"] = str(profile_path) + + def _capture_statement_profile(self, executed_sql: str) -> "dict[str, Any] | None": + """Read the latest DuckDB JSON profile and stage it for StatementEvent emission.""" + profile_path = self.driver_features.get("_duckdb_profile_path") + profile = read_duckdb_profile_json(profile_path, executed_sql) + self._pending_statement_profile = profile + return profile + def _materialize_select_rows(self, cursor: Any, compiled_sql: str) -> "tuple[list[Any], list[str], int, RowFormat]": """Materialize SELECT rows using DuckDB's native Python type contract.""" description = cursor.description diff --git a/sqlspec/adapters/duckdb/pool.py b/sqlspec/adapters/duckdb/pool.py index 1882f1b33..edfc92869 100644 --- a/sqlspec/adapters/duckdb/pool.py +++ b/sqlspec/adapters/duckdb/pool.py @@ -9,11 +9,17 @@ from sqlspec.adapters.duckdb._typing import DuckDBConnection from sqlspec.adapters.duckdb._typing import duckdb_module as duckdb +from sqlspec.adapters.duckdb.core import ( + cleanup_duckdb_profile_path, + configure_duckdb_statement_profiling, + create_duckdb_profile_path, +) from sqlspec.utils.logging import POOL_LOGGER_NAME, get_logger, log_with_context from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: from collections.abc import Callable, Generator + from pathlib import Path __all__ = ("DuckDBConnectionPool",) @@ -72,6 +78,7 @@ class DuckDBConnectionPool: __slots__ = ( "_connection_config", "_connection_registry", + "_enable_statement_profiling", "_extensions", "_generation", "_health_check_interval", @@ -80,6 +87,7 @@ class DuckDBConnectionPool: "_lock", "_on_connection_create", "_pool_id", + "_profile_paths", "_recycle", "_secrets", "_thread_local", @@ -93,6 +101,7 @@ def __init__( extensions: "list[dict[str, Any]] | None" = None, secrets: "list[dict[str, Any]] | None" = None, on_connection_create: "Callable[[DuckDBConnection], DuckDBConnection | None] | None" = None, + enable_statement_profiling: bool = False, ) -> None: """Initialize the thread-local connection manager. @@ -103,6 +112,7 @@ def __init__( extensions: List of extensions to install/load secrets: List of secrets to create on_connection_create: Callback executed when connection is created + enable_statement_profiling: Whether to enable per-connection JSON statement profiling """ self._connection_config = connection_config self._recycle = pool_recycle_seconds @@ -110,6 +120,8 @@ def __init__( self._extensions = extensions or [] self._secrets = secrets or [] self._on_connection_create = on_connection_create + self._enable_statement_profiling = enable_statement_profiling + self._profile_paths: dict[int, Path] = {} self._installed_signatures: set[tuple[Any, ...]] = set() self._thread_local = threading.local() self._connection_registry: set[DuckDBConnection] = set() @@ -229,11 +241,18 @@ def new_connection(self, record_thread_state: bool = False) -> DuckDBConnection: modified = self._on_connection_create(connection) if modified is not None: connection = modified + profile_path: Path | None = None + if self._enable_statement_profiling: + profile_path = create_duckdb_profile_path() + configure_duckdb_statement_profiling(connection, profile_path) + with self._lock: + self._profile_paths[id(connection)] = profile_path storage_protocols.update(connection.list_filesystems()) if record_thread_state: self._thread_local.storage_extensions = frozenset(loaded_extensions) self._thread_local.storage_secrets = tuple(created_secrets) self._thread_local.storage_protocols = frozenset(storage_protocols) + self._thread_local.profile_path = profile_path return connection @contextmanager @@ -280,10 +299,14 @@ def close(self) -> None: with self._lock: orphaned = list(self._connection_registry) self._connection_registry.clear() + profile_paths = list(self._profile_paths.values()) + self._profile_paths.clear() self._generation += 1 for connection in orphaned: with suppress(Exception): connection.close() + for profile_path in profile_paths: + cleanup_duckdb_profile_path(profile_path) def size(self) -> int: """Get current pool size (always 1 for thread-local).""" @@ -322,12 +345,20 @@ def _storage_settings(self, connection: DuckDBConnection) -> "dict[str, Any]": """Return storage setup completed for this thread's current connection.""" state = self._thread_local.__dict__ if state.get("connection") is not connection: + with self._lock: + profile_path = self._profile_paths.get(id(connection)) + if profile_path is not None: + return {"_duckdb_profile_path": str(profile_path)} return {} - return { + settings: dict[str, Any] = { "_duckdb_storage_extensions": state.get("storage_extensions", frozenset()), "_duckdb_storage_secrets": state.get("storage_secrets", ()), "_duckdb_storage_protocols": state.get("storage_protocols", frozenset()), } + profile_path = state.get("profile_path") + if profile_path is not None: + settings["_duckdb_profile_path"] = str(profile_path) + return settings def _install_extension_once( self, @@ -416,8 +447,10 @@ def _retire_connection(self, connection: DuckDBConnection) -> None: """Close a pool-owned connection and drop it from the shutdown registry.""" with self._lock: self._connection_registry.discard(connection) + profile_path = self._profile_paths.pop(id(connection), None) with suppress(Exception): connection.close() + cleanup_duckdb_profile_path(profile_path) def _close_thread_connection(self) -> None: """Close the connection for the current thread.""" @@ -425,6 +458,7 @@ def _close_thread_connection(self) -> None: thread_state.pop("storage_extensions", None) thread_state.pop("storage_secrets", None) thread_state.pop("storage_protocols", None) + thread_state.pop("profile_path", None) if "connection" in thread_state: self._retire_connection(cast("DuckDBConnection", self._thread_local.connection)) del self._thread_local.connection diff --git a/sqlspec/adapters/oracledb/_typing.py b/sqlspec/adapters/oracledb/_typing.py index 353dbc5e2..80eb93ba5 100644 --- a/sqlspec/adapters/oracledb/_typing.py +++ b/sqlspec/adapters/oracledb/_typing.py @@ -90,17 +90,24 @@ class OracleSyncCursor: """Sync context manager for Oracle cursor management.""" - __slots__ = ("arraysize", "connection", "cursor", "prefetchrows") + __slots__ = ("arraysize", "connection", "cursor", "on_enter", "prefetchrows") def __init__( - self, connection: OracleSyncConnection, arraysize: "int | None" = None, prefetchrows: "int | None" = None + self, + connection: OracleSyncConnection, + arraysize: "int | None" = None, + prefetchrows: "int | None" = None, + on_enter: "Callable[[Any], None] | None" = None, ) -> None: self.connection = connection self.cursor: OracleSyncRawCursor | None = None self.arraysize = arraysize self.prefetchrows = prefetchrows + self.on_enter = on_enter def __enter__(self) -> "OracleSyncRawCursor": + if self.on_enter is not None: + self.on_enter(self.connection) self.cursor = self.connection.cursor() if self.arraysize is not None: self.cursor.arraysize = self.arraysize @@ -117,17 +124,24 @@ def __exit__(self, *_: object) -> None: class OracleAsyncCursor: """Async context manager for Oracle cursor management.""" - __slots__ = ("arraysize", "connection", "cursor", "prefetchrows") + __slots__ = ("arraysize", "connection", "cursor", "on_enter", "prefetchrows") def __init__( - self, connection: OracleAsyncConnection, arraysize: "int | None" = None, prefetchrows: "int | None" = None + self, + connection: OracleAsyncConnection, + arraysize: "int | None" = None, + prefetchrows: "int | None" = None, + on_enter: "Callable[[Any], None] | None" = None, ) -> None: self.connection = connection self.cursor: OracleAsyncRawCursor | None = None self.arraysize = arraysize self.prefetchrows = prefetchrows + self.on_enter = on_enter async def __aenter__(self) -> "OracleAsyncRawCursor": + if self.on_enter is not None: + self.on_enter(self.connection) self.cursor = self.connection.cursor() if self.arraysize is not None: self.cursor.arraysize = self.arraysize diff --git a/sqlspec/adapters/oracledb/config.py b/sqlspec/adapters/oracledb/config.py index c3b187235..e4297a387 100644 --- a/sqlspec/adapters/oracledb/config.py +++ b/sqlspec/adapters/oracledb/config.py @@ -261,6 +261,7 @@ class OracleDriverFeatures(TypedDict): events_backend: NotRequired[OracleEventsBackend] enable_direct_path_load: NotRequired[bool] session_context_hook: NotRequired[OracleSessionContextConfig | dict[str, Any]] + enable_session_tracing: NotRequired[bool] class OracleSyncConnectionContext(SyncPoolConnectionContext): diff --git a/sqlspec/adapters/oracledb/core.py b/sqlspec/adapters/oracledb/core.py index 7e29a888c..a1e6f0b2f 100644 --- a/sqlspec/adapters/oracledb/core.py +++ b/sqlspec/adapters/oracledb/core.py @@ -26,6 +26,7 @@ build_statement_config_from_profile, create_sql_result, ) +from sqlspec.core.sqlcommenter import SQLCommenterContext from sqlspec.data_dictionary.dialects.oracle import ORACLE_JSON_STORAGE_NATIVE, resolve_oracle_json_storage from sqlspec.driver import rows_to_dicts from sqlspec.exceptions import ( @@ -46,7 +47,9 @@ TransactionError, UniqueViolationError, ) +from sqlspec.observability import get_trace_context from sqlspec.typing import NUMPY_INSTALLED +from sqlspec.utils.correlation import CorrelationContext from sqlspec.utils.module_loader import ensure_pyarrow, import_optional from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.type_converters import build_uuid_coercions @@ -71,6 +74,10 @@ "DB_TYPE_CLOB", "ORACLEDB_SUPPORTS_SPARSE_VECTORS", "ORACLEDB_VERSION", + "ORACLE_ACTION_MAX_BYTES", + "ORACLE_CLIENT_IDENTIFIER_MAX_BYTES", + "ORACLE_ECID_MAX_BYTES", + "ORACLE_MODULE_MAX_BYTES", "SPARSE_VECTOR_MIN_DATABASE_MAJOR", "VALID_ORACLE_DRIVER_FEATURE_KEYS", "OracleAsyncStreamSource", @@ -82,6 +89,7 @@ "OracleSyncStreamSource", "OracleTopologyConfig", "apply_driver_features", + "apply_oracle_session_trace_attributes", "build_arrow_fetch_kwargs", "build_connection_config", "build_duality_etag_update_sql", @@ -126,6 +134,7 @@ "sanitize_oracle_text_query", "supports_df_batches", "supports_direct_path_load", + "truncate_utf8_bytes", ) @@ -319,6 +328,7 @@ "enable_lob_fetch", "enable_lowercase_column_names", "enable_numpy_vectors", + "enable_session_tracing", "enable_uuid_binary", "events_backend", "fetch_decimals", @@ -1155,6 +1165,85 @@ def resolve_rowcount(cursor: Any) -> int: return 0 +ORACLE_MODULE_MAX_BYTES: int = 48 +ORACLE_ACTION_MAX_BYTES: int = 64 +ORACLE_CLIENT_IDENTIFIER_MAX_BYTES: int = 64 +ORACLE_ECID_MAX_BYTES: int = 64 + + +def truncate_utf8_bytes(value: str, max_bytes: int) -> str: + """Truncate a string to at most max_bytes in UTF-8 without splitting multi-byte codepoints. + + Args: + value: Input string to truncate. + max_bytes: Maximum UTF-8 byte length allowed. + + Returns: + UTF-8 safe truncated string. + """ + encoded = value.encode("utf-8") + if len(encoded) <= max_bytes: + return value + return encoded[:max_bytes].decode("utf-8", errors="ignore") + + +def apply_oracle_session_trace_attributes( + connection: Any, + *, + module: "str | None" = None, + action: "str | None" = None, + client_identifier: "str | None" = None, + econtext_id: "str | None" = None, + ecid: "str | None" = None, +) -> None: + """Set Oracle end-to-end tracing attributes on connection for zero-RTT OALL8 piggybacking. + + Args: + connection: Oracle sync or async connection object. + module: Optional explicit V$SESSION.MODULE value (truncated to 48 UTF-8 bytes). + action: Optional explicit V$SESSION.ACTION value (truncated to 64 UTF-8 bytes). + client_identifier: Optional explicit V$SESSION.CLIENT_IDENTIFIER value (truncated to 64 UTF-8 bytes). + econtext_id: Optional explicit V$SESSION.ECID value (truncated to 64 UTF-8 bytes). + ecid: Optional alias for ``econtext_id``. + """ + ctx = SQLCommenterContext.get() or {} + resolved_module = ( + module if module is not None else (ctx.get("module") or ctx.get("framework") or ctx.get("controller")) + ) + resolved_action = action if action is not None else (ctx.get("action") or ctx.get("route")) + resolved_client_id = ( + client_identifier + if client_identifier is not None + else (CorrelationContext.get() or ctx.get("correlation_id") or ctx.get("client_identifier")) + ) + resolved_ecid = econtext_id if econtext_id is not None else ecid + if resolved_ecid is None: + trace_id, span_id = get_trace_context() + if trace_id and span_id: + resolved_ecid = f"00-{trace_id}-{span_id}-01" + elif trace_id: + resolved_ecid = trace_id + else: + resolved_ecid = ctx.get("traceparent") or ctx.get("econtext_id") or ctx.get("ecid") + + if resolved_module is not None: + with contextlib.suppress(Exception): + connection.module = truncate_utf8_bytes(resolved_module, ORACLE_MODULE_MAX_BYTES) + if resolved_action is not None: + with contextlib.suppress(Exception): + connection.action = truncate_utf8_bytes(resolved_action, ORACLE_ACTION_MAX_BYTES) + if resolved_client_id is not None: + with contextlib.suppress(Exception): + connection.client_identifier = truncate_utf8_bytes(resolved_client_id, ORACLE_CLIENT_IDENTIFIER_MAX_BYTES) + if resolved_ecid is not None: + truncated_ecid = truncate_utf8_bytes(resolved_ecid, ORACLE_ECID_MAX_BYTES) + with contextlib.suppress(Exception): + try: + connection.econtext_id = truncated_ecid + except AttributeError: + connection.ecid = truncated_ecid + + def apply_driver_features( statement_config: "StatementConfig", driver_features: "Mapping[str, Any] | None" ) -> "tuple[StatementConfig, dict[str, Any]]": @@ -1188,6 +1277,7 @@ def apply_driver_features( session_hook = features.get("session_context_hook") if isinstance(session_hook, Mapping): features["session_context_hook"] = OracleSessionContextConfig(**session_hook) + features.setdefault("enable_session_tracing", False) return statement_config, features diff --git a/sqlspec/adapters/oracledb/driver.py b/sqlspec/adapters/oracledb/driver.py index 592b27e4d..8cfe9f024 100644 --- a/sqlspec/adapters/oracledb/driver.py +++ b/sqlspec/adapters/oracledb/driver.py @@ -20,6 +20,7 @@ ORACLEDB_VERSION, OracleAsyncStreamSource, OracleSyncStreamSource, + apply_oracle_session_trace_attributes, build_arrow_fetch_kwargs, build_fetch_kwargs, build_insert_statement, @@ -64,6 +65,7 @@ ORACLE_EXPLAIN_XPLAN_FORMAT_META_KEY, ORACLE_MANAGED_EXPLAIN_META_KEY, ) +from sqlspec.core.sqlcommenter import SQLCommenterContext, format_wire_sqlcomment, strip_sqlcomment from sqlspec.driver import ( AsyncDriverAdapterBase, AsyncRowStream, @@ -151,6 +153,52 @@ class OraclePipelineMixin: __slots__ = () + def _session_tracing_enabled(self) -> bool: + driver = cast("OraclePipelineDriver", self) + return bool( + driver.driver_features.get("enable_session_tracing", False) or driver.statement_config.enable_sqlcommenter + ) + + def _apply_session_tracing(self, connection: Any = None) -> None: + if not self._session_tracing_enabled(): + return + driver = cast("OraclePipelineDriver", self) + target_conn = connection if connection is not None else getattr(self, "connection", None) + if target_conn is None: + return + ctx = SQLCommenterContext.get() or {} + attrs = driver.statement_config.sqlcommenter_attributes or {} + default_module = ( + ctx.get("module") + or ctx.get("framework") + or ctx.get("controller") + or attrs.get("module") + or attrs.get("framework") + or attrs.get("controller") + or "sqlspec" + ) + default_action = ctx.get("action") or ctx.get("route") or attrs.get("action") or attrs.get("route") + apply_oracle_session_trace_attributes(target_conn, module=default_module, action=default_action) + + def _format_wire_sql(self, sql: str, config: "StatementConfig | None" = None) -> str: + driver = cast("OraclePipelineDriver", self) + effective_config = config or driver.statement_config + if self._session_tracing_enabled(): + return strip_sqlcomment(sql) + if not effective_config.enable_sqlcommenter: + return sql + return format_wire_sqlcomment(sql, effective_config) + + def _compiled_sql( + self, statement: "SQL", statement_config: "StatementConfig", flatten_single_parameters: bool = False + ) -> "tuple[str, Any]": + compiled_sql, prepared_parameters = super()._compiled_sql( # type: ignore[misc] + statement, statement_config, flatten_single_parameters=flatten_single_parameters + ) + if self._session_tracing_enabled(): + return strip_sqlcomment(compiled_sql), prepared_parameters + return compiled_sql, prepared_parameters + def _stack_native_blocker(self, stack: "StatementStack") -> "str | None": for operation in stack.operations: if operation.method == "execute_arrow": @@ -583,6 +631,7 @@ def with_cursor(self, connection: OracleSyncConnection) -> OracleSyncCursor: connection, arraysize=self.driver_features.get("arraysize"), prefetchrows=self.driver_features.get("prefetchrows"), + on_enter=self._apply_session_tracing if self._session_tracing_enabled() else None, ) @overload @@ -704,6 +753,7 @@ def select_to_arrow( prepared_parameters = cast("list[Any] | tuple[Any, ...] | dict[Any, Any] | None", prepared_parameters) try: + self._apply_session_tracing(self.connection) if return_format in {"batches", "reader"} and supports_df_batches(self.connection): record_batches, batch_schema = self._fetch_arrow_record_batches( sql, prepared_parameters, batch_size, arrow_schema @@ -1312,6 +1362,7 @@ def with_cursor(self, connection: OracleAsyncConnection) -> OracleAsyncCursor: connection, arraysize=self.driver_features.get("arraysize"), prefetchrows=self.driver_features.get("prefetchrows"), + on_enter=self._apply_session_tracing if self._session_tracing_enabled() else None, ) @overload @@ -1433,6 +1484,7 @@ async def select_to_arrow( prepared_parameters = cast("list[Any] | tuple[Any, ...] | dict[Any, Any] | None", prepared_parameters) try: + self._apply_session_tracing(self.connection) if return_format in {"batches", "reader"} and supports_df_batches(self.connection): record_batches, batch_schema = await self._fetch_arrow_record_batches( sql, prepared_parameters, batch_size, arrow_schema @@ -1818,6 +1870,7 @@ async def _execute_stack_native( if owns_transaction and not continue_on_error: await self.begin() + self._apply_session_tracing(self.connection) pipeline_results = await self.connection.run_pipeline(pipeline, continue_on_error=True) results = self._build_stack_results_from_pipeline( compiled_operations, pipeline_results, continue_on_error, observer diff --git a/sqlspec/adapters/psqlpy/driver.py b/sqlspec/adapters/psqlpy/driver.py index f5b722c5c..5e10fef47 100644 --- a/sqlspec/adapters/psqlpy/driver.py +++ b/sqlspec/adapters/psqlpy/driver.py @@ -32,6 +32,7 @@ ) from sqlspec.adapters.psqlpy.data_dictionary import PsqlpyDataDictionary from sqlspec.core import SQL, StatementConfig, get_cache_config, register_driver_profile +from sqlspec.core.sqlcommenter import split_sqlcomment from sqlspec.driver import AsyncDriverAdapterBase, AsyncRowStream, BaseAsyncExceptionHandler from sqlspec.exceptions import SQLSpecError from sqlspec.utils.serializers import from_json @@ -140,12 +141,14 @@ async def dispatch_execute(self, cursor: "PsqlpyConnection", statement: SQL) -> ) if statement.operation_type in {"INSERT", "UPDATE", "DELETE"}: - count_sql = _dml_count_query(sql) + base_sql, comment_suffix = split_sqlcomment(sql) + count_sql = _dml_count_query(base_sql) if count_sql is not None: + wire_count_sql = f"{count_sql} {comment_suffix}" if comment_suffix else count_sql count_result = ( - await cursor.fetch(count_sql, params, prepared=False) + await cursor.fetch(wire_count_sql, params, prepared=False) if self._is_transaction_pooler_mode() - else await cursor.fetch(count_sql, params) + else await cursor.fetch(wire_count_sql, params) ) count_rows, _ = collect_rows(count_result) if len(count_rows) != 1 or set(count_rows[0]) != {_DML_COUNT_COLUMN}: diff --git a/sqlspec/builder/_explain.py b/sqlspec/builder/_explain.py index d46d9ab95..db131024d 100644 --- a/sqlspec/builder/_explain.py +++ b/sqlspec/builder/_explain.py @@ -165,6 +165,7 @@ def build_duckdb_explain(statement_sql: str, options: "ExplainOptions") -> str: - EXPLAIN statement - EXPLAIN ANALYZE statement - EXPLAIN (FORMAT JSON) statement + - EXPLAIN (ANALYZE, FORMAT JSON) statement Args: statement_sql: The SQL statement to explain @@ -173,6 +174,9 @@ def build_duckdb_explain(statement_sql: str, options: "ExplainOptions") -> str: Returns: Complete EXPLAIN SQL string """ + if options.analyze and options.format == ExplainFormat.JSON: + return f"EXPLAIN (ANALYZE, FORMAT JSON) {statement_sql}" + if options.analyze: return f"EXPLAIN ANALYZE {statement_sql}" diff --git a/sqlspec/core/sqlcommenter.py b/sqlspec/core/sqlcommenter.py index dc1d4565f..918eb2bda 100644 --- a/sqlspec/core/sqlcommenter.py +++ b/sqlspec/core/sqlcommenter.py @@ -20,9 +20,14 @@ "SQLCommenterAttributes", "SQLCommenterContext", "append_comment", + "append_sqlcomment", + "append_wire_sqlcomment", "create_sqlcommenter_statement_transformer", + "format_wire_sqlcomment", "generate_comment", "parse_comment", + "split_sqlcomment", + "strip_sqlcomment", ) @@ -155,14 +160,161 @@ def parse_comment(expression: exp.Expr) -> tuple[exp.Expr, dict[str, str]]: return expression, attrs +def split_sqlcomment(sql: str) -> tuple[str, str]: + """Split a leading or trailing sqlcommenter block comment from SQL text. + + Args: + sql: SQL text that may contain a sqlcommenter block comment. + + Returns: + Tuple of ``(base_sql, comment_block)`` where ``comment_block`` is the + full ``/* ... */`` comment or ``""`` when no sqlcommenter comment is present. + """ + if "/*" not in sql: + return sql, "" + + stripped = sql.rstrip() + trailing_semi = stripped.endswith(";") + candidate = stripped[:-1].rstrip() if trailing_semi else stripped + if candidate.endswith("*/"): + open_idx = candidate.rfind("/*") + if open_idx != -1: + inner = candidate[open_idx + 2 : -2].strip() + if _is_sqlcommenter_comment(inner): + comment_block = candidate[open_idx:].strip() + prefix = candidate[:open_idx].rstrip() + base_sql = f"{prefix};" if trailing_semi else prefix + return base_sql, comment_block + + lstripped = sql.lstrip() + if lstripped.startswith("/*"): + close_idx = lstripped.find("*/") + if close_idx != -1: + inner = lstripped[2:close_idx].strip() + if _is_sqlcommenter_comment(inner): + comment_block = lstripped[: close_idx + 2].strip() + rest = lstripped[close_idx + 2 :].lstrip() + return rest, comment_block + + return sql, "" + + +def strip_sqlcomment(sql: str) -> str: + """Remove a leading or trailing sqlcommenter block comment from SQL text. + + Args: + sql: SQL text that may contain a sqlcommenter block comment. + + Returns: + Clean SQL text without the sqlcommenter block comment. + """ + base_sql, _ = split_sqlcomment(sql) + return base_sql + + +def append_sqlcomment(sql: str, attrs: Mapping[str, str | None]) -> str: + """Append or replace a sqlcommenter block comment on rendered SQL text. + + Args: + sql: Rendered SQL text. + attrs: Attributes to serialize into the comment. + + Returns: + SQL text with the formatted sqlcommenter comment appended. + """ + base_sql = strip_sqlcomment(sql) + return _append_comment(base_sql, attrs) + + +def format_wire_sqlcomment( + sql: str, + config: Any = None, + *, + attributes: Mapping[str, str | None] | None = None, + enable_traceparent: bool = False, + enable_context: bool = False, + enable_correlation_id: bool = False, +) -> str: + """Format wire-level SQL with static and dynamic sqlcommenter attributes. + + Args: + sql: Compiled SQL string. + config: Optional ``StatementConfig`` carrying sqlcommenter settings. + attributes: Optional static attributes when ``config`` is omitted. + enable_traceparent: Include OpenTelemetry ``traceparent`` when ``config`` is omitted. + enable_context: Include ``SQLCommenterContext`` and ``CorrelationContext`` when ``config`` is omitted. + enable_correlation_id: Include ``CorrelationContext`` when ``config`` is omitted. + + Returns: + SQL string with wire-level sqlcommenter comment applied. + """ + if config is not None: + if not getattr(config, "enable_sqlcommenter", False): + return sql + static_attrs: Mapping[str, str | None] = getattr(config, "sqlcommenter_attributes", None) or {} + use_traceparent = bool(getattr(config, "sqlcommenter_enable_traceparent", False)) + use_context = bool(getattr(config, "sqlcommenter_enable_context", True)) + use_correlation_id = bool(getattr(config, "sqlcommenter_enable_correlation_id", False)) + else: + static_attrs = attributes or {} + use_traceparent = enable_traceparent + use_context = enable_context + use_correlation_id = enable_correlation_id + + merged = _comment_attributes( + static_attrs, + enable_traceparent=use_traceparent, + enable_context=use_context, + enable_correlation_id=use_correlation_id, + ) + base_sql = strip_sqlcomment(sql) + return _append_comment(base_sql, merged) + + +def append_wire_sqlcomment( + sql: str, + config: Any = None, + *, + attributes: Mapping[str, str | None] | None = None, + enable_traceparent: bool = False, + enable_context: bool = False, + enable_correlation_id: bool = False, +) -> str: + """Append wire-level sqlcommenter comment to compiled SQL text. + + Args: + sql: Compiled SQL string. + config: Optional ``StatementConfig`` carrying sqlcommenter settings. + attributes: Optional static attributes when ``config`` is omitted. + enable_traceparent: Include OpenTelemetry ``traceparent`` when ``config`` is omitted. + enable_context: Include ``SQLCommenterContext`` and ``CorrelationContext`` when ``config`` is omitted. + enable_correlation_id: Include ``CorrelationContext`` when ``config`` is omitted. + + Returns: + SQL string with wire-level sqlcommenter comment applied. + """ + return format_wire_sqlcomment( + sql, + config, + attributes=attributes, + enable_traceparent=enable_traceparent, + enable_context=enable_context, + enable_correlation_id=enable_correlation_id, + ) + + def create_sqlcommenter_statement_transformer( - *, attributes: dict[str, str | None] | None = None, enable_traceparent: bool = False, enable_context: bool = False + *, + attributes: dict[str, str | None] | None = None, + enable_traceparent: bool = False, + enable_context: bool = False, + enable_correlation_id: bool = False, ) -> Callable[[exp.Expr, Any], tuple[exp.Expr, Any]]: """Create a ``statement_transformer`` that adds sqlcommenter comments to the AST. Static attributes are pre-serialized at creation time. When - ``enable_traceparent`` or ``enable_context`` is True, dynamic attributes - are resolved per invocation. + ``enable_traceparent``, ``enable_context``, or ``enable_correlation_id`` is True, + dynamic attributes are resolved per invocation. Args: attributes: Static key-value pairs to include in every comment. @@ -170,12 +322,14 @@ def create_sqlcommenter_statement_transformer( current OpenTelemetry span context on each invocation. enable_context: If True, read request-scoped attributes from :class:`SQLCommenterContext` and merge them with static attributes. + enable_correlation_id: If True, read ``correlation_id`` from + :class:`CorrelationContext` when not already present. Returns: A callable suitable for ``StatementConfig(statement_transformers=[...])``. """ static_attrs: dict[str, str | None] = dict(attributes) if attributes else {} - is_dynamic = enable_traceparent or enable_context + is_dynamic = enable_traceparent or enable_context or enable_correlation_id if not is_dynamic and not static_attrs: return _NOOP_SQLCOMMENTER_TRANSFORMER @@ -184,7 +338,10 @@ def create_sqlcommenter_statement_transformer( return _StaticSQLCommenterTransformer(static_attrs) return _DynamicSQLCommenterTransformer( - static_attrs, enable_traceparent=enable_traceparent, enable_context=enable_context + static_attrs, + enable_traceparent=enable_traceparent, + enable_context=enable_context, + enable_correlation_id=enable_correlation_id, ) @@ -231,7 +388,11 @@ def _append_comment(sql: str, attrs: Mapping[str, str | None]) -> str: def _comment_attributes( - static_attrs: Mapping[str, str | None], *, enable_traceparent: bool, enable_context: bool + static_attrs: Mapping[str, str | None] | None = None, + *, + enable_traceparent: bool = False, + enable_context: bool = False, + enable_correlation_id: bool = False, ) -> dict[str, str | None]: """Resolve static and dynamic sqlcommenter attributes for the current call.""" merged: dict[str, str | None] = {} @@ -239,10 +400,12 @@ def _comment_attributes( ctx_attrs = SQLCommenterContext.get() if ctx_attrs: merged.update(ctx_attrs) + if enable_context or enable_correlation_id: correlation_id = CorrelationContext.get() if correlation_id and "correlation_id" not in merged: merged["correlation_id"] = correlation_id - merged.update(static_attrs) + if static_attrs: + merged.update(static_attrs) if enable_traceparent: trace_id, span_id = get_trace_context() if trace_id and span_id: @@ -278,18 +441,27 @@ def __call__(self, expression: exp.Expr, params: Any) -> tuple[exp.Expr, Any]: @final class _DynamicSQLCommenterTransformer: - __slots__ = ("_enable_context", "_enable_traceparent", "_static_attrs") + __slots__ = ("_enable_context", "_enable_correlation_id", "_enable_traceparent", "_static_attrs") def __init__( - self, static_attrs: Mapping[str, str | None], *, enable_traceparent: bool, enable_context: bool + self, + static_attrs: Mapping[str, str | None], + *, + enable_traceparent: bool, + enable_context: bool, + enable_correlation_id: bool = False, ) -> None: self._static_attrs = dict(static_attrs) self._enable_traceparent = enable_traceparent self._enable_context = enable_context + self._enable_correlation_id = enable_correlation_id def __call__(self, expression: exp.Expr, params: Any) -> tuple[exp.Expr, Any]: merged = _comment_attributes( - self._static_attrs, enable_traceparent=self._enable_traceparent, enable_context=self._enable_context + self._static_attrs, + enable_traceparent=self._enable_traceparent, + enable_context=self._enable_context, + enable_correlation_id=self._enable_correlation_id, ) return append_comment(expression, merged), params diff --git a/sqlspec/core/statement.py b/sqlspec/core/statement.py index 9fbc98501..a3c928212 100644 --- a/sqlspec/core/statement.py +++ b/sqlspec/core/statement.py @@ -53,7 +53,7 @@ safe_modify_with_cte, wrap_as_subquery, ) -from sqlspec.core.sqlcommenter import create_sqlcommenter_statement_transformer +from sqlspec.core.sqlcommenter import create_sqlcommenter_statement_transformer, format_wire_sqlcomment from sqlspec.observability import resolve_db_system from sqlspec.typing import Empty, EmptyEnum from sqlspec.utils.logging import get_logger @@ -1398,6 +1398,8 @@ def _rebind_cached_parameters(self, state: "ProcessedState") -> "tuple[str, Any] apply_wrap_types=state.applied_wrap_types, ) compiled_sql = state.compiled_sql + if self._statement_config.enable_sqlcommenter: + compiled_sql = format_wire_sqlcomment(compiled_sql, self._statement_config) output_transformer = self._statement_config.output_transformer if output_transformer: compiled_sql, rebound_params = output_transformer(compiled_sql, rebound_params) @@ -1815,7 +1817,7 @@ def __init__( self._fingerprint_cache: str | None = None self._hash_cache: int | None = None self._is_frozen = False - self._has_transformers = bool(self.statement_transformers) + self._has_transformers = bool(self._user_statement_transformers) self._has_output_transformer = self.output_transformer is not None def __hash__(self) -> int: diff --git a/sqlspec/driver/_async.py b/sqlspec/driver/_async.py index 30516b8b6..6c35a694d 100644 --- a/sqlspec/driver/_async.py +++ b/sqlspec/driver/_async.py @@ -269,7 +269,7 @@ async def dispatch_statement_execution(self, statement: "SQL", connection: "Any" return await self._execute_with_observability_async( compiled_sql=compiled_sql, execution_parameters=execution_parameters, - operation=statement.operation_type, + operation="SCRIPT" if statement.is_script else statement.operation_type, is_many=statement.is_many, is_script=statement.is_script, runner=self._dispatch_statement_observed_async, @@ -326,15 +326,35 @@ async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> Execut statement_count: int = len(statements) successful_count: int = 0 + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers - for stmt in statements: + for idx, stmt in enumerate(statements): single_stmt = self._sub_statement(stmt, prepared_parameters) + stmt_started = perf_counter() if has_observers else 0.0 + self._pending_statement_profile = None try: - await self.dispatch_execute(cursor, single_stmt) + sub_result = await self.dispatch_execute(cursor, single_stmt) finally: self._release_pooled_statement(single_stmt) + if has_observers: + stmt_duration = perf_counter() - stmt_started + stmt_profile: dict[str, Any] | None = self._pending_statement_profile + self._pending_statement_profile = None + self._emit_script_statement_event( + sql=stmt, + parameters=prepared_parameters, + operation=single_stmt.operation_type, + rows_affected=sub_result.rowcount_override, + duration_s=stmt_duration, + started_at=stmt_started, + profile=stmt_profile, + script_ordinal=idx, + script_statement_count=statement_count, + ) successful_count += 1 + self._pending_statement_profile = None return self.create_execution_result( cursor, statement_count=statement_count, successful_statements=successful_count, is_script_result=True ) @@ -1510,12 +1530,16 @@ async def _execute_with_observability_async( sql_hash = observability_runtime.compute_sql_hash(compiled_sql) span = runtime.start_query_span(compiled_sql, operation, driver_name, sql_hash=sql_hash) started = perf_counter() + self._pending_statement_profile = None try: result = await runner(**runner_kwargs) except Exception as exc: + self._pending_statement_profile = None runtime.span_manager.end_span(span, error=exc) await runtime.emit_error_async(exc, **query_context) raise + profile: dict[str, Any] | None = self._pending_statement_profile + self._pending_statement_profile = None runtime.span_manager.end_span(span) duration = perf_counter() - started await runtime.emit_query_complete_async(**{**query_context, "rows_affected": result.rows_affected}) @@ -1532,6 +1556,7 @@ async def _execute_with_observability_async( storage_backend=(result.metadata or {}).get("storage_backend"), started_at=started, precomputed_sql_hash=sql_hash, + profile=profile, ) return result @@ -1599,7 +1624,7 @@ async def _execute_cache_hit_observed( """Execute an async fast-path cache hit wrapped in observability instrumentation.""" obs_params = list(params) if isinstance(params, tuple) else params return await self._execute_with_observability_async( - compiled_sql=cached.compiled_sql, + compiled_sql=self._format_wire_sql(cached.compiled_sql), execution_parameters=obs_params, operation=cached.operation_type, is_many=False, @@ -1779,6 +1804,7 @@ async def _execute_cache_hit_with_cursor( result: SQLResult | None = None try: async with self.with_cursor(self.connection) as cursor: + wire_sql = self._format_wire_sql(cached.compiled_sql) returns_rows = cached.operation_profile.returns_rows execute = getattr(cursor, "execute", None) if cursor is not self.connection else None fetchall = getattr(cursor, "fetchall", None) if execute is not None and returns_rows else None @@ -1787,7 +1813,7 @@ async def _execute_cache_hit_with_cursor( ) if can_use_cursor_fast_path: assert execute is not None - execute_result = execute(cached.compiled_sql, params) if params else execute(cached.compiled_sql) + execute_result = execute(wire_sql, params) if params else execute(wire_sql) if isawaitable(execute_result): await execute_result if returns_rows: @@ -1804,13 +1830,14 @@ async def _execute_cache_hit_with_cursor( column_names=column_names, data_row_count=row_count, row_format=resolve_row_format(data), + compiled_sql=wire_sql, ) else: affected_rows = self.resolve_rowcount(cursor) result = DMLResult(cached.operation_type, affected_rows) else: direct_statement = self._cached_statement( - sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql + sql, params, cached, params, params_are_simple=True, compiled_sql=wire_sql ) execution_result = await self.dispatch_execute(cursor, direct_statement) if returns_rows: diff --git a/sqlspec/driver/_common.py b/sqlspec/driver/_common.py index 1303a43d1..5f39046e2 100644 --- a/sqlspec/driver/_common.py +++ b/sqlspec/driver/_common.py @@ -39,6 +39,7 @@ value_fingerprint, ) from sqlspec.core.result._base import RowFormat +from sqlspec.core.sqlcommenter import format_wire_sqlcomment, strip_sqlcomment from sqlspec.core.statement import ProcessedState from sqlspec.data_dictionary import ( ForeignKeyMetadata, @@ -545,6 +546,7 @@ class CommonDriverAttributesMixin: __slots__ = ( "_cache_holder", "_observability", + "_pending_statement_profile", "_processed_state_pool", "_statement_cache", "_statement_pool", @@ -583,6 +585,7 @@ def __init__( self.statement_config = statement_config self.driver_features = driver_features or {} self._observability = observability + self._pending_statement_profile: dict[str, Any] | None = None self._transaction_depth = 0 self._stmt_cache_max_size = self._statement_cache_size() shared_holder = self.driver_features.get("_statement_cache_holder") @@ -970,13 +973,14 @@ def build_cached_select_result( data_row_count: int, row_format: "RowFormat" = "dict", last_inserted_id: int | str | None = None, + compiled_sql: str | None = None, ) -> "SQLResult": """Build a fast-path SELECT result with deferred statement materialization.""" return SQLResult.create_cached_select( sql=sql, config=self.statement_config, cached_state=cached.processed_state, - compiled_sql=cached.compiled_sql, + compiled_sql=cached.compiled_sql if compiled_sql is None else compiled_sql, params=params, data=selected_data, column_names=column_names, @@ -985,6 +989,12 @@ def build_cached_select_result( last_inserted_id=last_inserted_id, ) + def _format_wire_sql(self, sql: str) -> str: + """Apply wire-level sqlcommenter formatting when enabled on statement_config.""" + if not self.statement_config.enable_sqlcommenter: + return sql + return format_wire_sqlcomment(sql, self.statement_config) + def _statement_cache_size(self) -> int: """Return the configured statement-cache capacity.""" return resolve_statement_cache_size(self.driver_features) @@ -1067,6 +1077,41 @@ def _sub_statement(self, sql: str, execution_parameters: "ConvertedParameters") ProcessedState.__init__(direct_state, compiled_sql=sql, execution_parameters=execution_parameters) return SQL._create_cached_direct(sql, self.statement_config, direct_state) + def _emit_script_statement_event( + self, + *, + sql: str, + parameters: Any, + operation: "OperationType", + rows_affected: int | None, + duration_s: float, + started_at: float | None, + profile: "dict[str, Any] | None", + script_ordinal: int, + script_statement_count: int, + ) -> None: + """Emit a per-statement StatementEvent during multi-statement script execution.""" + runtime = self._observability + if runtime is None or not runtime.has_statement_observers: + return + runtime.emit_statement_event( + sql=sql, + parameters=parameters, + driver=type(self).__name__, + operation=operation, + execution_mode=self.statement_config.execution_mode, + is_many=False, + is_script=True, + rows_affected=rows_affected, + duration_s=duration_s, + storage_backend=None, + started_at=started_at, + profile=profile, + script_ordinal=script_ordinal, + script_statement_index=script_ordinal, + script_statement_count=script_statement_count, + ) + def _cached_execution( self, statement: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]" ) -> "SQLResult | Awaitable[SQLResult | None] | None": @@ -1102,10 +1147,26 @@ def _cached_execution( needs_rebind = parameter_values_need_processing(param_values, param_config.type_coercion_map) if not needs_rebind and not config._has_output_transformer: + effective_cached = cached + if config.enable_sqlcommenter: + wire_sql = self._format_wire_sql(cached.compiled_sql) + if wire_sql != cached.compiled_sql: + effective_cached = CachedQuery( + compiled_sql=wire_sql, + parameter_profile=cached.parameter_profile, + input_named_parameters=cached.input_named_parameters, + applied_wrap_types=cached.applied_wrap_types, + parameter_casts=cached.parameter_casts, + operation_type=cached.operation_type, + operation_profile=cached.operation_profile, + param_count=cached.param_count, + processed_state=cached.processed_state, + needs_style_remap=cached.needs_style_remap, + ) runtime = self._observability if runtime is not None and not runtime.is_idle: - return self._execute_cache_hit_observed(statement, params, cached) - return self._execute_cache_hit(statement, params, cached) + return self._execute_cache_hit_observed(statement, params, effective_cached) + return self._execute_cache_hit(statement, params, effective_cached) if needs_rebind: rebound_params = self.stmt_cache_rebind(params, cached) @@ -1114,7 +1175,7 @@ def _cached_execution( rebound_params = params params_are_simple = True - compiled_sql = cached.compiled_sql + compiled_sql = self._format_wire_sql(cached.compiled_sql) output_transformer = config.output_transformer if output_transformer is not None: compiled_sql, rebound_params = output_transformer(compiled_sql, rebound_params) @@ -1148,7 +1209,7 @@ def _execute_cache_hit( implementations that bypass SQL object construction entirely. """ prepared = self._cached_statement( - sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql + sql, params, cached, params, params_are_simple=True, compiled_sql=self._format_wire_sql(cached.compiled_sql) ) return self._execute_cached_statement(prepared) @@ -1202,9 +1263,14 @@ def _cache_statement(self, statement: "SQL") -> None: operation_profile = processed.operation_profile if not operation_profile.returns_rows and statement.returns_rows(): operation_profile = OperationProfile(returns_rows=True, modifies_rows=operation_profile.modifies_rows) - cached_state = _clone_processed_state(processed, operation_profile) + clean_sql = ( + strip_sqlcomment(processed.compiled_sql) + if self.statement_config.enable_sqlcommenter + else processed.compiled_sql + ) + cached_state = _clone_processed_state(processed, operation_profile, compiled_sql=clean_sql) cached = CachedQuery( - compiled_sql=processed.compiled_sql, + compiled_sql=clean_sql, parameter_profile=param_profile, input_named_parameters=processed.input_named_parameters, applied_wrap_types=processed.applied_wrap_types, @@ -1611,8 +1677,13 @@ def _compiled_statement( prepared_parameters = self.prepare_driver_parameters( statement.parameters, statement_config, is_many=statement.is_many, prepared_statement=statement ) + wire_sql = ( + format_wire_sqlcomment(cached_result.compiled_sql, statement_config) + if statement_config.enable_sqlcommenter + else cached_result.compiled_sql + ) updated_cached = CachedStatement( - compiled_sql=cached_result.compiled_sql, + compiled_sql=wire_sql, parameters=cast("tuple[Any, ...] | list[Any] | dict[str, Any] | None", prepared_parameters), expression=cached_result.expression, input_named_parameters=cached_result.input_named_parameters, @@ -1637,6 +1708,7 @@ def _compiled_statement( ) cached_parameters = tuple(prepared_parameters) if isinstance(prepared_parameters, list) else prepared_parameters + stored_sql = strip_sqlcomment(compiled_sql) if statement_config.enable_sqlcommenter else compiled_sql cached_statement = CachedStatement( compiled_sql=compiled_sql, parameters=cast("tuple[Any, ...] | list[Any] | dict[str, Any] | None", cached_parameters), @@ -1647,7 +1719,19 @@ def _compiled_statement( ) if cache_key is not None and cache is not None: - cache.put_statement(cache_key, cached_statement, dialect_key) + stored_statement = ( + CachedStatement( + compiled_sql=stored_sql, + parameters=cached_statement.parameters, + expression=cached_statement.expression, + input_named_parameters=cached_statement.input_named_parameters, + parameter_profile=cached_statement.parameter_profile, + applied_wrap_types=cached_statement.applied_wrap_types, + ) + if stored_sql != compiled_sql + else cached_statement + ) + cache.put_statement(cache_key, stored_statement, dialect_key) self._cache_statement(statement) return cached_statement, prepared_parameters @@ -2048,10 +2132,10 @@ def _cache_param_values(params: "tuple[Any, ...] | list[Any] | dict[str, Any]") def _clone_processed_state( - processed: "ProcessedState", operation_profile: "OperationProfile | None" = None + processed: "ProcessedState", operation_profile: "OperationProfile | None" = None, compiled_sql: str | None = None ) -> "ProcessedState": return ProcessedState( - compiled_sql=processed.compiled_sql, + compiled_sql=processed.compiled_sql if compiled_sql is None else compiled_sql, execution_parameters=processed.execution_parameters, parsed_expression=processed.parsed_expression, operation_type=processed.operation_type, diff --git a/sqlspec/driver/_sync.py b/sqlspec/driver/_sync.py index e8b9411c5..8a71eea4d 100644 --- a/sqlspec/driver/_sync.py +++ b/sqlspec/driver/_sync.py @@ -233,7 +233,7 @@ def dispatch_statement_execution(self, statement: "SQL", connection: "Any") -> " return self._execute_with_observability_sync( compiled_sql=compiled_sql, execution_parameters=execution_parameters, - operation=statement.operation_type, + operation="SCRIPT" if statement.is_script else statement.operation_type, is_many=statement.is_many, is_script=statement.is_script, runner=self._dispatch_statement_observed_sync, @@ -290,15 +290,35 @@ def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> ExecutionRes statement_count: int = len(statements) successful_count: int = 0 + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers - for stmt in statements: + for idx, stmt in enumerate(statements): single_stmt = self._sub_statement(stmt, prepared_parameters) + stmt_started = perf_counter() if has_observers else 0.0 + self._pending_statement_profile = None try: - self.dispatch_execute(cursor, single_stmt) + sub_result = self.dispatch_execute(cursor, single_stmt) finally: self._release_pooled_statement(single_stmt) + if has_observers: + stmt_duration = perf_counter() - stmt_started + stmt_profile: dict[str, Any] | None = self._pending_statement_profile + self._pending_statement_profile = None + self._emit_script_statement_event( + sql=stmt, + parameters=prepared_parameters, + operation=single_stmt.operation_type, + rows_affected=sub_result.rowcount_override, + duration_s=stmt_duration, + started_at=stmt_started, + profile=stmt_profile, + script_ordinal=idx, + script_statement_count=statement_count, + ) successful_count += 1 + self._pending_statement_profile = None return self.create_execution_result( cursor, statement_count=statement_count, successful_statements=successful_count, is_script_result=True ) @@ -1491,12 +1511,16 @@ def _execute_with_observability_sync( sql_hash = observability_runtime.compute_sql_hash(compiled_sql) span = runtime.start_query_span(compiled_sql, operation, driver_name, sql_hash=sql_hash) started = perf_counter() + self._pending_statement_profile = None try: result = runner(**runner_kwargs) except Exception as exc: + self._pending_statement_profile = None runtime.span_manager.end_span(span, error=exc) runtime.emit_error_sync(exc, **query_context) raise + profile: dict[str, Any] | None = self._pending_statement_profile + self._pending_statement_profile = None runtime.span_manager.end_span(span) duration = perf_counter() - started runtime.emit_query_complete_sync(**{**query_context, "rows_affected": result.rows_affected}) @@ -1513,6 +1537,7 @@ def _execute_with_observability_sync( storage_backend=(result.metadata or {}).get("storage_backend"), started_at=started, precomputed_sql_hash=sql_hash, + profile=profile, ) return result @@ -1579,7 +1604,7 @@ def _execute_cache_hit_observed( """Execute a fast-path cache hit wrapped in observability instrumentation.""" obs_params = list(params) if isinstance(params, tuple) else params return self._execute_with_observability_sync( - compiled_sql=cached.compiled_sql, + compiled_sql=self._format_wire_sql(cached.compiled_sql), execution_parameters=obs_params, operation=cached.operation_type, is_many=False, @@ -1676,6 +1701,7 @@ def _execute_cache_hit_with_cursor( result: SQLResult | None = None try: with self.with_cursor(self.connection) as cursor: + wire_sql = self._format_wire_sql(cached.compiled_sql) returns_rows = cached.operation_profile.returns_rows execute = getattr(cursor, "execute", None) if cursor is not self.connection else None fetchall = getattr(cursor, "fetchall", None) if execute is not None and returns_rows else None @@ -1685,9 +1711,9 @@ def _execute_cache_hit_with_cursor( if can_use_cursor_fast_path: assert execute is not None if params: - execute(cached.compiled_sql, params) + execute(wire_sql, params) else: - execute(cached.compiled_sql) + execute(wire_sql) if returns_rows: assert fetchall is not None fetched_data = fetchall() @@ -1700,13 +1726,14 @@ def _execute_cache_hit_with_cursor( column_names=column_names, data_row_count=row_count, row_format=resolve_row_format(data), + compiled_sql=wire_sql, ) else: affected_rows = self.resolve_rowcount(cursor) result = DMLResult(cached.operation_type, affected_rows) else: direct_statement = self._cached_statement( - sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql + sql, params, cached, params, params_are_simple=True, compiled_sql=wire_sql ) execution_result = self.dispatch_execute(cursor, direct_statement) if returns_rows: diff --git a/sqlspec/observability/_observer.py b/sqlspec/observability/_observer.py index 931277a96..9b354cca6 100644 --- a/sqlspec/observability/_observer.py +++ b/sqlspec/observability/_observer.py @@ -44,8 +44,12 @@ class StatementEvent: "operation", "parameters", "prepared_statement", + "profile", "rows_affected", "sampled", + "script_ordinal", + "script_statement_count", + "script_statement_index", "span_id", "sql", "sql_hash", @@ -83,7 +87,12 @@ def __init__( trace_id: "str | None", span_id: "str | None", sampled: bool = True, + profile: "dict[str, Any] | None" = None, + script_ordinal: "int | None" = None, + script_statement_index: "int | None" = None, + script_statement_count: "int | None" = None, ) -> None: + resolved_ordinal = script_ordinal if script_ordinal is not None else script_statement_index self.sql = sql self.parameters = parameters self.driver = driver @@ -107,6 +116,10 @@ def __init__( self.trace_id = trace_id self.span_id = span_id self.sampled = sampled + self.profile = profile + self.script_ordinal = resolved_ordinal + self.script_statement_index = resolved_ordinal + self.script_statement_count = script_statement_count def __hash__(self) -> int: # pragma: no cover msg = "StatementEvent objects are mutable and unhashable" @@ -117,7 +130,8 @@ def __repr__(self) -> str: f"StatementEvent(sql={self.sql!r}, parameters={self.parameters!r}, driver={self.driver!r}, adapter={self.adapter!r}, bind_key={self.bind_key!r}, db_system={self.db_system!r}, " f"operation={self.operation!r}, execution_mode={self.execution_mode!r}, is_many={self.is_many!r}, is_script={self.is_script!r}, rows_affected={self.rows_affected!r}, duration_s={self.duration_s!r}, " f"started_at={self.started_at!r}, correlation_id={self.correlation_id!r}, storage_backend={self.storage_backend!r}, sql_hash={self.sql_hash!r}, sql_truncated={self.sql_truncated!r}, " - f"sql_original_length={self.sql_original_length!r}, transaction_state={self.transaction_state!r}, prepared_statement={self.prepared_statement!r}, trace_id={self.trace_id!r}, span_id={self.span_id!r}, sampled={self.sampled!r})" + f"sql_original_length={self.sql_original_length!r}, transaction_state={self.transaction_state!r}, prepared_statement={self.prepared_statement!r}, trace_id={self.trace_id!r}, span_id={self.span_id!r}, sampled={self.sampled!r}, " + f"profile={self.profile!r}, script_ordinal={self.script_ordinal!r}, script_statement_index={self.script_statement_index!r}, script_statement_count={self.script_statement_count!r})" ) def __eq__(self, other: object) -> bool: @@ -147,6 +161,10 @@ def __eq__(self, other: object) -> bool: and self.trace_id == other.trace_id and self.span_id == other.span_id and self.sampled == other.sampled + and self.profile == other.profile + and self.script_ordinal == other.script_ordinal + and self.script_statement_index == other.script_statement_index + and self.script_statement_count == other.script_statement_count ) def as_dict(self) -> "dict[str, Any]": @@ -176,6 +194,10 @@ def as_dict(self) -> "dict[str, Any]": "trace_id": self.trace_id, "span_id": self.span_id, "sampled": self.sampled, + "profile": self.profile, + "script_ordinal": self.script_ordinal, + "script_statement_index": self.script_statement_index, + "script_statement_count": self.script_statement_count, } @@ -238,6 +260,10 @@ def create_event( trace_id: "str | None" = None, span_id: "str | None" = None, sampled: bool = True, + profile: "dict[str, Any] | None" = None, + script_ordinal: "int | None" = None, + script_statement_index: "int | None" = None, + script_statement_count: "int | None" = None, ) -> StatementEvent: """Factory helper for runtime statement events.""" @@ -265,6 +291,10 @@ def create_event( trace_id=trace_id, span_id=span_id, sampled=sampled, + profile=profile, + script_ordinal=script_ordinal, + script_statement_index=script_statement_index, + script_statement_count=script_statement_count, ) @@ -304,6 +334,16 @@ def _emit_otel_statement_log( extra["correlation_id"] = event.correlation_id if event.sql_hash: extra["db.statement.hash"] = event.sql_hash + if event.profile is not None: + extra["db.statement.profile"] = event.profile + if event.script_ordinal is not None: + extra["sqlspec.script.ordinal"] = event.script_ordinal + extra["sqlspec.script.statement_index"] = event.script_ordinal + elif event.script_statement_index is not None: + extra["sqlspec.script.ordinal"] = event.script_statement_index + extra["sqlspec.script.statement_index"] = event.script_statement_index + if event.script_statement_count is not None: + extra["sqlspec.script.statement_count"] = event.script_statement_count if logging_config.include_driver_name: extra["sqlspec.driver"] = event.driver diff --git a/sqlspec/observability/_runtime.py b/sqlspec/observability/_runtime.py index 5b34da347..2941c2a22 100644 --- a/sqlspec/observability/_runtime.py +++ b/sqlspec/observability/_runtime.py @@ -8,7 +8,7 @@ from sqlspec.observability._common import compute_sql_hash, get_trace_context, redact_sql_literals, resolve_db_system from sqlspec.observability._config import LoggingConfig, ObservabilityConfig, StatementObserver from sqlspec.observability._dispatcher import LifecycleDispatcher, LifecycleHook -from sqlspec.observability._observer import create_event, create_statement_observer +from sqlspec.observability._observer import StatementEvent, create_event, create_statement_observer from sqlspec.observability._spans import SpanManager from sqlspec.utils.correlation import CorrelationContext from sqlspec.utils.logging import get_logger @@ -348,37 +348,31 @@ async def emit_error_async(self, exception: Exception, **extras: Any) -> None: def emit_error(self, exception: Exception, **extras: Any) -> None: self.emit_error_sync(exception, **extras) - def emit_statement_event( + def create_statement_event( self, *, sql: str, parameters: Any, driver: str, operation: str, - execution_mode: str | None, - is_many: bool, - is_script: bool, - rows_affected: int | None, - duration_s: float, - storage_backend: str | None, + execution_mode: str | None = None, + is_many: bool = False, + is_script: bool = False, + rows_affected: int | None = None, + duration_s: float = 0.0, + storage_backend: str | None = None, started_at: float | None = None, precomputed_sql_hash: str | None = None, transaction_state: str | None = None, prepared_statement: bool | None = None, - is_error: bool = False, - ) -> None: - """Emit a statement event to all registered observers.""" - - if not self._statement_observers: - return + sampled: bool = True, + profile: dict[str, Any] | None = None, + script_ordinal: int | None = None, + script_statement_index: int | None = None, + script_statement_count: int | None = None, + ) -> StatementEvent: + """Create a redacted and enriched StatementEvent instance.""" correlation_id = CorrelationContext.get() - sampled = True - if self.config.sampling is not None: - sampled = self.config.sampling.should_sample( - correlation_id=correlation_id, duration_ms=duration_s * 1000, is_error=is_error - ) - if not sampled: - return sanitized_sql = self._redact_sql(sql) sanitized_params = self._redact_parameters(parameters) logging_config = self.config.logging @@ -401,7 +395,7 @@ def emit_statement_event( resolved_started_at = ( started_at if (started_at is not None and started_at >= _MIN_WALL_CLOCK_EPOCH_S) else (time() - duration_s) ) - event = create_event( + return create_event( sql=sanitized_sql, parameters=sanitized_params, driver=driver, @@ -425,6 +419,67 @@ def emit_statement_event( trace_id=trace_id, span_id=span_id, sampled=sampled, + profile=profile, + script_ordinal=script_ordinal, + script_statement_index=script_statement_index, + script_statement_count=script_statement_count, + ) + + def emit_statement_event( + self, + *, + sql: str, + parameters: Any, + driver: str, + operation: str, + execution_mode: str | None, + is_many: bool, + is_script: bool, + rows_affected: int | None, + duration_s: float, + storage_backend: str | None, + started_at: float | None = None, + precomputed_sql_hash: str | None = None, + transaction_state: str | None = None, + prepared_statement: bool | None = None, + is_error: bool = False, + profile: dict[str, Any] | None = None, + script_ordinal: int | None = None, + script_statement_index: int | None = None, + script_statement_count: int | None = None, + ) -> None: + """Emit a statement event to all registered observers.""" + + if not self._statement_observers: + return + correlation_id = CorrelationContext.get() + sampled = True + if self.config.sampling is not None: + sampled = self.config.sampling.should_sample( + correlation_id=correlation_id, duration_ms=duration_s * 1000, is_error=is_error + ) + if not sampled: + return + event = self.create_statement_event( + sql=sql, + parameters=parameters, + driver=driver, + operation=operation, + execution_mode=execution_mode, + is_many=is_many, + is_script=is_script, + rows_affected=rows_affected, + duration_s=duration_s, + storage_backend=storage_backend, + started_at=started_at, + precomputed_sql_hash=precomputed_sql_hash, + transaction_state=transaction_state, + prepared_statement=prepared_statement, + sampled=sampled, + profile=profile, + script_ordinal=script_ordinal, + script_statement_index=script_statement_index, + script_statement_count=script_statement_count, ) for observer in self._statement_observers: try: diff --git a/tests/integration/adapters/_shared/_driver_type_system.py b/tests/integration/adapters/_shared/_driver_type_system.py index 13ff1ca94..042c34544 100644 --- a/tests/integration/adapters/_shared/_driver_type_system.py +++ b/tests/integration/adapters/_shared/_driver_type_system.py @@ -107,6 +107,7 @@ class SourceEquivalenceCase: "enable_events", "on_connection_create", "events_backend", + "enable_statement_profiling", ), "aiomysql": ( "json_serializer", @@ -257,6 +258,7 @@ class SourceEquivalenceCase: "extension_flags", "enable_events", "events_backend", + "enable_statement_profiling", ), "mssql_python": ( "use_pool", @@ -291,6 +293,7 @@ class SourceEquivalenceCase: "events_backend", "enable_direct_path_load", "session_context_hook", + "enable_session_tracing", ), "psqlpy": ( "enable_cast_detection", diff --git a/tests/unit/adapters/test_adbc/test_profiling.py b/tests/unit/adapters/test_adbc/test_profiling.py new file mode 100644 index 000000000..3e6724719 --- /dev/null +++ b/tests/unit/adapters/test_adbc/test_profiling.py @@ -0,0 +1,57 @@ +"""Unit tests for ADBC DuckDB JSON statement profiling and per-statement script events.""" + +from pathlib import Path + +from sqlspec.adapters.adbc import AdbcConfig +from sqlspec.observability import ObservabilityConfig, StatementEvent + + +def test_adbc_duckdb_enable_statement_profiling_and_execute_script_events() -> None: + """AdbcConfig with driver_name='duckdb' and enable_statement_profiling=True populates StatementEvent.profile and emits per-statement script events.""" + events: list[StatementEvent] = [] + config = AdbcConfig( + connection_config={"driver_name": "duckdb"}, + driver_features={"enable_statement_profiling": True}, + observability_config=ObservabilityConfig(statement_observers=(events.append,)), + ) + profile_path: Path | None = None + script = "CREATE TABLE t (x INT); INSERT INTO t VALUES (1), (2); SELECT sum(x) FROM t;" + + with config.provide_session() as session: + result = session.execute("SELECT sum(i) AS total FROM range(10) t(i)") + assert result.one()["total"] == 45 + + raw_path = session.driver_features.get("_duckdb_profile_path") + assert isinstance(raw_path, str) + profile_path = Path(raw_path) + assert profile_path.exists() + + script_result = session.execute_script(script) + assert script_result.total_statements == 3 + assert script_result.successful_statements == 3 + + assert profile_path is not None + assert not profile_path.exists() + + assert len(events) == 5 + first_event = events[0] + assert "enable_profiling" not in first_event.sql.lower() + assert "profiling_output" not in first_event.sql.lower() + assert isinstance(first_event.profile, dict) + assert "cpu_time" in first_event.profile + assert "children" in first_event.profile + + sub_events = events[1:4] + outer_event = events[4] + assert [e.script_ordinal for e in sub_events] == [0, 1, 2] + assert [e.script_statement_index for e in sub_events] == [0, 1, 2] + assert [e.script_statement_count for e in sub_events] == [3, 3, 3] + assert all(e.is_script for e in sub_events) + for sub_event in sub_events: + assert isinstance(sub_event.profile, dict) + assert "cpu_time" in sub_event.profile + assert "children" in sub_event.profile + + assert outer_event.operation == "SCRIPT" + assert outer_event.is_script is True + assert outer_event.script_ordinal is None diff --git a/tests/unit/adapters/test_duckdb/test_profiling.py b/tests/unit/adapters/test_duckdb/test_profiling.py new file mode 100644 index 000000000..d0b882fd6 --- /dev/null +++ b/tests/unit/adapters/test_duckdb/test_profiling.py @@ -0,0 +1,88 @@ +"""Unit tests for DuckDB JSON statement profiling and per-statement script events.""" + +from pathlib import Path + +from sqlspec.adapters.duckdb import DuckDBConfig +from sqlspec.observability import ObservabilityConfig, StatementEvent + + +def test_duckdb_enable_statement_profiling_attaches_json_profile_to_statement_event(tmp_path: Path) -> None: + """enable_statement_profiling=True attaches parsed JSON operator trees to StatementEvent.profile.""" + events: list[StatementEvent] = [] + config = DuckDBConfig( + connection_config={"database": str(tmp_path / "profiling.duckdb")}, + driver_features={"enable_statement_profiling": True}, + observability_config=ObservabilityConfig(statement_observers=(events.append,)), + ) + try: + with config.provide_session() as session: + result = session.execute("SELECT sum(i) AS total FROM range(10) t(i)") + assert result.one()["total"] == 45 + session.execute("SELECT ?::INTEGER AS val", (1,)) + session.execute("SELECT ?::INTEGER AS val", (2,)) + finally: + config.close_pool() + + assert len(events) == 3 + for event in events: + assert "enable_profiling" not in event.sql.lower() + assert "profiling_output" not in event.sql.lower() + assert isinstance(event.profile, dict) + assert "cpu_time" in event.profile + assert "children" in event.profile + + +def test_duckdb_execute_script_emits_per_statement_events_with_ordinals_and_profiles(tmp_path: Path) -> None: + """execute_script emits per-statement StatementEvents with ordinals, counts, and profiles before the outer event.""" + events: list[StatementEvent] = [] + config = DuckDBConfig( + connection_config={"database": str(tmp_path / "script_profiling.duckdb")}, + driver_features={"enable_statement_profiling": True}, + observability_config=ObservabilityConfig(statement_observers=(events.append,)), + ) + script = "CREATE TABLE t (x INT); INSERT INTO t VALUES (1), (2); SELECT sum(x) FROM t;" + try: + with config.provide_session() as session: + result = session.execute_script(script) + assert result.total_statements == 3 + assert result.successful_statements == 3 + finally: + config.close_pool() + + assert len(events) == 4 + sub_events = events[:3] + outer_event = events[3] + + assert [e.script_ordinal for e in sub_events] == [0, 1, 2] + assert [e.script_statement_index for e in sub_events] == [0, 1, 2] + assert [e.script_statement_count for e in sub_events] == [3, 3, 3] + assert all(e.is_script for e in sub_events) + for sub_event in sub_events: + assert isinstance(sub_event.profile, dict) + assert "cpu_time" in sub_event.profile + assert "children" in sub_event.profile + + assert outer_event.operation == "SCRIPT" + assert outer_event.is_script is True + assert outer_event.script_ordinal is None + + +def test_duckdb_profiling_scratch_file_cleaned_up_on_close_pool(tmp_path: Path) -> None: + """Temporary JSON profile files created by DuckDBConnectionPool are unlinked on close_pool.""" + config = DuckDBConfig( + connection_config={"database": str(tmp_path / "cleanup.duckdb")}, + driver_features={"enable_statement_profiling": True}, + ) + profile_path: Path | None = None + try: + with config.provide_session() as session: + session.execute("SELECT 1") + raw_path = session.driver_features.get("_duckdb_profile_path") + assert raw_path is not None + profile_path = Path(raw_path) + assert profile_path.exists() + finally: + config.close_pool() + + assert profile_path is not None + assert not profile_path.exists() diff --git a/tests/unit/adapters/test_oracledb/test_driver.py b/tests/unit/adapters/test_oracledb/test_driver.py new file mode 100644 index 000000000..7f938d684 --- /dev/null +++ b/tests/unit/adapters/test_oracledb/test_driver.py @@ -0,0 +1,205 @@ +"""Unit tests for Oracle zero-RTT session trace attribute piggybacking and UTF-8 truncation.""" + +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import patch + +import pytest + +from sqlspec.adapters.oracledb.core import ( + ORACLE_ACTION_MAX_BYTES, + ORACLE_CLIENT_IDENTIFIER_MAX_BYTES, + ORACLE_ECID_MAX_BYTES, + ORACLE_MODULE_MAX_BYTES, + apply_driver_features, + apply_oracle_session_trace_attributes, + default_statement_config, + truncate_utf8_bytes, +) +from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver +from sqlspec.core.sqlcommenter import SQLCommenterContext +from sqlspec.utils.correlation import CorrelationContext + + +def test_truncate_utf8_bytes_preserves_multibyte_codepoints() -> None: + """truncate_utf8_bytes slices cleanly on multi-byte UTF-8 codepoint boundaries.""" + multibyte = "é" * 30 + truncated = truncate_utf8_bytes(multibyte, ORACLE_MODULE_MAX_BYTES) + assert truncated == "é" * 24 + assert len(truncated.encode("utf-8")) == 48 + + emoji_text = "🚀" * 20 + truncated_emoji = truncate_utf8_bytes(emoji_text, 15) + assert truncated_emoji == "🚀" * 3 + assert len(truncated_emoji.encode("utf-8")) == 12 + + +def test_apply_oracle_session_trace_attributes_ecid_fallback() -> None: + """apply_oracle_session_trace_attributes falls back to conn.ecid when conn.econtext_id is absent.""" + + class _StrictConn: + __slots__ = ("action", "client_identifier", "ecid", "module") + + def __init__(self) -> None: + self.module: str | None = None + self.action: str | None = None + self.client_identifier: str | None = None + self.ecid: str | None = None + + conn = _StrictConn() + apply_oracle_session_trace_attributes( + conn, module="m" * 100, action="a" * 100, client_identifier="c" * 100, econtext_id="e" * 100 + ) + assert conn.module == "m" * ORACLE_MODULE_MAX_BYTES + assert conn.action == "a" * ORACLE_ACTION_MAX_BYTES + assert conn.client_identifier == "c" * ORACLE_CLIENT_IDENTIFIER_MAX_BYTES + assert conn.ecid == "e" * ORACLE_ECID_MAX_BYTES + + +@pytest.mark.anyio +async def test_oracle_sync_and_async_driver_piggyback_trace_attributes_on_slow_and_fast_paths() -> None: + """OracleSyncDriver and OracleAsyncDriver piggyback trace attributes on slow and cached fast paths without SQL comments.""" + sync_executed: list[tuple[str, dict[str, str | None]]] = [] + sync_conn: Any = SimpleNamespace( + module=None, action=None, client_identifier=None, econtext_id=None, transaction_in_progress=False + ) + + class _SyncFakeCursor: + def __init__(self) -> None: + self.description = [("VAL", "NUMBER")] + self.rowcount = 1 + self.arraysize = 100 + self.prefetchrows = 2 + + def execute(self, sql: str, _params: Any = None, **_kwargs: Any) -> None: + sync_executed.append(( + sql, + { + "module": sync_conn.module, + "action": sync_conn.action, + "client_identifier": sync_conn.client_identifier, + "econtext_id": sync_conn.econtext_id, + }, + )) + + def fetchall(self) -> list[tuple[int]]: + return [(1,)] + + def close(self) -> None: + return None + + sync_conn.cursor = _SyncFakeCursor + base_config = default_statement_config.replace( + enable_caching=True, + enable_sqlcommenter=True, + sqlcommenter_enable_context=True, + sqlcommenter_enable_traceparent=True, + ) + stmt_config, features = apply_driver_features(base_config, {"enable_session_tracing": True}) + sync_driver = OracleSyncDriver( + connection=cast("Any", sync_conn), statement_config=stmt_config, driver_features=features + ) + assert sync_driver._stmt_cache_enabled is True + + with ( + patch( + "sqlspec.adapters.oracledb.core.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "1111111111111111"), + ), + SQLCommenterContext.scope({"framework": "litestar", "route": "/orders/1"}), + CorrelationContext.context("corr-sync-1"), + ): + res1 = sync_driver.execute("SELECT 1 FROM DUAL") + + with ( + patch( + "sqlspec.adapters.oracledb.core.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "2222222222222222"), + ), + SQLCommenterContext.scope({"framework": "litestar", "route": "/orders/2"}), + CorrelationContext.context("corr-sync-2"), + ): + res2 = sync_driver.execute("SELECT 1 FROM DUAL") + + assert res1.data is not None + assert res2.data is not None + assert len(sync_executed) == 2 + assert "/*" not in sync_executed[0][0] + assert "/*" not in sync_executed[1][0] + assert sync_executed[0][1]["module"] == "litestar" + assert sync_executed[0][1]["action"] == "/orders/1" + assert sync_executed[0][1]["client_identifier"] == "corr-sync-1" + assert sync_executed[0][1]["econtext_id"] == "00-0af7651916cd43dd8448eb211c80319c-1111111111111111-01" + assert sync_executed[1][1]["module"] == "litestar" + assert sync_executed[1][1]["action"] == "/orders/2" + assert sync_executed[1][1]["client_identifier"] == "corr-sync-2" + assert sync_executed[1][1]["econtext_id"] == "00-0af7651916cd43dd8448eb211c80319c-2222222222222222-01" + + async_executed: list[tuple[str, dict[str, str | None]]] = [] + async_conn: Any = SimpleNamespace( + module=None, action=None, client_identifier=None, econtext_id=None, transaction_in_progress=False + ) + + class _AsyncFakeCursor: + def __init__(self) -> None: + self.description = [("VAL", "NUMBER")] + self.rowcount = 1 + self.arraysize = 100 + self.prefetchrows = 2 + + async def execute(self, sql: str, _params: Any = None, **_kwargs: Any) -> None: + async_executed.append(( + sql, + { + "module": async_conn.module, + "action": async_conn.action, + "client_identifier": async_conn.client_identifier, + "econtext_id": async_conn.econtext_id, + }, + )) + + async def fetchall(self) -> list[tuple[int]]: + return [(1,)] + + def close(self) -> None: + return None + + async_conn.cursor = _AsyncFakeCursor + async_driver = OracleAsyncDriver( + connection=cast("Any", async_conn), statement_config=stmt_config, driver_features=features + ) + assert async_driver._stmt_cache_enabled is True + + with ( + patch( + "sqlspec.adapters.oracledb.core.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "aaaaaaaaaaaaaaaa"), + ), + SQLCommenterContext.scope({"framework": "litestar", "route": "/async/1"}), + CorrelationContext.context("corr-async-1"), + ): + ares1 = await async_driver.execute("SELECT 1 FROM DUAL") + + with ( + patch( + "sqlspec.adapters.oracledb.core.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "bbbbbbbbbbbbbbbb"), + ), + SQLCommenterContext.scope({"framework": "litestar", "route": "/async/2"}), + CorrelationContext.context("corr-async-2"), + ): + ares2 = await async_driver.execute("SELECT 1 FROM DUAL") + + assert ares1.data is not None + assert ares2.data is not None + assert len(async_executed) == 2 + assert "/*" not in async_executed[0][0] + assert "/*" not in async_executed[1][0] + assert async_executed[0][1]["module"] == "litestar" + assert async_executed[0][1]["action"] == "/async/1" + assert async_executed[0][1]["client_identifier"] == "corr-async-1" + assert async_executed[0][1]["econtext_id"] == "00-0af7651916cd43dd8448eb211c80319c-aaaaaaaaaaaaaaaa-01" + assert async_executed[1][1]["module"] == "litestar" + assert async_executed[1][1]["action"] == "/async/2" + assert async_executed[1][1]["client_identifier"] == "corr-async-2" + assert async_executed[1][1]["econtext_id"] == "00-0af7651916cd43dd8448eb211c80319c-bbbbbbbbbbbbbbbb-01" diff --git a/tests/unit/adapters/test_psqlpy/test_core.py b/tests/unit/adapters/test_psqlpy/test_core.py index 0816849b9..d8e03eaeb 100644 --- a/tests/unit/adapters/test_psqlpy/test_core.py +++ b/tests/unit/adapters/test_psqlpy/test_core.py @@ -5,6 +5,7 @@ from decimal import Decimal from types import SimpleNamespace from typing import Any, cast +from unittest.mock import AsyncMock, patch import pytest from psqlpy import exceptions as psqlpy_exceptions @@ -19,7 +20,9 @@ get_parameter_casts, prepare_parameters_with_casts, ) +from sqlspec.adapters.psqlpy.driver import PsqlpyDriver from sqlspec.core import SQL +from sqlspec.core.sqlcommenter import SQLCommenterContext from sqlspec.driver import CachedQuery from sqlspec.exceptions import ( ConnectionTimeoutError, @@ -346,3 +349,52 @@ def test_create_mapped_exception_maps_connection_timeout_and_data_parse_errors() data_err = psqlpy_exceptions.DataError("invalid input syntax for type json: failed to parse value") assert type(psqlpy_core.create_mapped_exception(data_err)) is DataError + + +@pytest.mark.anyio +async def test_psqlpy_dispatch_execute_strips_sqlcommenter_before_dml_count_cache_lookup() -> None: + """PsqlpyDriver.dispatch_execute strips trailing sqlcommenter comments before _dml_count_query cache lookup.""" + psqlpy_core._dml_count_query.cache_clear() + fetched_queries: list[str] = [] + + class _FakeResult: + def result(self) -> list[dict[str, int]]: + return [{psqlpy_core._DML_COUNT_COLUMN: 1}] + + def _fake_fetch(query: str, _params: Any) -> _FakeResult: + fetched_queries.append(query) + return _FakeResult() + + connection = SimpleNamespace(fetch=AsyncMock(side_effect=_fake_fetch), execute=AsyncMock()) + stmt_config = build_statement_config().replace( + enable_sqlcommenter=True, sqlcommenter_enable_context=True, sqlcommenter_enable_traceparent=True + ) + driver = PsqlpyDriver(connection=cast("Any", connection), statement_config=stmt_config) + + with ( + patch( + "sqlspec.core.sqlcommenter.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "1111111111111111"), + ), + SQLCommenterContext.scope({"route": "/items/1"}), + ): + stmt1 = driver.prepare_statement("UPDATE events SET payload = $1 WHERE id = $2", ("a", 1)) + await driver.dispatch_execute(cast("Any", connection), stmt1) + + with ( + patch( + "sqlspec.core.sqlcommenter.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "2222222222222222"), + ), + SQLCommenterContext.scope({"route": "/items/2"}), + ): + stmt2 = driver.prepare_statement("UPDATE events SET payload = $1 WHERE id = $2", ("b", 2)) + await driver.dispatch_execute(cast("Any", connection), stmt2) + + cache_info = psqlpy_core._dml_count_query.cache_info() + assert cache_info.misses == 1 + assert cache_info.hits >= 1 + assert len(fetched_queries) == 2 + assert "1111111111111111" in fetched_queries[0] + assert "2222222222222222" in fetched_queries[1] + assert "1111111111111111" not in fetched_queries[1] diff --git a/tests/unit/adapters/test_psqlpy/test_driver.py b/tests/unit/adapters/test_psqlpy/test_driver.py new file mode 100644 index 000000000..cd7b47b85 --- /dev/null +++ b/tests/unit/adapters/test_psqlpy/test_driver.py @@ -0,0 +1,52 @@ +"""Unit tests for PsqlpyDriver wire-level sqlcommenter and _dml_count_query cache preservation.""" + +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import AsyncMock, patch + +import pytest + +from sqlspec.adapters.psqlpy import core as psqlpy_core +from sqlspec.adapters.psqlpy.core import build_statement_config +from sqlspec.adapters.psqlpy.driver import PsqlpyDriver +from sqlspec.core.sqlcommenter import SQLCommenterContext + + +@pytest.mark.anyio +async def test_psqlpy_dispatch_execute_preserves_dml_count_query_lru_cache_under_sqlcommenter() -> None: + """Executing the same DML across 5 traceparent contexts causes only 1 _dml_count_query cache miss.""" + psqlpy_core._dml_count_query.cache_clear() + fetched_queries: list[str] = [] + + class _FakeResult: + def result(self) -> list[dict[str, int]]: + return [{psqlpy_core._DML_COUNT_COLUMN: 1}] + + def _fake_fetch(query: str, _params: Any) -> _FakeResult: + fetched_queries.append(query) + return _FakeResult() + + connection = SimpleNamespace(fetch=AsyncMock(side_effect=_fake_fetch), execute=AsyncMock()) + stmt_config = build_statement_config().replace( + enable_sqlcommenter=True, sqlcommenter_enable_context=True, sqlcommenter_enable_traceparent=True + ) + driver = PsqlpyDriver(connection=cast("Any", connection), statement_config=stmt_config) + + for idx in range(5): + span_id = f"{idx + 1:016x}" + with ( + patch( + "sqlspec.core.sqlcommenter.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", span_id), + ), + SQLCommenterContext.scope({"route": f"/items/{idx}"}), + ): + stmt = driver.prepare_statement("UPDATE items SET qty = $1", (idx,)) + await driver.dispatch_execute(cast("Any", connection), stmt) + + cache_info = psqlpy_core._dml_count_query.cache_info() + assert cache_info.misses == 1 + assert cache_info.hits == 4 + assert len(fetched_queries) == 5 + assert "0000000000000001" in fetched_queries[0] + assert "0000000000000005" in fetched_queries[4] diff --git a/tests/unit/builder/test_explain.py b/tests/unit/builder/test_explain.py index 14b12d2f4..fc9123a18 100644 --- a/tests/unit/builder/test_explain.py +++ b/tests/unit/builder/test_explain.py @@ -1546,3 +1546,14 @@ class MockSQLObject: explain = Explain(mock_obj) # pyright: ignore[reportArgumentType] assert "SELECT * FROM mock_table" in explain._statement_sql + + +def test_build_duckdb_explain_combines_analyze_and_format_json() -> None: + """build_duckdb_explain combines ANALYZE and FORMAT JSON into EXPLAIN (ANALYZE, FORMAT JSON).""" + options = ExplainOptions(analyze=True, format=ExplainFormat.JSON) + sql_str = build_duckdb_explain("SELECT * FROM users", options) + assert sql_str == "EXPLAIN (ANALYZE, FORMAT JSON) SELECT * FROM users" + + stmt = SQL("SELECT * FROM users", statement_config=StatementConfig(dialect="duckdb")) + explained = stmt.explain(analyze=True, format="json") + assert explained.raw_sql == "EXPLAIN (ANALYZE, FORMAT JSON) SELECT * FROM users" diff --git a/tests/unit/core/test_sqlcommenter.py b/tests/unit/core/test_sqlcommenter.py index 2d49e712f..cbd30455c 100644 --- a/tests/unit/core/test_sqlcommenter.py +++ b/tests/unit/core/test_sqlcommenter.py @@ -12,9 +12,14 @@ from sqlspec.core.sqlcommenter import ( SQLCommenterContext, append_comment, + append_sqlcomment, + append_wire_sqlcomment, create_sqlcommenter_statement_transformer, + format_wire_sqlcomment, generate_comment, parse_comment, + split_sqlcomment, + strip_sqlcomment, ) from sqlspec.utils.correlation import CorrelationContext @@ -480,3 +485,140 @@ def test_statement_config_replace_preserves_sqlcommenter() -> None: config = StatementConfig(enable_sqlcommenter=True, sqlcommenter_attributes={"db_driver": "sqlite"}) replaced = config.replace(enable_caching=False) assert len(replaced.statement_transformers) == 1 + + +def test_statement_config_enable_sqlcommenter_keeps_has_transformers_false_without_user_transformers() -> None: + """Enabling sqlcommenter without user statement_transformers must keep _has_transformers False.""" + config = StatementConfig(enable_sqlcommenter=True, sqlcommenter_attributes={"db_driver": "sqlite"}) + assert config._has_transformers is False + + def user_transformer(expr: exp.Expr, params: Any) -> tuple[exp.Expr, Any]: + return expr, params + + config_with_user = StatementConfig(enable_sqlcommenter=True, statement_transformers=[user_transformer]) + assert config_with_user._has_transformers is True + + +def test_split_and_strip_sqlcomment_wire_helpers() -> None: + """Wire-level sqlcommenter helpers split, strip, and append comments without mutating base SQL.""" + commented = append_sqlcomment("SELECT * FROM users;", {"db_driver": "asyncpg", "route": "/users"}) + assert commented == "SELECT * FROM users /* db_driver='asyncpg',route='%2Fusers' */;" + base_sql, suffix = split_sqlcomment(commented) + assert base_sql == "SELECT * FROM users;" + assert suffix == "/* db_driver='asyncpg',route='%2Fusers' */" + assert strip_sqlcomment(commented) == "SELECT * FROM users;" + assert split_sqlcomment("SELECT /* not sqlcommenter */ 1") == ("SELECT /* not sqlcommenter */ 1", "") + + +def test_format_and_append_wire_sqlcomment_refreshes_dynamic_context() -> None: + """format_wire_sqlcomment and append_wire_sqlcomment replace stale comments with current context.""" + config = StatementConfig( + enable_sqlcommenter=True, + sqlcommenter_enable_context=True, + sqlcommenter_enable_traceparent=True, + sqlcommenter_attributes={"db_driver": "sqlite"}, + ) + with ( + patch( + "sqlspec.core.sqlcommenter.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "1111111111111111"), + ), + SQLCommenterContext.scope({"route": "/one"}), + ): + first = format_wire_sqlcomment("SELECT 1", config) + + with ( + patch( + "sqlspec.core.sqlcommenter.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "2222222222222222"), + ), + SQLCommenterContext.scope({"route": "/two"}), + ): + second = format_wire_sqlcomment(first, config) + direct = append_wire_sqlcomment( + first, attributes={"db_driver": "sqlite"}, enable_traceparent=True, enable_context=True + ) + + assert "1111111111111111" in first + assert "route='%2Fone'" in first + assert "2222222222222222" in second + assert "1111111111111111" not in second + assert "route='%2Ftwo'" in second + assert "route='%2Fone'" not in second + assert direct == second + + +def test_rebind_cached_parameters_refreshes_sqlcommenter_comment() -> None: + """SQL._rebind_cached_parameters refreshes dynamic sqlcommenter comments across parameter rebinds.""" + config = StatementConfig( + enable_sqlcommenter=True, + sqlcommenter_enable_context=True, + sqlcommenter_enable_traceparent=True, + sqlcommenter_attributes={"db_driver": "sqlite"}, + ) + base_stmt = SQL("SELECT * FROM users WHERE id = :id", {"id": 1}, statement_config=config) + with ( + patch( + "sqlspec.core.sqlcommenter.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "aaaaaaaaaaaaaaaa"), + ), + SQLCommenterContext.scope({"route": "/alpha"}), + ): + first_sql, _ = base_stmt.compile() + + rebound_stmt = base_stmt.copy(parameters={"id": 2}) + with ( + patch( + "sqlspec.core.sqlcommenter.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "bbbbbbbbbbbbbbbb"), + ), + SQLCommenterContext.scope({"route": "/beta"}), + ): + second_sql, _ = rebound_stmt.compile() + + assert "aaaaaaaaaaaaaaaa" in first_sql + assert "route='%2Falpha'" in first_sql + assert "bbbbbbbbbbbbbbbb" in second_sql + assert "aaaaaaaaaaaaaaaa" not in second_sql + assert "route='%2Fbeta'" in second_sql + assert "route='%2Falpha'" not in second_sql + + +def test_sqlcommenter_reads_ambient_correlation_context() -> None: + """CorrelationContext.get() populates correlation_id when enable_correlation_id or enable_context is True.""" + with CorrelationContext.context("req-123"): + commented = append_wire_sqlcomment("SELECT 1", attributes={"db_driver": "sqlite"}, enable_correlation_id=True) + assert "correlation_id='req-123'" in commented + + +def test_enable_sqlcommenter_preserves_stmt_cache_enabled_and_updates_traceparent_per_call() -> None: + """enable_sqlcommenter=True preserves _stmt_cache_enabled and refreshes traceparent per execution.""" + config = StatementConfig( + enable_sqlcommenter=True, + sqlcommenter_enable_context=True, + sqlcommenter_enable_traceparent=True, + sqlcommenter_attributes={"db_driver": "sqlite"}, + ) + assert config._has_transformers is False + stmt = SQL("SELECT 1", statement_config=config) + with ( + patch( + "sqlspec.core.sqlcommenter.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "1111111111111111"), + ), + CorrelationContext.context("req-one"), + ): + sql_one = format_wire_sqlcomment(stmt.compile()[0], config) + with ( + patch( + "sqlspec.core.sqlcommenter.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "2222222222222222"), + ), + CorrelationContext.context("req-two"), + ): + sql_two = format_wire_sqlcomment(stmt.compile()[0], config) + assert "1111111111111111" in sql_one + assert "req-one" in sql_one + assert "2222222222222222" in sql_two + assert "req-two" in sql_two + assert "1111111111111111" not in sql_two diff --git a/tests/unit/observability/test_observability_core.py b/tests/unit/observability/test_observability_core.py index 49bf65fec..100133b47 100644 --- a/tests/unit/observability/test_observability_core.py +++ b/tests/unit/observability/test_observability_core.py @@ -19,6 +19,7 @@ from sqlspec.adapters.sqlite import SqliteConfig from sqlspec.config import LifecycleConfig from sqlspec.core import SQL, ArrowResult, StatementConfig +from sqlspec.core.sqlcommenter import SQLCommenterContext from sqlspec.data_dictionary import ColumnMetadata, ForeignKeyMetadata, IndexMetadata, TableMetadata from sqlspec.driver import SyncDataDictionaryBase, SyncDriverAdapterBase from sqlspec.exceptions import SQLSpecError @@ -1062,3 +1063,102 @@ def test_fast_path_with_observability_spans_and_error_recording(tmp_path: Path) assert isinstance(span_manager.finished[2].exception, SQLSpecError) finally: config.close_pool() + + +def test_fast_path_statement_cache_preserved_with_sqlcommenter_and_refreshes_wire_comment(tmp_path: Path) -> None: + """enable_sqlcommenter=True preserves _stmt_cache_enabled and refreshes dynamic comments on cache hits.""" + observed_events: list[Any] = [] + stmt_config = StatementConfig( + dialect="sqlite", + enable_sqlcommenter=True, + sqlcommenter_enable_context=True, + sqlcommenter_enable_traceparent=True, + ) + spec = SQLSpec() + config = SqliteConfig( + connection_config={"database": str(tmp_path / "obs_sqlcommenter.db")}, + statement_config=stmt_config, + observability_config=ObservabilityConfig(statement_observers=(observed_events.append,)), + ) + spec.add_config(config) + + try: + with spec.provide_session(config) as session: + assert session._stmt_cache_enabled is True + session.execute("CREATE TABLE t_comment (id INTEGER PRIMARY KEY, val TEXT)") + observed_events.clear() + + with ( + patch( + "sqlspec.core.sqlcommenter.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "1111111111111111"), + ), + SQLCommenterContext.scope({"route": "/first"}), + ): + res1 = session.execute("SELECT val FROM t_comment WHERE id = ?", (1,)) + + cached_entry = session._stmt_cache.get("SELECT val FROM t_comment WHERE id = ?") + assert cached_entry is not None + assert "traceparent=" not in cached_entry.compiled_sql + assert "route=" not in cached_entry.compiled_sql + + with ( + patch( + "sqlspec.core.sqlcommenter.get_trace_context", + return_value=("0af7651916cd43dd8448eb211c80319c", "2222222222222222"), + ), + SQLCommenterContext.scope({"route": "/second"}), + ): + res2 = session.execute("SELECT val FROM t_comment WHERE id = ?", (1,)) + + assert len(observed_events) == 2 + assert "1111111111111111" in observed_events[0].sql + assert "route='%2Ffirst'" in observed_events[0].sql + assert "2222222222222222" in observed_events[1].sql + assert "1111111111111111" not in observed_events[1].sql + assert "route='%2Fsecond'" in observed_events[1].sql + assert "route='%2Ffirst'" not in observed_events[1].sql + assert "1111111111111111" in res1.statement.compile()[0] + assert "2222222222222222" in res2.statement.compile()[0] + finally: + config.close_pool() + + +def test_statement_event_profile_and_script_index_fields(caplog: Any) -> None: + """StatementEvent exposes profile, script_statement_index, and script_statement_count fields.""" + observed: list[Any] = [] + caplog.set_level(logging.INFO, logger="sqlspec.sql") + runtime = ObservabilityRuntime( + ObservabilityConfig(print_sql=True, statement_observers=(observed.append,)), config_name="DuckDBConfig" + ) + sample_profile = {"format": "json", "latency": 0.0012, "children": []} + runtime.emit_statement_event( + sql="SELECT 1", + parameters=None, + driver="DuckDBDriver", + operation="SCRIPT", + execution_mode="single", + is_many=False, + is_script=True, + rows_affected=0, + duration_s=0.002, + storage_backend=None, + profile=sample_profile, + script_statement_index=0, + script_statement_count=2, + ) + assert len(observed) == 1 + ev = observed[0] + assert ev.profile == sample_profile + assert ev.script_statement_index == 0 + assert ev.script_statement_count == 2 + ev_dict = ev.as_dict() + assert ev_dict["profile"] == sample_profile + assert ev_dict["script_statement_index"] == 0 + assert ev_dict["script_statement_count"] == 2 + sql_records = [r for r in caplog.records if r.name == "sqlspec.sql"] + assert sql_records + last_record = sql_records[-1] + assert last_record.__dict__.get("db.statement.profile") == sample_profile + assert last_record.__dict__.get("sqlspec.script.statement_index") == 0 + assert last_record.__dict__.get("sqlspec.script.statement_count") == 2 diff --git a/tests/unit/observability/test_observer.py b/tests/unit/observability/test_observer.py new file mode 100644 index 000000000..06241c73f --- /dev/null +++ b/tests/unit/observability/test_observer.py @@ -0,0 +1,68 @@ +"""Unit tests for StatementEvent and ObservabilityRuntime profile and script ordinal fields.""" + +import logging + +import pytest + +from sqlspec.observability import LoggingConfig, ObservabilityConfig, ObservabilityRuntime, StatementEvent +from sqlspec.observability._observer import SQL_LOGGER_NAME, create_statement_observer + + +def test_statement_event_and_runtime_support_profile_and_script_ordinals(caplog: pytest.LogCaptureFixture) -> None: + """StatementEvent and ObservabilityRuntime preserve profile, script_ordinal, and script_statement_count.""" + captured: list[StatementEvent] = [] + config = ObservabilityConfig(statement_observers=(captured.append,)) + runtime = ObservabilityRuntime(config=config, bind_key="duck", config_name="DuckDBConfig") + profile_payload = {"cpu_time": 0.01, "children": [{"operator_type": "SEQ_SCAN"}]} + + direct_event = runtime.create_statement_event( + sql="SELECT 1", + parameters=(), + driver="DuckDBDriver", + operation="SELECT", + execution_mode=None, + is_many=False, + is_script=True, + rows_affected=1, + duration_s=0.002, + storage_backend=None, + profile=profile_payload, + script_ordinal=0, + script_statement_count=3, + ) + assert direct_event.profile == profile_payload + assert direct_event.script_ordinal == 0 + assert direct_event.script_statement_index == 0 + assert direct_event.script_statement_count == 3 + + runtime.emit_statement_event( + sql="SELECT 1", + parameters=(), + driver="DuckDBDriver", + operation="SELECT", + execution_mode=None, + is_many=False, + is_script=True, + rows_affected=1, + duration_s=0.002, + storage_backend=None, + profile=profile_payload, + script_ordinal=0, + script_statement_count=3, + ) + assert len(captured) == 1 + assert captured[0].profile == profile_payload + assert captured[0].script_ordinal == 0 + assert captured[0].script_statement_count == 3 + as_dict = captured[0].as_dict() + assert as_dict["profile"] == profile_payload + assert as_dict["script_ordinal"] == 0 + assert as_dict["script_statement_count"] == 3 + + observer = create_statement_observer(LoggingConfig()) + with caplog.at_level(logging.INFO, logger=SQL_LOGGER_NAME): + observer(captured[0]) + record = caplog.records[-1] + assert getattr(record, "db.statement.profile") == profile_payload + assert getattr(record, "sqlspec.script.ordinal") == 0 + assert getattr(record, "sqlspec.script.statement_count") == 3