openmetadata-ingestion 1.6.9.0__py3-none-any.whl → 1.7.0.0__py3-none-any.whl

metadata/sampler/models.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -41,6 +41,7 @@ class BaseProfileConfig(ConfigModel):
     profileSampleType: Optional[ProfileSampleType] = None
     samplingMethodType: Optional[SamplingMethodType] = None
     sampleDataCount: Optional[int] = 100
+    randomizedSample: Optional[bool] = True
 
 
 class ColumnConfig(ConfigModel):
@@ -56,6 +57,7 @@ class TableConfig(BaseProfileConfig):
     profileQuery: Optional[str] = None
     partitionConfig: Optional[PartitionProfilerConfig] = None
     columnConfig: Optional[ColumnConfig] = None
+    randomizedSample: Optional[bool] = True
 
     @classmethod
     def from_database_and_schema_config(
@@ -104,3 +106,4 @@ class SampleConfig(ConfigModel):
     profileSample: Optional[Union[float, int]] = None
     profileSampleType: Optional[ProfileSampleType] = ProfileSampleType.PERCENTAGE
     samplingMethodType: Optional[SamplingMethodType] = None
+    randomizedSample: Optional[bool] = True

metadata/sampler/nosql/sampler.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -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

metadata/sampler/pandas/sampler.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -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):

metadata/sampler/partition.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -26,6 +26,9 @@ from metadata.generated.schema.entity.services.databaseService import (
     DatabaseServiceType,
 )
 from metadata.sampler.models import TableConfig
