Skip to content

Commit b2f68c5

Browse files
dpxccdpxcc
committed
Support list_tables
1 parent 2b6fb7e commit b2f68c5

2 files changed

Lines changed: 56 additions & 24 deletions

File tree

‎moonlink‎

Submodule moonlink updated 63 files

‎src/table/functions.rs‎

Lines changed: 55 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::utils::{block_on, get_database_id, get_stream};
22
use core::ffi::CStr;
3-
use pgrx::{direct_function_call, prelude::*};
3+
use pgrx::{direct_function_call, pg_sys::Oid, prelude::*};
44
use postgres::{Client, NoTls};
55
use regex::Regex;
66

@@ -9,7 +9,30 @@ extern "C" {
99
}
1010

1111
#[pg_extern(sql = "
12-
CREATE PROCEDURE mooncake.create_table(dst TEXT, src TEXT, src_uri TEXT DEFAULT NULL) LANGUAGE c AS 'MODULE_PATHNAME', '@FUNCTION_NAME@';
12+
CREATE PROCEDURE mooncake.create_snapshot(dst text) LANGUAGE c AS 'MODULE_PATHNAME', '@FUNCTION_NAME@';
13+
")]
14+
fn create_snapshot(dst: &str) {
15+
let dst = parse_table(dst);
16+
let dst_uri = get_loopback_uri();
17+
let mut client = Client::connect(&dst_uri, NoTls)
18+
.unwrap_or_else(|_| panic!("error connecting to server: {dst_uri}"));
19+
let get_table_id_query = format!("SELECT '{}'::regclass::oid", dst.replace("'", "''"));
20+
let table_id: u32 = client
21+
.query_one(&get_table_id_query, &[])
22+
.unwrap_or_else(|_| panic!("relation does not exist: {dst}"))
23+
.get(0);
24+
let lsn = unsafe { GetActiveLsn() };
25+
block_on(moonlink_rpc::create_snapshot(
26+
&mut *get_stream(),
27+
get_database_id(),
28+
table_id,
29+
lsn,
30+
))
31+
.expect("create_snapshot failed");
32+
}
33+
34+
#[pg_extern(sql = "
35+
CREATE PROCEDURE mooncake.create_table(dst text, src text, src_uri text DEFAULT NULL) LANGUAGE c AS 'MODULE_PATHNAME', '@FUNCTION_NAME@';
1336
")]
1437
fn create_table(dst: &str, src: &str, src_uri: Option<&str>) {
1538
let dst = parse_table(dst);
@@ -59,30 +82,39 @@ fn drop_trigger() {
5982
}
6083

6184
#[pg_extern(sql = "
62-
CREATE PROCEDURE mooncake.create_snapshot(dst TEXT) LANGUAGE c AS 'MODULE_PATHNAME', '@FUNCTION_NAME@';
85+
CREATE FUNCTION mooncake.list_tables() RETURNS TABLE (
86+
database_id oid,
87+
table_id oid,
88+
commit_lsn pg_lsn,
89+
flush_lsn pg_lsn,
90+
iceberg_warehouse_location text
91+
) LANGUAGE c AS 'MODULE_PATHNAME', '@FUNCTION_NAME@';
6392
")]
64-
fn create_snapshot(dst: &str) {
65-
let dst = parse_table(dst);
66-
let dst_uri = get_loopback_uri();
67-
let mut client = Client::connect(&dst_uri, NoTls)
68-
.unwrap_or_else(|_| panic!("error connecting to server: {dst_uri}"));
69-
let get_table_id_query = format!("SELECT '{}'::regclass::oid", dst.replace("'", "''"));
70-
let table_id: u32 = client
71-
.query_one(&get_table_id_query, &[])
72-
.unwrap_or_else(|_| panic!("relation does not exist: {dst}"))
73-
.get(0);
74-
let lsn = unsafe { GetActiveLsn() };
75-
block_on(moonlink_rpc::create_snapshot(
76-
&mut *get_stream(),
77-
get_database_id(),
78-
table_id,
79-
lsn,
80-
))
81-
.expect("create_snapshot failed");
93+
fn list_tables() -> TableIterator<
94+
'static,
95+
(
96+
name!(database_id, Oid),
97+
name!(table_id, Oid),
98+
name!(commit_lsn, i64),
99+
name!(flush_lsn, Option<i64>),
100+
name!(iceberg_warehouse_location, String),
101+
),
102+
> {
103+
let tables =
104+
block_on(moonlink_rpc::list_tables(&mut *get_stream())).expect("list_tables failed");
105+
TableIterator::new(tables.into_iter().map(|table| {
106+
(
107+
Oid::from_u32(table.database_id),
108+
Oid::from_u32(table.table_id),
109+
table.commit_lsn as i64,
110+
table.flush_lsn.map(|lsn| lsn as i64),
111+
table.iceberg_warehouse_location,
112+
)
113+
}))
82114
}
83115

84116
#[pg_extern(sql = "
85-
CREATE PROCEDURE mooncake.optimize_table(dst TEXT, mode TEXT) LANGUAGE c AS 'MODULE_PATHNAME', '@FUNCTION_NAME@';
117+
CREATE PROCEDURE mooncake.optimize_table(dst text, mode text) LANGUAGE c AS 'MODULE_PATHNAME', '@FUNCTION_NAME@';
86118
")]
87119
fn optimize_table(dst: &str, mode: &str) {
88120
let dst = parse_table(dst);
@@ -122,7 +154,7 @@ fn parse_table(table: &str) -> String {
122154
spi::quote_qualified_identifier(schema, &caps["table"])
123155
}
124156

125-
pub(crate) fn get_loopback_uri() -> String {
157+
fn get_loopback_uri() -> String {
126158
let hosts = unsafe { CStr::from_ptr(pg_sys::Unix_socket_directories) };
127159
let host = hosts.to_str().unwrap().split(",").next().unwrap().trim();
128160
let port: i32 = unsafe { pg_sys::PostPortNumber };

0 commit comments

Comments
 (0)