openmetadata-ingestion 1.9.12.0__py3-none-any.whl → 1.9.13.0__py3-none-any.whl

metadata/ingestion/source/dashboard/powerbi/metadata.py

@@ -54,6 +54,8 @@ from metadata.generated.schema.type.basic import (
 from metadata.generated.schema.type.entityReferenceList import EntityReferenceList
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.lineage.models import Dialect
+from metadata.ingestion.lineage.parser import LineageParser
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.ometa.utils import model_str
 from metadata.ingestion.source.dashboard.dashboard_service import DashboardServiceSource
@@ -80,6 +82,7 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 OWNER_ACCESS_RIGHTS_KEYWORDS = ["owner", "write", "admin"]
+SNOWFLAKE_QUERY_EXPRESSION_KW = "Value.NativeQuery(Snowflake.Databases("
 
 
 class PowerbiSource(DashboardServiceSource):
@@ -302,6 +305,24 @@ class PowerbiSource(DashboardServiceSource):
             f"{workspace_id}/reports/{dashboard_id}/ReportSection?experience=power-bi"
         )
 
+    def _get_dataset_url(self, workspace_id: str, dataset_id: str) -> str:
+        """
+        Method to build the dataset url
+        """
+        return (
+            f"{clean_uri(self.service_connection.hostPort)}/groups/"
+            f"{workspace_id}/datasets/{dataset_id}?experience=power-bi"
+        )
+
+    def _get_dataflow_url(self, workspace_id: str, dataflow_id: str) -> str:
+        """
+        Method to build the dataset url
+        """
+        return (
+            f"{clean_uri(self.service_connection.hostPort)}/groups/"
+            f"{workspace_id}/dataflows/{dataflow_id}?experience=power-bi"
+        )
+
     def _get_chart_url(
         self, report_id: Optional[str], workspace_id: str, dashboard_id: str
     ) -> str:
@@ -545,9 +566,17 @@ class PowerbiSource(DashboardServiceSource):
                     if isinstance(dataset, Dataset):
                         data_model_type = DataModelType.PowerBIDataModel.value
                         datamodel_columns = self._get_column_info(dataset)
+                        source_url = self._get_dataset_url(
+                            workspace_id=self.context.get().workspace.id,
+                            dataset_id=dataset.id,
+                        )
                     elif isinstance(dataset, Dataflow):
                         data_model_type = DataModelType.PowerBIDataFlow.value
                         datamodel_columns = []
+                        source_url = self._get_dataflow_url(
+                            workspace_id=self.context.get().workspace.id,
+                            dataflow_id=dataset.id,
+                        )
                     else:
                         logger.warning(
                             f"Unknown dataset type: {type(dataset)}, name: {dataset.name}"
@@ -569,6 +598,7 @@ class PowerbiSource(DashboardServiceSource):
                         columns=datamodel_columns,
                         project=self.get_project_name(dashboard_details=dataset),
                         owners=self.get_owner_ref(dashboard_details=dataset),
+                        sourceUrl=SourceUrl(source_url),
                     )
                     yield Either(right=data_model_request)
                     self.register_record_datamodel(datamodel_request=data_model_request)
@@ -761,7 +791,7 @@ class PowerbiSource(DashboardServiceSource):
             logger.debug(traceback.format_exc())
         return None
 
