openmetadata-ingestion 1.7.1.4__py3-none-any.whl → 1.7.2.0__py3-none-any.whl

metadata/ingestion/source/database/databricks/metadata.py

@@ -50,6 +50,7 @@ from metadata.ingestion.source.database.databricks.queries import (
     DATABRICKS_GET_CATALOGS,
     DATABRICKS_GET_CATALOGS_TAGS,
     DATABRICKS_GET_COLUMN_TAGS,
+    DATABRICKS_GET_SCHEMA_COMMENTS,
     DATABRICKS_GET_SCHEMA_TAGS,
     DATABRICKS_GET_TABLE_COMMENTS,
     DATABRICKS_GET_TABLE_TAGS,
@@ -65,6 +66,8 @@ from metadata.utils.filters import filter_by_database
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqlalchemy_utils import (
     get_all_view_definitions,
+    get_schema_comment_result_wrapper,
+    get_schema_comment_results,
     get_table_comment_result_wrapper,
     get_table_comment_results,
     get_view_definition_wrapper,
@@ -325,6 +328,24 @@ def get_table_comment_result(
     )
 
 
+@reflection.cache
+def get_schema_comment_result(
+    self,
+    connection,
+    query,
+    database,
+    schema,
+    **kw,  # pylint: disable=unused-argument
+):
+    return get_schema_comment_result_wrapper(
+        self,
+        connection,
+        query=query,
+        database=database,
+        schema=schema,
+    )
+
+
 @reflection.cache
 def get_table_ddl(
     self, connection, table_name, schema=None, **kw
@@ -413,6 +434,8 @@ DatabricksDialect.get_view_definition = get_view_definition
 DatabricksDialect.get_table_names = get_table_names
 DatabricksDialect.get_all_view_definitions = get_all_view_definitions
 DatabricksDialect.get_table_comment_results = get_table_comment_results
+DatabricksDialect.get_schema_comment_results = get_schema_comment_results
+DatabricksDialect.get_schema_comment_result = get_schema_comment_result
 DatabricksDialect.get_table_comment_result = get_table_comment_result
 reflection.Inspector.get_schema_names = get_schema_names_reflection
 reflection.Inspector.get_table_ddl = get_table_ddl
@@ -756,6 +779,33 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
                 )
             )
 
+    def get_schema_description(self, schema_name: str) -> str:
+        description = None
+        try:
+            query = DATABRICKS_GET_SCHEMA_COMMENTS.format(
+                database_name=self.context.get().database,
+                schema_name=schema_name,
+            )
+            cursor = self.inspector.dialect.get_schema_comment_result(
+                connection=self.connection,
+                query=query,
+                database=self.context.get().database,
+                schema=schema_name,
+            )
+            for result in list(cursor):
+                data = result.values()
+                if data[0] and data[0].strip() == "Comment":
+                    description = data[1] if data and data[1] else None
+                    return description
+
+        # Catch any exception without breaking the ingestion
+        except Exception as exep:  # pylint: disable=broad-except
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Schema description error for schema [{schema_name}]: {exep}"
+            )
+        return description
+
     def get_table_description(
         self, schema_name: str, table_name: str, inspector: Inspector
     ) -> str:

metadata/ingestion/source/database/databricks/queries.py

@@ -52,6 +52,10 @@ DATABRICKS_GET_TABLE_COMMENTS = (
     "DESCRIBE TABLE EXTENDED `{database_name}`.`{schema_name}`.`{table_name}`"
 )
 
