use wrapt_timeout_decorator library for process-based timeouts

Author: hsheth2

metadata-ingestion/setup.py

@@ -95,6 +95,7 @@
     # Using an Acryl fork of sqlglot.
     # https://github.com/tobymao/sqlglot/compare/main...hsheth2:sqlglot:hsheth?expand=1
     "acryl-sqlglot==21.1.2.dev9",
+    "wrapt_timeout_decorator",
 }
 
 sql_common = (

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

@@ -12,6 +12,7 @@
 import sqlglot.optimizer.annotate_types
 import sqlglot.optimizer.optimizer
 import sqlglot.optimizer.qualify
+import wrapt_timeout_decorator
 
 from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.ingestion.graph.client import DataHubGraph
@@ -45,7 +46,6 @@
     is_dialect_instance,
     parse_statement,
 )
-from datahub.utilities.timeout import timeout
 
 logger = logging.getLogger(__name__)
 
@@ -311,6 +311,9 @@ class SqlUnderstandingError(Exception):
 
 
 # TODO: Break this up into smaller functions.
+@wrapt_timeout_decorator.timeout(
+    SQL_LINEAGE_TIMEOUT_SECONDS if SQL_LINEAGE_TIMEOUT_ENABLED else None,
+)
 def _column_level_lineage(  # noqa: C901
     statement: sqlglot.exp.Expression,
     dialect: sqlglot.Dialect,
@@ -766,7 +769,6 @@ def _translate_internal_column_lineage(
     )
 
 
-@timeout(enabled=SQL_LINEAGE_TIMEOUT_ENABLED, seconds=SQL_LINEAGE_TIMEOUT_SECONDS)
 def _sqlglot_lineage_inner(
     sql: sqlglot.exp.ExpOrStr,
     schema_resolver: SchemaResolverInterface,