1use serde::{Deserialize, Serialize};
13use serde_json::{Value, json};
14
15#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
17pub struct DatasetDescriptor {
18 pub name: String,
21 pub kind: String,
25 #[serde(default, skip_serializing_if = "Option::is_none")]
31 pub schema: Option<Value>,
32 #[serde(default, skip_serializing_if = "Option::is_none")]
35 pub estimated_rows: Option<u64>,
36 pub config_patch: Value,
42 #[serde(default, skip_serializing_if = "Option::is_none")]
48 pub sink_patch: Option<Value>,
49 #[serde(default, skip_serializing_if = "Option::is_none")]
53 pub primary_key: Option<Vec<String>>,
54}
55
56impl DatasetDescriptor {
57 pub fn new(name: impl Into<String>, kind: impl Into<String>, config_patch: Value) -> Self {
59 Self {
60 name: name.into(),
61 kind: kind.into(),
62 schema: None,
63 estimated_rows: None,
64 config_patch,
65 sink_patch: None,
66 primary_key: None,
67 }
68 }
69
70 pub fn with_primary_key(mut self, columns: Vec<String>) -> Self {
73 self.primary_key = (!columns.is_empty()).then_some(columns);
74 self
75 }
76
77 pub fn with_schema(mut self, schema: Value) -> Self {
79 self.schema = Some(schema);
80 self
81 }
82
83 pub fn with_sink_patch(mut self, sink_patch: Value) -> Self {
85 self.sink_patch = Some(sink_patch);
86 self
87 }
88
89 pub fn with_estimated_rows(mut self, rows: u64) -> Self {
91 self.estimated_rows = Some(rows);
92 self
93 }
94}
95
96pub fn attach_primary_keys(
101 descriptors: Vec<DatasetDescriptor>,
102 key_columns: impl IntoIterator<Item = (String, String)>,
103) -> Vec<DatasetDescriptor> {
104 let mut keys: std::collections::HashMap<String, Vec<String>> = Default::default();
105 for (name, column) in key_columns {
106 keys.entry(name).or_default().push(column);
107 }
108 descriptors
109 .into_iter()
110 .map(|d| match keys.remove(&d.name) {
111 Some(cols) => d.with_primary_key(cols),
112 None => d,
113 })
114 .collect()
115}
116
117pub fn sql_type_to_json_schema(data_type: &str) -> Value {
127 let t = data_type.to_ascii_lowercase();
128 let ty = if t.contains("bool") || t == "bit" {
129 "boolean"
130 } else if t.contains("json") || t.contains("variant") || t == "object" || t.contains("struct") {
131 "object"
133 } else if t.contains("array") || t.starts_with('_') {
134 "array"
137 } else if (t.contains("int") && !t.contains("interval") && !t.contains("point"))
138 || t == "serial"
139 || t == "bigserial"
140 || t == "smallserial"
141 {
142 "integer"
145 } else if t.contains("double")
146 || t.contains("float")
147 || t.contains("real")
148 || t.contains("decimal")
149 || t.contains("numeric")
150 || t.contains("money")
151 || t == "number"
152 {
153 "number"
154 } else {
155 "string"
158 };
159 json!({ "type": ty })
160}
161
162pub fn nullable_type(mut fragment: Value) -> Value {
165 if let Some(Value::Array(types)) = fragment.get_mut("type") {
166 if !types.iter().any(|t| t == "null") {
167 types.push(json!("null"));
168 }
169 return fragment;
170 }
171 let single = match fragment.get("type") {
172 Some(Value::String(t)) => t.clone(),
173 _ => return fragment,
174 };
175 if let Some(obj) = fragment.as_object_mut() {
178 obj.insert("type".into(), json!([single, "null"]));
179 }
180 fragment
181}
182
183pub fn columns_to_schema(columns: impl IntoIterator<Item = (String, Value)>) -> Value {
187 let mut properties = serde_json::Map::new();
188 for (name, fragment) in columns {
189 properties.insert(name, fragment);
190 }
191 json!({ "type": "object", "properties": Value::Object(properties) })
192}
193
194#[cfg(test)]
195mod tests {
196 use super::*;
197
198 #[test]
199 fn descriptor_builder_round_trips() {
200 let d = DatasetDescriptor::new("public.orders", "table", json!({"query": "SELECT 1"}))
201 .with_schema(json!({"type": "object", "properties": {}}))
202 .with_estimated_rows(42);
203 assert_eq!(d.name, "public.orders");
204 assert_eq!(d.kind, "table");
205 assert_eq!(d.estimated_rows, Some(42));
206 let v = serde_json::to_value(&d).unwrap();
207 let back: DatasetDescriptor = serde_json::from_value(v).unwrap();
208 assert_eq!(back, d);
209 }
210
211 #[test]
212 fn primary_key_builder_keeps_order_and_treats_empty_as_none() {
213 let d = DatasetDescriptor::new("public.orders", "table", json!({}))
214 .with_primary_key(vec!["tenant".into(), "id".into()]);
215 assert_eq!(d.primary_key, Some(vec!["tenant".into(), "id".into()]));
216 let v = serde_json::to_value(&d).unwrap();
217 assert_eq!(v["primary_key"], json!(["tenant", "id"]));
218 let none = DatasetDescriptor::new("t", "table", json!({})).with_primary_key(vec![]);
219 assert_eq!(none.primary_key, None);
220 assert!(
221 serde_json::to_value(&none)
222 .unwrap()
223 .get("primary_key")
224 .is_none()
225 );
226 }
227
228 #[test]
229 fn attach_primary_keys_matches_by_name_in_key_order() {
230 let ds = vec![
231 DatasetDescriptor::new("public.a", "table", json!({})),
232 DatasetDescriptor::new("public.b", "table", json!({})),
233 ];
234 let out = attach_primary_keys(
235 ds,
236 [
237 ("public.a".to_string(), "tenant".to_string()),
238 ("public.a".to_string(), "id".to_string()),
239 ("public.zzz".to_string(), "id".to_string()),
240 ],
241 );
242 assert_eq!(out[0].primary_key, Some(vec!["tenant".into(), "id".into()]));
243 assert_eq!(out[1].primary_key, None);
244 }
245
246 #[test]
247 fn descriptor_omits_absent_optionals_in_json() {
248 let d = DatasetDescriptor::new("t", "table", json!({}));
249 let v = serde_json::to_value(&d).unwrap();
250 assert!(v.get("schema").is_none());
251 assert!(v.get("estimated_rows").is_none());
252 }
253
254 #[test]
255 fn sql_types_map_to_json_types() {
256 for (sql, want) in [
257 ("integer", "integer"),
258 ("BIGINT", "integer"),
259 ("smallint", "integer"),
260 ("tinyint(1)", "integer"),
261 ("serial", "integer"),
262 ("double precision", "number"),
263 ("NUMERIC(10,2)", "number"),
264 ("decimal", "number"),
265 ("float8", "number"),
266 ("money", "number"),
267 ("boolean", "boolean"),
268 ("BOOL", "boolean"),
269 ("bit", "boolean"),
270 ("json", "object"),
271 ("JSONB", "object"),
272 ("VARIANT", "object"),
273 ("STRUCT<a INT>", "object"),
274 ("ARRAY", "array"),
275 ("_int4", "array"),
276 ("text", "string"),
277 ("character varying", "string"),
278 ("uuid", "string"),
279 ("timestamp with time zone", "string"),
280 ("date", "string"),
281 ("bytea", "string"),
282 ("some_exotic_type", "string"),
283 ] {
284 assert_eq!(
285 sql_type_to_json_schema(sql),
286 json!({ "type": want }),
287 "for SQL type {sql:?}"
288 );
289 }
290 }
291
292 #[test]
293 fn nullable_wraps_scalar_type() {
294 assert_eq!(
295 nullable_type(json!({"type": "integer"})),
296 json!({"type": ["integer", "null"]})
297 );
298 let complex = json!({"type": ["string", "null"]});
300 assert_eq!(nullable_type(complex.clone()), complex);
301 assert_eq!(
302 nullable_type(json!({"type": ["number", "string"]})),
303 json!({"type": ["number", "string", "null"]}),
304 "a type union gains null"
305 );
306 assert_eq!(
309 nullable_type(json!({"type": "string", "format": "date-time"})),
310 json!({"type": ["string", "null"], "format": "date-time"})
311 );
312 }
313
314 #[test]
315 fn columns_to_schema_shapes_like_infer_schema() {
316 let schema = columns_to_schema(vec![
317 ("id".to_string(), json!({"type": "integer"})),
318 ("name".to_string(), json!({"type": ["string", "null"]})),
319 ]);
320 assert_eq!(schema["type"], "object");
321 assert_eq!(schema["properties"]["id"]["type"], "integer");
322 assert_eq!(schema["properties"]["name"]["type"][1], "null");
323 }
324}