acryl-datahub 1.2.0.2rc3__py3-none-any.whl → 1.2.0.3rc1__py3-none-any.whl

{acryl_datahub-1.2.0.2rc3.dist-info → acryl_datahub-1.2.0.3rc1.dist-info}/RECORD

@@ -1,7 +1,7 @@
-acryl_datahub-1.2.0.2rc3.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
+acryl_datahub-1.2.0.3rc1.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
 datahub/__init__.py,sha256=aq_i5lVREmoLfYIqcx_pEQicO855YlhD19tWc1eZZNI,59
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
-datahub/_version.py,sha256=cyLeAlWzIql-O65xt-MALgl-0jJ-qmzJ-gFq-t6FkdE,323
+datahub/_version.py,sha256=BAFY2OaLzkEm6Hs8RwoQ69XvJQdk3iPPrc9oCg1xGAE,323
 datahub/entrypoints.py,sha256=9Qf-37rNnTzbGlx8S75OCDazIclFp6zWNcCEL1zCZto,9015
 datahub/errors.py,sha256=p5rFAdAGVCk4Lqolol1YvthceadUSwpaCxLXRcyCCFQ,676
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -227,7 +227,7 @@ datahub/ingestion/source/pulsar.py,sha256=u5F8QnCLJsht5-7XCiUTsnfhCPIpKVB_l32CgM
 datahub/ingestion/source/redash.py,sha256=YxjSad-X_wPmxYH8dJmFz_VCFhiLTCTSlK99WdvcYiA,30653
 datahub/ingestion/source/salesforce.py,sha256=Pa_w1XszxFd8fyhpSWOfc2nOnevHwwstIvnRrQT4R9M,40584
 datahub/ingestion/source/source_registry.py,sha256=a2mLjJPLkSI-gYCTb_7U7Jo4D8jGknNQ_yScPIihXFk,1208
-datahub/ingestion/source/sql_queries.py,sha256=1SdEZGAmojfWbU1xKbezH6uqwRd2G0rgpK_Sh9MRj1U,11791
+datahub/ingestion/source/sql_queries.py,sha256=9ICUC6tpXpxvtwfK-9lytJzFcLe8MrNlciwy9DIFM-4,13764
 datahub/ingestion/source/superset.py,sha256=oi7F2jlvkVr9ItJ_r1Jm4bYfXHYu4vPAFPMPaGJKB84,50608
 datahub/ingestion/source/abs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/abs/config.py,sha256=mBQe0JTaP-Rcv4HnMUUySoYbSr4r3jDEMioxaXHnxXU,6709
@@ -524,7 +524,7 @@ datahub/ingestion/source/sql/sqlalchemy_uri_mapper.py,sha256=KOpbmDIE2h1hyYEsbVH
 datahub/ingestion/source/sql/teradata.py,sha256=xL_c_UEM_JT-xoMw5Nb5UvSNBUfTGol5CpOkgK5Bsjk,65412
 datahub/ingestion/source/sql/trino.py,sha256=zIfQ6GvW8Sbw4sxqsTcnibT51STka_nzNYvmld6HfHw,18947
 datahub/ingestion/source/sql/two_tier_sql_source.py,sha256=AB3Gtx4omAy_08zadHQpmUGmIGufkZ6o_ihWNnfvzYc,5783
-datahub/ingestion/source/sql/vertica.py,sha256=MeohL8j68ISES1RhrBXLQlkT_YqgT-AvHRxuVCJSMbE,33458
+datahub/ingestion/source/sql/vertica.py,sha256=blnu1-H7vnSQD3ZD5QTotoQ2DQJWJeR0uxz_clxiPGo,33518
 datahub/ingestion/source/sql/mssql/__init__.py,sha256=1agpl8S_uDW40olkhCX_W19dbr5GO9qgjS3R7pLRZSk,87
 datahub/ingestion/source/sql/mssql/job_models.py,sha256=nAo3rciu-w2-dXCz6_ekDEbGMEjCMEfh8WvSfXoF2l0,9359
 datahub/ingestion/source/sql/mssql/source.py,sha256=Uise_u6yXKU__9B_U3D3yObWNIVDzrz2AgEDZOlk6bQ,43101
