Skip to main content

ursa_engine/
lib.rs

1//! Ursa's SQL engine session.
2//!
3//! [`new_session`] builds a DataFusion [`SessionContext`] on which a file is a
4//! table: `SELECT * FROM 'edges.parquet'`, `'edges.csv.gz'`, `'edges/'`,
5//! `'edges/*.parquet'`, and the same over `s3://`, `gs://`, `az://` and
6//! `https://`, with object stores registered the first time a URL is read.
7//!
8//! Graphs are defined per session ([`graph`]). A session may have a `.ursa`
9//! database ([`database`]), which keeps its catalog between processes.
10
11use std::sync::Arc;
12
13use datafusion::catalog::DynamicFileCatalog;
14use datafusion::prelude::{SessionConfig, SessionContext};
15
16mod blocking;
17pub mod database;
18pub mod error;
19mod file_table;
20pub mod graph;
21pub mod options;
22pub mod paths;
23pub mod serial;
24pub mod stores;
25
26pub use blocking::{run_blocking, BlockingError};
27pub use error::user_message;
28pub use file_table::open_file_table;
29pub use stores::{ensure_object_store, ensure_object_store_for};
30
31/// A session with Ursa's defaults.
32pub fn new_session() -> SessionContext {
33    new_session_with_config(SessionConfig::new())
34}
35
36/// A session with Ursa's file tables over `config`. `information_schema` is
37/// turned on, so `SHOW TABLES` and `SHOW COLUMNS` work.
38///
39/// In the browser (wasm32) the session plans for one partition. DataFusion runs
40/// each extra partition as a spawned Tokio task, and the browser build has no
41/// runtime to spawn on; its futures are driven from JavaScript's event loop, as
42/// in DataFusion's own `wasmtest` crate (#278).
43pub fn new_session_with_config(config: SessionConfig) -> SessionContext {
44    #[cfg(target_arch = "wasm32")]
45    let config = config.with_target_partitions(1);
46    let config = graph::with_algorithms(graph::with_registry(config.with_information_schema(true)));
47    let config = database::with_database_slot(options::with_options(config));
48    let ctx = SessionContext::new_with_config(config);
49    // Graph schemas are added to and dropped from the unwrapped catalogs.
50    graph::registry_of(&ctx).set_catalogs(Arc::clone(ctx.state().catalog_list()));
51    // As `SessionContext::enable_url_table` does, with Ursa's factory in place
52    // of DataFusion's.
53    let factory = Arc::new(file_table::FileTableFactory::default());
54    let catalogs = Arc::new(DynamicFileCatalog::new(
55        Arc::clone(ctx.state().catalog_list()),
56        Arc::clone(&factory) as _,
57    ));
58    let session_id = ctx.session_id();
59    let ctx = ctx
60        .into_state_builder()
61        .with_session_id(session_id)
62        .with_catalog_list(catalogs)
63        .with_query_planner(graph::query_planner())
64        .with_physical_optimizer_rule(graph::edge_order_rule());
65    // In the browser, last: every plan runs on the task that polls it.
66    #[cfg(target_arch = "wasm32")]
67    let ctx = ctx.with_physical_optimizer_rule(Arc::new(serial::SingleTask));
68    let ctx: SessionContext = ctx.build().into();
69    graph::register_functions(&ctx).expect("registering a relation planner cannot fail");
70    factory.attach(ctx.state_weak_ref());
71    ctx
72}