acryl-datahub 1.2.0.11rc1__py3-none-any.whl → 1.2.0.11rc3__py3-none-any.whl

datahub/ingestion/source/looker/looker_lib_wrapper.py

@@ -2,6 +2,7 @@
 import json
 import logging
 import os
+from enum import Enum
 from functools import lru_cache
 from typing import Dict, List, MutableMapping, Optional, Sequence, Set, Union, cast
 
@@ -31,6 +32,14 @@ from datahub.configuration.common import ConfigurationError
 logger = logging.getLogger(__name__)
 
 
+class LookerQueryResponseFormat(Enum):
+    # result_format - Ref: https://cloud.google.com/looker/docs/reference/looker-api/latest/methods/Query/run_inline_query
+    JSON = "json"
+    SQL = (
+        "sql"  # Note: This does not execute the query, it only generates the SQL query.
+    )
+
+
 class TransportOptionsConfig(ConfigModel):
     timeout: int
     headers: MutableMapping[str, str]
@@ -69,6 +78,7 @@ class LookerAPIStats(BaseModel):
     search_looks_calls: int = 0
     search_dashboards_calls: int = 0
     all_user_calls: int = 0
+    generate_sql_query_calls: int = 0
 
 
 class LookerAPI:
@@ -170,17 +180,40 @@ class LookerAPI:
         logger.debug(f"Executing query {write_query}")
         self.client_stats.query_calls += 1
 