@@ -558,11 +558,11 @@ datahub/ingestion/source/unity/config.py,sha256=7QosoBthg9kirHfXev_vhefkobUxYnp1
 datahub/ingestion/source/unity/connection_test.py,sha256=B143Wb28fS0V4GhygU9hzKqiArWBjsQO54IUCPf23dc,2586
 datahub/ingestion/source/unity/ge_profiler.py,sha256=NBRHZceq-f95iUn7u0h7cgcd9nAc48Aa-lmp_BqE0As,8409
 datahub/ingestion/source/unity/hive_metastore_proxy.py,sha256=IAWWJjaW0si_UF52Se2D7wmdYRY_afUG4QlVmQu6xaw,15351
-datahub/ingestion/source/unity/proxy.py,sha256=jfQ1N8Xrp08zeYN2j74YTweusygXtK4Q-5_FBbwCVTE,22803
+datahub/ingestion/source/unity/proxy.py,sha256=iZ2ftKOXkxpFr0_2bEYEm31ci9OZJWFYgna3DNLCXrQ,26706
 datahub/ingestion/source/unity/proxy_profiling.py,sha256=WLqvYP6MziaisA4LYL4T_GA-kPt6Xdde7bfaYsjYw40,9663
 datahub/ingestion/source/unity/proxy_types.py,sha256=qrvHiwPzl5cPX-KRvcIGGeJVdr0I8XUQmoAI6ErZ-v8,9371
 datahub/ingestion/source/unity/report.py,sha256=XFT9oQfvEB4RkTvWGgFOoQuLPUN_AIoPXZ79xeDhGHQ,2831
-datahub/ingestion/source/unity/source.py,sha256=47sWCYb3pd21RjIqs9NH4h6VMtF_YMvqS9-6cegGi1w,49980
+datahub/ingestion/source/unity/source.py,sha256=udK1_WI7RO2Uzx9xis0mx264nHXDiMeP4Kah3-CFXis,49832
 datahub/ingestion/source/unity/tag_entities.py,sha256=iWl6nRAWSye1hoFDx_Xh4aT53PN0sGzlX7n1-oTVUv8,11568
 datahub/ingestion/source/unity/usage.py,sha256=0wETBAaZvHI_EGgBlxX3bKsVHEAdnUV8_bKI_lbyWjY,11500
 datahub/ingestion/source/usage/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -627,7 +627,7 @@ datahub/lite/lite_server.py,sha256=p9Oa2nNs65mqcssSIVOr7VOzWqfVstz6ZQEdT4f82S0,1
 datahub/lite/lite_util.py,sha256=G0LQHKkyEb1pc_q183g6hflShclGx7kikgMaOxtVVcs,4545
 datahub/metadata/__init__.py,sha256=AjhXPjI6cnpdcrBRrE5gOWo15vv2TTl2ctU4UAnUN7A,238
 datahub/metadata/_internal_schema_classes.py,sha256=Zh2volhvkUCENRavXDwPsmwfRe62k8_O6f5QT8_bh-g,1051205
-datahub/metadata/schema.avsc,sha256=4X6Jx5TFcOGY8Qxdm-FSgbGkzG2wND992brsxwgumSU,737966
+datahub/metadata/schema.avsc,sha256=u5iUlz9AnFfJijjJ9xcZx3MoiBfxWSmhr7pZIOg0tMo,738363
 datahub/metadata/schema_classes.py,sha256=tPT8iHCak4IsZi_oL0nirbPpI8ETTPTZzapqLRpeKU4,1326
 datahub/metadata/urns.py,sha256=nfrCTExR-k2P9w272WVtWSN3xW1VUJngPwP3xnvULjU,1217
 datahub/metadata/_urns/__init__.py,sha256=cOF3GHMDgPhmbLKbN02NPpuLGHSu0qNgQyBRv08eqF0,243
