add timeout to SQL parsing

metaphor/common/sql/table_level_lineage/table_level_lineage.py

@@ -1,5 +1,6 @@
 import logging
 import re
+import signal
 from collections import defaultdict
 from typing import Dict, Optional, Set
 
@@ -128,6 +129,10 @@
     return sources
 
 
+def _handle_parser_timeout(signum, frame):
+    raise TimeoutError("Parser timeout")
+
+
 def extract_table_level_lineage(
     sql: str,
     platform: DataPlatform,
@@ -137,25 +142,32 @@
     default_database: Optional[str] = None,
     default_schema: Optional[str] = None,
 ) -> Result:
-
     if statement_type and statement_type.upper() not in _VALID_STATEMENT_TYPES:
         # No target, no TLL possible
         return Result()
 
+    signal.signal(signal.SIGALRM, _handle_parser_timeout)
+    signal.alarm(10)  # set timeout to 10 seconds
     try:
         expression: Expression = maybe_parse(
             sql, dialect=PLATFORM_TO_DIALECT.get(platform)
         )
     except (sqlglot.errors.ParseError, sqlglot.errors.TokenError):
         if not _is_truncated_insert_into_with_values(sql) and query_id:
-            logger.warning(f"Cannot parse sql with SQLGlot, query_id = {query_id}")
+            logger.warning(f"Cannot parse sql with query_id = {query_id}")
         return Result()
     except RecursionError:
         if query_id:
             logger.warning(
                 f"Cannot parse sql with SQLGlot (max recursion level exceeded), query_id = {query_id}"
             )
         return Result()
+    except TimeoutError:
+        if query_id:
+            logger.warning(f"Parser timeout, query_id = {query_id}")
+        return Result()
+    finally:
+        signal.alarm(0)
 
     try:
         return Result(

metaphor/quick_sight/lineage.py

@@ -57,6 +57,7 @@ class Column:
 
 def _get_source_from_custom_sql(
     resources: Dict[str, ResourceType],
+    table_id: str,
     custom_sql: TypeCustomSql,
 ) -> Tuple[Optional[List[str]], Optional[VirtualViewQuery]]:
     data_source = resources.get(custom_sql.DataSourceArn)
@@ -80,7 +81,7 @@ def _get_source_from_custom_sql(
         query,
         platform=data_platform,
         account=account,
-        query_id="",
+        query_id=table_id,
         default_database=database,
     )
 
@@ -198,7 +199,7 @@ def _process_physical_table_map(
 
             if physical_table.CustomSql:
                 source_entities, query = _get_source_from_custom_sql(
-                    self._resources, physical_table.CustomSql
+                    self._resources, table_id, physical_table.CustomSql
                 )
 
                 # CLL of custom sql is not supported

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "metaphor-connectors"
-version = "0.14.173"
+version = "0.14.174"
 license = "Apache-2.0"
 description = "A collection of Python-based 'connectors' that extract metadata from various sources to ingest into the Metaphor app."
 authors = ["Metaphor <dev@metaphor.io>"]