Add multi-team lookup to Azure Key Vault backend (#65692)

Author: PrithviBadiga

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
 """"""""""""""

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
@@ -36,6 +37,8 @@
 from airflow.secrets import BaseSecretsBackend
 from airflow.utils.log.logging_mixin import LoggingMixin
 
+TEAM_SEP = "--"
+
 
 class AzureKeyVaultBackend(BaseSecretsBackend, LoggingMixin):
     """
@@ -154,7 +157,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 +173,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 +209,34 @@ 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:
+    def _build_team_secret_name(self, path_prefix: str, team_name: str, secret_id: str) -> str:
+        """Build a team-scoped secret name using a dedicated separator before the secret id."""
+        team_prefix = self.build_path(path_prefix, team_name, self.sep)
+        normalized_secret_id = secret_id.replace("_", self.sep)
+        return f"{team_prefix}{TEAM_SEP}{normalized_secret_id}"
+
+    def _is_team_specific_accessed_as_global(self, secret_id: str, team_name: str | None = None) -> bool:
+        normalized_secret_id = self.build_path("", secret_id, self.sep)
+        return team_name is None and bool(re.fullmatch(rf".+{re.escape(TEAM_SEP)}.+", normalized_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_secret = self._get_secret_value(
+                path_prefix, self._build_team_secret_name("", team_name, 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)

providers/microsoft/azure/tests/unit/microsoft/azure/secrets/test_key_vault.py

@@ -73,6 +73,46 @@ 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_uses_team_secret_with_custom_prefix(self, mock_client):
+        mock_client.get_secret.return_value = mock.Mock(value="team-value")
+
+        backend = AzureKeyVaultBackend(variables_prefix="custom-variables")
+        secret_val = backend.get_variable("hello", team_name="team_a")
+
+        assert secret_val == "team-value"
+        mock_client.get_secret.assert_called_once_with(name="custom-variables-team-a--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):
         """