Merge pull request #103 from lsst/tickets/DM-51061

DM-51061: Add CLI command for replica chunk management.

Author: andy-slac

python/lsst/dax/apdb/cli/apdb_cli.py

@@ -31,7 +31,7 @@
 from .logging_cli import LoggingCli
 
 
-def main(args: Sequence[str] | None = None) -> None:
+def main(args: Sequence[str] | None = None) -> int | None:
     """APDB command line tools."""
     parser = argparse.ArgumentParser(description="APDB command line tools")
     log_cli = LoggingCli(parser)
@@ -45,14 +45,15 @@ def main(args: Sequence[str] | None = None) -> None:
     _metadata_subcommand(subparsers)
     _convert_legacy_config_subcommand(subparsers)
     _metrics_subcommand(subparsers)
+    _replication_subcommand(subparsers)
 
     parsed_args = parser.parse_args(args)
     log_cli.process_args(parsed_args)
 
     kwargs = vars(parsed_args)
     # Strip keywords not understood by scripts.
     method = kwargs.pop("method")
-    method(**kwargs)
+    return method(**kwargs)
 
 
 def _create_sql_subcommand(subparsers: argparse._SubParsersAction) -> None:
@@ -239,3 +240,38 @@ def _metrics_log_to_influx(subparsers: argparse._SubParsersAction) -> None:
         default=False,
     )
     parser.set_defaults(method=scripts.metrics_log_to_influx)
+
+
+def _replication_subcommand(subparsers: argparse._SubParsersAction) -> None:
+    parser = subparsers.add_parser("replication", help="Operations with replication tables produced by APDB.")
+    subparsers = parser.add_subparsers(title="available subcommands", required=True)
+    _replication_list_chunks_subcommand(subparsers)
+    _replication_delete_chunks_subcommand(subparsers)
+
+
+def _replication_list_chunks_subcommand(subparsers: argparse._SubParsersAction) -> None:
+    parser = subparsers.add_parser("list-chunks", help="Print full list of replication chunks in APDB.")
+    parser.add_argument("apdb_config", help="Path to the APDB configuration.")
+    parser.set_defaults(method=scripts.replication_list_chunks)
+
+
+def _replication_delete_chunks_subcommand(subparsers: argparse._SubParsersAction) -> None:
+    parser = subparsers.add_parser("delete-chunks", help="Delete replication chunks from APDB.")
+    parser.add_argument("apdb_config", help="Path to the APDB configuration.")
+    parser.add_argument(
+        "chunk_id", type=int, help="Chunk ID to delete, all earlier chunks are deleted as well."
+    )
+    parser.add_argument(
+        "-p",
+        "--print-only",
+        help="Only print the list ochunks that will be deleted, but do not delete them.",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "--force",
+        help="Do not ask for confirmation before deleting chunks.",
+        action="store_true",
+        default=False,
+    )
+    parser.set_defaults(method=scripts.replication_delete_chunks)

python/lsst/dax/apdb/scripts/__init__.py

@@ -27,3 +27,4 @@
 from .list_index import list_index
 from .metadata import metadata_delete, metadata_get, metadata_set, metadata_show
 from .metrics import metrics_log_to_influx
+from .replication import replication_delete_chunks, replication_list_chunks

python/lsst/dax/apdb/scripts/replication.py

@@ -0,0 +1,109 @@
+# This file is part of dax_ppdb
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (https://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from __future__ import annotations
+
+__all__ = ["replication_delete_chunks", "replication_list_chunks"]
+
+import sys
+from collections.abc import Collection
+
+from ..apdbReplica import ApdbReplica, ReplicaChunk
+
+
+def replication_list_chunks(apdb_config: str) -> None:
+    """Print full list of replica chunks existing in APDB.
+
+    Parameters
+    ----------
+    apdb_config : `str`
+        URL for APDB configuration file.
+    """
+    apdb = ApdbReplica.from_uri(apdb_config)
+    chunks = apdb.getReplicaChunks()
+    if chunks is not None:
+        chunks = sorted(chunks, key=lambda chunk: chunk.id)
+        _print_chunks(chunks)
+    else:
+        print("APDB instance does not support replication")
+
+
+def replication_delete_chunks(apdb_config: str, chunk_id: int, force: bool, print_only: bool) -> int:
+    """Delete replication chunks from APDB.
+
+    Parameters
+    ----------
+    apdb_config : `str`
+        URL for APDB configuration file.
+    chunk_id : `int`
+        Chunk to delete, all earlier chunks are deleted as well.
+    force : `bool`
+        If `True` do not ask confirmation.
+    print_only : `bool`
+        If `True` print the list of chunks, but do not delete anything.
+    """
+    apdb = ApdbReplica.from_uri(apdb_config)
+    chunks = apdb.getReplicaChunks()
+    if chunks is None:
+        print("APDB instance does not support replications")
+    else:
+        chunks = sorted(chunks, key=lambda chunk: chunk.id)
+        # Chunks to delete.
+        chunks = [chunk for chunk in chunks if chunk.id <= chunk_id]
+        # Check that given chunk ID actually exists.
+        if not chunks or chunks[-1].id != chunk_id:
+            print(f"ERROR: Replication chunk with ID={chunk_id} does not exist", file=sys.stderr)
+            return 1
+
+        if print_only:
+            print("Following chunks will be deleted:")
+            _print_chunks(chunks)
+            return 0
+
+        if not force:
+            try:
+                response = input(f"{len(chunks)} chunks will be removed, y[n]? ")
+            except EOFError:
+                response = ""
+            if response not in ("y", "Y"):
+                return 0
+
+        apdb.deleteReplicaChunks(chunk.id for chunk in chunks)
+
+    return 0
+
+
+def _print_chunks(chunks: Collection[ReplicaChunk]) -> None:
+    """Print the list of chunks.
+
+    Parameters
+    ----------
+    chunks : `~collections.abc.Collection` [`ReplicaChunk`]
+        Chunks to print.
+    """
+    print(" Chunk Id            Update time                      Unique Id")
+    sep = "-" * 77
+    print(sep)
+    for chunk in chunks:
+        insert_time = chunk.last_update_time
+        print(f"{chunk.id:10d}  {insert_time.tai.isot}/tai  {chunk.unique_id}")
+    print(sep)
+    print(f"Total: {len(chunks)}")