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

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

@@ -12,10 +12,15 @@
 """
 Source connection handler
 """
+import base64
+import io
 import os
+import shutil
 import sys
+import tempfile
+import zipfile
 from copy import deepcopy
-from typing import Optional
+from typing import Any, Optional
 from urllib.parse import quote_plus
 
 import oracledb
@@ -30,6 +35,7 @@
     OracleConnection as OracleConnectionConfig,
 )
 from metadata.generated.schema.entity.services.connections.database.oracleConnection import (
+    OracleAutonomousConnection,
     OracleDatabaseSchema,
     OracleServiceName,
     OracleTNSConnection,
@@ -41,6 +47,7 @@
     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
@@ -67,22 +74,104 @@
 class OracleConnection(BaseConnection[OracleConnectionConfig, Engine]):
     def __init__(self, connection: OracleConnectionConfig):
         super().__init__(connection)
+        self._wallet_temp_dir: Optional[str] = None
+
+    def __del__(self):
+        self._cleanup_wallet_temp_dir()
+
+    def _cleanup_wallet_temp_dir(self) -> None:
+        if self._wallet_temp_dir:
+            shutil.rmtree(self._wallet_temp_dir, ignore_errors=True)
+            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,
+    ) -> Any:
+        return connection_type.root
+
+    def _extract_wallet_content(self, wallet_content: SecretStr) -> str:
+        try:
+            decoded_wallet = base64.b64decode(wallet_content.get_secret_value())
+        except (ValueError, TypeError) as exc:
+            raise ValueError(
+                "Invalid walletContent. Expected a base64-encoded wallet zip."
+            ) from exc
+
+        self._cleanup_wallet_temp_dir()
+        self._wallet_temp_dir = tempfile.mkdtemp(prefix="oracle_wallet_")
+
+        try:
+            with zipfile.ZipFile(io.BytesIO(decoded_wallet)) as zip_ref:
+                zip_ref.extractall(self._wallet_temp_dir)
+        except zipfile.BadZipFile as exc:
+            self._cleanup_wallet_temp_dir()
+            raise ValueError(
+                "Invalid walletContent. Expected a valid zip archive."
+            ) from exc
+
+        return self._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)
+
+        wallet_path = autonomous_connection.walletPath
+        if autonomous_connection.walletContent:
+            wallet_path = self._extract_wallet_content(
+                autonomous_connection.walletContent
+            )
+
+        if not wallet_path:
+            raise ValueError(
+                "Oracle Autonomous connections require either walletPath or walletContent."
+            )
+
+        if not self.service_connection.connectionArguments:
+            self.service_connection.connectionArguments = (
+                init_empty_connection_arguments()
+            )
+        elif self.service_connection.connectionArguments.root is None:
+            self.service_connection.connectionArguments.root = {}
+
+        self.service_connection.connectionArguments.root["config_dir"] = wallet_path
+        self.service_connection.connectionArguments.root[
+            "wallet_location"
+        ] = wallet_path
+
+        if autonomous_connection.walletPassword:
+            self.service_connection.connectionArguments.root[
+                "wallet_password"
+            ] = autonomous_connection.walletPassword.get_secret_value()
 
     def _get_client(self) -> Engine:
         """
         Create 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}")
+        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}")
 
         return create_generic_db_connection(
             connection=self.service_connection,
@@ -150,6 +239,13 @@ def get_connection_dict(self) -> dict:
             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:
@@ -209,6 +305,13 @@ 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
 

ingestion/tests/unit/test_source_connection.py

@@ -9,7 +9,12 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import base64
+import io
+import os
+import zipfile
 from unittest import TestCase
+from unittest.mock import patch
 
 from trino.auth import BasicAuthentication, JWTAuthentication, OAuth2Authentication
 
@@ -78,6 +83,7 @@
     OracleConnection as OracleConnectionConfig,
 )
 from metadata.generated.schema.entity.services.connections.database.oracleConnection import (
+    OracleAutonomousConnection,
     OracleDatabaseSchema,
     OracleScheme,
     OracleServiceName,
@@ -1227,6 +1233,109 @@ def test_oracle_url(self):
         )
         assert OracleConnection.get_connection_url(oracle_conn_obj) == expected_url
 
+        expected_url = "oracle+cx_oracle://admin:password@myadb_high"
+        oracle_conn_obj = OracleConnectionConfig(
+            username="admin",
+            password="password",
+            oracleConnectionType=OracleAutonomousConnection(
+                tnsAlias="myadb_high",
+                walletPath="/tmp/my_wallet",
+            ),
+        )
+        assert OracleConnection.get_connection_url(oracle_conn_obj) == expected_url
+
+        expected_url = [
+            "oracle+cx_oracle://admin:password@myadb_high?test_key_2=test_value_2&test_key_1=test_value_1",
+            "oracle+cx_oracle://admin:password@myadb_high?test_key_1=test_value_1&test_key_2=test_value_2",
+        ]
+        oracle_conn_obj = OracleConnectionConfig(
+            username="admin",
+            password="password",
+            oracleConnectionType=OracleAutonomousConnection(
+                tnsAlias="myadb_high",
+                walletPath="/tmp/my_wallet",
+            ),
+            connectionOptions=dict(
+                test_key_1="test_value_1", test_key_2="test_value_2"
+            ),
+        )
+        assert OracleConnection.get_connection_url(oracle_conn_obj) in expected_url
+
+    @patch(
+        "metadata.ingestion.source.database.oracle.connection.oracledb.init_oracle_client"
+    )
+    @patch(
+        "metadata.ingestion.source.database.oracle.connection.create_generic_db_connection"
+    )
+    def test_oracle_autonomous_wallet_path_args(
+        self, mock_create_generic_db_connection, mock_init_oracle_client
+    ):
+        connection = OracleConnectionConfig(
+            username="admin",
+            password="password",
+            instantClientDirectory="/instantclient",
+            oracleConnectionType=OracleAutonomousConnection(
+                tnsAlias="myadb_high",
+                walletPath="/tmp/my_wallet",
+                walletPassword="wallet_password",
+            ),
+        )
+        oracle_connection = OracleConnection(connection)
+        mock_create_generic_db_connection.return_value = "dummy_engine"
+
+        oracle_connection._get_client()
+
+        assert mock_init_oracle_client.call_count == 0
+        assert (
+            oracle_connection.service_connection.connectionArguments.root["config_dir"]
+            == "/tmp/my_wallet"
+        )
+        assert (
+            oracle_connection.service_connection.connectionArguments.root[
+                "wallet_location"
+            ]
+            == "/tmp/my_wallet"
+        )
+        assert (
+            oracle_connection.service_connection.connectionArguments.root[
+                "wallet_password"
+            ]
+            == "wallet_password"
+        )
+
+    @patch(
+        "metadata.ingestion.source.database.oracle.connection.create_generic_db_connection"
+    )
+    def test_oracle_autonomous_wallet_content_args(
+        self, mock_create_generic_db_connection
+    ):
+        wallet_bytes = io.BytesIO()
+        with zipfile.ZipFile(wallet_bytes, "w", zipfile.ZIP_DEFLATED) as zip_file:
+            zip_file.writestr("tnsnames.ora", "MYADB_HIGH=(DESCRIPTION=...)")
+
+        encoded_wallet = base64.b64encode(wallet_bytes.getvalue()).decode("utf-8")
+
+        connection = OracleConnectionConfig(
+            username="admin",
+            password="password",
+            oracleConnectionType=OracleAutonomousConnection(
+                tnsAlias="myadb_high",
+                walletContent=encoded_wallet,
+            ),
+        )
+        oracle_connection = OracleConnection(connection)
+        mock_create_generic_db_connection.return_value = "dummy_engine"
+
+        oracle_connection._get_client()
+
+        wallet_dir = oracle_connection.service_connection.connectionArguments.root[
+            "config_dir"
+        ]
+        assert os.path.isdir(wallet_dir)
+        assert os.path.exists(os.path.join(wallet_dir, "tnsnames.ora"))
+        oracle_connection._cleanup_wallet_temp_dir()
+        assert not os.path.exists(wallet_dir)
+
     def test_exasol_url(self):
         from metadata.ingestion.source.database.exasol.connection import (
             get_connection_url,

openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/oracleConnection.json

@@ -63,6 +63,49 @@
       "required": [
         "oracleTNSConnection"
       ]
+    },
+    "OracleAutonomousConnection": {
+      "title": "Oracle Autonomous Connection",
+      "type": "object",
+      "properties": {
+        "tnsAlias": {
+          "title": "TNS Alias",
+          "description": "Service alias defined in the wallet tnsnames.ora file, such as myadb_high.",
+          "type": "string"
+        },
+        "walletPath": {
+          "title": "Wallet Path",
+          "description": "Path to the extracted Oracle wallet directory on the ingestion host.",
+          "type": "string"
+        },
+        "walletContent": {
+          "title": "Wallet Content",
+          "description": "Base64-encoded Oracle wallet zip content. If provided, OpenMetadata extracts it at runtime.",
+          "type": "string",
+          "format": "password"
+        },
+        "walletPassword": {
+          "title": "Wallet Password",
+          "description": "Wallet password for Oracle Autonomous mTLS connections, if required.",
+          "type": "string",
+          "format": "password"
+        }
+      },
+      "required": [
+        "tnsAlias"
+      ],
+      "anyOf": [
+        {
+          "required": [
+            "walletPath"
+          ]
+        },
+        {
+          "required": [
+            "walletContent"
+          ]
+        }
+      ]
     }
   },
   "properties": {
@@ -97,7 +140,7 @@
     "oracleConnectionType": {
       "title": "Oracle Connection Type",
       "type": "object",
-      "description": "Connect with oracle by either passing service name or database schema name.",
+      "description": "Connect with Oracle by using schema, service name, TNS connection string, or Oracle Autonomous wallet configuration.",
       "oneOf": [
         {
           "$ref": "#/definitions/OracleDatabaseSchema"
@@ -107,6 +150,9 @@
         },
         {
           "$ref": "#/definitions/OracleTNSConnection"
+        },
+        {
+          "$ref": "#/definitions/OracleAutonomousConnection"
         }
       ]
     },