acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.2.0.1rc1__py3-none-any.whl

datahub/ingestion/source/kafka_connect/sink_connectors.py

@@ -1,3 +1,4 @@
+import logging
 import re
 from dataclasses import dataclass
 from typing import Dict, Iterable, List, Optional, Tuple
@@ -9,6 +10,81 @@ from datahub.ingestion.source.kafka_connect.common import (
     KafkaConnectLineage,
 )
 
+logger = logging.getLogger(__name__)
+
+
+class RegexRouterTransform:
+    """Helper class to handle RegexRouter transformations for topic/table names."""
+
+    def __init__(self, config: Dict[str, str]) -> None:
+        self.transforms = self._parse_transforms(config)
+
+    def _parse_transforms(self, config: Dict[str, str]) -> List[Dict[str, str]]:
+        """Parse transforms configuration from connector config."""
+        transforms_list: List[Dict[str, str]] = []
+
+        # Get the transforms parameter
+        transforms_param: str = config.get("transforms", "")
+        if not transforms_param:
+            return transforms_list
+
+        # Parse individual transforms
+        transform_names: List[str] = [
+            name.strip() for name in transforms_param.split(",")
+        ]
+
+        for transform_name in transform_names:
+            if not transform_name:
+                continue
+            transform_config: Dict[str, str] = {}
+            transform_prefix: str = f"transforms.{transform_name}."
+
+            # Extract transform configuration
+            for key, value in config.items():
+                if key.startswith(transform_prefix):
+                    config_key: str = key[len(transform_prefix) :]
+                    transform_config[config_key] = value
+
+            # Only process RegexRouter transforms
+            if (
+                transform_config.get("type")
+                == "org.apache.kafka.connect.transforms.RegexRouter"
+            ):
+                transform_config["name"] = transform_name
+                transforms_list.append(transform_config)
+
+        return transforms_list
+
+    def apply_transforms(self, topic_name: str) -> str:
+        """Apply RegexRouter transforms to the topic name using Java regex."""
+        result: str = topic_name
+
+        for transform in self.transforms:
+            regex_pattern: Optional[str] = transform.get("regex")
+            replacement: str = transform.get("replacement", "")
+
+            if regex_pattern:
+                try:
+                    # Use Java Pattern and Matcher for exact Kafka Connect compatibility
+                    from java.util.regex import Pattern
+
+                    pattern = Pattern.compile(regex_pattern)
+                    matcher = pattern.matcher(result)
+
+                    if matcher.find():
+                        # Reset matcher to beginning for replaceFirst
+                        matcher.reset()
+                        result = matcher.replaceFirst(replacement)
+                        logger.debug(
+                            f"Applied transform {transform['name']}: {topic_name} -> {result}"
+                        )
+                except Exception as e:
+                    logger.warning(
+                        f"Invalid regex pattern in transform {transform['name']}: {e}"
+                    )
+
+        return str(result)
+
 
 @dataclass
 class ConfluentS3SinkConnector(BaseConnector):
@@ -18,28 +94,35 @@ class ConfluentS3SinkConnector(BaseConnector):
         bucket: str
         topics_dir: str
         topics: Iterable[str]
+        regex_router: RegexRouterTransform
 
     def _get_parser(self, connector_manifest: ConnectorManifest) -> S3SinkParser:
         # https://docs.confluent.io/kafka-connectors/s3-sink/current/configuration_options.html#s3
-        bucket = connector_manifest.config.get("s3.bucket.name")
+        bucket: Optional[str] = connector_manifest.config.get("s3.bucket.name")
         if not bucket:
             raise ValueError(
                 "Could not find 's3.bucket.name' in connector configuration"
             )
 
         # https://docs.confluent.io/kafka-connectors/s3-sink/current/configuration_options.html#storage
-        topics_dir = connector_manifest.config.get("topics.dir", "topics")
+        topics_dir: str = connector_manifest.config.get("topics.dir", "topics")
+
+        # Create RegexRouterTransform instance
+        regex_router: RegexRouterTransform = RegexRouterTransform(
+            connector_manifest.config
+        )
 
         return self.S3SinkParser(
             target_platform="s3",
             bucket=bucket,
             topics_dir=topics_dir,
             topics=connector_manifest.topic_names,
+            regex_router=regex_router,
         )
 
     def extract_flow_property_bag(self) -> Dict[str, str]:
         # Mask/Remove properties that may reveal credentials
