fix(ingest): tighten Source.create type annotations

Supersedes #8620

metadata-ingestion/src/datahub/ingestion/api/source.py

@@ -23,7 +23,7 @@
 )
 
 from pydantic import BaseModel
-from typing_extensions import LiteralString
+from typing_extensions import LiteralString, Self
 
 from datahub.configuration.common import ConfigModel
 from datahub.configuration.source_common import PlatformInstanceConfigMixin
@@ -400,7 +400,7 @@ class Source(Closeable, metaclass=ABCMeta):
     ctx: PipelineContext
 
     @classmethod
-    def create(cls, config_dict: dict, ctx: PipelineContext) -> "Source":
+    def create(cls, config_dict: dict, ctx: PipelineContext) -> Self:
         # Technically, this method should be abstract. However, the @config_class
         # decorator automatically generates a create method at runtime if one is
         # not defined. Python still treats the class as abstract because it thinks

metadata-ingestion/src/datahub/ingestion/source/delta_lake/source.py

@@ -122,11 +122,6 @@ def __init__(self, config: DeltaLakeSourceConfig, ctx: PipelineContext):
             config_report,
         )
 
-    @classmethod
-    def create(cls, config_dict: dict, ctx: PipelineContext) -> "Source":
-        config = DeltaLakeSourceConfig.parse_obj(config_dict)
-        return cls(config, ctx)
-
     def _parse_datatype(self, raw_field_json_str: str) -> List[SchemaFieldClass]:
         raw_field_json = json.loads(raw_field_json_str)
 

metadata-ingestion/src/datahub/ingestion/source/demo_data.py

@@ -29,7 +29,7 @@ class DemoDataSource(Source):
 
     def __init__(self, ctx: PipelineContext, config: DemoDataConfig):
         file_config = FileSourceConfig(path=str(download_sample_data()))
-        self.file_source = GenericFileSource(ctx, file_config)
+        self.file_source: GenericFileSource = GenericFileSource(ctx, file_config)
 
     def get_workunits(self) -> Iterable[MetadataWorkUnit]:
         yield from self.file_source.get_workunits()

metadata-ingestion/src/datahub/ingestion/source/fivetran/fivetran.py

@@ -16,7 +16,7 @@
     platform_name,
     support_status,
 )
