Make S3 request failures more verbose and add token expiry test

Author: trevorbonas

tools/python/liveanalytics_influxdb3_migration_plugin/liveanalytics_migration_client/liveanalytics_influxdb3_migration_client.py

@@ -11,16 +11,17 @@
 """
 
 import argparse
-import sys
+import json
+import logging
 import os
+import sys
 import time
+from datetime import datetime, timezone
+
 import boto3
-from botocore.exceptions import ClientError, BotoCoreError
-from influxdb_client_3 import InfluxDBClient3, WritePrecision
-import json
-import logging
 import requests
-from datetime import datetime, timezone
+from botocore.exceptions import BotoCoreError, ClientError
+from influxdb_client_3 import InfluxDBClient3, WritePrecision
 
 MIGRATION_METADATA_TABLE: str = "liveanalytics_migration_metadata"
 TRIGGER_NAME: str = "migration_trigger"
@@ -38,6 +39,7 @@ def __init__(
         timeout_seconds: int = 360,  # 6 minutes.
         region: str = "us-west-2",
         max_parquet_files: int | None = None,
+        presigned_url_expiry_seconds: int | None = None,
     ) -> None:
         """
         Initialize
@@ -49,6 +51,7 @@ def __init__(
             timeout_seconds (int): The number of seconds to wait for each migration request.
             region (str): The AWS Region to use.
             max_parquet_files (int | None): The maximum number of Parquet files to migrate.
+            presigned_url_expiry_seconds (int | None): The expiry time for all presigned URLs.
 
         Returns:
             None
@@ -59,6 +62,7 @@ def __init__(
         self.timeout_seconds: int = timeout_seconds
         self.region: str = region
         self.max_parquet_files: int | None = max_parquet_files
+        self.presigned_url_expiry_seconds: int | None = presigned_url_expiry_seconds
 
         logging.basicConfig(
             level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
@@ -582,7 +586,7 @@ def get_s3_objects_list(
         return list(files_to_migrate)
 
     def generate_metadata(
-        self, parquet_file_names: list[str], expiration: int = 604_800
+        self, parquet_file_names: list[str], expiration: int | None = 604_800
     ) -> dict[str, dict[str, str]]:
         """
         Generates metadata, including presigned URLs for S3 objects.
@@ -675,7 +679,7 @@ def setup_influxdb_metadata(self) -> bool:
                     )
                     return False
             except requests.exceptions.RequestException as e:
-                self.error(f"Error creating database: ", str(e))
+                self.error("Error creating database: ", str(e))
                 return False
 
             self.delete_metadata_table(silence_errors=True)
@@ -707,12 +711,12 @@ def setup_influxdb_metadata(self) -> bool:
                     )
                     return False
             except requests.exceptions.RequestException as e:
-                self.error(f"Error creating table: ", str(e))
+                self.error("Error creating table: ", str(e))
                 return False
             return True
 
         except Exception as e:
-            self.error(f"Failed to setup InfluxDB metadata: ", str(e))
+            self.error("Failed to setup InfluxDB metadata: ", str(e))
             sys.exit(1)
 
     def write_metadata_to_influxdb(self, metadata):
@@ -783,7 +787,7 @@ def write_metadata_to_influxdb(self, metadata):
             )
 
         except Exception as e:
-            self.error(f"Failed to write metadata to InfluxDB: ", str(e))
+            self.error("Failed to write metadata to InfluxDB: ", str(e))
             sys.exit(1)
 
     def bulk_invoke_http_trigger(self, metadata):
@@ -1012,7 +1016,7 @@ def create_processing_engine_trigger(self):
             None
         """
         try:
-            self.info(f"Creating processing engine HTTP trigger")
+            self.info("Creating processing engine HTTP trigger")
 
             headers = {
                 "Authorization": f"Bearer {self.influx_token}",
@@ -1081,7 +1085,8 @@ def run(self):
         parquet_file_names: list[str] = self.get_s3_objects_list()
 
         metadata: dict[str, dict[str, str]] = self.generate_metadata(
-            parquet_file_names=parquet_file_names
+            parquet_file_names=parquet_file_names,
+            expiration=self.presigned_url_expiry_seconds,
         )
 
         # Setup InfluxDB metadata.
@@ -1160,7 +1165,7 @@ def delete_unloaded_data(self):
                         Key=marker["Key"],
                         VersionId=marker["VersionId"],
                     )
-                except Exception as e:
+                except Exception:
                     self.error(
                         f"Failed to delete deletion marker for object {marker['Key']} with version ID {marker['VersionId']}"
                     )
@@ -1417,6 +1422,12 @@ def main(input_args) -> int:
         type=int,
         help="Optional. The maximum number of Parquet files to migrate. Default behaviour is to migrate all Parquet files.",
     )
+    parser.add_argument(
+        "--presigned-url-expiry-seconds",
+        required=False,
+        type=int,
+        help="Optional. The number of seconds that all presigned URLs will be valid for. Default is 604800 (7 days, max).",
+    )
 
     args = parser.parse_args(input_args)
 
@@ -1426,6 +1437,7 @@ def main(input_args) -> int:
     timeout_seconds: int = args.timeout_seconds
     region: str = args.region
     max_parquet_files: int | None = args.max_parquet_files
+    presigned_url_expiry_seconds: int | None = args.presigned_url_expiry_seconds
 
     required_env_vars = [
         "INFLUXDB3_HOST_URL",
@@ -1448,6 +1460,7 @@ def main(input_args) -> int:
         timeout_seconds=timeout_seconds,
         region=region,
         max_parquet_files=max_parquet_files,
+        presigned_url_expiry_seconds=presigned_url_expiry_seconds,
     )
 
     try:

tools/python/liveanalytics_influxdb3_migration_plugin/liveanalytics_migration_client/tests/integration/test_liveanalytics_influxdb3_migration_client.py

@@ -1,25 +1,27 @@
-from concurrent.futures import ThreadPoolExecutor
-from datetime import datetime, timezone
-import os
-from pathlib import Path
 import json
-import re
-import time
-from typing import Any
-import unittest
+import logging
+import os
 import random
+import re
 import string
 import sys
-import logging
+import time
+import unittest
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
 
 import boto3
-from botocore.exceptions import ClientError
 import pandas
+import pytest
 import requests
+from botocore.exceptions import ClientError
 from influxdb_client_3 import InfluxDBClient3
-from testcontainers.core.wait_strategies import LogMessageWaitStrategy
 from testcontainers.core.container import DockerContainer, ExecResult
-import pytest
+from testcontainers.core.wait_strategies import LogMessageWaitStrategy
+
+logger = logging.getLogger()
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../")))
 
@@ -352,7 +354,7 @@ def delete_s3_bucket(self, bucket_name: str):
                         Key=marker["Key"],
                         VersionId=marker["VersionId"],
                     )
-                except Exception as e:
+                except Exception:
                     logging.error(
                         f"Failed to delete deletion marker for object {marker['Key']} with version ID {marker['VersionId']}"
                     )
@@ -1023,6 +1025,54 @@ def test_migration_different_database_names(self):
             influxdb_v3_table_count,
         )
 
+    def test_migration_expired_presigned_urls(self):
+        """
+        Tests migrating a single record where all presigned URLs are expired.
+        """
+        current_time: pandas.Timestamp = pandas.Timestamp.now()
+
+        start_time = current_time - pandas.Timedelta(days=30)
+        assert isinstance(start_time, pandas.Timestamp)
+        end_time = start_time + pandas.Timedelta(days=1)
+        assert isinstance(end_time, pandas.Timestamp)
+
+        dimensions = [
+            {"Name": "hostname", "Value": "hostname1", "DimensionValueType": "VARCHAR"},
+            {"Name": "region", "Value": "us-west-2", "DimensionValueType": "VARCHAR"},
+        ]
+
+        record = {
+            "Dimensions": dimensions,
+            "MeasureName": "cpu_utilization",
+            "MeasureValue": "13.5",
+            "MeasureValueType": "DOUBLE",
+            "Time": str(start_time.value),
+            "TimeUnit": "NANOSECONDS",
+        }
+        self.put_records([record])
+
+        with self.assertLogs(
+            "liveanalytics_influxdb3_migration_client", level="INFO"
+        ) as captured_logs:
+            return_code = liveanalytics_influxdb3_migration_client.main(
+                [
+                    "--live-analytics-database-name",
+                    self.la_database_name,
+                    "--s3-bucket-name",
+                    self.s3_bucket_name,
+                    "--presigned-url-expiry-seconds",
+                    "0",
+                ]
+            )
+
+            self.assertEqual(return_code, 1)
+
+            expected_error = "403 Client Error: Forbidden for url: *****"
+            self.assertTrue(
+                any(expected_error in log for log in captured_logs.output),
+                f"Expected error not found in logs: {expected_error}",
+            )
+
 
 if __name__ == "__main__":
     unittest.main()

tools/python/liveanalytics_influxdb3_migration_plugin/liveanalytics_migration_client/tests/unit/test_liveanalytics_migration_plugin.py

@@ -1,6 +1,6 @@
-import unittest
 import os
 import sys
+import unittest
 
 sys.path.insert(
     0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))