fix(etl): Typo"

[csit.git] / csit.infra.etl / stats.py

diff --git a/csit.infra.etl/stats.py b/csit.infra.etl/stats.py
index ab8bcaf..d7b4c23 100644 (file)
--- a/csit.infra.etl/stats.py
+++ b/csit.infra.etl/stats.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright (c) 2022 Cisco and/or its affiliates.
+# Copyright (c) 2024 Cisco and/or its affiliates.
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at:
@@ -28,8 +28,9 @@ from pyspark.context import SparkContext
 from pyspark.sql.functions import lit
 from pyspark.sql.types import StructType
 
-S3_LOGS_BUCKET="fdio-logs-s3-cloudfront-index"
-S3_DOCS_BUCKET="fdio-docs-s3-cloudfront-index"
+
+S3_LOGS_BUCKET=environ.get("S3_LOGS_BUCKET", "fdio-logs-s3-cloudfront-index")
+S3_DOCS_BUCKET=environ.get("S3_DOCS_BUCKET", "fdio-docs-s3-cloudfront-index")
 PATH=f"s3://{S3_LOGS_BUCKET}/vex-yul-rot-jenkins-1/csit-*-perf-*"
 SUFFIX="suite.info.json.gz"
 IGNORE_SUFFIX=[]
@@ -106,7 +107,6 @@ paths = wr.s3.list_objects(
 
 for schema_name in ["sra"]:
     out_sdf = process_json_to_dataframe(schema_name, paths)
-    out_sdf.show(truncate=False)
     out_sdf.printSchema()
     out_sdf = out_sdf \
         .withColumn("year", lit(datetime.now().year)) \
@@ -114,6 +114,15 @@ for schema_name in ["sra"]:
         .withColumn("day", lit(datetime.now().day)) \
         .repartition(1)
 
+    try:
+        boto3_session = session.Session(
+            aws_access_key_id=environ["OUT_AWS_ACCESS_KEY_ID"],
+            aws_secret_access_key=environ["OUT_AWS_SECRET_ACCESS_KEY"],
+            region_name=environ["OUT_AWS_DEFAULT_REGION"]
+        )
+    except KeyError:
+        boto3_session = session.Session()
+
     try:
         wr.s3.to_parquet(
             df=out_sdf.toPandas(),
@@ -123,11 +132,7 @@ for schema_name in ["sra"]:
             compression="snappy",
             use_threads=True,
             mode="overwrite_partitions",
-            boto3_session=session.Session(
-                aws_access_key_id=environ["OUT_AWS_ACCESS_KEY_ID"],
-                aws_secret_access_key=environ["OUT_AWS_SECRET_ACCESS_KEY"],
-                region_name=environ["OUT_AWS_DEFAULT_REGION"]
-            )
+            boto3_session=boto3_session
         )
     except EmptyDataFrame:
         pass