Query、QueryAs 与 SqlSafeStr

pub struct Query<'q, DB: Database, A> {
    pub(crate) statement: Either<SqlStr, &'q DB::Statement>,
    pub(crate) arguments: Option<Result<A, BoxDynError>>,
    pub(crate) database: PhantomData<DB>,
    pub(crate) persistent: bool,
}

pub fn query<'a, DB>(sql: impl SqlSafeStr) -> Query<'a, DB, <DB as Database>::Arguments>
where
    DB: Database,
{
    Query {
        database: PhantomData,
        arguments: Some(Ok(Default::default())),
        statement: Either::Left(sql.into_sql_str()),
        persistent: true,
    }
}

    /// If encoding the value fails, the error is stored and later surfaced when executing the query.
    pub fn bind<'t, T: Encode<'t, DB> + Type<DB>>(mut self, value: T) -> Self {
        let Ok(arguments) = self.get_arguments() else {
            return self;
        };

        let argument_number = arguments.len() + 1;
        if let Err(error) = arguments.add(value) {
            self.arguments = Some(Err(format!(
                "Encoding argument ${argument_number} failed: {error}"
            )
            .into()));
        }

        self
    }

    /// Like [`Query::bind`] but immediately returns an error if encoding a value failed.
    pub fn try_bind<'t, T: Encode<'t, DB> + Type<DB>>(
        &mut self,
        value: T,
    ) -> Result<(), BoxDynError> {
        let arguments = self.get_arguments()?;

        arguments.add(value)
    }

    pub(crate) fn add<'q, T>(&mut self, value: T) -> Result<(), BoxDynError>
    where
        T: Encode<'q, Postgres> + Type<Postgres>,
    {
        let type_info = value.produces().unwrap_or_else(T::type_info);

        let buffer_snapshot = self.buffer.snapshot();

        // encode the value into our buffer
        if let Err(error) = self.buffer.encode(value) {
            // reset the value buffer to its previous value if encoding failed,
            // so we don't leave a half-encoded value behind
            self.buffer.reset_to_snapshot(buffer_snapshot);
            return Err(error);
        };

        // remember the type information for this value
        self.types.push(type_info);
        // increment the number of arguments we are tracking
        self.buffer.count += 1;

        Ok(())
    }

    pub fn bind<T: 'q + Encode<'q, DB> + Type<DB>>(mut self, value: T) -> Self {

// Query::bind: the value only has to live for the call; it is encoded right away.
fn by_query(name: &str) -> Query<'static, Postgres, PgArguments> {
  sqlx::query("select $1::text").bind(name)
}

// QueryAs::bind still requires T: 'q.
fn by_query_as(name: &str) -> QueryAs<'static, Postgres, (String,), PgArguments> {
  sqlx::query_as("select $1::text").bind(name)
}

error: lifetime may not live long enough
  --> src/bin/bind_lifetime.rs:12:3
   |
11 | fn by_query_as(name: &str) -> QueryAs<'static, Postgres, (String,), PgArguments> {
   |                      - let's call the lifetime of this reference `'1`
12 |   sqlx::query_as("select $1::text").bind(name)
   |   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ returning this value requires that `'1` must outlive `'static`

use std::time::Duration;

use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
use sqlx::{Postgres, QueryBuilder};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
  tracing_subscriber::fmt()
    .with_env_filter("sqlx::query=debug")
    .without_time()
    .with_target(false)
    .with_ansi(false)
    .init();
  let url = std::env::var("DATABASE_URL")?;
  let opts: PgConnectOptions = url.parse()?;
  let pool = PgPoolOptions::new()
    .max_connections(1)
    .acquire_timeout(Duration::from_secs(1))
    .connect_with(opts)
    .await?;

  // INTERVAL has microsecond precision; 1 ns cannot be encoded.
  let q = sqlx::query("select $1::interval, $2::int4")
    .bind(Duration::from_nanos(1))
    .bind(7_i32);
  println!("bind returned a Query, no error yet");
  let err = q.fetch_one(&pool).await.unwrap_err();
  println!("fetch_one: {err}");
  println!("debug    : {err:?}");

  // The only connection is checked out: the same bad query on &pool.
  let held = pool.acquire().await?;
  let err = sqlx::query("select $1::interval")
    .bind(Duration::from_nanos(1))
    .fetch_one(&pool).await.unwrap_err();
  println!("while held: {err}");
  drop(held);

  // try_bind reports the same failure right away.
  let mut q = sqlx::query("select $1::interval");
  let e = q.try_bind(Duration::from_nanos(1)).unwrap_err();
  println!("try_bind : {e}");

  // A value that does encode: 1 µs.
  let s: String = sqlx::query_scalar("select $1::interval::text")
    .bind(Duration::from_micros(1))
    .fetch_one(&pool).await?;
  println!("1 µs ok  : {s}");

  // QueryBuilder::push_bind does not defer: it panics.
  let r = std::panic::catch_unwind(|| {
    let mut qb = QueryBuilder::<Postgres>::new("select ");
    qb.push_bind(Duration::from_nanos(1));
  });
  println!("QueryBuilder::push_bind panicked: {}", r.is_err());
  Ok(())
}

bind returned a Query, no error yet
fetch_one: error occurred while encoding a value: Encoding argument $1 failed: PostgreSQL `INTERVAL` does not support nanoseconds precision
debug    : Encode("Encoding argument $1 failed: PostgreSQL `INTERVAL` does not support nanoseconds precision")
while held: pool timed out while waiting for an open connection
try_bind : PostgreSQL `INTERVAL` does not support nanoseconds precision
DEBUG summary="select $1::interval::text" db.statement="" rows_affected=1 rows_returned=1 elapsed=5.504042ms elapsed_secs=0.005504042
1 µs ok  : 00:00:00.000001

thread 'main' (3084215) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/sqlx-core-0.9.0/src/query_builder.rs:165:30:
Failed to add argument: "PostgreSQL `INTERVAL` does not support nanoseconds precision"
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
QueryBuilder::push_bind panicked: true

    /// If `false`, the prepared statement will be closed after execution.
    ///
    /// Default: `true`.
    pub fn persistent(mut self, value: bool) -> Self {
        self.persistent = value;
        self
    }

pub fn query<DB>(sql: &str) -> Query<'_, DB, <DB as Database>::Arguments<'_>>

#[diagnostic::on_unimplemented(
    label = "dynamic SQL string",
    message = "dynamic SQL strings should be audited for possible injections",
    note = "prefer literal SQL strings with bind parameters or `QueryBuilder` to add dynamic data to a query.

To bypass this error, manually audit for potential injection vulnerabilities and wrap with `AssertSqlSafe()`.
For details, see the docs for `SqlSafeStr`.\n",
    note = "this trait is only implemented for `&'static str`, not all `&str` like the compiler error may suggest"
)]
pub trait SqlSafeStr {
    /// Convert `self` to a [`SqlStr`].
    fn into_sql_str(self) -> SqlStr;
}

