Skip to content

Commit 5fccd65

Browse files
authored
feat: omnishard system catalogs by default (#723)
Make most system catalog tables omnisharded by default so `\d+` and other similar commands work better (no duplicate table names).
1 parent eaacfc4 commit 5fccd65

File tree

2 files changed

+111
-0
lines changed

2 files changed

+111
-0
lines changed

pgdog-config/src/core.rs

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -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]
726766
host = "0.0.0.0"
727767
port = 6432
768+
system_catalogs_omnisharded = false
728769
729770
[[databases]]
730771
name = "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
}

pgdog-config/src/general.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,9 @@ pub struct General {
194194
/// Minimum ID for unique ID generator.
195195
#[serde(default)]
196196
pub unique_id_min: u64,
197+
/// System catalogs are omnisharded?
198+
#[serde(default = "General::default_system_catalogs_omnisharded")]
199+
pub system_catalogs_omnisharded: bool,
197200
}
198201

199202
impl Default for General {
@@ -262,6 +265,7 @@ impl Default for General {
262265
lsn_check_timeout: Self::lsn_check_timeout(),
263266
lsn_check_delay: Self::lsn_check_delay(),
264267
unique_id_min: u64::default(),
268+
system_catalogs_omnisharded: Self::default_system_catalogs_omnisharded(),
265269
}
266270
}
267271
}
@@ -421,6 +425,10 @@ impl General {
421425
Self::env_or_default("PGDOG_SHUTDOWN_TIMEOUT", 60_000)
422426
}
423427

428+
fn default_system_catalogs_omnisharded() -> bool {
429+
Self::env_or_default("PGDOG_SYSTEM_CATALOGS_OMNISHARDED", true)
430+
}
431+
424432
fn default_shutdown_termination_timeout() -> Option<u64> {
425433
Self::env_option("PGDOG_SHUTDOWN_TERMINATION_TIMEOUT")
426434
}

0 commit comments

Comments
 (0)