f1-lake/sender.py at main · TeoMeWhy/f1-lake · GitHub

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
# %%

import argparse
import os

import dotenv
import boto3
from tqdm import tqdm

dotenv.load_dotenv()

AWS_KEY = os.getenv("AWS_KEY")
AWS_SECRET_KEY = os.getenv("AWS_SECRET_KEY")

class Sender:

    def __init__(self, bucket_name, bucket_folder):
        self.bucket_name = bucket_name
        self.bucket_folder = bucket_folder

        self.s3 = boto3.client("s3",
                        aws_access_key_id=AWS_KEY,
                        aws_secret_access_key=AWS_SECRET_KEY,
                        region_name="us-east-1")


    def process_file(self, filename):

        file = filename.split("/")[-1]
        bucket_path = os.path.join(self.bucket_folder, file)

        try:
            self.s3.upload_file(filename,
                                self.bucket_name,
                                bucket_path)

        except Exception as err:
            print(err)
            return False

        os.remove(filename)
        return True


    def process_folder(self, folder):
        files = [i for i in os.listdir(folder) if i.endswith(".parquet")]
        for f in tqdm(files):
            self.process_file(os.path.join(folder, f))


# %%

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--bucket", type=str)
    parser.add_argument("--bucket_path",default="f1/results", type=str)
    parser.add_argument("--folder", default="data", type=str)
    args = parser.parse_args()

    if args.bucket:
        send = Sender(args.bucket, args.bucket_path)
        send.process_folder(args.folder)

    else:
        print("sem bucket definido")