openmetadata-ingestion 1.9.3.0__py3-none-any.whl → 1.9.5.0__py3-none-any.whl

metadata/ingestion/source/database/saphana/cdata_parser.py

@@ -105,6 +105,7 @@ class CDATAKeys(Enum):
     BASE_MEASURES = "baseMeasures"
     MEASURE = "measure"
     MEASURE_MAPPING = "measureMapping"
+    CALCULATED_MEASURES = "calculatedMeasures"
     PRIVATE_MEASURE_GROUP = "privateMeasureGroup"
     LOGICAL_MODEL = "logicalModel"
     DATA_SOURCES = "dataSources"
@@ -136,6 +137,9 @@ class DataSourceMapping(BaseModel):
     parents: Annotated[
         List[ParentSource], Field(..., description="Parent Sources for a target col")
     ]
+    formula: Annotated[
+        Optional[str], Field(None, description="Formula used to derive the column")
+    ]
 
 
 class DataSource(BaseModel):
@@ -372,10 +376,10 @@ def _get_column_datasources(
 
 def _get_column_datasources_with_names(
     entry: ET.Element, datasource_map: Optional[DataSourceMap] = None
-) -> List[Tuple[DataSource, str]]:
+) -> List[Tuple[DataSource, str, Optional[str]]]:
     """
     Get the DataSource and the actual source column name after traversal.
-    Returns a list of tuples (DataSource, column_name).
+    Returns a list of tuples (DataSource, column_name, formula).
     """
     if (
         datasource_map
@@ -387,6 +391,7 @@ def _get_column_datasources_with_names(
             ds_origin_list=[],
             current_ds=datasource_map[entry.get(CDATAKeys.COLUMN_OBJECT_NAME.value)],
             datasource_map=datasource_map,
+            formula=None,
         )
         return ds_col_pairs
 
@@ -399,6 +404,7 @@ def _get_column_datasources_with_names(
                 source_type=ViewType.DATA_BASE_TABLE,
             ),
             entry.get(CDATAKeys.COLUMN_NAME.value),
+            None,  # No formula for direct table sources
         )
     ]
 
