MINOR: Pydantic fixes for redshift & kafka (#16638)

ingestion/setup.py

@@ -328,6 +328,8 @@ test = {
     "minio==7.2.5",
     *plugins["mlflow"],
     *plugins["datalake-s3"],
+    *plugins["kafka"],
+    "kafka-python==2.0.2",
     *plugins["pii-processor"],
     "requests==2.31.0",
 }

ingestion/src/metadata/ingestion/source/messaging/common_broker_source.py

@@ -33,7 +33,7 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.generated.schema.type.basic import EntityName
+from metadata.generated.schema.type.basic import EntityName, FullyQualifiedEntityName
 from metadata.generated.schema.type.schema import SchemaType, Topic
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.models.ometa_topic_data import OMetaTopicSampleData

ingestion/src/metadata/ingestion/source/messaging/kafka/connection.py

@@ -95,7 +95,7 @@ def get_connection(
     schema_registry_client = None
     consumer_client = None
     if connection.schemaRegistryURL:
-        schema_registry_config["url"] = connection.schemaRegistryURL
+        schema_registry_config["url"] = str(connection.schemaRegistryURL)
         schema_registry_client = SchemaRegistryClient(schema_registry_config)
         consumer_config["bootstrap.servers"] = connection.bootstrapServers
         if "group.id" not in consumer_config:

ingestion/tests/integration/kafka/test_kafka.py

@@ -0,0 +1,109 @@
+#  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.
+
+"""
+Test Kafka using the topology
+"""
+
+
+from unittest import TestCase
+
+from confluent_kafka import Producer
+from testcontainers.kafka import KafkaContainer
+
+from metadata.generated.schema.entity.data.topic import Topic
+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
+from metadata.workflow.metadata import MetadataWorkflow
+
+
+def produce_and_consume_kafka_message(container):
+    topic = "test-topic"
+    bootstrap_server = container.get_bootstrap_server()
+
+    producer = Producer({"bootstrap.servers": bootstrap_server})
+    producer.produce(topic, b"verification message")
+    producer.flush()
+
+
+OM_JWT = "eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg"
+
+
+def get_ingestion_config(port: str):
+    return {
+        "source": {
+            "type": "kafka",
+            "serviceName": "TEST_KAFKA",
+            "serviceConnection": {
+                "config": {
+                    "type": "Kafka",
+                    "bootstrapServers": f"localhost:{port}",
+                }
+            },
+            "sourceConfig": {"config": {"type": "MessagingMetadata"}},
+        },
+        "sink": {"type": "metadata-rest", "config": {}},
+        "workflowConfig": {
+            "loggerLevel": "DEBUG",
+            "openMetadataServerConfig": {
+                "hostPort": "http://localhost:8585/api",
+                "authProvider": "openmetadata",
+                "securityConfig": {
+                    "jwtToken": "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
+
+
+class KafkaUnitTest(TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.metadata = int_admin_ometa()
+        cls.kafka_container = KafkaContainer()
+        cls.kafka_container.start()
+        cls.ingestion_config = get_ingestion_config(
+            cls.kafka_container.get_exposed_port(9093)
+        )
+        produce_and_consume_kafka_message(cls.kafka_container)
+
+        ingestion_workflow = MetadataWorkflow.create(
+            cls.ingestion_config,
+        )
+        ingestion_workflow.execute()
+        ingestion_workflow.raise_from_status()
+        ingestion_workflow.stop()
+
+    def test_topic(self):
+        topic = self.metadata.get_by_name(
+            entity=Topic,
+            fqn="TEST_KAFKA.test-topic",
+        )
+        self.assertIsNotNone(topic)

openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/common.py

@@ -202,7 +202,7 @@ def metadata_ingestion_workflow(workflow_config: OpenMetadataWorkflowConfig):
 
     set_operator_logger(workflow_config)
 
-    config = json.loads(workflow_config.model_dump_json())
+    config = json.loads(workflow_config.model_dump_json(exclude_defaults=False))
     workflow = MetadataWorkflow.create(config)
 
     workflow.execute()
@@ -313,7 +313,8 @@ def send_failed_status_callback(workflow_config: OpenMetadataWorkflowConfig, *_,
         )
 
         pipeline_status = metadata.get_pipeline_status(
-            workflow_config.ingestionPipelineFQN, str(workflow_config.pipelineRunId)
+            workflow_config.ingestionPipelineFQN,
+            str(workflow_config.pipelineRunId.root),
         )
         pipeline_status.endDate = Timestamp(int(datetime.now().timestamp() * 1000))
         pipeline_status.pipelineState = PipelineState.failed