Add a couple of fixes to Debezium source

treff7es · treff7es · commit 2ede18b78cc0 · 2025-11-10T10:14:02.000+01:00
diff --git a/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/kafka_connect.py b/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/kafka_connect.py
@@ -1,6 +1,5 @@
 import base64
 import logging
-from functools import lru_cache
 from typing import Dict, Iterable, List, Optional
 
 import jpype
@@ -114,6 +113,9 @@ def __init__(self, config: KafkaConnectSourceConfig, ctx: PipelineContext):
             self.session, self.config, self.report
         )
 
+        # Cache for all Kafka topics (single ingestion run)
+        self._all_kafka_topics_cache: Optional[List[str]] = None
+
         if not jpype.isJVMStarted():
             jpype.startJVM()
 
@@ -445,18 +447,23 @@ def _get_topics_confluent_cloud_from_manifest(
             # Final fallback to config-based approach
             return self._get_topics_from_connector_config(connector_manifest)
 
-    @lru_cache(maxsize=1)
     def _get_all_topics_from_kafka_api(self) -> List[str]:
         """
         Get all topics from Confluent Cloud Kafka REST API v3.
 
         This provides the comprehensive topic list needed for the reverse transform
         pipeline strategy to work effectively.
 
+        Uses instance variable caching to avoid repeated API calls during a single ingestion run.
+
         Returns:
             List of all topic names from the Kafka cluster.
             Empty list if API is not accessible or fails.
         """
+        # Return cached result if available
+        if self._all_kafka_topics_cache is not None:
+            return self._all_kafka_topics_cache
+
         try:
             # Extract cluster information from Connect URI
             kafka_rest_endpoint, cluster_id = self._parse_confluent_cloud_info()
@@ -507,13 +514,19 @@ def _get_all_topics_from_kafka_api(self) -> List[str]:
                 logger.info(
                     f"Retrieved {len(all_topics)} topics from Confluent Cloud Kafka REST API v3"
                 )
+                # Cache the result for subsequent calls
+                self._all_kafka_topics_cache = all_topics
                 return all_topics
             else:
                 logger.warning("Unexpected response format from Kafka REST API")
+                # Cache empty result to avoid repeated failures
+                self._all_kafka_topics_cache = []
                 return []
 
         except Exception as e:
             logger.debug(f"Failed to get topics from Kafka REST API: {e}")
+            # Cache empty result to avoid repeated API calls on failure
+            self._all_kafka_topics_cache = []
             return []
 
     def _parse_confluent_cloud_info(self) -> tuple[Optional[str], Optional[str]]:
diff --git a/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/source_connectors.py b/metadata-ingestion/src/datahub/ingestion/source/kafka_connect/source_connectors.py
@@ -1692,15 +1692,26 @@ def extract_lineages(self) -> List[KafkaConnectLineage]:
             source_platform = parser.source_platform
             server_name = parser.server_name
             database_name = parser.database_name
+
+            if not self.connector_manifest.topic_names:
+                return lineages
+
+            # Check for EventRouter transform - requires special handling
+            if self._has_event_router_transform():
+                logger.debug(
+                    f"Connector {self.connector_manifest.name} uses EventRouter transform - using table-based lineage extraction"
+                )
+                return self._extract_lineages_for_event_router(
+                    source_platform, database_name
+                )
+
+            # Standard Debezium topic processing
             # Escape server_name to handle cases where topic.prefix contains dots
             # Some users configure topic.prefix like "my.server" which breaks the regex
             server_name = server_name or ""
             # Regex pattern (\w+\.\w+(?:\.\w+)?) supports BOTH 2-part and 3-part table names
             topic_naming_pattern = rf"({re.escape(server_name)})\.(\w+\.\w+(?:\.\w+)?)"
 
-            if not self.connector_manifest.topic_names:
-                return lineages
-
             # Handle connectors with 2-level container (database + schema) in topic pattern
             connector_class = self.connector_manifest.config.get(CONNECTOR_CLASS, "")
             maybe_duplicated_database_name = (
@@ -1749,6 +1760,137 @@ def extract_lineages(self) -> List[KafkaConnectLineage]:
 
         return []
 
+    def _has_event_router_transform(self) -> bool:
+        """Check if connector uses Debezium EventRouter transform."""
+        transforms_config = self.connector_manifest.config.get("transforms", "")
+        if not transforms_config:
+            return False
+
+        transform_names = parse_comma_separated_list(transforms_config)
+        for name in transform_names:
+            transform_type = self.connector_manifest.config.get(
+                f"transforms.{name}.type", ""
+            )
+            if transform_type == "io.debezium.transforms.outbox.EventRouter":
+                return True
+
+        return False
+
+    def _extract_lineages_for_event_router(
+        self, source_platform: str, database_name: Optional[str]
+    ) -> List[KafkaConnectLineage]:
+        """
+        Extract lineages for connectors using EventRouter transform.
+
+        EventRouter is a data-dependent transform that reads fields from row data
+        to determine output topics. We cannot predict output topics from configuration alone,
+        so we extract source tables from table.include.list and try to match them to
+        actual topics using RegexRouter patterns.
+
+        Reference: https://debezium.io/documentation/reference/transformations/outbox-event-router.html
+        """
+        lineages: List[KafkaConnectLineage] = []
+
+        # Extract source tables from configuration
+        table_config = self.connector_manifest.config.get(
+            "table.include.list"
+        ) or self.connector_manifest.config.get("table.whitelist")
+
+        if not table_config:
+            logger.warning(
+                f"EventRouter connector {self.connector_manifest.name} has no table.include.list config"
+            )
+            return lineages
+
+        table_names = parse_comma_separated_list(table_config)
+
+        # Try to filter topics using RegexRouter replacement pattern (if available)
+        filtered_topics = self._filter_topics_for_event_router()
+
+        # For each source table, create lineages to filtered topics
+        for table_name in table_names:
+            # Clean quoted table names
+            clean_table = table_name.strip('"')
+
+            # Apply database name if present
+            if database_name:
+                source_dataset = get_dataset_name(database_name, clean_table)
+            else:
+                source_dataset = clean_table
+
+            # Create lineages from this source table to filtered topics
+            for topic in filtered_topics:
+                lineage = KafkaConnectLineage(
+                    source_dataset=source_dataset,
+                    source_platform=source_platform,
+                    target_dataset=topic,
+                    target_platform=KAFKA,
+                )
+                lineages.append(lineage)
+
+        logger.info(
+            f"Created {len(lineages)} EventRouter lineages from {len(table_names)} source tables "
+            f"to {len(filtered_topics)} topics for connector {self.connector_manifest.name}"
+        )
+
+        return lineages
+
+    def _filter_topics_for_event_router(self) -> List[str]:
+        """
+        Filter topics for EventRouter connectors using RegexRouter replacement pattern.
+
+        EventRouter often works with RegexRouter to rename output topics. We can use
+        the RegexRouter replacement pattern to identify which topics belong to this connector.
+        """
+        # Look for RegexRouter transform configuration
+        transforms_config = self.connector_manifest.config.get("transforms", "")
+        if not transforms_config:
+            return list(self.connector_manifest.topic_names)
+
+        transform_names = parse_comma_separated_list(transforms_config)
+
+        # Find RegexRouter configuration
+        regex_replacement = None
+        for name in transform_names:
+            transform_type = self.connector_manifest.config.get(
+                f"transforms.{name}.type", ""
+            )
+            if transform_type in [
+                "org.apache.kafka.connect.transforms.RegexRouter",
+                "io.confluent.connect.cloud.transforms.TopicRegexRouter",
+            ]:
+                # Extract the replacement pattern
+                # Example: "dev.ern.cashout.$1" -> we want topics starting with "dev.ern.cashout."
+                replacement = self.connector_manifest.config.get(
+                    f"transforms.{name}.replacement", ""
+                )
+                if replacement:
+                    # Extract prefix from replacement pattern (before first $)
+                    # "dev.ern.cashout.$1" -> "dev.ern.cashout."
+                    if "$" in replacement:
+                        regex_replacement = replacement.split("$")[0]
+                    else:
+                        regex_replacement = replacement
+                    break
+
+        # Filter topics using the replacement prefix
+        if regex_replacement:
+            filtered_topics = [
+                topic
+                for topic in self.connector_manifest.topic_names
+                if topic.startswith(regex_replacement)
+            ]
+            logger.debug(
+                f"Filtered EventRouter topics to {len(filtered_topics)} topics matching prefix '{regex_replacement}'"
+            )
+            return filtered_topics
+
+        # No RegexRouter found - use all topics (risky but best effort)
+        logger.warning(
+            f"EventRouter connector {self.connector_manifest.name} has no RegexRouter - cannot filter topics accurately"
+        )
+        return list(self.connector_manifest.topic_names)
+
 
 @dataclass
 class ConfigDrivenSourceConnector(BaseConnector):
diff --git a/metadata-ingestion/tests/integration/kafka-connect-confluent-cloud/confluent_cloud_mock_api.py b/metadata-ingestion/tests/integration/kafka-connect-confluent-cloud/confluent_cloud_mock_api.py
@@ -2,6 +2,7 @@
 import logging
 import os
 from functools import wraps
+from typing import Any, Dict
 
 from flask import Flask, jsonify, request
 
@@ -149,9 +150,64 @@ def log_request_info():
         "id": None,
         "extensions": {},
     },
+    "outbox-source-connector": {
+        "status": None,
+        "info": {
+            "name": "outbox-source-connector",
+            "type": "source",
+            "config": {
+                "after.state.only": "false",
+                "cloud.environment": "prod",
+                "cloud.provider": "aws",
+                "connector.class": "MySqlCdcSource",
+                "database.connectionTimeZone": "UTC",
+                "database.dbname": "outbox_db",
+                "database.hostname": "mysql-primary.us-west-2.rds.amazonaws.com",
+                "database.password": "****************",
+                "database.port": "3306",
+                "database.server.name": "outbox-db-server",
+                "database.ssl.mode": "required",
+                "database.sslmode": "require",
+                "database.user": "outbox_connector",
+                "hstore.handling.mode": "json",
+                "interval.handling.mode": "numeric",
+                "kafka.auth.mode": "SERVICE_ACCOUNT",
+                "kafka.endpoint": "SASL_SSL://pkc-abc123.us-west-2.aws.confluent.cloud:9092",
+                "kafka.region": "us-west-2",
+                "kafka.service.account.id": "sa-123456",
+                "max.batch.size": "2000",
+                "name": "outbox-source-connector",
+                "output.data.format": "JSON",
+                "output.key.format": "JSON",
+                "poll.interval.ms": "500",
+                "provide.transaction.metadata": "true",
+                "snapshot.locking.mode": "none",
+                "snapshot.mode": "schema_only",
+                "table.include.list": "outbox_db.outbox",
+                "tasks.max": "1",
+                "tombstones.on.delete": "false",
+                "transforms": "EventRouter,RegexRouter",
+                "transforms.EventRouter.route.by.field": "_route_suffix",
+                "transforms.EventRouter.table.expand.json.payload": "true",
+                "transforms.EventRouter.table.field.event.id": "event_id",
+                "transforms.EventRouter.table.field.event.key": "aggregate_id",
+                "transforms.EventRouter.table.field.event.payload": "payload",
+                "transforms.EventRouter.table.field.event.type": "event_type",
+                "transforms.EventRouter.table.fields.additional.placement": "event_timestamp:header,event_version:header,aggregate_type:header",
+                "transforms.EventRouter.type": "io.debezium.transforms.outbox.EventRouter",
+                "transforms.RegexRouter.regex": "outbox.event.(.*)",
+                "transforms.RegexRouter.replacement": "app.events.$1",
+                "transforms.RegexRouter.type": "io.confluent.connect.cloud.transforms.TopicRegexRouter",
+            },
+            "tasks": [{"connector": "outbox-source-connector", "task": 0}],
+        },
+        "id": None,
+        "extensions": {},
+    },
 }
 