-    def _parse_redshift_source(self, source_expression: str) -> Optional[dict]:
+    def _parse_redshift_source(self, source_expression: str) -> Optional[List[dict]]:
         try:
             db_match = re.search(
                 r'AmazonRedshift\.Database\("[^"]+","([^"]+)"\)', source_expression
@@ -778,20 +808,167 @@ class PowerbiSource(DashboardServiceSource):
                 table = schema_table_match[1] if len(schema_table_match) > 1 else None
 
             if table:  # atlease table should be fetched
-                return {"database": database, "schema": schema, "table": table}
+                return [{"database": database, "schema": schema, "table": table}]
             return None
         except Exception as exc:
             logger.debug(f"Error to parse redshift table source: {exc}")
             logger.debug(traceback.format_exc())
         return None
 
+    def _parse_snowflake_query_source(
+        self, source_expression: str
+    ) -> Optional[List[dict]]:
+        """
+        Parse snowflake query source
+        source expressions like `Value.NativeQuery(Snowflake.Databases())`
+        """
+        try:
+            logger.debug(
+                f"parsing source expression through query parser: {source_expression[:100]}"
+            )
+
+            # Look for SQL query after [Data],
+            # The pattern needs to handle the concatenated strings with & operators
+            m = re.search(
+                r'\[Data\],\s*"(.+?)"(?:,\s*null|\s*\))',
+                source_expression,
+                re.IGNORECASE | re.DOTALL,
+            )
+
+            if not m:
+                logger.debug("sql query not found in source expression")
+                return None
+
+            # Extract and clean the SQL query
+            sql_query = m.group(1).replace('""', '"')
+
+            # Handle PowerBI parameter concatenation (e.g., "& Database &")
+            # For now, replace parameter references with wildcards for parsing
+            sql_query_cleaned = re.sub(r'"?\s*&\s*\w+\s*&\s*"?\.?', "%.", sql_query)
+
+            logger.debug(f"Extracted SQL query: {sql_query}")
+            logger.debug(f"Cleaned SQL query: {sql_query_cleaned}")
+
+            if not sql_query_cleaned:
+                logger.debug("Empty SQL query after extraction")
+                return None
+
+            # Clean the query for parser
+            # 1. Replace % with a placeholder database name
+            parser_query = sql_query_cleaned.replace("%", "PLACEHOLDER_DB")
+
+            # 2. Remove PowerBI line feed markers #(lf) and clean up the query
+            parser_query = parser_query.replace("#(lf)", "\n")
+
+            # 3. Remove SQL comments that might cause issues (// style comments)
+            parser_query = re.sub(r"//[^\n]*", "", parser_query)
+
+            # 4. Clean up excessive whitespace
+            parser_query = re.sub(r"\s+", " ", parser_query).strip()
+
+            logger.debug(
+                f"Attempting LineageParser with cleaned query: {parser_query[:200]}"
+            )
+
+            try:
+                parser = LineageParser(
+                    parser_query, dialect=Dialect.SNOWFLAKE, timeout_seconds=30
+                )
+            except Exception as parser_exc:
+                logger.debug(f"LineageParser failed with error: {parser_exc}")
+                logger.debug(f"Failed query was: {parser_query[:200]}...")
+                return None
+
+            if parser.source_tables:
+                logger.debug(
+                    f"LineageParser found {len(parser.source_tables)} source table(s)"
+                )
+                for table in parser.source_tables:
+                    logger.debug(
+                        f"source table: {table.raw_name}, schema: {table.schema if hasattr(table, 'schema') else 'N/A'}"
+                    )
+                lineage_tables_list = []
+                for source_table in parser.source_tables:
+                    # source_table = parser.source_tables[0]
+
+                    # Extract database from schema's parent if it exists
+                    database = None
+                    schema = None
+
+                    if hasattr(source_table, "schema") and source_table.schema:
+                        # Log what we have in the schema object
+                        logger.debug(
+                            f"Schema object type: {type(source_table.schema)}, value: {source_table.schema}"
+                        )
+
+                        # Get schema as string first
+                        schema_str = (
+                            source_table.schema.raw_name
+                            if hasattr(source_table.schema, "raw_name")
+                            else str(source_table.schema)
+                        )
+
+                        # If schema contains dots, it might be database.schema format
+                        if "." in schema_str:
+                            parts = schema_str.split(".")
+                            if len(parts) == 2:
+                                # Format: database.schema
+                                # Check for placeholder (case insensitive)
+                                database = (
+                                    parts[0]
+                                    if parts[0].upper() != "PLACEHOLDER_DB"
+                                    else None
+                                )
+                                schema = parts[1]
+                            else:
+                                # Just use as is
+                                schema = schema_str
+                        else:
+                            schema = schema_str
+                            # Check if schema has a parent (database)
+                            if (
+                                hasattr(source_table.schema, "parent")
+                                and source_table.schema.parent
+                            ):
+                                database = (
+                                    source_table.schema.parent.raw_name
+                                    if hasattr(source_table.schema.parent, "raw_name")
+                                    else str(source_table.schema.parent)
+                                )
+
+                    # Filter out placeholder values (case insensitive)
+                    if database and database.upper() == "PLACEHOLDER_DB":
+                        database = None
+
+                    table = source_table.raw_name
+
+                    if table:
+                        logger.debug(f"tables found = {database}.{schema}.{table}")
+                        lineage_tables_list.append(
+                            {
+                                "database": database,
+                                "schema": schema,
+                                "table": table,
+                            }
+                        )
+                return lineage_tables_list
+            logger.debug("tables in query not found through parser")
+            return None
+        except Exception as exc:
+            logger.debug(f"Error parsing snowflake query source: {exc}")
+            logger.debug(traceback.format_exc())
+        return None
+
     def _parse_snowflake_source(
         self, source_expression: str, datamodel_entity: DashboardDataModel
-    ) -> Optional[dict]:
+    ) -> Optional[List[dict]]:
         try:
             if "Snowflake.Databases" not in source_expression:
                 # Not a snowflake valid expression
                 return None
