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

airflow_provider_openmetadata/lineage/runner.py

@@ -45,10 +45,16 @@ from metadata.generated.schema.entity.services.pipelineService import (
 )
 from metadata.generated.schema.type.entityLineage import EntitiesEdge, LineageDetails
 from metadata.generated.schema.type.entityReference import EntityReference
+from metadata.ingestion.models.patch_request import (
+    ALLOWED_COMMON_PATCH_FIELDS,
+    RESTRICT_UPDATE_LIST,
+)
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.airflow.lineage_parser import XLets
+from metadata.utils import fqn
 from metadata.utils.constants import ENTITY_REFERENCE_TYPE_MAP
 from metadata.utils.helpers import clean_uri, datetime_to_ts
+from metadata.utils.source_hash import generate_source_hash
 
 
 class SimpleEdge(BaseModel):
@@ -151,11 +157,24 @@ class AirflowLineageRunner:
             for task in self.dag.tasks or []
         ]
 
-    def create_pipeline_entity(self, pipeline_service: PipelineService) -> Pipeline:
+    def create_or_update_pipeline_entity(
+        self, pipeline_service: PipelineService
+    ) -> Pipeline:
         """
-        Create the Pipeline Entity
+        Create the Pipeline Entity if it does not exist, or PATCH it
+        if there have been changes.
         """
-        self.dag.log.info("Creating or updating Pipeline Entity from DAG...")
+        pipeline: Pipeline = self.metadata.get_by_name(
+            entity=Pipeline,
+            fqn=fqn.build(
+                self.metadata,
+                Pipeline,
+                service_name=self.service_name,
+                pipeline_name=self.dag.dag_id,
+            ),
+            fields=["*"],
+        )
+
         pipeline_request = CreatePipelineRequest(
             name=self.dag.dag_id,
             description=self.dag.description,
@@ -167,7 +186,27 @@ class AirflowLineageRunner:
             service=pipeline_service.fullyQualifiedName,
         )
 
-        return self.metadata.create_or_update(pipeline_request)
+        create_entity_request_hash = generate_source_hash(
+            create_request=pipeline_request
+        )
+        pipeline_request.sourceHash = create_entity_request_hash
+
+        if pipeline is None:
+            self.dag.log.info("Creating Pipeline Entity from DAG...")
+            return self.metadata.create_or_update(pipeline_request)
+
+        if create_entity_request_hash != pipeline.sourceHash:
+            self.dag.log.info("Updating Pipeline Entity from DAG...")
+            return self.metadata.patch(
+                entity=Pipeline,
+                source=pipeline,
+                destination=pipeline.copy(update=pipeline_request.__dict__),
+                allowed_fields=ALLOWED_COMMON_PATCH_FIELDS,
+                restrict_update_fields=RESTRICT_UPDATE_LIST,
+            )
+
+        self.dag.log.info("DAG has not changed since last run")
+        return pipeline
 
     def get_all_pipeline_status(self) -> List[PipelineStatus]:
         """
@@ -368,7 +407,7 @@ class AirflowLineageRunner:
         """
         self.dag.log.info("Executing Airflow Lineage Runner...")
         pipeline_service = self.get_or_create_pipeline_service()
-        pipeline = self.create_pipeline_entity(pipeline_service)
+        pipeline = self.create_or_update_pipeline_entity(pipeline_service)
         self.add_all_pipeline_status(pipeline)
 
         self.dag.log.info(f"Processing XLet data {self.xlets}")

airflow_provider_openmetadata/lineage/status.py

@@ -113,7 +113,7 @@ def add_status(
     ]
 
     updated_status = PipelineStatus(
-        timestamp=datetime_to_ts(execution_date),
+        timestamp=execution_date,
         executionStatus=get_dag_status(
             all_tasks=dag.task_ids,
             task_status=updated_task_status,

metadata/antlr/split_listener.py

@@ -35,7 +35,7 @@ class EntityLinkSplitListener(EntityLinkListener):
     def __init__(self):
         self._list = []
 
-    def enterEntityAttribute(self, ctx: EntityLinkParser.EntityAttributeContext):
+    def enterNameOrFQN(self, ctx: EntityLinkParser.NameOrFQNContext):
         self._list.append(ctx.getText())
 
     def enterEntityType(self, ctx: EntityLinkParser.EntityTypeContext):
@@ -44,8 +44,5 @@ class EntityLinkSplitListener(EntityLinkListener):
     def enterEntityField(self, ctx: EntityLinkParser.EntityFieldContext):
         self._list.append(ctx.getText())
 
-    def enterEntityFqn(self, ctx: EntityLinkParser.EntityFqnContext):
-        self._list.append(ctx.getText())
-
     def split(self):
         return self._list

metadata/applications/auto_tagger.py

@@ -0,0 +1,212 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  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.
+"""
+PII application
+"""
+import traceback
+from typing import Iterable, List, Optional
+
+from metadata.generated.schema.entity.applications.configuration.external.autoTaggerAppConfig import (
+    AutoTaggerAppConfig,
+)
+from metadata.generated.schema.entity.data.table import Column, Table, TableData
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
+from metadata.generated.schema.metadataIngestion.application import (
+    OpenMetadataApplicationConfig,
+)
+from metadata.generated.schema.type.tagLabel import (
+    LabelType,
+    State,
+    TagFQN,
+    TagLabel,
+    TagSource,
+)
+from metadata.ingestion.models.table_metadata import ColumnTag
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.pii.constants import PII
+from metadata.pii.scanners.column_name_scanner import ColumnNameScanner
+from metadata.pii.scanners.ner_scanner import NERScanner
+from metadata.utils.logger import app_logger
+from metadata.workflow.application import AppRunner, InvalidAppConfiguration
+
+logger = app_logger()
+
+DEFAULT_CONFIDENCE = 80
+
+
+class AutoTaggerApp(AppRunner):
+    """
+    PII Application
+    You can execute it with `metadata app -c <path-to-yaml>`
+    with a YAML file like:
+
+    sourcePythonClass: metadata.applications.auto_tagger.AutoTaggerApp
+    appConfig:
+      type: AutoTagger
+      confidenceLevel: 80
+    workflowConfig:
+      loggerLevel: INFO
+      openMetadataServerConfig:
+        hostPort: http://localhost:8585/api
+        authProvider: openmetadata
+        securityConfig:
+          jwtToken: "..."
+    """
+
+    def __init__(self, config: OpenMetadataApplicationConfig, metadata: OpenMetadata):
+        super().__init__(config, metadata)
+
+        if not isinstance(self.app_config, AutoTaggerAppConfig):
+            raise InvalidAppConfiguration(
+                f"AutoTagger Runner expects an AutoTaggerAppConfig, we got [{config}]"
+            )
+
+        self._ner_scanner = None
+        self.confidence_threshold = (
+            self.app_config.confidenceLevel or DEFAULT_CONFIDENCE
+        )
+
+    @property
+    def name(self) -> str:
+        return "Auto Tagger"
+
+    @staticmethod
+    def build_column_tag(tag_fqn: str, column_fqn: str) -> ColumnTag:
+        """
+        Build the tag and run the PATCH
+        """
+        tag_label = TagLabel(
+            tagFQN=TagFQN(__root__=tag_fqn),
+            source=TagSource.Classification,
+            state=State.Suggested,
+            labelType=LabelType.Automated,
+        )
+
+        return ColumnTag(column_fqn=column_fqn, tag_label=tag_label)
+
+    @property
+    def ner_scanner(self) -> NERScanner:
+        """Load the NER Scanner only if called"""
+        if self._ner_scanner is None:
+            self._ner_scanner = NERScanner()
+
+        return self._ner_scanner
+
+    def process_column(
+        self,
+        idx: int,
+        column: Column,
+        table_data: Optional[TableData],
+        confidence_threshold: float,
+    ) -> Optional[List[ColumnTag]]:
+        """
+        Tag a column with PII if we find it using our scanners
+        """
+
+        # First, check if the column we are about to process
+        # already has PII tags or not
+        column_has_pii_tag = any(
+            (PII in tag.tagFQN.__root__ for tag in column.tags or [])
+        )
+
+        # If it has PII tags, we skip the processing
+        # for the column
+        if column_has_pii_tag is True:
+            return None
+
+        # Scan by column name. If no results there, check the sample data, if any
+        tag_and_confidence = ColumnNameScanner.scan(column.name.__root__) or (
+            self.ner_scanner.scan([row[idx] for row in table_data.rows])
+            if table_data
+            else None
+        )
+
+        if (
+            tag_and_confidence
+            and tag_and_confidence.tag_fqn
+            and tag_and_confidence.confidence >= confidence_threshold / 100
+        ):
+            # We support returning +1 tags for a single column in _run
+            return [
+                self.build_column_tag(
+                    tag_fqn=tag_and_confidence.tag_fqn,
+                    column_fqn=column.fullyQualifiedName.__root__,
+                )
+            ]
+
+        return None
+
+    def process_table(self, table: Table) -> Optional[List[ColumnTag]]:
+        """Run the patching of the table"""
+        column_tags = []
+        for idx, column in enumerate(table.columns):
+            try:
+                col_tags = self.process_column(
+                    idx=idx,
+                    column=column,
+                    table_data=table.sampleData,
+                    confidence_threshold=self.confidence_threshold,
+                )
+                if col_tags:
+                    column_tags.extend(col_tags)
+            except Exception as err:
+                self.status.failed(
+                    StackTraceError(
+                        name=table.fullyQualifiedName.__root__,
+                        error=f"Error computing PII tags for [{column}] - [{err}]",
+                        stackTrace=traceback.format_exc(),
+                    )
+                )
+
+        if column_tags:
+            return column_tags
+
+        return None
+
+    def patch_columns(self, table: Table, column_tags: List[ColumnTag]) -> None:
+        """Patch columns with PII"""
+        patched = self.metadata.patch_column_tags(table=table, column_tags=column_tags)
+        if not patched:
+            self.status.failed(
+                StackTraceError(
+                    name=table.fullyQualifiedName.__root__,
+                    error="Error patching tags for table",
+                )
+            )
+        else:
+            self.status.scanned(table)
+            logger.debug(
+                f"Successfully patched tag {column_tags} for {table.fullyQualifiedName.__root__}"
+            )
+
+    def run(self) -> None:
+        """
+        The PII Application will:
+        1. List tables
+        2. Check their column names and sample data (if any)
+        3. PATCH PII tags when needed
+        """
+        tables: Iterable[Table] = self.metadata.list_all_entities(
+            entity=Table, fields=["sampleData", "tags"]
+        )
+        for table in tables:
+            column_tags = self.process_table(table)
+            if column_tags:
+                self.patch_columns(table=table, column_tags=column_tags)
+            else:
+                self.status.filter(
+                    key=table.fullyQualifiedName.__root__, reason="No PII found"
+                )
+
+    def close(self) -> None:
+        """Nothing to close"""

metadata/cli/app.py

@@ -0,0 +1,47 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  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.
+
+"""
+Profiler utility for the metadata CLI
+"""
+import sys
+import traceback
+from pathlib import Path
+
+from metadata.config.common import load_config_file
+from metadata.utils.logger import cli_logger
+from metadata.workflow.application import ApplicationWorkflow
+from metadata.workflow.application_output_handler import print_status
+
+logger = cli_logger()
+
+
+def run_app(config_path: Path) -> None:
+    """
+    Run the application workflow from a config path
+    to a JSON or YAML file
+    :param config_path: Path to load JSON config
+    """
+
+    config_dict = None
+    try:
+        config_dict = load_config_file(config_path)
+        workflow = ApplicationWorkflow.create(config_dict)
+        logger.debug(f"Using config: {workflow.config}")
+    except Exception as exc:
+        logger.error(f"Error running the application {exc}")
+        logger.debug(traceback.format_exc())
+        sys.exit(1)
+
+    workflow.execute()
+    workflow.stop()
+    print_status(workflow)
+    workflow.raise_from_status()

metadata/cli/db_dump.py

@@ -31,6 +31,7 @@ TABLES_DUMP_ALL = {
     "field_relationship",
     "tag_usage",
     "openmetadata_settings",
+    "profiler_data_time_series",
 }
 
 CUSTOM_TABLES = {"entity_extension_time_series": {"exclude_columns": ["timestamp"]}}
@@ -123,9 +124,12 @@ def get_hash_column_name(engine: Engine, table_name: str) -> Optional[str]:
 
 def run_query_iter(engine: Engine, query: str) -> Iterable[Row]:
     """Return a generator of rows, one row at a time, with a limit of 100 in-mem rows"""
-
-    for row in engine.execute(text(query)).yield_per(100):
-        yield row
+    with engine.connect() as conn:
+        result = conn.execution_options(
+            stream_results=True, max_row_buffer=100
+        ).execute(text(query))
+        for row in result:
+            yield row
 
 
 def dump_json(tables: List[str], engine: Engine, output: Path) -> None:

metadata/clients/aws_client.py

@@ -35,6 +35,7 @@ class AWSServices(Enum):
     QUICKSIGHT = "quicksight"
     ATHENA = "athena"
     RDS = "rds"
+    LAKE_FORMATION = "lakeformation"
 
 
 class AWSAssumeRoleException(Exception):
@@ -189,3 +190,6 @@ class AWSClient:
 
     def get_athena_client(self):
         return self.get_client(AWSServices.ATHENA.value)
+
+    def get_lake_formation_client(self):
+        return self.get_client(AWSServices.LAKE_FORMATION.value)

metadata/clients/azure_client.py

@@ -0,0 +1,85 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  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.
+"""
+Module containing Azure Client
+"""
+
+from metadata.generated.schema.security.credentials.azureCredentials import (
+    AzureCredentials,
+)
+from metadata.utils.logger import utils_logger
+
+logger = utils_logger()
+
+
+class AzureClient:
+    """
+    AzureClient based on AzureCredentials.
+    """
+
+    def __init__(self, credentials: "AzureCredentials"):
+        self.credentials = credentials
+        if not isinstance(credentials, AzureCredentials):
+            self.credentials = AzureCredentials.parse_obj(credentials)
+
+    def create_client(
+        self,
+    ):
+        from azure.identity import ClientSecretCredential, DefaultAzureCredential
+
+        try:
+            if (
+                getattr(self.credentials, "tenantId", None)
+                and getattr(self.credentials, "clientId", None)
+                and getattr(self.credentials, "clientSecret", None)
+            ):
+                logger.info("Using Client Secret Credentials")
+                return ClientSecretCredential(
+                    tenant_id=self.credentials.tenantId,
+                    client_id=self.credentials.clientId,
+                    client_secret=self.credentials.clientSecret.get_secret_value(),
+                )
+            else:
+                logger.info("Using Default Azure Credentials")
+                return DefaultAzureCredential()
+        except Exception as e:
+            logger.error(f"Error creating Azure Client: {e}")
+            raise e
+
+    def create_blob_client(self):
+        from azure.storage.blob import BlobServiceClient
+
+        try:
+            logger.info("Creating Blob Service Client")
+            if self.credentials.accountName:
+                return BlobServiceClient(
+                    account_url=f"https://{self.credentials.accountName}.blob.core.windows.net/",
+                    credential=self.create_client(),
+                )
+            raise ValueError("Account Name is required to create Blob Service Client")
+        except Exception as e:
+            logger.error(f"Error creating Blob Service Client: {e}")
+            raise e
+
+    def create_secret_client(self):
+        from azure.keyvault.secrets import SecretClient
+
+        try:
+            if self.credentials.vaultName:
+                logger.info("Creating Secret Client")
+                return SecretClient(
+                    vault_url=f"https://{self.credentials.vaultName}.vault.azure.net/",
+                    credential=self.create_client(),
+                )
+            raise ValueError("Vault Name is required to create a Secret Client")
+        except Exception as e:
+            logger.error(f"Error creating Secret Client: {e}")
+            raise e

metadata/clients/domo_client.py

@@ -146,8 +146,6 @@ class DomoClient:
 
         return None
 
-    # def get_owner_details(self, owner_id) -> dict:
-
     def get_pipelines(self):
         try:
             response = self.client._request(  # pylint: disable=protected-access

metadata/cmd.py

@@ -18,6 +18,7 @@ from http.server import BaseHTTPRequestHandler, HTTPServer
 from pathlib import Path
 
 from metadata.__version__ import get_metadata_version
+from metadata.cli.app import run_app
 from metadata.cli.backup import UploadDestinationType, run_backup
 from metadata.cli.dataquality import run_test
 from metadata.cli.docker import BACKEND_DATABASES, DockerActions, run_docker
@@ -50,10 +51,22 @@ class MetadataCommands(Enum):
     WEBHOOK = "webhook"
     INSIGHT = "insight"
     LINEAGE = "lineage"
+    APP = "app"
     OPENMETADATA_IMPORTS_MIGRATION = "openmetadata_imports_migration"
     OPENMETADATA_DAG_CONFIG_MIGRATION = "openmetadata_dag_config_migration"
 
 
+RUN_PATH_METHODS = {
+    MetadataCommands.INGEST.value: run_ingest,
+    MetadataCommands.USAGE.value: run_usage,
+    MetadataCommands.LINEAGE.value: run_lineage,
+    MetadataCommands.INSIGHT.value: run_insight,
+    MetadataCommands.PROFILE.value: run_profiler,
+    MetadataCommands.TEST.value: run_test,
+    MetadataCommands.APP.value: run_app,
+}
+
+
 OM_IMPORTS_MIGRATION = """
     Update DAG files generated after creating workflow in 0.11 and before.
     In 0.12 the airflow managed API package name changed from `openmetadata` to 
@@ -365,6 +378,12 @@ def get_parser(args=None):
             MetadataCommands.TEST.value, help="Workflow for running test suites"
         )
     )
+    create_common_config_parser_args(
+        sub_parser.add_parser(
+            MetadataCommands.APP.value,
+            help="Workflow for running external applications",
+        )
+    )
     create_openmetadata_imports_migration_args(
         sub_parser.add_parser(
             MetadataCommands.OPENMETADATA_IMPORTS_MIGRATION.value,
@@ -409,13 +428,14 @@ def get_parser(args=None):
     return parser.parse_args(args)
 
 
-def metadata(args=None):  # pylint: disable=too-many-branches
+def metadata(args=None):
     """
     This method implements parsing of the arguments passed from CLI
     """
     contains_args = vars(get_parser(args))
     metadata_workflow = contains_args.get("command")
     config_file = contains_args.get("config")
+    path = None
     if config_file:
         path = Path(config_file).expanduser()
     if contains_args.get("debug"):
@@ -425,18 +445,9 @@ def metadata(args=None):  # pylint: disable=too-many-branches
     else:
         set_loggers_level(logging.INFO)
 
-    if metadata_workflow == MetadataCommands.INGEST.value:
-        run_ingest(config_path=path)
-    if metadata_workflow == MetadataCommands.USAGE.value:
-        run_usage(config_path=path)
-    if metadata_workflow == MetadataCommands.LINEAGE.value:
-        run_lineage(config_path=path)
-    if metadata_workflow == MetadataCommands.INSIGHT.value:
-        run_insight(config_path=path)
-    if metadata_workflow == MetadataCommands.PROFILE.value:
-        run_profiler(config_path=path)
-    if metadata_workflow == MetadataCommands.TEST.value:
-        run_test(config_path=path)
+    if metadata_workflow in RUN_PATH_METHODS:
+        RUN_PATH_METHODS[metadata_workflow](path)
+
     if metadata_workflow == MetadataCommands.BACKUP.value:
         run_backup(
             common_backup_obj_instance=BackupRestoreArgs(