include/exclude at generation time rather than fetch time

datahub-project · Jul 15, 2024 · 70b1b11 · 70b1b11
1 parent df2e04c
commit 70b1b11
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 19 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_schema.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_schema.py
@@ -132,9 +132,9 @@ def __init__(self) -> None:
         )
 
         # self._table_tags[<database_name>][<schema_name>][<table_name>] = list of tags applied to table
-        self._table_tags: Dict[
-            str, Dict[str, Dict[str, List[SnowflakeTag]]]
-        ] = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
+        self._table_tags: Dict[str, Dict[str, Dict[str, List[SnowflakeTag]]]] = (
+            defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
+        )
 
         # self._column_tags[<database_name>][<schema_name>][<table_name>][<column_name>] = list of tags applied to column
         self._column_tags: Dict[
@@ -320,9 +320,7 @@ def get_tables_for_schema(
         return tables
 
     @serialized_lru_cache(maxsize=1)
-    def get_views_for_database(
-        self, db_name: str, include_view_definitions: bool
-    ) -> Dict[str, List[SnowflakeView]]:
+    def get_views_for_database(self, db_name: str) -> Dict[str, List[SnowflakeView]]:
         page_limit = SHOW_VIEWS_MAX_PAGE_SIZE
 
         views: Dict[str, List[SnowflakeView]] = {}
@@ -355,9 +353,7 @@ def get_views_for_database(
                         created=view["created_on"],
                         # last_altered=table["last_altered"],
                         comment=view["comment"],
-                        view_definition=(
-                            view["text"] if include_view_definitions else None
-                        ),
+                        view_definition=view["text"],
                         last_altered=view["created_on"],
                         materialized=(
                             view.get("is_materialized", "false").lower() == "true"

diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_schema_gen.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_schema_gen.py
@@ -171,9 +171,9 @@ def __init__(
             config, self.data_dictionary, self.report
         )
         self.profiler: Optional[SnowflakeProfiler] = profiler
-        self.snowsight_url_builder: Optional[
-            SnowsightUrlBuilder
-        ] = snowsight_url_builder
+        self.snowsight_url_builder: Optional[SnowsightUrlBuilder] = (
+            snowsight_url_builder
+        )
 
         # These are populated as side-effects of get_workunits_internal.
         self.databases: List[SnowflakeDatabase] = []
@@ -228,9 +228,9 @@ def get_databases(self) -> Optional[List[SnowflakeDatabase]]:
             )
             return None
         else:
-            ischema_databases: List[
-                SnowflakeDatabase
-            ] = self.get_databases_from_ischema(databases)
+            ischema_databases: List[SnowflakeDatabase] = (
+                self.get_databases_from_ischema(databases)
+            )
 
             if len(ischema_databases) == 0:
                 self.report_error(
@@ -753,7 +753,11 @@ def gen_dataset_workunits(
             view_properties_aspect = ViewProperties(
                 materialized=table.materialized,
                 viewLanguage="SQL",
-                viewLogic=table.view_definition,
+                viewLogic=(
+                    table.view_definition
+                    if self.config.include_view_definitions
+                    else None
+                ),
             )
 
             yield MetadataChangeProposalWrapper(
@@ -1017,9 +1021,7 @@ def get_tables_for_schema(
     def get_views_for_schema(
         self, schema_name: str, db_name: str
     ) -> List[SnowflakeView]:
-        views = self.data_dictionary.get_views_for_database(
-            db_name, self.config.include_view_definitions
-        )
+        views = self.data_dictionary.get_views_for_database(db_name)
 
         # Some schema may not have any table
         return views.get(schema_name, [])