openmetadata-ingestion 1.4.8.0__py3-none-any.whl → 1.5.0.0rc2__py3-none-any.whl

_openmetadata_testutils/helpers/assumption.py

@@ -0,0 +1,118 @@
+from typing import Callable, List
+
+import pytest
+from pandas import DataFrame, Series
+
+
+class Assumption:
+    def __init__(self):
+        self.not_ = False
+
+    def validate(self, df: DataFrame):
+        if self.not_:
+            self.assume_false(df)
+        else:
+            self.assume_positive(df)
+
+    def __invert__(self):
+        return self.copy(not_=not self.not_)
+
+    def assume_positive(self, df: DataFrame):
+        raise NotImplementedError
+
+    def assume_false(self, df: DataFrame):
+        with pytest.raises(AssertionError):
+            self.assume_positive(df)
+
+    def copy(self, **kwargs):
+        new = self.__class__.__new__(self.__class__)
+        new.__dict__.update(self.__dict__)
+        new.__dict__.update(kwargs)
+        return new
+
+
+class AssumeUnqiue(Assumption):
+    def __init__(self, column: str):
+        super().__init__()
+        self.column = column
+
+    def assume_positive(self, df: DataFrame):
+        assert df[self.column].nunique() == len(df)
+
+
+class AssumeNotNull(Assumption):
+    def __init__(self, column: str):
+        super().__init__()
+        self.column = column
+
+    def assume_positive(self, df: DataFrame):
+        assert df[self.column].notnull().all()
+
+
+class AssumeMatchRegex(Assumption):
+    def __init__(self, column: str, regex: str):
+        super().__init__()
+        self.column = column
+        self.regex = regex
+
+    def assume_positive(self, df: DataFrame):
+        assert df[self.column].str.match(self.regex).all()
+
+
+class AssumeBetween(Assumption):
+    def __init__(self, column: str, min_value: int, max_value: int):
+        super().__init__()
+        self.column = column
+        self.min_value = min_value
+        self.max_value = max_value
+
+    def assume_positive(self, df: DataFrame):
+        assert (df[self.column] >= self.min_value).all()
+        assert (df[self.column] <= self.max_value).all()
+
+
+class AssumeLengthBetween(Assumption):
+    def __init__(self, column: str, min_length: int, max_length: int):
+        super().__init__()
+        self.column = column
+        self.min_length = min_length
+        self.max_length = max_length
+
+    def assume_positive(self, df: DataFrame):
+        assert (df[self.column].str.len() >= self.min_length).all()
+        assert (df[self.column].str.len() <= self.max_length).all()
+
+
+class AssumeColumnValuesIn(Assumption):
+    def __init__(self, column: str, allowed_values: List[str]):
+        super().__init__()
+        self.column = column
+        self.allowed_values = allowed_values
+
+    def assume_positive(self, df: DataFrame):
+        assert df[self.column].isin(self.allowed_values).all()
+
+
+class AssumeArbitrary(Assumption):
+    def __init__(self, column: str, fn: Callable[[Series], Series]):
+        super().__init__()
+        self.column = column
+        self.fn = fn
+
+    def assume_positive(self, df: DataFrame):
+        assert self.fn(
+            df[self.column]
+        ).all(), f"failed test {self.__class__.__name__} for column {self.column}"
+
+
+class Assumptions:
+    notnull = AssumeNotNull
+    unique = AssumeUnqiue
+    match_regex = AssumeMatchRegex
+    between = AssumeBetween
+    length_between = AssumeLengthBetween
+    column_values_in = AssumeColumnValuesIn
+    arbitrary = AssumeArbitrary
+
+
+assume = Assumptions()

_openmetadata_testutils/helpers/docker.py

