@@ -241,6 +241,46 @@ impl Config {
241241 }
242242 }
243243
244+ let databases = self
245+ . databases
246+ . iter ( )
247+ . map ( |database| database. name . clone ( ) )
248+ . collect :: < HashSet < _ > > ( ) ;
249+
250+ // Automatically configure system catalogs
251+ // as omnisharded.
252+ if self . general . system_catalogs_omnisharded {
253+ for database in databases {
254+ let entry = tables. entry ( database) . or_insert_with ( Vec :: new) ;
255+
256+ for table in [
257+ "pg_class" ,
258+ "pg_attribute" ,
259+ "pg_attrdef" ,
260+ "pg_index" ,
261+ "pg_constraint" ,
262+ "pg_namespace" ,
263+ "pg_database" ,
264+ "pg_tablespace" ,
265+ "pg_type" ,
266+ "pg_proc" ,
267+ "pg_operator" ,
268+ "pg_cast" ,
269+ "pg_enum" ,
270+ "pg_range" ,
271+ "pg_authid" ,
272+ "pg_am" ,
273+ ] {
274+ if entry. iter ( ) . find ( |t| t. name == table) . is_none ( ) {
275+ entry. push ( OmnishardedTable {
276+ name : table. to_string ( ) ,
277+ sticky_routing : true ,
278+ } ) ;
279+ }
280+ }
281+ }
282+ }
283+
244284 tables
245285 }
246286
@@ -725,6 +765,7 @@ password = "users_admin_password"
725765[general]
726766host = "0.0.0.0"
727767port = 6432
768+ system_catalogs_omnisharded = false
728769
729770[[databases]]
730771name = "db1"
@@ -772,4 +813,66 @@ tables = ["table_x"]
772813 assert_eq ! ( db2_tables[ 0 ] . name, "table_x" ) ;
773814 assert ! ( !db2_tables[ 0 ] . sticky_routing) ;
774815 }
816+
817+ #[ test]
818+ fn test_omnisharded_tables_system_catalogs ( ) {
819+ // Test with system_catalogs_omnisharded = true
820+ let source_enabled = r#"
821+ [general]
822+ host = "0.0.0.0"
823+ port = 6432
824+ system_catalogs_omnisharded = true
825+
826+ [[databases]]
827+ name = "db1"
828+ host = "127.0.0.1"
829+ port = 5432
830+
831+ [[omnisharded_tables]]
832+ database = "db1"
833+ tables = ["my_table"]
834+ "# ;
835+
836+ let config: Config = toml:: from_str ( source_enabled) . unwrap ( ) ;
837+ let tables = config. omnisharded_tables ( ) ;
838+ let db1_tables = tables. get ( "db1" ) . unwrap ( ) ;
839+
840+ // Should include my_table plus system catalogs
841+ assert ! ( db1_tables. iter( ) . any( |t| t. name == "my_table" ) ) ;
842+ assert ! ( db1_tables. iter( ) . any( |t| t. name == "pg_class" ) ) ;
843+ assert ! ( db1_tables. iter( ) . any( |t| t. name == "pg_attribute" ) ) ;
844+ assert ! ( db1_tables. iter( ) . any( |t| t. name == "pg_namespace" ) ) ;
845+ assert ! ( db1_tables. iter( ) . any( |t| t. name == "pg_type" ) ) ;
846+
847+ // System catalogs should have sticky_routing = true
848+ let pg_class = db1_tables. iter ( ) . find ( |t| t. name == "pg_class" ) . unwrap ( ) ;
849+ assert ! ( pg_class. sticky_routing) ;
850+
851+ // Test with system_catalogs_omnisharded = false
852+ let source_disabled = r#"
853+ [general]
854+ host = "0.0.0.0"
855+ port = 6432
856+ system_catalogs_omnisharded = false
857+
858+ [[databases]]
859+ name = "db1"
860+ host = "127.0.0.1"
861+ port = 5432
862+
863+ [[omnisharded_tables]]
864+ database = "db1"
865+ tables = ["my_table"]
866+ "# ;
867+
868+ let config: Config = toml:: from_str ( source_disabled) . unwrap ( ) ;
869+ let tables = config. omnisharded_tables ( ) ;
870+ let db1_tables = tables. get ( "db1" ) . unwrap ( ) ;
871+
872+ // Should only include my_table, no system catalogs
873+ assert_eq ! ( db1_tables. len( ) , 1 ) ;
874+ assert_eq ! ( db1_tables[ 0 ] . name, "my_table" ) ;
875+ assert ! ( !db1_tables. iter( ) . any( |t| t. name == "pg_class" ) ) ;
876+ assert ! ( !db1_tables. iter( ) . any( |t| t. name == "pg_attribute" ) ) ;
877+ }
775878}
0 commit comments