update augment_csv to add a creation time column

Author: luizirber

machine/archive/scripts/augment_csv.py

@@ -5,6 +5,8 @@
 # ]
 # ///
 
+from datetime import datetime
+
 import polars as pl
 
 
@@ -16,12 +18,43 @@ def main(args):
         separator=" ",
         new_columns=["sha256", "sep", "path"],
     ).drop("sep")
+
+    s3_metadata = {
+        "creation_date": [],
+        "size": [],
+        "filename": [],
+    }
+    with args.s3_metadata.open() as f:
+        _ = next(f)
+
+        for line in f:
+            date, time, size, filename = line.strip().split()
+            dt = datetime(
+                *[int(v) for v in date.split("/")], *[int(v) for v in time.split(":")]
+            )
+            s3_metadata["creation_date"].append(dt)
+            s3_metadata["size"].append(int(size))
+            s3_metadata["filename"].append(f"sigs/{filename}")
+
+    s3_metadata_df = pl.from_dict(
+        s3_metadata,
+        schema={
+            "creation_date": pl.Datetime,
+            "size": pl.Int64,
+            "filename": pl.String,
+        },
+    ).lazy()
+
+    print(s3_metadata_df.head(5).collect())
+
     df = manifest_df.join(sha256_df, left_on="internal_location", right_on="path")
     if args.basepath:
         df = df.with_columns(
             pl.col("internal_location").str.strip_prefix(args.basepath)
         )
 
+    df = df.join(s3_metadata_df, left_on="internal_location", right_on="filename")
+
     match args.format:
         case "parquet":
             df.sink_parquet(
@@ -38,12 +71,14 @@ def main(args):
 
 if __name__ == "__main__":
     import argparse
+    import pathlib
 
     parser = argparse.ArgumentParser()
     parser.add_argument("-b", "--basepath", default=None)
     parser.add_argument("-F", "--format", choices=["parquet", "csv"], default="parquet")
     parser.add_argument("manifest")
     parser.add_argument("sha256")
+    parser.add_argument("s3_metadata", type=pathlib.Path)
     parser.add_argument("output", type=argparse.FileType("w", encoding="utf-8"))
 
     args = parser.parse_args()

machine/mirror/wort-mirror.py

@@ -13,8 +13,8 @@
 
 import aiofiles
 import httpx
-from httpx_retries import RetryTransport
 import polars as pl
+from httpx_retries import RetryTransport
 
 
 ARCHIVE_URL = "https://farm.cse.ucdavis.edu/~irber"
@@ -75,6 +75,8 @@ async def main(args):
 
     print(to_mirror_df.collect())
 
+    (args.basedir / "sigs").mkdir(parents=True, exist_ok=True)
+
     async with httpx.AsyncClient(
         timeout=30.0,
         # limits=httpx.Limits(max_connections=args.max_downloaders),
@@ -95,22 +97,25 @@ async def main(args):
                         )
                     )
         except* Exception as eg:
-            print(*[str(e)[:50] for e in eg.exceptions])
-
-        # copy manifest
-        if args.dry_run:
-            print(f"download: {manifest_url}")
-            return
-
-        async with client.stream("GET", "SOURMASH-MANIFEST.parquet") as response:
-            async with aiofiles.tempfile.NamedTemporaryFile() as f:
-                async for chnk in response.aiter_raw(1024 * 1024):
-                    await f.write(chnk)
-                await f.flush()
-
-                await asyncio.to_thread(
-                    shutil.copyfile, f.name, args.basedir / "SOURMASH-MANIFEST.parquet"
-                )
+            print(*[str(e)[:80] for e in eg.exceptions])
+            print(len(eg.exceptions))
+        else:
+            # copy manifest
+            if args.dry_run:
+                print(f"download: {manifest_url}")
+                return
+
+            async with client.stream("GET", "SOURMASH-MANIFEST.parquet") as response:
+                async with aiofiles.tempfile.NamedTemporaryFile() as f:
+                    async for chnk in response.aiter_raw(1024 * 1024):
+                        await f.write(chnk)
+                    await f.flush()
+
+                    await asyncio.to_thread(
+                        shutil.copyfile,
+                        f.name,
+                        args.basedir / "SOURMASH-MANIFEST.parquet",
+                    )
 
 
 async def download_sig(location, sha256, basedir, client, limiter, dry_run):
@@ -175,9 +180,17 @@ async def download_sig(location, sha256, basedir, client, limiter, dry_run):
         help="Calculate sha256 for local files, instead of depending only on filename",
     )
     parser.add_argument(
-        "database", default="img", choices=DATABASES, metavar="database", help=f"Which database to download. Available databases: {', '.join(DATABASES)}"
+        "database",
+        default="img",
+        choices=DATABASES,
+        metavar="database",
+        help=f"Which database to download. Available databases: {', '.join(DATABASES)}",
+    )
+    parser.add_argument(
+        "basedir",
+        type=pathlib.Path,
+        help="base directory for the mirror (existing or new)",
     )
-    parser.add_argument("basedir", type=pathlib.Path, help="base directory for the mirror (existing or new)")
 
     args = parser.parse_args()