Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add connection Credentials to StorageLocation #743

Closed
wants to merge 5 commits into from
Closed

Add connection Credentials to StorageLocation #743

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
af4821f
Add connection credentials to StorageLocation and update related func…
lizardoluis Nov 21, 2024
a0fb6bc
Fix test to use get_credentials method
lizardoluis Nov 21, 2024
91ffb6e
Rename 'location' field to 'url' in StorageLocation message and updat…
lizardoluis Nov 21, 2024
00c79cf
Format
lizardoluis Nov 21, 2024
bbaeb87
Remove unused to_hashmap method from ObjectStoreConfig implementation
lizardoluis Nov 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
985 changes: 385 additions & 600 deletions Cargo.lock
View file
Open in desktop

Large diffs are not rendered by default.

4 changes: 3 additions & 1 deletion clade/proto/schema.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,13 @@ message TableObject {
// A single root storage location, hosting many individual tables
message StorageLocation {
// URL of the storage location root
string location = 1;
string url = 1;
// Connection options for the object store client
map<string, string> options = 2;
// Unique storage location identifier
string name = 3;
// Connection credentials for the object store client
map<string, string> credentials = 4;
}

message ListSchemaRequest {
Expand Down
250 changes: 154 additions & 96 deletions object_store_factory/src/aws.rs
View file
Open in desktop

Large diffs are not rendered by default.

91 changes: 69 additions & 22 deletions object_store_factory/src/google.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ impl GCSConfig {
})
}