+DATABRICKS_GET_SCHEMA_COMMENTS = (
+    "DESCRIBE SCHEMA EXTENDED `{database_name}`.`{schema_name}`"
+)
+
 DATABRICKS_GET_CATALOGS = "SHOW CATALOGS"
 
 DATABRICKS_GET_CATALOGS_TAGS = textwrap.dedent(

metadata/ingestion/source/database/databricks/service_spec.py

@@ -9,6 +9,7 @@ from metadata.ingestion.source.database.databricks.usage import DatabricksUsageS
 from metadata.profiler.interface.sqlalchemy.databricks.profiler_interface import (
     DatabricksProfilerInterface,
 )
+from metadata.sampler.sqlalchemy.databricks.sampler import DatabricksSamplerInterface
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
@@ -17,4 +18,5 @@ ServiceSpec = DefaultDatabaseSpec(
     usage_source_class=DatabricksUsageSource,
     profiler_class=DatabricksProfilerInterface,
     test_suite_class=DatabricksTestSuiteInterface,
+    sampler_class=DatabricksSamplerInterface,
 )

metadata/ingestion/source/database/db2/connection.py

@@ -33,13 +33,28 @@ from metadata.ingestion.connections.builders import (
 )
 from metadata.ingestion.connections.test_connections import test_connection_db_common
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.db2.utils import (
+    check_clidriver_version,
+    install_clidriver,
+)
 from metadata.utils.constants import THREE_MIN, UTF_8
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 def get_connection(connection: Db2Connection) -> Engine:
     """
     Create connection
     """
+    # Install ibm_db with specific version
+    clidriver_version = connection.clidriverVersion
+
+    if clidriver_version:
+        clidriver_version = check_clidriver_version(clidriver_version)
+        if clidriver_version:
+            install_clidriver(clidriver_version.value)
+
     # prepare license
     # pylint: disable=import-outside-toplevel
     if connection.license and connection.licenseFileName:

metadata/ingestion/source/database/db2/utils.py

@@ -12,9 +12,33 @@
 """
 Module to define overriden dialect methods
 """
+from enum import Enum
+
 from sqlalchemy import and_, join, sql
 from sqlalchemy.engine import reflection
 
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+BASE_CLIDRIVER_URL = (
+    "https://public.dhe.ibm.com/ibmdl/export/pub/software/data/db2/drivers/odbc_cli"
+)
+
+
+class DB2CLIDriverVersions(Enum):
+    """
+    Enum for the DB2 CLI Driver versions
+    """
+
+    V11_1_4 = "11.1.4"
+    V11_5_4 = "11.5.4"
+    V11_5_5 = "11.5.5"
+    V11_5_6 = "11.5.6"
+    V11_5_8 = "11.5.8"
+    V11_5_9 = "11.5.9"
+    V12_1_0 = "12.1.0"
+
 
 @reflection.cache
 def get_unique_constraints(
@@ -61,3 +85,109 @@ def get_unique_constraints(
                 }
             )
     return unique_consts
+
+
+def check_clidriver_version(clidriver_version: str):
+    """
+    Check if the CLI Driver version is valid
+    """
+    if clidriver_version not in [v.value for v in DB2CLIDriverVersions]:
+        logger.warning(f"Invalid CLI Driver version provided: {clidriver_version}")
+        return None
+    return DB2CLIDriverVersions(clidriver_version)
+
+
+# pylint: disable=too-many-statements,too-many-branches
+def install_clidriver(clidriver_version: str) -> None:
+    """
+    Install the CLI Driver for DB2
+    """
+    # pylint: disable=import-outside-toplevel
+    import os
+    import platform
+    import subprocess
+    import sys
+    from urllib.request import URLError, urlopen
+
+    import pkg_resources
+
+    clidriver_version = f"v{clidriver_version}"
+    system = platform.system().lower()
+    is_64bits = platform.architecture()[0] == "64bit"
+    clidriver_url = None
+    default_clidriver_url = None
+
+    def is_valid_url(url: str) -> bool:
+        """Check if the URL is valid and accessible"""
+        try:
+            with urlopen(url) as _:
+                return True
+        except URLError:
+            return False
+
+    if system == "darwin":  # macOS
+        machine = platform.machine().lower()
+        if machine == "arm64":  # Apple Silicon
+            default_clidriver_url = f"{BASE_CLIDRIVER_URL}/macarm64_odbc_cli.tar.gz"
+            clidriver_url = f"{BASE_CLIDRIVER_URL}/macarm64_odbc_cli.tar.gz"
+        elif machine == "x86_64":  # Intel
+            default_clidriver_url = f"{BASE_CLIDRIVER_URL}/macos64_odbc_cli.tar.gz"
+            clidriver_url = (
+                f"{BASE_CLIDRIVER_URL}/{str(clidriver_version)}/macos64_odbc_cli.tar.gz"
+            )
+    elif system == "linux":
+        if is_64bits:
+            default_clidriver_url = f"{BASE_CLIDRIVER_URL}/linuxx64_odbc_cli.tar.gz"
+            clidriver_url = f"{BASE_CLIDRIVER_URL}/{str(clidriver_version)}/linuxx64_odbc_cli.tar.gz"
+        else:
+            default_clidriver_url = f"{BASE_CLIDRIVER_URL}/linuxia32_odbc_cli.tar.gz"
+            clidriver_url = f"{BASE_CLIDRIVER_URL}/{str(clidriver_version)}/linuxia32_odbc_cli.tar.gz"
+    elif system == "windows":
+        if is_64bits:
+            default_clidriver_url = f"{BASE_CLIDRIVER_URL}/ntx64_odbc_cli.zip"
+            clidriver_url = (
+                f"{BASE_CLIDRIVER_URL}/{str(clidriver_version)}/ntx64_odbc_cli.zip"
+            )
+        else:
+            default_clidriver_url = f"{BASE_CLIDRIVER_URL}/nt32_odbc_cli.zip"
+            clidriver_url = (
+                f"{BASE_CLIDRIVER_URL}/{str(clidriver_version)}/nt32_odbc_cli.zip"
+            )
+    else:
+        logger.error(
+            f"Unsupported operating system for db2 driver installation: {system}"
+        )
+        return None
+
+    # set env variables for CLIDRIVER_VERSION and IBM_DB_INSTALLER_URL
+    os.environ["CLIDRIVER_VERSION"] = clidriver_version
+    if is_valid_url(clidriver_url):
+        os.environ["IBM_DB_INSTALLER_URL"] = clidriver_url
+    else:
+        os.environ["IBM_DB_INSTALLER_URL"] = default_clidriver_url
+    logger.info(f"Set IBM_DB_INSTALLER_URL to {os.environ['IBM_DB_INSTALLER_URL']}")
+    logger.info(f"Set CLIDRIVER_VERSION to {os.environ['CLIDRIVER_VERSION']}")
+    # Uninstall ibm_db if it is already installed
+    try:
+        pkg_resources.get_distribution("ibm_db")
+        # If we get here, ibm_db is installed, so uninstall it first
+        subprocess.check_call(
+            [sys.executable, "-m", "pip", "uninstall", "-y", "ibm_db"]
+        )
+    except pkg_resources.DistributionNotFound:
+        # ibm_db is not installed, proceed with installation
+        pass
+    # Install ibm_db with specific flags
+    subprocess.check_call(
+        [
+            sys.executable,
+            "-m",
+            "pip",
+            "install",
+            "ibm_db~=3.2.6",
+            "--no-binary",
+            ":all:",
+            "--no-cache-dir",
+        ]
+    )
+    return None

metadata/ingestion/source/database/dbt/metadata.py

@@ -97,6 +97,7 @@ from metadata.ingestion.source.database.dbt.models import DbtMeta
 from metadata.utils import fqn
 from metadata.utils.elasticsearch import get_entity_from_es_result
 from metadata.utils.entity_link import get_table_fqn
+from metadata.utils.filters import filter_by_tag
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.tag_utils import get_ometa_tag_and_classification, get_tag_labels
 from metadata.utils.time_utils import datetime_to_timestamp
@@ -245,6 +246,18 @@ class DbtSource(DbtServiceSource):
                             f"Unable to find the node or columns in the catalog file for dbt node: {key}"
                         )
 
