use graph_storage_sdk::models::EmbeddingSpaceId;
use sea_orm::{ActiveValue, ColumnTrait, Condition, EntityTrait};
use toolkit_db::secure::{Db, SecureEntityExt, SecureInsertExt};
use toolkit_security::AccessScope;
use uuid::Uuid;
use crate::infra::storage::entity::embedding_space;
use crate::infra::storage::migrations::m0003_embedding_space::{FIRST_EPOCH, STATE_ACTIVE};
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SpaceResolution {
Active { epoch: i64 },
Mismatched {
recorded_identity: String,
recorded_epoch: i64,
},
}
fn deployment_scope() -> AccessScope {
AccessScope::for_tenant(Uuid::nil())
}
pub async fn resolve(db: &Db, space: &EmbeddingSpaceId) -> anyhow::Result<SpaceResolution> {
let scope = deployment_scope();
let conn = db.conn()?;
if let Some(row) = active_row(&conn, &scope).await? {
return Ok(decide(row, space));
}
if let Err(error) = open_first_space(&conn, &scope, space).await {
return match active_row(&conn, &scope).await? {
Some(row) => Ok(decide(row, space)),
None => Err(error),
};
}
Ok(SpaceResolution::Active { epoch: FIRST_EPOCH })
}
fn decide(row: embedding_space::Model, space: &EmbeddingSpaceId) -> SpaceResolution {
if row.identity_hash == space.identity_hash {
SpaceResolution::Active { epoch: row.epoch }
} else {
SpaceResolution::Mismatched {
recorded_identity: row.identity_hash,
recorded_epoch: row.epoch,
}
}
}
async fn active_row(
conn: &impl toolkit_db::secure::DBRunner,
scope: &AccessScope,
) -> anyhow::Result<Option<embedding_space::Model>> {
Ok(embedding_space::Entity::find()
.secure()
.scope_with(scope)
.filter(Condition::all().add(embedding_space::Column::State.eq(STATE_ACTIVE)))
.one(conn)
.await?)
}
async fn open_first_space(
conn: &impl toolkit_db::secure::DBRunner,
scope: &AccessScope,
space: &EmbeddingSpaceId,
) -> anyhow::Result<()> {
let now = time::OffsetDateTime::now_utc();
let active = embedding_space::ActiveModel {
tenant_id: ActiveValue::Set(Uuid::nil()),
epoch: ActiveValue::Set(FIRST_EPOCH),
identity_hash: ActiveValue::Set(space.identity_hash.clone()),
model_artifact: ActiveValue::Set(space.model_artifact.clone()),
tokenizer_artifact: ActiveValue::Set(space.tokenizer_artifact.clone()),
preprocessing: ActiveValue::Set(space.preprocessing.clone()),
pooling: ActiveValue::Set(space.pooling.clone()),
normalization: ActiveValue::Set(space.normalization.clone()),
dimension: ActiveValue::Set(i32::try_from(space.dimension).unwrap_or(i32::MAX)),
state: ActiveValue::Set(STATE_ACTIVE.to_owned()),
created_at: ActiveValue::Set(now),
activated_at: ActiveValue::Set(Some(now)),
};
embedding_space::Entity::insert(active)
.secure()
.scope_unchecked(scope)?
.exec(conn)
.await?;
Ok(())
}