From 3ff236b1884eee190c60b69aced90307486264ba Mon Sep 17 00:00:00 2001 From: Tom Brooks Date: Fri, 26 Jun 2026 11:14:40 +0100 Subject: [PATCH] fix: free disk space in assemble-and-bake --- digital_land/cli.py | 7 ++++++ digital_land/commands.py | 33 +++++++++++++++++++++++++ digital_land/package/dataset_parquet.py | 6 +++-- tests/acceptance/test_dataset_create.py | 1 + tests/unit/test_commands.py | 19 ++++++++++++++ tests/unit/test_dataset_parquet.py | 27 ++++++++++++++++++++ 6 files changed, 91 insertions(+), 2 deletions(-) create mode 100644 tests/unit/test_dataset_parquet.py diff --git a/digital_land/cli.py b/digital_land/cli.py index 609516ec..a0ef184c 100644 --- a/digital_land/cli.py +++ b/digital_land/cli.py @@ -213,6 +213,11 @@ def convert_cmd(input_path, output_path): default="collection/resource.csv", help="link to where the resource list is stored", ) +@click.option( + "--cleanup-provenance-cache/--no-cleanup-provenance-cache", + default=True, + help="delete the intermediate provenance parquet cache before indexing to save disk (default: enabled)", +) @click.argument("input-dir", nargs=1, type=click.Path(exists=True)) @click.pass_context def dataset_create_cmd( @@ -225,6 +230,7 @@ def dataset_create_cmd( issue_dir, cache_dir, resource_path, + cleanup_provenance_cache, ): return dataset_create( input_dir=input_dir, @@ -238,6 +244,7 @@ def dataset_create_cmd( issue_dir=issue_dir, cache_dir=cache_dir, resource_path=resource_path, + cleanup_provenance_cache=cleanup_provenance_cache, ) diff --git a/digital_land/commands.py b/digital_land/commands.py index 4ec7d172..4c215744 100644 --- a/digital_land/commands.py +++ b/digital_land/commands.py @@ -322,6 +322,7 @@ def dataset_create( dataset_resource_dir="var/dataset-resource", cache_dir="var/cache", resource_path="collection/resource.csv", + cleanup_provenance_cache=True, ): """ Create a dataset package from transformed parquet files. @@ -434,6 +435,14 @@ def dataset_create( logger.info("loading fact,fact_resource and entity into {output_path}") pqpackage.load_to_sqlite(output_path) + # The provenance parquet has now been copied into the sqlite and is not + # needed again. Releasing it lowers peak disk usage on the Fargate task + # before the disk-heavy index build. Disabled in tests that need to + # inspect the intermediate parquet (e.g. parquet-vs-sqlite row checks). + if cleanup_provenance_cache: + pqpackage.close_conn() # release the duckdb overflow file handle first + _free_parquet_cache(dataset_parquet_path) + logger.info(f"add indexes to {output_path}") package.connect() package.create_cursor() @@ -524,6 +533,30 @@ def dataset_update( logging.warning("No directory for this dataset in the provided issue_directory") +def _free_parquet_cache(parquet_dir): + """Delete the intermediate provenance parquet cache to reclaim disk. + + Called after the parquet data has been loaded into the sqlite package and + before index creation, to lower peak disk usage on the (200 GiB, capped) + Fargate ephemeral volume. A failure here is logged but not fatal - we would + rather attempt the index build than abort on a cleanup error. + """ + parquet_dir = Path(parquet_dir) + if not parquet_dir.exists(): + return + try: + before = shutil.disk_usage(parquet_dir).free + shutil.rmtree(parquet_dir) + after = shutil.disk_usage(parquet_dir.parent).free + freed_gb = (after - before) / 1024**3 + logger.info( + f"removed intermediate parquet cache {parquet_dir}, " + f"freed ~{freed_gb:.1f} GiB (now {after / 1024**3:.1f} GiB free)" + ) + except OSError as e: + logger.warning(f"could not remove parquet cache {parquet_dir}: {e}") + + def dataset_dump(input_path, output_path): cmd = f"sqlite3 -header -csv {input_path} 'select * from entity;' > {output_path}" logging.info(cmd) diff --git a/digital_land/package/dataset_parquet.py b/digital_land/package/dataset_parquet.py index f2b983ca..be407967 100644 --- a/digital_land/package/dataset_parquet.py +++ b/digital_land/package/dataset_parquet.py @@ -823,8 +823,10 @@ def close_conn(self): logging.info("Close connection to duckdb database in session") if self.conn is not None: self.conn.close() - if os.path.exists(self.duckdb_file): - os.remove(self.duckdb_file) + # duckdb_path is only set when a file-backed db is used (not in-memory) + duckdb_path = getattr(self, "duckdb_path", None) + if duckdb_path is not None and os.path.exists(duckdb_path): + os.remove(duckdb_path) def load(self): pass diff --git a/tests/acceptance/test_dataset_create.py b/tests/acceptance/test_dataset_create.py index 09d5c78a..d19870b3 100644 --- a/tests/acceptance/test_dataset_create.py +++ b/tests/acceptance/test_dataset_create.py @@ -296,6 +296,7 @@ def test_acceptance_dataset_create( str(cache_path), "--resource-path", str(resource_path), + "--no-cleanup-provenance-cache", str(input_dir), ], catch_exceptions=False, diff --git a/tests/unit/test_commands.py b/tests/unit/test_commands.py index e4f0a05b..cdebf741 100644 --- a/tests/unit/test_commands.py +++ b/tests/unit/test_commands.py @@ -197,3 +197,22 @@ def test_validate_and_add_data_input_error_thrown_when_no_resource_downloaded( ) assert "Failed to collect from URL" in str(error.value) assert "log status: 200" in str(error.value) + + +def test_free_parquet_cache_removes_dir(tmp_path): + from digital_land.commands import _free_parquet_cache + + cache = tmp_path / "provenance" + (cache / "fact" / "dataset=x").mkdir(parents=True) + (cache / "fact" / "dataset=x" / "fact.parquet").write_text("data") + + _free_parquet_cache(cache) + + assert not cache.exists() + + +def test_free_parquet_cache_missing_dir_is_noop(tmp_path): + from digital_land.commands import _free_parquet_cache + + # should not raise when the directory was never created + _free_parquet_cache(tmp_path / "does-not-exist") diff --git a/tests/unit/test_dataset_parquet.py b/tests/unit/test_dataset_parquet.py new file mode 100644 index 00000000..ab3ab448 --- /dev/null +++ b/tests/unit/test_dataset_parquet.py @@ -0,0 +1,27 @@ +def test_close_conn_removes_duckdb_file(tmp_path): + import duckdb + from digital_land.package.dataset_parquet import DatasetParquetPackage + + duckdb_path = tmp_path / "overflow.duckdb" + + # build a minimal instance, bypassing the spec-driven __init__ so this stays + # a true unit test of close_conn (which only needs .conn and .duckdb_path) + pkg = DatasetParquetPackage.__new__(DatasetParquetPackage) + pkg.duckdb_path = duckdb_path + pkg.conn = duckdb.connect(str(duckdb_path)) + assert duckdb_path.exists() + + pkg.close_conn() + + assert not duckdb_path.exists() + + +def test_close_conn_handles_in_memory_db(): + import duckdb + from digital_land.package.dataset_parquet import DatasetParquetPackage + + # in-memory mode never sets duckdb_path; close_conn must not raise + pkg = DatasetParquetPackage.__new__(DatasetParquetPackage) + pkg.conn = duckdb.connect() # in-memory + + pkg.close_conn() # should not raise AttributeError