openmetadata-ingestion 1.2.4.0__py3-none-any.whl → 1.3.2.0rc1__py3-none-any.whl

metadata/profiler/source/mariadb/metrics/window/first_quartile.py

@@ -0,0 +1,10 @@
+"""Override first quartile metric definition for MariaDB"""
+
+from metadata.profiler.metrics.window.first_quartile import FirstQuartile
+from metadata.profiler.source.mariadb.functions.median import MariaDBMedianFn
+
+
+class MariaDBFirstQuartile(FirstQuartile):
+    def _compute_sqa_fn(self, column, table, percentile):
+        """Generic method to compute the quartile using sqlalchemy"""
+        return MariaDBMedianFn(column, table, percentile)

metadata/profiler/source/mariadb/metrics/window/median.py

@@ -0,0 +1,10 @@
+"""Override first quartile metric definition for MariaDB"""
+
+from metadata.profiler.metrics.window.median import Median
+from metadata.profiler.source.mariadb.functions.median import MariaDBMedianFn
+
+
+class MariaDBMedian(Median):
+    def _compute_sqa_fn(self, column, table, percentile):
+        """Generic method to compute the quartile using sqlalchemy"""
+        return MariaDBMedianFn(column, table, percentile)

metadata/profiler/source/mariadb/metrics/window/third_quartile.py

@@ -0,0 +1,10 @@
+"""Override first quartile metric definition for MariaDB"""
+
+from metadata.profiler.metrics.window.third_quartile import ThirdQuartile
+from metadata.profiler.source.mariadb.functions.median import MariaDBMedianFn
+
+
+class MariaDBThirdQuartile(ThirdQuartile):
+    def _compute_sqa_fn(self, column, table, percentile):
+        """Generic method to compute the quartile using sqlalchemy"""
+        return MariaDBMedianFn(column, table, percentile)

metadata/profiler/source/metadata.py

@@ -17,15 +17,19 @@ from typing import Iterable, Optional, cast
 from pydantic import BaseModel
 
 from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.table import Table, TableType
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
     DatabaseServiceProfilerPipeline,
 )
 from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.parser import parse_workflow_config_gracefully
 from metadata.ingestion.api.step import Step
 from metadata.ingestion.api.steps import Source
@@ -39,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"""
 
@@ -67,6 +75,10 @@ class OpenMetadataSource(Source):
     def init_steps(self):
         super().__init__()
 
+    @property
+    def name(self) -> str:
+        return "OpenMetadata Service"
+
     # pylint: disable=super-init-not-called
     def __init__(
         self,
@@ -114,7 +126,6 @@ class OpenMetadataSource(Source):
         self.metadata.health_check()
 
     def _iter(self, *_, **__) -> Iterable[Either[ProfilerSourceAndEntity]]:
-
         for database in self.get_database_entities():
             try:
                 profiler_source = profiler_source_factory.create(
@@ -135,7 +146,7 @@ class OpenMetadataSource(Source):
                     left=StackTraceError(
                         name=database.fullyQualifiedName.__root__,
                         error=f"Error listing source and entities for database due to [{exc}]",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 
@@ -146,9 +157,17 @@ class OpenMetadataSource(Source):
 
     def filter_databases(self, database: Database) -> Optional[Database]:
         """Returns filtered database entities"""
+        database_fqn = fqn.build(
+            self.metadata,
+            entity_type=Database,
+            service_name=self.config.source.serviceName,
+            database_name=database.name.__root__,
+        )
         if filter_by_database(
             self.source_config.databaseFilterPattern,
-            database.name.__root__,
+            database_fqn
+            if self.source_config.useFqnForFiltering
+            else database.name.__root__,
         ):
             self.status.filter(database.name.__root__, "Database pattern not allowed")
             return None
@@ -163,18 +182,38 @@ class OpenMetadataSource(Source):
         """
         for table in tables:
             try:
+                schema_fqn = fqn.build(
+                    self.metadata,
+                    entity_type=DatabaseSchema,
+                    service_name=self.config.source.serviceName,
+                    database_name=table.database.name,
+                    schema_name=table.databaseSchema.name,
+                )
                 if filter_by_schema(
                     self.source_config.schemaFilterPattern,
-                    table.databaseSchema.name,  # type: ignore
+                    schema_fqn
+                    if self.source_config.useFqnForFiltering
+                    else table.databaseSchema.name,  # type: ignore
                 ):
                     self.status.filter(
                         f"Schema pattern not allowed: {table.fullyQualifiedName.__root__}",
                         "Schema pattern not allowed",
                     )
                     continue