+            if SNOWFLAKE_QUERY_EXPRESSION_KW in source_expression:
+                # snowflake query source identified
+                return self._parse_snowflake_query_source(source_expression)
             db_match = re.search(
                 r'\[Name=(?:"([^"]+)"|([^,]+)),Kind="Database"\]', source_expression
             )
@@ -811,11 +988,13 @@ class PowerbiSource(DashboardServiceSource):
             view = self._parse_snowflake_regex_exp(view_match, datamodel_entity)
 
             if table or view:  # atlease table or view should be fetched
-                return {
-                    "database": database,
-                    "schema": schema,
-                    "table": table if table else view,
-                }
+                return [
+                    {
+                        "database": database,
+                        "schema": schema,
+                        "table": table if table else view,
+                    }
+                ]
             return None
         except Exception as exc:
             logger.debug(f"Error to parse snowflake table source: {exc}")
@@ -824,29 +1003,29 @@ class PowerbiSource(DashboardServiceSource):
 
     def _parse_table_info_from_source_exp(
         self, table: PowerBiTable, datamodel_entity: DashboardDataModel
-    ) -> dict:
+    ) -> Optional[List[dict]]:
         try:
             if not isinstance(table.source, list):
-                return {}
+                return None
             source_expression = table.source[0].expression
             if not source_expression:
                 logger.debug(f"No source expression found for table: {table.name}")
-                return {}
+                return None
             # parse snowflake source
