openmetadata-ingestion 1.5.0.0rc1__py3-none-any.whl → 1.5.0.0rc2__py3-none-any.whl

_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/postgres/conftest.py

@@ -6,6 +6,7 @@ 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
@@ -67,4 +68,8 @@ def postgres_container(tmp_path_factory):
             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

@@ -43,7 +43,6 @@ def assert_equal_pydantic_objects(
     """
     errors = []
     queue = deque([(expected, actual, "")])
-
     while queue:
         expected, actual, current_key_prefix = queue.popleft()
         if not isinstance(expected, actual.__class__):
@@ -52,11 +51,13 @@ def assert_equal_pydantic_objects(
                 f"expected: [{type(expected).__name__}], actual: [{type(actual).__name__}]"
             )
             continue
-        if issubclass(expected.__class__, BaseModel):
-            for key, expected_value in expected.dict().items():
+        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.dict().get(key)
+                actual_value = actual.model_dump().get(key)
                 new_key_prefix = (
                     f"{current_key_prefix}.{key}" if current_key_prefix else key
                 )

metadata/cli/app.py

@@ -33,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())

metadata/cli/dataquality.py

@@ -37,7 +37,6 @@ 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())

metadata/cli/ingest.py

@@ -38,7 +38,6 @@ def run_ingest(config_path: Path) -> None:
     try:
         config_dict = load_config_file(config_path)
         workflow = MetadataWorkflow.create(config_dict)
-        logger.debug(f"Using config: {workflow.config}")
     except Exception as exc:
         logger.debug(traceback.format_exc())
         WorkflowInitErrorHandler.print_init_error(

metadata/cli/profile.py

@@ -37,7 +37,6 @@ def run_profiler(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 = ProfilerWorkflow.create(workflow_config_dict)
     except Exception as exc:
         logger.debug(traceback.format_exc())

metadata/cli/usage.py

@@ -38,7 +38,6 @@ def run_usage(config_path: Path) -> None:
     try:
         config_dict = load_config_file(config_path)
         workflow = UsageWorkflow.create(config_dict)
-        logger.debug(f"Using config: {workflow.config}")
     except Exception as exc:
         logger.debug(traceback.format_exc())
         WorkflowInitErrorHandler.print_init_error(exc, config_dict, PipelineType.usage)

metadata/clients/domo_client.py

@@ -126,7 +126,7 @@ class DomoClient:
 
             if isinstance(response, list) and len(response) > 0:
                 return DomoChartDetails(
-                    id=response[0]["id"],
+                    id=str(response[0]["id"]),
                     name=response[0]["title"],
                     metadata=DomoChartMetadataDetails(
                         chartType=response[0].get("metadata", {}).get("chartType", "")

metadata/data_quality/validations/column/sqlalchemy/columnValueLengthsToBeBetween.py

@@ -12,8 +12,7 @@
 """
 Validator for column value length to be between test case
 """
-
-
+import math
 from typing import Optional
 
 from sqlalchemy import Column, inspect
@@ -65,14 +64,16 @@ class ColumnValueLengthsToBeBetweenValidator(
             NotImplementedError:
         """
         row_count = self._compute_row_count(self.runner, column)
+        filters = []
+        if min_bound > -math.inf:
+            filters.append((LenFn(column), "lt", min_bound))
+        if max_bound < math.inf:
+            filters.append((LenFn(column), "gt", max_bound))
         failed_rows = self._compute_row_count_between(
             self.runner,
             column,
             {
-                "filters": [
-                    (LenFn(column), "gt", max_bound),
-                    (LenFn(column), "lt", min_bound),
-                ],
+                "filters": filters,
                 "or_filter": True,
             },
         )

metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeBetween.py

@@ -12,7 +12,7 @@
 """
 Validator for column values to be between test case
 """
-
+import math
 from typing import Optional
 
 from sqlalchemy import Column, inspect
@@ -63,11 +63,16 @@ class ColumnValuesToBeBetweenValidator(
             NotImplementedError:
         """
         row_count = self._compute_row_count(self.runner, column)
+        filters = []
+        if not isinstance(min_bound, (int, float)) or min_bound > -math.inf:
+            filters.append((column, "lt", min_bound))
+        if not isinstance(min_bound, (int, float)) or max_bound < math.inf:
+            filters.append((column, "gt", max_bound))
         failed_rows = self._compute_row_count_between(
             self.runner,
             column,
             {
-                "filters": [(column, "gt", max_bound), (column, "lt", min_bound)],
+                "filters": filters,
                 "or_filter": True,
             },
         )

metadata/data_quality/validations/table/sqlalchemy/tableDiff.py

@@ -9,6 +9,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 # pylint: disable=missing-module-docstring
+import logging
 import traceback
 from itertools import islice
 from typing import Dict, Iterable, List, Optional, Tuple
@@ -59,6 +60,21 @@ class UnsupportedDialectError(Exception):
         super().__init__(f"Unsupported dialect in param {param}: {dialect}")
 
 
+def masked(s: str, mask: bool = True) -> str:
+    """Mask a string if masked is True otherwise return the string.
+    Only for development purposes, do not use in production.
+    Change it False if you want to see the data in the logs.
+
+    Args:
+        s: string to mask
+        mask: boolean to mask the string
+
+    Returns:
+        masked string if mask is True otherwise return the string
+    """
+    return "***" if mask else s
+
+
 class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
     """
     Compare two tables and fail if the number of differences exceeds a threshold
@@ -110,12 +126,14 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
             stats = table_diff_iter.get_stats_dict()
             if stats["total"] > 0:
                 logger.debug("Sample of failed rows:")
-                for s in islice(self.get_table_diff(), 10):
-                    # since the data can contiant sensitive information, we don't want to log it
-                    # we can uncomment this line if we must see the data in the logs
-                    # logger.debug(s)
-                    # by default we will log the data masked
-                    logger.debug([s[0], ["*" for _ in s[1]]])
+                # depending on the data, this require scanning a lot of data
+                # so we only log the sample in debug mode. data can be sensitive
+                # so it is masked by default
+                for s in islice(
+                    self.safe_table_diff_iterator(),
+                    10 if logger.level <= logging.DEBUG else 0,
+                ):
+                    logger.debug("%s", str([s[0]] + [masked(st) for st in s[1]]))
             test_case_result = self.get_row_diff_test_case_result(
                 threshold,
                 stats["total"],
@@ -222,12 +240,10 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
             "where": self.get_where(),
         }
         logger.debug(
-            "Calling table diff with parameters:"  # pylint: disable=consider-using-f-string
-            " table1={}, table2={}, kwargs={}".format(
-                table1.table_path,
-                table2.table_path,
-                ",".join(f"{k}={v}" for k, v in data_diff_kwargs.items()),
-            )
+            "Calling table diff with parameters: table1=%s, table2=%s, kwargs=%s",
+            table1.table_path,
+            table2.table_path,
+            ",".join(f"{k}={v}" for k, v in data_diff_kwargs.items()),
         )
         return data_diff.diff_tables(table1, table2, **data_diff_kwargs)  # type: ignore
 
@@ -400,3 +416,19 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
             if len(key_set) > limit:
                 len(key_set)
         return len(key_set)
+
+    def safe_table_diff_iterator(self) -> DiffResultWrapper:
+        """A safe iterator object which properly closes the diff object when the generator is exhausted.
+        Otherwise the data_diff library will continue to hold the connection open and eventually
+        raise a KeyError.
+        """
+        gen = self.get_table_diff()
+        try:
+            yield from gen
+        finally:
+            try:
+                gen.diff.close()
+            except KeyError as ex:
+                if str(ex) == "2":
+                    # This is a known issue in data_diff where the diff object is closed
+                    pass