@@ -781,7 +781,7 @@ datahub/metadata/schemas/DataHubUpgradeResult.avsc,sha256=VydVb4yqjIviR73-T6TooF
 datahub/metadata/schemas/DataHubViewInfo.avsc,sha256=U3fBIoG9ietLUpOknfQGNekqBdPQYwvhhv9RQv6gEeg,11642
 datahub/metadata/schemas/DataHubViewKey.avsc,sha256=p53axIdSVbubo3r23Vpsed7NqRcQBMGveVikEHAVAok,424
 datahub/metadata/schemas/DataJobInfo.avsc,sha256=Bc9qdDcXI0GQdEgNTpgHaBbnrppDKQ-1xR26diOSVIQ,7488
-datahub/metadata/schemas/DataJobInputOutput.avsc,sha256=H1O8eAzZV34tvULdu67iBSWkdn08rt7wS208b8Nisbk,15268
+datahub/metadata/schemas/DataJobInputOutput.avsc,sha256=BYKImZ8kQQHqWbSBMKXWD0tGi96yzUt8zJFW3_twVVM,15575
 datahub/metadata/schemas/DataJobKey.avsc,sha256=S7egH8jWjKW52MG6Pg7plDoP15XfTTiMde5V6nR6ycE,1624
 datahub/metadata/schemas/DataPlatformInfo.avsc,sha256=WGPFumBNHbR75vsLrivnRCbBc8vSCuxDw2UlylMieh4,2686
 datahub/metadata/schemas/DataPlatformInstance.avsc,sha256=SNd3v_YyyLaDflv8Rd5cQR9GrVuky_cDTkYM6FqJiM8,1058
@@ -884,7 +884,7 @@ datahub/metadata/schemas/MLModelProperties.avsc,sha256=hDCBHxGe-cmCBeU1k0ANuQlKj
 datahub/metadata/schemas/MLPrimaryKeyKey.avsc,sha256=F3lgpMnHBhXsqGncHE9x06P-0RiNCrzbUUWlMkPJxFI,1132
 datahub/metadata/schemas/MLPrimaryKeyProperties.avsc,sha256=URIuOpS93RVk8MZVcbZ-dmTwu_cN3KSOKxSR8fm-eTo,6744
 datahub/metadata/schemas/MLTrainingRunProperties.avsc,sha256=WGgj0MuQrGD4UgvyHCJHzTnHja2LlJTOr1gLu8SySj0,4269
-datahub/metadata/schemas/MetadataChangeEvent.avsc,sha256=l3tVuQces7sKrwWsaIJrn3nMRUiCl3MHqCJJHcw7Ylc,377705
+datahub/metadata/schemas/MetadataChangeEvent.avsc,sha256=oNK0N8WrBsM_AoZkdYAMJQlhYzbao_QWaAMOjqEvPBw,378228
 datahub/metadata/schemas/MetadataChangeLog.avsc,sha256=soCmgrcEBE5yS-mQIm-RIefhb74ONj9Fqayxa0-59KE,13254
 datahub/metadata/schemas/MetadataChangeProposal.avsc,sha256=pT14vUmpj7VJ8hinQ0pcCUtRKx6RAGHWh1eJixkqaE8,12647
 datahub/metadata/schemas/Metrics.avsc,sha256=O7DJGjOwmHbb1x_Zj7AuM_HaHKjBvkfJKfUsX8icXD4,690
@@ -933,7 +933,7 @@ datahub/metadata/schemas/TestInfo.avsc,sha256=rye90gdY_lxZt_1gpa_Xum923CJgDU6i_e
 datahub/metadata/schemas/TestKey.avsc,sha256=eL-S4Z8EuN1JEXV1t4fy3LwmdA2dJURasFcKygP2rLY,421
 datahub/metadata/schemas/TestResults.avsc,sha256=uspC95AzRvz2_AgHVb5-fxELm5u8NmBTaFVJvGunmh0,5178
 datahub/metadata/schemas/TrainingData.avsc,sha256=7p7sFBA_UyV5IbNU5qLgS3vVu70yevKCfJKSGmTzVTg,2069
