toasty_driver_postgresql/
lib.rs

1#![warn(missing_docs)]
2
3//! Toasty driver for [PostgreSQL](https://www.postgresql.org/) using
4//! [`tokio-postgres`](https://docs.rs/tokio-postgres).
5//!
6//! # Examples
7//!
8//! ```no_run
9//! use toasty_driver_postgresql::PostgreSQL;
10//!
11//! let driver = PostgreSQL::new("postgresql://localhost/mydb").unwrap();
12//! ```
13
14mod oid_cache;
15mod statement_cache;
16#[cfg(feature = "tls")]
17mod tls;
18mod r#type;
19mod value;
20
21pub(crate) use value::Value;
22
23use async_trait::async_trait;
24use percent_encoding::percent_decode_str;
25use std::{borrow::Cow, sync::Arc};
26use toasty_core::{
27    Result, Schema,
28    driver::{
29        Capability, Driver, ExecResponse, Operation,
30        operation::{RawSqlRet, Transaction, TransactionMode, TypedValue},
31    },
32    schema::{
33        db::{self, Migration, Table},
34        diff,
35    },
36    stmt,
37    stmt::ValueRecord,
38};
39use toasty_sql::{self as sql};
40use tokio_postgres::{Client, Config, Socket, tls::MakeTlsConnect, types::ToSql};
41use url::Url;
42
43enum SqlReturn {
44    Count,
45    Infer,
46    Types(Vec<stmt::Type>),
47}
48
49use crate::{oid_cache::OidCache, statement_cache::StatementCache};
50
51/// Classifies a `tokio_postgres::Error` into a Toasty error.
52///
53/// Errors that carry a server-side `DbError` are mapped to typed
54/// variants where one exists (`SerializationFailure`,
55/// `ReadOnlyTransaction`); everything else with a `DbError` becomes
56/// `DriverOperationFailed`. Errors *without* a `DbError` are
57/// classified as `ConnectionLost`: per `tokio-postgres`, those
58/// originate from the underlying socket or protocol layer (closed
59/// socket, IO error, end-of-stream during handshake), which the
60/// pool treats as evictable.
61fn classify_pg_error(e: tokio_postgres::Error) -> toasty_core::Error {
62    if let Some(db_err) = e.as_db_error() {
63        match db_err.code().code() {
64            "40001" => toasty_core::Error::serialization_failure(db_err.message()),
65            "25006" => toasty_core::Error::read_only_transaction(db_err.message()),
66            _ => toasty_core::Error::driver_operation_failed(e),
67        }
68    } else {
69        toasty_core::Error::connection_lost(e)
70    }
71}
72
73/// A PostgreSQL [`Driver`] that connects via `tokio-postgres`.
74///
75/// # Examples
76///
77/// ```no_run
78/// use toasty_driver_postgresql::PostgreSQL;
79///
80/// let driver = PostgreSQL::new("postgresql://localhost/mydb").unwrap();
81/// ```
82pub struct PostgreSQL {
83    url: String,
84    config: Config,
85    #[cfg(feature = "tls")]
86    tls: Option<tokio_postgres_rustls::MakeRustlsConnect>,
87}
88
89impl std::fmt::Debug for PostgreSQL {
90    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91        let Self {
92            url,
93            config,
94            #[cfg(feature = "tls")]
95            tls,
96        } = self;
97        let mut s = f.debug_struct("PostgreSQL");
98        s.field("url", url);
99        s.field("config", config);
100        #[cfg(feature = "tls")]
101        s.field("tls", &tls.as_ref().map(|_| "MakeRustlsConnect"));
102        s.finish()
103    }
104}
105
106impl PostgreSQL {
107    /// Create a new PostgreSQL driver from a connection URL
108    pub fn new(url: impl Into<String>) -> Result<Self> {
109        let url_str = url.into();
110        let url = Url::parse(&url_str).map_err(toasty_core::Error::driver_operation_failed)?;
111
112        if !matches!(url.scheme(), "postgresql" | "postgres") {
113            return Err(toasty_core::Error::invalid_connection_url(format!(
114                "connection URL does not have a `postgresql` scheme; url={}",
115                url
116            )));
117        }
118
119        if url.path().is_empty() {
120            return Err(toasty_core::Error::invalid_connection_url(format!(
121                "no database specified - missing path in connection URL; url={}",
122                url
123            )));
124        }
125
126        let mut config = Config::new();
127
128        let dbname = percent_decode_str(url.path().trim_start_matches('/'))
129            .decode_utf8()
130            .map_err(|_| {
131                toasty_core::Error::invalid_connection_url("database name is not valid UTF-8")
132            })?;
133        config.dbname(&*dbname);
134
135        if !url.username().is_empty() {
136            let user = percent_decode_str(url.username())
137                .decode_utf8()
138                .map_err(|_| {
139                    toasty_core::Error::invalid_connection_url("username is not valid UTF-8")
140                })?;
141            config.user(&*user);
142        }
143
144        if let Some(password) = url.password() {
145            config.password(percent_decode_str(password).collect::<Vec<u8>>());
146        }
147
148        // libpq lets standard connection parameters appear in the query
149        // string; honor the ones a Toasty user can reasonably set so that
150        // `postgresql:///mydb?host=/tmp&user=alice` reaches the server.
151        // Single-valued setters (user, password, dbname, application_name)
152        // replace earlier calls, so we can apply them inline; host and
153        // port are list-valued — staged into Options below so a query
154        // parameter cleanly overrides the URL component instead of being
155        // appended as a fallback tokio-postgres would try first.
156        let mut host: Option<String> = None;
157        let mut port: Option<u16> = None;
158
159        for (key, value) in url.query_pairs() {
160            match key.as_ref() {
161                "host" => host = Some(value.into_owned()),
162                "port" => {
163                    port = Some(value.parse::<u16>().map_err(|_| {
164                        toasty_core::Error::invalid_connection_url(format!(
165                            "invalid port in connection URL query parameter: {value}"
166                        ))
167                    })?);
168                }
169                "user" => {
170                    config.user(&*value);
171                }
172                "password" => {
173                    config.password(value.as_bytes());
174                }
175                "dbname" => {
176                    config.dbname(&*value);
177                }
178                "application_name" => {
179                    config.application_name(&*value);
180                }
181                _ => {}
182            }
183        }
184
185        let host = host
186            .or_else(|| url.host_str().filter(|h| !h.is_empty()).map(String::from))
187            .ok_or_else(|| {
188                toasty_core::Error::invalid_connection_url(format!(
189                    "missing host in connection URL; url={}",
190                    url
191                ))
192            })?;
193        config.host(&host);
194
195        if let Some(port) = port.or_else(|| url.port()) {
196            config.port(port);
197        }
198
199        #[cfg(feature = "tls")]
200        let tls = tls::configure_tls(&url, &mut config)?;
201
202        #[cfg(not(feature = "tls"))]
203        for (key, value) in url.query_pairs() {
204            if key == "sslmode" && value != "disable" {
205                return Err(toasty_core::Error::invalid_connection_url(
206                    "TLS not available: compile with the `tls` feature",
207                ));
208            }
209        }
210
211        Ok(Self {
212            url: url_str,
213            config,
214            #[cfg(feature = "tls")]
215            tls,
216        })
217    }
218
219    async fn connect_with_config(&self, config: Config) -> Result<Connection> {
220        #[cfg(feature = "tls")]
221        if let Some(ref tls) = self.tls {
222            return Connection::connect(config, tls.clone()).await;
223        }
224        Connection::connect(config, tokio_postgres::NoTls).await
225    }
226}
227
228#[async_trait]
229impl Driver for PostgreSQL {
230    fn url(&self) -> Cow<'_, str> {
231        Cow::Borrowed(&self.url)
232    }
233
234    fn capability(&self) -> &'static Capability {
235        &Capability::POSTGRESQL
236    }
237
238    async fn connect(&self) -> toasty_core::Result<Box<dyn toasty_core::driver::Connection>> {
239        Ok(Box::new(
240            self.connect_with_config(self.config.clone()).await?,
241        ))
242    }
243
244    fn generate_migration(&self, schema_diff: &diff::Schema<'_>) -> Migration {
245        let statements = sql::MigrationStatement::from_diff(schema_diff, &Capability::POSTGRESQL);
246
247        let sql_strings: Vec<String> = statements
248            .iter()
249            .map(|stmt| sql::Serializer::postgresql(stmt.schema()).serialize(stmt.statement()))
250            .collect();
251
252        Migration::new_sql(sql_strings.join("\n"))
253    }
254
255    async fn reset_db(&self) -> toasty_core::Result<()> {
256        let dbname = self
257            .config
258            .get_dbname()
259            .ok_or_else(|| {
260                toasty_core::Error::invalid_connection_url("no database name configured")
261            })?
262            .to_string();
263
264        // We cannot drop a database we are currently connected to, so we need a temp database.
265        let temp_dbname = "__toasty_reset_temp";
266
267        let connect = |dbname: &str| {
268            let mut config = self.config.clone();
269            config.dbname(dbname);
270            self.connect_with_config(config)
271        };
272
273        // Step 1: Connect to the target DB and create a temp DB
274        let conn = connect(&dbname).await?;
275        conn.client
276            .execute(&format!("DROP DATABASE IF EXISTS \"{}\"", temp_dbname), &[])
277            .await
278            .map_err(classify_pg_error)?;
279        conn.client
280            .execute(&format!("CREATE DATABASE \"{}\"", temp_dbname), &[])
281            .await
282            .map_err(classify_pg_error)?;
283        drop(conn);
284
285        // Step 2: Connect to the temp DB, drop and recreate the target
286        let conn = connect(temp_dbname).await?;
287        conn.client
288            .execute(
289                "SELECT pg_terminate_backend(pid) \
290                 FROM pg_stat_activity \
291                 WHERE datname = $1 AND pid <> pg_backend_pid()",
292                &[&dbname],
293            )
294            .await
295            .map_err(classify_pg_error)?;
296        conn.client
297            .execute(&format!("DROP DATABASE IF EXISTS \"{}\"", dbname), &[])
298            .await
299            .map_err(classify_pg_error)?;
300        conn.client
301            .execute(&format!("CREATE DATABASE \"{}\"", dbname), &[])
302            .await
303            .map_err(classify_pg_error)?;
304        drop(conn);
305
306        // Step 3: Connect back to the target and clean up the temp DB
307        let conn = connect(&dbname).await?;
308        conn.client
309            .execute(&format!("DROP DATABASE IF EXISTS \"{}\"", temp_dbname), &[])
310            .await
311            .map_err(classify_pg_error)?;
312
313        Ok(())
314    }
315}
316
317/// An open connection to a PostgreSQL database.
318#[derive(Debug)]
319pub struct Connection {
320    client: Client,
321    statement_cache: StatementCache,
322    oid_cache: OidCache,
323}
324
325impl Connection {
326    /// Initialize a Toasty PostgreSQL connection using an initialized client.
327    pub fn new(client: Client) -> Self {
328        Self {
329            client,
330            statement_cache: StatementCache::new(100),
331            oid_cache: OidCache::new(),
332        }
333    }
334
335    /// Connects to a PostgreSQL database using a [`postgres::Config`].
336    ///
337    /// See [`postgres::Client::configure`] for more information.
338    pub async fn connect<T>(config: Config, tls: T) -> Result<Self>
339    where
340        T: MakeTlsConnect<Socket> + 'static,
341        T::Stream: Send,
342    {
343        let (client, connection) = config.connect(tls).await.map_err(classify_pg_error)?;
344
345        tokio::spawn(async move {
346            if let Err(e) = connection.await {
347                eprintln!("connection error: {e}");
348            }
349        });
350
351        Ok(Self::new(client))
352    }
353
354    async fn exec_sql(
355        &mut self,
356        sql_as_str: &str,
357        typed_params: Vec<TypedValue>,
358        ret: SqlReturn,
359    ) -> Result<ExecResponse> {
360        tracing::debug!(db.system = "postgresql", db.statement = %sql_as_str, params = typed_params.len(), "executing SQL");
361
362        self.oid_cache
363            .preload(&self.client, typed_params.iter().map(|tv| &tv.ty))
364            .await?;
365        let param_types: Vec<_> = typed_params
366            .iter()
367            .map(|tv| self.oid_cache.get(&tv.ty).clone())
368            .collect();
369
370        let values: Vec<_> = typed_params
371            .into_iter()
372            .map(|tv| Value::from(tv.value))
373            .collect();
374        let params = values
375            .iter()
376            .map(|param| param as &(dyn ToSql + Sync))
377            .collect::<Vec<_>>();
378
379        let statement = self
380            .statement_cache
381            .prepare_typed(&mut self.client, sql_as_str, &param_types)
382            .await
383            .map_err(classify_pg_error)?;
384
385        if matches!(ret, SqlReturn::Count) {
386            let count = self
387                .client
388                .execute(&statement, &params)
389                .await
390                .map_err(classify_pg_error)?;
391            return Ok(ExecResponse::count(count));
392        }
393
394        let rows = self
395            .client
396            .query(&statement, &params)
397            .await
398            .map_err(classify_pg_error)?;
399
400        let results = rows.into_iter().map(move |row| {
401            let mut results = Vec::new();
402
403            match &ret {
404                SqlReturn::Count => unreachable!(),
405                SqlReturn::Infer => {
406                    for (i, column) in row.columns().iter().enumerate() {
407                        results.push(Value::from_sql_infer(i, &row, column).into_inner());
408                    }
409                }
410                SqlReturn::Types(ret_tys) => {
411                    for (i, column) in row.columns().iter().enumerate() {
412                        results.push(Value::from_sql(i, &row, column, &ret_tys[i]).into_inner());
413                    }
414                }
415            }
416
417            Ok(ValueRecord::from_vec(results))
418        });
419
420        Ok(ExecResponse::value_stream(stmt::ValueStream::from_iter(
421            results,
422        )))
423    }
424
425    /// Creates a table.
426    pub async fn create_table(&mut self, schema: &db::Schema, table: &Table) -> Result<()> {
427        let serializer = sql::Serializer::postgresql(schema);
428
429        let sql = serializer.serialize(&sql::Statement::create_table(
430            table,
431            &Capability::POSTGRESQL,
432        ));
433
434        self.client
435            .execute(&sql, &[])
436            .await
437            .map_err(classify_pg_error)?;
438
439        for index in &table.indices {
440            if index.primary_key {
441                continue;
442            }
443
444            let sql = serializer.serialize(&sql::Statement::create_index(index));
445
446            self.client
447                .execute(&sql, &[])
448                .await
449                .map_err(classify_pg_error)?;
450        }
451
452        Ok(())
453    }
454}
455
456impl From<Client> for Connection {
457    fn from(client: Client) -> Self {
458        Self::new(client)
459    }
460}
461
462#[async_trait]
463impl toasty_core::driver::Connection for Connection {
464    async fn exec(&mut self, schema: &Arc<Schema>, op: Operation) -> Result<ExecResponse> {
465        tracing::trace!(driver = "postgresql", op = %op.name(), "driver exec");
466
467        if let Operation::Transaction(ref t) = op {
468            // PostgreSQL has no `BEGIN IMMEDIATE` / `BEGIN EXCLUSIVE`
469            // analogue; reject non-Default modes loudly rather than
470            // silently dropping them at the serializer.
471            if let Transaction::Start {
472                mode: mode @ (TransactionMode::Immediate | TransactionMode::Exclusive),
473                ..
474            } = t
475            {
476                return Err(toasty_core::Error::unsupported_feature(format!(
477                    "PostgreSQL does not support TransactionMode::{mode:?}"
478                )));
479            }
480            let sql = sql::Serializer::postgresql(&schema.db).serialize_transaction(t);
481            self.client
482                .batch_execute(&sql)
483                .await
484                .map_err(classify_pg_error)?;
485            return Ok(ExecResponse::count(0));
486        }
487
488        let (sql, typed_params, ret_tys) = match op {
489            Operation::Insert(op) => (sql::Statement::from(op.stmt), op.params, None),
490            Operation::QuerySql(query) => {
491                assert!(
492                    query.last_insert_id_hack.is_none(),
493                    "last_insert_id_hack is MySQL-specific and should not be set for PostgreSQL"
494                );
495                (sql::Statement::from(query.stmt), query.params, query.ret)
496            }
497            Operation::RawSql(op) => {
498                let ret = match op.ret {
499                    RawSqlRet::None => SqlReturn::Count,
500                    RawSqlRet::Infer => SqlReturn::Infer,
501                    RawSqlRet::Types(types) => SqlReturn::Types(types),
502                };
503                return self.exec_sql(&op.sql, op.params, ret).await;
504            }
505            op => todo!("op={:#?}", op),
506        };
507
508        let sql_as_str = sql::Serializer::postgresql(&schema.db).serialize(&sql);
509
510        let ret = if sql.returning_len().is_some() {
511            SqlReturn::Types(ret_tys.unwrap())
512        } else {
513            SqlReturn::Count
514        };
515
516        self.exec_sql(&sql_as_str, typed_params, ret).await
517    }
518
519    async fn push_schema(&mut self, schema: &Schema) -> Result<()> {
520        let serializer = sql::Serializer::postgresql(&schema.db);
521
522        // Create PostgreSQL enum types before creating tables.
523        // Collect unique enum types across all columns.
524        let mut created_enum_types = hashbrown::HashSet::new();
525        for table in &schema.db.tables {
526            for column in &table.columns {
527                if let toasty_core::schema::db::Type::Enum(type_enum) = &column.storage_ty
528                    && created_enum_types.insert(type_enum.name.clone())
529                {
530                    let sql = serializer.serialize(&sql::Statement::create_enum_type(type_enum));
531
532                    tracing::debug!(enum_type = ?type_enum.name, "creating enum type");
533                    self.client
534                        .execute(&sql, &[])
535                        .await
536                        .map_err(classify_pg_error)?;
537                }
538            }
539        }
540
541        for table in &schema.db.tables {
542            tracing::debug!(table = %table.name, "creating table");
543            self.create_table(&schema.db, table).await?;
544        }
545        Ok(())
546    }
547
548    async fn applied_migrations(
549        &mut self,
550    ) -> Result<Vec<toasty_core::schema::db::AppliedMigration>> {
551        // Ensure the migrations table exists
552        self.client
553            .execute(
554                "CREATE TABLE IF NOT EXISTS __toasty_migrations (
555                id BIGINT PRIMARY KEY,
556                name TEXT NOT NULL,
557                applied_at TIMESTAMP NOT NULL
558            )",
559                &[],
560            )
561            .await
562            .map_err(classify_pg_error)?;
563
564        // Query all applied migrations
565        let rows = self
566            .client
567            .query(
568                "SELECT id FROM __toasty_migrations ORDER BY applied_at",
569                &[],
570            )
571            .await
572            .map_err(classify_pg_error)?;
573
574        Ok(rows
575            .iter()
576            .map(|row| {
577                let id: i64 = row.get(0);
578                toasty_core::schema::db::AppliedMigration::new(id as u64)
579            })
580            .collect())
581    }
582
583    async fn apply_migration(
584        &mut self,
585        id: u64,
586        name: &str,
587        migration: &toasty_core::schema::db::Migration,
588    ) -> Result<()> {
589        tracing::info!(id = id, name = %name, "applying migration");
590        // Ensure the migrations table exists
591        self.client
592            .execute(
593                "CREATE TABLE IF NOT EXISTS __toasty_migrations (
594                id BIGINT PRIMARY KEY,
595                name TEXT NOT NULL,
596                applied_at TIMESTAMP NOT NULL
597            )",
598                &[],
599            )
600            .await
601            .map_err(classify_pg_error)?;
602
603        // Start transaction
604        let transaction = self.client.transaction().await.map_err(classify_pg_error)?;
605
606        // Execute each migration statement
607        for statement in migration.statements() {
608            if let Err(e) = transaction
609                .batch_execute(statement)
610                .await
611                .map_err(classify_pg_error)
612            {
613                transaction.rollback().await.map_err(classify_pg_error)?;
614                return Err(e);
615            }
616        }
617
618        // Record the migration
619        if let Err(e) = transaction
620            .execute(
621                "INSERT INTO __toasty_migrations (id, name, applied_at) VALUES ($1, $2, NOW())",
622                &[&(id as i64), &name],
623            )
624            .await
625            .map_err(classify_pg_error)
626        {
627            transaction.rollback().await.map_err(classify_pg_error)?;
628            return Err(e);
629        }
630
631        // Commit transaction
632        transaction.commit().await.map_err(classify_pg_error)?;
633        Ok(())
634    }
635
636    fn is_valid(&self) -> bool {
637        !self.client.is_closed()
638    }
639
640    async fn ping(&mut self) -> Result<()> {
641        // An empty `simple_query` is the lightest sync round-trip in
642        // the PG protocol — it skips parsing entirely. Any failure is
643        // surfaced as `connection_lost`: the only meaningful outcome
644        // of a ping is "the connection is alive" or "evict it."
645        self.client
646            .simple_query("")
647            .await
648            .map(|_| ())
649            .map_err(toasty_core::Error::connection_lost)
650    }
651}
652
653#[cfg(test)]
654mod tests {
655    use super::*;
656    use tokio_postgres::config::Host;
657
658    fn cfg(url: &str) -> Config {
659        PostgreSQL::new(url).expect("valid URL").config
660    }
661
662    #[test]
663    fn host_in_url_authority() {
664        let c = cfg("postgresql://example.com/mydb");
665        assert_eq!(c.get_hosts(), &[Host::Tcp("example.com".into())]);
666        assert_eq!(c.get_dbname(), Some("mydb"));
667    }
668
669    // `Host::Unix` only exists on Unix targets in tokio-postgres, so the
670    // tests that assert socket-path resolution are gated to those
671    // platforms. Windows builds still exercise the URL-parsing path via
672    // the non-Unix tests below.
673    #[cfg(unix)]
674    mod unix_socket {
675        use super::*;
676        use std::path::PathBuf;
677
678        #[test]
679        fn unix_socket_via_host_query_param() {
680            // Regression for #984: libpq lets a Unix-socket directory be
681            // supplied via `?host=/path`, since URL syntax cannot put a
682            // filesystem path in the authority.
683            let c = cfg("postgresql:///mydb?host=/tmp&user=myuser");
684            assert_eq!(c.get_hosts(), &[Host::Unix(PathBuf::from("/tmp"))]);
685            assert_eq!(c.get_user(), Some("myuser"));
686            assert_eq!(c.get_dbname(), Some("mydb"));
687        }
688
689        #[test]
690        fn query_param_host_overrides_url_authority() {
691            // libpq semantics: a `host=` query parameter replaces (not
692            // appends to) the URL authority host. tokio-postgres's
693            // `Config::host` is additive across calls, so an authority
694            // host would otherwise be tried first and the configured
695            // socket reached only as a fallback.
696            let c = cfg("postgresql://example.com/mydb?host=/var/run/postgresql");
697            assert_eq!(
698                c.get_hosts(),
699                &[Host::Unix(PathBuf::from("/var/run/postgresql"))]
700            );
701        }
702
703        #[test]
704        fn query_param_port_user_password_dbname() {
705            let c = cfg(
706                "postgresql:///placeholder?host=/tmp&port=5433&user=alice&password=s3cret&dbname=real",
707            );
708            assert_eq!(c.get_hosts(), &[Host::Unix(PathBuf::from("/tmp"))]);
709            assert_eq!(c.get_ports(), &[5433]);
710            assert_eq!(c.get_user(), Some("alice"));
711            assert_eq!(c.get_password(), Some(&b"s3cret"[..]));
712            assert_eq!(c.get_dbname(), Some("real"));
713        }
714    }
715
716    #[test]
717    fn query_param_port_overrides_url_port() {
718        // `Config::port` is additive too, so a `port=` query parameter
719        // must replace the URL authority port for the same reason.
720        let c = cfg("postgresql://example.com:5432/mydb?port=5433");
721        assert_eq!(c.get_ports(), &[5433]);
722    }
723
724    #[test]
725    fn application_name_query_param() {
726        let c = cfg("postgresql://localhost/mydb?application_name=my_app");
727        assert_eq!(c.get_application_name(), Some("my_app"));
728    }
729
730    #[test]
731    fn missing_host_rejected() {
732        let err = PostgreSQL::new("postgresql:///mydb").unwrap_err();
733        assert!(
734            err.to_string().contains("missing host"),
735            "expected missing-host error, got: {err}"
736        );
737    }
738
739    #[test]
740    fn invalid_port_query_param_rejected() {
741        let err = PostgreSQL::new("postgresql:///mydb?host=/tmp&port=not-a-number").unwrap_err();
742        assert!(
743            err.to_string().contains("invalid port"),
744            "expected invalid-port error, got: {err}"
745        );
746    }
747}
toasty_driver_postgresql/lib.rs

toasty_driver_postgresql/
lib.rs
