openmetadata-ingestion 1.3.1.3__py3-none-any.whl → 1.3.2.0rc2__py3-none-any.whl

metadata/profiler/processor/sample_data_handler.py

@@ -17,8 +17,13 @@ from datetime import datetime
 from functools import singledispatch
 from io import BytesIO
 
+from pydantic.json import ENCODERS_BY_TYPE
+
 from metadata.clients.aws_client import AWSClient
 from metadata.generated.schema.entity.data.table import Table, TableData
+from metadata.generated.schema.entity.services.connections.connectionBasicType import (
+    DataStorageConfig,
+)
 from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.utils.helpers import clean_uri
@@ -27,15 +32,45 @@ from metadata.utils.logger import profiler_logger
 logger = profiler_logger()
 
 
-def _get_object_key(table: Table, prefix: str, overwrite_data: bool) -> str:
+class PathPatternException(Exception):
+    """
+    Exception class need to validate the file path pattern
+    """
+
+
+def validate_path_pattern(file_path_format: str) -> None:
+    if not (
+        "{service_name}" in file_path_format
+        and "{database_name}" in file_path_format
+        and "{database_schema_name}" in file_path_format
+        and "{table_name}" in file_path_format
+        and file_path_format.endswith(".parquet")
+    ):
+        raise PathPatternException(
+            "Please provide a valid path pattern, "
+            "the pattern should include these components {service_name}, "
+            "{database_name}, {database_schema_name}, {table_name} and "
+            "it should end with extension .parquet"
+        )
+
+
+def _get_object_key(
+    table: Table, prefix: str, overwrite_data: bool, file_path_format: str
+) -> str:
+    validate_path_pattern(file_path_format)
+    file_name = file_path_format.format(
+        service_name=table.service.name,
+        database_name=table.database.name,
+        database_schema_name=table.databaseSchema.name,
+        table_name=table.name.__root__,
+    )
     if not overwrite_data:
-        file_name = f"sample_data_{datetime.now().strftime('%Y_%m_%d')}.parquet"
-    else:
-        file_name = "sample_data.parquet"
-    path = str(table.fullyQualifiedName.__root__).replace(".", "/")
+        file_name = file_name.replace(
+            ".parquet", f"_{datetime.now().strftime('%Y_%m_%d')}.parquet"
+        )
     if prefix:
-        return f"{clean_uri(prefix)}/{path}/{file_name}"
-    return f"{path}/{file_name}"
+        return f"{clean_uri(prefix)}/{file_name}"
+    return file_name
 
 
 def upload_sample_data(data: TableData, profiler_interface: ProfilerInterface) -> None:
@@ -45,9 +80,10 @@ def upload_sample_data(data: TableData, profiler_interface: ProfilerInterface) -
     import pandas as pd  # pylint: disable=import-outside-toplevel
 
     try:
-        sample_storage_config = profiler_interface.storage_config
+        sample_storage_config: DataStorageConfig = profiler_interface.storage_config
         if not sample_storage_config:
             return
