lance_namespace_impls/
context.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Dynamic context provider for per-request context overrides.
5//!
6//! This module provides the [`DynamicContextProvider`] trait that enables
7//! per-request context injection (e.g., dynamic authentication headers).
8//!
9//! ## Usage
10//!
11//! Implement the trait and pass to namespace builders:
12//!
13//! ```ignore
14//! use lance_namespace_impls::{RestNamespaceBuilder, DynamicContextProvider, OperationInfo};
15//! use std::collections::HashMap;
16//! use std::sync::Arc;
17//!
18//! #[derive(Debug)]
19//! struct MyProvider;
20//!
21//! impl DynamicContextProvider for MyProvider {
22//!     fn provide_context(&self, info: &OperationInfo) -> HashMap<String, String> {
23//!         let mut context = HashMap::new();
24//!         context.insert("headers.Authorization".to_string(), format!("Bearer {}", get_current_token()));
25//!         context.insert("headers.X-Request-Id".to_string(), generate_request_id());
26//!         context
27//!     }
28//! }
29//!
30//! let namespace = RestNamespaceBuilder::new("https://api.example.com")
31//!     .context_provider(Arc::new(MyProvider))
32//!     .build();
33//! ```
34//!
35//! For RestNamespace, context keys that start with `headers.` are converted to HTTP headers
36//! by stripping the prefix. For example, `{"headers.Authorization": "Bearer abc123"}`
37//! becomes the `Authorization: Bearer abc123` header. Keys without the `headers.` prefix
38//! are ignored for HTTP headers but may be used for other purposes.
39
40use std::collections::HashMap;
41
42/// Information about the namespace operation being executed.
43///
44/// This is passed to the [`DynamicContextProvider`] to allow it to make
45/// context decisions based on the operation.
46#[derive(Debug, Clone)]
47pub struct OperationInfo {
48    /// The operation name (e.g., "list_tables", "describe_table", "create_namespace")
49    pub operation: String,
50    /// The object ID for the operation (namespace or table identifier).
51    /// This is the delimited string form, e.g., "workspace$table_name".
52    pub object_id: String,
53}
54
55impl OperationInfo {
56    /// Create a new OperationInfo.
57    pub fn new(operation: impl Into<String>, object_id: impl Into<String>) -> Self {
58        Self {
59            operation: operation.into(),
60            object_id: object_id.into(),
61        }
62    }
63}
64
65/// Trait for providing dynamic request context.
66///
67/// Implementations can generate per-request context (e.g., authentication headers)
68/// based on the operation being performed. The provider is called synchronously
69/// before each namespace operation.
70///
71/// For RestNamespace, context keys that start with `headers.` are converted to
72/// HTTP headers by stripping the prefix. For example, `{"headers.Authorization": "Bearer token"}`
73/// becomes the `Authorization: Bearer token` header.
74///
75/// ## Thread Safety
76///
77/// Implementations must be `Send + Sync` as the provider may be called from
78/// multiple threads concurrently.
79///
80/// ## Error Handling
81///
82/// If the provider needs to signal an error, it should return an empty HashMap
83/// and log the error. The namespace operation will proceed without the
84/// additional context.
85pub trait DynamicContextProvider: Send + Sync + std::fmt::Debug {
86    /// Provide context for a namespace operation.
87    ///
88    /// # Arguments
89    ///
90    /// * `info` - Information about the operation being performed
91    ///
92    /// # Returns
93    ///
94    /// Returns a HashMap of context key-value pairs. For HTTP headers, use keys
95    /// with the `headers.` prefix (e.g., `headers.Authorization`).
96    /// Returns an empty HashMap if no additional context is needed.
97    fn provide_context(&self, info: &OperationInfo) -> HashMap<String, String>;
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[derive(Debug)]
105    struct MockContextProvider {
106        prefix: String,
107    }
108
109    impl DynamicContextProvider for MockContextProvider {
110        fn provide_context(&self, info: &OperationInfo) -> HashMap<String, String> {
111            let mut context = HashMap::new();
112            context.insert(
113                "test-header".to_string(),
114                format!("{}-{}", self.prefix, info.operation),
115            );
116            context.insert("object-id".to_string(), info.object_id.clone());
117            context
118        }
119    }
120
121    #[test]
122    fn test_operation_info_creation() {
123        let info = OperationInfo::new("describe_table", "workspace$my_table");
124        assert_eq!(info.operation, "describe_table");
125        assert_eq!(info.object_id, "workspace$my_table");
126    }
127
128    #[test]
129    fn test_context_provider_basic() {
130        let provider = MockContextProvider {
131            prefix: "test".to_string(),
132        };
133
134        let info = OperationInfo::new("list_tables", "workspace$ns");
135
136        let context = provider.provide_context(&info);
137        assert_eq!(
138            context.get("test-header"),
139            Some(&"test-list_tables".to_string())
140        );
141        assert_eq!(context.get("object-id"), Some(&"workspace$ns".to_string()));
142    }
143
144    #[test]
145    fn test_empty_context() {
146        #[derive(Debug)]
147        struct EmptyProvider;
148
149        impl DynamicContextProvider for EmptyProvider {
150            fn provide_context(&self, _info: &OperationInfo) -> HashMap<String, String> {
151                HashMap::new()
152            }
153        }
154
155        let provider = EmptyProvider;
156        let info = OperationInfo::new("list_tables", "ns");
157
158        let context = provider.provide_context(&info);
159        assert!(context.is_empty());
160    }
161}
lance_namespace_impls/context.rs

lance_namespace_impls/
context.rs
