openmetadata-ingestion 1.6.7.1__py3-none-any.whl → 1.7.0.0rc1__py3-none-any.whl

metadata/sampler/sqlalchemy/sampler.py

@@ -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

@@ -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

@@ -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/utils/constants.py

@@ -85,6 +85,7 @@ THREE_MIN = 3 * 60
 UTF_8 = "utf-8"
 CHUNKSIZE = 200000
 DEFAULT_DATABASE = "default"
+DEFAULT_DASHBAORD = "default"
 BUILDER_PASSWORD_ATTR = "password"
 TIMEDELTA = "timedelta"
 COMPLEX_COLUMN_SEPARATOR = "_##"

metadata/utils/credentials.py

@@ -22,6 +22,7 @@ from google import auth
 from google.auth import impersonated_credentials
 
 from metadata.generated.schema.security.credentials.gcpCredentials import (
+    GcpADC,
     GCPCredentials,
     GcpCredentialsPath,
 )
@@ -89,7 +90,8 @@ def create_credential_tmp_file(credentials: dict) -> str:
 
 
 def build_google_credentials_dict(
-    gcp_values: Union[GcpCredentialsValues, GcpExternalAccount]
+    gcp_values: Union[GcpCredentialsValues, GcpExternalAccount],
+    single_project: bool = False,
 ) -> Dict[str, str]:
     """
     Given GcPCredentialsValues, build a dictionary as the JSON file
@@ -103,9 +105,18 @@ def build_google_credentials_dict(
         private_key_str = private_key_str.replace("\\n", "\n")
         validate_private_key(private_key_str)
 
+        if single_project:
+            project_id = (
+                gcp_values.projectId.root
+                if isinstance(gcp_values.projectId.root, str)
+                else gcp_values.projectId.root[0]
+            )
+        else:
+            project_id = gcp_values.projectId.root
+
         return {
             "type": gcp_values.type,
-            "project_id": gcp_values.projectId.root,
+            "project_id": project_id,
             "private_key_id": gcp_values.privateKeyId,
             "private_key": private_key_str,
             "client_email": gcp_values.clientEmail,
@@ -128,7 +139,9 @@ def build_google_credentials_dict(
     )
 
 
-def set_google_credentials(gcp_credentials: GCPCredentials) -> None:
+def set_google_credentials(
+    gcp_credentials: GCPCredentials, single_project: bool = False
+) -> None:
     """
     Set GCP credentials environment variable
     :param gcp_credentials: GCPCredentials
@@ -162,11 +175,19 @@ def set_google_credentials(gcp_credentials: GCPCredentials) -> None:
             )
             return
 
-        credentials_dict = build_google_credentials_dict(gcp_credentials.gcpConfig)
+        credentials_dict = build_google_credentials_dict(
+            gcp_credentials.gcpConfig, single_project
+        )
         tmp_credentials_file = create_credential_tmp_file(credentials=credentials_dict)
         os.environ[GOOGLE_CREDENTIALS] = tmp_credentials_file
         return
 
