From b94927b3ea81b4382472e4ffe001563d6eda554b Mon Sep 17 00:00:00 2001
From: Nahuel <nahuel@getcollate.io>
Date: Fri, 10 Mar 2023 10:55:56 +0100
Subject: [PATCH] Bug: fix flaky timeout test (#10509)

---
 .../src/metadata/ingestion/lineage/parser.py  | 28 +++++++++++--------
 ingestion/src/metadata/utils/timeout.py       |  2 --
 ingestion/tests/unit/test_sql_lineage.py      |  8 ++++--
 3 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/ingestion/src/metadata/ingestion/lineage/parser.py b/ingestion/src/metadata/ingestion/lineage/parser.py
index 756dc2bba88..640e1b39092 100644
--- a/ingestion/src/metadata/ingestion/lineage/parser.py
+++ b/ingestion/src/metadata/ingestion/lineage/parser.py
@@ -51,13 +51,6 @@ logger = ingestion_logger()
 LINEAGE_PARSING_TIMEOUT = 10
 
 
-@timeout(seconds=LINEAGE_PARSING_TIMEOUT)
-def get_sqlfluff_lineage_runner(query: str, dialect: str) -> LineageRunner:
-    lr_sqlfluff = LineageRunner(query, dialect=dialect)
-    lr_sqlfluff.get_column_lineage()
-    return lr_sqlfluff
-
-
 class LineageParser:
     """
     Class that acts like a wrapper for the LineageRunner library usage
@@ -67,10 +60,17 @@ class LineageParser:
     query: str
     _clean_query: str
 
-    def __init__(self, query: str, dialect: Dialect = Dialect.ANSI):
+    def __init__(
+        self,
+        query: str,
+        dialect: Dialect = Dialect.ANSI,
+        timeout_seconds: int = LINEAGE_PARSING_TIMEOUT,
+    ):
         self.query = query
         self._clean_query = self.clean_raw_query(query)
-        self.parser = self._evaluate_best_parser(self._clean_query, dialect=dialect)
+        self.parser = self._evaluate_best_parser(
+            self._clean_query, dialect=dialect, timeout_seconds=timeout_seconds
+        )
 
     @cached_property
     def involved_tables(self) -> Optional[List[Table]]:
@@ -367,8 +367,14 @@ class LineageParser:
 
     @staticmethod
     def _evaluate_best_parser(
-        query: str, dialect: Dialect = Dialect.ANSI
+        query: str, dialect: Dialect, timeout_seconds: int
     ) -> LineageRunner:
+        @timeout(seconds=timeout_seconds)
+        def get_sqlfluff_lineage_runner(qry: str, dlct: str) -> LineageRunner:
+            lr_dialect = LineageRunner(qry, dialect=dlct)
+            lr_dialect.get_column_lineage()
+            return lr_dialect
+
         sqlfluff_count = 0
         try:
             lr_sqlfluff = get_sqlfluff_lineage_runner(query, dialect.value)
@@ -382,7 +388,7 @@ class LineageParser:
         except TimeoutError:
             logger.debug(
                 f"Lineage with SqlFluff failed for the [{dialect.value}] query: [{query}]: "
-                f"Parser has been running for more than {LINEAGE_PARSING_TIMEOUT} seconds."
+                f"Parser has been running for more than {timeout_seconds} seconds."
             )
             lr_sqlfluff = None
         except Exception:
diff --git a/ingestion/src/metadata/utils/timeout.py b/ingestion/src/metadata/utils/timeout.py
index f8268416082..effd93a65ff 100644
--- a/ingestion/src/metadata/utils/timeout.py
+++ b/ingestion/src/metadata/utils/timeout.py
@@ -18,7 +18,6 @@ import inspect
 import os
 import platform
 import signal
-import traceback
 from typing import Callable
 
 from metadata.utils.constants import TEN_MIN
@@ -31,7 +30,6 @@ def _handle_timeout(signum, frame):
     """
     Handler for signal timeout
     """
-    logger.debug(traceback.print_stack(frame))
     raise TimeoutError(f"[SIGNUM {signum}] {os.strerror(errno.ETIME)}")
 
 
diff --git a/ingestion/tests/unit/test_sql_lineage.py b/ingestion/tests/unit/test_sql_lineage.py
index 01875b872d9..93948444e65 100644
--- a/ingestion/tests/unit/test_sql_lineage.py
+++ b/ingestion/tests/unit/test_sql_lineage.py
@@ -164,13 +164,15 @@ class SqlLineageTest(TestCase):
         # When
         with self.assertLogs(Loggers.INGESTION.value, level="DEBUG") as logger:
             LineageParser(
-                query.format(values="\n".join(values)), dialect=Dialect.SNOWFLAKE
+                query.format(values="\n".join(values)),
+                dialect=Dialect.SNOWFLAKE,
+                timeout_seconds=1,
             )
             # Then
             self.assertTrue(
                 any(
-                    "Parser has been running for more than 10 seconds." in log
+                    "Parser has been running for more than 1 seconds." in log
                     for log in logger.output
                 ),
-                "Parser finished before the 10 expected seconds!",
+                "Parser finished before the 1 expected seconds!",
             )