feat: support PoolConnection via AsExecutor trait

README.md

@@ -11,12 +11,18 @@ For each SQL query annotated with a sqlc command, the plugin emits:
 - An optional params struct (`QueryNameParams`) when a query has 2+ parameters.
 - A `&mut self` method on `pub struct Queries<E>` that executes the query.
 
-`Queries<E>` wraps an executor once at construction. Pass `&pool` for pool usage or `&mut tx` for transactions:
+`Queries<E>` wraps anything implementing the generated `AsExecutor` trait. `AsExecutor` is implemented for `PgPool`, `&PgPool`, `PgConnection`, `Transaction<'_, Postgres>`, `PoolConnection<Postgres>`, and `&mut T` of each:
 
 ```rust
+// From a pool:
 let mut q = Queries::new(&pool);
 let author = q.get_author(1).await?;
 
+// Borrowed or owned pool connection:
+let mut conn = pool.acquire().await?;
+let mut q = Queries::new(&mut conn);
+// ...or Queries::new(conn) to take ownership.
+
 // Transactions:
 let mut tx = pool.begin().await?;
 let mut q = Queries::new(&mut tx);
@@ -126,7 +132,7 @@ Array types (`type[]`) become `Vec<T>`. Nullable columns become `Option<T>`.
 | `:batchmany` | `impl Stream<Item = Result<Vec<QueryRow>, sqlx::Error>>` | Lazily fetch all rows per item |
 | `:copyfrom` | `Result<u64, sqlx::Error>` | Chunked bulk insert from any `IntoIterator` |
 
-All functions are `&mut self` methods on `Queries<E>`. The bound `for<'c> &'c mut E: sqlx::Executor<'c, Database = sqlx::Postgres>` is satisfied by `&PgPool`, `&mut PgConnection`, and `&mut Transaction<'_, Postgres>`.
+All functions are `&mut self` methods on `Queries<E>`. The bound is `E: AsExecutor`, where `AsExecutor` is the trait emitted in each generated file. Impls cover `PgPool`, `&PgPool`, `PgConnection`, `Transaction<'_, Postgres>`, `PoolConnection<Postgres>`, and `&mut T` of each.
 
 Batch methods generate `Stream`-returning APIs and reference `futures_core` and `futures_util` directly. Consumer crates should include those dependencies alongside `sqlx`.
 

examples/advanced-types/src/queries.rs

@@ -1,54 +1,105 @@
-// Code generated by sqlc-gen-sqlx. DO NOT EDIT.
+// Code generated by sqlc-gen-sqlx v0.1.0. DO NOT EDIT.
+// sqlc version: v1.30.0
 
 #![allow(
     dead_code,
     reason = "generated queries may expose items a caller does not use"
 )]
 
 const GET_EVENT: &str = "SELECT id, name, flags, event_window FROM events WHERE id = $1";
-
 #[derive(Debug, Clone, sqlx::FromRow)]
 pub struct GetEventRow {
     pub id: i64,
     pub name: String,
     pub flags: bit_vec::BitVec,
     pub event_window: sqlx::postgres::types::PgRange<chrono::DateTime<chrono::Utc>>,
 }
-
 const LIST_EVENTS: &str = "SELECT id, name, flags, event_window FROM events";
-
 #[derive(Debug, Clone, sqlx::FromRow)]
 pub struct ListEventsRow {
     pub id: i64,
     pub name: String,
     pub flags: bit_vec::BitVec,
     pub event_window: sqlx::postgres::types::PgRange<chrono::DateTime<chrono::Utc>>,
 }
-
+pub trait AsExecutor {
+    type Exec<'c>: sqlx::Executor<'c, Database = sqlx::Postgres>
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_>;
+}
+impl AsExecutor for sqlx::PgPool {
+    type Exec<'c>
+        = &'c sqlx::PgPool
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_> {
+        &*self
+    }
+}
+impl AsExecutor for &sqlx::PgPool {
+    type Exec<'c>
+        = &'c sqlx::PgPool
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_> {
+        *self
+    }
+}
+impl AsExecutor for sqlx::PgConnection {
+    type Exec<'c>
+        = &'c mut sqlx::PgConnection
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_> {
+        self
+    }
+}
+impl AsExecutor for sqlx::Transaction<'_, sqlx::Postgres> {
+    type Exec<'c>
+        = &'c mut sqlx::PgConnection
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_> {
+        &mut **self
+    }
+}
+impl AsExecutor for sqlx::pool::PoolConnection<sqlx::Postgres> {
+    type Exec<'c>
+        = &'c mut sqlx::PgConnection
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_> {
+        &mut **self
+    }
+}
+impl<T: AsExecutor + ?Sized> AsExecutor for &mut T {
+    type Exec<'c>
+        = T::Exec<'c>
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_> {
+        (**self).as_executor()
+    }
+}
 pub struct Queries<E> {
     db: E,
 }
-
 impl<E> Queries<E> {
     pub fn new(db: E) -> Self {
         Self { db }
     }
 }
