chore: import upstream snapshot with attribution

ingestion/tests/unit/ometa/test_lineage_sql_query_patch.py

@@ -0,0 +1,141 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Lineage sqlQuery patch tests (issue #29520).
+
+Patching an existing lineage edge must add or update the incoming sqlQuery, never
+drop an already-stored query, and never crash build_patch on an edge that already
+carries a query. Covered through both public entry points, add_lineage (by id) and
+add_lineage_by_name (by fqn).
+"""
+
+import json
+from unittest.mock import MagicMock
+
+import pytest
+
+from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
+from metadata.generated.schema.type.entityLineage import EntitiesEdge, LineageDetails
+from metadata.generated.schema.type.entityLineage import Source as LineageSource
+from metadata.generated.schema.type.entityReference import EntityReference
+from metadata.ingestion.ometa.mixins.lineage_mixin import OMetaLineageMixin
+
+FROM_ID = "d311bdf2-c4a9-4be3-9937-3b26309759af"
+TO_ID = "abea43f7-ccc2-4daf-9dfb-115549461244"
+STORED_QUERY = "SELECT 1 FROM source_table"
+INCOMING_QUERY = "SELECT 2 FROM source_table"
+
+
+class StubbedLineage(OMetaLineageMixin):
+    """OMetaLineageMixin with only the HTTP calls stubbed, so add_lineage and
+    add_lineage_by_name exercise the real patch reconstruction and build_patch.
+    `existing_edge` is what the edge lookup returns."""
+
+    def __init__(self, existing_edge):
+        self.client = MagicMock()
+        self._existing_edge = existing_edge
+
+    def _get_lineage_edge_for_references(self, from_entity, to_entity):
+        return self._existing_edge
+
+    def get_lineage_edge_by_name(self, *args, **kwargs):
+        return self._existing_edge
+
+    def get_suffix(self, entity):
+        return "/lineage"
+
+    def get_lineage_by_id(self, *args, **kwargs):
+        return {}
+
+    def get_lineage_by_name(self, *args, **kwargs):
+        return {}
+
+    def _update_cache(self, *args, **kwargs):
+        return None
+
+
+def edge_lookup(stored_query=None):
+    """Server edge-lookup payload, optionally carrying an already-stored sqlQuery."""
+    details = {"columnsLineage": []}
+    if stored_query is not None:
+        details["sqlQuery"] = stored_query
+    return {"edge": details}
+
+
+def incoming_details(sql_query):
+    return (
+        LineageDetails(source=LineageSource.QueryLineage, sqlQuery=sql_query)
+        if sql_query is not None
+        else LineageDetails(source=LineageSource.QueryLineage)
+    )
+
+
+def sql_query_ops(client):
+    """JSON-patch ops targeting /sqlQuery from the single client.patch call."""
+    if not client.patch.called:
+        return []
+    patch = json.loads(client.patch.call_args.kwargs["data"])
+    return [op for op in patch if op.get("path") == "/sqlQuery"]
+
+
+# (stored query on the existing edge, incoming query, expected query after patch)
+PATCH_CASES = [
+    pytest.param(None, INCOMING_QUERY, INCOMING_QUERY, id="backfill-onto-query-less-edge"),
+    pytest.param(STORED_QUERY, INCOMING_QUERY, INCOMING_QUERY, id="update-existing-query"),
+    pytest.param(STORED_QUERY, None, None, id="keep-stored-query-when-incoming-empty"),
+]
+
+
+def assert_patched(client, expected_query):
+    assert not client.put.called
+    if expected_query is None:
+        assert sql_query_ops(client) == []
+    else:
+        assert sql_query_ops(client) == [{"op": "add", "path": "/sqlQuery", "value": expected_query}]
+
+
+class TestAddLineageSqlQuery:
+    """add_lineage (by id) reconstruction adds/updates the incoming query, keeps a
+    stored one when the run has none, and never crashes on an edge with a query."""
+
+    @pytest.mark.parametrize("stored, incoming, expected", PATCH_CASES)
+    def test_patches_sql_query(self, stored, incoming, expected):
+        service = StubbedLineage(edge_lookup(stored))
+        request = AddLineageRequest(
+            edge=EntitiesEdge(
+                fromEntity=EntityReference(id=FROM_ID, type="table"),
+                toEntity=EntityReference(id=TO_ID, type="table"),
+                lineageDetails=incoming_details(incoming),
+            )
+        )
+
+        service.add_lineage(request, check_patch=True)
+
+        assert_patched(service.client, expected)
+
+
+class TestAddLineageByNameSqlQuery:
+    """Same reconstruction through the add_lineage_by_name (by fqn) entry point."""
+
+    @pytest.mark.parametrize("stored, incoming, expected", PATCH_CASES)
+    def test_patches_sql_query(self, stored, incoming, expected):
+        service = StubbedLineage(edge_lookup(stored))
+
+        service.add_lineage_by_name(
+            from_entity_fqn="svc.db.sch.source",
+            from_entity_type="table",
+            to_entity_fqn="svc.db.sch.target",
+            to_entity_type="table",
+            lineage_details=incoming_details(incoming),
+            check_patch=True,
+        )
+
+        assert_patched(service.client, expected)