@@ -0,0 +1,38 @@
+import contextlib
+import logging
+import os
+import tarfile
+
+import docker
+from docker.models.containers import Container
+
+
+@contextlib.contextmanager
+def try_bind(container, container_port, host_port):
+    """Try to bind a port to the container, if it is already in use try another port."""
+    try:
+        with container.with_bind_ports(container_port, host_port) as container:
+            yield container
+    except docker.errors.APIError:
+        logging.warning("Port %s is already in use, trying another port", host_port)
+        with container.with_bind_ports(container_port, None) as container:
+            yield container
+
+
+def copy_dir_to_container(dir_path: str, container: Container, container_path: str):
+    """Copy the contents of a directory to a path in a container. If the path
+    does not exist it will be created.
+
+    Args:
+        dir_path (str): Path to the directory to copy.
+        container (DockerContainer): The container to copy the directory to.
+        container_path (str): The path to copy the directory to in the container.
+    """
+    tar_path = dir_path + ".tar"
+    with tarfile.open(tar_path, "w") as tar:
+        for item in os.listdir(dir_path):
+            tar.add(os.path.join(dir_path, item), arcname=item)
+    container.exec_run(["mkdir", "-p", container_path])
+    with open(tar_path, "rb") as tar_file:
+        container.put_archive(container_path, tar_file)
+    os.remove(tar_path)

_openmetadata_testutils/helpers/markers.py

@@ -0,0 +1,5 @@
+import pytest
+
+
+def xfail_param(param, reason):
+    return pytest.param(param, marks=pytest.mark.xfail(reason=reason, strict=True))

_openmetadata_testutils/kafka/load_csv_data.py

@@ -0,0 +1,133 @@
+"""
+pip install confluent_kafka pandas requests avro-python3
+"""
+
+import json
+import os
+from typing import Dict, List
+
+import pandas as pd
+import requests
+from confluent_kafka import Producer
+from confluent_kafka.schema_registry import SchemaRegistryClient
+from confluent_kafka.schema_registry.avro import AvroSerializer
+from confluent_kafka.serialization import MessageField, SerializationContext
+
+
+class Kafka:
+    def __init__(self, broker: str):
+        self.broker = broker
+
+    def get_producer(self):
+        return Producer({"bootstrap.servers": self.broker})
+
+
+class SchemaRegistry:
+    def __init__(self, url: str):
+        self.url = url
+
+    def register_schema(self, schema: str, topic: str):
+        """Register the schema with the Schema Registry"""
+        url = f"{self.url}/subjects/{sanitize_name(topic)}-value/versions"
+        headers = {"Content-Type": "application/vnd.schemaregistry.v1+json"}
+        data = json.dumps({"schema": schema})
+        response = requests.post(url, headers=headers, data=data)
+        if response.status_code == 200:
+            print(f"Schema registered for topic {topic}")
+        else:
+            print(f"Failed to register schema for topic {topic}: {response.text}")
+
+    def get_avro_serializer(self, schema: str) -> AvroSerializer:
+        schema_registry_conf = {"url": self.url}
+        schema_registry_client = SchemaRegistryClient(schema_registry_conf)
+        return AvroSerializer(schema_registry_client, schema_str=schema)
+
+
+def delivery_report(err, msg):
+    """Called once for each message produced to indicate delivery result.
+    Triggered by poll() or flush()."""
+    if err is not None:
+        print(f"Message delivery failed: {err}")
+    else:
+        print(f"Message delivered to {msg.topic()} [{msg.partition()}]")
+
+
+def sanitize_name(name):
+    """Sanitize the name to conform to Avro naming rules"""
+    return name.replace("-", "_").replace(" ", "_")
+
+
+def generate_avro_schema(df: pd.DataFrame, topic: str) -> str:
+    """Generate an Avro schema from a pandas DataFrame"""
+    fields: List[Dict[str, str]] = []
+    for column in df.columns:
+        fields.append(
+            {"name": sanitize_name(column), "type": "string"}
+        )  # Assuming all columns are of type string
+    schema_dict: Dict[str, any] = {
+        "namespace": "example.avro",
+        "type": "record",
+        "name": sanitize_name(topic),
+        "fields": fields,
+    }
+    return json.dumps(schema_dict)
+
+
+def send_csv_to_kafka(kafka: Kafka, schema_registry: SchemaRegistry, file_path: str):
+    # Read the CSV file
+    df = pd.read_csv(file_path)
+
+    # Sanitize column names
+    df.columns = [sanitize_name(col) for col in df.columns]
+
+    # Convert all fields to string
+    df = df.astype(str)
+
+    # Get the file name without extension to use as the topic name
+    topic = os.path.splitext(os.path.basename(file_path))[0]
+
+    # Generate and register the Avro schema
+    schema = generate_avro_schema(df, topic)
+    schema_registry.register_schema(schema, topic)
+
+    # Create an Avro serializer with the generated schema
+    avro_serializer = schema_registry.get_avro_serializer(schema)
+
+    # Create a Kafka producer
+    producer = kafka.get_producer()
+
+    # Iterate over the rows of the DataFrame and send each row as an Avro message
+    for _, row in df.iterrows():
+        message = row.to_dict()
+        try:
+            producer.produce(
+                topic=topic,
+                value=avro_serializer(
+                    message, SerializationContext(topic, MessageField.VALUE)
+                ),
+                callback=delivery_report,
+            )
+        except Exception as e:
+            print(f"Message serialization failed: {e}")
+            break
+
+    # Wait for any outstanding messages to be delivered and delivery reports to be received
+    producer.flush()
+
+
+def main(kafka_broker: str, schema_registry_url: str, csv_directory: str):
+    # Iterate over all files in the directory
+    kafka = Kafka(kafka_broker)
+    schema_registry = SchemaRegistry(schema_registry_url)
+    for file_name in os.listdir(csv_directory):
+        if file_name.endswith(".csv"):
+            file_path = os.path.join(csv_directory, file_name)
+            send_csv_to_kafka(kafka, schema_registry, file_path)
+
+
+if __name__ == "__main__":
+    main(
+        kafka_broker="localhost:9092",
+        schema_registry_url="http://localhost:8081",
+        csv_directory="./data",
+    )

