1use crate::{
5 profiles::{Entry, Role},
6 project::Project,
7};
8use serde::Deserialize;
9use serde_json::{Value, json};
10use sqlparser::{
11 ast::{SetExpr, Statement, TableFactor},
12 dialect::GenericDialect,
13 parser::Parser,
14};
15
16#[derive(Default, Deserialize)]
17#[serde(deny_unknown_fields)]
18pub struct Request {
19 #[serde(default)]
20 pub tables: Vec<String>,
21 #[serde(default)]
22 pub schemas: Vec<String>,
23}
24
25fn parts(name: &str, kind: &str) -> Result<Vec<String>, String> {
26 let statements = Parser::parse_sql(&GenericDialect {}, &format!("select * from {name}"))
27 .map_err(|e| format!("invalid relation `{name}`: {e}"))?;
28 let [Statement::Query(query)] = statements.as_slice() else {
29 return Err("expected a relation name".into());
30 };
31 let SetExpr::Select(select) = query.body.as_ref() else {
32 return Err("expected a relation name".into());
33 };
34 if select.from.len() != 1
35 || !select.from[0].joins.is_empty()
36 || select.selection.is_some()
37 || query.limit_clause.is_some()
38 {
39 return Err(format!("`{name}` must be a relation name, not SQL"));
40 }
41 let TableFactor::Table {
42 name: object,
43 alias: None,
44 args: None,
45 ..
46 } = &select.from[0].relation
47 else {
48 return Err(format!(
49 "`{name}` must be a relation name, not a function or alias"
50 ));
51 };
52 if query.to_string() != format!("SELECT * FROM {}", select.from[0].relation) {
53 return Err(format!("`{name}` must contain only a relation name"));
54 }
55 object
56 .0
57 .iter()
58 .map(|part| {
59 part.as_ident()
60 .map(|id| {
61 if kind == "postgres" && id.quote_style.is_none() {
62 id.value.to_ascii_lowercase()
63 } else {
64 id.value.clone()
65 }
66 })
67 .ok_or_else(|| "relation identifiers must be literal".into())
68 })
69 .collect()
70}
71
72fn lit(value: &str, kind: &str) -> String {
73 let hex = value
74 .as_bytes()
75 .iter()
76 .map(|byte| format!("{byte:02x}"))
77 .collect::<String>();
78 match kind {
79 "databricks" => format!("decode(unhex('{hex}'), 'UTF-8')"),
80 "postgres" => format!("convert_from(decode('{hex}', 'hex'), 'UTF8')"),
81 _ => format!("'{}'", value.replace('\'', "''")),
82 }
83}
84fn rows(project: &Project, profile: &str, sql: &str) -> Result<Vec<Value>, String> {
85 let value = crate::query::execute(project, profile, sql, 100_000)?;
86 if value["truncated"] == true {
87 return Err("catalogue exceeds 100000 rows; narrow the requested schema".into());
88 }
89 Ok(value["rows"].as_array().cloned().unwrap_or_default())
90}
91
92pub fn build(project: &Project, profile: &str, request: &Request) -> Result<Value, String> {
93 let Entry::Use(target) = project.profiles.entry(Role::Connection, profile) else {
94 return Err(format!("connection `{profile}` is unavailable"));
95 };
96 let defaults = match target.kind.as_str() {
97 "duckdb" | "postgres" => "select current_database(), current_schema()",
98 "databricks" => "select current_catalog(), current_schema()",
99 _ => {
100 return Err(format!(
101 "live catalogue comments are not yet supported for `{}`; supported: duckdb, postgres, databricks",
102 target.kind
103 ));
104 }
105 };
106 let defaults = rows(project, profile, defaults)?;
107 let catalog = defaults
108 .first()
109 .and_then(|r| r[0].as_str())
110 .ok_or("the connection returned no current catalog")?;
111 let schema = defaults[0][1]
112 .as_str()
113 .ok_or("the connection returned no current schema")?;
114 let mut wanted = Vec::new();
115 for name in &request.tables {
116 let p = parts(name, &target.kind)?;
117 let (c, s, n) = match p.as_slice() {
118 [n] => (catalog.to_owned(), schema.to_owned(), n.clone()),
119 [s, n] => (catalog.to_owned(), s.clone(), n.clone()),
120 [c, s, n] => (c.clone(), s.clone(), n.clone()),
121 _ => return Err(format!("`{name}` needs one to three name parts")),
122 };
123 let (c, s, n) = if target.kind == "postgres" {
124 if c != catalog {
125 return Err("Postgres does not support cross-database relation context".into());
126 }
127 let resolved = rows(
129 project,
130 profile,
131 &format!(
132 "select current_database(), ns.nspname, cl.relname from pg_class cl join pg_namespace ns on ns.oid=cl.relnamespace where cl.oid=to_regclass({})",
133 lit(name, &target.kind)
134 ),
135 )?;
136 let row = resolved
137 .first()
138 .ok_or_else(|| format!("table `{name}` is missing or not visible"))?;
139 (
140 row[0].as_str().ok_or("catalog missing")?.to_owned(),
141 row[1].as_str().ok_or("schema missing")?.to_owned(),
142 row[2].as_str().ok_or("name missing")?.to_owned(),
143 )
144 } else {
145 (c, s, n)
146 };
147 wanted.push((name.clone(), c, s, n));
148 }
149 let mut filters = Vec::new();
150 let comparison = |column: &str, value: &str| {
151 if target.kind == "duckdb" {
152 format!("lower({column})=lower({})", lit(value, &target.kind))
153 } else {
154 format!("{column}={}", lit(value, &target.kind))
155 }
156 };
157 let equal = |left: &str, right: &str| {
158 if target.kind == "duckdb" {
159 left.eq_ignore_ascii_case(right)
160 } else {
161 left == right
162 }
163 };
164 for (_, c, s, n) in &wanted {
165 filters.push(format!(
166 "({} and {} and {})",
167 comparison("table_catalog", c),
168 comparison("table_schema", s),
169 comparison("table_name", n)
170 ));
171 }
172 for name in &request.schemas {
173 let p = parts(name, &target.kind)?;
174 let (c, s) = match p.as_slice() {
175 [s] => (catalog, s.as_str()),
176 [c, s] => (c.as_str(), s.as_str()),
177 _ => return Err(format!("schema `{name}` needs one or two name parts")),
178 };
179 filters.push(format!(
180 "({} and {})",
181 comparison("table_catalog", c),
182 comparison("table_schema", s)
183 ));
184 }
185 let tables_view = match target.kind.as_str() {
186 "duckdb" => "(select database_name as table_catalog, schema_name as table_schema, table_name, comment as description from duckdb_tables() union all select database_name, schema_name, view_name, comment from duckdb_views()) t".to_owned(),
187 "postgres" => "(select current_database() as table_catalog, n.nspname as table_schema, c.relname as table_name, obj_description(c.oid, 'pg_class') as description from pg_class c join pg_namespace n on n.oid=c.relnamespace where c.relkind in ('r','v','m','f','p')) t".to_owned(),
188 "databricks" => "(select table_catalog, table_schema, table_name, comment as description from system.information_schema.tables) t".to_owned(),
189 _ => unreachable!(),
190 };
191 let filter = if filters.is_empty() {
192 "false".into()
193 } else {
194 filters.join(" or ")
195 };
196 let tables = rows(
197 project,
198 profile,
199 &format!(
200 "select table_catalog, table_schema, table_name, description from {tables_view} where {filter} order by table_catalog, table_schema, table_name"
201 ),
202 )?;
203 let mut out = Vec::new();
204 for table in tables {
205 let c = table[0].as_str().ok_or("catalog is missing")?;
206 let s = table[1].as_str().ok_or("schema is missing")?;
207 let n = table[2].as_str().ok_or("table name is missing")?;
208 let sql = match target.kind.as_str() {
209 "duckdb" => format!(
210 "select column_name, data_type, comment from duckdb_columns() where database_name={} and schema_name={} and table_name={} order by column_index",
211 lit(c, &target.kind),
212 lit(s, &target.kind),
213 lit(n, &target.kind)
214 ),
215 "postgres" => format!(
216 "select a.attname, format_type(a.atttypid,a.atttypmod), col_description(c.oid,a.attnum) from pg_attribute a join pg_class c on c.oid=a.attrelid join pg_namespace n on n.oid=c.relnamespace where current_database()={} and n.nspname={} and c.relname={} and a.attnum>0 and not a.attisdropped order by a.attnum",
217 lit(c, &target.kind),
218 lit(s, &target.kind),
219 lit(n, &target.kind)
220 ),
221 "databricks" => format!(
222 "select column_name, full_data_type, comment from system.information_schema.columns where table_catalog={} and table_schema={} and table_name={} order by ordinal_position",
223 lit(c, &target.kind),
224 lit(s, &target.kind),
225 lit(n, &target.kind)
226 ),
227 _ => unreachable!(),
228 };
229 let columns: Vec<Value> = rows(project, profile, &sql)?
230 .iter()
231 .map(|r| json!({"name":r[0], "type":r[1], "description":r[2]}))
232 .collect();
233 let requested: Vec<_> = wanted
234 .iter()
235 .filter(|(_, wc, ws, wn)| equal(wc, c) && equal(ws, s) && equal(wn, n))
236 .map(|(name, _, _, _)| name.clone())
237 .collect();
238 out.push(json!({"catalog":c,"schema":s,"name":n,"qualified_name":[c,s,n],"description":table[3], "description_source":"database", "columns":columns,"requested":requested}));
239 }
240 for (name, c, s, n) in &wanted {
241 if !out.iter().any(|t| {
242 equal(t["catalog"].as_str().unwrap_or(""), c)
243 && equal(t["schema"].as_str().unwrap_or(""), s)
244 && equal(t["name"].as_str().unwrap_or(""), n)
245 }) {
246 return Err(format!(
247 "table `{name}` resolved to {c}.{s}.{n}, but is missing or not visible on connection `{profile}`"
248 ));
249 }
250 }
251 Ok(
252 json!({"schema_version":"dre/catalog/v1","connection":profile,"target":project.target_name,
253 "source_type":target.kind,"identifier_quote":if target.kind=="databricks" {"`"} else {"\""},
254 "current_catalog":catalog,"current_schema":schema,"tables":out}),
255 )
256}