Refactor platform extraction

Author: treff7es

metadata-ingestion/src/datahub/ingestion/source/kafka_connect/common.py

@@ -653,9 +653,8 @@ def supports_connector_class(connector_class: str) -> bool:
         """Check if this connector handles the given class. Override in subclasses."""
         return False
 
-    @staticmethod
-    def get_platform(connector_class: str) -> str:
-        """Get the platform for this connector type. Override in subclasses."""
+    def get_platform(self) -> str:
+        """Get the platform for this connector instance. Override in subclasses."""
         return "unknown"
 
     def _extract_fine_grained_lineage(

metadata-ingestion/src/datahub/ingestion/source/kafka_connect/connector_registry.py

@@ -43,15 +43,15 @@ class ConnectorRegistry:
     def create_schema_resolver(
         ctx: Optional["PipelineContext"],
         config: KafkaConnectSourceConfig,
-        connector_manifest: ConnectorManifest,
+        connector: BaseConnector,
     ) -> Optional["SchemaResolver"]:
         """
         Create SchemaResolver for enhanced lineage extraction if enabled.
 
         Args:
             ctx: Pipeline context (contains graph connection)
             config: Kafka Connect source configuration
-            connector_manifest: Connector manifest to determine platform
+            connector: Connector instance to get platform from
 
         Returns:
             SchemaResolver instance if feature is enabled and graph is available, None otherwise
@@ -62,19 +62,16 @@ def create_schema_resolver(
         try:
             from datahub.sql_parsing.schema_resolver import SchemaResolver
 
-            # Determine platform from connector class
-            connector_class = connector_manifest.config.get("connector.class", "")
-            platform = ConnectorRegistry._infer_platform_from_connector(
-                connector_class, connector_manifest
-            )
+            # Get platform from connector instance (single source of truth)
+            platform = connector.get_platform()
 
             # Get platform instance if configured
             platform_instance = get_platform_instance(
-                config, connector_manifest.name, platform
+                config, connector.connector_manifest.name, platform
             )
 
             logger.info(
-                f"Creating SchemaResolver for connector {connector_manifest.name} "
+                f"Creating SchemaResolver for connector {connector.connector_manifest.name} "
                 f"with platform={platform}, instance={platform_instance}"
             )
 
@@ -86,46 +83,11 @@ def create_schema_resolver(
             )
         except Exception as e:
             logger.warning(
-                f"Failed to create SchemaResolver for connector {connector_manifest.name}: {e}. "
+                f"Failed to create SchemaResolver for connector {connector.connector_manifest.name}: {e}. "
                 "Falling back to standard lineage extraction."
             )
             return None
 
-    @staticmethod
-    def _infer_platform_from_connector(
-        connector_class: str, manifest: ConnectorManifest
-    ) -> str:
-        """Infer the source platform from connector class."""
-        from datahub.ingestion.source.kafka_connect.source_connectors import (
-            DebeziumSourceConnector,
-        )
-
-        # Try to get platform from Debezium source connector
-        if (
-            "debezium" in connector_class.lower()
-            or connector_class in CLOUD_JDBC_SOURCE_CLASSES
-        ):
-            return DebeziumSourceConnector._get_platform_from_connector_class(
-                connector_class
-            )
-
-        # Default platform inference
-        connector_lower = connector_class.lower()
-        if "postgres" in connector_lower:
-            return "postgres"
-        elif "mysql" in connector_lower:
-            return "mysql"
-        elif "sqlserver" in connector_lower or "mssql" in connector_lower:
-            return "mssql"
-        elif "oracle" in connector_lower:
-            return "oracle"
-        elif "mongodb" in connector_lower or "mongo" in connector_lower:
-            return "mongodb"
-        elif "snowflake" in connector_lower:
-            return "snowflake"
-        else:
-            return "unknown"
-
     @staticmethod
     def get_connector_for_manifest(
         manifest: ConnectorManifest,
@@ -147,12 +109,7 @@ def get_connector_for_manifest(
         """
         connector_class_value = manifest.config.get("connector.class", "")
 
-        # Create schema resolver if enabled
-        schema_resolver = ConnectorRegistry.create_schema_resolver(
-            ctx, config, manifest
-        )
-
-        # Determine connector type based on manifest type
+        # Create connector instance first
         if manifest.type == SOURCE:
             connector = ConnectorRegistry._get_source_connector(
                 connector_class_value, manifest, config, report
@@ -164,9 +121,13 @@ def get_connector_for_manifest(
         else:
             return None
 
-        # Attach schema resolver to connector if created
-        if connector and schema_resolver:
-            connector.schema_resolver = schema_resolver
+        # Create and attach schema resolver using connector's platform
+        if connector:
+            schema_resolver = ConnectorRegistry.create_schema_resolver(
+                ctx, config, connector
+            )
+            if schema_resolver:
+                connector.schema_resolver = schema_resolver
 
         return connector
 
@@ -347,7 +308,6 @@ def supports_connector_class(connector_class: str) -> bool:
         """Generic connector supports any unknown class."""
         return True
 
-    @staticmethod
-    def get_platform(connector_class: str) -> str:
+    def get_platform(self) -> str:
         """Generic connectors have configurable platforms."""
         return "unknown"

metadata-ingestion/src/datahub/ingestion/source/kafka_connect/source_connectors.py

@@ -1502,6 +1502,23 @@ def _find_topics_by_advanced_prefix_patterns(
 
         return matching_topics
 
+    def get_platform(self) -> str:
+        """
+        Get platform for JDBC connector.
+
+        JDBC connectors can connect to multiple databases, so platform is inferred from
+        the connection URL in the connector configuration.
+        """
+        try:
+            parser = self.get_parser(self.connector_manifest)
+            return parser.source_platform
+        except Exception:
+            # If parser fails, try to infer from JDBC URL directly
+            jdbc_url = self.connector_manifest.config.get("connection.url", "")
+            if jdbc_url:
+                return self._extract_platform_from_jdbc_url(jdbc_url)
+            return "unknown"
+
 
 @dataclass
 class SnowflakeSourceConnector(BaseConnector):
@@ -1931,8 +1948,7 @@ def supports_connector_class(connector_class: str) -> bool:
         """Check if this connector handles Snowflake Source."""
         return connector_class == SNOWFLAKE_SOURCE_CLOUD
 
-    @staticmethod
-    def get_platform(connector_class: str) -> str:
+    def get_platform(self) -> str:
         """Get the platform for Snowflake Source connector."""
         return "snowflake"
 
@@ -1998,6 +2014,10 @@ def extract_lineages(self) -> List[KafkaConnectLineage]:
                 lineages.append(lineage)
         return lineages
 
+    def get_platform(self) -> str:
+        """Get the platform for Mongo Source connector."""
+        return "mongodb"
+
 
 @dataclass
 class DebeziumSourceConnector(BaseConnector):
@@ -2036,10 +2056,8 @@ def get_parser(
         self,
         connector_manifest: ConnectorManifest,
     ) -> DebeziumParser:
-        connector_class = connector_manifest.config.get(CONNECTOR_CLASS, "")
-
         # Map connector class to platform
-        platform = self._get_platform_from_connector_class(connector_class)
+        platform = self.get_platform()
 
         # Map handler platform to parser platform (handler uses "sqlserver", parser expects "mssql")
         parser_platform = "mssql" if platform == "sqlserver" else platform
@@ -2108,9 +2126,9 @@ def _get_database_name_for_platform(
             # postgres, oracle, db2 use database.dbname
             return config.get("database.dbname")
 
-    @staticmethod
-    def _get_platform_from_connector_class(connector_class: str) -> str:
+    def get_platform(self) -> str:
         """Map Debezium connector class to platform name."""
+        connector_class = self.connector_manifest.config.get(CONNECTOR_CLASS, "")
         # Map based on well-known Debezium connector classes
         if "mysql" in connector_class.lower():
             return "mysql"