sqlx-macros-core/src/database/mod.rs - main - germanubis/sqlx

sqlx-macros-core/src/database/mod.rs
106 строк3 KB
Sean Lynch
Drop cached db connections in macros upon hitting an error (#4009)
29 авг 2025, 06:37
29 авг 2025, 06:3769bb595
use sqlx_core::config;
use sqlx_core::connection::Connection;
use sqlx_core::database::Database;
use sqlx_core::describe::Describe;
use sqlx_core::executor::Executor;
use sqlx_core::sql_str::AssertSqlSafe;
use sqlx_core::sql_str::SqlSafeStr;
use sqlx_core::type_checking::TypeChecking;
use std::collections::hash_map;
use std::collections::HashMap;
use std::sync::{LazyLock, Mutex};

#[cfg(any(feature = "postgres", feature = "mysql", feature = "_sqlite"))]
mod impls;

pub trait DatabaseExt: Database + TypeChecking {
    const DATABASE_PATH: &'static str;
    const ROW_PATH: &'static str;

    fn db_path() -> syn::Path {
        syn::parse_str(Self::DATABASE_PATH).unwrap()
    }

    fn row_path() -> syn::Path {
        syn::parse_str(Self::ROW_PATH).unwrap()
    }

    fn describe_blocking(
        query: &str,
        database_url: &str,
        driver_config: &config::drivers::Config,
    ) -> sqlx_core::Result<Describe<Self>>;
}

#[allow(dead_code)]
pub struct CachingDescribeBlocking<DB: DatabaseExt> {
    connections: LazyLock<Mutex<HashMap<String, DB::Connection>>>,
}

#[allow(dead_code)]
impl<DB: DatabaseExt> CachingDescribeBlocking<DB> {
    #[allow(clippy::new_without_default, reason = "internal API")]
    pub const fn new() -> Self {
        CachingDescribeBlocking {
            connections: LazyLock::new(|| Mutex::new(HashMap::new())),
        }
    }

    pub fn describe(
        &self,
        query: &str,
        database_url: &str,
        _driver_config: &config::drivers::Config,
    ) -> sqlx_core::Result<Describe<DB>>
    where
        for<'a> &'a mut DB::Connection: Executor<'a, Database = DB>,
    {
        let mut cache = self
            .connections
            .lock()
            .expect("previous panic in describe call");

        crate::block_on(async {
            let conn = match cache.entry(database_url.to_string()) {
                hash_map::Entry::Occupied(hit) => hit.into_mut(),
                hash_map::Entry::Vacant(miss) => {
                    let conn = miss.insert(DB::Connection::connect(database_url).await?);

                    #[cfg(feature = "postgres")]
                    if DB::NAME == sqlx_postgres::Postgres::NAME {
                        conn.execute(
                            "
                            DO $$
                            BEGIN
                                IF EXISTS (
                                    SELECT 1
                                    FROM pg_settings
                                    WHERE name = 'plan_cache_mode'
                                ) THEN
                                    SET SESSION plan_cache_mode = 'force_generic_plan';
                                END IF;
                            END $$;
                        ",
                        )
                        .await?;
                    }
                    conn
                }
            };

            match conn
                .describe(AssertSqlSafe(query.to_string()).into_sql_str())
                .await
            {
                Ok(describe) => Ok(describe),
                Err(e) => {
                    if matches!(e, sqlx_core::Error::Io(_) | sqlx_core::Error::Protocol(_)) {
                        cache.remove(database_url);
                    }

                    Err(e)
                }
            }
        })
    }
}