little-durable-objects 0.1.2

use anyhow::{Context, Result};
use async_trait::async_trait;

use super::{HostLease, HostLeaseRegistry, HostLeaseRequest, HostLeaseStatus, HostLeaseStore};
use crate::{host::HostId, postgres::PostgresDatabase};

const REGISTRY_NOW_MS: &str = "(extract(epoch FROM clock_timestamp()) * 1000)::bigint";

pub struct PostgresHostLeaseStore {
    database: PostgresDatabase,
}

impl PostgresHostLeaseStore {
    pub async fn connect(url: &str) -> Result<Self> {
        Ok(Self::from_database(PostgresDatabase::connect(url).await?))
    }

    pub(crate) fn from_database(database: PostgresDatabase) -> Self {
        Self { database }
    }
}

#[async_trait]
impl HostLeaseRegistry for PostgresHostLeaseStore {
    async fn register(&self, request: &HostLeaseRequest) -> Result<HostLease> {
        request.validate_duration()?;
        let duration_ms = i64::try_from(request.duration_ms)
            .expect("validated host lease duration must fit PostgreSQL BIGINT");
        let row = self
            .database
            .client()
            .query_opt(
                &format!(
                    "INSERT INTO durable_object_host_leases (host_id, session_id, route, expires_at_ms) \
                     VALUES ($1, $2, $3, {REGISTRY_NOW_MS} + $4) \
                     ON CONFLICT (host_id) DO UPDATE \
                     SET session_id = EXCLUDED.session_id, route = EXCLUDED.route, \
                         expires_at_ms = EXCLUDED.expires_at_ms \
                     WHERE durable_object_host_leases.session_id = EXCLUDED.session_id \
                        OR durable_object_host_leases.expires_at_ms <= {REGISTRY_NOW_MS} \
                     RETURNING expires_at_ms"
                ),
                &[
                    &request.id.as_str(),
                    &request.session_id,
                    &request.route,
                    &duration_ms,
                ],
            )
            .await
            .context("register PostgreSQL host lease")?
            .context("host lease is held by a different active session")?;
        Ok(HostLease {
            id: request.id.clone(),
            session_id: request.session_id.clone(),
            route: request.route.clone(),
            expires_at_ms: u64::try_from(row.get::<_, i64>(0))
                .context("PostgreSQL host lease expiration is negative")?,
        })
    }

    async fn unregister(&self, id: &HostId, session_id: &str) -> Result<()> {
        self.database
            .client()
            .execute(
                "DELETE FROM durable_object_host_leases WHERE host_id = $1 AND session_id = $2",
                &[&id.as_str(), &session_id],
            )
            .await
            .context("unregister PostgreSQL host lease")?;
        Ok(())
    }
}

#[async_trait]
impl HostLeaseStore for PostgresHostLeaseStore {
    async fn lease_status(&self, id: &HostId) -> Result<HostLeaseStatus> {
        let row = self
            .database
            .client()
            .query_one(
                &format!(
                    "SELECT {REGISTRY_NOW_MS} AS store_now_ms, lease.session_id, \
                            lease.route, lease.expires_at_ms \
                     FROM (SELECT 1) AS clock_row \
                     LEFT JOIN durable_object_host_leases AS lease ON lease.host_id = $1"
                ),
                &[&id.as_str()],
            )
            .await
            .context("read PostgreSQL host lease status")?;
        let store_now_ms = u64::try_from(row.get::<_, i64>(0))
            .context("PostgreSQL lease-store clock is before the Unix epoch")?;
        let lease = row
            .get::<_, Option<String>>(1)
            .map(|session_id| {
                Ok::<_, anyhow::Error>(HostLease {
                    id: id.clone(),
                    session_id,
                    route: row
                        .get::<_, Option<String>>(2)
                        .context("PostgreSQL host lease row is missing its route")?,
                    expires_at_ms: u64::try_from(
                        row.get::<_, Option<i64>>(3)
                            .context("PostgreSQL host lease row is missing its expiration")?,
                    )
                    .context("PostgreSQL host lease expiration is negative")?,
                })
            })
            .transpose()?;
        Ok(HostLeaseStatus {
            lease,
            store_now_ms,
        })
    }
}