-        flow_property_bag = {
+        flow_property_bag: Dict[str, str] = {
             k: v
             for k, v in self.connector_manifest.config.items()
             if k
@@ -54,11 +137,17 @@ class ConfluentS3SinkConnector(BaseConnector):
 
     def extract_lineages(self) -> List[KafkaConnectLineage]:
         try:
-            parser = self._get_parser(self.connector_manifest)
+            parser: ConfluentS3SinkConnector.S3SinkParser = self._get_parser(
+                self.connector_manifest
+            )
 
             lineages: List[KafkaConnectLineage] = list()
             for topic in parser.topics:
-                target_dataset = f"{parser.bucket}/{parser.topics_dir}/{topic}"
+                # Apply RegexRouter transformations using the RegexRouterTransform class
+                transformed_topic: str = parser.regex_router.apply_transforms(topic)
+                target_dataset: str = (
+                    f"{parser.bucket}/{parser.topics_dir}/{transformed_topic}"
+                )
 
                 lineages.append(
                     KafkaConnectLineage(
@@ -86,6 +175,7 @@ class SnowflakeSinkConnector(BaseConnector):
         database_name: str
         schema_name: str
         topics_to_tables: Dict[str, str]
+        regex_router: RegexRouterTransform
 
     def get_table_name_from_topic_name(self, topic_name: str) -> str:
         """
@@ -93,7 +183,7 @@ class SnowflakeSinkConnector(BaseConnector):
         Refer below link for more info
         https://docs.snowflake.com/en/user-guide/kafka-connector-overview#target-tables-for-kafka-topics
         """
-        table_name = re.sub("[^a-zA-Z0-9_]", "_", topic_name)
+        table_name: str = re.sub("[^a-zA-Z0-9_]", "_", topic_name)
         if re.match("^[^a-zA-Z_].*", table_name):
             table_name = "_" + table_name
         # Connector  may append original topic's hash code as suffix for conflict resolution
@@ -106,8 +196,13 @@ class SnowflakeSinkConnector(BaseConnector):
         self,
         connector_manifest: ConnectorManifest,
     ) -> SnowflakeParser:
-        database_name = connector_manifest.config["snowflake.database.name"]
-        schema_name = connector_manifest.config["snowflake.schema.name"]
+        database_name: str = connector_manifest.config["snowflake.database.name"]
+        schema_name: str = connector_manifest.config["snowflake.schema.name"]
+
+        # Create RegexRouterTransform instance
+        regex_router: RegexRouterTransform = RegexRouterTransform(
+            connector_manifest.config
+        )
 
         # Fetch user provided topic to table map
         provided_topics_to_tables: Dict[str, str] = {}
@@ -121,24 +216,30 @@ class SnowflakeSinkConnector(BaseConnector):
         topics_to_tables: Dict[str, str] = {}
         # Extract lineage for only those topics whose data ingestion started
         for topic in connector_manifest.topic_names:
+            # Apply transforms first to get the transformed topic name
+            transformed_topic: str = regex_router.apply_transforms(topic)
+
             if topic in provided_topics_to_tables:
                 # If user provided which table to get mapped with this topic
                 topics_to_tables[topic] = provided_topics_to_tables[topic]
             else:
-                # Else connector converts topic name to a valid Snowflake table name.
-                topics_to_tables[topic] = self.get_table_name_from_topic_name(topic)
+                # Use the transformed topic name to generate table name
+                topics_to_tables[topic] = self.get_table_name_from_topic_name(
+                    transformed_topic
+                )
 
         return self.SnowflakeParser(
             database_name=database_name,
             schema_name=schema_name,
             topics_to_tables=topics_to_tables,
+            regex_router=regex_router,
         )
 
     def extract_flow_property_bag(self) -> Dict[str, str]:
         # For all snowflake sink connector properties, refer below link
         # https://docs.snowflake.com/en/user-guide/kafka-connector-install#configuring-the-kafka-connector
         # remove private keys, secrets from properties
-        flow_property_bag = {
+        flow_property_bag: Dict[str, str] = {
             k: v
             for k, v in self.connector_manifest.config.items()
             if k
@@ -153,10 +254,12 @@ class SnowflakeSinkConnector(BaseConnector):
 
     def extract_lineages(self) -> List[KafkaConnectLineage]:
         lineages: List[KafkaConnectLineage] = list()
-        parser = self.get_parser(self.connector_manifest)
+        parser: SnowflakeSinkConnector.SnowflakeParser = self.get_parser(
+            self.connector_manifest
+        )
 
         for topic, table in parser.topics_to_tables.items():
-            target_dataset = f"{parser.database_name}.{parser.schema_name}.{table}"
+            target_dataset: str = f"{parser.database_name}.{parser.schema_name}.{table}"
             lineages.append(
                 KafkaConnectLineage(
                     source_dataset=topic,
@@ -176,7 +279,8 @@ class BigQuerySinkConnector(BaseConnector):
         project: str
         target_platform: str
         sanitizeTopics: bool
-        transforms: list
+        transforms: List[Dict[str, str]]
+        regex_router: RegexRouterTransform
         topicsToTables: Optional[str] = None
         datasets: Optional[str] = None
         defaultDataset: Optional[str] = None
@@ -186,16 +290,18 @@ class BigQuerySinkConnector(BaseConnector):
         self,
         connector_manifest: ConnectorManifest,
     ) -> BQParser:
-        project = connector_manifest.config["project"]
-        sanitizeTopics = connector_manifest.config.get("sanitizeTopics") or "false"
-        transform_names = (
+        project: str = connector_manifest.config["project"]
+        sanitizeTopics: str = connector_manifest.config.get("sanitizeTopics") or "false"
+
+        # Parse ALL transforms (original BigQuery logic)
+        transform_names: List[str] = (
             self.connector_manifest.config.get("transforms", "").split(",")
             if self.connector_manifest.config.get("transforms")
             else []
         )
-        transforms = []
+        transforms: List[Dict[str, str]] = []
         for name in transform_names:
-            transform = {"name": name}
+            transform: Dict[str, str] = {"name": name}
             transforms.append(transform)
             for key in self.connector_manifest.config:
                 if key.startswith(f"transforms.{name}."):
@@ -203,8 +309,13 @@ class BigQuerySinkConnector(BaseConnector):
                         self.connector_manifest.config[key]
                     )
 
+        # Create RegexRouterTransform instance for RegexRouter-specific handling
+        regex_router: RegexRouterTransform = RegexRouterTransform(
+            connector_manifest.config
+        )
+
         if "defaultDataset" in connector_manifest.config:
-            defaultDataset = connector_manifest.config["defaultDataset"]
+            defaultDataset: str = connector_manifest.config["defaultDataset"]
             return self.BQParser(
                 project=project,
                 defaultDataset=defaultDataset,
@@ -212,11 +323,14 @@ class BigQuerySinkConnector(BaseConnector):
                 sanitizeTopics=sanitizeTopics.lower() == "true",
                 version="v2",
                 transforms=transforms,
+                regex_router=regex_router,
             )
         else:
             # version 1.6.x and similar configs supported
-            datasets = connector_manifest.config["datasets"]
-            topicsToTables = connector_manifest.config.get("topicsToTables")
+            datasets: str = connector_manifest.config["datasets"]
+            topicsToTables: Optional[str] = connector_manifest.config.get(
+                "topicsToTables"
+            )
 
             return self.BQParser(
                 project=project,
@@ -225,10 +339,11 @@ class BigQuerySinkConnector(BaseConnector):
                 target_platform="bigquery",
                 sanitizeTopics=sanitizeTopics.lower() == "true",
                 transforms=transforms,
+                regex_router=regex_router,
             )
 
     def get_list(self, property: str) -> Iterable[Tuple[str, str]]:
-        entries = property.split(",")
+        entries: List[str] = property.split(",")
         for entry in entries:
             key, val = entry.rsplit("=")
             yield (key.strip(), val.strip())
@@ -243,7 +358,7 @@ class BigQuerySinkConnector(BaseConnector):
                 return dataset
         return None
 
-    def sanitize_table_name(self, table_name):
+    def sanitize_table_name(self, table_name: str) -> str:
         table_name = re.sub("[^a-zA-Z0-9_]", "_", table_name)
         if re.match("^[^a-zA-Z_].*", table_name):
             table_name = "_" + table_name
@@ -254,8 +369,8 @@ class BigQuerySinkConnector(BaseConnector):
         self, topic: str, parser: BQParser
     ) -> Optional[str]:
         if parser.version == "v2":
-            dataset = parser.defaultDataset
-            parts = topic.split(":")
+            dataset: Optional[str] = parser.defaultDataset
+            parts: List[str] = topic.split(":")
             if len(parts) == 2:
                 dataset = parts[0]
                 table = parts[1]
@@ -283,21 +398,9 @@ class BigQuerySinkConnector(BaseConnector):
             table = self.sanitize_table_name(table)
         return f"{dataset}.{table}"
 
-    def apply_transformations(
-        self, topic: str, transforms: List[Dict[str, str]]
-    ) -> str:
-        for transform in transforms:
-            if transform["type"] == "org.apache.kafka.connect.transforms.RegexRouter":
-                regex = transform["regex"]
-                replacement = transform["replacement"]
-                pattern = re.compile(regex)
-                if pattern.match(topic):
-                    topic = pattern.sub(replacement, topic, count=1)
-        return topic
-
     def extract_flow_property_bag(self) -> Dict[str, str]:
         # Mask/Remove properties that may reveal credentials
-        flow_property_bag = {
+        flow_property_bag: Dict[str, str] = {
             k: v
             for k, v in self.connector_manifest.config.items()
             if k not in ["keyfile"]
@@ -307,27 +410,33 @@ class BigQuerySinkConnector(BaseConnector):
 
     def extract_lineages(self) -> List[KafkaConnectLineage]:
         lineages: List[KafkaConnectLineage] = list()
-        parser = self.get_parser(self.connector_manifest)
+        parser: BigQuerySinkConnector.BQParser = self.get_parser(
+            self.connector_manifest
+        )
         if not parser:
             return lineages
-        target_platform = parser.target_platform
-        project = parser.project
-        transforms = parser.transforms
+        target_platform: str = parser.target_platform
+        project: str = parser.project
 
         for topic in self.connector_manifest.topic_names:
-            transformed_topic = self.apply_transformations(topic, transforms)
-            dataset_table = self.get_dataset_table_for_topic(transformed_topic, parser)
+            # Apply RegexRouter transformations using the RegexRouterTransform class
+            transformed_topic: str = parser.regex_router.apply_transforms(topic)
+
+            # Use the transformed topic to determine dataset/table
+            dataset_table: Optional[str] = self.get_dataset_table_for_topic(
+                transformed_topic, parser
+            )
             if dataset_table is None:
                 self.report.warning(
                     "Could not find target dataset for topic, please check your connector configuration"
                     f"{self.connector_manifest.name} : {transformed_topic} ",
                 )
                 continue
-            target_dataset = f"{project}.{dataset_table}"
+            target_dataset: str = f"{project}.{dataset_table}"
 
             lineages.append(
                 KafkaConnectLineage(
-                    source_dataset=transformed_topic,
+                    source_dataset=topic,  # Keep original topic as source
                     source_platform=KAFKA,
                     target_dataset=target_dataset,
                     target_platform=target_platform,

datahub/ingestion/source/kafka_connect/source_connectors.py

@@ -20,6 +20,8 @@ from datahub.ingestion.source.sql.sqlalchemy_uri_mapper import (
     get_platform_from_sqlalchemy_uri,
 )
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class ConfluentJDBCSourceConnector(BaseConnector):
@@ -392,7 +394,7 @@ class MongoSourceConnector(BaseConnector):
             db_connection_url=connector_manifest.config.get("connection.uri"),
             source_platform="mongodb",
             database_name=connector_manifest.config.get("database"),
-            topic_prefix=connector_manifest.config.get("topic_prefix"),
+            topic_prefix=connector_manifest.config.get("topic.prefix"),
             transforms=(
                 connector_manifest.config["transforms"].split(",")
                 if "transforms" in connector_manifest.config
@@ -406,7 +408,11 @@ class MongoSourceConnector(BaseConnector):
         lineages: List[KafkaConnectLineage] = list()
         parser = self.get_parser(self.connector_manifest)
         source_platform = parser.source_platform
-        topic_naming_pattern = r"mongodb\.(\w+)\.(\w+)"
+        topic_prefix = parser.topic_prefix or ""
+
+        # Escape topic_prefix to handle cases where it contains dots
+        # Some users configure topic.prefix like "my.mongodb" which breaks the regex
+        topic_naming_pattern = rf"{re.escape(topic_prefix)}\.(\w+)\.(\w+)"
 
         if not self.connector_manifest.topic_names:
             return lineages
@@ -429,6 +435,26 @@ class MongoSourceConnector(BaseConnector):
 
 @dataclass
 class DebeziumSourceConnector(BaseConnector):
+    # Debezium topic naming patterns by connector type
+    # - MySQL: {topic.prefix}.{database}.{table}
+    # - PostgreSQL: {topic.prefix}.{schema}.{table}
+    # - SQL Server: {topic.prefix}.{database}.{schema}.{table}
+    # - Oracle: {topic.prefix}.{schema}.{table}
+    # - DB2: {topic.prefix}.{schema}.{table}
+    # - MongoDB: {topic.prefix}.{database}.{collection}
+    # - Vitess: {topic.prefix}.{keyspace}.{table}
+
+    # Note SQL Server allows for "database.names" (multiple databases) config,
+    # and so database is in the topic naming pattern.
+    # However, others have "database.dbname" which is a single database name. For these connectors,
+    # additional databases would require a different connector instance
+
+    # Connectors with 2-level container in pattern (database + schema)
+    # Others have either database XOR schema, but not both
+    DEBEZIUM_CONNECTORS_WITH_2_LEVEL_CONTAINER_IN_PATTERN = {
+        "io.debezium.connector.sqlserver.SqlServerConnector",
+    }
+
     @dataclass
     class DebeziumParser:
         source_platform: str
@@ -514,16 +540,45 @@ class DebeziumSourceConnector(BaseConnector):
             source_platform = parser.source_platform
             server_name = parser.server_name
             database_name = parser.database_name
-            topic_naming_pattern = rf"({server_name})\.(\w+\.\w+)"
+            # 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 = (
+                connector_class
+                in self.DEBEZIUM_CONNECTORS_WITH_2_LEVEL_CONTAINER_IN_PATTERN
+            )
+
             for topic in self.connector_manifest.topic_names:
                 found = re.search(re.compile(topic_naming_pattern), topic)
+                logger.debug(
+                    f"Processing topic: '{topic}' with regex pattern '{topic_naming_pattern}', found: {found}"
+                )
 
                 if found:
-                    table_name = get_dataset_name(database_name, found.group(2))
+                    # Extract the table part after server_name
+                    table_part = found.group(2)
+
+                    if (
+                        maybe_duplicated_database_name
+                        and database_name
+                        and table_part.startswith(f"{database_name}.")
+                    ):
+                        table_part = table_part[len(database_name) + 1 :]
+
+                    logger.debug(
+                        f"Extracted table part: '{table_part}' from topic '{topic}'"
+                    )
+                    # Apply database name to create final dataset name
+                    table_name = get_dataset_name(database_name, table_part)
+                    logger.debug(f"Final table name: '{table_name}'")
 
                     lineage = KafkaConnectLineage(
                         source_dataset=table_name,

datahub/ingestion/source/looker/looker_source.py

@@ -126,6 +126,7 @@ logger = logging.getLogger(__name__)
     SourceCapability.USAGE_STATS,
     "Enabled by default, configured using `extract_usage_history`",
 )
+@capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
 class LookerDashboardSource(TestableSource, StatefulIngestionSourceBase):
     """
     This plugin extracts the following:

datahub/ingestion/source/mlflow.py

@@ -33,7 +33,10 @@ from datahub.ingestion.api.source import (
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.data_platforms import KNOWN_VALID_PLATFORM_NAMES
-from datahub.ingestion.source.common.subtypes import MLAssetSubTypes
+from datahub.ingestion.source.common.subtypes import (
+    MLAssetSubTypes,
+    SourceCapabilityModifier,
+)
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
     StaleEntityRemovalSourceReport,
@@ -138,6 +141,13 @@ class MLflowRegisteredModelStageInfo:
     SourceCapability.DESCRIPTIONS,
     "Extract descriptions for MLflow Registered Models and Model Versions",
 )
+@capability(
+    SourceCapability.CONTAINERS,
+    "Extract ML experiments",
+    subtype_modifier=[
+        SourceCapabilityModifier.MLFLOW_EXPERIMENT,
+    ],
+)
 @capability(SourceCapability.TAGS, "Extract tags for MLflow Registered Model Stages")
 class MLflowSource(StatefulIngestionSourceBase):
     platform = "mlflow"