Skip to main content

ursa_engine/graph/
define.rs

1//! Defining and dropping graphs: what `CREATE PROPERTY GRAPH` and
2//! `DROP PROPERTY GRAPH` do once parsed.
3//!
4//! A definition is checked when it is created: its tables exist, the named
5//! columns exist in them, the endpoints have one node id type (an integer of
6//! any width, or a string), and the vertex key's type is compatible with it.
7//! A table given as a file path is opened then (see the [module
8//! docs](super)).
9//!
10//! The tables are read afresh by every query. [`check_ddl`] keeps them from
11//! being dropped, or replaced by ones the definition no longer fits, by SQL
12//! while a graph uses them.
13
14use 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
30/// The schema DataFusion reserves for `information_schema` views.
31const INFORMATION_SCHEMA: &str = "information_schema";
32
33/// Define `def` in `ctx`, as `CREATE [OR REPLACE] PROPERTY GRAPH [IF NOT EXISTS]`
34/// does. Follows DataFusion's `CREATE TABLE`: an existing graph is an error,
35/// unless `or_replace` (replace it) or `if_not_exists` (keep it, do nothing).
36pub 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    // Plan both tables once, so their columns are known to listings even if a
87    // table under the graph later disappears.
88    for table in [NODES, EDGES] {
89        schema.table(table).await?;
90    }
91    Ok(())
92}
93
94/// The table `table` names: a table or view in the session, or else, when it
95/// looks like a file path or URL, that file opened as a listing table (added
96/// to `files`).
97async 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
118/// The table or view `table` names in the session's catalogs, not opening
119/// files.
120async 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
137/// The error for a graph that does not exist, with a suggestion.
138pub 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
150/// Refuse a DDL statement, planned as `plan`, that would break a graph:
151///
152/// - `DROP TABLE` or `DROP VIEW` of a table a graph uses;
153/// - `DROP SCHEMA` of a graph's schema, or of a schema holding a table a graph
154///   uses;
155/// - `CREATE OR REPLACE` of a table a graph uses, by one whose columns the
156///   graph's definition no longer fits (see [`check_replace`]).
157///
158/// Every other plan passes.
159pub 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            // An empty schema is inferred from the files when the statement
177            // runs, so it cannot be checked here.
178            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
186/// Refuse to drop `table` (a `kind`: "table" or "view") while a graph uses it.
187pub 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
200/// Refuse to drop the schema `schema` while it is a graph's, or holds a table
201/// a graph uses.
202pub 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
231/// Refuse to replace `table` (a `kind`) by one with the columns `schema`
232/// while a graph uses it and its definition does not fit those columns.
233/// `None`: the new columns are not known yet, so the replacement is refused
234/// whenever a graph uses the table.
235pub 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
280/// The graphs with a session table (not a file path) for which `uses` holds.
281fn 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
297/// "graph 'g' uses it" or "graphs 'g', 'h' use it".
298fn 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
307/// "property graph g" or "property graphs g, h".
308fn 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
317/// `table` with the session's default catalog and schema filled in.
318fn 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
325/// Drop the graph called `name` and its schema, as `DROP PROPERTY GRAPH
326/// [IF EXISTS]` does.
327pub 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
346/// The closest of `candidates` to `name` by edit distance, if any is close
347/// enough to be a likely typo: for "did you mean" suggestions.
348pub 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
361/// "; did you mean 'x'?" for the closest candidate, or nothing.
362fn 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
368/// Check that `def`'s columns exist in its tables, whose columns are `edges`
369/// and `vertices`, and have node id types.
370fn 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/// Node ids are integers of any width, or strings.
421#[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
441/// The table `table` names, or an error with a suggestion.
442async 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
458/// Every table in the session, named as a query in the default schema would
459/// name it.
460async 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
511/// The default catalog, as queries see it.
512fn 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
520/// The default catalog without the file-table wrapper (which cannot
521/// deregister a schema), when the session was built by `new_session`.
522fn 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}