-datahub/metadata/schemas/UpstreamLineage.avsc,sha256=iaeFRbL2aVSYFwj-HQHyfIVaHRrK3kLbkkLXgIfJTsk,10639
+datahub/metadata/schemas/UpstreamLineage.avsc,sha256=dtpI7KUv9kYyGZmIlKfR2zLwgqsHO5P20egvIeup1EU,11000
 datahub/metadata/schemas/UsageAggregation.avsc,sha256=QaF6lyWGUq8IlRel2h4qIXOXCMxBhrwjoaUELsd-I6g,4538
 datahub/metadata/schemas/VersionInfo.avsc,sha256=9gMcZ8tjuhgcZiq2gOAp_EOV9q9jvuOgfph6m6v_X7c,1189
 datahub/metadata/schemas/VersionProperties.avsc,sha256=ME8V01JzG8lEsLXgYWnSYCehmpPcvv1UbE5Y8-8Ys9k,8022
@@ -1101,8 +1101,8 @@ datahub_provider/operators/datahub_assertion_operator.py,sha256=uvTQ-jk2F0sbqqxp
 datahub_provider/operators/datahub_assertion_sensor.py,sha256=lCBj_3x1cf5GMNpHdfkpHuyHfVxsm6ff5x2Z5iizcAo,140
 datahub_provider/operators/datahub_operation_operator.py,sha256=aevDp2FzX7FxGlXrR0khoHNbxbhKR2qPEX5e8O2Jyzw,174
 datahub_provider/operators/datahub_operation_sensor.py,sha256=8fcdVBCEPgqy1etTXgLoiHoJrRt_nzFZQMdSzHqSG7M,168
-acryl_datahub-1.2.0.2rc3.dist-info/METADATA,sha256=ixYC_JxEXERi_Ik-6tXxNn5W0EMN_poTxliZ99bCS6o,182014
-acryl_datahub-1.2.0.2rc3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-acryl_datahub-1.2.0.2rc3.dist-info/entry_points.txt,sha256=bnGf6eX9UhiW8yVHtt6MJCVcmLErvrVQxTJAayA-PKc,9885
-acryl_datahub-1.2.0.2rc3.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-1.2.0.2rc3.dist-info/RECORD,,
+acryl_datahub-1.2.0.3rc1.dist-info/METADATA,sha256=zo0PhZMaumsiXe8Vq8ud1VQPSZWVspaAuUIx1FoCk9s,182014
+acryl_datahub-1.2.0.3rc1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+acryl_datahub-1.2.0.3rc1.dist-info/entry_points.txt,sha256=bnGf6eX9UhiW8yVHtt6MJCVcmLErvrVQxTJAayA-PKc,9885
+acryl_datahub-1.2.0.3rc1.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-1.2.0.3rc1.dist-info/RECORD,,

datahub/_version.py

@@ -1,6 +1,6 @@
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "1.2.0.2rc3"
+__version__ = "1.2.0.3rc1"
 
 
 def is_dev_mode() -> bool:

datahub/ingestion/source/sql/vertica.py

@@ -4,6 +4,7 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Tuple, Union
 
 import pydantic
+import pytest
 from pydantic.class_validators import validator
 from vertica_sqlalchemy_dialect.base import VerticaInspector
 
@@ -55,6 +56,8 @@ from datahub.utilities import config_clean
 
 if TYPE_CHECKING:
     from datahub.ingestion.source.ge_data_profiler import GEProfilerRequest
+
+pytestmark = pytest.mark.integration_batch_4
 logger: logging.Logger = logging.getLogger(__name__)
 
 

datahub/ingestion/source/sql_queries.py

@@ -2,12 +2,13 @@ import json
 import logging
 import os
 from dataclasses import dataclass
-from datetime import datetime, timezone
+from datetime import datetime
 from functools import partial
-from typing import Iterable, List, Optional, Union
+from typing import ClassVar, Iterable, List, Optional, Union
 
-from pydantic import Field
+from pydantic import BaseModel, Field, validator
 
