Add multi team secret support to VaultBackend

providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py

@@ -21,6 +21,7 @@
 
 from typing import TYPE_CHECKING
 
+from airflow.providers.common.compat.sdk import conf
 from airflow.providers.hashicorp._internal_client.vault_client import _VaultClient
 from airflow.secrets import BaseSecretsBackend
 from airflow.utils.log.logging_mixin import LoggingMixin
@@ -52,6 +53,10 @@ class VaultBackend(BaseSecretsBackend, LoggingMixin):
         (default: 'variables'). If set to None (null), requests for variables will not be sent to Vault.
     :param config_path: Specifies the path of the secret to read Airflow Configurations
         (default: 'config'). If set to None (null), requests for configurations will not be sent to Vault.
+    :param use_team_secrets_path: Flag to enable team scoped secret retrieval from {base_path}/{team_name}/{key}
+        in multi team deployments. (default: true)
+    :param global_secrets_path: Path prefix to add to global scoped connections and variables in multi team deployments.
+        (default: No prefix)
     :param url: Base URL for the Vault instance being addressed.
     :param auth_type: Authentication Type for Vault. Default is ``token``. Available values are:
         ('approle', 'aws_iam', 'azure', 'github', 'gcp', 'jwt', 'kubernetes', 'ldap', 'radius', 'token', 'userpass')
