|
16 | 16 | * |
17 | 17 | */ |
18 | 18 |
|
| 19 | +use std::{ |
| 20 | + collections::{BTreeMap, HashSet}, |
| 21 | + fmt::Display, |
| 22 | + path::Path, |
| 23 | + str::FromStr, |
| 24 | + sync::Arc, |
| 25 | + time::{Duration, Instant}, |
| 26 | +}; |
| 27 | + |
19 | 28 | use async_trait::async_trait; |
20 | 29 | use bytes::Bytes; |
21 | | -use datafusion::datasource::listing::ListingTableUrl; |
22 | | -use datafusion::datasource::object_store::{ |
23 | | - DefaultObjectStoreRegistry, ObjectStoreRegistry, ObjectStoreUrl, |
| 30 | +use datafusion::{ |
| 31 | + datasource::listing::ListingTableUrl, |
| 32 | + execution::{ |
| 33 | + object_store::{DefaultObjectStoreRegistry, ObjectStoreRegistry, ObjectStoreUrl}, |
| 34 | + runtime_env::RuntimeEnvBuilder, |
| 35 | + }, |
| 36 | +}; |
| 37 | +use futures::{stream::FuturesUnordered, StreamExt, TryStreamExt}; |
| 38 | +use object_store::{ |
| 39 | + aws::{AmazonS3, AmazonS3Builder, AmazonS3ConfigKey, Checksum}, |
| 40 | + limit::LimitStore, |
| 41 | + path::Path as StorePath, |
| 42 | + BackoffConfig, ClientOptions, ObjectStore, PutPayload, RetryConfig, |
24 | 43 | }; |
25 | | -use datafusion::execution::runtime_env::RuntimeEnvBuilder; |
26 | | -use futures::stream::FuturesUnordered; |
27 | | -use futures::{StreamExt, TryStreamExt}; |
28 | | -use object_store::aws::{AmazonS3, AmazonS3Builder, AmazonS3ConfigKey, Checksum}; |
29 | | -use object_store::limit::LimitStore; |
30 | | -use object_store::path::Path as StorePath; |
31 | | -use object_store::{BackoffConfig, ClientOptions, ObjectStore, PutPayload, RetryConfig}; |
32 | 44 | use relative_path::{RelativePath, RelativePathBuf}; |
33 | 45 | use tracing::{error, info}; |
34 | 46 |
|
35 | | -use std::collections::{BTreeMap, HashSet}; |
36 | | -use std::fmt::Display; |
37 | | -use std::iter::Iterator; |
38 | | -use std::path::Path as StdPath; |
39 | | -use std::str::FromStr; |
40 | | -use std::sync::Arc; |
41 | | -use std::time::{Duration, Instant}; |
| 47 | +use crate::{ |
| 48 | + handlers::http::users::USERS_ROOT_DIR, |
| 49 | + metrics::storage::{azureblob::REQUEST_RESPONSE_TIME, StorageMetrics}, |
| 50 | +}; |
42 | 51 |
|
43 | | -use super::metrics_layer::MetricLayer; |
44 | | -use super::object_storage::parseable_json_path; |
45 | 52 | use super::{ |
46 | | - to_object_store_path, LogStream, ObjectStorageProvider, SCHEMA_FILE_NAME, |
47 | | - STREAM_METADATA_FILE_NAME, STREAM_ROOT_DIRECTORY, |
| 53 | + metrics_layer::MetricLayer, object_storage::parseable_json_path, to_object_store_path, LogStream, ObjectStorage, ObjectStorageError, ObjectStorageProvider, CONNECT_TIMEOUT_SECS, PARSEABLE_ROOT_DIRECTORY, REQUEST_TIMEOUT_SECS, SCHEMA_FILE_NAME, STREAM_METADATA_FILE_NAME, STREAM_ROOT_DIRECTORY |
48 | 54 | }; |
49 | | -use crate::handlers::http::users::USERS_ROOT_DIR; |
50 | | -use crate::metrics::storage::{s3::REQUEST_RESPONSE_TIME, StorageMetrics}; |
51 | | -use crate::storage::{ObjectStorage, ObjectStorageError, PARSEABLE_ROOT_DIRECTORY}; |
52 | 55 |
|
53 | 56 | // in bytes |
54 | 57 | // const MULTIPART_UPLOAD_SIZE: usize = 1024 * 1024 * 100; |
55 | | -const CONNECT_TIMEOUT_SECS: u64 = 5; |
56 | | -const REQUEST_TIMEOUT_SECS: u64 = 300; |
57 | 58 | const AWS_CONTAINER_CREDENTIALS_RELATIVE_URI: &str = "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"; |
58 | 59 |
|
59 | 60 | #[derive(Debug, Clone, clap::Args)] |
@@ -291,7 +292,7 @@ impl ObjectStorageProvider for S3Config { |
291 | 292 | let s3 = LimitStore::new(s3, super::MAX_OBJECT_STORE_REQUESTS); |
292 | 293 | let s3 = MetricLayer::new(s3); |
293 | 294 |
|
294 | | - let object_store_registry: DefaultObjectStoreRegistry = DefaultObjectStoreRegistry::new(); |
| 295 | + let object_store_registry = DefaultObjectStoreRegistry::new(); |
295 | 296 | let url = ObjectStoreUrl::parse(format!("s3://{}", &self.bucket_name)).unwrap(); |
296 | 297 | object_store_registry.register_store(url.as_ref(), Arc::new(s3)); |
297 | 298 |
|
@@ -474,7 +475,7 @@ impl S3 { |
474 | 475 | } |
475 | 476 | Ok(result_file_list) |
476 | 477 | } |
477 | | - async fn _upload_file(&self, key: &str, path: &StdPath) -> Result<(), ObjectStorageError> { |
| 478 | + async fn _upload_file(&self, key: &str, path: &Path) -> Result<(), ObjectStorageError> { |
478 | 479 | let instant = Instant::now(); |
479 | 480 |
|
480 | 481 | // // TODO: Uncomment this when multipart is fixed |
@@ -503,7 +504,7 @@ impl S3 { |
503 | 504 | } |
504 | 505 |
|
505 | 506 | // TODO: introduce parallel, multipart-uploads if required |
506 | | - // async fn _upload_multipart(&self, key: &str, path: &StdPath) -> Result<(), ObjectStorageError> { |
| 507 | + // async fn _upload_multipart(&self, key: &str, path: &Path) -> Result<(), ObjectStorageError> { |
507 | 508 | // let mut buf = vec![0u8; MULTIPART_UPLOAD_SIZE / 2]; |
508 | 509 | // let mut file = OpenOptions::new().read(true).open(path).await?; |
509 | 510 |
|
@@ -750,7 +751,7 @@ impl ObjectStorage for S3 { |
750 | 751 | Ok(files) |
751 | 752 | } |
752 | 753 |
|
753 | | - async fn upload_file(&self, key: &str, path: &StdPath) -> Result<(), ObjectStorageError> { |
| 754 | + async fn upload_file(&self, key: &str, path: &Path) -> Result<(), ObjectStorageError> { |
754 | 755 | self._upload_file(key, path).await?; |
755 | 756 |
|
756 | 757 | Ok(()) |
|
0 commit comments