1#![deny(missing_docs)]
2use std::future::Future;
5use std::marker::PhantomData;
6use std::pin::Pin;
7use std::sync::Arc;
8
9use type_bridge_contract::query_remote::RemoteCapabilities;
10use type_bridge_contract::query_remote_v2::RemoteLimitsV2;
11use type_bridge_orm::query_v2_prepared::QueryAuthority;
12use type_bridge_orm::{
13 DescriptorRegistry, InstalledRuntimeProjection, RemoteModelQueryV2Error, ValidatedMatchRequest,
14 ValidatedMatchResult, prepare_remote_model_query_v2,
15};
16
17use crate::Result;
18use crate::error::Error;
19use crate::query::QuerySession;
20use crate::schema::{Schema, SchemaPackage, Unbound};
21
22pub trait RemoteQueryTransport: Send + Sync + 'static {
27 fn capabilities(&self) -> Pin<Box<dyn Future<Output = Result<Vec<u8>>> + Send + '_>>;
29
30 fn exchange<'a>(
32 &'a self,
33 request: &'a [u8],
34 ) -> Pin<Box<dyn Future<Output = Result<Vec<u8>>> + Send + 'a>>;
35}
36
37#[derive(Clone, Copy, Debug, Eq, PartialEq)]
39pub struct RemoteQueryLimits {
40 limits: RemoteLimitsV2,
41}
42
43impl RemoteQueryLimits {
44 #[must_use]
46 pub const fn new(
47 max_items: u64,
48 max_bytes: u64,
49 max_collection_members: u64,
50 max_graph_nodes: u64,
51 max_attribute_values: u64,
52 max_role_players: u64,
53 ) -> Self {
54 Self {
55 limits: RemoteLimitsV2 {
56 deadline_ms: None,
57 max_bytes,
58 max_items,
59 max_collection_members,
60 max_graph_nodes,
61 max_attribute_values,
62 max_role_players,
63 },
64 }
65 }
66
67 #[must_use]
69 pub const fn deadline_ms(mut self, deadline_ms: u64) -> Self {
70 self.limits.deadline_ms = Some(deadline_ms);
71 self
72 }
73}
74
75pub struct RemoteConnectionOptions {
77 scope: String,
78 semantic_profile: String,
79 limits: RemoteQueryLimits,
80 transport: Arc<dyn RemoteQueryTransport>,
81 advertisement: Option<Vec<u8>>,
82}
83
84impl RemoteConnectionOptions {
85 #[must_use]
88 pub fn new(
89 scope: impl Into<String>,
90 semantic_profile: impl Into<String>,
91 limits: RemoteQueryLimits,
92 transport: impl RemoteQueryTransport,
93 ) -> Self {
94 Self {
95 scope: scope.into(),
96 semantic_profile: semantic_profile.into(),
97 limits,
98 transport: Arc::new(transport),
99 advertisement: None,
100 }
101 }
102}
103
104struct RemoteRuntime {
105 advertisement: Vec<u8>,
106 authority: Arc<QueryAuthority>,
107 limits: RemoteLimitsV2,
108 transport: Arc<dyn RemoteQueryTransport>,
109}
110
111pub struct RemoteDatabase<S: Schema = Unbound> {
113 options: Option<RemoteConnectionOptions>,
114 runtime: Option<Arc<RemoteRuntime>>,
115 installed: Option<Arc<InstalledRuntimeProjection>>,
116 registry: Option<Arc<DescriptorRegistry>>,
117 marker: PhantomData<fn() -> S>,
118}
119
120impl<S: Schema> std::fmt::Debug for RemoteDatabase<S> {
121 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
122 formatter
123 .debug_struct("RemoteDatabase")
124 .field("schema_bound", &self.installed.is_some())
125 .finish_non_exhaustive()
126 }
127}
128
129impl RemoteDatabase<Unbound> {
130 pub async fn connect(mut options: RemoteConnectionOptions) -> Result<Self> {
132 let advertisement = options.transport.capabilities().await?;
133 RemoteCapabilities::decode(&advertisement).map_err(remote_diagnostic)?;
134 options.advertisement = Some(advertisement);
135 Ok(Self {
136 options: Some(options),
137 runtime: None,
138 installed: None,
139 registry: None,
140 marker: PhantomData,
141 })
142 }
143
144 pub fn with_schema<S: Schema>(mut self, schema: SchemaPackage<S>) -> Result<RemoteDatabase<S>> {
146 let declared = schema
147 .declared_schema_json()
148 .ok_or_else(|| Error::SchemaVerification {
149 message: "generated schema package omits remote declared-schema authority".into(),
150 source: None,
151 })?;
152 let installed = schema.verify_and_install()?;
153 let registry = Arc::new(installed.match_registry().map_err(Error::from_orm)?);
154 let options = self.options.take().ok_or_else(|| Error::Other {
155 message: "remote connection options are unavailable".into(),
156 source: None,
157 })?;
158 let authority = QueryAuthority::from_declared_bytes(
159 declared.as_bytes(),
160 &options.scope,
161 &options.semantic_profile,
162 )
163 .map_err(remote_diagnostic)?;
164 if !authority.matches_semantic_fingerprint(installed.projection().semantic_fingerprint()) {
165 return Err(Error::SchemaVerification {
166 message: "remote declared-schema authority does not match the generated projection"
167 .into(),
168 source: None,
169 });
170 }
171 let runtime = Arc::new(RemoteRuntime {
172 advertisement: options.advertisement.ok_or_else(|| Error::Other {
173 message: "remote capability advertisement is unavailable".into(),
174 source: None,
175 })?,
176 authority: Arc::new(authority),
177 limits: options.limits.limits,
178 transport: options.transport,
179 });
180 Ok(RemoteDatabase {
181 options: None,
182 runtime: Some(runtime),
183 installed: Some(installed),
184 registry: Some(registry),
185 marker: PhantomData,
186 })
187 }
188}
189
190impl<S: Schema> RemoteDatabase<S> {
191 pub fn query(&self) -> Result<QuerySession<'_, S>> {
193 let installed = self.installed.as_deref().ok_or_else(remote_not_bound)?;
194 let registry = self.registry.as_ref().ok_or_else(remote_not_bound)?;
195 Ok(QuerySession::remote(installed, Arc::clone(registry), self))
196 }
197
198 pub(crate) async fn execute_match(
199 &self,
200 registry: &DescriptorRegistry,
201 validated: ValidatedMatchRequest,
202 ) -> Result<(ValidatedMatchRequest, ValidatedMatchResult)> {
203 let runtime = self.runtime.as_ref().ok_or_else(remote_not_bound)?;
204 let pending = prepare_remote_model_query_v2(
205 &runtime.authority,
206 registry,
207 validated,
208 &runtime.advertisement,
209 runtime.limits,
210 )
211 .map_err(remote_model_input_error)?;
212 let request = pending.request_bytes().to_vec();
213 let response = runtime.transport.exchange(&request).await?;
214 let claimed = pending
215 .claim_reply()
216 .map_err(remote_model_hydration_error)?;
217 if response.len() > claimed.response_snapshot_limit() {
218 return Err(Error::classified(
219 crate::ErrorCategory::ResourceLimit,
220 None,
221 "remote_response_limit",
222 Vec::new(),
223 "remote query reply exceeds the authenticated response ceiling",
224 None,
225 ));
226 }
227 let (request, result, _registry) = claimed
228 .decode(&response)
229 .map_err(remote_model_hydration_error)?;
230 Ok((request, result))
231 }
232}
233
234fn remote_not_bound() -> Error {
235 Error::ModelValidation {
236 phase: crate::ModelValidationPhase::Input,
237 code: "schema_not_bound".into(),
238 path: vec![],
239 message: "remote database is not schema-bound".into(),
240 source: None,
241 }
242}
243
244fn remote_diagnostic(error: type_bridge_contract::diagnostic::Diagnostic) -> Error {
245 Error::from_remote_diagnostic(error)
246}
247
248fn remote_model_input_error(error: RemoteModelQueryV2Error) -> Error {
249 match error {
250 RemoteModelQueryV2Error::Diagnostic(error) => Error::from_remote_diagnostic(error),
251 RemoteModelQueryV2Error::Match(error) => {
252 Error::from_match(error, crate::ModelValidationPhase::Input)
253 }
254 }
255}
256
257fn remote_model_hydration_error(error: RemoteModelQueryV2Error) -> Error {
258 match error {
259 RemoteModelQueryV2Error::Diagnostic(error) => Error::from_remote_diagnostic(error),
260 RemoteModelQueryV2Error::Match(error) => {
261 Error::from_match(error, crate::ModelValidationPhase::Hydration)
262 }
263 }
264}
265
266#[cfg(test)]
267mod tests {
268 use std::sync::Mutex;
269
270 use type_bridge_contract::codec::to_canonical_json;
271 use type_bridge_contract::diagnostic::{
272 Diagnostic, DiagnosticCategory, DiagnosticCode, DiagnosticPathSegment,
273 };
274 use type_bridge_contract::fingerprint::SemanticProfileId;
275 use type_bridge_contract::migration_assertion::BindingId;
276 use type_bridge_contract::projection::{BindingTarget, ProjectionConfig};
277 use type_bridge_contract::query_plan::query_plan_v2_capability_vocabulary;
278 use type_bridge_contract::query_remote::RemoteExecutorBinding;
279 use type_bridge_contract::query_remote_v2::{
280 HydrationGraphV2, RemoteOutcomeV2, RemoteQueryRequestV2, RemoteQueryResponseV2,
281 RemoteResultKindV2, query_remote_v2_required_capabilities,
282 };
283 use type_bridge_contract::schema::{DocumentId, encode_declared_schema};
284 use type_bridge_orm::OrmError;
285 use type_bridge_orm::match_request::SessionHandle;
286 use type_bridge_orm::query_v2_remote::RemoteReplySigningKey;
287 use type_bridge_schema::{SchemaDocumentSet, normalize_documents, project, resolve};
288 use type_bridge_schema_codegen::RustEmitter;
289
290 use super::*;
291 use crate::__codegen::{
292 self, CompleteModel, EncodedCreate, EntityModel, HydratedRow, HydrationCapability,
293 IntoEncodedCreate, MaterializeModel, Model, ThingModel, ValidationError,
294 };
295 use crate::schema::sealed;
296
297 struct TestSchema;
298 impl sealed::Sealed for TestSchema {}
299 impl Schema for TestSchema {}
300
301 struct Person;
302 impl sealed::Sealed for Person {}
303 impl Model for Person {
304 type Schema = TestSchema;
305 const TYPE_ID_JSON: &'static str = r#"{"kind":"entity","label":"person"}"#;
306 }
307 impl ThingModel for Person {
308 fn thing_kind() -> __codegen::ThingKind {
309 __codegen::ThingKind::Entity
310 }
311 }
312 impl EntityModel for Person {}
313 impl CompleteModel for Person {
314 type Create = PersonCreate;
315
316 fn iid(&self) -> &str {
317 unreachable!()
318 }
319 }
320 impl MaterializeModel for Person {
321 fn materialize(
322 _: &HydratedRow,
323 _: &HydrationCapability,
324 ) -> std::result::Result<Self, ValidationError> {
325 Ok(Self)
326 }
327 }
328
329 struct PersonCreate;
330 impl sealed::Sealed for PersonCreate {}
331 impl IntoEncodedCreate for PersonCreate {
332 fn into_encoded_create(self) -> std::result::Result<EncodedCreate, ValidationError> {
333 Ok(EncodedCreate::new(Person::TYPE_ID_JSON, vec![], vec![]))
334 }
335 }
336
337 #[test]
338 fn local_and_remote_failures_preserve_classification_codes_and_paths() {
339 let session = SessionHandle::new(Arc::new(DescriptorRegistry::new()));
340 let match_error = match session.exact("missing") {
341 Err(OrmError::Match(error)) => error,
342 Err(other) => panic!("unexpected ORM error: {other:?}"),
343 Ok(_) => panic!("missing descriptor unexpectedly resolved"),
344 };
345 let local = Error::from_orm(OrmError::Match(match_error.clone()));
346 let remote = remote_model_input_error(RemoteModelQueryV2Error::Match(match_error));
347
348 assert_eq!(local.category(), crate::ErrorCategory::QueryAuthoring);
349 assert_eq!(remote.category(), local.category());
350 assert_eq!(remote.code(), Some("unknown_descriptor"));
351 assert_eq!(remote.code(), local.code());
352 assert_eq!(remote.path(), local.path());
353 assert_eq!(
354 remote.model_validation_phase(),
355 local.model_validation_phase()
356 );
357
358 let diagnostic = Diagnostic::new(
359 DiagnosticCategory::UnsupportedCapability,
360 DiagnosticCode::new("missing_remote_capability").unwrap(),
361 "the remote executor does not advertise one required capability",
362 )
363 .at(DiagnosticPathSegment::Field("capabilities".into()))
364 .at(DiagnosticPathSegment::Index(2));
365 let classified = remote_diagnostic(diagnostic);
366
367 assert_eq!(classified.category(), crate::ErrorCategory::Capability);
368 assert_eq!(classified.code(), Some("missing_remote_capability"));
369 assert_eq!(
370 classified.path(),
371 Some(&["capabilities".to_owned(), "[2]".to_owned()][..])
372 );
373 assert_eq!(classified.model_validation_phase(), None);
374 }
375
376 fn package() -> SchemaPackage<TestSchema> {
377 let documents = SchemaDocumentSet::parse([(
378 DocumentId::new("remote.yaml").unwrap(),
379 "format: typebridge.schema/v2\nattributes:\n name: { value: string }\nentities:\n person:\n owns: { name: { key: true } }\n",
380 )])
381 .unwrap();
382 let declared = normalize_documents(&documents).unwrap();
383 let resolved = resolve(
384 &declared,
385 &SemanticProfileId::new("typedb-3.12.1/v1").unwrap(),
386 )
387 .unwrap();
388 let emitter = RustEmitter::new();
389 let projection = project(
390 &resolved,
391 BindingTarget::Rust,
392 &ProjectionConfig::rust(),
393 &emitter.generator_handlers(),
394 &emitter.code_resources().unwrap(),
395 )
396 .unwrap();
397 let leak = |bytes: Vec<u8>| {
398 Box::leak(String::from_utf8(bytes).unwrap().into_boxed_str()) as &'static str
399 };
400 SchemaPackage::new_with_declared(
401 leak(to_canonical_json(projection.semantic_fingerprint()).unwrap()),
402 leak(to_canonical_json(projection.projection_fingerprint()).unwrap()),
403 leak(to_canonical_json(&projection).unwrap()),
404 leak(encode_declared_schema(&declared).unwrap()),
405 )
406 }
407
408 struct Transport {
409 advertisement_contract: RemoteCapabilities,
410 advertisement: Vec<u8>,
411 capabilities: Arc<Mutex<usize>>,
412 exchanges: Arc<Mutex<Vec<Vec<u8>>>>,
413 signer: RemoteReplySigningKey,
414 }
415
416 impl RemoteQueryTransport for Transport {
417 fn capabilities(&self) -> Pin<Box<dyn Future<Output = Result<Vec<u8>>> + Send + '_>> {
418 *self.capabilities.lock().unwrap() += 1;
419 let bytes = self.advertisement.clone();
420 Box::pin(async move { Ok(bytes) })
421 }
422
423 fn exchange<'a>(
424 &'a self,
425 request: &'a [u8],
426 ) -> Pin<Box<dyn Future<Output = Result<Vec<u8>>> + Send + 'a>> {
427 self.exchanges.lock().unwrap().push(request.to_vec());
428 let response = (|| {
429 let request = RemoteQueryRequestV2::decode(request).map_err(remote_diagnostic)?;
430 request
431 .validate_advertisement(&self.advertisement_contract)
432 .map_err(remote_diagnostic)?;
433 let plan = request.plan().map_err(remote_diagnostic)?;
434 let root = BindingId::new(0).map_err(remote_diagnostic)?;
435 let outcome = match request.result_kind() {
436 RemoteResultKindV2::DistinctCount => {
437 RemoteOutcomeV2::DistinctCount { root, value: 7 }
438 }
439 RemoteResultKindV2::DistinctExists => {
440 RemoteOutcomeV2::DistinctExists { root, value: true }
441 }
442 RemoteResultKindV2::HydratedRows => RemoteOutcomeV2::HydratedRows {
443 graph: HydrationGraphV2::new(vec![]).map_err(remote_diagnostic)?,
444 rows: vec![],
445 },
446 RemoteResultKindV2::HydratedPage => RemoteOutcomeV2::HydratedPage {
447 entries: vec![],
448 graph: HydrationGraphV2::new(vec![]).map_err(remote_diagnostic)?,
449 limit: 2,
450 offset: 0,
451 root,
452 total: Some(0),
453 },
454 _ => {
455 return Err(Error::Other {
456 message: "test transport received an unexpected terminal".into(),
457 source: None,
458 });
459 }
460 };
461 RemoteQueryResponseV2::new(
462 request.nonce(),
463 &plan,
464 &request.fingerprint().map_err(remote_diagnostic)?,
465 request.result_kind(),
466 outcome,
467 )
468 .and_then(|response| {
469 response
470 .encode_signed(&self.advertisement_contract.fingerprint()?, &self.signer)
471 })
472 .map_err(remote_diagnostic)
473 })();
474 Box::pin(async move { response })
475 }
476 }
477
478 #[tokio::test]
479 async fn remote_database_fetches_capabilities_once_and_exchanges_once_per_terminal() {
480 let signer = RemoteReplySigningKey::from_secret_bytes([0x31; 32]);
481 let mut capabilities = query_plan_v2_capability_vocabulary();
482 for capability in query_remote_v2_required_capabilities(true) {
483 capabilities.insert(capability);
484 }
485 let advertisement_contract = RemoteCapabilities::new(
486 capabilities,
487 RemoteExecutorBinding::new("rust-client-test", "epoch-00000000001").unwrap(),
488 signer.public_key(),
489 );
490 let advertisement = advertisement_contract.encode().unwrap();
491 let capability_calls = Arc::new(Mutex::new(0));
492 let exchanges = Arc::new(Mutex::new(Vec::new()));
493 let transport = Transport {
494 advertisement_contract,
495 advertisement,
496 capabilities: Arc::clone(&capability_calls),
497 exchanges: Arc::clone(&exchanges),
498 signer,
499 };
500 let options = RemoteConnectionOptions::new(
501 "rust-client-test",
502 "typedb-3.12.1/v1",
503 RemoteQueryLimits::new(10, 1 << 20, 10, 100, 100, 100),
504 transport,
505 );
506 let remote = RemoteDatabase::connect(options)
507 .await
508 .unwrap()
509 .with_schema(package())
510 .unwrap();
511 let mut session = remote.query().unwrap();
512 let person = session.exact::<Person>().unwrap();
513 let query = session.query(person).unwrap();
514
515 assert_eq!(query.count().await.unwrap(), 7);
516 assert!(query.exists().await.unwrap());
517 assert!(
518 query
519 .rows(crate::RowsOptions::new(2))
520 .await
521 .unwrap()
522 .is_empty()
523 );
524 let page = query
525 .page_by(person, crate::PageOptions::new(2).include_total(true))
526 .await
527 .unwrap();
528 assert!(page.items().is_empty());
529 assert_eq!(page.total(), Some(0));
530 let error = query
531 .aggregate((crate::aggregate::count(),))
532 .await
533 .expect_err("native-only reductions fail before transport exchange");
534 assert!(
535 error
536 .to_string()
537 .contains("query_remote_v2_native_only_operation")
538 );
539 assert_eq!(*capability_calls.lock().unwrap(), 1);
540 let requests = exchanges.lock().unwrap();
541 assert_eq!(requests.len(), 4);
542 assert!(
543 std::str::from_utf8(&requests[0])
544 .unwrap()
545 .contains("\"format\":\"typebridge.query-remote-request/v2\"")
546 );
547 }
548}