+    if isinstance(gcp_credentials.gcpConfig, GcpADC):
+        logger.info(
+            "Using Application Default Credentials to authenticate with GCP services."
+        )
+        return
+
     raise InvalidGcpConfigException(
         f"Error trying to set GCP credentials with {gcp_credentials}."
         " Check https://docs.open-metadata.org/connectors/database/bigquery "

metadata/utils/db_utils.py

@@ -18,9 +18,6 @@ from typing import Iterable
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.table import Table
-from metadata.generated.schema.entity.services.databaseService import (
-    DatabaseServiceType,
-)
 from metadata.generated.schema.type.entityLineage import Source as LineageSource
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper
@@ -36,8 +33,6 @@ from metadata.utils.logger import utils_logger
 
 logger = utils_logger()
 
-PUBLIC_SCHEMA = "public"
-
 
 def get_host_from_host_port(uri: str) -> str:
     """
@@ -69,7 +64,7 @@ def get_view_lineage(
         schema_name=schema_name,
         table_name=table_name,
     )
-    table_entity: Table = metadata.get_by_name(
+    table_entity = metadata.get_by_name(
         entity=Table,
         fqn=table_fqn,
     )
@@ -84,11 +79,6 @@ def get_view_lineage(
         lineage_parser = LineageParser(
             view_definition, dialect, timeout_seconds=timeout_seconds
         )
-
-        if table_entity.serviceType == DatabaseServiceType.Postgres:
-            # For Postgres, if schema is not defined, we need to use the public schema
-            schema_name = PUBLIC_SCHEMA
-
         if lineage_parser.source_tables and lineage_parser.target_tables:
             yield from get_lineage_by_query(
                 metadata,

metadata/utils/elasticsearch.py

@@ -17,6 +17,9 @@ from typing import List, Optional, TypeVar
 from pydantic import BaseModel
 
 from metadata.generated.schema.analytics.reportData import ReportData
+from metadata.generated.schema.entity.classification.classification import (
+    Classification,
+)
 from metadata.generated.schema.entity.classification.tag import Tag
 from metadata.generated.schema.entity.data.apiCollection import APICollection
 from metadata.generated.schema.entity.data.apiEndpoint import APIEndpoint
@@ -66,6 +69,7 @@ ES_INDEX_MAP = {
     GlossaryTerm.__name__: "glossary_term_search_index",
     MlModel.__name__: "mlmodel_search_index",
     Tag.__name__: "tag_search_index",
+    Classification.__name__: "classification_search_index",
     Container.__name__: "container_search_index",
     Query.__name__: "query_search_index",
     ReportData.__name__: "entity_report_data_index",

metadata/utils/helpers.py

@@ -15,6 +15,7 @@ Helpers module for ingestion related methods
 
 from __future__ import annotations
 
+import hashlib
 import itertools
 import re
 import shutil
@@ -32,6 +33,9 @@ from metadata.generated.schema.entity.data.chart import ChartType
 from metadata.generated.schema.entity.data.table import Column, Table
 from metadata.generated.schema.entity.feed.suggestion import Suggestion, SuggestionType
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
 from metadata.generated.schema.type.basic import EntityLink
 from metadata.generated.schema.type.tagLabel import TagLabel
 from metadata.utils.constants import DEFAULT_DATABASE
@@ -476,3 +480,90 @@ def init_staging_dir(directory: str) -> None:
     location = Path(directory)
     logger.info(f"Creating the directory to store staging data in {location}")
     location.mkdir(parents=True, exist_ok=True)
+
+
+def retry_with_docker_host(config: Optional[WorkflowSource] = None):
+    """
+    Retries the function on exception, replacing "localhost" with "host.docker.internal"
+    in the `hostPort` config if applicable. Raises the original exception if no `config` is found.
+    """
+
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            nonlocal config
+            try:
+                func(*args, **kwargs)
+            except Exception as error:
+                config = config or kwargs.get("config")
+                if not config:
+                    for argument in args:
+                        if isinstance(argument, WorkflowSource):
+                            config = argument
+                            break
+                    else:
+                        raise error
+
+                host_port_str = str(
+                    getattr(config.serviceConnection.root.config, "hostPort", None)
+                    or ""
+                )
+                if "localhost" not in host_port_str:
+                    raise error
+
+                host_port_type = type(config.serviceConnection.root.config.hostPort)
+                docker_host_port_str = host_port_str.replace(
+                    "localhost", "host.docker.internal"
+                )
+                config.serviceConnection.root.config.hostPort = host_port_type(
+                    docker_host_port_str
+                )
+                func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def get_query_hash(query: str) -> str:
+    result = hashlib.md5(query.encode())
+    return str(result.hexdigest())
+
+
+def evaluate_threshold(threshold: int, operator: str, result: int) -> bool:
+    """Evaluate the threshold against the result.
+
+    Args:
+        threshold: A string representing a comparison threshold (e.g., "< 5", ">= 10").
+        result: The integer value to compare against the threshold.
+
+    Returns:
+        True if the result satisfies the threshold condition, False otherwise.
+        If no comparison operator is provided, it defaults to less than or equal to comparison.
+        Returns False for invalid threshold formats.
+    """
+    import operator as op  # pylint: disable=import-outside-toplevel
+
+    operators = {
+        "<": op.lt,
+        "<=": op.le,
+        ">": op.gt,
+        ">=": op.ge,
+        "==": op.eq,
+        "!=": op.ne,
+    }
+    op_func = operators.get(operator, op.le)
+    try:
+        if op_func:
+            return op_func(result, threshold)
+    except ValueError:
+        return False
+
+    # Fallback:
+    logger.error(
+        f"Invalid threshold: {threshold}, "
+        "Allowed format: <, >, <=, >=, ==, !=. Example: >5"
+    )
+    raise ValueError(
+        f"Invalid threshold: {threshold}, "
+        "Allowed format: <, >, <=, >=, ==, !=. Example: >5"
+    )

metadata/utils/logger.py

@@ -60,6 +60,7 @@ class Loggers(Enum):
     TEST_SUITE = "TestSuite"
     QUERY_RUNNER = "QueryRunner"
     APP = "App"
+    REVERSE_INGESTION = "ReverseIngestion"
 
     @DynamicClassAttribute
     def value(self):
@@ -143,6 +144,14 @@ def ingestion_logger():
     return logging.getLogger(Loggers.INGESTION.value)
 
 
+def reverse_ingestion_logger():
+    """
+    Method to get the REVERSE INGESTION logger
+    """
+
+    return logging.getLogger(Loggers.REVERSE_INGESTION.value)
+
+
 def utils_logger():
     """
     Method to get the UTILS logger

metadata/utils/ssl_manager.py

@@ -17,6 +17,7 @@ import os
 import tempfile
 import traceback
 from functools import singledispatch, singledispatchmethod
+from ssl import CERT_REQUIRED, SSLContext
 from typing import List, Optional, Union, cast
 
 from pydantic import SecretStr
@@ -27,6 +28,9 @@ from metadata.generated.schema.entity.services.connections.connectionBasicType i
 from metadata.generated.schema.entity.services.connections.dashboard.qlikSenseConnection import (
     QlikSenseConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.cassandraConnection import (
+    CassandraConnection,
+)
 from metadata.generated.schema.entity.services.connections.database.dorisConnection import (
     DorisConnection,
 )
@@ -55,6 +59,7 @@ from metadata.generated.schema.entity.services.connections.pipeline.matillionCon
     MatillionConnection,
 )
 from metadata.generated.schema.security.ssl import verifySSLConfig
+from metadata.generated.schema.security.ssl.verifySSLConfig import SslMode
 from metadata.ingestion.connections.builders import init_empty_connection_arguments
 from metadata.ingestion.models.custom_pydantic import CustomSecretStr
 from metadata.ingestion.source.connections import get_connection
@@ -223,6 +228,25 @@ class SSLManager:
         ] = self.cert_file_path
         return connection
 