+    def filter_tags(self, tags: List[str]) -> List[str]:
+        """
+        Filter tags based on tag filter pattern if configured
+        """
+        if self.source_config.tagFilterPattern:
+            return [
+                tag
+                for tag in tags
+                if not filter_by_tag(self.source_config.tagFilterPattern, tag)
+            ]
+        return tags
+
     def yield_dbt_tags(
         self, dbt_objects: DbtObjects
     ) -> Iterable[Either[OMetaTagAndClassification]]:
@@ -272,13 +285,13 @@ class DbtSource(DbtServiceSource):
                     # Add the tags from the model
                     model_tags = manifest_node.tags
                     if model_tags:
-                        dbt_tags_list.extend(model_tags)
+                        dbt_tags_list.extend(self.filter_tags(model_tags))
 
                     # Add the tags from the columns
                     for _, column in manifest_node.columns.items():
                         column_tags = column.tags
                         if column_tags:
-                            dbt_tags_list.extend(column_tags)
+                            dbt_tags_list.extend(self.filter_tags(column_tags))
                 except Exception as exc:
                     yield Either(
                         left=StackTraceError(
@@ -509,6 +522,7 @@ class DbtSource(DbtServiceSource):
 
                     dbt_table_tags_list = []
                     if manifest_node.tags:
+                        manifest_node.tags = self.filter_tags(manifest_node.tags)
                         dbt_table_tags_list = (
                             get_tag_labels(
                                 metadata=self.metadata,
@@ -669,6 +683,7 @@ class DbtSource(DbtServiceSource):
                     column_description = catalog_column.comment
 
                 dbt_column_tag_list = []
+                manifest_column.tags = self.filter_tags(manifest_column.tags)
                 dbt_column_tag_list.extend(
                     get_tag_labels(
                         metadata=self.metadata,

metadata/ingestion/source/database/doris/metadata.py

@@ -60,6 +60,7 @@ RELKIND_MAP = {
     "Doris": TableType.Regular,
     "View": TableType.View,
     "MEMORY": TableType.View,
+    "OLAP": TableType.MaterializedView,
 }
 
 DorisDialect.get_table_names_and_type = get_table_names_and_type
@@ -179,7 +180,9 @@ class DorisSource(CommonDbSourceService):
         logic on how to handle table types, e.g., external, foreign,...
         """
         tables = [
-            TableNameAndType(name=name, type_=RELKIND_MAP.get(engine))
+            TableNameAndType(
+                name=name, type_=RELKIND_MAP.get(engine, TableType.Regular)
+            )
             for name, engine in self.connection.execute(
                 sql.text(DORIS_GET_TABLE_NAMES), {"schema": schema_name}
             )

metadata/ingestion/source/database/snowflake/metadata.py

@@ -52,6 +52,7 @@ from metadata.generated.schema.type.basic import (
     FullyQualifiedEntityName,
     SourceUrl,
 )
+from metadata.generated.schema.type.entityReferenceList import EntityReferenceList
 from metadata.ingestion.api.delete import delete_entity_by_name
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
@@ -908,3 +909,21 @@ class SnowflakeSource(
             schema_name=self.context.get().database_schema,
             account_usage=self.service_connection.accountUsageSchema,
         )
+
+    def get_owner_ref(self, table_name: str) -> Optional[EntityReferenceList]:
+        """
+        Method to process the table owners
+
+        Snowflake uses a role-based ownership model, not a user-based one.
+        This means that ownership of database objects, such as tables, is assigned
+        to roles rather than individual users.
+
+        As OpenMetadata currently does not support role-based ownership assignment,
+        we are unable to retrieve or associate a meaningful table owner using this method.
+        Therefore, this function will return `None` or a placeholder, and ownership
+        metadata will not be populated in the OpenMetadata ingestion process.
+        """
+        logger.debug(
+            f"Processing ownership is not supported for {self.service_connection.type.name}"
+        )
+        return None

metadata/ingestion/source/database/trino/queries.py

@@ -20,7 +20,7 @@ TRINO_SQL_STATEMENT = textwrap.dedent(
       "user" as user_name,
       "started" as start_time,
       "end" as end_time
-    from "system"."runtime"."queries"
+    from {query_history_table}
     WHERE "query" NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
     AND "query" NOT LIKE '/* {{"app": "dbt", %%}} */%%'
     AND CAST("started" AS date)  >= date_parse('{start_time}', '%Y-%m-%d %H:%i:%s')

metadata/ingestion/source/database/trino/query_parser.py

@@ -12,6 +12,7 @@
 Trino usage module
 """
 from abc import ABC
+from datetime import datetime
 from typing import Optional
 
 from metadata.generated.schema.entity.services.connections.database.trinoConnection import (
@@ -44,3 +45,17 @@ class TrinoQueryParserSource(QueryParserSource, ABC):
                 f"Expected TrinoConnection, but got {connection}"
             )
         return cls(config, metadata)
+
+    def get_sql_statement(self, start_time: datetime, end_time: datetime) -> str:
+        """
+        returns sql statement to fetch query logs.
+
+        Override if we have specific parameters
+        """
+        return self.sql_stmt.format(
+            start_time=start_time,
+            end_time=end_time,
+            filters=self.get_filters(),
+            result_limit=self.source_config.resultLimit,
+            query_history_table=self.service_connection.queryHistoryTable,
+        )

metadata/ingestion/source/database/unitycatalog/service_spec.py

@@ -11,7 +11,7 @@ from metadata.ingestion.source.database.unitycatalog.usage import (
 from metadata.profiler.interface.sqlalchemy.unity_catalog.profiler_interface import (
     UnityCatalogProfilerInterface,
 )
-from metadata.profiler.interface.sqlalchemy.unity_catalog.sampler_interface import (
+from metadata.sampler.sqlalchemy.unitycatalog.sampler import (
     UnityCatalogSamplerInterface,
 )
 from metadata.utils.service_spec.default import DefaultDatabaseSpec

metadata/ingestion/source/pipeline/dbtcloud/metadata.py

@@ -168,13 +168,6 @@ class DbtcloudSource(PipelineServiceSource):
                     entity=Pipeline, fqn=pipeline_fqn
                 )
 
-                lineage_details = LineageDetails(
-                    pipeline=EntityReference(
-                        id=pipeline_entity.id.root, type="pipeline"
-                    ),
-                    source=LineageSource.PipelineLineage,
-                )
-
                 dbt_models = self.client.get_model_details(
                     job_id=pipeline_details.id, run_id=self.context.get().latest_run_id
                 )
@@ -222,6 +215,13 @@ class DbtcloudSource(PipelineServiceSource):
                                 if from_entity is None:
                                     continue
 
+                                lineage_details = LineageDetails(
+                                    pipeline=EntityReference(
+                                        id=pipeline_entity.id.root, type="pipeline"
+                                    ),
+                                    source=LineageSource.PipelineLineage,
+                                )
+
                                 yield Either(
                                     right=AddLineageRequest(
                                         edge=EntitiesEdge(

metadata/pii/algorithms/classifiers.py

@@ -38,7 +38,11 @@ from metadata.pii.algorithms.feature_extraction import (
     split_column_name,
 )
 from metadata.pii.algorithms.preprocessing import preprocess_values
-from metadata.pii.algorithms.presidio_patches import url_patcher
+from metadata.pii.algorithms.presidio_patches import (
+    combine_patchers,
+    date_time_patcher,
+    url_patcher,
+)
 from metadata.pii.algorithms.presidio_utils import (
     build_analyzer_engine,
     set_presidio_logger_level,
@@ -119,7 +123,7 @@ class HeuristicPIIClassifier(ColumnClassifier[PIITag]):
             self._presidio_analyzer,
             str_values,
             context=context,
-            recognizer_result_patcher=url_patcher,
+            recognizer_result_patcher=combine_patchers(date_time_patcher, url_patcher),
         )
 
         column_name_matches: Set[PIITag] = set()

metadata/pii/algorithms/presidio_patches.py

@@ -13,6 +13,7 @@ Patch the Presidio recognizer results to make adapt them to specific use cases.
 """
 from typing import List, Protocol, Sequence
 
+from dateutil.parser import parse
 from presidio_analyzer import RecognizerResult
 
 
@@ -29,6 +30,24 @@ class PresidioRecognizerResultPatcher(Protocol):
         ...
 
 
+def combine_patchers(
+    *patchers: PresidioRecognizerResultPatcher,
+) -> PresidioRecognizerResultPatcher:
+    """
+    Combine multiple patchers into one.
+    This allows us to apply multiple patches in sequence.
+    """
+
+    def combined_patcher(
+        recognizer_results: Sequence[RecognizerResult], text: str
+    ) -> Sequence[RecognizerResult]:
+        for patcher in patchers:
+            recognizer_results = patcher(recognizer_results, text)
+        return recognizer_results
+
+    return combined_patcher
+
+
 def url_patcher(
     recognizer_results: Sequence[RecognizerResult], text: str
 ) -> Sequence[RecognizerResult]:
@@ -43,3 +62,22 @@ def url_patcher(
                 continue
         patched_result.append(result)
     return patched_result
+
+
+def date_time_patcher(
+    recognizer_results: Sequence[RecognizerResult], text: str
+) -> Sequence[RecognizerResult]:
+    """
+    Patch the recognizer result to remove date time false positive with date.
+    """
+    patched_result: List[RecognizerResult] = []
+    for result in recognizer_results:
+        if result.entity_type == "DATE_TIME":
+            # try to parse using dateutils, if it fails, skip the result
+            try:
+                _ = parse(text[result.start : result.end])
+            except ValueError:
+                # if parsing fails, skip the result
+                continue
+        patched_result.append(result)
+    return patched_result

metadata/sampler/sqlalchemy/azuresql/sampler.py

@@ -15,7 +15,7 @@ for the profiler
 from typing import List, Optional
 
 from sqlalchemy import Column, Table, text
-from sqlalchemy.orm import Query
+from sqlalchemy.sql.selectable import CTE
 
 from metadata.generated.schema.entity.data.table import TableData, TableType
 from metadata.sampler.sqlalchemy.sampler import ProfileSampleType, SQASampler
@@ -49,13 +49,12 @@ class AzureSQLSampler(SQASampler):
 
         return selectable
 
-    def get_sample_query(self, *, column=None) -> Query:
-        """get query for sample data"""
+    def get_sample_query(self, *, column=None) -> CTE:
+        """Override the base method as ROWS or PERCENT sampling handled through the tablesample clause"""
         rnd = self._base_sample_query(column).cte(
             f"{self.get_sampler_table_name()}_rnd"
         )
-        with self.get_client() as client:
-            query = client.query(rnd)
+        query = self.get_client().query(rnd)
         return query.cte(f"{self.get_sampler_table_name()}_sample")
 
     def fetch_sample_data(self, columns: Optional[List[Column]] = None) -> TableData:

metadata/sampler/sqlalchemy/databricks/sampler.py

@@ -0,0 +1,26 @@
+#  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
+#  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.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Helper module to handle data sampling for the profiler
+"""
+from metadata.ingestion.source.database.databricks.connection import (
+    get_connection as databricks_get_connection,
+)
+from metadata.sampler.sqlalchemy.sampler import SQASampler
+
+
+class DatabricksSamplerInterface(SQASampler):
+    def get_client(self):
+        """client is the session for SQA"""
+        self.connection = databricks_get_connection(self.service_connection_config)
+        client = super().get_client()
+        self.set_catalog(client)
+        return client

metadata/sampler/sqlalchemy/mssql/sampler.py

@@ -44,10 +44,9 @@ class MssqlSampler(SQASampler):
         return selectable
 
     def get_sample_query(self, *, column=None) -> CTE:
-        """get query for sample data"""
+        """Override the base method as ROWS or PERCENT sampling handled through the tablesample clause"""
         rnd = self._base_sample_query(column).cte(
             f"{self.get_sampler_table_name()}_rnd"
         )
-        with self.get_client() as client:
-            query = client.query(rnd)
+        query = self.get_client().query(rnd)
         return query.cte(f"{self.get_sampler_table_name()}_sample")

metadata/sampler/sqlalchemy/sampler.py

@@ -136,6 +136,8 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
             ).cte(f"{self.get_sampler_table_name()}_sample")
 
         table_query = client.query(self.raw_dataset)
+        if self.partition_details:
+            table_query = self.get_partitioned_query(table_query)
         session_query = self._base_sample_query(
             column,
             (ModuloFn(RandomNumFn(), table_query.count())).label(RANDOM_LABEL)

metadata/sampler/sqlalchemy/snowflake/sampler.py

@@ -87,10 +87,9 @@ class SnowflakeSampler(SQASampler):
         )
 
     def get_sample_query(self, *, column=None) -> CTE:
-        """get query for sample data"""
+        """Override the base method as ROWS or PERCENT sampling handled through the tablesample clause"""
         rnd = self._base_sample_query(column).cte(
             f"{self.get_sampler_table_name()}_rnd"
         )
-        with self.get_client() as client:
-            query = client.query(rnd)
+        query = self.get_client().query(rnd)
         return query.cte(f"{self.get_sampler_table_name()}_sample")

metadata/utils/filters.py

@@ -310,3 +310,16 @@ def filter_by_collection(
     :return: True for filtering, False otherwise
     """
     return _filter(collection_pattern, collection_name)
+
+
+def filter_by_tag(tag_pattern: Optional[FilterPattern], tag_name: str) -> bool:
+    """
+    Return True if the models needs to be filtered, False otherwise
+
+    Include takes precedence over exclude
+
+    :param tag_pattern: Model defining tag filtering logic
+    :param tag_name: tag name
+    :return: True for filtering, False otherwise
+    """
+    return _filter(tag_pattern, tag_name)

metadata/utils/helpers.py

@@ -94,6 +94,7 @@ om_chart_type_dict = {
     "levelTable": ChartType.Table,
     "dist_bar": ChartType.Bar,
     "bar": ChartType.Bar,
+    "vertical_bar": ChartType.Bar,
     "box_plot": ChartType.BoxPlot,
     "box": ChartType.BoxPlot,
     "boxplot": ChartType.BoxPlot,