refactor(native): cache & job runtime (#15139)

This commit is contained in:
DarkSky
2026-06-22 11:48:37 +08:00
committed by GitHub
parent 7ea8800c99
commit bb8454e7e1
39 changed files with 6390 additions and 62 deletions
Generated
+625 -54
View File
File diff suppressed because it is too large Load Diff
+1
View File
@@ -21,6 +21,7 @@ resolver = "3"
anyhow = "1"
arbitrary = { version = "1.3", features = ["derive"] }
assert-json-diff = "2.0"
base64 = "0.22.1"
base64-simd = "0.8"
bitvec = "1.0"
block2 = "0.6"
+14 -1
View File
@@ -16,10 +16,13 @@ affine_common = { workspace = true, features = [
"ydoc-loader",
] }
anyhow = { workspace = true }
aws-sdk-s3 = "1.115"
base64 = { workspace = true }
chrono = { workspace = true }
doc_extractor = { workspace = true }
file-format = { workspace = true }
hex = { workspace = true }
homedir = { workspace = true }
image = { workspace = true }
infer = { workspace = true }
jsonschema = "0.46"
@@ -39,8 +42,18 @@ serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
sha2 = { workspace = true }
sha3 = { workspace = true }
sqlx = { workspace = true, default-features = false, features = [
"chrono",
"json",
"macros",
"migrate",
"postgres",
"runtime-tokio",
] }
tiktoken-rs = { workspace = true }
tokio = { workspace = true, features = ["rt-multi-thread", "sync"] }
url = { workspace = true }
uuid = { workspace = true, features = ["v4"] }
v_htmlescape = { workspace = true }
y-octo = { workspace = true, features = ["large_refs"] }
@@ -52,7 +65,7 @@ mimalloc = { workspace = true, features = ["local_dynamic_tls"] }
[dev-dependencies]
rayon = { workspace = true }
tokio = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
[build-dependencies]
napi-build = { workspace = true }
+214
View File
@@ -1,5 +1,66 @@
/* auto-generated by NAPI-RS */
/* eslint-disable */
export declare class BackendRuntime {
completeBlobUpload(workspaceId: string, key: string, expectedSize: number, expectedMime: string): Promise<RuntimeBlobCompleteResult>
completeFsBlobUpload(root: string, bucket: string, workspaceId: string, key: string, expectedSize: number, expectedMime: string): Promise<RuntimeBlobCompleteResult>
cleanupExpiredPendingBlobs(cutoffMs: number, limit: number): Promise<RuntimeBlobCleanupResult>
releaseDeletedBlobs(workspaceId: string, limit: number): Promise<RuntimeBlobCleanupResult>
acquireCoordinationLease(key: string, owner: string, ttlMs: number): Promise<CoordinationLeaseGrant | null>
releaseCoordinationLease(key: string, owner: string, fencingToken: bigint | number): Promise<boolean>
renewCoordinationLease(key: string, owner: string, fencingToken: bigint | number, ttlMs: number): Promise<boolean>
/**
* Merge pending doc updates with y-octo and persist the merged snapshot.
*
* Do not use this for snapshots that will be sent back to yjs clients until
* the y-octo/yjs round-trip compatibility issue is resolved.
*/
compactPendingDocUpdates(workspaceId: string, docId: string, batchLimit: number, historyMinIntervalMs: number, owner: string, leaseTtlMs: number): Promise<RuntimeDocCompactionResult>
upsertDocSnapshot(workspaceId: string, docId: string, blob: Buffer, timestampMs: number, editorId?: string | undefined | null): Promise<boolean>
createDocHistory(input: RuntimeDocHistoryInput): Promise<boolean>
deleteDocStorage(workspaceId: string, docId: string): Promise<void>
putRuntimeGateIfAbsent(key: string, ttlMs: number): Promise<boolean>
cleanupExpiredRuntimeGates(limit: number): Promise<number>
cleanupExpiredUserSessions(limit: number): Promise<number>
cleanupExpiredSnapshotHistories(limit: number): Promise<number>
objectStorageHealth(): RuntimeObjectStorageHealth
objectStoragePut(key: string, body: Buffer, metadata?: RuntimeObjectStoragePutOptions | undefined | null): Promise<void>
objectStoragePresignPut(key: string, metadata?: RuntimeObjectStoragePutOptions | undefined | null): Promise<RuntimePresignedObjectRequest>
objectStorageCreateMultipartUpload(key: string, metadata?: RuntimeObjectStoragePutOptions | undefined | null): Promise<RuntimeMultipartUploadInit | null>
objectStoragePresignUploadPart(key: string, uploadId: string, partNumber: number): Promise<RuntimePresignedObjectRequest>
objectStorageListMultipartUploadParts(key: string, uploadId: string): Promise<Array<RuntimeMultipartUploadPart>>
objectStorageCompleteMultipartUpload(key: string, uploadId: string, parts: Array<RuntimeMultipartUploadPart>): Promise<void>
objectStorageAbortMultipartUpload(key: string, uploadId: string): Promise<void>
objectStorageHead(key: string): Promise<RuntimeObjectMetadata | null>
objectStorageGet(key: string): Promise<RuntimeObjectGetResult | null>
objectStorageList(prefix?: string | undefined | null): Promise<Array<RuntimeObjectListEntry>>
objectStorageDelete(key: string): Promise<void>
createAuthChallenge(purpose: string, token: string, payload: any, ttlMs: number): Promise<boolean>
getAuthChallenge(purpose: string, token: string): Promise<any | null>
consumeAuthChallenge(purpose: string, token: string): Promise<any | null>
createVerificationToken(tokenType: number, credential: string | undefined | null, ttlMs: number): Promise<string>
getVerificationToken(tokenType: number, token: string, keep?: boolean | undefined | null): Promise<RuntimeVerificationTokenRecord | null>
verifyVerificationToken(tokenType: number, token: string, credential?: string | undefined | null, keep?: boolean | undefined | null): Promise<RuntimeVerificationTokenRecord | null>
cleanupExpiredVerificationTokens(limit: number): Promise<number>
upsertMagicLinkOtp(email: string, otpHash: string, token: string, clientNonce: string | undefined | null, ttlMs: number): Promise<void>
consumeMagicLinkOtp(email: string, otpHash: string, clientNonce?: string | undefined | null): Promise<RuntimeMagicLinkOtpConsumeResult>
createWorkspaceInviteLink(workspaceId: string, inviteId: string, inviterUserId: string, ttlMs: number): Promise<RuntimeWorkspaceInviteLinkRecord>
getWorkspaceInviteLink(workspaceId: string): Promise<RuntimeWorkspaceInviteLinkRecord | null>
getWorkspaceInviteLinkById(inviteId: string): Promise<RuntimeWorkspaceInviteLinkRecord | null>
revokeWorkspaceInviteLink(workspaceId: string): Promise<boolean>
createByokLocalLease(activeKey: string, leaseId: string, payload: any, ttlMs: number): Promise<RuntimeByokLocalLeaseRecord>
getByokLocalLease(leaseId: string): Promise<RuntimeByokLocalLeaseRecord | null>
cleanupExpiredRuntimeStates(limit: number): Promise<number>
refreshWorkspaceAdminStatsDirty(batchLimit: number, owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsRefreshResult>
recalibrateWorkspaceAdminStats(lastSid: number, batchLimit: number, owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsRecalibrationResult>
writeWorkspaceAdminStatsDailySnapshot(owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsSnapshotResult>
recalibrateWorkspaceAdminStatsDaily(batchLimit: number, owner: string, leaseTtlMs: number, lockRetryTimes: number, lockRetryDelayMs: number): Promise<RuntimeWorkspaceStatsDailyRecalibrationResult>
constructor()
start(): Promise<void>
stop(): Promise<void>
health(): Promise<BackendRuntimeHealth>
runMigrations(): Promise<void>
}
export declare class LlmStreamHandle {
abort(): void
}
@@ -74,6 +135,12 @@ export interface AssertSafeUrlRequest {
url: string
}
export interface BackendRuntimeHealth {
started: boolean
databaseConnected: boolean
objectStorageConfigured: boolean
}
export declare function buildPublicRootDoc(rootDocBin: Buffer, docMetas: Array<PublicDocMetaInput>): Buffer
export interface BuiltInPromptRenderContract {
@@ -164,6 +231,12 @@ export interface CommandResponse {
error?: LicenseError
}
export interface CoordinationLeaseGrant {
key: string
owner: string
fencingToken: bigint | number
}
/**
* Converts markdown content to AFFiNE-compatible y-octo document binary.
*
@@ -738,6 +811,147 @@ export declare function resolveEntitlementV1(input: ResolveEntitlementInput): Re
export declare function runNativeActionRecipePreparedStream(input: ActionRuntimeInput, callback: ((err: Error | null, arg: string) => void)): LlmStreamHandle
export interface RuntimeBlobCleanupResult {
scanned: number
deleted: number
abortedMultipart: number
workspaceIds: Array<string>
}
export interface RuntimeBlobCompleteResult {
ok: boolean
reason?: string
contentType?: string
contentLength?: number
lastModifiedMs?: number
}
export interface RuntimeByokLocalLeaseRecord {
leaseId: string
payload: any
expiresAtMs: number
}
export interface RuntimeDocCompactionResult {
leaseAcquired: boolean
merged: boolean
workspaceId: string
docId: string
updatesMerged: number
historyCreated: boolean
}
export interface RuntimeDocHistoryInput {
workspaceId: string
docId: string
blob: Buffer
timestampMs: number
editorId?: string
force: boolean
historyMinIntervalMs: number
historyMaxAgeMs: number
}
export interface RuntimeMagicLinkOtpConsumeResult {
ok: boolean
token?: string
reason?: string
}
export interface RuntimeMultipartUploadInit {
uploadId: string
expiresAtMs: number
}
export interface RuntimeMultipartUploadPart {
partNumber: number
etag: string
}
export interface RuntimeObjectGetResult {
body: Buffer
metadata: RuntimeObjectMetadata
}
export interface RuntimeObjectListEntry {
key: string
contentLength: number
lastModifiedMs: number
}
export interface RuntimeObjectMetadata {
contentType: string
contentLength: number
lastModifiedMs: number
checksumCrc32?: string
}
export interface RuntimeObjectStorageHealth {
configured: boolean
provider?: string
bucket?: string
endpoint?: string
region?: string
hasCredentials: boolean
forcePathStyle: boolean
requestTimeoutMs?: number
minPartSize?: number
presignExpiresInSeconds?: number
presignSignContentTypeForPut?: boolean
usePresignedUrl: boolean
clientBuildable: boolean
}
export interface RuntimeObjectStoragePutOptions {
contentType?: string
contentLength?: number
checksumCrc32?: string
}
export interface RuntimePresignedObjectRequest {
url: string
headersJson: string
expiresAtMs: number
}
export interface RuntimeVerificationTokenRecord {
tokenType: number
token: string
credential?: string
expiresAtMs: number
}
export interface RuntimeWorkspaceInviteLinkRecord {
workspaceId: string
inviteId: string
inviterUserId: string
expiresAtMs: number
}
export interface RuntimeWorkspaceStatsDailyRecalibrationResult {
processed: number
lastSid: number
snapshotted: number
skipped: boolean
}
export interface RuntimeWorkspaceStatsRecalibrationResult {
processed: number
lastSid: number
skipped: boolean
}
export interface RuntimeWorkspaceStatsRefreshResult {
processed: number
backlog: number
skipped: boolean
}
export interface RuntimeWorkspaceStatsSnapshotResult {
snapshotted: number
skipped: boolean
}
export declare function safeFetch(request: SafeFetchRequest): Promise<SafeFetchResponse>
export type SafeFetchMethod = 'get'|
+1 -1
View File
@@ -29,7 +29,7 @@
"test": "node --test ./__tests__/**/*.spec.js",
"bench": "node ./benchmark/index.js",
"build": "napi build --release --strip --no-const-enum",
"build:debug": "napi build"
"build:debug": "napi build --no-const-enum"
},
"devDependencies": {
"@napi-rs/cli": "3.5.0",
@@ -0,0 +1,296 @@
use std::{
fs,
path::{Path, PathBuf},
};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use napi::Result;
use serde::Deserialize;
use sha2::{Digest, Sha256};
use super::{BackendRuntime, error::napi_error, types::RuntimeBlobCompleteResult};
const MAX_BLOB_SIZE: i64 = i32::MAX as i64;
fn object_missing_error(err: &napi::Error) -> bool {
let message = err.to_string();
message.contains("NoSuchKey") || message.contains("NotFound") || message.contains("not found")
}
fn blob_complete_failure(reason: &str) -> RuntimeBlobCompleteResult {
RuntimeBlobCompleteResult {
ok: false,
reason: Some(reason.to_string()),
content_type: None,
content_length: None,
last_modified_ms: None,
}
}
fn blob_complete_success(
content_type: String,
content_length: i64,
last_modified_ms: i64,
) -> RuntimeBlobCompleteResult {
RuntimeBlobCompleteResult {
ok: true,
reason: None,
content_type: Some(content_type),
content_length: Some(content_length),
last_modified_ms: Some(last_modified_ms),
}
}
fn normalize_base64_url_key(key: &str) -> &str {
key.trim_end_matches('=')
}
fn sha256_base64_url(body: &[u8]) -> String {
URL_SAFE_NO_PAD.encode(Sha256::digest(body))
}
fn sha256_base64_url_matches(body: &[u8], key: &str) -> bool {
sha256_base64_url(body) == normalize_base64_url_key(key)
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct FsBlobMetadata {
content_type: String,
content_length: i64,
last_modified: i64,
}
fn normalize_storage_key(key: &str) -> Result<Vec<String>> {
let normalized = key.replace('\\', "/");
let segments = normalized.split('/').map(ToString::to_string).collect::<Vec<_>>();
if normalized.is_empty()
|| normalized.starts_with('/')
|| segments
.iter()
.any(|segment| segment.is_empty() || segment == "." || segment == "..")
{
return Err(napi_error(format!("Invalid storage key: {key}")));
}
Ok(segments)
}
fn fs_bucket_path(root: &str, bucket: &str) -> PathBuf {
if let Some(stripped) = root.strip_prefix("~/")
&& let Ok(Some(home)) = homedir::my_home()
{
return home.join(stripped).join(bucket);
}
Path::new(root).join(bucket)
}
fn fs_object_path(root: &str, bucket: &str, key: &str) -> Result<PathBuf> {
let mut path = fs_bucket_path(root, bucket);
for segment in normalize_storage_key(key)? {
path.push(segment);
}
Ok(path)
}
fn read_fs_metadata(path: &Path) -> Result<Option<FsBlobMetadata>> {
let metadata_path = PathBuf::from(format!("{}.metadata.json", path.display()));
let raw = match fs::read_to_string(metadata_path) {
Ok(raw) => raw,
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None),
Err(err) => {
return Err(napi_error(format!("BlobComplete read fs metadata failed: {err}")));
}
};
serde_json::from_str(&raw).map(Some).map_err(|err| {
napi_error(format!(
"BlobComplete parse fs metadata failed for {}: {err}",
path.display()
))
})
}
async fn upsert_completed_blob(
runtime: &BackendRuntime,
workspace_id: &str,
key: &str,
mime: &str,
size: i64,
) -> Result<()> {
if !(0..=MAX_BLOB_SIZE).contains(&size) {
return Err(napi_error("BlobComplete size exceeds limit"));
}
let size = i32::try_from(size).map_err(|_| napi_error("BlobComplete size exceeds limit"))?;
sqlx::query(
r#"
INSERT INTO blobs (workspace_id, key, mime, size, status, upload_id)
VALUES ($1, $2, $3, $4, 'completed', NULL)
ON CONFLICT (workspace_id, key)
DO UPDATE SET
mime = EXCLUDED.mime,
size = EXCLUDED.size,
status = EXCLUDED.status,
upload_id = NULL
"#,
)
.bind(workspace_id)
.bind(key)
.bind(mime)
.bind(size)
.execute(&runtime.pool().await?)
.await
.map_err(|err| napi_error(format!("BlobComplete upsert metadata failed: {err}")))?;
Ok(())
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn complete_blob_upload(
&self,
workspace_id: String,
key: String,
expected_size: i64,
expected_mime: String,
) -> Result<RuntimeBlobCompleteResult> {
if !(0..=MAX_BLOB_SIZE).contains(&expected_size) {
return Ok(blob_complete_failure("size_too_large"));
}
let object_key = format!("{workspace_id}/{key}");
let object = match self.object_storage_get(object_key.clone()).await {
Ok(Some(object)) => object,
Ok(None) => return Ok(blob_complete_failure("not_found")),
Err(err) if object_missing_error(&err) => return Ok(blob_complete_failure("not_found")),
Err(err) => return Err(err),
};
if !(0..=MAX_BLOB_SIZE).contains(&object.metadata.content_length) {
match self.object_storage_delete(object_key).await {
Ok(()) => {}
Err(err) if object_missing_error(&err) => {}
Err(err) => return Err(err),
}
return Ok(blob_complete_failure("size_too_large"));
}
if object.metadata.content_length != expected_size {
return Ok(blob_complete_failure("size_mismatch"));
}
if !expected_mime.is_empty() && object.metadata.content_type != expected_mime {
return Ok(blob_complete_failure("mime_mismatch"));
}
if !sha256_base64_url_matches(&object.body, &key) {
match self.object_storage_delete(object_key).await {
Ok(()) => {}
Err(err) if object_missing_error(&err) => {}
Err(err) => return Err(err),
}
return Ok(blob_complete_failure("checksum_mismatch"));
}
upsert_completed_blob(
self,
&workspace_id,
&key,
&object.metadata.content_type,
object.metadata.content_length,
)
.await?;
Ok(blob_complete_success(
object.metadata.content_type,
object.metadata.content_length,
object.metadata.last_modified_ms,
))
}
#[napi]
pub async fn complete_fs_blob_upload(
&self,
root: String,
bucket: String,
workspace_id: String,
key: String,
expected_size: i64,
expected_mime: String,
) -> Result<RuntimeBlobCompleteResult> {
if !(0..=MAX_BLOB_SIZE).contains(&expected_size) {
return Ok(blob_complete_failure("size_too_large"));
}
let storage_key = format!("{workspace_id}/{key}");
let path = fs_object_path(&root, &bucket, &storage_key)?;
let metadata = match read_fs_metadata(&path)? {
Some(metadata) => metadata,
None => return Ok(blob_complete_failure("not_found")),
};
if !(0..=MAX_BLOB_SIZE).contains(&metadata.content_length) {
let _ = fs::remove_file(&path);
let _ = fs::remove_file(PathBuf::from(format!("{}.metadata.json", path.display())));
return Ok(blob_complete_failure("size_too_large"));
}
if metadata.content_length != expected_size {
return Ok(blob_complete_failure("size_mismatch"));
}
if !expected_mime.is_empty() && metadata.content_type != expected_mime {
return Ok(blob_complete_failure("mime_mismatch"));
}
let body = match fs::read(&path) {
Ok(body) => body,
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(blob_complete_failure("not_found")),
Err(err) => return Err(napi_error(format!("BlobComplete read fs object failed: {err}"))),
};
if !sha256_base64_url_matches(&body, &key) {
let _ = fs::remove_file(&path);
let _ = fs::remove_file(PathBuf::from(format!("{}.metadata.json", path.display())));
return Ok(blob_complete_failure("checksum_mismatch"));
}
upsert_completed_blob(
self,
&workspace_id,
&key,
&metadata.content_type,
metadata.content_length,
)
.await?;
Ok(blob_complete_success(
metadata.content_type,
metadata.content_length,
metadata.last_modified,
))
}
}
#[cfg(test)]
mod tests {
use super::{sha256_base64_url, sha256_base64_url_matches};
#[test]
fn sha256_base64_url_omits_padding() {
assert_eq!(
sha256_base64_url(b"hello"),
"LPJNul-wow4m6DsqxbninhsWHlwfp0JecwQzYpOLmCQ"
);
}
#[test]
fn sha256_base64_url_matches_legacy_padding() {
assert!(sha256_base64_url_matches(
b"hello",
"LPJNul-wow4m6DsqxbninhsWHlwfp0JecwQzYpOLmCQ="
));
}
}
@@ -0,0 +1,190 @@
use chrono::{DateTime, Utc};
use napi::Result;
use sqlx::{FromRow, PgPool};
use super::{BackendRuntime, error::napi_error, types::RuntimeBlobCleanupResult};
#[derive(FromRow)]
struct BlobRow {
workspace_id: String,
key: String,
upload_id: Option<String>,
}
struct BlobReclaimerStore {
pool: PgPool,
}
impl BlobReclaimerStore {
fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn load_expired_pending(&self, cutoff: DateTime<Utc>, limit: i64) -> Result<Vec<BlobRow>> {
sqlx::query_as::<_, BlobRow>(
r#"
SELECT workspace_id, key, upload_id
FROM blobs
WHERE status = 'pending'
AND deleted_at IS NULL
AND created_at < $1
ORDER BY created_at ASC
LIMIT $2
"#,
)
.bind(cutoff)
.bind(limit)
.fetch_all(&self.pool)
.await
.map_err(|err| napi_error(format!("BlobReclaimer load pending blobs failed: {err}")))
}
async fn load_deleted(&self, workspace_id: &str, limit: i64) -> Result<Vec<BlobRow>> {
sqlx::query_as::<_, BlobRow>(
r#"
SELECT workspace_id, key, upload_id
FROM blobs
WHERE workspace_id = $1
AND deleted_at IS NOT NULL
ORDER BY deleted_at ASC
LIMIT $2
"#,
)
.bind(workspace_id)
.bind(limit)
.fetch_all(&self.pool)
.await
.map_err(|err| napi_error(format!("BlobReclaimer load deleted blobs failed: {err}")))
}
async fn delete_pending_metadata(&self, workspace_id: &str, key: &str) -> Result<i64> {
let result = sqlx::query(
r#"
DELETE FROM blobs
WHERE workspace_id = $1 AND key = $2
AND status = 'pending'
AND deleted_at IS NULL
"#,
)
.bind(workspace_id)
.bind(key)
.execute(&self.pool)
.await
.map_err(|err| napi_error(format!("BlobReclaimer delete pending blob metadata failed: {err}")))?;
Ok(result.rows_affected() as i64)
}
async fn delete_released_metadata(&self, workspace_id: &str, key: &str) -> Result<i64> {
let result = sqlx::query(
r#"
DELETE FROM blobs
WHERE workspace_id = $1 AND key = $2
AND deleted_at IS NOT NULL
"#,
)
.bind(workspace_id)
.bind(key)
.execute(&self.pool)
.await
.map_err(|err| napi_error(format!("BlobReclaimer delete blob metadata failed: {err}")))?;
Ok(result.rows_affected() as i64)
}
}
fn object_missing_error(err: &napi::Error) -> bool {
let message = err.to_string();
message.contains("NoSuchKey")
|| message.contains("NoSuchUpload")
|| message.contains("NotFound")
|| message.contains("not found")
}
async fn delete_object_idempotent(runtime: &BackendRuntime, key: &str) -> Result<()> {
match runtime.object_storage_delete_object(key).await {
Ok(()) => Ok(()),
Err(err) if object_missing_error(&err) => Ok(()),
Err(err) => Err(err),
}
}
async fn abort_upload_idempotent(runtime: &BackendRuntime, key: &str, upload_id: &str) -> Result<()> {
match runtime.object_storage_abort_upload(key, upload_id).await {
Ok(()) => Ok(()),
Err(err) if object_missing_error(&err) => Ok(()),
Err(err) => Err(err),
}
}
fn push_workspace_once(workspace_ids: &mut Vec<String>, workspace_id: &str) {
if !workspace_ids.iter().any(|id| id == workspace_id) {
workspace_ids.push(workspace_id.to_string());
}
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn cleanup_expired_pending_blobs(&self, cutoff_ms: i64, limit: i64) -> Result<RuntimeBlobCleanupResult> {
if limit <= 0 {
return Err(napi_error("pending blob cleanup limit must be positive"));
}
let cutoff = DateTime::<Utc>::from_timestamp_millis(cutoff_ms)
.ok_or_else(|| napi_error("pending blob cleanup cutoff is invalid"))?;
let store = BlobReclaimerStore::new(self.pool().await?);
let rows = store.load_expired_pending(cutoff, limit).await?;
let mut deleted = 0;
let mut aborted_multipart = 0;
let mut workspace_ids = Vec::new();
for row in &rows {
let object_key = format!("{}/{}", row.workspace_id, row.key);
if let Some(upload_id) = row.upload_id.as_deref() {
abort_upload_idempotent(self, &object_key, upload_id).await?;
aborted_multipart += 1;
}
delete_object_idempotent(self, &object_key).await?;
let affected = store.delete_pending_metadata(&row.workspace_id, &row.key).await?;
if affected > 0 {
deleted += affected;
push_workspace_once(&mut workspace_ids, &row.workspace_id);
}
}
Ok(RuntimeBlobCleanupResult {
scanned: rows.len() as i64,
deleted,
aborted_multipart,
workspace_ids,
})
}
#[napi]
pub async fn release_deleted_blobs(&self, workspace_id: String, limit: i64) -> Result<RuntimeBlobCleanupResult> {
if limit <= 0 {
return Err(napi_error("deleted blob release limit must be positive"));
}
let store = BlobReclaimerStore::new(self.pool().await?);
let rows = store.load_deleted(&workspace_id, limit).await?;
let mut deleted = 0;
let mut workspace_ids = Vec::new();
for row in &rows {
let object_key = format!("{}/{}", row.workspace_id, row.key);
delete_object_idempotent(self, &object_key).await?;
let affected = store.delete_released_metadata(&row.workspace_id, &row.key).await?;
if affected > 0 {
deleted += affected;
push_workspace_once(&mut workspace_ids, &row.workspace_id);
}
}
Ok(RuntimeBlobCleanupResult {
scanned: rows.len() as i64,
deleted,
aborted_multipart: 0,
workspace_ids,
})
}
}
@@ -0,0 +1,128 @@
use std::{
collections::HashMap,
env, fs,
path::{Path, PathBuf},
};
use napi::Result;
use serde::Deserialize;
use super::{
error::napi_error,
object_storage::{ObjectStorageConfig, StorageProviderConfig},
};
#[derive(Clone, Debug)]
pub(super) struct RuntimeConfig {
pub(super) database_url: String,
pub(super) storage: Option<ObjectStorageConfig>,
}
impl RuntimeConfig {
pub(super) fn from_config_files() -> Result<Self> {
let database_url =
database_url_from_config_files()?.unwrap_or_else(|| "postgresql://localhost:5432/affine".to_string());
let storage = ObjectStorageConfig::from_config_files()?;
Ok(Self { database_url, storage })
}
}
#[derive(Debug, Deserialize)]
struct AppConfigFile {
db: Option<DbConfigFile>,
storages: Option<HashMap<String, StorageProviderConfig>>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct DbConfigFile {
datasource_url: Option<String>,
}
fn database_url_from_config_files() -> Result<Option<String>> {
let mut database_url = None;
for path in config_json_paths() {
if !path.exists() {
continue;
}
let raw = fs::read_to_string(&path)
.map_err(|err| napi_error(format!("failed to read config file {}: {err}", path.display())))?;
let config: AppConfigFile = serde_json::from_str(&raw)
.map_err(|err| napi_error(format!("failed to parse config file {}: {err}", path.display())))?;
if let Some(next) = config.db.and_then(|db| db.datasource_url)
&& !next.trim().is_empty()
{
database_url = Some(next);
}
}
Ok(database_url)
}
pub(super) fn blob_storage_config_from_config_files() -> Result<Option<StorageProviderConfig>> {
let mut storage = None;
for path in config_json_paths() {
if !path.exists() {
continue;
}
let raw = fs::read_to_string(&path)
.map_err(|err| napi_error(format!("failed to read config file {}: {err}", path.display())))?;
let config: AppConfigFile = serde_json::from_str(&raw)
.map_err(|err| napi_error(format!("failed to parse config file {}: {err}", path.display())))?;
if let Some(next) = config.storages.and_then(|mut storages| storages.remove("blob.storage")) {
storage = Some(next);
}
}
Ok(storage)
}
pub(super) fn config_json_paths() -> Vec<PathBuf> {
let mut paths = Vec::new();
if let Ok(exe) = env::current_exe()
&& let Some(dir) = exe.parent()
{
paths.push(config_in(dir));
}
if let Ok(cwd) = env::current_dir() {
paths.push(config_in(&cwd));
}
dedupe_paths(paths)
}
fn config_in(dir: &Path) -> PathBuf {
dir.join("config.json")
}
fn dedupe_paths(paths: Vec<PathBuf>) -> Vec<PathBuf> {
let mut deduped = Vec::new();
for path in paths {
if !deduped.contains(&path) {
deduped.push(path);
}
}
deduped
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn config_paths_are_limited_to_executable_dir_and_cwd() {
let paths = config_json_paths();
assert!(!paths.is_empty());
assert!(paths.len() <= 2);
assert!(
paths
.iter()
.all(|path| path.file_name().is_some_and(|name| name == "config.json"))
);
assert!(paths.iter().all(|path| !path.to_string_lossy().contains(".affine")));
assert!(
paths
.iter()
.all(|path| !path.to_string_lossy().contains("packages/backend/server"))
);
}
}
@@ -0,0 +1,11 @@
pub(super) const DEFAULT_HISTORY_PERIOD_SECONDS: i32 = 7 * 24 * 60 * 60;
pub(super) const BYOK_LOCAL_LEASE_ACTIVE_PURPOSE: &str = "copilot_byok_local_lease:active";
pub(super) const BYOK_LOCAL_LEASE_PURPOSE: &str = "copilot_byok_local_lease";
pub(super) const MAGIC_LINK_OTP_PURPOSE: &str = "magic_link_otp";
pub(super) const MAX_MAGIC_LINK_OTP_ATTEMPTS: i32 = 10;
pub(super) const WORKSPACE_INVITE_LINK_ID_PURPOSE: &str = "workspace_invite_link:id";
pub(super) const WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE: &str = "workspace_invite_link:workspace";
pub(super) const WORKSPACE_STATS_LEASE_KEY: &str = "workspace:admin-stats:refresh";
pub(super) const WORKSPACE_STATS_LOCK_NAMESPACE: i64 = 97_301;
pub(super) const WORKSPACE_STATS_REFRESH_LOCK_KEY: i64 = 1;
pub(super) const RUNTIME_MIGRATIONS: &str = include_str!("sql/runtime_migrations.sql");
@@ -0,0 +1,138 @@
use napi::Result;
use sqlx::{FromRow, PgPool};
use super::{BackendRuntime, error::napi_error, types::CoordinationLeaseGrant};
#[derive(FromRow)]
struct LeaseGrantRow {
fencing_token: i64,
}
struct CoordinationLeaseStore {
pool: PgPool,
}
impl CoordinationLeaseStore {
fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn acquire(&self, key: String, owner: String, ttl_ms: i64) -> Result<Option<CoordinationLeaseGrant>> {
let row = sqlx::query_as::<_, LeaseGrantRow>(
r#"
INSERT INTO runtime_leases (key, owner, fencing_token, expires_at)
VALUES ($1, $2, 1, CURRENT_TIMESTAMP + ($3 * INTERVAL '1 millisecond'))
ON CONFLICT (key) DO UPDATE
SET owner = EXCLUDED.owner,
fencing_token = runtime_leases.fencing_token + 1,
expires_at = EXCLUDED.expires_at,
updated_at = CURRENT_TIMESTAMP
WHERE runtime_leases.expires_at <= CURRENT_TIMESTAMP
RETURNING fencing_token
"#,
)
.bind(&key)
.bind(&owner)
.bind(ttl_ms as f64)
.fetch_optional(&self.pool)
.await
.map_err(|err| napi_error(format!("CoordinationLease acquire failed: {err}")))?;
Ok(row.map(|row| CoordinationLeaseGrant {
key,
owner,
fencing_token: row.fencing_token,
}))
}
async fn release(&self, key: &str, owner: &str, fencing_token: i64) -> Result<bool> {
let result = sqlx::query(
r#"
DELETE FROM runtime_leases
WHERE key = $1 AND owner = $2 AND fencing_token = $3
"#,
)
.bind(key)
.bind(owner)
.bind(fencing_token)
.execute(&self.pool)
.await
.map_err(|err| napi_error(format!("CoordinationLease release failed: {err}")))?;
Ok(result.rows_affected() == 1)
}
async fn renew(&self, key: &str, owner: &str, fencing_token: i64, ttl_ms: i64) -> Result<bool> {
let result = sqlx::query(
r#"
UPDATE runtime_leases
SET expires_at = CURRENT_TIMESTAMP + ($4 * INTERVAL '1 millisecond'),
updated_at = CURRENT_TIMESTAMP
WHERE key = $1
AND owner = $2
AND fencing_token = $3
AND expires_at > CURRENT_TIMESTAMP
"#,
)
.bind(key)
.bind(owner)
.bind(fencing_token)
.bind(ttl_ms as f64)
.execute(&self.pool)
.await
.map_err(|err| napi_error(format!("CoordinationLease renew failed: {err}")))?;
Ok(result.rows_affected() == 1)
}
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn acquire_coordination_lease(
&self,
key: String,
owner: String,
ttl_ms: i64,
) -> Result<Option<CoordinationLeaseGrant>> {
if ttl_ms <= 0 {
return Err(napi_error("coordination lease ttl must be positive"));
}
if owner.is_empty() {
return Err(napi_error("coordination lease owner is required"));
}
CoordinationLeaseStore::new(self.pool().await?)
.acquire(key, owner, ttl_ms)
.await
}
#[napi]
pub async fn release_coordination_lease(
&self,
key: String,
owner: String,
#[napi(ts_arg_type = "bigint | number")] fencing_token: i64,
) -> Result<bool> {
CoordinationLeaseStore::new(self.pool().await?)
.release(&key, &owner, fencing_token)
.await
}
#[napi]
pub async fn renew_coordination_lease(
&self,
key: String,
owner: String,
#[napi(ts_arg_type = "bigint | number")] fencing_token: i64,
ttl_ms: i64,
) -> Result<bool> {
if ttl_ms <= 0 {
return Err(napi_error("coordination lease ttl must be positive"));
}
CoordinationLeaseStore::new(self.pool().await?)
.renew(&key, &owner, fencing_token, ttl_ms)
.await
}
}
@@ -0,0 +1,389 @@
use chrono::{DateTime, Duration, Utc};
use napi::Result;
use sqlx::{FromRow, PgPool, Postgres, Row, Transaction};
use y_octo::Doc;
use super::{
BackendRuntime, constants::DEFAULT_HISTORY_PERIOD_SECONDS, error::napi_error, types::RuntimeDocCompactionResult,
};
#[derive(FromRow)]
struct SnapshotRow {
blob: Vec<u8>,
updated_at: DateTime<Utc>,
updated_by: Option<String>,
}
#[derive(FromRow)]
struct UpdateRow {
blob: Vec<u8>,
created_at: DateTime<Utc>,
created_by: Option<String>,
}
struct DocCompactorStore {
pool: PgPool,
}
impl DocCompactorStore {
fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn compact_doc(
&self,
workspace_id: &str,
doc_id: &str,
batch_limit: i64,
history_min_interval_ms: i64,
) -> Result<(i64, bool)> {
compact_doc(
self.pool.clone(),
workspace_id,
doc_id,
batch_limit,
history_min_interval_ms,
)
.await
}
}
fn is_empty_doc(bin: &[u8]) -> bool {
bin.is_empty() || (bin.len() == 1 && bin[0] == 0) || (bin.len() == 2 && bin[0] == 0 && bin[1] == 0)
}
fn apply_updates(updates: impl IntoIterator<Item = Vec<u8>>) -> Result<Vec<u8>> {
let mut doc = Doc::default();
for update in updates {
doc
.apply_update_from_binary_v1(&update)
.map_err(|err| napi_error(format!("DocCompactor merge failed: {err}")))?;
}
doc
.encode_update_v1()
.map_err(|err| napi_error(format!("DocCompactor encode failed: {err}")))
}
async fn load_snapshot(
tx: &mut Transaction<'_, Postgres>,
workspace_id: &str,
doc_id: &str,
) -> Result<Option<SnapshotRow>> {
sqlx::query_as::<_, SnapshotRow>(
r#"
SELECT blob, updated_at, updated_by
FROM snapshots
WHERE workspace_id = $1 AND guid = $2
FOR UPDATE
"#,
)
.bind(workspace_id)
.bind(doc_id)
.fetch_optional(&mut **tx)
.await
.map_err(|err| napi_error(format!("DocCompactor load snapshot failed: {err}")))
}
async fn load_updates(
tx: &mut Transaction<'_, Postgres>,
workspace_id: &str,
doc_id: &str,
batch_limit: i64,
) -> Result<Vec<UpdateRow>> {
sqlx::query_as::<_, UpdateRow>(
r#"
SELECT blob, created_at, created_by
FROM updates
WHERE workspace_id = $1 AND guid = $2
ORDER BY created_at ASC
LIMIT $3
FOR UPDATE
"#,
)
.bind(workspace_id)
.bind(doc_id)
.bind(batch_limit)
.fetch_all(&mut **tx)
.await
.map_err(|err| napi_error(format!("DocCompactor load updates failed: {err}")))
}
async fn upsert_snapshot(
tx: &mut Transaction<'_, Postgres>,
workspace_id: &str,
doc_id: &str,
blob: &[u8],
timestamp: DateTime<Utc>,
editor: Option<&str>,
) -> Result<bool> {
if is_empty_doc(blob) {
return Ok(false);
}
let row = sqlx::query(
r#"
INSERT INTO snapshots
(workspace_id, guid, blob, size, created_at, updated_at, created_by, updated_by)
VALUES
($1, $2, $3, $4, $5, $5, $6, $6)
ON CONFLICT (workspace_id, guid)
DO UPDATE SET
blob = $3,
size = $4,
updated_at = $5,
updated_by = $6
WHERE snapshots.workspace_id = $1
AND snapshots.guid = $2
AND snapshots.updated_at <= $5
RETURNING updated_at
"#,
)
.bind(workspace_id)
.bind(doc_id)
.bind(blob)
.bind(blob.len() as i64)
.bind(timestamp)
.bind(editor)
.fetch_optional(&mut **tx)
.await
.map_err(|err| napi_error(format!("DocCompactor upsert snapshot failed: {err}")))?;
Ok(row.is_some())
}
async fn should_create_history(
tx: &mut Transaction<'_, Postgres>,
snapshot: &SnapshotRow,
workspace_id: &str,
doc_id: &str,
history_min_interval_ms: i64,
) -> Result<bool> {
if is_empty_doc(&snapshot.blob) {
return Ok(false);
}
let row = sqlx::query(
r#"
SELECT timestamp
FROM snapshot_histories
WHERE workspace_id = $1 AND guid = $2
ORDER BY timestamp DESC
LIMIT 1
"#,
)
.bind(workspace_id)
.bind(doc_id)
.fetch_optional(&mut **tx)
.await
.map_err(|err| napi_error(format!("DocCompactor load latest history failed: {err}")))?;
let Some(row) = row else {
return Ok(true);
};
let last_timestamp: DateTime<Utc> = row.get("timestamp");
if last_timestamp == snapshot.updated_at {
return Ok(false);
}
Ok(last_timestamp < snapshot.updated_at - Duration::milliseconds(history_min_interval_ms))
}
async fn history_max_age_seconds(tx: &mut Transaction<'_, Postgres>, workspace_id: &str) -> Result<i32> {
let row = sqlx::query(
r#"
SELECT history_period_seconds
FROM effective_workspace_quota_states
WHERE workspace_id = $1
"#,
)
.bind(workspace_id)
.fetch_optional(&mut **tx)
.await
.map_err(|err| napi_error(format!("DocCompactor load history quota failed: {err}")))?;
Ok(
row
.map(|row| row.get("history_period_seconds"))
.unwrap_or(DEFAULT_HISTORY_PERIOD_SECONDS),
)
}
async fn create_history(
tx: &mut Transaction<'_, Postgres>,
workspace_id: &str,
doc_id: &str,
snapshot: &SnapshotRow,
) -> Result<bool> {
let max_age_seconds = history_max_age_seconds(tx, workspace_id).await?;
if max_age_seconds <= 0 {
return Ok(false);
}
let expired_at = Utc::now() + Duration::seconds(max_age_seconds as i64);
sqlx::query(
r#"
INSERT INTO snapshot_histories
(workspace_id, guid, timestamp, blob, expired_at, created_by)
VALUES
($1, $2, $3, $4, $5, $6)
ON CONFLICT (workspace_id, guid, timestamp)
DO UPDATE SET expired_at = EXCLUDED.expired_at
"#,
)
.bind(workspace_id)
.bind(doc_id)
.bind(snapshot.updated_at)
.bind(&snapshot.blob)
.bind(expired_at)
.bind(snapshot.updated_by.as_deref())
.execute(&mut **tx)
.await
.map_err(|err| napi_error(format!("DocCompactor create history failed: {err}")))?;
Ok(true)
}
async fn delete_updates(
tx: &mut Transaction<'_, Postgres>,
workspace_id: &str,
doc_id: &str,
timestamps: &[DateTime<Utc>],
) -> Result<i64> {
let result = sqlx::query(
r#"
DELETE FROM updates
WHERE workspace_id = $1
AND guid = $2
AND created_at = ANY($3)
"#,
)
.bind(workspace_id)
.bind(doc_id)
.bind(timestamps)
.execute(&mut **tx)
.await
.map_err(|err| napi_error(format!("DocCompactor delete updates failed: {err}")))?;
Ok(result.rows_affected() as i64)
}
async fn compact_doc(
pool: PgPool,
workspace_id: &str,
doc_id: &str,
batch_limit: i64,
history_min_interval_ms: i64,
) -> Result<(i64, bool)> {
let mut tx = pool
.begin()
.await
.map_err(|err| napi_error(format!("DocCompactor begin transaction failed: {err}")))?;
let snapshot = load_snapshot(&mut tx, workspace_id, doc_id).await?;
let updates = load_updates(&mut tx, workspace_id, doc_id, batch_limit).await?;
if updates.is_empty() {
tx.commit()
.await
.map_err(|err| napi_error(format!("DocCompactor commit transaction failed: {err}")))?;
return Ok((0, false));
}
let last = updates.last().expect("updates is not empty");
let mut merge_inputs = Vec::with_capacity(updates.len() + usize::from(snapshot.is_some()));
if let Some(snapshot) = &snapshot {
merge_inputs.push(snapshot.blob.clone());
}
merge_inputs.extend(updates.iter().map(|update| update.blob.clone()));
let final_blob = if merge_inputs.len() == 1 {
merge_inputs.remove(0)
} else {
apply_updates(merge_inputs)?
};
let snapshot_updated = upsert_snapshot(
&mut tx,
workspace_id,
doc_id,
&final_blob,
last.created_at,
last.created_by.as_deref(),
)
.await?;
let mut history_created = false;
if snapshot_updated
&& let Some(snapshot) = &snapshot
&& should_create_history(&mut tx, snapshot, workspace_id, doc_id, history_min_interval_ms).await?
{
history_created = create_history(&mut tx, workspace_id, doc_id, snapshot).await?;
}
let timestamps = updates.iter().map(|update| update.created_at).collect::<Vec<_>>();
let deleted = delete_updates(&mut tx, workspace_id, doc_id, &timestamps).await?;
tx.commit()
.await
.map_err(|err| napi_error(format!("DocCompactor commit transaction failed: {err}")))?;
Ok((deleted, history_created))
}
#[napi_derive::napi]
impl BackendRuntime {
/// Merge pending doc updates with y-octo and persist the merged snapshot.
///
/// Do not use this for snapshots that will be sent back to yjs clients until
/// the y-octo/yjs round-trip compatibility issue is resolved.
#[napi]
pub async fn compact_pending_doc_updates(
&self,
workspace_id: String,
doc_id: String,
batch_limit: i64,
history_min_interval_ms: i64,
owner: String,
lease_ttl_ms: i64,
) -> Result<RuntimeDocCompactionResult> {
if batch_limit <= 0 {
return Err(napi_error("doc compactor batch limit must be positive"));
}
if history_min_interval_ms < 0 {
return Err(napi_error("doc compactor history interval must be non-negative"));
}
let lease_key = format!("doc:update:{workspace_id}:{doc_id}");
let Some(lease) = self.acquire_coordination_lease(lease_key, owner, lease_ttl_ms).await? else {
return Ok(RuntimeDocCompactionResult {
lease_acquired: false,
merged: false,
workspace_id,
doc_id,
updates_merged: 0,
history_created: false,
});
};
let result = DocCompactorStore::new(self.pool().await?)
.compact_doc(&workspace_id, &doc_id, batch_limit, history_min_interval_ms)
.await;
let released = self
.release_coordination_lease(lease.key, lease.owner, lease.fencing_token)
.await?;
if !released {
return Err(napi_error("DocCompactor failed to release coordination lease"));
}
let (updates_merged, history_created) = result?;
Ok(RuntimeDocCompactionResult {
lease_acquired: true,
merged: updates_merged > 0,
workspace_id,
doc_id,
updates_merged,
history_created,
})
}
}
@@ -0,0 +1,158 @@
use chrono::{DateTime, Duration, Utc};
use napi::{Result, bindgen_prelude::Buffer};
use sqlx::{PgPool, Row};
use super::{BackendRuntime, error::napi_error, types::RuntimeDocHistoryInput};
fn is_empty_doc(bin: &[u8]) -> bool {
bin.is_empty() || (bin.len() == 1 && bin[0] == 0) || (bin.len() == 2 && bin[0] == 0 && bin[1] == 0)
}
async fn latest_history_timestamp(pool: &PgPool, workspace_id: &str, doc_id: &str) -> Result<Option<DateTime<Utc>>> {
sqlx::query(
r#"
SELECT timestamp
FROM snapshot_histories
WHERE workspace_id = $1 AND guid = $2
ORDER BY timestamp DESC
LIMIT 1
"#,
)
.bind(workspace_id)
.bind(doc_id)
.fetch_optional(pool)
.await
.map(|row| row.map(|row| row.get("timestamp")))
.map_err(|err| napi_error(format!("DocStorage load latest history failed: {err}")))
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn upsert_doc_snapshot(
&self,
workspace_id: String,
doc_id: String,
blob: Buffer,
timestamp_ms: i64,
editor_id: Option<String>,
) -> Result<bool> {
if is_empty_doc(blob.as_ref()) {
return Ok(false);
}
let timestamp = DateTime::<Utc>::from_timestamp_millis(timestamp_ms)
.ok_or_else(|| napi_error(format!("Invalid doc snapshot timestamp: {timestamp_ms}")))?;
let pool = self.pool().await?;
let row = sqlx::query(
r#"
INSERT INTO snapshots
(workspace_id, guid, blob, size, created_at, updated_at, created_by, updated_by)
VALUES
($1, $2, $3, $4, $5, $5, $6, $6)
ON CONFLICT (workspace_id, guid)
DO UPDATE SET
blob = $3,
size = $4,
updated_at = $5,
updated_by = $6
WHERE snapshots.workspace_id = $1
AND snapshots.guid = $2
AND snapshots.updated_at <= $5
RETURNING updated_at
"#,
)
.bind(&workspace_id)
.bind(&doc_id)
.bind(blob.as_ref())
.bind(blob.len() as i64)
.bind(timestamp)
.bind(editor_id.as_deref())
.fetch_optional(&pool)
.await
.map_err(|err| napi_error(format!("DocStorage upsert snapshot failed: {err}")))?;
Ok(row.is_some())
}
#[napi]
pub async fn create_doc_history(&self, input: RuntimeDocHistoryInput) -> Result<bool> {
if input.history_min_interval_ms < 0 {
return Err(napi_error("doc history interval must be non-negative"));
}
if input.history_max_age_ms <= 0 || is_empty_doc(input.blob.as_ref()) {
return Ok(false);
}
let timestamp = DateTime::<Utc>::from_timestamp_millis(input.timestamp_ms)
.ok_or_else(|| napi_error(format!("Invalid doc history timestamp: {}", input.timestamp_ms)))?;
let pool = self.pool().await?;
let should_create = match latest_history_timestamp(&pool, &input.workspace_id, &input.doc_id).await? {
None => true,
Some(last_timestamp) if last_timestamp == timestamp => false,
Some(last_timestamp) => {
input.force || last_timestamp < timestamp - Duration::milliseconds(input.history_min_interval_ms)
}
};
if !should_create {
return Ok(false);
}
let expired_at = Utc::now() + Duration::milliseconds(input.history_max_age_ms);
sqlx::query(
r#"
INSERT INTO snapshot_histories
(workspace_id, guid, timestamp, blob, expired_at, created_by)
VALUES
($1, $2, $3, $4, $5, $6)
ON CONFLICT (workspace_id, guid, timestamp)
DO UPDATE SET expired_at = EXCLUDED.expired_at
"#,
)
.bind(&input.workspace_id)
.bind(&input.doc_id)
.bind(timestamp)
.bind(input.blob.as_ref())
.bind(expired_at)
.bind(input.editor_id.as_deref())
.execute(&pool)
.await
.map_err(|err| napi_error(format!("DocStorage create history failed: {err}")))?;
Ok(true)
}
#[napi]
pub async fn delete_doc_storage(&self, workspace_id: String, doc_id: String) -> Result<()> {
let pool = self.pool().await?;
let mut tx = pool
.begin()
.await
.map_err(|err| napi_error(format!("DocStorage delete begin transaction failed: {err}")))?;
sqlx::query("DELETE FROM snapshots WHERE workspace_id = $1 AND guid = $2")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&mut *tx)
.await
.map_err(|err| napi_error(format!("DocStorage delete snapshot failed: {err}")))?;
sqlx::query("DELETE FROM updates WHERE workspace_id = $1 AND guid = $2")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&mut *tx)
.await
.map_err(|err| napi_error(format!("DocStorage delete updates failed: {err}")))?;
sqlx::query("DELETE FROM snapshot_histories WHERE workspace_id = $1 AND guid = $2")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&mut *tx)
.await
.map_err(|err| napi_error(format!("DocStorage delete histories failed: {err}")))?;
tx.commit()
.await
.map_err(|err| napi_error(format!("DocStorage delete commit failed: {err}")))?;
Ok(())
}
}
@@ -0,0 +1,5 @@
use napi::{Error, Status};
pub(super) fn napi_error(message: impl Into<String>) -> Error {
Error::new(Status::GenericFailure, message.into())
}
@@ -0,0 +1,90 @@
use napi::Result;
use sqlx::PgPool;
use super::{BackendRuntime, error::napi_error};
struct RuntimeGateStore {
pool: PgPool,
}
impl RuntimeGateStore {
fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn put_if_absent(&self, key: &str, ttl_ms: i64) -> Result<bool> {
let mut tx = self
.pool
.begin()
.await
.map_err(|err| napi_error(format!("RuntimeGate transaction failed: {err}")))?;
sqlx::query("DELETE FROM runtime_gates WHERE key = $1 AND expires_at <= CURRENT_TIMESTAMP")
.bind(key)
.execute(&mut *tx)
.await
.map_err(|err| napi_error(format!("RuntimeGate expired cleanup failed: {err}")))?;
let inserted = sqlx::query(
r#"
INSERT INTO runtime_gates (key, expires_at)
VALUES ($1, CURRENT_TIMESTAMP + ($2 * INTERVAL '1 millisecond'))
ON CONFLICT (key) DO NOTHING
"#,
)
.bind(key)
.bind(ttl_ms as f64)
.execute(&mut *tx)
.await
.map_err(|err| napi_error(format!("RuntimeGate put_if_absent failed: {err}")))?
.rows_affected()
== 1;
tx.commit()
.await
.map_err(|err| napi_error(format!("RuntimeGate transaction commit failed: {err}")))?;
Ok(inserted)
}
async fn cleanup_expired(&self, limit: i64) -> Result<i64> {
let result = sqlx::query(
r#"
DELETE FROM runtime_gates
WHERE key IN (
SELECT key FROM runtime_gates
WHERE expires_at <= CURRENT_TIMESTAMP
ORDER BY expires_at ASC
LIMIT $1
)
"#,
)
.bind(limit)
.execute(&self.pool)
.await
.map_err(|err| napi_error(format!("RuntimeGate cleanup failed: {err}")))?;
Ok(result.rows_affected() as i64)
}
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn put_runtime_gate_if_absent(&self, key: String, ttl_ms: i64) -> Result<bool> {
if ttl_ms <= 0 {
return Err(napi_error("runtime gate ttl must be positive"));
}
RuntimeGateStore::new(self.pool().await?)
.put_if_absent(&key, ttl_ms)
.await
}
#[napi]
pub async fn cleanup_expired_runtime_gates(&self, limit: i64) -> Result<i64> {
if limit <= 0 {
return Err(napi_error("runtime gate cleanup limit must be positive"));
}
RuntimeGateStore::new(self.pool().await?).cleanup_expired(limit).await
}
}
@@ -0,0 +1,80 @@
use napi::Result;
use sqlx::PgPool;
use super::{BackendRuntime, error::napi_error};
struct HousekeepingStore {
pool: PgPool,
}
impl HousekeepingStore {
fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn cleanup_expired_user_sessions(&self, limit: i64) -> Result<i64> {
let result = sqlx::query(
r#"
DELETE FROM user_sessions
WHERE id IN (
SELECT id FROM user_sessions
WHERE expires_at <= CURRENT_TIMESTAMP
ORDER BY expires_at ASC
LIMIT $1
)
"#,
)
.bind(limit)
.execute(&self.pool)
.await
.map_err(|err| napi_error(format!("Housekeeping user sessions cleanup failed: {err}")))?;
Ok(result.rows_affected() as i64)
}
async fn cleanup_expired_snapshot_histories(&self, limit: i64) -> Result<i64> {
let result = sqlx::query(
r#"
DELETE FROM snapshot_histories
WHERE (workspace_id, guid, timestamp) IN (
SELECT workspace_id, guid, timestamp
FROM snapshot_histories
WHERE expired_at <= CURRENT_TIMESTAMP
ORDER BY expired_at ASC
LIMIT $1
)
"#,
)
.bind(limit)
.execute(&self.pool)
.await
.map_err(|err| napi_error(format!("Housekeeping snapshot histories cleanup failed: {err}")))?;
Ok(result.rows_affected() as i64)
}
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn cleanup_expired_user_sessions(&self, limit: i64) -> Result<i64> {
if limit <= 0 {
return Err(napi_error("user sessions cleanup limit must be positive"));
}
HousekeepingStore::new(self.pool().await?)
.cleanup_expired_user_sessions(limit)
.await
}
#[napi]
pub async fn cleanup_expired_snapshot_histories(&self, limit: i64) -> Result<i64> {
if limit <= 0 {
return Err(napi_error("snapshot histories cleanup limit must be positive"));
}
HousekeepingStore::new(self.pool().await?)
.cleanup_expired_snapshot_histories(limit)
.await
}
}
@@ -0,0 +1,128 @@
mod blob_complete;
mod blob_reclaimer;
mod config;
mod constants;
mod coordination_lease;
mod doc_compactor;
mod doc_storage;
mod error;
mod gate;
mod housekeeping;
mod object_storage;
mod runtime_state;
#[cfg(test)]
mod tests;
mod types;
mod workspace_stats;
use std::time::Duration;
use napi::Result;
use sha2::{Digest, Sha256};
use sqlx::{PgPool, Row, postgres::PgPoolOptions};
use tokio::sync::Mutex;
use self::{config::RuntimeConfig, constants::RUNTIME_MIGRATIONS, error::napi_error, types::BackendRuntimeHealth};
pub(super) fn token_hash(token: &str) -> String {
hex::encode(Sha256::digest(token.as_bytes()))
}
#[napi_derive::napi]
pub struct BackendRuntime {
config: RuntimeConfig,
pool: Mutex<Option<PgPool>>,
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi(constructor)]
pub fn new() -> Result<Self> {
Ok(Self {
config: RuntimeConfig::from_config_files()?,
pool: Mutex::new(None),
})
}
#[napi]
pub async fn start(&self) -> Result<()> {
let mut guard = self.pool.lock().await;
if guard.is_some() {
return Ok(());
}
let pool = PgPoolOptions::new()
.max_connections(5)
.acquire_timeout(Duration::from_secs(5))
.connect(&self.config.database_url)
.await
.map_err(|err| napi_error(format!("BackendRuntime failed to connect postgres: {err}")))?;
sqlx::query("SELECT 1")
.execute(&pool)
.await
.map_err(|err| napi_error(format!("BackendRuntime postgres health check failed: {err}")))?;
*guard = Some(pool);
Ok(())
}
#[napi]
pub async fn stop(&self) -> Result<()> {
let pool = self.pool.lock().await.take();
if let Some(pool) = pool {
pool.close().await;
}
Ok(())
}
#[napi]
pub async fn health(&self) -> Result<BackendRuntimeHealth> {
let pool = self.pool.lock().await.as_ref().cloned();
let database_connected = match pool.as_ref() {
Some(pool) => sqlx::query("SELECT 1")
.fetch_one(pool)
.await
.map(|row| row.try_get::<i32, _>(0).unwrap_or(0) == 1)
.unwrap_or(false),
None => false,
};
Ok(BackendRuntimeHealth {
started: pool.is_some(),
database_connected,
object_storage_configured: self.config.storage.is_some(),
})
}
#[napi]
pub async fn run_migrations(&self) -> Result<()> {
let pool = self.pool().await?;
migrate_runtime_tables(&pool).await
}
async fn pool(&self) -> Result<PgPool> {
self
.pool
.lock()
.await
.as_ref()
.cloned()
.ok_or_else(|| napi_error("BackendRuntime must be started before using postgres operations"))
}
}
async fn migrate_runtime_tables(pool: &PgPool) -> Result<()> {
for statement in RUNTIME_MIGRATIONS
.split(';')
.map(str::trim)
.filter(|statement| !statement.is_empty())
{
sqlx::query(statement)
.execute(pool)
.await
.map_err(|err| napi_error(format!("BackendRuntime migration failed: {err}")))?;
}
Ok(())
}
@@ -0,0 +1,353 @@
use std::{
collections::HashMap,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use aws_sdk_s3::{
Client as S3Client, presigning::PresigningConfig, primitives::ByteStream, types::CompletedMultipartUpload,
};
use napi::Result;
use super::types::{
MultipartUploadInitResult, MultipartUploadPart, ObjectGetResult, ObjectListEntry, ObjectMetadata, ObjectPutMetadata,
PresignedObjectRequest, completed_multipart_parts, trim_etag,
};
use crate::backend_runtime::error::napi_error;
#[derive(Clone)]
pub(super) struct ObjectStorageClient {
client: S3Client,
bucket: String,
presign_expires_in_seconds: u64,
presign_sign_content_type_for_put: bool,
}
impl ObjectStorageClient {
pub(super) fn new(
config: aws_sdk_s3::Config,
bucket: String,
presign_expires_in_seconds: u64,
presign_sign_content_type_for_put: bool,
) -> Self {
Self {
client: S3Client::from_conf(config),
bucket,
presign_expires_in_seconds,
presign_sign_content_type_for_put,
}
}
pub(super) fn non_destructive_health(&self) -> bool {
let _ = &self.client;
!self.bucket.is_empty()
}
pub(super) async fn put(&self, key: &str, body: Vec<u8>, metadata: ObjectPutMetadata) -> Result<()> {
let content_length = metadata.content_length.unwrap_or(body.len() as i64);
let content_type = metadata
.content_type
.unwrap_or_else(|| "application/octet-stream".to_string());
let mut request = self
.client
.put_object()
.bucket(&self.bucket)
.key(key)
.body(ByteStream::from(body))
.content_type(content_type)
.content_length(content_length);
if let Some(checksum) = metadata.checksum_crc32 {
request = request.checksum_crc32(checksum);
}
request
.send()
.await
.map_err(|err| napi_error(format!("ObjectStorage put failed for {key}: {err:?}")))?;
Ok(())
}
pub(super) async fn presign_put(&self, key: &str, metadata: ObjectPutMetadata) -> Result<PresignedObjectRequest> {
let content_type = metadata
.content_type
.unwrap_or_else(|| "application/octet-stream".to_string());
let expires_at_ms = expires_at_ms(self.presign_expires_in_seconds)?;
let config = PresigningConfig::expires_in(Duration::from_secs(self.presign_expires_in_seconds))
.map_err(|err| napi_error(format!("ObjectStorage presign config failed: {err}")))?;
let mut request = self.client.put_object().bucket(&self.bucket).key(key);
if self.presign_sign_content_type_for_put {
request = request.content_type(content_type.clone());
}
if let Some(content_length) = metadata.content_length {
request = request.content_length(content_length);
}
let presigned = request
.presigned(config)
.await
.map_err(|err| napi_error(format!("ObjectStorage presign put failed for {key}: {err}")))?;
let mut headers = presigned_headers(&presigned);
headers.insert("Content-Type".to_string(), content_type);
Ok(PresignedObjectRequest {
url: presigned.uri().to_string(),
headers,
expires_at_ms,
})
}
pub(super) async fn create_multipart_upload(
&self,
key: &str,
metadata: ObjectPutMetadata,
) -> Result<Option<MultipartUploadInitResult>> {
let content_type = metadata
.content_type
.unwrap_or_else(|| "application/octet-stream".to_string());
let result = self
.client
.create_multipart_upload()
.bucket(&self.bucket)
.key(key)
.content_type(content_type)
.send()
.await
.map_err(|err| {
napi_error(format!(
"ObjectStorage create multipart upload failed for {key}: {err:?}"
))
})?;
let expires_at_ms = expires_at_ms(self.presign_expires_in_seconds)?;
Ok(result.upload_id.map(|upload_id| MultipartUploadInitResult {
upload_id,
expires_at_ms,
}))
}
pub(super) async fn presign_upload_part(
&self,
key: &str,
upload_id: &str,
part_number: i32,
) -> Result<PresignedObjectRequest> {
let expires_at_ms = expires_at_ms(self.presign_expires_in_seconds)?;
let config = PresigningConfig::expires_in(Duration::from_secs(self.presign_expires_in_seconds))
.map_err(|err| napi_error(format!("ObjectStorage presign config failed: {err}")))?;
let presigned = self
.client
.upload_part()
.bucket(&self.bucket)
.key(key)
.upload_id(upload_id)
.part_number(part_number)
.presigned(config)
.await
.map_err(|err| napi_error(format!("ObjectStorage presign upload part failed for {key}: {err}")))?;
Ok(PresignedObjectRequest {
url: presigned.uri().to_string(),
headers: presigned_headers(&presigned),
expires_at_ms,
})
}
pub(super) async fn list_multipart_upload_parts(
&self,
key: &str,
upload_id: &str,
) -> Result<Vec<MultipartUploadPart>> {
let result = self
.client
.list_parts()
.bucket(&self.bucket)
.key(key)
.upload_id(upload_id)
.send()
.await
.map_err(|err| {
napi_error(format!(
"ObjectStorage list multipart upload parts failed for {key}: {err}"
))
})?;
Ok(
result
.parts()
.iter()
.filter_map(|part| {
Some(MultipartUploadPart {
part_number: part.part_number?,
etag: trim_etag(part.e_tag.as_deref().unwrap_or_default()),
})
})
.collect(),
)
}
pub(super) async fn complete_multipart_upload(
&self,
key: &str,
upload_id: &str,
parts: Vec<MultipartUploadPart>,
) -> Result<()> {
let ordered_parts = completed_multipart_parts(parts);
self
.client
.complete_multipart_upload()
.bucket(&self.bucket)
.key(key)
.upload_id(upload_id)
.multipart_upload(
CompletedMultipartUpload::builder()
.set_parts(Some(ordered_parts))
.build(),
)
.send()
.await
.map_err(|err| {
napi_error(format!(
"ObjectStorage complete multipart upload failed for {key}: {err}"
))
})?;
Ok(())
}
pub(super) async fn abort_multipart_upload(&self, key: &str, upload_id: &str) -> Result<()> {
self
.client
.abort_multipart_upload()
.bucket(&self.bucket)
.key(key)
.upload_id(upload_id)
.send()
.await
.map_err(|err| {
napi_error(format!(
"ObjectStorage abort multipart upload failed for {key}: {err:?}"
))
})?;
Ok(())
}
pub(super) async fn head(&self, key: &str) -> Result<Option<ObjectMetadata>> {
let result = self
.client
.head_object()
.bucket(&self.bucket)
.key(key)
.send()
.await
.map_err(|err| napi_error(format!("ObjectStorage head failed for {key}: {err:?}")))?;
Ok(Some(ObjectMetadata {
content_type: result
.content_type
.unwrap_or_else(|| "application/octet-stream".to_string()),
content_length: result.content_length.unwrap_or(0),
last_modified_ms: optional_datetime_ms(result.last_modified),
checksum_crc32: result.checksum_crc32,
}))
}
pub(super) async fn get(&self, key: &str) -> Result<Option<ObjectGetResult>> {
let result = self
.client
.get_object()
.bucket(&self.bucket)
.key(key)
.send()
.await
.map_err(|err| napi_error(format!("ObjectStorage get failed for {key}: {err:?}")))?;
let metadata = ObjectMetadata {
content_type: result
.content_type
.unwrap_or_else(|| "application/octet-stream".to_string()),
content_length: result.content_length.unwrap_or(0),
last_modified_ms: optional_datetime_ms(result.last_modified),
checksum_crc32: result.checksum_crc32,
};
let body = result
.body
.collect()
.await
.map_err(|err| napi_error(format!("ObjectStorage read body failed for {key}: {err}")))?
.into_bytes()
.to_vec();
Ok(Some(ObjectGetResult { body, metadata }))
}
pub(super) async fn list(&self, prefix: Option<String>) -> Result<Vec<ObjectListEntry>> {
let mut entries = Vec::new();
let mut token = None;
loop {
let mut request = self.client.list_objects_v2().bucket(&self.bucket);
if let Some(prefix) = &prefix {
request = request.prefix(prefix);
}
if let Some(next_token) = token {
request = request.continuation_token(next_token);
}
let result = request
.send()
.await
.map_err(|err| napi_error(format!("ObjectStorage list failed: {err:?}")))?;
entries.extend(result.contents().iter().filter_map(|object| {
Some(ObjectListEntry {
key: object.key.as_ref()?.clone(),
content_length: object.size.unwrap_or(0),
last_modified_ms: optional_datetime_ms(object.last_modified),
})
}));
if result.is_truncated.unwrap_or(false) {
token = result.next_continuation_token;
} else {
break;
}
}
Ok(entries)
}
pub(super) async fn delete(&self, key: &str) -> Result<()> {
self
.client
.delete_object()
.bucket(&self.bucket)
.key(key)
.send()
.await
.map_err(|err| napi_error(format!("ObjectStorage delete failed for {key}: {err:?}")))?;
Ok(())
}
}
fn expires_at_ms(expires_in_seconds: u64) -> Result<i64> {
let expires_at = SystemTime::now()
.checked_add(Duration::from_secs(expires_in_seconds))
.ok_or_else(|| napi_error("ObjectStorage presign expiration overflow"))?;
system_time_ms(expires_at)
}
fn system_time_ms(time: SystemTime) -> Result<i64> {
let duration = time
.duration_since(UNIX_EPOCH)
.map_err(|err| napi_error(format!("system time before unix epoch: {err}")))?;
Ok(duration.as_millis() as i64)
}
fn optional_datetime_ms(time: Option<aws_sdk_s3::primitives::DateTime>) -> i64 {
time.and_then(|value| value.to_millis().ok()).unwrap_or(0)
}
fn presigned_headers(request: &aws_sdk_s3::presigning::PresignedRequest) -> HashMap<String, String> {
request
.headers()
.map(|(key, value)| (key.to_string(), value.to_string()))
.collect()
}
@@ -0,0 +1,225 @@
use aws_sdk_s3::config::{
BehaviorVersion, Credentials, Region, RequestChecksumCalculation, ResponseChecksumValidation, timeout::TimeoutConfig,
};
use napi::Result;
use serde::Deserialize;
use super::{client::ObjectStorageClient, types::StorageProviderConfig};
use crate::backend_runtime::{
config::blob_storage_config_from_config_files, error::napi_error, types::RuntimeObjectStorageHealth,
};
#[derive(Clone, Debug)]
pub(in crate::backend_runtime) struct ObjectStorageConfig {
pub(super) provider: String,
pub(super) bucket: String,
pub(super) endpoint: Option<String>,
pub(super) region: Option<String>,
pub(super) access_key_id: Option<String>,
pub(super) secret_access_key: Option<String>,
pub(super) session_token: Option<String>,
pub(super) force_path_style: bool,
pub(super) request_timeout_ms: Option<u64>,
pub(super) min_part_size: Option<u64>,
pub(super) presign_expires_in_seconds: Option<u64>,
pub(super) presign_sign_content_type_for_put: Option<bool>,
pub(super) use_presigned_url: bool,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct S3ConfigFile {
endpoint: Option<String>,
region: Option<String>,
credentials: Option<S3CredentialsConfigFile>,
force_path_style: Option<bool>,
request_timeout_ms: Option<u64>,
min_part_size: Option<u64>,
presign: Option<S3PresignConfigFile>,
#[serde(rename = "usePresignedURL")]
use_presigned_url: Option<UsePresignedUrlConfigFile>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct R2ConfigFile {
account_id: String,
jurisdiction: Option<String>,
region: Option<String>,
credentials: Option<S3CredentialsConfigFile>,
request_timeout_ms: Option<u64>,
min_part_size: Option<u64>,
presign: Option<S3PresignConfigFile>,
#[serde(rename = "usePresignedURL")]
use_presigned_url: Option<UsePresignedUrlConfigFile>,
}
#[derive(Debug, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
struct S3CredentialsConfigFile {
access_key_id: Option<String>,
secret_access_key: Option<String>,
session_token: Option<String>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct S3PresignConfigFile {
expires_in_seconds: Option<u64>,
sign_content_type_for_put: Option<bool>,
}
#[derive(Debug, Deserialize)]
struct UsePresignedUrlConfigFile {
enabled: bool,
}
impl ObjectStorageConfig {
pub(in crate::backend_runtime) fn from_config_files() -> Result<Option<Self>> {
let Some(storage) = blob_storage_config_from_config_files()? else {
return Ok(None);
};
match storage.provider.as_str() {
"aws-s3" => Self::from_s3_config(storage),
"cloudflare-r2" => Self::from_r2_config(storage),
"fs" => Ok(None),
provider => Err(napi_error(format!(
"unsupported blob storage provider for BackendRuntime: {provider}"
))),
}
}
pub(super) fn from_s3_config(storage: StorageProviderConfig) -> Result<Option<Self>> {
let config: S3ConfigFile = serde_json::from_value(storage.config)
.map_err(|err| napi_error(format!("invalid aws-s3 blob storage config: {err}")))?;
let region = config
.region
.ok_or_else(|| napi_error("aws-s3 blob storage config requires region"))?;
let endpoint = config.endpoint.or_else(|| Some(resolve_s3_endpoint(&region)));
let credentials = config.credentials.unwrap_or_default();
Ok(Some(Self {
provider: storage.provider,
bucket: storage.bucket,
endpoint,
region: Some(region),
access_key_id: credentials.access_key_id,
secret_access_key: credentials.secret_access_key,
session_token: credentials.session_token,
force_path_style: config.force_path_style.unwrap_or(false),
request_timeout_ms: config.request_timeout_ms,
min_part_size: config.min_part_size,
presign_expires_in_seconds: config.presign.as_ref().and_then(|v| v.expires_in_seconds),
presign_sign_content_type_for_put: config.presign.as_ref().and_then(|v| v.sign_content_type_for_put),
use_presigned_url: config.use_presigned_url.map(|v| v.enabled).unwrap_or(false),
}))
}
pub(super) fn from_r2_config(storage: StorageProviderConfig) -> Result<Option<Self>> {
let config: R2ConfigFile = serde_json::from_value(storage.config)
.map_err(|err| napi_error(format!("invalid cloudflare-r2 blob storage config: {err}")))?;
let account = match config.jurisdiction {
Some(jurisdiction) => format!("{}.{}", config.account_id, jurisdiction),
None => config.account_id,
};
let credentials = config.credentials.unwrap_or_default();
Ok(Some(Self {
provider: storage.provider,
bucket: storage.bucket,
endpoint: Some(format!("https://{account}.r2.cloudflarestorage.com")),
region: Some(config.region.unwrap_or_else(|| "auto".to_string())),
access_key_id: credentials.access_key_id,
secret_access_key: credentials.secret_access_key,
session_token: credentials.session_token,
force_path_style: true,
request_timeout_ms: config.request_timeout_ms,
min_part_size: config.min_part_size,
presign_expires_in_seconds: config.presign.as_ref().and_then(|v| v.expires_in_seconds),
presign_sign_content_type_for_put: config.presign.as_ref().and_then(|v| v.sign_content_type_for_put),
use_presigned_url: config.use_presigned_url.map(|v| v.enabled).unwrap_or(false),
}))
}
pub(super) fn build_client(&self) -> Result<ObjectStorageClient> {
let region = self
.region
.clone()
.ok_or_else(|| napi_error("object storage region is required"))?;
let access_key_id = self
.access_key_id
.clone()
.ok_or_else(|| napi_error("object storage accessKeyId is required"))?;
let secret_access_key = self
.secret_access_key
.clone()
.ok_or_else(|| napi_error("object storage secretAccessKey is required"))?;
let credentials = Credentials::new(
access_key_id,
secret_access_key,
self.session_token.clone(),
None,
"affine-server-config-json",
);
let mut builder = aws_sdk_s3::Config::builder()
.behavior_version(BehaviorVersion::latest())
.region(Region::new(region))
.credentials_provider(credentials)
.force_path_style(self.force_path_style)
.request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
.response_checksum_validation(ResponseChecksumValidation::WhenRequired);
if let Some(endpoint) = &self.endpoint {
builder = builder.endpoint_url(endpoint);
}
if let Some(request_timeout_ms) = self.request_timeout_ms {
builder = builder.timeout_config(
TimeoutConfig::builder()
.operation_timeout(std::time::Duration::from_millis(request_timeout_ms))
.build(),
);
}
Ok(ObjectStorageClient::new(
builder.build(),
self.bucket.clone(),
self.presign_expires_in_seconds.unwrap_or(60),
self.presign_sign_content_type_for_put.unwrap_or(true),
))
}
pub(super) fn health(&self) -> RuntimeObjectStorageHealth {
let client_buildable = self
.build_client()
.map(|client| client.non_destructive_health())
.unwrap_or(false);
RuntimeObjectStorageHealth {
configured: true,
provider: Some(self.provider.clone()),
bucket: Some(self.bucket.clone()),
endpoint: self.endpoint.clone(),
region: self.region.clone(),
has_credentials: self.access_key_id.is_some()
&& self.secret_access_key.is_some()
&& self.session_token.as_ref().map(|v| !v.is_empty()).unwrap_or(true),
force_path_style: self.force_path_style,
request_timeout_ms: self.request_timeout_ms.map(|v| v as i64),
min_part_size: self.min_part_size.map(|v| v as i64),
presign_expires_in_seconds: self.presign_expires_in_seconds.map(|v| v as i64),
presign_sign_content_type_for_put: self.presign_sign_content_type_for_put,
use_presigned_url: self.use_presigned_url,
client_buildable,
}
}
}
fn resolve_s3_endpoint(region: &str) -> String {
if region == "us-east-1" {
"https://s3.amazonaws.com".to_string()
} else {
format!("https://s3.{region}.amazonaws.com")
}
}
@@ -0,0 +1,184 @@
mod client;
mod config;
#[cfg(test)]
mod tests;
mod types;
use client::ObjectStorageClient;
pub(super) use config::ObjectStorageConfig;
use napi::{Result, bindgen_prelude::Buffer};
pub(super) use types::StorageProviderConfig;
use super::{
BackendRuntime,
types::{
RuntimeMultipartUploadInit, RuntimeMultipartUploadPart, RuntimeObjectGetResult, RuntimeObjectListEntry,
RuntimeObjectMetadata, RuntimeObjectStorageHealth, RuntimeObjectStoragePutOptions, RuntimePresignedObjectRequest,
},
};
#[napi_derive::napi]
impl BackendRuntime {
fn object_storage_client(&self) -> Result<ObjectStorageClient> {
self
.config
.storage
.as_ref()
.ok_or_else(|| super::error::napi_error("ObjectStorageClient is not configured"))?
.build_client()
}
pub(super) async fn object_storage_delete_object(&self, key: &str) -> Result<()> {
self.object_storage_client()?.delete(key).await
}
pub(super) async fn object_storage_abort_upload(&self, key: &str, upload_id: &str) -> Result<()> {
self
.object_storage_client()?
.abort_multipart_upload(key, upload_id)
.await
}
#[napi]
pub fn object_storage_health(&self) -> RuntimeObjectStorageHealth {
match &self.config.storage {
Some(storage) => storage.health(),
None => RuntimeObjectStorageHealth {
configured: false,
provider: None,
bucket: None,
endpoint: None,
region: None,
has_credentials: false,
force_path_style: false,
request_timeout_ms: None,
min_part_size: None,
presign_expires_in_seconds: None,
presign_sign_content_type_for_put: None,
use_presigned_url: false,
client_buildable: false,
},
}
}
#[napi]
pub async fn object_storage_put(
&self,
key: String,
body: Buffer,
metadata: Option<RuntimeObjectStoragePutOptions>,
) -> Result<()> {
self
.object_storage_client()?
.put(&key, body.to_vec(), metadata.map(Into::into).unwrap_or_default())
.await
}
#[napi]
pub async fn object_storage_presign_put(
&self,
key: String,
metadata: Option<RuntimeObjectStoragePutOptions>,
) -> Result<RuntimePresignedObjectRequest> {
self
.object_storage_client()?
.presign_put(&key, metadata.map(Into::into).unwrap_or_default())
.await?
.try_into()
}
#[napi]
pub async fn object_storage_create_multipart_upload(
&self,
key: String,
metadata: Option<RuntimeObjectStoragePutOptions>,
) -> Result<Option<RuntimeMultipartUploadInit>> {
Ok(
self
.object_storage_client()?
.create_multipart_upload(&key, metadata.map(Into::into).unwrap_or_default())
.await?
.map(Into::into),
)
}
#[napi]
pub async fn object_storage_presign_upload_part(
&self,
key: String,
upload_id: String,
part_number: i32,
) -> Result<RuntimePresignedObjectRequest> {
self
.object_storage_client()?
.presign_upload_part(&key, &upload_id, part_number)
.await?
.try_into()
}
#[napi]
pub async fn object_storage_list_multipart_upload_parts(
&self,
key: String,
upload_id: String,
) -> Result<Vec<RuntimeMultipartUploadPart>> {
Ok(
self
.object_storage_client()?
.list_multipart_upload_parts(&key, &upload_id)
.await?
.into_iter()
.map(Into::into)
.collect(),
)
}
#[napi]
pub async fn object_storage_complete_multipart_upload(
&self,
key: String,
upload_id: String,
parts: Vec<RuntimeMultipartUploadPart>,
) -> Result<()> {
self
.object_storage_client()?
.complete_multipart_upload(&key, &upload_id, parts.into_iter().map(Into::into).collect())
.await
}
#[napi]
pub async fn object_storage_abort_multipart_upload(&self, key: String, upload_id: String) -> Result<()> {
self
.object_storage_client()?
.abort_multipart_upload(&key, &upload_id)
.await
}
#[napi]
pub async fn object_storage_head(&self, key: String) -> Result<Option<RuntimeObjectMetadata>> {
Ok(self.object_storage_client()?.head(&key).await?.map(Into::into))
}
#[napi]
pub async fn object_storage_get(&self, key: String) -> Result<Option<RuntimeObjectGetResult>> {
Ok(self.object_storage_client()?.get(&key).await?.map(Into::into))
}
#[napi]
pub async fn object_storage_list(&self, prefix: Option<String>) -> Result<Vec<RuntimeObjectListEntry>> {
Ok(
self
.object_storage_client()?
.list(prefix)
.await?
.into_iter()
.map(Into::into)
.collect(),
)
}
#[napi]
pub async fn object_storage_delete(&self, key: String) -> Result<()> {
self.object_storage_client()?.delete(&key).await
}
}
@@ -0,0 +1,129 @@
use super::{
config::ObjectStorageConfig,
types::{MultipartUploadPart, ObjectPutMetadata, StorageProviderConfig, completed_multipart_parts, trim_etag},
};
#[test]
fn resolves_r2_config_from_config_json_shape() {
let storage = StorageProviderConfig {
provider: "cloudflare-r2".to_string(),
bucket: "workspace-blobs".to_string(),
config: serde_json::json!({
"accountId": "account",
"jurisdiction": "eu",
"credentials": {
"accessKeyId": "key",
"secretAccessKey": "secret"
},
"usePresignedURL": {
"enabled": true
}
}),
};
let config = ObjectStorageConfig::from_r2_config(storage).unwrap().unwrap();
assert_eq!(config.provider, "cloudflare-r2");
assert_eq!(config.bucket, "workspace-blobs");
assert_eq!(
config.endpoint.as_deref(),
Some("https://account.eu.r2.cloudflarestorage.com")
);
assert_eq!(config.region.as_deref(), Some("auto"));
assert!(config.force_path_style);
assert!(config.use_presigned_url);
assert_eq!(config.access_key_id.as_deref(), Some("key"));
}
#[test]
fn resolves_s3_config_from_config_json_shape() {
let storage = StorageProviderConfig {
provider: "aws-s3".to_string(),
bucket: "workspace-blobs".to_string(),
config: serde_json::json!({
"region": "us-west-2",
"credentials": {
"accessKeyId": "key",
"secretAccessKey": "secret",
"sessionToken": "session"
},
"forcePathStyle": true,
"requestTimeoutMs": 1000,
"minPartSize": 1024,
"presign": {
"expiresInSeconds": 60,
"signContentTypeForPut": false
}
}),
};
let config = ObjectStorageConfig::from_s3_config(storage).unwrap().unwrap();
assert_eq!(config.provider, "aws-s3");
assert_eq!(config.endpoint.as_deref(), Some("https://s3.us-west-2.amazonaws.com"));
assert_eq!(config.session_token.as_deref(), Some("session"));
assert!(config.force_path_style);
assert_eq!(config.request_timeout_ms, Some(1000));
assert_eq!(config.min_part_size, Some(1024));
assert_eq!(config.presign_expires_in_seconds, Some(60));
assert_eq!(config.presign_sign_content_type_for_put, Some(false));
}
#[tokio::test]
async fn object_storage_presign_put_returns_sigv4_url_and_headers() {
let storage = StorageProviderConfig {
provider: "aws-s3".to_string(),
bucket: "test-bucket".to_string(),
config: serde_json::json!({
"region": "us-east-1",
"endpoint": "https://s3.us-east-1.amazonaws.com",
"credentials": {
"accessKeyId": "key",
"secretAccessKey": "secret"
},
"presign": {
"expiresInSeconds": 60
}
}),
};
let config = ObjectStorageConfig::from_s3_config(storage).unwrap().unwrap();
let Ok(Ok(client)) = std::panic::catch_unwind(|| config.build_client()) else {
eprintln!("skipping object storage presign test: S3 client cannot be built in this environment");
return;
};
let result = client
.presign_put(
"key",
ObjectPutMetadata {
content_type: Some("text/plain".to_string()),
..Default::default()
},
)
.await
.unwrap();
assert!(result.url.contains("X-Amz-Algorithm=AWS4-HMAC-SHA256"));
assert!(result.url.contains("X-Amz-SignedHeaders="));
assert_eq!(
result.headers.get("Content-Type").map(String::as_str),
Some("text/plain")
);
assert!(result.expires_at_ms > 0);
}
#[test]
fn object_storage_orders_completed_multipart_parts_and_trims_etags() {
let parts = completed_multipart_parts(vec![
MultipartUploadPart {
part_number: 2,
etag: trim_etag("\"b\""),
},
MultipartUploadPart {
part_number: 1,
etag: trim_etag("a"),
},
]);
assert_eq!(parts[0].part_number, Some(1));
assert_eq!(parts[0].e_tag.as_deref(), Some("a"));
assert_eq!(parts[1].part_number, Some(2));
assert_eq!(parts[1].e_tag.as_deref(), Some("b"));
}
@@ -0,0 +1,165 @@
use std::collections::HashMap;
use aws_sdk_s3::types::CompletedPart;
use napi::Result;
use serde::Deserialize;
use crate::backend_runtime::{
error::napi_error,
types::{
RuntimeMultipartUploadInit, RuntimeMultipartUploadPart, RuntimeObjectGetResult, RuntimeObjectListEntry,
RuntimeObjectMetadata, RuntimeObjectStoragePutOptions, RuntimePresignedObjectRequest,
},
};
#[derive(Clone, Debug, Default)]
pub(super) struct ObjectPutMetadata {
pub(super) content_type: Option<String>,
pub(super) content_length: Option<i64>,
pub(super) checksum_crc32: Option<String>,
}
#[derive(Clone, Debug, PartialEq)]
pub(super) struct ObjectMetadata {
pub(super) content_type: String,
pub(super) content_length: i64,
pub(super) last_modified_ms: i64,
pub(super) checksum_crc32: Option<String>,
}
#[derive(Clone, Debug, PartialEq)]
pub(super) struct ObjectListEntry {
pub(super) key: String,
pub(super) content_length: i64,
pub(super) last_modified_ms: i64,
}
#[derive(Clone, Debug, PartialEq)]
pub(super) struct ObjectGetResult {
pub(super) body: Vec<u8>,
pub(super) metadata: ObjectMetadata,
}
#[derive(Clone, Debug, PartialEq)]
pub(super) struct PresignedObjectRequest {
pub(super) url: String,
pub(super) headers: HashMap<String, String>,
pub(super) expires_at_ms: i64,
}
#[derive(Clone, Debug, PartialEq)]
pub(super) struct MultipartUploadInitResult {
pub(super) upload_id: String,
pub(super) expires_at_ms: i64,
}
#[derive(Clone, Debug, PartialEq)]
pub(super) struct MultipartUploadPart {
pub(super) part_number: i32,
pub(super) etag: String,
}
#[derive(Clone, Debug, Deserialize)]
pub(in crate::backend_runtime) struct StorageProviderConfig {
pub(super) provider: String,
pub(super) bucket: String,
#[serde(default)]
pub(super) config: serde_json::Value,
}
pub(super) fn trim_etag(etag: &str) -> String {
etag.trim_matches('"').to_string()
}
pub(super) fn completed_multipart_parts(mut parts: Vec<MultipartUploadPart>) -> Vec<CompletedPart> {
parts.sort_by_key(|part| part.part_number);
parts
.into_iter()
.map(|part| {
CompletedPart::builder()
.part_number(part.part_number)
.e_tag(part.etag)
.build()
})
.collect()
}
impl From<RuntimeObjectStoragePutOptions> for ObjectPutMetadata {
fn from(options: RuntimeObjectStoragePutOptions) -> Self {
Self {
content_type: options.content_type,
content_length: options.content_length,
checksum_crc32: options.checksum_crc32,
}
}
}
impl From<ObjectMetadata> for RuntimeObjectMetadata {
fn from(metadata: ObjectMetadata) -> Self {
Self {
content_type: metadata.content_type,
content_length: metadata.content_length,
last_modified_ms: metadata.last_modified_ms,
checksum_crc32: metadata.checksum_crc32,
}
}
}
impl From<ObjectListEntry> for RuntimeObjectListEntry {
fn from(entry: ObjectListEntry) -> Self {
Self {
key: entry.key,
content_length: entry.content_length,
last_modified_ms: entry.last_modified_ms,
}
}
}
impl TryFrom<PresignedObjectRequest> for RuntimePresignedObjectRequest {
type Error = napi::Error;
fn try_from(request: PresignedObjectRequest) -> Result<Self> {
Ok(Self {
url: request.url,
headers_json: serde_json::to_string(&request.headers)
.map_err(|err| napi_error(format!("ObjectStorage headers serialization failed: {err}")))?,
expires_at_ms: request.expires_at_ms,
})
}
}
impl From<ObjectGetResult> for RuntimeObjectGetResult {
fn from(result: ObjectGetResult) -> Self {
Self {
body: result.body.into(),
metadata: result.metadata.into(),
}
}
}
impl From<MultipartUploadInitResult> for RuntimeMultipartUploadInit {
fn from(init: MultipartUploadInitResult) -> Self {
Self {
upload_id: init.upload_id,
expires_at_ms: init.expires_at_ms,
}
}
}
impl From<RuntimeMultipartUploadPart> for MultipartUploadPart {
fn from(part: RuntimeMultipartUploadPart) -> Self {
Self {
part_number: part.part_number,
etag: part.etag,
}
}
}
impl From<MultipartUploadPart> for RuntimeMultipartUploadPart {
fn from(part: MultipartUploadPart) -> Self {
Self {
part_number: part.part_number,
etag: part.etag,
}
}
}
@@ -0,0 +1,42 @@
use napi::Result;
use super::{auth_challenge_purpose, dto::RuntimeStateRows};
pub(super) async fn create(
rows: &RuntimeStateRows,
purpose: &str,
token: &str,
payload: serde_json::Value,
ttl_ms: i64,
) -> Result<bool> {
rows
.insert_payload_if_absent(
&auth_challenge_purpose(purpose),
token,
None,
payload,
ttl_ms,
"RuntimeState auth challenge create",
)
.await
}
pub(super) async fn get(rows: &RuntimeStateRows, purpose: &str, token: &str) -> Result<Option<serde_json::Value>> {
rows
.active_payload(
&auth_challenge_purpose(purpose),
token,
"RuntimeState auth challenge get",
)
.await
}
pub(super) async fn consume(rows: &RuntimeStateRows, purpose: &str, token: &str) -> Result<Option<serde_json::Value>> {
rows
.consume_payload(
&auth_challenge_purpose(purpose),
token,
"RuntimeState auth challenge consume",
)
.await
}
@@ -0,0 +1,136 @@
use napi::Result;
use super::dto::{RuntimeStateInsertPayload, RuntimeStatePayloadRow, RuntimeStateRows};
use crate::backend_runtime::{
constants::{BYOK_LOCAL_LEASE_ACTIVE_PURPOSE, BYOK_LOCAL_LEASE_PURPOSE},
error::napi_error,
types::RuntimeByokLocalLeaseRecord,
};
pub(super) async fn get(rows: &RuntimeStateRows, lease_id: String) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
get_lease_by_id(rows, &lease_id).await
}
pub(super) async fn create(
rows: &RuntimeStateRows,
active_key: String,
lease_id: String,
payload: serde_json::Value,
ttl_ms: i64,
) -> Result<RuntimeByokLocalLeaseRecord> {
if ttl_ms <= 0 {
return Err(napi_error("BYOK local lease ttl must be positive"));
}
let mut tx = rows.begin("RuntimeState BYOK local lease").await?;
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
.bind(&active_key)
.execute(&mut *tx)
.await
.map_err(|err| napi_error(format!("RuntimeState BYOK local lease active lock failed: {err}")))?;
if let Some(active) = rows
.active_payload_with_expires_for_update_in_tx(
&mut tx,
BYOK_LOCAL_LEASE_ACTIVE_PURPOSE,
&active_key,
"RuntimeState BYOK local lease active get",
)
.await?
{
let existing_lease = match active.payload.get("leaseId").and_then(serde_json::Value::as_str) {
Some(existing_lease_id) => get_lease_by_id_in_tx(rows, &mut tx, existing_lease_id).await?,
None => None,
};
if let Some(lease) = existing_lease {
tx.commit().await.map_err(|err| {
napi_error(format!(
"RuntimeState BYOK local lease transaction commit failed: {err}"
))
})?;
return Ok(lease);
}
rows
.delete_by_key_in_tx(
&mut tx,
BYOK_LOCAL_LEASE_ACTIVE_PURPOSE,
&active_key,
"RuntimeState BYOK local lease stale active delete",
)
.await?;
}
let expires_at_ms = rows
.insert_payload_returning_expires_in_tx(
&mut tx,
RuntimeStateInsertPayload {
purpose: BYOK_LOCAL_LEASE_PURPOSE,
token: &lease_id,
lookup_key: &active_key,
payload: &payload,
ttl_ms,
context: "RuntimeState BYOK local lease create",
},
)
.await?;
let active_payload = serde_json::json!({ "leaseId": lease_id });
rows
.insert_payload_returning_expires_in_tx(
&mut tx,
RuntimeStateInsertPayload {
purpose: BYOK_LOCAL_LEASE_ACTIVE_PURPOSE,
token: &active_key,
lookup_key: &active_key,
payload: &active_payload,
ttl_ms,
context: "RuntimeState BYOK local lease active create",
},
)
.await?;
tx.commit().await.map_err(|err| {
napi_error(format!(
"RuntimeState BYOK local lease transaction commit failed: {err}"
))
})?;
Ok(RuntimeByokLocalLeaseRecord {
lease_id,
payload,
expires_at_ms,
})
}
async fn get_lease_by_id(rows: &RuntimeStateRows, lease_id: &str) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
rows
.active_payload_with_expires(BYOK_LOCAL_LEASE_PURPOSE, lease_id, "RuntimeState BYOK local lease get")
.await?
.map(|row| record_from_row(lease_id, row))
.transpose()
}
async fn get_lease_by_id_in_tx(
rows: &RuntimeStateRows,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
lease_id: &str,
) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
rows
.active_payload_with_expires_for_update_in_tx(
tx,
BYOK_LOCAL_LEASE_PURPOSE,
lease_id,
"RuntimeState BYOK local lease get",
)
.await?
.map(|row| record_from_row(lease_id, row))
.transpose()
}
fn record_from_row(lease_id: &str, row: RuntimeStatePayloadRow) -> Result<RuntimeByokLocalLeaseRecord> {
Ok(RuntimeByokLocalLeaseRecord {
lease_id: lease_id.to_string(),
payload: row.payload,
expires_at_ms: row.expires_at_ms,
})
}
@@ -0,0 +1,454 @@
use napi::Result;
use sqlx::{PgPool, Row};
use crate::backend_runtime::{error::napi_error, token_hash};
pub(super) struct RuntimeStatePayloadRow {
pub(super) payload: serde_json::Value,
pub(super) expires_at_ms: i64,
}
pub(super) struct RuntimeStateLockedRow {
pub(super) payload: serde_json::Value,
pub(super) attempts: i32,
pub(super) expires_at: chrono::DateTime<chrono::Utc>,
}
pub(super) struct RuntimeStateInsertPayload<'a> {
pub(super) purpose: &'a str,
pub(super) token: &'a str,
pub(super) lookup_key: &'a str,
pub(super) payload: &'a serde_json::Value,
pub(super) ttl_ms: i64,
pub(super) context: &'a str,
}
#[derive(Clone)]
pub(super) struct RuntimeStateRows {
pub(super) pool: PgPool,
}
impl RuntimeStateRows {
pub(super) fn new(pool: PgPool) -> Self {
Self { pool }
}
pub(super) fn pool(&self) -> &PgPool {
&self.pool
}
pub(super) async fn begin(&self, context: &str) -> Result<sqlx::Transaction<'_, sqlx::Postgres>> {
self
.pool
.begin()
.await
.map_err(|err| napi_error(format!("{context} transaction failed: {err}")))
}
pub(super) async fn insert_payload(
&self,
purpose: &str,
token: &str,
lookup_key: Option<&str>,
payload: serde_json::Value,
ttl_ms: i64,
context: &str,
) -> Result<()> {
sqlx::query(
r#"
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, expires_at)
VALUES ($1, $2, $3, $4, CURRENT_TIMESTAMP + ($5 * INTERVAL '1 millisecond'))
"#,
)
.bind(purpose)
.bind(token_hash(token))
.bind(lookup_key)
.bind(payload)
.bind(ttl_ms as f64)
.execute(&self.pool)
.await
.map_err(|err| napi_error(format!("{context} failed: {err}")))?;
Ok(())
}
pub(super) async fn insert_payload_if_absent(
&self,
purpose: &str,
token: &str,
lookup_key: Option<&str>,
payload: serde_json::Value,
ttl_ms: i64,
context: &str,
) -> Result<bool> {
let inserted = sqlx::query(
r#"
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, expires_at)
VALUES ($1, $2, $3, $4, CURRENT_TIMESTAMP + ($5 * INTERVAL '1 millisecond'))
ON CONFLICT (purpose, token_hash) DO NOTHING
"#,
)
.bind(purpose)
.bind(token_hash(token))
.bind(lookup_key)
.bind(payload)
.bind(ttl_ms as f64)
.execute(&self.pool)
.await
.map_err(|err| napi_error(format!("{context} failed: {err}")))?
.rows_affected()
== 1;
Ok(inserted)
}
pub(super) async fn upsert_payload_reset_attempts(
&self,
purpose: &str,
token: &str,
lookup_key: &str,
payload: serde_json::Value,
ttl_ms: i64,
context: &str,
) -> Result<()> {
sqlx::query(
r#"
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, attempts, consumed_at, expires_at)
VALUES ($1, $2, $3, $4, 0, NULL, CURRENT_TIMESTAMP + ($5 * INTERVAL '1 millisecond'))
ON CONFLICT (purpose, token_hash) DO UPDATE
SET lookup_key = EXCLUDED.lookup_key,
payload = EXCLUDED.payload,
attempts = 0,
consumed_at = NULL,
expires_at = EXCLUDED.expires_at,
updated_at = CURRENT_TIMESTAMP
"#,
)
.bind(purpose)
.bind(token_hash(token))
.bind(lookup_key)
.bind(payload)
.bind(ttl_ms as f64)
.execute(&self.pool)
.await
.map_err(|err| napi_error(format!("{context} failed: {err}")))?;
Ok(())
}
pub(super) async fn active_payload(
&self,
purpose: &str,
token: &str,
context: &str,
) -> Result<Option<serde_json::Value>> {
let row = sqlx::query(
r#"
SELECT payload
FROM runtime_states
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > CURRENT_TIMESTAMP
"#,
)
.bind(purpose)
.bind(token_hash(token))
.fetch_optional(&self.pool)
.await
.map_err(|err| napi_error(format!("{context} failed: {err}")))?;
Ok(row.map(|row| row.get::<serde_json::Value, _>("payload")))
}
pub(super) async fn active_payload_with_expires(
&self,
purpose: &str,
token: &str,
context: &str,
) -> Result<Option<RuntimeStatePayloadRow>> {
let row = sqlx::query(
r#"
SELECT payload, (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
FROM runtime_states
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > CURRENT_TIMESTAMP
"#,
)
.bind(purpose)
.bind(token_hash(token))
.fetch_optional(&self.pool)
.await
.map_err(|err| napi_error(format!("{context} failed: {err}")))?;
Ok(row.map(payload_row))
}
pub(super) async fn consume_payload(
&self,
purpose: &str,
token: &str,
context: &str,
) -> Result<Option<serde_json::Value>> {
let row = sqlx::query(
r#"
UPDATE runtime_states
SET consumed_at = CURRENT_TIMESTAMP,
updated_at = CURRENT_TIMESTAMP
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > CURRENT_TIMESTAMP
RETURNING payload
"#,
)
.bind(purpose)
.bind(token_hash(token))
.fetch_optional(&self.pool)
.await
.map_err(|err| napi_error(format!("{context} failed: {err}")))?;
Ok(row.map(|row| row.get::<serde_json::Value, _>("payload")))
}
pub(super) async fn consume_payload_with_expires(
&self,
purpose: &str,
token: &str,
context: &str,
) -> Result<Option<RuntimeStatePayloadRow>> {
let row = sqlx::query(
r#"
UPDATE runtime_states
SET consumed_at = CURRENT_TIMESTAMP,
updated_at = CURRENT_TIMESTAMP
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > CURRENT_TIMESTAMP
RETURNING payload, (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
"#,
)
.bind(purpose)
.bind(token_hash(token))
.fetch_optional(&self.pool)
.await
.map_err(|err| napi_error(format!("{context} failed: {err}")))?;
Ok(row.map(payload_row))
}
pub(super) async fn active_payload_with_expires_for_update_in_tx(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
purpose: &str,
token: &str,
context: &str,
) -> Result<Option<RuntimeStatePayloadRow>> {
let row = sqlx::query(
r#"
SELECT payload, (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
FROM runtime_states
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > clock_timestamp()
FOR UPDATE
"#,
)
.bind(purpose)
.bind(token_hash(token))
.fetch_optional(&mut **tx)
.await
.map_err(|err| napi_error(format!("{context} failed: {err}")))?;
Ok(row.map(payload_row))
}
pub(super) async fn unconsumed_row_for_update_in_tx(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
purpose: &str,
token: &str,
context: &str,
) -> Result<Option<RuntimeStateLockedRow>> {
let row = sqlx::query(
r#"
SELECT payload, attempts, expires_at
FROM runtime_states
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
FOR UPDATE
"#,
)
.bind(purpose)
.bind(token_hash(token))
.fetch_optional(&mut **tx)
.await
.map_err(|err| napi_error(format!("{context} failed: {err}")))?;
Ok(row.map(|row| RuntimeStateLockedRow {
payload: row.get("payload"),
attempts: row.get("attempts"),
expires_at: row.get("expires_at"),
}))
}
pub(super) async fn insert_payload_returning_expires_in_tx(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
input: RuntimeStateInsertPayload<'_>,
) -> Result<i64> {
let row = sqlx::query(
r#"
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, expires_at)
VALUES ($1, $2, $3, $4, CURRENT_TIMESTAMP + ($5 * INTERVAL '1 millisecond'))
RETURNING (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
"#,
)
.bind(input.purpose)
.bind(token_hash(input.token))
.bind(input.lookup_key)
.bind(input.payload)
.bind(input.ttl_ms as f64)
.fetch_one(&mut **tx)
.await
.map_err(|err| napi_error(format!("{} failed: {err}", input.context)))?;
Ok(row.get::<i64, _>("expires_at_ms"))
}
pub(super) async fn upsert_expired_or_consumed_payload_returning_expires_in_tx(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
input: RuntimeStateInsertPayload<'_>,
) -> Result<Option<i64>> {
let row = sqlx::query(
r#"
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, expires_at)
VALUES ($1, $2, $3, $4, clock_timestamp() + ($5 * INTERVAL '1 millisecond'))
ON CONFLICT (purpose, token_hash) DO UPDATE
SET lookup_key = EXCLUDED.lookup_key,
payload = EXCLUDED.payload,
attempts = 0,
consumed_at = NULL,
expires_at = clock_timestamp() + ($5 * INTERVAL '1 millisecond')
WHERE runtime_states.consumed_at IS NOT NULL
OR runtime_states.expires_at <= clock_timestamp()
RETURNING (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
"#,
)
.bind(input.purpose)
.bind(token_hash(input.token))
.bind(input.lookup_key)
.bind(input.payload)
.bind(input.ttl_ms as f64)
.fetch_optional(&mut **tx)
.await
.map_err(|err| napi_error(format!("{} failed: {err}", input.context)))?;
Ok(row.map(|row| row.get::<i64, _>("expires_at_ms")))
}
pub(super) async fn update_attempts_in_tx(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
purpose: &str,
token: &str,
attempts: i32,
context: &str,
) -> Result<()> {
sqlx::query(
r#"
UPDATE runtime_states
SET attempts = $3,
updated_at = CURRENT_TIMESTAMP
WHERE purpose = $1
AND token_hash = $2
"#,
)
.bind(purpose)
.bind(token_hash(token))
.bind(attempts)
.execute(&mut **tx)
.await
.map_err(|err| napi_error(format!("{context} failed: {err}")))?;
Ok(())
}
pub(super) async fn delete_by_key_in_tx(
&self,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
purpose: &str,
token: &str,
context: &str,
) -> Result<()> {
sqlx::query("DELETE FROM runtime_states WHERE purpose = $1 AND token_hash = $2")
.bind(purpose)
.bind(token_hash(token))
.execute(&mut **tx)
.await
.map_err(|err| napi_error(format!("{context} failed: {err}")))?;
Ok(())
}
pub(super) async fn cleanup_expired_or_consumed(&self, limit: i64, context: &str) -> Result<i64> {
let result = sqlx::query(
r#"
DELETE FROM runtime_states
WHERE (purpose, token_hash) IN (
SELECT purpose, token_hash FROM runtime_states
WHERE expires_at <= CURRENT_TIMESTAMP
OR consumed_at IS NOT NULL
ORDER BY expires_at ASC
LIMIT $1
)
"#,
)
.bind(limit)
.execute(&self.pool)
.await
.map_err(|err| napi_error(format!("{context} failed: {err}")))?;
Ok(result.rows_affected() as i64)
}
pub(super) async fn cleanup_expired_by_purpose_prefix(
&self,
purpose_prefix: &str,
limit: i64,
context: &str,
) -> Result<i64> {
let result = sqlx::query(
r#"
DELETE FROM runtime_states
WHERE (purpose, token_hash) IN (
SELECT purpose, token_hash FROM runtime_states
WHERE purpose LIKE $1
AND expires_at <= CURRENT_TIMESTAMP
ORDER BY expires_at ASC
LIMIT $2
)
"#,
)
.bind(format!("{purpose_prefix}%"))
.bind(limit)
.execute(&self.pool)
.await
.map_err(|err| napi_error(format!("{context} failed: {err}")))?;
Ok(result.rows_affected() as i64)
}
}
fn payload_row(row: sqlx::postgres::PgRow) -> RuntimeStatePayloadRow {
RuntimeStatePayloadRow {
payload: row.get("payload"),
expires_at_ms: row.get("expires_at_ms"),
}
}
@@ -0,0 +1,194 @@
use napi::Result;
use super::dto::{RuntimeStateInsertPayload, RuntimeStatePayloadRow, RuntimeStateRows};
use crate::backend_runtime::{
constants::{WORKSPACE_INVITE_LINK_ID_PURPOSE, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE},
error::napi_error,
types::RuntimeWorkspaceInviteLinkRecord,
};
pub(super) async fn get_by_workspace(
rows: &RuntimeStateRows,
workspace_id: String,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
get_by_key(rows, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE, &workspace_id).await
}
pub(super) async fn get_by_invite_id(
rows: &RuntimeStateRows,
invite_id: String,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
get_by_key(rows, WORKSPACE_INVITE_LINK_ID_PURPOSE, &invite_id).await
}
pub(super) async fn create(
rows: &RuntimeStateRows,
workspace_id: String,
invite_id: String,
inviter_user_id: String,
ttl_ms: i64,
) -> Result<RuntimeWorkspaceInviteLinkRecord> {
if ttl_ms <= 0 {
return Err(napi_error("workspace invite link ttl must be positive"));
}
let mut tx = rows.begin("RuntimeState workspace invite link").await?;
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
.bind(&workspace_id)
.execute(&mut *tx)
.await
.map_err(|err| napi_error(format!("RuntimeState workspace invite link active lock failed: {err}")))?;
if let Some(existing) =
get_by_key_in_tx(rows, &mut tx, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE, &workspace_id).await?
{
tx.commit().await.map_err(|err| {
napi_error(format!(
"RuntimeState workspace invite link transaction commit failed: {err}"
))
})?;
return Ok(existing);
}
let payload = serde_json::json!({
"workspaceId": workspace_id,
"inviteId": invite_id,
"inviterUserId": inviter_user_id,
});
let Some(expires_at_ms) = rows
.upsert_expired_or_consumed_payload_returning_expires_in_tx(
&mut tx,
RuntimeStateInsertPayload {
purpose: WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE,
token: &workspace_id,
lookup_key: &workspace_id,
payload: &payload,
ttl_ms,
context: "RuntimeState workspace invite link create",
},
)
.await?
else {
let existing = get_by_key_in_tx(rows, &mut tx, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE, &workspace_id).await?;
tx.commit().await.map_err(|err| {
napi_error(format!(
"RuntimeState workspace invite link transaction commit failed: {err}"
))
})?;
return existing.ok_or_else(|| napi_error("RuntimeState workspace invite link active conflict missing row"));
};
rows
.insert_payload_returning_expires_in_tx(
&mut tx,
RuntimeStateInsertPayload {
purpose: WORKSPACE_INVITE_LINK_ID_PURPOSE,
token: &invite_id,
lookup_key: &invite_id,
payload: &payload,
ttl_ms,
context: "RuntimeState workspace invite link create",
},
)
.await?;
tx.commit().await.map_err(|err| {
napi_error(format!(
"RuntimeState workspace invite link transaction commit failed: {err}"
))
})?;
Ok(RuntimeWorkspaceInviteLinkRecord {
workspace_id,
invite_id,
inviter_user_id,
expires_at_ms,
})
}
pub(super) async fn revoke(rows: &RuntimeStateRows, workspace_id: String) -> Result<bool> {
let mut tx = rows.begin("RuntimeState workspace invite link").await?;
let existing = get_by_key_in_tx(rows, &mut tx, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE, &workspace_id).await?;
let Some(existing) = existing else {
tx.commit().await.map_err(|err| {
napi_error(format!(
"RuntimeState workspace invite link transaction commit failed: {err}"
))
})?;
return Ok(false);
};
rows
.delete_by_key_in_tx(
&mut tx,
WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE,
&workspace_id,
"RuntimeState workspace invite link revoke",
)
.await?;
rows
.delete_by_key_in_tx(
&mut tx,
WORKSPACE_INVITE_LINK_ID_PURPOSE,
&existing.invite_id,
"RuntimeState workspace invite link revoke",
)
.await?;
tx.commit().await.map_err(|err| {
napi_error(format!(
"RuntimeState workspace invite link transaction commit failed: {err}"
))
})?;
Ok(true)
}
async fn get_by_key(
rows: &RuntimeStateRows,
purpose: &str,
key: &str,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
rows
.active_payload_with_expires(purpose, key, "RuntimeState workspace invite link get")
.await?
.map(record_from_row)
.transpose()
}
async fn get_by_key_in_tx(
rows: &RuntimeStateRows,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
purpose: &str,
key: &str,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
rows
.active_payload_with_expires_for_update_in_tx(tx, purpose, key, "RuntimeState workspace invite link get")
.await?
.map(record_from_row)
.transpose()
}
fn record_from_row(row: RuntimeStatePayloadRow) -> Result<RuntimeWorkspaceInviteLinkRecord> {
Ok(RuntimeWorkspaceInviteLinkRecord {
workspace_id: row
.payload
.get("workspaceId")
.and_then(serde_json::Value::as_str)
.ok_or_else(|| napi_error("RuntimeState workspace invite link payload missing workspaceId"))?
.to_string(),
invite_id: row
.payload
.get("inviteId")
.and_then(serde_json::Value::as_str)
.ok_or_else(|| napi_error("RuntimeState workspace invite link payload missing inviteId"))?
.to_string(),
inviter_user_id: row
.payload
.get("inviterUserId")
.and_then(serde_json::Value::as_str)
.ok_or_else(|| napi_error("RuntimeState workspace invite link payload missing inviterUserId"))?
.to_string(),
expires_at_ms: row.expires_at_ms,
})
}
@@ -0,0 +1,178 @@
use napi::Result;
use super::dto::RuntimeStateRows;
use crate::backend_runtime::{
constants::{MAGIC_LINK_OTP_PURPOSE, MAX_MAGIC_LINK_OTP_ATTEMPTS},
error::napi_error,
types::RuntimeMagicLinkOtpConsumeResult,
};
impl RuntimeMagicLinkOtpConsumeResult {
fn ok(token: String) -> Self {
Self {
ok: true,
token: Some(token),
reason: None,
}
}
fn fail(reason: &'static str) -> Self {
Self {
ok: false,
token: None,
reason: Some(reason.to_string()),
}
}
}
pub(super) async fn upsert(
rows: &RuntimeStateRows,
email: String,
otp_hash: String,
token: String,
client_nonce: Option<String>,
ttl_ms: i64,
) -> Result<()> {
if ttl_ms <= 0 {
return Err(napi_error("magic link otp ttl must be positive"));
}
let payload = serde_json::json!({
"otpHash": otp_hash,
"token": token,
"clientNonce": client_nonce,
});
rows
.upsert_payload_reset_attempts(
MAGIC_LINK_OTP_PURPOSE,
&email,
&email,
payload,
ttl_ms,
"RuntimeState magic link otp upsert",
)
.await
}
pub(super) async fn consume(
rows: &RuntimeStateRows,
email: String,
otp_hash: String,
client_nonce: Option<String>,
) -> Result<RuntimeMagicLinkOtpConsumeResult> {
let mut tx = rows.begin("RuntimeState magic link otp").await?;
let row = rows
.unconsumed_row_for_update_in_tx(
&mut tx,
MAGIC_LINK_OTP_PURPOSE,
&email,
"RuntimeState magic link otp lookup",
)
.await?;
let Some(row) = row else {
tx.rollback().await.map_err(|err| {
napi_error(format!(
"RuntimeState magic link otp transaction rollback failed: {err}"
))
})?;
return Ok(RuntimeMagicLinkOtpConsumeResult::fail("not_found"));
};
let payload = row.payload;
let attempts = row.attempts;
let expires_at = row.expires_at;
if expires_at <= chrono::Utc::now() {
rows
.delete_by_key_in_tx(
&mut tx,
MAGIC_LINK_OTP_PURPOSE,
&email,
"RuntimeState magic link otp delete",
)
.await?;
tx.commit()
.await
.map_err(|err| napi_error(format!("RuntimeState magic link otp transaction commit failed: {err}")))?;
return Ok(RuntimeMagicLinkOtpConsumeResult::fail("expired"));
}
let stored_client_nonce = payload.get("clientNonce").and_then(serde_json::Value::as_str);
if stored_client_nonce.is_some() && stored_client_nonce != client_nonce.as_deref() {
tx.commit()
.await
.map_err(|err| napi_error(format!("RuntimeState magic link otp transaction commit failed: {err}")))?;
return Ok(RuntimeMagicLinkOtpConsumeResult::fail("nonce_mismatch"));
}
if attempts >= MAX_MAGIC_LINK_OTP_ATTEMPTS {
rows
.delete_by_key_in_tx(
&mut tx,
MAGIC_LINK_OTP_PURPOSE,
&email,
"RuntimeState magic link otp delete",
)
.await?;
tx.commit()
.await
.map_err(|err| napi_error(format!("RuntimeState magic link otp transaction commit failed: {err}")))?;
return Ok(RuntimeMagicLinkOtpConsumeResult::fail("locked"));
}
let stored_otp_hash = payload.get("otpHash").and_then(serde_json::Value::as_str);
if stored_otp_hash != Some(otp_hash.as_str()) {
let attempts = attempts + 1;
if attempts >= MAX_MAGIC_LINK_OTP_ATTEMPTS {
rows
.delete_by_key_in_tx(
&mut tx,
MAGIC_LINK_OTP_PURPOSE,
&email,
"RuntimeState magic link otp delete",
)
.await?;
tx.commit()
.await
.map_err(|err| napi_error(format!("RuntimeState magic link otp transaction commit failed: {err}")))?;
return Ok(RuntimeMagicLinkOtpConsumeResult::fail("locked"));
}
rows
.update_attempts_in_tx(
&mut tx,
MAGIC_LINK_OTP_PURPOSE,
&email,
attempts,
"RuntimeState magic link otp attempts update",
)
.await?;
tx.commit()
.await
.map_err(|err| napi_error(format!("RuntimeState magic link otp transaction commit failed: {err}")))?;
return Ok(RuntimeMagicLinkOtpConsumeResult::fail("invalid_otp"));
}
let token = payload
.get("token")
.and_then(serde_json::Value::as_str)
.ok_or_else(|| napi_error("RuntimeState magic link otp payload missing token"))?
.to_string();
rows
.delete_by_key_in_tx(
&mut tx,
MAGIC_LINK_OTP_PURPOSE,
&email,
"RuntimeState magic link otp delete",
)
.await?;
tx.commit()
.await
.map_err(|err| napi_error(format!("RuntimeState magic link otp transaction commit failed: {err}")))?;
Ok(RuntimeMagicLinkOtpConsumeResult::ok(token))
}
@@ -0,0 +1,235 @@
use napi::Result;
use super::{
BackendRuntime,
error::napi_error,
types::{
RuntimeByokLocalLeaseRecord, RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord,
RuntimeWorkspaceInviteLinkRecord,
},
};
mod auth_challenge;
mod byok_local_lease;
mod dto;
mod invite_link;
mod magic_link_otp;
mod store;
mod verification_token;
use store::RuntimeStateStore;
pub(super) fn auth_challenge_purpose(purpose: &str) -> String {
format!("auth_challenge:{purpose}")
}
pub(super) fn verification_token_purpose(token_type: i32) -> String {
format!("verification_token:{token_type}")
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn create_auth_challenge(
&self,
purpose: String,
token: String,
payload: serde_json::Value,
ttl_ms: i64,
) -> Result<bool> {
if ttl_ms <= 0 {
return Err(napi_error("auth challenge ttl must be positive"));
}
RuntimeStateStore::new(self.pool().await?)
.create_auth_challenge(&purpose, &token, payload, ttl_ms)
.await
}
#[napi]
pub async fn get_auth_challenge(&self, purpose: String, token: String) -> Result<Option<serde_json::Value>> {
RuntimeStateStore::new(self.pool().await?)
.get_auth_challenge(&purpose, &token)
.await
}
#[napi]
pub async fn consume_auth_challenge(&self, purpose: String, token: String) -> Result<Option<serde_json::Value>> {
RuntimeStateStore::new(self.pool().await?)
.consume_auth_challenge(&purpose, &token)
.await
}
#[napi]
pub async fn create_verification_token(
&self,
token_type: i32,
credential: Option<String>,
ttl_ms: i64,
) -> Result<String> {
if ttl_ms <= 0 {
return Err(napi_error("verification token ttl must be positive"));
}
RuntimeStateStore::new(self.pool().await?)
.create_verification_token(token_type, credential, ttl_ms)
.await
}
#[napi]
pub async fn get_verification_token(
&self,
token_type: i32,
token: String,
keep: Option<bool>,
) -> Result<Option<RuntimeVerificationTokenRecord>> {
let keep = keep.unwrap_or(false);
RuntimeStateStore::new(self.pool().await?)
.get_verification_token(token_type, token, keep)
.await
}
#[napi]
pub async fn verify_verification_token(
&self,
token_type: i32,
token: String,
credential: Option<String>,
keep: Option<bool>,
) -> Result<Option<RuntimeVerificationTokenRecord>> {
let keep = keep.unwrap_or(false);
RuntimeStateStore::new(self.pool().await?)
.verify_verification_token(token_type, token, credential, keep)
.await
}
#[napi]
pub async fn cleanup_expired_verification_tokens(&self, limit: i64) -> Result<i64> {
if limit <= 0 {
return Err(napi_error("verification token cleanup limit must be positive"));
}
RuntimeStateStore::new(self.pool().await?)
.cleanup_expired_verification_tokens(limit)
.await
}
#[napi]
pub async fn upsert_magic_link_otp(
&self,
email: String,
otp_hash: String,
token: String,
client_nonce: Option<String>,
ttl_ms: i64,
) -> Result<()> {
RuntimeStateStore::new(self.pool().await?)
.upsert_magic_link_otp(email, otp_hash, token, client_nonce, ttl_ms)
.await
}
#[napi]
pub async fn consume_magic_link_otp(
&self,
email: String,
otp_hash: String,
client_nonce: Option<String>,
) -> Result<RuntimeMagicLinkOtpConsumeResult> {
RuntimeStateStore::new(self.pool().await?)
.consume_magic_link_otp(email, otp_hash, client_nonce)
.await
}
#[napi]
pub async fn create_workspace_invite_link(
&self,
workspace_id: String,
invite_id: String,
inviter_user_id: String,
ttl_ms: i64,
) -> Result<RuntimeWorkspaceInviteLinkRecord> {
RuntimeStateStore::new(self.pool().await?)
.create_workspace_invite_link(workspace_id, invite_id, inviter_user_id, ttl_ms)
.await
}
#[napi]
pub async fn get_workspace_invite_link(
&self,
workspace_id: String,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
RuntimeStateStore::new(self.pool().await?)
.get_workspace_invite_link(workspace_id)
.await
}
#[napi]
pub async fn get_workspace_invite_link_by_id(
&self,
invite_id: String,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
RuntimeStateStore::new(self.pool().await?)
.get_workspace_invite_link_by_id(invite_id)
.await
}
#[napi]
pub async fn revoke_workspace_invite_link(&self, workspace_id: String) -> Result<bool> {
RuntimeStateStore::new(self.pool().await?)
.revoke_workspace_invite_link(workspace_id)
.await
}
#[napi]
pub async fn create_byok_local_lease(
&self,
active_key: String,
lease_id: String,
payload: serde_json::Value,
ttl_ms: i64,
) -> Result<RuntimeByokLocalLeaseRecord> {
RuntimeStateStore::new(self.pool().await?)
.create_byok_local_lease(active_key, lease_id, payload, ttl_ms)
.await
}
#[napi]
pub async fn get_byok_local_lease(&self, lease_id: String) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
RuntimeStateStore::new(self.pool().await?)
.get_byok_local_lease(lease_id)
.await
}
#[napi]
pub async fn cleanup_expired_runtime_states(&self, limit: i64) -> Result<i64> {
if limit <= 0 {
return Err(napi_error("runtime state cleanup limit must be positive"));
}
RuntimeStateStore::new(self.pool().await?)
.cleanup_expired_runtime_states(limit)
.await
}
}
#[cfg(test)]
mod tests {
use crate::backend_runtime::{
constants::{MAGIC_LINK_OTP_PURPOSE, WORKSPACE_INVITE_LINK_ID_PURPOSE, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE},
token_hash,
};
#[test]
fn magic_link_otp_uses_scoped_purpose_and_email_hash() {
assert_eq!(MAGIC_LINK_OTP_PURPOSE, "magic_link_otp");
assert_ne!(token_hash("user@affine.test"), "user@affine.test");
assert_eq!(token_hash("user@affine.test"), token_hash("user@affine.test"));
assert_ne!(token_hash("user@affine.test"), token_hash("other@affine.test"));
}
#[test]
fn workspace_invite_link_uses_scoped_purposes_and_hashes() {
assert_eq!(
WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE,
"workspace_invite_link:workspace"
);
assert_eq!(WORKSPACE_INVITE_LINK_ID_PURPOSE, "workspace_invite_link:id");
assert_ne!(token_hash("workspace-id"), "workspace-id");
assert_ne!(token_hash("invite-id"), "invite-id");
}
}
@@ -0,0 +1,139 @@
use napi::Result;
use sqlx::PgPool;
use super::{auth_challenge, byok_local_lease, dto::RuntimeStateRows, invite_link, magic_link_otp, verification_token};
use crate::backend_runtime::types::{
RuntimeByokLocalLeaseRecord, RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord,
RuntimeWorkspaceInviteLinkRecord,
};
pub(super) struct RuntimeStateStore {
rows: RuntimeStateRows,
}
impl RuntimeStateStore {
pub(super) fn new(pool: PgPool) -> Self {
Self {
rows: RuntimeStateRows::new(pool),
}
}
pub(super) async fn create_auth_challenge(
&self,
purpose: &str,
token: &str,
payload: serde_json::Value,
ttl_ms: i64,
) -> Result<bool> {
auth_challenge::create(&self.rows, purpose, token, payload, ttl_ms).await
}
pub(super) async fn get_auth_challenge(&self, purpose: &str, token: &str) -> Result<Option<serde_json::Value>> {
auth_challenge::get(&self.rows, purpose, token).await
}
pub(super) async fn consume_auth_challenge(&self, purpose: &str, token: &str) -> Result<Option<serde_json::Value>> {
auth_challenge::consume(&self.rows, purpose, token).await
}
pub(super) async fn create_verification_token(
&self,
token_type: i32,
credential: Option<String>,
ttl_ms: i64,
) -> Result<String> {
verification_token::create(&self.rows, token_type, credential, ttl_ms).await
}
pub(super) async fn get_verification_token(
&self,
token_type: i32,
token: String,
keep: bool,
) -> Result<Option<RuntimeVerificationTokenRecord>> {
verification_token::get(&self.rows, token_type, token, keep).await
}
pub(super) async fn verify_verification_token(
&self,
token_type: i32,
token: String,
credential: Option<String>,
keep: bool,
) -> Result<Option<RuntimeVerificationTokenRecord>> {
verification_token::verify(&self.rows, token_type, token, credential, keep).await
}
pub(super) async fn cleanup_expired_verification_tokens(&self, limit: i64) -> Result<i64> {
verification_token::cleanup_expired(&self.rows, limit).await
}
pub(super) async fn cleanup_expired_runtime_states(&self, limit: i64) -> Result<i64> {
self
.rows
.cleanup_expired_or_consumed(limit, "RuntimeState cleanup")
.await
}
pub(super) async fn upsert_magic_link_otp(
&self,
email: String,
otp_hash: String,
token: String,
client_nonce: Option<String>,
ttl_ms: i64,
) -> Result<()> {
magic_link_otp::upsert(&self.rows, email, otp_hash, token, client_nonce, ttl_ms).await
}
pub(super) async fn consume_magic_link_otp(
&self,
email: String,
otp_hash: String,
client_nonce: Option<String>,
) -> Result<RuntimeMagicLinkOtpConsumeResult> {
magic_link_otp::consume(&self.rows, email, otp_hash, client_nonce).await
}
pub(super) async fn create_workspace_invite_link(
&self,
workspace_id: String,
invite_id: String,
inviter_user_id: String,
ttl_ms: i64,
) -> Result<RuntimeWorkspaceInviteLinkRecord> {
invite_link::create(&self.rows, workspace_id, invite_id, inviter_user_id, ttl_ms).await
}
pub(super) async fn get_workspace_invite_link(
&self,
workspace_id: String,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
invite_link::get_by_workspace(&self.rows, workspace_id).await
}
pub(super) async fn get_workspace_invite_link_by_id(
&self,
invite_id: String,
) -> Result<Option<RuntimeWorkspaceInviteLinkRecord>> {
invite_link::get_by_invite_id(&self.rows, invite_id).await
}
pub(super) async fn revoke_workspace_invite_link(&self, workspace_id: String) -> Result<bool> {
invite_link::revoke(&self.rows, workspace_id).await
}
pub(super) async fn create_byok_local_lease(
&self,
active_key: String,
lease_id: String,
payload: serde_json::Value,
ttl_ms: i64,
) -> Result<RuntimeByokLocalLeaseRecord> {
byok_local_lease::create(&self.rows, active_key, lease_id, payload, ttl_ms).await
}
pub(super) async fn get_byok_local_lease(&self, lease_id: String) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
byok_local_lease::get(&self.rows, lease_id).await
}
}
@@ -0,0 +1,150 @@
use napi::Result;
use sqlx::{PgPool, Row};
use uuid::Uuid;
use super::{
dto::{RuntimeStatePayloadRow, RuntimeStateRows},
verification_token_purpose,
};
use crate::backend_runtime::{error::napi_error, token_hash, types::RuntimeVerificationTokenRecord};
pub(super) async fn create(
rows: &RuntimeStateRows,
token_type: i32,
credential: Option<String>,
ttl_ms: i64,
) -> Result<String> {
let token = Uuid::new_v4().to_string();
let payload = serde_json::json!({ "credential": credential });
rows
.insert_payload(
&verification_token_purpose(token_type),
&token,
credential.as_deref(),
payload,
ttl_ms,
"RuntimeState verification token create",
)
.await?;
Ok(token)
}
pub(super) async fn get(
rows: &RuntimeStateRows,
token_type: i32,
token: String,
keep: bool,
) -> Result<Option<RuntimeVerificationTokenRecord>> {
let purpose = verification_token_purpose(token_type);
let row = if keep {
rows
.active_payload_with_expires(&purpose, &token, "RuntimeState verification token get")
.await?
} else {
rows
.consume_payload_with_expires(&purpose, &token, "RuntimeState verification token get")
.await?
};
Ok(row.map(|row| record_from_row(token_type, token, row)))
}
pub(super) async fn verify(
rows: &RuntimeStateRows,
token_type: i32,
token: String,
credential: Option<String>,
keep: bool,
) -> Result<Option<RuntimeVerificationTokenRecord>> {
let purpose = verification_token_purpose(token_type);
let row = if keep {
active_payload_with_credential(rows.pool(), &purpose, &token, credential.as_deref()).await
} else {
consume_payload_with_credential(rows.pool(), &purpose, &token, credential.as_deref()).await
}
.map_err(|err| napi_error(format!("RuntimeState verification token verify failed: {err}")))?;
Ok(row.map(|row| record_from_row(token_type, token, row)))
}
pub(super) async fn cleanup_expired(rows: &RuntimeStateRows, limit: i64) -> Result<i64> {
rows
.cleanup_expired_by_purpose_prefix("verification_token:", limit, "RuntimeState verification token cleanup")
.await
}
async fn active_payload_with_credential(
pool: &PgPool,
purpose: &str,
token: &str,
credential: Option<&str>,
) -> sqlx::Result<Option<RuntimeStatePayloadRow>> {
let row = sqlx::query(
r#"
SELECT payload, (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
FROM runtime_states
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > CURRENT_TIMESTAMP
AND (payload->>'credential' IS NULL OR payload->>'credential' = $3)
"#,
)
.bind(purpose)
.bind(token_hash(token))
.bind(credential)
.fetch_optional(pool)
.await?;
Ok(row.map(payload_row))
}
async fn consume_payload_with_credential(
pool: &PgPool,
purpose: &str,
token: &str,
credential: Option<&str>,
) -> sqlx::Result<Option<RuntimeStatePayloadRow>> {
let row = sqlx::query(
r#"
UPDATE runtime_states
SET consumed_at = CURRENT_TIMESTAMP,
updated_at = CURRENT_TIMESTAMP
WHERE purpose = $1
AND token_hash = $2
AND consumed_at IS NULL
AND expires_at > CURRENT_TIMESTAMP
AND (payload->>'credential' IS NULL OR payload->>'credential' = $3)
RETURNING payload, (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT AS expires_at_ms
"#,
)
.bind(purpose)
.bind(token_hash(token))
.bind(credential)
.fetch_optional(pool)
.await?;
Ok(row.map(payload_row))
}
fn payload_row(row: sqlx::postgres::PgRow) -> RuntimeStatePayloadRow {
RuntimeStatePayloadRow {
payload: row.get("payload"),
expires_at_ms: row.get("expires_at_ms"),
}
}
fn record_from_row(token_type: i32, token: String, row: RuntimeStatePayloadRow) -> RuntimeVerificationTokenRecord {
RuntimeVerificationTokenRecord {
token_type,
token,
credential: row
.payload
.get("credential")
.and_then(serde_json::Value::as_str)
.map(ToString::to_string),
expires_at_ms: row.expires_at_ms,
}
}
@@ -0,0 +1,40 @@
CREATE TABLE IF NOT EXISTS runtime_states (
purpose TEXT NOT NULL,
token_hash TEXT NOT NULL,
lookup_key TEXT,
payload JSONB NOT NULL,
attempts INTEGER NOT NULL DEFAULT 0,
consumed_at TIMESTAMPTZ(3),
expires_at TIMESTAMPTZ(3) NOT NULL,
created_at TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
PRIMARY KEY (purpose, token_hash)
);
CREATE INDEX IF NOT EXISTS runtime_states_lookup_idx
ON runtime_states (purpose, lookup_key)
WHERE lookup_key IS NOT NULL AND consumed_at IS NULL;
CREATE INDEX IF NOT EXISTS runtime_states_expires_at_idx
ON runtime_states (expires_at);
CREATE TABLE IF NOT EXISTS runtime_gates (
key TEXT PRIMARY KEY,
expires_at TIMESTAMPTZ(3) NOT NULL,
created_at TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS runtime_gates_expires_at_idx
ON runtime_gates (expires_at);
CREATE TABLE IF NOT EXISTS runtime_leases (
key TEXT PRIMARY KEY,
owner TEXT NOT NULL,
fencing_token BIGINT NOT NULL,
expires_at TIMESTAMPTZ(3) NOT NULL,
created_at TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS runtime_leases_expires_at_idx
ON runtime_leases (expires_at);
@@ -0,0 +1,85 @@
WITH targets AS (
SELECT UNNEST($1::varchar[]) AS workspace_id
),
snapshot_stats AS (
SELECT workspace_id,
COUNT(*) AS snapshot_count,
COALESCE(SUM(COALESCE(size, octet_length(blob))), 0) AS snapshot_size
FROM snapshots
WHERE workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
blob_stats AS (
SELECT workspace_id,
COUNT(*) FILTER (WHERE deleted_at IS NULL AND status = 'completed') AS blob_count,
COALESCE(SUM(size) FILTER (WHERE deleted_at IS NULL AND status = 'completed'), 0) AS blob_size
FROM blobs
WHERE workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
member_stats AS (
SELECT workspace_id, COUNT(*) AS member_count
FROM workspace_user_permissions
WHERE workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
public_page_stats AS (
SELECT workspace_id, COUNT(*) AS public_page_count
FROM workspace_pages
WHERE public = TRUE AND workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
feature_stats AS (
SELECT workspace_id,
ARRAY_AGG(DISTINCT name ORDER BY name) FILTER (WHERE activated) AS features
FROM workspace_features
WHERE workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
aggregated AS (
SELECT t.workspace_id,
COALESCE(ss.snapshot_count, 0) AS snapshot_count,
COALESCE(ss.snapshot_size, 0) AS snapshot_size,
COALESCE(bs.blob_count, 0) AS blob_count,
COALESCE(bs.blob_size, 0) AS blob_size,
COALESCE(ms.member_count, 0) AS member_count,
COALESCE(pp.public_page_count, 0) AS public_page_count,
COALESCE(fs.features, ARRAY[]::text[]) AS features
FROM targets t
LEFT JOIN snapshot_stats ss ON ss.workspace_id = t.workspace_id
LEFT JOIN blob_stats bs ON bs.workspace_id = t.workspace_id
LEFT JOIN member_stats ms ON ms.workspace_id = t.workspace_id
LEFT JOIN public_page_stats pp ON pp.workspace_id = t.workspace_id
LEFT JOIN feature_stats fs ON fs.workspace_id = t.workspace_id
)
INSERT INTO workspace_admin_stats (
workspace_id,
snapshot_count,
snapshot_size,
blob_count,
blob_size,
member_count,
public_page_count,
features,
updated_at
)
SELECT
workspace_id,
snapshot_count,
snapshot_size,
blob_count,
blob_size,
member_count,
public_page_count,
features,
NOW()
FROM aggregated
ON CONFLICT (workspace_id) DO UPDATE SET
snapshot_count = EXCLUDED.snapshot_count,
snapshot_size = EXCLUDED.snapshot_size,
blob_count = EXCLUDED.blob_count,
blob_size = EXCLUDED.blob_size,
member_count = EXCLUDED.member_count,
public_page_count = EXCLUDED.public_page_count,
features = EXCLUDED.features,
updated_at = EXCLUDED.updated_at
@@ -0,0 +1,431 @@
use anyhow::{Context, Result as AnyResult, anyhow};
use super::{runtime_state::*, *};
static PG_TEST_LOCK: std::sync::OnceLock<tokio::sync::Mutex<()>> = std::sync::OnceLock::new();
const TEST_VERIFICATION_TOKEN_TYPE: i32 = 99_999;
fn pg_test_lock() -> &'static tokio::sync::Mutex<()> {
PG_TEST_LOCK.get_or_init(|| tokio::sync::Mutex::new(()))
}
#[test]
fn migrations_include_runtime_tables_without_worker_heartbeats() {
assert!(RUNTIME_MIGRATIONS.contains("runtime_states"));
assert!(RUNTIME_MIGRATIONS.contains("runtime_gates"));
assert!(RUNTIME_MIGRATIONS.contains("runtime_leases"));
assert!(!RUNTIME_MIGRATIONS.contains("runtime_worker_heartbeats"));
}
#[test]
fn auth_challenge_state_uses_scoped_purpose_and_token_hash() {
assert_eq!(auth_challenge_purpose("oauth_state"), "auth_challenge:oauth_state");
assert_ne!(token_hash("plain-token"), "plain-token");
assert_eq!(token_hash("plain-token"), token_hash("plain-token"));
assert_ne!(token_hash("plain-token"), token_hash("other-token"));
}
#[test]
fn verification_token_state_uses_typed_purpose_and_token_hash() {
assert_eq!(verification_token_purpose(0), "verification_token:0");
assert_ne!(token_hash("verification-token"), "verification-token");
assert_eq!(token_hash("verification-token"), token_hash("verification-token"));
assert_ne!(token_hash("verification-token"), token_hash("other-token"));
}
async fn runtime_from_database_url() -> AnyResult<Option<BackendRuntime>> {
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return Ok(None);
};
let pool = PgPoolOptions::new()
.max_connections(5)
.connect(&database_url)
.await
.context("connect postgres for backend runtime tests")?;
migrate_runtime_tables(&pool)
.await
.map_err(|err| anyhow!(err.to_string()))?;
sqlx::query(
r#"
DELETE FROM runtime_states
WHERE purpose LIKE 'rust_test:%'
OR purpose LIKE 'auth_challenge:rust_test:%'
OR purpose = 'verification_token:99999'
"#,
)
.execute(&pool)
.await
.context("cleanup runtime_states for backend runtime tests")?;
sqlx::query("DELETE FROM runtime_gates WHERE key LIKE 'rust-test:%'")
.execute(&pool)
.await
.context("cleanup runtime_gates for backend runtime tests")?;
sqlx::query("DELETE FROM runtime_leases WHERE key LIKE 'rust-test:%'")
.execute(&pool)
.await
.context("cleanup runtime_leases for backend runtime tests")?;
Ok(Some(BackendRuntime {
config: RuntimeConfig {
database_url,
storage: None,
},
pool: Mutex::new(Some(pool)),
}))
}
#[tokio::test]
async fn runtime_gate_sql_semantics_are_atomic_and_ttl_bound() {
let _guard = pg_test_lock().lock().await;
let Some(runtime) = runtime_from_database_url().await.unwrap() else {
eprintln!("skipping postgres integration test: DATABASE_URL is not set");
return;
};
struct Case {
key: &'static str,
first_ttl_ms: i64,
wait_ms: Option<u64>,
second_expected: bool,
}
for case in [
Case {
key: "rust-test:gate:same-key",
first_ttl_ms: 30_000,
wait_ms: None,
second_expected: false,
},
Case {
key: "rust-test:gate:expired-key",
first_ttl_ms: 1,
wait_ms: Some(20),
second_expected: true,
},
] {
assert!(
runtime
.put_runtime_gate_if_absent(case.key.to_string(), case.first_ttl_ms)
.await
.unwrap()
);
if let Some(wait_ms) = case.wait_ms {
tokio::time::sleep(Duration::from_millis(wait_ms)).await;
}
assert_eq!(
runtime
.put_runtime_gate_if_absent(case.key.to_string(), 30_000)
.await
.unwrap(),
case.second_expected,
"{}",
case.key
);
}
let mut tasks = Vec::new();
for _ in 0..16 {
let runtime = BackendRuntime {
config: runtime.config.clone(),
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
};
tasks.push(tokio::spawn(async move {
runtime
.put_runtime_gate_if_absent("rust-test:gate:concurrent".to_string(), 30_000)
.await
.unwrap()
}));
}
let mut successful = 0;
for task in tasks {
if task.await.unwrap() {
successful += 1;
}
}
assert_eq!(successful, 1);
assert!(
runtime
.put_runtime_gate_if_absent("rust-test:gate:cleanup".to_string(), 1)
.await
.unwrap()
);
tokio::time::sleep(Duration::from_millis(20)).await;
assert_eq!(runtime.cleanup_expired_runtime_gates(100).await.unwrap(), 1);
assert_eq!(runtime.cleanup_expired_runtime_gates(100).await.unwrap(), 0);
}
#[tokio::test]
async fn coordination_lease_sql_semantics_are_fenced_and_ttl_bound() {
let _guard = pg_test_lock().lock().await;
let Some(runtime) = runtime_from_database_url().await.unwrap() else {
eprintln!("skipping postgres integration test: DATABASE_URL is not set");
return;
};
let lease = runtime
.acquire_coordination_lease("rust-test:lease:basic".to_string(), "owner-1".to_string(), 30_000)
.await
.unwrap()
.expect("first owner should acquire lease");
assert_eq!(lease.fencing_token, 1);
assert!(
!runtime
.release_coordination_lease(lease.key.clone(), "owner-2".to_string(), lease.fencing_token)
.await
.unwrap()
);
assert!(
runtime
.release_coordination_lease(lease.key.clone(), lease.owner.clone(), lease.fencing_token)
.await
.unwrap()
);
let mut tasks = Vec::new();
for index in 0..16 {
let runtime = BackendRuntime {
config: runtime.config.clone(),
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
};
tasks.push(tokio::spawn(async move {
runtime
.acquire_coordination_lease(
"rust-test:lease:concurrent".to_string(),
format!("owner-{index}"),
30_000,
)
.await
.unwrap()
.is_some()
}));
}
let mut successful = 0;
for task in tasks {
if task.await.unwrap() {
successful += 1;
}
}
assert_eq!(successful, 1);
let stale = runtime
.acquire_coordination_lease("rust-test:lease:stale".to_string(), "owner-1".to_string(), 1)
.await
.unwrap()
.expect("stale lease owner should acquire");
tokio::time::sleep(Duration::from_millis(20)).await;
let takeover = runtime
.acquire_coordination_lease("rust-test:lease:stale".to_string(), "owner-2".to_string(), 30_000)
.await
.unwrap()
.expect("expired lease should be taken over");
assert_eq!(takeover.fencing_token, stale.fencing_token + 1);
assert!(
!runtime
.release_coordination_lease(stale.key.clone(), stale.owner.clone(), stale.fencing_token)
.await
.unwrap()
);
let renew = runtime
.acquire_coordination_lease("rust-test:lease:renew".to_string(), "owner-1".to_string(), 30_000)
.await
.unwrap()
.expect("renew lease owner should acquire");
assert!(
!runtime
.renew_coordination_lease(renew.key.clone(), "owner-2".to_string(), renew.fencing_token, 30_000)
.await
.unwrap()
);
assert!(
!runtime
.renew_coordination_lease(renew.key.clone(), renew.owner.clone(), renew.fencing_token + 1, 30_000)
.await
.unwrap()
);
assert!(
runtime
.renew_coordination_lease(renew.key.clone(), renew.owner.clone(), renew.fencing_token, 30_000)
.await
.unwrap()
);
}
#[tokio::test]
async fn runtime_state_cleanup_deletes_expired_and_consumed_rows() {
let _guard = pg_test_lock().lock().await;
let Some(runtime) = runtime_from_database_url().await.unwrap() else {
eprintln!("skipping postgres integration test: DATABASE_URL is not set");
return;
};
assert!(
runtime
.create_auth_challenge(
"rust_test:cleanup".to_string(),
"expired".to_string(),
serde_json::json!({}),
1
)
.await
.unwrap()
);
assert!(
runtime
.create_auth_challenge(
"rust_test:cleanup".to_string(),
"consumed".to_string(),
serde_json::json!({}),
30_000,
)
.await
.unwrap()
);
assert!(
runtime
.consume_auth_challenge("rust_test:cleanup".to_string(), "consumed".to_string())
.await
.unwrap()
.is_some()
);
tokio::time::sleep(Duration::from_millis(20)).await;
assert_eq!(runtime.cleanup_expired_runtime_states(100).await.unwrap(), 2);
assert_eq!(runtime.cleanup_expired_runtime_states(100).await.unwrap(), 0);
}
#[tokio::test]
async fn verification_token_sql_state_machine_handles_keep_verify_and_cleanup() {
let _guard = pg_test_lock().lock().await;
let Some(runtime) = runtime_from_database_url().await.unwrap() else {
eprintln!("skipping postgres integration test: DATABASE_URL is not set");
return;
};
let mismatch_token = runtime
.create_verification_token(
TEST_VERIFICATION_TOKEN_TYPE,
Some("user@affine.test".to_string()),
30_000,
)
.await
.unwrap();
assert!(
runtime
.verify_verification_token(
TEST_VERIFICATION_TOKEN_TYPE,
mismatch_token.clone(),
Some("wrong@affine.test".to_string()),
None,
)
.await
.unwrap()
.is_none()
);
assert!(
runtime
.verify_verification_token(
TEST_VERIFICATION_TOKEN_TYPE,
mismatch_token.clone(),
Some("user@affine.test".to_string()),
None,
)
.await
.unwrap()
.is_some()
);
assert!(
runtime
.verify_verification_token(
TEST_VERIFICATION_TOKEN_TYPE,
mismatch_token.clone(),
Some("user@affine.test".to_string()),
None,
)
.await
.unwrap()
.is_none()
);
let keep_token = runtime
.create_verification_token(
TEST_VERIFICATION_TOKEN_TYPE,
Some("keep@affine.test".to_string()),
30_000,
)
.await
.unwrap();
assert!(
runtime
.get_verification_token(TEST_VERIFICATION_TOKEN_TYPE, keep_token.clone(), Some(true))
.await
.unwrap()
.is_some()
);
assert!(
runtime
.get_verification_token(TEST_VERIFICATION_TOKEN_TYPE, keep_token.clone(), None)
.await
.unwrap()
.is_some()
);
assert!(
runtime
.get_verification_token(TEST_VERIFICATION_TOKEN_TYPE, keep_token.clone(), None)
.await
.unwrap()
.is_none()
);
let concurrent_token = runtime
.create_verification_token(
TEST_VERIFICATION_TOKEN_TYPE,
Some("concurrent@affine.test".to_string()),
30_000,
)
.await
.unwrap();
let mut tasks = Vec::new();
for _ in 0..16 {
let runtime = BackendRuntime {
config: runtime.config.clone(),
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
};
let token = concurrent_token.clone();
tasks.push(tokio::spawn(async move {
runtime
.verify_verification_token(
TEST_VERIFICATION_TOKEN_TYPE,
token,
Some("concurrent@affine.test".to_string()),
None,
)
.await
.unwrap()
.is_some()
}));
}
let mut successful = 0;
for task in tasks {
if task.await.unwrap() {
successful += 1;
}
}
assert_eq!(successful, 1);
let expired_token = runtime
.create_verification_token(TEST_VERIFICATION_TOKEN_TYPE, Some("expired@affine.test".to_string()), 1)
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(20)).await;
assert!(
runtime
.get_verification_token(TEST_VERIFICATION_TOKEN_TYPE, expired_token.clone(), None)
.await
.unwrap()
.is_none()
);
assert_eq!(runtime.cleanup_expired_verification_tokens(100).await.unwrap(), 1);
assert_eq!(runtime.cleanup_expired_verification_tokens(100).await.unwrap(), 0);
}
@@ -0,0 +1,177 @@
use napi::bindgen_prelude::Buffer;
#[napi_derive::napi(object)]
pub struct RuntimeVerificationTokenRecord {
pub token_type: i32,
pub token: String,
pub credential: Option<String>,
pub expires_at_ms: i64,
}
#[napi_derive::napi(object)]
pub struct BackendRuntimeHealth {
pub started: bool,
pub database_connected: bool,
pub object_storage_configured: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeObjectStorageHealth {
pub configured: bool,
pub provider: Option<String>,
pub bucket: Option<String>,
pub endpoint: Option<String>,
pub region: Option<String>,
pub has_credentials: bool,
pub force_path_style: bool,
pub request_timeout_ms: Option<i64>,
pub min_part_size: Option<i64>,
pub presign_expires_in_seconds: Option<i64>,
pub presign_sign_content_type_for_put: Option<bool>,
pub use_presigned_url: bool,
pub client_buildable: bool,
}
#[napi_derive::napi(object)]
pub struct CoordinationLeaseGrant {
pub key: String,
pub owner: String,
#[napi(ts_type = "bigint | number")]
pub fencing_token: i64,
}
#[napi_derive::napi(object)]
pub struct RuntimeMagicLinkOtpConsumeResult {
pub ok: bool,
pub token: Option<String>,
pub reason: Option<String>,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceInviteLinkRecord {
pub workspace_id: String,
pub invite_id: String,
pub inviter_user_id: String,
pub expires_at_ms: i64,
}
#[napi_derive::napi(object)]
pub struct RuntimeByokLocalLeaseRecord {
pub lease_id: String,
pub payload: serde_json::Value,
pub expires_at_ms: i64,
}
#[napi_derive::napi(object)]
pub struct RuntimeDocHistoryInput {
pub workspace_id: String,
pub doc_id: String,
pub blob: Buffer,
pub timestamp_ms: i64,
pub editor_id: Option<String>,
pub force: bool,
pub history_min_interval_ms: i64,
pub history_max_age_ms: i64,
}
#[napi_derive::napi(object)]
pub struct RuntimeObjectStoragePutOptions {
pub content_type: Option<String>,
pub content_length: Option<i64>,
pub checksum_crc32: Option<String>,
}
#[napi_derive::napi(object)]
pub struct RuntimeObjectMetadata {
pub content_type: String,
pub content_length: i64,
pub last_modified_ms: i64,
pub checksum_crc32: Option<String>,
}
#[napi_derive::napi(object)]
pub struct RuntimeObjectListEntry {
pub key: String,
pub content_length: i64,
pub last_modified_ms: i64,
}
#[napi_derive::napi(object)]
pub struct RuntimeObjectGetResult {
pub body: Buffer,
pub metadata: RuntimeObjectMetadata,
}
#[napi_derive::napi(object)]
pub struct RuntimePresignedObjectRequest {
pub url: String,
pub headers_json: String,
pub expires_at_ms: i64,
}
#[napi_derive::napi(object)]
pub struct RuntimeMultipartUploadInit {
pub upload_id: String,
pub expires_at_ms: i64,
}
#[napi_derive::napi(object)]
pub struct RuntimeMultipartUploadPart {
pub part_number: i32,
pub etag: String,
}
#[napi_derive::napi(object)]
pub struct RuntimeBlobCleanupResult {
pub scanned: i64,
pub deleted: i64,
pub aborted_multipart: i64,
pub workspace_ids: Vec<String>,
}
#[napi_derive::napi(object)]
pub struct RuntimeBlobCompleteResult {
pub ok: bool,
pub reason: Option<String>,
pub content_type: Option<String>,
pub content_length: Option<i64>,
pub last_modified_ms: Option<i64>,
}
#[napi_derive::napi(object)]
pub struct RuntimeDocCompactionResult {
pub lease_acquired: bool,
pub merged: bool,
pub workspace_id: String,
pub doc_id: String,
pub updates_merged: i64,
pub history_created: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceStatsRefreshResult {
pub processed: i64,
pub backlog: i64,
pub skipped: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceStatsRecalibrationResult {
pub processed: i64,
pub last_sid: i64,
pub skipped: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceStatsSnapshotResult {
pub snapshotted: i64,
pub skipped: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceStatsDailyRecalibrationResult {
pub processed: i64,
pub last_sid: i64,
pub snapshotted: i64,
pub skipped: bool,
}
@@ -0,0 +1,527 @@
use napi::Result;
use sqlx::{FromRow, PgPool, Postgres, Row, Transaction};
use tokio::time::{Duration as TokioDuration, sleep};
use super::{
BackendRuntime,
constants::{WORKSPACE_STATS_LEASE_KEY, WORKSPACE_STATS_LOCK_NAMESPACE, WORKSPACE_STATS_REFRESH_LOCK_KEY},
error::napi_error,
types::{
CoordinationLeaseGrant, RuntimeWorkspaceStatsDailyRecalibrationResult, RuntimeWorkspaceStatsRecalibrationResult,
RuntimeWorkspaceStatsRefreshResult, RuntimeWorkspaceStatsSnapshotResult,
},
};
const UPSERT_WORKSPACE_ADMIN_STATS_SQL: &str = include_str!("sql/upsert_workspace_admin_stats.sql");
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn refresh_workspace_admin_stats_dirty(
&self,
batch_limit: i64,
owner: String,
lease_ttl_ms: i64,
) -> Result<RuntimeWorkspaceStatsRefreshResult> {
if batch_limit <= 0 {
return Err(napi_error("workspace stats dirty refresh limit must be positive"));
}
let Some(lease) = self
.acquire_coordination_lease(WORKSPACE_STATS_LEASE_KEY.to_string(), owner, lease_ttl_ms)
.await?
else {
return Ok(RuntimeWorkspaceStatsRefreshResult {
processed: 0,
backlog: 0,
skipped: true,
});
};
let result = async {
WorkspaceStatsStore::new(self.pool().await?)
.refresh_dirty(batch_limit)
.await
}
.await;
release_workspace_stats_lease(self, lease).await?;
result
}
#[napi]
pub async fn recalibrate_workspace_admin_stats(
&self,
last_sid: i64,
batch_limit: i64,
owner: String,
lease_ttl_ms: i64,
) -> Result<RuntimeWorkspaceStatsRecalibrationResult> {
if batch_limit <= 0 {
return Err(napi_error("workspace stats recalibration limit must be positive"));
}
let Some(lease) = self
.acquire_coordination_lease(WORKSPACE_STATS_LEASE_KEY.to_string(), owner, lease_ttl_ms)
.await?
else {
return Ok(RuntimeWorkspaceStatsRecalibrationResult {
processed: 0,
last_sid,
skipped: true,
});
};
let result = async {
WorkspaceStatsStore::new(self.pool().await?)
.recalibrate(last_sid, batch_limit)
.await
}
.await;
release_workspace_stats_lease(self, lease).await?;
result
}
#[napi]
pub async fn write_workspace_admin_stats_daily_snapshot(
&self,
owner: String,
lease_ttl_ms: i64,
) -> Result<RuntimeWorkspaceStatsSnapshotResult> {
let Some(lease) = self
.acquire_coordination_lease(WORKSPACE_STATS_LEASE_KEY.to_string(), owner, lease_ttl_ms)
.await?
else {
return Ok(RuntimeWorkspaceStatsSnapshotResult {
snapshotted: 0,
skipped: true,
});
};
let result = async {
WorkspaceStatsStore::new(self.pool().await?)
.write_daily_snapshot()
.await
}
.await;
release_workspace_stats_lease(self, lease).await?;
result
}
#[napi]
pub async fn recalibrate_workspace_admin_stats_daily(
&self,
batch_limit: i64,
owner: String,
lease_ttl_ms: i64,
lock_retry_times: i64,
lock_retry_delay_ms: i64,
) -> Result<RuntimeWorkspaceStatsDailyRecalibrationResult> {
if batch_limit <= 0 {
return Err(napi_error("workspace stats daily recalibration limit must be positive"));
}
if lock_retry_times <= 0 {
return Err(napi_error(
"workspace stats daily recalibration retry times must be positive",
));
}
if lock_retry_delay_ms < 0 {
return Err(napi_error(
"workspace stats daily recalibration retry delay must be non-negative",
));
}
let Some(lease) = acquire_workspace_stats_lease_with_retry(
self,
owner.clone(),
lease_ttl_ms,
lock_retry_times,
lock_retry_delay_ms,
)
.await?
else {
return Ok(RuntimeWorkspaceStatsDailyRecalibrationResult {
processed: 0,
last_sid: 0,
snapshotted: 0,
skipped: true,
});
};
let result = async {
let store = WorkspaceStatsStore::new(self.pool().await?);
let mut processed = 0;
let mut last_sid = 0;
loop {
let batch = retry_workspace_stats_operation(lock_retry_times, lock_retry_delay_ms, || {
store.recalibrate(last_sid, batch_limit)
})
.await?;
if batch.skipped {
return Ok(RuntimeWorkspaceStatsDailyRecalibrationResult {
processed,
last_sid,
snapshotted: 0,
skipped: true,
});
}
if batch.processed == 0 {
break;
}
processed += batch.processed;
last_sid = batch.last_sid;
if batch.processed < batch_limit {
break;
}
}
let snapshot =
retry_workspace_stats_operation(lock_retry_times, lock_retry_delay_ms, || store.write_daily_snapshot()).await?;
Ok(RuntimeWorkspaceStatsDailyRecalibrationResult {
processed,
last_sid,
snapshotted: snapshot.snapshotted,
skipped: snapshot.skipped,
})
}
.await;
release_workspace_stats_lease(self, lease).await?;
result
}
}
#[derive(FromRow)]
struct WorkspaceSid {
id: String,
sid: i32,
}
struct WorkspaceStatsStore {
pool: PgPool,
}
impl WorkspaceStatsStore {
fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn refresh_dirty(&self, batch_limit: i64) -> Result<RuntimeWorkspaceStatsRefreshResult> {
let mut tx = self
.pool
.begin()
.await
.map_err(|err| napi_error(format!("WorkspaceStats dirty refresh transaction failed: {err}")))?;
if !try_transaction_lock(&mut tx).await? {
tx.commit()
.await
.map_err(|err| napi_error(format!("WorkspaceStats dirty refresh commit failed: {err}")))?;
return Ok(RuntimeWorkspaceStatsRefreshResult {
processed: 0,
backlog: 0,
skipped: true,
});
}
let backlog = count_dirty(&mut tx).await?;
let dirty = load_dirty(&mut tx, batch_limit).await?;
if dirty.is_empty() {
tx.commit()
.await
.map_err(|err| napi_error(format!("WorkspaceStats dirty refresh commit failed: {err}")))?;
return Ok(RuntimeWorkspaceStatsRefreshResult {
processed: 0,
backlog,
skipped: false,
});
}
upsert_stats(&mut tx, &dirty).await?;
clear_dirty(&mut tx, &dirty).await?;
tx.commit()
.await
.map_err(|err| napi_error(format!("WorkspaceStats dirty refresh commit failed: {err}")))?;
Ok(RuntimeWorkspaceStatsRefreshResult {
processed: dirty.len() as i64,
backlog,
skipped: false,
})
}
async fn recalibrate(&self, last_sid: i64, batch_limit: i64) -> Result<RuntimeWorkspaceStatsRecalibrationResult> {
let mut tx = self
.pool
.begin()
.await
.map_err(|err| napi_error(format!("WorkspaceStats recalibration transaction failed: {err}")))?;
if !try_transaction_lock(&mut tx).await? {
tx.commit()
.await
.map_err(|err| napi_error(format!("WorkspaceStats recalibration commit failed: {err}")))?;
return Ok(RuntimeWorkspaceStatsRecalibrationResult {
processed: 0,
last_sid,
skipped: true,
});
}
let workspaces = fetch_workspace_batch(&mut tx, last_sid, batch_limit).await?;
if workspaces.is_empty() {
tx.commit()
.await
.map_err(|err| napi_error(format!("WorkspaceStats recalibration commit failed: {err}")))?;
return Ok(RuntimeWorkspaceStatsRecalibrationResult {
processed: 0,
last_sid,
skipped: false,
});
}
let ids = workspaces
.iter()
.map(|workspace| workspace.id.clone())
.collect::<Vec<_>>();
let next_sid = workspaces
.last()
.map(|workspace| workspace.sid as i64)
.unwrap_or(last_sid);
upsert_stats(&mut tx, &ids).await?;
tx.commit()
.await
.map_err(|err| napi_error(format!("WorkspaceStats recalibration commit failed: {err}")))?;
Ok(RuntimeWorkspaceStatsRecalibrationResult {
processed: ids.len() as i64,
last_sid: next_sid,
skipped: false,
})
}
async fn write_daily_snapshot(&self) -> Result<RuntimeWorkspaceStatsSnapshotResult> {
let mut tx = self
.pool
.begin()
.await
.map_err(|err| napi_error(format!("WorkspaceStats daily snapshot transaction failed: {err}")))?;
if !try_transaction_lock(&mut tx).await? {
tx.commit()
.await
.map_err(|err| napi_error(format!("WorkspaceStats daily snapshot commit failed: {err}")))?;
return Ok(RuntimeWorkspaceStatsSnapshotResult {
snapshotted: 0,
skipped: true,
});
}
let snapshotted = write_daily_snapshot(&mut tx).await?;
tx.commit()
.await
.map_err(|err| napi_error(format!("WorkspaceStats daily snapshot commit failed: {err}")))?;
Ok(RuntimeWorkspaceStatsSnapshotResult {
snapshotted,
skipped: false,
})
}
}
async fn release_workspace_stats_lease(runtime: &BackendRuntime, lease: CoordinationLeaseGrant) -> Result<()> {
let _ = runtime
.release_coordination_lease(lease.key, lease.owner, lease.fencing_token)
.await?;
Ok(())
}
async fn acquire_workspace_stats_lease_with_retry(
runtime: &BackendRuntime,
owner: String,
lease_ttl_ms: i64,
retry_times: i64,
retry_delay_ms: i64,
) -> Result<Option<CoordinationLeaseGrant>> {
for attempt in 0..retry_times {
let lease = runtime
.acquire_coordination_lease(WORKSPACE_STATS_LEASE_KEY.to_string(), owner.clone(), lease_ttl_ms)
.await?;
if lease.is_some() {
return Ok(lease);
}
if attempt < retry_times - 1 && retry_delay_ms > 0 {
sleep(TokioDuration::from_millis(retry_delay_ms as u64)).await;
}
}
Ok(None)
}
async fn retry_workspace_stats_operation<T, F, Fut>(
retry_times: i64,
retry_delay_ms: i64,
mut operation: F,
) -> Result<T>
where
T: WorkspaceStatsSkippable,
F: FnMut() -> Fut,
Fut: std::future::Future<Output = Result<T>>,
{
for attempt in 0..retry_times {
let result = operation().await?;
if !result.skipped() || attempt == retry_times - 1 {
return Ok(result);
}
if retry_delay_ms > 0 {
sleep(TokioDuration::from_millis(retry_delay_ms as u64)).await;
}
}
unreachable!("workspace stats retry loop validates retry_times > 0")
}
trait WorkspaceStatsSkippable {
fn skipped(&self) -> bool;
}
impl WorkspaceStatsSkippable for RuntimeWorkspaceStatsRecalibrationResult {
fn skipped(&self) -> bool {
self.skipped
}
}
impl WorkspaceStatsSkippable for RuntimeWorkspaceStatsSnapshotResult {
fn skipped(&self) -> bool {
self.skipped
}
}
async fn try_transaction_lock(tx: &mut Transaction<'_, Postgres>) -> Result<bool> {
let row = sqlx::query(
r#"
SELECT pg_try_advisory_xact_lock(($1::bigint << 32) + $2::bigint) AS locked
"#,
)
.bind(WORKSPACE_STATS_LOCK_NAMESPACE)
.bind(WORKSPACE_STATS_REFRESH_LOCK_KEY)
.fetch_one(&mut **tx)
.await
.map_err(|err| napi_error(format!("WorkspaceStats transaction lock failed: {err}")))?;
Ok(row.get::<bool, _>("locked"))
}
async fn load_dirty(tx: &mut Transaction<'_, Postgres>, limit: i64) -> Result<Vec<String>> {
let rows = sqlx::query(
r#"
SELECT workspace_id
FROM workspace_admin_stats_dirty
ORDER BY updated_at ASC
LIMIT $1
FOR UPDATE SKIP LOCKED
"#,
)
.bind(limit)
.fetch_all(&mut **tx)
.await
.map_err(|err| napi_error(format!("WorkspaceStats load dirty workspaces failed: {err}")))?;
Ok(rows.into_iter().map(|row| row.get("workspace_id")).collect())
}
async fn count_dirty(tx: &mut Transaction<'_, Postgres>) -> Result<i64> {
let row = sqlx::query("SELECT COUNT(*) AS total FROM workspace_admin_stats_dirty")
.fetch_one(&mut **tx)
.await
.map_err(|err| napi_error(format!("WorkspaceStats count dirty workspaces failed: {err}")))?;
Ok(row.get::<i64, _>("total"))
}
async fn clear_dirty(tx: &mut Transaction<'_, Postgres>, workspace_ids: &[String]) -> Result<()> {
sqlx::query(
r#"
DELETE FROM workspace_admin_stats_dirty
WHERE workspace_id = ANY($1::varchar[])
"#,
)
.bind(workspace_ids)
.execute(&mut **tx)
.await
.map_err(|err| napi_error(format!("WorkspaceStats clear dirty workspaces failed: {err}")))?;
Ok(())
}
async fn upsert_stats(tx: &mut Transaction<'_, Postgres>, workspace_ids: &[String]) -> Result<()> {
if workspace_ids.is_empty() {
return Ok(());
}
sqlx::query(UPSERT_WORKSPACE_ADMIN_STATS_SQL)
.bind(workspace_ids)
.execute(&mut **tx)
.await
.map_err(|err| napi_error(format!("WorkspaceStats upsert stats failed: {err}")))?;
Ok(())
}
async fn fetch_workspace_batch(
tx: &mut Transaction<'_, Postgres>,
last_sid: i64,
limit: i64,
) -> Result<Vec<WorkspaceSid>> {
sqlx::query_as::<_, WorkspaceSid>(
r#"
SELECT id, sid
FROM workspaces
WHERE sid > $1
ORDER BY sid
LIMIT $2
"#,
)
.bind(last_sid)
.bind(limit)
.fetch_all(&mut **tx)
.await
.map_err(|err| napi_error(format!("WorkspaceStats fetch workspace batch failed: {err}")))
}
async fn write_daily_snapshot(tx: &mut Transaction<'_, Postgres>) -> Result<i64> {
let result = sqlx::query(
r#"
INSERT INTO workspace_admin_stats_daily (
workspace_id,
date,
snapshot_size,
blob_size,
member_count,
updated_at
)
SELECT
workspace_id,
CURRENT_DATE,
snapshot_size,
blob_size,
member_count,
NOW()
FROM workspace_admin_stats
ON CONFLICT (workspace_id, date)
DO UPDATE SET
snapshot_size = EXCLUDED.snapshot_size,
blob_size = EXCLUDED.blob_size,
member_count = EXCLUDED.member_count,
updated_at = EXCLUDED.updated_at
"#,
)
.execute(&mut **tx)
.await
.map_err(|err| napi_error(format!("WorkspaceStats daily snapshot failed: {err}")))?;
Ok(result.rows_affected() as i64)
}
+1
View File
@@ -2,6 +2,7 @@
mod utils;
pub mod backend_runtime;
pub mod doc;
pub mod doc_loader;
pub mod entitlement;
+2 -5
View File
@@ -26,11 +26,8 @@ fn try_remove_label(s: &str, i: usize) -> Option<usize> {
return None;
}
if let Some(ch) = s[next_idx..].chars().next() {
if !ch.is_whitespace() {
return None;
}
} else {
let ch = s[next_idx..].chars().next()?;
if !ch.is_whitespace() {
return None;
}
@@ -0,0 +1,39 @@
CREATE TABLE "runtime_states" (
"purpose" TEXT NOT NULL,
"token_hash" TEXT NOT NULL,
"lookup_key" TEXT,
"payload" JSONB NOT NULL,
"attempts" INTEGER NOT NULL DEFAULT 0,
"consumed_at" TIMESTAMPTZ(3),
"expires_at" TIMESTAMPTZ(3) NOT NULL,
"created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "runtime_states_pkey" PRIMARY KEY ("purpose", "token_hash")
);
CREATE INDEX "runtime_states_lookup_idx" ON "runtime_states"("purpose", "lookup_key") WHERE "lookup_key" IS NOT NULL AND "consumed_at" IS NULL;
CREATE INDEX "runtime_states_expires_at_idx" ON "runtime_states"("expires_at");
CREATE TABLE "runtime_gates" (
"key" TEXT NOT NULL,
"expires_at" TIMESTAMPTZ(3) NOT NULL,
"created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "runtime_gates_pkey" PRIMARY KEY ("key")
);
CREATE INDEX "runtime_gates_expires_at_idx" ON "runtime_gates"("expires_at");
CREATE TABLE "runtime_leases" (
"key" TEXT NOT NULL,
"owner" TEXT NOT NULL,
"fencing_token" BIGINT NOT NULL,
"expires_at" TIMESTAMPTZ(3) NOT NULL,
"created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "runtime_leases_pkey" PRIMARY KEY ("key")
);
CREATE INDEX "runtime_leases_expires_at_idx" ON "runtime_leases"("expires_at");
+35
View File
@@ -2,6 +2,7 @@ import serverNativeModule, {
type ActionEvent as NativeActionEventContract,
type ActionRuntimeInput as NativeActionRuntimeInputContract,
type AssertSafeUrlRequest,
type BackendRuntimeHealth,
type BuiltInPromptRenderContract,
type BuiltInPromptSessionContract,
type BuiltInPromptSpec,
@@ -45,6 +46,22 @@ import serverNativeModule, {
type RequestedModelMatchResponse,
type ResolvedEntitlement,
type ResolveEntitlementInput,
type RuntimeBlobCleanupResult,
type RuntimeBlobCompleteResult,
type RuntimeByokLocalLeaseRecord,
type RuntimeDocCompactionResult,
type RuntimeMagicLinkOtpConsumeResult,
type RuntimeMultipartUploadInit,
type RuntimeMultipartUploadPart,
type RuntimeObjectGetResult,
type RuntimeObjectListEntry,
type RuntimeObjectMetadata,
type RuntimeObjectStorageHealth,
type RuntimeObjectStoragePutOptions,
type RuntimePresignedObjectRequest,
type RuntimeVerificationTokenRecord,
type RuntimeWorkspaceInviteLinkRecord,
type RuntimeWorkspaceStatsDailyRecalibrationResult,
type SafeFetchRequest,
type SafeFetchResponse,
type Tokenizer,
@@ -52,6 +69,7 @@ import serverNativeModule, {
export type {
AssertSafeUrlRequest,
BackendRuntimeHealth,
CapabilityAttachmentContract,
CapabilityModelCapability,
CommandResponse,
@@ -73,6 +91,22 @@ export type {
RemoteMimeTypeRequest,
ResolvedEntitlement,
ResolveEntitlementInput,
RuntimeBlobCleanupResult,
RuntimeBlobCompleteResult,
RuntimeByokLocalLeaseRecord,
RuntimeDocCompactionResult,
RuntimeMagicLinkOtpConsumeResult,
RuntimeMultipartUploadInit,
RuntimeMultipartUploadPart,
RuntimeObjectGetResult,
RuntimeObjectListEntry,
RuntimeObjectMetadata,
RuntimeObjectStorageHealth,
RuntimeObjectStoragePutOptions,
RuntimePresignedObjectRequest,
RuntimeVerificationTokenRecord,
RuntimeWorkspaceInviteLinkRecord,
RuntimeWorkspaceStatsDailyRecalibrationResult,
SafeFetchRequest,
SafeFetchResponse,
};
@@ -180,6 +214,7 @@ export const readAllDocIdsFromRootDoc =
export const AFFINE_PRO_PUBLIC_KEY = serverNativeModule.AFFINE_PRO_PUBLIC_KEY;
export const AFFINE_PRO_LICENSE_AES_KEY =
serverNativeModule.AFFINE_PRO_LICENSE_AES_KEY;
export const BackendRuntime = serverNativeModule.BackendRuntime;
export type PermissionWorkspaceRole = 'external' | 'member' | 'admin' | 'owner';
export type PermissionDocRole =
+1 -1
View File
@@ -61,7 +61,7 @@ impl Stamp {
let ts = now.format("%Y%m%d%H%M%S");
let bits = bits.unwrap_or(20);
let rand = String::from_iter(Alphanumeric.sample_iter(rng()).take(SALT_LENGTH).map(char::from));
let challenge = format!("{}:{}:{}:{}:{}:{}", version, bits, ts, &resource, "", rand);
let challenge = format!("{}:{}:{}:{}:{}:{}", version, bits, ts, resource, "", rand);
Stamp {
version: version.to_string(),