_openmetadata_testutils/kafka/schema_registry_container.py

@@ -0,0 +1,31 @@
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.waiting_utils import wait_for_logs
+
+
+class SchemaRegistryContainer(DockerContainer):
+    def __init__(
+        self,
+        image: str = "confluentinc/cp-schema-registry:7.6.0",
+        port: int = 8081,
+        schema_registry_kafkastore_bootstrap_servers="PLAINTEXT://localhost:9092",
+        schema_registry_host_name="localhost",
+        **kwargs,
+    ) -> None:
+        super().__init__(image, **kwargs)
+        self.with_env(
+            "SCHEMA_REGISTRY_KAFKASTORE_BOOTSTRAP_SERVERS",
+            schema_registry_kafkastore_bootstrap_servers,
+        )
+        self.with_env("SCHEMA_REGISTRY_HOST_NAME", schema_registry_host_name)
+        self.port = port
+        self.with_exposed_ports(port)
+
+    def start(self, timeout=30) -> "SchemaRegistryContainer":
+        super().start()
+        wait_for_logs(self, r".*\Server started.*", timeout=timeout)
+        return self
+
+    def get_connection_url(self):
+        return (
+            f"http://{self.get_container_host_ip()}:{self.get_exposed_port(self.port)}"
+        )

_openmetadata_testutils/ometa.py

@@ -0,0 +1,23 @@
+from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
+    AuthProvider,
+    OpenMetadataConnection,
+)
+from metadata.generated.schema.security.client.openMetadataJWTClientConfig import (
+    OpenMetadataJWTClientConfig,
+)
+from metadata.ingestion.models.custom_pydantic import CustomSecretStr
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+
+OM_JWT = "eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg"
+
+
+def int_admin_ometa(url: str = "http://localhost:8585/api") -> OpenMetadata:
+    """Initialize the ometa connection with default admin:admin creds"""
+    server_config = OpenMetadataConnection(
+        hostPort=url,
+        authProvider=AuthProvider.openmetadata,
+        securityConfig=OpenMetadataJWTClientConfig(jwtToken=CustomSecretStr(OM_JWT)),
+    )
+    metadata = OpenMetadata(server_config)
+    assert metadata.health_check()
+    return metadata

_openmetadata_testutils/postgres/conftest.py

