save

Author: levkk

integration/rust/tests/integration/mod.rs

@@ -20,3 +20,4 @@ pub mod timestamp_sorting;
 pub mod tls_enforced;
 pub mod tls_reload;
 pub mod transaction_state;
+pub mod unique_id;

integration/rust/tests/integration/unique_id.rs

@@ -0,0 +1,37 @@
+use rust::setup::connections_sqlx;
+use sqlx::{Executor, Row};
+
+#[tokio::test]
+async fn unique_id_returns_bigint() -> Result<(), Box<dyn std::error::Error>> {
+    let conns = connections_sqlx().await;
+    let sharded = conns.get(1).cloned().unwrap();
+
+    // Simple query
+    let row = sharded.fetch_one("SELECT pgdog.unique_id() AS id").await?;
+    let mut id: i64 = row.get("id");
+
+    assert!(
+        id > 0,
+        "unique_id should return a positive bigint, got {id}"
+    );
+
+    for _ in 0..100 {
+        // Prepared statement
+        let row = sqlx::query("SELECT pgdog.unique_id() AS id")
+            .fetch_one(&sharded)
+            .await?;
+        let prepared_id: i64 = row.get("id");
+        assert!(
+            prepared_id > 0,
+            "prepared unique_id should return a positive bigint, got {prepared_id}"
+        );
+
+        assert!(
+            prepared_id > id,
+            "prepared id should be greater than simple query id"
+        );
+        id = prepared_id;
+    }
+
+    Ok(())
+}

pgdog/src/backend/databases.rs

