diff --git a/providers/cncf/kubernetes/docs/secrets-backends/kubernetes-secrets-backend.rst b/providers/cncf/kubernetes/docs/secrets-backends/kubernetes-secrets-backend.rst index 4f68de3128dc2..9ee9b123dcaf5 100644 --- a/providers/cncf/kubernetes/docs/secrets-backends/kubernetes-secrets-backend.rst +++ b/providers/cncf/kubernetes/docs/secrets-backends/kubernetes-secrets-backend.rst @@ -81,6 +81,7 @@ The following parameters can be passed via ``backend_kwargs`` as a JSON dictiona * ``connections_label``: Label key used to discover connection secrets. Default: ``"airflow.apache.org/connection-id"`` * ``variables_label``: Label key used to discover variable secrets. Default: ``"airflow.apache.org/variable-key"`` * ``config_label``: Label key used to discover config secrets. Default: ``"airflow.apache.org/config-key"`` +* ``team_label``: Label key used to discover team-scoped secrets in multi-team mode. Default: ``"airflow.apache.org/team"`` * ``connections_data_key``: The data key in the Kubernetes secret that holds the connection value. Default: ``"value"`` * ``variables_data_key``: The data key in the Kubernetes secret that holds the variable value. Default: ``"value"`` * ``config_data_key``: The data key in the Kubernetes secret that holds the config value. Default: ``"value"`` @@ -207,6 +208,33 @@ You can create a variable secret with ``kubectl``: airflow.apache.org/variable-key=my_var \ --namespace=airflow +Multi-team lookup +""""""""""""""""" + +In multi-team mode, this backend first looks for a secret whose identifier label matches the requested +connection or variable and whose ``team_label`` matches the current team. If no team-scoped secret is +found, it falls back to a global secret with the same identifier label and no team label. + +For example, with ``team_label="airflow.apache.org/team"``, ``team_name="team_a"``, and +``conn_id="my_db"``, the backend queries: + +* Team-scoped: ``airflow.apache.org/connection-id=my_db,airflow.apache.org/team=team_a`` +* Global fallback: ``airflow.apache.org/connection-id=my_db,!airflow.apache.org/team`` + +Example team-scoped connection secret: + +.. code-block:: yaml + + apiVersion: v1 + kind: Secret + metadata: + name: my-team-db-secret + labels: + airflow.apache.org/connection-id: my_db + airflow.apache.org/team: team_a + data: + value: + Using with External Secrets Operator """"""""""""""""""""""""""""""""""""" diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/secrets/kubernetes_secrets_backend.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/secrets/kubernetes_secrets_backend.py index ea67f7c3da6b2..68a58dda0ef29 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/secrets/kubernetes_secrets_backend.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/secrets/kubernetes_secrets_backend.py @@ -20,6 +20,7 @@ from __future__ import annotations import base64 +import re from functools import cached_property from pathlib import Path @@ -96,6 +97,7 @@ class KubernetesSecretsBackend(BaseSecretsBackend, LoggingMixin): DEFAULT_CONNECTIONS_LABEL = "airflow.apache.org/connection-id" DEFAULT_VARIABLES_LABEL = "airflow.apache.org/variable-key" DEFAULT_CONFIG_LABEL = "airflow.apache.org/config-key" + DEFAULT_TEAM_LABEL = "airflow.apache.org/team" SERVICE_ACCOUNT_NAMESPACE_PATH = "/var/run/secrets/kubernetes.io/serviceaccount/namespace" def __init__( @@ -104,6 +106,7 @@ def __init__( connections_label: str = DEFAULT_CONNECTIONS_LABEL, variables_label: str = DEFAULT_VARIABLES_LABEL, config_label: str = DEFAULT_CONFIG_LABEL, + team_label: str | None = DEFAULT_TEAM_LABEL, connections_data_key: str = "value", variables_data_key: str = "value", config_data_key: str = "value", @@ -114,6 +117,7 @@ def __init__( self.connections_label = connections_label self.variables_label = variables_label self.config_label = config_label + self.team_label = team_label self.connections_data_key = connections_data_key self.variables_data_key = variables_data_key self.config_data_key = config_data_key @@ -143,26 +147,28 @@ def get_conn_value(self, conn_id: str, team_name: str | None = None) -> str | No """ Get serialized representation of Connection from a Kubernetes secret. - Multi-team isolation is not currently supported; ``team_name`` is accepted - for API compatibility but ignored. - :param conn_id: connection id - :param team_name: Team name (unused — multi-team is not currently supported) + :param team_name: Team name associated to the task trying to access the connection (if any) """ - return self._get_secret(self.connections_label, conn_id, self.connections_data_key) + if self._is_team_specific_accessed_as_global(conn_id, team_name): + return None + + return self._get_secret( + self.connections_label, conn_id, self.connections_data_key, team_name=team_name + ) def get_variable(self, key: str, team_name: str | None = None) -> str | None: """ Get Airflow Variable from a Kubernetes secret. - Multi-team isolation is not currently supported; ``team_name`` is accepted - for API compatibility but ignored. - :param key: Variable Key - :param team_name: Team name (unused — multi-team is not currently supported) + :param team_name: Team name associated to the task trying to access the variable (if any) :return: Variable Value """ - return self._get_secret(self.variables_label, key, self.variables_data_key) + if self._is_team_specific_accessed_as_global(key, team_name): + return None + + return self._get_secret(self.variables_label, key, self.variables_data_key, team_name=team_name) def get_config(self, key: str) -> str | None: """ @@ -173,7 +179,13 @@ def get_config(self, key: str) -> str | None: """ return self._get_secret(self.config_label, key, self.config_data_key) - def _get_secret(self, label_key: str | None, label_value: str, data_key: str) -> str | None: + @staticmethod + def _is_team_specific_accessed_as_global(secret_id: str, team_name: str | None = None) -> bool: + return team_name is None and bool(re.fullmatch(r"_[^_]+___.+", secret_id)) + + def _get_secret( + self, label_key: str | None, label_value: str, data_key: str, team_name: str | None = None + ) -> str | None: """ Get secret value from Kubernetes by label selector. @@ -188,18 +200,43 @@ def _get_secret(self, label_key: str | None, label_value: str, data_key: str) -> """ if label_key is None: return None - label_selector = f"{label_key}={label_value}" + + if team_name and self.team_label: + team_secret = self._get_secret_by_selector( + label_key, label_value, data_key, f"{self.team_label}={team_name}", warn_if_missing=False + ) + if team_secret is not None: + return team_secret + + team_selector = f"!{self.team_label}" if self.team_label else None + return self._get_secret_by_selector(label_key, label_value, data_key, team_selector) + + def _get_secret_by_selector( + self, + label_key: str, + label_value: str, + data_key: str, + extra_selector: str | None, + *, + warn_if_missing: bool = True, + ) -> str | None: + """Get secret value from Kubernetes by the given base and optional extra selectors.""" + selectors = [f"{label_key}={label_value}"] + if extra_selector: + selectors.append(extra_selector) + label_selector = ",".join(selectors) secret_list = self.client.list_namespaced_secret( self.namespace, label_selector=label_selector, resource_version="0", ) if not secret_list.items: - self.log.warning( - "No secret found with label %s in namespace %s.", - label_selector, - self.namespace, - ) + if warn_if_missing: + self.log.warning( + "No secret found with label %s in namespace %s.", + label_selector, + self.namespace, + ) return None if len(secret_list.items) > 1: self.log.warning( diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/secrets/test_kubernetes_secrets_backend.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/secrets/test_kubernetes_secrets_backend.py index 6dced3d340bcc..68b53903534c3 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/secrets/test_kubernetes_secrets_backend.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/secrets/test_kubernetes_secrets_backend.py @@ -63,7 +63,7 @@ def test_get_conn_value_uri(self, mock_client, mock_namespace): assert result == uri mock_client.return_value.list_namespaced_secret.assert_called_once_with( "default", - label_selector="airflow.apache.org/connection-id=my_db", + label_selector="airflow.apache.org/connection-id=my_db,!airflow.apache.org/team", resource_version="0", ) @@ -103,6 +103,33 @@ def test_get_conn_value_not_found(self, mock_client, mock_namespace): assert result is None + @mock.patch(f"{MODULE_PATH}.namespace", new_callable=mock.PropertyMock, return_value="default") + @mock.patch(f"{MODULE_PATH}.client", new_callable=mock.PropertyMock) + def test_get_conn_value_uses_team_specific_secret_first(self, mock_client, mock_namespace): + mock_client.return_value.list_namespaced_secret.side_effect = [ + _make_secret_list([_make_secret({"value": "team-conn"})]), + ] + + backend = KubernetesSecretsBackend() + result = backend.get_conn_value("my_db", team_name="team_a") + + assert result == "team-conn" + mock_client.return_value.list_namespaced_secret.assert_called_once_with( + "default", + label_selector="airflow.apache.org/connection-id=my_db,airflow.apache.org/team=team_a", + resource_version="0", + ) + + @mock.patch(f"{MODULE_PATH}.namespace", new_callable=mock.PropertyMock, return_value="default") + @mock.patch(f"{MODULE_PATH}.client", new_callable=mock.PropertyMock) + def test_get_conn_value_returns_none_for_team_scoped_id_without_team_name( + self, mock_client, mock_namespace + ): + backend = KubernetesSecretsBackend() + + assert backend.get_conn_value("_teama___my_db") is None + mock_client.return_value.list_namespaced_secret.assert_not_called() + class TestKubernetesSecretsBackendVariables: @mock.patch(f"{MODULE_PATH}.namespace", new_callable=mock.PropertyMock, return_value="default") @@ -119,7 +146,7 @@ def test_get_variable(self, mock_client, mock_namespace): assert result == "my-value" mock_client.return_value.list_namespaced_secret.assert_called_once_with( "default", - label_selector="airflow.apache.org/variable-key=api_key", + label_selector="airflow.apache.org/variable-key=api_key,!airflow.apache.org/team", resource_version="0", ) @@ -134,6 +161,33 @@ def test_get_variable_not_found(self, mock_client, mock_namespace): assert result is None + @mock.patch(f"{MODULE_PATH}.namespace", new_callable=mock.PropertyMock, return_value="default") + @mock.patch(f"{MODULE_PATH}.client", new_callable=mock.PropertyMock) + def test_get_variable_falls_back_to_global_secret_when_team_secret_is_missing( + self, mock_client, mock_namespace + ): + mock_client.return_value.list_namespaced_secret.side_effect = [ + _make_secret_list([]), + _make_secret_list([_make_secret({"value": "global-value"})]), + ] + + backend = KubernetesSecretsBackend() + result = backend.get_variable("api_key", team_name="team_a") + + assert result == "global-value" + assert mock_client.return_value.list_namespaced_secret.call_args_list == [ + mock.call( + "default", + label_selector="airflow.apache.org/variable-key=api_key,airflow.apache.org/team=team_a", + resource_version="0", + ), + mock.call( + "default", + label_selector="airflow.apache.org/variable-key=api_key,!airflow.apache.org/team", + resource_version="0", + ), + ] + class TestKubernetesSecretsBackendConfig: @mock.patch(f"{MODULE_PATH}.namespace", new_callable=mock.PropertyMock, return_value="default") @@ -150,7 +204,7 @@ def test_get_config(self, mock_client, mock_namespace): assert result == "sqlite:///airflow.db" mock_client.return_value.list_namespaced_secret.assert_called_once_with( "default", - label_selector="airflow.apache.org/config-key=sql_alchemy_conn", + label_selector="airflow.apache.org/config-key=sql_alchemy_conn,!airflow.apache.org/team", resource_version="0", ) @@ -181,7 +235,7 @@ def test_custom_label(self, mock_client, mock_namespace): assert result == "postgresql://localhost/db" mock_client.return_value.list_namespaced_secret.assert_called_once_with( "default", - label_selector="my-org/conn=my_db", + label_selector="my-org/conn=my_db,!airflow.apache.org/team", resource_version="0", ) @@ -225,6 +279,23 @@ def test_secret_with_none_data_returns_none(self, mock_client, mock_namespace): assert result is None + @mock.patch(f"{MODULE_PATH}.namespace", new_callable=mock.PropertyMock, return_value="default") + @mock.patch(f"{MODULE_PATH}.client", new_callable=mock.PropertyMock) + def test_team_specific_lookup_uses_custom_team_label(self, mock_client, mock_namespace): + mock_client.return_value.list_namespaced_secret.return_value = _make_secret_list( + [_make_secret({"value": "team-conn"})] + ) + + backend = KubernetesSecretsBackend(team_label="my-org.io/team") + result = backend.get_conn_value("my_db", team_name="team_a") + + assert result == "team-conn" + mock_client.return_value.list_namespaced_secret.assert_called_once_with( + "default", + label_selector="airflow.apache.org/connection-id=my_db,my-org.io/team=team_a", + resource_version="0", + ) + class TestKubernetesSecretsBackendLabelNone: @mock.patch(f"{MODULE_PATH}.client", new_callable=mock.PropertyMock) @@ -332,7 +403,7 @@ def test_namespace_used_in_api_calls(self, mock_client, mock_namespace): mock_client.return_value.list_namespaced_secret.assert_called_once_with( "airflow", - label_selector="airflow.apache.org/connection-id=my_db", + label_selector="airflow.apache.org/connection-id=my_db,!airflow.apache.org/team", resource_version="0", ) diff --git a/providers/microsoft/azure/docs/secrets-backends/azure-key-vault.rst b/providers/microsoft/azure/docs/secrets-backends/azure-key-vault.rst index 942d6d806f9e0..4a2ead30bef0a 100644 --- a/providers/microsoft/azure/docs/secrets-backends/azure-key-vault.rst +++ b/providers/microsoft/azure/docs/secrets-backends/azure-key-vault.rst @@ -67,6 +67,26 @@ Storing and Retrieving Variables If you have set ``variables_prefix`` as ``airflow-variables``, then for an Variable key of ``hello``, you would want to store your Variable at ``airflow-variables-hello``. +Multi-team lookup +""""""""""""""""" + +In multi-team mode, this backend looks for team-scoped secrets first and falls back to the global +secret name when a team-scoped secret is not found. + +For connections: + +* Team-scoped: ``{connections_prefix}-{team_name}-{conn_id}`` +* Global fallback: ``{connections_prefix}-{conn_id}`` + +For variables: + +* Team-scoped: ``{variables_prefix}-{team_name}-{key}`` +* Global fallback: ``{variables_prefix}-{key}`` + +Underscores are normalized to the configured separator, so with ``connections_prefix="airflow-connections"``, +``team_name="team_a"``, and ``conn_id="my_db"``, the backend looks up +``airflow-connections-team-a-my-db`` before falling back to ``airflow-connections-my-db``. + Authentication """""""""""""" diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/secrets/key_vault.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/secrets/key_vault.py index 613dc3310159d..93677ed9ecd9e 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/secrets/key_vault.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/secrets/key_vault.py @@ -26,6 +26,7 @@ import logging import os +import re from functools import cached_property from azure.core.exceptions import ResourceNotFoundError @@ -154,7 +155,10 @@ def get_conn_value(self, conn_id: str, team_name: str | None = None) -> str | No if self.connections_prefix is None: return None - return self._get_secret(self.connections_prefix, conn_id) + if self._is_team_specific_accessed_as_global(conn_id, team_name): + return None + + return self._get_secret(self.connections_prefix, conn_id, team_name=team_name) def get_variable(self, key: str, team_name: str | None = None) -> str | None: """ @@ -167,7 +171,10 @@ def get_variable(self, key: str, team_name: str | None = None) -> str | None: if self.variables_prefix is None: return None - return self._get_secret(self.variables_prefix, key) + if self._is_team_specific_accessed_as_global(key, team_name): + return None + + return self._get_secret(self.variables_prefix, key, team_name=team_name) def get_config(self, key: str) -> str | None: """ @@ -200,13 +207,27 @@ def build_path(path_prefix: str, secret_id: str, sep: str = "-") -> str: path = f"{path_prefix}{sep}{secret_id}" return path.replace("_", sep) - def _get_secret(self, path_prefix: str, secret_id: str) -> str | None: + @staticmethod + def _is_team_specific_accessed_as_global(secret_id: str, team_name: str | None = None) -> bool: + return team_name is None and bool(re.fullmatch(r"_[^_]+___.+", secret_id)) + + def _get_secret(self, path_prefix: str, secret_id: str, team_name: str | None = None) -> str | None: """ Get an Azure Key Vault secret value. :param path_prefix: Prefix for the Path to get Secret :param secret_id: Secret Key """ + if team_name: + team_prefix = self.build_path(path_prefix, team_name, self.sep) + team_secret = self._get_secret_value(team_prefix, secret_id) + if team_secret is not None: + return team_secret + + return self._get_secret_value(path_prefix, secret_id) + + def _get_secret_value(self, path_prefix: str, secret_id: str) -> str | None: + """Get an Azure Key Vault secret value for the given prefix and key.""" name = self.build_path(path_prefix, secret_id, self.sep) try: secret = self.client.get_secret(name=name) diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/secrets/test_key_vault.py b/providers/microsoft/azure/tests/unit/microsoft/azure/secrets/test_key_vault.py index 6ad744f21968d..03e73a4c119e2 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/secrets/test_key_vault.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/secrets/test_key_vault.py @@ -73,6 +73,36 @@ def test_get_secret_value(self, mock_client): mock_client.get_secret.assert_called_with(name="af-secrets-test-mysql-password") assert secret_val == "super-secret" + @mock.patch(f"{KEY_VAULT_MODULE}.AzureKeyVaultBackend.client") + def test_get_conn_value_uses_team_specific_secret_first(self, mock_client): + mock_client.get_secret.return_value = mock.Mock(value="team-secret") + + backend = AzureKeyVaultBackend() + secret_val = backend.get_conn_value("my_db", team_name="team_a") + + assert secret_val == "team-secret" + mock_client.get_secret.assert_called_once_with(name="airflow-connections-team-a-my-db") + + @mock.patch(f"{KEY_VAULT_MODULE}.AzureKeyVaultBackend.client") + def test_get_variable_falls_back_to_global_secret_when_team_secret_is_missing(self, mock_client): + mock_client.get_secret.side_effect = [ResourceNotFoundError, mock.Mock(value="global-value")] + + backend = AzureKeyVaultBackend() + secret_val = backend.get_variable("hello", team_name="team_a") + + assert secret_val == "global-value" + assert mock_client.get_secret.call_args_list == [ + mock.call(name="airflow-variables-team-a-hello"), + mock.call(name="airflow-variables-hello"), + ] + + @mock.patch(f"{KEY_VAULT_MODULE}.AzureKeyVaultBackend.client") + def test_get_variable_returns_none_for_team_scoped_key_without_team_name(self, mock_client): + backend = AzureKeyVaultBackend() + + assert backend.get_variable("_teama___hello") is None + mock_client.get_secret.assert_not_called() + @mock.patch(f"{KEY_VAULT_MODULE}.AzureKeyVaultBackend._get_secret") def test_variable_prefix_none_value(self, mock_get_secret): """ diff --git a/providers/yandex/docs/secrets-backends/yandex-cloud-lockbox-secret-backend.rst b/providers/yandex/docs/secrets-backends/yandex-cloud-lockbox-secret-backend.rst index f0baea493f1ac..b3c69ed048fca 100644 --- a/providers/yandex/docs/secrets-backends/yandex-cloud-lockbox-secret-backend.rst +++ b/providers/yandex/docs/secrets-backends/yandex-cloud-lockbox-secret-backend.rst @@ -251,6 +251,26 @@ To check the variable is correctly read from the Lockbox Secret Backend, you can $ airflow variables get my_variable some_secret_data +Multi-team lookup +----------------- + +In multi-team mode, this backend looks for team-scoped secret names first and falls back to the +global secret name when a team-scoped secret is not found. + +For connections: + +* Team-scoped: ``{connections_prefix}/{team_name}/{conn_id}`` +* Global fallback: ``{connections_prefix}/{conn_id}`` + +For variables: + +* Team-scoped: ``{variables_prefix}/{team_name}/{key}`` +* Global fallback: ``{variables_prefix}/{key}`` + +For example, with ``connections_prefix="airflow/connections"``, ``team_name="team_a"``, and +``conn_id="my_db"``, the backend looks up ``airflow/connections/team_a/my_db`` before falling back +to ``airflow/connections/my_db``. + Storing and retrieving configs ------------------------------ diff --git a/providers/yandex/src/airflow/providers/yandex/secrets/lockbox.py b/providers/yandex/src/airflow/providers/yandex/secrets/lockbox.py index 47717f1876fa3..a74d789e755a8 100644 --- a/providers/yandex/src/airflow/providers/yandex/secrets/lockbox.py +++ b/providers/yandex/src/airflow/providers/yandex/secrets/lockbox.py @@ -18,6 +18,7 @@ from __future__ import annotations +import re from functools import cached_property from typing import Any @@ -159,7 +160,10 @@ def get_conn_value(self, conn_id: str, team_name: str | None = None) -> str | No if conn_id == self.yc_connection_id: return None - return self._get_secret_value(self.connections_prefix, conn_id) + if self._is_team_specific_accessed_as_global(conn_id, team_name): + return None + + return self._get_secret_value(self.connections_prefix, conn_id, team_name=team_name) def get_variable(self, key: str, team_name: str | None = None) -> str | None: """ @@ -172,7 +176,10 @@ def get_variable(self, key: str, team_name: str | None = None) -> str | None: if self.variables_prefix is None: return None - return self._get_secret_value(self.variables_prefix, key) + if self._is_team_specific_accessed_as_global(key, team_name): + return None + + return self._get_secret_value(self.variables_prefix, key, team_name=team_name) def get_config(self, key: str) -> str | None: """ @@ -243,12 +250,20 @@ def _build_secret_name(self, prefix: str, key: str): return key return f"{prefix}{self.sep}{key}" - def _get_secret_value(self, prefix: str, key: str) -> str | None: + @staticmethod + def _is_team_specific_accessed_as_global(secret_id: str, team_name: str | None = None) -> bool: + return team_name is None and bool(re.fullmatch(r"_[^_]+___.+", secret_id)) + + def _get_secret_value(self, prefix: str, key: str, team_name: str | None = None) -> str | None: + secrets = self._get_secrets() secret: secret_pb.Secret | None = None - for s in self._get_secrets(): - if s.name == self._build_secret_name(prefix=prefix, key=key): - secret = s - break + if team_name: + team_prefix = self._build_secret_name(prefix, team_name) + secret = self._find_secret(secrets, team_prefix, key) + + if not secret: + secret = self._find_secret(secrets, prefix, key) + if not secret: return None @@ -259,6 +274,12 @@ def _get_secret_value(self, prefix: str, key: str) -> str | None: return None return sorted(entries.values())[0] + def _find_secret(self, secrets: list[secret_pb.Secret], prefix: str, key: str) -> secret_pb.Secret | None: + for s in secrets: + if s.name == self._build_secret_name(prefix=prefix, key=key): + return s + return None + def _get_secrets(self) -> list[secret_pb.Secret]: # generate client if not exists, to load folder_id from connections _ = self._client diff --git a/providers/yandex/tests/unit/yandex/secrets/test_lockbox.py b/providers/yandex/tests/unit/yandex/secrets/test_lockbox.py index 7b9ab19d3e87a..f5602c31c66f4 100644 --- a/providers/yandex/tests/unit/yandex/secrets/test_lockbox.py +++ b/providers/yandex/tests/unit/yandex/secrets/test_lockbox.py @@ -17,7 +17,7 @@ from __future__ import annotations import json -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import ANY, MagicMock, Mock, patch import pytest @@ -261,6 +261,57 @@ def test_yandex_lockbox_secret_backend__build_secret_name_custom_sep(self): assert res == expected + @patch("airflow.providers.yandex.secrets.lockbox.LockboxSecretBackend._get_secrets") + @patch("airflow.providers.yandex.secrets.lockbox.LockboxSecretBackend._get_payload") + def test_get_conn_value_uses_team_specific_secret_first(self, mock_get_payload, mock_get_secrets): + mock_get_secrets.return_value = [ + secret_pb.Secret( + id="123", + name="airflow/connections/team_a/my_db", + ), + secret_pb.Secret( + id="456", + name="airflow/connections/my_db", + ), + ] + mock_get_payload.return_value = payload_pb.Payload( + entries=[payload_pb.Payload.Entry(text_value="team-conn")] + ) + + backend = LockboxSecretBackend() + result = backend.get_conn_value("my_db", team_name="team_a") + + assert result == "team-conn" + mock_get_payload.assert_called_once_with("123", ANY) + + @patch("airflow.providers.yandex.secrets.lockbox.LockboxSecretBackend._get_secrets") + @patch("airflow.providers.yandex.secrets.lockbox.LockboxSecretBackend._get_payload") + def test_get_variable_falls_back_to_global_secret_when_team_secret_is_missing( + self, mock_get_payload, mock_get_secrets + ): + mock_get_secrets.return_value = [ + secret_pb.Secret( + id="456", + name="airflow/variables/hello", + ), + ] + mock_get_payload.return_value = payload_pb.Payload( + entries=[payload_pb.Payload.Entry(text_value="global-value")] + ) + + backend = LockboxSecretBackend() + result = backend.get_variable("hello", team_name="team_a") + + assert result == "global-value" + mock_get_payload.assert_called_once_with("456", ANY) + + @patch("airflow.providers.yandex.secrets.lockbox.LockboxSecretBackend._get_secrets") + def test_get_variable_returns_none_for_team_scoped_key_without_team_name(self, mock_get_secrets): + backend = LockboxSecretBackend() + + assert backend.get_variable("_teama___hello") is None + mock_get_secrets.assert_not_called() + @patch("airflow.providers.yandex.secrets.lockbox.LockboxSecretBackend._get_secrets") @patch("airflow.providers.yandex.secrets.lockbox.LockboxSecretBackend._get_payload") def test_yandex_lockbox_secret_backend__get_secret_value(self, mock_get_payload, mock_get_secrets):