Merge pull request #151 from 3DBAG/fix/ahn-adjustments

balazsdukai · web-flow · commit 4bbcdd80d8fe · 2025-02-10T15:28:24.000+01:00
Adjustments to the AHN assets
diff --git a/README.md b/README.md
@@ -46,14 +46,15 @@ You can read about the details on how can you deploy it in the [deployment secti
 ## Integration as a library
 
 The 3dbag-pipeline can be used as a library in other projects.
-The packages can be installed directly from GitHub using specific release versions:
+The packages can be installed directly from GitHub using specific release versions.
+Note that you must use the [uv](https://docs.astral.sh/uv/) package manager to install the packages, because *pip* cannot resolve the relative package paths within this repository.
 
 ```bash
 # Install specific release version of the common package
-pip install "bag3d-common @ git+https://github.com/3DBAG/3dbag-pipeline.git@v2024.12.16#egg=bag3d-common&subdirectory=packages/common"
+uv pip install "bag3d-common @ git+https://github.com/3DBAG/3dbag-pipeline.git@v2024.12.16#egg=bag3d-common&subdirectory=packages/common"
 
 # Install specific commit of the common package
-pip install "bag3d-common @ git+https://github.com/3DBAG/3dbag-pipeline.git@<commit-hash>#egg=bag3d-common&subdirectory=packages/common"
+uv pip install "bag3d-common @ git+https://github.com/3DBAG/3dbag-pipeline.git@<commit-hash>#egg=bag3d-common&subdirectory=packages/common"
 ```
 
 ## License
diff --git a/packages/core/src/bag3d/core/assets/ahn/core.py b/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:
diff --git a/packages/core/src/bag3d/core/assets/ahn/download.py b/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,12 +322,17 @@ 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
+    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.
 
     Args:
+        nr_retries: The number of retries to download the file.
         fpath: Path to the LAZ file that may exist locally. If not it will be downloaded.
         url_laz: Complete URL of the file to download. If provided, 'url_base' is
             ignored.
@@ -340,36 +345,50 @@ def download_ahn_laz(
 
     url = url_laz if url_laz is not None else "/".join([url_base, fpath.name])
 
+    success = False
+    file_size = 0.0
+    is_new = False
     if not fpath.is_file():
         logger.info(format_laz_log(fpath, "Not found. Downloading..."))
-        fpath = download_file(
-            url=url, target_path=fpath, chunk_size=1024 * 1024, verify=verify_ssl
-        )
-        if fpath is None:
-            # Download failed
-            logger.warning(format_laz_log(fpath, "Downloading failed!"))
-            return LAZDownload(
-                url=None,
-                path=Path(),
-                success=False,
-                hash_name=None,
-                hash_hexdigest=None,
-                new=False,
-                size=0.0,
-            )
-        else:
-            is_new = True
+        for i in range(nr_retries):
+            try:
+                fpath = download_file(
+                    url=url,
+                    target_path=fpath,
+                    chunk_size=1024 * 1024,
+                    verify=verify_ssl,
+                )
+                if fpath is None:
+                    # Download failed
+                    logger.warning(format_laz_log(fpath, "Downloading failed!"))
+                    url_laz = None
+                    fpath = Path()
+                    success = False
+                    is_new = False
+                    file_size = 0.0
+                else:
+                    success = True
+                    is_new = True
+                    file_size = round(fpath.stat().st_size / 1e6, 2)
+                    break
+            except ConnectionError as e:
+                if i == 4:
+                    raise e
+                else:
+                    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
+        file_size = round(fpath.stat().st_size / 1e6, 2)
         is_new = False
     return LAZDownload(
         url=url_laz,
         path=fpath,
-        success=True,
+        success=success,
         hash_name=None,
         hash_hexdigest=None,
         new=is_new,
-        size=round(fpath.stat().st_size / 1e6, 2),
+        size=file_size,
     )
 
 
diff --git a/packages/core/src/bag3d/core/assets/ahn/index.py b/packages/core/src/bag3d/core/assets/ahn/index.py
@@ -0,0 +1,111 @@
+from dagster import asset, Field, get_dagster_logger
+
+from bag3d.core.assets.ahn.core import partition_definition_ahn
+
+
+logger = get_dagster_logger("ahn.index")
+
+
+@asset(
+    config_schema={
+        "tile_size": Field(
+            int,
+            default_value=250,
+            description="Set smallest spatial area indexed to tile_size by tile_size units.",
+        ),
+        "force": Field(
+            bool,
+            default_value=False,
+            description="Force the re-index the file, even if it is already indexed.",
+        ),
+    },
+    required_resource_keys={"lastools"},
+    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`.
+
+    See https://lastools.osgeo.org/download/lasindex_README.txt.
+    """
+    cmd_list = [
+        "{exe}",
+        "-i {local_path}",
+        "-append",
+        "-tile_size",
+        str(context.op_execution_context.op_config["tile_size"]),
+    ]
+    if context.op_execution_context.op_config["force"] is False:
+        cmd_list.append("-dont_reindex")
+    context.resources.lastools.app.execute(
+        "lasindex", " ".join(cmd_list), local_path=laz_files_ahn3.path
+    )
+
+
+@asset(
+    config_schema={
+        "tile_size": Field(
+            int,
+            default_value=250,
+            description="Set smallest spatial area indexed to tile_size by tile_size units.",
+        ),
+        "force": Field(
+            bool,
+            default_value=False,
+            description="Force the re-index the file, even if it is already indexed.",
+        ),
+    },
+    required_resource_keys={"lastools"},
+    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`.
+
+    See https://lastools.osgeo.org/download/lasindex_README.txt.
+    """
+    cmd_list = [
+        "{exe}",
+        "-i {local_path}",
+        "-append",
+        "-tile_size",
+        str(context.op_execution_context.op_config["tile_size"]),
+    ]
+    if context.op_execution_context.op_config["force"] is False:
+        cmd_list.append("-dont_reindex")
+    context.resources.lastools.app.execute(
+        "lasindex", " ".join(cmd_list), local_path=laz_files_ahn4.path
+    )
+
+
+@asset(
+    config_schema={
+        "tile_size": Field(
+            int,
+            default_value=250,
+            description="Set smallest spatial area indexed to tile_size by tile_size units.",
+        ),
+        "force": Field(
+            bool,
+            default_value=False,
+            description="Force the re-index the file, even if it is already indexed.",
+        ),
+    },
+    required_resource_keys={"lastools"},
+    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`.
+
+    See https://lastools.osgeo.org/download/lasindex_README.txt.
+    """
+    cmd_list = [
+        "{exe}",
+        "-i {local_path}",
+        "-append",
+        "-tile_size",
+        str(context.op_execution_context.op_config["tile_size"]),
+    ]
+    if context.op_execution_context.op_config["force"] is False:
+        cmd_list.append("-dont_reindex")
+    context.resources.lastools.app.execute(
+        "lasindex", " ".join(cmd_list), local_path=laz_files_ahn5.path
+    )
diff --git a/packages/core/src/bag3d/core/assets/ahn/metadata.py b/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
diff --git a/packages/core/src/bag3d/core/assets/ahn/tile.py b/packages/core/src/bag3d/core/assets/ahn/tile.py
