unitycatalog_server/api/
temporary_credentials.rs

1use unitycatalog_common::models::credentials::v1::GetCredentialRequest;
2use unitycatalog_common::models::model_versions::v1::ModelVersion;
3use unitycatalog_common::models::tables::v1::Table;
4use unitycatalog_common::models::temporary_credentials::v1::*;
5use unitycatalog_common::models::volumes::v1::Volume;
6use unitycatalog_common::models::{ResourceIdent, ResourceName, ResourceRef};
7
8use super::RequestContext;
9use crate::api::CredentialHandler;
10use crate::api::credentials::CredentialHandlerExt;
11pub use crate::codegen::temporary_credentials::TemporaryCredentialHandler;
12use crate::policy::{Permission, Policy};
13use crate::services::location::StorageLocationUrl;
14use crate::services::object_store::find_external_location_for_url;
15use crate::services::vending::{mint, to_temporary_credential};
16use crate::store::ResourceStore;
17use crate::{Error, Result};
18
19use buffa::Enumeration;
20use olai_storage_credentials::AccessLevel;
21
22/// Vend a temporary credential for a securable stored at `location`.
23///
24/// Local (`file://`) storage needs neither an external location nor a cloud
25/// credential: the client builds a `LocalFileSystem` addressed by full path. Any
26/// other location is minted from the credential of the external location governing
27/// it.
28pub(crate) async fn vend_for_location<T>(
29    handler: &T,
30    location: &StorageLocationUrl,
31    access: AccessLevel,
32) -> Result<TemporaryCredential>
33where
34    T: ResourceStore + CredentialHandlerExt + ?Sized,
35{
36    let credential = if location.is_local() {
37        None
38    } else {
39        let ext_loc = find_external_location_for_url(location, handler).await?;
40        Some(
41            handler
42                .get_credential_internal(GetCredentialRequest {
43                    name: ext_loc.credential_name.clone(),
44                    ..Default::default()
45                })
46                .await?,
47        )
48    };
49    let vended = mint(credential.as_ref(), location, access).await?;
50    Ok(to_temporary_credential(&vended, location))
51}
52
53/// The permission a vend operation requires from the policy.
54///
55/// Read-only vending requires [`Permission::Read`]; read-write vending requires
56/// [`Permission::Write`] so the policy can deny write access independently.
57fn required_permission(access: AccessLevel) -> Permission {
58    match access {
59        AccessLevel::ReadWrite => Permission::Write,
60        _ => Permission::Read,
61    }
62}
63
64/// Map the proto `operation` integer to an [`AccessLevel`].
65///
66/// The path, table, and volume request enums all encode the same semantics with
67/// the same numeric values (read = 1, write/read-write = 2; path adds
68/// create-table = 3), so a single mapping serves all three:
69///
70/// - `PATH_CREATE_TABLE`, `READ_WRITE`, and `WRITE_VOLUME` are treated as `ReadWrite`.
71/// - `READ`, `READ_VOLUME`, and `Unspecified` default to `Read` (least privilege).
72fn to_access_level(operation: i32) -> AccessLevel {
73    use generate_temporary_path_credentials_request::Operation as PathOp;
74    use generate_temporary_table_credentials_request::Operation as TableOp;
75    use generate_temporary_volume_credentials_request::Operation as VolumeOp;
76
77    // Check path operations first (values 0–3 are defined for both enums,
78    // but semantically we just need to distinguish read from read-write).
79    match PathOp::from_i32(operation) {
80        Some(PathOp::PATH_READ_WRITE | PathOp::PATH_CREATE_TABLE) => {
81            return AccessLevel::ReadWrite;
82        }
83        Some(PathOp::PATH_READ | PathOp::UNSPECIFIED) | None => {}
84    }
85    // The table (READ_WRITE = 2), volume (WRITE_VOLUME = 2), and model-version
86    // (READ_WRITE_MODEL_VERSION = 2) write operations all share the same numeric
87    // value, so any of these enums resolves write access identically.
88    match (TableOp::from_i32(operation), VolumeOp::from_i32(operation)) {
89        (Some(TableOp::READ_WRITE), _) | (_, Some(VolumeOp::WRITE_VOLUME)) => {
90            AccessLevel::ReadWrite
91        }
92        _ => AccessLevel::Read,
93    }
94}
95
96#[async_trait::async_trait]
97impl<
98    T: ResourceStore
99        + Policy<RequestContext>
100        + CredentialHandler<RequestContext>
101        + CredentialHandlerExt,
102> TemporaryCredentialHandler<RequestContext> for T
103{
104    #[tracing::instrument(skip(self, context))]
105    async fn generate_temporary_path_credentials(
106        &self,
107        request: GenerateTemporaryPathCredentialsRequest,
108        context: RequestContext,
109    ) -> Result<TemporaryCredential> {
110        let operation = to_access_level(request.operation.to_i32());
111        let storage_url = StorageLocationUrl::parse(&request.url)?;
112
113        // Local (`file://`) storage is governed by no external location, so there is
114        // nothing to authorize against; vending it is what lets managed tables be
115        // created on a purely-local dev server.
116        if storage_url.is_local() {
117            return vend_for_location(self, &storage_url, operation).await;
118        }
119
120        let ext_loc = find_external_location_for_url(&storage_url, self).await?;
121        // Authorize against the concrete external location and the operation
122        // actually requested, rather than the unscoped `SecuredAction` default.
123        self.authorize_checked(
124            &(&ext_loc).into(),
125            &required_permission(operation),
126            &context,
127        )
128        .await?;
129        vend_for_location(self, &storage_url, operation).await
130    }
131
132    /// Generate temporary credentials for a volume.
133    ///
134    /// Resolves the volume by its UUID, authorizes the requested operation
135    /// against the concrete volume, then vends credentials scoped to the
136    /// volume's storage location — mirroring
137    /// `generate_temporary_table_credentials`.
138    ///
139    /// See: <https://docs.databricks.com/api/workspace/temporaryvolumecredentials/generatetemporaryvolumecredentials>
140    #[tracing::instrument(skip(self, context))]
141    async fn generate_temporary_volume_credentials(
142        &self,
143        request: GenerateTemporaryVolumeCredentialsRequest,
144        context: RequestContext,
145    ) -> Result<TemporaryCredential> {
146        let operation = to_access_level(request.operation.to_i32());
147        let volume_id = uuid::Uuid::parse_str(&request.volume_id)
148            .map_err(|_| Error::invalid_argument("volume_id is not a valid UUID"))?;
149        // Authorize against the concrete volume and the operation actually
150        // requested, rather than the unscoped `SecuredAction` default.
151        let volume_ident = ResourceIdent::Volume(ResourceRef::Uuid(volume_id));
152        self.authorize_checked(&volume_ident, &required_permission(operation), &context)
153            .await?;
154        let (resource, _) = self.get(&volume_ident).await?;
155        let volume: Volume = resource.try_into()?;
156        let location = volume.storage_location;
157        let storage_url = StorageLocationUrl::parse(&location)?;
158        vend_for_location(self, &storage_url, operation).await
159    }
160
161    /// Generate temporary credentials for a model version.
162    ///
163    /// Resolves the model version by its `(catalog, schema, model, version)`
164    /// composite name, authorizes the requested operation against the concrete
165    /// version, then vends credentials scoped to the version's storage location —
166    /// mirroring `generate_temporary_volume_credentials`.
167    ///
168    /// See: <https://docs.databricks.com/api/workspace/temporarymodelversioncredentials/generatetemporarymodelversioncredentials>
169    #[tracing::instrument(skip(self, context))]
170    async fn generate_temporary_model_version_credentials(
171        &self,
172        request: GenerateTemporaryModelVersionCredentialsRequest,
173        context: RequestContext,
174    ) -> Result<TemporaryCredential> {
175        let operation = to_access_level(request.operation.to_i32());
176        // The model version is keyed by the composite name that
177        // `ModelVersion::resource_name` produces.
178        let ident = ResourceIdent::model_version(ResourceName::new([
179            request.catalog_name.clone(),
180            request.schema_name.clone(),
181            request.model_name.clone(),
182            request.version.to_string(),
183        ]));
184        // Authorize against the concrete model version and the operation actually
185        // requested, rather than the unscoped `SecuredAction` default.
186        self.authorize_checked(&ident, &required_permission(operation), &context)
187            .await?;
188        let (resource, _) = self.get(&ident).await?;
189        let model_version: ModelVersion = resource.try_into()?;
190        let location = model_version.storage_location.ok_or_else(|| {
191            Error::invalid_argument("model version does not have a storage location")
192        })?;
193        let storage_url = StorageLocationUrl::parse(&location)?;
194        vend_for_location(self, &storage_url, operation).await
195    }
196
197    #[tracing::instrument(skip(self, context))]
198    async fn generate_temporary_table_credentials(
199        &self,
200        request: GenerateTemporaryTableCredentialsRequest,
201        context: RequestContext,
202    ) -> Result<TemporaryCredential> {
203        let operation = to_access_level(request.operation.to_i32());
204        let table_id = uuid::Uuid::parse_str(&request.table_id)
205            .map_err(|_| Error::invalid_argument("table_id is not a valid UUID"))?;
206        // Authorize against the concrete table and the operation actually
207        // requested, rather than the unscoped `SecuredAction` default.
208        let table_ident = ResourceIdent::Table(ResourceRef::Uuid(table_id));
209        self.authorize_checked(&table_ident, &required_permission(operation), &context)
210            .await?;
211        let (resource, _) = self.get(&table_ident).await?;
212        let table: Table = resource.try_into()?;
213        let location = table
214            .storage_location
215            .ok_or_else(|| Error::invalid_argument("table does not have a storage location"))?;
216        let storage_url = StorageLocationUrl::parse(&location)?;
217        vend_for_location(self, &storage_url, operation).await
218    }
219}
220// NOTE: These request types intentionally do not implement `SecuredAction`.
221// Authorization is performed inside the handlers against the *concrete* resolved
222// resource (external location / table) and the *requested* operation
223// (read vs. read-write), which a static `SecuredAction` impl cannot express.
224// See `generate_temporary_path_credentials` / `generate_temporary_table_credentials`.
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229    use generate_temporary_path_credentials_request::Operation as PathOp;
230    use generate_temporary_table_credentials_request::Operation as TableOp;
231    use generate_temporary_volume_credentials_request::Operation as VolumeOp;
232
233    #[test]
234    fn vend_operation_volume_read_is_read() {
235        assert_eq!(
236            to_access_level(VolumeOp::ReadVolume as i32),
237            AccessLevel::Read
238        );
239    }
240
241    #[test]
242    fn vend_operation_volume_write_is_read_write() {
243        assert_eq!(
244            to_access_level(VolumeOp::WriteVolume as i32),
245            AccessLevel::ReadWrite
246        );
247    }
248
249    #[test]
250    fn vend_operation_unspecified_defaults_to_read() {
251        assert_eq!(
252            to_access_level(VolumeOp::Unspecified as i32),
253            AccessLevel::Read
254        );
255        // Out-of-range values are treated as least-privilege read, not an error.
256        assert_eq!(to_access_level(99), AccessLevel::Read);
257    }
258
259    #[test]
260    fn vend_operation_table_and_path_unchanged() {
261        assert_eq!(to_access_level(TableOp::Read as i32), AccessLevel::Read);
262        assert_eq!(
263            to_access_level(TableOp::ReadWrite as i32),
264            AccessLevel::ReadWrite
265        );
266        assert_eq!(
267            to_access_level(PathOp::PathCreateTable as i32),
268            AccessLevel::ReadWrite
269        );
270    }
271
272    /// Local storage is vended without an external location or a credential, the
273    /// same way for every securable kind, and the response names the location as
274    /// given.
275    #[cfg(feature = "memory")]
276    #[tokio::test]
277    async fn local_locations_vend_without_a_credential() {
278        use std::sync::Arc;
279
280        use unitycatalog_common::services::encryption::{EnvelopeEncryptor, LocalKeyProvider};
281
282        use crate::memory::InMemoryResourceStore;
283        use crate::policy::ConstantPolicy;
284        use crate::services::ServerHandler;
285
286        let encryptor =
287            EnvelopeEncryptor::local(LocalKeyProvider::single("test", vec![0x42; 32]).unwrap());
288        let store = Arc::new(InMemoryResourceStore::new(encryptor));
289        let policy: Arc<dyn Policy<RequestContext>> = Arc::new(ConstantPolicy::default());
290        let handler = ServerHandler::try_new_tokio(policy, store).unwrap();
291
292        let location = StorageLocationUrl::parse("file:///tmp/uc/tables/t1").unwrap();
293        let credential = vend_for_location(&handler, &location, AccessLevel::ReadWrite)
294            .await
295            .unwrap();
296        assert!(credential.credentials.is_none());
297        assert_eq!(credential.url, "file:///tmp/uc/tables/t1");
298        assert_eq!(credential.expiration_time, 0);
299
300        // Cloud storage still needs a governing external location.
301        let location = StorageLocationUrl::parse("s3://bucket/tables/t1").unwrap();
302        assert!(matches!(
303            vend_for_location(&handler, &location, AccessLevel::Read).await,
304            Err(Error::NotFound)
305        ));
306    }
307}
unitycatalog_server/api/temporary_credentials.rs

unitycatalog_server/api/
temporary_credentials.rs
