MINOR: fix: close client after query (#19711)

* fix: close client after query

use context clients in SQL sampler to close the connection once the query is complete

* use self.context_client in all sql sampler implementations

* use sqlalchemy's built-in session management

* format

* format

* use get_client directly

ingestion/src/metadata/profiler/interface/sqlalchemy/unity_catalog/sampler_interface.py

@@ -23,7 +23,6 @@ class UnityCatalogSamplerInterface(SQASampler):
     def get_client(self):
         """client is the session for SQA"""
         self.connection = databricks_get_connection(self.service_connection_config)
-        self.client = super().get_client()
+        client = super().get_client()
-        self.set_catalog(self.client)
+        self.set_catalog(client)
-
+        return client
-        return self.client

ingestion/src/metadata/sampler/nosql/sampler.py

@@ -24,6 +24,10 @@ class NoSQLSampler(SamplerInterface):
 
     client: NoSQLAdaptor
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.client = self.get_client()
+
     @property
     def raw_dataset(self):
         return self.entity

ingestion/src/metadata/sampler/pandas/sampler.py

@@ -46,6 +46,7 @@ class DatalakeSampler(SamplerInterface, PandasInterfaceMixin):
         super().__init__(*args, **kwargs)
         self.partition_details = cast(PartitionProfilerConfig, self.partition_details)
         self._table = None
+        self.client = self.get_client()
 
     @property
     def raw_dataset(self):

ingestion/src/metadata/sampler/sampler_interface.py

@@ -86,7 +86,6 @@ class SamplerInterface(ABC):
 
         self.service_connection_config = service_connection_config
         self.connection = get_ssl_connection(self.service_connection_config)
-        self.client = self.get_client()
 
     # pylint: disable=too-many-arguments, too-many-locals
     @classmethod

ingestion/src/metadata/sampler/sqlalchemy/azuresql/sampler.py

@@ -54,7 +54,8 @@ class AzureSQLSampler(SQASampler):
         rnd = self._base_sample_query(column).cte(
             f"{self.get_sampler_table_name()}_rnd"
         )
-        query = self.client.query(rnd)
+        with self.get_client() as client:
+            query = client.query(rnd)
         return query.cte(f"{self.get_sampler_table_name()}_sample")
 
     def fetch_sample_data(self, columns: Optional[List[Column]] = None) -> TableData:

ingestion/src/metadata/sampler/sqlalchemy/mssql/sampler.py

@@ -48,5 +48,6 @@ class MssqlSampler(SQASampler):
         rnd = self._base_sample_query(column).cte(
             f"{self.get_sampler_table_name()}_rnd"
         )
-        query = self.client.query(rnd)
+        with self.get_client() as client:
+            query = client.query(rnd)
         return query.cte(f"{self.get_sampler_table_name()}_sample")

ingestion/src/metadata/sampler/sqlalchemy/sampler.py

@@ -102,10 +102,11 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         selectable = self.set_tablesample(self.raw_dataset.__table__)
 
         entity = selectable if column is None else selectable.c.get(column.key)
+        with self.get_client() as client:
             if label is not None:
-            query = self.client.query(entity, label)
+                query = client.query(entity, label)
             else:
-            query = self.client.query(entity)
+                query = client.query(entity)
 
             if self.partition_details:
                 query = self.get_partitioned_query(query)
@@ -129,12 +130,14 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
                 column,
                 (ModuloFn(RandomNumFn(), 100)).label(RANDOM_LABEL),
             ).cte(f"{self.get_sampler_table_name()}_rnd")
-            session_query = self.client.query(rnd)
+            with self.get_client() as client:
+                session_query = client.query(rnd)
             return session_query.where(
                 rnd.c.random <= self.sample_config.profileSample
             ).cte(f"{self.get_sampler_table_name()}_sample")
 
-        table_query = self.client.query(self.raw_dataset)
+        with self.get_client() as client:
+            table_query = client.query(self.raw_dataset)
         session_query = self._base_sample_query(
             column,
             (ModuloFn(RandomNumFn(), table_query.count())).label(RANDOM_LABEL)
@@ -195,10 +198,10 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
                 for col in inspect(ds).c
                 if col.name != RANDOM_LABEL and col.name in names
             ]
-
+        with self.get_client() as client:
             try:
                 sqa_sample = (
-                self.client.query(*sqa_columns)
+                    client.query(*sqa_columns)
                     .select_from(ds)
                     .limit(self.sample_limit)
                     .all()
@@ -210,12 +213,11 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
                 logger.debug(traceback.format_exc())
                 sqa_columns = list(inspect(self.raw_dataset).c)
                 sqa_sample = (
-                self.client.query(*sqa_columns)
+                    client.query(*sqa_columns)
                     .select_from(self.raw_dataset)
                     .limit(100)
                     .all()
                 )
-
             return TableData(
                 columns=[column.name for column in sqa_columns],
                 rows=[list(row) for row in sqa_sample],
@@ -226,7 +228,8 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         if not is_safe_sql_query(self.sample_query):
             raise RuntimeError(f"SQL expression is not safe\n\n{self.sample_query}")
 
-        rnd = self.client.execute(f"{self.sample_query}")
+        with self.get_client() as client:
+            rnd = client.execute(f"{self.sample_query}")
         try:
             columns = [col.name for col in rnd.cursor.description]
         except AttributeError:
@@ -243,8 +246,8 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
 
         stmt = text(f"{self.sample_query}")
         stmt = stmt.columns(*list(inspect(self.raw_dataset).c))
-
+        with self.get_client() as client:
-        return self.client.query(stmt.subquery()).cte(
+            return client.query(stmt.subquery()).cte(
                 f"{self.get_sampler_table_name()}_user_sampled"
             )
 
@@ -263,7 +266,8 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         )
         if query is not None:
             return query.filter(partition_filter)
-        return self.client.query(self.raw_dataset).filter(partition_filter)
+        with self.get_client() as client:
+            return client.query(self.raw_dataset).filter(partition_filter)
 
     def get_columns(self):
         """get columns from entity"""

ingestion/src/metadata/sampler/sqlalchemy/snowflake/sampler.py

@@ -91,5 +91,6 @@ class SnowflakeSampler(SQASampler):
         rnd = self._base_sample_query(column).cte(
             f"{self.get_sampler_table_name()}_rnd"
         )
-        query = self.client.query(rnd)
+        with self.get_client() as client:
+            query = client.query(rnd)
         return query.cte(f"{self.get_sampler_table_name()}_sample")

ingestion/src/metadata/sampler/sqlalchemy/trino/sampler.py

@@ -38,7 +38,8 @@ class TrinoSampler(SQASampler):
             col for col in inspect(self.raw_dataset).c if col.name != RANDOM_LABEL
         ]
         entity = self.raw_dataset if column is None else column
-        return self.client.query(entity, label).where(
+        with self.get_client() as client:
+            return client.query(entity, label).where(
                 or_(
                     *[
                         text(f'is_nan("{cols.name}") = False')