acryl-datahub 0.15.0rc25__py3-none-any.whl → 0.15.0.1rc2__py3-none-any.whl

datahub/ingestion/source/kafka_connect/common.py

@@ -0,0 +1,202 @@
+import logging
+from dataclasses import dataclass, field
+from typing import Dict, Iterable, List, Optional
+
+from pydantic.fields import Field
+
+from datahub.configuration.common import AllowDenyPattern, ConfigModel
+from datahub.configuration.source_common import (
+    DatasetLineageProviderConfigBase,
+    PlatformInstanceConfigMixin,
+)
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+    StaleEntityRemovalSourceReport,
+    StatefulStaleMetadataRemovalConfig,
+)
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulIngestionConfigBase,
+)
+
+logger = logging.getLogger(__name__)
+
+KAFKA = "kafka"
+SOURCE = "source"
+SINK = "sink"
+CONNECTOR_CLASS = "connector.class"
+
+
+class ProvidedConfig(ConfigModel):
+    provider: str
+    path_key: str
+    value: str
+
+
+class GenericConnectorConfig(ConfigModel):
+    connector_name: str
+    source_dataset: str
+    source_platform: str
+
+
+class KafkaConnectSourceConfig(
+    PlatformInstanceConfigMixin,
+    DatasetLineageProviderConfigBase,
+    StatefulIngestionConfigBase,
+):
+    # See the Connect REST Interface for details
+    # https://docs.confluent.io/platform/current/connect/references/restapi.html#
+    connect_uri: str = Field(
+        default="http://localhost:8083/", description="URI to connect to."
+    )
+    username: Optional[str] = Field(default=None, description="Kafka Connect username.")
+    password: Optional[str] = Field(default=None, description="Kafka Connect password.")
+    cluster_name: Optional[str] = Field(
+        default="connect-cluster", description="Cluster to ingest from."
+    )
+    # convert lineage dataset's urns to lowercase
+    convert_lineage_urns_to_lowercase: bool = Field(
+        default=False,
+        description="Whether to convert the urns of ingested lineage dataset to lowercase",
+    )
+    connector_patterns: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="regex patterns for connectors to filter for ingestion.",
+    )
+    provided_configs: Optional[List[ProvidedConfig]] = Field(
+        default=None, description="Provided Configurations"
+    )
+    connect_to_platform_map: Optional[Dict[str, Dict[str, str]]] = Field(
+        default=None,
+        description='Platform instance mapping when multiple instances for a platform is available. Entry for a platform should be in either `platform_instance_map` or `connect_to_platform_map`. e.g.`connect_to_platform_map: { "postgres-connector-finance-db": "postgres": "core_finance_instance" }`',
+    )
+    platform_instance_map: Optional[Dict[str, str]] = Field(
+        default=None,
+        description='Platform instance mapping to use when constructing URNs. e.g.`platform_instance_map: { "hive": "warehouse" }`',
+    )
+    generic_connectors: List[GenericConnectorConfig] = Field(
+        default=[],
+        description="Provide lineage graph for sources connectors other than Confluent JDBC Source Connector, Debezium Source Connector, and Mongo Source Connector",
+    )
+
+    stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = None
+
+
+@dataclass
+class KafkaConnectSourceReport(StaleEntityRemovalSourceReport):
+    connectors_scanned: int = 0
+    filtered: List[str] = field(default_factory=list)
+
+    def report_connector_scanned(self, connector: str) -> None:
+        self.connectors_scanned += 1
+
+    def report_dropped(self, connector: str) -> None:
+        self.filtered.append(connector)
+
+
+@dataclass
+class KafkaConnectLineage:
+    """Class to store Kafka Connect lineage mapping, Each instance is potential DataJob"""
+
+    source_platform: str
+    target_dataset: str
+    target_platform: str
+    job_property_bag: Optional[Dict[str, str]] = None
+    source_dataset: Optional[str] = None
+
+
+@dataclass
+class ConnectorManifest:
+    """Each instance is potential DataFlow"""
+
+    name: str
+    type: str
+    config: Dict
+    tasks: Dict
+    url: Optional[str] = None
+    flow_property_bag: Optional[Dict[str, str]] = None
+    lineages: List[KafkaConnectLineage] = field(default_factory=list)
+    topic_names: Iterable[str] = field(default_factory=list)
+
+
+def remove_prefix(text: str, prefix: str) -> str:
+    if text.startswith(prefix):
+        index = len(prefix)
+        return text[index:]
+    return text
+
+
+def unquote(
+    string: str, leading_quote: str = '"', trailing_quote: Optional[str] = None
+) -> str:
+    """
+    If string starts and ends with a quote, unquote it
+    """
+    trailing_quote = trailing_quote if trailing_quote else leading_quote
+    if string.startswith(leading_quote) and string.endswith(trailing_quote):
+        string = string[1:-1]
+    return string
+
+
+def get_dataset_name(
+    database_name: Optional[str],
+    source_table: str,
+) -> str:
+    if database_name:
+        dataset_name = database_name + "." + source_table
+    else:
+        dataset_name = source_table
+
+    return dataset_name
+
+
+def get_platform_instance(
+    config: KafkaConnectSourceConfig, connector_name: str, platform: str
+) -> Optional[str]:
+    instance_name = None
+    if (
+        config.connect_to_platform_map
+        and config.connect_to_platform_map.get(connector_name)
+        and config.connect_to_platform_map[connector_name].get(platform)
+    ):
+        instance_name = config.connect_to_platform_map[connector_name][platform]
+        if config.platform_instance_map and config.platform_instance_map.get(platform):
+            logger.warning(
+                f"Same source platform {platform} configured in both platform_instance_map and connect_to_platform_map."
+                "Will prefer connector specific platform instance from connect_to_platform_map."
+            )
+    elif config.platform_instance_map and config.platform_instance_map.get(platform):
+        instance_name = config.platform_instance_map[platform]
+    logger.info(
+        f"Instance name assigned is: {instance_name} for Connector Name {connector_name} and platform {platform}"
+    )
+    return instance_name
+
+
+def transform_connector_config(
+    connector_config: Dict, provided_configs: List[ProvidedConfig]
+) -> None:
+    """This method will update provided configs in connector config values, if any"""
+    lookupsByProvider = {}
+    for pconfig in provided_configs:
+        lookupsByProvider[f"${{{pconfig.provider}:{pconfig.path_key}}}"] = pconfig.value
+    for k, v in connector_config.items():
+        for key, value in lookupsByProvider.items():
+            if key in v:
+                connector_config[k] = connector_config[k].replace(key, value)
+
+
+# TODO: Find a more automated way to discover new platforms with 3 level naming hierarchy.
+def has_three_level_hierarchy(platform: str) -> bool:
+    return platform in ["postgres", "trino", "redshift", "snowflake"]
+
+
+@dataclass
+class BaseConnector:
+    connector_manifest: ConnectorManifest
+    config: KafkaConnectSourceConfig
+    report: KafkaConnectSourceReport
+
+    def extract_lineages(self) -> List[KafkaConnectLineage]:
+        return []
+
+    def extract_flow_property_bag(self) -> Optional[Dict[str, str]]:
+        return None