+from metadata.utils.logger import sampler_logger
+
+logger = sampler_logger()
 
 
 def validate_athena_injected_partitioning(
@@ -164,8 +167,12 @@ def _handle_bigquery_partition(
                 partitionIntegerRangeEnd=10000,
             )
         # TODO: Allow External Hive Partitioning for profiler
-        raise TypeError(
-            f"Unsupported partition type {partition.intervalType}. Skipping table"
+        logger.info(
+            f"Unknown partition type {partition.intervalType}. Cannot infer partition. Partitionning "
+            "will be set to None. If partitionning is required for your table, you can set it "
+            "directly in the table config. More information can be found at "
+            "https://docs.open-metadata.org/latest/how-to-guides/data-quality-observability/profiler/"
+            "workflow#4.-updating-profiler-setting-at-the-table-level"
         )
 
     return None

metadata/sampler/processor.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -17,6 +17,9 @@ from typing import Optional, cast
 
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
+    BigQueryConnection,
+)
 from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
@@ -38,6 +41,7 @@ from metadata.profiler.source.metadata import ProfilerSourceAndEntity
 from metadata.sampler.config import get_config_for_table
 from metadata.sampler.models import SampleConfig, SampleData, SamplerResponse
 from metadata.sampler.sampler_interface import SamplerInterface
+from metadata.utils.bigquery_utils import copy_service_config
 from metadata.utils.profiler_utils import get_context_entities
 from metadata.utils.service_spec.service_spec import import_sampler_class
 
@@ -133,6 +137,9 @@ class SamplerProcessor(Processor):
         Returns:
             DatabaseService.__config__
         """
+        if isinstance(config.source.serviceConnection.root.config, BigQueryConnection):
+            return copy_service_config(config, database.name.root)
+
         config_copy = deepcopy(
             config.source.serviceConnection.root.config  # type: ignore
         )

metadata/sampler/sampler_interface.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -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

metadata/sampler/sqlalchemy/azuresql/sampler.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -14,10 +14,11 @@ for the profiler
 """
 from typing import List, Optional
 
-from sqlalchemy import Column
+from sqlalchemy import Column, Table, text
+from sqlalchemy.orm import Query
 
-from metadata.generated.schema.entity.data.table import TableData
-from metadata.sampler.sqlalchemy.sampler import SQASampler
+from metadata.generated.schema.entity.data.table import TableData, TableType
+from metadata.sampler.sqlalchemy.sampler import ProfileSampleType, SQASampler
 
 
 class AzureSQLSampler(SQASampler):
@@ -31,6 +32,32 @@ class AzureSQLSampler(SQASampler):
     # pyodbc.ProgrammingError: ('ODBC SQL type -151 is not yet supported.  column-index=x  type=-151', 'HY106')
     NOT_COMPUTE_PYODBC = {"SQASGeography", "UndeterminedType"}
 
+    def set_tablesample(self, selectable: Table):
+        """Set the TABLESAMPLE clause for MSSQL
+        Args:
+            selectable (Table): _description_
+        """
+        if self.entity.tableType != TableType.View:
+            if self.sample_config.profileSampleType == ProfileSampleType.PERCENTAGE:
+                return selectable.tablesample(
+                    text(f"{self.sample_config.profileSample or 100} PERCENT")
+                )
+
+            return selectable.tablesample(
+                text(f"{int(self.sample_config.profileSample or 100)} ROWS")
+            )
+
+        return selectable
+
+    def get_sample_query(self, *, column=None) -> Query:
+        """get query for sample data"""
+        rnd = self._base_sample_query(column).cte(
+            f"{self.get_sampler_table_name()}_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:
         sqa_columns = []
         if columns:

metadata/sampler/sqlalchemy/bigquery/sampler.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -12,16 +12,18 @@
 Helper module to handle data sampling
 for the profiler
 """
-from typing import Dict, Optional, Union
+import traceback
+from typing import Dict, List, Optional, Union
 
 from sqlalchemy import Column
 from sqlalchemy import Table as SqaTable
-from sqlalchemy import text
+from sqlalchemy import inspect, text
 from sqlalchemy.orm import Query
 
 from metadata.generated.schema.entity.data.table import (
     ProfileSampleType,
     Table,
+    TableData,
     TableType,
 )
 from metadata.generated.schema.entity.services.connections.connectionBasicType import (
@@ -35,6 +37,11 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.sampler.models import SampleConfig
 from metadata.sampler.sqlalchemy.sampler import SQASampler
 from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
+from metadata.utils.logger import profiler_interface_registry_logger
+
+logger = profiler_interface_registry_logger()
+
+RANDOM_LABEL = "random"
 
 
 class BigQuerySampler(SQASampler):
@@ -123,3 +130,62 @@ class BigQuerySampler(SQASampler):
             )
 
         return super().get_sample_query(column=column)
+
+    def fetch_sample_data(self, columns: Optional[List[Column]] = None) -> TableData:
+        """
+        Use the sampler to retrieve sample data rows as per limit given by user
+
+        Args:
+            columns (Optional[List]): List of columns to fetch
+        Returns:
+            TableData to be added to the Table Entity
+        """
+        if self.sample_query:
+            return self._fetch_sample_data_from_user_query()
+
+        # Add new RandomNumFn column
+        ds = self.get_dataset()
+        if not columns:
+            sqa_columns = [col for col in inspect(ds).c if col.name != RANDOM_LABEL]
+        else:
+            # we can't directly use columns as it is bound to self.raw_dataset and not the rnd table.
+            # If we use it, it will result in a cross join between self.raw_dataset and rnd table
+            names = [col.name for col in columns]
+            sqa_columns = [
+                col
+                for col in inspect(ds).c
+                if col.name != RANDOM_LABEL and col.name in names
+            ]
+        with self.get_client() as client:
+            try:
+                query = (
+                    str(client.query(*sqa_columns)).replace(
+                        f"`{ds.__table_args__['schema']}`.`{ds.__tablename__}`",
+                        f"`{self.entity.database.name}`.`{ds.__table_args__['schema']}`.`{ds.__tablename__}`",
+                    )
+                    + f" LIMIT {self.sample_limit}"
+                )
+                sqa_sample = client.execute(text(query))
+            except Exception:
+                logger.debug(
+                    "Cannot fetch sample data with random sampling. Falling back to 100 rows."
+                )
+                logger.debug(traceback.format_exc())
+                ds_inspect = inspect(self.raw_dataset)
+                sqa_columns = list(ds_inspect.c)
+
+                schema = ds_inspect.class_.__table_args__["schema"]
+                table = ds_inspect.class_.__tablename__
+
+                query = (
+                    str(client.query(*sqa_columns)).replace(
+                        f"`{schema}`.`{table}`",
+                        f"`{self.entity.database.name}`.`{schema}`.`{table}`",
+                    )
+                    + f" LIMIT {self.sample_limit}"
+                )
+                sqa_sample = client.execute(text(query))
+            return TableData(
+                columns=[column.name for column in sqa_columns],
+                rows=[list(row) for row in sqa_sample],
+            )

metadata/sampler/sqlalchemy/mssql/sampler.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -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")

metadata/sampler/sqlalchemy/postgres/sampler.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

metadata/sampler/sqlalchemy/sampler.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -102,13 +102,14 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         selectable = self.set_tablesample(self.raw_dataset.__table__)
 
         entity = selectable if column is None else selectable.c.get(column.key)
-        if label is not None:
-            query = self.client.query(entity, label)
-        else:
-            query = self.client.query(entity)
+        with self.get_client() as client:
+            if label is not None:
+                query = client.query(entity, label)
+            else:
+                query = client.query(entity)
 
-        if self.partition_details:
-            query = self.get_partitioned_query(query)
+            if self.partition_details:
+                query = self.get_partitioned_query(query)
         return query
 
     def get_sampler_table_name(self) -> str:
@@ -129,20 +130,27 @@ 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),
+            (ModuloFn(RandomNumFn(), table_query.count())).label(RANDOM_LABEL)
+            if self.sample_config.randomizedSample
+            else None,
         )
-        return (
+        query = (
             session_query.order_by(RANDOM_LABEL)
-            .limit(self.sample_config.profileSample)
-            .cte(f"{self.get_sampler_table_name()}_rnd")
+            if self.sample_config.randomizedSample
+            else session_query
+        )
+        return query.limit(self.sample_config.profileSample).cte(
+            f"{self.get_sampler_table_name()}_rnd"
         )
 
     def get_dataset(self, column=None, **__) -> Union[DeclarativeMeta, AliasedClass]:
@@ -190,38 +198,38 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
                 for col in inspect(ds).c
                 if col.name != RANDOM_LABEL and col.name in names
             ]
-
-        try:
-            sqa_sample = (
-                self.client.query(*sqa_columns)
-                .select_from(ds)
-                .limit(self.sample_limit)
-                .all()
-            )
-        except Exception:
-            logger.debug(
-                "Cannot fetch sample data with random sampling. Falling back to 100 rows."
-            )
-            logger.debug(traceback.format_exc())
-            sqa_columns = list(inspect(self.raw_dataset).c)
-            sqa_sample = (
-                self.client.query(*sqa_columns)
-                .select_from(self.raw_dataset)
-                .limit(100)
-                .all()
+        with self.get_client() as client:
+            try:
+                sqa_sample = (
+                    client.query(*sqa_columns)
+                    .select_from(ds)
+                    .limit(self.sample_limit)
+                    .all()
+                )
+            except Exception:
+                logger.debug(
+                    "Cannot fetch sample data with random sampling. Falling back to 100 rows."
+                )
+                logger.debug(traceback.format_exc())
+                sqa_columns = list(inspect(self.raw_dataset).c)
+                sqa_sample = (
+                    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],
             )
 
-        return TableData(
-            columns=[column.name for column in sqa_columns],
-            rows=[list(row) for row in sqa_sample],
-        )
-
     def _fetch_sample_data_from_user_query(self) -> TableData:
         """Returns a table data object using results from query execution"""
         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:
@@ -238,10 +246,10 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
 
         stmt = text(f"{self.sample_query}")
         stmt = stmt.columns(*list(inspect(self.raw_dataset).c))
-
-        return self.client.query(stmt.subquery()).cte(
-            f"{self.get_sampler_table_name()}_user_sampled"
-        )
+        with self.get_client() as client:
+            return client.query(stmt.subquery()).cte(
+                f"{self.get_sampler_table_name()}_user_sampled"
+            )
 
     def _partitioned_table(self) -> Query:
         """Return the Query object for partitioned tables"""
@@ -258,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"""

metadata/sampler/sqlalchemy/snowflake/sampler.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -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")

metadata/sampler/sqlalchemy/trino/sampler.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -38,12 +38,13 @@ 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(
-            or_(
-                *[
-                    text(f'is_nan("{cols.name}") = False')
-                    for cols in sqa_columns
-                    if type(cols.type) in FLOAT_SET
-                ]
+        with self.get_client() as client:
+            return client.query(entity, label).where(
+                or_(
+                    *[
+                        text(f'is_nan("{cols.name}") = False')
+                        for cols in sqa_columns
+                        if type(cols.type) in FLOAT_SET
+                    ]
+                )
             )
-        )

metadata/timer/repeated_timer.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.