common_function/system/
procedure_state.rs

1// Copyright 2023 Greptime Team
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 api::v1::meta::ProcedureStatus;
16use arrow::datatypes::DataType as ArrowDataType;
17use common_macro::admin_fn;
18use common_meta::rpc::procedure::ProcedureStateResponse;
19use common_query::error::{
20    InvalidFuncArgsSnafu, MissingProcedureServiceHandlerSnafu, Result,
21    UnsupportedInputDataTypeSnafu,
22};
23use datafusion_expr::{Signature, TypeSignature, Volatility};
24use datatypes::prelude::*;
25use serde::Serialize;
26use session::context::QueryContextRef;
27use snafu::{OptionExt, ensure};
28
29use crate::handlers::ProcedureServiceHandlerRef;
30
31#[derive(Serialize)]
32struct ProcedureStateJson {
33    status: String,
34    #[serde(skip_serializing_if = "Option::is_none")]
35    error: Option<String>,
36}
37
38/// Queries procedure state using an explicit procedure ID, or the latest
39/// procedure ID recorded in the Session.
40/// The zero-argument form reads the procedure ID from the Session at execution
41/// time. State lookups do not update the recorded procedure ID.
42#[admin_fn(
43    name = ProcedureStateFunction,
44    display_name = procedure_state,
45    sig_fn = signature,
46    ret = string
47)]
48pub(crate) async fn procedure_state(
49    procedure_service_handler: &ProcedureServiceHandlerRef,
50    ctx: &QueryContextRef,
51    params: &[ValueRef<'_>],
52) -> Result<Value> {
53    ensure!(
54        params.len() <= 1,
55        InvalidFuncArgsSnafu {
56            err_msg: format!(
57                "The length of the args is not correct, expect 0 or 1, have: {}",
58                params.len()
59            ),
60        }
61    );
62
63    let last_procedure_id;
64    let pid = match params {
65        [] => {
66            last_procedure_id = ctx.last_procedure_id().context(InvalidFuncArgsSnafu {
67                err_msg: "No procedure ID has been returned in this session; pass an explicit procedure ID",
68            })?;
69            last_procedure_id.as_str()
70        }
71        [ValueRef::String(pid)] => *pid,
72        _ => {
73            return UnsupportedInputDataTypeSnafu {
74                function: "procedure_state",
75                datatypes: params.iter().map(|v| v.data_type()).collect::<Vec<_>>(),
76            }
77            .fail();
78        }
79    };
80
81    let ProcedureStateResponse { status, error, .. } =
82        procedure_service_handler.query_procedure_state(pid).await?;
83    let status = ProcedureStatus::try_from(status)
84        .map(|v| v.as_str_name())
85        .unwrap_or("Unknown");
86
87    let state = ProcedureStateJson {
88        status: status.to_string(),
89        error: if error.is_empty() { None } else { Some(error) },
90    };
91    let json = serde_json::to_string(&state).unwrap_or_default();
92
93    Ok(Value::from(json))
94}
95
96fn signature() -> Signature {
97    Signature::one_of(
98        vec![
99            TypeSignature::Exact(vec![]),
100            TypeSignature::Exact(vec![ArrowDataType::Utf8]),
101        ],
102        Volatility::Volatile,
103    )
104}
105
106#[cfg(test)]
107mod tests {
108    use std::sync::Arc;
109    use std::sync::atomic::{AtomicUsize, Ordering};
110
111    use arrow::array::StringArray;
112    use arrow::datatypes::{DataType, Field};
113    use datafusion_expr::ColumnarValue;
114    use datafusion_expr::type_coercion::functions::fields_with_udf;
115
116    use super::*;
117    use crate::function::FunctionContext;
118    use crate::function_factory::ScalarFunctionFactory;
119    use crate::state::FunctionState;
120
121    #[test]
122    fn test_procedure_state_misc() {
123        let factory: ScalarFunctionFactory = ProcedureStateFunction::factory().into();
124        let f = factory.provide(FunctionContext::mock());
125        assert_eq!("procedure_state", f.name());
126        assert_eq!(DataType::Utf8, f.return_type(&[]).unwrap());
127        assert_eq!(
128            f.signature(),
129            &Signature::one_of(
130                vec![
131                    TypeSignature::Exact(vec![]),
132                    TypeSignature::Exact(vec![ArrowDataType::Utf8]),
133                ],
134                Volatility::Volatile,
135            )
136        );
137        let field = Arc::new(Field::new("pid", DataType::Utf8, false));
138        assert!(fields_with_udf(&[field.clone(), field], &f).is_err());
139    }
140
141    #[tokio::test]
142    async fn test_procedure_state_args() {
143        let context = FunctionContext::mock();
144        let handler = context.state.procedure_service_handler.as_ref().unwrap();
145        let ctx = &context.query_ctx;
146        let err = procedure_state(handler, ctx, &[]).await.unwrap_err();
147        assert!(
148            err.to_string()
149                .contains("No procedure ID has been returned")
150        );
151
152        ctx.set_last_procedure_id("session-id".to_string());
153        for (params, succeeds) in [
154            (vec![], true),
155            (vec![ValueRef::String("explicit-id")], true),
156            (vec![ValueRef::Null], false),
157            (vec![ValueRef::String("a"), ValueRef::String("b")], false),
158        ] {
159            let result = procedure_state(handler, ctx, &params).await;
160            assert_eq!(result.is_ok(), succeeds, "{params:?}: {result:?}");
161            assert_eq!(ctx.last_procedure_id().as_deref(), Some("session-id"));
162        }
163    }
164
165    #[tokio::test]
166    async fn test_procedure_state_zero_arg_rows() {
167        let calls = Arc::new(AtomicUsize::new(0));
168        let context = FunctionContext {
169            state: Arc::new(FunctionState::mock_with_procedure_state_counter(
170                calls.clone(),
171            )),
172            ..Default::default()
173        };
174        context
175            .query_ctx
176            .set_last_procedure_id("session-id".to_string());
177        let factory: ScalarFunctionFactory = ProcedureStateFunction::factory().into();
178        let function = factory.provide(context);
179        let f = function.as_async().unwrap();
180
181        for number_rows in [0, 1, 2] {
182            calls.store(0, Ordering::Relaxed);
183            let args = datafusion_expr::ScalarFunctionArgs {
184                args: vec![],
185                arg_fields: vec![],
186                return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
187                number_rows,
188                config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
189            };
190            let ColumnarValue::Array(result) = f.invoke_async_with_args(args).await.unwrap() else {
191                panic!("expected array output");
192            };
193            assert_eq!(result.len(), number_rows);
194            assert_eq!(calls.load(Ordering::Relaxed), number_rows);
195        }
196    }
197
198    #[tokio::test]
199    async fn test_missing_procedure_service() {
200        let factory: ScalarFunctionFactory = ProcedureStateFunction::factory().into();
201        let binding = factory.provide(FunctionContext::default());
202        let f = binding.as_async().unwrap();
203
204        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
205            args: vec![ColumnarValue::Array(Arc::new(StringArray::from(vec![
206                "pid",
207            ])))],
208            arg_fields: vec![Arc::new(Field::new("arg_0", DataType::Utf8, false))],
209            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
210            number_rows: 1,
211            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
212        };
213        let result = f.invoke_async_with_args(func_args).await;
214        assert!(result.is_err());
215    }
216
217    #[tokio::test]
218    async fn test_procedure_state() {
219        let factory: ScalarFunctionFactory = ProcedureStateFunction::factory().into();
220        let provider = factory.provide(FunctionContext::mock());
221        let f = provider.as_async().unwrap();
222
223        let func_args = datafusion::logical_expr::ScalarFunctionArgs {
224            args: vec![ColumnarValue::Array(Arc::new(StringArray::from(vec![
225                "pid",
226            ])))],
227            arg_fields: vec![Arc::new(Field::new("arg_0", DataType::Utf8, false))],
228            return_field: Arc::new(Field::new("result", DataType::Utf8, true)),
229            number_rows: 1,
230            config_options: Arc::new(datafusion_common::config::ConfigOptions::default()),
231        };
232        let result = f.invoke_async_with_args(func_args).await.unwrap();
233
234        match result {
235            ColumnarValue::Array(array) => {
236                let result_array = array.as_any().downcast_ref::<StringArray>().unwrap();
237                assert_eq!(
238                    result_array.value(0),
239                    "{\"status\":\"Done\",\"error\":\"OK\"}"
240                );
241            }
242            ColumnarValue::Scalar(scalar) => {
243                assert_eq!(
244                    scalar,
245                    datafusion_common::ScalarValue::Utf8(Some(
246                        "{\"status\":\"Done\",\"error\":\"OK\"}".to_string()
247                    ))
248                );
249            }
250        }
251    }
252}
common_function/system/procedure_state.rs

common_function/system/
procedure_state.rs
