udb_client/
client.rs

1//! The data-plane client.
2//!
3//! Every method routes through [`UdbClient::request`], which is the single point
4//! that applies connection metadata. That is deliberate: if applying tenant scope
5//! were the caller's job, forgetting it once would be a cross-tenant read, and
6//! nothing in the type system would object.
7
8use tonic::transport::{Channel, Endpoint};
9use tonic::{Request, Status};
10
11use crate::error::{CallPolicy, UdbError};
12use crate::metadata::Metadata;
13use crate::proto::udb::entity::v1 as entity;
14use crate::proto::udb::services::v1::data_broker_client::DataBrokerClient;
15
16/// A connected UDB data-plane client bound to one identity.
17#[derive(Clone, Debug)]
18pub struct UdbClient {
19    inner: DataBrokerClient<Channel>,
20    meta: Metadata,
21    /// Overrides the per-RPC default when set. Reads default to a retrying
22    /// policy and mutations to a single attempt, so an override is only needed
23    /// to tighten a deadline or to opt a mutation into retries the CALLER knows
24    /// are safe (an idempotency key, say).
25    policy: Option<CallPolicy>,
26}
27
28/// Generate an RPC wrapper with identical retry and deadline handling.
29///
30/// A macro rather than a generic helper because the closure would have to hand a
31/// `&mut` client into an async block per attempt, and the lifetime dance costs
32/// more clarity than the repetition saves. `$default` is the policy used when the
33/// client carries no override.
34macro_rules! rpc {
35    ($(#[$m:meta])* $name:ident, $rpc:ident, $req:ty, $resp:ty, $path:expr) => {
36        $(#[$m])*
37        pub async fn $name(&mut self, req: $req) -> Result<$resp, UdbError> {
38            let policy = self.policy.unwrap_or_else(|| CallPolicy::from_contract($path));
39            let mut attempt: u32 = 1;
40            loop {
41                let mut request = self.request(req.clone())?;
42                if let Some(deadline) = policy.deadline {
43                    request.set_timeout(deadline);
44                }
45                match self.inner.$rpc(request).await {
46                    Ok(response) => return Ok(response.into_inner()),
47                    Err(status) => {
48                        let err = UdbError::from_status(status);
49                        if policy.should_retry(attempt, &err) {
50                            tokio::time::sleep(policy.backoff_for(attempt, &err)).await;
51                            attempt += 1;
52                            continue;
53                        }
54                        return Err(err);
55                    }
56                }
57            }
58        }
59    };
60}
61
62impl UdbClient {
63    /// Connect to a broker's data-plane listener (`:50051` by default).
64    ///
65    /// Note the port: UDB runs SEPARATE listeners with different authorization
66    /// models — the data plane authorizes through Casbin, while the native
67    /// service listener uses scope-based endpoint security. A credential accepted
68    /// by one is not automatically accepted by the other, and the mismatch
69    /// presents as a permissions error rather than a wrong-address error.
70    pub async fn connect(
71        endpoint: impl Into<String>,
72        meta: Metadata,
73    ) -> Result<Self, tonic::transport::Error> {
74        let channel = Endpoint::from_shared(endpoint.into())?.connect().await?;
75        Ok(Self::with_channel(channel, meta))
76    }
77
78    /// Wrap an already-configured channel — use this for TLS, custom timeouts,
79    /// load balancing, or interceptors.
80    pub fn with_channel(channel: Channel, meta: Metadata) -> Self {
81        Self {
82            inner: DataBrokerClient::new(channel),
83            meta,
84            policy: None,
85        }
86    }
87
88    /// The connection's metadata.
89    pub fn metadata(&self) -> &Metadata {
90        &self.meta
91    }
92
93    /// A copy of this client carrying request-scoped audit fields.
94    ///
95    /// Identity is intentionally not settable here; see [`Metadata`].
96    pub fn with_audit(
97        &self,
98        purpose: impl Into<String>,
99        correlation_id: impl Into<String>,
100    ) -> Self {
101        Self {
102            inner: self.inner.clone(),
103            meta: self.meta.clone().with_audit(purpose, correlation_id),
104            policy: self.policy,
105        }
106    }
107
108    /// Replace the bearer credential, e.g. after a token refresh.
109    pub fn with_bearer_token(&self, token: impl Into<String>) -> Self {
110        Self {
111            inner: self.inner.clone(),
112            meta: self.meta.clone().with_bearer_token(token),
113            policy: self.policy,
114        }
115    }
116
117    /// Wrap a message in a request carrying this connection's metadata.
118    ///
119    /// Public so callers can reach an RPC this wrapper does not expose yet
120    /// without hand-rolling — and without hand-rolling the headers.
121    pub fn request<T>(&self, message: T) -> Result<Request<T>, Status> {
122        let mut req = Request::new(message);
123        self.meta.apply(&mut req)?;
124        Ok(req)
125    }
126
127    /// The raw generated client, for RPCs this wrapper does not cover.
128    ///
129    /// Pair it with [`UdbClient::request`] so metadata is still applied.
130    pub fn raw(&mut self) -> &mut DataBrokerClient<Channel> {
131        &mut self.inner
132    }
133
134    /// Apply a deadline/retry policy to every call this client makes.
135    pub fn with_policy(&self, policy: CallPolicy) -> Self {
136        Self {
137            inner: self.inner.clone(),
138            meta: self.meta.clone(),
139            policy: Some(policy),
140        }
141    }
142
143    rpc!(
144        /// Read. Retry policy comes from the contract, not this wrapper.
145        select,
146        select,
147        entity::SelectRequest,
148        entity::RecordSet,
149        "/udb.services.v1.DataBroker/Select"
150    );
151    rpc!(
152        /// Write. The contract declares it replayable, so it retries.
153        upsert,
154        upsert,
155        entity::UpsertRequest,
156        entity::MutationResponse,
157        "/udb.services.v1.DataBroker/Upsert"
158    );
159    rpc!(
160        /// Write. Contract-declared replayable.
161        update,
162        update,
163        entity::UpdateRequest,
164        entity::MutationResponse,
165        "/udb.services.v1.DataBroker/Update"
166    );
167    rpc!(
168        /// Write. Contract-declared replayable.
169        delete,
170        delete,
171        entity::DeleteRequest,
172        entity::MutationResponse,
173        "/udb.services.v1.DataBroker/Delete"
174    );
175    rpc!(
176        /// Compare-and-swap. The contract does NOT declare it replayable: a CAS
177        /// that timed out may have committed, and a repeat would compare against
178        /// state its own first attempt wrote.
179        bulk_cas,
180        bulk_cas,
181        entity::BulkCasRequest,
182        entity::BulkCasResponse,
183        "/udb.services.v1.DataBroker/BulkCas"
184    );
185    rpc!(
186        /// Read.
187        vector_search,
188        vector_search,
189        entity::VectorSearchRequest,
190        entity::VectorSet,
191        "/udb.services.v1.DataBroker/VectorSearch"
192    );
193    rpc!(
194        /// Write. Not contract-declared replayable.
195        vector_upsert,
196        vector_upsert,
197        entity::VectorUpsertRequest,
198        entity::MutationResponse,
199        "/udb.services.v1.DataBroker/VectorUpsert"
200    );
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use crate::metadata::headers;
207
208    // `connect_lazy` registers with the Tokio reactor even though it dials
209    // nothing, so these must run inside a runtime. They exercise metadata
210    // assembly, not transport: nothing needs to be listening.
211    fn client() -> UdbClient {
212        let channel = Endpoint::from_static("http://127.0.0.1:50051").connect_lazy();
213        UdbClient::with_channel(
214            channel,
215            Metadata::new("tenant-1")
216                .with_project("proj-9")
217                .with_bearer_token("tok"),
218        )
219    }
220
221    #[tokio::test]
222    async fn request_carries_connection_identity() {
223        let c = client();
224        let req = c.request(()).expect("metadata applies");
225        let md = req.metadata();
226        assert_eq!(md.get(headers::TENANT_ID).unwrap(), "tenant-1");
227        assert_eq!(md.get(headers::PROJECT_ID).unwrap(), "proj-9");
228        assert_eq!(md.get(headers::AUTHORIZATION).unwrap(), "Bearer tok");
229    }
230
231    #[tokio::test]
232    async fn with_audit_keeps_identity_and_adds_audit() {
233        let c = client().with_audit("billing", "corr-7");
234        let req = c.request(()).expect("metadata applies");
235        let md = req.metadata();
236        assert_eq!(md.get(headers::TENANT_ID).unwrap(), "tenant-1");
237        assert_eq!(md.get(headers::PURPOSE).unwrap(), "billing");
238        assert_eq!(md.get(headers::CORRELATION_ID).unwrap(), "corr-7");
239    }
240
241    #[test]
242    fn wrapped_rpc_paths_exist_in_the_registry() {
243        // A typo in a path literal would make `from_contract` fall through to the
244        // conservative default and silently stop retrying that RPC — a change
245        // nothing else would catch.
246        for path in [
247            "/udb.services.v1.DataBroker/Select",
248            "/udb.services.v1.DataBroker/Upsert",
249            "/udb.services.v1.DataBroker/Update",
250            "/udb.services.v1.DataBroker/Delete",
251            "/udb.services.v1.DataBroker/BulkCas",
252            "/udb.services.v1.DataBroker/VectorSearch",
253            "/udb.services.v1.DataBroker/VectorUpsert",
254        ] {
255            assert!(
256                crate::generated_rpcs::spec_for_path(path).is_some(),
257                "{path} is not in the generated registry"
258            );
259        }
260    }
261
262    #[test]
263    fn contract_drives_retry_and_disagrees_with_naive_naming() {
264        use crate::generated_rpcs::is_retry_safe;
265        // Reads: obviously repeatable.
266        assert!(is_retry_safe("/udb.services.v1.DataBroker/Select"));
267        // Mutations the CONTRACT declares replayable. A name-based guess would
268        // refuse these, which is exactly the mistake this replaced.
269        assert!(is_retry_safe("/udb.services.v1.DataBroker/Upsert"));
270        assert!(is_retry_safe("/udb.services.v1.DataBroker/Delete"));
271        // And ones it does not.
272        assert!(!is_retry_safe("/udb.services.v1.DataBroker/BulkCas"));
273        assert!(!is_retry_safe("/udb.services.v1.DataBroker/VectorUpsert"));
274    }
275
276    #[tokio::test]
277    async fn with_bearer_token_replaces_only_the_credential() {
278        let c = client().with_bearer_token("rotated");
279        assert_eq!(c.metadata().tenant_id, "tenant-1");
280        let req = c.request(()).expect("metadata applies");
281        assert_eq!(
282            req.metadata().get(headers::AUTHORIZATION).unwrap(),
283            "Bearer rotated"
284        );
285    }
286}
udb_client/client.rs

udb_client/
client.rs
