This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new e009e5302f6 AIP-84 Add user id when updating notes in TI (#48582)
e009e5302f6 is described below
commit e009e5302f607806c996d00467b2723cfc4118e4
Author: Kalyan R <[email protected]>
AuthorDate: Tue Apr 1 21:01:22 2025 +0530
AIP-84 Add user id when updating notes in TI (#48582)
* add user id to ti note
* refresh ti_note to fix test
* fix failing test
* add tests to patch dry run
* add explanation for adding session.refresh
---
.../core_api/routes/public/task_instances.py | 9 ++---
.../core_api/routes/public/test_task_instances.py | 44 ++++++++++++++++++----
.../src/tests_common/test_utils/api_fastapi.py | 14 +++++++
3 files changed, 54 insertions(+), 13 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
index 5defd2b08ff..6a3afa74b4f 100644
---
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
+++
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
@@ -61,7 +61,7 @@ from airflow.api_fastapi.core_api.datamodels.task_instances
import (
TaskInstancesBatchBody,
)
from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
-from airflow.api_fastapi.core_api.security import ReadableTIFilterDep,
requires_access_dag
+from airflow.api_fastapi.core_api.security import GetUserDep,
ReadableTIFilterDep, requires_access_dag
from airflow.api_fastapi.logging.decorators import action_logging
from airflow.exceptions import TaskNotFound
from airflow.models import Base, DagRun
@@ -853,6 +853,7 @@ def patch_task_instance(
task_id: str,
request: Request,
body: PatchTaskInstanceBody,
+ user: GetUserDep,
session: SessionDep,
map_index: int = -1,
update_mask: list[str] | None = Query(None),
@@ -883,13 +884,11 @@ def patch_task_instance(
ti = tis[0] if isinstance(tis, list) else tis
elif key == "note":
if update_mask or body.note is not None:
- # @TODO: replace None passed for user_id with actual user id
when
- # permissions and auth is in place.
if ti.task_instance_note is None:
- ti.note = (body.note, None)
+ ti.note = (body.note, user.get_id())
else:
ti.task_instance_note.content = body.note
- ti.task_instance_note.user_id = None
+ ti.task_instance_note.user_id = user.get_id()
session.commit()
return TaskInstanceResponse.model_validate(ti)
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
index 486c6eaa8a7..228b6aa7f9f 100644
---
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
+++
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
@@ -42,6 +42,7 @@ from airflow.utils.state import DagRunState, State,
TaskInstanceState
from airflow.utils.timezone import datetime
from airflow.utils.types import DagRunType
+from tests_common.test_utils.api_fastapi import _check_task_instance_note
from tests_common.test_utils.db import (
clear_db_runs,
clear_rendered_ti_fields,
@@ -3355,13 +3356,21 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
assert mock_set_ti_state.call_count == set_ti_state_call_count
@pytest.mark.parametrize(
- "new_note_value",
+ "new_note_value,ti_note_data",
[
- "My super cool TaskInstance note.",
- None,
+ (
+ "My super cool TaskInstance note.",
+ {"content": "My super cool TaskInstance note.", "user_id":
"test"},
+ ),
+ (
+ None,
+ {"content": None, "user_id": "test"},
+ ),
],
)
- def test_update_mask_set_note_should_respond_200(self, test_client,
session, new_note_value):
+ def test_update_mask_set_note_should_respond_200(
+ self, test_client, session, new_note_value, ti_note_data
+ ):
self.create_task_instances(session)
response = test_client.patch(
"/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context",
@@ -3369,7 +3378,8 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
json={"note": new_note_value},
)
assert response.status_code == 200, response.text
- assert response.json() == {
+ response_data = response.json()
+ assert response_data == {
"dag_id": self.DAG_ID,
"dag_version": None,
"duration": 10000.0,
@@ -3403,6 +3413,7 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
"trigger": None,
"triggerer_job": None,
}
+ _check_task_instance_note(session, response_data["id"], ti_note_data)
def test_set_note_should_respond_200(self, test_client, session):
self.create_task_instances(session)
@@ -3412,7 +3423,8 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
json={"note": new_note_value},
)
assert response.status_code == 200, response.text
- assert response.json() == {
+ response_data = response.json()
+ assert response_data == {
"dag_id": self.DAG_ID,
"dag_version": None,
"duration": 10000.0,
@@ -3447,6 +3459,10 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
"triggerer_job": None,
}
+ _check_task_instance_note(
+ session, response_data["id"], {"content": new_note_value,
"user_id": "test"}
+ )
+
def test_set_note_should_respond_200_mapped_task_instance_with_rtif(self,
test_client, session):
"""Verify we don't duplicate rows through join to RTIF"""
tis = self.create_task_instances(session)
@@ -3468,8 +3484,9 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
json={"note": new_note_value},
)
assert response.status_code == 200, response.text
+ response_data = response.json()
- assert response.json() == {
+ assert response_data == {
"dag_id": self.DAG_ID,
"dag_version": None,
"duration": 10000.0,
@@ -3504,6 +3521,10 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
"triggerer_job": None,
}
+ _check_task_instance_note(
+ session, response_data["id"], {"content": new_note_value,
"user_id": "test"}
+ )
+
def test_set_note_should_respond_200_when_note_is_empty(self, test_client,
session):
tis = self.create_task_instances(session)
for ti in tis:
@@ -3516,7 +3537,11 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
json={"note": new_note_value},
)
assert response.status_code == 200, response.text
- assert response.json()["note"] == new_note_value
+ response_data = response.json()
+ assert response_data["note"] == new_note_value
+ _check_task_instance_note(
+ session, response_data["id"], {"content": new_note_value,
"user_id": "test"}
+ )
@mock.patch("airflow.models.dag.DAG.set_task_instance_state")
def test_should_raise_409_for_updating_same_task_instance_state(
@@ -3650,6 +3675,8 @@ class
TestPatchTaskInstanceDryRun(TestTaskInstanceEndpoint):
assert task_before == task_after
+ _check_task_instance_note(session, task_after["id"], {"content":
"placeholder-note", "user_id": None})
+
def test_should_respond_401(self, unauthenticated_test_client):
response = unauthenticated_test_client.patch(
f"{self.ENDPOINT_URL}/dry_run",
@@ -3687,6 +3714,7 @@ class
TestPatchTaskInstanceDryRun(TestTaskInstanceEndpoint):
task_after = test_client.get(f"{self.ENDPOINT_URL}/{map_index}").json()
assert task_before == task_after
+ _check_task_instance_note(session, task_after["id"], None)
@pytest.mark.parametrize(
"error, code, payload",
diff --git a/devel-common/src/tests_common/test_utils/api_fastapi.py
b/devel-common/src/tests_common/test_utils/api_fastapi.py
index 5c99ceeb9ee..10399f919ca 100644
--- a/devel-common/src/tests_common/test_utils/api_fastapi.py
+++ b/devel-common/src/tests_common/test_utils/api_fastapi.py
@@ -20,6 +20,7 @@ import json
from airflow.models import DagRun, Log
from airflow.models.dagrun import DagRunNote
+from airflow.models.taskinstance import TaskInstanceNote
from airflow.sdk.execution_time.secrets_masker import DEFAULT_SENSITIVE_FIELDS
as sensitive_fields
@@ -81,3 +82,16 @@ def _check_dag_run_note(session, dr_id, note_data):
else:
assert dr_note.user_id == note_data.get("user_id")
assert dr_note.content == note_data.get("content")
+
+
+def _check_task_instance_note(session, ti_id, note_data):
+ ti_note =
session.query(TaskInstanceNote).filter_by(ti_id=ti_id).one_or_none()
+ if note_data is None:
+ assert ti_note is None
+ else:
+ # Had to add this refresh because
TestPatchTaskInstance::test_set_note_should_respond_200_mapped_task_instance_with_rtif
+ # was failing for map index = 2. Unless I force refresh ti_note, it
was returning the old value.
+ # Even if I reverse the order of map indexes, only the map index 2 was
returning older value.
+ session.refresh(ti_note)
+ assert ti_note.content == note_data["content"]
+ assert ti_note.user_id == note_data["user_id"]
