Skip to content

Commit be70c7e

Browse files
levkklevkk
committed
feat: fetch fk constraints when loading schema
1 parent 2a02d04 commit be70c7e

File tree

6 files changed

+199
-1
lines changed

6 files changed

+199
-1
lines changed

‎pgdog-stats/src/schema.rs‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,40 @@ use std::{collections::HashMap, hash::Hash, ops::Deref, sync::Arc};
55
/// Schema name -> Table name -> Relation
66
pub type Relations = HashMap<String, HashMap<String, Relation>>;
77

8+
/// The action to take when a referenced row is deleted or updated.
9+
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq, Hash)]
10+
pub enum ForeignKeyAction {
11+
#[default]
12+
NoAction,
13+
Restrict,
14+
Cascade,
15+
SetNull,
16+
SetDefault,
17+
}
18+
19+
impl ForeignKeyAction {
20+
/// Parse from PostgreSQL's information_schema representation.
21+
pub fn from_pg_string(s: &str) -> Self {
22+
match s {
23+
"CASCADE" => Self::Cascade,
24+
"SET NULL" => Self::SetNull,
25+
"SET DEFAULT" => Self::SetDefault,
26+
"RESTRICT" => Self::Restrict,
27+
"NO ACTION" | _ => Self::NoAction,
28+
}
29+
}
30+
}
31+
32+
/// A foreign key reference to another table's column.
33+
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Hash)]
34+
pub struct ForeignKey {
35+
pub schema: String,
36+
pub table: String,
37+
pub column: String,
38+
pub on_delete: ForeignKeyAction,
39+
pub on_update: ForeignKeyAction,
40+
}
41+
842
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Hash)]
943
pub struct Column {
1044
pub table_catalog: String,
@@ -16,6 +50,7 @@ pub struct Column {
1650
pub data_type: String,
1751
pub ordinal_position: i32,
1852
pub is_primary_key: bool,
53+
pub foreign_keys: Vec<ForeignKey>,
1954
}
2055

2156
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]

‎pgdog/src/backend/schema/columns.rs‎

Lines changed: 129 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
//! Get all table definitions.
22
pub use pgdog_stats::Column as StatsColumn;
3+
pub use pgdog_stats::ForeignKey;
4+
pub use pgdog_stats::ForeignKeyAction;
35
use serde::{Deserialize, Serialize};
46

57
use super::Error;
@@ -10,6 +12,34 @@ use std::{
1012
};
1113

1214
static COLUMNS: &str = include_str!("columns.sql");
15+
static FOREIGN_KEYS: &str = include_str!("foreign_keys.sql");
16+
17+
/// Represents a row from the foreign_keys.sql query.
18+
struct ForeignKeyRow {
19+
source_schema: String,
20+
source_table: String,
21+
source_column: String,
22+
ref_schema: String,
23+
ref_table: String,
24+
ref_column: String,
25+
on_delete: ForeignKeyAction,
26+
on_update: ForeignKeyAction,
27+
}
28+
29+
impl From<DataRow> for ForeignKeyRow {
30+
fn from(value: DataRow) -> Self {
31+
Self {
32+
source_schema: value.get_text(0).unwrap_or_default(),
33+
source_table: value.get_text(1).unwrap_or_default(),
34+
source_column: value.get_text(2).unwrap_or_default(),
35+
ref_schema: value.get_text(3).unwrap_or_default(),
36+
ref_table: value.get_text(4).unwrap_or_default(),
37+
ref_column: value.get_text(5).unwrap_or_default(),
38+
on_delete: ForeignKeyAction::from_pg_string(&value.get_text(6).unwrap_or_default()),
39+
on_update: ForeignKeyAction::from_pg_string(&value.get_text(7).unwrap_or_default()),
40+
}
41+
}
42+
}
1343

