[Postgres] Add support for LISTEN/NOTIFY

.github/workflows/rust.yml

@@ -59,6 +59,16 @@ jobs:
     # unit test: tokio
     - run: cargo test --manifest-path sqlx-core/Cargo.toml --no-default-features --features 'chrono uuid postgres mysql tls runtime-tokio'
 
+    # integration test: sqlite + async-std
+    - run: cargo test --no-default-features --features 'runtime-async-std sqlite macros uuid chrono tls'
+      env:
+        DATABASE_URL: "sqlite::memory:"
+
+    # integration test: sqlite + tokio
+    - run: cargo test --no-default-features --features 'runtime-tokio sqlite macros uuid chrono tls'
+      env:
+        DATABASE_URL: "sqlite::memory:"
+
     # Rust ------------------------------------------------
 
     - name: Prepare build directory for cache

Cargo.toml

@@ -4,6 +4,7 @@ members = [
     "sqlx-core",
     "sqlx-macros",
     "sqlx-test",
+    "examples/listen-postgres",
     "examples/realworld-postgres",
     "examples/todos-postgres",
 ]
@@ -42,6 +43,7 @@ runtime-tokio = [ "sqlx-core/runtime-tokio", "sqlx-macros/runtime-tokio" ]
 # database
 postgres = [ "sqlx-core/postgres", "sqlx-macros/postgres" ]
 mysql = [ "sqlx-core/mysql", "sqlx-macros/mysql" ]
+sqlite = [ "sqlx-core/sqlite", "sqlx-macros/sqlite" ]
 
 # types
 chrono = [ "sqlx-core/chrono", "sqlx-macros/chrono" ]
@@ -70,6 +72,18 @@ required-features = [ "postgres", "macros" ]
 name = "mysql-macros"
 required-features = [ "mysql", "macros" ]
 
+[[test]]
+name = "sqlite"
+required-features = [ "sqlite" ]
+
+[[test]]
+name = "sqlite-raw"
+required-features = [ "sqlite" ]
+
+[[test]]
+name = "sqlite-types"
+required-features = [ "sqlite" ]
+
 [[test]]
 name = "mysql"
 required-features = [ "mysql" ]

examples/listen-postgres/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "sqlx-example-listen-postgres"
+version = "0.1.0"
+edition = "2018"
+workspace = "../.."
+
+[dependencies]
+async-std = { version = "1.4.0", features = [ "attributes", "unstable" ] }
+sqlx = { path = "../..", features = [ "postgres", "tls" ] }
+futures = "0.3.1"

examples/listen-postgres/README.md

@@ -0,0 +1,18 @@
+Postgres LISTEN/NOTIFY
+======================
+
+## Usage
+
+Declare the database URL. This example does not include any reading or writing of data.
+
+```
+export DATABASE_URL="postgres://postgres@localhost/postgres"
+```
+
+Run.
+
+```
+cargo run
+```
+
+The example program should connect to the database, and create a LISTEN loop on a predefined set of channels. A NOTIFY task will be spawned which will connect to the same database and will emit notifications on a 5 second interval.

examples/listen-postgres/src/main.rs

@@ -0,0 +1,68 @@
+use async_std::stream;
+use futures::StreamExt;
+use futures::TryStreamExt;
+use sqlx::postgres::PgListener;
+use sqlx::{Executor, PgPool};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::time::Duration;
+
+#[async_std::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    println!("Building PG pool.");
+    let conn_str =
+        std::env::var("DATABASE_URL").expect("Env var DATABASE_URL is required for this example.");
+    let pool = sqlx::PgPool::new(&conn_str).await?;
+
+    let mut listener = PgListener::new(&conn_str).await?;
+
+    // let notify_pool = pool.clone();
+    let _t = async_std::task::spawn(async move {
+        stream::interval(Duration::from_secs(2))
+            .for_each(|_| notify(&pool))
+            .await
+    });
+
+    println!("Starting LISTEN loop.");
+
+    listener.listen_all(&["chan0", "chan1", "chan2"]).await?;
+
+    let mut counter = 0usize;
+    loop {
+        let notification = listener.recv().await?;
+        println!("[from recv]: {:?}", notification);
+
+        counter += 1;
+        if counter >= 3 {
+            break;
+        }
+    }
+
+    // Prove that we are buffering messages by waiting for 6 seconds
+    listener.execute("SELECT pg_sleep(6)").await?;
+
+    let mut stream = listener.into_stream();
+    while let Some(notification) = stream.try_next().await? {
+        println!("[from stream]: {:?}", notification);
+    }
+
+    Ok(())
+}
+
+async fn notify(mut pool: &PgPool) {
+    static COUNTER: AtomicUsize = AtomicUsize::new(0);
+
+    let res = pool
+        .execute(&*format!(
+            r#"
+NOTIFY "chan0", '{{"payload": {}}}';
+NOTIFY "chan1", '{{"payload": {}}}';
+NOTIFY "chan2", '{{"payload": {}}}';
+                "#,
+            COUNTER.fetch_add(1, Ordering::SeqCst),
+            COUNTER.fetch_add(1, Ordering::SeqCst),
+            COUNTER.fetch_add(1, Ordering::SeqCst)
+        ))
+        .await;
+
+    println!("[from notify]: {:?}", res);
+}

