Executor 与 Acquire

pub trait Executor<'c>: Send + Debug + Sized {
    type Database: Database;
    // ...
    fn fetch_many<'e, 'q: 'e, E>(
        self,
        query: E,
    ) -> BoxStream<
        'e,
        Result<
            Either<<Self::Database as Database>::QueryResult, <Self::Database as Database>::Row>,
            Error,
        >,
    >
    where
        'c: 'e,
        E: 'q + Execute<'q, Self::Database>;
    // ...
    fn fetch_optional<'e, 'q: 'e, E>(
        self,
        query: E,
    ) -> BoxFuture<'e, Result<Option<<Self::Database as Database>::Row>, Error>>
    where
        'c: 'e,
        E: 'q + Execute<'q, Self::Database>;

use futures::future::BoxFuture;
use futures::stream::BoxStream;
use sqlx::Either;
use sqlx::postgres::{PgQueryResult, PgRow};
use sqlx::{Error, Execute, Executor, Postgres};

// An Executor that implements only the two methods the plan calls "required".
#[derive(Debug)]
struct Fake;

impl<'c> Executor<'c> for Fake {
  type Database = Postgres;

  fn fetch_many<'e, 'q: 'e, E>(self, _query: E) -> BoxStream<'e, Result<Either<PgQueryResult, PgRow>, Error>>
  where
    'c: 'e,
    E: 'q + Execute<'q, Postgres>,
  {
    Box::pin(futures::stream::empty())
  }

  fn fetch_optional<'e, 'q: 'e, E>(self, _query: E) -> BoxFuture<'e, Result<Option<PgRow>, Error>>
  where
    'c: 'e,
    E: 'q + Execute<'q, Postgres>,
  {
    Box::pin(async { Ok(None) })
  }
}

fn main() {}

error[E0046]: not all trait items implemented, missing: `prepare_with`, `describe`
  --> src/bin/exec_required.rs:11:1
   |
