openmetadata-ingestion 1.7.5.2__py3-none-any.whl → 1.7.5.3__py3-none-any.whl

metadata/ingestion/source/api/rest/metadata.py

@@ -300,38 +300,80 @@ class RestSource(ApiServiceSource):
             logger.warning(f"Error while parsing response schema: {err}")
         return None
 
-    def process_schema_fields(self, schema_ref: str) -> Optional[List[FieldModel]]:
+    def process_schema_fields(
+        self, schema_ref: str, parent_refs: Optional[List[str]] = None
+    ) -> Optional[List[FieldModel]]:
         try:
+            if parent_refs is None:
+                parent_refs = []
             schema_name = schema_ref.split("/")[-1]
             schema_fields = (
                 self.json_response.get("components").get("schemas").get(schema_name)
             )
-
+            parent_refs.append(schema_ref)
             fetched_fields = []
             for key, val in schema_fields.get("properties", {}).items():
                 dtype = val.get("type")
                 if dtype:
+                    dtype = "INT" if dtype.upper() == "INTEGER" else dtype
                     parsed_dtype = (
                         DataTypeTopic[dtype.upper()]
                         if dtype.upper() in DataTypeTopic.__members__
                         else DataTypeTopic.UNKNOWN
                     )
-                    fetched_fields.append(FieldModel(name=key, dataType=parsed_dtype))
+                    children = None
+                    if parsed_dtype.value == DataTypeTopic.ARRAY.value:
+                        # If field of array type then parse children
+                        children_ref = val.get("items", {}).get("$ref")
+                        if children_ref:
+                            # check infinite recursion by checking pre-processed schemas(parent_refs)
+                            if children_ref not in parent_refs:
+                                logger.debug(
+                                    f"Processing array fields inside schema: {children_ref}"
+                                )
+                                children = self.process_schema_fields(
+                                    children_ref, parent_refs
+                                )
+                                logger.debug(
+                                    f"Completed processing array fields inside schema: {children_ref}"
+                                )
+                            else:
+                                logger.debug(
+                                    f"Skipping array fields inside schema: {children_ref} to avoid infinite recursion"
+                                )
+                    fetched_fields.append(
+                        FieldModel(name=key, dataType=parsed_dtype, children=children)
+                    )
                 else:
                     # If type of field is not defined then check for sub-schema
                     # Check if it's `object` type field
-                    # check infinite recrusrion by comparing with parent(schema_ref)
-                    object_children = None
-                    if val.get("$ref") and val.get("$ref") != schema_ref:
-                        object_children = self.process_schema_fields(val.get("$ref"))
+                    children = None
+                    if val.get("$ref"):
+                        # check infinite recursion by checking pre-processed schemas(parent_refs)
+                        if val.get("$ref") not in parent_refs:
+                            children = self.process_schema_fields(
+                                val.get("$ref"), parent_refs
+                            )
+                        else:
+                            logger.debug(
+                                f"Skipping object fields inside schema: {val.get('$ref')} to avoid infinite recursion"
+                            )
                     fetched_fields.append(
                         FieldModel(
                             name=key,
                             dataType=DataTypeTopic.UNKNOWN,
-                            children=object_children,
+                            dataTypeDisplay="OBJECT",
+                            children=children,
                         )
                     )
+            if parent_refs and (schema_ref in parent_refs):
+                parent_refs.pop()
             return fetched_fields
         except Exception as err:
             logger.warning(f"Error while processing schema fields: {err}")
+            if parent_refs and (schema_ref in parent_refs):
+                parent_refs.pop()
+                logger.debug(
+                    f"Popping {schema_ref} from parent_refs due to processing error"
+                )
         return None

metadata/ingestion/source/dashboard/powerbi/metadata.py

@@ -455,7 +455,7 @@ class PowerbiSource(DashboardServiceSource):
                 measures.append(Column(**parsed_measure.model_dump()))
             except Exception as err:
                 logger.debug(traceback.format_exc())
-                logger.warning(f"Error processing datamodel nested measure: {err}")
+                logger.debug(f"Error processing datamodel nested measure: {err}")
         return measures
 
     def _get_child_columns(self, table: PowerBiTable) -> List[Column]:
