Skip to content

Commit 5ff0d48

Browse files
authored
Add dsn attribute to ADWSecretKeeper (#986)
1 parent 4607e07 commit 5ff0d48

File tree

2 files changed

+98
-10
lines changed

2 files changed

+98
-10
lines changed

ads/secrets/adb.py

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,18 @@
11
#!/usr/bin/env python
2-
# -*- coding: utf-8 -*--
32

4-
# Copyright (c) 2021, 2022 Oracle and/or its affiliates.
3+
# Copyright (c) 2021, 2024 Oracle and/or its affiliates.
54
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
65

7-
import ads
8-
from ads.secrets import SecretKeeper, Secret
96
import json
107
import os
118
import tempfile
129
import zipfile
10+
1311
from tqdm.auto import tqdm
1412

13+
import ads
14+
from ads.secrets import Secret, SecretKeeper
15+
1516
logger = ads.getLogger("ads.secrets")
1617

1718
from dataclasses import dataclass, field
@@ -25,7 +26,7 @@ class ADBSecret(Secret):
2526

2627
user_name: str
2728
password: str
28-
service_name: str
29+
service_name: str = field(default=None)
2930
wallet_location: str = field(
3031
default=None, metadata={"serializable": False}
3132
) # Not saved in vault
@@ -40,6 +41,7 @@ class ADBSecret(Secret):
4041
wallet_secret_ids: list = field(
4142
repr=False, default_factory=list
4243
) # Not exposed through environment or `to_dict` function
44+
dsn: str = field(default=None)
4345

4446
def __post_init__(self):
4547
self.wallet_file_name = (
@@ -76,6 +78,22 @@ class ADBSecretKeeper(SecretKeeper):
7678
>>> print(adw_keeper.secret_id) # Prints the secret_id of the stored credentials
7779
>>> adw_keeper.export_vault_details("adw_employee_att.json", format="json") # Save the secret id and vault info to a json file
7880
81+
82+
>>> # Saving credentials for TLS connection
83+
>>> from ads.secrets.adw import ADBSecretKeeper
84+
>>> vault_id = "ocid1.vault.oc1..<unique_ID>"
85+
>>> kid = "ocid1.ke..<unique_ID>"
86+
87+
>>> import ads
88+
>>> ads.set_auth("resource_principal") # If using resource principal for authentication
89+
>>> connection_parameters={
90+
... "user_name":"admin",
91+
... "password":"<your password>",
92+
... "dsn":"<dsn string>"
93+
... }
94+
>>> adw_keeper = ADBSecretKeeper(vault_id=vault_id, key_id=kid, **connection_parameters)
95+
>>> adw_keeper.save("adw_employee", "My DB credentials", freeform_tags={"schema":"emp"})
96+
7997
>>> # Loading credentails
8098
>>> import ads
8199
>>> ads.set_auth("resource_principal") # If using resource principal for authentication
@@ -133,6 +151,7 @@ def __init__(
133151
wallet_dir: str = None,
134152
repository_path: str = None,
135153
repository_key: str = None,
154+
dsn: str = None,
136155
**kwargs,
137156
):
138157
"""
@@ -152,6 +171,8 @@ def __init__(
152171
Path to credentials repository. For more details refer `ads.database.connection`
153172
repository_key: (str, optional). Default None.
154173
Configuration key for loading the right configuration from repository. For more details refer `ads.database.connection`
174+
dsn: (str, optional). Default None.
175+
dsn string copied from the OCI console for TLS connection
155176
kwargs:
156177
vault_id: str. OCID of the vault where the secret is stored. Required for saving secret.
157178
key_id: str. OCID of the key used for encrypting the secret. Required for saving secret.
@@ -180,6 +201,7 @@ def __init__(
180201
password=password,
181202
service_name=service_name,
182203
wallet_location=wallet_location,
204+
dsn=dsn,
183205
)
184206
self.wallet_dir = wallet_dir
185207

@@ -252,7 +274,7 @@ def decode(self) -> "ads.secrets.adb.ADBSecretKeeper":
252274
logger.debug(f"Setting wallet file to {self.data.wallet_location}")
253275
data.wallet_location = self.data.wallet_location
254276
elif data.wallet_secret_ids and len(data.wallet_secret_ids) > 0:
255-
logger.debug(f"Secret ids corresponding to the wallet files found.")
277+
logger.debug("Secret ids corresponding to the wallet files found.")
256278
# If the secret ids for wallet files are available in secret, then we
257279
# can generate the wallet file.
258280

tests/unitary/default_setup/secret/test_secretkeeper_adw.py

Lines changed: 70 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#!/usr/bin/env python
22

3-
# Copyright (c) 2021, 2023 Oracle and/or its affiliates.
3+
# Copyright (c) 2021, 2024 Oracle and/or its affiliates.
44
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
55

66
from ads.secrets import ADBSecretKeeper
@@ -36,6 +36,24 @@ def key_encoding():
3636
)
3737

3838

39+
@pytest.fixture
40+
def key_encoding_dsn():
41+
user_name = "myuser"
42+
password = "this-is-not-the-secret"
43+
dsn = "my long dsn string....................."
44+
secret_dict = {
45+
"user_name": user_name,
46+
"password": password,
47+
"dsn": dsn,
48+
}
49+
encoded = b64encode(json.dumps(secret_dict).encode("utf-8")).decode("utf-8")
50+
return (
51+
(user_name, password, dsn),
52+
secret_dict,
53+
encoded,
54+
)
55+
56+
3957
def generate_wallet_data(wallet_zip_path, wallet_dir_path):
4058
files = 4
4159
file_content = {}
@@ -133,9 +151,9 @@ def test_encode(mock_client, mock_signer, key_encoding):
133151

134152
@patch("ads.common.auth.default_signer")
135153
@patch("ads.common.oci_client.OCIClientFactory")
136-
def test_adw_save(mock_client, mock_signer, key_encoding, tmpdir):
154+
def test_adw_tls_save(mock_client, mock_signer, key_encoding_dsn, tmpdir):
137155
adwsecretkeeper = ADBSecretKeeper(
138-
*key_encoding[0],
156+
**key_encoding_dsn[1],
139157
vault_id="ocid.vault",
140158
key_id="ocid.key",
141159
compartment_id="dummy",
@@ -211,6 +229,7 @@ def test_adw_context(mock_client, mock_signer, key_encoding):
211229
assert adwsecretkeeper == {
212230
**key_encoding[1],
213231
"wallet_location": "/this/is/mywallet.zip",
232+
"dsn": None,
214233
}
215234
assert os.environ.get("user_name") == key_encoding[0][0]
216235
assert os.environ.get("password") == key_encoding[0][1]
@@ -221,13 +240,44 @@ def test_adw_context(mock_client, mock_signer, key_encoding):
221240
"password": None,
222241
"service_name": None,
223242
"wallet_location": None,
243+
"dsn": None,
224244
}
225245
assert os.environ.get("user_name") is None
226246
assert os.environ.get("password") is None
227247
assert os.environ.get("service_name") is None
228248
assert os.environ.get("wallet_location") is None
229249

230250

251+
@patch("ads.common.auth.default_signer")
252+
@patch("ads.common.oci_client.OCIClientFactory")
253+
def test_adw_context_tls(mock_client, mock_signer, key_encoding_dsn):
254+
with mock.patch(
255+
"ads.vault.Vault.get_secret", return_value=key_encoding_dsn[2]
256+
) as mocked_getsecret:
257+
with ADBSecretKeeper.load_secret(
258+
source="ocid.secret.id",
259+
export_env=True,
260+
) as adwsecretkeeper:
261+
assert adwsecretkeeper == {
262+
**key_encoding_dsn[1],
263+
"service_name": None,
264+
"wallet_location": None,
265+
}
266+
assert os.environ.get("user_name") == key_encoding_dsn[0][0]
267+
assert os.environ.get("password") == key_encoding_dsn[0][1]
268+
assert os.environ.get("dsn") == key_encoding_dsn[0][2]
269+
assert adwsecretkeeper == {
270+
"user_name": None,
271+
"password": None,
272+
"service_name": None,
273+
"wallet_location": None,
274+
"dsn": None,
275+
}
276+
assert os.environ.get("user_name") is None
277+
assert os.environ.get("password") is None
278+
assert os.environ.get("dsn") is None
279+
280+
231281
@patch("ads.common.auth.default_signer")
232282
@patch("ads.common.oci_client.OCIClientFactory")
233283
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):
240290
assert adwsecretkeeper == {
241291
**key_encoding[1],
242292
"wallet_location": None,
293+
"dsn": None,
243294
}
244295

245296

246297
@patch("ads.common.auth.default_signer")
247298
@patch("ads.common.oci_client.OCIClientFactory")
248299
def test_adw_keeper_with_repository(mock_client, mock_signer, key_encoding, tmpdir):
249-
expected = {**key_encoding[1], "wallet_location": key_encoding[3]}
300+
expected = {**key_encoding[1], "wallet_location": key_encoding[3], "dsn": None}
250301
os.makedirs(os.path.join(tmpdir, "testdb"))
251302
with open(os.path.join(tmpdir, "testdb", "config.json"), "w") as conffile:
252303
json.dump(expected, conffile)
@@ -270,6 +321,7 @@ def test_adw_context_namespace(mock_client, mock_signer, key_encoding):
270321
assert adwsecretkeeper == {
271322
**key_encoding[1],
272323
"wallet_location": "/this/is/mywallet.zip",
324+
"dsn": None,
273325
}
274326
assert os.environ.get("myapp.user_name") == key_encoding[0][0]
275327
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):
280332
"password": None,
281333
"service_name": None,
282334
"wallet_location": None,
335+
"dsn": None,
283336
}
284337
assert os.environ.get("myapp.user_name") is None
285338
assert os.environ.get("myapp.password") is None
@@ -300,6 +353,7 @@ def test_adw_context_noexport(mock_client, mock_signer, key_encoding):
300353
assert adwsecretkeeper == {
301354
**key_encoding[1],
302355
"wallet_location": "/this/is/mywallet.zip",
356+
"dsn": None,
303357
}
304358

305359
assert os.environ.get("user_name") is None
@@ -312,6 +366,7 @@ def test_adw_context_noexport(mock_client, mock_signer, key_encoding):
312366
"password": None,
313367
"service_name": None,
314368
"wallet_location": None,
369+
"dsn": None,
315370
}
316371

317372

@@ -413,6 +468,7 @@ def mock_get_secret_id(
413468
"password": key_encoding_with_wallet.credentials.password,
414469
"service_name": key_encoding_with_wallet.credentials.service_name,
415470
"wallet_location": f"{os.path.join(wallet_dir,'wallet.zip')}",
471+
"dsn": None,
416472
}
417473

418474
# with open(key_encoding_with_wallet[3], "rb") as orgfile:
@@ -449,6 +505,7 @@ def mock_get_secret_id(
449505
"password": key_encoding_with_wallet.credentials.password,
450506
"service_name": key_encoding_with_wallet.credentials.service_name,
451507
"wallet_location": f"{os.path.join(wallet_dir,'wallet.zip')}",
508+
"dsn": None,
452509
}
453510
assert (
454511
os.environ.get("user_name")
@@ -472,6 +529,7 @@ def mock_get_secret_id(
472529
"password": None,
473530
"service_name": None,
474531
"wallet_location": None,
532+
"dsn": None,
475533
}
476534
assert os.environ.get("user_name") is None
477535
assert os.environ.get("password") is None
@@ -508,6 +566,7 @@ def mock_get_secret_id(
508566
"password": key_encoding_with_wallet.credentials.password,
509567
"service_name": key_encoding_with_wallet.credentials.service_name,
510568
"wallet_location": f"{os.path.join(wallet_dir,'wallet.zip')}",
569+
"dsn": None,
511570
}
512571
assert (
513572
os.environ.get("myapp.user_name")
@@ -531,6 +590,7 @@ def mock_get_secret_id(
531590
"password": None,
532591
"service_name": None,
533592
"wallet_location": None,
593+
"dsn": None,
534594
}
535595
assert os.environ.get("myapp.user_name") is None
536596
assert os.environ.get("myapp.password") is None
@@ -565,6 +625,7 @@ def mock_get_secret_id(
565625
"password": key_encoding_with_wallet.credentials.password,
566626
"service_name": key_encoding_with_wallet.credentials.service_name,
567627
"wallet_location": f"{os.path.join(wallet_dir,'wallet.zip')}",
628+
"dsn": None,
568629
}
569630
assert os.environ.get("user_name") is None
570631
assert os.environ.get("password") is None
@@ -576,6 +637,7 @@ def mock_get_secret_id(
576637
"password": None,
577638
"service_name": None,
578639
"wallet_location": None,
640+
"dsn": None,
579641
}
580642

581643

@@ -730,6 +792,7 @@ def mock_get_secret_id(
730792
"password": key_encoding_with_wallet.credentials.password,
731793
"service_name": key_encoding_with_wallet.credentials.service_name,
732794
"wallet_location": f"{os.path.join(wallet_dir,'wallet.zip')}",
795+
"dsn": None,
733796
}
734797
assert (
735798
os.environ.get("user_name")
@@ -753,6 +816,7 @@ def mock_get_secret_id(
753816
"password": None,
754817
"service_name": None,
755818
"wallet_location": None,
819+
"dsn": None,
756820
}
757821
assert os.environ.get("user_name") is None
758822
assert os.environ.get("password") is None
@@ -786,6 +850,7 @@ def mock_get_secret_id(
786850
"password": key_encoding_with_wallet.credentials.password,
787851
"service_name": key_encoding_with_wallet.credentials.service_name,
788852
"wallet_location": f"{os.path.join(wallet_dir,'wallet.zip')}",
853+
"dsn": None,
789854
}
790855
assert (
791856
os.environ.get("user_name")
@@ -809,6 +874,7 @@ def mock_get_secret_id(
809874
"password": None,
810875
"service_name": None,
811876
"wallet_location": None,
877+
"dsn": None,
812878
}
813879
assert os.environ.get("user_name") is None
814880
assert os.environ.get("password") is None

0 commit comments

Comments
 (0)