File size: 5,831 Bytes
4384839
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
// Note: To use this module, you'll need to install the AWS SDK:
// npm install @aws-sdk/client-s3 @aws-sdk/lib-storage

import { S3Client, HeadBucketCommand } from "@aws-sdk/client-s3";
import { Upload } from "@aws-sdk/lib-storage";

// Define ContentSource type locally to avoid HuggingFace dependency
type ContentSource = Blob | ArrayBuffer | Uint8Array | string;
type FileArray = Array<URL | File | { path: string; content: ContentSource }>;

/**
 * Uploads a leRobot dataset to Amazon S3
 */
export class LeRobotS3Uploader extends EventTarget {
    private _bucketName: string;
    private _region: string;
    private _uploaded: boolean;
    private _bucket_exists: boolean;
    private _s3Client: S3Client | null;
    
    constructor(bucketName: string, region: string = "us-east-1") {
        super();
        this._bucketName = bucketName;
        this._region = region;
        this._uploaded = false;
        this._bucket_exists = false;
        this._s3Client = null;
    }
    
    /**
     * Returns whether the bucket has been successfully checked/created
     */
    get bucketExists(): boolean {
        return this._bucket_exists;
    }

    get uploaded(): boolean {
        return this._uploaded;
    }

    /**
     * Initialize the S3 client with credentials
     * 
     * @param accessKeyId AWS access key ID
     * @param secretAccessKey AWS secret access key
     */
    initializeClient(accessKeyId: string, secretAccessKey: string): void {
        this._s3Client = new S3Client({
            region: this._region,
            credentials: {
                accessKeyId,
                secretAccessKey
            }
        });
    }

    /**
     * Checks if the bucket exists and uploads files to it
     * 
     * @param files The files to upload
     * @param accessKeyId AWS access key ID
     * @param secretAccessKey AWS secret access key
     * @param prefix Optional prefix (folder) to upload files to within the bucket
     * @param referenceId The reference id for the upload, to track it (optional)
     */
    async checkBucketAndUploadFiles(
        files: FileArray, 
        accessKeyId: string, 
        secretAccessKey: string, 
        prefix: string = "",
        referenceId: string = ""
    ): Promise<void> {
        // Initialize the client if not already done
        if (!this._s3Client) {
            this.initializeClient(accessKeyId, secretAccessKey);
        }

        // Check if bucket exists
        try {
            await this._s3Client!.send(new HeadBucketCommand({ Bucket: this._bucketName }));
            this._bucket_exists = true;
            this.dispatchEvent(new CustomEvent("bucketExists", { 
                detail: { bucketName: this._bucketName } 
            }));
        } catch (error) {
            throw new Error(`Bucket ${this._bucketName} does not exist or you don't have permission to access it`);
        }

        // Upload files
        const uploadPromises: Promise<void>[] = [];
        for (const file of files) {
            uploadPromises.push(this.uploadFileWithProgress([file], prefix, referenceId));
        }

        await Promise.all(uploadPromises);
        this._uploaded = true;
    }

    /**
     * Uploads files to S3 with progress events
     * 
     * @param files The files to upload
     * @param prefix Optional prefix (folder) to upload files to within the bucket
     * @param referenceId The reference id for the upload, to track it (optional)
     */
    async uploadFileWithProgress(
        files: FileArray, 
        prefix: string = "", 
        referenceId: string = ""
    ): Promise<void> {
        if (!this._s3Client) {
            throw new Error("S3 client not initialized. Call initializeClient first.");
        }

        for (const file of files) {
            let key: string;
            let body: any;

            if (file instanceof URL) {
                const response = await fetch(file);
                body = await response.blob();
                key = `${prefix}${prefix ? '/' : ''}${file.pathname.split('/').pop()}`;
            } else if (file instanceof File) {
                body = file;
                key = `${prefix}${prefix ? '/' : ''}${file.name}`;
            } else {
                body = file.content;
                key = `${prefix}${prefix ? '/' : ''}${file.path}`;
            }

            const upload = new Upload({
                client: this._s3Client,
                params: {
                    Bucket: this._bucketName,
                    Key: key,
                    Body: body,
                },
            });

            // Set up progress tracking
            upload.on('httpUploadProgress', (progress) => {
                this.dispatchEvent(new CustomEvent("progress", { 
                    detail: {
                        progressEvent: progress,
                        bucketName: this._bucketName,
                        key,
                        referenceId
                    } 
                }));
            });

            await upload.done();
        }
    }

    /**
     * Generates a pre-signed URL for downloading a file from S3
     * 
     * @param _key The key (path) of the file in the S3 bucket
     * @param _expiresIn The number of seconds until the URL expires (default: 3600 = 1 hour)
     * @returns A pre-signed URL for downloading the file
     */
    async generatePresignedUrl(_key: string, _expiresIn: number = 3600): Promise<string> {
        if (!this._s3Client) {
            throw new Error("S3 client not initialized. Call initializeClient first.");
        }

        // Note: This requires the @aws-sdk/s3-request-presigner package
        // Implementation would go here
        throw new Error("generatePresignedUrl not implemented. Requires @aws-sdk/s3-request-presigner package.");
    }
}