Name loose S3 objects relative to their common prefix

grobid_client/grobid_client.py

@@ -824,6 +824,30 @@ def _s3_basename(self, uri: str) -> str:
         """Return the last path component of an s3:// key."""
         return self._split_s3(uri)[1].rsplit("/", 1)[-1]
 
+    def _s3_common_prefix(self, uris: list) -> str:
+        """Return the key prefix (up to a '/') shared by all the given objects.
+
+        Plays the role _common_base plays for local files: loose remote
+        objects are named relative to it, so pdfs/a/x.pdf and pdfs/b/x.pdf
+        listed under pdfs/ become a/x.pdf and b/x.pdf and do not collide on
+        the same output file.
+        """
+        keys = [self._split_s3(uri)[1] for uri in uris]
+        if not keys:
+            return ""
+        if len(keys) == 1:
+            prefix = keys[0]
+        else:
+            prefix = os.path.commonprefix(keys)
+        return prefix[:prefix.rfind("/") + 1] if "/" in prefix else ""
+
+    def _s3_relative_name(self, uri: str, prefix: str) -> str:
+        """Name a remote object by its key relative to prefix (see above)."""
+        key = self._split_s3(uri)[1]
+        name = key[len(prefix):] if prefix and key.startswith(prefix) else key
+        safe = self._safe_member_path("", name)
+        return safe if safe else self._s3_basename(uri)
+
     def _import_smart_open(self) -> Any:
         try:
             import smart_open  # noqa: F401
