From 7135e626956267684881d13564cb8a08eb5769d1 Mon Sep 17 00:00:00 2001
From: Matteo De Wint <matteo@gorilla.co>
Date: Mon, 1 Jan 2024 14:57:26 +0100
Subject: [PATCH] feat: add `s3pypi force-unlock` command

---
 CHANGELOG.md                   |  3 ++-
 s3pypi/__main__.py             | 23 +++++++++++++++++++++--
 s3pypi/core.py                 |  9 +++++++++
 s3pypi/locking.py              |  3 +--
 tests/integration/test_main.py |  4 ++++
 5 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f326261..7dc698f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,7 +11,8 @@ and this project adheres to [PEP 440](https://www.python.org/dev/peps/pep-0440/)
 ### Added
 
 - `s3pypi delete` command to delete packages from S3.
-- `--locks-table` to customise the DynamoDB table name used for locking.
+- `s3pypi force-unlock` command to release a stuck lock in DynamoDB.
+- `--locks-table` option to customise the DynamoDB table name used for locking.
 
 ### Changed
 
diff --git a/s3pypi/__main__.py b/s3pypi/__main__.py
index 2bb43d2..81c2aa9 100644
--- a/s3pypi/__main__.py
+++ b/s3pypi/__main__.py
@@ -59,15 +59,24 @@ def add_command(
     d.add_argument("version", help="Package version.")
     build_s3_args(d)
 
+    ul = add_command(force_unlock, help="Release a stuck lock in DynamoDB.")
+    ul.add_argument("table", help="DynamoDB table.")
+    ul.add_argument("lock_id", help="ID of the lock to release.")
+    build_aws_args(ul)
+
     return p
 
 
+def build_aws_args(p: ArgumentParser) -> None:
+    p.add_argument("--profile", help="Optional AWS profile to use.")
+    p.add_argument("--region", help="Optional AWS region to target.")
+
+
 def build_s3_args(p: ArgumentParser) -> None:
     p.add_argument("-b", "--bucket", required=True, help="The S3 bucket to upload to.")
     p.add_argument("--prefix", help="Optional prefix to use for S3 object names.")
 
-    p.add_argument("--profile", help="Optional AWS profile to use.")
-    p.add_argument("--region", help="Optional AWS region to target.")
+    build_aws_args(p)
     p.add_argument(
         "--no-sign-request",
         action="store_true",
@@ -116,6 +125,10 @@ def delete(cfg: core.Config, args: Namespace) -> None:
     core.delete_package(cfg, name=args.name, version=args.version)
 
 
+def force_unlock(cfg: core.Config, args: Namespace) -> None:
+    core.force_unlock(cfg, args.table, args.lock_id)
+
+
 def main(*raw_args: str) -> None:
     args = build_arg_parser().parse_args(raw_args or sys.argv[1:])
     log.setLevel(logging.DEBUG if args.verbose else logging.INFO)
@@ -131,6 +144,12 @@ def main(*raw_args: str) -> None:
             put_kwargs=args.s3_put_args,
             index_html=args.index_html,
             locks_table=args.locks_table,
+        )
+        if hasattr(args, "bucket")
+        else core.S3Config(
+            bucket="",
+            profile=args.profile,
+            region=args.region,
         ),
     )
 
diff --git a/s3pypi/core.py b/s3pypi/core.py
index 78440ed..2ef1f95 100644
--- a/s3pypi/core.py
+++ b/s3pypi/core.py
@@ -9,9 +9,12 @@
 from typing import List
 from zipfile import ZipFile
 
+import boto3
+
 from s3pypi import __prog__
 from s3pypi.exceptions import S3PyPiError
 from s3pypi.index import Hash
+from s3pypi.locking import DynamoDBLocker
 from s3pypi.storage import S3Config, S3Storage
 
 log = logging.getLogger(__prog__)
@@ -138,3 +141,9 @@ def delete_package(cfg: Config, name: str, version: str) -> None:
     if not index.filenames:
         with storage.locked_index(storage.root) as root_index:
             root_index.filenames.pop(directory, None)
+
+
+def force_unlock(cfg: Config, table: str, lock_id: str) -> None:
+    session = boto3.Session(profile_name=cfg.s3.profile, region_name=cfg.s3.region)
+    DynamoDBLocker.build(session, table)._unlock(lock_id)
+    log.info("Released lock %s", lock_id)
diff --git a/s3pypi/locking.py b/s3pypi/locking.py
index 1b25073..aa43a44 100644
--- a/s3pypi/locking.py
+++ b/s3pypi/locking.py
@@ -103,10 +103,9 @@ def _unlock(self, lock_id: str) -> None:
 
 class DynamoDBLockTimeoutError(exc.S3PyPiError):
     def __init__(self, table: str, item: dict):
-        key = json.dumps({"LockID": {"S": item["LockID"]}})
         super().__init__(
             f"Timed out trying to acquire lock:\n\n{json.dumps(item, indent=2)}\n\n"
             "Another instance of s3pypi may currently be holding the lock.\n"
             "If this is not the case, you may release the lock as follows:\n\n"
-            f"$ aws dynamodb delete-item --table-name {table} --key '{key}'\n"
+            f"$ s3pypi force-unlock {table} {item['LockID']}\n"
         )
diff --git a/tests/integration/test_main.py b/tests/integration/test_main.py
index 7078d7d..d4ed2e2 100644
--- a/tests/integration/test_main.py
+++ b/tests/integration/test_main.py
@@ -134,3 +134,7 @@ def assert_pkg_exists(pkg: str, filename: str):
     assert ">hello-world</a>" not in root_index
     assert_pkg_exists("foo", "foo-0.1.0.tar.gz")
     assert_pkg_exists("xyz", "xyz-0.1.0.zip")
+
+
+def test_main_force_unlock(dynamodb_table):
+    s3pypi("force-unlock", dynamodb_table.name, "12345")