@@ -0,0 +1,75 @@
+"""Test fixtures for PostgreSQL. It uses a testcontainer to start a PostgreSQL instance with the dvdrental database."""
+
+import logging
+import os
+import zipfile
+from subprocess import CalledProcessError
+
+import pytest
+from sqlalchemy import create_engine
+from testcontainers.postgres import PostgresContainer
+
+from _openmetadata_testutils.helpers.docker import copy_dir_to_container, try_bind
+
+
+@pytest.fixture(autouse=True, scope="session")
+def config_logging():
+    """Set the log level of sqlfluff to CRITICAL."""
+    logging.getLogger("sqlfluff").setLevel(logging.CRITICAL)
+
+
+@pytest.fixture(scope="module")
+def postgres_container(tmp_path_factory):
+    """Start a PostgreSQL container with the dvdrental database."""
+    data_dir = tmp_path_factory.mktemp("data")
+    dvd_rental_zip = os.path.join(
+        os.path.dirname(os.path.dirname(__file__)), "data", "dvdrental.zip"
+    )
+    zipfile.ZipFile(dvd_rental_zip, "r").extractall(str(data_dir))
+    container = PostgresContainer("postgres:15", dbname="dvdrental")
+    container._command = [
+        "-c",
+        "shared_preload_libraries=pg_stat_statements",
+        "-c",
+        "pg_stat_statements.max=10000",
+        "-c",
+        "pg_stat_statements.track=all",
+        "-c",
+        "track_commit_timestamp=on",
+    ]
+
+    with (
+        try_bind(container, 5432, 5432) if not os.getenv("CI") else container
+    ) as container:
+        docker_container = container.get_wrapped_container()
+        copy_dir_to_container(str(data_dir), docker_container, "/data")
+        for query in (
+            "CREATE USER postgres SUPERUSER;",
+            "CREATE EXTENSION pg_stat_statements;",
+        ):
+            res = docker_container.exec_run(
+                ["psql", "-U", container.username, "-d", container.dbname, "-c", query]
+            )
+            if res[0] != 0:
+                raise CalledProcessError(
+                    returncode=res[0], cmd=res, output=res[1].decode("utf-8")
+                )
+        res = docker_container.exec_run(
+            [
+                "pg_restore",
+                "-U",
+                container.username,
+                "-d",
+                container.dbname,
+                "/data/dvdrental.tar",
+            ]
+        )
+        if res[0] != 0:
+            raise CalledProcessError(
+                returncode=res[0], cmd=res, output=res[1].decode("utf-8")
+            )
+        engine = create_engine(container.get_connection_url())
+        engine.execute(
+            "ALTER TABLE customer ADD COLUMN json_field JSONB DEFAULT '{}'::JSONB;"
+        )
+        yield container

_openmetadata_testutils/pydantic/test_utils.py