-            table_info = self._parse_snowflake_source(
+            table_info_list = self._parse_snowflake_source(
                 source_expression, datamodel_entity
             )
-            if isinstance(table_info, dict):
-                return table_info
+            if isinstance(table_info_list, List):
+                return table_info_list
             # parse redshift source
-            table_info = self._parse_redshift_source(source_expression)
-            if isinstance(table_info, dict):
-                return table_info
-            return {}
+            table_info_list = self._parse_redshift_source(source_expression)
+            if isinstance(table_info_list, List):
+                return table_info_list
+            return None
         except Exception as exc:
             logger.debug(f"Error to parse table source: {exc}")
             logger.debug(traceback.format_exc())
-        return {}
+        return None
 
     def _get_table_and_datamodel_lineage(
         self,
@@ -865,60 +1044,64 @@ class PowerbiSource(DashboardServiceSource):
         ) = self.parse_db_service_prefix(db_service_prefix)
 
         try:
-            table_info = self._parse_table_info_from_source_exp(table, datamodel_entity)
-            table_name = table_info.get("table") or table.name
-            schema_name = table_info.get("schema")
-            database_name = table_info.get("database")
-            if (
-                prefix_table_name
-                and table_name
-                and prefix_table_name.lower() != table_name.lower()
-            ):
-                logger.debug(
-                    f"Table {table_name} does not match prefix {prefix_table_name}"
-                )
-                return
+            table_info_list = self._parse_table_info_from_source_exp(
+                table, datamodel_entity
+            )
+            if isinstance(table_info_list, List):
+                for table_info in table_info_list:
+                    table_name = table_info.get("table") or table.name
+                    schema_name = table_info.get("schema")
+                    database_name = table_info.get("database")
+                    if (
+                        prefix_table_name
+                        and table_name
+                        and prefix_table_name.lower() != table_name.lower()
+                    ):
+                        logger.debug(
+                            f"Table {table_name} does not match prefix {prefix_table_name}"
+                        )
+                        return
 
-            if (
-                prefix_schema_name
-                and schema_name
-                and prefix_schema_name.lower() != schema_name.lower()
-            ):
-                logger.debug(
-                    f"Schema {table_info.get('schema')} does not match prefix {prefix_schema_name}"
-                )
-                return
+                    if (
+                        prefix_schema_name
+                        and schema_name
+                        and prefix_schema_name.lower() != schema_name.lower()
+                    ):
+                        logger.debug(
+                            f"Schema {table_info.get('schema')} does not match prefix {prefix_schema_name}"
+                        )
+                        return
 
-            if (
-                prefix_database_name
-                and database_name
-                and prefix_database_name.lower() != database_name.lower()
-            ):
-                logger.debug(
-                    f"Database {table_info.get('database')} does not match prefix {prefix_database_name}"
-                )
-                return
+                    if (
+                        prefix_database_name
+                        and database_name
+                        and prefix_database_name.lower() != database_name.lower()
+                    ):
+                        logger.debug(
+                            f"Database {table_info.get('database')} does not match prefix {prefix_database_name}"
+                        )
+                        return
 
-            fqn_search_string = build_es_fqn_search_string(
-                service_name=prefix_service_name or "*",
-                table_name=(prefix_table_name or table_name),
-                schema_name=(prefix_schema_name or schema_name),
-                database_name=(prefix_database_name or database_name),
-            )
-            table_entity = self.metadata.search_in_any_service(
-                entity_type=Table,
-                fqn_search_string=fqn_search_string,
-            )
-            if table_entity and datamodel_entity:
-                columns_list = [column.name for column in table.columns]
-                column_lineage = self._get_column_lineage(
-                    table_entity, datamodel_entity, columns_list
-                )
-                yield self._get_add_lineage_request(
-                    to_entity=datamodel_entity,
-                    from_entity=table_entity,
-                    column_lineage=column_lineage,
-                )
+                    fqn_search_string = build_es_fqn_search_string(
+                        service_name=prefix_service_name or "*",
+                        table_name=(prefix_table_name or table_name),
+                        schema_name=(prefix_schema_name or schema_name),
+                        database_name=(prefix_database_name or database_name),
+                    )
+                    table_entity = self.metadata.search_in_any_service(
+                        entity_type=Table,
+                        fqn_search_string=fqn_search_string,
+                    )
+                    if table_entity and datamodel_entity:
+                        columns_list = [column.name for column in table.columns]
+                        column_lineage = self._get_column_lineage(
+                            table_entity, datamodel_entity, columns_list
+                        )
+                        yield self._get_add_lineage_request(
+                            to_entity=datamodel_entity,
+                            from_entity=table_entity,
+                            column_lineage=column_lineage,
+                        )
         except Exception as exc:  # pylint: disable=broad-except
             yield Either(
                 left=StackTraceError(

metadata/ingestion/source/database/iceberg/metadata.py

@@ -11,6 +11,7 @@
 """
 Iceberg source methods.
 """
+import time
 import traceback
 from typing import Any, Iterable, Optional, Tuple
 
@@ -169,6 +170,43 @@ class IcebergSource(DatabaseServiceSource):
         yield Either(right=schema_request)
         self.register_record_schema_request(schema_request=schema_request)
 
+    def _load_iceberg_table(self, table_identifier: str):
+        """
+        load iceberg table properly with handling
+        network connection error
+        """
+        try:
+            from botocore.exceptions import EndpointConnectionError
+
+            # Add retry logic for transient network issues
+            max_retries = 3
+            retry_delay = 1
+
+            for attempt in range(max_retries):
+                try:
+                    table = self.iceberg.load_table(table_identifier)
+                    # Success, exit retry loop
+                    return table
+                except (OSError, EndpointConnectionError) as e:
+                    if "Couldn't resolve host name" in str(
+                        e
+                    ) or "NETWORK_CONNECTION" in str(e):
+                        if attempt < max_retries - 1:
+                            logger.warning(
+                                f"Network error loading table {table_identifier}, "
+                                f"retrying in {retry_delay} seconds (attempt {attempt + 1}/{max_retries}): {e}"
+                            )
+                            time.sleep(retry_delay)
+                            retry_delay *= 2  # Exponential backoff
+                        else:
+                            logger.warning(f"Maximum retries reached::: {exc}")
+                    else:
+                        logger.warning(f"Other error than host connection::: {exc}")
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Could not load iceberg table properly {exc}")
+        return None
+
     def get_tables_name_and_type(self) -> Optional[Iterable[Tuple[str, str]]]:
         """
         Prepares the table name to be sent to stage.
@@ -178,9 +216,14 @@ class IcebergSource(DatabaseServiceSource):
 
         for table_identifier in self.iceberg.list_tables(namespace):
             try:
-                table = self.iceberg.load_table(table_identifier)
+                table = self._load_iceberg_table(table_identifier)
                 # extract table name from table identifier, which does not include catalog name
                 table_name = get_table_name_as_str(table_identifier)
+                if not table:
+                    logger.debug(
+                        f"iceberg Table could not be fetched for table name = {table_name}"
+                    )
+                    continue
                 table_fqn = fqn.build(
                     self.metadata,
                     entity_type=Table,

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

@@ -22,6 +22,7 @@ from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_SQL_STATEMENT,
 )
 from metadata.ingestion.source.database.snowflake.query_parser import (
+    SNOWFLAKE_QUERY_BATCH_SIZE,
     SnowflakeQueryParserSource,
 )
 from metadata.ingestion.source.database.stored_procedures_mixin import (
@@ -72,26 +73,48 @@ class SnowflakeLineageSource(
         yield a TableQuery with query parsing info
         """
         for engine in self.get_engine():
-            rows = []
-            with engine.connect() as conn:
-                rows = conn.execution_options(
-                    stream_results=True, max_row_buffer=100
-                ).execute(
-                    self.get_sql_statement(start_time=self.start, end_time=self.end)
+            offset = 0
+            total_fetched = 0
+            max_results = self.source_config.resultLimit
+            while total_fetched < max_results:
+                batch_size = min(
+                    SNOWFLAKE_QUERY_BATCH_SIZE, max_results - total_fetched
                 )
-            # exit from active connection after fetching rows & during
-            # further process of `yield_query_lineage`
-            for row in rows:
-                query_dict = dict(row)
-                query_dict.update({k.lower(): v for k, v in query_dict.items()})
-                try:
-                    yield TableQuery(
-                        dialect=self.dialect.value,
-                        query=query_dict["query_text"],
-                        databaseName=self.get_database_name(query_dict),
-                        serviceName=self.config.serviceName,
-                        databaseSchema=self.get_schema_name(query_dict),
+                rows = []
+                row_count = 0
+                with engine.connect() as conn:
+                    rows = conn.execution_options(
+                        stream_results=True, max_row_buffer=100
+                    ).execute(
+                        self.get_sql_statement(
+                            start_time=self.start,
+                            end_time=self.end,
+                            offset=offset,
+                            limit=batch_size,
+                        )
                     )
-                except Exception as exc:
-                    logger.debug(traceback.format_exc())
-                    logger.warning(f"Error processing query_dict {query_dict}: {exc}")
+                for row in rows:
+                    query_dict = dict(row)
+                    query_dict.update({k.lower(): v for k, v in query_dict.items()})
+                    row_count += 1
+                    try:
+                        yield TableQuery(
+                            dialect=self.dialect.value,
+                            query=query_dict["query_text"],
+                            databaseName=self.get_database_name(query_dict),
+                            serviceName=self.config.serviceName,
+                            databaseSchema=self.get_schema_name(query_dict),
+                        )
+                    except Exception as exc:
+                        logger.debug(traceback.format_exc())
+                        logger.warning(
+                            f"Error processing query_dict {query_dict}: {exc}"
+                        )
+                total_fetched += row_count
+                if row_count < batch_size:
+                    break
+                offset += batch_size
+                logger.info(
+                    f"Fetching next page with offset {offset} (fetched {total_fetched}/{max_results}) "
+                    f"for lineage queries"
+                )

metadata/ingestion/source/database/snowflake/metadata.py

@@ -20,6 +20,7 @@ import sqlalchemy.types as sqltypes
 import sqlparse
 from snowflake.sqlalchemy.custom_types import VARIANT, StructuredType
 from snowflake.sqlalchemy.snowdialect import SnowflakeDialect, ischema_names
+from sqlalchemy import exc as sa_exc
 from sqlalchemy.engine.reflection import Inspector
 from sqlparse.sql import Function, Identifier, Token
 
@@ -948,17 +949,26 @@ class SnowflakeSource(
                     table_name = result[6].split(".")[-1]
                     # Can't fetch source of stream is source is dropped or no priviledge
                     if table_name == "No privilege or table dropped":
-                        logger.debug(
-                            f"Couldn't fetch columns of stream [{result and result[1]}],"
-                            f" due to error on source: {table_name}. Result: {result}"
+                        logger.warning(
+                            f"Couldn't fetch columns of stream [{result and result[1]}] "
+                            f"(schema: '{schema_name}', db: '{db_name}') due to error on"
+                            f" source: [{table_name}]. Result: {result}"
                         )
                         return []
             except Exception:
                 pass
 
-        columns = inspector.get_columns(
-            table_name, schema_name, table_type=table_type, db_name=db_name
-        )
+        try:
+            columns = inspector.get_columns(
+                table_name, schema_name, table_type=table_type, db_name=db_name
+            )
+        except sa_exc.NoSuchTableError:
+            logger.warning(
+                f"Table [{table_name}] (schema: '{schema_name}', db: '{db_name}') not found."
+                " Unable to fetch columns. Please check if the configured Snowflake user has"
+                " necessary grants on this table."
+            )
+            return []
 
         if table_type == TableType.Stream:
             columns = [*columns, *DEFAULT_STREAM_COLUMNS]

metadata/ingestion/source/database/snowflake/queries.py

@@ -31,7 +31,9 @@ SNOWFLAKE_SQL_STATEMENT = textwrap.dedent(
     AND query_text NOT LIKE '/* {{"app": "dbt", %%}} */%%'
     AND start_time between to_timestamp_ltz('{start_time}') and to_timestamp_ltz('{end_time}')
     {filters}
+    ORDER BY start_time
     LIMIT {result_limit}
+    OFFSET {offset}
     """
 )
 

metadata/ingestion/source/database/snowflake/query_parser.py

@@ -32,6 +32,7 @@ from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 SNOWFLAKE_ABORTED_CODE = "1969"
+SNOWFLAKE_QUERY_BATCH_SIZE = 1000
 
 
 class SnowflakeQueryParserSource(QueryParserSource, ABC):
@@ -51,18 +52,27 @@ class SnowflakeQueryParserSource(QueryParserSource, ABC):
             )
         return cls(config, metadata)
 
-    def get_sql_statement(self, start_time: datetime, end_time: datetime) -> str:
+    def get_sql_statement(
+        self,
+        start_time: datetime,
+        end_time: datetime,
+        offset: int = 0,
+        limit: int = None,
+    ) -> str:
         """
         returns sql statement to fetch query logs
         """
+        if limit is None:
+            limit = self.config.sourceConfig.config.resultLimit
         return self.sql_stmt.format(
             start_time=start_time,
             end_time=end_time,
-            result_limit=self.config.sourceConfig.config.resultLimit,
+            result_limit=limit,
             filters=self.get_filters(),
             account_usage=self.service_connection.accountUsageSchema,
             credit_cost=self.service_connection.creditCost
             * self.service_connection.creditCost,
+            offset=offset,
         )
 
     def check_life_cycle_query(