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`.
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 `generate_temporary_volume_credentials`.
155    ///
156    /// See: <https://docs.databricks.com/api/workspace/temporarymodelversioncredentials/generatetemporarymodelversioncredentials>
157    #[tracing::instrument(skip(self, context))]
158    async fn generate_temporary_model_version_credentials(
159        &self,
160        request: GenerateTemporaryModelVersionCredentialsRequest,
161        context: RequestContext,
162    ) -> Result<TemporaryCredential> {
163        let operation = to_vend_operation(request.operation.to_i32());
164        // The model version is keyed by the composite name that
165        // `ModelVersion::resource_name` produces.
166        let ident = ResourceIdent::model_version(ResourceName::new([
167            request.catalog_name.clone(),
168            request.schema_name.clone(),
169            request.model_name.clone(),
170            request.version.to_string(),
171        ]));
172        // Authorize against the concrete model version and the operation actually
173        // requested, rather than the unscoped `SecuredAction` default.
174        self.authorize_checked(&ident, &required_permission(operation), &context)
175            .await?;
176        let (resource, _) = self.get(&ident).await?;
177        let model_version: ModelVersion = resource.try_into()?;
178        let location = model_version.storage_location.ok_or_else(|| {
179            Error::invalid_argument("model version does not have a storage location")
180        })?;
181        let storage_url = StorageLocationUrl::parse(&location)?;
182        let ext_loc = find_external_location_for_url(&storage_url, self).await?;
183        let credential = self
184            .get_credential_internal(GetCredentialRequest {
185                name: ext_loc.credential_name.clone(),
186                ..Default::default()
187            })
188            .await?;
189        vend_credential(&credential, &location, operation).await
190    }
191
192    #[tracing::instrument(skip(self, context))]
193    async fn generate_temporary_table_credentials(
194        &self,
195        request: GenerateTemporaryTableCredentialsRequest,
196        context: RequestContext,
197    ) -> Result<TemporaryCredential> {
198        let operation = to_vend_operation(request.operation.to_i32());
199        let table_id = uuid::Uuid::parse_str(&request.table_id)
200            .map_err(|_| Error::invalid_argument("table_id is not a valid UUID"))?;
201        // Authorize against the concrete table and the operation actually
202        // requested, rather than the unscoped `SecuredAction` default.
203        let table_ident = ResourceIdent::Table(ResourceRef::Uuid(table_id));
204        self.authorize_checked(&table_ident, &required_permission(operation), &context)
205            .await?;
206        let (resource, _) = self.get(&table_ident).await?;
207        let table: Table = resource.try_into()?;
208        let location = table
209            .storage_location
210            .ok_or_else(|| Error::invalid_argument("table does not have a storage location"))?;
211        let storage_url = StorageLocationUrl::parse(&location)?;
212        let ext_loc = find_external_location_for_url(&storage_url, self).await?;
213        let credential = self
214            .get_credential_internal(GetCredentialRequest {
215                name: ext_loc.credential_name.clone(),
216                ..Default::default()
217            })
218            .await?;
219        vend_credential(&credential, &location, operation).await
220    }
221}
222// NOTE: These request types intentionally do not implement `SecuredAction`.
223// Authorization is performed inside the handlers against the *concrete* resolved
224// resource (external location / table) and the *requested* operation
225// (read vs. read-write), which a static `SecuredAction` impl cannot express.
226// See `generate_temporary_path_credentials` / `generate_temporary_table_credentials`.
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231    use generate_temporary_path_credentials_request::Operation as PathOp;
232    use generate_temporary_table_credentials_request::Operation as TableOp;
233    use generate_temporary_volume_credentials_request::Operation as VolumeOp;
234
235    #[test]
236    fn vend_operation_volume_read_is_read() {
237        assert_eq!(
238            to_vend_operation(VolumeOp::ReadVolume as i32),
239            VendOperation::Read
240        );
241    }
242
243    #[test]
244    fn vend_operation_volume_write_is_read_write() {
245        assert_eq!(
246            to_vend_operation(VolumeOp::WriteVolume as i32),
247            VendOperation::ReadWrite
248        );
249    }
250
251    #[test]
252    fn vend_operation_unspecified_defaults_to_read() {
253        assert_eq!(
254            to_vend_operation(VolumeOp::Unspecified as i32),
255            VendOperation::Read
256        );
257        // Out-of-range values are treated as least-privilege read, not an error.
258        assert_eq!(to_vend_operation(99), VendOperation::Read);
259    }
260
261    #[test]
262    fn vend_operation_table_and_path_unchanged() {
263        assert_eq!(to_vend_operation(TableOp::Read as i32), VendOperation::Read);
264        assert_eq!(
265            to_vend_operation(TableOp::ReadWrite as i32),
266            VendOperation::ReadWrite
267        );
268        assert_eq!(
269            to_vend_operation(PathOp::PathCreateTable as i32),
270            VendOperation::ReadWrite
271        );
272    }
273}
unitycatalog_server/api/temporary_credentials.rs

unitycatalog_server/api/
temporary_credentials.rs