@@ -0,0 +1,79 @@
+from collections import deque
+
+from pydantic import BaseModel
+
+
+def assert_equal_pydantic_objects(
+    expected: BaseModel, actual: BaseModel, ignore_none=True
+):
+    """Compare 2 pydantic objects recursively and raise an AssertionError if they are not equal along with all
+    the differences by field. If `ignore_none` is set to True, expected None values will be ignored. This can be
+    useful when comparing objects that are partially filled.
+
+    Example:
+        >>> from pydantic import BaseModel
+        >>> class A(BaseModel):
+        ...     a: int
+        >>> class B(BaseModel):
+        ...     b: A
+        >>> a1 = A(a=1)
+        >>> a2 = A(a=2)
+        >>> b1 = B(b=a1)
+        >>> b2 = B(b=a2)
+        >>> assert_equal_pydantic_objects(a1, b1)
+        Traceback (most recent call last):
+        ```
+        AssertionError: objects mismatched on type at : expected: [A], actual: [B]
+        >>> assert_equal_pydantic_objects(a1, a2)
+        Traceback (most recent call last):
+        ```
+        AssertionError: objects mismatched on field: [a], expected: [1], actual: [2]
+        >>> assert_equal_pydantic_objects(b1, b2)
+        Traceback (most recent call last):
+        ```
+        AssertionError: objects mismatched on field: [b.a], expected: [1], actual: [2]
+
+    Args:
+        expected (BaseModel): The expected pydantic object.
+        actual (BaseModel): The actual pydantic object.
+        ignore_none (bool, optional): Whether to ignore None values. Defaults to True.
+
+    Raises:
+        AssertionError: If the objects are not equal. The error message will contain all the differences.
+    """
+    errors = []
+    queue = deque([(expected, actual, "")])
+    while queue:
+        expected, actual, current_key_prefix = queue.popleft()
+        if not isinstance(expected, actual.__class__):
+            errors.append(
+                f"objects mismatched on type at {current_key_prefix}: "
+                f"expected: [{type(expected).__name__}], actual: [{type(actual).__name__}]"
+            )
+            continue
+        if issubclass(expected.__class__, BaseModel) and isinstance(
+            expected.model_dump(), dict
+        ):
+            for key, expected_value in expected.model_dump().items():
+                if expected_value is None and ignore_none:
+                    continue
+                actual_value = actual.model_dump().get(key)
+                new_key_prefix = (
+                    f"{current_key_prefix}.{key}" if current_key_prefix else key
+                )
+                if issubclass(getattr(expected, key).__class__, BaseModel):
+                    queue.append(
+                        (getattr(expected, key), getattr(actual, key), new_key_prefix)
+                    )
+                elif expected_value != actual_value:
+                    errors.append(
+                        f"objects mismatched on field: [{new_key_prefix}], expected: [{expected_value}], actual: [{actual_value}]"
+                    )
+        else:
+            if expected != actual:
+                errors.append(
+                    f"mismatch at {current_key_prefix}: expected: [{expected}], actual: [{actual}]"
+                )
+
+    if errors:
+        raise AssertionError("\n".join(errors))

airflow_provider_openmetadata/lineage/callback.py

@@ -48,7 +48,7 @@ def failure_callback(context: Dict[str, str]) -> None:
         )
         pipeline: Pipeline = metadata.get_by_name(
             entity=Pipeline,
-            fqn=f"{airflow_service_entity.name.__root__}.{dag.dag_id}",
+            fqn=f"{airflow_service_entity.name.root}.{dag.dag_id}",
         )
 
         if pipeline:
@@ -60,7 +60,7 @@ def failure_callback(context: Dict[str, str]) -> None:
             )
         else:
             logging.warning(
-                f"Pipeline {airflow_service_entity.name.__root__}.{dag.dag_id} not found. Skipping status update."
+                f"Pipeline {airflow_service_entity.name.root}.{dag.dag_id} not found. Skipping status update."
             )
 
     except Exception as exc:  # pylint: disable=broad-except