+    @setup_ssl.register(CassandraConnection)
+    def _(self, connection):
+        connection = cast(CassandraConnection, connection)
+
+        ssl_context = None
+        if connection.sslMode != SslMode.disable:
+            ssl_context = SSLContext()
+            ssl_context.load_verify_locations(cafile=self.ca_file_path)
+            ssl_context.verify_mode = CERT_REQUIRED
+            ssl_context.load_cert_chain(
+                certfile=self.cert_file_path, keyfile=self.key_file_path
+            )
+
+        connection.connectionArguments = (
+            connection.connectionArguments or init_empty_connection_arguments()
+        )
+        connection.connectionArguments.root["ssl_context"] = ssl_context
+        return connection
+
 
 @singledispatch
 def check_ssl_and_init(
@@ -340,6 +364,17 @@ def _(connection):
     return None
 
 
+@check_ssl_and_init.register(CassandraConnection)
+def _(connection):
+    service_connection = cast(CassandraConnection, connection)
+    ssl: Optional[verifySSLConfig.SslConfig] = service_connection.sslConfig
+    if ssl and (ssl.root.caCertificate or ssl.root.sslCertificate or ssl.root.sslKey):
+        return SSLManager(
+            ca=ssl.root.caCertificate, cert=ssl.root.sslCertificate, key=ssl.root.sslKey
+        )
+    return None
+
+
 def get_ssl_connection(service_config):
     try:
         # To be cleaned up as part of https://github.com/open-metadata/OpenMetadata/issues/15913

metadata/utils/tag_utils.py

@@ -20,6 +20,9 @@ from metadata.generated.schema.api.classification.createClassification import (
     CreateClassificationRequest,
 )
 from metadata.generated.schema.api.classification.createTag import CreateTagRequest
+from metadata.generated.schema.entity.classification.classification import (
+    Classification,
+)
 from metadata.generated.schema.entity.classification.tag import Tag
 from metadata.generated.schema.entity.data.glossaryTerm import GlossaryTerm
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
@@ -29,6 +32,7 @@ from metadata.generated.schema.type.basic import (
     EntityName,
     FullyQualifiedEntityName,
     Markdown,
+    ProviderType,
 )
 from metadata.generated.schema.type.tagLabel import (
     LabelType,
@@ -46,6 +50,7 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 
+# pylint: disable=too-many-arguments
 def get_ometa_tag_and_classification(
     tags: List[str],
     classification_name: str,
@@ -53,39 +58,79 @@ def get_ometa_tag_and_classification(
     classification_description: str,
     include_tags: bool = True,
     tag_fqn: Optional[FullyQualifiedEntityName] = None,
+    metadata: Optional[OpenMetadata] = None,
+    system_tags: bool = False,
 ) -> Iterable[Either[OMetaTagAndClassification]]:
     """
     Returns the OMetaTagAndClassification object
     """
-    if include_tags:
-        for tag in tags:
-            try:
-                classification = OMetaTagAndClassification(
-                    fqn=tag_fqn,
-                    classification_request=CreateClassificationRequest(
-                        name=EntityName(classification_name),
-                        description=Markdown(classification_description),
-                    ),
-                    tag_request=CreateTagRequest(
-                        classification=FullyQualifiedEntityName(classification_name),
-                        name=EntityName(tag),
-                        description=Markdown(tag_description)
-                        if tag_description
-                        else None,
-                    ),
-                )
-                yield Either(right=classification)
-                logger.debug(
-                    f"Classification {classification_name}, Tag {tag} Ingested"
-                )
-            except Exception as err:
-                yield Either(
-                    left=StackTraceError(
-                        name=tag,
-                        error=f"Error yielding tag [{tag}]: [{err}]",
-                        stackTrace=traceback.format_exc(),
+    if not include_tags:
+        return
+
+    if system_tags:
+        # Checking for system classification
+        for classification_entity in (
+            metadata.es_search_from_fqn(
+                entity_type=Classification, fqn_search_string=classification_name
+            )
+            or []
+        ):
+            if (
+                classification_entity.provider == ProviderType.system
+                and classification_entity.name.root.lower()
+                == classification_name.lower()
+            ):
+                classification_name = classification_entity.name.root
+                classification_description = classification_entity.description.root
+                break
+
+    for tag in tags:
+        specific_tag_description = tag_description
+        try:
+            if system_tags:
+                # Checking for system tag
+                for tag_entity in (
+                    metadata.es_search_from_fqn(
+                        entity_type=Tag,
+                        fqn_search_string=f"{classification_name}.{tag}",
                     )
+                    or []
+                ):
+                    if (
+                        tag_entity.provider == ProviderType.system
+                        and tag_entity.classification.name == classification_name
+                        and tag_entity.name.root.lower() == tag.lower()
+                    ):
+                        tag = tag_entity.name.root
+                        specific_tag_description = tag_entity.description.root
+                        break
+
+            classification = OMetaTagAndClassification(
+                fqn=tag_fqn,
+                classification_request=CreateClassificationRequest(
+                    name=EntityName(classification_name),
+                    description=Markdown(classification_description),
+                ),
+                tag_request=CreateTagRequest(
+                    classification=FullyQualifiedEntityName(classification_name),
+                    name=EntityName(tag),
+                    description=(
+                        Markdown(specific_tag_description)
+                        if specific_tag_description
+                        else None
+                    ),
+                ),
+            )
+            yield Either(right=classification)
+            logger.debug(f"Classification {classification_name}, Tag {tag} Ingested")
+        except Exception as err:
+            yield Either(
+                left=StackTraceError(
+                    name=tag,
+                    error=f"Error yielding tag [{tag}]: [{err}]",
+                    stackTrace=traceback.format_exc(),
                 )
+            )
 
 
 @functools.lru_cache(maxsize=512)

metadata/workflow/classification.py

@@ -29,6 +29,9 @@ class AutoClassificationWorkflow(ProfilerWorkflow):
     """Auto Classification workflow implementation. Based on the Profiler logic with different steps"""
 
     def set_steps(self):
+        # NOTE: Call test_connection to update host value before creating the source class
+        self.test_connection()
+
         source_class = self._get_source_class()
         self.source = source_class.create(self.config.model_dump(), self.metadata)