+from datahub.configuration.datetimes import parse_user_datetime
 from datahub.configuration.source_common import (
     EnvConfigMixin,
     PlatformInstanceConfigMixin,
@@ -35,7 +36,7 @@ from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
 from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
-from datahub.metadata.urns import CorpUserUrn
+from datahub.metadata.urns import CorpUserUrn, DatasetUrn
 from datahub.sql_parsing.schema_resolver import SchemaResolver
 from datahub.sql_parsing.sql_parsing_aggregator import (
     KnownQueryLineageInfo,
@@ -208,19 +209,40 @@ class SqlQueriesSource(Source):
     def _add_query_to_aggregator(self, query_entry: "QueryEntry") -> None:
         """Add a query to the SQL parsing aggregator."""
         try:
-            # If we have explicit lineage, use it directly
-            if query_entry.upstream_tables or query_entry.downstream_tables:
+            # If we have both upstream and downstream tables, use explicit lineage
+            if query_entry.upstream_tables and query_entry.downstream_tables:
                 logger.debug("Using explicit lineage from query file")
                 for downstream_table in query_entry.downstream_tables:
                     known_lineage = KnownQueryLineageInfo(
                         query_text=query_entry.query,
-                        downstream=downstream_table,
-                        upstreams=query_entry.upstream_tables,
+                        downstream=str(downstream_table),
+                        upstreams=[str(urn) for urn in query_entry.upstream_tables],
                         timestamp=query_entry.timestamp,
                         session_id=query_entry.session_id,
                     )
                     self.aggregator.add_known_query_lineage(known_lineage)
             else:
+                # Warn if only partial lineage information is provided
+                # XOR: true if exactly one of upstream_tables or downstream_tables is provided
+                if bool(query_entry.upstream_tables) ^ bool(
+                    query_entry.downstream_tables
+                ):
+                    query_preview = (
+                        query_entry.query[:150] + "..."
+                        if len(query_entry.query) > 150
+                        else query_entry.query
+                    )
+                    missing_upstream = (
+                        "Missing upstream. " if not query_entry.upstream_tables else ""
+                    )
+                    missing_downstream = (
+                        "Missing downstream. "
+                        if not query_entry.downstream_tables
+                        else ""
+                    )
+                    logger.info(
+                        f"Only partial lineage information provided, falling back to SQL parsing for complete lineage detection. {missing_upstream}{missing_downstream}Query: {query_preview}"
+                    )
                 # No explicit lineage, rely on parsing
                 observed_query = ObservedQuery(
                     query=query_entry.query,
@@ -243,46 +265,66 @@ class SqlQueriesSource(Source):
             )
 
 
-@dataclass
-class QueryEntry:
+class QueryEntry(BaseModel):
     query: str
-    timestamp: Optional[datetime]
-    user: Optional[CorpUserUrn]
-    operation_type: Optional[str]
-    downstream_tables: List[str]
-    upstream_tables: List[str]
+    timestamp: Optional[datetime] = None
+    user: Optional[CorpUserUrn] = None
+    operation_type: Optional[str] = None
+    downstream_tables: List[DatasetUrn] = Field(default_factory=list)
+    upstream_tables: List[DatasetUrn] = Field(default_factory=list)
     session_id: Optional[str] = None
 
+    # Validation context for URN creation
+    _validation_context: ClassVar[Optional[SqlQueriesSourceConfig]] = None
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @validator("timestamp", pre=True)
+    def parse_timestamp(cls, v):
+        return None if v is None else parse_user_datetime(str(v))
+
+    @validator("user", pre=True)
+    def parse_user(cls, v):
+        if v is None:
+            return None
+
+        return v if isinstance(v, CorpUserUrn) else CorpUserUrn(v)
+
+    @validator("downstream_tables", "upstream_tables", pre=True)
+    def parse_tables(cls, v):
+        if not v:
+            return []
+
+        result = []
+        for item in v:
+            if isinstance(item, DatasetUrn):
+                result.append(item)
+            elif isinstance(item, str):
+                # Skip empty/whitespace-only strings
+                if item and item.strip():
+                    # Convert to URN using validation context
+                    assert cls._validation_context, (
+                        "Validation context must be set for URN creation"
+                    )
+                    urn_string = make_dataset_urn_with_platform_instance(
+                        name=item,
+                        platform=cls._validation_context.platform,
+                        platform_instance=cls._validation_context.platform_instance,
+                        env=cls._validation_context.env,
+                    )
+                    result.append(DatasetUrn.from_string(urn_string))
+
+        return result
+
     @classmethod
     def create(
         cls, entry_dict: dict, *, config: SqlQueriesSourceConfig
     ) -> "QueryEntry":
-        return cls(
-            query=entry_dict["query"],
-            timestamp=(
-                datetime.fromtimestamp(entry_dict["timestamp"], tz=timezone.utc)
-                if "timestamp" in entry_dict
-                else None
-            ),
-            user=CorpUserUrn(entry_dict["user"]) if "user" in entry_dict else None,
-            operation_type=entry_dict.get("operation_type"),
-            downstream_tables=[
-                make_dataset_urn_with_platform_instance(
-                    name=table,
-                    platform=config.platform,
-                    platform_instance=config.platform_instance,
-                    env=config.env,
-                )
-                for table in entry_dict.get("downstream_tables", [])
-            ],
-            upstream_tables=[
-                make_dataset_urn_with_platform_instance(
-                    name=table,
-                    platform=config.platform,
-                    platform_instance=config.platform_instance,
-                    env=config.env,
-                )
-                for table in entry_dict.get("upstream_tables", [])
-            ],
-            session_id=entry_dict.get("session_id"),
-        )
+        """Create QueryEntry from dict with config context."""
+        # Set validation context for URN creation
+        cls._validation_context = config
+        try:
+            return cls.parse_obj(entry_dict)
+        finally:
+            cls._validation_context = None

datahub/ingestion/source/unity/proxy.py

@@ -4,8 +4,9 @@ Manage the communication with DataBricks Server and provide equivalent dataclass
 
 import dataclasses
 import logging
+from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
-from typing import Any, Dict, Iterable, List, Optional, Union, cast
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Union, cast
 from unittest.mock import patch
 
 import cachetools
@@ -28,6 +29,7 @@ from databricks.sdk.service.sql import (
 )
 from databricks.sdk.service.workspace import ObjectType
 from databricks.sql import connect
+from databricks.sql.types import Row
 
 from datahub._version import nice_version_name
 from datahub.api.entities.external.unity_catalog_external_entites import UnityCatalogTag
@@ -291,10 +293,59 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                 method, path, body={**body, "page_token": response["next_page_token"]}
             )
 
+    @cached(cachetools.FIFOCache(maxsize=100))
+    def get_catalog_column_lineage(self, catalog: str) -> Dict[str, Dict[str, dict]]:
+        """Get column lineage for all tables in a catalog."""
+        logger.info(f"Fetching column lineage for catalog: {catalog}")
+        try:
+            query = """
+                SELECT
+                    source_table_catalog, source_table_schema, source_table_name, source_column_name, source_type,
+                    target_table_schema, target_table_name, target_column_name,
+                    max(event_time)
+                FROM system.access.column_lineage
+                WHERE
+                    target_table_catalog = %s
+                    AND target_table_schema IS NOT NULL
+                    AND target_table_name IS NOT NULL
+                    AND target_column_name IS NOT NULL
+                    AND source_table_catalog IS NOT NULL
+                    AND source_table_schema IS NOT NULL
+                    AND source_table_name IS NOT NULL
+                    AND source_column_name IS NOT NULL
+                GROUP BY
+                    source_table_catalog, source_table_schema, source_table_name, source_column_name,  source_type,
+                    target_table_schema, target_table_name, target_column_name
+                """
+            rows = self._execute_sql_query(query, (catalog,))
+
+            result_dict: Dict[str, Dict[str, dict]] = {}
+            for row in rows:
+                result_dict.setdefault(row["target_table_schema"], {}).setdefault(
+                    row["target_table_name"], {}
+                ).setdefault(row["target_column_name"], []).append(
+                    # make fields look like the response from the older HTTP API
+                    {
+                        "catalog_name": row["source_table_catalog"],
+                        "schema_name": row["source_table_schema"],
+                        "table_name": row["source_table_name"],
+                        "name": row["source_column_name"],
+                    }
+                )
+
+            return result_dict
+        except Exception as e:
+            logger.warning(
+                f"Error getting column lineage for catalog {catalog}: {e}",
+                exc_info=True,
+            )
+            return {}
+
     def list_lineages_by_table(
         self, table_name: str, include_entity_lineage: bool
     ) -> dict:
         """List table lineage by table name."""
+        logger.debug(f"Getting table lineage for {table_name}")
         return self._workspace_client.api_client.do(  # type: ignore
             method="GET",
             path="/api/2.0/lineage-tracking/table-lineage",
@@ -304,13 +355,24 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             },
         )
 
