duckdb / duckdb-wasm

WebAssembly version of DuckDB
https://shell.duckdb.org
MIT License
1.02k stars 110 forks source link

Support parquet range requests for AWS S3 pre-signed URLs #1702

Closed carstonhernke closed 2 months ago

carstonhernke commented 2 months ago

Resolves #1367

AWS S3 allows for creating pre-signed URLs. However, these pre-signed URLs only allow a single HTTP operation/verb. So a single URL cannot support both HEAD and GET requests.

This causes problems for the current logic which detects whether range requests are supported, because it looks for a successful HEAD request for the Range bytes=0-. For an AWS pre-signed URL, this request will fail. However, AWS pre-signed URLs do support range requests.

This PR addresses the issue by adding additional logic which checks whether a GET request with a Range 0-0 header returns 206. It uses the Content-Range response header to determine the remote file size.

NOTE: the remote file must have CORS settings that expose the Content-Range header.

carlopi commented 2 months ago

Hi @carstonhernke, thanks for the PR, I think this is better unified with the logic within the if (file.allowFullHttpReads) bracket, since when we issue a GET request we might still get the whole file back, depending on the server.

Or it can special cased for S3 specifically, possibly via a "HasContentRange()" property that can be false by default but true for S3.

My preference would be for integrating the code you propose with the code below, something like:

-                        if ((contentLength !== null) && (+contentLength > 1)) {
+                        if (((contentLength !== null) && (+contentLength > 1)) || file.HasContentRange()) {
                            // 2. Send a dummy GET range request querying the first byte of the file
                            //          -> good IFF status is 206 and contentLenght2 is 1
                            //          -> otherwise, iff 200 and contentLenght2 == contentLenght
                            //                 we just downloaded the file, save it and move further
                            const xhr = new XMLHttpRequest();
                            if (file.dataProtocol == DuckDBDataProtocol.S3) {
                                xhr.open('GET', getHTTPUrl(file.s3Config, file.dataUrl!), false);
                                addS3Headers(xhr, file.s3Config, file.dataUrl!, 'GET');
                            } else {
                                xhr.open('GET', file.dataUrl!, false);
                            }
                            xhr.responseType = 'arraybuffer';
                            xhr.setRequestHeader('Range', `bytes=0-0`);
                            xhr.send(null);
                            const contentLength2 = xhr.getResponseHeader('Content-Length');
+                           const contentRange = xhr.getResponseHeader('Content-Range')?.split('/')[1];
-                            if (xhr.status == 206 && contentLength2 !== null && +contentLength2 == 1) {
+                           if (xhr.status == 206 && contentLength2 !== null && +contentLength2 == 1 && contentLength !== null) {
                                const result = mod._malloc(2 * 8);
                                mod.HEAPF64[(result >> 3) + 0] = +contentLength;
                                mod.HEAPF64[(result >> 3) + 1] = 0;
                                return result;
                            }
+.                         // This can also be before, unsure who should get precedence
+.                         // Actually this has likely to be before
+                           if (xhr.status == 206 && contentRange !== null) {
+                               const result = mod._malloc(2 * 8);
+                               mod.HEAPF64[(result >> 3) + 0] = +contentRange;
+                               mod.HEAPF64[(result >> 3) + 1] = 0;
+                               return result;
+                           }
                            if (xhr.status == 200 && contentLength2 !== null && +contentLength2 == +contentLength) {
                                console.warn(`fall back to full HTTP read for: ${file.dataUrl}`);
                                const data = mod._malloc(xhr.response.byteLength);
                                const src = new Uint8Array(xhr.response, 0, xhr.response.byteLength);
                                mod.HEAPU8.set(src, data);
                                const result = mod._malloc(2 * 8);
                                mod.HEAPF64[(result >> 3) + 0] = xhr.response.byteLength;
                                mod.HEAPF64[(result >> 3) + 1] = data;
                                return result;
                            }
                        }

mostly to avoid throwing all informations away when we happen to have just downloaded the whole file (that given we don't know it might as well happen) and to keep the logic together / more well tested.

carlopi commented 2 months ago

I filed a similar PR at https://github.com/duckdb/duckdb-wasm/pull/1717, mostly so it should be easier to check the diff and since it was simpler to just go ahead.

Could you possibly take a look and test with presigned URLs? I think moving in that direction is better, but very open to other ideas.

carstonhernke commented 2 months ago

Sounds good. I will leave some comments in https://github.com/duckdb/duckdb-wasm/pull/1717