rullst-orm 13.0.0-alpha.2

//! SurrealDB HTTP adapter.
//!
//! The adapter intentionally uses the public HTTP protocol instead of
//! embedding the SurrealDB SDK. This keeps the MIT-licensed framework from
//! implicitly redistributing the database engine's BSL dependency.

use std::{marker::PhantomData, time::Duration};

use async_trait::async_trait;
use reqwest::{Client, Method, Url, redirect::Policy};
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::Value;

use super::{
    Backend, BackendCapabilities, Capability, CollectionName, DocumentEntry, DocumentId,
    DocumentInventory, DocumentPage, DocumentRepository, MockDocumentStore, PolyglotError,
};

mod config;
mod graph;
mod live;
#[cfg(test)]
mod tests;

pub use config::{SurrealAuth, SurrealConfig};
use config::{decode_surreal_record_key, is_mock_credential, validate_endpoint};
pub use graph::{GraphQuery, GraphRepository};

struct LiveSurreal {
    client: Client,
    endpoint: Url,
    namespace: CollectionName,
    database: CollectionName,
    auth: SurrealAuth,
    response_limit: usize,
}

#[async_trait]
impl<T> DocumentInventory<T> for SurrealDbStore<T>
where
    T: Serialize + DeserializeOwned + Send + Sync,
{
    async fn list_entries(
        &self,
        collection: &CollectionName,
        page: DocumentPage,
    ) -> Result<Vec<DocumentEntry<T>>, PolyglotError> {
        let SurrealDbInner::Live { live, .. } = &self.inner else {
            return surreal_mock(self)?.list_entries(collection, page).await;
        };
        let mut url = live.route(&["sql"])?;
        url.query_pairs_mut()
            .append_pair("table", collection.as_str())
            .append_pair("start", &page.offset().to_string())
            .append_pair("limit", &page.limit().to_string());
        let envelopes = live
            .send_text(
                Method::POST,
                &url,
                "SELECT * FROM type::table($table) ORDER BY id START type::int($start) LIMIT type::int($limit)",
            )
            .await?;
        let values = result_values(statement_result(envelopes, false)?)?;
        if values.len() > page.limit() as usize {
            return Err(invalid_surreal_response(
                "server exceeded the requested page limit",
            ));
        }
        values
            .into_iter()
            .map(|value| decode_document_entry(collection, value))
            .collect()
    }
}

enum SurrealDbInner<T> {
    Live {
        live: LiveSurreal,
        marker: PhantomData<fn() -> T>,
    },
    Mock(MockDocumentStore<T>),
}

/// SurrealDB multi-model adapter with deterministic mock fallback.
pub struct SurrealDbStore<T> {
    inner: SurrealDbInner<T>,
}

impl<T> SurrealDbStore<T> {
    /// Builds the HTTP adapter or selects the deterministic fallback.
    pub fn connect_or_mock(config: SurrealConfig) -> Result<Self, PolyglotError> {
        let namespace = CollectionName::new(config.namespace)?;
        let database = CollectionName::new(config.database)?;
        if is_mock_credential(&config.endpoint) || config.auth.requests_mock() {
            return Ok(Self {
                inner: SurrealDbInner::Mock(MockDocumentStore::new()),
            });
        }
        let endpoint = validate_endpoint(&config.endpoint, config.allow_insecure_http)?;
        let client = Client::builder()
            .redirect(Policy::none())
            .timeout(Duration::from_secs(30))
            .build()
            .map_err(|error| PolyglotError::driver("SurrealDB", error))?;
        Ok(Self {
            inner: SurrealDbInner::Live {
                live: LiveSurreal {
                    client,
                    endpoint,
                    namespace,
                    database,
                    auth: config.auth,
                    response_limit: config.response_limit,
                },
                marker: PhantomData,
            },
        })
    }

    /// Declares the adapter's bounded behavior.
    pub const fn capabilities() -> BackendCapabilities {
        BackendCapabilities::new(
            Backend::SurrealDb,
            &[Capability::Documents, Capability::Graph],
        )
    }

    /// Returns whether this instance is using the deterministic offline store.
    pub const fn is_mock(&self) -> bool {
        matches!(&self.inner, SurrealDbInner::Mock(_))
    }
}

