Simplicity

Signed-off-by: Harm Weites <[email protected]>
amundsen-io · Feb 10, 2022 · 06b9fc3 · 06b9fc3
1 parent ad0280e
commit 06b9fc3
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/databuilder/databuilder/extractor/delta_lake_metadata_extractor.py b/databuilder/databuilder/extractor/delta_lake_metadata_extractor.py
@@ -210,8 +210,7 @@ def _get_extract_iter(self) -> Iterator[Union[TableMetadata, List[Tuple[Watermar
                 watermarks = self.create_table_watermarks(scraped_table)
                 if watermarks:
                     for watermark in watermarks:
-                        yield watermark[0]
-                        yield watermark[1]
+                        yield watermark
                 last_updated = self.create_table_last_updated(scraped_table)
                 if last_updated:
                     yield last_updated
@@ -432,8 +431,7 @@ def is_array_type(self, delta_type: Any) -> bool:
     def is_map_type(self, delta_type: Any) -> bool:
         return isinstance(delta_type, MapType)
 
-    def create_table_watermarks(self, table: ScrapedTableMetadata) -> Union[List[Tuple[Optional[Watermark],
-                                                                                       Optional[Watermark]]], None]:
+    def create_table_watermarks(self, table: ScrapedTableMetadata) -> Optional[List[Watermark]]:
         """
         Creates the watermark objects that reflect the highest and lowest values in the partition columns
         """

diff --git a/databuilder/tests/unit/extractor/test_deltalake_extractor.py b/databuilder/tests/unit/extractor/test_deltalake_extractor.py
@@ -218,7 +218,7 @@ def test_extract(self) -> None:
         while data is not None:
             ret.append(data)
             data = self.dExtractor.extract()
-        self.assertEqual(len(ret), 40)
+        self.assertEqual(len(ret), 37)
 
     def test_extract_with_only_specific_schemas(self) -> None:
         self.config_dict = {
@@ -235,7 +235,7 @@ def test_extract_with_only_specific_schemas(self) -> None:
         while data is not None:
             ret.append(data)
             data = self.dExtractor.extract()
-        self.assertEqual(len(ret), 12)
+        self.assertEqual(len(ret), 9)
 
     def test_extract_when_excluding(self) -> None:
         self.config_dict = {