use std::net::SocketAddr;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::net::TcpStream;
use crate::error::{BrainError, Result};
use crate::mux::HandshakeOutcome;
use crate::mux::{MuxConnection, Subscription};
use crate::wire::cbor::{f32_slice_to_le_bytes, from_cbor_bytes, to_cbor_bytes};
use crate::wire::opcode::Opcode;
use crate::wire::types::{
AnswerKindWire, AuthCredentials, AuthMethod, AuthPayload, CancelStreamAck, CancelStreamRequest,
CancellationReason, EncodeRequest, EncodeResponse, EncodeVectorDirectRequest,
EntityCreateRequest, EntityCreateResponse, EntityGetRequest, EntityGetResponse, EntityListItem,
EntityListRequest, EntityListResponseFrame, EntityMergeRequest, EntityMergeResponse,
EntityRenameRequest, EntityRenameResponse, EntityResolveRequest, EntityResolveResponse,
EntityTombstoneRequest, EntityTombstoneResponse, EntityUnmergeRequest, EntityUnmergeResponse,
EntityUpdateRequest, EntityUpdateResponse, ExtractorListRequest, ExtractorListResponseFrame,
ForgetRequest, ForgetResponse, GetCapabilitiesRequest, GetCapabilitiesResponse, GraphEdge,
GraphFetchRequest, GraphFetchResponseFrame, GraphNode, HelloCapabilities, HelloPayload,
InferenceStep, LinkRequest, LinkResponse, MaterializeProceduralRequest,
MaterializeProceduralResponse, MemoryInspectRequest, MemoryInspectResponse, MemoryListItem,
MemoryListRequest, MemoryListResponseFrame, MemoryResult, MtlsClaim, PingRequest, PlanRequest,
PlanResponseFrame, PlanStep, PongResponse, QueryExplainRequest, QueryExplainResponse,
QueryTraceRequest, QueryTraceResponse, ReasonRequest, ReasonResponseFrame, RecallRequest,
RecallResponseFrame, RelationCreateRequest, RelationCreateResponse, RelationGetRequest,
RelationGetResponse, RelationListFromRequest, RelationListFromResponseFrame,
RelationListToRequest, RelationListToResponseFrame, RelationSupersedeRequest,
RelationSupersedeResponse, RelationTombstoneRequest, RelationTombstoneResponse,
RelationTraverseRequest, RelationTraverseResponseFrame, RelationView, SchemaDropRequest,
SchemaDropResponse, SchemaGetRequest, SchemaGetResponse, SchemaListItemWire, SchemaListRequest,
SchemaListResponseFrame, SchemaReplaceRequest, SchemaReplaceResponse, SchemaUploadRequest,
SchemaUploadResponse, SchemaValidateRequest, SchemaValidateResponse, ServerFeatures,
SessionCreateRequest, SessionCreateResponse, SessionDeleteRequest, SessionDeleteResponse,
SessionListRequest, SessionListResponse, SpaceCreateRequest, SpaceCreateResponse,
SpaceDeleteRequest, SpaceDeleteResponse, SpaceListRequest, SpaceListResponse, SpacePermissions,
StatementCreateRequest, StatementCreateResponse, StatementGetRequest, StatementGetResponse,
StatementHistoryRequest, StatementHistoryResponseFrame, StatementListRequest,
StatementListResponseFrame, StatementRetractRequest, StatementRetractResponse,
StatementSupersedeRequest, StatementSupersedeResponse, StatementTombstoneRequest,
StatementTombstoneResponse, StatementView, SubscribeRequest, TraversalPathWire,
TxnAbortRequest, TxnAbortResponse, TxnBeginRequest, TxnBeginResponse, TxnCommitRequest,
TxnCommitResponse, UnlinkRequest, UnlinkResponse,
};
const DEFAULT_CLIENT_ID: &str = "brain-db-sdk-rust";
#[derive(Clone, Debug)]
pub enum Auth {
Token(Vec<u8>),
Mtls(MtlsClaim),
}
impl Auth {
fn method(&self) -> AuthMethod {
match self {
Auth::Token(_) => AuthMethod::Token,
Auth::Mtls(_) => AuthMethod::Mtls,
}
}
fn credentials(self) -> AuthCredentials {
match self {
Auth::Token(t) => AuthCredentials::Token(t),
Auth::Mtls(c) => AuthCredentials::Mtls(c),
}
}
}
#[derive(Clone, Debug)]
pub struct ClientConfig {
pub client_id: String,
pub supported_versions: Vec<u8>,
pub capabilities: HelloCapabilities,
pub auth: Auth,
pub connect_timeout: Option<Duration>,
pub request_timeout: Option<Duration>,
}
impl ClientConfig {
#[must_use]
pub fn new(auth: Auth) -> Self {
Self {
client_id: DEFAULT_CLIENT_ID.to_string(),
supported_versions: vec![1],
capabilities: HelloCapabilities {
streaming: true,
compression_zstd: false,
server_push: false,
},
auth,
connect_timeout: Some(Duration::from_secs(10)),
request_timeout: Some(Duration::from_secs(30)),
}
}
}
#[derive(Clone, Debug)]
pub struct ConnectionInfo {
pub space_id: [u8; 16],
pub server_id: String,
pub chosen_version: u8,
pub connection_id: [u8; 16],
pub bound_shard_id: u16,
pub permissions: SpacePermissions,
pub namespace: String,
pub server_features: ServerFeatures,
}
#[derive(Clone, Debug, PartialEq)]
pub struct RecallAnswer {
pub answer_kind: AnswerKindWire,
pub memories: Vec<MemoryResult>,
}
impl RecallAnswer {
#[must_use]
pub fn is_empty(&self) -> bool {
self.memories.is_empty()
}
#[must_use]
pub fn memories(&self) -> &[MemoryResult] {
&self.memories
}
}
pub struct BrainClient {
conn: MuxConnection<TcpStream>,
connection: ConnectionInfo,
}
impl std::fmt::Debug for BrainClient {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BrainClient")
.field("space_id", &hex16(&self.connection.space_id))
.field("namespace", &self.connection.namespace)
.field("chosen_version", &self.connection.chosen_version)
.finish_non_exhaustive()
}
}
fn hex16(id: &[u8; 16]) -> String {
id.iter().fold(String::with_capacity(32), |mut s, b| {
use std::fmt::Write as _;
let _ = write!(s, "{b:02x}");
s
})
}
impl BrainClient {
pub async fn connect(addr: SocketAddr, auth: Auth) -> Result<Self> {
Self::connect_with(addr, ClientConfig::new(auth)).await
}
pub async fn connect_with(addr: SocketAddr, config: ClientConfig) -> Result<Self> {
let stream = match config.connect_timeout {
Some(timeout) => tokio::time::timeout(timeout, TcpStream::connect(addr))
.await
.map_err(|_| BrainError::Timeout(timeout))??,
None => TcpStream::connect(addr).await?,
};
stream.set_nodelay(true)?;
Self::handshake_over(stream, config).await
}
async fn handshake_over(stream: TcpStream, config: ClientConfig) -> Result<Self> {
let hello = HelloPayload {
client_id: config.client_id,
supported_versions: config.supported_versions,
capabilities: config.capabilities,
client_connection_token: None,
};
let auth = AuthPayload {
method: config.auth.method(),
credentials: config.auth.credentials(),
};
let (conn, HandshakeOutcome { welcome, auth_ok }) =
MuxConnection::handshake(stream, hello, auth, config.request_timeout).await?;
let connection = ConnectionInfo {
space_id: auth_ok.space_id,
server_id: welcome.server_id,
chosen_version: welcome.chosen_version,
connection_id: welcome.connection_id,
bound_shard_id: auth_ok.bound_shard_id,
permissions: auth_ok.permissions,
namespace: auth_ok.namespace,
server_features: welcome.server_features,
};
Ok(Self { conn, connection })
}
#[must_use]
pub fn is_closed(&self) -> bool {
self.conn.is_closed()
}
#[must_use]
pub fn connection(&self) -> &ConnectionInfo {
&self.connection
}
#[must_use]
pub fn space_id(&self) -> [u8; 16] {
self.connection.space_id
}
#[must_use]
pub fn namespace(&self) -> &str {
&self.connection.namespace
}
pub async fn encode(&self, request: &EncodeRequest) -> Result<EncodeResponse> {
self.unary(
Opcode::EncodeReq,
Opcode::EncodeResp,
"ENCODE_RESP",
request,
)
.await
}
pub async fn recall(&self, request: &RecallRequest) -> Result<RecallAnswer> {
let frames = self.recall_frames(request).await?;
let mut memories = Vec::new();
let mut answer_kind = AnswerKindWire::None;
for frame in frames {
answer_kind = frame.answer_kind;
memories.extend(frame.memories);
}
Ok(RecallAnswer {
answer_kind,
memories,
})
}
pub async fn recall_frames(&self, request: &RecallRequest) -> Result<Vec<RecallResponseFrame>> {
let frames = self
.conn
.request(Opcode::RecallReq, to_cbor_bytes(request))
.await?;
let mut out = Vec::with_capacity(frames.len());
for frame in frames {
if frame.opcode != Opcode::RecallResp as u16 {
return Err(BrainError::Protocol(format!(
"expected RECALL_RESP ({:#06x}), got {:#06x}",
Opcode::RecallResp as u16,
frame.opcode
)));
}
out.push(from_cbor_bytes(&frame.payload)?);
}
Ok(out)
}
pub async fn forget(&self, request: &ForgetRequest) -> Result<ForgetResponse> {
self.unary(
Opcode::ForgetReq,
Opcode::ForgetResp,
"FORGET_RESP",
request,
)
.await
}
pub async fn link(&self, request: &LinkRequest) -> Result<LinkResponse> {
self.unary(Opcode::LinkReq, Opcode::LinkResp, "LINK_RESP", request)
.await
}
pub async fn unlink(&self, request: &UnlinkRequest) -> Result<UnlinkResponse> {
self.unary(
Opcode::UnlinkReq,
Opcode::UnlinkResp,
"UNLINK_RESP",
request,
)
.await
}
pub async fn plan(&self, request: &PlanRequest) -> Result<Vec<PlanStep>> {
let frames = self.plan_frames(request).await?;
let mut steps = Vec::new();
for frame in frames {
steps.extend(frame.steps);
}
Ok(steps)
}
pub async fn plan_frames(&self, request: &PlanRequest) -> Result<Vec<PlanResponseFrame>> {
self.streamed(Opcode::PlanReq, Opcode::PlanResp, "PLAN_RESP", request)
.await
}
pub async fn reason(&self, request: &ReasonRequest) -> Result<Vec<InferenceStep>> {
let frames = self.reason_frames(request).await?;
let mut inferences = Vec::new();
for frame in frames {
inferences.extend(frame.inferences);
}
Ok(inferences)
}
pub async fn reason_frames(&self, request: &ReasonRequest) -> Result<Vec<ReasonResponseFrame>> {
self.streamed(
Opcode::ReasonReq,
Opcode::ReasonResp,
"REASON_RESP",
request,
)
.await
}
pub async fn create_entity(
&self,
request: &EntityCreateRequest,
) -> Result<EntityCreateResponse> {
self.unary(
Opcode::EntityCreateReq,
Opcode::EntityCreateResp,
"ENTITY_CREATE_RESP",
request,
)
.await
}
pub async fn update_entity(
&self,
request: &EntityUpdateRequest,
) -> Result<EntityUpdateResponse> {
self.unary(
Opcode::EntityUpdateReq,
Opcode::EntityUpdateResp,
"ENTITY_UPDATE_RESP",
request,
)
.await
}
pub async fn rename_entity(
&self,
request: &EntityRenameRequest,
) -> Result<EntityRenameResponse> {
self.unary(
Opcode::EntityRenameReq,
Opcode::EntityRenameResp,
"ENTITY_RENAME_RESP",
request,
)
.await
}
pub async fn merge_entities(
&self,
request: &EntityMergeRequest,
) -> Result<EntityMergeResponse> {
self.unary(
Opcode::EntityMergeReq,
Opcode::EntityMergeResp,
"ENTITY_MERGE_RESP",
request,
)
.await
}
pub async fn unmerge_entity(
&self,
request: &EntityUnmergeRequest,
) -> Result<EntityUnmergeResponse> {
self.unary(
Opcode::EntityUnmergeReq,
Opcode::EntityUnmergeResp,
"ENTITY_UNMERGE_RESP",
request,
)
.await
}
pub async fn tombstone_entity(
&self,
request: &EntityTombstoneRequest,
) -> Result<EntityTombstoneResponse> {
self.unary(
Opcode::EntityTombstoneReq,
Opcode::EntityTombstoneResp,
"ENTITY_TOMBSTONE_RESP",
request,
)
.await
}
pub async fn create_statement(
&self,
request: &StatementCreateRequest,
) -> Result<StatementCreateResponse> {
self.unary(
Opcode::StatementCreateReq,
Opcode::StatementCreateResp,
"STATEMENT_CREATE_RESP",
request,
)
.await
}
pub async fn create_relation(
&self,
request: &RelationCreateRequest,
) -> Result<RelationCreateResponse> {
self.unary(
Opcode::RelationCreateReq,
Opcode::RelationCreateResp,
"RELATION_CREATE_RESP",
request,
)
.await
}
pub async fn upload_schema(
&self,
request: &SchemaUploadRequest,
) -> Result<SchemaUploadResponse> {
self.unary(
Opcode::SchemaUploadReq,
Opcode::SchemaUploadResp,
"SCHEMA_UPLOAD_RESP",
request,
)
.await
}
pub async fn get_relation(&self, request: &RelationGetRequest) -> Result<RelationGetResponse> {
self.unary(
Opcode::RelationGetReq,
Opcode::RelationGetResp,
"RELATION_GET_RESP",
request,
)
.await
}
pub async fn supersede_relation(
&self,
request: &RelationSupersedeRequest,
) -> Result<RelationSupersedeResponse> {
self.unary(
Opcode::RelationSupersedeReq,
Opcode::RelationSupersedeResp,
"RELATION_SUPERSEDE_RESP",
request,
)
.await
}
pub async fn tombstone_relation(
&self,
request: &RelationTombstoneRequest,
) -> Result<RelationTombstoneResponse> {
self.unary(
Opcode::RelationTombstoneReq,
Opcode::RelationTombstoneResp,
"RELATION_TOMBSTONE_RESP",
request,
)
.await
}
pub async fn traverse_relations(
&self,
request: &RelationTraverseRequest,
) -> Result<Vec<TraversalPathWire>> {
let frames = self.traverse_relations_frames(request).await?;
let mut paths = Vec::new();
for frame in frames {
paths.extend(frame.paths);
}
Ok(paths)
}
pub async fn traverse_relations_frames(
&self,
request: &RelationTraverseRequest,
) -> Result<Vec<RelationTraverseResponseFrame>> {
self.streamed(
Opcode::RelationTraverseReq,
Opcode::RelationTraverseResp,
"RELATION_TRAVERSE_RESP",
request,
)
.await
}
pub async fn query_explain(
&self,
request: &QueryExplainRequest,
) -> Result<QueryExplainResponse> {
self.unary(
Opcode::QueryExplainReq,
Opcode::QueryExplainResp,
"QUERY_EXPLAIN_RESP",
request,
)
.await
}
pub async fn query_trace(&self, request: &QueryTraceRequest) -> Result<QueryTraceResponse> {
self.unary(
Opcode::QueryTraceReq,
Opcode::QueryTraceResp,
"QUERY_TRACE_RESP",
request,
)
.await
}
pub async fn encode_vector_direct(
&self,
request: &EncodeVectorDirectRequest,
) -> Result<EncodeResponse> {
let mut payload = to_cbor_bytes(request);
payload.extend_from_slice(&f32_slice_to_le_bytes(&request.vector));
let frame = self
.conn
.request_one(Opcode::EncodeVectorDirectReq, payload)
.await?;
if frame.opcode != Opcode::EncodeVectorDirectResp as u16 {
return Err(BrainError::Protocol(format!(
"expected ENCODE_VECTOR_DIRECT_RESP ({:#06x}), got {:#06x}",
Opcode::EncodeVectorDirectResp as u16,
frame.opcode
)));
}
Ok(from_cbor_bytes(&frame.payload)?)
}
pub async fn materialize_procedural(
&self,
request: &MaterializeProceduralRequest,
) -> Result<MaterializeProceduralResponse> {
self.unary(
Opcode::MaterializeProceduralReq,
Opcode::MaterializeProceduralResp,
"MATERIALIZE_PROCEDURAL_RESP",
request,
)
.await
}
pub async fn capabilities(
&self,
request: &GetCapabilitiesRequest,
) -> Result<GetCapabilitiesResponse> {
self.unary(
Opcode::GetCapabilitiesReq,
Opcode::GetCapabilitiesResp,
"GET_CAPABILITIES_RESP",
request,
)
.await
}
pub async fn extractor_list(
&self,
request: &ExtractorListRequest,
) -> Result<ExtractorListResponseFrame> {
self.unary(
Opcode::ExtractorListReq,
Opcode::ExtractorListResp,
"EXTRACTOR_LIST_RESP",
request,
)
.await
}
pub async fn get_entity(&self, request: &EntityGetRequest) -> Result<EntityGetResponse> {
self.unary(
Opcode::EntityGetReq,
Opcode::EntityGetResp,
"ENTITY_GET_RESP",
request,
)
.await
}
pub async fn memory_inspect(
&self,
request: &MemoryInspectRequest,
) -> Result<MemoryInspectResponse> {
self.unary(
Opcode::MemoryInspectReq,
Opcode::MemoryInspectResp,
"MEMORY_INSPECT_RESP",
request,
)
.await
}
pub async fn resolve_entity(
&self,
request: &EntityResolveRequest,
) -> Result<EntityResolveResponse> {
self.unary(
Opcode::EntityResolveReq,
Opcode::EntityResolveResp,
"ENTITY_RESOLVE_RESP",
request,
)
.await
}
pub async fn get_statement(
&self,
request: &StatementGetRequest,
) -> Result<StatementGetResponse> {
self.unary(
Opcode::StatementGetReq,
Opcode::StatementGetResp,
"STATEMENT_GET_RESP",
request,
)
.await
}
pub async fn get_schema(&self, request: &SchemaGetRequest) -> Result<SchemaGetResponse> {
self.unary(
Opcode::SchemaGetReq,
Opcode::SchemaGetResp,
"SCHEMA_GET_RESP",
request,
)
.await
}
pub async fn validate_schema(
&self,
request: &SchemaValidateRequest,
) -> Result<SchemaValidateResponse> {
self.unary(
Opcode::SchemaValidateReq,
Opcode::SchemaValidateResp,
"SCHEMA_VALIDATE_RESP",
request,
)
.await
}
pub async fn replace_schema(
&self,
request: &SchemaReplaceRequest,
) -> Result<SchemaReplaceResponse> {
self.unary(
Opcode::SchemaReplaceReq,
Opcode::SchemaReplaceResp,
"SCHEMA_REPLACE_RESP",
request,
)
.await
}
pub async fn drop_schema(&self, request: &SchemaDropRequest) -> Result<SchemaDropResponse> {
self.unary(
Opcode::SchemaDropReq,
Opcode::SchemaDropResp,
"SCHEMA_DROP_RESP",
request,
)
.await
}
pub async fn ping(&self, client_timestamp_unix_nanos: Option<u64>) -> Result<PongResponse> {
let ts = client_timestamp_unix_nanos.unwrap_or_else(|| {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |d| u64::try_from(d.as_nanos()).unwrap_or(u64::MAX))
});
let request = PingRequest {
client_timestamp_unix_nanos: ts,
};
self.unary(Opcode::Ping, Opcode::Pong, "PONG", &request)
.await
}
pub async fn cancel_stream(
&self,
target_stream_id: u32,
reason: CancellationReason,
) -> Result<CancelStreamAck> {
let request = CancelStreamRequest {
target_stream_id,
reason,
};
self.unary(
Opcode::CancelStream,
Opcode::CancelStreamAck,
"CANCEL_STREAM_ACK",
&request,
)
.await
}
pub async fn list_entities(&self, request: &EntityListRequest) -> Result<Vec<EntityListItem>> {
let frames = self.list_entities_frames(request).await?;
let mut items = Vec::new();
for frame in frames {
items.extend(frame.items);
}
Ok(items)
}
pub async fn list_entities_frames(
&self,
request: &EntityListRequest,
) -> Result<Vec<EntityListResponseFrame>> {
self.streamed(
Opcode::EntityListReq,
Opcode::EntityListResp,
"ENTITY_LIST_RESP",
request,
)
.await
}
pub async fn memory_list(&self, request: &MemoryListRequest) -> Result<Vec<MemoryListItem>> {
let frames = self.memory_list_frames(request).await?;
let mut items = Vec::new();
for frame in frames {
items.extend(frame.items);
}
Ok(items)
}
pub async fn memory_list_frames(
&self,
request: &MemoryListRequest,
) -> Result<Vec<MemoryListResponseFrame>> {
self.streamed(
Opcode::MemoryListReq,
Opcode::MemoryListResp,
"MEMORY_LIST_RESP",
request,
)
.await
}
pub async fn graph_fetch(
&self,
request: &GraphFetchRequest,
) -> Result<(Vec<GraphNode>, Vec<GraphEdge>)> {
let frames = self.graph_fetch_frames(request).await?;
let mut nodes = Vec::new();
let mut edges = Vec::new();
for frame in frames {
nodes.extend(frame.nodes);
edges.extend(frame.edges);
}
Ok((nodes, edges))
}
pub async fn graph_fetch_frames(
&self,
request: &GraphFetchRequest,
) -> Result<Vec<GraphFetchResponseFrame>> {
self.streamed(
Opcode::GraphFetchReq,
Opcode::GraphFetchResp,
"GRAPH_FETCH_RESP",
request,
)
.await
}
pub async fn list_statements(
&self,
request: &StatementListRequest,
) -> Result<Vec<StatementView>> {
let frames = self.list_statements_frames(request).await?;
let mut items = Vec::new();
for frame in frames {
items.extend(frame.items);
}
Ok(items)
}
pub async fn list_statements_frames(
&self,
request: &StatementListRequest,
) -> Result<Vec<StatementListResponseFrame>> {
self.streamed(
Opcode::StatementListReq,
Opcode::StatementListResp,
"STATEMENT_LIST_RESP",
request,
)
.await
}
pub async fn supersede_statement(
&self,
request: &StatementSupersedeRequest,
) -> Result<StatementSupersedeResponse> {
self.unary(
Opcode::StatementSupersedeReq,
Opcode::StatementSupersedeResp,
"STATEMENT_SUPERSEDE_RESP",
request,
)
.await
}
pub async fn tombstone_statement(
&self,
request: &StatementTombstoneRequest,
) -> Result<StatementTombstoneResponse> {
self.unary(
Opcode::StatementTombstoneReq,
Opcode::StatementTombstoneResp,
"STATEMENT_TOMBSTONE_RESP",
request,
)
.await
}
pub async fn retract_statement(
&self,
request: &StatementRetractRequest,
) -> Result<StatementRetractResponse> {
self.unary(
Opcode::StatementRetractReq,
Opcode::StatementRetractResp,
"STATEMENT_RETRACT_RESP",
request,
)
.await
}
pub async fn statement_history(
&self,
request: &StatementHistoryRequest,
) -> Result<Vec<StatementView>> {
let frames = self.statement_history_frames(request).await?;
let mut items = Vec::new();
for frame in frames {
items.extend(frame.items);
}
Ok(items)
}
pub async fn statement_history_frames(
&self,
request: &StatementHistoryRequest,
) -> Result<Vec<StatementHistoryResponseFrame>> {
self.streamed(
Opcode::StatementHistoryReq,
Opcode::StatementHistoryResp,
"STATEMENT_HISTORY_RESP",
request,
)
.await
}
pub async fn list_relations_from(
&self,
request: &RelationListFromRequest,
) -> Result<Vec<RelationView>> {
let frames = self.list_relations_from_frames(request).await?;
let mut items = Vec::new();
for frame in frames {
items.extend(frame.items);
}
Ok(items)
}
pub async fn list_relations_from_frames(
&self,
request: &RelationListFromRequest,
) -> Result<Vec<RelationListFromResponseFrame>> {
self.streamed(
Opcode::RelationListFromReq,
Opcode::RelationListFromResp,
"RELATION_LIST_FROM_RESP",
request,
)
.await
}
pub async fn list_relations_to(
&self,
request: &RelationListToRequest,
) -> Result<Vec<RelationView>> {
let frames = self.list_relations_to_frames(request).await?;
let mut items = Vec::new();
for frame in frames {
items.extend(frame.items);
}
Ok(items)
}
pub async fn list_relations_to_frames(
&self,
request: &RelationListToRequest,
) -> Result<Vec<RelationListToResponseFrame>> {
self.streamed(
Opcode::RelationListToReq,
Opcode::RelationListToResp,
"RELATION_LIST_TO_RESP",
request,
)
.await
}
pub async fn list_schemas(
&self,
request: &SchemaListRequest,
) -> Result<Vec<SchemaListItemWire>> {
let frames = self.list_schemas_frames(request).await?;
let mut items = Vec::new();
for frame in frames {
items.extend(frame.items);
}
Ok(items)
}
pub async fn list_schemas_frames(
&self,
request: &SchemaListRequest,
) -> Result<Vec<SchemaListResponseFrame>> {
self.streamed(
Opcode::SchemaListReq,
Opcode::SchemaListResp,
"SCHEMA_LIST_RESP",
request,
)
.await
}
pub async fn txn_begin(&self, request: &TxnBeginRequest) -> Result<TxnBeginResponse> {
self.unary(
Opcode::TxnBegin,
Opcode::TxnBeginResp,
"TXN_BEGIN_RESP",
request,
)
.await
}
pub async fn txn_commit(&self, request: &TxnCommitRequest) -> Result<TxnCommitResponse> {
self.unary(
Opcode::TxnCommit,
Opcode::TxnCommitResp,
"TXN_COMMIT_RESP",
request,
)
.await
}
pub async fn txn_abort(&self, request: &TxnAbortRequest) -> Result<TxnAbortResponse> {
self.unary(
Opcode::TxnAbort,
Opcode::TxnAbortResp,
"TXN_ABORT_RESP",
request,
)
.await
}
pub async fn create_space(&self, request: &SpaceCreateRequest) -> Result<SpaceCreateResponse> {
self.unary(
Opcode::SpaceCreateReq,
Opcode::SpaceCreateResp,
"SPACE_CREATE_RESP",
request,
)
.await
}
pub async fn list_spaces(&self, request: &SpaceListRequest) -> Result<SpaceListResponse> {
self.unary(
Opcode::SpaceListReq,
Opcode::SpaceListResp,
"SPACE_LIST_RESP",
request,
)
.await
}
pub async fn delete_space(&self, request: &SpaceDeleteRequest) -> Result<SpaceDeleteResponse> {
self.unary(
Opcode::SpaceDeleteReq,
Opcode::SpaceDeleteResp,
"SPACE_DELETE_RESP",
request,
)
.await
}
pub async fn create_session(
&self,
request: &SessionCreateRequest,
) -> Result<SessionCreateResponse> {
self.unary(
Opcode::SessionCreateReq,
Opcode::SessionCreateResp,
"SESSION_CREATE_RESP",
request,
)
.await
}
pub async fn list_sessions(&self, request: &SessionListRequest) -> Result<SessionListResponse> {
self.unary(
Opcode::SessionListReq,
Opcode::SessionListResp,
"SESSION_LIST_RESP",
request,
)
.await
}
pub async fn delete_session(
&self,
request: &SessionDeleteRequest,
) -> Result<SessionDeleteResponse> {
self.unary(
Opcode::SessionDeleteReq,
Opcode::SessionDeleteResp,
"SESSION_DELETE_RESP",
request,
)
.await
}
pub async fn subscribe(&self, request: &SubscribeRequest) -> Result<Subscription<TcpStream>> {
self.conn.subscribe(request).await
}
async fn streamed<Req, Frame>(
&self,
req_opcode: Opcode,
resp_opcode: Opcode,
resp_name: &str,
request: &Req,
) -> Result<Vec<Frame>>
where
Req: serde::Serialize,
Frame: serde::de::DeserializeOwned,
{
let frames = self
.conn
.request(req_opcode, to_cbor_bytes(request))
.await?;
let mut out = Vec::with_capacity(frames.len());
for frame in frames {
if frame.opcode != resp_opcode as u16 {
return Err(BrainError::Protocol(format!(
"expected {resp_name} ({:#06x}), got {:#06x}",
resp_opcode as u16, frame.opcode
)));
}
out.push(from_cbor_bytes(&frame.payload)?);
}
Ok(out)
}
async fn unary<Req, Resp>(
&self,
req_opcode: Opcode,
resp_opcode: Opcode,
resp_name: &str,
request: &Req,
) -> Result<Resp>
where
Req: serde::Serialize,
Resp: serde::de::DeserializeOwned,
{
let frame = self
.conn
.request_one(req_opcode, to_cbor_bytes(request))
.await?;
if frame.opcode != resp_opcode as u16 {
return Err(BrainError::Protocol(format!(
"expected {resp_name} ({:#06x}), got {:#06x}",
resp_opcode as u16, frame.opcode
)));
}
Ok(from_cbor_bytes(&frame.payload)?)
}
pub async fn close(self) -> Result<()> {
self.conn.send_bye().await
}
}
#[must_use]
pub fn new_id() -> [u8; 16] {
*uuid::Uuid::now_v7().as_bytes()
}
const BRAIN_ROOT_NAMESPACE_UUID: uuid::Uuid = uuid::Uuid::from_bytes([
0x6b, 0x72, 0x61, 0x69, 0x6e, 0x2d, 0x73, 0x70, 0x61, 0x63, 0x65, 0x2d, 0x72, 0x6f, 0x6f, 0x74,
]);
#[must_use]
pub fn derive_space_id(namespace: &str, space: &str) -> [u8; 16] {
let ns = uuid::Uuid::new_v5(&BRAIN_ROOT_NAMESPACE_UUID, namespace.as_bytes());
*uuid::Uuid::new_v5(&ns, space.as_bytes()).as_bytes()
}
#[cfg(test)]
mod space_id_tests {
use super::derive_space_id;
#[test]
fn derive_space_id_matches_server_golden() {
let id = derive_space_id("acme", "support-bot:user123");
assert_eq!(
uuid::Uuid::from_bytes(id).to_string(),
"119061cc-b6db-5cde-8edd-0cefa33452eb"
);
}
#[test]
fn same_space_diverges_across_namespaces() {
assert_ne!(derive_space_id("ns_a", "u1"), derive_space_id("ns_b", "u1"));
}
}