diff --git a/ads/secrets/adb.py b/ads/secrets/adb.py index a3ff08034..90772d7d2 100644 --- a/ads/secrets/adb.py +++ b/ads/secrets/adb.py @@ -1,17 +1,18 @@ #!/usr/bin/env python -# -*- coding: utf-8 -*-- -# Copyright (c) 2021, 2022 Oracle and/or its affiliates. +# Copyright (c) 2021, 2024 Oracle and/or its affiliates. # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ -import ads -from ads.secrets import SecretKeeper, Secret import json import os import tempfile import zipfile + from tqdm.auto import tqdm +import ads +from ads.secrets import Secret, SecretKeeper + logger = ads.getLogger("ads.secrets") from dataclasses import dataclass, field @@ -25,7 +26,7 @@ class ADBSecret(Secret): user_name: str password: str - service_name: str + service_name: str = field(default=None) wallet_location: str = field( default=None, metadata={"serializable": False} ) # Not saved in vault @@ -40,6 +41,7 @@ class ADBSecret(Secret): wallet_secret_ids: list = field( repr=False, default_factory=list ) # Not exposed through environment or `to_dict` function + dsn: str = field(default=None) def __post_init__(self): self.wallet_file_name = ( @@ -76,6 +78,22 @@ class ADBSecretKeeper(SecretKeeper): >>> print(adw_keeper.secret_id) # Prints the secret_id of the stored credentials >>> adw_keeper.export_vault_details("adw_employee_att.json", format="json") # Save the secret id and vault info to a json file + + >>> # Saving credentials for TLS connection + >>> from ads.secrets.adw import ADBSecretKeeper + >>> vault_id = "ocid1.vault.oc1.." + >>> kid = "ocid1.ke.." + + >>> import ads + >>> ads.set_auth("resource_principal") # If using resource principal for authentication + >>> connection_parameters={ + ... "user_name":"admin", + ... "password":"", + ... "dsn":"" + ... } + >>> adw_keeper = ADBSecretKeeper(vault_id=vault_id, key_id=kid, **connection_parameters) + >>> adw_keeper.save("adw_employee", "My DB credentials", freeform_tags={"schema":"emp"}) + >>> # Loading credentails >>> import ads >>> ads.set_auth("resource_principal") # If using resource principal for authentication @@ -133,6 +151,7 @@ def __init__( wallet_dir: str = None, repository_path: str = None, repository_key: str = None, + dsn: str = None, **kwargs, ): """ @@ -152,6 +171,8 @@ def __init__( Path to credentials repository. For more details refer `ads.database.connection` repository_key: (str, optional). Default None. Configuration key for loading the right configuration from repository. For more details refer `ads.database.connection` + dsn: (str, optional). Default None. + dsn string copied from the OCI console for TLS connection kwargs: vault_id: str. OCID of the vault where the secret is stored. Required for saving secret. key_id: str. OCID of the key used for encrypting the secret. Required for saving secret. @@ -180,6 +201,7 @@ def __init__( password=password, service_name=service_name, wallet_location=wallet_location, + dsn=dsn, ) self.wallet_dir = wallet_dir @@ -252,7 +274,7 @@ def decode(self) -> "ads.secrets.adb.ADBSecretKeeper": logger.debug(f"Setting wallet file to {self.data.wallet_location}") data.wallet_location = self.data.wallet_location elif data.wallet_secret_ids and len(data.wallet_secret_ids) > 0: - logger.debug(f"Secret ids corresponding to the wallet files found.") + logger.debug("Secret ids corresponding to the wallet files found.") # If the secret ids for wallet files are available in secret, then we # can generate the wallet file. diff --git a/tests/unitary/default_setup/secret/test_secretkeeper_adw.py b/tests/unitary/default_setup/secret/test_secretkeeper_adw.py index 2002ffde3..1d5d40e63 100644 --- a/tests/unitary/default_setup/secret/test_secretkeeper_adw.py +++ b/tests/unitary/default_setup/secret/test_secretkeeper_adw.py @@ -1,6 +1,6 @@ #!/usr/bin/env python -# Copyright (c) 2021, 2023 Oracle and/or its affiliates. +# Copyright (c) 2021, 2024 Oracle and/or its affiliates. # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ from ads.secrets import ADBSecretKeeper @@ -36,6 +36,24 @@ def key_encoding(): ) +@pytest.fixture +def key_encoding_dsn(): + user_name = "myuser" + password = "this-is-not-the-secret" + dsn = "my long dsn string....................." + secret_dict = { + "user_name": user_name, + "password": password, + "dsn": dsn, + } + encoded = b64encode(json.dumps(secret_dict).encode("utf-8")).decode("utf-8") + return ( + (user_name, password, dsn), + secret_dict, + encoded, + ) + + def generate_wallet_data(wallet_zip_path, wallet_dir_path): files = 4 file_content = {} @@ -133,9 +151,9 @@ def test_encode(mock_client, mock_signer, key_encoding): @patch("ads.common.auth.default_signer") @patch("ads.common.oci_client.OCIClientFactory") -def test_adw_save(mock_client, mock_signer, key_encoding, tmpdir): +def test_adw_tls_save(mock_client, mock_signer, key_encoding_dsn, tmpdir): adwsecretkeeper = ADBSecretKeeper( - *key_encoding[0], + **key_encoding_dsn[1], vault_id="ocid.vault", key_id="ocid.key", compartment_id="dummy", @@ -211,6 +229,7 @@ def test_adw_context(mock_client, mock_signer, key_encoding): assert adwsecretkeeper == { **key_encoding[1], "wallet_location": "/this/is/mywallet.zip", + "dsn": None, } assert os.environ.get("user_name") == key_encoding[0][0] assert os.environ.get("password") == key_encoding[0][1] @@ -221,6 +240,7 @@ def test_adw_context(mock_client, mock_signer, key_encoding): "password": None, "service_name": None, "wallet_location": None, + "dsn": None, } assert os.environ.get("user_name") is None assert os.environ.get("password") is None @@ -228,6 +248,36 @@ def test_adw_context(mock_client, mock_signer, key_encoding): assert os.environ.get("wallet_location") is None +@patch("ads.common.auth.default_signer") +@patch("ads.common.oci_client.OCIClientFactory") +def test_adw_context_tls(mock_client, mock_signer, key_encoding_dsn): + with mock.patch( + "ads.vault.Vault.get_secret", return_value=key_encoding_dsn[2] + ) as mocked_getsecret: + with ADBSecretKeeper.load_secret( + source="ocid.secret.id", + export_env=True, + ) as adwsecretkeeper: + assert adwsecretkeeper == { + **key_encoding_dsn[1], + "service_name": None, + "wallet_location": None, + } + assert os.environ.get("user_name") == key_encoding_dsn[0][0] + assert os.environ.get("password") == key_encoding_dsn[0][1] + assert os.environ.get("dsn") == key_encoding_dsn[0][2] + assert adwsecretkeeper == { + "user_name": None, + "password": None, + "service_name": None, + "wallet_location": None, + "dsn": None, + } + assert os.environ.get("user_name") is None + assert os.environ.get("password") is None + assert os.environ.get("dsn") is None + + @patch("ads.common.auth.default_signer") @patch("ads.common.oci_client.OCIClientFactory") def test_adw_keeper_no_wallet(mock_client, mock_signer, key_encoding): @@ -240,13 +290,14 @@ def test_adw_keeper_no_wallet(mock_client, mock_signer, key_encoding): assert adwsecretkeeper == { **key_encoding[1], "wallet_location": None, + "dsn": None, } @patch("ads.common.auth.default_signer") @patch("ads.common.oci_client.OCIClientFactory") def test_adw_keeper_with_repository(mock_client, mock_signer, key_encoding, tmpdir): - expected = {**key_encoding[1], "wallet_location": key_encoding[3]} + expected = {**key_encoding[1], "wallet_location": key_encoding[3], "dsn": None} os.makedirs(os.path.join(tmpdir, "testdb")) with open(os.path.join(tmpdir, "testdb", "config.json"), "w") as conffile: json.dump(expected, conffile) @@ -270,6 +321,7 @@ def test_adw_context_namespace(mock_client, mock_signer, key_encoding): assert adwsecretkeeper == { **key_encoding[1], "wallet_location": "/this/is/mywallet.zip", + "dsn": None, } assert os.environ.get("myapp.user_name") == key_encoding[0][0] assert os.environ.get("myapp.password") == key_encoding[0][1] @@ -280,6 +332,7 @@ def test_adw_context_namespace(mock_client, mock_signer, key_encoding): "password": None, "service_name": None, "wallet_location": None, + "dsn": None, } assert os.environ.get("myapp.user_name") is None assert os.environ.get("myapp.password") is None @@ -300,6 +353,7 @@ def test_adw_context_noexport(mock_client, mock_signer, key_encoding): assert adwsecretkeeper == { **key_encoding[1], "wallet_location": "/this/is/mywallet.zip", + "dsn": None, } assert os.environ.get("user_name") is None @@ -312,6 +366,7 @@ def test_adw_context_noexport(mock_client, mock_signer, key_encoding): "password": None, "service_name": None, "wallet_location": None, + "dsn": None, } @@ -413,6 +468,7 @@ def mock_get_secret_id( "password": key_encoding_with_wallet.credentials.password, "service_name": key_encoding_with_wallet.credentials.service_name, "wallet_location": f"{os.path.join(wallet_dir,'wallet.zip')}", + "dsn": None, } # with open(key_encoding_with_wallet[3], "rb") as orgfile: @@ -449,6 +505,7 @@ def mock_get_secret_id( "password": key_encoding_with_wallet.credentials.password, "service_name": key_encoding_with_wallet.credentials.service_name, "wallet_location": f"{os.path.join(wallet_dir,'wallet.zip')}", + "dsn": None, } assert ( os.environ.get("user_name") @@ -472,6 +529,7 @@ def mock_get_secret_id( "password": None, "service_name": None, "wallet_location": None, + "dsn": None, } assert os.environ.get("user_name") is None assert os.environ.get("password") is None @@ -508,6 +566,7 @@ def mock_get_secret_id( "password": key_encoding_with_wallet.credentials.password, "service_name": key_encoding_with_wallet.credentials.service_name, "wallet_location": f"{os.path.join(wallet_dir,'wallet.zip')}", + "dsn": None, } assert ( os.environ.get("myapp.user_name") @@ -531,6 +590,7 @@ def mock_get_secret_id( "password": None, "service_name": None, "wallet_location": None, + "dsn": None, } assert os.environ.get("myapp.user_name") is None assert os.environ.get("myapp.password") is None @@ -565,6 +625,7 @@ def mock_get_secret_id( "password": key_encoding_with_wallet.credentials.password, "service_name": key_encoding_with_wallet.credentials.service_name, "wallet_location": f"{os.path.join(wallet_dir,'wallet.zip')}", + "dsn": None, } assert os.environ.get("user_name") is None assert os.environ.get("password") is None @@ -576,6 +637,7 @@ def mock_get_secret_id( "password": None, "service_name": None, "wallet_location": None, + "dsn": None, } @@ -730,6 +792,7 @@ def mock_get_secret_id( "password": key_encoding_with_wallet.credentials.password, "service_name": key_encoding_with_wallet.credentials.service_name, "wallet_location": f"{os.path.join(wallet_dir,'wallet.zip')}", + "dsn": None, } assert ( os.environ.get("user_name") @@ -753,6 +816,7 @@ def mock_get_secret_id( "password": None, "service_name": None, "wallet_location": None, + "dsn": None, } assert os.environ.get("user_name") is None assert os.environ.get("password") is None @@ -786,6 +850,7 @@ def mock_get_secret_id( "password": key_encoding_with_wallet.credentials.password, "service_name": key_encoding_with_wallet.credentials.service_name, "wallet_location": f"{os.path.join(wallet_dir,'wallet.zip')}", + "dsn": None, } assert ( os.environ.get("user_name") @@ -809,6 +874,7 @@ def mock_get_secret_id( "password": None, "service_name": None, "wallet_location": None, + "dsn": None, } assert os.environ.get("user_name") is None assert os.environ.get("password") is None