# Copyright (c) Microsoft. All rights reserved. # Copyright (c) 2025, Oracle Corporation. All rights reserved. # Standard Library import array import ast import datetime import logging import re import sys import uuid from collections.abc import AsyncIterable, Mapping, Sequence from typing import Any, ClassVar, Final, Generic, TypeVar # Third-party Libraries import numpy as np import oracledb from pydantic import Field, PrivateAttr, SecretStr # Semantic Kernel AI and Data Abstractions from semantic_kernel.connectors.ai.embedding_generator_base import ( EmbeddingGeneratorBase, ) from semantic_kernel.data.vector import ( DistanceFunction, GetFilteredRecordOptions, IndexKind, KernelSearchResults, SearchType, VectorSearch, VectorSearchOptions, VectorSearchResult, VectorStore, VectorStoreCollection, VectorStoreCollectionDefinition, VectorStoreField, ) # Semantic Kernel Exceptions from semantic_kernel.exceptions import ( MemoryConnectorConnectionException, VectorSearchExecutionException, VectorStoreOperationException, ) # Semantic Kernel Utilities & Config from semantic_kernel.kernel_pydantic import KernelBaseSettings from semantic_kernel.kernel_types import OneOrMany from semantic_kernel.utils.feature_stage_decorator import release_candidate oracledb.defaults.fetch_lobs = False # Compatibility: @override decorator # Python 3.12+ has typing.override natively, # for older versions use typing_extensions. if sys.version_info >= (3, 12): from typing import override else: from typing_extensions import override # Type variables for generics TKey = TypeVar("TKey", bound=str | int | uuid.UUID) TModel = TypeVar("TModel") # Module-level logger logger = logging.getLogger(__name__) # Explicit module exports: # Only expose high-level entry points; keep helpers internal. __all__ = [ "OracleCollection", "OracleSettings", "OracleStore", ] # Environment Variable keys for Oracle DB configuration POOL_MIN_ENV_VAR: Final[str] = "ORACLE_POOL_MIN" POOL_MAX_ENV_VAR: Final[str] = "ORACLE_POOL_MAX" POOL_INCREMENT_ENV_VAR: Final[str] = "ORACLE_POOL_INCREMENT" # Maps Semantic Kernel enums to Oracle SQL keywords DISTANCE_FUNCTION_MAP: Final[dict[DistanceFunction, str]] = { DistanceFunction.COSINE_DISTANCE: "COSINE", DistanceFunction.EUCLIDEAN_DISTANCE: "EUCLIDEAN", DistanceFunction.EUCLIDEAN_SQUARED_DISTANCE: "EUCLIDEAN_SQUARED", DistanceFunction.DOT_PROD: "DOT", DistanceFunction.HAMMING: "HAMMING", DistanceFunction.MANHATTAN: "MANHATTAN", DistanceFunction.DEFAULT: "COSINE", } # Maps index kind enums to Oracle SQL keywords INDEX_KIND_MAP: Final[dict[IndexKind, str]] = {IndexKind.HNSW: "HNSW", IndexKind.IVF_FLAT: "IVF"} # Maps dtype strings to NumPy types and array codes KIND_MAP = { "float32": (np.float32, "f"), "float": (np.float64, "d"), "float64": (np.float64, "d"), "int8": (np.int8, "b"), "uint8": (np.uint8, "B"), "binary": (np.uint8, "B"), } VECTOR_TYPE_MAPPING: dict[str, str] = { "uint8": "BINARY", "int8": "INT8", "float": "FLOAT64", "float32": "FLOAT32", "float64": "FLOAT64", "binary": "BINARY", } # helper methods def _map_scalar_field_type_to_oracle(field_type_str: str) -> str | None: """Map a Semantic Kernel record *scalar field* type (key or data). This is used when generating table DDL for all non-vector fields, including primary keys and data fields. Args: field_type_str: The field type as a string (e.g., 'int', 'float', 'uuid', 'dict', 'bytes'). Returns: The corresponding Oracle SQL column type as a string, or None if there is no known mapping. """ type_mapping = { "bool": "BOOLEAN", "byte": "NUMBER(3)", "int": "NUMBER(10)", "long": "NUMBER(19)", "float": "BINARY_FLOAT", "double": "BINARY_DOUBLE", "Decimal": "NUMBER", "UUID": "RAW(16)", "date": "DATE", "datetime": "TIMESTAMP", "timedelta": "INTERVAL DAY TO SECOND", "bytes": "RAW(2000)", "dict": "JSON", "clob": "CLOB", "blob": "BLOB", } list_pattern = re.compile(r"list\[(.*)\]") if list_pattern.match(field_type_str): return "JSON" dict_pattern = re.compile(r"dict\[(.*?),\s*(.*?)\]") if dict_pattern.match(field_type_str): return "JSON" str_match = re.match(r"str(?:\((\d+)\))?$", field_type_str) if str_match: size = str_match.group(1) or "4000" return f"VARCHAR2({size})" if field_type_str not in type_mapping: raise VectorStoreOperationException(f"Unsupported scalar field type: {field_type_str}") return type_mapping.get(field_type_str) def _sk_vector_element_to_oracle(field_type_str: str) -> str | None: """Convert a Semantic Kernel vector element type string to an Oracle VECTOR element type string.""" list_pattern = re.compile(r"(?i)^list\[(.*)\]$") field_type = field_type_str.strip() # Iteratively unwrap list[...] until no longer matches while True: match = list_pattern.match(field_type) if not match: break field_type = match.group(1).strip() # Return final mapped type if available return VECTOR_TYPE_MAPPING.get(field_type) class BindCounter: """Helper class to generate unique bind variable names for SQL queries.""" def __init__(self, start: int = 1): self._index = start def next_bind(self) -> str: name = f"bind_val{self._index}" self._index += 1 return name # region: Oracle Settings @release_candidate class OracleSettings(KernelBaseSettings): """Oracle connector settings. This class is used to configure the Oracle connection pool and related options for the Oracle vector or memory store connectors. It supports both standard username/password authentication and wallet-based secure connections, and is compatible with Oracle's native async driver (python-oracledb) for efficient async operations. The settings align with common Oracle client environment variables such as ORACLE_USER, ORACLE_PASSWORD, and ORACLE_CONNECTION_STRING, while following the Semantic Kernel convention for configuration through environment variables or explicit keyword arguments. Args: user: Oracle database username. (Env var ORACLE_USER) password: Oracle database password. (Env var ORACLE_PASSWORD) connection_string: Full Oracle connection string, for example: "host:port/service_name". (Env var ORACLE_CONNECTION_STRING) min: Minimum number of connections in the pool. (Env var ORACLE_MIN) max: Maximum number of connections in the pool. (Env var ORACLE_MAX) increment: Number of connections to add when the pool grows. (Env var ORACLE_INCREMENT) wallet_location: Path to the Oracle wallet directory for wallet-based authentication. (Env var ORACLE_WALLET_LOCATION) wallet_password: Password for the Oracle wallet. (Env var ORACLE_WALLET_PASSWORD) connection_pool: Optional preconfigured AsyncConnectionPool instance. """ env_prefix: ClassVar[str] = "ORACLE_" user: str | None = None password: SecretStr | None = None connection_string: str | None = None min: int | None = Field(default=None, validation_alias=POOL_MIN_ENV_VAR) max: int | None = Field(default=None, validation_alias=POOL_MAX_ENV_VAR) increment: int | None = Field(default=None, validation_alias=POOL_INCREMENT_ENV_VAR) wallet_location: str | None = None wallet_password: SecretStr | None = None _connection_pool: oracledb.AsyncConnectionPool | None = PrivateAttr(default=None) async def create_connection_pool(self, **kwargs: Any) -> oracledb.AsyncConnectionPool: """Creates an async Oracle connection pool.""" try: # Create pool with extra user-supplied kwargs self._connection_pool = oracledb.create_pool_async( user=self.user, password=self.password.get_secret_value() if self.password else None, dsn=self.connection_string, wallet_location=self.wallet_location, wallet_password=self.wallet_password.get_secret_value() if self.wallet_password else None, min=self.min, max=self.max, increment=self.increment, **kwargs, # extra pool params ) except Exception as err: raise MemoryConnectorConnectionException("Error creating Oracle connection pool.") from err return self._connection_pool # region: Oracle Collections @release_candidate class OracleCollection( VectorStoreCollection[TKey, TModel], VectorSearch[TKey, TModel], Generic[TKey, TModel], ): """Oracle implementation of VectorStoreCollection + VectorSearch.""" connection_pool: oracledb.AsyncConnectionPool | None = None db_schema: str | None = None pool_args: dict[str, Any] | None = None supported_key_types: ClassVar[set[str] | None] = {"str", "int", "UUID"} supported_vector_types: ClassVar[set[str] | None] = set(VECTOR_TYPE_MAPPING.keys()) supported_search_types: ClassVar[set[SearchType]] = {SearchType.VECTOR} _distance_column_name: str = "SIMILARITY_SKOVS" def __init__( self, record_type: type[TModel], collection_name: str | None = None, definition: VectorStoreCollectionDefinition | None = None, embedding_generator: EmbeddingGeneratorBase | None = None, connection_pool: oracledb.AsyncConnectionPool | None = None, db_schema: str | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None, settings: OracleSettings | None = None, pool_args: dict[str, Any] | None = None, **kwargs: Any, ): """Initialize the collection. Args: record_type: The type of the data model. definition: The data model definition. collection_name: The name of the collection, which corresponds to the table name. embedding_generator: The embedding generator. connection_pool: The connection pool. db_schema: The database schema. env_file_path: Use the environment settings file as a fallback to environment variables. env_file_encoding: The encoding of the environment settings file. settings: The settings for creating a new connection pool. If not provided, the settings will be created from the environment. pool_args: Optional dict of additional arguments to configure the connection pool (e.g., timeout, ping_interval). **kwargs: Additional arguments. """ # Build settings from env if we need to manage our own pool settings = settings or OracleSettings( env_file_path=env_file_path, env_file_encoding=env_file_encoding, ) managed_client = False # Create pool if missing if connection_pool is None: try: pool_args = pool_args or {} connection_pool = oracledb.create_pool_async( user=settings.user, password=settings.password.get_secret_value() if settings.password else None, dsn=settings.connection_string, wallet_location=settings.wallet_location, wallet_password=settings.wallet_password.get_secret_value() if settings.wallet_password else None, min=settings.min, max=settings.max, increment=settings.increment, **pool_args, ) except Exception as err: raise MemoryConnectorConnectionException("Error creating Oracle connection pool.") from err managed_client = True else: managed_client = False super().__init__( collection_name=collection_name, # type: ignore record_type=record_type, definition=definition, # type: ignore embedding_generator=embedding_generator, connection_pool=connection_pool, # type: ignore db_schema=db_schema, # type: ignore settings=settings, # type: ignore # This controls whether the connection pool is managed by the collection # in the __aenter__ and __aexit__ methods. managed_client=managed_client, ) # Compute UUID field names once self._uuid_fields = [ field.storage_name or field.name for field in (*self.definition.data_fields, self.definition.key_field) if field.type_ == "UUID" ] # Validate key/data/vector field once per life-cycle key_field = self.definition.key_field key_field_name = key_field.storage_name or key_field.name self._validate_identifiers(key_field_name) for field in self.definition.data_fields: data_field_name = field.storage_name or field.name self._validate_identifiers(data_field_name) for field in self.definition.vector_fields: vector_field_name = field.storage_name or field.name self._validate_identifiers(vector_field_name) dtype = field.type_ or "float32" if dtype not in KIND_MAP: raise VectorStoreOperationException( f"Unsupported dtype '{dtype}' for field '{field.name}'. " f"Supported dtypes: {', '.join(KIND_MAP.keys())}" ) @override async def __aenter__(self) -> "OracleCollection": return self @override async def __aexit__(self, *args: Any) -> None: # Only close the connection pool if it was created by the collection itself. if self.managed_client and self.connection_pool: try: await self.connection_pool.close() except Exception as e: logger.warning("Error closing Oracle connection pool: %s", e) finally: self.connection_pool = None self.managed_client = False def _check_pool(self) -> oracledb.AsyncConnectionPool: """Ensure that the connection pool is available, otherwise raise a consistent error.""" if self.connection_pool is None: raise VectorStoreOperationException("Collection has no connection pool.") return self.connection_pool @override def _deserialize_store_models_to_dicts(self, records: Sequence[Any], **kwargs: Any) -> Sequence[dict[str, Any]]: """Deserialize the store models to a list of dicts. Pass the records through without modification.""" return records def _full_table_name(self) -> str: """Return the fully qualified table name with optional schema prefix, quoted.""" self._validate_identifiers(self.collection_name) if self.db_schema: self._validate_identifiers(self.db_schema) return f'"{self.db_schema}"."{self.collection_name}"' return f'"{self.collection_name}"' async def _get_connection(self): """Acquire a connection from the pool, ensuring input/output type handlers are always set. Raises VectorStoreOperationException if no connection pool is configured. """ pool = self._check_pool() conn = await pool.acquire() # if you only want to set these in certain circumstances you can add a parameter to # the function for that purpose but it should be safe to set at all times conn.inputtypehandler = self._input_type_handler conn.outputtypehandler = self._output_type_handler return conn def _input_type_handler(self, cursor, value, arraysize): """Map Python types to Oracle bind variables with correct DB types.""" if isinstance(value, np.ndarray): return cursor.var(oracledb.DB_TYPE_VECTOR, arraysize=arraysize, inconverter=self._numpy_converter_in) if isinstance(value, uuid.UUID): return cursor.var(oracledb.DB_TYPE_RAW, arraysize=arraysize, inconverter=lambda v: v.bytes) if isinstance(value, (dict, list)): return cursor.var(oracledb.DB_TYPE_JSON, arraysize=arraysize) if isinstance(value, bytes): return cursor.var(oracledb.DB_TYPE_RAW, arraysize=arraysize) return None def _numpy_converter_in(self, value): """Convert a NumPy array into a Python array.array compatible with Oracle DB_TYPE_VECTOR.""" dtype_name = value.dtype.name np_dtype, code = KIND_MAP[dtype_name] value = value.astype(np_dtype, copy=False) return array.array(code, value) def _output_type_handler(self, cursor, metadata): """Map Oracle DB column types to Python-native objects during fetch operations.""" # VECTOR columns to list if metadata.type_code == oracledb.DB_TYPE_VECTOR: return cursor.var(oracledb.DB_TYPE_VECTOR, arraysize=cursor.arraysize, outconverter=lambda arr: list(arr)) # RAW to UUID if metadata.type_code == oracledb.DB_TYPE_RAW and metadata.name in self._uuid_fields: return cursor.var( oracledb.DB_TYPE_RAW, arraysize=cursor.arraysize, outconverter=lambda b: uuid.UUID(bytes=b) if b is not None else None, ) return None @override def _serialize_dicts_to_store_models(self, records: Sequence[dict[str, Any]], **kwargs: Any) -> Sequence[Any]: """Serialize a list of dicts of the data to the store model. Pass the records through without modification.""" return records def _validate_identifiers(self, name: str) -> None: """Validate Oracle identifier to disallow embedded double quotes. Since quoted identifiers are not allowed, any double quote is invalid. """ if not name: raise VectorStoreOperationException("Identifier cannot be empty") if '"' in name: raise VectorStoreOperationException(f"Invalid identifier with quotes: {name}") def _build_check_table_exists_query(self) -> tuple[str, dict[str, str]]: """Build SQL + bind variables for checking table existence. - If schema is provided, query ALL_TABLES. - If no schema, query USER_TABLES. """ if self.db_schema: sql = """ SELECT 1 FROM all_tables WHERE owner = :owner AND table_name = :tbl """ bind_vars = { "owner": self.db_schema, "tbl": self.collection_name, } else: sql = """ SELECT 1 FROM user_tables WHERE table_name = :tbl """ bind_vars = {"tbl": self.collection_name} return sql, bind_vars @override async def collection_exists(self, **kwargs: Any) -> bool: """Return True if the table backing this collection exists.""" pool = self._check_pool() sql, binds = self._build_check_table_exists_query() async with pool.acquire() as conn: result = await conn.fetchone(sql, binds) return result is not None @override async def ensure_collection_deleted(self, **kwargs: Any) -> None: """Deletes collection if it exists.""" pool = self._check_pool() tbl = self._full_table_name() drop_sql = f"DROP TABLE IF EXISTS {tbl} CASCADE CONSTRAINTS PURGE" async with pool.acquire() as conn: await conn.execute(drop_sql) logger.debug(f"Oracle table '{self.collection_name}' deleted successfully.") def _build_create_table_query( self, *, table: str, key_field: VectorStoreField, data_fields: list[VectorStoreField], vector_fields: list[VectorStoreField], ) -> str: col_lines: list[str] = [] if not key_field.type_: raise VectorStoreOperationException(f"Type missing for key field '{key_field.name}'") pk_name = key_field.storage_name or key_field.name col_lines.append(f'"{pk_name}" {_map_scalar_field_type_to_oracle(key_field.type_)} PRIMARY KEY') for f in data_fields: if not f.type_: raise VectorStoreOperationException(f"Type missing for data field '{f.name}'") sql_type = _map_scalar_field_type_to_oracle(f.type_) if sql_type is None: raise VectorStoreOperationException(f'Unsupported Oracle type for field "{f.name}" ({f.type_})') col_name = f.storage_name or f.name col_lines.append(f'"{col_name}" {sql_type}') for f in vector_fields: if not f.type_ or f.dimensions is None: raise VectorStoreOperationException(f"Vector field '{f.name}' missing type or dimensions") col_name = f.storage_name or f.name col_lines.append(f'"{col_name}" VECTOR({f.dimensions} , {_sk_vector_element_to_oracle(f.type_)})') columns_sql = ",\n ".join(col_lines) return f"CREATE TABLE IF NOT EXISTS {table} (\n {columns_sql}\n)" def _create_vector_index(self, table_name: str, vector_field: VectorStoreField) -> str | None: """Build a CREATE VECTOR INDEX statement for an Oracle vector column using HNSW or IVF indexing.""" if vector_field.index_kind not in INDEX_KIND_MAP: logger.warning( f"Index kind '{vector_field.index_kind}' is not supported. " "Please set the index kind in the vector field definition." ) return None if not vector_field.distance_function or vector_field.distance_function not in DISTANCE_FUNCTION_MAP: raise VectorStoreOperationException( f"Distance function '{vector_field.distance_function}' is not supported. " "Please set the distance function in the vector field definition." ) distance = DISTANCE_FUNCTION_MAP[vector_field.distance_function] column = vector_field.storage_name or vector_field.name last_token = table_name.split(".")[-1] base_table = last_token.strip('"') index_name = f'"{base_table}_{column}_idx"' index_kind = vector_field.index_kind if index_kind == IndexKind.HNSW or index_kind == IndexKind.DEFAULT: index_subtype = f"ORGANIZATION INMEMORY NEIGHBOR GRAPH DISTANCE {distance} " else: # IndexKind.IVF_FLAT index_subtype = f"ORGANIZATION NEIGHBOR PARTITIONS DISTANCE {distance} " return f'CREATE VECTOR INDEX IF NOT EXISTS {index_name}\nON {table_name} ("{column}")\n{index_subtype}' def _create_data_index(self, table_name: str, field) -> str | None: """Build a CREATE INDEX statement for a single data field if it is indexable. Returns the SQL string or None if no index should be created. """ if not getattr(field, "is_indexed", False): return None oracle_type = _map_scalar_field_type_to_oracle(field.type_) if field.type_ else None if oracle_type and oracle_type.lower() != "json": col = field.storage_name or field.name last_token = table_name.split(".")[-1] base_table = last_token.strip('"') index_name = f'"{base_table}_{col}_idx"' return f'CREATE INDEX {index_name} ON {table_name} ("{col}" ASC)' return None @override async def ensure_collection_exists(self, **kwargs: Any) -> None: """Create the table (and vector indexes) if not existing.""" pool = self._check_pool() tbl = self._full_table_name() create_sql = self._build_create_table_query( table=tbl, key_field=self.definition.key_field, data_fields=self.definition.data_fields, vector_fields=self.definition.vector_fields, ) # Combine create table + vector + data indexes statements = ( [create_sql] + [stmt for vf in self.definition.vector_fields if (stmt := self._create_vector_index(tbl, vf))] + [stmt for field in self.definition.data_fields if (stmt := self._create_data_index(tbl, field))] ) async with pool.acquire() as conn: for statement in statements: await conn.execute(statement) await conn.commit() logger.info(f"Oracle table '{self.collection_name}' created successfully.") @override async def _inner_delete(self, keys: Sequence[TKey], **kwargs: Any) -> None: """Delete the records whose primary keys are in ``keys``.""" pool = self._check_pool() if not keys: return tbl = self._full_table_name() pk_col = self.definition.key_field.storage_name or self.definition.key_field.name delete_sql = f'DELETE FROM {tbl} WHERE "{pk_col}" = :1' # nosec B608 async with pool.acquire() as conn: binds = [(k,) for k in keys] await conn.executemany(delete_sql, binds) await conn.commit() def _build_select_query( self, table: str, key_field: VectorStoreField, data_fields: list[VectorStoreField], vector_fields: list[VectorStoreField], keys: Sequence[Any] | None, options: GetFilteredRecordOptions | None, include_vectors: bool = False, ) -> tuple[str, list[Any]]: # SELECT clause all_fields = [key_field, *data_fields, *vector_fields] if include_vectors else [key_field, *data_fields] field_lookup = {f.name: f for f in all_fields} select_clause = ", ".join(f'"{f.storage_name or f.name}" AS "{f.name}"' for f in all_fields) sql = f"SELECT {select_clause} FROM {table}" # nosec B608 bind_values: list[Any] = [] # WHERE clause by keys if keys: placeholders = ", ".join(f":{i + 1}" for i in range(len(keys))) sql += f' WHERE "{key_field.storage_name or key_field.name}" IN ({placeholders})' bind_values.extend(keys) # ORDER BY if options and options.order_by: parts: list[str] = [] for logical_name, asc in options.order_by.items(): field = field_lookup[logical_name] field_name = field.storage_name or field.name direction = "ASC" if asc else "DESC" parts.append(f'"{field_name}" {direction}') sql += " ORDER BY " + ", ".join(parts) # Pagination if options: if options.skip is not None: sql += f" OFFSET {options.skip} ROWS" if options.top is not None: sql += f" FETCH NEXT {options.top} ROWS ONLY" return sql, bind_values @override async def _inner_get( self, keys: Sequence[TKey] | None = None, options: GetFilteredRecordOptions | None = None, **kwargs: Any, ) -> OneOrMany[dict[str, Any]] | None: """Retrieve one or more records from the Oracle table. Returns: OneOrMany[dict[str, Any]] | None - A single record (dict) when exactly one row matches - A list of dicts when multiple rows match - ``None`` when no rows match """ pool = self._check_pool() if not keys and options is None: return None table = self._full_table_name() include_vectors = kwargs.get("include_vectors", True) q, binds = self._build_select_query( table, self.definition.key_field, self.definition.data_fields, self.definition.vector_fields, keys, options, include_vectors, ) async with pool.acquire() as conn: conn.outputtypehandler = self._output_type_handler rows = await conn.fetchall(q, binds) # Build column list once: key, data, then vector fields columns = [self.definition.key_field.name] + [f.name for f in self.definition.data_fields] if include_vectors: columns = columns + [f.name for f in self.definition.vector_fields] if not rows: return None # build list of dict records records: list[dict[str, Any]] = [] for row in rows: record = {col: val for col, val in zip(columns, row)} records.append(record) return records[0] if len(records) == 1 else records def _convert_dict_to_row( self, record: Mapping[str, Any], fields: Sequence[VectorStoreField], ) -> tuple[Any, ...]: """Convert an in-memory record (dict) into a positional tuple ready for executemany() with Oracle.""" row: list[Any] = [] for field in fields: column_name = field.name value = record.get(column_name) if value is None: row.append(None) continue if field.field_type == "vector" and isinstance(value, (list)): _, code = KIND_MAP[field.type_] # type: ignore[index] value = array.array(code, value) row.append(value) continue row.append(value) return tuple(row) def _build_single_merge_query( self, table_name: str, key_field: VectorStoreField, data_fields: list[VectorStoreField], vector_fields: list[VectorStoreField], ) -> str: """Build a parameterised MERGE statement for Oracle. One executemany() call executes this MERGE once per record. """ all_fields = [key_field, *data_fields, *vector_fields] src_bindings = ",\n ".join( f':{idx + 1} AS "{field.storage_name or field.name}"' for idx, field in enumerate(all_fields) ) # When matched then update data update_clause = ",\n ".join( f't."{field.storage_name or field.name}" = s."{field.storage_name or field.name}"' for field in data_fields + vector_fields ) # When not matched then insert data insert_columns = ", ".join(f'"{field.storage_name or field.name}"' for field in all_fields) insert_values = ", ".join(f's."{field.storage_name or field.name}"' for field in all_fields) merge_sql = f""" MERGE INTO {table_name} t USING ( SELECT {src_bindings} FROM dual ) s ON ( t."{key_field.storage_name or key_field.name}" = s."{key_field.storage_name or key_field.name}" ) WHEN MATCHED THEN UPDATE SET {update_clause} WHEN NOT MATCHED THEN INSERT ({insert_columns}) VALUES ({insert_values}) """ # nosec B608 return merge_sql.strip() @override async def _inner_upsert( self, records: Sequence[dict[str, Any]], **kwargs: Any, ) -> Sequence[TKey]: pool = self._check_pool() if not records: return [] key_field = self.definition.key_field data_fields = self.definition.data_fields vector_fields = self.definition.vector_fields table_name = self._full_table_name() ordered_fields = [key_field, *data_fields, *vector_fields] query = self._build_single_merge_query(table_name, key_field, data_fields, vector_fields) async with pool.acquire() as conn: conn.inputtypehandler = self._input_type_handler binds = [self._convert_dict_to_row(record, ordered_fields) for record in records] await conn.executemany(query, binds) await conn.commit() return [record[key_field.name] for record in records] def _build_search_query( self, table: str, key_field: VectorStoreField, data_fields: list[VectorStoreField], vector_fields: list[VectorStoreField], vector: Sequence[float | int], vector_field: VectorStoreField, options: VectorSearchOptions, filter_clause: str | None, ) -> tuple[str, list[Any], list[str]]: bind_values = [] bind_index = 1 # Vector bind placeholder vector_placeholder = f":{bind_index}" bind_values.append(vector) bind_index += 1 # Build TO_VECTOR() expression from field metadata dim = vector_field.dimensions or len(vector) # Normalize user-provided dtype and map to Oracle-supported VECTOR types. raw_dtype = (vector_field.type_ or "float32").lower() dtype = VECTOR_TYPE_MAPPING.get(raw_dtype, "FLOAT32") to_vector_expr = f"TO_VECTOR({vector_placeholder}, {dim}, {dtype})" # Fields to SELECT select_fields = [f'"{key_field.storage_name or key_field.name}"'] + [ f'"{field.storage_name or field.name}"' for field in data_fields ] if options.include_vectors: select_fields += [f'"{f.storage_name or f.name}"' for f in vector_fields] select_clause = ", ".join(select_fields) # Choose distance function if not vector_field.distance_function or vector_field.distance_function not in DISTANCE_FUNCTION_MAP: raise VectorStoreOperationException( f"Distance function '{vector_field.distance_function}' is not supported. " "Please set the distance function in the vector field definition." ) distance_fn = DISTANCE_FUNCTION_MAP[vector_field.distance_function] distance_expr = ( f'VECTOR_DISTANCE("{vector_field.storage_name or vector_field.name}", ' f'{to_vector_expr}, {distance_fn}) AS "{self._distance_column_name}"' ) # Final SQL assembly sql = f""" SELECT {select_clause}, {distance_expr} FROM {table} """.strip() # nosec B608 if filter_clause: sql += f"\nWHERE {filter_clause}" sql += f'\nORDER BY "{self._distance_column_name}" ASC' if options: if options.skip is not None: sql += f" OFFSET {options.skip} ROWS" if options.top is not None: sql += f" FETCH NEXT {options.top} ROWS ONLY" return sql, bind_values, select_fields async def _fetch_records(self, sql: str, binds: list[Any]) -> AsyncIterable[dict[str, Any]]: """Execute the SQL with binds and yield rows as dictionaries mapping column name to value. Uses zip() for clean row-to-dict mapping. """ async with await self._get_connection() as conn: conn.inputtypehandler = self._input_type_handler conn.outputtypehandler = self._output_type_handler with conn.cursor() as cur: await cur.execute(sql, binds) col_names = [d.name for d in cur.description] async for row in cur: yield dict(zip(col_names, row)) @override def _get_record_from_result(self, result: dict[str, Any]) -> dict[str, Any]: """Pass-through: assumes result is already normalized and cleaned.""" return result @override def _get_score_from_result(self, result: Any) -> float | None: return result.get("SIMILARITY_SKOVS", None) @override async def _inner_search( self, search_type: SearchType, options: VectorSearchOptions, values: Any | None = None, vector: Sequence[float | int] | None = None, **kwargs: Any, ) -> KernelSearchResults[VectorSearchResult[TModel]]: # Oracle does not support accurate total_count if options.include_total_count: logger.warning("`include_total_count=True` is not supported in OracleVectorStore and will be ignored.") # Build SQL & bind parameters query, bind, _ = await self._inner_search_vector(options, values, vector, **kwargs) # Always run streaming search (even if include_total_count=True) stream: AsyncIterable[dict[str, Any]] = self._fetch_records(query, bind) return KernelSearchResults( results=self._get_vector_search_results_from_results(stream, options), total_count=None, # always None in Oracle ) async def _inner_search_vector( self, options: VectorSearchOptions, values: Any | None, vector: Sequence[float | int] | None, **kwargs: Any, ) -> tuple[str, list[Any], list[str]]: if vector is None: vector = await self._generate_vector_from_values(values, options) if vector is None or len(vector) == 0: raise VectorSearchExecutionException("Vector cannot be None or empty.") if options.vector_property_name is None: raise VectorStoreOperationException("vector_property_name cannot be None") vector_field = next( field for field in self.definition.vector_fields if field.name == options.vector_property_name ) dtype = vector_field.type_ if vector_field.type_ else "float32" if isinstance(vector, (list)): _, code = KIND_MAP[dtype] vector = array.array(code, vector) table = self._full_table_name() # When building filter: parsed_filter = self._build_filter(options.filter) filter_clause: str | None = None filter_binds: list[Any] = [] if parsed_filter is not None: filter_clause = parsed_filter[0] filter_binds.extend(parsed_filter[1].values()) query, bind_values, columns = self._build_search_query( table, self.definition.key_field, self.definition.data_fields, self.definition.vector_fields, vector, vector_field, options, filter_clause, ) # Append filter binds after vector bind_values.extend(filter_binds) return query, bind_values, columns @override def _lambda_parser(self, node: ast.AST, bind_counter: BindCounter | None = None) -> Any: """Parse a lambda AST node and return a tuple: (sql_expression, bind_values_dict). Uses bind variables for all scalar values, including dates. """ if bind_counter is None: bind_counter = BindCounter() bind_dict: dict[str, Any] = {} match node: # Comparisons case ast.Compare(): # IS / IS NOT NULL if isinstance(node.ops[0], (ast.Is, ast.IsNot)): left_sql, left_bind = self._lambda_parser(node.left, bind_counter) bind_dict.update(left_bind) right = node.comparators[0] if isinstance(right, ast.Constant) and right.value is None: op_str = "IS" if isinstance(node.ops[0], ast.Is) else "IS NOT" return (f"{left_sql} {op_str} NULL", bind_dict) raise NotImplementedError("IS comparison only supports None/NULL checks") # Chained comparisons: a <= b < c if len(node.ops) > 1: values = [] for idx in range(len(node.ops)): left_node = node.left if idx == 0 else node.comparators[idx - 1] right_node = node.comparators[idx] op = node.ops[idx] expr, binds = self._lambda_parser( ast.Compare(left=left_node, ops=[op], comparators=[right_node]), bind_counter ) values.append(expr) bind_dict.update(binds) return (f"({' AND '.join(values)})", bind_dict) # Single comparison left_sql, left_bind = self._lambda_parser(node.left, bind_counter) right_sql, right_bind = self._lambda_parser(node.comparators[0], bind_counter) bind_dict.update(left_bind) bind_dict.update(right_bind) op = node.ops[0] if isinstance(op, ast.Eq): return (f"{left_sql} = {right_sql}", bind_dict) if isinstance(op, ast.NotEq): return (f"{left_sql} <> {right_sql}", bind_dict) if isinstance(op, ast.Gt): return (f"{left_sql} > {right_sql}", bind_dict) if isinstance(op, ast.GtE): return (f"{left_sql} >= {right_sql}", bind_dict) if isinstance(op, ast.Lt): return (f"{left_sql} < {right_sql}", bind_dict) if isinstance(op, ast.LtE): return (f"{left_sql} <= {right_sql}", bind_dict) if isinstance(op, ast.In): return (f"{left_sql} IN {right_sql}", bind_dict) if isinstance(op, ast.NotIn): return (f"{left_sql} NOT IN {right_sql}", bind_dict) raise NotImplementedError(f"Unsupported comparison operator: {type(op)}") # Boolean operations case ast.BoolOp(): parts = [] for v in node.values: sql, binds = self._lambda_parser(v, bind_counter) parts.append(sql) bind_dict.update(binds) if isinstance(node.op, ast.And): return (f"({' AND '.join(parts)})", bind_dict) if isinstance(node.op, ast.Or): return (f"({' OR '.join(parts)})", bind_dict) raise NotImplementedError(f"Unsupported BoolOp: {type(node.op)}") # Unary operations case ast.UnaryOp(): if isinstance(node.op, ast.Not): operand_sql, operand_bind = self._lambda_parser(node.operand, bind_counter) bind_dict.update(operand_bind) return (f"NOT ({operand_sql})", bind_dict) raise NotImplementedError(f"Unsupported UnaryOp: {type(node.op)}") # Handling attribute or name nodes (fields) case ast.Attribute(): if node.attr not in self.definition.storage_names: raise VectorStoreOperationException(f"Field '{node.attr}' not in data model.") return (f'"{node.attr}"', {}) case ast.Name(): if node.id not in self.definition.storage_names: raise VectorStoreOperationException(f"Field '{node.id}' not in data model.") return (f'"{node.id}"', {}) # Constants (scalar values) used in this module case ast.Constant(): val = node.value if val is None: return ("NULL", {}) bind_name = bind_counter.next_bind() return (f":{bind_name}", {bind_name: val}) # Lists (for IN operator) case ast.List(): parts = [] for elt in node.elts: elt_sql, elt_bind = self._lambda_parser(elt, bind_counter) parts.append(elt_sql) bind_dict.update(elt_bind) return (f"({', '.join(parts)})", bind_dict) # Function calls case ast.Call(): # Supported methods in this block: contains, startswith, endswith, between if isinstance(node.func, ast.Attribute): obj_sql, obj_bind = self._lambda_parser(node.func.value, bind_counter) bind_dict.update(obj_bind) sql_args = [] for arg in node.args: arg_sql, arg_bind = self._lambda_parser(arg, bind_counter) sql_args.append(arg_sql) bind_dict.update(arg_bind) method = node.func.attr if method == "contains" and len(sql_args) == 1: return (f"LOWER({obj_sql}) LIKE LOWER('%' || {sql_args[0]} || '%')", bind_dict) if method == "startswith" and len(sql_args) == 1: return (f"{obj_sql} LIKE {sql_args[0]} || '%'", bind_dict) if method == "endswith" and len(sql_args) == 1: return (f"{obj_sql} LIKE '%' || {sql_args[0]}", bind_dict) if method == "between" and len(sql_args) == 2: return (f"{obj_sql} BETWEEN {sql_args[0]} AND {sql_args[1]}", bind_dict) # Handle datetime function with arguments (year, month, day) if isinstance(node.func, ast.Name) and node.func.id == "datetime": if not (3 <= len(node.args) <= 6): raise NotImplementedError("datetime() only supports between 3 and 6 integer arguments") def get_const(arg: ast.AST) -> int: if isinstance(arg, ast.Constant) and isinstance(arg.value, int): return arg.value raise NotImplementedError("datetime() arguments must be int constants") year = get_const(node.args[0]) month = get_const(node.args[1]) day = get_const(node.args[2]) hour = get_const(node.args[3]) if len(node.args) > 3 else 0 minute = get_const(node.args[4]) if len(node.args) > 4 else 0 second = get_const(node.args[5]) if len(node.args) > 5 else 0 dt = datetime.datetime(year, month, day, hour, minute, second) bind_name = bind_counter.next_bind() bind_dict[bind_name] = dt return (f":{bind_name}", bind_dict) # Handle date function with arguments (year, month, day) if isinstance(node.func, ast.Name) and node.func.id == "date": if len(node.args) != 3: raise NotImplementedError("date() only supports year, month, day as int constants") def get_const(arg: ast.AST) -> int: if isinstance(arg, ast.Constant) and isinstance(arg.value, int): return arg.value raise NotImplementedError("date() arguments must be int constants") year = get_const(node.args[0]) month = get_const(node.args[1]) day = get_const(node.args[2]) dt = datetime.date(year, month, day) # type: ignore[assignment] bind_name = bind_counter.next_bind() bind_dict[bind_name] = dt return (f":{bind_name}", bind_dict) raise NotImplementedError(f"Unsupported AST node: {type(node)}") # region: Oracle Store @release_candidate class OracleStore(VectorStore): """VectorStore wrapper holding a shared Oracle connection-pool.""" connection_pool: oracledb.AsyncConnectionPool | None = None db_schema: str | None = None env_file_path: str | None = None env_file_encoding: str | None = None def _build_select_table_names_query(self) -> tuple[str, dict[str, str]]: if self.db_schema: sql = """ SELECT table_name FROM all_tables WHERE owner = :schema ORDER BY table_name """ bind_vars = {"schema": self.db_schema} else: sql = """ SELECT table_name FROM user_tables ORDER BY table_name """ bind_vars = {} return sql, bind_vars @override def get_collection( self, record_type: type[TModel], *, collection_name: str | None = None, definition: VectorStoreCollectionDefinition | None = None, embedding_generator: EmbeddingGeneratorBase | None = None, pool_args: dict[str, Any] | None = None, **kwargs: Any, ) -> OracleCollection: """Return an OracleCollection that shares this store's pool. Args: record_type: The type of the records that will be used. collection_name: Name of the table (optional if `definition` supplied). definition: VectorStoreCollectionDefinition describing schema/PK. embedding_generator: Overrides store's default generator. pool_args: Dict of connection-pool overrides (user, min, max, …). **kwargs: Additional keyword arguments passed to OracleCollection. Returns: OracleCollection ready for use (optionally as an async context manager). """ return OracleCollection( record_type=record_type, collection_name=collection_name, db_schema=self.db_schema, definition=definition, connection_pool=self.connection_pool, env_file_path=self.env_file_path, env_file_encoding=self.env_file_encoding, embedding_generator=embedding_generator or self.embedding_generator, pool_args=pool_args, **kwargs, ) @override async def list_collection_names(self, **kwargs: Any) -> Sequence[str]: """Get the names of all collections.""" if self.connection_pool is None: raise VectorStoreOperationException("Store has no connection pool.") sql, binds = self._build_select_table_names_query() async with self.connection_pool.acquire() as conn: rows = await conn.fetchall(sql, binds) return [row[0] for row in rows]