-TOPICS_DATA = {
+TOPICS_DATA: Dict[str, Any] = {
+    "kind": "KafkaTopicList",
     "data": [
         {
             "topic_name": "public.customer_profiles",
@@ -201,7 +257,31 @@ def log_request_info():
             "is_internal": False,
             "cluster_id": "4k0R9d1GTS5tI9f4Y2xZ0Q",
         },
-    ]
+        {
+            "topic_name": "app.events.order_created",
+            "partitions": {
+                "related": "https://api.confluent.cloud/kafka/v3/clusters/4k0R9d1GTS5tI9f4Y2xZ0Q/topics/app.events.order_created/partitions"
+            },
+            "configs": {
+                "related": "https://api.confluent.cloud/kafka/v3/clusters/4k0R9d1GTS5tI9f4Y2xZ0Q/topics/app.events.order_created/configs"
+            },
+            "replication_factor": 3,
+            "is_internal": False,
+            "cluster_id": "4k0R9d1GTS5tI9f4Y2xZ0Q",
+        },
+        {
+            "topic_name": "app.events.order_updated",
+            "partitions": {
+                "related": "https://api.confluent.cloud/kafka/v3/clusters/4k0R9d1GTS5tI9f4Y2xZ0Q/topics/app.events.order_updated/partitions"
+            },
+            "configs": {
+                "related": "https://api.confluent.cloud/kafka/v3/clusters/4k0R9d1GTS5tI9f4Y2xZ0Q/topics/app.events.order_updated/configs"
+            },
+            "replication_factor": 3,
+            "is_internal": False,
+            "cluster_id": "4k0R9d1GTS5tI9f4Y2xZ0Q",
+        },
+    ],
 }
 
 
diff --git a/metadata-ingestion/tests/integration/kafka-connect-confluent-cloud/kafka_connect_confluent_cloud_mces_golden.json b/metadata-ingestion/tests/integration/kafka-connect-confluent-cloud/kafka_connect_confluent_cloud_mces_golden.json
