1use std::sync::Arc;
15
16use datafusion::arrow::datatypes::{DataType, Schema};
17use datafusion::catalog::{CatalogProvider, TableProvider};
18use datafusion::common::utils::datafusion_strsim::levenshtein;
19use datafusion::common::{ResolvedTableReference, SchemaReference, TableReference};
20use datafusion::error::{DataFusionError, Result};
21use datafusion::execution::context::SessionState;
22use datafusion::logical_expr::{DdlStatement, LogicalPlan};
23use datafusion::prelude::SessionContext;
24
25use super::schema::{graph_table, lookup_table, GraphSchema, EDGES, ID, NODES};
26use super::{registry, FileTable, GraphDef, GraphRegistry};
27use crate::file_table::open_file_table;
28use crate::paths::looks_like_location;
29
30const INFORMATION_SCHEMA: &str = "information_schema";
32
33pub async fn create_graph(
37 ctx: &SessionContext,
38 def: GraphDef,
39 or_replace: bool,
40 if_not_exists: bool,
41) -> Result<()> {
42 if or_replace && if_not_exists {
43 return plan_err("'IF NOT EXISTS' cannot coexist with 'REPLACE'");
44 }
45 let state = ctx.state();
46 let graphs = registry(&state);
47 if graphs.get(&def.name).is_some() {
48 if if_not_exists {
49 return Ok(());
50 }
51 if !or_replace {
52 return plan_err(format!("graph '{}' already exists", def.name));
53 }
54 } else if def.name == INFORMATION_SCHEMA || default_catalog(&state)?.schema(&def.name).is_some()
55 {
56 return plan_err(format!(
57 "cannot name a graph '{}': a schema of that name already exists, and a graph's \
58 name is the name of the schema holding its nodes and edges tables",
59 def.name
60 ));
61 }
62
63 let mut files = Vec::new();
64 let edges = open_table(&state, &def.edges.table, &mut files).await?;
65 let vertices = match &def.vertices {
66 Some(v) => Some(open_table(&state, &v.table, &mut files).await?),
67 None => None,
68 };
69 validate(
70 &def,
71 &edges.schema(),
72 vertices.as_ref().map(|v| v.schema()).as_deref(),
73 )?;
74
75 let name = def.name.clone();
76 graphs.put_with_files(def, files);
77 let catalog = default_catalog_inner(&state)?;
78 let schema = match catalog.schema(&name) {
79 Some(schema) => schema,
80 None => {
81 let schema = Arc::new(GraphSchema::new(name.clone(), ctx.state_weak_ref()));
82 catalog.register_schema(&name, Arc::clone(&schema) as _)?;
83 schema
84 }
85 };
86 for table in [NODES, EDGES] {
89 schema.table(table).await?;
90 }
91 Ok(())
92}
93
94async fn open_table(
98 state: &SessionState,
99 table: &TableReference,
100 files: &mut Vec<FileTable>,
101) -> Result<Arc<dyn TableProvider>> {
102 if let Some(provider) = session_table(state, table).await? {
103 return Ok(provider);
104 }
105 if let TableReference::Bare { table: location } = table {
106 if looks_like_location(location) {
107 let provider = open_file_table(state, location).await?;
108 files.push(FileTable {
109 table: table.clone(),
110 provider: Arc::clone(&provider),
111 });
112 return Ok(provider);
113 }
114 }
115 find_table(state, table).await
116}
117
118async fn session_table(
121 state: &SessionState,
122 table: &TableReference,
123) -> Result<Option<Arc<dyn TableProvider>>> {
124 let Some(catalogs) = registry(state).catalogs().cloned() else {
125 return lookup_table(state, table).await;
126 };
127 let resolved = resolve(state, table);
128 match catalogs
129 .catalog(&resolved.catalog)
130 .and_then(|c| c.schema(&resolved.schema))
131 {
132 Some(schema) => schema.table(&resolved.table).await,
133 None => Ok(None),
134 }
135}
136
137pub fn unknown_graph(name: &str, graphs: &GraphRegistry) -> DataFusionError {
139 let names = graphs.names();
140 let hint = match closest_name(name, names.iter().map(String::as_str)) {
141 Some(s) => format!("; did you mean '{s}'?"),
142 None if names.is_empty() => {
143 "; no graphs are defined (define one with CREATE PROPERTY GRAPH)".to_string()
144 }
145 None => format!(". The graphs are: {}", names.join(", ")),
146 };
147 DataFusionError::Plan(format!("graph '{name}' does not exist{hint}"))
148}
149
150pub async fn check_ddl(state: &SessionState, plan: &LogicalPlan) -> Result<()> {
160 let LogicalPlan::Ddl(ddl) = plan else {
161 return Ok(());
162 };
163 match ddl {
164 DdlStatement::DropTable(drop) => check_drop(state, "table", &drop.name),
165 DdlStatement::DropView(drop) => check_drop(state, "view", &drop.name),
166 DdlStatement::DropCatalogSchema(drop) => check_drop_schema(state, &drop.name),
167 DdlStatement::CreateView(create) if create.or_replace => {
168 let schema = create.input.schema().as_arrow();
169 check_replace(state, "view", &create.name, Some(schema)).await
170 }
171 DdlStatement::CreateMemoryTable(create) if create.or_replace => {
172 let schema = create.input.schema().as_arrow();
173 check_replace(state, "table", &create.name, Some(schema)).await
174 }
175 DdlStatement::CreateExternalTable(create) if create.or_replace => {
176 let schema = create.schema.as_arrow();
179 let known = (!schema.fields().is_empty()).then_some(schema);
180 check_replace(state, "table", &create.name, known).await
181 }
182 _ => Ok(()),
183 }
184}
185
186pub fn check_drop(state: &SessionState, kind: &str, table: &TableReference) -> Result<()> {
188 let dropped = resolve(state, table);
189 let users = graphs_using(state, |t| t == &dropped);
190 if users.is_empty() {
191 return Ok(());
192 }
193 plan_err(format!(
194 "cannot drop {kind} '{table}': {}; drop {} first",
195 used_by(&users),
196 drop_statements(&users)
197 ))
198}
199
200pub fn check_drop_schema(state: &SessionState, schema: &SchemaReference) -> Result<()> {
203 let options = &state.config_options().catalog;
204 let (catalog, name) = match schema {
205 SchemaReference::Bare { schema } => (options.default_catalog.as_str(), schema.as_ref()),
206 SchemaReference::Full { schema, catalog } => (catalog.as_ref(), schema.as_ref()),
207 };
208 if catalog == options.default_catalog && registry(state).get(name).is_some() {
209 return plan_err(format!(
210 "cannot drop schema '{name}': it holds the tables of graph '{name}'; drop the \
211 graph with DROP PROPERTY GRAPH {name}"
212 ));
213 }
214 let users = graphs_using(state, |t| {
215 t.catalog.as_ref() == catalog && t.schema.as_ref() == name
216 });
217 if users.is_empty() {
218 return Ok(());
219 }
220 plan_err(format!(
221 "cannot drop schema '{name}': {} a table in it; drop {} first",
222 if users.len() == 1 {
223 format!("graph '{}' uses", users[0].name)
224 } else {
225 used_by(&users).replace(" use it", " use")
226 },
227 drop_statements(&users)
228 ))
229}
230
231pub async fn check_replace(
236 state: &SessionState,
237 kind: &str,
238 table: &TableReference,
239 schema: Option<&Schema>,
240) -> Result<()> {
241 let replaced = resolve(state, table);
242 let users = graphs_using(state, |t| t == &replaced);
243 if users.is_empty() {
244 return Ok(());
245 }
246 let Some(schema) = schema else {
247 return plan_err(format!(
248 "cannot replace {kind} '{table}' by one whose columns are not known yet: {}; \
249 drop {} first",
250 used_by(&users),
251 drop_statements(&users)
252 ));
253 };
254 let graphs = registry(state);
255 for def in &users {
256 let mut schemas = Vec::new();
257 for t in std::iter::once(&def.edges.table).chain(def.vertices.as_ref().map(|v| &v.table)) {
258 let schema =
259 if graphs.file_table(&def.name, t).is_none() && resolve(state, t) == replaced {
260 Arc::new(schema.clone())
261 } else {
262 graph_table(state, def, t).await?.schema()
263 };
264 schemas.push(schema);
265 }
266 if let Err(e) = validate(def, &schemas[0], schemas.get(1).map(|s| s.as_ref())) {
267 return plan_err(format!(
268 "cannot replace {kind} '{table}': graph '{}' uses it, and the new {kind} does \
269 not fit its definition ({}). Keep the columns the graph uses, or drop \
270 property graph {} first",
271 def.name,
272 crate::user_message(&e),
273 def.name
274 ));
275 }
276 }
277 Ok(())
278}
279
280fn graphs_using(
282 state: &SessionState,
283 uses: impl Fn(&ResolvedTableReference) -> bool,
284) -> Vec<Arc<GraphDef>> {
285 let graphs = registry(state);
286 graphs
287 .all()
288 .into_iter()
289 .filter(|def| {
290 std::iter::once(&def.edges.table)
291 .chain(def.vertices.as_ref().map(|v| &v.table))
292 .any(|t| graphs.file_table(&def.name, t).is_none() && uses(&resolve(state, t)))
293 })
294 .collect()
295}
296
297fn used_by(graphs: &[Arc<GraphDef>]) -> String {
299 let names: Vec<String> = graphs.iter().map(|g| format!("'{}'", g.name)).collect();
300 if names.len() == 1 {
301 format!("graph {} uses it", names[0])
302 } else {
303 format!("graphs {} use it", names.join(", "))
304 }
305}
306
307fn drop_statements(graphs: &[Arc<GraphDef>]) -> String {
309 let names: Vec<&str> = graphs.iter().map(|g| g.name.as_str()).collect();
310 if names.len() == 1 {
311 format!("property graph {}", names[0])
312 } else {
313 format!("property graphs {}", names.join(", "))
314 }
315}
316
317fn resolve(state: &SessionState, table: &TableReference) -> ResolvedTableReference {
319 let options = &state.config_options().catalog;
320 table
321 .clone()
322 .resolve(&options.default_catalog, &options.default_schema)
323}
324
325pub fn drop_graph(ctx: &SessionContext, name: &str, if_exists: bool) -> Result<()> {
328 let state = ctx.state();
329 let graphs = registry(&state);
330 if graphs.remove(name).is_none() {
331 if if_exists {
332 return Ok(());
333 }
334 return Err(unknown_graph(name, &graphs));
335 }
336 let catalog = default_catalog_inner(&state)?;
337 if catalog
338 .schema(name)
339 .is_some_and(|schema| schema.is::<GraphSchema>())
340 {
341 catalog.deregister_schema(name, true)?;
342 }
343 Ok(())
344}
345
346pub fn closest_name<'a>(
349 name: &str,
350 candidates: impl IntoIterator<Item = &'a str>,
351) -> Option<&'a str> {
352 let limit = (name.chars().count() / 3).max(2);
353 candidates
354 .into_iter()
355 .map(|c| (levenshtein(name, c), c))
356 .filter(|(d, _)| *d <= limit)
357 .min_by_key(|(d, c)| (*d, *c))
358 .map(|(_, c)| c)
359}
360
361fn did_you_mean<'a>(name: &str, candidates: impl IntoIterator<Item = &'a str>) -> String {
363 closest_name(name, candidates)
364 .map(|c| format!("; did you mean '{c}'?"))
365 .unwrap_or_default()
366}
367
368fn validate(def: &GraphDef, edges: &Schema, vertices: Option<&Schema>) -> Result<()> {
371 let table = &def.edges.table;
372 let source = column_type(edges, table, &def.edges.source, "SOURCE KEY")?;
373 let destination = column_type(edges, table, &def.edges.destination, "DESTINATION KEY")?;
374 if source != destination {
375 return plan_err(format!(
376 "the source and destination columns of edge table '{table}' must have the same \
377 type, but {} is {source} and {} is {destination}",
378 def.edges.source, def.edges.destination
379 ));
380 }
381 let Some(family) = IdFamily::of(&source) else {
382 return plan_err(format!(
383 "the endpoints of edge table '{table}' ({} and {}) are {source}; node ids must \
384 be integers or strings",
385 def.edges.source, def.edges.destination
386 ));
387 };
388
389 let (Some(v), Some(vertices)) = (&def.vertices, vertices) else {
390 return Ok(());
391 };
392 let key = column_type(vertices, &v.table, &v.key, "KEY")?;
393 match IdFamily::of(&key) {
394 None => {
395 return plan_err(format!(
396 "the key of vertex table '{}' ({}) is {key}; node ids must be integers or \
397 strings",
398 v.table, v.key
399 ))
400 }
401 Some(f) if f != family => {
402 return plan_err(format!(
403 "the key of vertex table '{}' ({}) is {key}, but the edge endpoints are \
404 {source}; both must be integers or both strings",
405 v.table, v.key
406 ))
407 }
408 Some(_) => {}
409 }
410 if v.key != ID && vertices.field_with_name(ID).is_ok() {
411 return plan_err(format!(
412 "vertex table '{}' has a column named '{ID}' that is not its key ({}); \
413 {}.nodes names the key '{ID}', so rename that column (in a view) first",
414 v.table, v.key, def.name
415 ));
416 }
417 Ok(())
418}
419
420#[derive(Debug, PartialEq, Eq, Clone, Copy)]
422enum IdFamily {
423 Integer,
424 String,
425}
426
427impl IdFamily {
428 fn of(data_type: &DataType) -> Option<Self> {
429 match data_type {
430 t if t.is_integer() => Some(Self::Integer),
431 DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Some(Self::String),
432 DataType::Dictionary(_, value) => match Self::of(value) {
433 Some(Self::String) => Some(Self::String),
434 _ => None,
435 },
436 _ => None,
437 }
438 }
439}
440
441async fn find_table(
443 state: &SessionState,
444 table: &TableReference,
445) -> Result<Arc<dyn TableProvider>> {
446 match lookup_table(state, table).await? {
447 Some(provider) => Ok(provider),
448 None => {
449 let names = table_names(state).await;
450 plan_err(format!(
451 "table '{table}' not found{}",
452 did_you_mean(&table.to_string(), names.iter().map(String::as_str))
453 ))
454 }
455 }
456}
457
458async fn table_names(state: &SessionState) -> Vec<String> {
461 let options = &state.config_options().catalog;
462 let catalogs = state.catalog_list();
463 let mut names = Vec::new();
464 for catalog_name in catalogs.catalog_names() {
465 let Some(catalog) = catalogs.catalog(&catalog_name) else {
466 continue;
467 };
468 for schema_name in catalog.schema_names() {
469 if schema_name == INFORMATION_SCHEMA {
470 continue;
471 }
472 let Some(schema) = catalog.schema(&schema_name) else {
473 continue;
474 };
475 for table in schema.table_names() {
476 let reference = if catalog_name != options.default_catalog {
477 TableReference::full(catalog_name.as_str(), schema_name.as_str(), table)
478 } else if schema_name != options.default_schema {
479 TableReference::partial(schema_name.as_str(), table)
480 } else {
481 TableReference::bare(table)
482 };
483 names.push(reference.to_string());
484 }
485 }
486 }
487 names
488}
489
490fn column_type(
491 schema: &Schema,
492 table: &TableReference,
493 column: &str,
494 clause: &str,
495) -> Result<DataType> {
496 match schema.field_with_name(column) {
497 Ok(field) => Ok(field.data_type().clone()),
498 Err(_) => {
499 let columns: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
500 let hint = match closest_name(column, columns.iter().copied()) {
501 Some(c) => format!("did you mean '{c}'?"),
502 None => format!("its columns are {}", columns.join(", ")),
503 };
504 plan_err(format!(
505 "{clause} column '{column}' not found in table '{table}'; {hint}"
506 ))
507 }
508 }
509}
510
511fn default_catalog(state: &SessionState) -> Result<Arc<dyn CatalogProvider>> {
513 let name = &state.config_options().catalog.default_catalog;
514 state
515 .catalog_list()
516 .catalog(name)
517 .ok_or_else(|| DataFusionError::Plan(format!("the default catalog '{name}' is missing")))
518}
519
520fn default_catalog_inner(state: &SessionState) -> Result<Arc<dyn CatalogProvider>> {
523 let name = &state.config_options().catalog.default_catalog;
524 match registry(state).catalogs() {
525 Some(catalogs) => catalogs.catalog(name).ok_or_else(|| {
526 DataFusionError::Plan(format!("the default catalog '{name}' is missing"))
527 }),
528 None => default_catalog(state),
529 }
530}
531
532fn plan_err<T>(message: impl Into<String>) -> Result<T> {
533 Err(DataFusionError::Plan(message.into()))
534}