@@ -441,7 +447,7 @@ def _traverse_ds(
         else:
             logger.info(
                 f"Can't find mapping for column [{current_column}] in [{current_ds}]. "
-                f"We still have to implement `calculatedViewAttributes`."
+                f"This might be a constant or derived column."
             )
 
     return ds_origin_list
@@ -449,19 +455,20 @@ def _traverse_ds(
 
 def _traverse_ds_with_columns(
     current_column: str,
-    ds_origin_list: List[Tuple[DataSource, str]],
+    ds_origin_list: List[Tuple[DataSource, str, Optional[str]]],
     current_ds: DataSource,
     datasource_map: Optional[DataSourceMap],
-) -> List[Tuple[DataSource, str]]:
+    formula: Optional[str] = None,
+) -> List[Tuple[DataSource, str, Optional[str]]]:
     """
     Traverse the ds dict jumping from target -> source columns and getting the right parent.
     We keep inspecting current datasources and will append to the origin list the ones
-    that are not LOGICAL, along with the final column name.
-    Returns a list of tuples (DataSource, column_name).
+    that are not LOGICAL, along with the final column name and formula.
+    Returns a list of tuples (DataSource, column_name, formula).
     """
     if current_ds.source_type != ViewType.LOGICAL:
-        # This is a final datasource, append it with the current column name
-        ds_origin_list.append((current_ds, current_column))
+        # This is a final datasource, append it with the current column name and formula
+        ds_origin_list.append((current_ds, current_column, formula))
 
     else:
         # Based on our current column, find the parents from the mappings in the current_ds
@@ -470,6 +477,10 @@ def _traverse_ds_with_columns(
         )
 
         if current_ds_mapping:
+            # Use this layer's formula if we don't have one yet
+            if current_ds_mapping.formula and not formula:
+                formula = current_ds_mapping.formula
+
             for parent in current_ds_mapping.parents:
                 parent_ds = datasource_map.get(parent.parent)
                 if not parent_ds:
@@ -484,17 +495,44 @@ def _traverse_ds_with_columns(
                     ds_origin_list=ds_origin_list,
                     current_ds=parent_ds,
                     datasource_map=datasource_map,
+                    formula=formula,
                 )
         else:
-            # Current column not in mapping. This can happen for calculated view attributes
             logger.info(
                 f"Can't find mapping for column [{current_column}] in [{current_ds}]. "
-                f"We still have to implement `calculatedViewAttributes`."
+                f"This might be a constant or derived column."
             )
 
     return ds_origin_list
 
 
+def _get_formula_from_logical_mapping(
+    entry: Optional[ET.Element], datasource_map: Optional[DataSourceMap]
+) -> Optional[str]:
+    """Extract formula from logical datasource mapping if it exists."""
+    if not entry or not datasource_map:
+        return None
+
+    column_object_name = entry.get(CDATAKeys.COLUMN_OBJECT_NAME.value)
+    column_name = entry.get(CDATAKeys.COLUMN_NAME.value)
+
+    if not column_object_name or not column_name:
+        return None
+
+    datasource = datasource_map.get(column_object_name)
+    if not datasource:
+        return None
+
+    if datasource.source_type != ViewType.LOGICAL or not datasource.mapping:
+        return None
+
+    mapping = datasource.mapping.get(column_name)
+    if not mapping:
+        return None
+
+    return mapping.formula
+
+
 def _read_attributes(
     tree: ET.Element, ns: dict, datasource_map: Optional[DataSourceMap] = None
 ) -> ParsedLineage:
@@ -506,8 +544,9 @@ def _read_attributes(
 
     for attribute in attribute_list.findall(CDATAKeys.ATTRIBUTE.value, ns):
         key_mapping = attribute.find(CDATAKeys.KEY_MAPPING.value, ns)
+        target_name = attribute.get(CDATAKeys.ID.value)
 
-        # Get the actual source datasources and their column names
+        # Get the actual source datasources, column names, and formulas
         data_sources_with_columns = _get_column_datasources_with_names(
             entry=key_mapping, datasource_map=datasource_map
         )
@@ -517,7 +556,8 @@ def _read_attributes(
                 ColumnMapping(
                     data_source=ds_info[0],  # The datasource
                     sources=[ds_info[1]],  # The actual source column name
-                    target=attribute.get(CDATAKeys.ID.value),
+                    target=target_name,
+                    formula=ds_info[2],  # Formula from traversal (if any)
                 )
                 for ds_info in data_sources_with_columns
             ]
@@ -541,16 +581,39 @@ def _read_calculated_attributes(
         return lineage
 
     for calculated_attr in calculated_attrs.findall(key.value, ns):
-        formula = (
-            calculated_attr.find(CDATAKeys.KEY_CALCULATION.value, ns)
-            .find(CDATAKeys.FORMULA.value, ns)
-            .text
-        )
-        lineage += _explode_formula(
-            target=calculated_attr.get(CDATAKeys.ID.value),
-            formula=formula,
-            base_lineage=base_lineage,
-        )
+        key_calc = calculated_attr.find(CDATAKeys.KEY_CALCULATION.value, ns)
+        if key_calc is not None:
+            formula_elem = key_calc.find(CDATAKeys.FORMULA.value, ns)
+            if formula_elem is not None and formula_elem.text:
+                lineage += _explode_formula(
+                    target=calculated_attr.get(CDATAKeys.ID.value),
+                    formula=formula_elem.text,
+                    base_lineage=base_lineage,
+                )
+
+    return lineage
+
+
+def _read_calculated_measures(
+    tree: ET.Element,
+    ns: dict,
+    base_lineage: ParsedLineage,
+) -> ParsedLineage:
+    """Compute the lineage based on the calculated measures"""
+    lineage = ParsedLineage()
+
+    calculated_measures = tree.find(CDATAKeys.CALCULATED_MEASURES.value, ns)
+    if not calculated_measures:
+        return lineage
+
+    for measure in calculated_measures.findall(CDATAKeys.MEASURE.value, ns):
+        formula_elem = measure.find(CDATAKeys.FORMULA.value, ns)
+        if formula_elem is not None and formula_elem.text:
+            lineage += _explode_formula(
+                target=measure.get(CDATAKeys.ID.value),
+                formula=formula_elem.text,
+                base_lineage=base_lineage,
+            )
 
     return lineage
 
@@ -573,8 +636,9 @@ def _read_base_measures(
 
     for measure in base_measures.findall(CDATAKeys.MEASURE.value, ns):
         measure_mapping = measure.find(CDATAKeys.MEASURE_MAPPING.value, ns)
+        target_name = measure.get(CDATAKeys.ID.value)
 
-        # Get the actual source datasources and their column names
+        # Get the actual source datasources, column names, and formulas
         data_sources_with_columns = _get_column_datasources_with_names(
             entry=measure_mapping, datasource_map=datasource_map
         )
@@ -584,7 +648,8 @@ def _read_base_measures(
                 ColumnMapping(
                     data_source=ds_info[0],  # The datasource
                     sources=[ds_info[1]],  # The actual source column name
-                    target=measure.get(CDATAKeys.ID.value),
+                    target=target_name,
+                    formula=ds_info[2],  # Formula from traversal (if any)
                 )
                 for ds_info in data_sources_with_columns
             ]
@@ -605,10 +670,16 @@ def _explode_formula(
     Returns:
         Parsed Lineage from the formula
     """
-    column_ds = {
-        match.group(1): base_lineage.find_target(match.group(1)).data_source
-        for match in FORMULA_PATTERN.finditer(formula)
-    }
+    column_ds = {}
+    for match in FORMULA_PATTERN.finditer(formula):
+        col_name = match.group(1)
+        mapping = base_lineage.find_target(col_name)
+        if mapping:
+            column_ds[col_name] = mapping.data_source
+
+    # If no columns found in base_lineage, it might be a constant formula
+    if not column_ds:
+        return ParsedLineage()
 
     # Group every datasource (key) with a list of the involved columns (values)
     ds_columns = defaultdict(list)
@@ -694,17 +765,42 @@ def _(cdata: str) -> ParsedLineage:
     attribute_lineage = _read_attributes(
         tree=logical_model, ns=ns, datasource_map=datasource_map
     )
-    calculated_attrs_lineage = _read_calculated_attributes(
+
+    base_measure_lineage = _read_base_measures(
+        tree=logical_model, ns=ns, datasource_map=datasource_map
+    )
+
+    # Combine base attributes and measures for calculated columns
+    combined_base_lineage = attribute_lineage + base_measure_lineage
+
+    # Read calculated attributes from calculationViews (if they exist)
+    cv_calculated_attrs_lineage = _read_calculated_attributes(
         tree=tree,
         ns=ns,
-        base_lineage=attribute_lineage,
+        base_lineage=combined_base_lineage,
         key=CalculatedAttrKey.CALCULATED_VIEW_ATTRIBUTE,
     )
-    base_measure_lineage = _read_base_measures(
-        tree=logical_model, ns=ns, datasource_map=datasource_map
+
+    # Read calculated attributes from logical model
+    logical_calculated_attrs_lineage = _read_calculated_attributes(
+        tree=logical_model,
+        ns=ns,
+        base_lineage=combined_base_lineage,
+        key=CalculatedAttrKey.CALCULATED_ATTRIBUTE,
     )
 
-    return attribute_lineage + calculated_attrs_lineage + base_measure_lineage
+    # Read calculated measures from logical model
+    calculated_measures_lineage = _read_calculated_measures(
+        tree=logical_model, ns=ns, base_lineage=combined_base_lineage
+    )
+
+    return (
+        attribute_lineage
+        + cv_calculated_attrs_lineage
+        + logical_calculated_attrs_lineage
+        + base_measure_lineage
+        + calculated_measures_lineage
+    )
 
 
 def _parse_cv_data_sources(tree: ET.Element, ns: dict) -> DataSourceMap:
@@ -812,11 +908,13 @@ def _build_mappings(calculation_view: ET.Element, ns: dict) -> List[DataSourceMa
     """
 
     input_mappings = _build_input_mappings(calculation_view=calculation_view, ns=ns)
-    # calculated_view_attrs = _build_cv_attributes(
-    #     calculation_view=calculation_view, ns=ns, input_mappings=input_mappings
-    # )
+    calculated_view_attrs = _build_cv_attributes(
+        calculation_view=calculation_view, ns=ns, input_mappings=input_mappings
+    )
 
-    return input_mappings
+    # Combine input mappings and calculated view attributes
+    all_mappings = input_mappings + calculated_view_attrs
+    return all_mappings
 
 
 def _build_input_mappings(
@@ -875,6 +973,8 @@ def _build_cv_attributes(
     if view_attrs is None:
         return mappings
 
+    cv_id = calculation_view.get(CDATAKeys.ID.value)
+
     for view_attr in view_attrs.findall(CDATAKeys.CALCULATION_VIEW_ATTRIBUTE.value, ns):
         formula = (
             view_attr.find(CDATAKeys.FORMULA.value, ns).text
@@ -886,25 +986,28 @@ def _build_cv_attributes(
             continue
 
         involved_columns = FORMULA_PATTERN.findall(formula)
+        # For calculated columns, all source columns should come from the same calculation view
+        # where the formula is defined
+        parents = []
         for col in involved_columns:
-            # Find the mapping for the involved column
-            mapping = next(
-                (mapping for mapping in input_mappings if mapping.target == col), None
-            )
-            if not mapping:
-                logger.debug(
-                    f"Can't find mapping for column [{col}] in [{input_mappings}]"
+            # The source columns for the formula are in the same calculation view
+            parents.append(
+                ParentSource(
+                    source=col,
+                    parent=cv_id,  # The parent is the current calculation view
                 )
-                continue
+            )
 
+        if parents:
             mappings.append(
                 DataSourceMapping(
                     target=view_attr.get(CDATAKeys.ID.value),
-                    parents=mapping.parents,
+                    parents=parents,
+                    formula=formula,
                 )
             )
 
-    return _group_mappings(mappings)
+    return mappings
 
 
 def _group_mappings(mappings: List[DataSourceMapping]) -> List[DataSourceMapping]:

metadata/ingestion/source/database/unitycatalog/lineage.py

@@ -93,29 +93,36 @@ class UnitycatalogLineageSource(Source):
     def _get_lineage_details(
         self, from_table: Table, to_table: Table, databricks_table_fqn: str
     ) -> Optional[LineageDetails]:
-        col_lineage = []
-        for column in to_table.columns:
-            column_streams = self.client.get_column_lineage(
-                databricks_table_fqn, column_name=column.name.root
-            )
-            from_columns = []
-            for col in column_streams.upstream_cols:
-                col_fqn = get_column_fqn(from_table, col.name)
-                if col_fqn:
-                    from_columns.append(col_fqn)
-
-            if from_columns:
-                col_lineage.append(
-                    ColumnLineage(
-                        fromColumns=from_columns,
-                        toColumn=column.fullyQualifiedName.root,
+        try:
+            col_lineage = []
+            for column in to_table.columns:
+                column_streams = self.client.get_column_lineage(
+                    databricks_table_fqn, column_name=column.name.root
+                )
+                from_columns = []
+                for col in column_streams.upstream_cols:
+                    col_fqn = get_column_fqn(from_table, col.name)
+                    if col_fqn:
+                        from_columns.append(col_fqn)
+
+                if from_columns:
+                    col_lineage.append(
+                        ColumnLineage(
+                            fromColumns=from_columns,
+                            toColumn=column.fullyQualifiedName.root,
+                        )
                     )
+            if col_lineage:
+                return LineageDetails(
+                    columnsLineage=col_lineage, source=LineageSource.QueryLineage
                 )
-        if col_lineage:
-            return LineageDetails(
-                columnsLineage=col_lineage, source=LineageSource.QueryLineage
+            return None
+        except Exception as exc:
+            logger.debug(
+                f"Error computing column lineage for {to_table.fullyQualifiedName.root} - {exc}"
             )
-        return None
+            logger.debug(traceback.format_exc())
+            return None
 
     def _handle_upstream_table(
         self,
@@ -157,6 +164,12 @@ class UnitycatalogLineageSource(Source):
                             )
                         ),
                     )
+                else:
+                    logger.debug(
+                        f"Unable to find upstream entity for "
+                        f"{upstream_table.catalog_name}.{upstream_table.schema_name}.{upstream_table.name}"
+                        f" -> {databricks_table_fqn}"
+                    )
             except Exception:
                 logger.debug(
                     "Error while processing lineage for "

metadata/profiler/interface/sqlalchemy/databricks/profiler_interface.py

@@ -85,10 +85,28 @@ class DatabricksProfilerInterface(SQAProfilerInterface):
             result += "`.`".join(splitted_result)
         return result
 
+    def visit_table(self, *args, **kwargs):
+        result = super(  # pylint: disable=bad-super-call
+            HiveCompiler, self
+        ).visit_table(*args, **kwargs)
+        # Handle table references with hyphens in database/schema names
+        # Format: `database`.`schema`.`table` for Unity Catalog/Databricks
+        if "." in result and not result.startswith("`"):
+            parts = result.split(".")
+            quoted_parts = []
+            for part in parts:
+                if "-" in part and not (part.startswith("`") and part.endswith("`")):
+                    quoted_parts.append(f"`{part}`")
+                else:
+                    quoted_parts.append(part)
+            result = ".".join(quoted_parts)
+        return result
+
     def __init__(self, service_connection_config, **kwargs):
         super().__init__(service_connection_config=service_connection_config, **kwargs)
         self.set_catalog(self.session)
         HiveCompiler.visit_column = DatabricksProfilerInterface.visit_column
+        HiveCompiler.visit_table = DatabricksProfilerInterface.visit_table
 
     def _get_struct_columns(self, columns: List[OMColumn], parent: str):
         """Get struct columns"""

metadata/profiler/interface/sqlalchemy/unity_catalog/profiler_interface.py

@@ -14,6 +14,8 @@ Interfaces with database for all database engine
 supporting sqlalchemy abstraction layer
 """
 
+from sqlalchemy import event
+from sqlalchemy.orm import scoped_session, sessionmaker
 
 from metadata.ingestion.source.database.databricks.connection import (
     get_connection as databricks_get_connection,
@@ -26,5 +28,17 @@ from metadata.profiler.interface.sqlalchemy.databricks.profiler_interface import
 class UnityCatalogProfilerInterface(DatabricksProfilerInterface):
     def create_session(self):
         self.connection = databricks_get_connection(self.service_connection_config)
-        super().create_session()
-        self.set_catalog(self.session)
+
+        # Create custom session factory with after_begin event to set catalog
+        session_maker = sessionmaker(bind=self.connection)
+
+        @event.listens_for(session_maker, "after_begin")
+        def set_catalog(session, transaction, connection):
+            # Safely quote the catalog name to prevent SQL injection
+            quoted_catalog = connection.dialect.identifier_preparer.quote(
+                self.service_connection_config.catalog
+            )
+            connection.execute(f"USE CATALOG {quoted_catalog};")
+
+        self.session_factory = scoped_session(session_maker)
+        self.session = self.session_factory()

metadata/profiler/metrics/system/databricks/system.py

@@ -24,7 +24,7 @@ SYSTEM_QUERY = textwrap.dedent(
         '{database}' AS database,
         '{schema}' AS schema,
         '{table}' AS table
-    FROM (DESCRIBE HISTORY {database}.{schema}.{table})
+    FROM (DESCRIBE HISTORY `{database}`.`{schema}`.`{table}`)
     WHERE operation IN ({operations}) AND timestamp > DATEADD(day, -1, CURRENT_TIMESTAMP())
     """
 )

metadata/sampler/sqlalchemy/databricks/sampler.py

@@ -11,7 +11,8 @@
 """
 Helper module to handle data sampling for the profiler
 """
-from sqlalchemy import Column, text
+from sqlalchemy import Column, event, text
+from sqlalchemy.orm import scoped_session, sessionmaker
 
 from metadata.ingestion.source.database.databricks.connection import (
     get_connection as databricks_get_connection,
@@ -25,6 +26,17 @@ class DatabricksSamplerInterface(SQASampler):
         """Initialize with a single Databricks connection"""
         super().__init__(*args, **kwargs)
         self.connection = databricks_get_connection(self.service_connection_config)
+        session_maker = sessionmaker(bind=self.connection)
+
+        @event.listens_for(session_maker, "after_begin")
+        def set_catalog(session, transaction, connection):
+            # Safely quote the catalog name to prevent SQL injection
+            quoted_catalog = connection.dialect.identifier_preparer.quote(
+                self.service_connection_config.catalog
+            )
+            connection.execute(f"USE CATALOG {quoted_catalog};")
+
+        self.session_factory = scoped_session(session_maker)
 
     def get_client(self):
         """client is the session for SQA"""

metadata/sampler/sqlalchemy/unitycatalog/sampler.py

@@ -14,47 +14,14 @@ Interfaces with database for all database engine
 supporting sqlalchemy abstraction layer
 """
 
-from sqlalchemy import Column, text
 
-from metadata.ingestion.source.database.databricks.connection import (
-    get_connection as databricks_get_connection,
-)
-from metadata.profiler.orm.types.custom_array import CustomArray
-from metadata.sampler.sqlalchemy.sampler import SQASampler
+from metadata.sampler.sqlalchemy.databricks.sampler import DatabricksSamplerInterface
 
 
-class UnityCatalogSamplerInterface(SQASampler):
+class UnityCatalogSamplerInterface(DatabricksSamplerInterface):
+    """
+    Unity Catalog Sampler Interface
+    """
+
     def __init__(self, *args, **kwargs):
-        """Initialize with a single Databricks connection"""
         super().__init__(*args, **kwargs)
-        self.connection = databricks_get_connection(self.service_connection_config)
-
-    def get_client(self):
-        """client is the session for SQA"""
-        client = super().get_client()
-        self.set_catalog(client)
-        return client
-
-    def _handle_array_column(self, column: Column) -> bool:
-        """Check if a column is an array type"""
-        return isinstance(column.type, CustomArray)
-
-    def _get_slice_expression(self, column: Column):
-        """Generate SQL expression to slice array elements at query level
-
-        Args:
-            column_name: Name of the column
-            max_elements: Maximum number of elements to extract
-
-        Returns:
-            SQL expression string for array slicing
-        """
-        max_elements = self._get_max_array_elements()
-        return text(
-            f"""
-        CASE 
-            WHEN `{column.name}` IS NULL THEN NULL
-            ELSE slice(`{column.name}`, 1, {max_elements})
-        END AS `{column._label}`
-        """
-        )

metadata/utils/bigquery_utils.py

@@ -14,9 +14,7 @@ Utils module of BigQuery
 """
 
 from copy import deepcopy
-from typing import List, Optional
-
-from google.cloud import bigquery
+from typing import TYPE_CHECKING, List, Optional
 
 from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
     BigQueryConnection,
@@ -34,6 +32,9 @@ from metadata.utils.credentials import (
     get_gcp_impersonate_credentials,
 )
 
+if TYPE_CHECKING:
+    from google.cloud import bigquery
+
 
 def get_bigquery_client(
     project_id: Optional[str] = None,
@@ -42,7 +43,7 @@ def get_bigquery_client(
     quota_project_id: Optional[str] = None,
     scopes: Optional[List[str]] = None,
     lifetime: Optional[int] = 3600,
-) -> bigquery.Client:
+) -> "bigquery.Client":
     """Get a BigQuery client
 
     Args:
@@ -65,6 +66,8 @@ def get_bigquery_client(
             scopes=scopes,
             lifetime=lifetime,
         )
+    from google.cloud import bigquery  # pylint: disable=import-outside-toplevel
+
     return bigquery.Client(
         credentials=credentials, project=project_id, location=location
     )

metadata/utils/logger.py

@@ -212,6 +212,8 @@ def get_log_name(record: Entity) -> Optional[str]:
     try:
         if hasattr(record, "name"):
             return f"{type(record).__name__} [{getattr(record, 'name').root}]"
+        if hasattr(record, "table") and hasattr(record.table, "name"):
+            return f"{type(record).__name__} [{record.table.name.root}]"
         return f"{type(record).__name__} [{record.entity.name.root}]"
     except Exception:
         return str(record)