iceberg_catalog_glue/
utils.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::collections::HashMap;
19
20use aws_config::{BehaviorVersion, Region, SdkConfig};
21use aws_sdk_glue::config::Credentials;
22use aws_sdk_glue::types::{Database, DatabaseInput, StorageDescriptor, TableInput};
23use iceberg::spec::TableMetadata;
24use iceberg::{Error, ErrorKind, Namespace, NamespaceIdent, Result};
25
26use crate::error::from_aws_build_error;
27use crate::schema::GlueSchemaBuilder;
28
29/// Property aws profile name
30pub const AWS_PROFILE_NAME: &str = "profile_name";
31/// Property aws region
32pub const AWS_REGION_NAME: &str = "region_name";
33/// Property aws access key
34pub const AWS_ACCESS_KEY_ID: &str = "aws_access_key_id";
35/// Property aws secret access key
36pub const AWS_SECRET_ACCESS_KEY: &str = "aws_secret_access_key";
37/// Property aws session token
38pub const AWS_SESSION_TOKEN: &str = "aws_session_token";
39/// Parameter namespace description
40const DESCRIPTION: &str = "description";
41/// Parameter namespace location uri
42const LOCATION: &str = "location_uri";
43/// Property `metadata_location` for `TableInput`
44const METADATA_LOCATION: &str = "metadata_location";
45/// Property `previous_metadata_location` for `TableInput`
46const PREV_METADATA_LOCATION: &str = "previous_metadata_location";
47/// Property external table for `TableInput`
48const EXTERNAL_TABLE: &str = "EXTERNAL_TABLE";
49/// Parameter key `table_type` for `TableInput`
50const TABLE_TYPE: &str = "table_type";
51/// Parameter value `table_type` for `TableInput`
52const ICEBERG: &str = "ICEBERG";
53
54/// Creates an aws sdk configuration based on
55/// provided properties and an optional endpoint URL.
56pub(crate) async fn create_sdk_config(
57    properties: &HashMap<String, String>,
58    endpoint_uri: Option<&String>,
59) -> SdkConfig {
60    let mut config = aws_config::defaults(BehaviorVersion::latest());
61
62    if let Some(endpoint) = endpoint_uri {
63        config = config.endpoint_url(endpoint)
64    };
65
66    if properties.is_empty() {
67        return config.load().await;
68    }
69
70    if let (Some(access_key), Some(secret_key)) = (
71        properties.get(AWS_ACCESS_KEY_ID),
72        properties.get(AWS_SECRET_ACCESS_KEY),
73    ) {
74        let session_token = properties.get(AWS_SESSION_TOKEN).cloned();
75        let credentials_provider =
76            Credentials::new(access_key, secret_key, session_token, None, "properties");
77
78        config = config.credentials_provider(credentials_provider)
79    };
80
81    if let Some(profile_name) = properties.get(AWS_PROFILE_NAME) {
82        config = config.profile_name(profile_name);
83    }
84
85    if let Some(region_name) = properties.get(AWS_REGION_NAME) {
86        let region = Region::new(region_name.clone());
87        config = config.region(region);
88    }
89
90    config.load().await
91}
92
93/// Create `DatabaseInput` from `NamespaceIdent` and properties
94pub(crate) fn convert_to_database(
95    namespace: &NamespaceIdent,
96    properties: &HashMap<String, String>,
97) -> Result<DatabaseInput> {
98    let db_name = validate_namespace(namespace)?;
99    let mut builder = DatabaseInput::builder().name(db_name);
100
101    for (k, v) in properties.iter() {
102        match k.as_ref() {
103            DESCRIPTION => {
104                builder = builder.description(v);
105            }
106            LOCATION => {
107                builder = builder.location_uri(v);
108            }
109            _ => {
110                builder = builder.parameters(k, v);
111            }
112        }
113    }
114
115    builder.build().map_err(from_aws_build_error)
116}
117
118/// Create `Namespace` from aws sdk glue `Database`
119pub(crate) fn convert_to_namespace(database: &Database) -> Namespace {
120    let db_name = database.name().to_string();
121    let mut properties = database
122        .parameters()
123        .map_or_else(HashMap::new, |p| p.clone());
124
125    if let Some(location_uri) = database.location_uri() {
126        properties.insert(LOCATION.to_string(), location_uri.to_string());
127    };
128
129    if let Some(description) = database.description() {
130        properties.insert(DESCRIPTION.to_string(), description.to_string());
131    }
132
133    Namespace::with_properties(NamespaceIdent::new(db_name), properties)
134}
135
136/// Converts Iceberg table metadata into an
137/// AWS Glue `TableInput` representation.
138///
139/// This function facilitates the integration of Iceberg tables with AWS Glue
140/// by converting Iceberg table metadata into a Glue-compatible `TableInput`
141/// structure.
142pub(crate) fn convert_to_glue_table(
143    table_name: impl Into<String>,
144    metadata_location: String,
145    metadata: &TableMetadata,
146    properties: &HashMap<String, String>,
147    prev_metadata_location: Option<String>,
148) -> Result<TableInput> {
149    let glue_schema = GlueSchemaBuilder::from_iceberg(metadata)?.build();
150
151    let storage_descriptor = StorageDescriptor::builder()
152        .set_columns(Some(glue_schema))
153        .location(metadata.location().to_string())
154        .build();
155
156    let mut parameters = HashMap::from([
157        (TABLE_TYPE.to_string(), ICEBERG.to_string()),
158        (METADATA_LOCATION.to_string(), metadata_location),
159    ]);
160
161    if let Some(prev) = prev_metadata_location {
162        parameters.insert(PREV_METADATA_LOCATION.to_string(), prev);
163    }
164
165    let mut table_input_builder = TableInput::builder()
166        .name(table_name)
167        .set_parameters(Some(parameters))
168        .storage_descriptor(storage_descriptor)
169        .table_type(EXTERNAL_TABLE);
170
171    if let Some(description) = properties.get(DESCRIPTION) {
172        table_input_builder = table_input_builder.description(description);
173    }
174
175    let table_input = table_input_builder.build().map_err(from_aws_build_error)?;
176
177    Ok(table_input)
178}
179
180/// Checks if provided `NamespaceIdent` is valid
181pub(crate) fn validate_namespace(namespace: &NamespaceIdent) -> Result<String> {
182    let name = namespace.as_ref();
183
184    if name.len() != 1 {
185        return Err(Error::new(
186            ErrorKind::DataInvalid,
187            format!(
188                "Invalid database name: {namespace:?}, hierarchical namespaces are not supported"
189            ),
190        ));
191    }
192
193    let name = name[0].clone();
194
195    if name.is_empty() {
196        return Err(Error::new(
197            ErrorKind::DataInvalid,
198            "Invalid database, provided namespace is empty.",
199        ));
200    }
201
202    Ok(name)
203}
204
205/// Get default table location from `Namespace` properties
206pub(crate) fn get_default_table_location(
207    namespace: &Namespace,
208    db_name: impl AsRef<str>,
209    table_name: impl AsRef<str>,
210    warehouse: impl AsRef<str>,
211) -> String {
212    let properties = namespace.properties();
213
214    match properties.get(LOCATION) {
215        Some(location) => format!("{}/{}", location, table_name.as_ref()),
216        None => {
217            let warehouse_location = warehouse.as_ref().trim_end_matches('/');
218
219            format!(
220                "{}/{}.db/{}",
221                warehouse_location,
222                db_name.as_ref(),
223                table_name.as_ref()
224            )
225        }
226    }
227}
228
229/// Returns `true` if the given Glue table is an Iceberg table.
230///
231/// Iceberg tables are identified by the `table_type=ICEBERG` parameter
232/// (case-insensitive).
233pub(crate) fn is_iceberg_table(parameters: &Option<HashMap<String, String>>) -> bool {
234    parameters
235        .as_ref()
236        .and_then(|p| p.get(TABLE_TYPE))
237        .is_some_and(|v| v.eq_ignore_ascii_case(ICEBERG))
238}
239
240/// Get metadata location from `GlueTable` parameters
241pub(crate) fn get_metadata_location(
242    parameters: &Option<HashMap<String, String>>,
243) -> Result<String> {
244    match parameters {
245        Some(properties) => match properties.get(METADATA_LOCATION) {
246            Some(location) => Ok(location.to_string()),
247            None => Err(Error::new(
248                ErrorKind::DataInvalid,
249                format!("No '{METADATA_LOCATION}' set on table"),
250            )),
251        },
252        None => Err(Error::new(
253            ErrorKind::DataInvalid,
254            "No 'parameters' set on table. Location of metadata is undefined",
255        )),
256    }
257}
258
259#[macro_export]
260/// Extends aws sdk builder with `catalog_id` if present
261macro_rules! with_catalog_id {
262    ($builder:expr, $config:expr) => {{
263        if let Some(catalog_id) = &$config.catalog_id {
264            $builder.catalog_id(catalog_id)
265        } else {
266            $builder
267        }
268    }};
269}
270
271#[cfg(test)]
272mod tests {
273    use aws_sdk_glue::config::ProvideCredentials;
274    use aws_sdk_glue::types::Column;
275    use iceberg::spec::{NestedField, PrimitiveType, Schema, TableMetadataBuilder, Type};
276    use iceberg::{MetadataLocation, Namespace, Result, TableCreation};
277
278    use super::*;
279    use crate::schema::{ICEBERG_FIELD_CURRENT, ICEBERG_FIELD_ID, ICEBERG_FIELD_OPTIONAL};
280
281    fn create_metadata(schema: Schema) -> Result<TableMetadata> {
282        let table_creation = TableCreation::builder()
283            .name("my_table".to_string())
284            .location("my_location".to_string())
285            .schema(schema)
286            .build();
287        let metadata = TableMetadataBuilder::from_table_creation(table_creation)?
288            .build()?
289            .metadata;
290
291        Ok(metadata)
292    }
293
294    #[test]
295    fn test_get_metadata_location() -> Result<()> {
296        let params_valid = Some(HashMap::from([(
297            METADATA_LOCATION.to_string(),
298            "my_location".to_string(),
299        )]));
300        let params_missing_key = Some(HashMap::from([(
301            "not_here".to_string(),
302            "my_location".to_string(),
303        )]));
304
305        let result_valid = get_metadata_location(&params_valid)?;
306        let result_missing_key = get_metadata_location(&params_missing_key);
307        let result_no_params = get_metadata_location(&None);
308
309        assert_eq!(result_valid, "my_location");
310        assert!(result_missing_key.is_err());
311        assert!(result_no_params.is_err());
312
313        Ok(())
314    }
315
316    #[test]
317    fn test_is_iceberg_table() {
318        // table_type=ICEBERG -> iceberg
319        let params_table_type = Some(HashMap::from([(
320            TABLE_TYPE.to_string(),
321            ICEBERG.to_string(),
322        )]));
323        assert!(is_iceberg_table(&params_table_type));
324
325        // table_type is case-insensitive
326        let params_table_type_lower = Some(HashMap::from([(
327            TABLE_TYPE.to_string(),
328            "iceberg".to_string(),
329        )]));
330        assert!(is_iceberg_table(&params_table_type_lower));
331
332        // Plain Hive table -> not iceberg
333        let params_hive = Some(HashMap::from([(
334            TABLE_TYPE.to_string(),
335            "EXTERNAL_TABLE".to_string(),
336        )]));
337        assert!(!is_iceberg_table(&params_hive));
338
339        // No parameters at all -> not iceberg
340        assert!(!is_iceberg_table(&None));
341
342        // Empty parameters -> not iceberg
343        assert!(!is_iceberg_table(&Some(HashMap::new())));
344    }
345
346    #[test]
347    fn test_convert_to_glue_table() -> Result<()> {
348        let table_name = "my_table".to_string();
349        let schema = Schema::builder()
350            .with_schema_id(1)
351            .with_fields(vec![
352                NestedField::required(1, "foo", Type::Primitive(PrimitiveType::Int)).into(),
353            ])
354            .build()?;
355
356        let metadata = create_metadata(schema)?;
357        let metadata_location = MetadataLocation::try_new_with_metadata(&metadata)?.to_string();
358
359        let parameters = HashMap::from([
360            (ICEBERG_FIELD_ID.to_string(), "1".to_string()),
361            (ICEBERG_FIELD_OPTIONAL.to_string(), "false".to_string()),
362            (ICEBERG_FIELD_CURRENT.to_string(), "true".to_string()),
363        ]);
364
365        let column = Column::builder()
366            .name("foo")
367            .r#type("int")
368            .set_parameters(Some(parameters))
369            .set_comment(None)
370            .build()
371            .map_err(from_aws_build_error)?;
372
373        let storage_descriptor = StorageDescriptor::builder()
374            .set_columns(Some(vec![column]))
375            .location(metadata.location())
376            .build();
377
378        let result = convert_to_glue_table(
379            &table_name,
380            metadata_location,
381            &metadata,
382            metadata.properties(),
383            None,
384        )?;
385
386        assert_eq!(result.name(), &table_name);
387        assert_eq!(result.description(), None);
388        assert_eq!(result.storage_descriptor, Some(storage_descriptor));
389
390        Ok(())
391    }
392
393    #[test]
394    fn test_get_default_table_location() -> Result<()> {
395        let properties = HashMap::from([(LOCATION.to_string(), "db_location".to_string())]);
396
397        let namespace =
398            Namespace::with_properties(NamespaceIdent::new("default".into()), properties);
399        let db_name = validate_namespace(namespace.name())?;
400        let table_name = "my_table";
401
402        let expected = "db_location/my_table";
403        let result =
404            get_default_table_location(&namespace, db_name, table_name, "warehouse_location");
405
406        assert_eq!(expected, result);
407
408        Ok(())
409    }
410
411    #[test]
412    fn test_get_default_table_location_warehouse() -> Result<()> {
413        let namespace = Namespace::new(NamespaceIdent::new("default".into()));
414        let db_name = validate_namespace(namespace.name())?;
415        let table_name = "my_table";
416
417        let expected = "warehouse_location/default.db/my_table";
418        let result =
419            get_default_table_location(&namespace, db_name, table_name, "warehouse_location");
420
421        assert_eq!(expected, result);
422
423        Ok(())
424    }
425
426    #[test]
427    fn test_convert_to_namespace() -> Result<()> {
428        let db = Database::builder()
429            .name("my_db")
430            .location_uri("my_location")
431            .description("my_description")
432            .build()
433            .map_err(from_aws_build_error)?;
434
435        let properties = HashMap::from([
436            (DESCRIPTION.to_string(), "my_description".to_string()),
437            (LOCATION.to_string(), "my_location".to_string()),
438        ]);
439
440        let expected =
441            Namespace::with_properties(NamespaceIdent::new("my_db".to_string()), properties);
442        let result = convert_to_namespace(&db);
443
444        assert_eq!(result, expected);
445
446        Ok(())
447    }
448
449    #[test]
450    fn test_convert_to_database() -> Result<()> {
451        let namespace = NamespaceIdent::new("my_database".to_string());
452        let properties = HashMap::from([(LOCATION.to_string(), "my_location".to_string())]);
453
454        let result = convert_to_database(&namespace, &properties)?;
455
456        assert_eq!("my_database", result.name());
457        assert_eq!(Some("my_location".to_string()), result.location_uri);
458
459        Ok(())
460    }
461
462    #[test]
463    fn test_validate_namespace() {
464        let valid_ns = Namespace::new(NamespaceIdent::new("ns".to_string()));
465        let empty_ns = Namespace::new(NamespaceIdent::new("".to_string()));
466        let hierarchical_ns = Namespace::new(
467            NamespaceIdent::from_vec(vec!["level1".to_string(), "level2".to_string()]).unwrap(),
468        );
469
470        let valid = validate_namespace(valid_ns.name());
471        let empty = validate_namespace(empty_ns.name());
472        let hierarchical = validate_namespace(hierarchical_ns.name());
473
474        assert!(valid.is_ok());
475        assert!(empty.is_err());
476        assert!(hierarchical.is_err());
477    }
478
479    #[tokio::test]
480    async fn test_config_with_custom_endpoint() {
481        let properties = HashMap::new();
482        let endpoint_url = "http://custom_url:5001";
483
484        let sdk_config = create_sdk_config(&properties, Some(&endpoint_url.to_string())).await;
485
486        let result = sdk_config.endpoint_url().unwrap();
487
488        assert_eq!(result, endpoint_url);
489    }
490
491    #[tokio::test]
492    async fn test_config_with_properties() {
493        let properties = HashMap::from([
494            (AWS_PROFILE_NAME.to_string(), "my_profile".to_string()),
495            (AWS_REGION_NAME.to_string(), "us-east-1".to_string()),
496            (AWS_ACCESS_KEY_ID.to_string(), "my-access-id".to_string()),
497            (
498                AWS_SECRET_ACCESS_KEY.to_string(),
499                "my-secret-key".to_string(),
500            ),
501            (AWS_SESSION_TOKEN.to_string(), "my-token".to_string()),
502        ]);
503
504        let sdk_config = create_sdk_config(&properties, None).await;
505
506        let region = sdk_config.region().unwrap().as_ref();
507        let credentials = sdk_config
508            .credentials_provider()
509            .unwrap()
510            .provide_credentials()
511            .await
512            .unwrap();
513
514        assert_eq!("us-east-1", region);
515        assert_eq!("my-access-id", credentials.access_key_id());
516        assert_eq!("my-secret-key", credentials.secret_access_key());
517        assert_eq!("my-token", credentials.session_token().unwrap());
518    }
519}
iceberg_catalog_glue/utils.rs

iceberg_catalog_glue/
utils.rs
