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::credential_vending::{VendOperation, local_path_credential, vend_credential};
14use crate::services::location::{StorageLocationScheme, StorageLocationUrl};
15use crate::services::object_store::find_external_location_for_url;
16use crate::store::ResourceStore;
17use crate::{Error, Result};
18
19use buffa::Enumeration;
20use object_store::ObjectStoreScheme;
21
22/// The permission a vend operation requires from the policy.
23///
24/// Read-only vending requires [`Permission::Read`]; read-write vending requires
25/// [`Permission::Write`] so the policy can deny write access independently.
26fn required_permission(operation: VendOperation) -> Permission {
27    match operation {
28        VendOperation::Read => Permission::Read,
29        VendOperation::ReadWrite => Permission::Write,
30    }
31}
32
33/// Map the proto `operation` integer to a `VendOperation`.
34///
35/// The path, table, and volume request enums all encode the same semantics with
36/// the same numeric values (read = 1, write/read-write = 2; path adds
37/// create-table = 3), so a single mapping serves all three:
38///
39/// - `PATH_CREATE_TABLE`, `READ_WRITE`, and `WRITE_VOLUME` are treated as `ReadWrite`.
40/// - `READ`, `READ_VOLUME`, and `Unspecified` default to `Read` (least privilege).
41fn to_vend_operation(operation: i32) -> VendOperation {
42    use generate_temporary_path_credentials_request::Operation as PathOp;
43    use generate_temporary_table_credentials_request::Operation as TableOp;
44    use generate_temporary_volume_credentials_request::Operation as VolumeOp;
45
46    // Check path operations first (values 0–3 are defined for both enums,
47    // but semantically we just need to distinguish read from read-write).
48    match PathOp::from_i32(operation) {
49        Some(PathOp::PATH_READ_WRITE | PathOp::PATH_CREATE_TABLE) => {
50            return VendOperation::ReadWrite;
51        }
52        Some(PathOp::PATH_READ | PathOp::UNSPECIFIED) | None => {}
53    }
54    // The table (READ_WRITE = 2), volume (WRITE_VOLUME = 2), and model-version
55    // (READ_WRITE_MODEL_VERSION = 2) write operations all share the same numeric
56    // value, so any of these enums resolves write access identically.
57    match (TableOp::from_i32(operation), VolumeOp::from_i32(operation)) {
58        (Some(TableOp::READ_WRITE), _) | (_, Some(VolumeOp::WRITE_VOLUME)) => {
59            VendOperation::ReadWrite
60        }
61        _ => VendOperation::Read,
62    }
63}
64
65#[async_trait::async_trait]
66impl<
67    T: ResourceStore
68        + Policy<RequestContext>
69        + CredentialHandler<RequestContext>
70        + CredentialHandlerExt,
71> TemporaryCredentialHandler<RequestContext> for T
72{
73    #[tracing::instrument(skip(self, context))]
74    async fn generate_temporary_path_credentials(
75        &self,
76        request: GenerateTemporaryPathCredentialsRequest,
77        context: RequestContext,
78    ) -> Result<TemporaryCredential> {
79        let operation = to_vend_operation(request.operation.to_i32());
80        let storage_url = StorageLocationUrl::parse(&request.url)?;
81
82        // Local (`file://`) storage needs neither an external location nor a vended
83        // cloud credential: the client builds a `LocalFileSystem` addressed by full
84        // path, exactly as the server's own `get_object_store` short-circuits local
85        // storage. Without this, vending for a local managed root 404s
86        // (`find_external_location_for_url` finds no covering external location),
87        // which breaks managed-table creation on a purely-local dev server.
88        if matches!(
89            storage_url.scheme(),
90            StorageLocationScheme::ObjectStore(ObjectStoreScheme::Local)
91        ) {
92            return Ok(local_path_credential(&request.url));
93        }
94
95        let ext_loc = find_external_location_for_url(&storage_url, self).await?;
96        // Authorize against the concrete external location and the operation
97        // actually requested, rather than the unscoped `SecuredAction` default.
98        self.authorize_checked(
99            &(&ext_loc).into(),
100            &required_permission(operation),
101            &context,
102        )
103        .await?;
104        let credential = self
105            .get_credential_internal(GetCredentialRequest {
106                name: ext_loc.credential_name.clone(),
107                ..Default::default()
108            })
109            .await?;
110        vend_credential(&credential, &request.url, operation).await
111    }
112
113    /// Generate temporary credentials for a volume.
114    ///
115    /// Resolves the volume by its UUID, authorizes the requested operation
116    /// against the concrete volume, then vends credentials scoped to the
117    /// volume's storage location — mirroring
118    /// [`generate_temporary_table_credentials`](Self::generate_temporary_table_credentials).
119    ///
120    /// See: <https://docs.databricks.com/api/workspace/temporaryvolumecredentials/generatetemporaryvolumecredentials>
121    #[tracing::instrument(skip(self, context))]
122    async fn generate_temporary_volume_credentials(
123        &self,
124        request: GenerateTemporaryVolumeCredentialsRequest,
125        context: RequestContext,
126    ) -> Result<TemporaryCredential> {
127        let operation = to_vend_operation(request.operation.to_i32());
128        let volume_id = uuid::Uuid::parse_str(&request.volume_id)
129            .map_err(|_| Error::invalid_argument("volume_id is not a valid UUID"))?;
130        // Authorize against the concrete volume and the operation actually
131        // requested, rather than the unscoped `SecuredAction` default.
132        let volume_ident = ResourceIdent::Volume(ResourceRef::Uuid(volume_id));
133        self.authorize_checked(&volume_ident, &required_permission(operation), &context)
134            .await?;
135        let (resource, _) = self.get(&volume_ident).await?;
136        let volume: Volume = resource.try_into()?;
137        let location = volume.storage_location;
138        let storage_url = StorageLocationUrl::parse(&location)?;
139        let ext_loc = find_external_location_for_url(&storage_url, self).await?;
140        let credential = self
141            .get_credential_internal(GetCredentialRequest {
142                name: ext_loc.credential_name.clone(),
143                ..Default::default()
144            })
145            .await?;
146        vend_credential(&credential, &location, operation).await
147    }
148
149    /// Generate temporary credentials for a model version.
150    ///
151    /// Resolves the model version by its `(catalog, schema, model, version)`
152    /// composite name, authorizes the requested operation against the concrete
153    /// version, then vends credentials scoped to the version's storage location —
154    /// mirroring
155    /// [`generate_temporary_volume_credentials`](Self::generate_temporary_volume_credentials).
156    ///
157    /// See: <https://docs.databricks.com/api/workspace/temporarymodelversioncredentials/generatetemporarymodelversioncredentials>
158    #[tracing::instrument(skip(self, context))]
159    async fn generate_temporary_model_version_credentials(
160        &self,
161        request: GenerateTemporaryModelVersionCredentialsRequest,
162        context: RequestContext,
163    ) -> Result<TemporaryCredential> {
164        let operation = to_vend_operation(request.operation.to_i32());
165        // The model version is keyed by the composite name that
166        // `ModelVersion::resource_name` produces.
167        let ident = ResourceIdent::model_version(ResourceName::new([
168            request.catalog_name.clone(),
169            request.schema_name.clone(),
170            request.model_name.clone(),
171            request.version.to_string(),
172        ]));
173        // Authorize against the concrete model version and the operation actually
174        // requested, rather than the unscoped `SecuredAction` default.
175        self.authorize_checked(&ident, &required_permission(operation), &context)
176            .await?;
177        let (resource, _) = self.get(&ident).await?;
178        let model_version: ModelVersion = resource.try_into()?;
179        let location = model_version.storage_location.ok_or_else(|| {
180            Error::invalid_argument("model version does not have a storage location")
181        })?;
182        let storage_url = StorageLocationUrl::parse(&location)?;
183        let ext_loc = find_external_location_for_url(&storage_url, self).await?;
184        let credential = self
185            .get_credential_internal(GetCredentialRequest {
186                name: ext_loc.credential_name.clone(),
187                ..Default::default()
188            })
189            .await?;
190        vend_credential(&credential, &location, operation).await
191    }
192
193    #[tracing::instrument(skip(self, context))]
194    async fn generate_temporary_table_credentials(
195        &self,
196        request: GenerateTemporaryTableCredentialsRequest,
197        context: RequestContext,
198    ) -> Result<TemporaryCredential> {
199        let operation = to_vend_operation(request.operation.to_i32());
200        let table_id = uuid::Uuid::parse_str(&request.table_id)
201            .map_err(|_| Error::invalid_argument("table_id is not a valid UUID"))?;
202        // Authorize against the concrete table and the operation actually
203        // requested, rather than the unscoped `SecuredAction` default.
204        let table_ident = ResourceIdent::Table(ResourceRef::Uuid(table_id));
205        self.authorize_checked(&table_ident, &required_permission(operation), &context)
206            .await?;
207        let (resource, _) = self.get(&table_ident).await?;
208        let table: Table = resource.try_into()?;
209        let location = table
210            .storage_location
211            .ok_or_else(|| Error::invalid_argument("table does not have a storage location"))?;
212        let storage_url = StorageLocationUrl::parse(&location)?;
213        let ext_loc = find_external_location_for_url(&storage_url, self).await?;
214        let credential = self
215            .get_credential_internal(GetCredentialRequest {
216                name: ext_loc.credential_name.clone(),
217                ..Default::default()
218            })
219            .await?;
220        vend_credential(&credential, &location, operation).await
221    }
222}
223// NOTE: These request types intentionally do not implement `SecuredAction`.
224// Authorization is performed inside the handlers against the *concrete* resolved
225// resource (external location / table) and the *requested* operation
226// (read vs. read-write), which a static `SecuredAction` impl cannot express.
227// See `generate_temporary_path_credentials` / `generate_temporary_table_credentials`.
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    use generate_temporary_path_credentials_request::Operation as PathOp;
233    use generate_temporary_table_credentials_request::Operation as TableOp;
234    use generate_temporary_volume_credentials_request::Operation as VolumeOp;
235
236    #[test]
237    fn vend_operation_volume_read_is_read() {
238        assert_eq!(
239            to_vend_operation(VolumeOp::ReadVolume as i32),
240            VendOperation::Read
241        );
242    }
243
244    #[test]
245    fn vend_operation_volume_write_is_read_write() {
246        assert_eq!(
247            to_vend_operation(VolumeOp::WriteVolume as i32),
248            VendOperation::ReadWrite
249        );
250    }
251
252    #[test]
253    fn vend_operation_unspecified_defaults_to_read() {
254        assert_eq!(
255            to_vend_operation(VolumeOp::Unspecified as i32),
256            VendOperation::Read
257        );
258        // Out-of-range values are treated as least-privilege read, not an error.
259        assert_eq!(to_vend_operation(99), VendOperation::Read);
260    }
261
262    #[test]
263    fn vend_operation_table_and_path_unchanged() {
264        assert_eq!(to_vend_operation(TableOp::Read as i32), VendOperation::Read);
265        assert_eq!(
266            to_vend_operation(TableOp::ReadWrite as i32),
267            VendOperation::ReadWrite
268        );
269        assert_eq!(
270            to_vend_operation(PathOp::PathCreateTable as i32),
271            VendOperation::ReadWrite
272        );
273    }
274}
unitycatalog_server/api/temporary_credentials.rs

unitycatalog_server/api/
temporary_credentials.rs
