Additional object fields

.vscode/settings.json

@@ -1,4 +1,5 @@
 {
     "ruff.lineLength": 120,
-    "editor.defaultFormatter": "charliermarsh.ruff"
+    "editor.defaultFormatter": "charliermarsh.ruff",
+    "ruff.format.preview": true
 }

README.md

@@ -16,10 +16,13 @@ positional arguments:
 options:
   -h, --help           show this help message and exit
   --bucket BUCKET      S3 bucket name (S3PUSHER_BUCKET environment variable can also be used)
-  --hostname HOSTNAME  Hostname to include in S3 object key (S3PUSHER_HOSTNAME environment variable can also be used)  --log-json           Log in JSON format
+  --fields FIELDS      Comma separated list of key=value pairs to include in S3 object key (S3PUSHER_FIELDS environment variable can also be used)
+  --log-json           Log in JSON format
   --debug              Enable debugging
 ```
 
+The environment variable `S3PUSHER_HOSTNAME` may also be used to set field `hostname`.
+
 
 ## Authentication
 
@@ -31,5 +34,11 @@ Environment variables used for authentication can be found in the [Boto3 documen
 Files will be uploaded to the specified bucket in the following format:
 
 ```
-year=YYYY/month=MM/day=DD/hour=HH/minute=MM/second=SS/hostname=HOSTNAME/uuid=UUID/FILENAME
+year=YYYY/month=MM/day=DD/hour=HH/minute=MM/second=SS/uuid=UUID/FILENAME
+```
+
+If `--fields provider=xyzzy,hostname=host` is specified, the format is:
+
+```
+year=YYYY/month=MM/day=DD/hour=HH/minute=MM/second=SS/provider=xyzzy/hostname=host/uuid=UUID/FILENAME
 ```

s3pusher.py

@@ -1,6 +1,7 @@
 import argparse
 import logging
 import os
+import re
 import time
 import uuid
 from datetime import UTC, datetime
@@ -16,11 +17,14 @@
 
 logger = structlog.get_logger()
 
+RESERVED_FIELD_NAMES = {"year", "month", "day", "hour", "minute", "second", "uuid"}
+FIELD_KV_RE = re.compile(r"^[a-zA-Z0-9_\-\.]+$")
+
 
 class ThePusher(FileSystemEventHandler):
-    def __init__(self, bucket: str | None, hostname: str | None) -> None:
+    def __init__(self, bucket: str | None, object_kvs: dict[str, str] | None = None) -> None:
         self.bucket = bucket
-        self.hostname = hostname
+        self.object_kvs = object_kvs or {}
         self.logger = structlog.get_logger()
 
     def on_modified(self, event: FileSystemEvent) -> None:
@@ -47,14 +51,18 @@ def upload_directory_to_s3(self, directory: Path) -> None:
     def upload_file_to_s3(self, filename: Path) -> None:
         """Upload file to S3 and delete it if upload is successful"""
 
-        with structlog.contextvars.bound_contextvars(filename=str(filename)):
-            if not self.wait_for_stable_file(filename):
+        if not self.wait_for_stable_file(filename):
+            return
+
+        s3_object_key = self.get_s3_object_key(filename=filename)
+
+        with structlog.contextvars.bound_contextvars(filename=str(filename), s3_object_key=s3_object_key):
+            if not self.bucket:
+                self.logger.warning("No bucket configured, skipping upload")
                 return
 
             try:
-                s3_client = boto3.client("s3")
-                s3_object_key = self.get_s3_object_key(hostname=self.hostname, filename=filename)
-                with structlog.contextvars.bound_contextvars(s3_bucket=self.bucket, s3_object_key=s3_object_key):
+                with structlog.contextvars.bound_contextvars(s3_bucket=self.bucket):
                     self.logger.debug("Uploading file")
                     s3_client = boto3.client("s3")
                     t1 = time.time()
@@ -68,9 +76,9 @@ def upload_file_to_s3(self, filename: Path) -> None:
                 self.logger.error("Failed to upload", exception=str(exc))
                 time.sleep(EXCEPTION_DELAY_SECONDS)
 
-    @staticmethod
-    def get_s3_object_key(filename: Path | None = None, hostname: str | None = None) -> str:
+    def get_s3_object_key(self, filename: Path | None = None) -> str:
         """Get S3 object key from filename"""
+
         dt = datetime.now(tz=UTC)
         fields_dict = {
             "year": f"{dt.year:04}",
@@ -79,7 +87,7 @@ def get_s3_object_key(filename: Path | None = None, hostname: str | None = None)
             "hour": f"{dt.hour:02}",
             "minute": f"{dt.minute:02}",
             "second": f"{dt.second:02}",
-            **({"hostname": hostname} if hostname else {}),
+            **self.object_kvs,
             "uuid": str(uuid.uuid4()),
         }
         fields_list = [f"{k}={v}" for k, v in fields_dict.items() if v is not None]
@@ -93,23 +101,60 @@ def wait_for_stable_file(self, filename: Path, timeout: int = 60) -> bool:
         start_time = time.time()
         last_size = -1
 
-        while time.time() - start_time < timeout:
-            try:
-                current_size = filename.stat().st_size
-                if current_size == last_size:
-                    self.logger.debug("File is stable")
-                    return True
-                last_size = current_size
-            except FileNotFoundError:
-                return False
-            self.logger.debug("Waiting for file to become stable", filename=str(filename))
-            time.sleep(STABLE_FILE_DELAY_SECONDS)
-
-        self.logger.warning("File not stable, timeout reached")
+        with structlog.contextvars.bound_contextvars(filename=str(filename)):
+            while time.time() - start_time < timeout:
+                try:
+                    current_size = filename.stat().st_size
+                    if current_size == last_size:
+                        self.logger.debug("File is stable")
+                        return True
+                    last_size = current_size
+                except FileNotFoundError:
+                    return False
+                self.logger.debug("Waiting for file to become stable", filename=str(filename))
+                time.sleep(STABLE_FILE_DELAY_SECONDS)
+
+            self.logger.warning("File not stable, timeout reached")
 
         return False
 
 
+def get_object_kvs(fields_str: str) -> dict[str, str]:
+    object_kvs: dict[str, str] = {}
+
+    for field in fields_str.split(","):
+        if "=" not in field:
+            raise ValueError(f"Invalid field format: '{field}' (expected 'key=value')")
+
+        k, v = field.split("=", 1)
+
+        k = k.strip()
+        v = v.strip()
+
+        if k in RESERVED_FIELD_NAMES:
+            raise ValueError(f"Field name '{k}' is reserved and cannot be used")
+
+        if not k:
+            raise ValueError("Field with empty key")
+        if not v:
+            raise ValueError(f"Field value for key '{k}' is empty")
+
+        if not FIELD_KV_RE.match(k):
+            raise ValueError(
+                f"Invalid characters in field key '{k}'"
+                + " (only letters, numbers, period, underscores and hyphens are allowed)"
+            )
+        if not FIELD_KV_RE.match(v):
+            raise ValueError(
+                f"Invalid characters in field value '{v}' for key '{k}'"
+                + " (only letters, numbers, period, underscores and hyphens are allowed)"
+            )
+
+        object_kvs[k] = v
+
+    return object_kvs
+
+
 def main():
 
     parser = argparse.ArgumentParser(description="S3 Pusher")
@@ -121,10 +166,10 @@ def main():
         help="S3 bucket name (S3PUSHER_BUCKET environment variable can also be used)",
     )
     parser.add_argument(
-        "--hostname",
+        "--fields",
         required=False,
         default=None,
-        help="Hostname to include in S3 object key (S3PUSHER_HOSTNAME environment variable can also be used)",
+        help="Fields to include in S3 object key (S3PUSHER_FIELDS environment variable can also be used)",
     )
     parser.add_argument("--log-json", action="store_true", help="Log in JSON format")
     parser.add_argument("--debug", action="store_true", help="Enable debugging")
@@ -151,20 +196,35 @@ def main():
     else:
         logger.warning("No bucket configured (file upload will be skipped)")
 
-    if hostname := args.hostname or os.getenv("S3PUSHER_HOSTNAME"):
-        logger.info("Hostname configured", hostname=hostname)
-    else:
-        logger.info("No hostname configured")
+    object_kvs: dict[str, str] = {}
+
+    try:
+        if fields_str := (args.fields or os.getenv("S3PUSHER_FIELDS")):
+            object_kvs = get_object_kvs(fields_str)
+
+        # Add hostname to object_kvs if configured via environment variable (for backwards compatibility)
+        if hostname := os.getenv("S3PUSHER_HOSTNAME"):
+            if FIELD_KV_RE.match(hostname):
+                if "hostname" not in object_kvs:
+                    object_kvs["hostname"] = hostname
+            else:
+                raise ValueError(f"Invalid hostname value '{hostname}'")
+    except ValueError as exc:
+        parser.error(str(exc))
+
+    if object_kvs:
+        logger.info("Configured with object fields", object_kvs=object_kvs)
 
     logger.info("Watching directories for changes", directories=args.directory)
 
-    event_handler = ThePusher(bucket=bucket, hostname=hostname)
+    event_handler = ThePusher(bucket=bucket, object_kvs=object_kvs)
 
     observer = Observer()
 
     for directory in args.directory:
         event_handler.upload_directory_to_s3(Path(directory))
         observer.schedule(event_handler, directory)
+
     observer.start()
 
     try: