'''
Add the following credentials to your Beam secrets store:

AWS_ACCESS_KEY
AWS_SECRET_ACCESS_KEY
AWS_REGION_NAME
'''

class S3Client:
    def __init__(self, bucket_name=None, key=None, file_path=None):
        self.bucket_name = bucket_name
        self.key = key
        self.file_path = file_path
        self.session = boto3.Session(
            region_name=os.environ["AWS_REGION_NAME"],
            aws_access_key_id=os.environ["AWS_ACCESS_KEY"],
            aws_secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"],
        )
        self.s3 = self.session.client("s3")

    def upload_object(self, file_path):
        with open(file_path, "rb") as data:
            return self.s3.upload_fileobj(data, "mybucket", "mykey")

    def download_object(self, file_path, bucket_name, key):
        object = self.s3.get_object(Bucket=bucket_name, Key=key)

        with open(file_path, "wb") as f:
            self.s3.download_fileobj(bucket_name, key, f)

        return object

Was this page helpful?