-    def list_lineages_by_column(self, table_name: str, column_name: str) -> dict:
+    def list_lineages_by_column(self, table_name: str, column_name: str) -> list:
         """List column lineage by table name and column name."""
-        return self._workspace_client.api_client.do(  # type: ignore
-            "GET",
-            "/api/2.0/lineage-tracking/column-lineage",
-            body={"table_name": table_name, "column_name": column_name},
-        )
+        logger.debug(f"Getting column lineage for {table_name}.{column_name}")
+        try:
+            return (
+                self._workspace_client.api_client.do(  # type: ignore
+                    "GET",
+                    "/api/2.0/lineage-tracking/column-lineage",
+                    body={"table_name": table_name, "column_name": column_name},
+                ).get("upstream_cols")
+                or []
+            )
+        except Exception as e:
+            logger.warning(
+                f"Error getting column lineage on table {table_name}, column {column_name}: {e}",
+                exc_info=True,
+            )
+            return []
 
     def table_lineage(self, table: Table, include_entity_lineage: bool) -> None:
         if table.schema.catalog.type == CustomCatalogType.HIVE_METASTORE_CATALOG:
@@ -348,23 +410,51 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                 f"Error getting lineage on table {table.ref}: {e}", exc_info=True
             )
 
-    def get_column_lineage(self, table: Table, column_name: str) -> None:
+    def get_column_lineage(
+        self,
+        table: Table,
+        column_names: List[str],
+        *,
+        max_workers: Optional[int] = None,
+    ) -> None:
         try:
