Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/common/src/param_value.rs
Line
Count
Source (jump to first uncovered line)
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements.  See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership.  The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License.  You may obtain a copy of the License at
8
//
9
//   http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied.  See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
18
use crate::error::{_plan_datafusion_err, _plan_err};
19
use crate::{Result, ScalarValue};
20
use arrow_schema::DataType;
21
use std::collections::HashMap;
22
23
/// The parameter value corresponding to the placeholder
24
#[derive(Debug, Clone)]
25
pub enum ParamValues {
26
    /// For positional query parameters, like `SELECT * FROM test WHERE a > $1 AND b = $2`
27
    List(Vec<ScalarValue>),
28
    /// For named query parameters, like `SELECT * FROM test WHERE a > $foo AND b = $goo`
29
    Map(HashMap<String, ScalarValue>),
30
}
31
32
impl ParamValues {
33
    /// Verify parameter list length and type
34
0
    pub fn verify(&self, expect: &[DataType]) -> Result<()> {
35
0
        match self {
36
0
            ParamValues::List(list) => {
37
0
                // Verify if the number of params matches the number of values
38
0
                if expect.len() != list.len() {
39
0
                    return _plan_err!(
40
0
                        "Expected {} parameters, got {}",
41
0
                        expect.len(),
42
0
                        list.len()
43
0
                    );
44
0
                }
45
0
46
0
                // Verify if the types of the params matches the types of the values
47
0
                let iter = expect.iter().zip(list.iter());
48
0
                for (i, (param_type, value)) in iter.enumerate() {
49
0
                    if *param_type != value.data_type() {
50
0
                        return _plan_err!(
51
0
                            "Expected parameter of type {:?}, got {:?} at index {}",
52
0
                            param_type,
53
0
                            value.data_type(),
54
0
                            i
55
0
                        );
56
0
                    }
57
                }
58
0
                Ok(())
59
            }
60
            ParamValues::Map(_) => {
61
                // If it is a named query, variables can be reused,
62
                // but the lengths are not necessarily equal
63
0
                Ok(())
64
            }
65
        }
66
0
    }
67
68
0
    pub fn get_placeholders_with_values(&self, id: &str) -> Result<ScalarValue> {
69
0
        match self {
70
0
            ParamValues::List(list) => {
71
0
                if id.is_empty() {
72
0
                    return _plan_err!("Empty placeholder id");
73
0
                }
74
                // convert id (in format $1, $2, ..) to idx (0, 1, ..)
75
0
                let idx = id[1..]
76
0
                    .parse::<usize>()
77
0
                    .map_err(|e| {
78
0
                        _plan_datafusion_err!("Failed to parse placeholder id: {e}")
79
0
                    })?
80
0
                    .checked_sub(1);
81
                // value at the idx-th position in param_values should be the value for the placeholder
82
0
                let value = idx.and_then(|idx| list.get(idx)).ok_or_else(|| {
83
0
                    _plan_datafusion_err!("No value found for placeholder with id {id}")
84
0
                })?;
85
0
                Ok(value.clone())
86
            }
87
0
            ParamValues::Map(map) => {
88
0
                // convert name (in format $a, $b, ..) to mapped values (a, b, ..)
89
0
                let name = &id[1..];
90
                // value at the name position in param_values should be the value for the placeholder
91
0
                let value = map.get(name).ok_or_else(|| {
92
0
                    _plan_datafusion_err!("No value found for placeholder with name {id}")
93
0
                })?;
94
0
                Ok(value.clone())
95
            }
96
        }
97
0
    }
98
}
99
100
impl From<Vec<ScalarValue>> for ParamValues {
101
0
    fn from(value: Vec<ScalarValue>) -> Self {
102
0
        Self::List(value)
103
0
    }
104
}
105
106
impl<K> From<Vec<(K, ScalarValue)>> for ParamValues
107
where
108
    K: Into<String>,
109
{
110
0
    fn from(value: Vec<(K, ScalarValue)>) -> Self {
111
0
        let value: HashMap<String, ScalarValue> =
112
0
            value.into_iter().map(|(k, v)| (k.into(), v)).collect();
113
0
        Self::Map(value)
114
0
    }
115
}
116
117
impl<K> From<HashMap<K, ScalarValue>> for ParamValues
118
where
119
    K: Into<String>,
120
{
121
0
    fn from(value: HashMap<K, ScalarValue>) -> Self {
122
0
        let value: HashMap<String, ScalarValue> =
123
0
            value.into_iter().map(|(k, v)| (k.into(), v)).collect();
124
0
        Self::Map(value)
125
0
    }
126
}