fix(ingest/snowflake): explicit set schema if public schema is absent (…

…datahub-project#9526)

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

@@ -24,6 +24,8 @@
 
 logger = logging.getLogger(__name__)
 
+PUBLIC_SCHEMA = "PUBLIC"
+
 
 class SnowflakeProfiler(GenericProfiler, SnowflakeCommonMixin):
     def __init__(
@@ -36,6 +38,7 @@ def __init__(
         self.config: SnowflakeV2Config = config
         self.report: SnowflakeV2Report = report
         self.logger = logger
+        self.database_default_schema: Dict[str, str] = dict()
 
     def get_workunits(
         self, database: SnowflakeDatabase, db_tables: Dict[str, List[SnowflakeTable]]
@@ -47,6 +50,10 @@ def get_workunits(
                 "max_overflow", self.config.profiling.max_workers
             )
 
+        if PUBLIC_SCHEMA not in db_tables:
+            # If PUBLIC schema is absent, we use any one of schemas as default schema
+            self.database_default_schema[database.name] = list(db_tables.keys())[0]
+
         profile_requests = []
         for schema in database.schemas:
             for table in db_tables[schema.name]:
@@ -136,9 +143,16 @@ def get_profiler_instance(
         )
 
     def callable_for_db_connection(self, db_name: str) -> Callable:
+        schema_name = self.database_default_schema.get(db_name)
+
         def get_db_connection():
             conn = self.config.get_connection()
             conn.cursor().execute(SnowflakeQuery.use_database(db_name))
+
+            # As mentioned here - https://docs.snowflake.com/en/sql-reference/sql/use-database#usage-notes
+            # no schema is selected if PUBLIC schema is absent. We need to explicitly call `USE SCHEMA <schema>`
+            if schema_name:
+                conn.cursor().execute(SnowflakeQuery.use_schema(schema_name))
             return conn
 
         return get_db_connection

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

@@ -80,6 +80,10 @@ def show_tags() -> str:
     def use_database(db_name: str) -> str:
         return f'use database "{db_name}"'
 
+    @staticmethod
+    def use_schema(schema_name: str) -> str:
+        return f'use schema "{schema_name}"'
+
     @staticmethod
     def get_databases(db_name: Optional[str]) -> str:
         db_clause = f'"{db_name}".' if db_name is not None else ""

metadata-ingestion/src/datahub/ingestion/source/state/stateful_ingestion_base.py

@@ -98,7 +98,7 @@ class StatefulIngestionConfigBase(GenericModel, Generic[CustomConfig]):
     )
 
 
-class StatefulLineageConfigMixin:
+class StatefulLineageConfigMixin(ConfigModel):
     enable_stateful_lineage_ingestion: bool = Field(
         default=True,
         description="Enable stateful lineage ingestion."