@@ -99,6 +104,8 @@ def __init__(
         connections_path: str | None = "connections",
         variables_path: str | None = "variables",
         config_path: str | None = "config",
+        use_team_secrets_path: bool = True,
+        global_secrets_path: str | None = None,
         url: str | None = None,
         auth_type: str = "token",
         auth_mount_point: str | None = None,
@@ -132,6 +139,10 @@ def __init__(
         self.connections_path = connections_path.rstrip("/") if connections_path is not None else None
         self.variables_path = variables_path.rstrip("/") if variables_path is not None else None
         self.config_path = config_path.rstrip("/") if config_path is not None else None
+        self.use_team_secrets_path = use_team_secrets_path
+        self.global_secrets_path = (
+            global_secrets_path.rstrip("/") if global_secrets_path is not None else None
+        )
         self.mount_point = mount_point
         self.kv_engine_version = kv_engine_version
         self.vault_client = _VaultClient(
@@ -189,13 +200,29 @@ def _get_secret_with_base(self, base_path: str | None, key: str) -> dict | None:
             secret_path=(mount_point + "/" if mount_point else "") + secret_path
         )
 
-    def get_response(self, conn_id: str) -> dict | None:
+    def _get_team_or_global_secret(self, base_path: str | None, team_name: str | None, key: str):
         """
-        Get data from Vault.
+        Get a secret from a team specific path or the global path.
 
-        :return: The data from the Vault path if exists
+        If multi team is enabled, check {base_path}/{team_name}/{key}, then fallback to {base_path}/{global_path}/{key} or {base_path}/{key}.
         """
-        return self._get_secret_with_base(self.connections_path, conn_id)
+        if base_path is None:
+            return None
+        if (
+            conf.getboolean("core", "multi_team", fallback=False)
+            and self.use_team_secrets_path
+            and team_name is not None
+        ):
+            response = self._get_secret_with_base(self.build_path(base_path, team_name), key)
+            if response is not None:
+                return response
+        # Fallback to global secret
+        if conf.getboolean("core", "multi_team", fallback=False) and self.global_secrets_path is not None:
+            path = self.build_path(base_path, self.global_secrets_path)
+        else:
+            path = base_path
+
+        return self._get_secret_with_base(path, key)
 
     # Make sure connection is imported this way for type checking, otherwise when importing
     # the backend it will get a circular dependency and fail
@@ -214,7 +241,7 @@ def get_connection(self, conn_id: str, team_name: str | None = None) -> Connecti
         # problems when instantiating the backend during configuration
         from airflow.models.connection import Connection
 
-        response = self.get_response(conn_id)
+        response = self._get_team_or_global_secret(self.connections_path, team_name, conn_id)
         if response is None:
             return None
 
@@ -232,7 +259,7 @@ def get_variable(self, key: str, team_name: str | None = None) -> str | None:
         :param team_name: Team name associated to the task trying to access the variable (if any)
         :return: Variable Value retrieved from the vault
         """
-        response = self._get_secret_with_base(self.variables_path, key)
+        response = self._get_team_or_global_secret(self.variables_path, team_name, key)
 
         if not response:
             return None

providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py

@@ -23,8 +23,64 @@
 
 from airflow.providers.hashicorp.secrets.vault import VaultBackend
 
+from tests_common.test_utils.config import conf_vars
+
 
 class TestVaultSecrets:
+    @pytest.fixture
+    def secret_not_found(self):
+        return {"data": {"data": None}}
+
+    @pytest.fixture
+    def connection_result(self):
+        return {
+            "request_id": "94011e25-f8dc-ec29-221b-1f9c1d9ad2ae",
+            "lease_id": "",
+            "renewable": False,
+            "lease_duration": 0,
+            "data": {
+                "data": {
+                    "conn_type": "postgresql",
+                    "login": "airflow",
+                    "password": "airflow",
+                    "host": "host",
+                    "port": "5432",
+                    "schema": "airflow",
+                    "extra": '{"foo":"bar","baz":"taz"}',
+                },
+                "metadata": {
+                    "created_time": "2020-03-16T21:01:43.331126Z",
+                    "deletion_time": "",
+                    "destroyed": False,
+                    "version": 1,
+                },
+            },
+            "wrap_info": None,
+            "warnings": None,
+            "auth": None,
+        }
+
+    @pytest.fixture
+    def variable_result(self):
+        return {
+            "request_id": "2d48a2ad-6bcb-e5b6-429d-da35fdf31f56",
+            "lease_id": "",
+            "renewable": False,
+            "lease_duration": 0,
+            "data": {
+                "data": {"value": "world"},
+                "metadata": {
+                    "created_time": "2020-03-28T02:10:54.301784Z",
+                    "deletion_time": "",
+                    "destroyed": False,
+                    "version": 1,
+                },
+            },
+            "wrap_info": None,
+            "warnings": None,
+            "auth": None,
+        }
+
     @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
     def test_get_connection(self, mock_hvac):
         mock_client = mock.MagicMock()
@@ -68,6 +124,76 @@ def test_get_connection(self, mock_hvac):
         connection = test_client.get_connection(conn_id="test_postgres")
         assert connection.get_uri() == "postgresql://airflow:airflow@host:5432/airflow?foo=bar&baz=taz"
 
+    @pytest.mark.parametrize(
+        ("side_effects", "extra_kwargs", "exp_paths", "team_name"),
+        [
+            pytest.param(["connection_result"], {}, ["/foo/test_postgres"], "foo", id="team_conn"),
+            pytest.param(["connection_result"], {}, ["/test_postgres"], None, id="conn_no_team"),
+            pytest.param(
+                ["connection_result"],
+                {
+                    "use_team_secrets_path": False,
+                },
+                ["/test_postgres"],
+                "foo",
+                id="team_conn_no_separation",
+            ),
+            pytest.param(
+                ["secret_not_found", "connection_result"],
+                {"global_secrets_path": "global"},
+                ["/foo/test_postgres", "/global/test_postgres"],
+                "foo",
+                id="fallback_global_conn",
+            ),
+            pytest.param(
+                ["connection_result"],
+                {"global_secrets_path": "global"},
+                ["/global/test_postgres"],
+                None,
+                id="global_conn_no_team",
+            ),
+        ],
+    )
+    @conf_vars({("core", "multi_team"): "True"})
+    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
+    def test_get_connection_value_multi_team(
+        self, mock_hvac, side_effects, extra_kwargs, exp_paths, team_name, request
+    ):
+        read_secret_side_effects = []
+        for eft in side_effects:
+            # Populate side effects from fixtures
+            if eft is not None:
+                read_secret_side_effects.append(request.getfixturevalue(eft))
+            else:
+                read_secret_side_effects.append(None)
+        mock_client = mock.MagicMock()
+        mock_hvac.Client.return_value = mock_client
+        mock_client.secrets.kv.v2.read_secret_version.side_effect = read_secret_side_effects
+
+        kwargs = dict(
+            connections_path="connections",
+            mount_point="airflow",
+            auth_type="token",
+            url="http://127.0.0.1:8200",
+            token="s.7AU0I51yv1Q1lxOIg1F3ZRAS",
+            **extra_kwargs,
+        )
+
+        test_client = VaultBackend(**kwargs)
+        connection = test_client.get_connection(conn_id="test_postgres", team_name=team_name)
+        mock_client.secrets.kv.v2.read_secret_version.assert_has_calls(
+            [
+                mock.call(
+                    path=test_client.connections_path + path,
+                    mount_point="airflow",
+                    version=None,
+                    raise_on_deleted_version=True,
+                )
+                for path in exp_paths
+            ]
+        )
+        assert connection.get_uri() == "postgresql://airflow:airflow@host:5432/airflow?foo=bar&baz=taz"
+
     @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
     def test_get_connection_without_predefined_mount_point(self, mock_hvac):
         mock_client = mock.MagicMock()
@@ -154,6 +280,76 @@ def test_get_variable_value(self, mock_hvac):
         returned_uri = test_client.get_variable("hello")
         assert returned_uri == "world"
 
+    @pytest.mark.parametrize(
+        ("side_effects", "extra_kwargs", "exp_paths", "team_name"),
+        [
+            pytest.param(["variable_result"], {}, ["/foo/hello"], "foo", id="team_var"),
+            pytest.param(["variable_result"], {}, ["/hello"], None, id="no_team_var"),
+            pytest.param(
+                ["variable_result"],
+                {
+                    "use_team_secrets_path": False,
+                },
+                ["/hello"],
+                "foo",
+                id="team_var_no_separation",
+            ),
+            pytest.param(
+                ["secret_not_found", "variable_result"],
+                {"global_secrets_path": "global"},
+                ["/foo/hello", "/global/hello"],
+                "foo",
+                id="fallback_global_var",
+            ),
+            pytest.param(
+                ["variable_result"],
+                {"global_secrets_path": "global"},
+                ["/global/hello"],
+                None,
+                id="global_var_no_team",
+            ),
+        ],
+    )
+    @conf_vars({("core", "multi_team"): "True"})
+    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
+    def test_get_variable_value_multi_team(
+        self, mock_hvac, side_effects, extra_kwargs, exp_paths, team_name, request
+    ):
+        read_secret_side_effects = []
+        for eft in side_effects:
+            # Populate side effects from fixtures
+            if eft is not None:
+                read_secret_side_effects.append(request.getfixturevalue(eft))
+            else:
+                read_secret_side_effects.append(None)
+        mock_client = mock.MagicMock()
+        mock_hvac.Client.return_value = mock_client
+        mock_client.secrets.kv.v2.read_secret_version.side_effect = read_secret_side_effects
+
+        kwargs = dict(
+            connections_path="connections",
+            mount_point="airflow",
+            auth_type="token",
+            url="http://127.0.0.1:8200",
+            token="s.7AU0I51yv1Q1lxOIg1F3ZRAS",
+            **extra_kwargs,
+        )
+
+        test_client = VaultBackend(**kwargs)
+        returned_uri = test_client.get_variable("hello", team_name)
+        mock_client.secrets.kv.v2.read_secret_version.assert_has_calls(
+            [
+                mock.call(
+                    path=test_client.variables_path + path,
+                    mount_point="airflow",
+                    version=None,
+                    raise_on_deleted_version=True,
+                )
+                for path in exp_paths
+            ]
+        )
+        assert returned_uri == "world"
+
     @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
     def test_get_variable_value_without_predefined_mount_point(self, mock_hvac):
         mock_client = mock.MagicMock()