created _add_default_options

Author: brock-acryl

metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py

@@ -352,6 +352,13 @@ def __init__(self, config: SQLCommonConfig, ctx: PipelineContext, platform: str)
         )
         self.report.sql_aggregator = self.aggregator.report
 
+    def _add_default_options(self, sql_config: SQLCommonConfig) -> None:
+        """Add default SQLAlchemy options. Can be overridden by subclasses to add additional defaults."""
+        if sql_config.is_profiling_enabled():
+            sql_config.options.setdefault(
+                "max_overflow", sql_config.profiling.max_workers
+            )
+
     @classmethod
     def test_connection(cls, config_dict: dict) -> TestConnectionReport:
         test_report = TestConnectionReport()
@@ -521,10 +528,7 @@ def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit
 
         # Extra default SQLAlchemy option for better connection pooling and threading.
         # https://docs.sqlalchemy.org/en/14/core/pooling.html#sqlalchemy.pool.QueuePool.params.max_overflow
-        if sql_config.is_profiling_enabled():
-            sql_config.options.setdefault(
-                "max_overflow", sql_config.profiling.max_workers
-            )
+        self._add_default_options(sql_config)
 
         for inspector in self.get_inspectors():
             profiler = None

metadata-ingestion/src/datahub/ingestion/source/sql/teradata.py

@@ -679,6 +679,12 @@ def __init__(self, config: TeradataConfig, ctx: PipelineContext):
             if self.config.stateful_ingestion:
                 self.config.stateful_ingestion.remove_stale_metadata = False
 
+    def _add_default_options(self, sql_config: SQLCommonConfig) -> None:
+        """Add Teradata-specific default options"""
+        # Teradata does not support max_overflow, instead we use QueuePool when profiling
+        if sql_config.is_profiling_enabled():
+            sql_config.options.setdefault("poolclass", QueuePool)
+
     @classmethod
     def create(cls, config_dict, ctx):
         config = TeradataConfig.parse_obj(config_dict)
@@ -707,11 +713,6 @@ def get_inspectors(self):
         # run on multiple databases.
         url = self.config.get_sql_alchemy_url()
 
-        # Teradata does not support max_overflow, instead we use QueuePool when profiling data
-        if "max_overflow" in self.config.options:
-            self.config.options.pop("max_overflow")
-            self.config.options["poolclass"] = QueuePool
-
         logger.debug(f"sql_alchemy_url={url}")
         engine = create_engine(url, **self.config.options)
         with engine.connect() as conn: