openmetadata-ingestion 1.3.3.0rc2__py3-none-any.whl → 1.4.0.0__py3-none-any.whl

metadata/applications/auto_tagger.py

@@ -1,212 +0,0 @@
-#  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/backup.py

@@ -1,202 +0,0 @@
-#  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.
-
-"""
-Backup utility for the metadata CLI
-"""
-import traceback
-from datetime import datetime
-from enum import Enum
-from pathlib import Path
-from typing import Optional, Tuple
-
-from metadata.cli.db_dump import dump
-from metadata.cli.utils import get_engine
-from metadata.utils.helpers import BackupRestoreArgs
-from metadata.utils.logger import ANSI, cli_logger, log_ansi_encoded_string
-
-
-class UploadDestinationType(Enum):
-    AWS = "AWS"
-    AZURE = "Azure"
-
-
-logger = cli_logger()
-
-
-def get_output(output: Optional[str] = None, filename: Optional[str] = None) -> Path:
-    """
-    Helper function to prepare the output backup file
-    path and name.
-
-    It will create the output dir if it does not exist.
-
-    :param output: local path to store the backup
-    :param filename: name of the backup file
-    :return: backup file name
-    """
-    now = datetime.now().strftime("%Y%m%d%H%M")
-    name = filename if filename else f"openmetadata_{now}_backup.sql"
-
-    if output:
-        path = Path(output).expanduser()
-        # Create the output directory if it does not exist
-        if not path.is_dir():
-            path.mkdir(parents=True, exist_ok=True)
-
-        return path / name
-
-    return Path(name)
-
-
-def upload_backup_aws(endpoint: str, bucket: str, key: str, file: Path) -> None:
-    """
-    Upload the mysqldump backup file.
-    We will use boto3 to upload the file to the endpoint
-    and the key provided.
-
-    :param endpoint: S3 endpoint
-    :param bucket: S3 bucket to upload the file to
-    :param key: S3 key to upload the backup file
-    :param file: file to upload
-    """
-
-    try:
-        # We just want to force boto3 install if uploading backup
-        # pylint: disable=import-outside-toplevel
-        import boto3
-        from boto3.exceptions import S3UploadFailedError
-    except ModuleNotFoundError as err:
-        logger.debug(traceback.format_exc())
-        logger.error(
-            "Trying to import boto3 to run the backup upload."
-            + " Please install openmetadata-ingestion[backup]."
-        )
-        raise err
-
-    s3_key = Path(key) / file.name
-    log_ansi_encoded_string(
-        color=ANSI.GREEN,
-        bold=False,
-        message=f"Uploading {file} to {endpoint}/{bucket}/{str(s3_key)}...",
-    )
-
-    try:
-        resource = boto3.resource(service_name="s3", endpoint_url=endpoint)
-        resource.Object(bucket, str(s3_key)).upload_file(str(file.absolute()))
-
-    except ValueError as err:
-        logger.debug(traceback.format_exc())
-        logger.error("Revisit the values of --upload")
-        raise err
-    except S3UploadFailedError as err:
-        logger.debug(traceback.format_exc())
-        logger.error(
-            "Error when uploading the backup to S3. Revisit the config and permissions."
-            + " You should have set the environment values for AWS_ACCESS_KEY_ID"
-            + " and AWS_SECRET_ACCESS_KEY"
-        )
-        raise err
-
-
-def upload_backup_azure(account_url: str, container: str, file: Path) -> None:
-    """
-    Upload the mysqldump backup file.
-
-    :param account_url: Azure account url
-    :param container: Azure container to upload file to
-    :param file: file to upload
-    """
-
-    try:
-        # pylint: disable=import-outside-toplevel
-        from azure.identity import DefaultAzureCredential
-        from azure.storage.blob import BlobServiceClient
-
-        default_credential = DefaultAzureCredential()
-        # Create the BlobServiceClient object
-        blob_service_client = BlobServiceClient(
-            account_url, credential=default_credential
-        )
-    except ModuleNotFoundError as err:
-        logger.debug(traceback.format_exc())
-        logger.error(
-            "Trying to import DefaultAzureCredential to run the backup upload."
-        )
-        raise err
-
-    log_ansi_encoded_string(
-        color=ANSI.GREEN,
-        message=f"Uploading {file} to {account_url}/{container}...",
-    )
-
-    try:
-        # Create a blob client using the local file name as the name for the blob
-        blob_client = blob_service_client.get_blob_client(
-            container=container, blob=file.name
-        )
-
-        # Upload the created file
-        with open(file=file, mode="rb") as data:
-            blob_client.upload_blob(data)
-
-    except ValueError as err:
-        logger.debug(traceback.format_exc())
-        logger.error("Revisit the values of --upload")
-        raise err
-    except Exception as err:
-        logger.debug(traceback.format_exc())
-        logger.error(err)
-        raise err
-
-
-def run_backup(
-    common_backup_obj_instance: BackupRestoreArgs,
-    output: Optional[str],
-    filename: Optional[str],
-    upload_destination_type: Optional[UploadDestinationType],
-    upload: Optional[Tuple[str, str, str]],
-) -> None:
-    """
-    Run `mysqldump` to MySQL database and store the
-    output. Optionally, upload it to S3.
-
-    :param common_backup_obj_instance: cls instance to fetch common args
-    :param output: local path to store the backup
-    :param filename: filename to store the backup
-    :param upload_destination_type: Azure or AWS Destination Type
-    :param upload: URI to upload result file
-
-    """
-    log_ansi_encoded_string(
-        color=ANSI.GREEN,
-        bold=False,
-        message="Creating OpenMetadata backup for "
-        f"{common_backup_obj_instance.host}:{common_backup_obj_instance.port}/{common_backup_obj_instance.database}...",
-    )
-
-    out = get_output(output, filename)
-
-    engine = get_engine(common_args=common_backup_obj_instance)
-    dump(engine=engine, output=out, schema=common_backup_obj_instance.schema)
-
-    log_ansi_encoded_string(
-        color=ANSI.GREEN, bold=False, message=f"Backup stored locally under {out}"
-    )
-
-    if upload:
-        if upload_destination_type == UploadDestinationType.AWS.value:
-            endpoint, bucket, key = upload
-            upload_backup_aws(endpoint, bucket, key, out)
-        elif upload_destination_type.title() == UploadDestinationType.AZURE.value:
-            # only need two parameters from upload, key would be null
-            account_url, container, key = upload
-            upload_backup_azure(account_url, container, out)

metadata/cli/db_dump.py

@@ -1,239 +0,0 @@
-#  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.
-
-"""
-Database Dumping utility for the metadata CLI
-"""
-
-import json
-from functools import singledispatch
-from pathlib import Path
-from typing import Iterable, List, Optional, Union
-
-from sqlalchemy import inspect, text
-from sqlalchemy.engine import Engine, Row
-
-from metadata.utils.constants import UTF_8
-
-TABLES_DUMP_ALL = {
-    "task_sequence",
-    "entity_usage",
-    "entity_relationship",
-    "entity_extension",
-    "field_relationship",
-    "tag_usage",
-    "openmetadata_settings",
-    "profiler_data_time_series",
-}
-
-CUSTOM_TABLES = {"entity_extension_time_series": {"exclude_columns": ["timestamp"]}}
-NOT_MIGRATE = {"DATABASE_CHANGE_LOG", "SERVER_MIGRATION_SQL_LOGS", "SERVER_CHANGE_LOG"}
-
-STATEMENT_JSON = "SELECT json FROM {table}"
-STATEMENT_HASH_JSON = "SELECT json, {hash_column_name} FROM {table}"
-STATEMENT_ALL = "SELECT * FROM {table}"
-STATEMENT_TRUNCATE = "TRUNCATE TABLE {table};\n"
-STATEMENT_ALL_NEW = "SELECT {cols} FROM {table}"
-
-MYSQL_ENGINE_NAME = "mysql"
-FQN_HASH_COLUMN = "fqnHash"
-NAME_HASH_COLUMN = "nameHash"
-
-
-def single_quote_wrap(raw: str) -> str:
-    """
-    Add single quote wrap to string. From `str` to `'str'`
-    """
-    return f"'{raw}'"
-
-
-def clean_str(raw: str, engine: Engine) -> str:
-    """
-    String cleaning for SQL parsing.
-
-    Logic is different between MySQL and Postgres
-
-    - descriptions/comments with single quotes, e.g., `Mysql's data`.
-      get converted to `Mysql''s data`
-    - To insert a literal backlash in MySQL you need to escape with another one. This applies for `\n` and `\"` in
-      inner JSONs for a field. This is not required for postgres
-    """
-    quoted_str = raw.replace("'", "''")
-
-    if engine.name == MYSQL_ENGINE_NAME:
-        quoted_str = quoted_str.replace("\\", "\\\\")
-
-    return quoted_str
-
-
-@singledispatch
-def clean_col(column_raw: Optional[Union[dict, str]], engine: Engine) -> str:
-    return (
-        single_quote_wrap(clean_str(str(column_raw), engine))
-        if column_raw is not None
-        else "null"
-    )
-
-
-@clean_col.register(dict)
-@clean_col.register(list)
-def _(column_raw: Optional[Union[dict, list]], engine: Engine) -> str:
-    """
-    Prepare the JSON column to be inserted to MySQL
-
-    Handle:
-    - quotes
-    - True/False values
-    """
-    return (
-        single_quote_wrap(
-            clean_str(
-                json.dumps(
-                    column_raw,
-                    default=str,  # If we don't know how to serialize, convert to str
-                ),
-                engine,
-            )
-        )
-        if column_raw is not None
-        else "null"
-    )
-
-
-def get_hash_column_name(engine: Engine, table_name: str) -> Optional[str]:
-    """
-    Method to get name of the hash column (fqnHash or nameHash)
-    """
-    inspector = inspect(engine)
-    columns = inspector.get_columns(table_name)
-    for column in columns:
-        if column["name"].lower() == FQN_HASH_COLUMN.lower():
-            return column["name"]
-        if column["name"].lower() == NAME_HASH_COLUMN.lower():
-            return column["name"]
-    return None
-
-
-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"""
-    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:
-    """
-    Dumps JSON data.
-
-    Postgres: engine.name == "postgresql"
-    MySQL: engine.name == "mysql"
-    """
-    with open(output, "a", encoding=UTF_8) as file:
-        for table in tables:
-            truncate = STATEMENT_TRUNCATE.format(table=table)
-            file.write(truncate)
-
-            hash_column_name = get_hash_column_name(engine=engine, table_name=table)
-            if hash_column_name:
-                query = STATEMENT_HASH_JSON.format(
-                    table=table, hash_column_name=hash_column_name
-                )
-                for row in run_query_iter(engine=engine, query=query):
-                    insert = f"INSERT INTO {table} (json, {hash_column_name}) VALUES ({clean_col(row.json, engine)}, {clean_col(row[1], engine)});\n"  # pylint: disable=line-too-long
-                    file.write(insert)
-            else:
-                res = engine.execute(text(STATEMENT_JSON.format(table=table))).all()
-                for row in res:
-                    insert = f"INSERT INTO {table} (json) VALUES ({clean_col(row.json, engine)});\n"
-                    file.write(insert)
-
-
-def dump_all(tables: List[str], engine: Engine, output: Path) -> None:
-    """
-    Dump tables that need to store all data
-    """
-    with open(output, "a", encoding=UTF_8) as file:
-        for table in tables:
-            truncate = STATEMENT_TRUNCATE.format(table=table)
-            file.write(truncate)
-
-            query = STATEMENT_ALL.format(table=table)
-            for row in run_query_iter(engine=engine, query=query):
-                data = ",".join(clean_col(col, engine) for col in row)
-
-                insert = f"INSERT INTO {table} VALUES ({data});\n"
-                file.write(insert)
-
-
-def dump_entity_custom(engine: Engine, output: Path, inspector) -> None:
-    """
-    This function is used to dump entities with custom handling
-    """
-    with open(output, "a", encoding=UTF_8) as file:
-        for table, data in CUSTOM_TABLES.items():
-            truncate = STATEMENT_TRUNCATE.format(table=table)
-            file.write(truncate)
-
-            columns = inspector.get_columns(table_name=table)
-
-            query = STATEMENT_ALL_NEW.format(
-                cols=",".join(
-                    col["name"]
-                    for col in columns
-                    if col["name"] not in data["exclude_columns"]
-                ),
-                table=table,
-            )
-            for row in run_query_iter(engine=engine, query=query):
-                # Let's use .format here to not add more variables
-                # pylint: disable=consider-using-f-string
-                insert = "INSERT INTO {table} ({cols}) VALUES ({data});\n".format(
-                    table=table,
-                    data=",".join(clean_col(col, engine) for col in row),
-                    cols=",".join(
-                        col["name"]
-                        for col in columns
-                        if col["name"] not in data["exclude_columns"]
-                    ),
-                )
-                file.write(insert)
-
-
-def get_lower_table_names(tables):
-    return [table.lower() for table in tables]
-
-
-def dump(engine: Engine, output: Path, schema: str = None) -> None:
-    """
-    Get all tables from the database and dump
-    only the JSON column for the required tables
-    """
-    inspector = inspect(engine)
-    tables = (
-        inspector.get_table_names(schema) if schema else inspector.get_table_names()
-    )
-    lower_tables = get_lower_table_names(tables)
-    all_non_json_tables = (
-        get_lower_table_names(TABLES_DUMP_ALL)
-        + get_lower_table_names(NOT_MIGRATE)
-        + get_lower_table_names(CUSTOM_TABLES)
-    )
-
-    dump_json_tables = [
-        table for table in lower_tables if table not in all_non_json_tables
-    ]
-
-    dump_all(tables=list(TABLES_DUMP_ALL), engine=engine, output=output)
-    dump_json(tables=dump_json_tables, engine=engine, output=output)
-    dump_entity_custom(engine=engine, output=output, inspector=inspector)