Fixes #27443: add OCI Autonomous Database support for Oracle connector

ingestion/src/metadata/ingestion/source/database/oracle/connection.py

@@ -13,10 +13,18 @@
 Source connection handler
 """
 
+import base64
+import binascii
+import io
 import os
+import shutil
 import sys
+import tempfile
+import weakref
+import zipfile
 from copy import deepcopy
-from typing import Optional
+from pathlib import Path
+from typing import Any, Optional
 from urllib.parse import quote_plus
 
 import oracledb
@@ -28,20 +36,22 @@
     Workflow as AutomationWorkflow,
 )
 from metadata.generated.schema.entity.services.connections.database.oracleConnection import (
-    OracleConnection as OracleConnectionConfig,
-)
-from metadata.generated.schema.entity.services.connections.database.oracleConnection import (
+    OracleAutonomousConnection,
     OracleDatabaseSchema,
     OracleServiceName,
     OracleTNSConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.oracleConnection import (
+    OracleConnection as OracleConnectionConfig,
+)
 from metadata.generated.schema.entity.services.connections.testConnectionResult import (
     TestConnectionResult,
 )
 from metadata.ingestion.connections.builders import (
     create_generic_db_connection,
     get_connection_args_common,
     get_connection_options_dict,
+    init_empty_connection_arguments,
 )
 from metadata.ingestion.connections.connection import BaseConnection
 from metadata.ingestion.connections.secrets import connection_with_options_secrets
@@ -68,24 +78,161 @@
 class OracleConnection(BaseConnection[OracleConnectionConfig, Engine]):
     def __init__(self, connection: OracleConnectionConfig):
         super().__init__(connection)
+        self._wallet_temp_dir: str | None = None
+        self._wallet_cleanup_finalizer: Any = None
+
+    def _set_wallet_temp_dir(self, wallet_temp_dir: str) -> None:
+        self._cleanup_wallet_temp_dir()
+        self._wallet_temp_dir = wallet_temp_dir
+        self._wallet_cleanup_finalizer = weakref.finalize(
+            self,
+            shutil.rmtree,
+            wallet_temp_dir,
+            ignore_errors=True,
+        )
+
+    def _cleanup_wallet_temp_dir(self) -> None:
+        wallet_temp_dir = self._wallet_temp_dir
+        if self._wallet_cleanup_finalizer and self._wallet_cleanup_finalizer.alive:
+            self._wallet_cleanup_finalizer()
+        elif wallet_temp_dir:
+            shutil.rmtree(wallet_temp_dir, ignore_errors=True)
+
+        self._wallet_cleanup_finalizer = None
+        self._wallet_temp_dir = None
+
+    def _is_autonomous_connection(self) -> bool:
+        return isinstance(self.service_connection.oracleConnectionType, OracleAutonomousConnection)
+
+    @staticmethod
+    def _get_autonomous_connection_config(
+        connection_type: OracleAutonomousConnection,
+    ) -> OracleAutonomousConnection:
+        return connection_type
+
+    @staticmethod
+    def _safe_extract_wallet_archive(zip_ref: zipfile.ZipFile, target_dir: str) -> None:
+        target_root = Path(target_dir).resolve()
+
+        for member in zip_ref.infolist():
+            member_path = (target_root / member.filename).resolve()
+
+            if member_path != target_root and target_root not in member_path.parents:
+                raise ValueError("Invalid walletContent. Wallet zip contains unsafe file paths.")
+
+            if member.is_dir():
+                OracleConnection._mkdir_secure_within(member_path, target_root)
+                continue
+
+            OracleConnection._mkdir_secure_within(member_path.parent, target_root)
+            with (
+                zip_ref.open(member, "r") as source_file,
+                open(
+                    member_path,
+                    "wb",
+                    opener=lambda path, flags: os.open(path, flags, 0o600),
+                ) as target_file,
+            ):
+                shutil.copyfileobj(source_file, target_file)
+
+    @staticmethod
+    def _mkdir_secure_within(path: Path, root: Path) -> None:
+        """Create path and any intermediate dirs with 0o700, only within root."""
+        if path == root:
+            return
+        OracleConnection._mkdir_secure_within(path.parent, root)
+        try:
+            path.mkdir(mode=0o700, exist_ok=False)
+        except FileExistsError:
+            return
+        path.chmod(0o700)
+
+    def _extract_wallet_content(self, wallet_content: SecretStr) -> str:
+        # Strip whitespace/newlines so wrapped base64 (e.g. from `base64 -i` on macOS,
+        # which inserts line breaks every 76 chars) decodes the same as a single line.
+        sanitized = "".join(wallet_content.get_secret_value().split())
+        try:
+            decoded_wallet = base64.b64decode(sanitized, validate=True)
+        except (binascii.Error, TypeError) as exc:
+            raise ValueError("Invalid walletContent. Expected a base64-encoded wallet zip.") from exc
+
+        wallet_temp_dir = tempfile.mkdtemp(prefix="oracle_wallet_")
+        self._set_wallet_temp_dir(wallet_temp_dir)
+
+        try:
+            with zipfile.ZipFile(io.BytesIO(decoded_wallet)) as zip_ref:
+                self._safe_extract_wallet_archive(zip_ref, wallet_temp_dir)
+        except zipfile.BadZipFile as exc:
+            self._cleanup_wallet_temp_dir()
+            raise ValueError("Invalid walletContent. Expected a valid zip archive.") from exc
+        except Exception:
+            self._cleanup_wallet_temp_dir()
+            raise
+
+        return wallet_temp_dir
+
+    def _configure_autonomous_connection_arguments(self) -> None:
+        connection_type = self.service_connection.oracleConnectionType
+        if not isinstance(connection_type, OracleAutonomousConnection):
+            return
+
+        autonomous_connection = self._get_autonomous_connection_config(connection_type)
+        if not self.service_connection.connectionArguments:
+            self.service_connection.connectionArguments = init_empty_connection_arguments()
+        if self.service_connection.connectionArguments.root is None:
+            self.service_connection.connectionArguments.root = {}
+
+        connection_arguments: dict[str, Any] = self.service_connection.connectionArguments.root
+
+        wallet_path = autonomous_connection.walletPath
+        if autonomous_connection.walletContent:
+            if self._wallet_temp_dir and Path(self._wallet_temp_dir).is_dir():
+                wallet_path = self._wallet_temp_dir
+            else:
+                wallet_path = self._extract_wallet_content(autonomous_connection.walletContent)
+        else:
+            self._cleanup_wallet_temp_dir()
+
+        if not wallet_path:
+            raise ValueError("Oracle Autonomous connections require either walletPath or walletContent.")
+
+        connection_arguments["config_dir"] = wallet_path
+        connection_arguments["wallet_location"] = wallet_path
+
+        if autonomous_connection.walletPassword:
+            connection_arguments["wallet_password"] = autonomous_connection.walletPassword.get_secret_value()
+        else:
+            connection_arguments.pop("wallet_password", None)
+
+    def _uses_inline_wallet_content(self) -> bool:
+        connection_type = self.service_connection.oracleConnectionType
+        return bool(isinstance(connection_type, OracleAutonomousConnection) and connection_type.walletContent)
 
     def _get_client(self) -> Engine:
         """
         Create connection
         """
+        self._configure_autonomous_connection_arguments()
+
+        if not self._is_autonomous_connection():
+            try:
+                if self.service_connection.instantClientDirectory:
+                    logger.info(f"Initializing Oracle thick client at {self.service_connection.instantClientDirectory}")
+                    os.environ[LD_LIB_ENV] = self.service_connection.instantClientDirectory
+                    oracledb.init_oracle_client(lib_dir=self.service_connection.instantClientDirectory)
+            except DatabaseError as err:
+                logger.info(f"Could not initialize Oracle thick client: {err}")
+
         try:
-            if self.service_connection.instantClientDirectory:
-                logger.info(f"Initializing Oracle thick client at {self.service_connection.instantClientDirectory}")
-                os.environ[LD_LIB_ENV] = self.service_connection.instantClientDirectory
-                oracledb.init_oracle_client(lib_dir=self.service_connection.instantClientDirectory)
-        except DatabaseError as err:
-            logger.info(f"Could not initialize Oracle thick client: {err}")
-
-        return create_generic_db_connection(
-            connection=self.service_connection,
-            get_connection_url_fn=self.get_connection_url,
-            get_connection_args_fn=get_connection_args_common,
-        )
+            return create_generic_db_connection(
+                connection=self.service_connection,
+                get_connection_url_fn=self.get_connection_url,
+                get_connection_args_fn=get_connection_args_common,
+            )
+        except Exception:
+            if self._uses_inline_wallet_content():
+                self._cleanup_wallet_temp_dir()
+            raise
 
     def test_connection(
         self,
@@ -139,6 +286,9 @@ def get_connection_dict(self) -> dict:
             connection_dict["database"] = connection_copy.oracleConnectionType.oracleServiceName
         elif isinstance(connection_copy.oracleConnectionType, OracleTNSConnection):
             connection_dict["host"] = connection_copy.oracleConnectionType.oracleTNSConnection
+        elif isinstance(connection_copy.oracleConnectionType, OracleAutonomousConnection):
+            autonomous_connection = self._get_autonomous_connection_config(connection_copy.oracleConnectionType)
+            connection_dict["host"] = autonomous_connection.tnsAlias
 
         # Add connection options if present
         if connection_copy.connectionOptions and connection_copy.connectionOptions.root:
@@ -191,6 +341,11 @@ def _handle_connection_type(url: str, connection: OracleConnectionConfig) -> str
             url += connection.oracleConnectionType.oracleTNSConnection
             return url
 
+        if isinstance(connection.oracleConnectionType, OracleAutonomousConnection):
+            autonomous_connection = OracleConnection._get_autonomous_connection_config(connection.oracleConnectionType)
+            url += autonomous_connection.tnsAlias
+            return url
+
         # If not TNS, we add the hostPort
         url += connection.hostPort