This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit ff4ae16826174a2bbab6d188f663aa078a76987f
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Wed Apr 15 10:59:28 2026 +0530
[v3-2-test] Feat/check multi team enabled when team name provided api
(#63994) (#65226)
* Add API check to ensure multi team is enabled when team_name is provided
* remove unnecessary arguments in added tests
* add variable tests and add slight change to other tests to align with
variables test file
* Change error message, Modify tests, Add bulk tests, Fix CI issues
(cherry picked from commit 62711897b3be6c097fb6ccd2f3d4362a5beacd9d)
Co-authored-by: ahilashsasidharan
<[email protected]>
---
.../api_fastapi/core_api/datamodels/connections.py | 11 ++-
.../api_fastapi/core_api/datamodels/pools.py | 19 ++++-
.../api_fastapi/core_api/datamodels/variables.py | 9 +++
.../airflow/ui/src/queries/useEditConnection.tsx | 6 ++
.../src/airflow/ui/src/queries/useEditVariable.ts | 3 +-
.../core_api/routes/public/test_connections.py | 93 +++++++++++++++++++++-
.../core_api/routes/public/test_pools.py | 93 +++++++++++++++++++++-
.../core_api/routes/public/test_variables.py | 82 +++++++++++++++++++
8 files changed, 311 insertions(+), 5 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 0b69c47cfdf..680376cddac 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 @@ import json
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,13 @@ class ConnectionBody(StrictBaseModel):
)
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.
Please contact your administrator."
+ )
+ 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 a661308d443..55c7c3ad35e 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 __future__ import annotations
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,14 @@ 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.
Please contact your administrator."
+ )
+ return self
+
class PoolBody(BasePool, StrictBaseModel):
"""Pool serializer for post bodies."""
@@ -91,3 +100,11 @@ 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.
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 9dc7969b69b..001d8c70e95 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 pydantic import Field, JsonValue, model_validator
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,14 @@ 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.
Please contact your administrator."
+ )
+ return self
+
VariableBodyPartial = make_partial_model(VariableBody)
diff --git a/airflow-core/src/airflow/ui/src/queries/useEditConnection.tsx
b/airflow-core/src/airflow/ui/src/queries/useEditConnection.tsx
index 54f5cef48c2..0b7b9cefd0a 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 73a0fcab7dc..ae339930581 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 4ca21b0ff49..a37ae170676 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 airflow.utils.session import NEW_SESSION, provide_session
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
@@ -299,10 +300,15 @@ class TestPostConnection(TestConnectionEndpoint):
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):
response = test_client.post(
"/connections",
- json={"connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE,
"team_name": testing_team.name},
+ json={
+ "connection_id": TEST_CONN_ID,
+ "conn_type": TEST_CONN_TYPE,
+ "team_name": testing_team.name,
+ },
)
assert response.status_code == 201
assert response.json() == {
@@ -351,6 +357,22 @@ class TestPostConnection(TestConnectionEndpoint):
]
}
+ @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 (
+ "team_name cannot be set when multi_team mode is disabled. Please
contact your administrator."
+ in response.json()["detail"][0]["msg"]
+ )
+
@pytest.mark.parametrize(
"body",
[
@@ -615,6 +637,7 @@ class TestPatchConnection(TestConnectionEndpoint):
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()
@@ -979,6 +1002,23 @@ class TestPatchConnection(TestConnectionEndpoint):
)
assert response.status_code == 422
+ @conf_vars({("core", "multi_team"): "False"})
+ def test_patch_rejects_team_name_when_multi_team_disabled(self,
test_client):
+ self.create_connection()
+ 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 (
+ "team_name cannot be set when multi_team mode is disabled. Please
contact your administrator."
+ in response.json()["detail"][0]["msg"]
+ )
+
class TestConnection(TestConnectionEndpoint):
def setup_method(self):
@@ -1629,6 +1669,57 @@ class TestBulkConnections(TestConnectionEndpoint):
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 e2ef8720735..08aa24cd767 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.pool import Pool
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
@@ -416,8 +417,27 @@ class TestPatchPool(TestPoolsEndpoint):
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()
+ 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 (
+ "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"),
[
@@ -483,9 +503,10 @@ 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
+ 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 +544,22 @@ class TestPostPool(TestPoolsEndpoint):
)
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):
+ response = test_client.post(
+ "/pools",
+ json={
+ "name": "bad_team_pool",
+ "slots": 1,
+ "team_name": "test_team",
+ },
+ )
+ assert response.status_code == 422
+ assert (
+ "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):
response = unauthenticated_test_client.post("/pools", json={})
assert response.status_code == 401
@@ -590,6 +627,7 @@ class TestPostPool(TestPoolsEndpoint):
class TestBulkPools(TestPoolsEndpoint):
+ @conf_vars({("core", "multi_team"): "True"})
@pytest.mark.enable_redact
@pytest.mark.parametrize(
("actions", "expected_results"),
@@ -1045,7 +1083,9 @@ class TestBulkPools(TestPoolsEndpoint):
)
def test_bulk_pools(self, test_client, actions, expected_results, session):
self.create_pools()
+
response = test_client.patch("/pools", json=actions)
+
response_data = response.json()
for key, value in expected_results.items():
assert response_data[key] == value
@@ -1105,3 +1145,54 @@ class TestBulkPools(TestPoolsEndpoint):
},
)
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 50871879675..045b78acc56 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,21 @@ class TestPatchVariable(TestVariableEndpoint):
body = response.json()
assert f"The Variable with key: `{TEST_VARIABLE_KEY}` was not found"
== body["detail"]
+ @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": "test_team",
+ }
+ response = test_client.patch(f"/variables/{TEST_VARIABLE_KEY}",
json=body)
+ assert response.status_code == 422
+ assert (
+ "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
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 +701,21 @@ class TestPostVariable(TestVariableEndpoint):
]
}
+ @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": "test_team",
+ }
+ response = test_client.post("/variables", json=body)
+ assert response.status_code == 422
+ assert (
+ "team_name cannot be set when multi_team mode is disabled. Please
contact your administrator."
+ in response.json()["detail"][0]["msg"]
+ )
+
@pytest.mark.parametrize(
"body",
[
@@ -1366,3 +1396,55 @@ class TestBulkVariables(TestVariableEndpoint):
},
)
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"]