pub fn to_hashmap(&self) -> HashMap<String, String> {
pub fn get_options(&self) -> HashMap<String, String> {
let mut map = HashMap::new();
map.insert(
GoogleConfigKey::Bucket.as_ref().to_string(),
Expand All @@ -58,6 +58,18 @@ impl GCSConfig {
map
}

pub fn get_credentials(&self) -> HashMap<String, String> {
let mut map = HashMap::new();
if let Some(google_application_credentials) = &self.google_application_credentials
{
map.insert(
GoogleConfigKey::ApplicationCredentials.as_ref().to_string(),
google_application_credentials.clone(),
);
}
map
}

pub fn bucket_to_url(&self) -> String {
format!("gs://{}", &self.bucket)
}
Expand Down Expand Up @@ -85,12 +97,16 @@ impl GCSConfig {
}
}

pub fn map_options_into_google_config_keys(
pub fn map_options_and_credentials_into_google_config_keys(
input_options: HashMap<String, String>,
input_credentials: HashMap<String, String>,
) -> Result<HashMap<GoogleConfigKey, String>, object_store::Error> {
let mut mapped_keys = HashMap::new();

for (key, value) in input_options {
for (key, value) in input_options
.into_iter()
.chain(input_credentials.into_iter())
{
match GoogleConfigKey::from_str(&key) {
Ok(config_key) => {
mapped_keys.insert(config_key, value);
Expand Down Expand Up @@ -245,18 +261,21 @@ mod tests {

#[test]
fn test_map_options_into_google_config_keys_with_valid_keys() {
let mut input_options = HashMap::new();
input_options.insert(
let mut options = HashMap::new();
options.insert(
"google_service_account".to_string(),
"my_google_service_account".to_string(),
);
Comment on lines +265 to 268
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't this be in credentials now?

input_options.insert("bucket".to_string(), "my-bucket".to_string());
input_options.insert(
options.insert("bucket".to_string(), "my-bucket".to_string());

let mut credentials = HashMap::new();
credentials.insert(
"google_service_account_key".to_string(),
"my_google_service_account_key".to_string(),
);

let result = map_options_into_google_config_keys(input_options);
let result =
map_options_and_credentials_into_google_config_keys(options, credentials);
assert!(result.is_ok());

let mapped_keys = result.unwrap();
Expand All @@ -275,11 +294,14 @@ mod tests {
}

#[test]
fn test_map_options_into_google_config_keys_with_invalid_key() {
let mut input_options = HashMap::new();
input_options.insert("invalid_key".to_string(), "some_value".to_string());
fn test_map_options_and_credentials_into_google_config_keys_with_invalid_key() {
let mut options = HashMap::new();
options.insert("invalid_key".to_string(), "some_value".to_string());

let credentials = HashMap::new();

let result = map_options_into_google_config_keys(input_options);
let result =
map_options_and_credentials_into_google_config_keys(options, credentials);
assert!(result.is_err());

let error = result.err().unwrap();
Expand All @@ -291,12 +313,35 @@ mod tests {
}

#[test]
fn test_map_options_into_google_config_keys_with_mixed_keys() {
let mut input_options = HashMap::new();
input_options.insert("invalid_key".to_string(), "some_value".to_string());
input_options.insert("bucket".to_string(), "my-bucket".to_string());
fn test_map_options_and_credentials_into_google_config_keys_with_invalid_credential()
{
let options = HashMap::new();

let mut credentials = HashMap::new();
credentials.insert("invalid_credential".to_string(), "some_value".to_string());

let result =
map_options_and_credentials_into_google_config_keys(options, credentials);
assert!(result.is_err());

let error = result.err().unwrap();
print!("ERROR1: {:?}", error.to_string());
assert_eq!(
error.to_string(),
"Configuration key: 'invalid_credential' is not valid for store 'GCS'."
)
}

#[test]
fn test_map_options_and_credentials_into_google_config_keys_with_mixed_keys() {
let mut options = HashMap::new();
options.insert("invalid_key".to_string(), "some_value".to_string());
options.insert("bucket".to_string(), "my-bucket".to_string());

let credentials = HashMap::new();

let result = map_options_into_google_config_keys(input_options);
let result =
map_options_and_credentials_into_google_config_keys(options, credentials);
assert!(result.is_err());

let error = result.err().unwrap();
Expand All @@ -307,9 +352,11 @@ mod tests {
}

#[test]
fn test_map_options_into_google_config_keys_empty_input() {
let input_options = HashMap::new();
let result = map_options_into_google_config_keys(input_options);
fn test_map_options_and_credentials_into_google_config_keys_empty_input() {
let options = HashMap::new();
let credentials = HashMap::new();
let result =
map_options_and_credentials_into_google_config_keys(options, credentials);
assert!(result.is_ok());

let mapped_keys = result.unwrap();
Expand Down Expand Up @@ -362,7 +409,7 @@ mod tests {
google_application_credentials: Some("path/to/credentials.json".to_string()),
};

let hashmap = gcs_config.to_hashmap();
let hashmap = gcs_config.get_options();

assert_eq!(
hashmap.get(GoogleConfigKey::Bucket.as_ref()),
Expand All @@ -383,7 +430,7 @@ mod tests {
google_application_credentials: None,
};

let hashmap = gcs_config.to_hashmap();
let hashmap = gcs_config.get_options();

assert_eq!(
hashmap.get(GoogleConfigKey::Bucket.as_ref()),
Expand Down
45 changes: 33 additions & 12 deletions object_store_factory/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,9 @@ pub struct StorageLocationInfo {
// Actual object store for this location
pub object_store: Arc<DynObjectStore>,

// Options used to construct the object store (for gRPC consumers)
pub options: HashMap<String, String>,
pub url: String,
pub options: HashMap<String, String>,
pub credentials: HashMap<String, String>,
}

impl ObjectStoreConfig {
Expand Down Expand Up @@ -68,8 +68,8 @@ impl ObjectStoreConfig {
pub fn to_hashmap(&self) -> HashMap<String, String> {
match self {
ObjectStoreConfig::Local(config) => config.to_hashmap(),
ObjectStoreConfig::AmazonS3(config) => config.to_hashmap(),
ObjectStoreConfig::GoogleCloudStorage(config) => config.to_hashmap(),
ObjectStoreConfig::AmazonS3(config) => config.get_options(),
ObjectStoreConfig::GoogleCloudStorage(config) => config.get_options(),
ObjectStoreConfig::Memory => HashMap::new(), // Memory config has no associated data
}
}
Expand All @@ -93,14 +93,16 @@ impl ObjectStoreConfig {
match self {
ObjectStoreConfig::AmazonS3(aws_config) => Ok(StorageLocationInfo {
object_store: aws_config.build_amazon_s3()?,
options: aws_config.to_hashmap(),
url: aws_config.bucket_to_url(),
options: aws_config.get_options(),
credentials: aws_config.get_credentials(),
}),
ObjectStoreConfig::GoogleCloudStorage(google_config) => {
Ok(StorageLocationInfo {
object_store: google_config.build_google_cloud_storage()?,
options: google_config.to_hashmap(),
url: google_config.bucket_to_url(),
options: google_config.get_options(),
credentials: google_config.get_credentials(),
})
}
_ => {
Expand Down Expand Up @@ -130,9 +132,10 @@ impl ObjectStoreConfig {
}
}

pub async fn build_object_store_from_opts(
pub async fn build_object_store_from_options_and_credentials(
url: &Url,
options: HashMap<String, String>,
credentials: HashMap<String, String>,
) -> Result<Box<dyn ObjectStore>, object_store::Error> {
let (scheme, _) = ObjectStoreScheme::parse(url).unwrap();

Expand All @@ -155,7 +158,11 @@ pub async fn build_object_store_from_opts(
Ok(Box::new(store))
}
ObjectStoreScheme::AmazonS3 => {
let mut s3_options = aws::map_options_into_amazon_s3_config_keys(options)?;
let mut s3_options =
aws::map_options_and_credentials_into_amazon_s3_config_keys(
options,
credentials,
)?;
aws::add_amazon_s3_specific_options(url, &mut s3_options).await;
aws::add_amazon_s3_environment_variables(&mut s3_options);

Expand All @@ -166,7 +173,11 @@ pub async fn build_object_store_from_opts(
Ok(store)
}
ObjectStoreScheme::GoogleCloudStorage => {
let mut gcs_options = google::map_options_into_google_config_keys(options)?;
let mut gcs_options =
google::map_options_and_credentials_into_google_config_keys(
options,
credentials,
)?;
google::add_google_cloud_storage_environment_variables(&mut gcs_options);

let (mut store, _) = parse_url_opts(url, gcs_options)?;
Expand All @@ -188,10 +199,19 @@ pub async fn build_object_store_from_opts(
pub async fn build_storage_location_info_from_opts(
url: &Url,
options: &HashMap<String, String>,
credentials: &HashMap<String, String>,
) -> Result<StorageLocationInfo, object_store::Error> {
Ok(StorageLocationInfo {
object_store: Arc::new(build_object_store_from_opts(url, options.clone()).await?),
object_store: Arc::new(
build_object_store_from_options_and_credentials(
url,
options.clone(),
credentials.clone(),
)
.await?,
),
options: options.clone(),
credentials: credentials.clone(),
url: url.to_string(),
})
}
Expand Down Expand Up @@ -294,7 +314,8 @@ mod tests {
#[tokio::test]
async fn test_build_aws_object_store(#[values(true, false)] use_env: bool) {
let url = Url::parse("s3://my-bucket").unwrap();
let options: HashMap<String, String> = if use_env {
let options = HashMap::new();
let credentials: HashMap<String, String> = if use_env {
HashMap::new()
} else {
HashMap::from([
Expand All @@ -308,7 +329,7 @@ mod tests {
("AWS_ACCESS_KEY_ID", Some("env-key")),
("AWS_SECRET_ACCESS_KEY", Some("env-secret")),
],
build_object_store_from_opts(&url, options),
build_object_store_from_options_and_credentials(&url, options, credentials),
)
.await;

Expand Down
25 changes: 14 additions & 11 deletions src/catalog/metastore.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,8 @@ use url::Url;
use super::empty::EmptyStore;

// Root URL for a storage location alongside client connection options
type LocationAndOptions = (String, HashMap<String, String>);
type LocationOptionsAndCredentials =
(String, HashMap<String, String>, HashMap<String, String>);

// This is the main entrypoint to all individual catalogs for various objects types.
// The intention is to make it extensible and de-coupled from the underlying metastore
Expand Down Expand Up @@ -115,7 +116,7 @@ impl Metastore {
let store_options = catalog_schemas
.stores
.into_iter()
.map(|store| (store.name, (store.location, store.options)))
.map(|store| (store.name, (store.url, store.options, store.credentials)))
.collect();

// Turn the list of all collections, tables and their columns into a nested map.
Expand All @@ -137,12 +138,12 @@ impl Metastore {
async fn build_schema(
&self,
schema: SchemaObject,
store_options: &HashMap<String, LocationAndOptions>,
store_options_and_credentials: &HashMap<String, LocationOptionsAndCredentials>,
) -> CatalogResult<(Arc<str>, Arc<SeafowlSchema>)> {
let schema_name = schema.name;

let tables: DashMap<_, _> = stream::iter(schema.tables)
.then(|table| self.build_table(table, store_options))
.then(|table| self.build_table(table, store_options_and_credentials))
.try_collect()
.await?;

Expand All @@ -158,7 +159,7 @@ impl Metastore {
async fn build_table(
&self,
table: TableObject,
store_options: &HashMap<String, LocationAndOptions>,
store_options_and_credentials: &HashMap<String, LocationOptionsAndCredentials>,
) -> CatalogResult<(Arc<str>, Arc<dyn TableProvider>)> {
// Build a delta table but don't load it yet; we'll do that only for tables that are
// actually referenced in a statement, via the async `table` method of the schema provider.
Expand All @@ -169,17 +170,19 @@ impl Metastore {
let table_log_store = match table.store {
// Use the provided customized location
Some(name) => {
let (location, this_store_options) = store_options
.get(&name)
.ok_or(CatalogError::Generic {
reason: format!("Object store with name {name} not found"),
})?
.clone();
let (location, this_store_options, this_store_credentials) =
store_options_and_credentials
.get(&name)
.ok_or(CatalogError::Generic {
reason: format!("Object store with name {name} not found"),
})?
.clone();

self.object_stores
.get_log_store_for_table(
Url::parse(&location)?,
this_store_options,
this_store_credentials,
table.path,
)
.await?
Expand Down
3 changes: 2 additions & 1 deletion src/frontend/flight/handler.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,12 +172,13 @@ impl SeafowlFlightHandler {
.metastore
.object_stores
.get_log_store_for_table(
Url::parse(&store_loc.location).map_err(|e| {
Url::parse(&store_loc.url).map_err(|e| {
DataFusionError::Execution(format!(
"Couldn't parse sync location: {e}"
))
})?,
store_loc.options,
store_loc.credentials,
cmd.path,
)
.await?
Expand Down
Loading