datahub/ingestion/source/kafka_connect/kafka_connect.py

@@ -0,0 +1,367 @@
+import logging
+from typing import Iterable, List, Optional, Type
+
+import jpype
+import jpype.imports
+import requests
+
+import datahub.emitter.mce_builder as builder
+import datahub.metadata.schema_classes as models
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.api.decorators import (
+    SourceCapability,
+    SupportStatus,
+    capability,
+    config_class,
+    platform_name,
+    support_status,
+)
+from datahub.ingestion.api.source import MetadataWorkUnitProcessor, Source
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.kafka_connect.common import (
+    CONNECTOR_CLASS,
+    SINK,
+    SOURCE,
+    BaseConnector,
+    ConnectorManifest,
+    KafkaConnectLineage,
+    KafkaConnectSourceConfig,
+    KafkaConnectSourceReport,
+    get_platform_instance,
+    transform_connector_config,
+)
+from datahub.ingestion.source.kafka_connect.sink_connectors import (
+    BIGQUERY_SINK_CONNECTOR_CLASS,
+    S3_SINK_CONNECTOR_CLASS,
+    SNOWFLAKE_SINK_CONNECTOR_CLASS,
+    BigQuerySinkConnector,
+    ConfluentS3SinkConnector,
+    SnowflakeSinkConnector,
+)
+from datahub.ingestion.source.kafka_connect.source_connectors import (
+    DEBEZIUM_SOURCE_CONNECTOR_PREFIX,
+    JDBC_SOURCE_CONNECTOR_CLASS,
+    MONGO_SOURCE_CONNECTOR_CLASS,
+    ConfigDrivenSourceConnector,
+    ConfluentJDBCSourceConnector,
+    DebeziumSourceConnector,
+    MongoSourceConnector,
+)
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+    StaleEntityRemovalHandler,
+)
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulIngestionSourceBase,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@platform_name("Kafka Connect")
+@config_class(KafkaConnectSourceConfig)
+@support_status(SupportStatus.CERTIFIED)
+@capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
+@capability(SourceCapability.SCHEMA_METADATA, "Enabled by default")
+@capability(SourceCapability.LINEAGE_COARSE, "Enabled by default")
+class KafkaConnectSource(StatefulIngestionSourceBase):
+    config: KafkaConnectSourceConfig
+    report: KafkaConnectSourceReport
+    platform: str = "kafka-connect"
+
+    def __init__(self, config: KafkaConnectSourceConfig, ctx: PipelineContext):
+        super().__init__(config, ctx)
+        self.config = config
+        self.report = KafkaConnectSourceReport()
+        self.session = requests.Session()
+        self.session.headers.update(
+            {
+                "Accept": "application/json",
+                "Content-Type": "application/json",
+            }
+        )
+
+        # Test the connection
+        if self.config.username is not None and self.config.password is not None:
+            logger.info(
+                f"Connecting to {self.config.connect_uri} with Authentication..."
+            )
+            self.session.auth = (self.config.username, self.config.password)
+
+        test_response = self.session.get(f"{self.config.connect_uri}/connectors")
+        test_response.raise_for_status()
+        logger.info(f"Connection to {self.config.connect_uri} is ok")
+        if not jpype.isJVMStarted():
+            jpype.startJVM()
+
+    @classmethod
+    def create(cls, config_dict: dict, ctx: PipelineContext) -> Source:
+        config = KafkaConnectSourceConfig.parse_obj(config_dict)
+        return cls(config, ctx)
+
+    def get_connectors_manifest(self) -> Iterable[ConnectorManifest]:
+        """Get Kafka Connect connectors manifest using REST API.
+        Enrich with lineages metadata.
+        """
+
+        connector_response = self.session.get(
+            f"{self.config.connect_uri}/connectors",
+        )
+
+        payload = connector_response.json()
+
+        for connector_name in payload:
+            connector_url = f"{self.config.connect_uri}/connectors/{connector_name}"
+            connector_manifest = self._get_connector_manifest(
+                connector_name, connector_url
+            )
+            if (
+                connector_manifest is None
+                or not self.config.connector_patterns.allowed(connector_manifest.name)
+            ):
+                self.report.report_dropped(connector_name)
+                continue
+
+            if self.config.provided_configs:
+                transform_connector_config(
+                    connector_manifest.config, self.config.provided_configs
+                )
+            connector_manifest.url = connector_url
+            connector_manifest.topic_names = self._get_connector_topics(connector_name)
+            connector_class_value = connector_manifest.config.get(CONNECTOR_CLASS) or ""
+
+            class_type: Type[BaseConnector] = BaseConnector
+
+            # Populate Source Connector metadata
+            if connector_manifest.type == SOURCE:
+                connector_manifest.tasks = self._get_connector_tasks(connector_name)
+
+                # JDBC source connector lineages
+                if connector_class_value == JDBC_SOURCE_CONNECTOR_CLASS:
+                    class_type = ConfluentJDBCSourceConnector
+                elif connector_class_value.startswith(DEBEZIUM_SOURCE_CONNECTOR_PREFIX):
+                    class_type = DebeziumSourceConnector
+                elif connector_class_value == MONGO_SOURCE_CONNECTOR_CLASS:
+                    class_type = MongoSourceConnector
+                elif any(
+                    [
+                        connector.connector_name == connector_manifest.name
+                        for connector in self.config.generic_connectors
+                    ]
+                ):
+                    class_type = ConfigDrivenSourceConnector
+                else:
+                    self.report.report_dropped(connector_manifest.name)
+                    self.report.warning(
+                        "Lineage for Source Connector not supported. "
+                        "Please refer to Kafka Connect docs to use `generic_connectors` config.",
+                        context=f"{connector_manifest.name} of type {connector_class_value}",
+                    )
+                    continue
+            elif connector_manifest.type == SINK:
+                if connector_class_value == BIGQUERY_SINK_CONNECTOR_CLASS:
+                    class_type = BigQuerySinkConnector
+                elif connector_class_value == S3_SINK_CONNECTOR_CLASS:
+                    class_type = ConfluentS3SinkConnector
+                elif connector_class_value == SNOWFLAKE_SINK_CONNECTOR_CLASS:
+                    class_type = SnowflakeSinkConnector
+                else:
+                    self.report.report_dropped(connector_manifest.name)
+                    self.report.warning(
+                        "Lineage for Sink Connector not supported.",
+                        context=f"{connector_manifest.name} of type {connector_class_value}",
+                    )
+
+            connector_class = class_type(connector_manifest, self.config, self.report)
+            connector_manifest.lineages = connector_class.extract_lineages()
+            connector_manifest.flow_property_bag = (
+                connector_class.extract_flow_property_bag()
+            )
+
+            yield connector_manifest
+
+    def _get_connector_manifest(
+        self, connector_name: str, connector_url: str
+    ) -> Optional[ConnectorManifest]:
+        try:
+            connector_response = self.session.get(connector_url)
+            connector_response.raise_for_status()
+        except Exception as e:
+            self.report.warning(
+                "Failed to get connector details", connector_name, exc=e
+            )
+            return None
+        manifest = connector_response.json()
+        connector_manifest = ConnectorManifest(**manifest)
+        return connector_manifest
+
+    def _get_connector_tasks(self, connector_name: str) -> dict:
+        try:
+            response = self.session.get(
+                f"{self.config.connect_uri}/connectors/{connector_name}/tasks",
+            )
+            response.raise_for_status()
+        except Exception as e:
+            self.report.warning(
+                "Error getting connector tasks", context=connector_name, exc=e
+            )
+            return {}
+
+        return response.json()
+
+    def _get_connector_topics(self, connector_name: str) -> List[str]:
+        try:
+            response = self.session.get(
+                f"{self.config.connect_uri}/connectors/{connector_name}/topics",
+            )
+            response.raise_for_status()
+        except Exception as e:
+            self.report.warning(
+                "Error getting connector topics", context=connector_name, exc=e
+            )
+            return []
+
+        return response.json()[connector_name]["topics"]
+
+    def construct_flow_workunit(self, connector: ConnectorManifest) -> MetadataWorkUnit:
+        connector_name = connector.name
+        connector_type = connector.type
+        connector_class = connector.config.get(CONNECTOR_CLASS)
+        flow_property_bag = connector.flow_property_bag
+        # connector_url = connector.url  # NOTE: this will expose connector credential when used
+        flow_urn = builder.make_data_flow_urn(
+            self.platform,
+            connector_name,
+            self.config.env,
+            self.config.platform_instance,
+        )
+
+        return MetadataChangeProposalWrapper(
+            entityUrn=flow_urn,
+            aspect=models.DataFlowInfoClass(
+                name=connector_name,
+                description=f"{connector_type.capitalize()} connector using `{connector_class}` plugin.",
+                customProperties=flow_property_bag,
+                # externalUrl=connector_url, # NOTE: this will expose connector credential when used
+            ),
+        ).as_workunit()
+
+    def construct_job_workunits(
+        self, connector: ConnectorManifest
+    ) -> Iterable[MetadataWorkUnit]:
+        connector_name = connector.name
+        flow_urn = builder.make_data_flow_urn(
+            self.platform,
+            connector_name,
+            self.config.env,
+            self.config.platform_instance,
+        )
+
+        lineages = connector.lineages
+        if lineages:
+            for lineage in lineages:
+                source_dataset = lineage.source_dataset
+                source_platform = lineage.source_platform
+                target_dataset = lineage.target_dataset
+                target_platform = lineage.target_platform
+                job_property_bag = lineage.job_property_bag
+
+                source_platform_instance = get_platform_instance(
+                    self.config, connector_name, source_platform
+                )
+                target_platform_instance = get_platform_instance(
+                    self.config, connector_name, target_platform
+                )
+
+                job_id = self.get_job_id(lineage, connector, self.config)
+                job_urn = builder.make_data_job_urn_with_flow(flow_urn, job_id)
+
+                inlets = (
+                    [
+                        self.make_lineage_dataset_urn(
+                            source_platform, source_dataset, source_platform_instance
+                        )
+                    ]
+                    if source_dataset
+                    else []
+                )
+                outlets = [
+                    self.make_lineage_dataset_urn(
+                        target_platform, target_dataset, target_platform_instance
+                    )
+                ]
+
+                yield MetadataChangeProposalWrapper(
+                    entityUrn=job_urn,
+                    aspect=models.DataJobInfoClass(
+                        name=f"{connector_name}:{job_id}",
+                        type="COMMAND",
+                        customProperties=job_property_bag,
+                    ),
+                ).as_workunit()
+
+                yield MetadataChangeProposalWrapper(
+                    entityUrn=job_urn,
+                    aspect=models.DataJobInputOutputClass(
+                        inputDatasets=inlets,
+                        outputDatasets=outlets,
+                    ),
+                ).as_workunit()
+
+    def get_job_id(
+        self,
+        lineage: KafkaConnectLineage,
+        connector: ConnectorManifest,
+        config: KafkaConnectSourceConfig,
+    ) -> str:
+        connector_class = connector.config.get(CONNECTOR_CLASS)
+
+        # Note - This block is only to maintain backward compatibility of Job URN
+        if (
+            connector_class
+            and connector.type == SOURCE
+            and (
+                "JdbcSourceConnector" in connector_class
+                or connector_class.startswith("io.debezium.connector")
+            )
+            and lineage.source_dataset
+            and config.connect_to_platform_map
+            and config.connect_to_platform_map.get(connector.name)
+            and config.connect_to_platform_map[connector.name].get(
+                lineage.source_platform
+            )
+        ):
+            return f"{config.connect_to_platform_map[connector.name][lineage.source_platform]}.{lineage.source_dataset}"
+
+        return (
+            lineage.source_dataset
+            if lineage.source_dataset
+            else f"unknown_source.{lineage.target_dataset}"
+        )
+
+    def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
+        return [
+            *super().get_workunit_processors(),
+            StaleEntityRemovalHandler.create(
+                self, self.config, self.ctx
+            ).workunit_processor,
+        ]
+
+    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+        for connector in self.get_connectors_manifest():
+            yield self.construct_flow_workunit(connector)
+            yield from self.construct_job_workunits(connector)
+            self.report.report_connector_scanned(connector.name)
+
+    def get_report(self) -> KafkaConnectSourceReport:
+        return self.report
+
+    def make_lineage_dataset_urn(
+        self, platform: str, name: str, platform_instance: Optional[str]
+    ) -> str:
+        if self.config.convert_lineage_urns_to_lowercase:
+            name = name.lower()
+
+        return builder.make_dataset_urn_with_platform_instance(
+            platform, name, platform_instance, self.config.env
+        )