-from datahub.ingestion.api.source import MetadataWorkUnitProcessor, Source, SourceReport
+from datahub.ingestion.api.source import MetadataWorkUnitProcessor, SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.fivetran.config import (
     KNOWN_DATA_PLATFORM_MAPPING,
@@ -291,11 +291,6 @@ def _get_connector_workunits(
             dpi = self._generate_dpi_from_job(job, datajob)
             yield from self._get_dpi_workunits(job, dpi)
 
-    @classmethod
-    def create(cls, config_dict: dict, ctx: PipelineContext) -> Source:
-        config = FivetranSourceConfig.parse_obj(config_dict)
-        return cls(config, ctx)
-
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
         return [
             *super().get_workunit_processors(),

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

@@ -17,7 +17,7 @@
     platform_name,
     support_status,
 )
-from datahub.ingestion.api.source import MetadataWorkUnitProcessor, Source
+from datahub.ingestion.api.source import MetadataWorkUnitProcessor
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.kafka_connect.common import (
     CONNECTOR_CLASS,
@@ -94,11 +94,6 @@ def __init__(self, config: KafkaConnectSourceConfig, ctx: PipelineContext):
         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.

metadata-ingestion/src/datahub/ingestion/source/metabase.py

@@ -23,7 +23,7 @@
     platform_name,
     support_status,
 )
-from datahub.ingestion.api.source import MetadataWorkUnitProcessor, Source, SourceReport
+from datahub.ingestion.api.source import MetadataWorkUnitProcessor, SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
@@ -789,11 +789,6 @@ def get_datasource_from_id(
 
         return platform, dbname, schema, platform_instance
 
-    @classmethod
-    def create(cls, config_dict: dict, ctx: PipelineContext) -> Source:
-        config = MetabaseConfig.parse_obj(config_dict)
-        return cls(ctx, config)
-
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
         return [
             *super().get_workunit_processors(),

metadata-ingestion/src/datahub/ingestion/source/mlflow.py

@@ -333,8 +333,3 @@ def _get_global_tags_workunit(
             aspect=global_tags,
         )
         return wu
-
-    @classmethod
-    def create(cls, config_dict: dict, ctx: PipelineContext) -> Source:
-        config = MLflowConfig.parse_obj(config_dict)
-        return cls(ctx, config)

metadata-ingestion/src/datahub/ingestion/source/nifi.py

@@ -484,11 +484,6 @@ def __init__(self, config: NifiSourceConfig, ctx: PipelineContext) -> None:
     def rest_api_base_url(self):
         return self.config.site_url[: -len("nifi/")] + "nifi-api/"
 
-    @classmethod
-    def create(cls, config_dict: dict, ctx: PipelineContext) -> "Source":
-        config = NifiSourceConfig.parse_obj(config_dict)
-        return cls(config, ctx)
-
     def get_report(self) -> SourceReport:
         return self.report
 

metadata-ingestion/src/datahub/ingestion/source/redash.py

@@ -369,11 +369,6 @@ def validate_connection(self) -> None:
         else:
             raise ValueError(f"Failed to connect to {self.config.connect_uri}/api")
 
-    @classmethod
-    def create(cls, config_dict: dict, ctx: PipelineContext) -> Source:
-        config = RedashConfig.parse_obj(config_dict)
-        return cls(ctx, config)
-
     def _get_chart_data_source(self, data_source_id: Optional[int] = None) -> Dict:
         url = f"/api/data_sources/{data_source_id}"
         resp = self.client._get(url).json()

metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_v2.py

@@ -23,7 +23,6 @@
 from datahub.ingestion.api.source import (
     CapabilityReport,
     MetadataWorkUnitProcessor,
-    Source,
     SourceCapability,
     SourceReport,
     TestableSource,
@@ -251,11 +250,6 @@ def __init__(self, ctx: PipelineContext, config: SnowflakeV2Config):
 
         self.add_config_to_report()
 
-    @classmethod
-    def create(cls, config_dict: dict, ctx: PipelineContext) -> "Source":
-        config = SnowflakeV2Config.parse_obj(config_dict)
-        return cls(ctx, config)
-
     @staticmethod
     def test_connection(config_dict: dict) -> TestConnectionReport:
         test_report = TestConnectionReport()

metadata-ingestion/src/datahub/ingestion/source/superset.py

@@ -33,7 +33,7 @@
     platform_name,
     support_status,
 )
-from datahub.ingestion.api.source import MetadataWorkUnitProcessor, Source
+from datahub.ingestion.api.source import MetadataWorkUnitProcessor
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.sql.sql_types import resolve_sql_type
 from datahub.ingestion.source.sql.sqlalchemy_uri_mapper import (
@@ -265,11 +265,6 @@ def login(self) -> requests.Session:
             # TODO(Gabe): how should we message about this error?
         return requests_session
 
-    @classmethod
-    def create(cls, config_dict: dict, ctx: PipelineContext) -> Source:
-        config = SupersetConfig.parse_obj(config_dict)
-        return cls(ctx, config)
-
     def paginate_entity_api_results(self, entity_type, page_size=100):
         current_page = 0
         total_items = page_size

metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py

@@ -71,7 +71,6 @@
 from datahub.ingestion.api.source import (
     CapabilityReport,
     MetadataWorkUnitProcessor,
-    Source,
     StructuredLogLevel,
     TestableSource,
     TestConnectionReport,
@@ -804,11 +803,6 @@ def test_connection(config_dict: dict) -> TestConnectionReport:
     def get_report(self) -> TableauSourceReport:
         return self.report
 
-    @classmethod
-    def create(cls, config_dict: dict, ctx: PipelineContext) -> Source:
-        config = TableauConfig.parse_obj(config_dict)
-        return cls(config, ctx)
-
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
         return [
             *super().get_workunit_processors(),

metadata-ingestion/tests/unit/api/test_pipeline.py

@@ -4,6 +4,7 @@
 
 import pytest
 from freezegun import freeze_time
+from typing_extensions import Self
 
 from datahub.configuration.common import DynamicTypedConfig
 from datahub.ingestion.api.committable import CommitPolicy, Committable
@@ -440,7 +441,7 @@ def __init__(self, ctx: PipelineContext):
         ]
 
     @classmethod
-    def create(cls, config_dict: dict, ctx: PipelineContext) -> "Source":
+    def create(cls, config_dict: dict, ctx: PipelineContext) -> Self:
         assert not config_dict
         return cls(ctx)