From 910da17121a76cee034f4544452b972585996337 Mon Sep 17 00:00:00 2001 From: ahilashsasidharan Date: Thu, 19 Mar 2026 00:06:49 -0400 Subject: [PATCH 1/4] Add API check to ensure multi team is enabled when team_name is provided --- .../core_api/datamodels/connections.py | 9 ++- .../api_fastapi/core_api/datamodels/pools.py | 15 ++++- .../core_api/datamodels/variables.py | 7 ++ .../routes/public/test_connections.py | 64 ++++++++++++++++--- .../core_api/routes/public/test_pools.py | 45 ++++++++++++- 5 files changed, 127 insertions(+), 13 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py index a44edb5b91b43..aac7af789fa33 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py @@ -21,11 +21,12 @@ from collections.abc import Iterable, Mapping from typing import Annotated, Any -from pydantic import Field, field_validator +from pydantic import Field, field_validator, model_validator from pydantic_core.core_schema import ValidationInfo from airflow._shared.secrets_masker import redact, should_hide_value_for_key from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel, make_partial_model +from airflow.configuration import conf # Response Models @@ -199,5 +200,11 @@ def validate_extra(cls, v: str | None) -> str | None: ) return v + @model_validator(mode="after") + def validate_team_name(self) -> ConnectionBody: + if self.team_name is not None and not conf.getboolean("core", "multi_team"): + raise ValueError("team_name cannot be set when multi_team mode is disabled") + return self + ConnectionBodyPartial = make_partial_model(ConnectionBody) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/pools.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/pools.py index a661308d4434c..3c8853da0e2ba 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/pools.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/pools.py @@ -20,9 +20,10 @@ from collections.abc import Callable, Iterable from typing import Annotated -from pydantic import BeforeValidator, Field +from pydantic import BeforeValidator, Field, model_validator from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel +from airflow.configuration import conf def _call_function(function: Callable[[], int]) -> int: @@ -83,6 +84,12 @@ class PoolPatchBody(StrictBaseModel): include_deferred: bool | None = None team_name: str | None = Field(max_length=50, default=None) + @model_validator(mode="after") + def validate_team_name(self) -> PoolPatchBody: + if self.team_name is not None and not conf.getboolean("core", "multi_team"): + raise ValueError("team_name cannot be set when multi_team mode is disabled") + return self + class PoolBody(BasePool, StrictBaseModel): """Pool serializer for post bodies.""" @@ -91,3 +98,9 @@ class PoolBody(BasePool, StrictBaseModel): description: str | None = None include_deferred: bool = False team_name: str | None = Field(max_length=50, default=None) + + @model_validator(mode="after") + def validate_team_name(self) -> PoolBody: + if self.team_name is not None and not conf.getboolean("core", "multi_team"): + raise ValueError("team_name cannot be set when multi_team mode is disabled") + return self diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/variables.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/variables.py index 9dc7969b69b6a..a4f1ff334f178 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/variables.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/variables.py @@ -24,6 +24,7 @@ from airflow._shared.secrets_masker import redact from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel, make_partial_model +from airflow.configuration import conf from airflow.models.base import ID_LEN from airflow.typing_compat import Self @@ -60,6 +61,12 @@ class VariableBody(StrictBaseModel): description: str | None = Field(default=None) team_name: str | None = Field(max_length=50, default=None) + @model_validator(mode="after") + def validate_team_name(self) -> VariableBody: + if self.team_name is not None and not conf.getboolean("core", "multi_team"): + raise ValueError("team_name cannot be set when multi_team mode is disabled") + return self + VariableBodyPartial = make_partial_model(VariableBody) diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 1f63247cfa9ab..adbacd13e1833 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -287,10 +287,17 @@ def test_post_should_respond_201(self, test_client, session, body): _check_last_log(session, dag_id=None, event="post_connection", logical_date=None) def test_post_should_respond_201_with_team(self, test_client, session, testing_team): - response = test_client.post( - "/connections", - json={"connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, "team_name": testing_team.name}, - ) + with mock.patch( + "airflow.api_fastapi.core_api.datamodels.connections.conf.getboolean", return_value=True + ): + response = test_client.post( + "/connections", + json={ + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "team_name": testing_team.name, + }, + ) assert response.status_code == 201 assert response.json() == { "connection_id": TEST_CONN_ID, @@ -338,6 +345,24 @@ def test_post_should_respond_422_for_invalid_conn_id(self, test_client, body): ] } + def test_post_rejects_team_name_when_multi_team_disabled(self, test_client, session, testing_team): + with mock.patch( + "airflow.api_fastapi.core_api.datamodels.connections.conf.getboolean", return_value=False + ): + response = test_client.post( + "/connections", + json={ + "connection_id": TEST_CONN_ID_2, + "conn_type": TEST_CONN_TYPE_2, + "team_name": testing_team.name, + }, + ) + assert response.status_code == 422 + assert ( + response.json()["detail"][0]["msg"] + == "Value error, team_name cannot be set when multi_team mode is disabled" + ) + @pytest.mark.parametrize( "body", [ @@ -605,10 +630,13 @@ def test_patch_should_respond_200( def test_patch_with_team_should_respond_200(self, test_client, testing_team, session): self.create_connection() - response = test_client.patch( - f"/connections/{TEST_CONN_ID}", - json={"connection_id": TEST_CONN_ID, "conn_type": "new_type", "team_name": testing_team.name}, - ) + with mock.patch( + "airflow.api_fastapi.core_api.datamodels.connections.conf.getboolean", return_value=True + ): + response = test_client.patch( + f"/connections/{TEST_CONN_ID}", + json={"connection_id": TEST_CONN_ID, "conn_type": "new_type", "team_name": testing_team.name}, + ) assert response.status_code == 200 _check_last_log(session, dag_id=None, event="patch_connection", logical_date=None) @@ -966,6 +994,26 @@ def test_patch_with_update_mask_rejects_extra_fields(self, test_client): ) assert response.status_code == 422 + def test_patch_rejects_team_name_when_multi_team_disabled(self, test_client, testing_team, session): + self.create_connection() + + with mock.patch( + "airflow.api_fastapi.core_api.datamodels.connections.conf.getboolean", return_value=False + ): + response = test_client.patch( + f"/connections/{TEST_CONN_ID_2}", + json={ + "connection_id": TEST_CONN_ID_2, + "conn_type": "new_type", + "team_name": testing_team.name, + }, + ) + assert response.status_code == 422 + assert ( + response.json()["detail"][0]["msg"] + == "Value error, team_name cannot be set when multi_team mode is disabled" + ) + class TestConnection(TestConnectionEndpoint): def setup_method(self): diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py index e2ef872073530..c932279f0cc40 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py @@ -416,6 +416,24 @@ def test_patch_pool3_should_respond_200(self, test_client, session): assert response.json() == expected_response check_last_log(session, dag_id=None, event="patch_pool", logical_date=None) + def test_patch_pool_rejects_team_name_when_multi_team_disabled(self, test_client, session): + self.create_pools() + with mock.patch("airflow.api_fastapi.core_api.datamodels.pools.conf.getboolean", return_value=False): + response = test_client.patch( + f"/pools/{POOL2_NAME}", + json={ + "name": POOL2_NAME, + "slots": POOL2_SLOT, + "include_deferred": POOL2_INCLUDE_DEFERRED, + "team_name": "test", + }, + ) + assert response.status_code == 422 + assert ( + response.json()["detail"][0]["msg"] + == "Value error, team_name cannot be set when multi_team mode is disabled" + ) + class TestPostPool(TestPoolsEndpoint): @pytest.mark.parametrize( @@ -483,9 +501,11 @@ class TestPostPool(TestPoolsEndpoint): def test_should_respond_200(self, test_client, session, body, expected_status_code, expected_response): self.create_pools() n_pools = session.scalar(select(func.count()).select_from(Pool)) - response = test_client.post("/pools", json=body) - assert response.status_code == expected_status_code + with mock.patch("airflow.api_fastapi.core_api.datamodels.pools.conf.getboolean", return_value=True): + response = test_client.post("/pools", json=body) + + assert response.status_code == expected_status_code assert response.json() == expected_response assert session.scalar(select(func.count()).select_from(Pool)) == n_pools + 1 check_last_log(session, dag_id=None, event="post_pool", logical_date=None) @@ -523,6 +543,22 @@ def test_post_pool_rejects_infinity_string(self, test_client, session): ) assert response.status_code == 422 + def test_post_pool_rejects_team_name_when_multi_team_disabled(self, test_client, session): + with mock.patch("airflow.api_fastapi.core_api.datamodels.pools.conf.getboolean", return_value=False): + response = test_client.post( + "/pools", + json={ + "name": "bad_team_pool", + "slots": 1, + "team_name": "test", + }, + ) + assert response.status_code == 422 + assert ( + response.json()["detail"][0]["msg"] + == "Value error, team_name cannot be set when multi_team mode is disabled" + ) + def test_should_respond_401(self, unauthenticated_test_client): response = unauthenticated_test_client.post("/pools", json={}) assert response.status_code == 401 @@ -1045,7 +1081,10 @@ class TestBulkPools(TestPoolsEndpoint): ) def test_bulk_pools(self, test_client, actions, expected_results, session): self.create_pools() - response = test_client.patch("/pools", json=actions) + + with mock.patch("airflow.api_fastapi.core_api.datamodels.pools.conf.getboolean", return_value=True): + response = test_client.patch("/pools", json=actions) + response_data = response.json() for key, value in expected_results.items(): assert response_data[key] == value From ecfb9e9f8c2f6a06f2af9916e205e8d4e5fb0364 Mon Sep 17 00:00:00 2001 From: ahilashsasidharan Date: Thu, 19 Mar 2026 23:47:38 -0400 Subject: [PATCH 2/4] remove unnecessary arguments in added tests --- .../api_fastapi/core_api/routes/public/test_connections.py | 4 ++-- .../unit/api_fastapi/core_api/routes/public/test_pools.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index adbacd13e1833..2491a59581c84 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -345,7 +345,7 @@ def test_post_should_respond_422_for_invalid_conn_id(self, test_client, body): ] } - def test_post_rejects_team_name_when_multi_team_disabled(self, test_client, session, testing_team): + def test_post_rejects_team_name_when_multi_team_disabled(self, test_client, testing_team): with mock.patch( "airflow.api_fastapi.core_api.datamodels.connections.conf.getboolean", return_value=False ): @@ -994,7 +994,7 @@ def test_patch_with_update_mask_rejects_extra_fields(self, test_client): ) assert response.status_code == 422 - def test_patch_rejects_team_name_when_multi_team_disabled(self, test_client, testing_team, session): + def test_patch_rejects_team_name_when_multi_team_disabled(self, test_client, testing_team): self.create_connection() with mock.patch( diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py index c932279f0cc40..07ebebd497748 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py @@ -416,7 +416,7 @@ def test_patch_pool3_should_respond_200(self, test_client, session): assert response.json() == expected_response check_last_log(session, dag_id=None, event="patch_pool", logical_date=None) - def test_patch_pool_rejects_team_name_when_multi_team_disabled(self, test_client, session): + def test_patch_pool_rejects_team_name_when_multi_team_disabled(self, test_client): self.create_pools() with mock.patch("airflow.api_fastapi.core_api.datamodels.pools.conf.getboolean", return_value=False): response = test_client.patch( @@ -543,7 +543,7 @@ def test_post_pool_rejects_infinity_string(self, test_client, session): ) assert response.status_code == 422 - def test_post_pool_rejects_team_name_when_multi_team_disabled(self, test_client, session): + def test_post_pool_rejects_team_name_when_multi_team_disabled(self, test_client): with mock.patch("airflow.api_fastapi.core_api.datamodels.pools.conf.getboolean", return_value=False): response = test_client.post( "/pools", From 809f13ff744942206830d3b4e7c780dd64534aae Mon Sep 17 00:00:00 2001 From: ahilashsasidharan Date: Fri, 20 Mar 2026 08:45:37 -0400 Subject: [PATCH 3/4] add variable tests and add slight change to other tests to align with variables test file --- .../routes/public/test_connections.py | 17 +++------- .../core_api/routes/public/test_pools.py | 9 +++--- .../core_api/routes/public/test_variables.py | 32 +++++++++++++++++++ 3 files changed, 42 insertions(+), 16 deletions(-) diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 2491a59581c84..137ffee0e74c1 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -34,6 +34,7 @@ from tests_common.test_utils.api_fastapi import _check_last_log from tests_common.test_utils.asserts import assert_queries_count +from tests_common.test_utils.config import conf_vars from tests_common.test_utils.db import clear_db_connections, clear_db_logs, clear_test_connections from tests_common.test_utils.markers import skip_if_force_lowest_dependencies_marker @@ -287,9 +288,7 @@ def test_post_should_respond_201(self, test_client, session, body): _check_last_log(session, dag_id=None, event="post_connection", logical_date=None) def test_post_should_respond_201_with_team(self, test_client, session, testing_team): - with mock.patch( - "airflow.api_fastapi.core_api.datamodels.connections.conf.getboolean", return_value=True - ): + with conf_vars({("core", "multi_team"): "True"}): response = test_client.post( "/connections", json={ @@ -346,9 +345,7 @@ def test_post_should_respond_422_for_invalid_conn_id(self, test_client, body): } def test_post_rejects_team_name_when_multi_team_disabled(self, test_client, testing_team): - with mock.patch( - "airflow.api_fastapi.core_api.datamodels.connections.conf.getboolean", return_value=False - ): + with conf_vars({("core", "multi_team"): "False"}): response = test_client.post( "/connections", json={ @@ -630,9 +627,7 @@ def test_patch_should_respond_200( def test_patch_with_team_should_respond_200(self, test_client, testing_team, session): self.create_connection() - with mock.patch( - "airflow.api_fastapi.core_api.datamodels.connections.conf.getboolean", return_value=True - ): + with conf_vars({("core", "multi_team"): "True"}): response = test_client.patch( f"/connections/{TEST_CONN_ID}", json={"connection_id": TEST_CONN_ID, "conn_type": "new_type", "team_name": testing_team.name}, @@ -997,9 +992,7 @@ def test_patch_with_update_mask_rejects_extra_fields(self, test_client): def test_patch_rejects_team_name_when_multi_team_disabled(self, test_client, testing_team): self.create_connection() - with mock.patch( - "airflow.api_fastapi.core_api.datamodels.connections.conf.getboolean", return_value=False - ): + with conf_vars({("core", "multi_team"): "False"}): response = test_client.patch( f"/connections/{TEST_CONN_ID_2}", json={ diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py index 07ebebd497748..1459b75ecbd51 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py @@ -25,6 +25,7 @@ from airflow.models.team import Team from airflow.utils.session import provide_session +from tests_common.test_utils.config import conf_vars from tests_common.test_utils.db import clear_db_pools, clear_db_teams from tests_common.test_utils.logs import check_last_log @@ -418,7 +419,7 @@ def test_patch_pool3_should_respond_200(self, test_client, session): def test_patch_pool_rejects_team_name_when_multi_team_disabled(self, test_client): self.create_pools() - with mock.patch("airflow.api_fastapi.core_api.datamodels.pools.conf.getboolean", return_value=False): + with conf_vars({("core", "multi_team"): "False"}): response = test_client.patch( f"/pools/{POOL2_NAME}", json={ @@ -502,7 +503,7 @@ def test_should_respond_200(self, test_client, session, body, expected_status_co self.create_pools() n_pools = session.scalar(select(func.count()).select_from(Pool)) - with mock.patch("airflow.api_fastapi.core_api.datamodels.pools.conf.getboolean", return_value=True): + with conf_vars({("core", "multi_team"): "True"}): response = test_client.post("/pools", json=body) assert response.status_code == expected_status_code @@ -544,7 +545,7 @@ def test_post_pool_rejects_infinity_string(self, test_client, session): assert response.status_code == 422 def test_post_pool_rejects_team_name_when_multi_team_disabled(self, test_client): - with mock.patch("airflow.api_fastapi.core_api.datamodels.pools.conf.getboolean", return_value=False): + with conf_vars({("core", "multi_team"): "False"}): response = test_client.post( "/pools", json={ @@ -1082,7 +1083,7 @@ class TestBulkPools(TestPoolsEndpoint): def test_bulk_pools(self, test_client, actions, expected_results, session): self.create_pools() - with mock.patch("airflow.api_fastapi.core_api.datamodels.pools.conf.getboolean", return_value=True): + with conf_vars({("core", "multi_team"): "True"}): response = test_client.patch("/pools", json=actions) response_data = response.json() diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py index 508718796755a..2386bb441e628 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py @@ -521,6 +521,22 @@ def test_patch_should_respond_404(self, test_client): body = response.json() assert f"The Variable with key: `{TEST_VARIABLE_KEY}` was not found" == body["detail"] + def test_patch_rejects_team_name_when_multi_team_disabled(self, test_client, testing_team): + self.create_variables() + body = { + "key": TEST_VARIABLE_KEY, + "value": "The new value", + "description": "The new description", + "team_name": str(testing_team.name), + } + with conf_vars({("core", "multi_team"): "False"}): + response = test_client.patch(f"/variables/{TEST_VARIABLE_KEY}", json=body) + assert response.status_code == 422 + assert ( + response.json()["detail"][0]["msg"] + == "Value error, team_name cannot be set when multi_team mode is disabled" + ) + @pytest.mark.enable_redact def test_patch_with_update_mask_description_only(self, test_client, session): """PATCH with update_mask=['description'] should only update description, keeping value unchanged.""" @@ -686,6 +702,22 @@ def test_post_should_respond_422_when_key_too_large(self, test_client): ] } + def test_post_rejects_team_name_when_multi_team_disabled(self, test_client, testing_team): + self.create_variables() + body = { + "key": "new variable key", + "value": "new variable value", + "description": "new variable description", + "team_name": str(testing_team.name), + } + with conf_vars({("core", "multi_team"): "False"}): + response = test_client.post("/variables", json=body) + assert response.status_code == 422 + assert ( + response.json()["detail"][0]["msg"] + == "Value error, team_name cannot be set when multi_team mode is disabled" + ) + @pytest.mark.parametrize( "body", [ From 0659255daf5b3b5801d0a37e99041120afaea49b Mon Sep 17 00:00:00 2001 From: ahilashsasidharan Date: Sun, 12 Apr 2026 07:25:01 -0400 Subject: [PATCH 4/4] Change error message, Modify tests, Add bulk tests, Fix CI issues --- .../core_api/datamodels/connections.py | 4 +- .../api_fastapi/core_api/datamodels/pools.py | 8 +- .../core_api/datamodels/variables.py | 4 +- .../ui/src/queries/useEditConnection.tsx | 6 + .../airflow/ui/src/queries/useEditVariable.ts | 3 +- .../routes/public/test_connections.py | 128 ++++++++++++------ .../core_api/routes/public/test_pools.py | 105 ++++++++++---- .../core_api/routes/public/test_variables.py | 78 +++++++++-- 8 files changed, 251 insertions(+), 85 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py index aac7af789fa33..6b4154bb9b525 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py @@ -203,7 +203,9 @@ def validate_extra(cls, v: str | None) -> str | None: @model_validator(mode="after") def validate_team_name(self) -> ConnectionBody: if self.team_name is not None and not conf.getboolean("core", "multi_team"): - raise ValueError("team_name cannot be set when multi_team mode is disabled") + raise ValueError( + "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." + ) return self diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/pools.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/pools.py index 3c8853da0e2ba..55c7c3ad35e56 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/pools.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/pools.py @@ -87,7 +87,9 @@ class PoolPatchBody(StrictBaseModel): @model_validator(mode="after") def validate_team_name(self) -> PoolPatchBody: if self.team_name is not None and not conf.getboolean("core", "multi_team"): - raise ValueError("team_name cannot be set when multi_team mode is disabled") + raise ValueError( + "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." + ) return self @@ -102,5 +104,7 @@ class PoolBody(BasePool, StrictBaseModel): @model_validator(mode="after") def validate_team_name(self) -> PoolBody: if self.team_name is not None and not conf.getboolean("core", "multi_team"): - raise ValueError("team_name cannot be set when multi_team mode is disabled") + raise ValueError( + "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." + ) return self diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/variables.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/variables.py index a4f1ff334f178..001d8c70e95ad 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/variables.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/variables.py @@ -64,7 +64,9 @@ class VariableBody(StrictBaseModel): @model_validator(mode="after") def validate_team_name(self) -> VariableBody: if self.team_name is not None and not conf.getboolean("core", "multi_team"): - raise ValueError("team_name cannot be set when multi_team mode is disabled") + raise ValueError( + "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." + ) return self diff --git a/airflow-core/src/airflow/ui/src/queries/useEditConnection.tsx b/airflow-core/src/airflow/ui/src/queries/useEditConnection.tsx index 54f5cef48c210..0b7b9cefd0af4 100644 --- a/airflow-core/src/airflow/ui/src/queries/useEditConnection.tsx +++ b/airflow-core/src/airflow/ui/src/queries/useEditConnection.tsx @@ -89,6 +89,11 @@ export const useEditConnection = ( if (requestBody.schema !== initialConnection.schema) { updateMask.push("schema"); } + if (requestBody.team_name !== initialConnection.team_name) { + updateMask.push("team_name"); + } + + const teamName = requestBody.team_name === "" ? undefined : requestBody.team_name; mutate({ connectionId: initialConnection.connection_id, @@ -99,6 +104,7 @@ export const useEditConnection = ( extra: requestBody.extra === "{}" ? undefined : requestBody.extra, // eslint-disable-next-line unicorn/no-null port: requestBody.port === "" ? null : Number(requestBody.port), + team_name: teamName, }, updateMask, }); diff --git a/airflow-core/src/airflow/ui/src/queries/useEditVariable.ts b/airflow-core/src/airflow/ui/src/queries/useEditVariable.ts index 73a0fcab7dca7..ae33993058181 100644 --- a/airflow-core/src/airflow/ui/src/queries/useEditVariable.ts +++ b/airflow-core/src/airflow/ui/src/queries/useEditVariable.ts @@ -78,12 +78,13 @@ export const useEditVariable = ( const parsedDescription = editVariableRequestBody.description === "" ? undefined : editVariableRequestBody.description; + const teamName = editVariableRequestBody.team_name === "" ? undefined : editVariableRequestBody.team_name; mutate({ requestBody: { description: parsedDescription, key: editVariableRequestBody.key, - team_name: editVariableRequestBody.team_name, + team_name: teamName, value: editVariableRequestBody.value, }, updateMask, diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 137ffee0e74c1..f631f197651cc 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -287,16 +287,16 @@ def test_post_should_respond_201(self, test_client, session, body): assert len(connection) == 1 _check_last_log(session, dag_id=None, event="post_connection", logical_date=None) + @conf_vars({("core", "multi_team"): "True"}) def test_post_should_respond_201_with_team(self, test_client, session, testing_team): - with conf_vars({("core", "multi_team"): "True"}): - response = test_client.post( - "/connections", - json={ - "connection_id": TEST_CONN_ID, - "conn_type": TEST_CONN_TYPE, - "team_name": testing_team.name, - }, - ) + response = test_client.post( + "/connections", + json={ + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "team_name": testing_team.name, + }, + ) assert response.status_code == 201 assert response.json() == { "connection_id": TEST_CONN_ID, @@ -344,20 +344,20 @@ def test_post_should_respond_422_for_invalid_conn_id(self, test_client, body): ] } - def test_post_rejects_team_name_when_multi_team_disabled(self, test_client, testing_team): - with conf_vars({("core", "multi_team"): "False"}): - response = test_client.post( - "/connections", - json={ - "connection_id": TEST_CONN_ID_2, - "conn_type": TEST_CONN_TYPE_2, - "team_name": testing_team.name, - }, - ) + @conf_vars({("core", "multi_team"): "False"}) + def test_post_rejects_team_name_when_multi_team_disabled(self, test_client): + response = test_client.post( + "/connections", + json={ + "connection_id": TEST_CONN_ID_2, + "conn_type": TEST_CONN_TYPE_2, + "team_name": "test_team", + }, + ) assert response.status_code == 422 assert ( - response.json()["detail"][0]["msg"] - == "Value error, team_name cannot be set when multi_team mode is disabled" + "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." + in response.json()["detail"][0]["msg"] ) @pytest.mark.parametrize( @@ -624,14 +624,14 @@ def test_patch_should_respond_200( assert response.json() == expected_result + @conf_vars({("core", "multi_team"): "True"}) def test_patch_with_team_should_respond_200(self, test_client, testing_team, session): self.create_connection() - with conf_vars({("core", "multi_team"): "True"}): - response = test_client.patch( - f"/connections/{TEST_CONN_ID}", - json={"connection_id": TEST_CONN_ID, "conn_type": "new_type", "team_name": testing_team.name}, - ) + response = test_client.patch( + f"/connections/{TEST_CONN_ID}", + json={"connection_id": TEST_CONN_ID, "conn_type": "new_type", "team_name": testing_team.name}, + ) assert response.status_code == 200 _check_last_log(session, dag_id=None, event="patch_connection", logical_date=None) @@ -989,22 +989,21 @@ def test_patch_with_update_mask_rejects_extra_fields(self, test_client): ) assert response.status_code == 422 - def test_patch_rejects_team_name_when_multi_team_disabled(self, test_client, testing_team): + @conf_vars({("core", "multi_team"): "False"}) + def test_patch_rejects_team_name_when_multi_team_disabled(self, test_client): self.create_connection() - - with conf_vars({("core", "multi_team"): "False"}): - response = test_client.patch( - f"/connections/{TEST_CONN_ID_2}", - json={ - "connection_id": TEST_CONN_ID_2, - "conn_type": "new_type", - "team_name": testing_team.name, - }, - ) + response = test_client.patch( + f"/connections/{TEST_CONN_ID_2}", + json={ + "connection_id": TEST_CONN_ID_2, + "conn_type": "new_type", + "team_name": "test_team", + }, + ) assert response.status_code == 422 assert ( - response.json()["detail"][0]["msg"] - == "Value error, team_name cannot be set when multi_team mode is disabled" + "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." + in response.json()["detail"][0]["msg"] ) @@ -1657,6 +1656,57 @@ def test_bulk_delete_avoids_n_plus_one_queries(self, session): assert sorted(results.success) == [TEST_CONN_ID, TEST_CONN_ID_2] + @conf_vars({("core", "multi_team"): "False"}) + def test_bulk_rejects_team_name_when_multi_team_is_disabled(self, test_client): + actions = { + "actions": [ + { + "action": "create", + "entities": [ + { + "connection_id": "test_conn_id_1", + "conn_type": TEST_CONN_TYPE, + "description": "description", + }, + { + "connection_id": "test_conn_id_2", + "conn_type": TEST_CONN_TYPE_2, + "description": "description_2", + "team_name": "test_team", + }, + ], + }, + { + "action": "update", + "entities": [ + { + "connection_id": "test_conn_id_3", + "conn_type": TEST_CONN_TYPE, + "description": "updated_description", + "team_name": "test_team", + }, + { + "connection_id": "test_conn_id_4", + "conn_type": TEST_CONN_TYPE_2, + "description": "updated_description_2", + }, + ], + }, + ] + } + response = test_client.patch("/connections", json=actions) + assert response.status_code == 422 + detail = response.json()["detail"] + + assert all( + "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." + in err["msg"] + for err in detail + ), f"Unexpected errors in detail: {detail}" + + expected_error_conn_ids = {err["input"]["connection_id"] for err in detail} + assert sorted(expected_error_conn_ids) == ["test_conn_id_2", "test_conn_id_3"] + class TestPostConnectionExtraBackwardCompatibility(TestConnectionEndpoint): def test_post_should_accept_empty_string_as_extra(self, test_client, session): diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py index 1459b75ecbd51..08aa24cd76721 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py @@ -417,26 +417,27 @@ def test_patch_pool3_should_respond_200(self, test_client, session): assert response.json() == expected_response check_last_log(session, dag_id=None, event="patch_pool", logical_date=None) + @conf_vars({("core", "multi_team"): "False"}) def test_patch_pool_rejects_team_name_when_multi_team_disabled(self, test_client): self.create_pools() - with conf_vars({("core", "multi_team"): "False"}): - response = test_client.patch( - f"/pools/{POOL2_NAME}", - json={ - "name": POOL2_NAME, - "slots": POOL2_SLOT, - "include_deferred": POOL2_INCLUDE_DEFERRED, - "team_name": "test", - }, - ) + response = test_client.patch( + f"/pools/{POOL2_NAME}", + json={ + "name": POOL2_NAME, + "slots": POOL2_SLOT, + "include_deferred": POOL2_INCLUDE_DEFERRED, + "team_name": "test_team", + }, + ) assert response.status_code == 422 assert ( - response.json()["detail"][0]["msg"] - == "Value error, team_name cannot be set when multi_team mode is disabled" + "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." + in response.json()["detail"][0]["msg"] ) class TestPostPool(TestPoolsEndpoint): + @conf_vars({("core", "multi_team"): "True"}) @pytest.mark.parametrize( ("body", "expected_status_code", "expected_response"), [ @@ -503,8 +504,7 @@ def test_should_respond_200(self, test_client, session, body, expected_status_co self.create_pools() n_pools = session.scalar(select(func.count()).select_from(Pool)) - with conf_vars({("core", "multi_team"): "True"}): - response = test_client.post("/pools", json=body) + response = test_client.post("/pools", json=body) assert response.status_code == expected_status_code assert response.json() == expected_response @@ -544,20 +544,20 @@ def test_post_pool_rejects_infinity_string(self, test_client, session): ) assert response.status_code == 422 + @conf_vars({("core", "multi_team"): "False"}) def test_post_pool_rejects_team_name_when_multi_team_disabled(self, test_client): - with conf_vars({("core", "multi_team"): "False"}): - response = test_client.post( - "/pools", - json={ - "name": "bad_team_pool", - "slots": 1, - "team_name": "test", - }, - ) + response = test_client.post( + "/pools", + json={ + "name": "bad_team_pool", + "slots": 1, + "team_name": "test_team", + }, + ) assert response.status_code == 422 assert ( - response.json()["detail"][0]["msg"] - == "Value error, team_name cannot be set when multi_team mode is disabled" + "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." + in response.json()["detail"][0]["msg"] ) def test_should_respond_401(self, unauthenticated_test_client): @@ -627,6 +627,7 @@ def test_should_response_409( class TestBulkPools(TestPoolsEndpoint): + @conf_vars({("core", "multi_team"): "True"}) @pytest.mark.enable_redact @pytest.mark.parametrize( ("actions", "expected_results"), @@ -1083,8 +1084,7 @@ class TestBulkPools(TestPoolsEndpoint): def test_bulk_pools(self, test_client, actions, expected_results, session): self.create_pools() - with conf_vars({("core", "multi_team"): "True"}): - response = test_client.patch("/pools", json=actions) + response = test_client.patch("/pools", json=actions) response_data = response.json() for key, value in expected_results.items(): @@ -1145,3 +1145,54 @@ def test_should_respond_403(self, unauthorized_test_client): }, ) assert response.status_code == 403 + + @conf_vars({("core", "multi_team"): "False"}) + def test_bulk_rejects_team_name_when_multi_team_is_disabled(self, test_client): + actions = { + "actions": [ + { + "action": "create", + "entities": [ + { + "name": "pool_1", + "slots": 1, + "description": "description", + }, + { + "name": "pool_2", + "slots": 2, + "description": "description_2", + "team_name": "test_team", + }, + ], + }, + { + "action": "update", + "entities": [ + { + "name": "pool_3", + "slots": 3, + "description": "updated_description", + "team_name": "test_team", + }, + { + "name": "pool_4", + "slots": 4, + "description": "updated_description_2", + }, + ], + }, + ] + } + response = test_client.patch("/pools", json=actions) + assert response.status_code == 422 + detail = response.json()["detail"] + + assert all( + "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." + in err["msg"] + for err in detail + ), f"Unexpected errors in detail: {detail}" + + expected_error_names = {err["input"]["name"] for err in detail} + assert sorted(expected_error_names) == ["pool_2", "pool_3"] diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py index 2386bb441e628..045b78acc5636 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py @@ -521,20 +521,19 @@ def test_patch_should_respond_404(self, test_client): body = response.json() assert f"The Variable with key: `{TEST_VARIABLE_KEY}` was not found" == body["detail"] - def test_patch_rejects_team_name_when_multi_team_disabled(self, test_client, testing_team): - self.create_variables() + @conf_vars({("core", "multi_team"): "False"}) + def test_patch_rejects_team_name_when_multi_team_disabled(self, test_client): body = { "key": TEST_VARIABLE_KEY, "value": "The new value", "description": "The new description", - "team_name": str(testing_team.name), + "team_name": "test_team", } - with conf_vars({("core", "multi_team"): "False"}): - response = test_client.patch(f"/variables/{TEST_VARIABLE_KEY}", json=body) + response = test_client.patch(f"/variables/{TEST_VARIABLE_KEY}", json=body) assert response.status_code == 422 assert ( - response.json()["detail"][0]["msg"] - == "Value error, team_name cannot be set when multi_team mode is disabled" + "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." + in response.json()["detail"][0]["msg"] ) @pytest.mark.enable_redact @@ -702,20 +701,19 @@ def test_post_should_respond_422_when_key_too_large(self, test_client): ] } - def test_post_rejects_team_name_when_multi_team_disabled(self, test_client, testing_team): - self.create_variables() + @conf_vars({("core", "multi_team"): "False"}) + def test_post_rejects_team_name_when_multi_team_disabled(self, test_client): body = { "key": "new variable key", "value": "new variable value", "description": "new variable description", - "team_name": str(testing_team.name), + "team_name": "test_team", } - with conf_vars({("core", "multi_team"): "False"}): - response = test_client.post("/variables", json=body) + response = test_client.post("/variables", json=body) assert response.status_code == 422 assert ( - response.json()["detail"][0]["msg"] - == "Value error, team_name cannot be set when multi_team mode is disabled" + "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." + in response.json()["detail"][0]["msg"] ) @pytest.mark.parametrize( @@ -1398,3 +1396,55 @@ def test_bulk_variables_should_respond_403(self, unauthorized_test_client): }, ) assert response.status_code == 403 + + @conf_vars({("core", "multi_team"): "False"}) + def test_bulk_rejects_team_name_when_multi_team_is_disabled(self, test_client): + actions = { + "actions": [ + { + "action": "create", + "entities": [ + { + "key": "var_1", + "value": "value_1", + "description": "description", + }, + { + "key": "var_2", + "value": "value_2", + "description": "description_2", + "team_name": "test_team", + }, + ], + }, + { + "action": "update", + "entities": [ + { + "key": "var_3", + "value": "value_3", + "description": "updated_description", + "team_name": "test_team", + }, + { + "key": "var_4", + "value": "value_4", + "description": "updated_description_2", + }, + ], + }, + ] + } + response = test_client.patch("/variables", json=actions) + + assert response.status_code == 422 + detail = response.json()["detail"] + + assert all( + "team_name cannot be set when multi_team mode is disabled. Please contact your administrator." + in err["msg"] + for err in detail + ), f"Unexpected errors in detail: {detail}" + + expected_error_keys = {err["input"]["key"] for err in detail} + assert sorted(expected_error_keys) == ["var_2", "var_3"]