#[async_trait]
impl<T> DocumentRepository<T> for SurrealDbStore<T>
where
    T: Serialize + DeserializeOwned + Send + Sync,
{
    async fn create(
        &self,
        collection: &CollectionName,
        id: &DocumentId,
        entity: &T,
    ) -> Result<(), PolyglotError> {
        let SurrealDbInner::Live { live, .. } = &self.inner else {
            return surreal_mock(self)?.create(collection, id, entity).await;
        };
        let body = encode_document(entity)?;
        let envelopes = live
            .send(
                Method::POST,
                &record_route(live, collection, id)?,
                Some(body),
            )
            .await?;
        statement_result(envelopes, true).map(|_| ())
    }

    async fn find(
        &self,
        collection: &CollectionName,
        id: &DocumentId,
    ) -> Result<Option<T>, PolyglotError> {
        let SurrealDbInner::Live { live, .. } = &self.inner else {
            return surreal_mock(self)?.find(collection, id).await;
        };
        let envelopes = live
            .send(Method::GET, &record_route(live, collection, id)?, None)
            .await?;
        let mut values = result_values(statement_result(envelopes, false)?)?;
        if values.len() > 1 {
            return Err(invalid_surreal_response(
                "single-record lookup returned multiple rows",
            ));
        }
        values.pop().map(decode_document).transpose()
    }

    async fn replace(
        &self,
        collection: &CollectionName,
        id: &DocumentId,
        entity: &T,
    ) -> Result<(), PolyglotError> {
        let SurrealDbInner::Live { live, .. } = &self.inner else {
            return surreal_mock(self)?.replace(collection, id, entity).await;
        };
        let body = encode_document(entity)?;
        // `PUT /key/...` upserts, which would recreate a deleted document.
        // `UPDATE` never creates a record; `type::record` (SurrealDB 3) with
        // the string `$id` addresses the same key as `record_route`.
        let envelopes = live
            .rpc_query(
                "UPDATE type::record($table, $id) CONTENT $data RETURN AFTER",
                serde_json::json!({
                    "table": collection.as_str(),
                    "id": id.as_str(),
                    "data": body,
                }),
            )
            .await?;
        if result_values(statement_result(envelopes, false)?)?.is_empty() {
            return Err(PolyglotError::NotFound);
        }
        Ok(())
    }

    async fn delete(
        &self,
        collection: &CollectionName,
        id: &DocumentId,
    ) -> Result<bool, PolyglotError> {
        let SurrealDbInner::Live { live, .. } = &self.inner else {
            return surreal_mock(self)?.delete(collection, id).await;
        };
        let envelopes = live
            .send(Method::DELETE, &record_route(live, collection, id)?, None)
            .await?;
        Ok(!result_values(statement_result(envelopes, false)?)?.is_empty())
    }

    async fn list(
        &self,
        collection: &CollectionName,
        page: DocumentPage,
    ) -> Result<Vec<T>, PolyglotError> {
        let SurrealDbInner::Live { live, .. } = &self.inner else {
            return surreal_mock(self)?.list(collection, page).await;
        };
        let mut url = live.route(&["sql"])?;
        url.query_pairs_mut()
            .append_pair("table", collection.as_str())
            .append_pair("start", &page.offset().to_string())
            .append_pair("limit", &page.limit().to_string());
        let envelopes = live
            .send_text(
                Method::POST,
                &url,
                "SELECT * FROM type::table($table) ORDER BY id START type::int($start) LIMIT type::int($limit)",
            )
            .await?;
        let values = result_values(statement_result(envelopes, false)?)?;
        if values.len() > page.limit() as usize {
            return Err(invalid_surreal_response(
                "server exceeded the requested page limit",
            ));
        }
        values.into_iter().map(decode_document).collect()
    }
}

#[async_trait]
impl<T> GraphRepository for SurrealDbStore<T>
where
    T: Send + Sync,
{
    async fn query_graph(&self, query: &GraphQuery) -> Result<Vec<Value>, PolyglotError> {
        let SurrealDbInner::Live { live, .. } = &self.inner else {
            return Ok(Vec::new());
        };
        let envelopes = live
            .send_text(Method::POST, &live.route(&["gql"])?, &query.bounded_query())
            .await?;
        let values = result_values(statement_result(envelopes, false)?)?;
        if values.len() > query.limit as usize {
            return Err(invalid_surreal_response(
                "server exceeded the graph row limit",
            ));
        }
        Ok(values)
    }
}

#[derive(Debug, Deserialize)]
struct StatementEnvelope {
    status: String,
    #[serde(default)]
    result: Value,
}