+        ENCODERS_BY_TYPE[bytes] = lambda v: v.decode("utf-8", "ignore")
         deserialized_data = json.loads(data.json())
         df = pd.DataFrame(
             data=deserialized_data.get("rows", []),
@@ -59,6 +95,7 @@ def upload_sample_data(data: TableData, profiler_interface: ProfilerInterface) -
             table=profiler_interface.table_entity,
             prefix=sample_storage_config.prefix,
             overwrite_data=sample_storage_config.overwriteData,
+            file_path_format=sample_storage_config.filePathPattern,
         )
         upload_to_storage(
             sample_storage_config.storageConfig,

metadata/profiler/source/databricks/profiler_source.py

@@ -0,0 +1,36 @@
+"""Extend the ProfilerSource class to add support for Databricks is_disconnect SQA method"""
+
+from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.metadataIngestion.workflow import (
+    OpenMetadataWorkflowConfig,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.profiler.source.base.profiler_source import ProfilerSource
+
+
+def is_disconnect(self, e, connection, cursor):
+    """is_disconnect method for the Databricks dialect"""
+    if "Invalid SessionHandle: SessionHandle" in str(e):
+        return True
+    return False
+
+
+class DataBricksProfilerSource(ProfilerSource):
+    """Databricks Profiler source"""
+
+    def __init__(
+        self,
+        config: OpenMetadataWorkflowConfig,
+        database: DatabaseService,
+        ometa_client: OpenMetadata,
+    ):
+        super().__init__(config, database, ometa_client)
+        self.set_is_disconnect()
+
+    def set_is_disconnect(self):
+        """Set the is_disconnect method for the Databricks dialect"""
+        from databricks.sqlalchemy import (
+            DatabricksDialect,  # pylint: disable=import-outside-toplevel
+        )
+
+        DatabricksDialect.is_disconnect = is_disconnect

metadata/profiler/source/metadata.py

@@ -43,6 +43,10 @@ from metadata.utils.logger import profiler_logger
 logger = profiler_logger()
 
 
+TABLE_FIELDS = ["tableProfilerConfig", "columns", "customMetrics"]
+TAGS_FIELD = ["tags"]
+
+
 class ProfilerSourceAndEntity(BaseModel):
     """Return class for the OpenMetadata Profiler Source"""
 
@@ -273,7 +277,9 @@ class OpenMetadataSource(Source):
         """
         tables = self.metadata.list_all_entities(
             entity=Table,
-            fields=["tableProfilerConfig", "columns", "customMetrics"],
+            fields=TABLE_FIELDS
+            if not self.source_config.processPiiSensitive
+            else TABLE_FIELDS + TAGS_FIELD,
             params={
                 "service": self.config.source.serviceName,
                 "database": fqn.build(

metadata/profiler/source/profiler_source_factory.py

@@ -16,8 +16,12 @@ Factory class for creating profiler source objects
 from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
     BigqueryType,
 )
+from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
+    DatabricksType,
+)
 from metadata.profiler.source.base.profiler_source import ProfilerSource
 from metadata.profiler.source.bigquery.profiler_source import BigQueryProfilerSource
+from metadata.profiler.source.databricks.profiler_source import DataBricksProfilerSource
 
 
 class ProfilerSourceFactory:
@@ -44,3 +48,7 @@ profiler_source_factory.register_source(
     BigqueryType.BigQuery.value.lower(),
     BigQueryProfilerSource,
 )
+profiler_source_factory.register_source(
+    DatabricksType.Databricks.value.lower(),
+    DataBricksProfilerSource,
+)

metadata/readers/dataframe/json.py

@@ -16,7 +16,7 @@ import gzip
 import io
 import json
 import zipfile
-from typing import List, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from metadata.readers.dataframe.base import DataFrameReader
 from metadata.readers.dataframe.common import dataframe_to_chunks
@@ -47,7 +47,7 @@ class JSONDataFrameReader(DataFrameReader):
     @staticmethod
     def read_from_json(
         key: str, json_text: bytes, decode: bool = False, **__
-    ) -> List["DataFrame"]:
+    ) -> Tuple[List["DataFrame"], Optional[Dict[str, Any]]]:
         """
         Decompress a JSON file (if needed) and read its contents
         as a dataframe.
@@ -60,20 +60,25 @@ class JSONDataFrameReader(DataFrameReader):
         import pandas as pd
 
         json_text = _get_json_text(key=key, text=json_text, decode=decode)
+        raw_data = None
         try:
             data = json.loads(json_text)
+            if isinstance(data, dict) and data.get("$schema"):
+                raw_data = json_text
         except json.decoder.JSONDecodeError:
             logger.debug("Failed to read as JSON object. Trying to read as JSON Lines")
             data = [json.loads(json_obj) for json_obj in json_text.strip().split("\n")]
 
         # if we get a scalar value (e.g. {"a":"b"}) then we need to specify the index
         data = data if not isinstance(data, dict) else [data]
-        return dataframe_to_chunks(pd.DataFrame.from_records(data))
+        return dataframe_to_chunks(pd.DataFrame.from_records(data)), raw_data
 
     def _read(self, *, key: str, bucket_name: str, **kwargs) -> DatalakeColumnWrapper:
         text = self.reader.read(key, bucket_name=bucket_name)
+        dataframes, raw_data = self.read_from_json(
+            key=key, json_text=text, decode=True, **kwargs
+        )
         return DatalakeColumnWrapper(
-            dataframes=self.read_from_json(
-                key=key, json_text=text, decode=True, **kwargs
-            )
+            dataframes=dataframes,
+            raw_data=raw_data,
         )

metadata/readers/dataframe/models.py

@@ -29,6 +29,7 @@ class DatalakeColumnWrapper(BaseModel):
 
     columns: Optional[List[Column]]
     dataframes: Optional[List[Any]]  # pandas.Dataframe does not have any validators
+    raw_data: Any  # in special cases like json schema, we need to store the raw data
 
 
 class DatalakeTableSchemaWrapper(BaseModel):

metadata/utils/credentials.py

@@ -15,7 +15,7 @@ import base64
 import json
 import os
 import tempfile
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 from cryptography.hazmat.primitives import serialization
 from google import auth
@@ -25,6 +25,9 @@ from metadata.generated.schema.security.credentials.gcpCredentials import (
     GCPCredentials,
     GcpCredentialsPath,
 )
+from metadata.generated.schema.security.credentials.gcpExternalAccount import (
+    GcpExternalAccount,
+)
 from metadata.generated.schema.security.credentials.gcpValues import (
     GcpCredentialsValues,
 )
@@ -85,30 +88,44 @@ def create_credential_tmp_file(credentials: dict) -> str:
         return temp_file_path
 
 
-def build_google_credentials_dict(gcp_values: GcpCredentialsValues) -> Dict[str, str]:
+def build_google_credentials_dict(
+    gcp_values: Union[GcpCredentialsValues, GcpExternalAccount]
+) -> Dict[str, str]:
     """
     Given GcPCredentialsValues, build a dictionary as the JSON file
     downloaded from GCP with the service_account
     :param gcp_values: GCP credentials
     :return: Dictionary with credentials
     """
-    private_key_str = gcp_values.privateKey.get_secret_value()
-    # adding the replace string here to escape line break if passed from env
-    private_key_str = private_key_str.replace("\\n", "\n")
-    validate_private_key(private_key_str)
-
-    return {
-        "type": gcp_values.type,
-        "project_id": gcp_values.projectId.__root__,
-        "private_key_id": gcp_values.privateKeyId,
-        "private_key": private_key_str,
-        "client_email": gcp_values.clientEmail,
-        "client_id": gcp_values.clientId,
-        "auth_uri": str(gcp_values.authUri),
-        "token_uri": str(gcp_values.tokenUri),
-        "auth_provider_x509_cert_url": str(gcp_values.authProviderX509CertUrl),
-        "client_x509_cert_url": str(gcp_values.clientX509CertUrl),
-    }
+    if isinstance(gcp_values, GcpCredentialsValues):
+        private_key_str = gcp_values.privateKey.get_secret_value()
+        # adding the replace string here to escape line break if passed from env
+        private_key_str = private_key_str.replace("\\n", "\n")
+        validate_private_key(private_key_str)
+
+        return {
+            "type": gcp_values.type,
+            "project_id": gcp_values.projectId.__root__,
+            "private_key_id": gcp_values.privateKeyId,
+            "private_key": private_key_str,
+            "client_email": gcp_values.clientEmail,
+            "client_id": gcp_values.clientId,
+            "auth_uri": str(gcp_values.authUri),
+            "token_uri": str(gcp_values.tokenUri),
+            "auth_provider_x509_cert_url": str(gcp_values.authProviderX509CertUrl),
+            "client_x509_cert_url": str(gcp_values.clientX509CertUrl),
+        }
+    if isinstance(gcp_values, GcpExternalAccount):
+        return {
+            "type": gcp_values.externalType,
+            "audience": gcp_values.audience,
+            "subject_token_type": gcp_values.subjectTokenType,
+            "token_url": gcp_values.tokenURL,
+            "credential_source": gcp_values.credentialSource,
+        }
+    raise InvalidGcpConfigException(
+        f"Error trying to build GCP credentials dict due to Invalid GCP config {type(gcp_values)}"
+    )
 
 
 def set_google_credentials(gcp_credentials: GCPCredentials) -> None:

metadata/utils/datalake/datalake_utils.py

@@ -17,10 +17,11 @@ import ast
 import json
 import random
 import traceback
-from typing import Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Optional, Union, cast
 
 from metadata.generated.schema.entity.data.table import Column, DataType
 from metadata.ingestion.source.database.column_helpers import truncate_column_name
+from metadata.parsers.json_schema_parser import parse_json_schema
 from metadata.readers.dataframe.models import (
     DatalakeColumnWrapper,
     DatalakeTableSchemaWrapper,
@@ -35,6 +36,7 @@ def fetch_dataframe(
     config_source,
     client,
     file_fqn: DatalakeTableSchemaWrapper,
+    fetch_raw_data: bool = False,
     **kwargs,
 ) -> Optional[List["DataFrame"]]:
     """
@@ -60,6 +62,8 @@ def fetch_dataframe(
                 df_wrapper: DatalakeColumnWrapper = df_reader.read(
                     key=key, bucket_name=bucket_name, **kwargs
                 )
+                if fetch_raw_data:
+                    return df_wrapper.dataframes, df_wrapper.raw_data
                 return df_wrapper.dataframes
             except Exception as err:
                 logger.error(
@@ -73,6 +77,8 @@ def fetch_dataframe(
         # Here we need to blow things up. Without the dataframe we cannot move forward
         raise err
 
+    if fetch_raw_data:
+        return None, None
     return None
 
 
@@ -112,6 +118,7 @@ class DataFrameColumnParser:
         file_type: Optional[SupportedTypes] = None,
         sample: bool = True,
         shuffle: bool = False,
+        raw_data: Any = None,
     ):
         """Instantiate a column parser object with the appropriate parser
 
@@ -126,8 +133,14 @@ class DataFrameColumnParser:
         data_frame = cls._get_data_frame(data_frame, sample, shuffle)
         if file_type == SupportedTypes.PARQUET:
             parser = ParquetDataFrameColumnParser(data_frame)
-            return cls(parser)
-        parser = GenericDataFrameColumnParser(data_frame)
+        elif file_type in {
+            SupportedTypes.JSON,
+            SupportedTypes.JSONGZ,
+            SupportedTypes.JSONZIP,
+        }:
+            parser = JsonDataFrameColumnParser(data_frame, raw_data=raw_data)
+        else:
+            parser = GenericDataFrameColumnParser(data_frame)
         return cls(parser)
 
     @staticmethod
@@ -172,8 +185,9 @@ class GenericDataFrameColumnParser:
         "bytes": DataType.BYTES,
     }
 
-    def __init__(self, data_frame: "DataFrame"):
+    def __init__(self, data_frame: "DataFrame", raw_data: Any = None):
         self.data_frame = data_frame
+        self.raw_data = raw_data
 
     def get_columns(self):
         """
@@ -472,3 +486,19 @@ class ParquetDataFrameColumnParser:
             data_type = self._data_formats.get(str(column.type), DataType.UNKNOWN)
 
         return data_type
+
+
+class JsonDataFrameColumnParser(GenericDataFrameColumnParser):
+    """Given a dataframe object generated from a json file, parse the columns and return a list of Column objects."""
+
+    def get_columns(self):
+        """
+        method to process column details for json files
+        """
+        if self.raw_data:
+            try:
+                return parse_json_schema(schema_text=self.raw_data, cls=Column)
+            except Exception as exc:
+                logger.warning(f"Unable to parse the json schema: {exc}")
+                logger.debug(traceback.format_exc())
+        return self._get_columns(self.data_frame)

metadata/utils/secrets/azure_kv_secrets_manager.py

@@ -17,9 +17,9 @@ import traceback
 from abc import ABC
 from typing import Optional
 
-from azure.identity import ClientSecretCredential, DefaultAzureCredential
-from azure.keyvault.secrets import KeyVaultSecret, SecretClient
+from azure.keyvault.secrets import KeyVaultSecret
 
+from metadata.clients.azure_client import AzureClient
 from metadata.generated.schema.security.secrets.secretsManagerClientLoader import (
     SecretsManagerClientLoader,
 )
@@ -105,23 +105,7 @@ class AzureKVSecretsManager(ExternalSecretsManager, ABC):
     ):
         super().__init__(provider=SecretsManagerProvider.azure_kv, loader=loader)
 
-        if (
-            self.credentials.tenantId
-            and self.credentials.clientId
-            and self.credentials.clientSecret
-        ):
-            azure_identity = ClientSecretCredential(
-                tenant_id=self.credentials.tenantId,
-                client_id=self.credentials.clientId,
-                client_secret=self.credentials.clientSecret.get_secret_value(),
-            )
-        else:
-            azure_identity = DefaultAzureCredential()
-
-        self.client = SecretClient(
-            vault_url=f"https://{self.credentials.vaultName}.vault.azure.net/",
-            credential=azure_identity,
-        )
+        self.client = AzureClient(self.credentials).create_secret_client()
 
     def get_string_value(self, secret_id: str) -> str:
         """

metadata/utils/source_hash.py

@@ -14,9 +14,14 @@ Source hash utils module
 """
 
 import hashlib
+import traceback
 from typing import Dict, Optional
 
 from metadata.ingestion.ometa.ometa_api import C
+from metadata.utils.logger import utils_logger
+
+logger = utils_logger()
+
 
 SOURCE_HASH_EXCLUDE_FIELDS = {
     "sourceHash": True,
@@ -25,19 +30,24 @@ SOURCE_HASH_EXCLUDE_FIELDS = {
 
 def generate_source_hash(
     create_request: C, exclude_fields: Optional[Dict] = None
-) -> str:
+) -> Optional[str]:
     """
     Given a create_request model convert it to json string and generate a hash value
     """
-
-    # We always want to exclude the sourceHash when generating the fingerprint
-    exclude_fields = (
-        SOURCE_HASH_EXCLUDE_FIELDS.update(exclude_fields)
-        if exclude_fields
-        else SOURCE_HASH_EXCLUDE_FIELDS
-    )
-
-    create_request_json = create_request.json(exclude=exclude_fields)
-
-    json_bytes = create_request_json.encode("utf-8")
-    return hashlib.md5(json_bytes).hexdigest()
+    try:
+        # We always want to exclude the sourceHash when generating the fingerprint
+        exclude_fields = (
+            SOURCE_HASH_EXCLUDE_FIELDS.update(exclude_fields)
+            if exclude_fields
+            else SOURCE_HASH_EXCLUDE_FIELDS
+        )
+
+        create_request_json = create_request.json(exclude=exclude_fields)
+
+        json_bytes = create_request_json.encode("utf-8")
+        return hashlib.md5(json_bytes).hexdigest()
+
+    except Exception as exc:
+        logger.warning(f"Failed to generate source hash due to - {exc}")
+        logger.debug(traceback.format_exc())
+    return None

metadata/utils/storage_metadata_config.py

@@ -17,6 +17,7 @@ from functools import singledispatch
 
 import requests
 
+from metadata.clients.azure_client import AzureClient
 from metadata.generated.schema.entity.services.connections.database.datalake.azureConfig import (
     AzureConfig,
 )
@@ -153,21 +154,7 @@ def _(config: StorageMetadataAdlsConfig) -> ManifestMetadataConfig:
             else STORAGE_METADATA_MANIFEST_FILE_NAME
         )
 
-        from azure.identity import (  # pylint: disable=import-outside-toplevel
-            ClientSecretCredential,
-        )
-        from azure.storage.blob import (  # pylint: disable=import-outside-toplevel
-            BlobServiceClient,
-        )
-
-        blob_client = BlobServiceClient(
-            account_url=f"https://{config.securityConfig.accountName}.blob.core.windows.net/",
-            credential=ClientSecretCredential(
-                config.securityConfig.tenantId,
-                config.securityConfig.clientId,
-                config.securityConfig.clientSecret.get_secret_value(),
-            ),
-        )
+        blob_client = AzureClient(config.securityConfig).create_blob_client()
 
         reader = get_reader(
             config_source=AzureConfig(securityConfig=config.securityConfig),

metadata/workflow/base.py

@@ -108,7 +108,7 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
     @property
     def ingestion_pipeline(self):
         """Get or create the Ingestion Pipeline from the configuration"""
-        if not self._ingestion_pipeline:
+        if not self._ingestion_pipeline and self.config.ingestionPipelineFQN:
             self._ingestion_pipeline = self.get_or_create_ingestion_pipeline()
 
         return self._ingestion_pipeline