use crate::connection::{DefaultLoadingMode, LoadConnection};
use crate::deserialize::FromSqlRow;
use crate::expression::QueryMetadata;
use crate::query_builder::{AstPass, Query, QueryFragment, QueryId};
use crate::result::{Error, QueryResult};
use crate::sql_types::SqlType;
use crate::sqlite::{AutoVacuumMode, Sqlite, SqliteConnection};
use alloc::string::ToString;
use alloc::vec::Vec;
use core::marker::PhantomData;
use core::num::NonZeroU32;
pub trait Pragma: Sized {
const NAME: &'static str;
type Scope: PragmaScope;
}
pub trait PragmaScope: private::Sealed {}
#[derive(Debug, Clone, Copy)]
pub enum PerSchema {}
#[derive(Debug, Clone, Copy)]
pub enum PerConnection {}
impl PragmaScope for PerSchema {}
impl PragmaScope for PerConnection {}
mod private {
use crate::query_builder::AstPass;
use crate::result::QueryResult;
use crate::sqlite::Sqlite;
pub trait Sealed {}
impl Sealed for super::PerSchema {}
impl Sealed for super::PerConnection {}
pub trait RenderArgument {
fn render(&self, out: &mut AstPass<'_, '_, Sqlite>) -> QueryResult<()>;
}
}
pub trait ReadPragma: Pragma {
type SqlType: SqlType;
}
pub trait WritePragma: Pragma {
type Argument<'a>: PragmaArgument;
}
pub trait PragmaArgument: private::RenderArgument {}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PragmaKeyword(&'static str);
impl PragmaKeyword {
pub const fn new(keyword: &'static str) -> Self {
Self(keyword)
}
}
impl private::RenderArgument for PragmaKeyword {
fn render(&self, out: &mut AstPass<'_, '_, Sqlite>) -> QueryResult<()> {
push_bare_word(out, "keyword", self.0)
}
}
impl PragmaArgument for PragmaKeyword {}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PragmaText<'a>(&'a str);
impl<'a> From<&'a str> for PragmaText<'a> {
fn from(text: &'a str) -> Self {
Self(text)
}
}
impl private::RenderArgument for PragmaText<'_> {
fn render(&self, out: &mut AstPass<'_, '_, Sqlite>) -> QueryResult<()> {
out.push_sql("'");
for piece in self.0.split_inclusive('\'') {
out.push_sql(piece);
if piece.ends_with('\'') {
out.push_sql("'");
}
}
out.push_sql("'");
Ok(())
}
}
impl PragmaArgument for PragmaText<'_> {}
macro_rules! integer_pragma_argument {
($($ty:ty),*) => {$(
impl private::RenderArgument for $ty {
fn render(&self, out: &mut AstPass<'_, '_, Sqlite>) -> QueryResult<()> {
out.push_sql(&i64::from(*self).to_string());
Ok(())
}
}
impl PragmaArgument for $ty {}
)*};
}
integer_pragma_argument!(i8, i16, i32, i64, u8, u16, u32, bool);
pub trait SchemaPragma: Pragma<Scope = PerSchema> {
fn get<Row>(schema: Option<&str>, conn: &mut SqliteConnection) -> QueryResult<Row>
where
Self: ReadPragma,
Sqlite: QueryMetadata<Self::SqlType>,
Row: FromSqlRow<Self::SqlType, Sqlite>,
{
first_row::<Self, Row>(conn, Some(schema.unwrap_or("main")))
}
fn get_all<Row>(schema: Option<&str>, conn: &mut SqliteConnection) -> QueryResult<Vec<Row>>
where
Self: ReadPragma,
Sqlite: QueryMetadata<Self::SqlType>,
Row: FromSqlRow<Self::SqlType, Sqlite>,
{
all_rows::<Self, Row>(conn, Some(schema.unwrap_or("main")))
}
fn set<'a>(
argument: impl Into<Self::Argument<'a>>,
schema: Option<&str>,
conn: &mut SqliteConnection,
) -> QueryResult<()>
where
Self: WritePragma,
{
assign::<Self>(conn, Some(schema.unwrap_or("main")), argument.into())
}
}
impl<P: Pragma<Scope = PerSchema>> SchemaPragma for P {}
pub trait ConnectionPragma: Pragma<Scope = PerConnection> {
fn get<Row>(conn: &mut SqliteConnection) -> QueryResult<Row>
where
Self: ReadPragma,
Sqlite: QueryMetadata<Self::SqlType>,
Row: FromSqlRow<Self::SqlType, Sqlite>,
{
first_row::<Self, Row>(conn, None)
}
fn get_all<Row>(conn: &mut SqliteConnection) -> QueryResult<Vec<Row>>
where
Self: ReadPragma,
Sqlite: QueryMetadata<Self::SqlType>,
Row: FromSqlRow<Self::SqlType, Sqlite>,
{
all_rows::<Self, Row>(conn, None)
}
fn set<'a>(
argument: impl Into<Self::Argument<'a>>,
conn: &mut SqliteConnection,
) -> QueryResult<()>
where
Self: WritePragma,
{
assign::<Self>(conn, None, argument.into())
}
}
impl<P: Pragma<Scope = PerConnection>> ConnectionPragma for P {}
#[derive(Debug, Clone, Copy)]
pub struct PageCount;
impl Pragma for PageCount {
const NAME: &'static str = "page_count";
type Scope = PerSchema;
}
impl ReadPragma for PageCount {
type SqlType = crate::sql_types::BigInt;
}
#[derive(Debug, Clone, Copy)]
pub struct FreelistCount;
impl Pragma for FreelistCount {
const NAME: &'static str = "freelist_count";
type Scope = PerSchema;
}
impl ReadPragma for FreelistCount {
type SqlType = crate::sql_types::BigInt;
}
#[derive(Debug, Clone, Copy)]
pub struct AutoVacuum;
impl Pragma for AutoVacuum {
const NAME: &'static str = "auto_vacuum";
type Scope = PerSchema;
}
impl ReadPragma for AutoVacuum {
type SqlType = crate::sql_types::Integer;
}
impl WritePragma for AutoVacuum {
type Argument<'a> = AutoVacuumMode;
}
impl private::RenderArgument for AutoVacuumMode {
fn render(&self, out: &mut AstPass<'_, '_, Sqlite>) -> QueryResult<()> {
out.push_sql(&(*self as i32).to_string());
Ok(())
}
}
impl PragmaArgument for AutoVacuumMode {}
#[derive(Debug, Clone, Copy)]
pub struct IncrementalVacuum;
impl Pragma for IncrementalVacuum {
const NAME: &'static str = "incremental_vacuum";
type Scope = PerSchema;
}
impl WritePragma for IncrementalVacuum {
type Argument<'a> = Option<NonZeroU32>;
}
impl private::RenderArgument for Option<NonZeroU32> {
fn render(&self, out: &mut AstPass<'_, '_, Sqlite>) -> QueryResult<()> {
out.push_sql(&self.map_or(0, NonZeroU32::get).to_string());
Ok(())
}
}
impl PragmaArgument for Option<NonZeroU32> {}
fn cursor<'conn, 'query, P>(
conn: &'conn mut SqliteConnection,
schema: Option<&'query str>,
) -> QueryResult<<SqliteConnection as LoadConnection<DefaultLoadingMode>>::Cursor<'conn, 'query>>
where
P: ReadPragma + 'query,
Sqlite: QueryMetadata<P::SqlType>,
{
LoadConnection::<DefaultLoadingMode>::load(conn, PragmaQuery::<P>::new(schema))
}
fn first_row<P, Row>(conn: &mut SqliteConnection, schema: Option<&str>) -> QueryResult<Row>
where
P: ReadPragma,
Sqlite: QueryMetadata<P::SqlType>,
Row: FromSqlRow<P::SqlType, Sqlite>,
{
let mut cursor = cursor::<P>(conn, schema)?;
let row = cursor.next().ok_or(Error::NotFound)??;
Row::build_from_row(&row).map_err(Error::DeserializationError)
}
fn all_rows<P, Row>(conn: &mut SqliteConnection, schema: Option<&str>) -> QueryResult<Vec<Row>>
where
P: ReadPragma,
Sqlite: QueryMetadata<P::SqlType>,
Row: FromSqlRow<P::SqlType, Sqlite>,
{
cursor::<P>(conn, schema)?
.map(|row| Row::build_from_row(&row?).map_err(Error::DeserializationError))
.collect()
}
fn assign<P: WritePragma>(
conn: &mut SqliteConnection,
schema: Option<&str>,
argument: P::Argument<'_>,
) -> QueryResult<()> {
use crate::connection::SimpleConnection;
use crate::query_builder::QueryBuilder;
use crate::sqlite::SqliteQueryBuilder;
let mut query = SqliteQueryBuilder::new();
PragmaAssignment::<P, _> {
schema,
argument,
pragma: PhantomData,
}
.to_sql(&mut query, &Sqlite)?;
conn.batch_execute(&query.finish())
}
#[derive(Debug)]
struct PragmaQuery<'a, P> {
schema: Option<&'a str>,
pragma: PhantomData<P>,
}
impl<'a, P: Pragma> PragmaQuery<'a, P> {
fn new(schema: Option<&'a str>) -> Self {
Self {
schema,
pragma: PhantomData,
}
}
}
impl<P: Pragma> QueryFragment<Sqlite> for PragmaQuery<'_, P> {
fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Sqlite>) -> QueryResult<()> {
push_pragma::<P>(&mut out, self.schema)
}
}
impl<P> QueryId for PragmaQuery<'_, P> {
type QueryId = ();
const HAS_STATIC_QUERY_ID: bool = false;
}
impl<P: ReadPragma> Query for PragmaQuery<'_, P> {
type SqlType = P::SqlType;
}
struct PragmaAssignment<'a, P, A> {
schema: Option<&'a str>,
argument: A,
pragma: PhantomData<P>,
}
impl<P: Pragma, A: PragmaArgument> QueryFragment<Sqlite> for PragmaAssignment<'_, P, A> {
fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Sqlite>) -> QueryResult<()> {
push_pragma::<P>(&mut out, self.schema)?;
out.push_sql(" = ");
self.argument.render(&mut out)
}
}
fn push_pragma<P: Pragma>(
out: &mut AstPass<'_, '_, Sqlite>,
schema: Option<&str>,
) -> QueryResult<()> {
out.push_sql("PRAGMA ");
if let Some(schema) = schema {
out.push_identifier(schema)?;
out.push_sql(".");
}
push_bare_word(out, "name", P::NAME)
}
fn push_bare_word(
out: &mut AstPass<'_, '_, Sqlite>,
kind: &str,
word: &'static str,
) -> QueryResult<()> {
if word.is_empty()
|| !word
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
{
return Err(Error::QueryBuilderError(
alloc::format!("`{word}` is not a valid pragma {kind}").into(),
));
}
out.push_sql(word);
Ok(())
}