wip

Author: GeigerJ2

src/aiida/cmdline/commands/cmd_archive.py

@@ -136,6 +136,10 @@ def inspect(ctx, archive, version, meta_data, database):
     help='Determine entities to export, but do not create the archive. Deprecated, please use `--dry-run` instead.',
 )
 @options.DRY_RUN(help='Determine entities to export, but do not create the archive.')
+@click.option(
+    '--base-tmp-dir',
+    help='Determine entities to export, but do not create the archive. Deprecated, please use `--dry-run` instead.',
+)
 @decorators.with_dbenv()
 def create(
     output_file,
@@ -158,6 +162,7 @@ def create(
     batch_size,
     test_run,
     dry_run,
+    temp_dir
 ):
     """Create an archive from all or part of a profiles's data.
 
@@ -209,6 +214,7 @@ def create(
         'compression': compress,
         'batch_size': batch_size,
         'test_run': dry_run,
+        'temp_dir': temp_dir
     }
 
     if AIIDA_LOGGER.level <= logging.REPORT:

src/aiida/tools/archive/create.py

@@ -59,6 +59,7 @@ def create_archive(
     compression: int = 6,
     test_run: bool = False,
     backend: Optional[StorageBackend] = None,
+    temp_dir: Optional[Union[str, Path]] = None,
     **traversal_rules: bool,
 ) -> Path:
     """Export AiiDA data to an archive file.
@@ -139,6 +140,12 @@ def create_archive(
 
     :param backend: the backend to export from. If not specified, the default backend is used.
 
+    :param temp_dir: Directory to use for temporary files during archive creation.
+        If not specified, a temporary directory will be created in the same directory as the output file
+        with a '.aiida-export-' prefix. This parameter is useful when the output directory has limited
+        space or when you want to use a specific filesystem (e.g., faster storage) for temporary operations.
+        The directory must exist and be writable.
+
     :param traversal_rules: graph traversal rules. See :const:`aiida.common.links.GraphTraversalRules`
         what rule names are toggleable and what the defaults are.
 
@@ -179,6 +186,25 @@ def querybuilder():
         name: traversal_rules.get(name, rule.default) for name, rule in GraphTraversalRules.EXPORT.value.items()
     }
 
+    # Handle temporary directory configuration
+    if temp_dir is not None:
+        temp_dir = Path(temp_dir)
+        if not temp_dir.exists():
+            msg = f"Specified temporary directory '{temp_dir}' does not exist"
+            raise ArchiveExportError(msg)
+        if not temp_dir.is_dir():
+            msg = f"Specified temporary directory '{temp_dir}' is not a directory"
+            raise ArchiveExportError(msg)
+        # Check if directory is writable
+        if not temp_dir.is_writable():
+            msg = f"Specified temporary directory '{temp_dir}' is not writable"
+            raise ArchiveExportError()
+        tmp_prefix = None  # Use default tempfile prefix
+    else:
+        # Create temporary directory in the same folder as the output file
+        temp_dir = filename.parent
+        tmp_prefix = '.aiida-export-'
+
     initial_summary = get_init_summary(
         archive_version=archive_format.latest_version,
         outfile=filename,
@@ -283,10 +309,10 @@ def querybuilder():
     # Create and open the archive for writing.
     # We create in a temp dir then move to final place at end,
     # so that the user cannot end up with a half written archive on errors
-    import ipdb; ipdb.set_trace()
-    base_temp_dir = '/mount'  # or whatever directory you want to use
-    with tempfile.TemporaryDirectory(dir=base_temp_dir) as tmpdir:
-    # with tempfile.TemporaryDirectory() as tmpdir:
+    # import ipdb; ipdb.set_trace()
+    temp_dir = Path('/mount')  # or whatever directory you want to use
+    with tempfile.TemporaryDirectory(dir=temp_dir, prefix=tmp_prefix) as tmpdir:
+        # NOTE: Add the `tmp_prefix` to the directory or file?
         tmp_filename = Path(tmpdir) / 'export.zip'
         with archive_format.open(tmp_filename, mode='x', compression=compression) as writer:
             # add metadata