fn statement_result(
    mut envelopes: Vec<StatementEnvelope>,
    conflict_on_duplicate: bool,
) -> Result<Value, PolyglotError> {
    if envelopes.len() != 1 {
        return Err(invalid_surreal_response(
            "expected exactly one statement result",
        ));
    }
    let envelope = envelopes.remove(0);
    if envelope.status != "OK" {
        let detail = envelope
            .result
            .as_str()
            .unwrap_or_default()
            .to_ascii_lowercase();
        if conflict_on_duplicate && detail.contains("already exists") {
            return Err(PolyglotError::Conflict);
        }
        return Err(PolyglotError::Driver {
            backend: "SurrealDB",
            message: format!("statement returned status {}", envelope.status),
        });
    }
    Ok(envelope.result)
}

fn result_values(result: Value) -> Result<Vec<Value>, PolyglotError> {
    match result {
        Value::Array(values) => Ok(values),
        Value::Null => Ok(Vec::new()),
        _ => Err(invalid_surreal_response(
            "statement result was not an array",
        )),
    }
}

fn encode_document<T: Serialize>(entity: &T) -> Result<Value, PolyglotError> {
    let value = serde_json::to_value(entity).map_err(PolyglotError::serialization)?;
    let Some(object) = value.as_object() else {
        return Err(invalid_surreal_response(
            "portable models must serialize as JSON objects",
        ));
    };
    if object.contains_key("id") {
        return Err(PolyglotError::InvalidIdentifier {
            kind: "SurrealDB model",
            reason: "the portable model must not define an id field",
        });
    }
    Ok(value)
}

fn decode_document<T: DeserializeOwned>(mut value: Value) -> Result<T, PolyglotError> {
    let Some(object) = value.as_object_mut() else {
        return Err(invalid_surreal_response(
            "document result was not a JSON object",
        ));
    };
    object.remove("id");
    serde_json::from_value(value).map_err(PolyglotError::serialization)
}

fn decode_document_entry<T: DeserializeOwned>(
    collection: &CollectionName,
    mut value: Value,
) -> Result<DocumentEntry<T>, PolyglotError> {
    let Some(object) = value.as_object_mut() else {
        return Err(invalid_surreal_response(
            "document inventory row was not a JSON object",
        ));
    };
    let record_id = object
        .remove("id")
        .and_then(|id| id.as_str().map(ToOwned::to_owned))
        .ok_or_else(|| {
            invalid_surreal_response("document inventory returned a non-portable identifier")
        })?;
    let expected_prefix = format!("{}:", collection.as_str());
    let encoded_id = record_id.strip_prefix(&expected_prefix).ok_or_else(|| {
        invalid_surreal_response("document inventory identifier crossed its collection boundary")
    })?;
    let id = decode_surreal_record_key(encoded_id)?;
    let entity = serde_json::from_value(value).map_err(PolyglotError::serialization)?;
    Ok(DocumentEntry::new(id, entity))
}

/// The `/key/:table/:id` route of one document. SurrealDB parses the `:id`
/// segment as a value, so an all-digit ID such as `9` would become a numeric
/// key that sorts numerically, unlike the portable string order of
/// `DocumentId`. The ID is therefore sent as a quoted string literal; the
/// `DocumentId` grammar contains no quote or backslash to escape.
fn record_route(
    live: &LiveSurreal,
    collection: &CollectionName,
    id: &DocumentId,
) -> Result<Url, PolyglotError> {
    let key = format!("\"{}\"", id.as_str());
    live.route(&["key", collection.as_str(), &key])
}

fn surreal_mock<T>(store: &SurrealDbStore<T>) -> Result<&MockDocumentStore<T>, PolyglotError> {
    match store {
        SurrealDbStore {
            inner: SurrealDbInner::Mock(store),
        } => Ok(store),
        SurrealDbStore {
            inner: SurrealDbInner::Live { .. },
        } => Err(invalid_surreal_response(
            "adapter state changed during operation",
        )),
    }
}

fn invalid_surreal_response(message: &'static str) -> PolyglotError {
    PolyglotError::Driver {
        backend: "SurrealDB",
        message: message.to_owned(),
    }
}

fn response_too_large(limit_bytes: usize) -> PolyglotError {
    PolyglotError::ResponseTooLarge {
        backend: "SurrealDB",
        limit_bytes,
    }
}