iceberg/io/storage/config/
hf.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
18//! HuggingFace Hub storage configuration.
19
20use serde::{Deserialize, Serialize};
21use typed_builder::TypedBuilder;
22
23use super::StorageConfig;
24use crate::Result;
25
26/// HuggingFace Hub authentication token.
27pub const HF_TOKEN: &str = "hf.token";
28/// HuggingFace Hub endpoint URL. Defaults to `https://huggingface.co`.
29pub const HF_ENDPOINT: &str = "hf.endpoint";
30/// Default git revision/branch for all paths that don't specify one. Defaults to `main`.
31pub const HF_REVISION: &str = "hf.revision";
32
33/// HuggingFace Hub storage configuration.
34///
35/// Repo type, repo ID, and revision are normally encoded in the file path URI
36/// (`hf://<repo_type>/<owner>/<repo>[@<revision>]/<path>`, where `<repo_type>`
37/// is one of `models`, `datasets`, `spaces`, or `buckets`).
38/// The fields here provide credentials and a default revision fallback.
39#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, TypedBuilder)]
40pub struct HfConfig {
41    /// HuggingFace Hub API token (required for private repos and write access).
42    #[builder(default, setter(strip_option, into))]
43    pub token: Option<String>,
44    /// HuggingFace Hub endpoint. Defaults to `https://huggingface.co`.
45    #[builder(default, setter(strip_option, into))]
46    pub endpoint: Option<String>,
47    /// Default revision to use when a path URI does not specify one.
48    #[builder(default, setter(strip_option, into))]
49    pub revision: Option<String>,
50}
51
52impl TryFrom<&StorageConfig> for HfConfig {
53    type Error = crate::Error;
54
55    fn try_from(config: &StorageConfig) -> Result<Self> {
56        let props = config.props();
57        let mut cfg = HfConfig::default();
58        if let Some(token) = props.get(HF_TOKEN) {
59            cfg.token = Some(token.clone());
60        }
61        if let Some(endpoint) = props.get(HF_ENDPOINT) {
62            cfg.endpoint = Some(endpoint.clone());
63        }
64        if let Some(revision) = props.get(HF_REVISION) {
65            cfg.revision = Some(revision.clone());
66        }
67        Ok(cfg)
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74
75    #[test]
76    fn test_hf_config_builder() {
77        let cfg = HfConfig::builder()
78            .token("hf_mytoken")
79            .endpoint("https://huggingface.co")
80            .revision("dev")
81            .build();
82        assert_eq!(cfg.token.as_deref(), Some("hf_mytoken"));
83        assert_eq!(cfg.endpoint.as_deref(), Some("https://huggingface.co"));
84        assert_eq!(cfg.revision.as_deref(), Some("dev"));
85    }
86
87    #[test]
88    fn test_hf_config_from_storage_config() {
89        let storage_config = StorageConfig::new()
90            .with_prop(HF_TOKEN, "hf_abc123")
91            .with_prop(HF_ENDPOINT, "https://huggingface.co");
92
93        let cfg = HfConfig::try_from(&storage_config).unwrap();
94        assert_eq!(cfg.token.as_deref(), Some("hf_abc123"));
95        assert_eq!(cfg.endpoint.as_deref(), Some("https://huggingface.co"));
96    }
97
98    #[test]
99    fn test_hf_config_empty() {
100        let cfg = HfConfig::try_from(&StorageConfig::new()).unwrap();
101        assert_eq!(cfg.token, None);
102        assert_eq!(cfg.endpoint, None);
103    }
104}
iceberg/io/storage/config/hf.rs

iceberg/io/storage/config/
hf.rs
