From 9bdaf6758ceaa40288eda257097d851266a6ac44 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Fri, 8 Nov 2024 23:31:49 +0000 Subject: [PATCH] Reviews --- .../api_fastapi/execution_api/routes/__init__.py | 13 +++++-------- .../api_fastapi/execution_api/routes/connections.py | 7 ++----- airflow/api_fastapi/execution_api/routes/health.py | 4 ++-- .../execution_api/routes/task_instance.py | 9 +++------ .../api_fastapi/execution_api/routes/variables.py | 7 ++----- .../{test_connection.py => test_connections.py} | 8 ++++---- .../routes/{test_variable.py => test_variables.py} | 11 ++++++----- 7 files changed, 24 insertions(+), 35 deletions(-) rename tests/api_fastapi/execution_api/routes/{test_connection.py => test_connections.py} (92%) rename tests/api_fastapi/execution_api/routes/{test_variable.py => test_variables.py} (89%) diff --git a/airflow/api_fastapi/execution_api/routes/__init__.py b/airflow/api_fastapi/execution_api/routes/__init__.py index 40ccced0e2f7..76479c96cf77 100644 --- a/airflow/api_fastapi/execution_api/routes/__init__.py +++ b/airflow/api_fastapi/execution_api/routes/__init__.py @@ -17,13 +17,10 @@ from __future__ import annotations from airflow.api_fastapi.common.router import AirflowRouter -from airflow.api_fastapi.execution_api.routes.connections import connection_router -from airflow.api_fastapi.execution_api.routes.health import health_router -from airflow.api_fastapi.execution_api.routes.task_instance import ti_router -from airflow.api_fastapi.execution_api.routes.variables import variable_router +from airflow.api_fastapi.execution_api.routes import connections, health, task_instance, variables execution_api_router = AirflowRouter() -execution_api_router.include_router(connection_router) -execution_api_router.include_router(health_router) -execution_api_router.include_router(ti_router) -execution_api_router.include_router(variable_router) +execution_api_router.include_router(connections.router, prefix="/connections", tags=["Connections"]) +execution_api_router.include_router(health.router, tags=["Health"]) +execution_api_router.include_router(task_instance.router, prefix="/task_instance", tags=["Task Instance"]) +execution_api_router.include_router(variables.router, prefix="/variables", tags=["Variables"]) diff --git a/airflow/api_fastapi/execution_api/routes/connections.py b/airflow/api_fastapi/execution_api/routes/connections.py index f65b189453b9..553cb0785d67 100644 --- a/airflow/api_fastapi/execution_api/routes/connections.py +++ b/airflow/api_fastapi/execution_api/routes/connections.py @@ -28,9 +28,7 @@ from airflow.models.connection import Connection # TODO: Add dependency on JWT token -connection_router = AirflowRouter( - prefix="/connection", - tags=["Connection"], +router = AirflowRouter( responses={status.HTTP_404_NOT_FOUND: {"description": "Connection not found"}}, ) @@ -42,9 +40,8 @@ def get_task_token() -> datamodels.TIToken: return datamodels.TIToken(ti_key="test_key") -@connection_router.get( +@router.get( "/{connection_id}", - status_code=status.HTTP_200_OK, responses={ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"}, status.HTTP_403_FORBIDDEN: {"description": "Task does not have access to the connection"}, diff --git a/airflow/api_fastapi/execution_api/routes/health.py b/airflow/api_fastapi/execution_api/routes/health.py index c8d903815dc6..7bf4c4a2de0a 100644 --- a/airflow/api_fastapi/execution_api/routes/health.py +++ b/airflow/api_fastapi/execution_api/routes/health.py @@ -19,9 +19,9 @@ from airflow.api_fastapi.common.router import AirflowRouter -health_router = AirflowRouter(tags=["Health"]) +router = AirflowRouter() -@health_router.get("/health") +@router.get("/health") def health() -> dict: return {"status": "healthy"} diff --git a/airflow/api_fastapi/execution_api/routes/task_instance.py b/airflow/api_fastapi/execution_api/routes/task_instance.py index 4612b0c0425b..3ef37013f89b 100644 --- a/airflow/api_fastapi/execution_api/routes/task_instance.py +++ b/airflow/api_fastapi/execution_api/routes/task_instance.py @@ -35,16 +35,13 @@ from airflow.utils.state import State # TODO: Add dependency on JWT token -ti_router = AirflowRouter( - prefix="/task_instance", - tags=["Task Instance"], -) +router = AirflowRouter() log = logging.getLogger(__name__) -@ti_router.patch( +@router.patch( "/{task_instance_id}/state", status_code=status.HTTP_204_NO_CONTENT, # TODO: Add description to the operation @@ -133,7 +130,7 @@ def ti_update_state( ) -@ti_router.put( +@router.put( "/{task_instance_id}/heartbeat", status_code=status.HTTP_204_NO_CONTENT, responses={ diff --git a/airflow/api_fastapi/execution_api/routes/variables.py b/airflow/api_fastapi/execution_api/routes/variables.py index ba8eefe9564a..79df5678aca4 100644 --- a/airflow/api_fastapi/execution_api/routes/variables.py +++ b/airflow/api_fastapi/execution_api/routes/variables.py @@ -27,9 +27,7 @@ from airflow.models.variable import Variable # TODO: Add dependency on JWT token -variable_router = AirflowRouter( - prefix="/variable", - tags=["Variable"], +router = AirflowRouter( responses={status.HTTP_404_NOT_FOUND: {"description": "Variable not found"}}, ) @@ -41,9 +39,8 @@ def get_task_token() -> datamodels.TIToken: return datamodels.TIToken(ti_key="test_key") -@variable_router.get( +@router.get( "/{variable_key}", - status_code=status.HTTP_200_OK, responses={ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"}, status.HTTP_403_FORBIDDEN: {"description": "Task does not have access to the variable"}, diff --git a/tests/api_fastapi/execution_api/routes/test_connection.py b/tests/api_fastapi/execution_api/routes/test_connections.py similarity index 92% rename from tests/api_fastapi/execution_api/routes/test_connection.py rename to tests/api_fastapi/execution_api/routes/test_connections.py index 107fb8741eac..287a5bac8019 100644 --- a/tests/api_fastapi/execution_api/routes/test_connection.py +++ b/tests/api_fastapi/execution_api/routes/test_connections.py @@ -43,7 +43,7 @@ def test_connection_get_from_db(self, client, session): session.add(connection) session.commit() - response = client.get("/execution/connection/test_conn") + response = client.get("/execution/connections/test_conn") assert response.status_code == 200 assert response.json() == { @@ -66,7 +66,7 @@ def test_connection_get_from_db(self, client, session): {"AIRFLOW_CONN_TEST_CONN2": '{"uri": "http://root:admin@localhost:8080/https?headers=header"}'}, ) def test_connection_get_from_env_var(self, client, session): - response = client.get("/execution/connection/test_conn2") + response = client.get("/execution/connections/test_conn2") assert response.status_code == 200 assert response.json() == { @@ -81,7 +81,7 @@ def test_connection_get_from_env_var(self, client, session): } def test_connection_get_not_found(self, client): - response = client.get("/execution/connection/non_existent_test_conn") + response = client.get("/execution/connections/non_existent_test_conn") assert response.status_code == 404 assert response.json() == { @@ -95,7 +95,7 @@ def test_connection_get_access_denied(self, client): with mock.patch( "airflow.api_fastapi.execution_api.routes.connections.has_connection_access", return_value=False ): - response = client.get("/execution/connection/test_conn") + response = client.get("/execution/connections/test_conn") # Assert response status code and detail for access denied assert response.status_code == 403 diff --git a/tests/api_fastapi/execution_api/routes/test_variable.py b/tests/api_fastapi/execution_api/routes/test_variables.py similarity index 89% rename from tests/api_fastapi/execution_api/routes/test_variable.py rename to tests/api_fastapi/execution_api/routes/test_variables.py index db837d53a3f5..67247e4adb95 100644 --- a/tests/api_fastapi/execution_api/routes/test_variable.py +++ b/tests/api_fastapi/execution_api/routes/test_variables.py @@ -23,14 +23,15 @@ from airflow.models.variable import Variable +pytestmark = pytest.mark.db_test + class TestGetVariable: - @pytest.mark.db_test def test_variable_get_from_db(self, client, session): Variable.set(key="var1", value="value", session=session) session.commit() - response = client.get("/execution/variable/var1") + response = client.get("/execution/variables/var1") assert response.status_code == 200 assert response.json() == {"key": "var1", "value": "value"} @@ -44,13 +45,13 @@ def test_variable_get_from_db(self, client, session): {"AIRFLOW_VAR_KEY1": "VALUE"}, ) def test_variable_get_from_env_var(self, client, session): - response = client.get("/execution/variable/key1") + response = client.get("/execution/variables/key1") assert response.status_code == 200 assert response.json() == {"key": "key1", "value": "VALUE"} def test_variable_get_not_found(self, client): - response = client.get("/execution/variable/non_existent_var") + response = client.get("/execution/variables/non_existent_var") assert response.status_code == 404 assert response.json() == { @@ -64,7 +65,7 @@ def test_variable_get_access_denied(self, client): with mock.patch( "airflow.api_fastapi.execution_api.routes.variables.has_variable_access", return_value=False ): - response = client.get("/execution/variable/key1") + response = client.get("/execution/variables/key1") # Assert response status code and detail for access denied assert response.status_code == 403