Skip to main content

type_bridge/
remote.rs

1#![deny(missing_docs)]
2//! Authenticated one-exchange remote execution for generated queries.
3
4use 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
22/// One caller-owned asynchronous transport for the authenticated V2 routes.
23///
24/// Implementations fetch the exact `/v2/capabilities` bytes once at connect
25/// time and perform exactly one `/v2/query` exchange per terminal.
26pub trait RemoteQueryTransport: Send + Sync + 'static {
27    /// Fetch the executor's exact signed capability advertisement.
28    fn capabilities(&self) -> Pin<Box<dyn Future<Output = Result<Vec<u8>>> + Send + '_>>;
29
30    /// Exchange one exact canonical request for one exact signed reply.
31    fn exchange<'a>(
32        &'a self,
33        request: &'a [u8],
34    ) -> Pin<Box<dyn Future<Output = Result<Vec<u8>>> + Send + 'a>>;
35}
36
37/// Explicit immutable budgets for one remote generated-query terminal.
38#[derive(Clone, Copy, Debug, Eq, PartialEq)]
39pub struct RemoteQueryLimits {
40    limits: RemoteLimitsV2,
41}
42
43impl RemoteQueryLimits {
44    /// Construct one explicit remote response and hydration budget.
45    #[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    /// Attach an optional executor deadline in milliseconds.
68    #[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
75/// Connection-time authority, transport, and limit configuration.
76pub 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    /// Construct remote options for one managed schema scope and semantic
86    /// profile.
87    #[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
111/// A client-owned remote generated-query database branded by schema `S`.
112pub 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    /// Fetch and validate one immutable executor advertisement.
131    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    /// Verify and bind one generated schema package and its remote authority.
145    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    /// Start one owner-branded query session over this remote executor.
192    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}