-
-impl<E> Queries<E>
-where
-    for<'c> &'c mut E: sqlx::Executor<'c, Database = sqlx::Postgres>,
-{
+impl<E: AsExecutor> Queries<E> {
     pub async fn get_event(&mut self, id: i64) -> Result<GetEventRow, sqlx::Error> {
         sqlx::query_as::<_, GetEventRow>(GET_EVENT)
             .bind(id)
-            .fetch_one(&mut self.db)
+            .fetch_one(self.db.as_executor())
             .await
     }
-
     pub async fn list_events(&mut self) -> Result<Vec<ListEventsRow>, sqlx::Error> {
         sqlx::query_as::<_, ListEventsRow>(LIST_EVENTS)
-            .fetch_all(&mut self.db)
+            .fetch_all(self.db.as_executor())
             .await
     }
 }

examples/basic/src/queries.rs

@@ -35,6 +35,66 @@ pub struct CreateAuthorRow {
 }
 const DELETE_AUTHOR: &str = "DELETE FROM authors WHERE id = $1";
 const DELETE_AUTHOR_ROWS: &str = "DELETE FROM authors WHERE id = $1";
+pub trait AsExecutor {
+    type Exec<'c>: sqlx::Executor<'c, Database = sqlx::Postgres>
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_>;
+}
+impl AsExecutor for sqlx::PgPool {
+    type Exec<'c>
+        = &'c sqlx::PgPool
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_> {
+        &*self
+    }
+}
+impl AsExecutor for &sqlx::PgPool {
+    type Exec<'c>
+        = &'c sqlx::PgPool
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_> {
+        *self
+    }
+}
+impl AsExecutor for sqlx::PgConnection {
+    type Exec<'c>
+        = &'c mut sqlx::PgConnection
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_> {
+        self
+    }
+}
+impl AsExecutor for sqlx::Transaction<'_, sqlx::Postgres> {
+    type Exec<'c>
+        = &'c mut sqlx::PgConnection
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_> {
+        &mut **self
+    }
+}
+impl AsExecutor for sqlx::pool::PoolConnection<sqlx::Postgres> {
+    type Exec<'c>
+        = &'c mut sqlx::PgConnection
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_> {
+        &mut **self
+    }
+}
+impl<T: AsExecutor + ?Sized> AsExecutor for &mut T {
+    type Exec<'c>
+        = T::Exec<'c>
+    where
+        Self: 'c;
+    fn as_executor(&mut self) -> Self::Exec<'_> {
+        (**self).as_executor()
+    }
+}
 pub struct Queries<E> {
     db: E,
 }
@@ -43,19 +103,16 @@ impl<E> Queries<E> {
         Self { db }
     }
 }
-impl<E> Queries<E>
-where
-    for<'c> &'c mut E: sqlx::Executor<'c, Database = sqlx::Postgres>,
-{
+impl<E: AsExecutor> Queries<E> {
     pub async fn get_author(&mut self, id: i64) -> Result<GetAuthorRow, sqlx::Error> {
         sqlx::query_as::<_, GetAuthorRow>(GET_AUTHOR)
             .bind(id)
-            .fetch_one(&mut self.db)
+            .fetch_one(self.db.as_executor())
             .await
     }
     pub async fn list_authors(&mut self) -> Result<Vec<ListAuthorsRow>, sqlx::Error> {
         sqlx::query_as::<_, ListAuthorsRow>(LIST_AUTHORS)
-            .fetch_all(&mut self.db)
+            .fetch_all(self.db.as_executor())
             .await
     }
     pub async fn create_author(
@@ -65,20 +122,20 @@ where
         sqlx::query_as::<_, CreateAuthorRow>(CREATE_AUTHOR)
             .bind(arg.name)
             .bind(arg.bio)
-            .fetch_one(&mut self.db)
+            .fetch_one(self.db.as_executor())
             .await
     }
     pub async fn delete_author(&mut self, id: i64) -> Result<(), sqlx::Error> {
         sqlx::query(DELETE_AUTHOR)
             .bind(id)
-            .execute(&mut self.db)
+            .execute(self.db.as_executor())
             .await?;
         Ok(())
     }
     pub async fn delete_author_rows(&mut self, id: i64) -> Result<u64, sqlx::Error> {
         let result = sqlx::query(DELETE_AUTHOR_ROWS)
             .bind(id)
-            .execute(&mut self.db)
+            .execute(self.db.as_executor())
             .await?;
         Ok(result.rows_affected())
     }

examples/batch/Cargo.toml

@@ -6,3 +6,5 @@ edition = "2024"
 [dependencies]
 sqlx = { workspace = true }
 tokio = { workspace = true }
+futures-core = "0.3"
+futures-util = "0.3"

examples/batch/src/main.rs

@@ -2,6 +2,8 @@
 #[path = "queries.rs"]
 mod queries;
 #[cfg(test)]
+use futures_util::TryStreamExt;
+#[cfg(test)]
 use queries::Queries;
 #[cfg(test)]
 use sqlx::{Connection as _, PgConnection};
@@ -34,12 +36,18 @@ async fn test_batch_roundtrip() {
 
     let mut q = Queries::new(conn);
 
-    let authors = q.batch_get_author(vec![1, 2]).await.expect("batch get");
+    let authors: Vec<_> = q
+        .batch_get_author(vec![1, 2])
+        .try_collect()
+        .await
+        .expect("batch get");
     assert_eq!(authors.len(), 2);
     assert_eq!(authors[0].name, "Alice");
     assert_eq!(authors[1].name, "Bob");
 
-    q.batch_delete_author(vec![1, 2])
+    let _: Vec<()> = q
+        .batch_delete_author(vec![1, 2])
+        .try_collect()
         .await
         .expect("batch delete");