@@ -1373,6 +1397,9 @@ def _process_remote_files(
 
         queue_size = self._effective_queue_size(n)
         print(f"Found {total} remote file(s) to process")
+        # Objects are named relative to their common prefix so that the
+        # output mirrors the listed "directory", as it does for local ones
+        prefix = self._s3_common_prefix(uris)
         processed_count = 0
         error_count = 0
         skipped_count = 0
@@ -1397,7 +1424,7 @@ def _process_remote_files(
                         self.logger.error(f"Failed to fetch {uri}: {str(e)}")
                         error_count += 1
                         continue
-                    document.name = self._s3_basename(uri)
+                    document.name = self._s3_relative_name(uri, prefix)
                     documents.append(document)
 
                 if not documents:
@@ -1434,8 +1461,9 @@ def _process_remote_files(
                 for uri in chunk:
                     if verbose:
                         self.logger.info(f"Fetching {uri}")
-                    dest = os.path.join(temp_dir, self._s3_basename(uri))
+                    dest = os.path.join(temp_dir, self._s3_relative_name(uri, prefix))
                     try:
+                        os.makedirs(os.path.dirname(dest), exist_ok=True)
                         with self._s3_open(uri) as src, open(dest, "wb") as out_file:
                             shutil.copyfileobj(src, out_file)
                         local_files.append(dest)

tests/test_s3_output_names.py

@@ -0,0 +1,104 @@
+"""
+Loose S3 objects are named relative to their common prefix, so two keys
+that share a basename (a/x.pdf, b/x.pdf) get two output files.
+"""
+import os
+
+import pytest
+from unittest.mock import Mock, patch
+
+boto3 = pytest.importorskip("boto3")
+pytest.importorskip("smart_open")
+pytest.importorskip("moto")
+try:
+    from moto import mock_aws
+except ImportError:  # moto < 5
+    from moto import mock_s3 as mock_aws
+
+from grobid_client.grobid_client import GrobidClient
+
+BUCKET = "test-bucket"
+REGION = "us-east-1"
+
+
+@pytest.fixture(autouse=True)
+def _aws_env(monkeypatch):
+    for var in ("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SECURITY_TOKEN", "AWS_SESSION_TOKEN"):
+        monkeypatch.setenv(var, "testing")
+    monkeypatch.setenv("AWS_DEFAULT_REGION", REGION)
+
+
+def _client():
+    with patch.object(GrobidClient, "_test_server_connection"):
+        with patch.object(GrobidClient, "_configure_logging"):
+            c = GrobidClient(check_server=False)
+    c.logger = Mock()
+    c.config["queue_size"] = 1          # one object per chunk: collisions across chunks too
+    return c
+
+
+def _tei(out):
+    found = []
+    for root, _, files in os.walk(out):
+        for f in files:
+            if f.endswith(".grobid.tei.xml"):
+                found.append(os.path.relpath(os.path.join(root, f), out).replace(os.sep, "/"))
+    return sorted(found)
+
+
+def test_common_prefix():
+    c = _client()
+    assert c._s3_common_prefix([f"s3://{BUCKET}/pdfs/a/x.pdf", f"s3://{BUCKET}/pdfs/b/x.pdf"]) == "pdfs/"
+    assert c._s3_common_prefix([f"s3://{BUCKET}/pdfs/x.pdf", f"s3://{BUCKET}/pdfs/y.pdf"]) == "pdfs/"
+    assert c._s3_common_prefix([f"s3://{BUCKET}/pdfs/x.pdf"]) == "pdfs/"
+    assert c._s3_common_prefix([f"s3://{BUCKET}/x.pdf", f"s3://{BUCKET}/y.pdf"]) == ""
+    assert c._s3_common_prefix([f"s3://{BUCKET}/pdfs/xa.pdf", f"s3://{BUCKET}/pdfs/xb.pdf"]) == "pdfs/"
+    assert c._s3_relative_name(f"s3://{BUCKET}/pdfs/a/x.pdf", "pdfs/") == os.path.join("a", "x.pdf")
+    assert c._s3_relative_name(f"s3://{BUCKET}/pdfs/x.pdf", "pdfs/") == "x.pdf"
+
+
+@mock_aws
+def test_same_basename_under_different_prefixes(tmp_path):
+    s3 = boto3.client("s3", region_name=REGION)
+    s3.create_bucket(Bucket=BUCKET)
+    for k in ["pdfs/a/x.pdf", "pdfs/b/x.pdf", "pdfs/top.pdf"]:
+        s3.put_object(Bucket=BUCKET, Key=k, Body=b"%PDF-" + k.encode())
+    c = _client()
+    out = str(tmp_path / "out")
+    posted = []
+
+    def fake_post(url, files=None, data=None, headers=None, timeout=None):
+        posted.append(files["input"][0])
+        resp = Mock()
+        resp.text = "<TEI>ok</TEI>"
+        return (resp, 200)
+
+    with patch.object(GrobidClient, "post", side_effect=fake_post):
+        c.process("processFulltextDocument", f"s3://{BUCKET}/pdfs/", output=out, force=True)
+    assert sorted(posted) == [os.path.join("a", "x.pdf"), os.path.join("b", "x.pdf"), "top.pdf"]
+    assert _tei(out) == ["a/x.grobid.tei.xml", "b/x.grobid.tei.xml", "top.grobid.tei.xml"]
+
+    # a rerun without --force skips all three, none is mistaken for another
+    posted.clear()
+    with patch.object(GrobidClient, "post", side_effect=fake_post):
+        c.process("processFulltextDocument", f"s3://{BUCKET}/pdfs/", output=out, force=False)
+    assert posted == []
+
+
+@mock_aws
+def test_citation_lists_keep_their_prefix_on_the_disk_route(tmp_path):
+    s3 = boto3.client("s3", region_name=REGION)
+    s3.create_bucket(Bucket=BUCKET)
+    for k in ["lists/a/refs.txt", "lists/b/refs.txt"]:
+        s3.put_object(Bucket=BUCKET, Key=k, Body=b"one ref\n")
+    c = _client()
+    out = str(tmp_path / "out")
+
+    def fake_post(url, files=None, data=None, headers=None, timeout=None):
+        resp = Mock()
+        resp.text = "<TEI>refs</TEI>"
+        return (resp, 200)
+
+    with patch.object(GrobidClient, "post", side_effect=fake_post):
+        c.process("processCitationList", f"s3://{BUCKET}/lists/", output=out, force=True)
+    assert _tei(out) == ["a/refs.grobid.tei.xml", "b/refs.grobid.tei.xml"]