-            response: dict = self.list_lineages_by_column(
-                table_name=table.ref.qualified_table_name,
-                column_name=column_name,
-            )
-            for item in response.get("upstream_cols") or []:
-                table_ref = TableReference.create_from_lineage(
-                    item, table.schema.catalog.metastore
+            # use the newer system tables if we have a SQL warehouse, otherwise fall back
+            # and use the older (and much slower) HTTP API.
+            if self.warehouse_id:
+                lineage = (
+                    self.get_catalog_column_lineage(table.ref.catalog)
+                    .get(table.ref.schema, {})
+                    .get(table.ref.table, {})
                 )
-                if table_ref:
-                    table.upstreams.setdefault(table_ref, {}).setdefault(
-                        column_name, []
-                    ).append(item["name"])
+            else:
+                with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                    futures = [
+                        executor.submit(
+                            self.list_lineages_by_column,
+                            table.ref.qualified_table_name,
+                            column_name,
+                        )
+                        for column_name in column_names
+                    ]
+                lineage = {
+                    column_name: future.result()
+                    for column_name, future in zip(column_names, futures)
+                }
+
+            for column_name in column_names:
+                for item in lineage.get(column_name) or []:
+                    table_ref = TableReference.create_from_lineage(
+                        item,
+                        table.schema.catalog.metastore,
+                    )
+                    if table_ref:
+                        table.upstreams.setdefault(table_ref, {}).setdefault(
+                            column_name, []
+                        ).append(item["name"])
+
         except Exception as e:
             logger.warning(
-                f"Error getting column lineage on table {table.ref}, column {column_name}: {e}",
+                f"Error getting column lineage on table {table.ref}: {e}",
                 exc_info=True,
             )
 
@@ -504,14 +594,14 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             executed_as_user_name=info.executed_as_user_name,
         )
 
-    def _execute_sql_query(self, query: str) -> List[List[str]]:
+    def _execute_sql_query(self, query: str, params: Sequence[Any] = ()) -> List[Row]:
         """Execute SQL query using databricks-sql connector for better performance"""
         try:
             with (
                 connect(**self._sql_connection_params) as connection,
                 connection.cursor() as cursor,
             ):