11 | impl<'c> Executor<'c> for Fake {
   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ missing `prepare_with`, `describe` in implementation
   |
   = help: implement the missing item: `fn prepare_with<'e>(self, _: SqlStr, _: &'e [<<Self as Executor<'c>>::Database as sqlx::Database>::TypeInfo]) -> Pin<Box<(dyn futures::Future<Output = Result<<<Self as Executor<'c>>::Database as sqlx::Database>::Statement, sqlx::Error>> + std::marker::Send + 'e)>> { todo!() }`
   = help: implement the missing item: `fn describe<'e>(self, _: SqlStr) -> Pin<Box<(dyn futures::Future<Output = Result<sqlx::Describe<<Self as Executor<'c>>::Database>, sqlx::Error>> + std::marker::Send + 'e)>> { todo!() }`

error[E0046]: not all trait items implemented, missing: `prepare_with`
  --> src/main.rs:11:1
   |
11 | impl<'c> Executor<'c> for Fake {
   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ missing `prepare_with` in implementation
...

pub trait Execute<'q, DB: Database>: Send + Sized {
    /// Gets the SQL that will be executed.
    fn sql(self) -> SqlStr;

    /// Gets the previously cached statement, if available.
    fn statement(&self) -> Option<&DB::Statement>;

    /// Returns the arguments to be bound against the query string.
    ///
    /// Returning `Ok(None)` for `Arguments` indicates to use a "simple" query protocol and to not
    /// prepare the query. Returning `Ok(Some(Default::default()))` is an empty arguments object that
    /// will be prepared (and cached) before execution.
    ///
    /// Returns `Err` if encoding any of the arguments failed.
    fn take_arguments(&mut self) -> Result<Option<<DB as Database>::Arguments>, BoxDynError>;

    /// Returns `true` if the statement should be cached.
    fn persistent(&self) -> bool;
}

    fn fetch_many<'e, 'q, E>(
        self,
        mut query: E,
    ) -> BoxStream<'e, Result<Either<PgQueryResult, PgRow>, Error>>
    where
        'c: 'e,
        E: Execute<'q, Self::Database>,
        'q: 'e,
        E: 'q,
    {
        // False positive: https://github.com/rust-lang/rust-clippy/issues/12560
        #[allow(clippy::map_clone)]
        let metadata = query.statement().map(|s| Arc::clone(&s.metadata));
        let arguments = query.take_arguments().map_err(Error::Encode);
        let persistent = query.persistent();
        let sql = query.sql();

        Box::pin(try_stream! {
            let arguments = arguments?;
            let mut s = pin!(self.run(sql, arguments, persistent, metadata).await?);

            while let Some(v) = s.try_next().await? {
                r#yield!(v);
            }

            Ok(())
        })
    }

use sqlx::{Connection, Executor, Row};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
  let pool = exp_sqlx_core1::pool(1).await?;
  let mut conn = pool.acquire().await?;
  conn.clear_cached_statements().await?;

  let cases: [(&str, &'static str); 5] = [
    ("&str", "select 1"),
    ("raw_sql", "select 2"),
    ("query, no bind", "select 3"),
    ("query, persistent(false)", "select 4"),
    ("query, one bind", "select $1::int4"),
  ];
  for (i, (label, sql)) in cases.into_iter().enumerate() {
    let row = match i {
      0 => conn.fetch_one(sql).await?,
      1 => sqlx::raw_sql(sql).fetch_one(&mut *conn).await?,
      2 => sqlx::query(sql).fetch_one(&mut *conn).await?,
      3 => sqlx::query(sql).persistent(false).fetch_one(&mut *conn).await?,
      _ => sqlx::query(sql).bind(5_i32).fetch_one(&mut *conn).await?,
    };
    let v: i32 = row.try_get(0)?;
    println!("{label:<25} value={v} cached_statements_size={}", conn.cached_statements_size());
  }
  let names: Vec<String> = sqlx::query_scalar("select name from pg_prepared_statements order by name")
    .persistent(false)
    .fetch_all(&mut *conn)
    .await?;
  println!("pg_prepared_statements: {names:?}");
  Ok(())
}

&str                      value=1 cached_statements_size=0
raw_sql                   value=2 cached_statements_size=0
query, no bind            value=3 cached_statements_size=1
query, persistent(false)  value=4 cached_statements_size=1
query, one bind           value=5 cached_statements_size=2
pg_prepared_statements: ["sqlx_s_1", "sqlx_s_2"]

    = help: the following other types implement trait `Executor<'c>`:
              `&'c mut AnyConnection` implements `Executor<'c>`
              `&'c mut PgConnection` implements `Executor<'c>`
              `&'c mut PgListener` implements `Executor<'c>`
              `&'c mut SqliteConnection` implements `Executor<'c>`
              `&Pool<DB>` implements `Executor<'p>`

impl<'p, DB: Database> Executor<'p> for &'_ Pool<DB>
where
    for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>,
{
    type Database = DB;

    fn fetch_many<'e, 'q: 'e, E>(
        self,
        query: E,
    ) -> BoxStream<'e, Result<Either<DB::QueryResult, DB::Row>, Error>>
    where
        E: 'q + Execute<'q, Self::Database>,
    {
        let pool = self.clone();

        Box::pin(try_stream! {
            let mut conn = pool.acquire().await?;
            let mut s = conn.fetch_many(query);

            while let Some(v) = s.try_next().await? {
                r#yield!(v);
            }

            Ok(())
        })
    }
    // ...

#[tokio::main]
async fn main() -> anyhow::Result<()> {
  let pool = exp_sqlx_core1::pool(1).await?;
  let mut tx = pool.begin().await?;
  let one: i32 = sqlx::query_scalar("select 1").fetch_one(&mut tx).await?;
  println!("{one}");
  Ok(())
}

error[E0277]: the trait bound `&mut Transaction<'_, Postgres>: Executor<'_>` is not satisfied
   --> src/bin/tx_noderef.rs:5:59
    |
  5 |   let one: i32 = sqlx::query_scalar("select 1").fetch_one(&mut tx).await?;
    |                                                 --------- ^^^^^^^ the trait `Executor<'_>` is not implemented for `&mut Transaction<'_, Postgres>`
    |                                                 |
    |                                                 required by a bound introduced by this call
    |
    = help: the following other types implement trait `Executor<'c>`:
              `&'c mut AnyConnection` implements `Executor<'c>`
              `&'c mut PgConnection` implements `Executor<'c>`
              `&'c mut PgListener` implements `Executor<'c>`
              `&'c mut SqliteConnection` implements `Executor<'c>`
              `&Pool<DB>` implements `Executor<'p>`
note: required by a bound in `QueryScalar::<'q, DB, O, A>::fetch_one`
...

use sqlx::PgExecutor;

async fn two_queries<'c, E>(e: E) -> sqlx::Result<i32>
where
  E: PgExecutor<'c>,
{
  let a: i32 = sqlx::query_scalar("select 1").fetch_one(e).await?;
  let b: i32 = sqlx::query_scalar("select 2").fetch_one(e).await?;
  Ok(a + b)
}

fn main() {
  let _ = two_queries::<&sqlx::PgPool>;
}

error[E0382]: use of moved value: `e`
 --> src/bin/executor_twice.rs:8:57
  |
3 | async fn two_queries<'c, E>(e: E) -> sqlx::Result<i32>
  |                             - move occurs because `e` has type `E`, which does not implement the `Copy` trait
...
7 |   let a: i32 = sqlx::query_scalar("select 1").fetch_one(e).await?;
  |                                                         - value moved here
8 |   let b: i32 = sqlx::query_scalar("select 2").fetch_one(e).await?;
  |                                                         ^ value used here after move
...

pub trait Acquire<'c> {
    type Database: Database;

    type Connection: Deref<Target = <Self::Database as Database>::Connection> + DerefMut + Send;

    fn acquire(self) -> BoxFuture<'c, Result<Self::Connection, Error>>;

    fn begin(self) -> BoxFuture<'c, Result<Transaction<'c, Self::Database>, Error>>;
}

