X-Git-Url: https://gerrit.fd.io/r/gitweb?a=blobdiff_plain;f=csit.infra.etl%2Ftrending_soak.py;h=40da521884c39b39ed387db54bde132c69f29261;hb=HEAD;hp=e54cf9f18afb0f726a36199e31e09fe08fd4ae6c;hpb=bff439b69ee71b654b1da92564ff62de7327fe71;p=csit.git

diff --git a/csit.infra.etl/trending_soak.py b/csit.infra.etl/trending_soak.py
index e54cf9f18a..086a1943ee 100644
--- a/csit.infra.etl/trending_soak.py
+++ b/csit.infra.etl/trending_soak.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:
@@ -29,8 +29,8 @@ from pyspark.sql.functions import col, lit, regexp_replace
 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="info.json.gz"
 IGNORE_SUFFIX=[
@@ -144,7 +144,6 @@ paths = wr.s3.list_objects(
 filtered_paths = [path for path in paths if "daily" in path or "weekly" in path]
 
 out_sdf = process_json_to_dataframe("soak", filtered_paths)
-out_sdf.show(truncate=False)
 out_sdf.printSchema()
 out_sdf = out_sdf \
     .withColumn("year", lit(datetime.now().year)) \
@@ -152,6 +151,15 @@ out_sdf = out_sdf \
     .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(),
@@ -161,11 +169,7 @@ try:
         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