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