Update PartitionDefinitionAHN

To new style of partitiondefintion

Author: balazsdukai

packages/core/src/bag3d/core/assets/ahn/core.py

@@ -8,10 +8,7 @@
 
 logger = get_dagster_logger("ahn")
 
-
-class PartitionDefinitionAHN(StaticPartitionsDefinition):
-    def __init__(self):
-        super().__init__(partition_keys=sorted(list(AHN_TILE_IDS)))
+partition_definition_ahn = StaticPartitionsDefinition(sorted(list(AHN_TILE_IDS)))
 
 
 def format_laz_log(fpath: Path, msg: str) -> str:

packages/core/src/bag3d/core/assets/ahn/download.py

@@ -7,10 +7,10 @@
 
 from bag3d.common.utils.requests import download_file, download_as_str
 from bag3d.core.assets.ahn.core import (
-    PartitionDefinitionAHN,
     format_laz_log,
     download_ahn_index,
     ahn_laz_dir,
+    partition_definition_ahn,
 )
 
 logger = get_dagster_logger("ahn.download")
@@ -141,7 +141,7 @@ def tile_index_ahn(context):
 
 @asset(
     required_resource_keys={"file_store"},
-    partitions_def=PartitionDefinitionAHN(),
+    partitions_def=partition_definition_ahn,
 )
 def laz_files_ahn3(context, md5_ahn3, tile_index_ahn):
     """AHN3 LAZ files as they are downloaded from PDOK.
@@ -196,7 +196,7 @@ def laz_files_ahn3(context, md5_ahn3, tile_index_ahn):
 
 @asset(
     required_resource_keys={"file_store"},
-    partitions_def=PartitionDefinitionAHN(),
+    partitions_def=partition_definition_ahn,
 )
 def laz_files_ahn4(context, md5_ahn4, tile_index_ahn):
     """AHN4 LAZ files as they are downloaded from PDOK.
@@ -254,7 +254,7 @@ def laz_files_ahn4(context, md5_ahn4, tile_index_ahn):
 
 @asset(
     required_resource_keys={"file_store"},
-    partitions_def=PartitionDefinitionAHN(),
+    partitions_def=partition_definition_ahn,
 )
 def laz_files_ahn5(context, sha256_ahn5, tile_index_ahn):
     """AHN5 LAZ files as they are downloaded from PDOK.
@@ -322,8 +322,11 @@ def get_checksums(url: str) -> Mapping[str, str]:
 
 
 def download_ahn_laz(
-    fpath: Path, url_laz: str = None, url_base: str = None, verify_ssl: bool = False,
-    nr_retries: int = 5
+    fpath: Path,
+    url_laz: str = None,
+    url_base: str = None,
+    verify_ssl: bool = False,
+    nr_retries: int = 5,
 ) -> LAZDownload:
     """Download an AHN LAZ file from the input url to the given path,
     if the file does not exists.
@@ -350,7 +353,10 @@ def download_ahn_laz(
         for i in range(nr_retries):
             try:
                 fpath = download_file(
-                    url=url, target_path=fpath, chunk_size=1024 * 1024, verify=verify_ssl
+                    url=url,
+                    target_path=fpath,
+                    chunk_size=1024 * 1024,
+                    verify=verify_ssl,
                 )
                 if fpath is None:
                     # Download failed
@@ -366,10 +372,10 @@ def download_ahn_laz(
                     file_size = round(fpath.stat().st_size / 1e6, 2)
                     break
             except ConnectionError as e:
-                if i==4:
+                if i == 4:
                     raise e
                 else:
-                    logger.warning(f"Retrying ({i+1}/5) due to {e}")
+                    logger.warning(f"Retrying ({i + 1}/5) due to {e}")
     else:  # pragma: no cover
         logger.info(format_laz_log(fpath, "File already downloaded"))
         success = True

packages/core/src/bag3d/core/assets/ahn/index.py

@@ -1,9 +1,11 @@
 from dagster import asset, Field, get_dagster_logger
 
-from bag3d.core.assets.ahn.core import PartitionDefinitionAHN
+from bag3d.core.assets.ahn.core import partition_definition_ahn
+
 
 logger = get_dagster_logger("ahn.index")
 
+
 @asset(
     config_schema={
         "tile_size": Field(
@@ -18,7 +20,7 @@
         ),
     },
     required_resource_keys={"lastools"},
-    partitions_def=PartitionDefinitionAHN(),
+    partitions_def=partition_definition_ahn,
 )
 def lasindex_ahn3(context, laz_files_ahn3):
     """Append a spatial index to the AHN3 LAZ file, using LASTools's `lasindex`.
@@ -53,7 +55,7 @@ def lasindex_ahn3(context, laz_files_ahn3):
         ),
     },
     required_resource_keys={"lastools"},
-    partitions_def=PartitionDefinitionAHN(),
+    partitions_def=partition_definition_ahn,
 )
 def lasindex_ahn4(context, laz_files_ahn4):
     """Append a spatial index to the AHN4 LAZ file, using LASTools's `lasindex`.
@@ -88,7 +90,7 @@ def lasindex_ahn4(context, laz_files_ahn4):
         ),
     },
     required_resource_keys={"lastools"},
-    partitions_def=PartitionDefinitionAHN(),
+    partitions_def=partition_definition_ahn,
 )
 def lasindex_ahn5(context, laz_files_ahn5):
     """Append a spatial index to the AHN5 LAZ file, using LASTools's `lasindex`.

packages/core/src/bag3d/core/assets/ahn/metadata.py

@@ -8,7 +8,7 @@
 
 from bag3d.common.utils.geodata import pdal_info
 from bag3d.common.utils.database import create_schema, load_sql
-from bag3d.core.assets.ahn.core import PartitionDefinitionAHN
+from bag3d.core.assets.ahn.core import partition_definition_ahn
 
 
 @asset(required_resource_keys={"db_connection"})
@@ -41,7 +41,7 @@ def metadata_table_ahn5(context):
         ),
     },
     required_resource_keys={"pdal", "db_connection"},
-    partitions_def=PartitionDefinitionAHN(),
+    partitions_def=partition_definition_ahn,
 )
 def metadata_ahn3(context, laz_files_ahn3, metadata_table_ahn3, tile_index_ahn):
     """Metadata of the AHN3 LAZ file, retrieved from the PDOK tile index and
@@ -64,7 +64,7 @@ def metadata_ahn3(context, laz_files_ahn3, metadata_table_ahn3, tile_index_ahn):
         ),
     },
     required_resource_keys={"pdal", "db_connection"},
-    partitions_def=PartitionDefinitionAHN(),
+    partitions_def=partition_definition_ahn,
 )
 def metadata_ahn4(context, laz_files_ahn4, metadata_table_ahn4, tile_index_ahn):
     """Metadata of the AHN4 LAZ file, retrieved from the PDOK tile index and
@@ -87,7 +87,7 @@ def metadata_ahn4(context, laz_files_ahn4, metadata_table_ahn4, tile_index_ahn):
         ),
     },
     required_resource_keys={"pdal", "db_connection"},
-    partitions_def=PartitionDefinitionAHN(),
+    partitions_def=partition_definition_ahn,
 )
 def metadata_ahn5(context, laz_files_ahn5, metadata_table_ahn5, tile_index_ahn):
     """Metadata of the AHN5 LAZ file, retrieved from the PDOK tile index and