diff options
Diffstat (limited to 'csit.infra.etl/trending_ndrpdr.py')
| -rw-r--r-- | csit.infra.etl/trending_ndrpdr.py | 22 |
1 files changed, 14 insertions, 8 deletions
diff --git a/csit.infra.etl/trending_ndrpdr.py b/csit.infra.etl/trending_ndrpdr.py index e35d27b0bf..96582f5928 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) 2023 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=[ @@ -153,6 +153,16 @@ out_sdf = out_sdf \ .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(), path=f"s3://{S3_DOCS_BUCKET}/csit/parquet/trending", @@ -161,11 +171,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 |