Omnisharded tables routing fix (#618)

* Fix omnishard tables detection

* Tests

* make explain work

Author: levkk

integration/d_plus/pgdog.toml

@@ -0,0 +1,25 @@
+[general]
+expanded_explain = true
+
+[[databases]]
+name = "pgdog"
+host = "127.0.0.1"
+shard = 0
+database_name = "shard_0"
+
+[[databases]]
+name = "pgdog"
+host = "127.0.0.1"
+shard = 1
+database_name = "shard_1"
+
+[admin]
+password = "pgdog"
+
+[[omnisharded_tables]]
+database = "pgdog"
+tables = [
+    "pg_class",
+    "pg_namespace",
+    "pg_am"
+]

integration/d_plus/users.toml

@@ -0,0 +1,4 @@
+[[users]]
+name = "pgdog"
+database = "pgdog"
+password = "pgdog"

pgdog/src/frontend/router/parser/from_clause.rs

@@ -3,7 +3,7 @@ use pg_query::{Node, NodeEnum};
 use super::*;
 
 /// Handle FROM <table/join> clause.
-#[derive(Copy, Clone)]
+#[derive(Copy, Clone, Debug)]
 pub struct FromClause<'a> {
     nodes: &'a [Node],
 }
@@ -64,4 +64,37 @@ impl<'a> FromClause<'a> {
 
         None
     }
+
+    /// Get all tables from the FROM clause.
+    pub fn tables(&'a self) -> Vec<Table<'a>> {
+        let mut tables = vec![];
+
+        fn tables_recursive(node: &Node) -> Vec<Table<'_>> {
+            let mut tables = vec![];
+            match node.node {
+                Some(NodeEnum::RangeVar(ref range_var)) => {
+                    tables.push(Table::from(range_var));
+                }
+
+                Some(NodeEnum::JoinExpr(ref join)) => {
+                    if let Some(ref node) = join.larg {
+                        tables.extend(tables_recursive(node));
+                    }
+                    if let Some(ref node) = join.rarg {
+                        tables.extend(tables_recursive(node));
+                    }
+                }
+
+                _ => (),
+            }
+
+            tables
+        }
+
+        for node in self.nodes {
+            tables.extend(tables_recursive(node));
+        }
+
+        tables
+    }
 }

pgdog/src/frontend/router/parser/query/select.rs

@@ -105,15 +105,36 @@ impl QueryParser {
 
         let mut query = Route::select(shard, order_by, aggregates, limit, distinct);
 
-        let mut omni = false;
+        // Omnisharded tables check.
         if query.is_all_shards() {
-            if let Some(name) = from_clause.table_name() {
-                omni = context.sharding_schema.tables.omnishards().contains(name);
-            }
-        }
+            let tables = from_clause.tables();
+            let omni = tables.iter().all(|table| {
+                context
+                    .sharding_schema
+                    .tables
+                    .omnishards()
+                    .contains(table.name)
+            });
+
+            if omni {
+                let shard = round_robin::next() % context.shards;
 
-        if omni {
-            query.set_shard_mut(round_robin::next() % context.shards);
+                query.set_shard_mut(shard);
+
+                if let Some(recorder) = self.recorder_mut() {
+                    recorder.record_entry(
+                        Some(shard.into()),
+                        format!(
+                            "SELECT matched omnisharded tables: {}",
+                            tables
+                                .iter()
+                                .map(|table| table.name)
+                                .collect::<Vec<_>>()
+                                .join(", ")
+                        ),
+                    );
+                }
+            }
         }
 
         // Only rewrite if query is cross-shard.

pgdog/src/frontend/router/parser/query/test.rs

@@ -329,6 +329,33 @@ fn test_omni() {
     assert!(matches!(route.shard(), Shard::Direct(_)));
     let (_, qp) = command!(q);
     assert!(!qp.in_transaction);
+
+    // Test that sharded tables take priority.
+    let q = "
+        SELECT
+            sharded_omni.*,
+            sharded.*
+        FROM
+            sharded_omni
+        INNER JOIN
+            sharded
+        ON sharded_omni.id = sharded.i
+    WHERE sharded.id = 5";
+
+    let route = query!(q);
+    let shard = route.shard().clone();
+
+    for _ in 0..5 {
+        let route = query!(q);
+        // Test that shard doesn't change (i.e. not round robin)
+        assert_eq!(&shard, route.shard());
+        assert!(matches!(shard, Shard::Direct(_)));
+    }
+
+    // Test that all tables have to be omnisharded.
+    let q = "SELECT * FROM sharded_omni INNER JOIN not_sharded ON sharded_omni.id = not_sharded.id WHERE sharded_omni = $1";
+    let route = query!(q);
+    assert!(matches!(route.shard(), Shard::All));
 }
 
 #[test]

pgdog/src/frontend/router/parser/where_clause.rs

@@ -10,7 +10,7 @@ use crate::frontend::router::parser::{from_clause::FromClause, Table};
 
 use super::Key;
 
-#[derive(Copy, Clone)]
+#[derive(Copy, Clone, Debug)]
 pub enum TablesSource<'a> {
     Table(Table<'a>),
     FromClause(FromClause<'a>),
@@ -54,6 +54,13 @@ impl<'a> TablesSource<'a> {
             Self::FromClause(fc) => fc.table_name(),
         }
     }
+
+    pub fn tables(&'a self) -> Vec<Table<'a>> {
+        match self {
+            Self::Table(table) => vec![*table],
+            Self::FromClause(from_clause) => from_clause.tables(),
+        }
+    }
 }
 
 #[derive(Debug)]