sqlx-core/Cargo.toml

@@ -17,13 +17,14 @@ default = [ "runtime-async-std" ]
 unstable = []
 postgres = [ "md-5", "sha2", "base64", "sha-1", "rand", "hmac" ]
 mysql = [ "sha-1", "sha2", "generic-array", "num-bigint", "base64", "digest", "rand" ]
+sqlite = [ "libsqlite3-sys" ]
 tls = [ "async-native-tls" ]
 runtime-async-std = [ "async-native-tls/runtime-async-std", "async-std" ]
 runtime-tokio = [ "async-native-tls/runtime-tokio", "tokio" ]
 
 [dependencies]
 async-native-tls = { version = "0.3.2", default-features = false, optional = true }
-async-std = { version = "1.5.0", optional = true }
+async-std = { version = "1.5.0", features = [ "unstable" ], optional = true }
 async-stream = { version = "0.2.1", default-features = false }
 base64 = { version = "0.11.0", default-features = false, optional = true, features = [ "std" ] }
 bitflags = { version = "1.2.1", default-features = false }
@@ -50,5 +51,12 @@ tokio = { version = "0.2.13", default-features = false, features = [ "dns", "fs"
 url = { version = "2.1.1", default-features = false }
 uuid = { version = "0.8.1", default-features = false, optional = true, features = [ "std" ] }
 
+# <https://github.com/jgallagher/rusqlite/tree/master/libsqlite3-sys>
+[dependencies.libsqlite3-sys]
+version = "0.17.1"
+optional = true
+default-features = false
+features = [ "pkg-config", "vcpkg", "bundled" ]
+
 [dev-dependencies]
-matches = "0.1.8"
+matches = "0.1.8"

sqlx-core/src/connection.rs

@@ -4,6 +4,7 @@ use futures_core::future::BoxFuture;
 
 use crate::executor::Executor;
 use crate::pool::{Pool, PoolConnection};
+use crate::transaction::Transaction;
 use crate::url::Url;
 
 /// Represents a single database connection rather than a pool of database connections.
@@ -15,6 +16,16 @@ where
     Self: Send + 'static,
     Self: Executor,
 {
+    /// Starts a transaction.
+    ///
+    /// Returns [`Transaction`](struct.Transaction.html).
+    fn begin(self) -> BoxFuture<'static, crate::Result<Transaction<Self>>>
+    where
+        Self: Sized,
+    {
+        Box::pin(Transaction::new(0, self))
+    }
+
     /// Close this database connection.
     fn close(self) -> BoxFuture<'static, crate::Result<()>>;
 
@@ -31,63 +42,61 @@ pub trait Connect: Connection {
         Self: Sized;
 }
 
-mod internal {
-    #[allow(dead_code)]
-    pub enum MaybeOwnedConnection<'c, C>
-    where
-        C: super::Connect,
-    {
-        Borrowed(&'c mut C),
-        Owned(super::PoolConnection<C>),
-    }
-
-    #[allow(dead_code)]
-    pub enum ConnectionSource<'c, C>
-    where
-        C: super::Connect,
-    {
-        Connection(MaybeOwnedConnection<'c, C>),
-        Pool(super::Pool<C>),
-    }
+#[allow(dead_code)]
+pub(crate) enum ConnectionSource<'c, C>
+where
+    C: Connect,
+{
+    ConnectionRef(&'c mut C),
+    Connection(C),
+    PoolConnection(Pool<C>, PoolConnection<C>),
+    Pool(Pool<C>),
 }
 
-pub(crate) use self::internal::{ConnectionSource, MaybeOwnedConnection};
-
 impl<'c, C> ConnectionSource<'c, C>
 where
     C: Connect,
 {
     #[allow(dead_code)]
-    pub(crate) async fn resolve_by_ref(&mut self) -> crate::Result<&'_ mut C> {
+    pub(crate) async fn resolve(&mut self) -> crate::Result<&'_ mut C> {
         if let ConnectionSource::Pool(pool) = self {
-            *self =
-                ConnectionSource::Connection(MaybeOwnedConnection::Owned(pool.acquire().await?));
+            let conn = pool.acquire().await?;
+
+            *self = ConnectionSource::PoolConnection(pool.clone(), conn);
         }
 
         Ok(match self {
-            ConnectionSource::Connection(conn) => match conn {
-                MaybeOwnedConnection::Borrowed(conn) => &mut *conn,
-                MaybeOwnedConnection::Owned(ref mut conn) => conn,
-            },
+            ConnectionSource::ConnectionRef(conn) => conn,
+            ConnectionSource::PoolConnection(_, ref mut conn) => conn,
+            ConnectionSource::Connection(ref mut conn) => conn,
             ConnectionSource::Pool(_) => unreachable!(),
         })
     }
 }
 
-impl<'c, C> From<&'c mut C> for MaybeOwnedConnection<'c, C>
+impl<'c, C> From<C> for ConnectionSource<'c, C>
+where
+    C: Connect,
+{
+    fn from(connection: C) -> Self {
+        ConnectionSource::Connection(connection)
+    }
+}
+
+impl<'c, C> From<PoolConnection<C>> for ConnectionSource<'c, C>
 where
     C: Connect,
 {
-    fn from(conn: &'c mut C) -> Self {
-        MaybeOwnedConnection::Borrowed(conn)
+    fn from(connection: PoolConnection<C>) -> Self {
+        ConnectionSource::PoolConnection(Pool(connection.pool.clone()), connection)
     }
 }
 
-impl<'c, C> From<PoolConnection<C>> for MaybeOwnedConnection<'c, C>
+impl<'c, C> From<Pool<C>> for ConnectionSource<'c, C>
 where
     C: Connect,
 {
-    fn from(conn: PoolConnection<C>) -> Self {
-        MaybeOwnedConnection::Owned(conn)
+    fn from(pool: Pool<C>) -> Self {
+        ConnectionSource::Pool(pool)
     }
 }

sqlx-core/src/cursor.rs

@@ -1,6 +1,5 @@
 use futures_core::future::BoxFuture;
 
-use crate::connection::{Connect, MaybeOwnedConnection};
 use crate::database::{Database, HasRow};
 use crate::executor::Execute;
 use crate::pool::Pool;
@@ -19,20 +18,21 @@ where
 {
     type Database: Database;
 
-    #[doc(hidden)]
     fn from_pool<E>(pool: &Pool<<Self::Database as Database>::Connection>, query: E) -> Self
     where
         Self: Sized,
         E: Execute<'q, Self::Database>;
 
-    #[doc(hidden)]
-    fn from_connection<E, C>(conn: C, query: E) -> Self
+    fn from_connection<E>(
+        connection: &'c mut <Self::Database as Database>::Connection,
+        query: E,
+    ) -> Self
     where
         Self: Sized,
-        <Self::Database as Database>::Connection: Connect,
-        C: Into<MaybeOwnedConnection<'c, <Self::Database as Database>::Connection>>,
         E: Execute<'q, Self::Database>;
 
     /// Fetch the next row in the result. Returns `None` if there are no more rows.
-    fn next(&mut self) -> BoxFuture<crate::Result<Option<<Self::Database as HasRow>::Row>>>;
+    fn next<'cur>(
+        &'cur mut self,
+    ) -> BoxFuture<'cur, crate::Result<Option<<Self::Database as HasRow<'cur>>::Row>>>;
 }

sqlx-core/src/database.rs

@@ -28,6 +28,8 @@ where
 
     /// The Rust type of table identifiers for this database.
     type TableId: Display + Clone;
+
+    type RawBuffer;
 }
 
 pub trait HasRawValue<'c> {