separate config file

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

@@ -0,0 +1,58 @@
+from typing import Dict, Optional
+
+from pydantic import Field, root_validator
+
+from datahub.ingestion.source.sql.sql_common import SQLAlchemyStatefulIngestionConfig
+from datahub.ingestion.source_config.sql.snowflake import (
+    SnowflakeConfig,
+    SnowflakeProvisionRoleConfig,
+)
+
+
+class SnowflakeV2Config(SnowflakeConfig):
+    _convert_urns_to_lowercase: bool = Field(
+        default=True,
+        exclude=True,
+        description="Not supported",
+    )
+
+    check_role_grants: bool = Field(
+        default=False,
+        exclude=True,
+        description="Not supported",
+    )
+
+    provision_role: Optional[SnowflakeProvisionRoleConfig] = Field(
+        default=None, exclude=True, description="Not supported"
+    )
+
+    stateful_ingestion: Optional[SQLAlchemyStatefulIngestionConfig] = Field(
+        default=None, exclude=True, description="Not supported"
+    )
+
+    @root_validator(pre=False)
+    def validate_unsupported_configs(cls, values: Dict) -> Dict:
+        value = values.get("stateful_ingestion")
+        if value is not None and value.enabled:
+            raise ValueError(
+                "Stateful ingestion is currently not supported. Set `stateful_ingestion.enabled` to False"
+            )
+
+        value = values.get("provision_role")
+        if value is not None and value.enabled:
+            raise ValueError(
+                "Provision role is currently not supported. Set `provision_role.enabled` to False."
+            )
+
+        value = values.get("profiling")
+        if value is not None and value.enabled and not value.profile_table_level_only:
+            raise ValueError(
+                "Only table level profiling is supported. Set `profiling.profile_table_level_only` to True.",
+            )
+
+        value = values.get("check_role_grants")
+        if value is not None and value:
+            raise ValueError(
+                "Check role grants is not supported. Set `check_role_grants` to False.",
+            )
+        return values

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

@@ -3,7 +3,6 @@
 from typing import Dict, Iterable, List, Optional, Tuple, Union, cast
 
 from avrogen.dict_wrapper import DictWrapper
-from pydantic import Field, root_validator
 from snowflake.connector import SnowflakeConnection
 
 from datahub.emitter.mce_builder import (
@@ -29,6 +28,7 @@
     TestConnectionReport,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.snowflake.snowflake_config import SnowflakeV2Config
 from datahub.ingestion.source.snowflake.snowflake_lineage import (
     SnowflakeLineageExtractor,
 )
@@ -44,14 +44,7 @@
     SnowflakeView,
 )
 from datahub.ingestion.source.sql.snowflake import SnowflakeSource
-from datahub.ingestion.source.sql.sql_common import (
-    SQLAlchemySource,
-    SQLAlchemyStatefulIngestionConfig,
-)
-from datahub.ingestion.source_config.sql.snowflake import (
-    SnowflakeConfig,
-    SnowflakeProvisionRoleConfig,
-)
+from datahub.ingestion.source.sql.sql_common import SQLAlchemySource
 from datahub.ingestion.source_report.sql.snowflake import SnowflakeReport
 from datahub.metadata.com.linkedin.pegasus2avro.common import Status, SubTypes
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
@@ -120,7 +113,7 @@
 
 
 @platform_name("Snowflake")
-@config_class(SnowflakeConfig)
+@config_class(SnowflakeV2Config)
 @support_status(SupportStatus.INCUBATING)
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
@@ -133,57 +126,6 @@
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
 @capability(SourceCapability.LINEAGE_COARSE, "Optionally enabled via configuration")
 @capability(SourceCapability.DELETION_DETECTION, "Coming soon", supported=False)
-class SnowflakeV2Config(SnowflakeConfig):
-    _convert_urns_to_lowercase: bool = Field(
-        default=True,
-        exclude=True,
-        description="Not supported",
-    )
-
-    check_role_grants: bool = Field(
-        default=False,
-        exclude=True,
-        description="Not supported",
-    )
-
-    provision_role: Optional[SnowflakeProvisionRoleConfig] = Field(
-        default=None, exclude=True, description="Not supported"
-    )
-
-    stateful_ingestion: Optional[SQLAlchemyStatefulIngestionConfig] = Field(
-        default=None, exclude=True, description="Not supported"
-    )
-
-    @root_validator(pre=False)
-    def validate_unsupported_configs(cls, values: Dict) -> Dict:
-        value = values.get("stateful_ingestion")
-        if value is not None and value.enabled:
-            raise ValueError(
-                "Stateful ingestion is currently not supported. Set `stateful_ingestion.enabled` to False"
-            )
-
-        value = values.get("provision_role")
-        if value is not None and value.enabled:
-            raise ValueError(
-                "Provision role is currently not supported. Set `provision_role.enabled` to False."
-            )
-
-        value = values.get("profiling")
-        if value is not None and value.enabled and not value.profile_table_level_only:
-            raise ValueError(
-                "Only table level profiling is supported. Set `profiling.profile_table_level_only` to True.",
-            )
-
-        value = values.get("check_role_grants")
-        if value is not None and value:
-            raise ValueError(
-                "Check role grants is not supported. Set `check_role_grants` to False.",
-            )
-        return values
-
-
-@platform_name("Snowflake")
-@config_class(SnowflakeV2Config)
 class SnowflakeV2Source(TestableSource):
     def __init__(self, ctx: PipelineContext, config: SnowflakeV2Config):
         super().__init__(ctx)