1444
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
1545
pub struct Column {
@@ -43,7 +73,7 @@ impl DerefMut for Column {
4373
}
4474

4575
impl Column {
46-
/// Load all columns from server.
76+
/// Load all columns from server, including foreign key information.
4777
pub async fn load(
4878
server: &mut Server,
4979
) -> Result<HashMap<(String, String), Vec<Column>>, Error> {
@@ -57,6 +87,26 @@ impl Column {
5787
entry.push(row);
5888
}
5989

90+
// Load foreign key constraints and merge into columns
91+
let fk_rows: Vec<ForeignKeyRow> = server.fetch_all(FOREIGN_KEYS).await?;
92+
for fk_row in fk_rows {
93+
let key = (fk_row.source_schema.clone(), fk_row.source_table.clone());
94+
if let Some(columns) = result.get_mut(&key) {
95+
if let Some(column) = columns
96+
.iter_mut()
97+
.find(|c| c.column_name == fk_row.source_column)
98+
{
99+
column.foreign_keys.push(ForeignKey {
100+
schema: fk_row.ref_schema,
101+
table: fk_row.ref_table,
102+
column: fk_row.ref_column,
103+
on_delete: fk_row.on_delete,
104+
on_update: fk_row.on_update,
105+
});
106+
}
107+
}
108+
}
109+
60110
Ok(result)
61111
}
62112
}
@@ -75,6 +125,7 @@ impl From<DataRow> for Column {
75125
data_type: value.get_text(6).unwrap_or_default(),
76126
ordinal_position: value.get::<i32>(7, Format::Text).unwrap_or(0),
77127
is_primary_key: value.get_text(8).unwrap_or_default() == "true",
128+
foreign_keys: Vec::new(),
78129
},
79130
}
80131
}
@@ -93,4 +144,81 @@ mod test {
93144
let columns = Column::load(&mut conn).await.unwrap();
94145
println!("{:#?}", columns);
95146
}
147+
148+
#[tokio::test]
149+
async fn test_load_foreign_keys() {
150+
use crate::backend::schema::columns::ForeignKeyAction;
151+
152+
let pool = pool();
153+
let mut conn = pool.get(&Request::default()).await.unwrap();
154+
155+
// Create test tables with FK relationships using different action types
156+
conn.execute("DROP TABLE IF EXISTS fk_test_orders CASCADE")
157+
.await
158+
.unwrap();
159+
conn.execute("DROP TABLE IF EXISTS fk_test_customers CASCADE")
160+
.await
161+
.unwrap();
162+
163+
conn.execute(
164+
"CREATE TABLE fk_test_customers (
165+
id SERIAL PRIMARY KEY,
166+
name TEXT NOT NULL
167+
)",
168+
)
169+
.await
170+
.unwrap();
171+
172+
conn.execute(
173+
"CREATE TABLE fk_test_orders (
174+
id SERIAL PRIMARY KEY,
175+
customer_id INTEGER NOT NULL REFERENCES fk_test_customers(id)
176+
ON DELETE CASCADE ON UPDATE RESTRICT,
177+
amount NUMERIC
178+
)",
179+
)
180+
.await
181+
.unwrap();
182+
183+
let columns = Column::load(&mut conn).await.unwrap();
184+
185+
// Find the customer_id column in fk_test_orders (uses "pgdog" schema in test db)
186+
let orders_columns = columns
187+
.get(&("pgdog".to_string(), "fk_test_orders".to_string()))
188+
.expect("fk_test_orders table should exist");
189+
190+
let customer_id_col = orders_columns
191+
.iter()
192+
.find(|c| c.column_name == "customer_id")
193+
.expect("customer_id column should exist");
194+
195+
// Verify FK info is captured
196+
assert!(
197+
!customer_id_col.foreign_keys.is_empty(),
198+
"customer_id should have foreign key info"
199+
);
200+
201+
let fk = &customer_id_col.foreign_keys[0];
202+
assert_eq!(fk.schema, "pgdog");
203+
assert_eq!(fk.table, "fk_test_customers");
204+
assert_eq!(fk.column, "id");
205+
assert_eq!(fk.on_delete, ForeignKeyAction::Cascade);
206+
assert_eq!(fk.on_update, ForeignKeyAction::Restrict);
207+
208+
// Verify the id column (PK) does NOT have FK info
209+
let id_col = orders_columns
210+
.iter()
211+
.find(|c| c.column_name == "id")
212+
.expect("id column should exist");
213+
assert!(id_col.is_primary_key);
214+
assert!(id_col.foreign_keys.is_empty());
215+
216+
// Clean up
217+
conn.execute("DROP TABLE IF EXISTS fk_test_orders CASCADE")
218+
.await
219+
.unwrap();
220+
conn.execute("DROP TABLE IF EXISTS fk_test_customers CASCADE")
221+
.await
222+
.unwrap();
223+
}
96224
}