+                table_fqn = fqn.build(
+                    self.metadata,
+                    entity_type=Table,
+                    service_name=self.config.source.serviceName,
+                    database_name=table.database.name,
+                    schema_name=table.databaseSchema.name,
+                    table_name=table.name.__root__,
+                )
+
                 if filter_by_table(
                     self.source_config.tableFilterPattern,
-                    table.name.__root__,
+                    table_fqn
+                    if self.source_config.useFqnForFiltering
+                    else table.name.__root__,
                 ):
                     self.status.filter(
                         f"Table pattern not allowed: {table.fullyQualifiedName.__root__}",
@@ -196,7 +235,7 @@ class OpenMetadataSource(Source):
                     StackTraceError(
                         name=table.fullyQualifiedName.__root__,
                         error=f"Unexpected error filtering entities for table [{table}]: {exc}",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 
@@ -238,9 +277,9 @@ class OpenMetadataSource(Source):
         """
         tables = self.metadata.list_all_entities(
             entity=Table,
-            fields=[
-                "tableProfilerConfig",
-            ],
+            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/metadata_ext.py

@@ -26,7 +26,10 @@ from typing import Iterable, cast
 from sqlalchemy.inspection import inspect
 
 from metadata.generated.schema.entity.data.database import Database
-from metadata.generated.schema.entity.data.table import Table, TableType
+from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
     DatabaseServiceMetadataPipeline,
 )
@@ -36,7 +39,7 @@ from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline
 from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection
 from metadata.profiler.source.metadata import (
@@ -159,7 +162,7 @@ class OpenMetadataSourceExt(OpenMetadataSource):
                     left=StackTraceError(
                         name=database_name,
                         error=f"Error listing source and entities for database due to [{exc}]",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 
@@ -194,8 +197,17 @@ class OpenMetadataSourceExt(OpenMetadataSource):
                 else:
                     database_names = self.source.get_database_names_raw()
                     for database in database_names:
+                        database_fqn = fqn.build(
+                            self.metadata,
+                            entity_type=Database,
+                            service_name=self.config.source.serviceName,
+                            database_name=database,
+                        )
                         if filter_by_database(
-                            self.source_config.databaseFilterPattern, database
+                            self.source_config.databaseFilterPattern,
+                            database_fqn
+                            if self.source_config.useFqnForFiltering
+                            else database,
                         ):
                             self.status.filter(database, "Database pattern not allowed")
                             continue
@@ -213,52 +225,6 @@ class OpenMetadataSourceExt(OpenMetadataSource):
             logger.debug(f"Failed to fetch database names {exc}")
             logger.debug(traceback.format_exc())
 
-    def filter_entities(self, tables: Iterable[Table]) -> Iterable[Table]:
-        """
-        From a list of tables, apply the SQLSourceConfig
-        filter patterns.
-
-        We will update the status on the SQLSource Status.
-        """
-        for table in tables:
-            try:
-                if filter_by_schema(
-                    self.source_config.schemaFilterPattern,
-                    table.databaseSchema.name,  # type: ignore
-                ):
-                    self.status.filter(
-                        f"Schema pattern not allowed: {table.fullyQualifiedName.__root__}",
-                        "Schema pattern not allowed",
-                    )
-                    continue
-                if filter_by_table(
-                    self.source_config.tableFilterPattern,
-                    table.name.__root__,
-                ):
-                    self.status.filter(
-                        f"Table pattern not allowed: {table.fullyQualifiedName.__root__}",
-                        "Table pattern not allowed",
-                    )
-                    continue
-                if (
-                    table.tableType == TableType.View
-                    and not self.source_config.includeViews
-                ):
-                    self.status.filter(
-                        table.fullyQualifiedName.__root__,
-                        "View filtered out",
-                    )
-                    continue
-                yield table
-            except Exception as exc:
-                self.status.failed(
-                    StackTraceError(
-                        name=table.fullyQualifiedName.__root__,
-                        error=f"Unexpected error filtering entities for table [{table}]: {exc}",
-                        stack_trace=traceback.format_exc(),
-                    )
-                )
-
     def get_table_entities(self, database):
         """
         List and filter OpenMetadata tables based on the

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

@@ -28,14 +28,15 @@ logger = ingestion_logger()
 
 
 def _get_json_text(key: str, text: bytes, decode: bool) -> Union[str, bytes]:
+    processed_text: Union[str, bytes] = text
     if key.endswith(".gz"):
-        return gzip.decompress(text)
+        processed_text = gzip.decompress(text)
     if key.endswith(".zip"):
         with zipfile.ZipFile(io.BytesIO(text)) as zip_file:
-            return zip_file.read(zip_file.infolist()[0]).decode(UTF_8)
+            processed_text = zip_file.read(zip_file.infolist()[0]).decode(UTF_8)
     if decode:
-        return text.decode(UTF_8) if isinstance(text, bytes) else text
-    return text
+        return processed_text.decode(UTF_8) if isinstance(text, bytes) else text
+    return processed_text
 
 
 class JSONDataFrameReader(DataFrameReader):

metadata/readers/file/api_reader.py

@@ -35,7 +35,6 @@ class ApiReader(Reader, ABC):
     """
 
     def __init__(self, credentials: ReadersCredentials):
-
         self._auth_headers = None
         self.credentials = credentials
 

metadata/utils/constants.py

@@ -32,6 +32,8 @@ from metadata.generated.schema.entity.services.mlmodelService import MlModelServ
 from metadata.generated.schema.entity.services.pipelineService import PipelineService
 from metadata.generated.schema.entity.services.searchService import SearchService
 from metadata.generated.schema.entity.services.storageService import StorageService
+from metadata.generated.schema.entity.teams.team import Team
+from metadata.generated.schema.entity.teams.user import User
 
 DOT = "_DOT_"
 TEN_MIN = 10 * 60
@@ -92,4 +94,7 @@ ENTITY_REFERENCE_TYPE_MAP = {
     SearchIndex.__name__: "searchIndex",
     MlModel.__name__: "mlmodel",
     Container.__name__: "container",
+    # User Entities
+    User.__name__: "user",
+    Team.__name__: "team",
 }

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: