From 53e6d8ead3a7e777867314198c283e601777b894 Mon Sep 17 00:00:00 2001 From: Bugra Ozturk Date: Sat, 26 Oct 2024 12:59:19 +0200 Subject: [PATCH 1/8] Migrate Create a Connection to FastAPI --- .../endpoints/connection_endpoint.py | 1 + .../core_api/routes/public/connections.py | 23 +++++ .../core_api/serializers/connections.py | 17 ++++ .../ui/openapi-gen/requests/schemas.gen.ts | 83 +++++++++++++++++++ .../routes/public/test_connections.py | 61 ++++++++++++++ 5 files changed, 185 insertions(+) diff --git a/airflow/api_connexion/endpoints/connection_endpoint.py b/airflow/api_connexion/endpoints/connection_endpoint.py index 37c91c44eb698..c0c2fcbf4610a 100644 --- a/airflow/api_connexion/endpoints/connection_endpoint.py +++ b/airflow/api_connexion/endpoints/connection_endpoint.py @@ -151,6 +151,7 @@ def patch_connection( return connection_schema.dump(connection) +@mark_fastapi_migration_done @security.requires_access_connection("POST") @provide_session @action_logging( diff --git a/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow/api_fastapi/core_api/routes/public/connections.py index 8d9f9ddb8ebfc..0f11d4716d0e6 100644 --- a/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow/api_fastapi/core_api/routes/public/connections.py @@ -26,10 +26,12 @@ from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc from airflow.api_fastapi.core_api.serializers.connections import ( + ConnectionBody, ConnectionCollectionResponse, ConnectionResponse, ) from airflow.models import Connection +from airflow.utils import helpers connections_router = AirflowRouter(tags=["Connection"], prefix="/connections") @@ -114,3 +116,24 @@ async def get_connections( ], total_entries=total_entries, ) + + +@connections_router.post("/", status_code=201, responses=create_openapi_http_exception_doc([401, 403, 409])) +async def post_connection( + post_body: ConnectionBody, + session: Annotated[Session, Depends(get_session)], +) -> ConnectionResponse: + """Create connection entry.""" + try: + helpers.validate_key(post_body.connection_id, max_length=200) + except Exception as e: + raise HTTPException(400, f"{e}") + + connection = session.scalar(select(Connection).filter_by(conn_id=post_body.connection_id)) + if connection is not None: + raise HTTPException(409, f"Connection with connection_id: `{post_body.connection_id}` already exists") + + connection = Connection(**post_body.model_dump(by_alias=True, exclude_none=True)) + session.add(connection) + + return ConnectionResponse.model_validate(connection, from_attributes=True) diff --git a/airflow/api_fastapi/core_api/serializers/connections.py b/airflow/api_fastapi/core_api/serializers/connections.py index 1cc069cac0cb3..81a436968764c 100644 --- a/airflow/api_fastapi/core_api/serializers/connections.py +++ b/airflow/api_fastapi/core_api/serializers/connections.py @@ -24,9 +24,12 @@ from airflow.utils.log.secrets_masker import redact +# Response Models class ConnectionResponse(BaseModel): """Connection serializer for responses.""" + """Connection serializer for responses.""" + connection_id: str = Field(serialization_alias="connection_id", validation_alias="conn_id") conn_type: str description: str | None @@ -55,3 +58,17 @@ class ConnectionCollectionResponse(BaseModel): connections: list[ConnectionResponse] total_entries: int + + +# Request Models +class ConnectionBody(BaseModel): + """Connection Serializer for requests body.""" + + connection_id: str = Field(serialization_alias="conn_id") + conn_type: str + description: str | None = Field(default=None) + host: str | None = Field(default=None) + login: str | None = Field(default=None) + schema_: str | None = Field(None, alias="schema") + port: int | None = Field(default=None) + extra: str | None = Field(default=None) diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index 517743af17f6c..35a6505205c16 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -151,6 +151,89 @@ export const $BaseInfoSchema = { description: "Base status field for metadatabase and scheduler.", } as const; +export const $ConnectionBody = { + properties: { + connection_id: { + type: "string", + title: "Connection Id", + }, + conn_type: { + type: "string", + title: "Conn Type", + }, + description: { + anyOf: [ + { + type: "string", + }, + { + type: "null", + }, + ], + title: "Description", + }, + host: { + anyOf: [ + { + type: "string", + }, + { + type: "null", + }, + ], + title: "Host", + }, + login: { + anyOf: [ + { + type: "string", + }, + { + type: "null", + }, + ], + title: "Login", + }, + schema: { + anyOf: [ + { + type: "string", + }, + { + type: "null", + }, + ], + title: "Schema", + }, + port: { + anyOf: [ + { + type: "integer", + }, + { + type: "null", + }, + ], + title: "Port", + }, + extra: { + anyOf: [ + { + type: "string", + }, + { + type: "null", + }, + ], + title: "Extra", + }, + }, + type: "object", + required: ["connection_id", "conn_type"], + title: "ConnectionBody", + description: "Connection Serializer for requests body.", +} as const; + export const $ConnectionCollectionResponse = { properties: { connections: { diff --git a/tests/api_fastapi/core_api/routes/public/test_connections.py b/tests/api_fastapi/core_api/routes/public/test_connections.py index ee9c80219eb18..9f25739b6df40 100644 --- a/tests/api_fastapi/core_api/routes/public/test_connections.py +++ b/tests/api_fastapi/core_api/routes/public/test_connections.py @@ -169,3 +169,64 @@ def test_should_respond_200( body = response.json() assert body["total_entries"] == expected_total_entries assert [connection["connection_id"] for connection in body["connections"]] == expected_ids + + +class TestPostConnection(TestConnectionEndpoint): + @pytest.mark.parametrize( + "body", + [ + {"connection_id": "test-connection-id", "conn_type": "test_type"}, + {"connection_id": "test-connection-id", "conn_type": "test_type", "extra": None}, + {"connection_id": "test-connection-id", "conn_type": "test_type", "extra": "{}"}, + {"connection_id": "test-connection-id", "conn_type": "test_type", "extra": '{"key": "value"}'}, + { + "connection_id": "test-connection-id", + "conn_type": "test_type", + "description": "test_description", + "host": "test_host", + "login": "test_login", + "schema": "test_schema", + "port": 8080, + "extra": '{"key": "value"}', + }, + ], + ) + def test_post_should_respond_200(self, test_client, session, body): + response = test_client.post("/public/connections", json=body) + assert response.status_code == 201 + connection = session.query(Connection).all() + assert len(connection) == 1 + + @pytest.mark.parametrize( + "body", + [ + {"connection_id": "****", "conn_type": "test_type"}, + {"connection_id": "test()", "conn_type": "test_type"}, + {"connection_id": "this_^$#is_invalid", "conn_type": "test_type"}, + {"connection_id": "iam_not@#$_connection_id", "conn_type": "test_type"}, + ], + ) + def test_post_should_respond_400_for_invalid_conn_id(self, test_client, body): + response = test_client.post("/public/connections", json=body) + assert response.status_code == 400 + connection_id = body["connection_id"] + assert response.json() == { + "detail": f"The key '{connection_id}' has to be made of " + "alphanumeric characters, dashes, dots and underscores exclusively", + } + + @pytest.mark.parametrize( + "body", + [ + {"connection_id": "test-connection-id", "conn_type": "test_type"}, + ], + ) + def test_post_should_respond_already_exist(self, test_client, body): + response = test_client.post("/public/connections", json=body) + assert response.status_code == 201 + # Another request + response = test_client.post("/public/connections", json=body) + assert response.status_code == 409 + assert response.json() == { + "detail": "Connection with connection_id: `test-connection-id` already exists", + } From 439c5c172cac67474e950853013761c7e1008279 Mon Sep 17 00:00:00 2001 From: Bugra Ozturk Date: Sun, 27 Oct 2024 02:53:12 +0200 Subject: [PATCH 2/8] Remove additional duplicate comment --- airflow/api_fastapi/core_api/serializers/connections.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/airflow/api_fastapi/core_api/serializers/connections.py b/airflow/api_fastapi/core_api/serializers/connections.py index 81a436968764c..3794d2490a7ad 100644 --- a/airflow/api_fastapi/core_api/serializers/connections.py +++ b/airflow/api_fastapi/core_api/serializers/connections.py @@ -28,8 +28,6 @@ class ConnectionResponse(BaseModel): """Connection serializer for responses.""" - """Connection serializer for responses.""" - connection_id: str = Field(serialization_alias="connection_id", validation_alias="conn_id") conn_type: str description: str | None From 85018368473416322b16207ff0b34009ab73c6d1 Mon Sep 17 00:00:00 2001 From: Bugra Ozturk Date: Wed, 30 Oct 2024 01:08:09 +0100 Subject: [PATCH 3/8] Include password in connection and move dashboard.py to serializers/ui/ --- .../core_api/routes/public/connections.py | 2 +- .../core_api/routes/ui/dashboard.py | 2 +- .../core_api/serializers/connections.py | 1 + .../serializers/{ => ui}/dashboard.py | 0 .../ui/openapi-gen/requests/schemas.gen.ts | 11 +++++++ .../routes/public/test_connections.py | 31 ++++++++++++++++--- 6 files changed, 41 insertions(+), 6 deletions(-) rename airflow/api_fastapi/core_api/serializers/{ => ui}/dashboard.py (100%) diff --git a/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow/api_fastapi/core_api/routes/public/connections.py index 0f11d4716d0e6..a31c97c91488f 100644 --- a/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow/api_fastapi/core_api/routes/public/connections.py @@ -133,7 +133,7 @@ async def post_connection( if connection is not None: raise HTTPException(409, f"Connection with connection_id: `{post_body.connection_id}` already exists") - connection = Connection(**post_body.model_dump(by_alias=True, exclude_none=True)) + connection = Connection(**post_body.model_dump(by_alias=True)) session.add(connection) return ConnectionResponse.model_validate(connection, from_attributes=True) diff --git a/airflow/api_fastapi/core_api/routes/ui/dashboard.py b/airflow/api_fastapi/core_api/routes/ui/dashboard.py index e101ca78be7d9..0eeea4d0dc15d 100644 --- a/airflow/api_fastapi/core_api/routes/ui/dashboard.py +++ b/airflow/api_fastapi/core_api/routes/ui/dashboard.py @@ -25,7 +25,7 @@ from airflow.api_fastapi.common.parameters import DateTimeQuery from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc -from airflow.api_fastapi.core_api.serializers.dashboard import HistoricalMetricDataResponse +from airflow.api_fastapi.core_api.serializers.ui.dashboard import HistoricalMetricDataResponse from airflow.models.dagrun import DagRun, DagRunType from airflow.models.taskinstance import TaskInstance from airflow.utils.state import DagRunState, TaskInstanceState diff --git a/airflow/api_fastapi/core_api/serializers/connections.py b/airflow/api_fastapi/core_api/serializers/connections.py index 3794d2490a7ad..79cd241b1ae2d 100644 --- a/airflow/api_fastapi/core_api/serializers/connections.py +++ b/airflow/api_fastapi/core_api/serializers/connections.py @@ -69,4 +69,5 @@ class ConnectionBody(BaseModel): login: str | None = Field(default=None) schema_: str | None = Field(None, alias="schema") port: int | None = Field(default=None) + password: str | None = Field(default=None) extra: str | None = Field(default=None) diff --git a/airflow/api_fastapi/core_api/serializers/dashboard.py b/airflow/api_fastapi/core_api/serializers/ui/dashboard.py similarity index 100% rename from airflow/api_fastapi/core_api/serializers/dashboard.py rename to airflow/api_fastapi/core_api/serializers/ui/dashboard.py diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index 35a6505205c16..8a72a45383c40 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -216,6 +216,17 @@ export const $ConnectionBody = { ], title: "Port", }, + password: { + anyOf: [ + { + type: "string", + }, + { + type: "null", + }, + ], + title: "Password", + }, extra: { anyOf: [ { diff --git a/tests/api_fastapi/core_api/routes/public/test_connections.py b/tests/api_fastapi/core_api/routes/public/test_connections.py index 9f25739b6df40..652dcc0ce9c87 100644 --- a/tests/api_fastapi/core_api/routes/public/test_connections.py +++ b/tests/api_fastapi/core_api/routes/public/test_connections.py @@ -30,6 +30,7 @@ TEST_CONN_DESCRIPTION = "some_description_a" TEST_CONN_HOST = "some_host_a" TEST_CONN_PORT = 8080 +TEST_PASS = "test-password" TEST_CONN_ID_2 = "test_connection_id_2" @@ -37,6 +38,7 @@ TEST_CONN_DESCRIPTION_2 = "some_description_b" TEST_CONN_HOST_2 = "some_host_b" TEST_CONN_PORT_2 = 8081 +TEST_PASS_2 = "test-password_2" @provide_session @@ -47,6 +49,7 @@ def _create_connection(session) -> None: description=TEST_CONN_DESCRIPTION, host=TEST_CONN_HOST, port=TEST_CONN_PORT, + password=TEST_PASS, ) session.add(connection_model) @@ -60,6 +63,7 @@ def _create_connections(session) -> None: description=TEST_CONN_DESCRIPTION_2, host=TEST_CONN_HOST_2, port=TEST_CONN_PORT_2, + password=TEST_PASS_2, ) session.add(connection_model_2) @@ -192,7 +196,7 @@ class TestPostConnection(TestConnectionEndpoint): ], ) def test_post_should_respond_200(self, test_client, session, body): - response = test_client.post("/public/connections", json=body) + response = test_client.post("/public/connections/", json=body) assert response.status_code == 201 connection = session.query(Connection).all() assert len(connection) == 1 @@ -207,7 +211,7 @@ def test_post_should_respond_200(self, test_client, session, body): ], ) def test_post_should_respond_400_for_invalid_conn_id(self, test_client, body): - response = test_client.post("/public/connections", json=body) + response = test_client.post("/public/connections/", json=body) assert response.status_code == 400 connection_id = body["connection_id"] assert response.json() == { @@ -222,11 +226,30 @@ def test_post_should_respond_400_for_invalid_conn_id(self, test_client, body): ], ) def test_post_should_respond_already_exist(self, test_client, body): - response = test_client.post("/public/connections", json=body) + response = test_client.post("/public/connections/", json=body) assert response.status_code == 201 # Another request - response = test_client.post("/public/connections", json=body) + response = test_client.post("/public/connections/", json=body) assert response.status_code == 409 assert response.json() == { "detail": "Connection with connection_id: `test-connection-id` already exists", } + + @pytest.mark.parametrize( + "body", + [ + {"connection_id": "test-connection-id", "conn_type": "test_type", "password": "test-password"}, + {"connection_id": "test-connection-id", "conn_type": "test_type", "password": "?>@#+!_%()#"}, + { + "connection_id": "test-connection-id", + "conn_type": "test_type", + "password": "A!rF|0wi$aw3s0m3", + "extra": '{"password": "test-password"}', + }, + ], + ) + def test_post_should_response_201_redact_password(self, test_client, body): + response = test_client.post("/public/connections/", json=body) + assert response.status_code == 201 + connection = response.json() + assert connection["password"] == "***" From b33d775044b1b7a2bc2d9975378dba9cf92e12af Mon Sep 17 00:00:00 2001 From: Bugra Ozturk Date: Wed, 30 Oct 2024 01:36:17 +0100 Subject: [PATCH 4/8] Fix test for password --- .../core_api/routes/public/test_connections.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/api_fastapi/core_api/routes/public/test_connections.py b/tests/api_fastapi/core_api/routes/public/test_connections.py index 652dcc0ce9c87..5d7bf54ade466 100644 --- a/tests/api_fastapi/core_api/routes/public/test_connections.py +++ b/tests/api_fastapi/core_api/routes/public/test_connections.py @@ -30,7 +30,6 @@ TEST_CONN_DESCRIPTION = "some_description_a" TEST_CONN_HOST = "some_host_a" TEST_CONN_PORT = 8080 -TEST_PASS = "test-password" TEST_CONN_ID_2 = "test_connection_id_2" @@ -38,7 +37,6 @@ TEST_CONN_DESCRIPTION_2 = "some_description_b" TEST_CONN_HOST_2 = "some_host_b" TEST_CONN_PORT_2 = 8081 -TEST_PASS_2 = "test-password_2" @provide_session @@ -49,7 +47,6 @@ def _create_connection(session) -> None: description=TEST_CONN_DESCRIPTION, host=TEST_CONN_HOST, port=TEST_CONN_PORT, - password=TEST_PASS, ) session.add(connection_model) @@ -63,7 +60,6 @@ def _create_connections(session) -> None: description=TEST_CONN_DESCRIPTION_2, host=TEST_CONN_HOST_2, port=TEST_CONN_PORT_2, - password=TEST_PASS_2, ) session.add(connection_model_2) @@ -248,8 +244,8 @@ def test_post_should_respond_already_exist(self, test_client, body): }, ], ) - def test_post_should_response_201_redact_password(self, test_client, body): + def test_post_should_response_201_redacted_password(self, test_client, body): response = test_client.post("/public/connections/", json=body) assert response.status_code == 201 connection = response.json() - assert connection["password"] == "***" + assert "password" not in connection From 04bd842c58af967327eef588c6c5843677ce53d5 Mon Sep 17 00:00:00 2001 From: Bugra Ozturk Date: Thu, 31 Oct 2024 21:48:43 +0100 Subject: [PATCH 5/8] Include password field to response and redact it, run pre-commit after rebase --- .../core_api/serializers/connections.py | 11 ++- .../ui/openapi-gen/requests/schemas.gen.ts | 16 +++- .../routes/public/test_connections.py | 85 ++++++++++++++----- 3 files changed, 86 insertions(+), 26 deletions(-) diff --git a/airflow/api_fastapi/core_api/serializers/connections.py b/airflow/api_fastapi/core_api/serializers/connections.py index 79cd241b1ae2d..dc2612fb15a3f 100644 --- a/airflow/api_fastapi/core_api/serializers/connections.py +++ b/airflow/api_fastapi/core_api/serializers/connections.py @@ -19,7 +19,8 @@ import json -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, Field, field_validator, model_validator +from typing_extensions import Self from airflow.utils.log.secrets_masker import redact @@ -35,8 +36,16 @@ class ConnectionResponse(BaseModel): login: str | None schema_: str | None = Field(alias="schema") port: int | None + password: str | None extra: str | None + @model_validator(mode="after") + def redact_password(self) -> Self: + if self.password is None: + return self + self.password = redact(self.password) + return self + @field_validator("extra", mode="before") @classmethod def redact_extra(cls, v: str | None) -> str | None: diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index 8a72a45383c40..c1dc8cd34576b 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -330,6 +330,17 @@ export const $ConnectionResponse = { ], title: "Port", }, + password: { + anyOf: [ + { + type: "string", + }, + { + type: "null", + }, + ], + title: "Password", + }, extra: { anyOf: [ { @@ -351,6 +362,7 @@ export const $ConnectionResponse = { "login", "schema", "port", + "password", "extra", ], title: "ConnectionResponse", @@ -1825,7 +1837,7 @@ export const $HistoricalMetricDataResponse = { $ref: "#/components/schemas/DAGRunStates", }, task_instance_states: { - $ref: "#/components/schemas/airflow__api_fastapi__core_api__serializers__dashboard__TaskInstanceState", + $ref: "#/components/schemas/airflow__api_fastapi__core_api__serializers__ui__dashboard__TaskInstanceState", }, }, type: "object", @@ -2860,7 +2872,7 @@ export const $VersionInfo = { description: "Version information serializer for responses.", } as const; -export const $airflow__api_fastapi__core_api__serializers__dashboard__TaskInstanceState = +export const $airflow__api_fastapi__core_api__serializers__ui__dashboard__TaskInstanceState = { properties: { no_status: { diff --git a/tests/api_fastapi/core_api/routes/public/test_connections.py b/tests/api_fastapi/core_api/routes/public/test_connections.py index 5d7bf54ade466..f70743dfd4fb8 100644 --- a/tests/api_fastapi/core_api/routes/public/test_connections.py +++ b/tests/api_fastapi/core_api/routes/public/test_connections.py @@ -175,13 +175,13 @@ class TestPostConnection(TestConnectionEndpoint): @pytest.mark.parametrize( "body", [ - {"connection_id": "test-connection-id", "conn_type": "test_type"}, - {"connection_id": "test-connection-id", "conn_type": "test_type", "extra": None}, - {"connection_id": "test-connection-id", "conn_type": "test_type", "extra": "{}"}, - {"connection_id": "test-connection-id", "conn_type": "test_type", "extra": '{"key": "value"}'}, + {"connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE}, + {"connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, "extra": None}, + {"connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, "extra": "{}"}, + {"connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, "extra": '{"key": "value"}'}, { - "connection_id": "test-connection-id", - "conn_type": "test_type", + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, "description": "test_description", "host": "test_host", "login": "test_login", @@ -200,10 +200,10 @@ def test_post_should_respond_200(self, test_client, session, body): @pytest.mark.parametrize( "body", [ - {"connection_id": "****", "conn_type": "test_type"}, - {"connection_id": "test()", "conn_type": "test_type"}, - {"connection_id": "this_^$#is_invalid", "conn_type": "test_type"}, - {"connection_id": "iam_not@#$_connection_id", "conn_type": "test_type"}, + {"connection_id": "****", "conn_type": TEST_CONN_TYPE}, + {"connection_id": "test()", "conn_type": TEST_CONN_TYPE}, + {"connection_id": "this_^$#is_invalid", "conn_type": TEST_CONN_TYPE}, + {"connection_id": "iam_not@#$_connection_id", "conn_type": TEST_CONN_TYPE}, ], ) def test_post_should_respond_400_for_invalid_conn_id(self, test_client, body): @@ -218,7 +218,7 @@ def test_post_should_respond_400_for_invalid_conn_id(self, test_client, body): @pytest.mark.parametrize( "body", [ - {"connection_id": "test-connection-id", "conn_type": "test_type"}, + {"connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE}, ], ) def test_post_should_respond_already_exist(self, test_client, body): @@ -231,21 +231,60 @@ def test_post_should_respond_already_exist(self, test_client, body): "detail": "Connection with connection_id: `test-connection-id` already exists", } + @pytest.mark.enable_redact @pytest.mark.parametrize( - "body", + "body, expected_response", [ - {"connection_id": "test-connection-id", "conn_type": "test_type", "password": "test-password"}, - {"connection_id": "test-connection-id", "conn_type": "test_type", "password": "?>@#+!_%()#"}, - { - "connection_id": "test-connection-id", - "conn_type": "test_type", - "password": "A!rF|0wi$aw3s0m3", - "extra": '{"password": "test-password"}', - }, + ( + {"connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, "password": "test-password"}, + { + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "description": None, + "extra": None, + "host": None, + "login": None, + "password": "***", + "port": None, + "schema": None, + }, + ), + ( + {"connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, "password": "?>@#+!_%()#"}, + { + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "description": None, + "extra": None, + "host": None, + "login": None, + "password": "***", + "port": None, + "schema": None, + }, + ), + ( + { + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "password": "A!rF|0wi$aw3s0m3", + "extra": '{"password": "test-password"}', + }, + { + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "description": None, + "extra": '{"password": "***"}', + "host": None, + "login": None, + "password": "***", + "port": None, + "schema": None, + }, + ), ], ) - def test_post_should_response_201_redacted_password(self, test_client, body): + def test_post_should_response_201_redacted_password(self, test_client, body, expected_response): response = test_client.post("/public/connections/", json=body) assert response.status_code == 201 - connection = response.json() - assert "password" not in connection + assert response.json() == expected_response From 94210799648b85ac5ab904a60c2d298a205a82b7 Mon Sep 17 00:00:00 2001 From: Bugra Ozturk Date: Thu, 31 Oct 2024 22:00:26 +0100 Subject: [PATCH 6/8] Convert redact to field_validator and fix tests --- .../core_api/serializers/connections.py | 15 +++++++-------- .../core_api/routes/public/test_connections.py | 2 +- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/airflow/api_fastapi/core_api/serializers/connections.py b/airflow/api_fastapi/core_api/serializers/connections.py index dc2612fb15a3f..0e141915437d5 100644 --- a/airflow/api_fastapi/core_api/serializers/connections.py +++ b/airflow/api_fastapi/core_api/serializers/connections.py @@ -19,8 +19,7 @@ import json -from pydantic import BaseModel, Field, field_validator, model_validator -from typing_extensions import Self +from pydantic import BaseModel, Field, field_validator from airflow.utils.log.secrets_masker import redact @@ -39,12 +38,12 @@ class ConnectionResponse(BaseModel): password: str | None extra: str | None - @model_validator(mode="after") - def redact_password(self) -> Self: - if self.password is None: - return self - self.password = redact(self.password) - return self + @field_validator("password", mode="after") + @classmethod + def redact_password(cls, v: str | None) -> str | None: + if v is None: + return None + return redact(v) @field_validator("extra", mode="before") @classmethod diff --git a/tests/api_fastapi/core_api/routes/public/test_connections.py b/tests/api_fastapi/core_api/routes/public/test_connections.py index f70743dfd4fb8..1dc3cf9d2cd4d 100644 --- a/tests/api_fastapi/core_api/routes/public/test_connections.py +++ b/tests/api_fastapi/core_api/routes/public/test_connections.py @@ -228,7 +228,7 @@ def test_post_should_respond_already_exist(self, test_client, body): response = test_client.post("/public/connections/", json=body) assert response.status_code == 409 assert response.json() == { - "detail": "Connection with connection_id: `test-connection-id` already exists", + "detail": f"Connection with connection_id: `{TEST_CONN_ID}` already exists", } @pytest.mark.enable_redact From 4c6098313b0e33a31f80f46f177482272bddd823 Mon Sep 17 00:00:00 2001 From: Bugra Ozturk Date: Mon, 4 Nov 2024 19:34:09 +0100 Subject: [PATCH 7/8] Pass field name into redact --- airflow/api_fastapi/core_api/serializers/connections.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/airflow/api_fastapi/core_api/serializers/connections.py b/airflow/api_fastapi/core_api/serializers/connections.py index 0e141915437d5..c5956b6ec517d 100644 --- a/airflow/api_fastapi/core_api/serializers/connections.py +++ b/airflow/api_fastapi/core_api/serializers/connections.py @@ -20,6 +20,7 @@ import json from pydantic import BaseModel, Field, field_validator +from pydantic_core.core_schema import ValidationInfo from airflow.utils.log.secrets_masker import redact @@ -40,10 +41,10 @@ class ConnectionResponse(BaseModel): @field_validator("password", mode="after") @classmethod - def redact_password(cls, v: str | None) -> str | None: + def redact_password(cls, v: str | None, field_info: ValidationInfo) -> str | None: if v is None: return None - return redact(v) + return redact(v, field_info.field_name) @field_validator("extra", mode="before") @classmethod From 5d50c7e539a154969e2d33ce3f934092d0c79529 Mon Sep 17 00:00:00 2001 From: Bugra Ozturk Date: Mon, 4 Nov 2024 19:37:48 +0100 Subject: [PATCH 8/8] run pre-commit after rebase --- .../core_api/openapi/v1-generated.yaml | 102 +++++++++++++++++- airflow/ui/openapi-gen/queries/common.ts | 3 + airflow/ui/openapi-gen/queries/queries.ts | 40 +++++++ .../ui/openapi-gen/requests/services.gen.ts | 27 +++++ airflow/ui/openapi-gen/requests/types.gen.ts | 51 ++++++++- 5 files changed, 219 insertions(+), 4 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 28e38884803a7..06a041b55daab 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -1104,6 +1104,49 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' + post: + tags: + - Connection + summary: Post Connection + description: Create connection entry. + operationId: post_connection + requestBody: + required: true + content: + application/json: + schema: + $ref: '#/components/schemas/ConnectionBody' + responses: + '201': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/ConnectionResponse' + '401': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Unauthorized + '403': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Forbidden + '409': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Conflict + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' /public/dags/{dag_id}/dagRuns/{dag_run_id}: get: tags: @@ -2515,6 +2558,55 @@ components: - status title: BaseInfoSchema description: Base status field for metadatabase and scheduler. + ConnectionBody: + properties: + connection_id: + type: string + title: Connection Id + conn_type: + type: string + title: Conn Type + description: + anyOf: + - type: string + - type: 'null' + title: Description + host: + anyOf: + - type: string + - type: 'null' + title: Host + login: + anyOf: + - type: string + - type: 'null' + title: Login + schema: + anyOf: + - type: string + - type: 'null' + title: Schema + port: + anyOf: + - type: integer + - type: 'null' + title: Port + password: + anyOf: + - type: string + - type: 'null' + title: Password + extra: + anyOf: + - type: string + - type: 'null' + title: Extra + type: object + required: + - connection_id + - conn_type + title: ConnectionBody + description: Connection Serializer for requests body. ConnectionCollectionResponse: properties: connections: @@ -2564,6 +2656,11 @@ components: - type: integer - type: 'null' title: Port + password: + anyOf: + - type: string + - type: 'null' + title: Password extra: anyOf: - type: string @@ -2578,6 +2675,7 @@ components: - login - schema - port + - password - extra title: ConnectionResponse description: Connection serializer for responses. @@ -3545,7 +3643,7 @@ components: dag_run_states: $ref: '#/components/schemas/DAGRunStates' task_instance_states: - $ref: '#/components/schemas/airflow__api_fastapi__core_api__serializers__dashboard__TaskInstanceState' + $ref: '#/components/schemas/airflow__api_fastapi__core_api__serializers__ui__dashboard__TaskInstanceState' type: object required: - dag_run_types @@ -4224,7 +4322,7 @@ components: - git_version title: VersionInfo description: Version information serializer for responses. - airflow__api_fastapi__core_api__serializers__dashboard__TaskInstanceState: + airflow__api_fastapi__core_api__serializers__ui__dashboard__TaskInstanceState: properties: no_status: type: integer diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 36ea524e0138f..2ed842201c2e9 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -682,6 +682,9 @@ export const UseVersionServiceGetVersionKeyFn = (queryKey?: Array) => [ export type BackfillServiceCreateBackfillMutationResult = Awaited< ReturnType >; +export type ConnectionServicePostConnectionMutationResult = Awaited< + ReturnType +>; export type PoolServicePostPoolMutationResult = Awaited< ReturnType >; diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index f4b7c41195f77..583f14f7711d9 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -28,6 +28,7 @@ import { } from "../requests/services.gen"; import { BackfillPostBody, + ConnectionBody, DAGPatchBody, DAGRunPatchBody, DagRunState, @@ -1130,6 +1131,45 @@ export const useBackfillServiceCreateBackfill = < }) as unknown as Promise, ...options, }); +/** + * Post Connection + * Create connection entry. + * @param data The data for the request. + * @param data.requestBody + * @returns ConnectionResponse Successful Response + * @throws ApiError + */ +export const useConnectionServicePostConnection = < + TData = Common.ConnectionServicePostConnectionMutationResult, + TError = unknown, + TContext = unknown, +>( + options?: Omit< + UseMutationOptions< + TData, + TError, + { + requestBody: ConnectionBody; + }, + TContext + >, + "mutationFn" + >, +) => + useMutation< + TData, + TError, + { + requestBody: ConnectionBody; + }, + TContext + >({ + mutationFn: ({ requestBody }) => + ConnectionService.postConnection({ + requestBody, + }) as unknown as Promise, + ...options, + }); /** * Post Pool * Create a Pool. diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 5597b0a6a9cfe..4eecb848a57ce 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -41,6 +41,8 @@ import type { GetConnectionResponse, GetConnectionsData, GetConnectionsResponse, + PostConnectionData, + PostConnectionResponse, GetDagRunData, GetDagRunResponse, DeleteDagRunData, @@ -661,6 +663,31 @@ export class ConnectionService { }, }); } + + /** + * Post Connection + * Create connection entry. + * @param data The data for the request. + * @param data.requestBody + * @returns ConnectionResponse Successful Response + * @throws ApiError + */ + public static postConnection( + data: PostConnectionData, + ): CancelablePromise { + return __request(OpenAPI, { + method: "POST", + url: "/public/connections/", + body: data.requestBody, + mediaType: "application/json", + errors: { + 401: "Unauthorized", + 403: "Forbidden", + 409: "Conflict", + 422: "Validation Error", + }, + }); + } } export class DagRunService { diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index e3071b64936a8..603a20d090032 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -43,6 +43,21 @@ export type BaseInfoSchema = { status: string | null; }; +/** + * Connection Serializer for requests body. + */ +export type ConnectionBody = { + connection_id: string; + conn_type: string; + description?: string | null; + host?: string | null; + login?: string | null; + schema?: string | null; + port?: number | null; + password?: string | null; + extra?: string | null; +}; + /** * Connection Collection serializer for responses. */ @@ -62,6 +77,7 @@ export type ConnectionResponse = { login: string | null; schema: string | null; port: number | null; + password: string | null; extra: string | null; }; @@ -413,7 +429,7 @@ export type HealthInfoSchema = { export type HistoricalMetricDataResponse = { dag_run_types: DAGRunTypes; dag_run_states: DAGRunStates; - task_instance_states: airflow__api_fastapi__core_api__serializers__dashboard__TaskInstanceState; + task_instance_states: airflow__api_fastapi__core_api__serializers__ui__dashboard__TaskInstanceState; }; /** @@ -651,7 +667,7 @@ export type VersionInfo = { /** * TaskInstance serializer for responses. */ -export type airflow__api_fastapi__core_api__serializers__dashboard__TaskInstanceState = +export type airflow__api_fastapi__core_api__serializers__ui__dashboard__TaskInstanceState = { no_status: number; removed: number; @@ -841,6 +857,12 @@ export type GetConnectionsData = { export type GetConnectionsResponse = ConnectionCollectionResponse; +export type PostConnectionData = { + requestBody: ConnectionBody; +}; + +export type PostConnectionResponse = ConnectionResponse; + export type GetDagRunData = { dagId: string; dagRunId: string; @@ -1512,6 +1534,31 @@ export type $OpenApiTs = { 422: HTTPValidationError; }; }; + post: { + req: PostConnectionData; + res: { + /** + * Successful Response + */ + 201: ConnectionResponse; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Conflict + */ + 409: HTTPExceptionResponse; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; }; "/public/dags/{dag_id}/dagRuns/{dag_run_id}": { get: {