refactor: move locker into S3Storage

CHANGELOG.md

@@ -6,7 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [PEP 440](https://www.python.org/dev/peps/pep-0440/).
 
 
-## 2.0.0rc1 - 2024-01-XX
+## 2.0.0rc1 - 2024-01-01
 
 ### Added
 

s3pypi/__main__.py

@@ -129,10 +129,10 @@ def main(*raw_args: str) -> None:
             put_kwargs=args.s3_put_args,
             unsafe_s3_website=args.unsafe_s3_website,
             no_sign_request=args.no_sign_request,
+            lock_indexes=args.lock_indexes,
+            profile=args.profile,
+            region=args.region,
         ),
-        lock_indexes=args.lock_indexes,
-        profile=args.profile,
-        region=args.region,
     )
 
     try:

s3pypi/core.py

@@ -6,15 +6,12 @@
 from itertools import groupby
 from operator import attrgetter
 from pathlib import Path
-from typing import List, Optional, Tuple
+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 DummyLocker, DynamoDBLocker, Locker
 from s3pypi.storage import S3Config, S3Storage
 
 log = logging.getLogger(__prog__)
@@ -25,9 +22,6 @@
 @dataclass
 class Config:
     s3: S3Config
-    lock_indexes: bool = False
-    profile: Optional[str] = None
-    region: Optional[str] = None
 
 
 @dataclass
@@ -41,35 +35,23 @@ def normalize_package_name(name: str) -> str:
     return re.sub(r"[-_.]+", "-", name.lower())
 
 
-def build_storage_and_locker(cfg: Config) -> Tuple[S3Storage, Locker]:
-    session = boto3.Session(profile_name=cfg.profile, region_name=cfg.region)
-    storage = S3Storage(session, cfg.s3)
-    lock = (
-        DynamoDBLocker(session, table=f"{cfg.s3.bucket}-locks")
-        if cfg.lock_indexes
-        else DummyLocker()
-    )
-    return storage, lock
-
-
 def upload_packages(
     cfg: Config,
     dist: List[Path],
     put_root_index: bool = False,
     strict: bool = False,
     force: bool = False,
 ) -> None:
-    storage, lock = build_storage_and_locker(cfg)
+    storage = S3Storage(cfg.s3)
     distributions = parse_distributions(dist)
 
     get_name = attrgetter("name")
     existing_files = []
 
     for name, group in groupby(sorted(distributions, key=get_name), get_name):
         directory = normalize_package_name(name)
-        with lock(directory):
-            index = storage.get_index(directory)
 
+        with storage.locked_index(directory) as index:
             for distr in group:
                 filename = distr.local_path.name
 
