jason810496 commented on code in PR #76457:
URL: https://github.com/apache/airflow/pull/76457#discussion_r4512688832
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_warning.py:
##########
@@ -73,9 +90,43 @@ def list_dag_warnings(
limit=limit,
session=session,
)
- dag_warnings = session.scalars(dag_warnings_select)
+ dag_warnings = session.scalars(dag_warnings_select).all()
+
+ # Being able to read the Dag a file-level warning is attached to does not
mean being able to
+ # read every Dag of that file, so its message is redacted like the source
of that file.
+ file_level_warnings = [
+ dag_warning
+ for dag_warning in dag_warnings
+ if dag_warning.warning_type not in PER_DAG_WARNING_TYPES and
dag_warning.dag_model.relative_fileloc
+ ]
Review Comment:
This is now the third per-file check (`dag_sources`, `import_error`, here),
each written differently. Could it be one helper in `security.py` next to
`REDACTED_FILE_CONTENT` (files and readable Dag ids in, files with unreadable
Dags out), so a fix like the stale filter lands once?
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_warning.py:
##########
@@ -73,9 +90,43 @@ def list_dag_warnings(
limit=limit,
session=session,
)
- dag_warnings = session.scalars(dag_warnings_select)
+ dag_warnings = session.scalars(dag_warnings_select).all()
+
+ # Being able to read the Dag a file-level warning is attached to does not
mean being able to
+ # read every Dag of that file, so its message is redacted like the source
of that file.
+ file_level_warnings = [
+ dag_warning
+ for dag_warning in dag_warnings
+ if dag_warning.warning_type not in PER_DAG_WARNING_TYPES and
dag_warning.dag_model.relative_fileloc
+ ]
+ redacted_warnings: set[DagWarning] = set()
+ if file_level_warnings:
+ files_with_unreadable_dags = set(
+ session.execute(
+ select(DagModel.bundle_name, DagModel.relative_fileloc)
+ .where(
+ tuple_(DagModel.bundle_name,
DagModel.relative_fileloc).in_(
+ {(w.dag_model.bundle_name,
w.dag_model.relative_fileloc) for w in file_level_warnings}
+ ),
+ DagModel.dag_id.not_in(readable_dag_warning_filter.value
or set()),
+ )
+ .distinct()
+ ).tuples()
+ )
Review Comment:
Two things here:
- Stale Dags still recorded against the file keep its warnings redacted
after they leave it, though the message is computed from the current file. The
duplicate Dag id check in `collection.py` filters `is_stale` for this reason.
- The caller's readable Dag ids are sent a third time as a `NOT IN` list,
which gets large on big deployments.
Could we fetch the file's Dag ids and check them in Python like
`dag_sources` does? (needs `false` from `sqlalchemy`)
```suggestion
readable_dag_ids = readable_dag_warning_filter.value or set()
file_dags = session.execute(
select(DagModel.bundle_name, DagModel.relative_fileloc,
DagModel.dag_id).where(
tuple_(DagModel.bundle_name, DagModel.relative_fileloc).in_(
{(w.dag_model.bundle_name, w.dag_model.relative_fileloc)
for w in file_level_warnings}
),
DagModel.is_stale == false(),
)
).tuples()
files_with_unreadable_dags = {
(bundle_name, fileloc) for bundle_name, fileloc, dag_id in
file_dags if dag_id not in readable_dag_ids
}
```
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_warning.py:
##########
@@ -73,9 +90,43 @@ def list_dag_warnings(
limit=limit,
session=session,
)
- dag_warnings = session.scalars(dag_warnings_select)
+ dag_warnings = session.scalars(dag_warnings_select).all()
+
+ # Being able to read the Dag a file-level warning is attached to does not
mean being able to
+ # read every Dag of that file, so its message is redacted like the source
of that file.
+ file_level_warnings = [
+ dag_warning
+ for dag_warning in dag_warnings
+ if dag_warning.warning_type not in PER_DAG_WARNING_TYPES and
dag_warning.dag_model.relative_fileloc
+ ]
+ redacted_warnings: set[DagWarning] = set()
+ if file_level_warnings:
+ files_with_unreadable_dags = set(
+ session.execute(
+ select(DagModel.bundle_name, DagModel.relative_fileloc)
+ .where(
+ tuple_(DagModel.bundle_name,
DagModel.relative_fileloc).in_(
+ {(w.dag_model.bundle_name,
w.dag_model.relative_fileloc) for w in file_level_warnings}
+ ),
+ DagModel.dag_id.not_in(readable_dag_warning_filter.value
or set()),
+ )
+ .distinct()
+ ).tuples()
+ )
+ redacted_warnings = {
+ w
+ for w in file_level_warnings
+ if (w.dag_model.bundle_name, w.dag_model.relative_fileloc) in
files_with_unreadable_dags
+ }
+
+ response_warnings = []
+ for dag_warning in dag_warnings:
+ response_warning = DAGWarningResponse.model_validate(dag_warning)
+ if dag_warning in redacted_warnings:
+ response_warning.message = REDACTED_FILE_CONTENT
Review Comment:
The message is redacted only in the response, but `order_by=message` still
sorts on the stored one, so where a redacted row lands among readable ones
hints at its content. Is that acceptable here, or should sorting by `message`
be limited?
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_warning.py:
##########
@@ -138,3 +143,78 @@ def test_get_dag_warnings_importer_defined_type(self,
test_client, session):
assert response_json["total_entries"] == 1
assert response_json["dag_warnings"][0]["warning_type"] ==
"yaml:schema_violation"
assert response_json["dag_warnings"][0]["message"] == "importer
message"
+
+
+class TestGetDagWarningsOfMultiDagFile:
+ FILE_LEVEL_MESSAGE = "Line 12, Col 4\nCode:
BashOperator(task_id='colocated_task')"
+
+ @pytest.fixture(autouse=True)
+ def colocate_dags(self, session):
+ """Put DAG1 and DAG2 in one file and DAG3 in a file of its own."""
+ session.get(DagModel, DAG1_ID).relative_fileloc = "multi_dag_file.py"
+ session.get(DagModel, DAG2_ID).relative_fileloc = "multi_dag_file.py"
+ session.get(DagModel, DAG3_ID).relative_fileloc = "single_dag_file.py"
+ session.commit()
+
+ @pytest.mark.parametrize("warning_type", ["runtime varying value",
"yaml:schema_violation"])
+ @pytest.mark.parametrize(
+ ("readable_dag_ids", "expected_messages"),
+ [
+ pytest.param({DAG1_ID}, {DAG1_ID: REDACTED_FILE_CONTENT},
id="colocated-dag-unreadable"),
+ pytest.param(
+ {DAG1_ID, DAG3_ID},
+ {DAG1_ID: REDACTED_FILE_CONTENT, DAG3_ID: FILE_LEVEL_MESSAGE},
+ id="only-file-with-unreadable-dag-redacted",
+ ),
+ pytest.param(
+ {DAG1_ID, DAG2_ID},
+ {DAG1_ID: FILE_LEVEL_MESSAGE, DAG2_ID: FILE_LEVEL_MESSAGE},
+ id="all-dags-in-file-readable",
+ ),
+ ],
+ )
+
@mock.patch("airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_dag_ids")
Review Comment:
`AGENTS.md` asks for `spec`/`autospec` when mocking (same at line 204).
```suggestion
@mock.patch(
"airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_dag_ids",
autospec=True,
)
```
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_warning.py:
##########
@@ -138,3 +143,78 @@ def test_get_dag_warnings_importer_defined_type(self,
test_client, session):
assert response_json["total_entries"] == 1
assert response_json["dag_warnings"][0]["warning_type"] ==
"yaml:schema_violation"
assert response_json["dag_warnings"][0]["message"] == "importer
message"
+
+
+class TestGetDagWarningsOfMultiDagFile:
+ FILE_LEVEL_MESSAGE = "Line 12, Col 4\nCode:
BashOperator(task_id='colocated_task')"
+
+ @pytest.fixture(autouse=True)
+ def colocate_dags(self, session):
+ """Put DAG1 and DAG2 in one file and DAG3 in a file of its own."""
+ session.get(DagModel, DAG1_ID).relative_fileloc = "multi_dag_file.py"
+ session.get(DagModel, DAG2_ID).relative_fileloc = "multi_dag_file.py"
+ session.get(DagModel, DAG3_ID).relative_fileloc = "single_dag_file.py"
+ session.commit()
+
+ @pytest.mark.parametrize("warning_type", ["runtime varying value",
"yaml:schema_violation"])
+ @pytest.mark.parametrize(
+ ("readable_dag_ids", "expected_messages"),
+ [
+ pytest.param({DAG1_ID}, {DAG1_ID: REDACTED_FILE_CONTENT},
id="colocated-dag-unreadable"),
+ pytest.param(
+ {DAG1_ID, DAG3_ID},
+ {DAG1_ID: REDACTED_FILE_CONTENT, DAG3_ID: FILE_LEVEL_MESSAGE},
+ id="only-file-with-unreadable-dag-redacted",
+ ),
+ pytest.param(
+ {DAG1_ID, DAG2_ID},
+ {DAG1_ID: FILE_LEVEL_MESSAGE, DAG2_ID: FILE_LEVEL_MESSAGE},
+ id="all-dags-in-file-readable",
+ ),
+ ],
+ )
+
@mock.patch("airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_dag_ids")
+ def
test_file_level_message_is_redacted_unless_caller_can_read_all_dags_in_file(
+ self,
+ mock_get_authorized_dag_ids,
+ test_client,
+ session,
+ readable_dag_ids,
+ expected_messages,
+ warning_type,
+ ):
+ mock_get_authorized_dag_ids.return_value = readable_dag_ids
+ for dag_id in (DAG1_ID, DAG2_ID, DAG3_ID):
+ session.add(DagWarning(dag_id, warning_type,
self.FILE_LEVEL_MESSAGE))
+ session.commit()
+
+ with assert_queries_count(4):
+ response = test_client.get("/dagWarnings", params={"warning_type":
warning_type})
+
+ assert response.status_code == 200
+ response_json = response.json()
+ assert response_json["total_entries"] == len(expected_messages)
+ assert {w["dag_id"]: w["message"] for w in
response_json["dag_warnings"]} == expected_messages
+ session.expire_all()
+ assert set(
+
session.scalars(select(DagWarning.message).where(DagWarning.warning_type ==
warning_type))
+ ) == {self.FILE_LEVEL_MESSAGE}
+
+ @pytest.mark.parametrize("warning_type", sorted(PER_DAG_WARNING_TYPES))
+
@mock.patch("airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_dag_ids")
Review Comment:
```suggestion
@mock.patch(
"airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_dag_ids",
autospec=True,
)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