@@ -90,7 +90,7 @@ def success_callback(context: Dict[str, str]) -> None:
         )
         pipeline: Pipeline = metadata.get_by_name(
             entity=Pipeline,
-            fqn=f"{airflow_service_entity.name.__root__}.{dag.dag_id}",
+            fqn=f"{airflow_service_entity.name.root}.{dag.dag_id}",
         )
 
         add_status(

airflow_provider_openmetadata/lineage/config/loader.py

@@ -90,7 +90,7 @@ def get_lineage_config() -> AirflowLineageConfig:
     if openmetadata_config_file:
         with open(openmetadata_config_file, encoding="utf-8") as config_file:
             config = json.load(config_file)
-            return AirflowLineageConfig.parse_obj(config)
+            return AirflowLineageConfig.model_validate(config)
 
     # If nothing is configured, raise
     raise ValueError("Missing lineage backend configuration")

airflow_provider_openmetadata/lineage/runner.py

@@ -200,7 +200,7 @@ class AirflowLineageRunner:
             return self.metadata.patch(
                 entity=Pipeline,
                 source=pipeline,
-                destination=pipeline.copy(update=pipeline_request.__dict__),
+                destination=pipeline.model_copy(update=pipeline_request.__dict__),
                 allowed_fields=ALLOWED_COMMON_PATCH_FIELDS,
                 restrict_update_fields=RESTRICT_UPDATE_LIST,
             )
@@ -282,7 +282,7 @@ class AirflowLineageRunner:
 
         for status in pipeline_status_list:
             self.metadata.add_pipeline_status(
-                fqn=pipeline.fullyQualifiedName.__root__, status=status
+                fqn=pipeline.fullyQualifiedName.root, status=status
             )
 
     def add_lineage(self, pipeline: Pipeline, xlets: XLets) -> None:
@@ -327,12 +327,12 @@ class AirflowLineageRunner:
                     else:
                         self.dag.log.warning(
                             f"Could not find [{to_xlet.entity.__name__}] [{to_xlet.fqn}] from "
-                            f"[{pipeline.fullyQualifiedName.__root__}] outlets"
+                            f"[{pipeline.fullyQualifiedName.root}] outlets"
                         )
             else:
                 self.dag.log.warning(
                     f"Could not find [{from_xlet.entity.__name__}] [{from_xlet.fqn}] from "
-                    f"[{pipeline.fullyQualifiedName.__root__}] inlets"
+                    f"[{pipeline.fullyQualifiedName.root}] inlets"
                 )
 
     def clean_lineage(self, pipeline: Pipeline, xlets: XLets):
@@ -343,7 +343,7 @@ class AirflowLineageRunner:
         """
         lineage_data = self.metadata.get_lineage_by_name(
             entity=Pipeline,
-            fqn=pipeline.fullyQualifiedName.__root__,
+            fqn=pipeline.fullyQualifiedName.root,
             up_depth=1,
             down_depth=1,
         )

airflow_provider_openmetadata/lineage/status.py

@@ -91,7 +91,7 @@ def add_status(
 
     task_status = []
     # We will append based on the current registered status
-    if pipeline_status and pipeline_status.timestamp.__root__ == execution_date:
+    if pipeline_status and pipeline_status.timestamp.root == execution_date:
         # If we are clearing a task, use the status of the new execution
         task_status = [
             task
@@ -123,5 +123,5 @@ def add_status(
 
     operator.log.info(f"Added status to DAG {updated_status}")
     metadata.add_pipeline_status(
-        fqn=pipeline.fullyQualifiedName.__root__, status=updated_status
+        fqn=pipeline.fullyQualifiedName.root, status=updated_status
     )

metadata/automations/runner.py

@@ -37,7 +37,7 @@ def execute(encrypted_automation_workflow: AutomationWorkflow) -> Any:
     )
 
     automation_workflow = metadata.get_by_name(
-        entity=AutomationWorkflow, fqn=encrypted_automation_workflow.name.__root__
+        entity=AutomationWorkflow, fqn=encrypted_automation_workflow.name.root
     )
 
     return run_workflow(automation_workflow.request, automation_workflow, metadata)

metadata/cli/app.py

@@ -19,7 +19,6 @@ 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()
 
@@ -34,7 +33,6 @@ def run_app(config_path: Path) -> 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())
@@ -42,5 +40,5 @@ def run_app(config_path: Path) -> None:
 
     workflow.execute()
     workflow.stop()
-    print_status(workflow)
+    workflow.print_status()
     workflow.raise_from_status()

metadata/cli/dataquality.py

@@ -17,13 +17,12 @@ import traceback
 from pathlib import Path
 
 from metadata.config.common import load_config_file
+from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
+    PipelineType,
+)
 from metadata.utils.logger import cli_logger
 from metadata.workflow.data_quality import TestSuiteWorkflow
-from metadata.workflow.workflow_output_handler import (
-    WorkflowType,
-    print_init_error,
-    print_status,
-)
+from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
 logger = cli_logger()
 
@@ -38,14 +37,15 @@ def run_test(config_path: Path) -> None:
     workflow_config_dict = None
     try:
         workflow_config_dict = load_config_file(config_path)
-        logger.debug(f"Using config: {workflow_config_dict}")
         workflow = TestSuiteWorkflow.create(workflow_config_dict)
     except Exception as exc:
         logger.debug(traceback.format_exc())
-        print_init_error(exc, workflow_config_dict, WorkflowType.TEST)
+        WorkflowInitErrorHandler.print_init_error(
+            exc, workflow_config_dict, PipelineType.TestSuite
+        )
         sys.exit(1)
 
     workflow.execute()
     workflow.stop()
-    print_status(workflow)
+    workflow.print_status()
     workflow.raise_from_status()