Skip to content

Commit

Permalink
Reviews
Browse files Browse the repository at this point in the history
  • Loading branch information
kaxil committed Nov 8, 2024
1 parent 4d56ec4 commit 9bdaf67
Show file tree
Hide file tree
Showing 7 changed files with 24 additions and 35 deletions.
13 changes: 5 additions & 8 deletions airflow/api_fastapi/execution_api/routes/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"])
7 changes: 2 additions & 5 deletions airflow/api_fastapi/execution_api/routes/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"}},
)

Expand All @@ -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"},
Expand Down
4 changes: 2 additions & 2 deletions airflow/api_fastapi/execution_api/routes/health.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"}
9 changes: 3 additions & 6 deletions airflow/api_fastapi/execution_api/routes/task_instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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={
Expand Down
7 changes: 2 additions & 5 deletions airflow/api_fastapi/execution_api/routes/variables.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"}},
)

Expand All @@ -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"},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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() == {
Expand All @@ -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() == {
Expand All @@ -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() == {
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"}
Expand All @@ -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() == {
Expand All @@ -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
Expand Down

0 comments on commit 9bdaf67

Please sign in to comment.