feat(api-public): expose actor kv api

Author: NathanFlurry

Cargo.lock

@@ -8,6 +8,7 @@ license.workspace = true
 [dependencies]
 anyhow.workspace = true
 axum.workspace = true
+base64.workspace = true
 gas.workspace = true
 epoxy.workspace = true
 futures-util.workspace = true
@@ -27,6 +28,7 @@ tokio.workspace = true
 tracing.workspace = true
 namespace.workspace = true
 pegboard.workspace = true
+pegboard-actor-kv.workspace = true
 universalpubsub.workspace = true
 uuid.workspace = true
 utoipa.workspace = true

engine/artifacts/errors/actor.kv_key_not_found.json

@@ -0,0 +1,67 @@
+use anyhow::*;
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
+use pegboard_actor_kv as actor_kv;
+use rivet_api_builder::ApiCtx;
+use rivet_util::Id;
+use serde::{Deserialize, Serialize};
+use utoipa::ToSchema;
+
+#[derive(Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct KvGetPath {
+	pub actor_id: Id,
+	pub key: String,
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(deny_unknown_fields)]
+pub struct KvGetQuery {}
+
+#[derive(Serialize, ToSchema)]
+#[schema(as = ActorsKvGetResponse)]
+pub struct KvGetResponse {
+	/// Value encoded in base 64.
+	pub value: String,
+	pub update_ts: i64,
+}
+
+#[utoipa::path(
+	get,
+	operation_id = "actors_kv_get",
+	path = "/actors/{actor_id}/kv/keys/{key}",
+	params(
+		("actor_id" = Id, Path),
+		("key" = String, Path),
+	),
+	responses(
+		(status = 200, body = KvGetResponse),
+	),
+)]
+#[tracing::instrument(skip_all)]
+pub async fn kv_get(ctx: ApiCtx, path: KvGetPath, _query: KvGetQuery) -> Result<KvGetResponse> {
+	// Decode base64 key
+	let key_bytes = BASE64_STANDARD
+		.decode(&path.key)
+		.context("failed to decode base64 key")?;
+
+	// Get the KV value
+	let udb = ctx.pools().udb()?;
+	let (keys, values, metadata) =
+		actor_kv::get(&*udb, path.actor_id, vec![key_bytes.clone()]).await?;
+
+	// Check if key was found
+	if keys.is_empty() {
+		return Err(pegboard::errors::Actor::KvKeyNotFound.build());
+	}
+
+	// Encode value as base64
+	let value_base64 = BASE64_STANDARD.encode(&values[0]);
+
+	Ok(KvGetResponse {
+		value: value_base64,
+		// NOTE: Intentionally uses different name in public API. `create_ts` is actually
+		// `update_ts`.
+		update_ts: metadata[0].create_ts,
+	})
+}

engine/artifacts/openapi.json

@@ -1,4 +1,5 @@
 pub mod create;
 pub mod delete;
+pub mod kv_get;
 pub mod list;
 pub mod list_names;

engine/packages/api-peer/Cargo.toml

@@ -26,6 +26,10 @@ pub async fn router(
 			.route("/actors", post(actors::create::create))
 			.route("/actors/{actor_id}", delete(actors::delete::delete))
 			.route("/actors/names", get(actors::list_names::list_names))
+			.route(
+				"/actors/{actor_id}/kv/keys/{key}",
+				get(actors::kv_get::kv_get),
+			)
 			// MARK: Runners
 			.route("/runners", get(runners::list))
 			.route("/runners/names", get(runners::list_names))

engine/packages/api-peer/src/actors/kv_get.rs

@@ -30,6 +30,7 @@ serde.workspace = true
 tokio.workspace = true
 tower-http.workspace = true
 tracing.workspace = true
+urlencoding.workspace = true
 utoipa.workspace = true
 
 [build-dependencies]

engine/packages/api-peer/src/actors/mod.rs

@@ -0,0 +1,75 @@
+use anyhow::Result;
+use axum::response::{IntoResponse, Response};
+use rivet_api_builder::{
+	ApiError,
+	extract::{Extension, Path},
+};
+use rivet_api_util::request_remote_datacenter_raw;
+use rivet_util::Id;
+use serde::{Deserialize, Serialize};
+use utoipa::ToSchema;
+
+use crate::ctx::ApiCtx;
+
+#[derive(Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct KvGetPath {
+	pub actor_id: Id,
+	pub key: String,
+}
+
+#[derive(Serialize, ToSchema)]
+#[schema(as = ActorsKvGetResponse)]
+pub struct KvGetResponse {
+	pub value: String,
+	pub update_ts: i64,
+}
+
+#[utoipa::path(
+	get,
+	operation_id = "actors_kv_get",
+	path = "/actors/{actor_id}/kv/keys/{key}",
+	params(
+		("actor_id" = Id, Path),
+		("key" = String, Path),
+	),
+	responses(
+		(status = 200, body = KvGetResponse),
+	),
+	security(("bearer_auth" = [])),
+)]
+#[tracing::instrument(skip_all)]
+pub async fn kv_get(Extension(ctx): Extension<ApiCtx>, Path(path): Path<KvGetPath>) -> Response {
+	match kv_get_inner(ctx, path).await {
+		Ok(response) => response,
+		Err(err) => ApiError::from(err).into_response(),
+	}
+}
+
+#[tracing::instrument(skip_all)]
+async fn kv_get_inner(ctx: ApiCtx, path: KvGetPath) -> Result<Response> {
+	use axum::Json;
+
+	ctx.auth().await?;
+
+	if path.actor_id.label() == ctx.config().dc_label() {
+		let peer_path = rivet_api_peer::actors::kv_get::KvGetPath {
+			actor_id: path.actor_id,
+			key: path.key,
+		};
+		let peer_query = rivet_api_peer::actors::kv_get::KvGetQuery {};
+		let res = rivet_api_peer::actors::kv_get::kv_get(ctx.into(), peer_path, peer_query).await?;
+
+		Ok(Json(res).into_response())
+	} else {
+		request_remote_datacenter_raw(
+			&ctx,
+			path.actor_id.label(),
+			&format!("/actors/{}/kv/keys/{}", path.actor_id, urlencoding::encode(&path.key)),
+			axum::http::Method::GET,
+			Option::<&()>::None,
+			Option::<&()>::None,
+		)
+		.await
+	}
+}

engine/packages/api-peer/src/router.rs

@@ -1,6 +1,7 @@
 pub mod create;
 pub mod delete;
 pub mod get_or_create;
+pub mod kv_get;
 pub mod list;
 pub mod list_names;
 pub mod utils;