This is an automated email from the ASF dual-hosted git repository.
taragolis pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new cea58c1111 fix: using endpoint from connection if not specified
(#37076)
cea58c1111 is described below
commit cea58c11118cb6776cc14b9460d674b251888bad
Author: Vadim Vladimirov
<[email protected]>
AuthorDate: Thu Feb 1 23:15:21 2024 +0300
fix: using endpoint from connection if not specified (#37076)
* fix: use endpoint from connection if not specified
* fix: generate client before get_secrets if not exists, to load folder_id
from connections
---
airflow/providers/yandex/secrets/lockbox.py | 7 +++-
.../yandex-cloud-lockbox-secret-backend.rst | 2 +-
tests/providers/yandex/secrets/test_lockbox.py | 39 ++++++++++++++++++++--
3 files changed, 44 insertions(+), 4 deletions(-)
diff --git a/airflow/providers/yandex/secrets/lockbox.py
b/airflow/providers/yandex/secrets/lockbox.py
index adbf994873..d9ade29646 100644
--- a/airflow/providers/yandex/secrets/lockbox.py
+++ b/airflow/providers/yandex/secrets/lockbox.py
@@ -103,7 +103,8 @@ class LockboxSecretBackend(BaseSecretsBackend,
LoggingMixin):
:param sep: Specifies the separator used to concatenate secret_prefix and
secret_id.
Default: "/"
:param endpoint: Specifies an API endpoint.
- Leave blank to use default.
+ If set to None (null in JSON), requests will use the connection
endpoint, if specified,
+ or the default endpoint.
"""
def __init__(
@@ -191,6 +192,7 @@ class LockboxSecretBackend(BaseSecretsBackend,
LoggingMixin):
self.yc_sa_key_json = self._get_field("service_account_json")
self.yc_sa_key_json_path =
self._get_field("service_account_json_path")
self.folder_id = self.folder_id or self._get_field("folder_id")
+ self.endpoint = self.endpoint or self._get_field("endpoint")
credentials = get_credentials(
oauth_token=self.yc_oauth_token,
@@ -251,6 +253,9 @@ class LockboxSecretBackend(BaseSecretsBackend,
LoggingMixin):
return sorted(entries.values())[0]
def _get_secrets(self) -> list[secret_pb.Secret]:
+ # generate client if not exists, to load folder_id from connections
+ _ = self._client
+
response = self._list_secrets(folder_id=self.folder_id)
secrets: list[secret_pb.Secret] = response.secrets[:]
diff --git
a/docs/apache-airflow-providers-yandex/secrets-backends/yandex-cloud-lockbox-secret-backend.rst
b/docs/apache-airflow-providers-yandex/secrets-backends/yandex-cloud-lockbox-secret-backend.rst
index 9403dad0ea..e0d3625df4 100644
---
a/docs/apache-airflow-providers-yandex/secrets-backends/yandex-cloud-lockbox-secret-backend.rst
+++
b/docs/apache-airflow-providers-yandex/secrets-backends/yandex-cloud-lockbox-secret-backend.rst
@@ -73,7 +73,7 @@ You can pass the following parameters:
* ``variables_prefix``: Specifies the prefix of the secret to read to get
Variables. If set to None (null in JSON), requests for variables will not be
sent to Yandex Lockbox. Default: "airflow/variables"
* ``config_prefix``: Specifies the prefix of the secret to read to get
Configurations. If set to None (null in JSON), requests for variables will not
be sent to Yandex Lockbox. Default: "airflow/config"
* ``sep``: Specifies the separator used to concatenate secret_prefix and
secret_id. Default: "/"
-* ``endpoint``: Specifies an API endpoint. Leave blank to use default.
+* ``endpoint``: Specifies an API endpoint. If set to None (null in JSON),
requests will use the connection endpoint, if specified, or the default
endpoint.
All options should be passed as a JSON dictionary.
diff --git a/tests/providers/yandex/secrets/test_lockbox.py
b/tests/providers/yandex/secrets/test_lockbox.py
index e51724f866..ae91ae637b 100644
--- a/tests/providers/yandex/secrets/test_lockbox.py
+++ b/tests/providers/yandex/secrets/test_lockbox.py
@@ -176,6 +176,37 @@ class TestLockboxSecretBackend:
assert sm._client is not None
+
@patch("airflow.providers.yandex.secrets.lockbox.LockboxSecretBackend._get_field")
+ def
test_yandex_lockbox_secret_backend__client_credentials_received_from_connection(self,
mock_get_field):
+ yc_oauth_token =
"y3_Vdheub7w9bIut67GHeL345gfb5GAnd3dZnf08FRbvjeUFvetYiohGvc"
+ yc_sa_key_json = "sa_key_json"
+ yc_sa_key_json_path = "sa_key_json_path"
+ folder_id = "folder_id123"
+ endpoint = "some-custom-api-endpoint.cloud.yandex.net"
+ yc_connection_id = "connection_id"
+
+ fields = {
+ "oauth": yc_oauth_token,
+ "service_account_json": yc_sa_key_json,
+ "service_account_json_path": yc_sa_key_json_path,
+ "folder_id": folder_id,
+ "endpoint": endpoint,
+ }
+ mock_get_field.side_effect = lambda key: fields[key]
+
+ sm = LockboxSecretBackend(
+ yc_connection_id=yc_connection_id,
+ )
+ client = sm._client
+
+ assert client is not None
+ assert sm.yc_oauth_token == yc_oauth_token
+ assert sm.yc_sa_key_json == yc_sa_key_json
+ assert sm.yc_sa_key_json_path == yc_sa_key_json_path
+ assert sm.folder_id == folder_id
+ assert sm.endpoint == endpoint
+ assert sm.yc_connection_id == yc_connection_id
+
def test_yandex_lockbox_secret_backedn__get_endpoint(self):
endpoint = "api.cloud.yandex.net"
expected = {
@@ -315,8 +346,9 @@ class TestLockboxSecretBackend:
assert res is None
+
@patch("airflow.providers.yandex.secrets.lockbox.LockboxSecretBackend._client")
@patch("airflow.providers.yandex.secrets.lockbox.LockboxSecretBackend._list_secrets")
- def test_yandex_lockbox_secret_backend__get_secrets(self,
mock_list_secrets):
+ def test_yandex_lockbox_secret_backend__get_secrets(self,
mock_list_secrets, mock_client):
secrets = secret_service_pb.ListSecretsResponse(
secrets=[
secret_pb.Secret(
@@ -329,6 +361,7 @@ class TestLockboxSecretBackend:
)
mock_list_secrets.return_value = secrets
+ mock_client.return_value = None
res = LockboxSecretBackend(
folder_id="someid",
@@ -336,8 +369,9 @@ class TestLockboxSecretBackend:
assert res == secrets.secrets
+
@patch("airflow.providers.yandex.secrets.lockbox.LockboxSecretBackend._client")
@patch("airflow.providers.yandex.secrets.lockbox.LockboxSecretBackend._list_secrets")
- def test_yandex_lockbox_secret_backend__get_secrets_page_token(self,
mock_list_secrets):
+ def test_yandex_lockbox_secret_backend__get_secrets_page_token(self,
mock_list_secrets, mock_client):
first_secrets = secret_service_pb.ListSecretsResponse(
secrets=[
secret_pb.Secret(
@@ -365,6 +399,7 @@ class TestLockboxSecretBackend:
first_secrets,
second_secrets,
]
+ mock_client.return_value = None
res = LockboxSecretBackend(
folder_id="someid",
