feat(ingest/sigma): add default_db/schema config and optimize SQL parsing performance

metadata-ingestion/src/datahub/ingestion/source/sigma/config.py

@@ -136,6 +136,16 @@ class PlatformDetail(PlatformInstanceConfigMixin, EnvConfigMixin):
     data_source_platform: str = pydantic.Field(
         description="A chart's data sources platform name.",
     )
+    default_db: Optional[str] = pydantic.Field(
+        default=None,
+        description="Default database name to use when parsing SQL queries. "
+        "Used to generate fully qualified table URNs (e.g., 'prod' for 'prod.public.table').",
+    )
+    default_schema: Optional[str] = pydantic.Field(
+        default=None,
+        description="Default schema name to use when parsing SQL queries. "
+        "Used to generate fully qualified table URNs (e.g., 'public' for 'prod.public.table').",
+    )
 
 
 class SigmaSourceConfig(
@@ -173,6 +183,12 @@ class SigmaSourceConfig(
         default={},
         description="A mapping of the sigma workspace/workbook/chart folder path to all chart's data sources platform details present inside that folder path.",
     )
+    generate_column_lineage: bool = pydantic.Field(
+        default=False,
+        description="Whether to generate column-level lineage. "
+        "Disabled by default as Sigma only uses table-level lineage and "
+        "column lineage computation can be expensive for complex queries.",
+    )
     stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = pydantic.Field(
         default=None, description="Sigma Stateful Ingestion Config."
     )

metadata-ingestion/src/datahub/ingestion/source/sigma/sigma.py

@@ -404,13 +404,15 @@ def _get_element_input_details(
             try:
                 sql_parser_in_tables = create_lineage_sql_parsed_result(
                     query=element.query.strip(),
-                    default_db=None,
+                    default_db=data_source_platform_details.default_db,
+                    default_schema=data_source_platform_details.default_schema,
                     platform=data_source_platform_details.data_source_platform,
                     env=data_source_platform_details.env,
                     platform_instance=data_source_platform_details.platform_instance,
+                    generate_column_lineage=self.config.generate_column_lineage,
                 ).in_tables
             except Exception:
-                logging.debug(f"Unable to parse query of element {element.name}")
+                logger.debug(f"Unable to parse query of element {element.name}")
 
         # Add sigma dataset as input of element if present
         # and its matched sql parser in_table as its upsteam dataset
@@ -671,6 +673,7 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
 
         for dataset in self.sigma_api.get_sigma_datasets():
             yield from self._gen_dataset_workunit(dataset)
+
         for workbook in self.sigma_api.get_sigma_workbooks():
             yield from self._gen_workbook_workunit(workbook)
 

metadata-ingestion/src/datahub/ingestion/source/sigma/sigma_api.py

@@ -1,6 +1,7 @@
 import functools
 import logging
 import sys
+import time
 from typing import Any, Dict, List, Optional
 
 import requests
@@ -22,6 +23,10 @@
 # Logger instance
 logger = logging.getLogger(__name__)
 
+# Retry constants for 429/503 errors
+RETRY_MAX_ATTEMPTS = 3
+RETRY_BASE_DELAY_SECONDS = 2.0  # Exponential backoff: 2s, 4s, 8s
+
 
 class SigmaAPI:
     def __init__(self, config: SigmaSourceConfig, report: SigmaSourceReport) -> None:
@@ -31,6 +36,7 @@ def __init__(self, config: SigmaSourceConfig, report: SigmaSourceReport) -> None
         self.users: Dict[str, str] = {}
         self.session = requests.Session()
         self.refresh_token: Optional[str] = None
+
         # Test connection by generating access token
         logger.info(f"Trying to connect to {self.config.api_url}")
         self._generate_token()
@@ -87,11 +93,35 @@ def _refresh_access_token(self):
             )
 
     def _get_api_call(self, url: str) -> requests.Response:
-        get_response = self.session.get(url)
-        if get_response.status_code == 401 and self.refresh_token:
-            logger.debug("Access token might expired. Refreshing access token.")
-            self._refresh_access_token()
+        """Make an API call with retry on 429/503 errors."""
+        get_response: requests.Response
+        for attempt in range(RETRY_MAX_ATTEMPTS):
             get_response = self.session.get(url)
+
+            # Handle token refresh on 401
+            if get_response.status_code == 401 and self.refresh_token:
+                logger.debug("Access token might expired. Refreshing access token.")
+                self._refresh_access_token()
+                get_response = self.session.get(url)
+
+            # Success or non-retryable error
+            if get_response.status_code not in (429, 503):
+                break
+
+            # Retry with exponential backoff, or give up on last attempt
+            if attempt < RETRY_MAX_ATTEMPTS - 1:
+                delay = RETRY_BASE_DELAY_SECONDS * (2**attempt)
+                logger.debug(
+                    f"Retryable error ({get_response.status_code}) on {url}, "
+                    f"retrying in {delay}s (attempt {attempt + 1}/{RETRY_MAX_ATTEMPTS})"
+                )
+                time.sleep(delay)
+            else:
+                logger.warning(
+                    f"Retryable error ({get_response.status_code}) on {url}, "
+                    f"max retries exceeded"
+                )
+
         return get_response
 
     def get_workspace(self, workspace_id: str) -> Optional[Workspace]:
@@ -299,6 +329,11 @@ def _get_element_upstream_sources(
                     f"Lineage metadata not accessible for element {element.name} of workbook '{workbook.name}'"
                 )
                 return upstream_sources
+            if response.status_code == 400:
+                logger.debug(
+                    f"Lineage not supported for element {element.name} of workbook '{workbook.name}' (400 Bad Request)"
+                )
+                return upstream_sources
 
             response.raise_for_status()
             response_dict = response.json()

metadata-ingestion/src/datahub/sql_parsing/sqlglot_lineage.py

@@ -1846,6 +1846,7 @@ def _sqlglot_lineage_inner(
     default_db: Optional[str] = None,
     default_schema: Optional[str] = None,
     override_dialect: Optional[DialectOrStr] = None,
+    generate_column_lineage: bool = True,
 ) -> SqlParsingResult:
     if override_dialect:
         dialect = get_dialect(override_dialect)
@@ -1957,31 +1958,32 @@ def _sqlglot_lineage_inner(
 
     column_lineage: Optional[List[_ColumnLineageInfo]] = None
     joins = None
-    try:
-        with cooperative_timeout(
-            timeout=(
-                SQL_LINEAGE_TIMEOUT_SECONDS if SQL_LINEAGE_TIMEOUT_ENABLED else None
-            )
-        ):
-            column_lineage_debug_info = _column_level_lineage(
-                statement,
-                dialect=dialect,
-                downstream_table=downstream_table,
-                table_name_schema_mapping=table_name_schema_mapping,
-                default_db=default_db,
-                default_schema=default_schema,
-            )
-            column_lineage = column_lineage_debug_info.column_lineage
-            joins = column_lineage_debug_info.joins
-    except CooperativeTimeoutError as e:
-        logger.debug(f"Timed out while generating column-level lineage: {e}")
-        debug_info.column_error = e
-    except UnsupportedStatementTypeError as e:
-        # For this known exception type, we assume the error is logged at the point of failure.
-        debug_info.column_error = e
-    except Exception as e:
-        logger.debug(f"Failed to generate column-level lineage: {e}", exc_info=True)
-        debug_info.column_error = e
+    if generate_column_lineage:
+        try:
+            with cooperative_timeout(
+                timeout=(
+                    SQL_LINEAGE_TIMEOUT_SECONDS if SQL_LINEAGE_TIMEOUT_ENABLED else None
+                )
+            ):
+                column_lineage_debug_info = _column_level_lineage(
+                    statement,
+                    dialect=dialect,
+                    downstream_table=downstream_table,
+                    table_name_schema_mapping=table_name_schema_mapping,
+                    default_db=default_db,
+                    default_schema=default_schema,
+                )
+                column_lineage = column_lineage_debug_info.column_lineage
+                joins = column_lineage_debug_info.joins
+        except CooperativeTimeoutError as e:
+            logger.debug(f"Timed out while generating column-level lineage: {e}")
+            debug_info.column_error = e
+        except UnsupportedStatementTypeError as e:
+            # For this known exception type, we assume the error is logged at the point of failure.
+            debug_info.column_error = e
+        except Exception as e:
+            logger.debug(f"Failed to generate column-level lineage: {e}", exc_info=True)
+            debug_info.column_error = e
 
     # TODO: Can we generate a common JOIN tables / keys section?
     # TODO: Can we generate a common WHERE clauses section?
@@ -2043,6 +2045,7 @@ def _sqlglot_lineage_nocache(
     default_db: Optional[str] = None,
     default_schema: Optional[str] = None,
     override_dialect: Optional[DialectOrStr] = None,
+    generate_column_lineage: bool = True,
 ) -> SqlParsingResult:
     """Parse a SQL statement and generate lineage information.
 
@@ -2102,6 +2105,7 @@ def _sqlglot_lineage_nocache(
             default_db=default_db,
             default_schema=default_schema,
             override_dialect=override_dialect,
+            generate_column_lineage=generate_column_lineage,
         )
     except Exception as e:
         return SqlParsingResult.make_from_error(e)
@@ -2140,14 +2144,25 @@ def sqlglot_lineage(
     default_db: Optional[str] = None,
     default_schema: Optional[str] = None,
     override_dialect: Optional[DialectOrStr] = None,
+    generate_column_lineage: bool = True,
 ) -> SqlParsingResult:
     if schema_resolver.includes_temp_tables():
         return _sqlglot_lineage_nocache(
-            sql, schema_resolver, default_db, default_schema, override_dialect
+            sql=sql,
+            schema_resolver=schema_resolver,
+            default_db=default_db,
+            default_schema=default_schema,
+            override_dialect=override_dialect,
+            generate_column_lineage=generate_column_lineage,
         )
     else:
         return _sqlglot_lineage_cached(
-            sql, schema_resolver, default_db, default_schema, override_dialect
+            sql=sql,
+            schema_resolver=schema_resolver,
+            default_db=default_db,
+            default_schema=default_schema,
+            override_dialect=override_dialect,
+            generate_column_lineage=generate_column_lineage,
         )
 
 
@@ -2200,6 +2215,7 @@ def create_lineage_sql_parsed_result(
     graph: Optional[DataHubGraph] = None,
     schema_aware: bool = True,
     override_dialect: Optional[DialectOrStr] = None,
+    generate_column_lineage: bool = True,
 ) -> SqlParsingResult:
     schema_resolver = create_schema_resolver(
         platform=platform,
@@ -2220,6 +2236,7 @@ def create_lineage_sql_parsed_result(
             default_db=default_db,
             default_schema=default_schema,
             override_dialect=override_dialect,
+            generate_column_lineage=generate_column_lineage,
         )
     except Exception as e:
         return SqlParsingResult.make_from_error(e)