@@ -359,9 +359,9 @@ impl Databases {
         for cluster in self.all().values() {
             cluster.launch();
 
-            if cluster.pooler_mode() == PoolerMode::Session && cluster.router_needed() {
+            if cluster.pooler_mode() == PoolerMode::Session && cluster.use_parser() {
                 warn!(
-                    r#"user "{}" for database "{}" requires transaction mode to route queries"#,
+                    r#"user "{}" for database "{}" requires transaction mode to parse and route queries"#,
                     cluster.user(),
                     cluster.name()
                 );

pgdog/src/backend/pool/cluster.rs

@@ -413,10 +413,14 @@ impl Cluster {
         self.stats.clone()
     }
 
-    /// We'll need the query router to figure out
-    /// where a query should go.
-    pub fn router_needed(&self) -> bool {
+    /// We need to parse the query using pg_query.
+    pub fn use_parser(&self) -> bool {
         !(self.shards().len() == 1 && (self.read_only() || self.write_only()))
+            || self.query_parser_enabled
+            || self.multi_tenant.is_some()
+            || self.pub_sub_enabled()
+            || self.prepared_statements() == &PreparedStatements::Full
+            || self.dry_run
     }
 
     /// Multi-tenant config.

pgdog/src/backend/replication/logical/subscriber/context.rs

@@ -57,6 +57,7 @@ impl<'a> StreamContext<'a> {
             &self.params,
             None,
             1,
+            None,
         )?)
     }
 }

pgdog/src/frontend/client/query_engine/context.rs

@@ -4,6 +4,7 @@ use crate::{
     backend::pool::{connection::mirror::Mirror, stats::MemoryStats},
     frontend::{
         client::{timeouts::Timeouts, TransactionType},
+        router::parser::cache::CachedAst,
         Client, ClientRequest, PreparedStatements,
     },
     net::{BackendKeyData, Parameters, Stream},
@@ -38,6 +39,8 @@ pub struct QueryEngineContext<'a> {
     pub(super) rollback: bool,
     /// Omnisharded modulo.
     pub(super) omni_sticky_index: usize,
+    /// Query AST.
+    pub(super) ast: Option<CachedAst>,
 }
 
 impl<'a> QueryEngineContext<'a> {
@@ -58,6 +61,7 @@ impl<'a> QueryEngineContext<'a> {
             requests_left: 0,
             rollback: false,
             omni_sticky_index: client.omni_sticky_index,
+            ast: None,
         }
     }
 
@@ -83,6 +87,7 @@ impl<'a> QueryEngineContext<'a> {
             requests_left: 0,
             rollback: false,
             omni_sticky_index: thread_rng().gen_range(1..usize::MAX),
+            ast: None,
         }
     }
 

pgdog/src/frontend/client/query_engine/mod.rs

@@ -3,7 +3,11 @@ use crate::{
     config::config,
     frontend::{
         client::query_engine::hooks::QueryEngineHooks,
-        router::{parser::Shard, Route},
+        router::{
+            parser::Shard,
+            rewrite::{self, RewriteRequest},
+            Route,
+        },
         BufferedQuery, Client, Command, Comms, Error, Router, RouterContext, Stats,
     },
     net::{BackendKeyData, ErrorResponse, Message, Parameters},
@@ -114,6 +118,28 @@ impl QueryEngine {
         &mut self,
         context: &mut QueryEngineContext<'_>,
     ) -> Result<QueryEngineOutput, Error> {
+        // Check that we have the latest version of the config.
+        if let Some(error) = self.ensure_cluster(context.in_transaction()).await {
+            self.error_response(context, error).await?;
+            return Ok(QueryEngineOutput::Executed);
+        }
+
+        if let Ok(cluster) = self.backend.cluster() {
+            if cluster.use_parser() && context.ast.is_none() {
+                // Execute request rewrite, if needed.
+                let mut rewrite = RewriteRequest::new(
+                    context.client_request,
+                    self.backend.cluster()?,
+                    context.prepared_statements,
+                );
+                match rewrite.execute() {
+                    Ok(ast) => context.ast = Some(ast),
+                    Err(rewrite::Error::EmptyQuery) => (),
+                    Err(err) => return Err(err.into()),
+                }
+            }
+        }
+
         self.stats
             .received(context.client_request.total_message_len());
         self.set_state(State::Active); // Client is active.

pgdog/src/frontend/client/query_engine/route_query.rs

@@ -1,9 +1,41 @@
 use pgdog_config::PoolerMode;
 use tracing::trace;
 
+use crate::backend::Cluster;
+
 use super::*;
 
 impl QueryEngine {
+    pub fn cluster(&self) -> Result<&Cluster, Error> {
+        Ok(self.backend.cluster()?)
+    }
+
+    /// Check that cluster still exists.
+    pub async fn ensure_cluster(&mut self, in_transaction: bool) -> Option<ErrorResponse> {
+        if let Ok(cluster) = self.backend.cluster() {
+            let identifier = cluster.identifier();
+
+            if !in_transaction && !cluster.online() {
+                // Reload cluster config.
+                if let Err(_) = self.backend.safe_reload().await {
+                    return Some(ErrorResponse::connection(
+                        &identifier.user,
+                        &identifier.database,
+                    ));
+                }
+
+                if let Err(_) = self.backend.cluster() {
+                    return Some(ErrorResponse::connection(
+                        &identifier.user,
+                        &identifier.database,
+                    ));
+                }
+            }
+        }
+
+        None
+    }
+
     pub(super) async fn route_transaction(
         &mut self,
         context: &mut QueryEngineContext<'_>,
@@ -18,28 +50,7 @@ impl QueryEngine {
 
         // Admin doesn't have a cluster.
         let cluster = if let Ok(cluster) = self.backend.cluster() {
-            if !context.in_transaction() && !cluster.online() {
-                let identifier = cluster.identifier();
-
-                // Reload cluster config.
-                self.backend.safe_reload().await?;
-
-                match self.backend.cluster() {
-                    Ok(cluster) => cluster,
-                    Err(_) => {
-                        // Cluster is gone.
-                        self.error_response(
-                            context,
-                            ErrorResponse::connection(&identifier.user, &identifier.database),
-                        )
-                        .await?;
-
-                        return Ok(false);
-                    }
-                }
-            } else {
-                cluster
-            }
+            cluster
         } else {
             return Ok(true);
         };
@@ -51,6 +62,7 @@ impl QueryEngine {
             context.params,
             context.transaction,
             context.omni_sticky_index,
+            context.ast.as_ref(),
         )?;
         match self.router.query(router_context) {
             Ok(cmd) => {

pgdog/src/frontend/error.rs

@@ -50,6 +50,9 @@ pub enum Error {
 
     #[error("unique id: {0}")]
     UniqueId(#[from] unique_id::Error),
+
+    #[error("rewrite: {0}")]
+    Rewrite(#[from] crate::frontend::router::rewrite::Error),
 }
 
 impl Error {

pgdog/src/frontend/prepared_statements/mod.rs

@@ -9,7 +9,7 @@ use tracing::debug;
 
 use crate::{
     config::{config, PreparedStatements as PreparedStatementsLevel},
-    frontend::router::parser::RewritePlan,
+    frontend::router::parser::{cache::CachedAst, RewritePlan},
     net::{Parse, ProtocolMessage},
     stats::memory::MemoryUsage,
 };
@@ -31,6 +31,7 @@ pub struct PreparedStatements {
     pub(super) local: HashMap<String, String>,
     pub(super) level: PreparedStatementsLevel,
     pub(super) memory_used: usize,
+    pub(super) rewrite: HashMap<String, CachedAst>,
 }
 
 impl MemoryUsage for PreparedStatements {
@@ -49,6 +50,7 @@ impl Default for PreparedStatements {
             local: HashMap::default(),
             level: PreparedStatementsLevel::Extended,
             memory_used: 0,
+            rewrite: HashMap::new(),
         }
     }
 }
@@ -97,15 +99,15 @@ impl PreparedStatements {
         parse.rename_fast(&name)
     }
 
-    /// Store a rewritten statement in the global cache forever.
-    pub fn cache_rewritten(parse: &Parse) -> String {
-        let exists = Self::global().read().name(parse);
-        if let Some(exists) = exists {
-            exists
-        } else {
-            let (_, name) = Self::global().write().insert(parse);
-            name
-        }
+    /// Get original AST for a prepared statement
+    /// we have rewritten.
+    pub fn get_original_ast(&self, name: &str) -> Option<&CachedAst> {
+        self.rewrite.get(name)
+    }
+
+    /// Save original AST for re-use by subsequent Bind messages.
+    pub fn save_original_ast(&mut self, name: &str, ast: &CachedAst) {
+        self.rewrite.insert(name.to_string(), ast.clone());
     }
 
     /// Retrieve stored rewrite plan for a prepared statement, if any.