6use support::mock_jev::MockJev;
7use support::mock_jev::Reply;
8use support::{jev_instance, noul};
9use tokio_postgres::error::SqlState;

13async fn sidecar(mock: &MockJev) -> (support::postgres::Instance, tokio_postgres::Client) {
14    let (pg, client) = jev_instance(mock).await;
15    let conn = pg.conn_str();
16    let host = conn.split_whitespace().find_map(|kv| kv.strip_prefix("host=")).expect("a host");
17    client
18        .batch_execute(&format!(
19            "CREATE EXTENSION postgres_fdw;
20             CREATE SERVER loopback FOREIGN DATA WRAPPER postgres_fdw
21               OPTIONS (host '{host}', dbname 'postgres');
22             CREATE USER MAPPING FOR postgres SERVER loopback OPTIONS (user 'postgres');
23             CREATE FOREIGN TABLE remote_tickets (id int, body text)
24               SERVER loopback OPTIONS (table_name 'tickets');"
25        ))
26        .await
27        .expect("loopback");
28    (pg, client)
29}

31#[tokio::test(flavor = "multi_thread")]
32async fn a_foreign_table_is_judged_by_the_scan() {
33    let mock = MockJev::start(|_| Reply::json(200, noul(0.9))).await;
34    let (_pg, client) = sidecar(&mock).await;
35
36    let plan: Vec<String> = client
37        .query("EXPLAIN (COSTS OFF) SELECT id FROM remote_tickets t WHERE jev(t, 'Urgent?')", &[])
38        .await
39        .unwrap()
40        .iter()
41        .map(|r| r.get(0))
42        .collect();
43    assert!(plan.iter().any(|l| l.contains("Custom Scan")), "{plan:#?}");
44    assert!(plan.iter().any(|l| l.contains("Foreign Scan")), "{plan:#?}");
45
46    let rows = client.query("SELECT id FROM remote_tickets t WHERE jev(t, 'Urgent?')", &[]).await.expect("judged");
47    assert_eq!(rows.len(), 1);
48    assert_eq!(rows[0].get::<_, i32>(0), 1);
49    assert_eq!(mock.requests().len(), 1);
50}
51
52async fn assert_refused(client: &tokio_postgres::Client, mock: &MockJev, query: &str) {
53    let err = client.query(query, &[]).await.expect_err("refused");
54    let db = err.as_db_error().expect("an ERROR");
55    assert_eq!(db.code(), &SqlState::INVALID_PARAMETER_VALUE, "{err:?}");
56    assert_eq!(db.message(), "foreign server \"loopback\" lists postjevsql in its extensions option");
57    assert!(mock.requests().is_empty(), "sent {:?}", mock.requests().len());
58}
59
60#[tokio::test(flavor = "multi_thread")]
61async fn a_server_listing_the_extension_is_refused_before_anything_is_sent() {
62    let mock = MockJev::start(|_| Reply::json(200, noul(0.9))).await;
63    let (_pg, client) = sidecar(&mock).await;
64    client.batch_execute("ALTER SERVER loopback OPTIONS (ADD extensions 'postjevsql')").await.unwrap();
65
66    assert_refused(&client, &mock, "SELECT id FROM remote_tickets t WHERE jev(t, 'Urgent?')").await;
67    assert_refused(&client, &mock, "SELECT jev_prob(t, 'Urgent?') FROM remote_tickets t").await;
68    // With ORDER BY the call is computed over the Sort, and the ForeignScan
69    // is below it.
70    assert_refused(&client, &mock, "SELECT id, jev_prob(t, 'Urgent?') FROM remote_tickets t ORDER BY id").await;
71    // EXPLAIN plans too, so it is refused as well.
72    assert_refused(&client, &mock, "EXPLAIN SELECT id FROM remote_tickets t WHERE jev(t, 'Urgent?')").await;
73}
74
75#[tokio::test(flavor = "multi_thread")]
76async fn the_extension_is_found_anywhere_in_the_list() {
77    let mock = MockJev::start(|_| Reply::json(200, noul(0.9))).await;
78    let (_pg, client) = sidecar(&mock).await;
79    client.batch_execute("ALTER SERVER loopback OPTIONS (ADD extensions 'plpgsql, \"postjevsql\"')").await.unwrap();
80
81    assert_refused(&client, &mock, "SELECT id FROM remote_tickets t WHERE jev(t, 'Urgent?')").await;
82}
83
84#[tokio::test(flavor = "multi_thread")]
85async fn other_extensions_are_allowed() {
86    let mock = MockJev::start(|_| Reply::json(200, noul(0.9))).await;
87    let (_pg, client) = sidecar(&mock).await;
88    client.batch_execute("ALTER SERVER loopback OPTIONS (ADD extensions 'plpgsql')").await.unwrap();
89
90    let rows = client.query("SELECT id FROM remote_tickets t WHERE jev(t, 'Urgent?')", &[]).await.expect("judged");
91    assert_eq!(rows.len(), 1);
92    assert_eq!(mock.requests().len(), 1);
93}
94
95#[tokio::test(flavor = "multi_thread")]
96async fn an_update_through_a_foreign_table_is_judged_by_the_scan() {
97    let mock = MockJev::start(|_| Reply::json(200, noul(0.9))).await;
98    let (_pg, client) = sidecar(&mock).await;
99
100    // The jev condition stays on the sidecar, so postgres_fdw cannot push
101    // the UPDATE down: it fetches the rows FOR UPDATE, and the scan over
102    // that ForeignScan judges them before each is updated remotely. The
103    // call reads a column list: a whole-row `t` here is not yet batched,
104    // since postgres_fdw's own whole-row identity column shares its number.
105    let plan: Vec<String> = client
106        .query("EXPLAIN (COSTS OFF) UPDATE remote_tickets t SET body = 'urgent' WHERE jev((t.id, t.body), 'Urgent?')", &[])
107        .await
108        .unwrap()
109        .iter()
110        .map(|r| r.get(0))
111        .collect();
112    assert!(plan.iter().any(|l| l.contains("Custom Scan")), "{plan:#?}");
113    assert!(plan.iter().any(|l| l.contains("Foreign Scan")), "{plan:#?}");
114
115    let n = client.execute("UPDATE remote_tickets t SET body = 'urgent' WHERE jev((t.id, t.body), 'Urgent?')", &[]).await.expect("judged");
116    assert_eq!(n, 1);
117    assert_eq!(mock.requests().len(), 1);
118    let body: String = client.query_one("SELECT body FROM tickets WHERE id = 1", &[]).await.unwrap().get(0);
119    assert_eq!(body, "urgent", "the target was updated");
120}