mirror of
https://github.com/element-hq/synapse.git
synced 2026-08-14 15:50:19 +00:00
Add backend-agnostic Rust-native query helpers over the pool
Introduce a small enum-based layer so simple queries can be written once regardless of backend: DbPool/DbConn dispatch over each backend's own pool, and DbConn::query(sql, params) binds positional `?` placeholders and returns rows as backend-agnostic DbValue cells (read out with DbRowExt::try_get). No trait/async_trait machinery — a closed enum keeps the pool.get().await? / conn.query(...).await? surface simple. The Postgres arm rewrites `?` to `$1, $2, ...` (matching convert_param_style), reuses PgValue for parameter binding, and adds a Rust-native DbValueFromSql decoder (the non-Python counterpart of PythonPgFromSql). The shared column-type list is factored into accepts_column_type so the three mappings can't drift. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Erik Johnston
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
31d55d7887
commit
2a4ad76bda
+139
-3
@@ -1,14 +1,68 @@
|
||||
//! DBAPI2-shaped Connection / Cursor types implemented in Rust.
|
||||
//! Rust database access.
|
||||
//!
|
||||
//! Currently this provides a single Postgres backend ([`tokio_postgres`]); see
|
||||
//! the [`postgres`] submodule.
|
||||
//! Two layers live here:
|
||||
//! - the DBAPI2-shaped `Connection`/`Cursor` shim (see the [`postgres`]
|
||||
//! submodule), which lets existing Python transaction functions run against a
|
||||
//! Rust connection unchanged, and
|
||||
//! - the backend-agnostic Rust-native query helpers ([`DbPool`]/[`DbConn`] plus
|
||||
//! [`value::DbValue`]), so simple queries can be written once and run against
|
||||
//! any backend.
|
||||
//!
|
||||
//! Currently the only backend is Postgres ([`tokio_postgres`]); SQLite is
|
||||
//! expected to follow as a sibling variant of [`DbPool`]/[`DbConn`].
|
||||
|
||||
pub mod postgres;
|
||||
pub mod runtime;
|
||||
pub mod value;
|
||||
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::PyModule;
|
||||
|
||||
use crate::database::value::{DbRow, DbValue};
|
||||
|
||||
/// A backend-agnostic connection pool.
|
||||
///
|
||||
/// A thin enum over each backend's own pool rather than a trait: with a small,
|
||||
/// closed set of backends this keeps the ergonomic `pool.get().await?` /
|
||||
/// `conn.query(...).await?` surface without `dyn`/`async_trait` machinery.
|
||||
pub enum DbPool {
|
||||
Postgres(postgres::pool::ConnectionPool),
|
||||
}
|
||||
|
||||
impl DbPool {
|
||||
/// Check a connection out of the pool.
|
||||
pub async fn get(&self) -> Result<DbConn, anyhow::Error> {
|
||||
match self {
|
||||
DbPool::Postgres(pool) => Ok(DbConn::Postgres(pool.get().await?)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A backend-agnostic connection checked out of a [`DbPool`].
|
||||
///
|
||||
/// Returned to the pool when dropped.
|
||||
pub enum DbConn {
|
||||
Postgres(postgres::pool::PooledConnection),
|
||||
}
|
||||
|
||||
impl DbConn {
|
||||
/// Run a query and return all resulting rows.
|
||||
///
|
||||
/// Parameters are bound positionally to `?` placeholders (Synapse's existing
|
||||
/// convention); each backend rewrites them to its own placeholder syntax.
|
||||
/// Result cells come back as [`DbValue`]s, read out with
|
||||
/// [`value::DbRowExt::try_get`].
|
||||
pub async fn query(
|
||||
&mut self,
|
||||
sql: &str,
|
||||
params: &[DbValue],
|
||||
) -> Result<Vec<DbRow>, anyhow::Error> {
|
||||
match self {
|
||||
DbConn::Postgres(conn) => postgres::query::query(conn, sql, params).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Register the `database` submodule (and its per-backend children) on the
|
||||
/// top-level `synapse_rust` module.
|
||||
pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
@@ -26,3 +80,85 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()>
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
//! These exercise the end-to-end query helper against a live Postgres, so
|
||||
//! they only run when `SYNAPSE_TEST_POSTGRES_DSN` is set (e.g. to
|
||||
//! `host=postgres user=postgres password=postgres dbname=postgres`);
|
||||
//! otherwise they no-op. The placeholder rewriting and value mapping have
|
||||
//! their own server-free unit tests.
|
||||
|
||||
use super::*;
|
||||
use crate::database::runtime::runtime;
|
||||
use crate::database::value::DbRowExt;
|
||||
|
||||
fn test_dsn() -> Option<String> {
|
||||
std::env::var("SYNAPSE_TEST_POSTGRES_DSN").ok()
|
||||
}
|
||||
|
||||
fn test_pool(dsn: &str) -> DbPool {
|
||||
DbPool::Postgres(postgres::pool::create_pool(dsn, 2).unwrap())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn query_binds_placeholders_and_decodes_scalar_rows() {
|
||||
let Some(dsn) = test_dsn() else {
|
||||
eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run");
|
||||
return;
|
||||
};
|
||||
let pool = test_pool(&dsn);
|
||||
|
||||
runtime().block_on(async {
|
||||
let mut conn = pool.get().await.unwrap();
|
||||
let rows = conn
|
||||
.query(
|
||||
"SELECT ?::int AS a, ?::text AS b, ?::bool AS c",
|
||||
&[
|
||||
DbValue::Int(7),
|
||||
DbValue::Text("hi".into()),
|
||||
DbValue::Bool(true),
|
||||
],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(rows[0].try_get::<i64>(0).unwrap(), 7);
|
||||
assert_eq!(rows[0].try_get::<String>(1).unwrap(), "hi");
|
||||
assert!(rows[0].try_get::<bool>(2).unwrap());
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn query_decodes_nulls_bytes_and_multiple_rows() {
|
||||
let Some(dsn) = test_dsn() else {
|
||||
eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run");
|
||||
return;
|
||||
};
|
||||
let pool = test_pool(&dsn);
|
||||
|
||||
runtime().block_on(async {
|
||||
let mut conn = pool.get().await.unwrap();
|
||||
|
||||
// A NULL cell reads back as `None`; bytea round-trips byte-for-byte.
|
||||
let rows = conn
|
||||
.query(
|
||||
"SELECT NULL::text AS a, ?::bytea AS b",
|
||||
&[DbValue::Bytes(vec![0, 255])],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(rows[0].try_get::<Option<String>>(0).unwrap(), None);
|
||||
assert_eq!(rows[0].try_get::<Vec<u8>>(1).unwrap(), vec![0u8, 255]);
|
||||
|
||||
// Several rows come back in order.
|
||||
let rows = conn
|
||||
.query("SELECT g FROM generate_series(1, 3) AS g ORDER BY g", &[])
|
||||
.await
|
||||
.unwrap();
|
||||
let got: Vec<i64> = rows.iter().map(|r| r.try_get::<i64>(0).unwrap()).collect();
|
||||
assert_eq!(got, vec![1, 2, 3]);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -18,6 +18,7 @@ mod errors;
|
||||
mod helpers;
|
||||
mod libpq;
|
||||
pub mod pool;
|
||||
pub(crate) mod query;
|
||||
mod value;
|
||||
|
||||
/// Register the `postgres` submodule (the `Connection` / `Cursor` classes, the
|
||||
|
||||
@@ -13,7 +13,7 @@
|
||||
//! *same* pool, so both share a single set of connections rather than running
|
||||
//! two pools that could exhaust the server's connection limit between them.
|
||||
|
||||
use deadpool::managed::{Manager, Metrics, Pool, RecycleError, RecycleResult};
|
||||
use deadpool::managed::{Manager, Metrics, Object, Pool, RecycleError, RecycleResult};
|
||||
use log::warn;
|
||||
use tokio_postgres::{Client, Config, NoTls};
|
||||
|
||||
@@ -71,6 +71,10 @@ impl Manager for ConnectionManager {
|
||||
/// A pool of [`tokio_postgres`] connections.
|
||||
pub type ConnectionPool = Pool<ConnectionManager>;
|
||||
|
||||
/// A connection checked out of a [`ConnectionPool`]. Dereferences to the
|
||||
/// underlying [`tokio_postgres::Client`]; returned to the pool when dropped.
|
||||
pub type PooledConnection = Object<ConnectionManager>;
|
||||
|
||||
/// Build a [`ConnectionPool`] from a libpq-style DSN, capped at `max_size`
|
||||
/// connections.
|
||||
pub fn create_pool(dsn: &str, max_size: usize) -> Result<ConnectionPool, anyhow::Error> {
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
//! The Rust-native `query` helper for the Postgres backend.
|
||||
//!
|
||||
//! This is the Postgres arm of [`crate::database::DbConn::query`]: it binds
|
||||
//! backend-agnostic [`DbValue`] parameters and decodes rows back into
|
||||
//! [`DbRow`]s, so simple queries can be written once regardless of backend.
|
||||
|
||||
use tokio_postgres::types::ToSql;
|
||||
use tokio_postgres::Client;
|
||||
|
||||
use crate::database::postgres::value::{pg_row_to_db_row, PgValue};
|
||||
use crate::database::value::{DbRow, DbValue};
|
||||
|
||||
/// Run `sql` (with `?` placeholders bound to `params`) and return all rows.
|
||||
///
|
||||
/// Parameters are bound positionally, in order, to the `?` placeholders — which
|
||||
/// are first rewritten to Postgres' `$1, $2, ...` (see [`convert_placeholders`]).
|
||||
pub(crate) async fn query(
|
||||
client: &Client,
|
||||
sql: &str,
|
||||
params: &[DbValue],
|
||||
) -> Result<Vec<DbRow>, anyhow::Error> {
|
||||
let sql = convert_placeholders(sql);
|
||||
|
||||
// Bind each agnostic `DbValue` as a Postgres parameter. `PgValue` is the
|
||||
// `ToSql` type the rest of the backend already uses.
|
||||
let pg_params: Vec<PgValue> = params.iter().map(PgValue::from_db_value).collect();
|
||||
let params_dyn: Vec<&(dyn ToSql + Sync)> =
|
||||
pg_params.iter().map(|v| v as &(dyn ToSql + Sync)).collect();
|
||||
|
||||
// Passing the SQL as a `&str` lets `tokio_postgres` prepare it (inferring
|
||||
// the parameter types our `ToSql` impl needs) and run it in one call.
|
||||
let rows = client.query(sql.as_str(), ¶ms_dyn).await?;
|
||||
|
||||
rows.iter().map(pg_row_to_db_row).collect()
|
||||
}
|
||||
|
||||
/// Rewrite `?` placeholders to Postgres' positional `$1, $2, ...`.
|
||||
///
|
||||
/// A naive left-to-right substitution, matching Synapse's existing
|
||||
/// `PostgresEngine.convert_param_style` (`?` -> `%s`): it does not skip a `?`
|
||||
/// inside a string literal, so callers should parameterise rather than embed a
|
||||
/// literal `?` in the SQL.
|
||||
fn convert_placeholders(sql: &str) -> String {
|
||||
let mut out = String::with_capacity(sql.len() + 8);
|
||||
let mut n = 0u32;
|
||||
for ch in sql.chars() {
|
||||
if ch == '?' {
|
||||
n += 1;
|
||||
out.push('$');
|
||||
out.push_str(&n.to_string());
|
||||
} else {
|
||||
out.push(ch);
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn convert_placeholders_numbers_each_question_mark() {
|
||||
assert_eq!(
|
||||
convert_placeholders("SELECT * FROM t WHERE a = ? AND b = ?"),
|
||||
"SELECT * FROM t WHERE a = $1 AND b = $2"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_placeholders_leaves_sql_without_params_untouched() {
|
||||
assert_eq!(convert_placeholders("SELECT 1"), "SELECT 1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_placeholders_counts_past_nine() {
|
||||
// The index is decimal, not a single digit.
|
||||
let sql = "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)";
|
||||
let converted = convert_placeholders(sql);
|
||||
assert!(converted.ends_with("$9, $10, $11)"), "{converted}");
|
||||
}
|
||||
}
|
||||
@@ -28,13 +28,16 @@ use std::error::Error;
|
||||
|
||||
use bytes::BytesMut;
|
||||
use postgres_protocol::types::{
|
||||
bool_to_sql, bytea_to_sql, float4_to_sql, float8_to_sql, int2_to_sql, int4_to_sql, int8_to_sql,
|
||||
text_to_sql,
|
||||
bool_from_sql, bool_to_sql, bytea_to_sql, float4_from_sql, float4_to_sql, float8_from_sql,
|
||||
float8_to_sql, int2_from_sql, int2_to_sql, int4_from_sql, int4_to_sql, int8_from_sql,
|
||||
int8_to_sql, text_to_sql,
|
||||
};
|
||||
use pyo3::exceptions::{PyTypeError, PyValueError};
|
||||
use pyo3::types::{PyBool, PyBytes, PyFloat, PyInt, PyString, PyTuple};
|
||||
use pyo3::{prelude::*, BoundObject};
|
||||
use tokio_postgres::types::{to_sql_checked, IsNull, ToSql, Type, WrongType};
|
||||
use tokio_postgres::types::{to_sql_checked, FromSql, IsNull, ToSql, Type, WrongType};
|
||||
|
||||
use crate::database::value::{DbRow, DbValue};
|
||||
|
||||
/// Owned representation of a Python value that we can hand to [`tokio_postgres`]
|
||||
/// as a [`ToSql`] parameter.
|
||||
@@ -86,6 +89,42 @@ impl PgValue {
|
||||
obj.get_type().name()?,
|
||||
)))
|
||||
}
|
||||
|
||||
/// Bind a backend-agnostic [`DbValue`] (from the Rust-native query helpers)
|
||||
/// as a Postgres parameter. This is the non-Python counterpart of
|
||||
/// [`PgValue::from_py`].
|
||||
pub(crate) fn from_db_value(value: &DbValue) -> Self {
|
||||
match value {
|
||||
DbValue::Null => PgValue::Null,
|
||||
DbValue::Bool(b) => PgValue::Bool(*b),
|
||||
DbValue::Int(i) => PgValue::Int(*i),
|
||||
DbValue::Float(f) => PgValue::Float(*f),
|
||||
DbValue::Text(s) => PgValue::Text(s.as_str().into()),
|
||||
DbValue::Bytes(b) => PgValue::Bytea(b.as_slice().into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The Postgres column types the value mapping supports, in both directions.
|
||||
///
|
||||
/// Shared by every `accepts` implementation here ([`PgValue`]'s `ToSql`,
|
||||
/// [`PythonPgFromSql`] and [`DbValueFromSql`]) and kept in sync with their
|
||||
/// `match` arms, so the supported set can't silently drift between them.
|
||||
pub(crate) fn accepts_column_type(ty: &Type) -> bool {
|
||||
matches!(
|
||||
*ty,
|
||||
Type::BOOL
|
||||
| Type::INT2
|
||||
| Type::INT4
|
||||
| Type::INT8
|
||||
| Type::FLOAT4
|
||||
| Type::FLOAT8
|
||||
| Type::TEXT
|
||||
| Type::VARCHAR
|
||||
| Type::NAME
|
||||
| Type::BPCHAR
|
||||
| Type::BYTEA
|
||||
)
|
||||
}
|
||||
|
||||
// Lets PyO3 extract a `PgValue` directly from a Python argument, e.g. when a
|
||||
@@ -168,20 +207,7 @@ impl ToSql for PgValue {
|
||||
}
|
||||
|
||||
fn accepts(ty: &Type) -> bool {
|
||||
matches!(
|
||||
*ty,
|
||||
Type::BOOL
|
||||
| Type::INT2
|
||||
| Type::INT4
|
||||
| Type::INT8
|
||||
| Type::FLOAT4
|
||||
| Type::FLOAT8
|
||||
| Type::TEXT
|
||||
| Type::VARCHAR
|
||||
| Type::NAME
|
||||
| Type::BPCHAR
|
||||
| Type::BYTEA
|
||||
)
|
||||
accepts_column_type(ty)
|
||||
}
|
||||
|
||||
to_sql_checked!();
|
||||
@@ -229,20 +255,7 @@ impl<'a> tokio_postgres::types::FromSql<'a> for PythonPgFromSql {
|
||||
}
|
||||
|
||||
fn accepts(ty: &Type) -> bool {
|
||||
matches!(
|
||||
*ty,
|
||||
Type::BOOL
|
||||
| Type::INT2
|
||||
| Type::INT4
|
||||
| Type::INT8
|
||||
| Type::FLOAT4
|
||||
| Type::FLOAT8
|
||||
| Type::TEXT
|
||||
| Type::VARCHAR
|
||||
| Type::NAME
|
||||
| Type::BPCHAR
|
||||
| Type::BYTEA
|
||||
)
|
||||
accepts_column_type(ty)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -294,6 +307,63 @@ impl PythonPgFromSql {
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a Postgres row into a backend-agnostic [`DbRow`] for the Rust-native
|
||||
/// query helpers, one [`DbValue`] per column.
|
||||
///
|
||||
/// The non-Python counterpart of [`pg_row_to_py`]; decodes each column via
|
||||
/// [`DbValueFromSql`]. Errors (with the column index and type) if a column can't
|
||||
/// be decoded.
|
||||
pub(crate) fn pg_row_to_db_row(row: &tokio_postgres::Row) -> Result<DbRow, anyhow::Error> {
|
||||
let mut out = Vec::with_capacity(row.len());
|
||||
for idx in 0..row.len() {
|
||||
let value: DbValueFromSql = row.try_get(idx).map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"failed to decode column {idx} (type {}): {e}",
|
||||
row.columns()[idx].type_()
|
||||
)
|
||||
})?;
|
||||
out.push(value.0);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// A column value decoded into a backend-agnostic [`DbValue`].
|
||||
///
|
||||
/// The non-Python counterpart of [`PythonPgFromSql`]: the same supported types
|
||||
/// (see [`accepts_column_type`]), but with no GIL and no Python objects — just
|
||||
/// the plain Rust value.
|
||||
pub(crate) struct DbValueFromSql(pub DbValue);
|
||||
|
||||
impl<'a> FromSql<'a> for DbValueFromSql {
|
||||
fn from_sql(ty: &Type, raw: &'a [u8]) -> Result<Self, Box<dyn Error + Sync + Send>> {
|
||||
let value = match *ty {
|
||||
Type::BOOL => DbValue::Bool(bool_from_sql(raw)?),
|
||||
Type::INT2 => DbValue::Int(int2_from_sql(raw)?.into()),
|
||||
Type::INT4 => DbValue::Int(int4_from_sql(raw)?.into()),
|
||||
Type::INT8 => DbValue::Int(int8_from_sql(raw)?),
|
||||
Type::FLOAT4 => DbValue::Float(float4_from_sql(raw)?.into()),
|
||||
Type::FLOAT8 => DbValue::Float(float8_from_sql(raw)?),
|
||||
Type::TEXT | Type::VARCHAR | Type::NAME | Type::BPCHAR => {
|
||||
DbValue::Text(std::str::from_utf8(raw)?.to_owned())
|
||||
}
|
||||
Type::BYTEA => DbValue::Bytes(raw.to_vec()),
|
||||
_ => {
|
||||
// Unreachable unless `accepts` drifts out of sync with this match.
|
||||
return Err(format!("unsupported column type for postgres: {ty}").into());
|
||||
}
|
||||
};
|
||||
Ok(DbValueFromSql(value))
|
||||
}
|
||||
|
||||
fn from_sql_null(_ty: &Type) -> Result<Self, Box<dyn Error + Sync + Send>> {
|
||||
Ok(DbValueFromSql(DbValue::Null))
|
||||
}
|
||||
|
||||
fn accepts(ty: &Type) -> bool {
|
||||
accepts_column_type(ty)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
//! These tests exercise the value mapping in isolation — no live Postgres
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
//! Backend-agnostic value and row types for the Rust-native query helpers.
|
||||
//!
|
||||
//! [`DbValue`] is the common currency for both query *parameters* and result
|
||||
//! *cells*, so simple queries can be written once and run against any backend
|
||||
//! (currently Postgres; SQLite to follow). Each backend maps its driver's own
|
||||
//! value types to and from [`DbValue`] at its edge — see
|
||||
//! `postgres::value` for the Postgres mapping.
|
||||
//!
|
||||
//! The set is deliberately small (the common scalar types); richer types can be
|
||||
//! added as the query helpers grow to need them.
|
||||
|
||||
/// A single backend-agnostic value, used both as a query parameter and as a
|
||||
/// cell in a returned [`DbRow`].
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum DbValue {
|
||||
/// SQL `NULL`.
|
||||
Null,
|
||||
Bool(bool),
|
||||
Int(i64),
|
||||
Float(f64),
|
||||
Text(String),
|
||||
Bytes(Vec<u8>),
|
||||
}
|
||||
|
||||
// Ergonomic conversions so callers can write `&[1i64.into(), "x".into()]` (or
|
||||
// pass typed values straight through) rather than spelling out `DbValue`.
|
||||
impl From<bool> for DbValue {
|
||||
fn from(v: bool) -> Self {
|
||||
DbValue::Bool(v)
|
||||
}
|
||||
}
|
||||
impl From<i64> for DbValue {
|
||||
fn from(v: i64) -> Self {
|
||||
DbValue::Int(v)
|
||||
}
|
||||
}
|
||||
impl From<i32> for DbValue {
|
||||
fn from(v: i32) -> Self {
|
||||
DbValue::Int(v as i64)
|
||||
}
|
||||
}
|
||||
impl From<f64> for DbValue {
|
||||
fn from(v: f64) -> Self {
|
||||
DbValue::Float(v)
|
||||
}
|
||||
}
|
||||
impl From<String> for DbValue {
|
||||
fn from(v: String) -> Self {
|
||||
DbValue::Text(v)
|
||||
}
|
||||
}
|
||||
impl From<&str> for DbValue {
|
||||
fn from(v: &str) -> Self {
|
||||
DbValue::Text(v.to_owned())
|
||||
}
|
||||
}
|
||||
impl From<Vec<u8>> for DbValue {
|
||||
fn from(v: Vec<u8>) -> Self {
|
||||
DbValue::Bytes(v)
|
||||
}
|
||||
}
|
||||
impl From<&[u8]> for DbValue {
|
||||
fn from(v: &[u8]) -> Self {
|
||||
DbValue::Bytes(v.to_vec())
|
||||
}
|
||||
}
|
||||
impl<T: Into<DbValue>> From<Option<T>> for DbValue {
|
||||
fn from(v: Option<T>) -> Self {
|
||||
match v {
|
||||
Some(v) => v.into(),
|
||||
None => DbValue::Null,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A row of data returned by a query: one [`DbValue`] per column, read out by
|
||||
/// numeric index with [`DbRowExt::try_get`].
|
||||
pub type DbRow = Vec<DbValue>;
|
||||
|
||||
/// Reads a typed value out of a [`DbValue`], analogous to `tokio_postgres`'s
|
||||
/// `FromSql`. `Option<T>` reads a nullable column (SQL `NULL` -> `None`).
|
||||
pub trait FromDbValue: Sized {
|
||||
fn from_value(value: DbValue) -> Result<Self, anyhow::Error>;
|
||||
}
|
||||
|
||||
impl FromDbValue for bool {
|
||||
fn from_value(value: DbValue) -> Result<Self, anyhow::Error> {
|
||||
match value {
|
||||
DbValue::Bool(b) => Ok(b),
|
||||
// SQLite has no native boolean type and stores them as integers.
|
||||
DbValue::Int(i) => Ok(i != 0),
|
||||
other => anyhow::bail!("cannot read {other:?} as bool"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FromDbValue for i64 {
|
||||
fn from_value(value: DbValue) -> Result<Self, anyhow::Error> {
|
||||
match value {
|
||||
DbValue::Int(i) => Ok(i),
|
||||
DbValue::Bool(b) => Ok(b as i64),
|
||||
other => anyhow::bail!("cannot read {other:?} as i64"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FromDbValue for f64 {
|
||||
fn from_value(value: DbValue) -> Result<Self, anyhow::Error> {
|
||||
match value {
|
||||
DbValue::Float(f) => Ok(f),
|
||||
DbValue::Int(i) => Ok(i as f64),
|
||||
other => anyhow::bail!("cannot read {other:?} as f64"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FromDbValue for String {
|
||||
fn from_value(value: DbValue) -> Result<Self, anyhow::Error> {
|
||||
match value {
|
||||
DbValue::Text(s) => Ok(s),
|
||||
other => anyhow::bail!("cannot read {other:?} as String"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FromDbValue for Vec<u8> {
|
||||
fn from_value(value: DbValue) -> Result<Self, anyhow::Error> {
|
||||
match value {
|
||||
DbValue::Bytes(b) => Ok(b),
|
||||
other => anyhow::bail!("cannot read {other:?} as bytes"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: FromDbValue> FromDbValue for Option<T> {
|
||||
fn from_value(value: DbValue) -> Result<Self, anyhow::Error> {
|
||||
match value {
|
||||
DbValue::Null => Ok(None),
|
||||
other => Ok(Some(T::from_value(other)?)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Extension methods for reading typed values out of a [`DbRow`].
|
||||
pub trait DbRowExt {
|
||||
/// Read the value at `index`, converting it into `T` via [`FromDbValue`].
|
||||
/// Errors if the index is out of bounds or the value can't become `T`.
|
||||
fn try_get<T: FromDbValue>(&self, index: usize) -> Result<T, anyhow::Error>;
|
||||
}
|
||||
|
||||
impl DbRowExt for DbRow {
|
||||
fn try_get<T: FromDbValue>(&self, index: usize) -> Result<T, anyhow::Error> {
|
||||
let value = self.get(index).cloned().ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"tried to read column {index} but the row only has {} column(s)",
|
||||
self.len()
|
||||
)
|
||||
})?;
|
||||
T::from_value(value)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn from_impls_classify_scalars() {
|
||||
assert_eq!(DbValue::from(true), DbValue::Bool(true));
|
||||
assert_eq!(DbValue::from(7i64), DbValue::Int(7));
|
||||
assert_eq!(DbValue::from(7i32), DbValue::Int(7));
|
||||
assert_eq!(DbValue::from(1.5f64), DbValue::Float(1.5));
|
||||
assert_eq!(DbValue::from("x"), DbValue::Text("x".to_owned()));
|
||||
assert_eq!(DbValue::from(vec![1u8, 2]), DbValue::Bytes(vec![1, 2]));
|
||||
// `Option` folds `None` to NULL and `Some` to the inner conversion.
|
||||
assert_eq!(DbValue::from(None::<i64>), DbValue::Null);
|
||||
assert_eq!(DbValue::from(Some(3i64)), DbValue::Int(3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn try_get_reads_typed_values() {
|
||||
let row: DbRow = vec![
|
||||
DbValue::Int(42),
|
||||
DbValue::Text("hi".to_owned()),
|
||||
DbValue::Null,
|
||||
];
|
||||
assert_eq!(row.try_get::<i64>(0).unwrap(), 42);
|
||||
assert_eq!(row.try_get::<String>(1).unwrap(), "hi");
|
||||
// A NULL read as `Option` is `None`...
|
||||
assert_eq!(row.try_get::<Option<String>>(2).unwrap(), None);
|
||||
// ...but read as a non-optional type it's an error, not a silent default.
|
||||
assert!(row.try_get::<String>(2).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn try_get_out_of_bounds_errors() {
|
||||
let row: DbRow = vec![DbValue::Int(1)];
|
||||
let err = row.try_get::<i64>(5).unwrap_err();
|
||||
assert!(err.to_string().contains("only has 1 column"), "{err}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bool_reads_from_sqlite_style_integer() {
|
||||
// SQLite returns booleans as integers; `FromDbValue for bool` bridges it.
|
||||
assert!(bool::from_value(DbValue::Int(1)).unwrap());
|
||||
assert!(!bool::from_value(DbValue::Int(0)).unwrap());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user