vibesql/catalog/
function.rs1use crate::types::SqlType;
4
5#[derive(Debug, Clone, PartialEq)]
7pub struct FunctionSignature {
8 pub name: String,
10 pub parameters: Vec<FunctionParameter>,
12 pub return_type: SqlType,
14 pub is_aggregate: bool,
16 pub is_window: bool,
18 pub is_deterministic: bool,
20 pub min_args: usize,
22 pub max_args: Option<usize>,
24}
25
26#[derive(Debug, Clone, PartialEq)]
28pub struct FunctionParameter {
29 pub name: Option<String>,
31 pub data_type: Option<SqlType>,
33 pub optional: bool,
35 pub variadic: bool,
37}
38
39impl FunctionSignature {
40 pub fn scalar(name: impl Into<String>, return_type: SqlType) -> Self {
42 Self {
43 name: name.into().to_uppercase(),
44 parameters: Vec::new(),
45 return_type,
46 is_aggregate: false,
47 is_window: false,
48 is_deterministic: true,
49 min_args: 0,
50 max_args: None,
51 }
52 }
53
54 pub fn aggregate(name: impl Into<String>, return_type: SqlType) -> Self {
56 Self {
57 name: name.into().to_uppercase(),
58 parameters: Vec::new(),
59 return_type,
60 is_aggregate: true,
61 is_window: false,
62 is_deterministic: true,
63 min_args: 0,
64 max_args: None,
65 }
66 }
67
68 pub fn window(name: impl Into<String>, return_type: SqlType) -> Self {
70 Self {
71 name: name.into().to_uppercase(),
72 parameters: Vec::new(),
73 return_type,
74 is_aggregate: false,
75 is_window: true,
76 is_deterministic: true,
77 min_args: 0,
78 max_args: None,
79 }
80 }
81
82 pub fn with_min_args(mut self, min: usize) -> Self {
84 self.min_args = min;
85 self
86 }
87
88 pub fn with_max_args(mut self, max: usize) -> Self {
90 self.max_args = Some(max);
91 self
92 }
93
94 pub fn with_args(mut self, count: usize) -> Self {
96 self.min_args = count;
97 self.max_args = Some(count);
98 self
99 }
100
101 pub fn with_param(mut self, param: FunctionParameter) -> Self {
103 self.parameters.push(param);
104 self
105 }
106
107 pub fn non_deterministic(mut self) -> Self {
109 self.is_deterministic = false;
110 self
111 }
112
113 pub fn accepts_arg_count(&self, count: usize) -> bool {
115 if count < self.min_args {
116 return false;
117 }
118 match self.max_args {
119 Some(max) => count <= max,
120 None => true,
121 }
122 }
123
124 pub fn can_be_window(&self) -> bool {
126 self.is_window || self.is_aggregate
127 }
128}
129
130impl FunctionParameter {
131 pub fn new(name: impl Into<String>, data_type: SqlType) -> Self {
133 Self {
134 name: Some(name.into()),
135 data_type: Some(data_type),
136 optional: false,
137 variadic: false,
138 }
139 }
140
141 pub fn any(name: impl Into<String>) -> Self {
143 Self {
144 name: Some(name.into()),
145 data_type: None,
146 optional: false,
147 variadic: false,
148 }
149 }
150
151 pub fn unnamed(data_type: SqlType) -> Self {
153 Self {
154 name: None,
155 data_type: Some(data_type),
156 optional: false,
157 variadic: false,
158 }
159 }
160
161 pub fn optional(mut self) -> Self {
163 self.optional = true;
164 self
165 }
166
167 pub fn variadic(mut self) -> Self {
169 self.variadic = true;
170 self
171 }
172}
173
174#[cfg(test)]
175mod tests {
176 use super::*;
177
178 #[test]
179 fn test_function_signature() {
180 let count = FunctionSignature::aggregate("COUNT", SqlType::Int64)
181 .with_min_args(1)
182 .with_max_args(1);
183
184 assert!(count.is_aggregate);
185 assert!(!count.is_window);
186 assert!(count.accepts_arg_count(1));
187 assert!(!count.accepts_arg_count(0));
188 assert!(!count.accepts_arg_count(2));
189 }
190
191 #[test]
192 fn test_window_function() {
193 let row_number = FunctionSignature::window("ROW_NUMBER", SqlType::Int64).with_args(0);
194
195 assert!(row_number.is_window);
196 assert!(row_number.can_be_window());
197 assert!(row_number.accepts_arg_count(0));
198 }
199}