diff --git a/docs/source/usage.rst b/docs/source/usage.rst index b2d4223a..8a434d0e 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -437,7 +437,10 @@ where * ``--hpss=`` specifies the destination path on the HPSS file system, * ``-l`` an optional argument to display more information. * ``--cache`` to use a cache other than the default of ``zstash``. -* ``--tars`` to list the tars in addition to the files. +* ``--tars`` to list the tars containing the matched files. + When combined with a ``[files]`` pattern, only the tars that hold those + matching files are shown — making it easy to identify the subset of tar + archives to download. * ``-v`` increases output verbosity. * ``[files]`` is a list of files to be listed (standard wildcards supported). @@ -476,6 +479,26 @@ Below is an example of using ``ls`` to look at the tars in addition to the files Tars: 000000.tar +When combined with a file pattern, ``--tars`` shows only the tars that contain +the matching files. This is useful for downloading a subset of the archive — +for example, to find which tars hold files from 1850–1900 of a historical +simulation:: + + $ zstash ls --hpss=hpss_archive --tars "*historical*185[0-9]*" "*historical*18[6-9][0-9]*" "*historical*190[0]*" + + archive/run/historical.cam.h0.1850-01.nc + archive/run/historical.cam.h0.1851-06.nc + ... + archive/run/historical.cam.h0.1900-12.nc + + Tars: + 000000.tar + 000001.tar + 000007.tar + +You can then pass these tar names to ``zstash extract --tars`` to download only +that subset. + .. warning:: Running ``zstash ls`` outside the source directory (the directory you're archiving) is not advised. ``zstash`` will only retrieve ``index.db`` from the HPSS archive diff --git a/tests/unit/test_ls.py b/tests/unit/test_ls.py new file mode 100644 index 00000000..d717a163 --- /dev/null +++ b/tests/unit/test_ls.py @@ -0,0 +1,81 @@ +import datetime +import os +import sqlite3 +import tempfile +import unittest +from unittest.mock import patch + +from zstash.ls import ls_tars_database +from zstash.settings import FilesRow + + +class TestLsTarsDatabase(unittest.TestCase): + def setUp(self): + # Create an in-memory SQLite database with files and tars tables + self.tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False) + self.db_path = self.tmp.name + self.tmp.close() + + con = sqlite3.connect(self.db_path) + cur = con.cursor() + cur.execute( + "CREATE TABLE tars (id INTEGER, name TEXT, size INTEGER, md5 TEXT)" + ) + cur.execute( + "INSERT INTO tars VALUES (1, '000000.tar', 10240, 'abc123')" + ) + cur.execute( + "INSERT INTO tars VALUES (2, '000001.tar', 20480, 'def456')" + ) + cur.execute( + "INSERT INTO tars VALUES (3, '000002.tar', 30720, 'ghi789')" + ) + con.commit() + con.close() + + def tearDown(self): + os.unlink(self.db_path) + + def _make_args(self, long: bool = False): + from unittest.mock import MagicMock + + args = MagicMock() + args.long = long + return args + + def test_ls_tars_database_filtered(self): + """Only tars containing matched files should be returned.""" + args = self._make_args() + with patch("zstash.ls.get_db_filename", return_value=self.db_path): + result = ls_tars_database(args, "zstash", ["000000.tar", "000002.tar"]) + + names = [r.name for r in result] + self.assertIn("000000.tar", names) + self.assertIn("000002.tar", names) + self.assertNotIn("000001.tar", names) + + def test_ls_tars_database_unfiltered(self): + """When tar_names is None, all tars should be returned.""" + args = self._make_args() + with patch("zstash.ls.get_db_filename", return_value=self.db_path): + result = ls_tars_database(args, "zstash", None) + + names = [r.name for r in result] + self.assertIn("000000.tar", names) + self.assertIn("000001.tar", names) + self.assertIn("000002.tar", names) + + def test_tar_names_from_file_matches(self): + """The unique tar names extracted from FilesRow matches are correct.""" + now = datetime.datetime(2024, 1, 1) + row1 = FilesRow((1, "archive/file_1850.nc", 100, now, "md5a", "000000.tar", 0)) + row2 = FilesRow((2, "archive/file_1851.nc", 100, now, "md5b", "000000.tar", 512)) + row3 = FilesRow((3, "archive/file_1900.nc", 100, now, "md5c", "000005.tar", 0)) + + matches = [row1, row2, row3] + tar_names = sorted(set(m.tar for m in matches)) + self.assertEqual(tar_names, ["000000.tar", "000005.tar"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/zstash/ls.py b/zstash/ls.py index 8b6ad6e4..49f06746 100644 --- a/zstash/ls.py +++ b/zstash/ls.py @@ -5,7 +5,7 @@ import os import sqlite3 import sys -from typing import List, Tuple, Union +from typing import List, Optional, Tuple, Union from .hpss import hpss_get from .settings import ( @@ -36,7 +36,8 @@ def ls(): print_matches(args, matches) if args.tars: - tar_matches: List[TarsRow] = ls_tars_database(args, cache) + tar_names: List[str] = sorted(set(m.tar for m in matches)) + tar_matches: List[TarsRow] = ls_tars_database(args, cache, tar_names) print_matches(args, tar_matches) @@ -165,7 +166,9 @@ def ls_database(args: argparse.Namespace, cache: str) -> List[FilesRow]: return matches -def ls_tars_database(args: argparse.Namespace, cache: str) -> List[TarsRow]: +def ls_tars_database( + args: argparse.Namespace, cache: str, tar_names: Optional[List[str]] = None +) -> List[TarsRow]: con: sqlite3.Connection = sqlite3.connect( get_db_filename(cache), detect_types=sqlite3.PARSE_DECLTYPES ) @@ -175,8 +178,14 @@ def ls_tars_database(args: argparse.Namespace, cache: str) -> List[TarsRow]: print("\ntars table does not exist") return [] - # Find matching files - cur.execute("select * from tars") + # Find matching tars + if tar_names is not None: + placeholders = ",".join("?" * len(tar_names)) + cur.execute( + "select * from tars where name in ({})".format(placeholders), tar_names + ) + else: + cur.execute("select * from tars") matches_: List[TupleTarsRow] = cur.fetchall() # Remove duplicates