work in progress

Author: rustyconover

src/query_farm_server_base/flight_inventory.py

@@ -13,9 +13,6 @@
     AirportSerializedSchema,
 )
 
-# SCHEMA_BASE_URL = "https://schemas.beta.database.flights"
-# SCHEMA_BUCKET_NAME = "schemas.beta.database.flights"
-
 # This is the level of ZStandard compression to use for individual FlightInfo
 # objects, since the schemas are pretty small, we can use a lower compression
 # preferring fast decompression.
@@ -72,6 +69,14 @@ def serialize(self) -> bytes:
 FlightInventoryWithMetadata = tuple[flight.FlightInfo, FlightSchemaMetadata]
 
 
+@dataclass
+class UploadParameters:
+    s3_client: S3Client
+    base_url: str
+    bucket_name: str
+    bucket_prefix: str | None = None
+
+
 def upload_and_generate_schema_list(
     *,
     flight_service_name: str,
@@ -80,10 +85,11 @@ def upload_and_generate_schema_list(
     skip_upload: bool,
     catalog_version: int,
     catalog_version_fixed: bool,
-    schema_base_url: str,
-    schema_bucket_name: str,
-    s3_client: S3Client,
-    s3_bucket_prefix: str | None = None,
+    upload_parameters: UploadParameters,
+    # schema_base_url: str,
+    # schema_bucket_name: str,
+    # s3_client: S3Client,
+    # s3_bucket_prefix: str | None = None,
     serialize_inline: bool = False,
 ) -> AirportSerializedCatalogRoot:
     serialized_schema_data: list[AirportSerializedSchema] = []
@@ -99,7 +105,12 @@ def upload_and_generate_schema_list(
     #
     # I think we can suffer with this problem for a bit longer.
     #
-    s3_bucket_prefix = s3_bucket_prefix.rstrip("/") + "/" if s3_bucket_prefix else ""
+    if upload_parameters and upload_parameters.bucket_prefix:
+        upload_parameters.bucket_prefix = (
+            upload_parameters.bucket_prefix.rstrip("/") + "/"
+            if upload_parameters.bucket_prefix
+            else ""
+        )
 
     for catalog_name, schema_names in flight_inventory.items():
         for schema_name, schema_items in schema_names.items():
@@ -110,16 +121,14 @@ def upload_and_generate_schema_list(
             assert packed_flight_info
 
             uploaded_schema_contents = schema_uploader.upload(
-                s3_client=s3_client,
+                s3_client=upload_parameters.s3_client,
                 data=packed_flight_info,
                 compression_level=SCHEMA_COMPRESSION_LEVEL,
-                key_prefix=f"{s3_bucket_prefix}schemas/{flight_service_name}/{catalog_name}",
-                bucket=schema_bucket_name,
+                key_prefix=f"{upload_parameters.bucket_prefix}schemas/{flight_service_name}/{catalog_name}",
+                bucket=upload_parameters.bucket_name,
                 skip_upload=skip_upload or serialize_inline,
             )
 
-            schema_path = f"{schema_base_url}/{uploaded_schema_contents.s3_path}"
-
             assert uploaded_schema_contents.compressed_data
 
             all_schema_flights_serialized.append(
@@ -136,7 +145,9 @@ def upload_and_generate_schema_list(
                     if schema_name in schema_details
                     else "",
                     contents=AirportSerializedContentsWithSHA256Hash(
-                        url=schema_path if not serialize_inline else None,
+                        url=f"{upload_parameters.base_url}/{uploaded_schema_contents.s3_path}"
+                        if not serialize_inline
+                        else None,
                         sha256=uploaded_schema_contents.sha256_hash,
                         serialized=None,
                     ),
@@ -148,20 +159,21 @@ def upload_and_generate_schema_list(
     assert all_packed
 
     all_schema_contents_upload = schema_uploader.upload(
-        s3_client=s3_client,
+        s3_client=upload_parameters.s3_client,
         data=all_packed,
-        key_prefix=f"{s3_bucket_prefix}schemas/{flight_service_name}",
-        bucket=schema_bucket_name,
+        key_prefix=f"{upload_parameters.bucket_prefix}schemas/{flight_service_name}",
+        bucket=upload_parameters.bucket_name,
         compression_level=None,  # Don't compress since all contained schemas are compressed
         skip_upload=skip_upload or serialize_inline,
     )
-    all_schema_path = f"{schema_base_url}/{all_schema_contents_upload.s3_path}"
 
     return AirportSerializedCatalogRoot(
         schemas=serialized_schema_data,
         contents=AirportSerializedContentsWithSHA256Hash(
             sha256=all_schema_contents_upload.sha256_hash,
-            url=all_schema_path if not serialize_inline else None,
+            url=f"{upload_parameters.base_url}/{all_schema_contents_upload.s3_path}"
+            if not serialize_inline
+            else None,
             serialized=all_schema_contents_upload.compressed_data if serialize_inline else None,
         ),
         version_info=server.GetCatalogVersionResult(

src/query_farm_server_base/schema_uploader.py

@@ -25,8 +25,12 @@ def _compress_and_prefix_with_length(data: bytes, compression_level: int) -> byt
     return result
 
 
+def _hash_value(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
 def _build_sha256_key_name_and_hash(key_prefix: str, data: bytes) -> tuple[str, str]:
-    sha256_hash = hashlib.sha256(data).hexdigest()
+    sha256_hash = _hash_value(data)
     return (
         f"{key_prefix}/{hex_to_url_safe_characters(sha256_hash)}",
         sha256_hash,
@@ -36,7 +40,7 @@ def _build_sha256_key_name_and_hash(key_prefix: str, data: bytes) -> tuple[str,
 def upload(
     *,
     compression_level: int | None,
-    s3_client: S3Client,
+    s3_client: S3Client | None,
     data: bytes,
     key_prefix: str,
     bucket: str,
@@ -62,6 +66,7 @@ def upload(
     s3_path, sha256_hash = _build_sha256_key_name_and_hash(key_prefix, data)
 
     if not skip_upload:
+        assert s3_client, "S3 client must be provided if not skipping upload"
         s3_client.put_object(
             Body=data,
             Bucket=bucket,