impl SqlSafeStr for &'static str {
    #[inline]
    fn into_sql_str(self) -> SqlStr {
        SqlStr(Repr::Static(self))
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
  let pool = exp_sqlx_core1::pool(1).await?;
  let table = "pg_class";
  let r = sqlx::query(&format!("select count(*) from {table}"))
    .execute(&pool)
    .await?;
  println!("rows_affected = {}", r.rows_affected());
  Ok(())
}

error[E0277]: dynamic SQL strings should be audited for possible injections
   --> src/bin/fmt_sql.rs:5:23
    |
  5 |   let r = sqlx::query(&format!("select count(*) from {table}"))
    |           ----------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ dynamic SQL string
    |           |
    |           required by a bound introduced by this call
    |
    = help: the trait `SqlSafeStr` is not implemented for `&std::string::String`
    = note: prefer literal SQL strings with bind parameters or `QueryBuilder` to add dynamic data to a query.
            
            To bypass this error, manually audit for potential injection vulnerabilities and wrap with `AssertSqlSafe()`.
            For details, see the docs for `SqlSafeStr`.
            
    = note: this trait is only implemented for `&'static str`, not all `&str` like the compiler error may suggest
help: the trait `SqlSafeStr` is implemented for `&str`
   --> ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/sqlx-core-0.9.0/src/sql_str.rs:52:1
    |
 52 | impl SqlSafeStr for &'static str {
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
note: required by a bound in `sqlx::query`
   --> ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/sqlx-core-0.9.0/src/query.rs:653:32
    |
653 | pub fn query<'a, DB>(sql: impl SqlSafeStr) -> Query<'a, DB, <DB as Database>::Arguments>
    |                                ^^^^^^^^^^ required by this bound in `query`

For more information about this error, try `rustc --explain E0277`.

rows_affected = 1

error[E0308]: mismatched types
 --> src/bin/fmt_sql_owned.rs:6:23
  |
6 |   let r = sqlx::query(format!("select count(*) from {table}"))
  |                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `&str`, found `String`

  let sql = String::from("select 1");
  let r = sqlx::query(sql.as_str()).execute(&pool).await?;

error[E0597]: `sql` does not live long enough
   --> src/bin/borrowed_sql.rs:5:23
    |
  4 |   let sql = String::from("select 1");
    |       --- binding `sql` declared here
  5 |   let r = sqlx::query(sql.as_str()).execute(&pool).await?;
    |           ------------^^^----------
    |           |           |
    |           |           borrowed value does not live long enough
    |           argument requires that `sql` is borrowed for `'static`
...
  8 | }
    | - `sql` dropped here while still borrowed
    |
note: requirement that the value outlives `'static` introduced here
   --> ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/sqlx-core-0.9.0/src/query.rs:653:32
    |
653 | pub fn query<'a, DB>(sql: impl SqlSafeStr) -> Query<'a, DB, <DB as Database>::Arguments>
    |                                ^^^^^^^^^^

pub struct SqlStr(Repr);

#[derive(Debug)]
enum Repr {
    /// We need a variant to memoize when we already have a static string, so we don't copy it.
    Static(&'static str),
    /// Thanks to the new niche in `String`, this doesn't increase the size beyond 3 words.
    /// We essentially get all these variants for free.
    Owned(String),
    Boxed(Box<str>),
    Arced(Arc<str>),
    /// Allows for dynamic shared ownership with `query_builder`.
    ArcString(Arc<String>),
}

use sqlx::postgres::PgArguments;
use sqlx::query::Query;
use sqlx::{AssertSqlSafe, Postgres, SqlStr};

// Table names cannot be bind parameters; this one comes from a fixed list.
fn count_rows(table: &str, min_oid: i64) -> Query<'static, Postgres, PgArguments> {
  assert!(["pg_class", "pg_type"].contains(&table));
  let sql = format!("select count(*) from {table} where oid::int8 >= $1");
  sqlx::query(AssertSqlSafe(sql)).bind(min_oid)
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
  println!("size_of::<SqlStr>()  = {}", size_of::<SqlStr>());
  println!("size_of::<String>()  = {}", size_of::<String>());
  println!("size_of::<Option<String>>() = {}", size_of::<Option<String>>());
  let pool = exp_sqlx_core1::pool(1).await?;
  let q = count_rows("pg_type", 0);
  let row = q.fetch_one(&pool).await?;
  let n: i64 = sqlx::Row::get(&row, 0);
  println!("pg_type rows > 0: {}", n > 0);
  Ok(())
}

error[E0515]: cannot return value referencing local variable `sql`
 --> src/bin/owned_query.rs:8:3
  |
8 |   sqlx::query(&sql).bind(min_oid)
  |   ^^^^^^^^^^^^----^^^^^^^^^^^^^^^
  |   |           |
  |   |           `sql` is borrowed here
  |   returns a value referencing data owned by the current function

impl<DB: Database> Execute<'_, DB> for RawSql {
    fn sql(self) -> SqlStr {
        self.0
    }

    fn statement(&self) -> Option<&<DB as Database>::Statement> {
        None
    }

    fn take_arguments(&mut self) -> Result<Option<<DB as Database>::Arguments>, BoxDynError> {
        Ok(None)
    }

    fn persistent(&self) -> bool {
        false
    }
}

pub struct QueryBuilder<DB>
where
    DB: Database,
{
    query: Arc<String>,
    init_len: usize,
    arguments: Option<<DB as Database>::Arguments>,
}

    pub fn push_values<I, F>(&mut self, tuples: I, mut push_tuple: F) -> &mut Self
    where
        I: IntoIterator,
        F: FnMut(Separated<'_, DB, &'static str>, I::Item),
    {
        self.sanity_check();

        self.push("VALUES ");

        let mut separated = self.separated(", ");

        for tuple in tuples {
            separated.push("(");

            // use a `Separated` with a separate (hah) internal state
            push_tuple(separated.query_builder.separated(", "), tuple);

            separated.push_unseparated(")");
        }

        debug_assert!(
            separated.push_separator,
            "No value being pushed. QueryBuilder may not build correct sql query!"
        );

        separated.query_builder
    }

use sqlx::{Postgres, QueryBuilder};

struct User {
  id: i64,
  name: &'static str,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
  let pool = exp_sqlx_core1::pool(1).await?;
  sqlx::raw_sql("drop table if exists qb_users; create table qb_users(id int8 primary key, name text not null)")
    .execute(&pool).await?;

  let users = [User { id: 1, name: "ann" }, User { id: 2, name: "bob" }, User { id: 3, name: "cy" }];
  let mut qb: QueryBuilder<Postgres> = QueryBuilder::new("insert into qb_users(id, name) ");
  qb.push_values(&users, |mut b, u| {
    b.push_bind(u.id).push_bind(u.name);
  });
  qb.push(" on conflict (id) do nothing");
  println!("push_values: {}", qb.sql().as_str());
  let r = qb.build().execute(&pool).await?;
  println!("inserted   : {}", r.rows_affected());

  let mut qb: QueryBuilder<Postgres> = QueryBuilder::new("select name from qb_users where id in (");
  let mut sep = qb.separated(", ");
  for id in [1_i64, 3] {
    sep.push_bind(id);
  }
  sep.push_unseparated(") order by id");
  println!("separated  : {}", qb.sql().as_str());
  let names: Vec<String> = qb.build_query_scalar().fetch_all(&pool).await?;
  println!("names      : {names:?}");

  sqlx::raw_sql("drop table qb_users").execute(&pool).await?;
  Ok(())
}

push_values: insert into qb_users(id, name) VALUES ($1, $2), ($3, $4), ($5, $6) on conflict (id) do nothing
inserted   : 3
separated  : select name from qb_users where id in ($1, $2) order by id
names      : ["ann", "cy"]

size_of::<SqlStr>()  = 24
size_of::<String>()  = 24
size_of::<Option<String>>() = 24
pg_type rows > 0: true