fix imports, separate snowflake constants

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

@@ -0,0 +1,18 @@
+from enum import Enum
+
+
+class SnowflakeCloudProvider(str, Enum):
+    AWS = "aws"
+    GCP = "gcp"
+    AZURE = "azure"
+
+
+SNOWFLAKE_DEFAULT_CLOUD_REGION_ID = "us-west-2"
+SNOWFLAKE_DEFAULT_CLOUD = SnowflakeCloudProvider.AWS
+
+
+class SnowflakeEdition(str, Enum):
+    STANDARD = "Standard"
+
+    # We use this to represent Enterprise Edition or higher
+    ENTERPRISE = "Enterprise or above"

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

@@ -11,6 +11,7 @@
 import datahub.emitter.mce_builder as builder
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.aws.s3_util import make_s3_urn
+from datahub.ingestion.source.snowflake.constants import SnowflakeEdition
 from datahub.ingestion.source.snowflake.snowflake_config import SnowflakeV2Config
 from datahub.ingestion.source.snowflake.snowflake_query import SnowflakeQuery
 from datahub.ingestion.source.snowflake.snowflake_report import SnowflakeV2Report
@@ -19,7 +20,6 @@
 )
 from datahub.ingestion.source.snowflake.snowflake_utils import (
     SnowflakeCommonMixin,
-    SnowflakeEdition,
     SnowflakePermissionError,
     SnowflakeQueryMixin,
 )
@@ -207,7 +207,7 @@ def _get_upstream_lineage_info(
         if self._lineage_map is None or self._external_lineage_map is None:
             conn = self.config.get_connection()
         if self._lineage_map is None:
-            if self.report.edition == SnowflakeEdition.STANDARD.value:
+            if self.report.edition == SnowflakeEdition.STANDARD:
                 logger.info(
                     "Snowflake Account is Standard Edition. Table to Table Lineage Feature is not supported."
                 )
@@ -354,7 +354,7 @@ def _populate_view_lineage(self, conn: SnowflakeConnection) -> None:
         with PerfTimer() as timer:
             self._populate_view_upstream_lineage(conn)
             self.report.view_upstream_lineage_query_secs = timer.elapsed_seconds()
-        if self.report.edition == SnowflakeEdition.STANDARD.value:
+        if self.report.edition == SnowflakeEdition.STANDARD:
             logger.info(
                 "Snowflake Account is Standard Edition. View to Table Lineage Feature is not supported."
             )
@@ -390,7 +390,7 @@ def _populate_external_lineage(self, conn: SnowflakeConnection) -> None:
                     f"ExternalLineage[Table(Down)={key}]:External(Up)={self._external_lineage_map[key]} via access_history"
                 )
         except SnowflakePermissionError:
-            if self.report.edition == SnowflakeEdition.STANDARD.value:
+            if self.report.edition == SnowflakeEdition.STANDARD:
                 logger.info(
                     "Snowflake Account is Standard Edition. External Lineage Feature is not supported."
                 )

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

@@ -1,5 +1,6 @@
 from typing import Optional
 
+from datahub.ingestion.source.snowflake.constants import SnowflakeEdition
 from datahub.ingestion.source.sql.sql_generic_profiler import ProfilingSqlReport
 from datahub.ingestion.source_report.sql.snowflake import SnowflakeReport
 from datahub.ingestion.source_report.usage.snowflake_usage import SnowflakeUsageReport

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

@@ -13,12 +13,12 @@
 )
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.snowflake.constants import SnowflakeEdition
 from datahub.ingestion.source.snowflake.snowflake_config import SnowflakeV2Config
 from datahub.ingestion.source.snowflake.snowflake_query import SnowflakeQuery
 from datahub.ingestion.source.snowflake.snowflake_report import SnowflakeV2Report
 from datahub.ingestion.source.snowflake.snowflake_utils import (
     SnowflakeCommonMixin,
-    SnowflakeEdition,
     SnowflakePermissionError,
     SnowflakeQueryMixin,
 )

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

@@ -1,5 +1,4 @@
 import logging
-from enum import Enum
 from functools import lru_cache
 from typing import Any, Optional
 
@@ -11,6 +10,11 @@
 from datahub.configuration.pattern_utils import is_schema_allowed
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.snowflake.constants import (
+    SNOWFLAKE_DEFAULT_CLOUD,
+    SNOWFLAKE_DEFAULT_CLOUD_REGION_ID,
+    SnowflakeCloudProvider,
+)
 from datahub.ingestion.source.snowflake.snowflake_config import SnowflakeV2Config
 from datahub.ingestion.source.snowflake.snowflake_report import SnowflakeV2Report
 from datahub.metadata.com.linkedin.pegasus2avro.events.metadata import ChangeType
@@ -23,23 +27,6 @@ class SnowflakePermissionError(MetaError):
     """A permission error has happened"""
 
 
-class SnowflakeEdition(str, Enum):
-    STANDARD = "Standard"
-
-    # We use this to represent Enterprise Edition or higher
-    ENTERPRISE = "Enterprise or above"
-
-
-class SnowflakeCloudProvider(str, Enum):
-    AWS = "aws"
-    GCP = "gcp"
-    AZURE = "azure"
-
-
-SNOWFLAKE_DEFAULT_CLOUD_REGION_ID = "us-west-2"
-SNOWFLAKE_DEFAULT_CLOUD = SnowflakeCloudProvider.AWS
-
-
 # Required only for mypy, since we are using mixin classes, and not inheritance.
 # Reference - https://mypy.readthedocs.io/en/latest/more_types.html#mixin-classes
 class SnowflakeLoggingProtocol(Protocol):

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

@@ -43,6 +43,7 @@
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.glossary.classification_mixin import ClassificationMixin
+from datahub.ingestion.source.snowflake.constants import SnowflakeEdition
 from datahub.ingestion.source.snowflake.snowflake_config import SnowflakeV2Config
 from datahub.ingestion.source.snowflake.snowflake_lineage import (
     SnowflakeLineageExtractor,
@@ -65,7 +66,6 @@
 )
 from datahub.ingestion.source.snowflake.snowflake_utils import (
     SnowflakeCommonMixin,
-    SnowflakeEdition,
     SnowflakePermissionError,
     SnowflakeQueryMixin,
 )
@@ -874,8 +874,7 @@ def get_schema_metadata(
                     else {},
                 )
             except Exception as e:
-                self.warn(
-                    self.logger,
+                self.report_warning(
                     dataset_name,
                     f"unable to classify table columns due to error -> {e}",
                 )
@@ -1185,9 +1184,9 @@ def inspect_session_metadata(self, conn: SnowflakeConnection) -> None:
         try:
             logger.info("Checking current edition")
             if self.is_standard_edition(conn):
-                self.report.edition = SnowflakeEdition.STANDARD.value
+                self.report.edition = SnowflakeEdition.STANDARD
             else:
-                self.report.edition = SnowflakeEdition.ENTERPRISE.value
+                self.report.edition = SnowflakeEdition.ENTERPRISE
         except Exception:
             self.report.edition = None