risingwave_frontend/handler/
fetch_cursor.rs

1// Copyright 2024 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use pgwire::pg_field_descriptor::PgFieldDescriptor;
16use pgwire::pg_response::{PgResponse, StatementType};
17use pgwire::types::{Format, Row};
18use risingwave_common::bail_not_implemented;
19use risingwave_common::catalog::Schema;
20use risingwave_common::types::DataType;
21use risingwave_sqlparser::ast::{FetchCursorStatement, Statement};
22
23use super::RwPgResponse;
24use super::extended_handle::{PortalResult, PrepareStatement, PreparedResult};
25use super::query::BoundResult;
26use super::util::convert_interval_to_u64_seconds;
27use crate::binder::BoundStatement;
28use crate::error::Result;
29use crate::handler::HandlerArgs;
30use crate::session::cursor_manager::FetchCursorCancelHandle;
31use crate::{Binder, PgResponseStream, WithOptions};
32
33pub async fn handle_fetch_cursor_execute(
34    handler_args: HandlerArgs,
35    portal_result: PortalResult,
36) -> Result<RwPgResponse> {
37    if let PortalResult {
38        statement: Statement::FetchCursor { stmt },
39        bound_result:
40            BoundResult {
41                bound: BoundStatement::FetchCursor(fetch_cursor),
42                ..
43            },
44        result_formats,
45        ..
46    } = portal_result
47    {
48        match fetch_cursor.returning_schema {
49            Some(_) => handle_fetch_cursor(handler_args, stmt, &result_formats).await,
50            None => Ok(build_fetch_cursor_response(vec![], vec![])),
51        }
52    } else {
53        bail_not_implemented!("unsupported portal {}", portal_result)
54    }
55}
56pub async fn handle_fetch_cursor(
57    handler_args: HandlerArgs,
58    stmt: FetchCursorStatement,
59    formats: &Vec<Format>,
60) -> Result<RwPgResponse> {
61    let session = handler_args.session.clone();
62    let cursor_name = stmt.cursor_name.real_value();
63    let with_options = WithOptions::try_from(stmt.with_properties.0.as_slice())?;
64
65    if with_options.len() > 1 {
66        bail_not_implemented!("only `timeout` is supported in with options")
67    }
68
69    let timeout_seconds = with_options
70        .get("timeout")
71        .map(convert_interval_to_u64_seconds)
72        .transpose()?;
73
74    if with_options.len() == 1 && timeout_seconds.is_none() {
75        bail_not_implemented!("only `timeout` is supported in with options")
76    }
77
78    let cursor_manager = session.get_cursor_manager();
79    let mut cancel_handle = FetchCursorCancelHandle::new();
80
81    let result = cursor_manager
82        .get_rows_with_cursor(
83            &cursor_name,
84            stmt.count,
85            handler_args,
86            formats,
87            timeout_seconds,
88            &mut cancel_handle,
89        )
90        .await;
91    session.clear_cancel_query_flag();
92    let (rows, pg_descs) = result?;
93    Ok(build_fetch_cursor_response(rows, pg_descs))
94}
95
96fn build_fetch_cursor_response(rows: Vec<Row>, pg_descs: Vec<PgFieldDescriptor>) -> RwPgResponse {
97    PgResponse::builder(StatementType::FETCH_CURSOR)
98        .row_cnt_opt(Some(rows.len() as i32))
99        .values(PgResponseStream::from(rows), pg_descs)
100        .into()
101}
102
103pub async fn handle_parse(
104    handler_args: HandlerArgs,
105    statement: Statement,
106    specified_param_types: Vec<Option<DataType>>,
107) -> Result<PrepareStatement> {
108    if let Statement::FetchCursor { stmt } = &statement {
109        let session = handler_args.session.clone();
110        let cursor_name = stmt.cursor_name.real_value();
111        let fields = session
112            .get_cursor_manager()
113            .get_fields_with_cursor(&cursor_name)
114            .await?;
115
116        let mut binder =
117            Binder::new_for_batch(&session).with_specified_params_types(specified_param_types);
118        let schema = Some(Schema::new(fields));
119
120        let bound = binder.bind_fetch_cursor(cursor_name, stmt.count, schema)?;
121        let bound_result = BoundResult {
122            stmt_type: StatementType::FETCH_CURSOR,
123            must_dist: false,
124            bound: BoundStatement::FetchCursor(Box::new(bound)),
125            param_types: binder.export_param_types()?,
126            parsed_params: None,
127            dependent_relations: binder.included_relations().clone(),
128            dependent_udfs: binder.included_udfs().clone(),
129            dependent_secrets: binder.included_secrets().clone(),
130        };
131        let result = PreparedResult {
132            statement,
133            bound_result,
134        };
135        Ok(PrepareStatement::Prepared(result))
136    } else {
137        bail_not_implemented!("unsupported statement {:?}", statement)
138    }
139}
risingwave_frontend/handler/fetch_cursor.rs

risingwave_frontend/handler/
fetch_cursor.rs
