feat: add cursor statements (#5094)

* feat: add sql parsers for cursor operations

* feat: cursor operator

* feat: implement RecordBatchStreamCursor

* feat: implement cursor storage and execution

* test: add tests

* chore: update docstring

* feat: add a temporary sql rewrite for cast in limit

this issue is described in #5097

* test: add more sql for cursor integration test

* feat: reject non-select query for cursor statement

* refactor: address review issues

* test: add empty result case

* feat: address review comments
This commit is contained in:
Ning Sun
2024-12-06 17:32:22 +08:00
committed by GitHub
parent 8b944268da
commit 3133f3fb4e
21 changed files with 786 additions and 5 deletions

View File

@@ -72,6 +72,7 @@ macro_rules! sql_tests {
test_postgres_parameter_inference,
test_postgres_array_types,
test_mysql_prepare_stmt_insert_timestamp,
test_declare_fetch_close_cursor,
);
)*
};
@@ -1198,3 +1199,66 @@ pub async fn test_postgres_array_types(store_type: StorageType) {
let _ = fe_pg_server.shutdown().await;
guard.remove_all().await;
}
pub async fn test_declare_fetch_close_cursor(store_type: StorageType) {
let (addr, mut guard, fe_pg_server) = setup_pg_server(store_type, "sql_inference").await;
let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls)
.await
.unwrap();
let (tx, rx) = tokio::sync::oneshot::channel();
tokio::spawn(async move {
connection.await.unwrap();
tx.send(()).unwrap();
});
client
.execute(
"DECLARE c1 CURSOR FOR SELECT * FROM numbers WHERE number > 2 LIMIT 50::bigint",
&[],
)
.await
.expect("declare cursor");
// duplicated cursor
assert!(client
.execute("DECLARE c1 CURSOR FOR SELECT 1", &[],)
.await
.is_err());
let rows = client.query("FETCH 5 FROM c1", &[]).await.unwrap();
assert_eq!(5, rows.len());
let rows = client.query("FETCH 100 FROM c1", &[]).await.unwrap();
assert_eq!(45, rows.len());
let rows = client.query("FETCH 100 FROM c1", &[]).await.unwrap();
assert_eq!(0, rows.len());
client.execute("CLOSE c1", &[]).await.expect("close cursor");
// cursor not found
let result = client.query("FETCH 100 FROM c1", &[]).await;
assert!(result.is_err());
client
.execute(
"DECLARE c2 CURSOR FOR SELECT * FROM numbers WHERE number < 0",
&[],
)
.await
.expect("declare cursor");
let rows = client.query("FETCH 5 FROM c2", &[]).await.unwrap();
assert_eq!(0, rows.len());
client.execute("CLOSE c2", &[]).await.expect("close cursor");
// Shutdown the client.
drop(client);
rx.await.unwrap();
let _ = fe_pg_server.shutdown().await;
guard.remove_all().await;
}