Endpoint for pretty print migration plan (#3137)

# Description of Changes
- Adds endpoint for for pretty printing migration plan.
- It also changes current `publish` endpoint to optionally provide
`MigrationToken` and `MigrationPolicy` to allow migration with breaking
clients.

# API and ABI breaking changes
Backward compatible change to existing API and new Api

# Expected complexity level and risk
2

# Testing
- Existing smoketest should cover changes for `publish` endpoint.
- For pretty print endpoint, smoketests can be written only after cli
changes.

---------

Signed-off-by: Shubham Mishra <shivam828787@gmail.com>
Co-authored-by: James Gilles <jameshgilles@gmail.com>
Co-authored-by: Phoebe Goldman <phoebe@clockworklabs.io>

Author: 3 people

Cargo.lock

@@ -106,6 +106,27 @@ pub enum PublishResult {
     PermissionDenied { name: DatabaseName },
 }
 
+#[derive(serde::Serialize, serde::Deserialize, Debug, Default)]
+pub enum MigrationPolicy {
+    #[default]
+    Compatible,
+    BreakClients,
+}
+
+#[derive(serde::Serialize, serde::Deserialize, Debug, Default)]
+pub enum PrettyPrintStyle {
+    #[default]
+    AnsiColor,
+    NoColor,
+}
+
+#[derive(serde::Serialize, serde::Deserialize, Debug)]
+pub struct PrintPlanResult {
+    pub migrate_plan: Box<str>,
+    pub break_clients: bool,
+    pub token: spacetimedb_lib::Hash,
+}
+
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub enum DnsLookupResponse {
     /// The lookup was successful and the domain and identity are returned.

crates/client-api-messages/src/name.rs

@@ -7,14 +7,15 @@ use http::StatusCode;
 
 use spacetimedb::client::ClientActorIndex;
 use spacetimedb::energy::{EnergyBalance, EnergyQuanta};
-use spacetimedb::host::{HostController, ModuleHost, NoSuchModule, UpdateDatabaseResult};
+use spacetimedb::host::{HostController, MigratePlanResult, ModuleHost, NoSuchModule, UpdateDatabaseResult};
 use spacetimedb::identity::{AuthCtx, Identity};
 use spacetimedb::messages::control_db::{Database, HostType, Node, Replica};
 use spacetimedb::sql;
 use spacetimedb_client_api_messages::http::{SqlStmtResult, SqlStmtStats};
 use spacetimedb_client_api_messages::name::{DomainName, InsertDomainResult, RegisterTldResult, SetDomainsResult, Tld};
 use spacetimedb_lib::{ProductTypeElement, ProductValue};
 use spacetimedb_paths::server::ModuleLogsDir;
+use spacetimedb_schema::auto_migrate::{MigrationPolicy, PrettyPrintStyle};
 use tokio::sync::watch;
 
 pub mod auth;
@@ -146,9 +147,10 @@ impl Host {
         database: Database,
         host_type: HostType,
         program_bytes: Box<[u8]>,
+        policy: MigrationPolicy,
     ) -> anyhow::Result<UpdateDatabaseResult> {
         self.host_controller
-            .update_module_host(database, host_type, self.replica_id, program_bytes)
+            .update_module_host(database, host_type, self.replica_id, program_bytes, policy)
             .await
     }
 }
@@ -231,8 +233,11 @@ pub trait ControlStateWriteAccess: Send + Sync {
         &self,
         publisher: &Identity,
         spec: DatabaseDef,
+        policy: MigrationPolicy,
     ) -> anyhow::Result<Option<UpdateDatabaseResult>>;
 
+    async fn migrate_plan(&self, spec: DatabaseDef, style: PrettyPrintStyle) -> anyhow::Result<MigratePlanResult>;
+
     async fn delete_database(&self, caller_identity: &Identity, database_identity: &Identity) -> anyhow::Result<()>;
 
     // Energy
@@ -321,8 +326,13 @@ impl<T: ControlStateWriteAccess + ?Sized> ControlStateWriteAccess for Arc<T> {
         &self,
         identity: &Identity,
         spec: DatabaseDef,
+        policy: MigrationPolicy,
     ) -> anyhow::Result<Option<UpdateDatabaseResult>> {
-        (**self).publish_database(identity, spec).await
+        (**self).publish_database(identity, spec, policy).await
+    }
+
+    async fn migrate_plan(&self, spec: DatabaseDef, style: PrettyPrintStyle) -> anyhow::Result<MigratePlanResult> {
+        (**self).migrate_plan(spec, style).await
     }
 
     async fn delete_database(&self, caller_identity: &Identity, database_identity: &Identity) -> anyhow::Result<()> {

crates/client-api/src/lib.rs

@@ -20,16 +20,21 @@ use http::StatusCode;
 use serde::Deserialize;
 use spacetimedb::database_logger::DatabaseLogger;
 use spacetimedb::host::module_host::ClientConnectedError;
-use spacetimedb::host::ReducerArgs;
 use spacetimedb::host::ReducerCallError;
 use spacetimedb::host::ReducerOutcome;
 use spacetimedb::host::UpdateDatabaseResult;
+use spacetimedb::host::{MigratePlanResult, ReducerArgs};
 use spacetimedb::identity::Identity;
 use spacetimedb::messages::control_db::{Database, HostType};
-use spacetimedb_client_api_messages::name::{self, DatabaseName, DomainName, PublishOp, PublishResult};
+use spacetimedb_client_api_messages::name::{
+    self, DatabaseName, DomainName, MigrationPolicy, PrettyPrintStyle, PrintPlanResult, PublishOp, PublishResult,
+};
 use spacetimedb_lib::db::raw_def::v9::RawModuleDefV9;
 use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_lib::{sats, Timestamp};
+use spacetimedb_schema::auto_migrate::{
+    MigrationPolicy as SchemaMigrationPolicy, MigrationToken, PrettyPrintStyle as AutoMigratePrettyPrintStyle,
+};
 
 use super::subscribe::{handle_websocket, HasWebSocketOptions};
 
@@ -474,6 +479,13 @@ pub struct PublishDatabaseQueryParams {
     #[serde(default)]
     clear: bool,
     num_replicas: Option<usize>,
+    /// [`Hash`] of [`MigrationToken`]` to be checked if `MigrationPolicy::BreakClients` is set.
+    ///
+    /// Users obtain such a hash via the `/database/:name_or_identity/pre-publish POST` route.
+    /// This is a safeguard to require explicit approval for updates which will break clients.
+    token: Option<spacetimedb_lib::Hash>,
+    #[serde(default)]
+    policy: MigrationPolicy,
 }
 
 use std::env;
@@ -501,7 +513,12 @@ fn allow_creation(auth: &SpacetimeAuth) -> Result<(), ErrorResponse> {
 pub async fn publish<S: NodeDelegate + ControlStateDelegate>(
     State(ctx): State<S>,
     Path(PublishDatabaseParams { name_or_identity }): Path<PublishDatabaseParams>,
-    Query(PublishDatabaseQueryParams { clear, num_replicas }): Query<PublishDatabaseQueryParams>,
+    Query(PublishDatabaseQueryParams {
+        clear,
+        num_replicas,
+        token,
+        policy,
+    }): Query<PublishDatabaseQueryParams>,
     Extension(auth): Extension<SpacetimeAuth>,
     body: Bytes,
 ) -> axum::response::Result<axum::Json<PublishResult>> {
@@ -551,6 +568,21 @@ pub async fn publish<S: NodeDelegate + ControlStateDelegate>(
         }
     };
 
+    let policy: SchemaMigrationPolicy = match policy {
+        MigrationPolicy::BreakClients => {
+            if let Some(token) = token {
+                Ok(SchemaMigrationPolicy::BreakClients(token))
+            } else {
+                Err((
+                    StatusCode::BAD_REQUEST,
+                    "Migration policy is set to `BreakClients`, but no migration token was provided.",
+                ))
+            }
+        }
+
+        MigrationPolicy::Compatible => Ok(SchemaMigrationPolicy::Compatible),
+    }?;
+
     log::trace!("Publishing to the identity: {}", database_identity.to_hex());
 
     let op = {
@@ -592,6 +624,7 @@ pub async fn publish<S: NodeDelegate + ControlStateDelegate>(
                 num_replicas,
                 host_type: HostType::Wasm,
             },
+            policy,
         )
         .await
         .map_err(log_and_500)?;
@@ -619,6 +652,101 @@ pub async fn publish<S: NodeDelegate + ControlStateDelegate>(
     }))
 }
 
