diff --git a/sqlspec/adapters/oracledb/__init__.py b/sqlspec/adapters/oracledb/__init__.py index 3d2c43c6f..87b8de277 100644 --- a/sqlspec/adapters/oracledb/__init__.py +++ b/sqlspec/adapters/oracledb/__init__.py @@ -8,11 +8,35 @@ OracleSyncConfig, ) from sqlspec.adapters.oracledb.core import ( + VALID_ORACLE_DRIVER_FEATURE_KEYS, OracleBlob, OracleClob, + OracleEndpoint, OracleJson, + OracleSessionContextConfig, + OracleTopologyConfig, + build_arrow_fetch_kwargs, build_connection_config, + build_duality_etag_update_sql, + build_session_checkout_plsql, + build_session_context_plsql, + build_session_release_plsql, + build_session_reset_plsql, + build_tns_description, + coerce_oracle_df_to_arrow_batches, + coerce_oracle_df_to_arrow_table, default_statement_config, + execute_session_checkout_async, + execute_session_checkout_sync, + execute_session_release_async, + execute_session_release_sync, + extract_duality_etag, + extract_duality_view_etag, + extract_oracle_wallet_zip, + inject_duality_view_etag, + oracle_df_batches_to_arrow_reader, + oracle_df_to_arrow_table, + sanitize_oracle_text_query, ) from sqlspec.adapters.oracledb.data_dictionary import ( OracledbAsyncDataDictionary, @@ -28,6 +52,7 @@ from sqlspec.adapters.oracledb.type_converter import OracleOutputConverter __all__ = ( + "VALID_ORACLE_DRIVER_FEATURE_KEYS", "OracleAsyncConfig", "OracleAsyncDriver", "OracleAsyncExceptionHandler", @@ -35,15 +60,38 @@ "OracleClob", "OracleConnectionParams", "OracleDriverFeatures", + "OracleEndpoint", "OracleJson", "OracleOutputConverter", "OraclePoolParams", + "OracleSessionContextConfig", "OracleSyncConfig", "OracleSyncDriver", "OracleSyncExceptionHandler", + "OracleTopologyConfig", "OracleVersionInfo", "OracledbAsyncDataDictionary", "OracledbSyncDataDictionary", + "build_arrow_fetch_kwargs", "build_connection_config", + "build_duality_etag_update_sql", + "build_session_checkout_plsql", + "build_session_context_plsql", + "build_session_release_plsql", + "build_session_reset_plsql", + "build_tns_description", + "coerce_oracle_df_to_arrow_batches", + "coerce_oracle_df_to_arrow_table", "default_statement_config", + "execute_session_checkout_async", + "execute_session_checkout_sync", + "execute_session_release_async", + "execute_session_release_sync", + "extract_duality_etag", + "extract_duality_view_etag", + "extract_oracle_wallet_zip", + "inject_duality_view_etag", + "oracle_df_batches_to_arrow_reader", + "oracle_df_to_arrow_table", + "sanitize_oracle_text_query", ) diff --git a/sqlspec/adapters/oracledb/adk/store.py b/sqlspec/adapters/oracledb/adk/store.py index 030a0b580..b5d85e187 100644 --- a/sqlspec/adapters/oracledb/adk/store.py +++ b/sqlspec/adapters/oracledb/adk/store.py @@ -8,8 +8,13 @@ from sqlspec import SQL from sqlspec.adapters.oracledb._storage import oracle_table_feature_report, validate_oracle_identifier from sqlspec.adapters.oracledb._typing import DatabaseError as OracleDatabaseError -from sqlspec.adapters.oracledb.core import coerce_large_parameters_async, coerce_large_parameters_sync +from sqlspec.adapters.oracledb.core import ( + coerce_large_parameters_async, + coerce_large_parameters_sync, + sanitize_oracle_text_query, +) from sqlspec.adapters.oracledb.data_dictionary import JSONStorageType, OracleVersionInfo, storage_type_from_version +from sqlspec.builder._vector_distance import build_oracle_hnsw_subquery, build_oracle_hybrid_rrf_sql from sqlspec.config import ADKConfig from sqlspec.extensions.adk import ( BaseAsyncADKStore, @@ -2463,6 +2468,10 @@ async def search_entries( effective_limit = limit if limit is not None else self._max_results try: + if embedding is not None and self._use_fts: + return await self._search_entries_hybrid_rrf( + query, embedding, app_name, user_id, effective_limit, scope_filter + ) if self._use_fts: return await self._search_entries_fts(query, app_name, user_id, effective_limit, scope_filter) return await self._search_entries_simple(query, app_name, user_id, effective_limit, scope_filter) @@ -2659,6 +2668,9 @@ async def _execute_insert_entry(self, cursor: Any, sql: str, params: "dict[str, async def _search_entries_fts( self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" ) -> "list[StoredMemory]": + sanitized_query = sanitize_oracle_text_query(query) + if not sanitized_query: + return [] where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, @@ -2674,7 +2686,74 @@ async def _search_entries_fts( ) WHERE ROWNUM <= :limit """ - params = {**scope_params, "query": query, "limit": limit} + params = {**scope_params, "query": sanitized_query, "limit": limit} + async with self._config.provide_connection() as conn: + with conn.cursor() as cursor: + await cursor.execute(sql, params) + rows = await cursor.fetchall() + return await self._rows_to_records(rows) + + async def _search_entries_hybrid_rrf( + self, + query: str, + embedding: "Sequence[float] | Any", + app_name: str, + user_id: str, + limit: int, + scope_filter: Literal["all", "user", "app"] = "all", + *, + vector_column: str = "embedding", + rrf_k: int = 60, + metric: str = "COSINE", + ) -> "list[StoredMemory]": + sanitized_query = sanitize_oracle_text_query(query) + where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) + select_cols = ( + "id", + "session_id", + "app_name", + "user_id", + "scope", + "event_id", + "author", + "timestamp", + "content_json", + "content_text", + "metadata_json", + "inserted_at", + ) + if not sanitized_query: + sql = build_oracle_hnsw_subquery( + table=self._memory_table, + vector_column=vector_column, + select_columns=select_cols, + vector_param="query_vector", + top_k_param="top_k", + metric=metric, + where_clause=where_scope, + ) + params: dict[str, Any] = {**scope_params, "query_vector": embedding, "top_k": limit} + else: + sql = build_oracle_hybrid_rrf_sql( + table=self._memory_table, + key_column="id", + vector_column=vector_column, + text_column="content_text", + select_columns=select_cols, + vector_param="query_vector", + text_param="text_query", + top_k_param="top_k", + rrf_k_param="rrf_k", + metric=metric, + where_clause=where_scope, + ) + params = { + **scope_params, + "query_vector": embedding, + "text_query": sanitized_query, + "top_k": limit, + "rrf_k": rrf_k, + } async with self._config.provide_connection() as conn: with conn.cursor() as cursor: await cursor.execute(sql, params) @@ -2823,6 +2902,10 @@ def search_entries( effective_limit = limit if limit is not None else self._max_results try: + if embedding is not None and self._use_fts: + return self._search_entries_hybrid_rrf( + query, embedding, app_name, user_id, effective_limit, scope_filter + ) if self._use_fts: return self._search_entries_fts(query, app_name, user_id, effective_limit, scope_filter) return self._search_entries_simple(query, app_name, user_id, effective_limit, scope_filter) @@ -3017,6 +3100,9 @@ def _execute_insert_entry(self, cursor: Any, sql: str, params: "dict[str, Any]") def _search_entries_fts( self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" ) -> "list[StoredMemory]": + sanitized_query = sanitize_oracle_text_query(query) + if not sanitized_query: + return [] where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, @@ -3032,7 +3118,73 @@ def _search_entries_fts( ) WHERE ROWNUM <= :limit """ - params = {**scope_params, "query": query, "limit": limit} + params = {**scope_params, "query": sanitized_query, "limit": limit} + with self._config.provide_connection() as conn, conn.cursor() as cursor: + cursor.execute(sql, params) + rows = cursor.fetchall() + return self._rows_to_records(rows) + + def _search_entries_hybrid_rrf( + self, + query: str, + embedding: "Sequence[float] | Any", + app_name: str, + user_id: str, + limit: int, + scope_filter: Literal["all", "user", "app"] = "all", + *, + vector_column: str = "embedding", + rrf_k: int = 60, + metric: str = "COSINE", + ) -> "list[StoredMemory]": + sanitized_query = sanitize_oracle_text_query(query) + where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) + select_cols = ( + "id", + "session_id", + "app_name", + "user_id", + "scope", + "event_id", + "author", + "timestamp", + "content_json", + "content_text", + "metadata_json", + "inserted_at", + ) + if not sanitized_query: + sql = build_oracle_hnsw_subquery( + table=self._memory_table, + vector_column=vector_column, + select_columns=select_cols, + vector_param="query_vector", + top_k_param="top_k", + metric=metric, + where_clause=where_scope, + ) + params: dict[str, Any] = {**scope_params, "query_vector": embedding, "top_k": limit} + else: + sql = build_oracle_hybrid_rrf_sql( + table=self._memory_table, + key_column="id", + vector_column=vector_column, + text_column="content_text", + select_columns=select_cols, + vector_param="query_vector", + text_param="text_query", + top_k_param="top_k", + rrf_k_param="rrf_k", + metric=metric, + where_clause=where_scope, + ) + params = { + **scope_params, + "query_vector": embedding, + "text_query": sanitized_query, + "top_k": limit, + "rrf_k": rrf_k, + } with self._config.provide_connection() as conn, conn.cursor() as cursor: cursor.execute(sql, params) rows = cursor.fetchall() diff --git a/sqlspec/adapters/oracledb/config.py b/sqlspec/adapters/oracledb/config.py index ff43a94bf..c3b187235 100644 --- a/sqlspec/adapters/oracledb/config.py +++ b/sqlspec/adapters/oracledb/config.py @@ -1,10 +1,12 @@ """OracleDB database configuration with direct field-based configuration.""" import contextlib -from collections.abc import Awaitable, Callable +import shutil +from collections.abc import Awaitable, Callable, Mapping from inspect import isawaitable +from pathlib import Path from ssl import TLSVersion -from typing import TYPE_CHECKING, Any, ClassVar, Literal, NotRequired, TypedDict, cast +from typing import TYPE_CHECKING, Any, ClassVar, Literal, NotRequired, Self, TypedDict, cast from sqlspec.adapters.oracledb._json_handlers import register_json_handlers from sqlspec.adapters.oracledb._typing import ( @@ -23,7 +25,20 @@ from sqlspec.adapters.oracledb._typing import oracledb_module as oracledb from sqlspec.adapters.oracledb._uuid_handlers import register_uuid_handlers from sqlspec.adapters.oracledb._vector_handlers import register_numpy_handlers -from sqlspec.adapters.oracledb.core import apply_driver_features, build_connection_config, default_statement_config +from sqlspec.adapters.oracledb.core import ( + VALID_ORACLE_DRIVER_FEATURE_KEYS, + OracleEndpoint, + OracleSessionContextConfig, + OracleTopologyConfig, + apply_driver_features, + build_connection_config, + default_statement_config, + execute_session_checkout_async, + execute_session_checkout_sync, + execute_session_release_async, + execute_session_release_sync, + extract_oracle_wallet_zip, +) from sqlspec.adapters.oracledb.data_dictionary import ( OracleVersionCache, resolve_oracle_connection_compatible_major, @@ -36,7 +51,7 @@ OracleSyncExceptionHandler, ) from sqlspec.adapters.oracledb.migrations import OracleAsyncMigrationTracker, OracleSyncMigrationTracker -from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs, SyncDatabaseConfig +from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs, MigrationConfig, SyncDatabaseConfig from sqlspec.core.capabilities import TypeCoercionCapabilities from sqlspec.driver import ( AsyncPoolConnectionContext, @@ -52,14 +67,18 @@ from types import TracebackType from sqlspec.core import StatementConfig + from sqlspec.observability import ObservabilityConfig __all__ = ( "OracleAsyncConfig", "OracleConnectionParams", "OracleDriverFeatures", + "OracleEndpoint", "OraclePoolParams", + "OracleSessionContextConfig", "OracleSyncConfig", + "OracleTopologyConfig", ) @@ -133,6 +152,8 @@ class OracleConnectionParams(TypedDict): on_connect_callback: NotRequired[Callable[..., Any]] handle: NotRequired[int] extra: NotRequired[dict[str, Any]] + topology: NotRequired[OracleTopologyConfig | dict[str, Any]] + wallet_zip: NotRequired[str | Path | bytes] class OraclePoolParams(OracleConnectionParams): @@ -217,6 +238,8 @@ class OracleDriverFeatures(TypedDict): Thin-mode only; falls back to executemany when the API is absent or the connection is in Thick mode. Defaults to True; set to False to force executemany. + session_context_hook: Optional declarative session context configuration + applied on checkout and cleared on release. Native pipeline execution is runtime-gated by async Thin-mode driver API support, Oracle Database 26ai or newer, and the SQLSPEC_ORACLE_DISABLE_PIPELINE environment override; there is no adapter config switch that can force-enable unsupported @@ -237,6 +260,7 @@ class OracleDriverFeatures(TypedDict): enable_events: NotRequired[bool] events_backend: NotRequired[OracleEventsBackend] enable_direct_path_load: NotRequired[bool] + session_context_hook: NotRequired[OracleSessionContextConfig | dict[str, Any]] class OracleSyncConnectionContext(SyncPoolConnectionContext): @@ -251,16 +275,31 @@ def __init__(self, config: "OracleSyncConfig") -> None: def __enter__(self) -> "OracleSyncConnection": if self._config.connection_instance is None: self._config.connection_instance = self._config.create_pool() - self._conn = self._config.connection_instance.acquire() - return cast("OracleSyncConnection", self._conn) + conn = self._config.connection_instance.acquire() + self._conn = conn + session_hook = self._config.driver_features.get("session_context_hook") + if session_hook is not None: + try: + execute_session_checkout_sync(conn, session_hook) + except BaseException: + self._conn = None + self._config.connection_instance.release(conn) + raise + return cast("OracleSyncConnection", conn) def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> bool | None: - if self._conn: - if self._config.connection_instance: - self._config.connection_instance.release(self._conn) + if self._conn is not None: + conn = self._conn self._conn = None + session_hook = self._config.driver_features.get("session_context_hook") + try: + if session_hook is not None: + execute_session_release_sync(conn, session_hook) + finally: + if self._config.connection_instance: + self._config.connection_instance.release(conn) return None @@ -274,21 +313,36 @@ def __init__(self, config: "OracleSyncConfig") -> None: def acquire_connection(self) -> "OracleSyncConnection": if self._config.connection_instance is None: self._config.connection_instance = self._config.create_pool() - self._conn = self._config.connection_instance.acquire() - return cast("OracleSyncConnection", self._conn) + conn = self._config.connection_instance.acquire() + self._conn = conn + session_hook = self._config.driver_features.get("session_context_hook") + if session_hook is not None: + try: + execute_session_checkout_sync(conn, session_hook) + except BaseException: + self._conn = None + self._config.connection_instance.release(conn) + raise + return cast("OracleSyncConnection", conn) def release_connection(self, _conn: "OracleSyncConnection", **kwargs: Any) -> None: if self._conn is None: return - if self._config.connection_instance: - self._config.connection_instance.release(self._conn) + conn = self._conn self._conn = None + session_hook = self._config.driver_features.get("session_context_hook") + try: + if session_hook is not None: + execute_session_release_sync(conn, session_hook) + finally: + if self._config.connection_instance: + self._config.connection_instance.release(conn) class OracleSyncConfig(SyncDatabaseConfig[OracleSyncConnection, "OracleSyncConnectionPool", OracleSyncDriver]): """Configuration for Oracle synchronous database connections.""" - __slots__ = ("_oracle_version_cache", "_pool_session_callback", "_user_connection_hook") + __slots__ = ("_extracted_wallet_dir", "_oracle_version_cache", "_pool_session_callback", "_user_connection_hook") driver_type: ClassVar[type[OracleSyncDriver]] = OracleSyncDriver connection_type: "ClassVar[type[OracleSyncConnection]]" = OracleSyncConnection @@ -333,7 +387,14 @@ def __init__( extension_config: Extension-specific configuration. **kwargs: Additional keyword arguments. """ - connection_config = build_connection_config(normalize_connection_config(connection_config)) + normalized_conn = dict(normalize_connection_config(connection_config)) + self._extracted_wallet_dir: Path | None = None + if normalized_conn.get("wallet_zip") is not None: + wallet_dir = extract_oracle_wallet_zip(normalized_conn.pop("wallet_zip")) + self._extracted_wallet_dir = wallet_dir + normalized_conn.setdefault("config_dir", str(wallet_dir)) + normalized_conn.setdefault("wallet_location", str(wallet_dir)) + connection_config = build_connection_config(normalized_conn) self._oracle_version_cache = OracleVersionCache() self._pool_session_callback = cast( "Callable[[OracleSyncConnection, str], None] | None", connection_config.pop("session_callback", None) @@ -376,6 +437,49 @@ def provide_pool(self) -> "OracleSyncConnectionPool": self.connection_instance = self.create_pool() return self.connection_instance + def derive_pool( + self, + *, + connection_config_overrides: "Mapping[str, Any] | None" = None, + driver_features_overrides: "Mapping[str, Any] | None" = None, + statement_config: "StatementConfig | None" = None, + migration_config: "dict[str, Any] | MigrationConfig | None" = None, + extension_config: "ExtensionConfigs | None" = None, + observability_config: "ObservabilityConfig | None" = None, + bind_key: "str | None" = None, + **overrides: Any, + ) -> Self: + """Clone this OracleSyncConfig with overridden connection/pool or feature settings and ``connection_instance=None``.""" + new_conn_cfg: dict[str, Any] = dict(self.connection_config) + new_driver_features: dict[str, Any] = dict(self.driver_features) + if self._pool_session_callback is not None and "session_callback" not in new_conn_cfg: + new_conn_cfg["session_callback"] = self._pool_session_callback + if self._user_connection_hook is not None and "on_connection_create" not in new_driver_features: + new_driver_features["on_connection_create"] = self._user_connection_hook + if connection_config_overrides: + new_conn_cfg.update(connection_config_overrides) + if driver_features_overrides: + new_driver_features.update(driver_features_overrides) + for key, val in overrides.items(): + if key in VALID_ORACLE_DRIVER_FEATURE_KEYS: + new_driver_features[key] = val + else: + new_conn_cfg[key] = val + return type(self)( + connection_config=new_conn_cfg, + connection_instance=None, + migration_config=cast("dict[str, Any] | None", dict(self.migration_config)) + if migration_config is None + else cast("dict[str, Any] | None", migration_config), + statement_config=self.statement_config if statement_config is None else statement_config, + driver_features=new_driver_features, + bind_key=self.bind_key if bind_key is None else bind_key, + extension_config=cast("ExtensionConfigs", dict(self.extension_config)) + if extension_config is None + else extension_config, + observability_config=self.observability_config if observability_config is None else observability_config, + ) + def get_signature_namespace(self) -> "dict[str, Any]": """Get the signature namespace for OracleDB types. @@ -394,7 +498,9 @@ def get_signature_namespace(self) -> "dict[str, Any]": "OracleAsyncExceptionHandler": OracleAsyncExceptionHandler, "OracleConnectionParams": OracleConnectionParams, "OracleDriverFeatures": OracleDriverFeatures, + "OracleEndpoint": OracleEndpoint, "OraclePoolParams": OraclePoolParams, + "OracleSessionContextConfig": OracleSessionContextConfig, "OracleSyncConnectionContext": OracleSyncConnectionContext, "OracleSyncConnection": OracleSyncConnection, "OracleSyncConnectionPool": OracleSyncConnectionPool, @@ -402,6 +508,7 @@ def get_signature_namespace(self) -> "dict[str, Any]": "OracleSyncDriver": OracleSyncDriver, "OracleSyncExceptionHandler": OracleSyncExceptionHandler, "OracleSyncSessionContext": OracleSyncSessionContext, + "OracleTopologyConfig": OracleTopologyConfig, }) return namespace @@ -470,11 +577,16 @@ def _prepare_driver(self, driver: "OracleSyncDriver") -> "OracleSyncDriver": return driver def _close_pool(self) -> None: - """Close the actual connection pool.""" - if self.connection_instance: - self.connection_instance.close() - self.connection_instance = None - self._oracle_version_cache.reset() + """Close the actual connection pool and remove any extracted temporary wallet directory.""" + try: + if self.connection_instance: + self.connection_instance.close() + self.connection_instance = None + self._oracle_version_cache.reset() + finally: + if self._extracted_wallet_dir is not None: + shutil.rmtree(self._extracted_wallet_dir, ignore_errors=True) + self._extracted_wallet_dir = None class OracleAsyncConnectionContext(AsyncPoolConnectionContext): @@ -482,15 +594,58 @@ class OracleAsyncConnectionContext(AsyncPoolConnectionContext): __slots__ = () + async def __aenter__(self) -> Any: + conn = await super().__aenter__() + session_hook = self._config.driver_features.get("session_context_hook") + if session_hook is not None: + try: + await execute_session_checkout_async(conn, session_hook) + except BaseException: + await super().__aexit__(None, None, None) + raise + return conn + + async def __aexit__( + self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: Any + ) -> "bool | None": + conn = self._connection + session_hook = self._config.driver_features.get("session_context_hook") if conn is not None else None + try: + if conn is not None and session_hook is not None: + await execute_session_release_async(conn, session_hook) + finally: + await super().__aexit__(exc_type, exc_val, exc_tb) + return None + class _OracleAsyncSessionConnectionHandler(AsyncPoolSessionFactory): __slots__ = () + async def acquire_connection(self) -> Any: + conn = await super().acquire_connection() + session_hook = self._config.driver_features.get("session_context_hook") + if session_hook is not None: + try: + await execute_session_checkout_async(conn, session_hook) + except BaseException: + await super().release_connection(conn) + raise + return conn + + async def release_connection(self, _conn: Any, **kwargs: Any) -> None: + conn = self._connection + session_hook = self._config.driver_features.get("session_context_hook") if conn is not None else None + try: + if conn is not None and session_hook is not None: + await execute_session_release_async(conn, session_hook) + finally: + await super().release_connection(_conn, **kwargs) + class OracleAsyncConfig(AsyncDatabaseConfig[OracleAsyncConnection, "OracleAsyncConnectionPool", OracleAsyncDriver]): """Configuration for Oracle asynchronous database connections.""" - __slots__ = ("_oracle_version_cache", "_pool_session_callback", "_user_connection_hook") + __slots__ = ("_extracted_wallet_dir", "_oracle_version_cache", "_pool_session_callback", "_user_connection_hook") connection_type: "ClassVar[type[OracleAsyncConnection]]" = OracleAsyncConnection driver_type: ClassVar[type[OracleAsyncDriver]] = OracleAsyncDriver @@ -537,7 +692,14 @@ def __init__( extension_config: Extension-specific configuration. **kwargs: Additional keyword arguments. """ - connection_config = build_connection_config(normalize_connection_config(connection_config)) + normalized_conn = dict(normalize_connection_config(connection_config)) + self._extracted_wallet_dir: Path | None = None + if normalized_conn.get("wallet_zip") is not None: + wallet_dir = extract_oracle_wallet_zip(normalized_conn.pop("wallet_zip")) + self._extracted_wallet_dir = wallet_dir + normalized_conn.setdefault("config_dir", str(wallet_dir)) + normalized_conn.setdefault("wallet_location", str(wallet_dir)) + connection_config = build_connection_config(normalized_conn) self._oracle_version_cache = OracleVersionCache() self._pool_session_callback = cast( "Callable[[OracleAsyncConnection, str], Any] | None", connection_config.pop("session_callback", None) @@ -581,6 +743,49 @@ async def provide_pool(self) -> "OracleAsyncConnectionPool": self.connection_instance = await self.create_pool() return self.connection_instance + def derive_pool( + self, + *, + connection_config_overrides: "Mapping[str, Any] | None" = None, + driver_features_overrides: "Mapping[str, Any] | None" = None, + statement_config: "StatementConfig | None" = None, + migration_config: "dict[str, Any] | MigrationConfig | None" = None, + extension_config: "ExtensionConfigs | None" = None, + observability_config: "ObservabilityConfig | None" = None, + bind_key: "str | None" = None, + **overrides: Any, + ) -> Self: + """Clone this OracleAsyncConfig with overridden connection/pool or feature settings and ``connection_instance=None``.""" + new_conn_cfg: dict[str, Any] = dict(self.connection_config) + new_driver_features: dict[str, Any] = dict(self.driver_features) + if self._pool_session_callback is not None and "session_callback" not in new_conn_cfg: + new_conn_cfg["session_callback"] = self._pool_session_callback + if self._user_connection_hook is not None and "on_connection_create" not in new_driver_features: + new_driver_features["on_connection_create"] = self._user_connection_hook + if connection_config_overrides: + new_conn_cfg.update(connection_config_overrides) + if driver_features_overrides: + new_driver_features.update(driver_features_overrides) + for key, val in overrides.items(): + if key in VALID_ORACLE_DRIVER_FEATURE_KEYS: + new_driver_features[key] = val + else: + new_conn_cfg[key] = val + return type(self)( + connection_config=new_conn_cfg, + connection_instance=None, + migration_config=cast("dict[str, Any] | None", dict(self.migration_config)) + if migration_config is None + else cast("dict[str, Any] | None", migration_config), + statement_config=self.statement_config if statement_config is None else statement_config, + driver_features=new_driver_features, + bind_key=self.bind_key if bind_key is None else bind_key, + extension_config=cast("ExtensionConfigs", dict(self.extension_config)) + if extension_config is None + else extension_config, + observability_config=self.observability_config if observability_config is None else observability_config, + ) + def get_signature_namespace(self) -> "dict[str, Any]": """Get the signature namespace for OracleAsyncConfig types. @@ -598,12 +803,15 @@ def get_signature_namespace(self) -> "dict[str, Any]": "OracleAsyncSessionContext": OracleAsyncSessionContext, "OracleConnectionParams": OracleConnectionParams, "OracleDriverFeatures": OracleDriverFeatures, + "OracleEndpoint": OracleEndpoint, "OraclePoolParams": OraclePoolParams, + "OracleSessionContextConfig": OracleSessionContextConfig, "OracleSyncConnection": OracleSyncConnection, "OracleSyncConnectionPool": OracleSyncConnectionPool, "OracleSyncCursor": OracleSyncCursor, "OracleSyncDriver": OracleSyncDriver, "OracleSyncExceptionHandler": OracleSyncExceptionHandler, + "OracleTopologyConfig": OracleTopologyConfig, }) return namespace @@ -672,8 +880,13 @@ def _prepare_driver(self, driver: "OracleAsyncDriver") -> "OracleAsyncDriver": return driver async def _close_pool(self) -> None: - """Close the actual async connection pool.""" - if self.connection_instance: - await self.connection_instance.close() - self.connection_instance = None - self._oracle_version_cache.reset() + """Close the actual async connection pool and remove any extracted temporary wallet directory.""" + try: + if self.connection_instance: + await self.connection_instance.close() + self.connection_instance = None + self._oracle_version_cache.reset() + finally: + if self._extracted_wallet_dir is not None: + shutil.rmtree(self._extracted_wallet_dir, ignore_errors=True) + self._extracted_wallet_dir = None diff --git a/sqlspec/adapters/oracledb/core.py b/sqlspec/adapters/oracledb/core.py index b8df3ddfe..7e29a888c 100644 --- a/sqlspec/adapters/oracledb/core.py +++ b/sqlspec/adapters/oracledb/core.py @@ -1,9 +1,14 @@ """OracleDB adapter compiled helpers.""" import contextlib +import io import re -from collections.abc import Iterable, Sized -from typing import TYPE_CHECKING, Any, Protocol, cast +import shutil +import tempfile +import zipfile +from collections.abc import Callable, Iterable, Mapping, Sequence, Sized +from pathlib import Path +from typing import TYPE_CHECKING, Any, Final, Protocol, cast from sqlspec.adapters.oracledb._json_handlers import is_json_payload from sqlspec.adapters.oracledb._param_types import OracleBlob, OracleClob, OracleJson @@ -30,6 +35,7 @@ DataError, DeadlockError, ForeignKeyViolationError, + ImproperConfigurationError, IntegrityError, NotNullViolationError, OperationalError, @@ -41,12 +47,13 @@ UniqueViolationError, ) from sqlspec.typing import NUMPY_INSTALLED -from sqlspec.utils.serializers import to_json +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 from sqlspec.utils.type_guards import has_rowcount, is_async_readable, is_readable if TYPE_CHECKING: - from collections.abc import Awaitable, Callable, Mapping + from collections.abc import Awaitable from sqlspec.adapters.oracledb._typing import ( OracleAsyncConnection, @@ -56,6 +63,7 @@ ) from sqlspec.core import SQL from sqlspec.driver import AsyncExceptionHandler, SyncExceptionHandler + from sqlspec.typing import ArrowRecordBatch, ArrowRecordBatchReader, ArrowTable __all__ = ( @@ -64,35 +72,58 @@ "ORACLEDB_SUPPORTS_SPARSE_VECTORS", "ORACLEDB_VERSION", "SPARSE_VECTOR_MIN_DATABASE_MAJOR", + "VALID_ORACLE_DRIVER_FEATURE_KEYS", "OracleAsyncStreamSource", "OracleBlob", "OracleClob", + "OracleEndpoint", "OracleJson", + "OracleSessionContextConfig", "OracleSyncStreamSource", + "OracleTopologyConfig", "apply_driver_features", "build_arrow_fetch_kwargs", "build_connection_config", + "build_duality_etag_update_sql", "build_fetch_kwargs", "build_insert_statement", "build_pipeline_stack_result", "build_profile", + "build_session_checkout_plsql", + "build_session_context_plsql", + "build_session_release_plsql", + "build_session_reset_plsql", "build_statement_config", + "build_tns_description", "build_truncate_statement", "coerce_large_parameters_async", "coerce_large_parameters_sync", "coerce_many_parameters_async", "coerce_many_parameters_sync", + "coerce_oracle_df_to_arrow_batches", + "coerce_oracle_df_to_arrow_table", "collect_async_rows", "collect_sync_rows", "connection_is_thin", "create_mapped_exception", "default_statement_config", "driver_profile", + "execute_session_checkout_async", + "execute_session_checkout_sync", + "execute_session_release_async", + "execute_session_release_sync", + "extract_duality_etag", + "extract_duality_view_etag", + "extract_oracle_wallet_zip", + "inject_duality_view_etag", "normalize_column_names", "normalize_execute_many_parameters_async", "normalize_execute_many_parameters_sync", + "oracle_df_batches_to_arrow_reader", + "oracle_df_to_arrow_table", "resolve_row_metadata", "resolve_rowcount", + "sanitize_oracle_text_query", "supports_df_batches", "supports_direct_path_load", ) @@ -212,6 +243,7 @@ }) _VERSION_COMPONENTS: int = 3 TYPE_CONVERTER = OracleOutputConverter() +_pa: Any = import_optional("pyarrow") _LOB_TYPE_NAME_MARKERS: "tuple[str, ...]" = ("LOB", "BFILE") _SCALAR_PASSTHROUGH_TYPES: "tuple[type[Any], ...]" = (bool, int, float, str, bytes, bytearray, type(None)) _BIND_PASSTHROUGH_TYPES: "tuple[type[Any], ...]" = (bool, int, float, type(None)) @@ -277,6 +309,398 @@ ORACLEDB_SUPPORTS_SPARSE_VECTORS: bool = True +VALID_ORACLE_DRIVER_FEATURE_KEYS: "Final[frozenset[str]]" = frozenset({ + "_statement_cache_holder", + "arraysize", + "arrow_requested_schema", + "custom", + "enable_direct_path_load", + "enable_events", + "enable_lob_fetch", + "enable_lowercase_column_names", + "enable_numpy_vectors", + "enable_uuid_binary", + "events_backend", + "fetch_decimals", + "fetch_lobs", + "json_deserializer", + "json_serializer", + "on_connection_create", + "on_connection_destroy", + "on_pool_create", + "on_pool_destroy", + "oracle_raw_byte_limit", + "oracle_varchar2_byte_limit", + "prefetchrows", + "session_context_hook", + "storage_capabilities", + "vector_return_format", +}) + +_ORACLE_TIME_ZONE_PATTERN: "re.Pattern[str]" = re.compile(r"^[A-Za-z0-9_+\-/:]{1,64}$") +_ORACLE_NLS_KEY_PATTERN: "re.Pattern[str]" = re.compile(r"^NLS_[A-Z0-9_]{1,60}$") +_ORACLE_TEXT_TOKEN_PATTERN: "re.Pattern[str]" = re.compile(r"[A-Za-z0-9_]+") +_ORACLE_TEXT_RESERVED_OPERATORS: "frozenset[str]" = frozenset({ + "ABOUT", + "ACCUM", + "AND", + "BT", + "DEFINEMERGE", + "DEFINESCORE", + "EQUIV", + "FUZZY", + "HASPATH", + "INPATH", + "MDATA", + "MINUS", + "NEAR", + "NOT", + "NT", + "OR", + "PT", + "RT", + "SOUNDEX", + "SQE", + "STEM", + "SYN", + "TR", + "TRSYN", + "WEIGHT", + "WILDCARD", + "WITHIN", +}) + + +_MAX_TCP_PORT: int = 65535 + + +class OracleEndpoint: + """Typed Oracle Net listener endpoint for multi-endpoint TNS descriptors.""" + + __slots__ = ("host", "port", "protocol") + + def __init__(self, host: str, port: int = 1521, protocol: str = "tcp") -> None: + if not host or not host.strip(): + msg = "OracleEndpoint host must be non-empty" + raise ImproperConfigurationError(msg) + cleaned_host = host.strip() + _validate_tns_token(cleaned_host, "host") + if isinstance(port, bool) or not isinstance(port, int) or port <= 0 or port > _MAX_TCP_PORT: + msg = "OracleEndpoint port must be an integer between 1 and 65535" + raise ImproperConfigurationError(msg) + normalized_protocol = protocol.strip().lower() + if normalized_protocol not in {"tcp", "tcps"}: + msg = f"OracleEndpoint protocol must be 'tcp' or 'tcps', got {protocol!r}" + raise ImproperConfigurationError(msg) + self.host: str = cleaned_host + self.port: int = port + self.protocol: str = normalized_protocol + + def __repr__(self) -> str: + return f"OracleEndpoint(host={self.host!r}, port={self.port!r}, protocol={self.protocol!r})" + + def __eq__(self, other: object) -> bool: + if not isinstance(other, OracleEndpoint): + return NotImplemented + return (self.host, self.port, self.protocol) == (other.host, other.port, other.protocol) + + def __hash__(self) -> int: + return hash((self.host, self.port, self.protocol)) + + +class OracleTopologyConfig: + """Typed Oracle Net multi-endpoint topology configuration for RAC, Data Guard, and CMAN.""" + + __slots__ = ( + "endpoints", + "failover", + "failover_delay", + "failover_mode_method", + "failover_mode_type", + "failover_retries", + "load_balance", + "retry_count", + "retry_delay", + "sdu", + "server", + "service_name", + "source_route", + "transport_connect_timeout", + ) + + def __init__( + self, + endpoints: "Sequence[OracleEndpoint | Mapping[str, Any]]", + service_name: str, + *, + load_balance: bool | None = True, + failover: bool | None = True, + retry_count: int | None = None, + retry_delay: int | None = None, + transport_connect_timeout: float | None = None, + sdu: int | None = None, + source_route: bool = False, + server: str | None = None, + failover_mode_type: str | None = None, + failover_mode_method: str | None = None, + failover_type: str | None = None, + failover_method: str | None = None, + failover_retries: int | None = None, + failover_delay: int | None = None, + ) -> None: + if not endpoints: + msg = "OracleTopologyConfig requires at least one endpoint" + raise ImproperConfigurationError(msg) + if not service_name or not service_name.strip(): + msg = "OracleTopologyConfig service_name must be non-empty" + raise ImproperConfigurationError(msg) + cleaned_service_name = service_name.strip() + _validate_tns_token(cleaned_service_name, "service_name") + normalized_endpoints: list[OracleEndpoint] = [] + for ep in endpoints: + if isinstance(ep, OracleEndpoint): + normalized_endpoints.append(ep) + elif isinstance(ep, Mapping): + normalized_endpoints.append( + OracleEndpoint( + host=str(ep.get("host", "")), + port=int(ep.get("port", 1521)), + protocol=str(ep.get("protocol", "tcp")), + ) + ) + else: + msg = f"Invalid endpoint in OracleTopologyConfig: {ep!r}" + raise ImproperConfigurationError(msg) + resolved_server = server.strip().upper() if server else None + if resolved_server is not None: + _validate_tns_token(resolved_server, "server") + raw_fm_type = failover_mode_type if failover_mode_type is not None else failover_type + raw_fm_method = failover_mode_method if failover_mode_method is not None else failover_method + resolved_fm_type = raw_fm_type.strip().upper() if raw_fm_type else None + if resolved_fm_type is not None: + _validate_tns_token(resolved_fm_type, "failover_mode_type") + resolved_fm_method = raw_fm_method.strip().upper() if raw_fm_method else None + if resolved_fm_method is not None: + _validate_tns_token(resolved_fm_method, "failover_mode_method") + self.endpoints: tuple[OracleEndpoint, ...] = tuple(normalized_endpoints) + self.service_name: str = cleaned_service_name + self.load_balance: bool | None = None if load_balance is None else bool(load_balance) + self.failover: bool | None = None if failover is None else bool(failover) + self.retry_count: int | None = retry_count + self.retry_delay: int | None = retry_delay + self.transport_connect_timeout: float | None = transport_connect_timeout + self.sdu: int | None = sdu + self.source_route: bool = bool(source_route) + self.server: str | None = resolved_server + self.failover_mode_type: str | None = resolved_fm_type + self.failover_mode_method: str | None = resolved_fm_method + self.failover_retries: int | None = failover_retries + self.failover_delay: int | None = failover_delay + + def __repr__(self) -> str: + return ( + f"OracleTopologyConfig(endpoints={self.endpoints!r}, service_name={self.service_name!r}, " + f"load_balance={self.load_balance!r}, failover={self.failover!r})" + ) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, OracleTopologyConfig): + return NotImplemented + return ( + self.endpoints, + self.service_name, + self.load_balance, + self.failover, + self.retry_count, + self.retry_delay, + self.transport_connect_timeout, + self.sdu, + self.source_route, + self.server, + self.failover_mode_type, + self.failover_mode_method, + self.failover_retries, + self.failover_delay, + ) == ( + other.endpoints, + other.service_name, + other.load_balance, + other.failover, + other.retry_count, + other.retry_delay, + other.transport_connect_timeout, + other.sdu, + other.source_route, + other.server, + other.failover_mode_type, + other.failover_mode_method, + other.failover_retries, + other.failover_delay, + ) + + def __hash__(self) -> int: + return hash(( + self.endpoints, + self.service_name, + self.load_balance, + self.failover, + self.retry_count, + self.retry_delay, + self.transport_connect_timeout, + self.sdu, + self.source_route, + self.server, + self.failover_mode_type, + self.failover_mode_method, + self.failover_retries, + self.failover_delay, + )) + + +class OracleSessionContextConfig: + """Declarative Oracle session context, EBR edition, NLS, EBS, APEX, and VPD configuration.""" + + __slots__ = ( + "action", + "apex_detach", + "apex_session", + "clear_context_namespaces", + "clear_identifier", + "client_identifier", + "context_attributes", + "contexts", + "ebs_apps_initialize", + "ecid", + "edition", + "mo_global_policy", + "module", + "nls_parameters", + "reset_on_release", + "reset_package", + "time_zone", + ) + + def __init__( + self, + *, + module: "str | Callable[[], str | None] | None" = None, + action: "str | Callable[[], str | None] | None" = None, + client_identifier: "str | Callable[[], str | None] | None" = None, + ecid: "str | Callable[[], str | None] | None" = None, + edition: "str | Callable[[], str | None] | None" = None, + time_zone: "str | Callable[[], str | None] | None" = None, + nls_parameters: "Mapping[str, str] | Callable[[], Mapping[str, str] | None] | None" = None, + contexts: "Mapping[str, Mapping[str, str | Callable[[], str | None]]] | Callable[[], Mapping[str, Mapping[str, str]] | None] | None" = None, + context_attributes: "Mapping[str, Mapping[str, str | Callable[[], str | None]]] | Callable[[], Mapping[str, Mapping[str, str]] | None] | None" = None, + ebs_apps_initialize: "tuple[int, int, int] | None" = None, + mo_global_policy: "tuple[str, int | None] | None" = None, + apex_session: "tuple[int, int, int] | None" = None, + reset_package: bool = False, + clear_identifier: bool = False, + clear_context_namespaces: "Sequence[str]" = (), + apex_detach: bool = False, + reset_on_release: bool = True, + ) -> None: + if isinstance(edition, str): + _validate_oracle_identifier(edition, "edition") + if isinstance(time_zone, str): + _validate_oracle_time_zone(time_zone) + if isinstance(nls_parameters, Mapping): + _validate_oracle_nls_parameters(nls_parameters) + resolved_contexts = contexts if contexts is not None else context_attributes + if isinstance(resolved_contexts, Mapping): + _validate_oracle_contexts(resolved_contexts) + for ns in clear_context_namespaces: + _validate_oracle_identifier(ns, "context namespace") + self.module = module + self.action = action + self.client_identifier = client_identifier + self.ecid = ecid + self.edition = edition + self.time_zone = time_zone + self.nls_parameters = dict(nls_parameters) if isinstance(nls_parameters, Mapping) else nls_parameters + normalized_ctx = ( + {ns: dict(attrs) for ns, attrs in resolved_contexts.items()} + if isinstance(resolved_contexts, Mapping) + else resolved_contexts + ) + self.contexts = normalized_ctx + self.context_attributes = normalized_ctx + self.ebs_apps_initialize = ebs_apps_initialize + self.mo_global_policy = mo_global_policy + self.apex_session = apex_session + self.reset_package = bool(reset_package) + self.clear_identifier = bool(clear_identifier) + self.clear_context_namespaces: tuple[str, ...] = tuple(clear_context_namespaces) + self.apex_detach = bool(apex_detach) + self.reset_on_release = bool(reset_on_release) + + def __repr__(self) -> str: + return ( + f"OracleSessionContextConfig(module={self.module!r}, action={self.action!r}, " + f"client_identifier={self.client_identifier!r}, ecid={self.ecid!r}, " + f"edition={self.edition!r}, time_zone={self.time_zone!r}, " + f"nls_parameters={self.nls_parameters!r}, contexts={self.contexts!r}, " + f"reset_on_release={self.reset_on_release!r})" + ) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, OracleSessionContextConfig): + return NotImplemented + return ( + self.module, + self.action, + self.client_identifier, + self.ecid, + self.edition, + self.time_zone, + self.nls_parameters, + self.contexts, + self.ebs_apps_initialize, + self.mo_global_policy, + self.apex_session, + self.reset_package, + self.clear_identifier, + self.clear_context_namespaces, + self.apex_detach, + self.reset_on_release, + ) == ( + other.module, + other.action, + other.client_identifier, + other.ecid, + other.edition, + other.time_zone, + other.nls_parameters, + other.contexts, + other.ebs_apps_initialize, + other.mo_global_policy, + other.apex_session, + other.reset_package, + other.clear_identifier, + other.clear_context_namespaces, + other.apex_detach, + other.reset_on_release, + ) + + def __hash__(self) -> int: + return hash(( + self.module, + self.action, + self.client_identifier, + self.ecid, + self.edition, + self.time_zone, + self.ebs_apps_initialize, + self.mo_global_policy, + self.apex_session, + self.reset_package, + self.clear_identifier, + self.clear_context_namespaces, + self.apex_detach, + self.reset_on_release, + )) + def connection_is_thin(connection: object) -> bool: """Return whether an Oracle connection is in Thin mode.""" @@ -329,9 +753,104 @@ def build_fetch_kwargs(driver_features: "dict[str, Any]") -> "dict[str, object]" return fetch_kwargs -def build_arrow_fetch_kwargs(driver_features: "dict[str, Any]") -> "dict[str, object]": - """Build Oracle Arrow/DataFrame fetch keyword arguments.""" - return build_fetch_kwargs(driver_features) +def build_arrow_fetch_kwargs( + driver_features_or_batch_size: "Mapping[str, Any] | int | None" = None, + driver_features: "Mapping[str, Any] | None" = None, + arrow_schema: Any = None, + *, + batch_size: int | None = None, +) -> "dict[str, object]": + """Build Oracle Arrow/DataFrame fetch keyword arguments. + + ``oracledb.Connection.fetch_df_all`` and ``fetch_df_batches`` accept only + ``statement``, ``parameters``, ``arraysize``, ``fetch_decimals``, and + ``requested_schema``; passing ``fetch_lobs`` raises ``TypeError``. + """ + if isinstance(driver_features_or_batch_size, Mapping): + resolved_features: Mapping[str, Any] | None = driver_features_or_batch_size + resolved_batch_size = batch_size + else: + resolved_features = driver_features + resolved_batch_size = batch_size if batch_size is not None else driver_features_or_batch_size + + fetch_kwargs: dict[str, object] = {} + if resolved_batch_size is not None: + fetch_kwargs["arraysize"] = resolved_batch_size + if resolved_features is not None: + fetch_decimals = resolved_features.get("fetch_decimals") + if fetch_decimals is not None: + fetch_kwargs["fetch_decimals"] = fetch_decimals + requested_schema = resolved_features.get("arrow_requested_schema") + if arrow_schema is None and requested_schema is not None: + fetch_kwargs["requested_schema"] = requested_schema + if arrow_schema is not None: + fetch_kwargs["requested_schema"] = arrow_schema + return fetch_kwargs + + +def oracle_df_to_arrow_table(oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None) -> "ArrowTable": + """Convert an Oracle DataFrame to a PyArrow Table using zero-copy PyCapsule when supported.""" + if pa_module is None: + ensure_pyarrow() + pa_module = _pa + + if hasattr(oracle_df, "__arrow_c_stream__"): + reader = ( + pa_module.RecordBatchReader.from_stream(oracle_df, schema=arrow_schema) + if arrow_schema is not None + else pa_module.RecordBatchReader.from_stream(oracle_df) + ) + return cast("ArrowTable", reader.read_all()) + table = pa_module.table(oracle_df) + if arrow_schema is not None: + table = table.cast(arrow_schema) + return cast("ArrowTable", table) + + +def coerce_oracle_df_to_arrow_table(oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None) -> "ArrowTable": + """Alias for :func:`oracle_df_to_arrow_table`.""" + return oracle_df_to_arrow_table(oracle_df, pa_module=pa_module, arrow_schema=arrow_schema) + + +def coerce_oracle_df_to_arrow_batches( + oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None +) -> "list[ArrowRecordBatch]": + """Convert an Oracle DataFrame batch into PyArrow RecordBatches via PyCapsule when supported.""" + if pa_module is None: + ensure_pyarrow() + pa_module = _pa + + if hasattr(oracle_df, "__arrow_c_stream__"): + reader = ( + pa_module.RecordBatchReader.from_stream(oracle_df, schema=arrow_schema) + if arrow_schema is not None + else pa_module.RecordBatchReader.from_stream(oracle_df) + ) + return cast("list[ArrowRecordBatch]", list(reader)) + table = pa_module.table(oracle_df) + if arrow_schema is not None: + table = table.cast(arrow_schema) + return cast("list[ArrowRecordBatch]", table.to_batches()) + + +def oracle_df_batches_to_arrow_reader( + batches: "Iterable[Any]", schema: Any = None, pa_module: Any = None +) -> "ArrowRecordBatchReader": + """Convert an iterable of Oracle DataFrame batches into a PyArrow RecordBatchReader.""" + if pa_module is None: + ensure_pyarrow() + pa_module = _pa + + collected_batches: list[Any] = [] + for df_batch in batches: + collected_batches.extend(coerce_oracle_df_to_arrow_batches(df_batch, pa_module=pa_module, arrow_schema=schema)) + if collected_batches: + resolved_schema = schema if schema is not None else collected_batches[0].schema + return cast( + "ArrowRecordBatchReader", pa_module.RecordBatchReader.from_batches(resolved_schema, collected_batches) + ) + resolved_schema = schema if schema is not None else pa_module.schema([]) + return cast("ArrowRecordBatchReader", pa_module.RecordBatchReader.from_batches(resolved_schema, [])) def normalize_column_names(column_names: "list[str]", driver_features: "dict[str, Any]") -> "list[str]": @@ -373,32 +892,6 @@ def normalize_execute_many_parameters_async(parameters: Any) -> Any: return normalize_execute_many_parameters_sync(parameters) -class _OracleJsonBindingState: - """Lazily resolve whether JSON-shaped values require BLOB locators.""" - - __slots__ = ("_connection", "_resolved", "_use_blob", "_version_cache") - - def __init__(self, connection: Any, version_cache: Any) -> None: - self._connection = connection - self._version_cache = version_cache - self._resolved = False - self._use_blob = False - - def uses_blob(self) -> bool: - """Return the operation-scoped JSON storage decision.""" - if not self._resolved: - server_major = resolve_oracle_connection_major(self._connection, self._version_cache) - compatible_major = resolve_oracle_connection_compatible_major(self._connection, self._version_cache) - self._use_blob = ( - isinstance(server_major, int) - and not isinstance(server_major, bool) - and resolve_oracle_json_storage(server_major, compatible_major=compatible_major) - != ORACLE_JSON_STORAGE_NATIVE - ) - self._resolved = True - return self._use_blob - - def coerce_large_parameters_sync( connection: Any, parameters: Any, @@ -665,16 +1158,23 @@ def resolve_rowcount(cursor: Any) -> int: def apply_driver_features( statement_config: "StatementConfig", driver_features: "Mapping[str, Any] | None" ) -> "tuple[StatementConfig, dict[str, Any]]": - """Apply OracleDB driver feature defaults. + """Apply OracleDB driver feature defaults and validate known keys. Defaults ``fetch_lobs`` to ``False`` so direct LOB reads under the 1 GB fetch ceiling return text/bytes by default. Locators remain available for OUT binds, PL/SQL, and explicit ``fetch_lobs=True``. """ + if driver_features: + unknown_keys = sorted(set(driver_features.keys()) - VALID_ORACLE_DRIVER_FEATURE_KEYS) + if unknown_keys: + msg = f"Unknown Oracle driver_features key(s): {', '.join(unknown_keys)}" + raise ImproperConfigurationError(msg) features: dict[str, Any] = dict(driver_features) if driver_features else {} features.setdefault("enable_numpy_vectors", NUMPY_INSTALLED) features.setdefault("enable_lowercase_column_names", True) features.setdefault("enable_uuid_binary", True) + if "enable_lob_fetch" in features and "fetch_lobs" not in features: + features["fetch_lobs"] = bool(features["enable_lob_fetch"]) features.setdefault("fetch_lobs", False) if "vector_return_format" not in features: if not NUMPY_INSTALLED: @@ -685,6 +1185,9 @@ def apply_driver_features( features["vector_return_format"] = "array" features.setdefault("oracle_varchar2_byte_limit", 4000) features.setdefault("oracle_raw_byte_limit", 2000) + session_hook = features.get("session_context_hook") + if isinstance(session_hook, Mapping): + features["session_context_hook"] = OracleSessionContextConfig(**session_hook) return statement_config, features @@ -1142,39 +1645,441 @@ def build_statement_config(*, json_serializer: "Callable[[Any], str] | None" = N ) -class _OracleSyncStreamDriver(Protocol): - connection: "OracleSyncConnection" - driver_features: dict[str, Any] +driver_profile = build_profile() - def handle_database_exceptions(self) -> "SyncExceptionHandler": ... +default_statement_config = build_statement_config() - def _check_pending_exception(self, exc_handler: "SyncExceptionHandler") -> None: ... - def _resolve_row_metadata(self, description: object) -> tuple[list[str], bool]: ... +def extract_oracle_wallet_zip( + wallet_zip: "Path | str | bytes", target_dir: "Path | None" = None, *, prefix: str = "sqlspec_oracle_wallet_" +) -> Path: + """Extract an Oracle mTLS wallet archive safely into a 0o700 directory. + Validates all archive entries against path traversal (ZipSlip), sets + ``0o600`` permissions on extracted files, and verifies expected Oracle + wallet artifacts (``ewallet.pem`` or ``tnsnames.ora``) exist. -class _OracleAsyncStreamDriver(Protocol): - connection: "OracleAsyncConnection" - driver_features: dict[str, Any] + Args: + wallet_zip: Path to a wallet ``.zip`` archive or raw archive bytes. + target_dir: Optional directory to extract into. When ``None``, creates + a temporary directory with ``prefix``. + prefix: Temporary directory prefix when ``target_dir`` is omitted. - def handle_database_exceptions(self) -> "AsyncExceptionHandler": ... + Returns: + Path to the extracted wallet directory. - def _check_pending_exception(self, exc_handler: "AsyncExceptionHandler") -> None: ... + Raises: + ImproperConfigurationError: If the archive is missing, corrupt, contains + unsafe paths, or lacks required wallet files. + """ + created_temp = target_dir is None + if target_dir is None: + extract_dir = Path(tempfile.mkdtemp(prefix=prefix)) + else: + extract_dir = Path(target_dir) + extract_dir.mkdir(parents=True, exist_ok=True) + extract_dir.chmod(0o700) - async def _run_with_exception_handler( - self, - exc_handler: "AsyncExceptionHandler", - operation: "Callable[..., Awaitable[Any]]", - *args: Any, - **kwargs: Any, - ) -> Any: ... + resolved_root = extract_dir.resolve() + try: + _extract_wallet_archive_into(wallet_zip, extract_dir, resolved_root) + except BaseException: + if created_temp: + shutil.rmtree(extract_dir, ignore_errors=True) + raise - def _resolve_row_metadata(self, description: object) -> tuple[list[str], bool]: ... + return extract_dir -driver_profile = build_profile() +def build_tns_description(topology: "OracleTopologyConfig | Mapping[str, Any]") -> str: + """Compile a multi-endpoint Oracle Net ``(DESCRIPTION=...)`` connect descriptor. -default_statement_config = build_statement_config() + Args: + topology: :class:`OracleTopologyConfig` or mapping of topology options. + + Returns: + Formatted TNS descriptor string. + """ + resolved = OracleTopologyConfig(**topology) if isinstance(topology, Mapping) else topology + if not resolved.endpoints: + msg = "OracleTopologyConfig requires at least one endpoint" + raise ImproperConfigurationError(msg) + _validate_tns_token(resolved.service_name, "service_name") + for ep in resolved.endpoints: + _validate_tns_token(ep.host, "host") + desc_parts: list[str] = [] + if resolved.sdu is not None: + desc_parts.append(f"(SDU={resolved.sdu})") + if resolved.retry_count is not None: + desc_parts.append(f"(RETRY_COUNT={resolved.retry_count})") + if resolved.retry_delay is not None: + desc_parts.append(f"(RETRY_DELAY={resolved.retry_delay})") + if resolved.transport_connect_timeout is not None: + desc_parts.append(f"(TRANSPORT_CONNECT_TIMEOUT={resolved.transport_connect_timeout})") + + addr_parts: list[str] = [] + if resolved.load_balance is not None: + addr_parts.append(f"(LOAD_BALANCE={'ON' if resolved.load_balance else 'OFF'})") + if resolved.failover is not None: + addr_parts.append(f"(FAILOVER={'ON' if resolved.failover else 'OFF'})") + if resolved.source_route: + addr_parts.append("(SOURCE_ROUTE=ON)") + addr_parts.extend( + f"(ADDRESS=(PROTOCOL={ep.protocol.upper()})(HOST={ep.host})(PORT={ep.port}))" for ep in resolved.endpoints + ) + desc_parts.append(f"(ADDRESS_LIST={''.join(addr_parts)})") + + connect_parts: list[str] = [f"(SERVICE_NAME={resolved.service_name})"] + if resolved.server is not None: + _validate_tns_token(resolved.server, "server") + connect_parts.append(f"(SERVER={resolved.server})") + if resolved.failover_mode_type is not None or resolved.failover_mode_method is not None: + fm_type = resolved.failover_mode_type or "SELECT" + fm_method = resolved.failover_mode_method or "BASIC" + _validate_tns_token(fm_type, "failover_mode_type") + _validate_tns_token(fm_method, "failover_mode_method") + fm_clauses = [f"(TYPE={fm_type})", f"(METHOD={fm_method})"] + if resolved.failover_retries is not None: + fm_clauses.append(f"(RETRIES={resolved.failover_retries})") + if resolved.failover_delay is not None: + fm_clauses.append(f"(DELAY={resolved.failover_delay})") + connect_parts.append(f"(FAILOVER_MODE={''.join(fm_clauses)})") + desc_parts.append(f"(CONNECT_DATA={''.join(connect_parts)})") + return f"(DESCRIPTION={''.join(desc_parts)})" + + +def _build_session_checkout_plan( + config: "OracleSessionContextConfig | Mapping[str, Any]", +) -> "tuple[str | None, dict[str, Any], dict[str, str | None]]": + resolved = OracleSessionContextConfig(**config) if isinstance(config, Mapping) else config + piggyback_attrs: dict[str, str | None] = {} + + client_id = _resolve_callable_or_str(resolved.client_identifier) + if client_id is not None: + piggyback_attrs["client_identifier"] = client_id + mod = _resolve_callable_or_str(resolved.module) + if mod is not None: + piggyback_attrs["module"] = mod + act = _resolve_callable_or_str(resolved.action) + if act is not None: + piggyback_attrs["action"] = act + ecid_val = _resolve_callable_or_str(resolved.ecid) + if ecid_val is not None: + piggyback_attrs["econtext_id"] = ecid_val + + alter_clauses: list[str] = [] + tz = _resolve_callable_or_str(resolved.time_zone) + if tz is not None: + _validate_oracle_time_zone(tz) + alter_clauses.append(f"TIME_ZONE = ''{tz}''") + + ed = _resolve_callable_or_str(resolved.edition) + if ed is not None: + _validate_oracle_identifier(ed, "edition") + alter_clauses.append(f"EDITION = {ed}") + + nls_raw = resolved.nls_parameters() if callable(resolved.nls_parameters) else resolved.nls_parameters + if nls_raw: + _validate_oracle_nls_parameters(nls_raw) + for key in sorted(nls_raw.keys()): + val = nls_raw[key] + alter_clauses.append(f"{key.upper()} = ''{val}''") + + plsql_statements: list[str] = [] + if alter_clauses: + joined_alter = " ".join(alter_clauses) + plsql_statements.append(f"EXECUTE IMMEDIATE 'ALTER SESSION SET {joined_alter}';") + + bind_params: dict[str, Any] = {} + uses_plsql_identifier = bool( + resolved.clear_identifier + or resolved.ebs_apps_initialize is not None + or resolved.mo_global_policy is not None + or resolved.apex_session is not None + or ( + client_id is not None + and resolved.module is None + and resolved.action is None + and resolved.ecid is None + and not alter_clauses + and not resolved.contexts + ) + ) + if client_id is not None and uses_plsql_identifier: + plsql_statements.append("DBMS_SESSION.SET_IDENTIFIER(:p_client_id);") + bind_params["p_client_id"] = client_id + + ctx_raw = resolved.contexts() if callable(resolved.contexts) else resolved.contexts + if ctx_raw: + _validate_oracle_contexts(ctx_raw) + idx = 0 + for namespace in sorted(ctx_raw.keys()): + attrs = ctx_raw[namespace] + for attr_key in sorted(attrs.keys()): + raw_val = attrs[attr_key] + attr_val = _resolve_callable_or_str(raw_val) + ns_param = f"ctx_ns_{idx}" + key_param = f"ctx_attr_{idx}" + val_param = f"ctx_val_{idx}" + bind_params[ns_param] = namespace + bind_params[key_param] = attr_key + bind_params[val_param] = attr_val + plsql_statements.append(f"DBMS_SESSION.SET_CONTEXT(:{ns_param}, :{key_param}, :{val_param});") + idx += 1 + + if resolved.ebs_apps_initialize is not None: + u_id, r_id, a_id = resolved.ebs_apps_initialize + plsql_statements.append( + "FND_GLOBAL.APPS_INITIALIZE(user_id => :p_ebs_uid, resp_id => :p_ebs_rid, resp_appl_id => :p_ebs_aid);" + ) + bind_params["p_ebs_uid"] = u_id + bind_params["p_ebs_rid"] = r_id + bind_params["p_ebs_aid"] = a_id + + if resolved.mo_global_policy is not None: + mode, org_id = resolved.mo_global_policy + plsql_statements.append("MO_GLOBAL.SET_POLICY_CONTEXT(:p_mo_mode, :p_mo_org);") + bind_params["p_mo_mode"] = mode + bind_params["p_mo_org"] = org_id + + if resolved.apex_session is not None: + app_id, page_id, sess_id = resolved.apex_session + plsql_statements.append( + "APEX_SESSION.ATTACH(p_app_id => :p_apex_app, p_page_id => :p_apex_page, p_session_id => :p_apex_sess);" + ) + bind_params["p_apex_app"] = app_id + bind_params["p_apex_page"] = page_id + bind_params["p_apex_sess"] = sess_id + + if not plsql_statements: + return None, bind_params, piggyback_attrs + return f"BEGIN {' '.join(plsql_statements)} END;", bind_params, piggyback_attrs + + +def build_session_checkout_plsql( + config: "OracleSessionContextConfig | Mapping[str, Any]", +) -> "tuple[str | None, dict[str, Any]]": + """Compile checkout session context into a single PL/SQL block and bind parameters.""" + plsql, bind_params, _ = _build_session_checkout_plan(config) + return plsql, bind_params + + +def build_session_context_plsql( + config: "OracleSessionContextConfig | Mapping[str, Any]", +) -> "tuple[str | None, dict[str, Any]]": + """Return checkout PL/SQL block and bind parameters for :class:`OracleSessionContextConfig`.""" + return build_session_checkout_plsql(config) + + +def _build_session_release_plan( + config: "OracleSessionContextConfig | Mapping[str, Any]", +) -> "tuple[str | None, dict[str, Any], dict[str, str | None]]": + resolved = OracleSessionContextConfig(**config) if isinstance(config, Mapping) else config + if not resolved.reset_on_release and not ( + resolved.apex_detach or resolved.clear_identifier or resolved.clear_context_namespaces or resolved.reset_package + ): + return None, {}, {} + + piggyback_attrs: dict[str, str | None] = {} + if resolved.client_identifier is not None or resolved.clear_identifier: + piggyback_attrs["client_identifier"] = "" + if resolved.module is not None: + piggyback_attrs["module"] = "" + if resolved.action is not None: + piggyback_attrs["action"] = "" + if resolved.ecid is not None: + piggyback_attrs["econtext_id"] = "" + + plsql_statements: list[str] = [] + bind_params: dict[str, Any] = {} + if resolved.apex_detach: + plsql_statements.append("APEX_SESSION.DETACH;") + if resolved.clear_identifier: + plsql_statements.append("DBMS_SESSION.CLEAR_IDENTIFIER;") + for idx, namespace in enumerate(resolved.clear_context_namespaces): + plsql_statements.append(f"DBMS_SESSION.CLEAR_ALL_CONTEXT(:p_clr_ns_{idx});") + bind_params[f"p_clr_ns_{idx}"] = namespace + if resolved.reset_on_release: + ctx_raw = resolved.contexts() if callable(resolved.contexts) else resolved.contexts + if ctx_raw: + _validate_oracle_contexts(ctx_raw) + plsql_statements.extend( + f"DBMS_SESSION.CLEAR_CONTEXT('{namespace}');" + for namespace in sorted(ctx_raw.keys()) + if namespace not in resolved.clear_context_namespaces + ) + if resolved.reset_package: + plsql_statements.append("DBMS_SESSION.RESET_PACKAGE;") + + if not plsql_statements: + return None, bind_params, piggyback_attrs + return f"BEGIN {' '.join(plsql_statements)} END;", bind_params, piggyback_attrs + + +def build_session_release_plsql( + config: "OracleSessionContextConfig | Mapping[str, Any]", +) -> "tuple[str | None, dict[str, Any]]": + """Compile release reset PL/SQL block and bind parameters.""" + plsql, bind_params, _ = _build_session_release_plan(config) + return plsql, bind_params + + +def build_session_reset_plsql( + config: "OracleSessionContextConfig | Mapping[str, Any]", +) -> "tuple[str | None, dict[str, Any]]": + """Return release reset PL/SQL block and bind parameters for :class:`OracleSessionContextConfig`.""" + return build_session_release_plsql(config) + + +def execute_session_checkout_sync(connection: Any, config: "OracleSessionContextConfig | Mapping[str, Any]") -> None: + """Apply TTC piggyback attributes and execute session checkout PL/SQL synchronously.""" + plsql, bind_params, piggyback_attrs = _build_session_checkout_plan(config) + for attr_name, attr_val in piggyback_attrs.items(): + with contextlib.suppress(AttributeError): + setattr(connection, attr_name, attr_val) + if plsql is not None: + with connection.cursor() as cursor: + if bind_params: + cursor.execute(plsql, bind_params) + else: + cursor.execute(plsql) + + +def execute_session_release_sync(connection: Any, config: "OracleSessionContextConfig | Mapping[str, Any]") -> None: + """Clear TTC piggyback attributes and execute session release PL/SQL synchronously.""" + plsql, bind_params, piggyback_attrs = _build_session_release_plan(config) + for attr_name, attr_val in piggyback_attrs.items(): + with contextlib.suppress(AttributeError): + setattr(connection, attr_name, attr_val) + if plsql is not None: + with connection.cursor() as cursor: + if bind_params: + cursor.execute(plsql, bind_params) + else: + cursor.execute(plsql) + + +async def execute_session_checkout_async( + connection: Any, config: "OracleSessionContextConfig | Mapping[str, Any]" +) -> None: + """Apply TTC piggyback attributes and execute session checkout PL/SQL asynchronously.""" + plsql, bind_params, piggyback_attrs = _build_session_checkout_plan(config) + for attr_name, attr_val in piggyback_attrs.items(): + with contextlib.suppress(AttributeError): + setattr(connection, attr_name, attr_val) + if plsql is not None: + with connection.cursor() as cursor: + if bind_params: + await cursor.execute(plsql, bind_params) + else: + await cursor.execute(plsql) + + +async def execute_session_release_async( + connection: Any, config: "OracleSessionContextConfig | Mapping[str, Any]" +) -> None: + """Clear TTC piggyback attributes and execute session release PL/SQL asynchronously.""" + plsql, bind_params, piggyback_attrs = _build_session_release_plan(config) + for attr_name, attr_val in piggyback_attrs.items(): + with contextlib.suppress(AttributeError): + setattr(connection, attr_name, attr_val) + if plsql is not None: + with connection.cursor() as cursor: + if bind_params: + await cursor.execute(plsql, bind_params) + else: + await cursor.execute(plsql) + + +def sanitize_oracle_text_query(raw_query: str, *, operator: str = "ACCUM") -> str: + """Sanitize a free-text search string for Oracle Text ``CONTAINS``. + + Extracts alphanumeric tokens and wraps each token in ``{token}`` literal + braces so user queries containing Oracle Text reserved keywords or syntax + punctuation never raise ``DRG-50901``. + + Args: + raw_query: Raw user search string. + operator: Operator used to join escaped tokens (defaults to ``"ACCUM"``). + + Returns: + Sanitized Oracle Text query string, or ``""`` when no tokens remain. + """ + if not raw_query or not raw_query.strip(): + return "" + tokens = [match.group(0) for match in _ORACLE_TEXT_TOKEN_PATTERN.finditer(raw_query)] + if not tokens: + return "" + joiner = f" {operator} " if operator else " " + return joiner.join(f"{{{tok}}}" for tok in tokens) + + +def extract_duality_view_etag(row_or_doc: "Mapping[str, Any] | str | bytes") -> str | None: + """Extract ``_metadata.etag`` from an Oracle 26ai JSON Relational Duality View document or row. + + Args: + row_or_doc: Duality view document mapping, row mapping containing a + ``data``/``DATA`` document, or raw JSON string/bytes. + + Returns: + The hex ETag string if present, otherwise ``None``. + """ + parsed: Any = row_or_doc + if isinstance(parsed, (str, bytes)): + try: + parsed = from_json(parsed) + except (ValueError, TypeError): + return None + if not isinstance(parsed, Mapping): + return None + metadata = parsed.get("_metadata") + if isinstance(metadata, Mapping): + etag = metadata.get("etag") + if isinstance(etag, str) and etag: + return etag + for key in ("data", "DATA"): + nested = parsed.get(key) + if isinstance(nested, (str, bytes)): + try: + nested = from_json(nested) + except (ValueError, TypeError): + continue + if isinstance(nested, Mapping): + nested_meta = nested.get("_metadata") + if isinstance(nested_meta, Mapping): + etag = nested_meta.get("etag") + if isinstance(etag, str) and etag: + return etag + return None + + +def extract_duality_etag(payload: "Mapping[str, Any] | str | bytes") -> str | None: + """Alias for :func:`extract_duality_view_etag`.""" + return extract_duality_view_etag(payload) + + +def inject_duality_view_etag(doc: "Mapping[str, Any]", etag: str) -> "dict[str, Any]": + """Return a copy of a Duality View document with ``_metadata.etag`` injected for optimistic concurrency.""" + result: dict[str, Any] = dict(doc) + existing_meta = result.get("_metadata") + merged_meta: dict[str, Any] = dict(existing_meta) if isinstance(existing_meta, Mapping) else {} + merged_meta["etag"] = etag + result["_metadata"] = merged_meta + return result + + +def build_duality_etag_update_sql( + view_name: str, id_column: str = "ID", data_column: str = "DATA", *, etag_param: str = "expected_etag" +) -> str: + """Build an optimistic-concurrency ``UPDATE`` statement for an Oracle 26ai JSON Relational Duality View.""" + _validate_oracle_identifier(view_name, "view_name") + _validate_oracle_identifier(id_column, "id_column") + _validate_oracle_identifier(data_column, "data_column") + _validate_oracle_identifier(etag_param, "etag_param") + return ( + f"UPDATE {view_name} SET {data_column} = :data " + f"WHERE {id_column} = :id AND JSON_VALUE({data_column}, '$._metadata.etag') = :{etag_param}" + ) def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str, Any]": @@ -1186,7 +2091,20 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str Returns: Normalized connection parameters. """ - config = dict(connection_config) + config = {k: v for k, v in connection_config.items() if v is not None} + if "wallet_zip" in config: + wallet_dir = extract_oracle_wallet_zip(config.pop("wallet_zip")) + config.setdefault("config_dir", str(wallet_dir)) + config.setdefault("wallet_location", str(wallet_dir)) + if "topology" in config: + topology_val = config.pop("topology") + resolved_topology = OracleTopologyConfig(**topology_val) if isinstance(topology_val, Mapping) else topology_val + if "dsn" not in config: + config["dsn"] = build_tns_description(resolved_topology) + if resolved_topology.retry_count is not None and "retry_count" not in config: + config["retry_count"] = resolved_topology.retry_count + if resolved_topology.retry_delay is not None and "retry_delay" not in config: + config["retry_delay"] = resolved_topology.retry_delay if "dsn" not in config: dsn = config.pop("url", None) or config.pop("connection_string", None) if dsn is not None: @@ -1200,6 +2118,154 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str return config +def _extract_wallet_archive_into(wallet_zip: "Path | str | bytes", extract_dir: Path, resolved_root: Path) -> None: + if isinstance(wallet_zip, bytes): + zip_source: Any = io.BytesIO(wallet_zip) + else: + zip_path = Path(wallet_zip) + if not zip_path.is_file(): + msg = f"Oracle wallet archive does not exist: {zip_path}" + raise ImproperConfigurationError(msg) + zip_source = zip_path + + try: + with zipfile.ZipFile(zip_source, "r") as archive: + for info in archive.infolist(): + raw_name = info.filename + member_path = Path(raw_name) + if ( + member_path.is_absolute() + or ".." in member_path.parts + or raw_name.startswith(("/", "\\")) + or ":" in raw_name + ): + msg = f"Unsafe path in Oracle wallet archive: {raw_name!r}" + raise ImproperConfigurationError(msg) + dest_path = (resolved_root / member_path).resolve() + if dest_path != resolved_root and resolved_root not in dest_path.parents: + msg = f"Unsafe path in Oracle wallet archive: {raw_name!r}" + raise ImproperConfigurationError(msg) + + for info in archive.infolist(): + dest_path = (resolved_root / Path(info.filename)).resolve() + if info.is_dir() or info.filename.endswith("/"): + dest_path.mkdir(parents=True, exist_ok=True) + dest_path.chmod(0o700) + continue + dest_path.parent.mkdir(parents=True, exist_ok=True) + dest_path.parent.chmod(0o700) + with archive.open(info, "r") as src_file, dest_path.open("wb") as dst_file: + shutil.copyfileobj(src_file, dst_file) + dest_path.chmod(0o600) + except zipfile.BadZipFile as exc: + msg = f"Invalid Oracle wallet zip archive: {exc}" + raise ImproperConfigurationError(msg) from exc + + if not (extract_dir / "ewallet.pem").exists() and not (extract_dir / "tnsnames.ora").exists(): + msg = "Oracle wallet archive must contain 'ewallet.pem' or 'tnsnames.ora'" + raise ImproperConfigurationError(msg) + + +def _validate_tns_token(value: str, field_name: str) -> None: + if any(ch in value for ch in ("(", ")", "=", "'", '"', "\r", "\n", "\x00")): + msg = f"Unsafe TNS descriptor character in {field_name}: {value!r}" + raise ImproperConfigurationError(msg) + + +def _validate_oracle_identifier(value: str, field_name: str) -> None: + if not value or not _UNQUOTED_ORACLE_IDENTIFIER_PATTERN.fullmatch(value): + msg = f"Invalid Oracle identifier for {field_name}: {value!r}" + raise ImproperConfigurationError(msg) + + +def _validate_oracle_time_zone(value: str) -> None: + if not value or not _ORACLE_TIME_ZONE_PATTERN.fullmatch(value) or "'" in value or ";" in value: + msg = f"Invalid Oracle session time_zone: {value!r}" + raise ImproperConfigurationError(msg) + + +def _validate_oracle_nls_parameters(nls_parameters: "Mapping[str, str]") -> None: + for key, val in nls_parameters.items(): + if not isinstance(key, str) or not _ORACLE_NLS_KEY_PATTERN.fullmatch(key.upper()): + msg = f"Invalid Oracle NLS parameter name: {key!r}" + raise ImproperConfigurationError(msg) + if not isinstance(val, str) or "'" in val or ";" in val or "--" in val or "/*" in val or "\n" in val: + msg = f"Unsafe Oracle NLS parameter value for {key!r}: {val!r}" + raise ImproperConfigurationError(msg) + + +def _validate_oracle_contexts(contexts: "Mapping[str, Mapping[str, Any]]") -> None: + for namespace, attrs in contexts.items(): + _validate_oracle_identifier(namespace, "context namespace") + if not isinstance(attrs, Mapping): + msg = f"Oracle context attributes for namespace {namespace!r} must be a mapping" + raise ImproperConfigurationError(msg) + for attr_key in attrs: + _validate_oracle_identifier(attr_key, "context attribute") + + +def _resolve_callable_or_str(value: "str | Callable[[], str | None] | None") -> str | None: + if callable(value): + return value() + return value + + +class _OracleSyncStreamDriver(Protocol): + connection: "OracleSyncConnection" + driver_features: dict[str, Any] + + def handle_database_exceptions(self) -> "SyncExceptionHandler": ... + + def _check_pending_exception(self, exc_handler: "SyncExceptionHandler") -> None: ... + + def _resolve_row_metadata(self, description: object) -> tuple[list[str], bool]: ... + + +class _OracleAsyncStreamDriver(Protocol): + connection: "OracleAsyncConnection" + driver_features: dict[str, Any] + + def handle_database_exceptions(self) -> "AsyncExceptionHandler": ... + + def _check_pending_exception(self, exc_handler: "AsyncExceptionHandler") -> None: ... + + async def _run_with_exception_handler( + self, + exc_handler: "AsyncExceptionHandler", + operation: "Callable[..., Awaitable[Any]]", + *args: Any, + **kwargs: Any, + ) -> Any: ... + + def _resolve_row_metadata(self, description: object) -> tuple[list[str], bool]: ... + + +class _OracleJsonBindingState: + """Lazily resolve whether JSON-shaped values require BLOB locators.""" + + __slots__ = ("_connection", "_resolved", "_use_blob", "_version_cache") + + def __init__(self, connection: Any, version_cache: Any) -> None: + self._connection = connection + self._version_cache = version_cache + self._resolved = False + self._use_blob = False + + def uses_blob(self) -> bool: + """Return the operation-scoped JSON storage decision.""" + if not self._resolved: + server_major = resolve_oracle_connection_major(self._connection, self._version_cache) + compatible_major = resolve_oracle_connection_compatible_major(self._connection, self._version_cache) + self._use_blob = ( + isinstance(server_major, int) + and not isinstance(server_major, bool) + and resolve_oracle_json_storage(server_major, compatible_major=compatible_major) + != ORACLE_JSON_STORAGE_NATIVE + ) + self._resolved = True + return self._use_blob + + def _parse_version_tuple(version: str) -> "tuple[int, int, int]": parts = [int(part) for part in version.split(".") if part.isdigit()] while len(parts) < _VERSION_COMPONENTS: diff --git a/sqlspec/adapters/oracledb/data_dictionary.py b/sqlspec/adapters/oracledb/data_dictionary.py index 415ef42f3..4808db035 100644 --- a/sqlspec/adapters/oracledb/data_dictionary.py +++ b/sqlspec/adapters/oracledb/data_dictionary.py @@ -536,6 +536,55 @@ def get_foreign_keys( schema_type=ForeignKeyMetadata, ) + def get_vector_memory_pool(self, driver: "OracleSyncDriver") -> "list[dict[str, Any]]": + """Return V$VECTOR_MEMORY_POOL allocation rows or empty list when unavailable.""" + try: + return driver.select(self.get_query("system", "vector_memory_pool")) + except Exception as error: + if _is_missing_or_inaccessible_oracle_object(error): + return [] + raise + + def has_cloud_ai(self, driver: "OracleSyncDriver") -> bool: + """Return True when DBMS_CLOUD_AI is installed and accessible.""" + try: + rows = driver.select(self.get_query("objects", "cloud_ai_available")) + except Exception as error: + if _is_missing_or_inaccessible_oracle_object(error): + return False + raise + if not rows: + return False + first = rows[0] + if isinstance(first, dict): + if "cnt" in first or "CNT" in first: + return int(first.get("cnt", 0) or first.get("CNT", 0) or 0) > 0 + return bool(first) + return bool(first) + + def get_apex_release(self, driver: "OracleSyncDriver") -> "dict[str, Any] | None": + """Return APEX_RELEASE metadata row or None when APEX is not installed.""" + try: + rows = driver.select(self.get_query("version", "apex_release")) + except Exception as error: + if _is_missing_or_inaccessible_oracle_object(error): + return None + raise + return rows[0] if rows else None + + def get_json_duality_views(self, driver: "OracleSyncDriver", schema: "str | None" = None) -> "list[dict[str, Any]]": + """Return Oracle 26ai JSON Relational Duality View catalog rows.""" + try: + if schema is None: + return driver.select(self.get_query("views", "json_duality_views")) + return driver.select( + self.get_query("views", "json_duality_views_by_schema"), schema_name=self.resolve_schema(schema) + ) + except Exception as error: + if _is_missing_or_inaccessible_oracle_object(error): + return [] + raise + def _build_version_info( self, version_value: "str | None", compatible: "str | None", is_autonomous: bool ) -> "OracleVersionInfo | None": @@ -910,6 +959,57 @@ async def get_foreign_keys( schema_type=ForeignKeyMetadata, ) + async def get_vector_memory_pool(self, driver: "OracleAsyncDriver") -> "list[dict[str, Any]]": + """Return V$VECTOR_MEMORY_POOL allocation rows or empty list when unavailable.""" + try: + return await driver.select(self.get_query("system", "vector_memory_pool")) + except Exception as error: + if _is_missing_or_inaccessible_oracle_object(error): + return [] + raise + + async def has_cloud_ai(self, driver: "OracleAsyncDriver") -> bool: + """Return True when DBMS_CLOUD_AI is installed and accessible.""" + try: + rows = await driver.select(self.get_query("objects", "cloud_ai_available")) + except Exception as error: + if _is_missing_or_inaccessible_oracle_object(error): + return False + raise + if not rows: + return False + first = rows[0] + if isinstance(first, dict): + if "cnt" in first or "CNT" in first: + return int(first.get("cnt", 0) or first.get("CNT", 0) or 0) > 0 + return bool(first) + return bool(first) + + async def get_apex_release(self, driver: "OracleAsyncDriver") -> "dict[str, Any] | None": + """Return APEX_RELEASE metadata row or None when APEX is not installed.""" + try: + rows = await driver.select(self.get_query("version", "apex_release")) + except Exception as error: + if _is_missing_or_inaccessible_oracle_object(error): + return None + raise + return rows[0] if rows else None + + async def get_json_duality_views( + self, driver: "OracleAsyncDriver", schema: "str | None" = None + ) -> "list[dict[str, Any]]": + """Return Oracle 26ai JSON Relational Duality View catalog rows.""" + try: + if schema is None: + return await driver.select(self.get_query("views", "json_duality_views")) + return await driver.select( + self.get_query("views", "json_duality_views_by_schema"), schema_name=self.resolve_schema(schema) + ) + except Exception as error: + if _is_missing_or_inaccessible_oracle_object(error): + return [] + raise + def _build_version_info( self, version_value: "str | None", compatible: "str | None", is_autonomous: bool ) -> "OracleVersionInfo | None": @@ -1150,6 +1250,18 @@ def _oracle_system_metadata_request( return ensure_system_metadata_request(request, **kwargs) +def _is_missing_or_inaccessible_oracle_object(error: BaseException) -> bool: + """Return True when an error corresponds to ORA-00942 or ORA-01031.""" + msg = str(error).upper() + if "ORA-00942" in msg or "ORA-01031" in msg: + return True + if error.args: + code = getattr(error.args[0], "code", None) + if code in {942, 1031}: + return True + return False + + def _storage_type_from_version(version_info: "OracleVersionInfo | None") -> JSONStorageType: """Determine the JSON storage rung for an Oracle version. diff --git a/sqlspec/adapters/oracledb/driver.py b/sqlspec/adapters/oracledb/driver.py index 288fd433e..592b27e4d 100644 --- a/sqlspec/adapters/oracledb/driver.py +++ b/sqlspec/adapters/oracledb/driver.py @@ -29,6 +29,7 @@ coerce_large_parameters_sync, coerce_many_parameters_async, coerce_many_parameters_sync, + coerce_oracle_df_to_arrow_table, collect_async_rows, collect_sync_rows, connection_is_thin, @@ -45,6 +46,7 @@ OracledbAsyncDataDictionary, OracledbSyncDataDictionary, OracleVersionCache, + resolve_oracle_connection_major, ) from sqlspec.core import ( SQL, @@ -56,7 +58,12 @@ get_cache_config, register_driver_profile, ) -from sqlspec.core.explain import ORACLE_EXPLAIN_PREFIX, ORACLE_MANAGED_EXPLAIN_META_KEY +from sqlspec.core.explain import ( + ORACLE_EXPLAIN_DISPLAY_CURSOR_META_KEY, + ORACLE_EXPLAIN_PREFIX, + ORACLE_EXPLAIN_XPLAN_FORMAT_META_KEY, + ORACLE_MANAGED_EXPLAIN_META_KEY, +) from sqlspec.driver import ( AsyncDriverAdapterBase, AsyncRowStream, @@ -394,7 +401,10 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": explained_sql = _managed_explain_statement(statement, sql) if explained_sql is not None: - return self._dispatch_explain_plan(cursor, explained_sql, prepared_parameters) + display_cursor, xplan_format = _managed_explain_options(statement) + return self._dispatch_explain_plan( + cursor, explained_sql, prepared_parameters, display_cursor=display_cursor, xplan_format=xplan_format + ) cursor.execute(sql, prepared_parameters or {}, **build_fetch_kwargs(self.driver_features)) @@ -722,7 +732,7 @@ def select_to_arrow( arrow_schema=arrow_schema, ) - arrow_table = pa.table(oracle_df) + arrow_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) column_names = normalize_column_names(arrow_table.column_names, self.driver_features) if column_names != arrow_table.column_names: arrow_table = arrow_table.rename_columns(column_names) @@ -881,27 +891,63 @@ def _detect_oracledb_version(self) -> "tuple[int, int, int]": def _resolve_row_metadata(self, description: Any) -> "tuple[list[str], bool]": return resolve_row_metadata(description, self.driver_features, self._row_metadata_cache) - def _dispatch_explain_plan(self, cursor: Any, explained_sql: str, parameters: Any) -> "ExecutionResult": + def _dispatch_explain_plan( + self, + cursor: Any, + explained_sql: str, + parameters: Any, + *, + display_cursor: bool = False, + xplan_format: str = "TYPICAL", + ) -> "ExecutionResult": """Request a plan, read it back, and release the plan table rows. Oracle writes the plan to a table instead of returning it, so the plan is tagged with a generated statement id, formatted through ``DBMS_XPLAN``, and deleted again. All three statements run in the caller's transaction, so the - plan is read before a rollback can discard it. + plan is read before a rollback can discard it. When ``display_cursor`` is + enabled, the statement is executed first and formatted through + ``DBMS_XPLAN.DISPLAY_CURSOR``. Args: cursor: Oracle cursor object. explained_sql: Statement whose plan is wanted. parameters: Bind values for the explained statement. + display_cursor: Whether to inspect actual cursor execution statistics via ``DBMS_XPLAN.DISPLAY_CURSOR``. + xplan_format: ``DBMS_XPLAN`` format string. Returns: Execution result carrying the formatted plan rows. """ + if display_cursor: + cursor.execute(explained_sql, parameters or {}, **build_fetch_kwargs(self.driver_features)) + if cursor.description is not None: + cursor.fetchall() + cursor.execute(_format_display_cursor_sql(xplan_format), {}, **build_fetch_kwargs(self.driver_features)) + fetched_data = cursor.fetchall() + column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) + data, column_names = collect_sync_rows( + cast("list[Any] | None", fetched_data), + cursor.description, + self.driver_features, + column_names=column_names, + requires_lob_coercion=requires_lob_coercion, + ) + return self.create_execution_result( + cursor, + selected_data=data, + column_names=column_names, + data_row_count=len(data), + is_select_result=True, + row_format="tuple", + ) + statement_id = _new_plan_statement_id() cursor.execute(_tagged_explain_plan_sql(statement_id, explained_sql), parameters or {}) + display_sql = PLAN_DISPLAY_SQL if xplan_format == "TYPICAL" else _format_plan_display_sql(xplan_format) try: - cursor.execute(PLAN_DISPLAY_SQL, {"statement_id": statement_id}, **build_fetch_kwargs(self.driver_features)) + cursor.execute(display_sql, {"statement_id": statement_id}, **build_fetch_kwargs(self.driver_features)) fetched_data = cursor.fetchall() column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) data, column_names = collect_sync_rows( @@ -977,7 +1023,7 @@ def _fetch_arrow_record_batches( statement=sql, parameters=params, size=batch_size or 1000, **retry_kwargs ) for oracle_df in oracle_dfs: - batch_table = pa.table(oracle_df) + batch_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) column_names = normalize_column_names(batch_table.column_names, self.driver_features) if column_names != batch_table.column_names: batch_table = batch_table.rename_columns(column_names) @@ -1079,7 +1125,10 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes explained_sql = _managed_explain_statement(statement, sql) if explained_sql is not None: - return await self._dispatch_explain_plan(cursor, explained_sql, prepared_parameters) + display_cursor, xplan_format = _managed_explain_options(statement) + return await self._dispatch_explain_plan( + cursor, explained_sql, prepared_parameters, display_cursor=display_cursor, xplan_format=xplan_format + ) await cursor.execute(sql, prepared_parameters or {}, **build_fetch_kwargs(self.driver_features)) @@ -1412,7 +1461,7 @@ async def select_to_arrow( arrow_schema=arrow_schema, ) - arrow_table = pa.table(oracle_df) + arrow_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) column_names = normalize_column_names(arrow_table.column_names, self.driver_features) if column_names != arrow_table.column_names: arrow_table = arrow_table.rename_columns(column_names) @@ -1580,28 +1629,66 @@ def _detect_oracledb_version(self) -> "tuple[int, int, int]": def _resolve_row_metadata(self, description: Any) -> "tuple[list[str], bool]": return resolve_row_metadata(description, self.driver_features, self._row_metadata_cache) - async def _dispatch_explain_plan(self, cursor: Any, explained_sql: str, parameters: Any) -> "ExecutionResult": + async def _dispatch_explain_plan( + self, + cursor: Any, + explained_sql: str, + parameters: Any, + *, + display_cursor: bool = False, + xplan_format: str = "TYPICAL", + ) -> "ExecutionResult": """Request a plan, read it back, and release the plan table rows. Oracle writes the plan to a table instead of returning it, so the plan is tagged with a generated statement id, formatted through ``DBMS_XPLAN``, and deleted again. All three statements run in the caller's transaction, so the - plan is read before a rollback can discard it. + plan is read before a rollback can discard it. When ``display_cursor`` is + enabled, the statement is executed first and formatted through + ``DBMS_XPLAN.DISPLAY_CURSOR``. Args: cursor: Oracle cursor object. explained_sql: Statement whose plan is wanted. parameters: Bind values for the explained statement. + display_cursor: Whether to inspect actual cursor execution statistics via ``DBMS_XPLAN.DISPLAY_CURSOR``. + xplan_format: ``DBMS_XPLAN`` format string. Returns: Execution result carrying the formatted plan rows. """ + if display_cursor: + await cursor.execute(explained_sql, parameters or {}, **build_fetch_kwargs(self.driver_features)) + if cursor.description is not None: + await cursor.fetchall() + await cursor.execute( + _format_display_cursor_sql(xplan_format), {}, **build_fetch_kwargs(self.driver_features) + ) + fetched_data = await cursor.fetchall() + column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) + data, column_names = await collect_async_rows( + cast("list[Any] | None", fetched_data), + cursor.description, + self.driver_features, + column_names=column_names, + requires_lob_coercion=requires_lob_coercion, + ) + return self.create_execution_result( + cursor, + selected_data=data, + column_names=column_names, + data_row_count=len(data), + is_select_result=True, + row_format="tuple", + ) + statement_id = _new_plan_statement_id() await cursor.execute(_tagged_explain_plan_sql(statement_id, explained_sql), parameters or {}) + display_sql = PLAN_DISPLAY_SQL if xplan_format == "TYPICAL" else _format_plan_display_sql(xplan_format) try: await cursor.execute( - PLAN_DISPLAY_SQL, {"statement_id": statement_id}, **build_fetch_kwargs(self.driver_features) + display_sql, {"statement_id": statement_id}, **build_fetch_kwargs(self.driver_features) ) fetched_data = await cursor.fetchall() column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) @@ -1678,7 +1765,7 @@ async def _fetch_arrow_record_batches( statement=sql, parameters=params, size=batch_size or 1000, **retry_kwargs ) async for oracle_df in oracle_dfs: - batch_table = pa.table(oracle_df) + batch_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) column_names = normalize_column_names(batch_table.column_names, self.driver_features) if column_names != batch_table.column_names: batch_table = batch_table.rename_columns(column_names) @@ -1774,6 +1861,16 @@ async def _pipeline_native_supported(self) -> bool: self._pipeline_support_reason = "thin_mode_required" return False + connection_major = resolve_oracle_connection_major(self.connection, self._oracle_version_cache) + if connection_major is not None: + if connection_major >= PIPELINE_MIN_DATABASE_MAJOR: + self._pipeline_support = True + self._pipeline_support_reason = None + return True + self._pipeline_support = False + self._pipeline_support_reason = "database_version" + return False + version_info = await self.data_dictionary.get_version(self) if version_info and version_info.major >= PIPELINE_MIN_DATABASE_MAJOR: self._pipeline_support = True @@ -1901,6 +1998,33 @@ def _tagged_explain_plan_sql(statement_id: str, explained_sql: str) -> str: return f"EXPLAIN PLAN SET STATEMENT_ID = '{statement_id}' FOR {explained_sql}" +def _format_plan_display_sql(xplan_format: str = "TYPICAL") -> str: + """Build a DBMS_XPLAN.DISPLAY query with a validated format string.""" + if "'" in xplan_format: + msg = f"Invalid Oracle XPLAN format string: {xplan_format!r}" + raise SQLSpecError(msg) + return f"SELECT plan_table_output FROM TABLE(DBMS_XPLAN.DISPLAY(NULL, :statement_id, '{xplan_format}'))" + + +def _format_display_cursor_sql(xplan_format: str = "TYPICAL") -> str: + """Build a DBMS_XPLAN.DISPLAY_CURSOR query with a validated format string.""" + if "'" in xplan_format: + msg = f"Invalid Oracle XPLAN format string: {xplan_format!r}" + raise SQLSpecError(msg) + return f"SELECT plan_table_output FROM TABLE(DBMS_XPLAN.DISPLAY_CURSOR(NULL, NULL, '{xplan_format}'))" + + +def _managed_explain_options(statement: "SQL") -> "tuple[bool, str]": + """Extract Oracle display_cursor and xplan_format options from a managed Explain statement.""" + expression = statement.raw_expression + if expression is None: + return False, "TYPICAL" + display_cursor = bool(expression.meta.get(ORACLE_EXPLAIN_DISPLAY_CURSOR_META_KEY, False)) + raw_format = expression.meta.get(ORACLE_EXPLAIN_XPLAN_FORMAT_META_KEY) + xplan_format = str(raw_format) if isinstance(raw_format, str) and raw_format else "TYPICAL" + return display_cursor, xplan_format + + def _retry_arrow_without_fetch_lobs(exc: TypeError, fetch_kwargs: "dict[str, object]") -> "dict[str, object] | None": if "fetch_lobs" not in fetch_kwargs or "fetch_lobs" not in str(exc): return None diff --git a/sqlspec/builder/__init__.py b/sqlspec/builder/__init__.py index aa481b83b..985b95847 100644 --- a/sqlspec/builder/__init__.py +++ b/sqlspec/builder/__init__.py @@ -52,6 +52,7 @@ build_explain_sql, build_generic_explain, build_mysql_explain, + build_oracle_display_cursor, build_oracle_explain, build_postgres_explain, build_sqlite_explain, @@ -102,7 +103,12 @@ from sqlspec.builder._temporal import create_temporal_table, register_version_generators from sqlspec.builder._update import Update from sqlspec.builder._values import Values - from sqlspec.builder._vector_distance import VectorDistance + from sqlspec.builder._vector_distance import ( + VectorDistance, + build_oracle_hnsw_subquery, + build_oracle_hybrid_rrf_query, + build_oracle_hybrid_rrf_sql, + ) from sqlspec.exceptions import SQLBuilderError __all__ = ( @@ -179,7 +185,11 @@ "build_explain_sql", "build_generic_explain", "build_mysql_explain", + "build_oracle_display_cursor", "build_oracle_explain", + "build_oracle_hnsw_subquery", + "build_oracle_hybrid_rrf_query", + "build_oracle_hybrid_rrf_sql", "build_postgres_explain", "build_sqlite_explain", "create_temporal_table", @@ -268,7 +278,11 @@ "build_explain_sql": ("sqlspec.builder._explain", "build_explain_sql"), "build_generic_explain": ("sqlspec.builder._explain", "build_generic_explain"), "build_mysql_explain": ("sqlspec.builder._explain", "build_mysql_explain"), + "build_oracle_display_cursor": ("sqlspec.builder._explain", "build_oracle_display_cursor"), "build_oracle_explain": ("sqlspec.builder._explain", "build_oracle_explain"), + "build_oracle_hnsw_subquery": ("sqlspec.builder._vector_distance", "build_oracle_hnsw_subquery"), + "build_oracle_hybrid_rrf_query": ("sqlspec.builder._vector_distance", "build_oracle_hybrid_rrf_query"), + "build_oracle_hybrid_rrf_sql": ("sqlspec.builder._vector_distance", "build_oracle_hybrid_rrf_sql"), "build_postgres_explain": ("sqlspec.builder._explain", "build_postgres_explain"), "build_sqlite_explain": ("sqlspec.builder._explain", "build_sqlite_explain"), "create_temporal_table": ("sqlspec.builder._temporal", "create_temporal_table"), diff --git a/sqlspec/builder/_explain.py b/sqlspec/builder/_explain.py index 61d9056c0..d46d9ab95 100644 --- a/sqlspec/builder/_explain.py +++ b/sqlspec/builder/_explain.py @@ -10,7 +10,14 @@ from sqlglot import Dialect, exp from sqlspec.core import SQL, StatementConfig -from sqlspec.core.explain import ORACLE_EXPLAIN_PREFIX, ORACLE_MANAGED_EXPLAIN_META_KEY, ExplainFormat, ExplainOptions +from sqlspec.core.explain import ( + ORACLE_EXPLAIN_DISPLAY_CURSOR_META_KEY, + ORACLE_EXPLAIN_PREFIX, + ORACLE_EXPLAIN_XPLAN_FORMAT_META_KEY, + ORACLE_MANAGED_EXPLAIN_META_KEY, + ExplainFormat, + ExplainOptions, +) from sqlspec.exceptions import SQLBuilderError from sqlspec.utils.type_guards import has_expression_and_sql, has_parameter_builder, is_expression @@ -28,6 +35,7 @@ "build_explain_sql", "build_generic_explain", "build_mysql_explain", + "build_oracle_display_cursor", "build_oracle_explain", "build_postgres_explain", "build_sqlite_explain", @@ -174,6 +182,30 @@ def build_duckdb_explain(statement_sql: str, options: "ExplainOptions") -> str: return f"EXPLAIN {statement_sql}" +def build_oracle_display_cursor( + sql_id: str | None = None, cursor_child_no: int | None = None, format: str = "TYPICAL" +) -> str: + """Build an Oracle DBMS_XPLAN.DISPLAY_CURSOR query. + + Args: + sql_id: Optional SQL_ID to inspect (defaults to NULL for last statement in session). + cursor_child_no: Optional cursor child number (defaults to NULL). + format: DBMS_XPLAN format string (e.g. 'TYPICAL +PREDICATE', 'ALLSTATS LAST +PEEKED_BINDS'). + + Returns: + SELECT statement querying TABLE(DBMS_XPLAN.DISPLAY_CURSOR(...)). + """ + if "'" in format: + msg = f"Invalid Oracle XPLAN format string: {format!r}" + raise SQLBuilderError(msg) + if sql_id is not None and "'" in sql_id: + msg = f"Invalid Oracle sql_id: {sql_id!r}" + raise SQLBuilderError(msg) + sql_id_arg = f"'{sql_id}'" if sql_id is not None else "NULL" + child_arg = str(cursor_child_no) if cursor_child_no is not None else "NULL" + return f"SELECT plan_table_output FROM TABLE(DBMS_XPLAN.DISPLAY_CURSOR({sql_id_arg}, {child_arg}, '{format}'))" + + def build_oracle_explain(statement_sql: str, options: "ExplainOptions") -> str: """Build Oracle EXPLAIN statement. @@ -439,6 +471,33 @@ def generic_plan(self, enabled: bool = True) -> Self: self._options = self._options.copy(generic_plan=enabled) return self + def display_cursor(self, enabled: bool = True) -> Self: + """Enable Oracle DBMS_XPLAN.DISPLAY_CURSOR post-execution plan inspection. + + Args: + enabled: Whether to use DISPLAY_CURSOR instead of EXPLAIN PLAN FOR. + + Returns: + Self for method chaining + """ + self._options = self._options.copy(display_cursor=enabled) + return self + + def xplan_format(self, fmt: str) -> Self: + """Set Oracle DBMS_XPLAN format string (e.g. 'TYPICAL +PREDICATE +NOTE'). + + Args: + fmt: Oracle DBMS_XPLAN format string. + + Returns: + Self for method chaining + """ + if "'" in fmt: + msg = f"Invalid Oracle XPLAN format string: {fmt!r}" + raise SQLBuilderError(msg) + self._options = self._options.copy(xplan_format=fmt) + return self + def with_options(self, options: "ExplainOptions") -> Self: """Replace all options with the provided ExplainOptions. @@ -465,6 +524,12 @@ def build(self, dialect: "DialectType | None" = None) -> "SQL": if normalize_dialect_name(target_dialect) in ORACLE_DIALECTS: oracle_statement = exp.Command(this="EXPLAIN PLAN FOR", expression=exp.var(self._statement_sql)) oracle_statement.meta[ORACLE_MANAGED_EXPLAIN_META_KEY] = True + if self._options.display_cursor: + oracle_statement.meta[ORACLE_EXPLAIN_DISPLAY_CURSOR_META_KEY] = True + if self._options.xplan_format is not None: + oracle_statement.meta[ORACLE_EXPLAIN_XPLAN_FORMAT_META_KEY] = self._options.xplan_format + elif self._options.analyze and self._options.display_cursor: + oracle_statement.meta[ORACLE_EXPLAIN_XPLAN_FORMAT_META_KEY] = "ALLSTATS LAST" statement: str | exp.Expr = oracle_statement else: statement = build_explain_sql(self._statement_sql, self._options, target_dialect) diff --git a/sqlspec/builder/_select.py b/sqlspec/builder/_select.py index 42d934dfb..66d550156 100644 --- a/sqlspec/builder/_select.py +++ b/sqlspec/builder/_select.py @@ -353,6 +353,13 @@ def limit(self, value: int) -> Self: builder.set_expression(select_expr.limit(exp.convert(value), copy=False)) return cast("Self", builder) + def approx_limit(self, value: int, *, target_accuracy: int | None = None) -> Self: + self.limit(value) + hints = getattr(self, "_hints", None) + if isinstance(hints, list): + hints.append({"location": "approx_limit", "target_accuracy": target_accuracy}) + return self + def offset(self, value: int) -> Self: builder = cast("SQLBuilderProtocol", self) select_expr = _ensure_select_expression(builder, error_message="OFFSET only valid for SELECT.") @@ -1248,6 +1255,17 @@ def _postprocess_built_sql(self, sql_string: str, dialect: "DialectType" = None) target_dialect = _normalize_dialect(dialect or self.dialect) modified_sql = sql_string for hint_dict in self._hints: + if hint_dict.get("location") == "approx_limit": + acc = hint_dict.get("target_accuracy") + replacement = ( + r"FETCH APPROX FIRST \1 ROWS ONLY" + if acc is None + else rf"FETCH APPROX FIRST \1 ROWS ONLY WITH TARGET ACCURACY {int(cast('int', acc))}" + ) + modified_sql = re.sub( + r"\bFETCH\s+FIRST\s+(\S+)\s+ROWS\s+ONLY\b", replacement, modified_sql, count=1, flags=re.IGNORECASE + ) + continue if ( hint_dict.get("location") == "table" and hint_dict.get("table") diff --git a/sqlspec/builder/_vector_distance.py b/sqlspec/builder/_vector_distance.py index f57304a0c..4cfe1f000 100644 --- a/sqlspec/builder/_vector_distance.py +++ b/sqlspec/builder/_vector_distance.py @@ -2,7 +2,7 @@ # ruff: noqa: N802 # pyright: ignore[reportConstantRedefinition] -from collections.abc import Callable, MutableMapping +from collections.abc import Callable, MutableMapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast from sqlglot import exp @@ -14,6 +14,9 @@ __all__ = ( "VectorDistance", + "build_oracle_hnsw_subquery", + "build_oracle_hybrid_rrf_query", + "build_oracle_hybrid_rrf_sql", "has_vector_distance_ancestor", "is_vector_distance_expression", "render_vector_distance_bigquery", @@ -54,6 +57,10 @@ _DUCKDB_OPERATOR_TRANSFORM: _OperatorTransform | None = None +_ORACLE_MIN_ACCURACY: Final[int] = 1 +_ORACLE_MAX_ACCURACY: Final[int] = 100 + + def is_vector_distance_expression(expression: object) -> bool: """Return True when an Operator node is a SQLSpec vector-distance expression.""" return isinstance(expression, exp.Operator) and _VECTOR_DISTANCE_META_KEY in expression.meta @@ -154,6 +161,183 @@ def render_vector_distance_generic(left: str, right: str, metric: str) -> str: return f"VECTOR_DISTANCE({left}, {right}, '{metric.upper()}')" +def build_oracle_hnsw_subquery( + table: str, + vector_column: str = "embedding", + vector_param: str = ":query_vec", + *, + id_column: str = "id", + select_columns: Sequence[str] = ("id",), + metric: str = "COSINE", + top_k_param: str = ":top_k", + query_vector_bind: str | None = None, + top_k_bind: str | None = None, + target_accuracy: int | None = None, + where_clause: str | None = None, + rank_alias: str = "vector_rank", + distance_alias: str = "vector_distance", +) -> str: + """Build an Oracle 26ai HNSW-preserving top-K subquery using outer ``ROWNUM``. + + Oracle 26ai CBO disables ``VECTOR INDEX HNSW SCAN`` when ``ROW_NUMBER() OVER (...)`` + appears in the same query block as ``ORDER BY VECTOR_DISTANCE(...) FETCH APPROX FIRST :k ROWS ONLY``. + This helper isolates the ``ORDER BY VECTOR_DISTANCE(...) FETCH APPROX FIRST`` clause inside + an inner derived table and assigns ``ROWNUM`` in the outer query block. + """ + oracle_metric = _ORACLE_METRIC_MAP.get(metric.lower(), metric.upper()) + vec_bind = _format_oracle_bind(query_vector_bind if query_vector_bind is not None else vector_param) + limit_bind = _format_oracle_bind(top_k_bind if top_k_bind is not None else top_k_param) + cols = (id_column,) if tuple(select_columns) == ("id",) and id_column != "id" else tuple(select_columns) + cols_sql = ", ".join(cols) + where_sql = f" WHERE {where_clause}" if where_clause else "" + accuracy_sql = "" + if target_accuracy is not None: + if ( + isinstance(target_accuracy, bool) + or not isinstance(target_accuracy, int) + or not (_ORACLE_MIN_ACCURACY <= target_accuracy <= _ORACLE_MAX_ACCURACY) + ): + msg = f"target_accuracy must be an integer between 1 and 100, got {target_accuracy!r}" + raise ValueError(msg) + accuracy_sql = f" WITH TARGET ACCURACY {target_accuracy}" + return ( + f"SELECT {cols_sql}, {distance_alias}, ROWNUM AS {rank_alias} " + f"FROM (" + f"SELECT {cols_sql}, VECTOR_DISTANCE({vector_column}, {vec_bind}, {oracle_metric}) AS {distance_alias} " + f"FROM {table}{where_sql} " + f"ORDER BY VECTOR_DISTANCE({vector_column}, {vec_bind}, {oracle_metric}) " + f"FETCH APPROX FIRST {limit_bind} ROWS ONLY{accuracy_sql}" + f") vec_top" + ) + + +def build_oracle_hybrid_rrf_sql( + table: str, + vector_column: str = "embedding", + text_column: str = "content_text", + *, + id_column: str = "id", + key_column: str | None = None, + select_columns: Sequence[str] | None = None, + vector_param: str = ":query_vec", + text_param: str = ":fts_query", + metric: str = "COSINE", + rrf_k: int = 60, + rrf_k_param: str | None = None, + candidate_limit_param: str = ":candidate_limit", + final_limit_param: str = ":final_limit", + top_k_param: str | None = None, + target_accuracy: int | None = None, + where_clause: str | None = None, +) -> str: + """Build an Oracle 26ai Hybrid Search query fusing HNSW vector similarity and Oracle Text FTS via RRF.""" + resolved_id_col = key_column if key_column is not None else id_column + if top_k_param is not None: + if candidate_limit_param == ":candidate_limit": + candidate_limit_param = top_k_param + if final_limit_param == ":final_limit": + final_limit_param = top_k_param + cand_bind = _format_oracle_bind(candidate_limit_param) + final_bind = _format_oracle_bind(final_limit_param) + txt_bind = _format_oracle_bind(text_param) + rrf_k_expr = _format_oracle_bind(rrf_k_param) if rrf_k_param is not None else str(rrf_k) + + vec_subquery = build_oracle_hnsw_subquery( + table=table, + vector_column=vector_column, + vector_param=vector_param, + id_column=resolved_id_col, + select_columns=(resolved_id_col,), + metric=metric, + top_k_param=cand_bind, + target_accuracy=target_accuracy, + where_clause=where_clause, + rank_alias="vector_rank", + distance_alias="vector_distance", + ) + fts_where_parts = [f"CONTAINS({text_column}, {txt_bind}, 1) > 0"] + if where_clause: + fts_where_parts.append(f"({where_clause})") + fts_where_sql = " AND ".join(fts_where_parts) + fts_subquery = ( + f"SELECT {resolved_id_col}, fts_score, ROWNUM AS fts_rank " + f"FROM (" + f"SELECT {resolved_id_col}, SCORE(1) AS fts_score " + f"FROM {table} " + f"WHERE {fts_where_sql} " + f"ORDER BY SCORE(1) DESC " + f"FETCH FIRST {cand_bind} ROWS ONLY" + f") fts_top" + ) + + if select_columns: + base_cols = ", ".join(col if "." in col or col == "*" else f"base.{col}" for col in select_columns) + else: + base_cols = "base.*" + + return ( + f"WITH vec_candidates AS ({vec_subquery}), " + f"fts_candidates AS ({fts_subquery}), " + f"rrf_fused AS (" + f"SELECT COALESCE(v.{resolved_id_col}, f.{resolved_id_col}) AS {resolved_id_col}, " + f"v.vector_distance, v.vector_rank, f.fts_score, f.fts_rank, " + f"(NVL(1.0 / ({rrf_k_expr} + v.vector_rank), 0) + NVL(1.0 / ({rrf_k_expr} + f.fts_rank), 0)) AS rrf_score " + f"FROM vec_candidates v " + f"FULL OUTER JOIN fts_candidates f ON v.{resolved_id_col} = f.{resolved_id_col}" + f") " + f"SELECT {base_cols}, rrf.vector_distance, rrf.vector_rank, rrf.fts_score, rrf.fts_rank, rrf.rrf_score " + f"FROM rrf_fused rrf " + f"JOIN {table} base ON base.{resolved_id_col} = rrf.{resolved_id_col} " + f"ORDER BY rrf.rrf_score DESC " + f"FETCH FIRST {final_bind} ROWS ONLY" + ) + + +def build_oracle_hybrid_rrf_query( + table: str, + vector_column: str = "embedding", + text_column: str = "content_text", + *, + id_column: str = "id", + key_column: str | None = None, + select_columns: Sequence[str] | None = None, + vector_param: str = ":query_vec", + text_param: str = ":fts_query", + metric: str = "COSINE", + rrf_k: int = 60, + rrf_k_param: str | None = None, + candidate_limit_param: str = ":candidate_limit", + final_limit_param: str = ":final_limit", + top_k_param: str | None = None, + where_clause: str | None = None, +) -> str: + """Alias for :func:`build_oracle_hybrid_rrf_sql`.""" + return build_oracle_hybrid_rrf_sql( + table, + vector_column, + text_column, + id_column=id_column, + key_column=key_column, + select_columns=select_columns, + vector_param=vector_param, + text_param=text_param, + metric=metric, + rrf_k=rrf_k, + rrf_k_param=rrf_k_param, + candidate_limit_param=candidate_limit_param, + final_limit_param=final_limit_param, + top_k_param=top_k_param, + where_clause=where_clause, + ) + + +def _format_oracle_bind(param: str) -> str: + stripped = param.strip() + if stripped.startswith(":") or stripped.isdigit() or "(" in stripped: + return stripped + return f":{stripped}" + + def _render_with_metric(generator: "Generator", expression: exp.Operator, dialect: str) -> str: left_sql = generator.sql(expression, "this") right_sql = generator.sql(expression, "expression") diff --git a/sqlspec/config.py b/sqlspec/config.py index 86cc89238..a4f165a1b 100644 --- a/sqlspec/config.py +++ b/sqlspec/config.py @@ -15,7 +15,7 @@ from difflib import get_close_matches from inspect import Signature, signature from pathlib import Path -from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, NotRequired, TypeAlias, TypeVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, NotRequired, Self, TypeAlias, TypeVar, cast from typing_extensions import TypedDict @@ -1797,6 +1797,49 @@ def provide_session( self._provide_session_impl(*args, statement_config=statement_config, **kwargs), ) + def derive_pool( + self, + *, + connection_config_overrides: "Mapping[str, Any] | None" = None, + driver_features_overrides: "Mapping[str, Any] | None" = None, + statement_config: "StatementConfig | None" = None, + migration_config: "dict[str, Any] | MigrationConfig | None" = None, + extension_config: "ExtensionConfigs | None" = None, + observability_config: "ObservabilityConfig | None" = None, + bind_key: "str | None" = None, + **overrides: Any, + ) -> Self: + """Clone this configuration with overridden connection/pool or feature settings and ``connection_instance=None``.""" + new_conn_cfg: dict[str, Any] = dict(self.connection_config) + new_driver_features: dict[str, Any] = dict(self.driver_features) + pool_session_cb = getattr(self, "_pool_session_callback", None) + if pool_session_cb is not None and "session_callback" not in new_conn_cfg: + new_conn_cfg["session_callback"] = pool_session_cb + user_conn_hook = getattr(self, "_user_connection_hook", None) + if user_conn_hook is not None and "on_connection_create" not in new_driver_features: + new_driver_features["on_connection_create"] = user_conn_hook + if connection_config_overrides: + new_conn_cfg.update(connection_config_overrides) + if driver_features_overrides: + new_driver_features.update(driver_features_overrides) + for key, val in overrides.items(): + if key in self.driver_features or key in DRIVER_FEATURE_LIFECYCLE_HOOKS: + new_driver_features[key] = val + else: + new_conn_cfg[key] = val + return type(self)( + connection_config=new_conn_cfg, + connection_instance=None, + migration_config=dict(self.migration_config) if migration_config is None else migration_config, + statement_config=self.statement_config if statement_config is None else statement_config, + driver_features=new_driver_features, + bind_key=self.bind_key if bind_key is None else bind_key, + extension_config=cast("ExtensionConfigs", dict(self.extension_config)) + if extension_config is None + else extension_config, + observability_config=self.observability_config if observability_config is None else observability_config, + ) + @abstractmethod def _create_pool(self) -> PoolT: """Actual pool creation implementation.""" @@ -1893,6 +1936,49 @@ def provide_session( self._provide_session_impl(*args, statement_config=statement_config, **kwargs), ) + def derive_pool( + self, + *, + connection_config_overrides: "Mapping[str, Any] | None" = None, + driver_features_overrides: "Mapping[str, Any] | None" = None, + statement_config: "StatementConfig | None" = None, + migration_config: "dict[str, Any] | MigrationConfig | None" = None, + extension_config: "ExtensionConfigs | None" = None, + observability_config: "ObservabilityConfig | None" = None, + bind_key: "str | None" = None, + **overrides: Any, + ) -> Self: + """Clone this configuration with overridden connection/pool or feature settings and ``connection_instance=None``.""" + new_conn_cfg: dict[str, Any] = dict(self.connection_config) + new_driver_features: dict[str, Any] = dict(self.driver_features) + pool_session_cb = getattr(self, "_pool_session_callback", None) + if pool_session_cb is not None and "session_callback" not in new_conn_cfg: + new_conn_cfg["session_callback"] = pool_session_cb + user_conn_hook = getattr(self, "_user_connection_hook", None) + if user_conn_hook is not None and "on_connection_create" not in new_driver_features: + new_driver_features["on_connection_create"] = user_conn_hook + if connection_config_overrides: + new_conn_cfg.update(connection_config_overrides) + if driver_features_overrides: + new_driver_features.update(driver_features_overrides) + for key, val in overrides.items(): + if key in self.driver_features or key in DRIVER_FEATURE_LIFECYCLE_HOOKS: + new_driver_features[key] = val + else: + new_conn_cfg[key] = val + return type(self)( + connection_config=new_conn_cfg, + connection_instance=None, + migration_config=dict(self.migration_config) if migration_config is None else migration_config, + statement_config=self.statement_config if statement_config is None else statement_config, + driver_features=new_driver_features, + bind_key=self.bind_key if bind_key is None else bind_key, + extension_config=cast("ExtensionConfigs", dict(self.extension_config)) + if extension_config is None + else extension_config, + observability_config=self.observability_config if observability_config is None else observability_config, + ) + @abstractmethod async def _create_pool(self) -> PoolT: """Actual async pool creation implementation.""" diff --git a/sqlspec/core/explain.py b/sqlspec/core/explain.py index 159ef2b5e..4a8a3126d 100644 --- a/sqlspec/core/explain.py +++ b/sqlspec/core/explain.py @@ -9,7 +9,14 @@ from mypy_extensions import mypyc_attr -__all__ = ("ExplainFormat", "ExplainOptions") +__all__ = ( + "ORACLE_EXPLAIN_DISPLAY_CURSOR_META_KEY", + "ORACLE_EXPLAIN_PREFIX", + "ORACLE_EXPLAIN_XPLAN_FORMAT_META_KEY", + "ORACLE_MANAGED_EXPLAIN_META_KEY", + "ExplainFormat", + "ExplainOptions", +) EXPLAIN_OPTION_FIELDS: Final[tuple[str, ...]] = ( @@ -24,9 +31,13 @@ "settings", "wal", "generic_plan", + "display_cursor", + "xplan_format", ) EXPLAIN_OPTIONS_SLOTS: Final = EXPLAIN_OPTION_FIELDS ORACLE_MANAGED_EXPLAIN_META_KEY: Final[str] = "sqlspec_managed_oracle_explain" +ORACLE_EXPLAIN_DISPLAY_CURSOR_META_KEY: Final[str] = "sqlspec_oracle_explain_display_cursor" +ORACLE_EXPLAIN_XPLAN_FORMAT_META_KEY: Final[str] = "sqlspec_oracle_explain_xplan_format" ORACLE_EXPLAIN_PREFIX: Final[str] = "EXPLAIN PLAN FOR " @@ -104,6 +115,8 @@ def __init__( settings: bool | None = None, wal: bool | None = None, generic_plan: bool | None = None, + display_cursor: bool | None = None, + xplan_format: str | None = None, ) -> None: """Initialize ExplainOptions. @@ -119,6 +132,8 @@ def __init__( settings: Include configuration parameter information wal: Include WAL usage information generic_plan: Generate a generic plan ignoring parameter values + display_cursor: Use Oracle DBMS_XPLAN.DISPLAY_CURSOR instead of PLAN_TABLE + xplan_format: Custom Oracle DBMS_XPLAN format string """ self.analyze = analyze self.verbose = verbose @@ -130,6 +145,8 @@ def __init__( self.settings = settings self.wal = wal self.generic_plan = generic_plan + self.display_cursor = display_cursor + self.xplan_format = xplan_format if format is not None: if isinstance(format, ExplainFormat): @@ -178,6 +195,8 @@ def copy( settings: "bool | None" = None, wal: "bool | None" = None, generic_plan: "bool | None" = None, + display_cursor: "bool | None" = None, + xplan_format: "str | None" = None, ) -> "ExplainOptions": """Create a copy with modifications. @@ -193,6 +212,8 @@ def copy( settings: Override settings setting wal: Override wal setting generic_plan: Override generic_plan setting + display_cursor: Override Oracle display_cursor setting + xplan_format: Override Oracle xplan_format setting Returns: New ExplainOptions instance with modifications applied @@ -209,6 +230,8 @@ def copy( settings=settings if settings is not None else self.settings, wal=wal if wal is not None else self.wal, generic_plan=generic_plan if generic_plan is not None else self.generic_plan, + display_cursor=display_cursor if display_cursor is not None else self.display_cursor, + xplan_format=xplan_format if xplan_format is not None else self.xplan_format, ) def to_dict(self) -> "dict[str, Any]": @@ -254,4 +277,8 @@ def _field_value(self, field_name: str) -> Any: return self.wal if field_name == "generic_plan": return self.generic_plan + if field_name == "display_cursor": + return self.display_cursor + if field_name == "xplan_format": + return self.xplan_format return None diff --git a/sqlspec/data_dictionary/dialects/oracle/sql/objects.sql b/sqlspec/data_dictionary/dialects/oracle/sql/objects.sql index f27eeb9fe..3ec05c56b 100644 --- a/sqlspec/data_dictionary/dialects/oracle/sql/objects.sql +++ b/sqlspec/data_dictionary/dialects/oracle/sql/objects.sql @@ -16,3 +16,11 @@ FROM all_objects WHERE owner = COALESCE(:schema_name, USER) AND (:object_name IS NULL OR object_name = :object_name) ORDER BY owner, object_type, object_name; + +-- name: cloud_ai_available +-- dialect: oracle +SELECT + COUNT(*) AS "cnt" +FROM all_objects +WHERE object_name = 'DBMS_CLOUD_AI' + AND object_type = 'PACKAGE'; diff --git a/sqlspec/data_dictionary/dialects/oracle/sql/system.sql b/sqlspec/data_dictionary/dialects/oracle/sql/system.sql index 9972e45c8..36b394551 100644 --- a/sqlspec/data_dictionary/dialects/oracle/sql/system.sql +++ b/sqlspec/data_dictionary/dialects/oracle/sql/system.sql @@ -27,3 +27,11 @@ SELECT FROM dual WHERE :include_diagnostics = 1 AND :acknowledge_license = 1; + +-- name: vector_memory_pool +-- dialect: oracle +SELECT + pool AS "pool", + alloc_bytes AS "alloc_bytes", + used_bytes AS "used_bytes" +FROM v$vector_memory_pool; diff --git a/sqlspec/data_dictionary/dialects/oracle/sql/version.sql b/sqlspec/data_dictionary/dialects/oracle/sql/version.sql index 5f6f36a5a..bbf3ec91c 100644 --- a/sqlspec/data_dictionary/dialects/oracle/sql/version.sql +++ b/sqlspec/data_dictionary/dialects/oracle/sql/version.sql @@ -18,3 +18,10 @@ SELECT sys_context('USERENV', 'CLOUD_SERVICE') AS "service" FROM dual; SELECT value AS compatible FROM v$parameter WHERE name = 'compatible'; + +-- name: apex_release +-- dialect: oracle +SELECT + version_no AS "version_no", + api_compatibility AS "api_compatibility" +FROM apex_release; diff --git a/sqlspec/data_dictionary/dialects/oracle/sql/views.sql b/sqlspec/data_dictionary/dialects/oracle/sql/views.sql index 2bce3c383..2537d5f01 100644 --- a/sqlspec/data_dictionary/dialects/oracle/sql/views.sql +++ b/sqlspec/data_dictionary/dialects/oracle/sql/views.sql @@ -21,3 +21,29 @@ FROM all_views WHERE owner = COALESCE(:schema_name, USER) AND (:view_name IS NULL OR view_name = :view_name) ORDER BY owner, view_name; + +-- name: json_duality_views +-- dialect: oracle +SELECT + view_name AS "view_name", + root_table_name AS "root_table_name", + root_table_owner AS "root_table_owner", + allow_insert AS "allow_insert", + allow_update AS "allow_update", + allow_delete AS "allow_delete" +FROM user_json_duality_views +ORDER BY view_name; + +-- name: json_duality_views_by_schema +-- dialect: oracle +SELECT + view_name AS "view_name", + root_table_name AS "root_table_name", + root_table_owner AS "root_table_owner", + allow_insert AS "allow_insert", + allow_update AS "allow_update", + allow_delete AS "allow_delete" +FROM all_json_duality_views +WHERE (:schema_name IS NULL AND view_owner = SYS_CONTEXT('USERENV', 'CURRENT_SCHEMA')) + OR (:schema_name IS NOT NULL AND view_owner = :schema_name) +ORDER BY view_name; diff --git a/tests/integration/adapters/_shared/_driver_type_system.py b/tests/integration/adapters/_shared/_driver_type_system.py index 7601c431d..13ff1ca94 100644 --- a/tests/integration/adapters/_shared/_driver_type_system.py +++ b/tests/integration/adapters/_shared/_driver_type_system.py @@ -290,6 +290,7 @@ class SourceEquivalenceCase: "enable_events", "events_backend", "enable_direct_path_load", + "session_context_hook", ), "psqlpy": ( "enable_cast_detection", diff --git a/tests/unit/adapters/test_oracledb/test_arrow_batches.py b/tests/unit/adapters/test_oracledb/test_arrow_batches.py index 95ccb4ffd..73545af34 100644 --- a/tests/unit/adapters/test_oracledb/test_arrow_batches.py +++ b/tests/unit/adapters/test_oracledb/test_arrow_batches.py @@ -1,11 +1,16 @@ """Unit tests for Oracle native Arrow batch export.""" from collections.abc import AsyncIterator, Iterator -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, cast import pyarrow as pa import pytest +from sqlspec.adapters.oracledb.core import ( + build_arrow_fetch_kwargs, + coerce_oracle_df_to_arrow_batches, + coerce_oracle_df_to_arrow_table, +) from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver if TYPE_CHECKING: @@ -73,10 +78,10 @@ def test_sync_select_to_arrow_batches_uses_fetch_df_batches() -> None: assert result.rows_affected == 3 assert connection.fetch_batch_calls[0]["size"] == 2 assert connection.fetch_batch_calls[0]["fetch_decimals"] is True - assert connection.fetch_batch_calls[0]["fetch_lobs"] is False + assert "fetch_lobs" not in connection.fetch_batch_calls[0] -def test_sync_select_to_arrow_table_forwards_fetch_lobs_to_fetch_df_all() -> None: +def test_sync_select_to_arrow_table_omits_fetch_lobs_from_fetch_df_all() -> None: connection = _OracleAllConnection() driver = OracleSyncDriver( cast("OracleSyncConnection", connection), @@ -88,7 +93,7 @@ def test_sync_select_to_arrow_table_forwards_fetch_lobs_to_fetch_df_all() -> Non assert result.get_data().to_pydict() == {"id": [1], "value": [10]} assert connection.fetch_all_calls[0]["arraysize"] == 2 assert connection.fetch_all_calls[0]["fetch_decimals"] is True - assert connection.fetch_all_calls[0]["fetch_lobs"] is False + assert "fetch_lobs" not in connection.fetch_all_calls[0] @pytest.mark.anyio @@ -105,11 +110,11 @@ async def test_async_select_to_arrow_reader_uses_fetch_df_batches() -> None: assert result.get_data().read_all().to_pydict() == {"id": [1, 2, 3], "value": [10, 20, 30]} assert connection.fetch_batch_calls[0]["size"] == 2 assert connection.fetch_batch_calls[0]["fetch_decimals"] is True - assert connection.fetch_batch_calls[0]["fetch_lobs"] is False + assert "fetch_lobs" not in connection.fetch_batch_calls[0] @pytest.mark.anyio -async def test_async_select_to_arrow_table_forwards_fetch_lobs_to_fetch_df_all() -> None: +async def test_async_select_to_arrow_table_omits_fetch_lobs_from_fetch_df_all() -> None: connection = _OracleAsyncAllConnection() driver = OracleAsyncDriver( cast("OracleAsyncConnection", connection), @@ -121,4 +126,66 @@ async def test_async_select_to_arrow_table_forwards_fetch_lobs_to_fetch_df_all() assert result.get_data().to_pydict() == {"id": [1], "value": [10]} assert connection.fetch_all_calls[0]["arraysize"] == 2 assert connection.fetch_all_calls[0]["fetch_decimals"] is True - assert connection.fetch_all_calls[0]["fetch_lobs"] is False + assert "fetch_lobs" not in connection.fetch_all_calls[0] + + +class _StreamOnlyOracleDataFrame: + """Fake OracleDataFrame that exposes only __arrow_c_stream__.""" + + def __init__(self, table: pa.Table) -> None: + self._table = table + self.stream_calls = 0 + + def __arrow_c_stream__(self, requested_schema: Any = None) -> Any: + self.stream_calls += 1 + return cast("Any", self._table).__arrow_c_stream__(requested_schema=requested_schema) + + def __arrow_c_array__(self, *_args: Any, **_kwargs: Any) -> Any: + msg = "select_to_arrow must use __arrow_c_stream__ rather than __arrow_c_array__" + raise AssertionError(msg) + + +def test_build_arrow_fetch_kwargs_excludes_fetch_lobs_and_passes_requested_schema() -> None: + schema = pa.schema([("id", pa.int64())]) + kwargs = build_arrow_fetch_kwargs({"fetch_lobs": True, "fetch_decimals": True}, arrow_schema=schema, batch_size=500) + assert "fetch_lobs" not in kwargs + assert kwargs["arraysize"] == 500 + assert kwargs["fetch_decimals"] is True + assert kwargs["requested_schema"] is schema + + +def test_select_to_arrow_consumes_arrow_c_stream_pycapsule_sync() -> None: + raw_table = pa.table({"ID": [1, 2], "VALUE": [10, 20]}) + stream_df = _StreamOnlyOracleDataFrame(raw_table) + + class _StreamConn: + def fetch_df_all(self, **_kwargs: object) -> _StreamOnlyOracleDataFrame: + return stream_df + + driver = OracleSyncDriver( + cast("OracleSyncConnection", _StreamConn()), driver_features={"enable_lowercase_column_names": True} + ) + res_table = driver.select_to_arrow("SELECT id, value FROM example", return_format="table") + assert res_table.get_data().to_pydict() == {"id": [1, 2], "value": [10, 20]} + assert stream_df.stream_calls == 1 + + batches = coerce_oracle_df_to_arrow_batches(stream_df, pa) + assert sum(b.num_rows for b in batches) == 2 + assert coerce_oracle_df_to_arrow_table(stream_df, pa).num_rows == 2 + + +@pytest.mark.anyio +async def test_select_to_arrow_consumes_arrow_c_stream_pycapsule_async() -> None: + raw_table = pa.table({"ID": [1, 2], "VALUE": [10, 20]}) + stream_df = _StreamOnlyOracleDataFrame(raw_table) + + class _AsyncStreamConn: + async def fetch_df_all(self, **_kwargs: object) -> _StreamOnlyOracleDataFrame: + return stream_df + + driver = OracleAsyncDriver( + cast("OracleAsyncConnection", _AsyncStreamConn()), driver_features={"enable_lowercase_column_names": True} + ) + res_table = await driver.select_to_arrow("SELECT id, value FROM example", return_format="table") + assert res_table.get_data().to_pydict() == {"id": [1, 2], "value": [10, 20]} + assert stream_df.stream_calls == 1 diff --git a/tests/unit/adapters/test_oracledb/test_config.py b/tests/unit/adapters/test_oracledb/test_config.py index d357c73db..97b792dfa 100644 --- a/tests/unit/adapters/test_oracledb/test_config.py +++ b/tests/unit/adapters/test_oracledb/test_config.py @@ -1,15 +1,25 @@ """OracleDB configuration tests covering driver kwargs and typed options.""" +import io +import zipfile from collections.abc import Awaitable, Callable from inspect import isawaitable +from pathlib import Path from ssl import TLSVersion from typing import Any, NotRequired, cast, get_args, get_origin, get_type_hints -from unittest.mock import AsyncMock, Mock, call +from unittest.mock import AsyncMock, MagicMock, Mock, call import pytest from oracledb import AuthMode, PoolGetMode, Purity -from sqlspec.adapters.oracledb import build_connection_config +from sqlspec.adapters.oracledb import ( + OracleEndpoint, + OracleSessionContextConfig, + OracleTopologyConfig, + build_connection_config, + build_session_checkout_plsql, + build_session_release_plsql, +) from sqlspec.adapters.oracledb import config as oracle_config_module from sqlspec.adapters.oracledb.config import ( OracleAsyncConfig, @@ -339,3 +349,246 @@ async def test_async_soda_option_does_not_initialize_thick_mode(monkeypatch: pyt initialize.assert_not_called() assert create_pool.call_args.kwargs["soda_metadata_cache"] is True + + +def test_oracle_config_wallet_zip_extracted_and_cleaned_up_on_close_pool() -> None: + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + zf.writestr("ewallet.pem", b"pem-bytes") + pool_mock = Mock() + config = OracleSyncConfig( + connection_instance=pool_mock, + connection_config={ + "user": "admin", + "password": "pw", + "wallet_zip": buf.getvalue(), + "topology": OracleTopologyConfig( + endpoints=(OracleEndpoint(host="adb.oraclecloud.com", port=1522, protocol="tcps"),), + service_name="adb_high.adb.oraclecloud.com", + ), + }, + ) + wallet_dir = config._extracted_wallet_dir + assert wallet_dir is not None + assert Path(wallet_dir).exists() + assert config.connection_config["wallet_location"] == str(wallet_dir) + assert config.connection_config["config_dir"] == str(wallet_dir) + assert "(SERVICE_NAME=adb_high.adb.oraclecloud.com)" in config.connection_config["dsn"] + + config._close_pool() + assert not Path(wallet_dir).exists() + + +def test_oracle_derive_pool_preserves_popped_callbacks_and_applies_overrides() -> None: + def _session_cb(conn: Any, tag: str) -> None: + return None + + def _conn_hook(conn: Any, tag: str) -> None: + return None + + session_cb: Any = _session_cb + conn_hook: Any = _conn_hook + + base_sync = OracleSyncConfig( + connection_config={"dsn": "tp", "session_callback": session_cb}, + driver_features={"on_connection_create": conn_hook, "arraysize": 100}, + extension_config={"litestar": {"commit_mode": "manual"}}, + ) + derived_sync = base_sync.derive_pool( + arraysize=500, cclass="BATCH", extension_config={"adk": {"session_table": "tenant_sessions"}} + ) + assert derived_sync.connection_instance is None + assert derived_sync._pool_session_callback is session_cb + assert derived_sync._user_connection_hook is conn_hook + assert derived_sync.driver_features["arraysize"] == 500 + assert derived_sync.connection_config["cclass"] == "BATCH" + assert base_sync.driver_features["arraysize"] == 100 + assert "cclass" not in base_sync.connection_config + assert "adk" in derived_sync.extension_config + assert "adk" not in base_sync.extension_config + + base_async = OracleAsyncConfig( + connection_config={"dsn": "tp", "session_callback": session_cb}, + driver_features={"on_connection_create": conn_hook, "fetch_lobs": False, "arraysize": 100}, + ) + batch_cfg = base_async.derive_pool(fetch_lobs=True, arraysize=250, min=1, max=4, stmtcachesize=0, cclass="BATCH") + assert batch_cfg.connection_instance is None + assert batch_cfg._pool_session_callback is session_cb + assert batch_cfg._user_connection_hook is conn_hook + assert batch_cfg.driver_features["fetch_lobs"] is True + assert batch_cfg.driver_features["arraysize"] == 250 + assert batch_cfg.connection_config["cclass"] == "BATCH" + assert batch_cfg.connection_config["min"] == 1 + assert batch_cfg.connection_config["max"] == 4 + assert batch_cfg.connection_config["stmtcachesize"] == 0 + assert base_async.driver_features["fetch_lobs"] is False + assert base_async.driver_features["arraysize"] == 100 + + +def test_oracle_session_context_hook_sync_checkout_and_release() -> None: + executed_sql: list[tuple[str, dict[str, Any]]] = [] + cursor_mock = MagicMock() + cursor_mock.__enter__.return_value = cursor_mock + cursor_mock.__exit__.return_value = False + cursor_mock.execute.side_effect = lambda sql, params=None: executed_sql.append((sql, dict(params or {}))) + + conn_mock = MagicMock() + conn_mock.__enter__.return_value = conn_mock + conn_mock.__exit__.return_value = False + conn_mock.cursor.return_value = cursor_mock + + pool_mock = MagicMock() + pool_mock.acquire.return_value = conn_mock + + hook = OracleSessionContextConfig( + module="billing", + action=lambda: "invoice", + client_identifier="tenant_42", + edition="ORA$BASE", + time_zone="UTC", + nls_parameters={"NLS_DATE_FORMAT": "YYYY-MM-DD"}, + contexts={"APP_CTX": {"TENANT_ID": "42"}}, + reset_on_release=True, + ) + config = OracleSyncConfig(connection_instance=pool_mock, driver_features={"session_context_hook": hook}) + with config.provide_connection(): + assert conn_mock.module == "billing" + assert conn_mock.action == "invoice" + assert conn_mock.client_identifier == "tenant_42" + + assert len(executed_sql) == 2 + checkout_sql, checkout_binds = executed_sql[0] + release_sql, _ = executed_sql[1] + assert "EDITION = ORA$BASE" in checkout_sql + assert "TIME_ZONE = ''UTC''" in checkout_sql + assert "DBMS_SESSION.SET_CONTEXT(:ctx_ns_0, :ctx_attr_0, :ctx_val_0);" in checkout_sql + assert checkout_binds == {"ctx_ns_0": "APP_CTX", "ctx_attr_0": "TENANT_ID", "ctx_val_0": "42"} + assert "DBMS_SESSION.CLEAR_CONTEXT('APP_CTX');" in release_sql + assert conn_mock.module == "" + assert conn_mock.action == "" + assert conn_mock.client_identifier == "" + + +@pytest.mark.anyio +async def test_oracle_session_context_hook_async_checkout_and_release() -> None: + executed_sql: list[tuple[str, dict[str, Any]]] = [] + cursor_mock = MagicMock() + cursor_mock.__enter__.return_value = cursor_mock + cursor_mock.__exit__.return_value = False + cursor_mock.__aenter__ = AsyncMock(return_value=cursor_mock) + cursor_mock.__aexit__ = AsyncMock(return_value=False) + + async def record_execute(sql: str, params: dict[str, Any] | None = None) -> None: + executed_sql.append((sql, dict(params or {}))) + + cursor_mock.execute = AsyncMock(side_effect=record_execute) + + conn_mock = MagicMock() + conn_mock.cursor.return_value = cursor_mock + + pool_mock = MagicMock() + pool_mock.acquire = AsyncMock(return_value=conn_mock) + pool_mock.release = AsyncMock() + + hook = OracleSessionContextConfig( + module="search", client_identifier="user_99", contexts={"CTX_NS": {"ORG_ID": "org_1"}}, reset_on_release=True + ) + config = OracleAsyncConfig(connection_instance=pool_mock, driver_features={"session_context_hook": hook}) + async with config.provide_connection(): + assert conn_mock.module == "search" + assert conn_mock.client_identifier == "user_99" + + assert len(executed_sql) == 2 + assert "DBMS_SESSION.SET_CONTEXT" in executed_sql[0][0] + assert "DBMS_SESSION.CLEAR_CONTEXT('CTX_NS');" in executed_sql[1][0] + pool_mock.release.assert_awaited_once_with(conn_mock) + + +def test_build_session_checkout_and_release_plsql() -> None: + cfg = OracleSessionContextConfig( + client_identifier=lambda: "tenant_user_1", + context_attributes={"APP_CTX": {"TENANT_ID": "t1"}}, + ebs_apps_initialize=(1010, 20420, 1), + mo_global_policy=("S", 84), + apex_session=(100, 1, 999999), + reset_package=True, + clear_identifier=True, + clear_context_namespaces=("APP_CTX",), + apex_detach=True, + ) + checkout_sql, checkout_binds = build_session_checkout_plsql(cfg) + assert checkout_sql is not None + assert "DBMS_SESSION.SET_IDENTIFIER(:p_client_id);" in checkout_sql + assert "DBMS_SESSION.SET_CONTEXT(:ctx_ns_0, :ctx_attr_0, :ctx_val_0);" in checkout_sql + assert ( + "FND_GLOBAL.APPS_INITIALIZE(user_id => :p_ebs_uid, resp_id => :p_ebs_rid, resp_appl_id => :p_ebs_aid);" + in checkout_sql + ) + assert "MO_GLOBAL.SET_POLICY_CONTEXT(:p_mo_mode, :p_mo_org);" in checkout_sql + assert ( + "APEX_SESSION.ATTACH(p_app_id => :p_apex_app, p_page_id => :p_apex_page, p_session_id => :p_apex_sess);" + in checkout_sql + ) + assert checkout_binds["p_client_id"] == "tenant_user_1" + assert checkout_binds["p_ebs_uid"] == 1010 + assert checkout_binds["p_mo_mode"] == "S" + assert checkout_binds["p_apex_app"] == 100 + + release_sql, release_binds = build_session_release_plsql(cfg) + assert release_sql is not None + assert "APEX_SESSION.DETACH;" in release_sql + assert "DBMS_SESSION.CLEAR_IDENTIFIER;" in release_sql + assert "DBMS_SESSION.CLEAR_ALL_CONTEXT(:p_clr_ns_0);" in release_sql + assert "DBMS_SESSION.RESET_PACKAGE;" in release_sql + assert release_binds == {"p_clr_ns_0": "APP_CTX"} + + +@pytest.mark.anyio +async def test_oracle_sync_and_async_provide_connection_and_session_run_checkout_and_release_hooks_on_exception() -> ( + None +): + executed_sync: list[str] = [] + cursor_sync = MagicMock() + cursor_sync.__enter__.return_value = cursor_sync + cursor_sync.__exit__.return_value = False + cursor_sync.execute.side_effect = lambda sql, params=None: executed_sync.append(sql) + + conn_sync = MagicMock() + conn_sync.cursor.return_value = cursor_sync + pool_sync = MagicMock() + pool_sync.acquire.return_value = conn_sync + + hook = OracleSessionContextConfig(client_identifier="u1", clear_identifier=True, reset_package=True) + sync_cfg = OracleSyncConfig(connection_instance=pool_sync, driver_features={"session_context_hook": hook}) + with pytest.raises(RuntimeError, match="boom_sync"), sync_cfg.provide_session(): + raise RuntimeError("boom_sync") + + assert len(executed_sync) == 2 + assert "DBMS_SESSION.SET_IDENTIFIER" in executed_sync[0] + assert "DBMS_SESSION.RESET_PACKAGE;" in executed_sync[1] + pool_sync.release.assert_called_once_with(conn_sync) + + executed_async: list[str] = [] + cursor_async = MagicMock() + cursor_async.__enter__.return_value = cursor_async + cursor_async.__exit__.return_value = False + + async def rec_async(sql: str, params: Any = None) -> None: + executed_async.append(sql) + + cursor_async.execute = AsyncMock(side_effect=rec_async) + conn_async = MagicMock() + conn_async.cursor.return_value = cursor_async + pool_async = MagicMock() + pool_async.acquire = AsyncMock(return_value=conn_async) + pool_async.release = AsyncMock() + + async_cfg = OracleAsyncConfig(connection_instance=pool_async, driver_features={"session_context_hook": hook}) + with pytest.raises(RuntimeError, match="boom_async"): + async with async_cfg.provide_session(): + raise RuntimeError("boom_async") + + assert len(executed_async) == 2 + assert "DBMS_SESSION.SET_IDENTIFIER" in executed_async[0] + assert "DBMS_SESSION.RESET_PACKAGE;" in executed_async[1] + pool_async.release.assert_awaited_once_with(conn_async) diff --git a/tests/unit/adapters/test_oracledb/test_core_driver_features.py b/tests/unit/adapters/test_oracledb/test_core_driver_features.py index d7b464e81..691849589 100644 --- a/tests/unit/adapters/test_oracledb/test_core_driver_features.py +++ b/tests/unit/adapters/test_oracledb/test_core_driver_features.py @@ -2,6 +2,10 @@ """Unit tests for OracleDB ``apply_driver_features`` defaults.""" import inspect +import io +import shutil +import zipfile +from pathlib import Path from typing import TYPE_CHECKING, cast import pytest @@ -9,9 +13,21 @@ import sqlspec.adapters.oracledb._typing as oracle_typing from sqlspec.adapters.oracledb import config from sqlspec.adapters.oracledb._typing import OracleAsyncCursor -from sqlspec.adapters.oracledb.config import OracleAsyncConfig, OracleSyncConfig -from sqlspec.adapters.oracledb.core import apply_driver_features +from sqlspec.adapters.oracledb.config import OracleAsyncConfig, OracleDriverFeatures, OracleSyncConfig +from sqlspec.adapters.oracledb.core import ( + OracleEndpoint, + OracleTopologyConfig, + apply_driver_features, + build_duality_etag_update_sql, + build_tns_description, + extract_duality_etag, + extract_duality_view_etag, + extract_oracle_wallet_zip, + inject_duality_view_etag, + sanitize_oracle_text_query, +) from sqlspec.core import StatementConfig +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.typing import NUMPY_INSTALLED if TYPE_CHECKING: @@ -195,3 +211,86 @@ def test_typing_exports_oracle_sync_cursor_still_importable() -> None: from sqlspec.adapters.oracledb._typing import OracleSyncCursor assert OracleSyncCursor is oracle_typing.OracleSyncCursor + + +def test_apply_driver_features_rejects_unknown_keys() -> None: + assert OracleDriverFeatures is not None + with pytest.raises(ImproperConfigurationError, match="Unknown Oracle driver_features"): + apply_driver_features(StatementConfig(), {"unknown_oracle_knob": True}) + + +def test_build_tns_description_multi_endpoint_and_rejects_injection() -> None: + topology = OracleTopologyConfig( + endpoints=( + OracleEndpoint(host="rac1.example.com", port=1521, protocol="tcp"), + OracleEndpoint(host="rac2.example.com", port=2484, protocol="tcps"), + ), + service_name="orcl_high.adb.oraclecloud.com", + load_balance=True, + failover=True, + retry_count=3, + retry_delay=2, + transport_connect_timeout=5, + sdu=65535, + source_route=True, + server="DEDICATED", + failover_mode_type="SELECT", + failover_mode_method="BASIC", + ) + dsn = build_tns_description(topology) + assert "(HOST=rac1.example.com)(PORT=1521)" in dsn + assert "(PROTOCOL=TCPS)(HOST=rac2.example.com)(PORT=2484)" in dsn + assert "(SERVICE_NAME=orcl_high.adb.oraclecloud.com)" in dsn + assert "(LOAD_BALANCE=ON)" in dsn + assert "(FAILOVER=ON)" in dsn + assert "(RETRY_COUNT=3)" in dsn + assert "(RETRY_DELAY=2)" in dsn + assert "(TRANSPORT_CONNECT_TIMEOUT=5)" in dsn + assert "(SDU=65535)" in dsn + assert "(SOURCE_ROUTE=ON)" in dsn + assert "(SERVER=DEDICATED)" in dsn + assert "(FAILOVER_MODE=(TYPE=SELECT)(METHOD=BASIC))" in dsn + + with pytest.raises(ImproperConfigurationError): + build_tns_description( + OracleTopologyConfig(endpoints=(OracleEndpoint(host="bad)host", port=1521),), service_name="svc") + ) + + +def test_extract_oracle_wallet_zip_validates_entries_and_permissions() -> None: + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + zf.writestr("ewallet.pem", b"pem-content") + zf.writestr("tnsnames.ora", b"orcl_high = (DESCRIPTION=...)") + extracted_dir = extract_oracle_wallet_zip(buf.getvalue()) + try: + wallet_path = Path(extracted_dir) + assert (wallet_path / "ewallet.pem").read_bytes() == b"pem-content" + assert (wallet_path / "tnsnames.ora").exists() + finally: + shutil.rmtree(extracted_dir, ignore_errors=True) + + slip_buf = io.BytesIO() + with zipfile.ZipFile(slip_buf, "w") as zf: + zf.writestr("../evil.pem", b"pwn") + with pytest.raises(ImproperConfigurationError, match="Unsafe path in Oracle wallet archive"): + extract_oracle_wallet_zip(slip_buf.getvalue()) + + +def test_sanitize_oracle_text_query_escapes_reserved_words_and_syntax() -> None: + sanitized = sanitize_oracle_text_query("NOT (foo & bar) OR NEAR") + assert sanitized == "{NOT} ACCUM {foo} ACCUM {bar} ACCUM {OR} ACCUM {NEAR}" + assert sanitize_oracle_text_query(" !@#$%^&*() ") == "" + + +def test_duality_view_etag_helpers_extract_inject_and_build_update_sql() -> None: + doc = {"_metadata": {"etag": "etag-abc-123"}, "id": 42, "name": "Order"} + assert extract_duality_etag(doc) == "etag-abc-123" + assert extract_duality_view_etag(doc) == "etag-abc-123" + assert extract_duality_etag({"id": 42}) is None + + injected = inject_duality_view_etag({"id": 42}, "etag-999") + assert injected["_metadata"]["etag"] == "etag-999" + + sql = build_duality_etag_update_sql("orders_dv") + assert "JSON_VALUE(DATA, '$._metadata.etag') = :expected_etag" in sql diff --git a/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py b/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py index 4fb7f0b40..ab228b9ca 100644 --- a/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py +++ b/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py @@ -481,3 +481,40 @@ async def test_async_oracle_replacement_methods_execute_domain_queries() -> None statement_text, kwargs = async_driver_stub.select_value_calls[0] assert "all_objects" in statement_text.lower() assert kwargs == {"schema_name": "APP", "object_name": None} + + +def test_oracle_26ai_sync_data_dictionary_introspection_methods() -> None: + data_dictionary = OracledbSyncDataDictionary() + sync_driver_stub = _RecordingSyncOracleDriver() + sync_driver = cast("OracleSyncDriver", sync_driver_stub) + + vpool = data_dictionary.get_vector_memory_pool(sync_driver) + assert isinstance(vpool, list) + + has_ai = data_dictionary.has_cloud_ai(sync_driver) + assert has_ai is True + + apex = data_dictionary.get_apex_release(sync_driver) + assert isinstance(apex, dict) + + dviews = data_dictionary.get_json_duality_views(sync_driver, schema="app") + assert isinstance(dviews, list) + + +@pytest.mark.anyio +async def test_oracle_26ai_async_data_dictionary_introspection_methods() -> None: + data_dictionary = OracledbAsyncDataDictionary() + async_driver_stub = _RecordingAsyncOracleDriver() + async_driver = cast("OracleAsyncDriver", async_driver_stub) + + vpool = await data_dictionary.get_vector_memory_pool(async_driver) + assert isinstance(vpool, list) + + has_ai = await data_dictionary.has_cloud_ai(async_driver) + assert has_ai is True + + apex = await data_dictionary.get_apex_release(async_driver) + assert isinstance(apex, dict) + + dviews = await data_dictionary.get_json_duality_views(async_driver, schema="app") + assert isinstance(dviews, list) diff --git a/tests/unit/adapters/test_oracledb/test_explain_plan.py b/tests/unit/adapters/test_oracledb/test_explain_plan.py index f0db07f3d..732499ae3 100644 --- a/tests/unit/adapters/test_oracledb/test_explain_plan.py +++ b/tests/unit/adapters/test_oracledb/test_explain_plan.py @@ -1,19 +1,21 @@ """Unit tests for the Oracle EXPLAIN PLAN retrieval flow.""" -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest from sqlspec.adapters.oracledb.driver import ( PLAN_STATEMENT_ID_MAX_LENGTH, + OracleAsyncDriver, OracleSyncDriver, + _managed_explain_options, _managed_explain_statement, _new_plan_statement_id, _tagged_explain_plan_sql, ) -from sqlspec.builder import Explain +from sqlspec.builder import Explain, build_oracle_display_cursor from sqlspec.core import SQL, StatementConfig -from sqlspec.exceptions import SQLSpecError +from sqlspec.exceptions import SQLBuilderError, SQLSpecError def test_managed_explain_plan_yields_the_explained_statement() -> None: @@ -74,3 +76,54 @@ def test_tagged_sql_preserves_the_explained_statement() -> None: tagged = _tagged_explain_plan_sql("sqlspec_abc123", "SELECT * FROM t WHERE a = :1") assert tagged == "EXPLAIN PLAN SET STATEMENT_ID = 'sqlspec_abc123' FOR SELECT * FROM t WHERE a = :1" + + +def test_oracle_explain_builder_display_cursor_and_xplan_format() -> None: + stmt = ( + Explain("SELECT * FROM t WHERE a = :1", dialect="oracle") + .display_cursor() + .xplan_format("ALLSTATS LAST +PEEKED_BINDS") + .build() + ) + display_cursor, xplan_format = _managed_explain_options(stmt) + assert display_cursor is True + assert xplan_format == "ALLSTATS LAST +PEEKED_BINDS" + + sql_cursor = build_oracle_display_cursor(format="TYPICAL +PREDICATE +NOTE") + assert ( + sql_cursor + == "SELECT plan_table_output FROM TABLE(DBMS_XPLAN.DISPLAY_CURSOR(NULL, NULL, 'TYPICAL +PREDICATE +NOTE'))" + ) + + with pytest.raises(SQLBuilderError): + Explain("SELECT 1 FROM dual", dialect="oracle").xplan_format("BAD'FMT") + + +def test_oracle_explain_plan_custom_xplan_format_and_display_cursor_sync() -> None: + driver = OracleSyncDriver(MagicMock()) + cursor = MagicMock() + cursor.description = [("PLAN_TABLE_OUTPUT", None, None, None, None, None, None)] + cursor.fetchall.return_value = [("Plan hash value: 123",)] + + driver._dispatch_explain_plan( + cursor, "SELECT 1 FROM dual", {}, display_cursor=True, xplan_format="ALLSTATS LAST +PEEKED_BINDS" + ) + executed_sqls = [call.args[0] for call in cursor.execute.call_args_list] + assert executed_sqls[0] == "SELECT 1 FROM dual" + assert "DBMS_XPLAN.DISPLAY_CURSOR(NULL, NULL, 'ALLSTATS LAST +PEEKED_BINDS')" in executed_sqls[1] + + +@pytest.mark.anyio +async def test_oracle_explain_plan_custom_xplan_format_and_display_cursor_async() -> None: + driver = OracleAsyncDriver(MagicMock()) + cursor = MagicMock() + cursor.description = [("PLAN_TABLE_OUTPUT", None, None, None, None, None, None)] + cursor.execute = AsyncMock() + cursor.fetchall = AsyncMock(return_value=[("Plan hash value: 123",)]) + + await driver._dispatch_explain_plan( + cursor, "SELECT 1 FROM dual", {}, display_cursor=True, xplan_format="ALLSTATS LAST +PEEKED_BINDS" + ) + executed_sqls = [call.args[0] for call in cursor.execute.call_args_list] + assert executed_sqls[0] == "SELECT 1 FROM dual" + assert "DBMS_XPLAN.DISPLAY_CURSOR(NULL, NULL, 'ALLSTATS LAST +PEEKED_BINDS')" in executed_sqls[1] diff --git a/tests/unit/adapters/test_oracledb/test_fetch_tuning.py b/tests/unit/adapters/test_oracledb/test_fetch_tuning.py index 3a656f3b8..324654c84 100644 --- a/tests/unit/adapters/test_oracledb/test_fetch_tuning.py +++ b/tests/unit/adapters/test_oracledb/test_fetch_tuning.py @@ -67,10 +67,7 @@ def test_build_fetch_kwargs_includes_explicit_fetch_options() -> None: def test_build_arrow_fetch_kwargs_matches_execute_fetch_options() -> None: - assert build_arrow_fetch_kwargs({"fetch_lobs": False, "fetch_decimals": True}) == { - "fetch_lobs": False, - "fetch_decimals": True, - } + assert build_arrow_fetch_kwargs({"fetch_lobs": False, "fetch_decimals": True}) == {"fetch_decimals": True} def test_oracle_sync_cursor_applies_arraysize_and_prefetchrows() -> None: diff --git a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py index 24bb54dc4..0a0419082 100644 --- a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py +++ b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py @@ -624,3 +624,48 @@ def test_oracle_adk_owner_id_column_ddl_escapes_single_quotes() -> None: assert "DEFAULT ''anon'' NOT NULL" in session_ddl assert "DEFAULT ''anon'' NOT NULL" in memory_ddl + + +def test_oracle_sync_adk_memory_fts_sanitizes_query_and_hybrid_rrf_fuses_candidates() -> None: + cursor = _RecordingCursor() + conn = _RecordingConnection(cursor) + config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config.extension_config = {"adk": {"memory_use_fts": True}} + config.provide_connection = lambda *_a, **_k: conn + store = OracleSyncADKMemoryStore(config) + + store.search_entries(query="NOT (alpha & beta)", app_name="app", user_id="u1", limit=5) + assert len(cursor.calls) == 1 + fts_sql, fts_params = cursor.calls[0] + assert "CONTAINS(content_text, :query, 1) > 0" in fts_sql + assert fts_params["query"] == "{NOT} ACCUM {alpha} ACCUM {beta}" + + cursor.calls.clear() + assert store.search_entries(query=" !@# ", app_name="app", user_id="u1", limit=5) == [] + assert cursor.calls == [] + + store.search_entries(query="alpha beta", app_name="app", user_id="u1", limit=5, embedding=[0.1, 0.2, 0.3]) + assert len(cursor.calls) == 1 + rrf_sql, rrf_params = cursor.calls[0] + assert "rrf_score" in rrf_sql + assert "FETCH APPROX FIRST :top_k ROWS ONLY" in rrf_sql + assert rrf_params["text_query"] == "{alpha} ACCUM {beta}" + assert rrf_params["top_k"] == 5 + + +@pytest.mark.anyio +async def test_oracle_async_adk_memory_fts_sanitizes_query_and_hybrid_rrf_fuses_candidates() -> None: + cursor = _AsyncRecordingCursor() + conn = _RecordingConnection(cursor) + config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config.extension_config = {"adk": {"memory_use_fts": True}} + config.provide_connection = lambda *_a, **_k: conn + store = OracleAsyncADKMemoryStore(config) + + await store.search_entries( + query="NOT (alpha & beta)", app_name="app", user_id="u1", limit=5, embedding=[0.1, 0.2, 0.3] + ) + assert len(cursor.calls) == 1 + rrf_sql, rrf_params = cursor.calls[0] + assert "rrf_score" in rrf_sql + assert rrf_params["text_query"] == "{NOT} ACCUM {alpha} ACCUM {beta}" diff --git a/tests/unit/adapters/test_oracledb/test_pipeline_helpers.py b/tests/unit/adapters/test_oracledb/test_pipeline_helpers.py index 7e51014e9..71b89f09f 100644 --- a/tests/unit/adapters/test_oracledb/test_pipeline_helpers.py +++ b/tests/unit/adapters/test_oracledb/test_pipeline_helpers.py @@ -48,9 +48,11 @@ def run_pipeline(self, *_args: object, **_kwargs: object) -> list[object]: class _AsyncPipelineConnection: """Async connection stub exposing the python-oracledb async pipeline API.""" - def __init__(self, *, thin: bool = True) -> None: + def __init__(self, *, thin: bool = True, version: str | None = None) -> None: self.in_transaction = False self.thin = thin + if version is not None: + self.version = version async def run_pipeline(self, *_args: object, **_kwargs: object) -> list[object]: return [] @@ -384,3 +386,21 @@ def test_dispatch_execute_script_filters_named_binds_per_statement() -> None: assert cursor.execute.call_count == 2 assert cursor.execute.call_args_list[0].args[1] == {"a": 1} assert cursor.execute.call_args_list[1].args[1] == {"b": 2} + + +async def test_oracle_async_driver_pipeline_native_supported_zero_rtt_from_connection_version( + monkeypatch: pytest.MonkeyPatch, +) -> None: + conn = _AsyncPipelineConnection(thin=True, version="23.4.0.24.05") + driver = OracleAsyncDriver( + connection=cast("OracleAsyncConnection", conn), statement_config=default_statement_config, driver_features={} + ) + monkeypatch.setattr(OracleAsyncDriver, "_detect_oracledb_version", lambda _self: (4, 1, 0)) + get_version_mock = AsyncMock( + side_effect=AssertionError("SQL get_version should not be called when connection.version is set") + ) + monkeypatch.setattr(OracledbAsyncDataDictionary, "get_version", get_version_mock) + + assert await driver._pipeline_native_supported() is True + assert driver._pipeline_support_reason is None + get_version_mock.assert_not_called() diff --git a/tests/unit/builder/test_vector_functions.py b/tests/unit/builder/test_vector_functions.py index 70f30a867..eb7ba9084 100644 --- a/tests/unit/builder/test_vector_functions.py +++ b/tests/unit/builder/test_vector_functions.py @@ -8,7 +8,13 @@ from sqlglot import exp from sqlspec import sql -from sqlspec.builder import Column, VectorDistance +from sqlspec.builder import ( + Column, + VectorDistance, + build_oracle_hnsw_subquery, + build_oracle_hybrid_rrf_query, + build_oracle_hybrid_rrf_sql, +) from sqlspec.builder._vector_distance import is_vector_distance_expression, vector_distance_metric @@ -567,3 +573,43 @@ def test_distance_with_null_handling() -> None: assert "IS NULL" in stmt.sql assert "<->" in stmt.sql + + +def test_oracle_select_approx_limit_and_hybrid_rrf_builders() -> None: + dist_col = Column("embedding").vector_distance([0.1, 0.2], metric="cosine") + query = ( + sql + .select("id", "content_text", dist_col.alias("dist")) + .from_("docs") + .order_by("dist") + .approx_limit(10, target_accuracy=95) + ) + stmt = query.build(dialect="oracle") + assert "FETCH APPROX FIRST 10 ROWS ONLY WITH TARGET ACCURACY 95" in stmt.sql + + hnsw_sql = build_oracle_hnsw_subquery( + "docs", + vector_column="embedding", + query_vector_bind="query_vec", + metric="COSINE", + top_k_bind="vec_top_k", + where_clause="app_name = :app_name", + target_accuracy=90, + ) + assert "VECTOR_DISTANCE(embedding, :query_vec, COSINE)" in hnsw_sql + assert "FETCH APPROX FIRST :vec_top_k ROWS ONLY WITH TARGET ACCURACY 90" in hnsw_sql + + rrf_sql = build_oracle_hybrid_rrf_sql( + "docs", + id_column="id", + text_column="content_text", + vector_column="embedding", + select_columns=("id", "content_text"), + where_clause="app_name = :app_name", + ) + assert "CONTAINS(content_text, :fts_query, 1) > 0" in rrf_sql + assert "FULL OUTER JOIN fts_candidates" in rrf_sql + assert "ORDER BY rrf.rrf_score DESC" in rrf_sql + + rrf_stmt = build_oracle_hybrid_rrf_query("docs", vector_column="embedding", text_column="content_text") + assert "rrf_score" in rrf_stmt