-                cursor.execute(query)
+                cursor.execute(query, list(params))
                 return cursor.fetchall()
 
         except Exception as e:

datahub/ingestion/source/unity/source.py

@@ -1,7 +1,6 @@
 import logging
 import re
 import time
-from concurrent.futures import ThreadPoolExecutor
 from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
 from urllib.parse import urljoin
 
@@ -657,15 +656,13 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             if len(table.columns) > self.config.column_lineage_column_limit:
                 self.report.num_column_lineage_skipped_column_count += 1
 
-            with ThreadPoolExecutor(
-                max_workers=self.config.lineage_max_workers
-            ) as executor:
-                for column in table.columns[: self.config.column_lineage_column_limit]:
-                    executor.submit(
-                        self.unity_catalog_api_proxy.get_column_lineage,
-                        table,
-                        column.name,
-                    )
+            column_names = [
+                column.name
+                for column in table.columns[: self.config.column_lineage_column_limit]
+            ]
+            self.unity_catalog_api_proxy.get_column_lineage(
+                table, column_names, max_workers=self.config.lineage_max_workers
+            )
 
         return self._generate_lineage_aspect(self.gen_dataset_urn(table.ref), table)
 

datahub/metadata/schema.avsc

@@ -4319,6 +4319,14 @@
                   "doc": "The type of upstream entity"
                 },
                 {
+                  "Searchable": {
+                    "/*": {
+                      "fieldName": "fineGrainedUpstreams",
+                      "fieldType": "URN",
+                      "hasValuesFieldName": "hasFineGrainedUpstreams",
+                      "queryByDefault": false
+                    }
+                  },
                   "Urn": "Urn",
                   "urn_is_array": true,
                   "type": [
@@ -12875,6 +12883,7 @@
                                   "Searchable": {
                                     "fieldName": "upstreams",
                                     "fieldType": "URN",
+                                    "hasValuesFieldName": "hasUpstreams",
                                     "queryByDefault": false
                                   },
                                   "java": {

datahub/metadata/schemas/DataJobInputOutput.avsc

@@ -375,6 +375,14 @@
                 "doc": "The type of upstream entity"
               },
               {
+                "Searchable": {
+                  "/*": {
+                    "fieldName": "fineGrainedUpstreams",
+                    "fieldType": "URN",
+                    "hasValuesFieldName": "hasFineGrainedUpstreams",
+                    "queryByDefault": false
+                  }
+                },
                 "type": [
                   "null",
                   {

datahub/metadata/schemas/MetadataChangeEvent.avsc

@@ -3070,6 +3070,14 @@
                                   "doc": "The type of upstream entity"
                                 },
                                 {
+                                  "Searchable": {
+                                    "/*": {
+                                      "fieldName": "fineGrainedUpstreams",
+                                      "fieldType": "URN",
+                                      "hasValuesFieldName": "hasFineGrainedUpstreams",
+                                      "queryByDefault": false
+                                    }
+                                  },
                                   "type": [
                                     "null",
                                     {
@@ -3691,6 +3699,7 @@
                                 "Searchable": {
                                   "fieldName": "upstreams",
                                   "fieldType": "URN",
+                                  "hasValuesFieldName": "hasUpstreams",
                                   "queryByDefault": false
                                 },
                                 "java": {

datahub/metadata/schemas/UpstreamLineage.avsc

@@ -94,6 +94,7 @@
               "Searchable": {
                 "fieldName": "upstreams",
                 "fieldType": "URN",
+                "hasValuesFieldName": "hasUpstreams",
                 "queryByDefault": false
               },
               "java": {
@@ -199,6 +200,14 @@
                 "doc": "The type of upstream entity"
               },
               {
+                "Searchable": {
+                  "/*": {
+                    "fieldName": "fineGrainedUpstreams",
+                    "fieldType": "URN",
+                    "hasValuesFieldName": "hasFineGrainedUpstreams",
+                    "queryByDefault": false
+                  }
+                },
                 "type": [
                   "null",
                   {