@@ -631,20 +631,26 @@ class PowerbiSource(DashboardServiceSource):
                 entity=Dashboard,
                 fqn=report_fqn,
             )
-            datamodel_fqn = fqn.build(
-                self.metadata,
-                entity_type=DashboardDataModel,
-                service_name=self.context.get().dashboard_service,
-                data_model_name=dashboard_details.datasetId,
-            )
-            datamodel_entity = self.metadata.get_by_name(
-                entity=DashboardDataModel,
-                fqn=datamodel_fqn,
-            )
+            if dashboard_details.datasetId:
+                datamodel_fqn = fqn.build(
+                    self.metadata,
+                    entity_type=DashboardDataModel,
+                    service_name=self.context.get().dashboard_service,
+                    data_model_name=dashboard_details.datasetId,
+                )
+                datamodel_entity = self.metadata.get_by_name(
+                    entity=DashboardDataModel,
+                    fqn=datamodel_fqn,
+                )
 
-            if datamodel_entity and report_entity:
-                yield self._get_add_lineage_request(
-                    to_entity=report_entity, from_entity=datamodel_entity
+                if datamodel_entity and report_entity:
+                    yield self._get_add_lineage_request(
+                        to_entity=report_entity, from_entity=datamodel_entity
+                    )
+            else:
+                logger.debug(
+                    f"Skipping datamodel and report lineage for"
+                    f" {dashboard_details.id} as datasetId is not found"
                 )
 
         except Exception as exc:  # pylint: disable=broad-except
@@ -1184,6 +1190,12 @@ class PowerbiSource(DashboardServiceSource):
         Method to process the dashboard owners
         """
         try:
+            if not self.source_config.includeOwners:
+                logger.debug(
+                    f"Skipping owner processing for {dashboard_details.id} "
+                    f"as includeOwners is False"
+                )
+                return None
             owner_ref_list = []  # to assign multiple owners to entity if they exist
             for owner in dashboard_details.users or []:
                 owner_ref = None
@@ -1198,7 +1210,7 @@ class PowerbiSource(DashboardServiceSource):
                     access_right = owner.dashboardUserAccessRight
 
                 if owner.userType != "Member":
-                    logger.warning(
+                    logger.debug(
                         f"User is not a member of {dashboard_details.id}:"
                         f" ({owner.displayName}, {owner.email})"
                     )
@@ -1211,7 +1223,7 @@ class PowerbiSource(DashboardServiceSource):
                         try:
                             owner_email = EmailStr._validate(owner.email)
                         except PydanticCustomError:
-                            logger.warning(f"Invalid email for owner: {owner.email}")
+                            logger.debug(f"Invalid email for owner: {owner.email}")
                             owner_email = None
                         if owner_email:
                             try:
@@ -1219,7 +1231,7 @@ class PowerbiSource(DashboardServiceSource):
                                     owner_email.lower()
                                 )
                             except Exception as err:
-                                logger.warning(
+                                logger.debug(
                                     f"Could not process owner data with email"
                                     f" {owner.email} in {dashboard_details.id}: {err}"
                                 )
@@ -1229,14 +1241,14 @@ class PowerbiSource(DashboardServiceSource):
                                 name=owner.displayName
                             )
                         except Exception as err:
-                            logger.warning(
+                            logger.debug(
                                 f"Could not process owner data with name"
                                 f" {owner.displayName} in {dashboard_details.id}: {err}"
                             )
                     if owner_ref:
                         owner_ref_list.append(owner_ref.root[0])
                 else:
-                    logger.warning(
+                    logger.debug(
                         f"User does not have owner, admin or write access to"
                         f" {dashboard_details.id}: ({owner.displayName}, {owner.email})"
                     )
@@ -1254,7 +1266,7 @@ class PowerbiSource(DashboardServiceSource):
                     if owner_ref and owner_ref.root[0] not in owner_ref_list:
                         owner_ref_list.append(owner_ref.root[0])
                 except Exception as err:
-                    logger.warning(f"Could not fetch owner data due to {err}")
+                    logger.debug(f"Could not fetch current active user due to {err}")
             if len(owner_ref_list) > 0:
                 logger.debug(
                     f"Successfully fetched owners data for {dashboard_details.id}"

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

@@ -13,7 +13,7 @@
 Source connection handler
 """
 from copy import deepcopy
-from typing import Optional
+from typing import Optional, cast
 from urllib.parse import quote_plus
 
 from requests import Session
@@ -25,6 +25,7 @@ from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
 from metadata.generated.schema.entity.services.connections.database.common import (
+    azureConfig,
     basicAuth,
     jwtAuth,
     noConfigAuthenticationTypes,
@@ -147,17 +148,11 @@ def get_connection(connection: TrinoConnection) -> Engine:
         )
         connection.connectionArguments.root["verify"] = {"verify": connection.verify}
     if hasattr(connection.authType, "azureConfig"):
-        azure_client = AzureClient(connection.authType.azureConfig).create_client()
-        if not connection.authType.azureConfig.scopes:
-            raise ValueError(
-                "Azure Scopes are missing, please refer https://learn.microsoft.com/en-gb/azure/mysql/flexible-server/how-to-azure-ad#2---retrieve-microsoft-entra-access-token and fetch the resource associated with it, for e.g. https://ossrdbms-aad.database.windows.net/.default"
-            )
-        access_token_obj = azure_client.get_token(
-            *connection.authType.azureConfig.scopes.split(",")
-        )
+        access_token_obj = get_azure_token(connection)
+
         if not connection.connectionOptions:
             connection.connectionOptions = init_empty_connection_options()
-        connection.connectionOptions.root["access_token"] = access_token_obj.token
+        connection.connectionOptions.root["access_token"] = access_token_obj
     return create_generic_db_connection(
         connection=connection_copy,
         get_connection_url_fn=get_connection_url,
@@ -165,6 +160,80 @@ def get_connection(connection: TrinoConnection) -> Engine:
     )
 
 
+def get_connection_dict(connection: TrinoConnection) -> dict:
+    """
+    Create connection
+    """
+    url = get_connection(connection).url
+    connection_copy = deepcopy(connection)
+
+    connection_dict = {
+        "driver": url.drivername,
+        "host": url.host,
+        "port": url.port,
+        "user": url.username,
+        "catalog": url.database,
+        "schema": url.query.get("schema"),
+    }
+    connection_dict.update(url.query)
+
+    if connection_copy.proxies:
+        connection_dict["http_session"] = connection_copy.proxies
+
+    if connection_copy.connectionArguments and connection_copy.connectionArguments.root:
+        connection_with_options_secrets(lambda: connection_copy)
+        connection_dict.update(get_connection_args_common(connection_copy))
+
+    if isinstance(connection_copy.authType, basicAuth.BasicAuth):
+        connection_dict["auth"] = {
+            "authType": "basic",
+            "username": connection.username,
+            "password": connection.authType.password.get_secret_value()
+            if connection.authType.password
+            else None,
+        }
+        connection_dict["http_scheme"] = "https"
+
+    elif isinstance(connection_copy.authType, jwtAuth.JwtAuth):
+        connection_dict["auth"] = {
+            "authType": "jwt",
+            "jwt": connection.authType.jwt.get_secret_value(),
+        }
+        connection_dict["http_scheme"] = "https"
+
+    elif hasattr(connection_copy.authType, "azureConfig"):
+        connection_dict["auth"] = {
+            "authType": "jwt",
+            "jwt": get_azure_token(connection_copy),
+        }
+        connection_dict["http_scheme"] = "https"
+
+    elif (
+        connection_copy.authType
+        == noConfigAuthenticationTypes.NoConfigAuthenticationTypes.OAuth2
+    ):
+        connection_dict["auth"] = {"authType": "oauth2"}
+        connection_dict["http_scheme"] = "https"
+
+    return connection_dict
+
+
+def get_azure_token(connection: TrinoConnection) -> str:
+    """
+    Get the azure token for the trino connection
+    """
+    auth_type = cast(azureConfig.AzureConfigurationSource, connection.authType)
+
+    if not auth_type.azureConfig.scopes:
+        raise ValueError(
+            "Azure Scopes are missing, please refer https://learn.microsoft.com/en-gb/azure/mysql/flexible-server/how-to-azure-ad#2---retrieve-microsoft-entra-access-token and fetch the resource associated with it, for e.g. https://ossrdbms-aad.database.windows.net/.default"
+        )
+
+    azure_client = AzureClient(auth_type.azureConfig).create_client()
+
+    return azure_client.get_token(*auth_type.azureConfig.scopes.split(",")).token
+
+
 def test_connection(
     metadata: OpenMetadata,
     engine: Engine,

metadata/pii/algorithms/classifiers.py

@@ -103,7 +103,6 @@ class HeuristicPIIClassifier(ColumnClassifier[PIITag]):
         column_name: Optional[str] = None,
         column_data_type: Optional[DataType] = None,
     ) -> Mapping[PIITag, float]:
-
         if column_data_type is not None and is_non_pii_datatype(column_data_type):
             return {}
 

metadata/pii/algorithms/presidio_patches.py

@@ -16,6 +16,10 @@ from typing import List, Protocol, Sequence
 from dateutil.parser import parse
 from presidio_analyzer import RecognizerResult
 
+from metadata.utils.logger import pii_logger
+
+logger = pii_logger()
+
 
 class PresidioRecognizerResultPatcher(Protocol):
     """
@@ -76,8 +80,11 @@ def date_time_patcher(
             # try to parse using dateutils, if it fails, skip the result
             try:
                 _ = parse(text[result.start : result.end])
-            except ValueError:
+            except (ValueError, OverflowError):
                 # if parsing fails, skip the result
                 continue
+            except Exception as e:
+                logger.info("Unexpected error while parsing date time: %s", e)
+                continue
         patched_result.append(result)
     return patched_result

metadata/sampler/sqlalchemy/databricks/sampler.py

@@ -18,9 +18,13 @@ from metadata.sampler.sqlalchemy.sampler import SQASampler
 
 
 class DatabricksSamplerInterface(SQASampler):
+    def __init__(self, *args, **kwargs):
+        """Initialize with a single Databricks connection"""
+        super().__init__(*args, **kwargs)
+        self.connection = databricks_get_connection(self.service_connection_config)
+
     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/sampler.py

@@ -13,6 +13,7 @@ Helper module to handle data sampling
 for the profiler
 """
 import hashlib
+from contextlib import contextmanager
 from typing import List, Optional, Union, cast
 
 from sqlalchemy import Column, inspect, text
@@ -71,6 +72,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         self._table = self.build_table_orm(
             self.entity, self.service_connection_config, self.ometa_client
         )
+        self._active_sessions = set()
 
     @property
     def raw_dataset(self):
@@ -79,7 +81,20 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
     def get_client(self):
         """Build the SQA Client"""
         session_factory = create_and_bind_thread_safe_session(self.connection)
-        return session_factory()
+        session = session_factory()
+        self._active_sessions.add(session)
+        return session
+
+    @contextmanager
+    def get_client_context(self):
+        """Get client as context manager for proper cleanup"""
+        session = self.get_client()
+        try:
+            yield session
+        finally:
+            if session in self._active_sessions:
+                self._active_sessions.remove(session)
+            session.close()
 
     def set_tablesample(self, selectable: Table):
         """Set the tablesample for the table. To be implemented by the child SQA sampler class
@@ -199,7 +214,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
                 if col.name != RANDOM_LABEL and col.name in names
             ]
 
-        with self.get_client() as client:
+        with self.get_client_context() as client:
             sqa_sample = (
                 client.query(*sqa_columns)
                 .select_from(ds)
@@ -217,7 +232,7 @@ 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}")
 
-        with self.get_client() as client:
+        with self.get_client_context() as client:
             rnd = client.execute(f"{self.sample_query}")
         try:
             columns = [col.name for col in rnd.cursor.description]
@@ -264,5 +279,6 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
 
     def close(self):
         """Close the connection"""
-        self.get_client().close()
+        for session in self._active_sessions:
+            session.close()
         self.connection.pool.dispose()

metadata/sampler/sqlalchemy/unitycatalog/sampler.py

@@ -20,9 +20,13 @@ from metadata.sampler.sqlalchemy.sampler import SQASampler
 
 
 class UnityCatalogSamplerInterface(SQASampler):
+    def __init__(self, *args, **kwargs):
+        """Initialize with a single Databricks connection"""
+        super().__init__(*args, **kwargs)
+        self.connection = databricks_get_connection(self.service_connection_config)
+
     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