@@ -82,12 +64,9 @@ def upload_packages(
                     storage.put_distribution(directory, distr.local_path)
                     index.filenames[filename] = Hash.of("sha256", distr.local_path)
 
-            storage.put_index(directory, index)
-
     if put_root_index:
-        with lock(storage.root):
-            index = storage.build_root_index()
-            storage.put_index(storage.root, index)
+        with storage.locked_index(storage.root) as index:
+            index.filenames = dict.fromkeys(storage.list_directories())
 
     if strict and existing_files:
         raise S3PyPiError(f"Found {len(existing_files)} existing files on S3")
@@ -143,12 +122,10 @@ def extract_wheel_metadata(path: Path) -> PackageMetadata:
 
 
 def delete_package(cfg: Config, name: str, version: str) -> None:
-    storage, lock = build_storage_and_locker(cfg)
+    storage = S3Storage(cfg.s3)
     directory = normalize_package_name(name)
 
-    with lock(directory):
-        index = storage.get_index(directory)
-
+    with storage.locked_index(directory) as index:
         filenames = [f for f in index.filenames if f.split("-", 2)[1] == version]
         if not filenames:
             raise S3PyPiError(f"Package not found: {name} {version}")
@@ -158,14 +135,6 @@ def delete_package(cfg: Config, name: str, version: str) -> None:
             storage.delete(directory, filename)
             del index.filenames[filename]
 
-        if not index.filenames:
-            storage.delete(directory, storage.index_name)
-        else:
-            storage.put_index(directory, index)
-
     if not index.filenames:
-        with lock(storage.root):
-            index = storage.get_index(storage.root)
-            if directory in index.filenames:
-                del index.filenames[directory]
-                storage.put_index(storage.root, index)
+        with storage.locked_index(storage.root) as index:
+            index.filenames.pop(directory, None)

s3pypi/storage.py

@@ -1,13 +1,15 @@
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, Iterator, List, Optional
 
 import boto3
 import botocore
 from botocore.config import Config as BotoConfig
 from mypy_boto3_s3.service_resource import Object
 
 from s3pypi.index import Index
+from s3pypi.locking import DummyLocker, DynamoDBLocker
 
 
 @dataclass
@@ -18,21 +20,32 @@ class S3Config:
     put_kwargs: Dict[str, str] = field(default_factory=dict)
     unsafe_s3_website: bool = False
     no_sign_request: bool = False
+    lock_indexes: bool = False
+    profile: Optional[str] = None
+    region: Optional[str] = None
 
 
 class S3Storage:
     root = "/"
     _index = "index.html"
 
-    def __init__(self, session: boto3.session.Session, cfg: S3Config):
-        _config = None
+    def __init__(self, cfg: S3Config):
+        session = boto3.Session(profile_name=cfg.profile, region_name=cfg.region)
+
+        config = None
         if cfg.no_sign_request:
-            _config = BotoConfig(signature_version=botocore.session.UNSIGNED)  # type: ignore
+            config = BotoConfig(signature_version=botocore.session.UNSIGNED)  # type: ignore
 
-        self.s3 = session.resource("s3", endpoint_url=cfg.endpoint_url, config=_config)
+        self.s3 = session.resource("s3", endpoint_url=cfg.endpoint_url, config=config)
         self.index_name = self._index if cfg.unsafe_s3_website else ""
         self.cfg = cfg
 
+        self.lock = (
+            DynamoDBLocker(session, table=f"{cfg.bucket}-locks")
+            if cfg.lock_indexes
+            else DummyLocker()
+        )
+
     def _object(self, directory: str, filename: str) -> Object:
         parts = [directory, filename]
         if parts == [self.root, self.index_name]:
@@ -48,10 +61,18 @@ def get_index(self, directory: str) -> Index:
             return Index()
         return Index.parse(html.decode())
 
-    def build_root_index(self) -> Index:
-        return Index(dict.fromkeys(self._list_dirs()))
+    @contextmanager
+    def locked_index(self, directory: str) -> Iterator[Index]:
+        with self.lock(directory):
+            index = self.get_index(directory)
+            yield index
+
+            if index.filenames:
+                self.put_index(directory, index)
+            else:
+                self.delete(directory, self.index_name)
 
-    def _list_dirs(self) -> List[str]:
+    def list_directories(self) -> List[str]:
         prefix = f"{p}/" if (p := self.cfg.prefix) else ""
         return [
             d[len(prefix) :]

tests/integration/test_storage.py

@@ -4,12 +4,12 @@
 from s3pypi.storage import S3Config, S3Storage
 
 
-def test_index_storage_roundtrip(boto3_session, s3_bucket):
+def test_index_storage_roundtrip(s3_bucket):
     directory = "foo"
     index = Index({"bar": None})
 
     cfg = S3Config(bucket=s3_bucket.name)
-    s = S3Storage(boto3_session, cfg)
+    s = S3Storage(cfg)
 
     s.put_index(directory, index)
     got = s.get_index(directory)
@@ -31,8 +31,8 @@ def test_index_storage_roundtrip(boto3_session, s3_bucket):
         (S3Config("", unsafe_s3_website=True), "/", index, "index.html"),
     ],
 )
-def test_s3_key(boto3_session, cfg, directory, filename, expected_key):
-    s = S3Storage(boto3_session, cfg)
+def test_s3_key(cfg, directory, filename, expected_key):
+    s = S3Storage(cfg)
     if filename is index:
         filename = s.index_name
 
@@ -41,23 +41,23 @@ def test_s3_key(boto3_session, cfg, directory, filename, expected_key):
     assert obj.key == expected_key
 
 
-def test_list_dirs(boto3_session, s3_bucket):
+def test_list_directories(s3_bucket):
     cfg = S3Config(bucket=s3_bucket.name, prefix="AA")
-    s = S3Storage(boto3_session, cfg)
+    s = S3Storage(cfg)
     s.put_index("one", Index())
     s.put_index("two", Index())
     s.put_index("three", Index())
 
-    assert s._list_dirs() == ["one/", "three/", "two/"]
+    assert s.list_directories() == ["one/", "three/", "two/"]
 
     cfg = S3Config(bucket=s3_bucket.name, prefix="BBBB")
-    s = S3Storage(boto3_session, cfg)
+    s = S3Storage(cfg)
     s.put_index("xxx", Index())
     s.put_index("yyy", Index())
 
-    assert s._list_dirs() == ["xxx/", "yyy/"]
+    assert s.list_directories() == ["xxx/", "yyy/"]
 
     cfg = S3Config(bucket=s3_bucket.name)
-    s = S3Storage(boto3_session, cfg)
+    s = S3Storage(cfg)
 
-    assert s._list_dirs() == ["AA/", "BBBB/"]
+    assert s.list_directories() == ["AA/", "BBBB/"]