+#[derive(serde::Deserialize)]
+pub struct PrePublishParams {
+    name_or_identity: NameOrIdentity,
+}
+
+#[derive(serde::Deserialize)]
+pub struct PrePublishQueryParams {
+    #[serde(default)]
+    style: PrettyPrintStyle,
+}
+
+pub async fn pre_publish<S: NodeDelegate + ControlStateDelegate>(
+    State(ctx): State<S>,
+    Path(PrePublishParams { name_or_identity }): Path<PrePublishParams>,
+    Query(PrePublishQueryParams { style }): Query<PrePublishQueryParams>,
+    Extension(auth): Extension<SpacetimeAuth>,
+    body: Bytes,
+) -> axum::response::Result<axum::Json<PrintPlanResult>> {
+    // User should not be able to print migration plans for a database that they do not own
+    let database_identity = resolve_and_authenticate(&ctx, &name_or_identity, &auth).await?;
+    let style = match style {
+        PrettyPrintStyle::NoColor => AutoMigratePrettyPrintStyle::NoColor,
+        PrettyPrintStyle::AnsiColor => AutoMigratePrettyPrintStyle::AnsiColor,
+    };
+
+    let migrate_plan = ctx
+        .migrate_plan(
+            DatabaseDef {
+                database_identity,
+                program_bytes: body.into(),
+                num_replicas: None,
+                host_type: HostType::Wasm,
+            },
+            style,
+        )
+        .await
+        .map_err(log_and_500)?;
+
+    match migrate_plan {
+        MigratePlanResult::Success {
+            old_module_hash,
+            new_module_hash,
+            breaks_client,
+            plan,
+        } => {
+            let token = MigrationToken {
+                database_identity,
+                old_module_hash,
+                new_module_hash,
+            }
+            .hash();
+
+            Ok(PrintPlanResult {
+                token,
+                migrate_plan: plan,
+                break_clients: breaks_client,
+            })
+        }
+        MigratePlanResult::AutoMigrationError(e) => Err((
+            StatusCode::BAD_REQUEST,
+            format!("Automatic migration is not possible: {e}"),
+        )
+            .into()),
+    }
+    .map(axum::Json)
+}
+
+/// Resolves the [`NameOrIdentity`] to a database identity and checks if the
+/// `auth` identity owns the database.
+async fn resolve_and_authenticate<S: ControlStateDelegate>(
+    ctx: &S,
+    name_or_identity: &NameOrIdentity,
+    auth: &SpacetimeAuth,
+) -> axum::response::Result<Identity> {
+    let database_identity = name_or_identity.resolve(ctx).await?;
+
+    let database = worker_ctx_find_database(ctx, &database_identity)
+        .await?
+        .ok_or(NO_SUCH_DATABASE)?;
+
+    if database.owner_identity != auth.identity {
+        return Err((
+            StatusCode::UNAUTHORIZED,
+            format!(
+                "Identity does not own database, expected: {} got: {}",
+                database.owner_identity.to_hex(),
+                auth.identity.to_hex()
+            ),
+        )
+            .into());
+    }
+
+    Ok(database_identity)
+}
+
 #[derive(Deserialize)]
 pub struct DeleteDatabaseParams {
     name_or_identity: NameOrIdentity,
@@ -788,7 +916,8 @@ pub struct DatabaseRoutes<S> {
     pub logs_get: MethodRouter<S>,
     /// POST: /database/:name_or_identity/sql
     pub sql_post: MethodRouter<S>,
-
+    /// POST: /database/:name_or_identity/pre-publish
+    pub pre_publish: MethodRouter<S>,
     /// GET: /database/: name_or_identity/unstable/timestamp
     pub timestamp_get: MethodRouter<S>,
 }
@@ -813,6 +942,7 @@ where
             schema_get: get(schema::<S>),
             logs_get: get(logs::<S>),
             sql_post: post(sql::<S>),
+            pre_publish: post(pre_publish::<S>),
             timestamp_get: get(get_timestamp::<S>),
         }
     }
@@ -836,7 +966,8 @@ where
             .route("/schema", self.schema_get)
             .route("/logs", self.logs_get)
             .route("/sql", self.sql_post)
-            .route("/unstable/timestamp", self.timestamp_get);
+            .route("/unstable/timestamp", self.timestamp_get)
+            .route("/pre-publish", self.pre_publish);
 
         axum::Router::new()
             .route("/", self.root_post)