kaxil commented on code in PR #70715:
URL: https://github.com/apache/airflow/pull/70715#discussion_r3686184256
##########
airflow-core/tests/unit/models/test_serialized_dag.py:
##########
@@ -1155,3 +1155,92 @@ def
test_deadline_reuse_skips_write_when_hash_matches(self, testing_dag_bundle,
alert = session.scalar(select(DAM).where(DAM.serialized_dag_id ==
orig_serdag.id))
assert alert is not None
assert alert.id == orig_alert.id
+
+
+class TestComputeHashAndStorageJson:
+ """Tests for SerializedDagModel._compute_hash_and_storage_json."""
+
+ @staticmethod
+ def _make_dag_data(fileloc="/tmp/test.py", bundle_name="test-bundle",
max_active_runs=16):
+ return {
+ "dag": {
+ "dag_id": "test_dag",
+ "fileloc": fileloc,
+ "bundle_name": bundle_name,
+ "max_active_runs": max_active_runs,
+ "tasks": [{"task_id": "t1"}, {"task_id": "t2"}],
+ },
+ "__version": 1,
+ }
+
+ def test_hash_matches_legacy_hash_method(self):
+ dag_data = self._make_dag_data()
+ new_hash, _, _ = SDM._compute_hash_and_storage_json(dag_data)
+
+ sorted_data = SDM._sort_serialized_dag_dict(dag_data)
+ legacy_copy = sorted_data.copy()
+ legacy_copy["dag"] = legacy_copy["dag"].copy()
+ legacy_copy["dag"].pop("fileloc", None)
+ legacy_copy["dag"].pop("bundle_name", None)
+ legacy_json = json.dumps(legacy_copy, sort_keys=True).encode("utf-8")
+ legacy_hash = md5(legacy_json).hexdigest()
+
+ assert new_hash == legacy_hash
+
+ def test_storage_json_contains_fileloc(self):
+ dag_data = self._make_dag_data(fileloc="/my/dag/file.py")
+ _, storage_json, _ = SDM._compute_hash_and_storage_json(dag_data)
+ parsed = json.loads(storage_json)
+ assert parsed["dag"]["fileloc"] == "/my/dag/file.py"
+
+ def test_hash_excludes_fileloc(self):
+ hash1, _, _ =
SDM._compute_hash_and_storage_json(self._make_dag_data(fileloc="/path/a.py"))
+ hash2, _, _ =
SDM._compute_hash_and_storage_json(self._make_dag_data(fileloc="/path/b.py"))
+ assert hash1 == hash2
+
+ def test_hash_excludes_bundle_name(self):
+ hash1, _, _ =
SDM._compute_hash_and_storage_json(self._make_dag_data(bundle_name="bundle-a"))
+ hash2, _, _ =
SDM._compute_hash_and_storage_json(self._make_dag_data(bundle_name="bundle-b"))
+ assert hash1 == hash2
+
+ def test_hash_changes_when_dag_content_changes(self):
+ hash1, _, _ =
SDM._compute_hash_and_storage_json(self._make_dag_data(max_active_runs=16))
+ hash2, _, _ =
SDM._compute_hash_and_storage_json(self._make_dag_data(max_active_runs=32))
+ assert hash1 != hash2
+
+ def test_original_dict_not_mutated(self):
+ import copy
+
+ dag_data = self._make_dag_data()
+ original = copy.deepcopy(dag_data)
+ SDM._compute_hash_and_storage_json(dag_data)
+ assert dag_data == original
+
+ @pytest.mark.parametrize(
+ "compress",
+ [
+ pytest.param(False, id="uncompressed"),
+ pytest.param(True, id="compressed"),
+ ],
+ )
+ def test_init_data_cache_and_storage(self, compress, monkeypatch):
+ import zlib
+
+
monkeypatch.setattr("airflow.models.serialized_dag._COMPRESS_SERIALIZED_DAGS",
compress)
+ dag_data = self._make_dag_data()
+ lazy_dag = mock.MagicMock()
Review Comment:
Repo convention is to always spec mocks:
`mock.MagicMock(spec=LazyDeserializedDAG)`, otherwise attribute typos pass
silently. The `import copy` / `import zlib` inside the tests should move to the
top of the file too, and `SDM(lazy_dag)` does the same as the `__new__` +
`__init__` pair below.
##########
airflow-core/src/airflow/models/serialized_dag.py:
##########
@@ -366,19 +364,37 @@ def __init__(self, dag: LazyDeserializedDAG) -> None:
def __repr__(self) -> str:
return f"<SerializedDag: {self.dag_id}>"
+ @classmethod
+ def _compute_hash_and_storage_json(cls, dag_data: dict) -> tuple[str,
bytes, dict]:
+ """
+ Compute the Dag hash and storage JSON in a single pass.
+
+ Sorts the serialized dict once, generates hash JSON (without
fileloc/bundle_name),
+ then restores those fields and generates storage JSON.
+
+ :return: (dag_hash, storage_json_bytes, sorted_data)
+ """
+ sorted_data = cls._sort_serialized_dag_dict(dag_data)
+ dag_section = sorted_data["dag"]
+ saved_fileloc = dag_section.pop("fileloc", None)
+ saved_bundle_name = dag_section.pop("bundle_name", None)
+
+ hash_json = json.dumps(sorted_data, sort_keys=True).encode("utf-8")
+ dag_hash = md5(hash_json).hexdigest()
+
+ if saved_fileloc is not None:
+ dag_section["fileloc"] = saved_fileloc
+ if saved_bundle_name is not None:
+ dag_section["bundle_name"] = saved_bundle_name
+
+ storage_json = json.dumps(sorted_data, sort_keys=True).encode("utf-8")
+ return dag_hash, storage_json, sorted_data
+
@classmethod
def hash(cls, dag_data):
"""Hash the data to get the dag_hash."""
- dag_data = cls._sort_serialized_dag_dict(dag_data)
- data_ = dag_data.copy()
- # Remove fileloc from the hash so changes to fileloc
- # does not affect the hash. In 3.0+, a combination of
- # bundle_path and relative fileloc more correctly determines the
- # dag file location.
- data_["dag"].pop("fileloc", None)
- data_["dag"].pop("bundle_name", None)
- data_json = json.dumps(data_, sort_keys=True).encode("utf-8")
- return md5(data_json).hexdigest()
+ dag_hash, _, _ = cls._compute_hash_and_storage_json(dag_data)
Review Comment:
This makes the common path slower. `write_dag` calls `hash()` on every
serialization pass to decide whether anything changed, and unchanged DAGs (the
usual case) now pay for building the storage JSON that gets thrown away here.
On a 50-task dict I measure standalone `hash()` about 26% slower. The
constructor doesn't get faster either: on main, `__init__` only sorts inside
`hash()`, the storage dump is a plain `json.dumps(dag_data, sort_keys=True)` on
the unsorted dict, so it was already one sort plus two dumps, exactly what this
helper does. The duplication that does exist is `write_dag` hashing and then
`cls(dag)` hashing again when the DAG changed; passing the precomputed hash
into the constructor would remove that without touching the unchanged-DAG path.
##########
airflow-core/src/airflow/models/serialized_dag.py:
##########
@@ -366,19 +364,37 @@ def __init__(self, dag: LazyDeserializedDAG) -> None:
def __repr__(self) -> str:
return f"<SerializedDag: {self.dag_id}>"
+ @classmethod
+ def _compute_hash_and_storage_json(cls, dag_data: dict) -> tuple[str,
bytes, dict]:
+ """
+ Compute the Dag hash and storage JSON in a single pass.
+
+ Sorts the serialized dict once, generates hash JSON (without
fileloc/bundle_name),
+ then restores those fields and generates storage JSON.
+
+ :return: (dag_hash, storage_json_bytes, sorted_data)
+ """
+ sorted_data = cls._sort_serialized_dag_dict(dag_data)
+ dag_section = sorted_data["dag"]
+ saved_fileloc = dag_section.pop("fileloc", None)
+ saved_bundle_name = dag_section.pop("bundle_name", None)
+
+ hash_json = json.dumps(sorted_data, sort_keys=True).encode("utf-8")
+ dag_hash = md5(hash_json).hexdigest()
+
+ if saved_fileloc is not None:
+ dag_section["fileloc"] = saved_fileloc
+ if saved_bundle_name is not None:
+ dag_section["bundle_name"] = saved_bundle_name
+
+ storage_json = json.dumps(sorted_data, sort_keys=True).encode("utf-8")
Review Comment:
Dumping `sorted_data` here changes what gets stored when
`compress_serialized_dags` is enabled. On main the stored JSON is dumped from
the unsorted dict, so list order is preserved; `_sort_serialized_dag_dict`
reorders tasks by task_id and alphabetizes every all-string list, e.g.
`op_args: ["zeta", "alpha"]` round-trips from compressed storage as `["alpha",
"zeta"]`, and `template_fields`/`tags` get the same treatment. I verified the
bytes differ on such a dict. It also means `_data_compressed` no longer
decompresses to the same lists as `__data_cache` or the uncompressed `_data`,
which keep the original order, so the "stored JSON is identical" claim only
holds for dicts whose lists are already sorted. (A `fileloc: None` entry would
also be dropped from storage entirely, since the restore checks `is not None`.)
--
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]
