重写

server-rs/crates/spacetime-client/src/lib.rs

@@ -0,0 +1,306 @@
+pub mod module_bindings;
+
+use std::{
+    error::Error,
+    fmt,
+    sync::{Arc, Mutex},
+    time::Duration,
+};
+
+use module_assets::{
+    AssetEntityBindingRecord, AssetObjectAccessPolicy, AssetObjectRecord,
+    build_asset_entity_binding_record, build_asset_object_record,
+};
+use spacetimedb_sdk::DbContext;
+use tokio::{sync::oneshot, time::timeout};
+
+use crate::module_bindings::{
+    AssetEntityBindingInput as BindingAssetEntityBindingInput,
+    AssetEntityBindingProcedureResult as BindingAssetEntityBindingProcedureResult,
+    AssetEntityBindingSnapshot as BindingAssetEntityBindingSnapshot,
+    AssetObjectProcedureResult as BindingAssetObjectProcedureResult,
+    AssetObjectUpsertInput as BindingAssetObjectUpsertInput,
+    AssetObjectUpsertSnapshot as BindingAssetObjectUpsertSnapshot, DbConnection,
+    bind_asset_object_to_entity_and_return_procedure::bind_asset_object_to_entity_and_return as _,
+    confirm_asset_object_and_return_procedure::confirm_asset_object_and_return as _,
+};
+
+#[derive(Clone, Debug)]
+pub struct SpacetimeClientConfig {
+    pub server_url: String,
+    pub database: String,
+    pub token: Option<String>,
+}
+
+#[derive(Clone, Debug)]
+pub struct SpacetimeClient {
+    config: SpacetimeClientConfig,
+}
+
+#[derive(Debug)]
+pub enum SpacetimeClientError {
+    Build(String),
+    ConnectDropped,
+    Procedure(String),
+    Runtime(String),
+    Timeout,
+}
+
+const CONFIRM_ASSET_OBJECT_TIMEOUT: Duration = Duration::from_secs(10);
+
+type ProcedureResultSender<T> =
+    Arc<Mutex<Option<oneshot::Sender<Result<T, SpacetimeClientError>>>>>;
+
+impl SpacetimeClient {
+    pub fn new(config: SpacetimeClientConfig) -> Self {
+        Self { config }
+    }
+
+    pub async fn confirm_asset_object(
+        &self,
+        input: module_assets::AssetObjectUpsertInput,
+    ) -> Result<AssetObjectRecord, SpacetimeClientError> {
+        let procedure_input = map_upsert_input(input);
+
+        self.call_after_connect(move |connection, sender| {
+            connection
+                .procedures()
+                .confirm_asset_object_and_return_then(procedure_input, move |_, result| {
+                    let mapped = result
+                        .map_err(|error| SpacetimeClientError::Procedure(error.to_string()))
+                        .and_then(map_procedure_result);
+                    send_once(&sender, mapped);
+                });
+        })
+        .await
+    }
+
+    pub async fn bind_asset_object_to_entity(
+        &self,
+        input: module_assets::AssetEntityBindingInput,
+    ) -> Result<AssetEntityBindingRecord, SpacetimeClientError> {
+        let procedure_input = map_entity_binding_input(input);
+
+        self.call_after_connect(move |connection, sender| {
+            connection
+                .procedures()
+                .bind_asset_object_to_entity_and_return_then(procedure_input, move |_, result| {
+                    let mapped = result
+                        .map_err(|error| SpacetimeClientError::Procedure(error.to_string()))
+                        .and_then(map_entity_binding_procedure_result);
+                    send_once(&sender, mapped);
+                });
+        })
+        .await
+    }
+
+    async fn call_after_connect<T>(
+        &self,
+        call: impl FnOnce(&DbConnection, ProcedureResultSender<T>) + Send + 'static,
+    ) -> Result<T, SpacetimeClientError>
+    where
+        T: Send + 'static,
+    {
+        let config = self.config.clone();
+        let (sender, receiver) = oneshot::channel();
+        let result_sender = Arc::new(Mutex::new(Some(sender)));
+        let connect_sender = result_sender.clone();
+        let disconnect_sender = result_sender.clone();
+
+        let connection = tokio::task::spawn_blocking(move || {
+            DbConnection::builder()
+                .with_uri(config.server_url)
+                .with_database_name(config.database)
+                .with_token(config.token)
+                .on_connect(move |connection, _, _| {
+                    // SDK 收到 IdentityToken 后才调用 procedure，避免 WebSocket 已建好但身份握手未完成时丢请求。
+                    call(connection, connect_sender);
+                })
+                .on_disconnect(move |_, error| {
+                    let message = error
+                        .map(|error| error.to_string())
+                        .unwrap_or_else(|| "SpacetimeDB 连接在 procedure 返回前断开".to_string());
+                    send_once(
+                        &disconnect_sender,
+                        Err(SpacetimeClientError::Procedure(message)),
+                    );
+                })
+                .build()
+                .map_err(|error| SpacetimeClientError::Build(error.to_string()))
+        })
+        .await
+        .map_err(|error| SpacetimeClientError::Runtime(error.to_string()))??;
+
+        let runner = connection.run_threaded();
+        let result = timeout(CONFIRM_ASSET_OBJECT_TIMEOUT, receiver).await;
+        let _ = connection.disconnect();
+        // SDK 线程会在断开消息被处理后自行退出；HTTP 请求不能同步等待该线程，否则 Windows 本地联调可能卡在收尾阶段。
+        drop(runner);
+
+        result
+            .map_err(|_| SpacetimeClientError::Timeout)?
+            .map_err(|_| SpacetimeClientError::ConnectDropped)?
+    }
+}
+
+fn send_once<T>(sender: &ProcedureResultSender<T>, result: Result<T, SpacetimeClientError>) {
+    if let Some(sender) = sender
+        .lock()
+        .expect("spacetime result sender should not poison")
+        .take()
+    {
+        let _ = sender.send(result);
+    }
+}
+
+fn map_entity_binding_input(
+    input: module_assets::AssetEntityBindingInput,
+) -> BindingAssetEntityBindingInput {
+    BindingAssetEntityBindingInput {
+        binding_id: input.binding_id,
+        asset_object_id: input.asset_object_id,
+        entity_kind: input.entity_kind,
+        entity_id: input.entity_id,
+        slot: input.slot,
+        asset_kind: input.asset_kind,
+        owner_user_id: input.owner_user_id,
+        profile_id: input.profile_id,
+        updated_at_micros: input.updated_at_micros,
+    }
+}
+
+fn map_upsert_input(input: module_assets::AssetObjectUpsertInput) -> BindingAssetObjectUpsertInput {
+    BindingAssetObjectUpsertInput {
+        asset_object_id: input.asset_object_id,
+        bucket: input.bucket,
+        object_key: input.object_key,
+        access_policy: map_access_policy(input.access_policy),
+        content_type: input.content_type,
+        content_length: input.content_length,
+        content_hash: input.content_hash,
+        version: input.version,
+        source_job_id: input.source_job_id,
+        owner_user_id: input.owner_user_id,
+        profile_id: input.profile_id,
+        entity_id: input.entity_id,
+        asset_kind: input.asset_kind,
+        updated_at_micros: input.updated_at_micros,
+    }
+}
+
+fn map_procedure_result(
+    result: BindingAssetObjectProcedureResult,
+) -> Result<AssetObjectRecord, SpacetimeClientError> {
+    if !result.ok {
+        return Err(SpacetimeClientError::Procedure(
+            result
+                .error_message
+                .unwrap_or_else(|| "SpacetimeDB procedure 返回未知错误".to_string()),
+        ));
+    }
+
+    let snapshot = result.record.ok_or_else(|| {
+        SpacetimeClientError::Procedure("SpacetimeDB procedure 未返回对象快照".to_string())
+    })?;
+
+    Ok(build_asset_object_record(map_snapshot(snapshot)))
+}
+
+fn map_entity_binding_procedure_result(
+    result: BindingAssetEntityBindingProcedureResult,
+) -> Result<AssetEntityBindingRecord, SpacetimeClientError> {
+    if !result.ok {
+        return Err(SpacetimeClientError::Procedure(
+            result
+                .error_message
+                .unwrap_or_else(|| "SpacetimeDB procedure 返回未知错误".to_string()),
+        ));
+    }
+
+    let snapshot = result.record.ok_or_else(|| {
+        SpacetimeClientError::Procedure("SpacetimeDB procedure 未返回绑定快照".to_string())
+    })?;
+
+    Ok(build_asset_entity_binding_record(
+        map_entity_binding_snapshot(snapshot),
+    ))
+}
+
+fn map_entity_binding_snapshot(
+    snapshot: BindingAssetEntityBindingSnapshot,
+) -> module_assets::AssetEntityBindingSnapshot {
+    module_assets::AssetEntityBindingSnapshot {
+        binding_id: snapshot.binding_id,
+        asset_object_id: snapshot.asset_object_id,
+        entity_kind: snapshot.entity_kind,
+        entity_id: snapshot.entity_id,
+        slot: snapshot.slot,
+        asset_kind: snapshot.asset_kind,
+        owner_user_id: snapshot.owner_user_id,
+        profile_id: snapshot.profile_id,
+        created_at_micros: snapshot.created_at_micros,
+        updated_at_micros: snapshot.updated_at_micros,
+    }
+}
+
+fn map_snapshot(
+    snapshot: BindingAssetObjectUpsertSnapshot,
+) -> module_assets::AssetObjectUpsertSnapshot {
+    module_assets::AssetObjectUpsertSnapshot {
+        asset_object_id: snapshot.asset_object_id,
+        bucket: snapshot.bucket,
+        object_key: snapshot.object_key,
+        access_policy: map_access_policy_back(snapshot.access_policy),
+        content_type: snapshot.content_type,
+        content_length: snapshot.content_length,
+        content_hash: snapshot.content_hash,
+        version: snapshot.version,
+        source_job_id: snapshot.source_job_id,
+        owner_user_id: snapshot.owner_user_id,
+        profile_id: snapshot.profile_id,
+        entity_id: snapshot.entity_id,
+        asset_kind: snapshot.asset_kind,
+        created_at_micros: snapshot.created_at_micros,
+        updated_at_micros: snapshot.updated_at_micros,
+    }
+}
+
+fn map_access_policy(
+    value: AssetObjectAccessPolicy,
+) -> crate::module_bindings::AssetObjectAccessPolicy {
+    match value {
+        AssetObjectAccessPolicy::Private => {
+            crate::module_bindings::AssetObjectAccessPolicy::Private
+        }
+        AssetObjectAccessPolicy::PublicRead => {
+            crate::module_bindings::AssetObjectAccessPolicy::PublicRead
+        }
+    }
+}
+
+fn map_access_policy_back(
+    value: crate::module_bindings::AssetObjectAccessPolicy,
+) -> AssetObjectAccessPolicy {
+    match value {
+        crate::module_bindings::AssetObjectAccessPolicy::Private => {
+            AssetObjectAccessPolicy::Private
+        }
+        crate::module_bindings::AssetObjectAccessPolicy::PublicRead => {
+            AssetObjectAccessPolicy::PublicRead
+        }
+    }
+}
+
+impl fmt::Display for SpacetimeClientError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Build(message) | Self::Procedure(message) | Self::Runtime(message) => {
+                f.write_str(message)
+            }
+            Self::ConnectDropped => f.write_str("SpacetimeDB 连接在返回结果前已断开"),
+            Self::Timeout => f.write_str("SpacetimeDB procedure 调用超时"),
+        }
+    }
+}
+
+impl Error for SpacetimeClientError {}