-        response_json = self.client.run_inline_query(
-            result_format="json",
+        response = self.client.run_inline_query(
+            result_format=LookerQueryResponseFormat.JSON.value,
             body=write_query,
             transport_options=self.transport_options,
         )
 
+        data = json.loads(response)
+
         logger.debug("=================Response=================")
-        data = json.loads(response_json)
         logger.debug("Length of response: %d", len(data))
         return data
 
+    def generate_sql_query(
+        self, write_query: WriteQuery, use_cache: bool = False
+    ) -> str:
+        """
+        Generates a SQL query string for a given WriteQuery.
+
+        Note: This does not execute the query, it only generates the SQL query.
+        """
+        logger.debug(f"Generating SQL query for {write_query}")
+        self.client_stats.generate_sql_query_calls += 1
+
+        response = self.client.run_inline_query(
+            result_format=LookerQueryResponseFormat.SQL.value,
+            body=write_query,
+            transport_options=self.transport_options,
+            cache=use_cache,
+        )
+
+        logger.debug("=================Response=================")
+        logger.debug("Length of SQL response: %d", len(response))
+        return str(response)
+
     def dashboard(self, dashboard_id: str, fields: Union[str, List[str]]) -> Dashboard:
         self.client_stats.dashboard_calls += 1
         return self.client.dashboard(

datahub/ingestion/source/looker/looker_view_id_cache.py

@@ -3,11 +3,11 @@ from typing import Dict, List, Optional
 
 from datahub.ingestion.source.looker.looker_common import LookerViewId, ViewFieldValue
 from datahub.ingestion.source.looker.looker_config import LookerConnectionDefinition
+from datahub.ingestion.source.looker.looker_constant import NAME
 from datahub.ingestion.source.looker.looker_dataclasses import LookerModel
 from datahub.ingestion.source.looker.looker_file_loader import LookerViewFileLoader
 from datahub.ingestion.source.looker.lookml_config import (
     BASE_PROJECT_NAME,
-    NAME,
     LookMLSourceReport,
 )
 

datahub/ingestion/source/looker/lookml_concept_context.py

@@ -12,12 +12,12 @@ from datahub.ingestion.source.looker.looker_constant import (
     DIMENSION_GROUPS,
     DIMENSIONS,
     MEASURES,
+    NAME,
 )
 from datahub.ingestion.source.looker.looker_dataclasses import LookerViewFile
 from datahub.ingestion.source.looker.looker_file_loader import LookerViewFileLoader
 from datahub.ingestion.source.looker.lookml_config import (
     DERIVED_VIEW_SUFFIX,
-    NAME,
     LookMLSourceReport,
 )
 from datahub.ingestion.source.looker.lookml_refinement import LookerRefinementResolver

datahub/ingestion/source/looker/lookml_config.py

@@ -28,11 +28,10 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionConfigBase,
 )
 from datahub.utilities.lossy_collections import LossyList
+from datahub.utilities.stats_collections import TopKDict, float_top_k_dict
 
 logger = logging.getLogger(__name__)
 
-NAME: str = "name"
-
 BASE_PROJECT_NAME = "__BASE"
 
 EXPLORE_FILE_EXTENSION = ".explore.lkml"
@@ -47,6 +46,9 @@ DERIVED_VIEW_PATTERN: str = r"\$\{([^}]*)\}"
 @dataclass
 class LookMLSourceReport(StaleEntityRemovalSourceReport):
     git_clone_latency: Optional[timedelta] = None
+    looker_query_api_latency_seconds: TopKDict[str, float] = dataclass_field(
+        default_factory=float_top_k_dict
+    )
     models_discovered: int = 0
     models_dropped: LossyList[str] = dataclass_field(default_factory=LossyList)
     views_discovered: int = 0
@@ -81,6 +83,11 @@ class LookMLSourceReport(StaleEntityRemovalSourceReport):
             self.api_stats = self._looker_api.compute_stats()
         return super().compute_stats()
 
+    def report_looker_query_api_latency(
+        self, view_urn: str, latency: timedelta
+    ) -> None:
+        self.looker_query_api_latency_seconds[view_urn] = latency.total_seconds()
+
 
 class LookMLSourceConfig(
     LookerCommonConfig, StatefulIngestionConfigBase, EnvConfigMixin
@@ -122,6 +129,16 @@ class LookMLSourceConfig(
         description="List of regex patterns for LookML views to include in the extraction.",
     )
     parse_table_names_from_sql: bool = Field(True, description="See note below.")
+    use_api_for_view_lineage: bool = Field(
+        False,
+        description="When enabled, uses Looker API to get SQL representation of views for lineage parsing instead of parsing LookML files directly. Requires 'api' configuration to be provided."
+        "Coverage of regex based lineage extraction has limitations, it only supportes ${TABLE}.column_name syntax, See (https://cloud.google.com/looker/docs/reference/param-field-sql#sql_for_dimensions) to"
+        "understand the other substitutions and cross-references allowed in LookML.",
+    )
+    use_api_cache_for_view_lineage: bool = Field(
+        False,
+        description="When enabled, uses Looker API server-side caching for query execution. Requires 'api' configuration to be provided.",
+    )
     api: Optional[LookerAPIConfig] = None
     project_name: Optional[str] = Field(
         None,
@@ -239,6 +256,17 @@ class LookMLSourceConfig(
             )
         return values
 
+    @root_validator(skip_on_failure=True)
+    def check_api_provided_for_view_lineage(cls, values):
+        """Validate that we must have an api credential to use Looker API for view's column lineage"""
+        if not values.get("api") and values.get("use_api_for_view_lineage"):
+            raise ValueError(
+                "API credential was not found. LookML source requires api credentials "
+                "for Looker to use Looker APIs for view's column lineage extraction."
+                "Set `use_api_for_view_lineage` to False to skip using Looker APIs."
+            )
+        return values
+
     @validator("base_folder", always=True)
     def check_base_folder_if_not_provided(
         cls, v: Optional[pydantic.DirectoryPath], values: Dict[str, Any]

datahub/ingestion/source/looker/lookml_refinement.py

@@ -4,10 +4,10 @@ import logging
 from typing import ClassVar, Dict, List, Set
 
 from datahub.ingestion.source.looker.looker_config import LookerConnectionDefinition
+from datahub.ingestion.source.looker.looker_constant import NAME
 from datahub.ingestion.source.looker.looker_dataclasses import LookerModel
 from datahub.ingestion.source.looker.looker_file_loader import LookerViewFileLoader
 from datahub.ingestion.source.looker.lookml_config import (
-    NAME,
     LookMLSourceConfig,
     LookMLSourceReport,
 )

datahub/ingestion/source/looker/lookml_source.py

@@ -142,6 +142,8 @@ class LookerView:
         ctx: PipelineContext,
         extract_col_level_lineage: bool = False,
         populate_sql_logic_in_descriptions: bool = False,
+        looker_client: Optional[LookerAPI] = None,
+        view_to_explore_map: Optional[Dict[str, str]] = None,
     ) -> Optional["LookerView"]:
         view_name = view_context.name()
 
@@ -160,6 +162,8 @@ class LookerView:
             config=config,
             ctx=ctx,
             reporter=reporter,
+            looker_client=looker_client,
+            view_to_explore_map=view_to_explore_map,
         )
 
         field_type_vs_raw_fields = OrderedDict(
@@ -705,6 +709,11 @@ class LookMLSource(StatefulIngestionSourceBase):
         # Value: Tuple(model file name, connection name)
         view_connection_map: Dict[str, Tuple[str, str]] = {}
 
+        # Map of view name to explore name for API-based view lineage
+        # A view can be referenced by multiple explores, we only need one of the explores to use Looker Query API
+        # Key: view_name, Value: explore_name
+        view_to_explore_map: Dict[str, str] = {}
+
         # The ** means "this directory and all subdirectories", and hence should
         # include all the files we want.
         model_files = sorted(
@@ -759,37 +768,37 @@ class LookMLSource(StatefulIngestionSourceBase):
                 )
             )
 
-            if self.source_config.emit_reachable_views_only:
-                model_explores_map = {d["name"]: d for d in model.explores}
-                for explore_dict in model.explores:
-                    try:
-                        if LookerRefinementResolver.is_refinement(explore_dict["name"]):
-                            continue
+            model_explores_map = {d["name"]: d for d in model.explores}
+            for explore_dict in model.explores:
+                try:
+                    if LookerRefinementResolver.is_refinement(explore_dict["name"]):
+                        continue
 
-                        explore_dict = (
-                            looker_refinement_resolver.apply_explore_refinement(
-                                explore_dict
-                            )
-                        )
-                        explore: LookerExplore = LookerExplore.from_dict(
-                            model_name,
-                            explore_dict,
-                            model.resolved_includes,
-                            viewfile_loader,
-                            self.reporter,
-                            model_explores_map,
-                        )
-                        if explore.upstream_views:
-                            for view_name in explore.upstream_views:
+                    explore_dict = looker_refinement_resolver.apply_explore_refinement(
+                        explore_dict
+                    )
+                    explore: LookerExplore = LookerExplore.from_dict(
+                        model_name,
+                        explore_dict,
+                        model.resolved_includes,
+                        viewfile_loader,
+                        self.reporter,
+                        model_explores_map,
+                    )
+                    if explore.upstream_views:
+                        for view_name in explore.upstream_views:
+                            if self.source_config.emit_reachable_views_only:
                                 explore_reachable_views.add(view_name.include)
-                    except Exception as e:
-                        self.reporter.report_warning(
-                            title="Failed to process explores",
-                            message="Failed to process explore dictionary.",
-                            context=f"Explore Details: {explore_dict}",
-                            exc=e,
-                        )
-                        logger.debug("Failed to process explore", exc_info=e)
+                            # Build view to explore mapping for API-based view lineage
+                            view_to_explore_map[view_name.include] = explore.name
+                except Exception as e:
+                    self.reporter.report_warning(
+                        title="Failed to process explores",
+                        message="Failed to process explore dictionary.",
+                        context=f"Explore Details: {explore_dict}",
+                        exc=e,
+                    )
+                    logger.debug("Failed to process explore", exc_info=e)
 
             processed_view_files = processed_view_map.setdefault(
                 model.connection, set()
@@ -878,6 +887,10 @@ class LookMLSource(StatefulIngestionSourceBase):
                                 populate_sql_logic_in_descriptions=self.source_config.populate_sql_logic_for_missing_descriptions,
                                 config=self.source_config,
                                 ctx=self.ctx,
+                                looker_client=self.looker_client,
+                                view_to_explore_map=view_to_explore_map
+                                if view_to_explore_map
+                                else None,
                             )
                         except Exception as e:
                             self.reporter.report_warning(