File size: 5,391 Bytes
4a190c2
 
 
 
b4066c5
4a190c2
 
 
 
b4066c5
4a190c2
 
a73f005
4a190c2
a73f005
4a190c2
 
96d818b
4a190c2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b4066c5
 
 
 
 
4a190c2
 
 
 
 
b4066c5
 
4a190c2
b4066c5
 
4a190c2
b4066c5
 
 
 
 
4a190c2
b4066c5
 
4a190c2
b4066c5
 
 
 
 
 
4a190c2
 
 
 
 
b4066c5
 
4a190c2
b4066c5
4a190c2
b4066c5
 
 
 
4a190c2
b4066c5
4a190c2
b4066c5
4a190c2
b4066c5
4a190c2
b4066c5
 
 
4a190c2
b4066c5
 
4a190c2
b4066c5
 
 
 
 
 
4a190c2
 
 
 
 
 
 
 
 
 
 
 
 
 
b4066c5
4a190c2
b4066c5
4a190c2
b4066c5
4a190c2
b4066c5
 
4a190c2
b4066c5
 
 
 
4a190c2
b4066c5
 
 
 
 
 
 
 
 
 
4a190c2
b4066c5
 
4a190c2
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
from typing import Type, List
import pandas as pd
import boto3
import os
from chatfuncs.config import AWS_REGION, RUN_AWS_FUNCTIONS, QA_CHATBOT_BUCKET

PandasDataFrame = Type[pd.DataFrame]

# Get AWS credentials if required
bucket_name=QA_CHATBOT_BUCKET


if RUN_AWS_FUNCTIONS == "1":
    try:
        bucket_name = os.environ['']
        session = boto3.Session() # profile_name="default"
    except Exception as e:
        print("Failed to start boto3 session due to:", e)

    def get_assumed_role_info():
        sts_endpoint = 'https://sts.' + AWS_REGION + '.amazonaws.com'
        sts = boto3.client('sts', region_name=AWS_REGION, endpoint_url=sts_endpoint)
        response = sts.get_caller_identity()

        # Extract ARN of the assumed role
        assumed_role_arn = response['Arn']
        
        # Extract the name of the assumed role from the ARN
        assumed_role_name = assumed_role_arn.split('/')[-1]
        
        return assumed_role_arn, assumed_role_name

    try:
        assumed_role_arn, assumed_role_name = get_assumed_role_info()

        print("Assumed Role ARN:", assumed_role_arn)
        print("Assumed Role Name:", assumed_role_name)

    except Exception as e:
        
        print(e)

# Download direct from S3 - requires login credentials
def download_file_from_s3(bucket_name, key, local_file_path):

    if RUN_AWS_FUNCTIONS == "1":

        s3 = boto3.client('s3')
        s3.download_file(bucket_name, key, local_file_path)
        print(f"File downloaded from S3: s3://{bucket_name}/{key} to {local_file_path}")
                         
def download_folder_from_s3(bucket_name, s3_folder, local_folder):
    """
    Download all files from an S3 folder to a local folder.
    """
    if RUN_AWS_FUNCTIONS == "1":
        s3 = boto3.client('s3')

        # List objects in the specified S3 folder
        response = s3.list_objects_v2(Bucket=bucket_name, Prefix=s3_folder)

        # Download each object
        for obj in response.get('Contents', []):
            # Extract object key and construct local file path
            object_key = obj['Key']
            local_file_path = os.path.join(local_folder, os.path.relpath(object_key, s3_folder))

            # Create directories if necessary
            os.makedirs(os.path.dirname(local_file_path), exist_ok=True)

            # Download the object
            try:
                s3.download_file(bucket_name, object_key, local_file_path)
                print(f"Downloaded 's3://{bucket_name}/{object_key}' to '{local_file_path}'")
            except Exception as e:
                print(f"Error downloading 's3://{bucket_name}/{object_key}':", e)

def download_files_from_s3(bucket_name, s3_folder, local_folder, filenames):
    """
    Download specific files from an S3 folder to a local folder.
    """
    if RUN_AWS_FUNCTIONS == "1":
        s3 = boto3.client('s3')

        print("Trying to download file: ", filenames)

        if filenames == '*':
            # List all objects in the S3 folder
            print("Trying to download all files in AWS folder: ", s3_folder)
            response = s3.list_objects_v2(Bucket=bucket_name, Prefix=s3_folder)

            print("Found files in AWS folder: ", response.get('Contents', []))

            filenames = [obj['Key'].split('/')[-1] for obj in response.get('Contents', [])]

            print("Found filenames in AWS folder: ", filenames)

        for filename in filenames:
            object_key = os.path.join(s3_folder, filename)
            local_file_path = os.path.join(local_folder, filename)

            # Create directories if necessary
            os.makedirs(os.path.dirname(local_file_path), exist_ok=True)

            # Download the object
            try:
                s3.download_file(bucket_name, object_key, local_file_path)
                print(f"Downloaded 's3://{bucket_name}/{object_key}' to '{local_file_path}'")
            except Exception as e:
                print(f"Error downloading 's3://{bucket_name}/{object_key}':", e)

def upload_file_to_s3(local_file_paths:List[str], s3_key:str, s3_bucket:str=bucket_name):
    """
    Uploads a file from local machine to Amazon S3.

    Args:
    - local_file_path: Local file path(s) of the file(s) to upload.
    - s3_key: Key (path) to the file in the S3 bucket.
    - s3_bucket: Name of the S3 bucket.

    Returns:
    - Message as variable/printed to console
    """
    final_out_message = []
    final_out_message_str = ""

    if RUN_AWS_FUNCTIONS == "1":

        s3_client = boto3.client('s3')

        if isinstance(local_file_paths, str):
            local_file_paths = [local_file_paths]

        for file in local_file_paths:
            try:
                # Get file name off file path
                file_name = os.path.basename(file)

                s3_key_full = s3_key + file_name
                print("S3 key: ", s3_key_full)

                s3_client.upload_file(file, s3_bucket, s3_key_full)
                out_message = "File " + file_name + " uploaded successfully!"
                print(out_message)
            
            except Exception as e:
                out_message = f"Error uploading file(s): {e}"
                print(out_message)

            final_out_message.append(out_message)
            final_out_message_str = '\n'.join(final_out_message)

    return final_out_message_str