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

diff --git a/csit.infra.etl/trending_ndrpdr.py b/csit.infra.etl/trending_ndrpdr.py
index d3c51ba757..f139f099f7 100644
--- a/csit.infra.etl/trending_ndrpdr.py
+++ b/csit.infra.etl/trending_ndrpdr.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=[
@@ -152,6 +152,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 +170,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