use sqlx::{Acquire, Postgres};

// One function, generic over anything that can hand out a Postgres connection.
async fn insert_and_count<'a, A>(conn: A, name: &str) -> anyhow::Result<(i32, i64)>
where
  A: Acquire<'a, Database = Postgres>,
{
  let mut conn = conn.acquire().await?;
  let pid: i32 = sqlx::query_scalar("select pg_backend_pid()").fetch_one(&mut *conn).await?;
  sqlx::query("insert into acq_demo(name) values ($1)").bind(name).execute(&mut *conn).await?;
  let n: i64 = sqlx::query_scalar("select count(*) from acq_demo").fetch_one(&mut *conn).await?;
  Ok((pid, n))
}

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

  let (pid, n) = insert_and_count(&pool, "via &Pool").await?;
  println!("&Pool            : pid={pid} count={n}");

  let mut tx = pool.begin().await?;
  let (pid, n) = insert_and_count(&mut tx, "via &mut Transaction").await?;
  println!("&mut Transaction : pid={pid} count={n}");
  tx.rollback().await?;

  let mut conn = pool.acquire().await?;
  let (pid, n) = insert_and_count(&mut conn, "via &mut PoolConnection").await?;
  println!("&mut PoolConn    : pid={pid} count={n}");
  let (pid, n) = insert_and_count(&mut *conn, "via &mut PgConnection").await?;
  println!("&mut PgConnection: pid={pid} count={n}");
  drop(conn);

  let names: Vec<String> = sqlx::query_scalar("select name from acq_demo order by id").fetch_all(&pool).await?;
  println!("after rollback   : {names:?}");
  sqlx::raw_sql("drop table acq_demo").execute(&pool).await?;
  Ok(())
}

&Pool            : pid=113 count=1
&mut Transaction : pid=114 count=2
&mut PoolConn    : pid=113 count=2
&mut PgConnection: pid=113 count=3
after rollback   : ["via &Pool", "via &mut PoolConnection", "via &mut PgConnection"]

fetch().try_next(): Ok(Some(5))
fetch_one()       : Err(Database(PgDatabaseError { severity: Error, code: "22012", message: "division by zero", detail: None, hint: None, position: None, where: None, schema: None, table: None, column: None, data_type: None, constraint: None, file: Some("int.c"), line: Some(870), routine: Some("int4div") }))

            // With deferred constraints we need to check all responses as we
            // could get a OK response (with uncommitted data), only to get an
            // error response after (when the deferred constraint is actually
            // checked).
            let mut ret = None;
            while let Some(result) = s.try_next().await? {
                match result {
                    Either::Right(r) if ret.is_none() => ret = Some(r),
                    _ => {}
                }
            }
            Ok(ret)