‎pgdog/src/backend/schema/foreign_keys.sql‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
SELECT
2+
kcu.table_schema::text AS source_schema,
3+
kcu.table_name::text AS source_table,
4+
kcu.column_name::text AS source_column,
5+
ccu.table_schema::text AS ref_schema,
6+
ccu.table_name::text AS ref_table,
7+
ccu.column_name::text AS ref_column,
8+
rc.delete_rule::text AS on_delete,
9+
rc.update_rule::text AS on_update
10+
FROM
11+
information_schema.table_constraints tc
12+
JOIN
13+
information_schema.key_column_usage kcu
14+
ON tc.constraint_name = kcu.constraint_name
15+
AND tc.table_schema = kcu.table_schema
16+
JOIN
17+
information_schema.constraint_column_usage ccu
18+
ON tc.constraint_name = ccu.constraint_name
19+
AND tc.table_schema = ccu.table_schema
20+
JOIN
21+
information_schema.referential_constraints rc
22+
ON tc.constraint_name = rc.constraint_name
23+
AND tc.table_schema = rc.constraint_schema
24+
WHERE
25+
tc.constraint_type = 'FOREIGN KEY'
26+
AND tc.table_schema NOT IN ('pg_catalog', 'information_schema')
27+
ORDER BY
28+
kcu.table_schema, kcu.table_name, kcu.column_name;

‎pgdog/src/frontend/router/parser/multi_tenant.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,7 @@ mod tests {
126126
data_type: "bigint".into(),
127127
ordinal_position: 1,
128128
is_primary_key: false,
129+
foreign_keys: Vec::new(),
129130
}
130131
.into(),
131132
);

‎pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,7 @@ mod tests {
267267
data_type: "bigint".into(),
268268
ordinal_position: 1,
269269
is_primary_key: true,
270+
foreign_keys: Vec::new(),
270271
}
271272
.into(),
272273
);
@@ -282,6 +283,7 @@ mod tests {
282283
data_type: "text".into(),
283284
ordinal_position: 2,
284285
is_primary_key: false,
286+
foreign_keys: Vec::new(),
285287
}
286288
.into(),
287289
);
@@ -304,6 +306,7 @@ mod tests {
304306
data_type: "uuid".into(),
305307
ordinal_position: 1,
306308
is_primary_key: true,
309+
foreign_keys: Vec::new(),
307310
}
308311
.into(),
309312
);
@@ -319,6 +322,7 @@ mod tests {
319322
data_type: "text".into(),
320323
ordinal_position: 2,
321324
is_primary_key: false,
325+
foreign_keys: Vec::new(),
322326
}
323327
.into(),
324328
);

‎pgdog/src/frontend/router/parser/statement.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2392,6 +2392,7 @@ mod test {
23922392
data_type: "bigint".into(),
23932393
ordinal_position: 1,
23942394
is_primary_key: true,
2395+
foreign_keys: Vec::new(),
23952396
}
23962397
.into(),
23972398
);
@@ -2407,6 +2408,7 @@ mod test {
24072408
data_type: "text".into(),
24082409
ordinal_position: 2,
24092410
is_primary_key: false,
2411+
foreign_keys: Vec::new(),
24102412
}
24112413
.into(),
24122414
);

0 commit comments

Comments
 (0)