mirror of
https://github.com/ModernRelay/omnigraph.git
synced 2026-06-30 02:49:39 +02:00
Merge branch 'main' into devin/mr-983-composite-unique
This commit is contained in:
commit
5133c0f88d
73 changed files with 7297 additions and 451 deletions
|
|
@ -1,6 +1,6 @@
|
|||
[package]
|
||||
name = "omnigraph-cli"
|
||||
version = "0.6.0"
|
||||
version = "0.6.1"
|
||||
edition = "2024"
|
||||
description = "CLI for the Omnigraph graph database."
|
||||
license = "MIT"
|
||||
|
|
@ -13,10 +13,10 @@ name = "omnigraph"
|
|||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
omnigraph = { package = "omnigraph-engine", path = "../omnigraph", version = "0.6.0" }
|
||||
omnigraph-compiler = { path = "../omnigraph-compiler", version = "0.6.0" }
|
||||
omnigraph-policy = { path = "../omnigraph-policy", version = "0.6.0" }
|
||||
omnigraph-server = { path = "../omnigraph-server", version = "0.6.0" }
|
||||
omnigraph = { package = "omnigraph-engine", path = "../omnigraph", version = "0.6.1" }
|
||||
omnigraph-compiler = { path = "../omnigraph-compiler", version = "0.6.1" }
|
||||
omnigraph-policy = { path = "../omnigraph-policy", version = "0.6.1" }
|
||||
omnigraph-server = { path = "../omnigraph-server", version = "0.6.1" }
|
||||
clap = { workspace = true }
|
||||
color-eyre = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ use clap::{Arg, ArgAction, Args, CommandFactory, FromArgMatches, Parser, Subcomm
|
|||
use color_eyre::eyre::{Result, bail};
|
||||
use omnigraph::db::{Omnigraph, ReadTarget, SnapshotId};
|
||||
use omnigraph::loader::LoadMode;
|
||||
use omnigraph::storage::normalize_root_uri;
|
||||
use omnigraph_compiler::query::parser::parse_query;
|
||||
use omnigraph_compiler::schema::parser::parse_schema;
|
||||
use omnigraph_compiler::{
|
||||
|
|
@ -24,9 +25,10 @@ use omnigraph_server::api::{
|
|||
SnapshotTableOutput, commit_output, ingest_output, read_output, schema_apply_output,
|
||||
snapshot_payload,
|
||||
};
|
||||
use omnigraph_server::queries::{QueryRegistry, check, format_check_breakages};
|
||||
use omnigraph_server::{
|
||||
AliasCommand, OmnigraphConfig, PolicyAction, PolicyDecision, PolicyEngine, PolicyRequest,
|
||||
PolicyTestConfig, ReadOutputFormat, load_config,
|
||||
PolicyTestConfig, ReadOutputFormat, graph_resource_id_for_selection, load_config,
|
||||
};
|
||||
use reqwest::Method;
|
||||
use reqwest::header::AUTHORIZATION;
|
||||
|
|
@ -153,6 +155,11 @@ enum Command {
|
|||
#[arg(long)]
|
||||
json: bool,
|
||||
},
|
||||
/// Operate on the server-side stored-query registry (`queries:`).
|
||||
Queries {
|
||||
#[command(subcommand)]
|
||||
command: QueriesCommand,
|
||||
},
|
||||
/// Show graph snapshot
|
||||
Snapshot {
|
||||
/// Graph URI
|
||||
|
|
@ -502,6 +509,35 @@ enum PolicyCommand {
|
|||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Subcommand)]
|
||||
enum QueriesCommand {
|
||||
/// Type-check the stored-query registry against the live schema.
|
||||
///
|
||||
/// Distinct from `omnigraph lint` (which lints one `.gq` file):
|
||||
/// this validates the whole `queries:` registry — opening the graph
|
||||
/// to read its schema and confirming every stored query still
|
||||
/// type-checks. Exits non-zero on any breakage.
|
||||
Validate {
|
||||
/// Graph URI
|
||||
uri: Option<String>,
|
||||
#[arg(long)]
|
||||
target: Option<String>,
|
||||
#[arg(long)]
|
||||
config: Option<PathBuf>,
|
||||
#[arg(long)]
|
||||
json: bool,
|
||||
},
|
||||
/// List the registered stored queries (name, MCP exposure, params).
|
||||
List {
|
||||
#[arg(long)]
|
||||
target: Option<String>,
|
||||
#[arg(long)]
|
||||
config: Option<PathBuf>,
|
||||
#[arg(long)]
|
||||
json: bool,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Args, Clone)]
|
||||
struct ParamsArgs {
|
||||
#[arg(long, conflicts_with = "params_file")]
|
||||
|
|
@ -743,25 +779,66 @@ fn load_cli_config(config_path: Option<&PathBuf>) -> Result<OmnigraphConfig> {
|
|||
Ok(config)
|
||||
}
|
||||
|
||||
fn resolve_policy_engine(config: &OmnigraphConfig) -> Result<PolicyEngine> {
|
||||
let policy_file = config
|
||||
.resolve_policy_file()
|
||||
.ok_or_else(|| color_eyre::eyre::eyre!("policy.file must be set in omnigraph.yaml"))?;
|
||||
PolicyEngine::load_graph(&policy_file, &policy_graph_id(config))
|
||||
#[derive(Debug, Clone)]
|
||||
struct ResolvedCliGraph {
|
||||
uri: String,
|
||||
selected: Option<String>,
|
||||
graph_id: String,
|
||||
policy_file: Option<PathBuf>,
|
||||
is_remote: bool,
|
||||
}
|
||||
|
||||
/// Open a local-URI graph and, when `policy.file` is configured in
|
||||
/// `omnigraph.yaml`, install the resolved `PolicyEngine` on the engine
|
||||
/// handle so every direct-engine write goes through
|
||||
/// `Omnigraph::enforce(...)` (MR-722). Without a configured policy this
|
||||
/// is identical to a bare `Omnigraph::open`.
|
||||
///
|
||||
/// Returns owned `Omnigraph`; chained on top of `Omnigraph::open(...)`'s
|
||||
/// existing future to keep call sites narrow.
|
||||
async fn open_local_db_with_policy(uri: &str, config: &OmnigraphConfig) -> Result<Omnigraph> {
|
||||
let db = Omnigraph::open(uri).await?;
|
||||
if config.resolve_policy_file().is_some() {
|
||||
let engine = Arc::new(resolve_policy_engine(config)?);
|
||||
impl ResolvedCliGraph {
|
||||
fn selected(&self) -> Option<&str> {
|
||||
self.selected.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
struct ResolvedPolicyContext {
|
||||
policy_file: PathBuf,
|
||||
graph_id: String,
|
||||
}
|
||||
|
||||
fn resolve_policy_context(config: &OmnigraphConfig) -> Result<ResolvedPolicyContext> {
|
||||
let selected = config.resolve_policy_tooling_graph_selection()?;
|
||||
let policy_file = config
|
||||
.resolve_policy_file_for(selected)
|
||||
.ok_or_else(|| {
|
||||
color_eyre::eyre::eyre!(
|
||||
"policy.file or graphs.<name>.policy.file must be set in omnigraph.yaml"
|
||||
)
|
||||
})?;
|
||||
let graph_id = match selected {
|
||||
Some(name) => graph_resource_id_for_selection(Some(name), ""),
|
||||
None => graph_resource_id_for_selection(None, "default"),
|
||||
};
|
||||
Ok(ResolvedPolicyContext {
|
||||
policy_file,
|
||||
graph_id,
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_policy_engine(context: &ResolvedPolicyContext) -> Result<PolicyEngine> {
|
||||
PolicyEngine::load_graph(&context.policy_file, &context.graph_id)
|
||||
}
|
||||
|
||||
fn resolve_policy_engine_for_graph(graph: &ResolvedCliGraph) -> Result<PolicyEngine> {
|
||||
let policy_file = graph.policy_file.as_ref().ok_or_else(|| {
|
||||
color_eyre::eyre::eyre!(
|
||||
"policy.file or graphs.<name>.policy.file must be set in omnigraph.yaml"
|
||||
)
|
||||
})?;
|
||||
PolicyEngine::load_graph(policy_file, &graph.graph_id)
|
||||
}
|
||||
|
||||
/// Open a local graph and install the policy resolved for the same graph
|
||||
/// identity that produced the URI. A named graph uses
|
||||
/// `graphs.<name>.policy.file`; an explicit positional URI is anonymous and
|
||||
/// uses the legacy top-level `policy.file`.
|
||||
async fn open_local_db_with_policy(graph: &ResolvedCliGraph) -> Result<Omnigraph> {
|
||||
let db = Omnigraph::open(&graph.uri).await?;
|
||||
if graph.policy_file.is_some() {
|
||||
let engine = Arc::new(resolve_policy_engine_for_graph(graph)?);
|
||||
Ok(db.with_policy(engine as Arc<dyn omnigraph_policy::PolicyChecker>))
|
||||
} else {
|
||||
Ok(db)
|
||||
|
|
@ -778,22 +855,16 @@ fn resolve_cli_actor<'a>(cli_as: Option<&'a str>, config: &'a OmnigraphConfig) -
|
|||
cli_as.or(config.cli.actor.as_deref())
|
||||
}
|
||||
|
||||
fn resolve_policy_tests_path(config: &OmnigraphConfig) -> Result<PathBuf> {
|
||||
config.resolve_policy_tests_file().ok_or_else(|| {
|
||||
color_eyre::eyre::eyre!(
|
||||
"policy.tests.yaml requires policy.file to be set in omnigraph.yaml"
|
||||
)
|
||||
})
|
||||
fn resolve_policy_tests_path(context: &ResolvedPolicyContext) -> PathBuf {
|
||||
context.policy_file.with_file_name("policy.tests.yaml")
|
||||
}
|
||||
|
||||
fn policy_graph_id(config: &OmnigraphConfig) -> String {
|
||||
if let Some(name) = &config.project.name {
|
||||
return name.clone();
|
||||
fn normalize_policy_graph_uri(uri: &str) -> Result<String> {
|
||||
if is_remote_uri(uri) {
|
||||
Ok(uri.trim_end_matches('/').to_string())
|
||||
} else {
|
||||
Ok(normalize_root_uri(uri)?)
|
||||
}
|
||||
config
|
||||
.resolve_target_uri(None, None, config.server_graph_name())
|
||||
.or_else(|_| config.resolve_target_uri(None, None, config.cli_graph_name()))
|
||||
.unwrap_or_else(|_| "default".to_string())
|
||||
}
|
||||
|
||||
fn resolve_remote_bearer_token(
|
||||
|
|
@ -877,6 +948,47 @@ fn resolve_uri(
|
|||
config.resolve_target_uri(cli_uri, cli_target, config.cli_graph_name())
|
||||
}
|
||||
|
||||
fn resolve_cli_graph(
|
||||
config: &OmnigraphConfig,
|
||||
cli_uri: Option<String>,
|
||||
cli_target: Option<&str>,
|
||||
) -> Result<ResolvedCliGraph> {
|
||||
let selected = if cli_uri.is_some() {
|
||||
None
|
||||
} else {
|
||||
cli_target
|
||||
.map(str::to_string)
|
||||
.or_else(|| config.cli_graph_name().map(str::to_string))
|
||||
};
|
||||
config.resolve_graph_selection(selected.as_deref())?;
|
||||
let uri = resolve_uri(config, cli_uri, cli_target)?;
|
||||
let normalized_uri = normalize_policy_graph_uri(&uri)?;
|
||||
let graph_id = graph_resource_id_for_selection(selected.as_deref(), &normalized_uri);
|
||||
Ok(ResolvedCliGraph {
|
||||
graph_id,
|
||||
is_remote: is_remote_uri(&uri),
|
||||
policy_file: config.resolve_policy_file_for(selected.as_deref()),
|
||||
selected,
|
||||
uri,
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_local_graph(
|
||||
config: &OmnigraphConfig,
|
||||
cli_uri: Option<String>,
|
||||
cli_target: Option<&str>,
|
||||
operation: &str,
|
||||
) -> Result<ResolvedCliGraph> {
|
||||
let graph = resolve_cli_graph(config, cli_uri, cli_target)?;
|
||||
if graph.is_remote {
|
||||
bail!(
|
||||
"{} is only supported against local graph URIs in this milestone",
|
||||
operation
|
||||
);
|
||||
}
|
||||
Ok(graph)
|
||||
}
|
||||
|
||||
/// Parse a Go-style compact duration: `7d`, `24h`, `30m`, `90s`, or a plain
|
||||
/// integer as seconds. Used by the `cleanup --older-than` flag.
|
||||
fn parse_duration_arg(s: &str) -> Result<std::time::Duration> {
|
||||
|
|
@ -915,14 +1027,7 @@ fn resolve_local_uri(
|
|||
cli_target: Option<&str>,
|
||||
operation: &str,
|
||||
) -> Result<String> {
|
||||
let uri = resolve_uri(config, cli_uri, cli_target)?;
|
||||
if is_remote_uri(&uri) {
|
||||
bail!(
|
||||
"{} is only supported against local graph URIs in this milestone",
|
||||
operation
|
||||
);
|
||||
}
|
||||
Ok(uri)
|
||||
Ok(resolve_local_graph(config, cli_uri, cli_target, operation)?.uri)
|
||||
}
|
||||
|
||||
fn resolve_branch(
|
||||
|
|
@ -1609,6 +1714,248 @@ async fn execute_query_lint(
|
|||
))
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct QueriesIssue {
|
||||
query: String,
|
||||
message: String,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct QueriesValidateOutput {
|
||||
ok: bool,
|
||||
breakages: Vec<QueriesIssue>,
|
||||
warnings: Vec<QueriesIssue>,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct QueriesParam {
|
||||
name: String,
|
||||
#[serde(rename = "type")]
|
||||
type_name: String,
|
||||
nullable: bool,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct QueriesListItem {
|
||||
name: String,
|
||||
mcp_expose: bool,
|
||||
tool_name: Option<String>,
|
||||
mutation: bool,
|
||||
params: Vec<QueriesParam>,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct QueriesListOutput {
|
||||
queries: Vec<QueriesListItem>,
|
||||
}
|
||||
|
||||
/// Resolve the selected graph to `(local URI, registry selection)` from one
|
||||
/// precedence, so a command's schema and its stored-query registry can never
|
||||
/// come from different graphs. A **positional URI is anonymous** (top-level
|
||||
/// registry, ignoring the configured default graph); otherwise `--target`
|
||||
/// or the configured `cli.graph` names the graph (its per-graph block).
|
||||
/// Mirrors the server's single-mode identity rule.
|
||||
fn resolve_selected_graph(
|
||||
config: &OmnigraphConfig,
|
||||
cli_uri: Option<String>,
|
||||
cli_target: Option<&str>,
|
||||
operation: &str,
|
||||
) -> Result<(String, Option<String>)> {
|
||||
let graph = resolve_local_graph(config, cli_uri, cli_target, operation)?;
|
||||
Ok((graph.uri, graph.selected))
|
||||
}
|
||||
|
||||
/// Load the stored-query registry for an already-resolved graph selection
|
||||
/// (`None` = anonymous → top-level; `Some(name)` = that graph's block).
|
||||
fn load_registry_or_report(
|
||||
config: &OmnigraphConfig,
|
||||
selected: Option<&str>,
|
||||
) -> Result<QueryRegistry> {
|
||||
QueryRegistry::load(config, config.query_entries_for(selected)).map_err(|errors| {
|
||||
color_eyre::eyre::eyre!(
|
||||
"stored-query registry failed to load:\n {}",
|
||||
errors
|
||||
.iter()
|
||||
.map(|e| e.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n ")
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn graph_query_registry_names(config: &OmnigraphConfig) -> Vec<&str> {
|
||||
config
|
||||
.graphs
|
||||
.iter()
|
||||
.filter_map(|(name, graph)| (!graph.queries.is_empty()).then_some(name.as_str()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn resolve_registry_selection_for_list(
|
||||
config: &OmnigraphConfig,
|
||||
target: Option<&str>,
|
||||
) -> Result<Option<String>> {
|
||||
let selected = target
|
||||
.map(str::to_string)
|
||||
.or_else(|| config.cli_graph_name().map(str::to_string));
|
||||
if let Some(name) = selected.as_deref() {
|
||||
config.resolve_graph_selection(Some(name))?;
|
||||
return Ok(selected);
|
||||
}
|
||||
|
||||
if !config.query_entries().is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let graph_names = graph_query_registry_names(config);
|
||||
if graph_names.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
bail!(
|
||||
"stored-query registries are configured for graph{} {} but no graph was selected. Pass `--target {}` or set `cli.graph`.",
|
||||
if graph_names.len() == 1 { "" } else { "s" },
|
||||
graph_names.join(", "),
|
||||
graph_names[0],
|
||||
)
|
||||
}
|
||||
|
||||
fn validate_registry_for_catalog(
|
||||
registry: &QueryRegistry,
|
||||
catalog: &omnigraph_compiler::catalog::Catalog,
|
||||
label: &str,
|
||||
) -> omnigraph::error::Result<()> {
|
||||
let report = check(registry, catalog);
|
||||
if report.has_breakages() {
|
||||
return Err(omnigraph::error::OmniError::manifest(
|
||||
format_check_breakages(label, &report),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn execute_queries_validate(
|
||||
uri: Option<String>,
|
||||
target: Option<String>,
|
||||
config_path: Option<&PathBuf>,
|
||||
json: bool,
|
||||
) -> Result<()> {
|
||||
let config = load_cli_config(config_path)?;
|
||||
// One selection drives both the schema URI and the registry, so a
|
||||
// positional URI and a `--target` can't validate different graphs.
|
||||
let (uri, selected) =
|
||||
resolve_selected_graph(&config, uri, target.as_deref(), "queries validate")?;
|
||||
let registry = load_registry_or_report(&config, selected.as_deref())?;
|
||||
let db = Omnigraph::open(&uri).await?;
|
||||
let report = check(®istry, &db.catalog());
|
||||
|
||||
let output = QueriesValidateOutput {
|
||||
ok: !report.has_breakages(),
|
||||
breakages: report
|
||||
.breakages
|
||||
.iter()
|
||||
.map(|b| QueriesIssue {
|
||||
query: b.query.clone(),
|
||||
message: b.message.clone(),
|
||||
})
|
||||
.collect(),
|
||||
warnings: report
|
||||
.warnings
|
||||
.iter()
|
||||
.map(|w| QueriesIssue {
|
||||
query: w.query.clone(),
|
||||
message: w.message.clone(),
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
|
||||
if json {
|
||||
print_json(&output)?;
|
||||
} else {
|
||||
if output.breakages.is_empty() {
|
||||
println!(
|
||||
"OK {} stored quer{} type-check against the schema",
|
||||
registry.len(),
|
||||
if registry.len() == 1 { "y" } else { "ies" }
|
||||
);
|
||||
}
|
||||
for issue in &output.breakages {
|
||||
println!("ERROR query '{}': {}", issue.query, issue.message);
|
||||
}
|
||||
for issue in &output.warnings {
|
||||
println!("WARN query '{}': {}", issue.query, issue.message);
|
||||
}
|
||||
}
|
||||
|
||||
if report.has_breakages() {
|
||||
io::stdout().flush()?;
|
||||
std::process::exit(1);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn execute_queries_list(
|
||||
target: Option<String>,
|
||||
config_path: Option<&PathBuf>,
|
||||
json: bool,
|
||||
) -> Result<()> {
|
||||
let config = load_cli_config(config_path)?;
|
||||
let selected = resolve_registry_selection_for_list(&config, target.as_deref())?;
|
||||
let registry = load_registry_or_report(&config, selected.as_deref())?;
|
||||
|
||||
let output = QueriesListOutput {
|
||||
queries: registry
|
||||
.iter()
|
||||
.map(|q| QueriesListItem {
|
||||
name: q.name.clone(),
|
||||
mcp_expose: q.expose,
|
||||
tool_name: q.tool_name.clone(),
|
||||
mutation: q.is_mutation(),
|
||||
params: q
|
||||
.decl
|
||||
.params
|
||||
.iter()
|
||||
.map(|p| QueriesParam {
|
||||
name: p.name.clone(),
|
||||
type_name: p.type_name.clone(),
|
||||
nullable: p.nullable,
|
||||
})
|
||||
.collect(),
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
|
||||
if json {
|
||||
print_json(&output)?;
|
||||
} else if output.queries.is_empty() {
|
||||
println!("(no stored queries registered)");
|
||||
} else {
|
||||
for q in &output.queries {
|
||||
let kind = if q.mutation { "mutation" } else { "read" };
|
||||
let params = q
|
||||
.params
|
||||
.iter()
|
||||
.map(|p| {
|
||||
format!(
|
||||
"${}: {}{}",
|
||||
p.name,
|
||||
p.type_name,
|
||||
if p.nullable { "?" } else { "" }
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
let mcp = if q.mcp_expose {
|
||||
format!(" [mcp: {}]", q.tool_name.as_deref().unwrap_or(&q.name))
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
println!("{kind} {}({params}){mcp}", q.name);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn execute_read(
|
||||
uri: &str,
|
||||
query_source: &str,
|
||||
|
|
@ -1655,7 +2002,7 @@ async fn execute_read_remote(
|
|||
}
|
||||
|
||||
async fn execute_change(
|
||||
uri: &str,
|
||||
graph: &ResolvedCliGraph,
|
||||
query_source: &str,
|
||||
query_name: Option<&str>,
|
||||
branch: &str,
|
||||
|
|
@ -1665,7 +2012,7 @@ async fn execute_change(
|
|||
) -> Result<ChangeOutput> {
|
||||
let (selected_name, query_params) = select_named_query(query_source, query_name)?;
|
||||
let params = query_params_from_json(&query_params, params_json)?;
|
||||
let db = open_local_db_with_policy(uri, config).await?;
|
||||
let db = open_local_db_with_policy(graph).await?;
|
||||
let actor = resolve_cli_actor(cli_as_actor, config);
|
||||
let result = db
|
||||
.mutate_as(branch, query_source, &selected_name, ¶ms, actor)
|
||||
|
|
@ -1893,9 +2240,10 @@ async fn main() -> Result<()> {
|
|||
json,
|
||||
} => {
|
||||
let config = load_cli_config(config.as_ref())?;
|
||||
let uri = resolve_local_uri(&config, uri, target.as_deref(), "load")?;
|
||||
let graph = resolve_local_graph(&config, uri, target.as_deref(), "load")?;
|
||||
let uri = graph.uri.clone();
|
||||
let branch = resolve_branch(&config, branch, None, "main");
|
||||
let db = open_local_db_with_policy(&uri, &config).await?;
|
||||
let db = open_local_db_with_policy(&graph).await?;
|
||||
let actor = resolve_cli_actor(cli.as_actor.as_deref(), &config);
|
||||
let result = db
|
||||
.load_file_as(&branch, &data.to_string_lossy(), mode.into(), actor)
|
||||
|
|
@ -1936,10 +2284,11 @@ async fn main() -> Result<()> {
|
|||
let config = load_cli_config(config.as_ref())?;
|
||||
let bearer_token =
|
||||
resolve_remote_bearer_token(&config, uri.as_deref(), target.as_deref())?;
|
||||
let uri = resolve_uri(&config, uri, target.as_deref())?;
|
||||
let graph = resolve_cli_graph(&config, uri, target.as_deref())?;
|
||||
let uri = graph.uri.clone();
|
||||
let branch = resolve_branch(&config, branch, None, "main");
|
||||
let from = resolve_branch(&config, from, None, "main");
|
||||
let payload = if is_remote_uri(&uri) {
|
||||
let payload = if graph.is_remote {
|
||||
let data = fs::read_to_string(&data)?;
|
||||
remote_json::<IngestOutput>(
|
||||
&http_client,
|
||||
|
|
@ -1955,7 +2304,7 @@ async fn main() -> Result<()> {
|
|||
)
|
||||
.await?
|
||||
} else {
|
||||
let db = open_local_db_with_policy(&uri, &config).await?;
|
||||
let db = open_local_db_with_policy(&graph).await?;
|
||||
let actor = resolve_cli_actor(cli.as_actor.as_deref(), &config);
|
||||
let result = db
|
||||
.ingest_file_as(
|
||||
|
|
@ -1986,9 +2335,10 @@ async fn main() -> Result<()> {
|
|||
let config = load_cli_config(config.as_ref())?;
|
||||
let bearer_token =
|
||||
resolve_remote_bearer_token(&config, uri.as_deref(), target.as_deref())?;
|
||||
let uri = resolve_uri(&config, uri, target.as_deref())?;
|
||||
let graph = resolve_cli_graph(&config, uri, target.as_deref())?;
|
||||
let uri = graph.uri.clone();
|
||||
let from = resolve_branch(&config, from, None, "main");
|
||||
let payload = if is_remote_uri(&uri) {
|
||||
let payload = if graph.is_remote {
|
||||
remote_json::<BranchCreateOutput>(
|
||||
&http_client,
|
||||
Method::POST,
|
||||
|
|
@ -2001,7 +2351,7 @@ async fn main() -> Result<()> {
|
|||
)
|
||||
.await?
|
||||
} else {
|
||||
let db = open_local_db_with_policy(&uri, &config).await?;
|
||||
let db = open_local_db_with_policy(&graph).await?;
|
||||
let actor = resolve_cli_actor(cli.as_actor.as_deref(), &config);
|
||||
db.branch_create_from_as(ReadTarget::branch(&from), &name, actor)
|
||||
.await?;
|
||||
|
|
@ -2027,8 +2377,9 @@ async fn main() -> Result<()> {
|
|||
let config = load_cli_config(config.as_ref())?;
|
||||
let bearer_token =
|
||||
resolve_remote_bearer_token(&config, uri.as_deref(), target.as_deref())?;
|
||||
let uri = resolve_uri(&config, uri, target.as_deref())?;
|
||||
let payload = if is_remote_uri(&uri) {
|
||||
let graph = resolve_cli_graph(&config, uri, target.as_deref())?;
|
||||
let uri = graph.uri.clone();
|
||||
let payload = if graph.is_remote {
|
||||
remote_json::<BranchListOutput>(
|
||||
&http_client,
|
||||
Method::GET,
|
||||
|
|
@ -2061,8 +2412,9 @@ async fn main() -> Result<()> {
|
|||
let config = load_cli_config(config.as_ref())?;
|
||||
let bearer_token =
|
||||
resolve_remote_bearer_token(&config, uri.as_deref(), target.as_deref())?;
|
||||
let uri = resolve_uri(&config, uri, target.as_deref())?;
|
||||
let payload = if is_remote_uri(&uri) {
|
||||
let graph = resolve_cli_graph(&config, uri, target.as_deref())?;
|
||||
let uri = graph.uri.clone();
|
||||
let payload = if graph.is_remote {
|
||||
remote_json::<BranchDeleteOutput>(
|
||||
&http_client,
|
||||
Method::DELETE,
|
||||
|
|
@ -2072,7 +2424,7 @@ async fn main() -> Result<()> {
|
|||
)
|
||||
.await?
|
||||
} else {
|
||||
let db = open_local_db_with_policy(&uri, &config).await?;
|
||||
let db = open_local_db_with_policy(&graph).await?;
|
||||
let actor = resolve_cli_actor(cli.as_actor.as_deref(), &config);
|
||||
db.branch_delete_as(&name, actor).await?;
|
||||
BranchDeleteOutput {
|
||||
|
|
@ -2098,9 +2450,10 @@ async fn main() -> Result<()> {
|
|||
let config = load_cli_config(config.as_ref())?;
|
||||
let bearer_token =
|
||||
resolve_remote_bearer_token(&config, uri.as_deref(), target.as_deref())?;
|
||||
let uri = resolve_uri(&config, uri, target.as_deref())?;
|
||||
let graph = resolve_cli_graph(&config, uri, target.as_deref())?;
|
||||
let uri = graph.uri.clone();
|
||||
let into = resolve_branch(&config, into, None, "main");
|
||||
let payload = if is_remote_uri(&uri) {
|
||||
let payload = if graph.is_remote {
|
||||
remote_json::<BranchMergeOutput>(
|
||||
&http_client,
|
||||
Method::POST,
|
||||
|
|
@ -2113,7 +2466,7 @@ async fn main() -> Result<()> {
|
|||
)
|
||||
.await?
|
||||
} else {
|
||||
let db = open_local_db_with_policy(&uri, &config).await?;
|
||||
let db = open_local_db_with_policy(&graph).await?;
|
||||
let actor = resolve_cli_actor(cli.as_actor.as_deref(), &config);
|
||||
let outcome = db.branch_merge_as(&source, &into, actor).await?;
|
||||
BranchMergeOutput {
|
||||
|
|
@ -2248,9 +2601,10 @@ async fn main() -> Result<()> {
|
|||
let config = load_cli_config(config.as_ref())?;
|
||||
let bearer_token =
|
||||
resolve_remote_bearer_token(&config, uri.as_deref(), target.as_deref())?;
|
||||
let uri = resolve_uri(&config, uri, target.as_deref())?;
|
||||
let graph = resolve_cli_graph(&config, uri, target.as_deref())?;
|
||||
let uri = graph.uri.clone();
|
||||
let schema_source = fs::read_to_string(&schema)?;
|
||||
let output = if is_remote_uri(&uri) {
|
||||
let output = if graph.is_remote {
|
||||
// MR-694 PR B: SchemaApplyRequest gained an
|
||||
// allow_data_loss field so Hard-mode drops are no
|
||||
// longer CLI-only. The previous bail is gone; the
|
||||
|
|
@ -2268,13 +2622,22 @@ async fn main() -> Result<()> {
|
|||
)
|
||||
.await?
|
||||
} else {
|
||||
let db = open_local_db_with_policy(&uri, &config).await?;
|
||||
let db = open_local_db_with_policy(&graph).await?;
|
||||
let actor = resolve_cli_actor(cli.as_actor.as_deref(), &config);
|
||||
let registry = load_registry_or_report(&config, graph.selected())?;
|
||||
let registry = (!registry.is_empty()).then_some(registry);
|
||||
let label = graph.selected().unwrap_or(&uri).to_string();
|
||||
let result = db
|
||||
.apply_schema_as(
|
||||
.apply_schema_as_with_catalog_check(
|
||||
&schema_source,
|
||||
omnigraph::db::SchemaApplyOptions { allow_data_loss },
|
||||
actor,
|
||||
|catalog| {
|
||||
if let Some(registry) = registry.as_ref() {
|
||||
validate_registry_for_catalog(registry, catalog, &label)?;
|
||||
}
|
||||
Ok(())
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
schema_apply_output(&uri, result)
|
||||
|
|
@ -2331,6 +2694,23 @@ async fn main() -> Result<()> {
|
|||
.await?;
|
||||
finish_query_lint(&output, json)?;
|
||||
}
|
||||
Command::Queries { command } => match command {
|
||||
QueriesCommand::Validate {
|
||||
uri,
|
||||
target,
|
||||
config,
|
||||
json,
|
||||
} => {
|
||||
execute_queries_validate(uri, target, config.as_ref(), json).await?;
|
||||
}
|
||||
QueriesCommand::List {
|
||||
target,
|
||||
config,
|
||||
json,
|
||||
} => {
|
||||
execute_queries_list(target, config.as_ref(), json)?;
|
||||
}
|
||||
},
|
||||
Command::Snapshot {
|
||||
uri,
|
||||
target,
|
||||
|
|
@ -2436,7 +2816,8 @@ async fn main() -> Result<()> {
|
|||
.as_deref()
|
||||
.or_else(|| alias_config.and_then(|alias| alias.graph.as_deref()));
|
||||
let bearer_token = resolve_remote_bearer_token(&config, uri.as_deref(), target_name)?;
|
||||
let uri = resolve_uri(&config, uri, target_name)?;
|
||||
let graph = resolve_cli_graph(&config, uri, target_name)?;
|
||||
let uri = graph.uri.clone();
|
||||
let query_source = resolve_query_source(
|
||||
&config,
|
||||
query.as_ref(),
|
||||
|
|
@ -2458,7 +2839,7 @@ async fn main() -> Result<()> {
|
|||
alias_config.and_then(|alias| alias.branch.clone()),
|
||||
)?;
|
||||
let query_name = name.or_else(|| alias_config.and_then(|alias| alias.name.clone()));
|
||||
let output = if is_remote_uri(&uri) {
|
||||
let output = if graph.is_remote {
|
||||
execute_read_remote(
|
||||
&http_client,
|
||||
&uri,
|
||||
|
|
@ -2521,7 +2902,8 @@ async fn main() -> Result<()> {
|
|||
.as_deref()
|
||||
.or_else(|| alias_config.and_then(|alias| alias.graph.as_deref()));
|
||||
let bearer_token = resolve_remote_bearer_token(&config, uri.as_deref(), target_name)?;
|
||||
let uri = resolve_uri(&config, uri, target_name)?;
|
||||
let graph = resolve_cli_graph(&config, uri, target_name)?;
|
||||
let uri = graph.uri.clone();
|
||||
let query_source = resolve_query_source(
|
||||
&config,
|
||||
query.as_ref(),
|
||||
|
|
@ -2543,7 +2925,7 @@ async fn main() -> Result<()> {
|
|||
"main",
|
||||
);
|
||||
let query_name = name.or_else(|| alias_config.and_then(|alias| alias.name.clone()));
|
||||
let output = if is_remote_uri(&uri) {
|
||||
let output = if graph.is_remote {
|
||||
execute_change_remote(
|
||||
&http_client,
|
||||
&uri,
|
||||
|
|
@ -2556,7 +2938,7 @@ async fn main() -> Result<()> {
|
|||
.await?
|
||||
} else {
|
||||
execute_change(
|
||||
&uri,
|
||||
&graph,
|
||||
&query_source,
|
||||
query_name.as_deref(),
|
||||
&branch,
|
||||
|
|
@ -2575,20 +2957,19 @@ async fn main() -> Result<()> {
|
|||
Command::Policy { command } => match command {
|
||||
PolicyCommand::Validate { config } => {
|
||||
let config = load_cli_config(config.as_ref())?;
|
||||
let engine = resolve_policy_engine(&config)?;
|
||||
let policy_file = config
|
||||
.resolve_policy_file()
|
||||
.expect("policy file should exist after resolve_policy_engine");
|
||||
let context = resolve_policy_context(&config)?;
|
||||
let engine = resolve_policy_engine(&context)?;
|
||||
println!(
|
||||
"policy valid: {} [{} actors]",
|
||||
policy_file.display(),
|
||||
context.policy_file.display(),
|
||||
engine.known_actor_count()
|
||||
);
|
||||
}
|
||||
PolicyCommand::Test { config } => {
|
||||
let config = load_cli_config(config.as_ref())?;
|
||||
let engine = resolve_policy_engine(&config)?;
|
||||
let tests_path = resolve_policy_tests_path(&config)?;
|
||||
let context = resolve_policy_context(&config)?;
|
||||
let engine = resolve_policy_engine(&context)?;
|
||||
let tests_path = resolve_policy_tests_path(&context);
|
||||
let tests = PolicyTestConfig::load(&tests_path)?;
|
||||
engine.run_tests(&tests)?;
|
||||
println!("policy tests passed: {} cases", tests.cases.len());
|
||||
|
|
@ -2601,7 +2982,8 @@ async fn main() -> Result<()> {
|
|||
target_branch,
|
||||
} => {
|
||||
let config = load_cli_config(config.as_ref())?;
|
||||
let engine = resolve_policy_engine(&config)?;
|
||||
let context = resolve_policy_context(&config)?;
|
||||
let engine = resolve_policy_engine(&context)?;
|
||||
let request = PolicyRequest {
|
||||
action,
|
||||
branch,
|
||||
|
|
@ -2629,18 +3011,19 @@ async fn main() -> Result<()> {
|
|||
"fragments_removed": s.fragments_removed,
|
||||
"fragments_added": s.fragments_added,
|
||||
"committed": s.committed,
|
||||
"skipped": s.skipped.map(|r| r.as_str()),
|
||||
})).collect::<Vec<_>>(),
|
||||
});
|
||||
print_json(&value)?;
|
||||
} else {
|
||||
println!("optimize {} — {} tables", uri, stats.len());
|
||||
for s in &stats {
|
||||
if s.committed {
|
||||
if let Some(reason) = s.skipped {
|
||||
println!(" {:<40} skipped ({reason})", s.table_key);
|
||||
} else if s.committed {
|
||||
println!(
|
||||
" {:<40} frags {} → {} ✓",
|
||||
s.table_key,
|
||||
s.fragments_removed + s.fragments_added - s.fragments_added,
|
||||
s.fragments_added
|
||||
s.table_key, s.fragments_removed, s.fragments_added
|
||||
);
|
||||
} else {
|
||||
println!(" {:<40} no-op", s.table_key);
|
||||
|
|
@ -2699,20 +3082,33 @@ async fn main() -> Result<()> {
|
|||
"table_key": s.table_key,
|
||||
"bytes_removed": s.bytes_removed,
|
||||
"old_versions_removed": s.old_versions_removed,
|
||||
"error": s.error,
|
||||
})).collect::<Vec<_>>(),
|
||||
});
|
||||
print_json(&value)?;
|
||||
} else {
|
||||
let total_bytes: u64 = stats.iter().map(|s| s.bytes_removed).sum();
|
||||
let total_versions: u64 = stats.iter().map(|s| s.old_versions_removed).sum();
|
||||
let failed: Vec<&str> = stats
|
||||
.iter()
|
||||
.filter(|s| s.error.is_some())
|
||||
.map(|s| s.table_key.as_str())
|
||||
.collect();
|
||||
println!(
|
||||
"cleanup {} ({}) — removed {} versions ({} bytes) across {} tables",
|
||||
uri,
|
||||
policy_desc,
|
||||
total_versions,
|
||||
total_bytes,
|
||||
stats.len()
|
||||
stats.len() - failed.len()
|
||||
);
|
||||
if !failed.is_empty() {
|
||||
println!(
|
||||
" {} table(s) failed and will be retried on the next cleanup: {}",
|
||||
failed.len(),
|
||||
failed.join(", ")
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Command::Graphs { command } => match command {
|
||||
|
|
@ -2761,7 +3157,8 @@ mod tests {
|
|||
use super::{
|
||||
DEFAULT_BEARER_TOKEN_ENV, apply_bearer_token, bearer_token_from_env_file,
|
||||
legacy_change_request_body, load_cli_config, load_env_file_into_process,
|
||||
normalize_bearer_token, parse_env_assignment, resolve_remote_bearer_token,
|
||||
normalize_bearer_token, parse_env_assignment, resolve_policy_context,
|
||||
resolve_cli_graph, resolve_remote_bearer_token,
|
||||
};
|
||||
use omnigraph_server::load_config;
|
||||
use reqwest::header::AUTHORIZATION;
|
||||
|
|
@ -3021,4 +3418,150 @@ graphs:
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn graph_identity_resolve_policy_context_named_cli_graph_uses_graph_key_not_project_name_or_uri() {
|
||||
let temp = tempdir().unwrap();
|
||||
let config_path = temp.path().join("omnigraph.yaml");
|
||||
fs::write(
|
||||
&config_path,
|
||||
r#"
|
||||
project:
|
||||
name: misleading-project
|
||||
graphs:
|
||||
local:
|
||||
uri: /tmp/local-policy-graph.omni
|
||||
policy:
|
||||
file: ./policy.yaml
|
||||
cli:
|
||||
graph: local
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = load_config(Some(&config_path)).unwrap();
|
||||
let context = resolve_policy_context(&config).unwrap();
|
||||
assert_eq!(context.graph_id, "local");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn graph_identity_resolve_policy_context_server_graph_uses_graph_key_when_cli_graph_absent() {
|
||||
let temp = tempdir().unwrap();
|
||||
let config_path = temp.path().join("omnigraph.yaml");
|
||||
fs::write(
|
||||
&config_path,
|
||||
r#"
|
||||
project:
|
||||
name: misleading-project
|
||||
graphs:
|
||||
local:
|
||||
uri: /tmp/local-policy-graph.omni
|
||||
policy:
|
||||
file: ./server-policy.yaml
|
||||
server:
|
||||
graph: local
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = load_config(Some(&config_path)).unwrap();
|
||||
let context = resolve_policy_context(&config).unwrap();
|
||||
assert_eq!(context.graph_id, "local");
|
||||
assert!(context.policy_file.ends_with("server-policy.yaml"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn graph_identity_resolve_policy_context_anonymous_uses_top_level_default_identity() {
|
||||
let temp = tempdir().unwrap();
|
||||
let config_path = temp.path().join("omnigraph.yaml");
|
||||
fs::write(
|
||||
&config_path,
|
||||
r#"
|
||||
project:
|
||||
name: misleading-project
|
||||
graphs:
|
||||
local:
|
||||
uri: /tmp/local-policy-graph.omni
|
||||
policy:
|
||||
file: ./top-policy.yaml
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = load_config(Some(&config_path)).unwrap();
|
||||
let context = resolve_policy_context(&config).unwrap();
|
||||
assert_eq!(context.graph_id, "default");
|
||||
assert!(context.policy_file.ends_with("top-policy.yaml"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn graph_identity_resolve_cli_graph_named_target_uses_graph_key_not_project_name_or_uri() {
|
||||
let temp = tempdir().unwrap();
|
||||
let config_path = temp.path().join("omnigraph.yaml");
|
||||
fs::write(
|
||||
&config_path,
|
||||
r#"
|
||||
project:
|
||||
name: misleading-project
|
||||
graphs:
|
||||
prod:
|
||||
uri: s3://bucket/prod-graph/
|
||||
policy:
|
||||
file: ./prod-policy.yaml
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = load_config(Some(&config_path)).unwrap();
|
||||
let graph = resolve_cli_graph(&config, None, Some("prod")).unwrap();
|
||||
assert_eq!(graph.selected(), Some("prod"));
|
||||
assert_eq!(graph.graph_id, "prod");
|
||||
assert_eq!(graph.uri, "s3://bucket/prod-graph/");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn graph_identity_resolve_cli_graph_positional_uri_uses_anonymous_normalized_uri() {
|
||||
let temp = tempdir().unwrap();
|
||||
let config_path = temp.path().join("omnigraph.yaml");
|
||||
fs::write(
|
||||
&config_path,
|
||||
r#"
|
||||
project:
|
||||
name: misleading-project
|
||||
graphs:
|
||||
local:
|
||||
uri: /tmp/configured-graph.omni
|
||||
policy:
|
||||
file: ./policy.yaml
|
||||
cli:
|
||||
graph: local
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = load_config(Some(&config_path)).unwrap();
|
||||
let local_graph_path = temp.path().join("explicit-graph.omni");
|
||||
let local_graph = resolve_cli_graph(
|
||||
&config,
|
||||
Some(format!("file://{}", local_graph_path.display())),
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(local_graph.selected(), None);
|
||||
assert_eq!(
|
||||
local_graph.graph_id,
|
||||
local_graph_path.to_string_lossy().as_ref()
|
||||
);
|
||||
assert_eq!(local_graph.policy_file, None);
|
||||
|
||||
let s3_graph = resolve_cli_graph(
|
||||
&config,
|
||||
Some("s3://bucket/anonymous-graph/".to_string()),
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(s3_graph.selected(), None);
|
||||
assert_eq!(s3_graph.graph_id, "s3://bucket/anonymous-graph");
|
||||
assert_eq!(s3_graph.policy_file, None);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2376,3 +2376,295 @@ fn graphs_list_against_local_uri_errors_with_remote_only_message() {
|
|||
"expected 'remote multi-graph server URL' rejection in stderr; got:\n{stderr}"
|
||||
);
|
||||
}
|
||||
|
||||
fn queries_test_config(graph_uri: &str, entry: &str, gq_file: &str) -> String {
|
||||
format!(
|
||||
"graphs:\n local:\n uri: '{}'\n queries:\n {entry}:\n file: ./{gq_file}\n\
|
||||
cli:\n graph: local\npolicy: {{}}\n",
|
||||
graph_uri.replace('\'', "''")
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queries_validate_exits_zero_on_clean_registry() {
|
||||
let graph = SystemGraph::loaded();
|
||||
graph.write_query(
|
||||
"find_person.gq",
|
||||
"query find_person($name: String) { match { $p: Person { name: $name } } return { $p.age } }",
|
||||
);
|
||||
let config = graph.write_config(
|
||||
"omnigraph.yaml",
|
||||
&queries_test_config(&graph.path().to_string_lossy(), "find_person", "find_person.gq"),
|
||||
);
|
||||
let output = output_success(cli().arg("queries").arg("validate").arg("--config").arg(&config));
|
||||
let stdout = stdout_string(&output);
|
||||
assert!(stdout.contains("OK"), "stdout:\n{stdout}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queries_validate_exits_nonzero_on_type_broken_query() {
|
||||
let graph = SystemGraph::loaded();
|
||||
// `Widget` is not in the fixture schema.
|
||||
graph.write_query("ghost.gq", "query ghost() { match { $w: Widget } return { $w.name } }");
|
||||
let config = graph.write_config(
|
||||
"omnigraph.yaml",
|
||||
&queries_test_config(&graph.path().to_string_lossy(), "ghost", "ghost.gq"),
|
||||
);
|
||||
let output = output_failure(cli().arg("queries").arg("validate").arg("--config").arg(&config));
|
||||
let stdout = stdout_string(&output);
|
||||
assert!(
|
||||
stdout.contains("ghost"),
|
||||
"validation should name the broken query; stdout:\n{stdout}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queries_list_prints_registered_query() {
|
||||
let graph = SystemGraph::loaded();
|
||||
graph.write_query(
|
||||
"find_person.gq",
|
||||
"query find_person($name: String) { match { $p: Person { name: $name } } return { $p.age } }",
|
||||
);
|
||||
// Exposed with an explicit tool name so the list shows the MCP suffix.
|
||||
let config = graph.write_config(
|
||||
"omnigraph.yaml",
|
||||
&format!(
|
||||
concat!(
|
||||
"graphs:\n",
|
||||
" local:\n",
|
||||
" uri: '{}'\n",
|
||||
" queries:\n",
|
||||
" find_person:\n",
|
||||
" file: ./find_person.gq\n",
|
||||
" mcp: {{ expose: true, tool_name: lookup_person }}\n",
|
||||
"cli:\n",
|
||||
" graph: local\n",
|
||||
"policy: {{}}\n",
|
||||
),
|
||||
graph.path().to_string_lossy().replace('\'', "''")
|
||||
),
|
||||
);
|
||||
let output = output_success(cli().arg("queries").arg("list").arg("--config").arg(&config));
|
||||
let stdout = stdout_string(&output);
|
||||
assert!(stdout.contains("find_person"), "stdout:\n{stdout}");
|
||||
assert!(
|
||||
stdout.contains("$name: String"),
|
||||
"list should show typed params; stdout:\n{stdout}"
|
||||
);
|
||||
assert!(
|
||||
stdout.contains("[mcp: lookup_person]"),
|
||||
"list should show the MCP tool name for exposed queries; stdout:\n{stdout}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queries_list_requires_graph_selection_for_per_graph_only_registries() {
|
||||
let graph = SystemGraph::loaded();
|
||||
graph.write_query(
|
||||
"find_person.gq",
|
||||
"query find_person($name: String) { match { $p: Person { name: $name } } return { $p.age } }",
|
||||
);
|
||||
let config = graph.write_config(
|
||||
"omnigraph.yaml",
|
||||
&format!(
|
||||
concat!(
|
||||
"graphs:\n",
|
||||
" local:\n",
|
||||
" uri: '{}'\n",
|
||||
" queries:\n",
|
||||
" find_person:\n",
|
||||
" file: ./find_person.gq\n",
|
||||
"policy: {{}}\n",
|
||||
),
|
||||
graph.path().to_string_lossy().replace('\'', "''")
|
||||
),
|
||||
);
|
||||
|
||||
let output = output_failure(cli().arg("queries").arg("list").arg("--config").arg(&config));
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
assert!(
|
||||
stderr.contains("local") && stderr.contains("--target local"),
|
||||
"error must name the graph and give a concrete selection hint; stderr:\n{stderr}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queries_list_without_graph_selection_lists_top_level_registry() {
|
||||
let graph = SystemGraph::loaded();
|
||||
graph.write_query(
|
||||
"top_find.gq",
|
||||
"query top_find($name: String) { match { $p: Person { name: $name } } return { $p.age } }",
|
||||
);
|
||||
let config = graph.write_config(
|
||||
"omnigraph.yaml",
|
||||
concat!(
|
||||
"queries:\n",
|
||||
" top_find:\n",
|
||||
" file: ./top_find.gq\n",
|
||||
"policy: {}\n",
|
||||
),
|
||||
);
|
||||
|
||||
let output = output_success(cli().arg("queries").arg("list").arg("--config").arg(&config));
|
||||
let stdout = stdout_string(&output);
|
||||
assert!(stdout.contains("top_find"), "stdout:\n{stdout}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queries_list_unknown_target_errors() {
|
||||
// `queries list` opens no graph URI, so unknown-graph validation can't ride
|
||||
// along on URI resolution the way it does for every other command. An
|
||||
// unknown `--target` must still error (naming the graph) instead of
|
||||
// silently falling back to the top-level registry and showing the wrong
|
||||
// (or empty) catalog.
|
||||
let graph = SystemGraph::loaded();
|
||||
graph.write_query(
|
||||
"find_person.gq",
|
||||
"query find_person($name: String) { match { $p: Person { name: $name } } return { $p.age } }",
|
||||
);
|
||||
let config = graph.write_config(
|
||||
"omnigraph.yaml",
|
||||
&queries_test_config(&graph.path().to_string_lossy(), "find_person", "find_person.gq"),
|
||||
);
|
||||
let output = output_failure(
|
||||
cli()
|
||||
.arg("queries")
|
||||
.arg("list")
|
||||
.arg("--target")
|
||||
.arg("nonexistent")
|
||||
.arg("--config")
|
||||
.arg(&config),
|
||||
);
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
assert!(
|
||||
stderr.contains("nonexistent"),
|
||||
"error must name the unknown graph; stderr:\n{stderr}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queries_commands_reject_named_graph_with_populated_top_level_block() {
|
||||
// A named graph (here via `cli.graph`) uses its own `graphs.<name>` block,
|
||||
// so a populated top-level `queries:` block would be silently ignored — a
|
||||
// config the server REFUSES to boot. `queries validate`/`list` must reject
|
||||
// it too (matching boot) instead of validating/listing the per-graph block
|
||||
// and giving a false green.
|
||||
let graph = SystemGraph::loaded();
|
||||
graph.write_query(
|
||||
"find_person.gq",
|
||||
"query find_person($name: String) { match { $p: Person { name: $name } } return { $p.age } }",
|
||||
);
|
||||
let config = graph.write_config(
|
||||
"omnigraph.yaml",
|
||||
&format!(
|
||||
concat!(
|
||||
"graphs:\n",
|
||||
" local:\n",
|
||||
" uri: '{}'\n",
|
||||
" queries:\n",
|
||||
" find_person:\n",
|
||||
" file: ./find_person.gq\n",
|
||||
"cli:\n",
|
||||
" graph: local\n",
|
||||
"queries:\n", // populated top-level block: the coherence violation
|
||||
" legacy:\n",
|
||||
" file: ./legacy.gq\n",
|
||||
"policy: {{}}\n",
|
||||
),
|
||||
graph.path().to_string_lossy().replace('\'', "''")
|
||||
),
|
||||
);
|
||||
// Both resolve `local` from cli.graph (no positional URI), so both must
|
||||
// error and name the graph + the ignored block — like server boot does.
|
||||
for sub in ["validate", "list"] {
|
||||
let output = output_failure(cli().arg("queries").arg(sub).arg("--config").arg(&config));
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
assert!(
|
||||
stderr.contains("local") && stderr.contains("queries"),
|
||||
"`queries {sub}` must reject a named graph with a populated top-level block; stderr:\n{stderr}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queries_validate_exits_nonzero_on_duplicate_tool_name() {
|
||||
// Two exposed queries claiming one MCP tool name is a load-time
|
||||
// collision — `queries validate` must fail (offline, before the engine
|
||||
// opens) and name both queries plus the contested tool.
|
||||
let graph = SystemGraph::loaded();
|
||||
graph.write_query("a.gq", "query a() { match { $p: Person } return { $p.name } }");
|
||||
graph.write_query("b.gq", "query b() { match { $p: Person } return { $p.name } }");
|
||||
let config = graph.write_config(
|
||||
"omnigraph.yaml",
|
||||
&format!(
|
||||
concat!(
|
||||
"graphs:\n",
|
||||
" local:\n",
|
||||
" uri: '{}'\n",
|
||||
" queries:\n",
|
||||
" a:\n",
|
||||
" file: ./a.gq\n",
|
||||
" mcp: {{ expose: true, tool_name: dup }}\n",
|
||||
" b:\n",
|
||||
" file: ./b.gq\n",
|
||||
" mcp: {{ expose: true, tool_name: dup }}\n",
|
||||
"cli:\n",
|
||||
" graph: local\n",
|
||||
"policy: {{}}\n",
|
||||
),
|
||||
graph.path().to_string_lossy().replace('\'', "''")
|
||||
),
|
||||
);
|
||||
let output = output_failure(cli().arg("queries").arg("validate").arg("--config").arg(&config));
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
assert!(
|
||||
stderr.contains("dup") && stderr.contains("'a'") && stderr.contains("'b'"),
|
||||
"duplicate tool name should be reported naming both queries; stderr:\n{stderr}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queries_validate_positional_uri_ignores_default_graph() {
|
||||
// A positional URI is anonymous → the schema AND the registry both come
|
||||
// from top-level, even when `cli.graph` names a graph whose per-graph
|
||||
// queries would fail. Pins that the URI and registry can't diverge.
|
||||
let graph = SystemGraph::loaded();
|
||||
graph.write_query(
|
||||
"clean.gq",
|
||||
"query clean($name: String) { match { $p: Person { name: $name } } return { $p.age } }",
|
||||
);
|
||||
// `Widget` is not in the fixture schema — the default graph's per-graph
|
||||
// query would break validate if it were (wrongly) selected.
|
||||
graph.write_query("broken.gq", "query broken() { match { $w: Widget } return { $w.name } }");
|
||||
let config = graph.write_config(
|
||||
"omnigraph.yaml",
|
||||
concat!(
|
||||
"cli:\n graph: prod\n",
|
||||
"graphs:\n",
|
||||
" prod:\n",
|
||||
" uri: /nonexistent-prod.omni\n",
|
||||
" queries:\n",
|
||||
" broken:\n",
|
||||
" file: ./broken.gq\n",
|
||||
"queries:\n",
|
||||
" clean:\n",
|
||||
" file: ./clean.gq\n",
|
||||
"policy: {}\n",
|
||||
),
|
||||
);
|
||||
// Positional URI = the real loaded graph; selection is anonymous, so the
|
||||
// CLEAN top-level registry validates (not prod's broken one).
|
||||
let output = output_success(
|
||||
cli()
|
||||
.arg("queries")
|
||||
.arg("validate")
|
||||
.arg(graph.path())
|
||||
.arg("--config")
|
||||
.arg(&config),
|
||||
);
|
||||
let stdout = stdout_string(&output);
|
||||
assert!(
|
||||
stdout.contains("OK"),
|
||||
"positional URI must validate the top-level registry, not the cli.graph default; stdout:\n{stdout}"
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -74,14 +74,36 @@ project:
|
|||
graphs:
|
||||
local:
|
||||
uri: {}
|
||||
policy:
|
||||
file: ./policy.yaml
|
||||
cli:
|
||||
graph: local
|
||||
branch: main
|
||||
query:
|
||||
roots:
|
||||
- .
|
||||
policy:
|
||||
file: ./policy.yaml
|
||||
",
|
||||
yaml_string(&graph.path().to_string_lossy())
|
||||
)
|
||||
}
|
||||
|
||||
fn local_policy_server_graph_config(graph: &SystemGraph) -> String {
|
||||
format!(
|
||||
"\
|
||||
project:
|
||||
name: policy-e2e-local
|
||||
graphs:
|
||||
local:
|
||||
uri: {}
|
||||
policy:
|
||||
file: ./policy.yaml
|
||||
server:
|
||||
graph: local
|
||||
cli:
|
||||
branch: main
|
||||
query:
|
||||
roots:
|
||||
- .
|
||||
",
|
||||
yaml_string(&graph.path().to_string_lossy())
|
||||
)
|
||||
|
|
@ -1000,49 +1022,55 @@ query vector_search($q: String) {
|
|||
#[test]
|
||||
fn local_cli_policy_tooling_is_end_to_end() {
|
||||
// Sanity check for the read-only policy CLI surfaces. These don't
|
||||
// mutate the graph — they just parse and evaluate the policy file —
|
||||
// so they don't depend on PR #4's engine-side enforcement.
|
||||
// mutate the graph; they parse and evaluate the effective policy for
|
||||
// named graph selections, including per-graph policy files.
|
||||
let graph = SystemGraph::loaded();
|
||||
let config = graph.write_config("omnigraph-policy.yaml", &local_policy_config(&graph));
|
||||
let server_graph_config = graph.write_config(
|
||||
"omnigraph-policy-server.yaml",
|
||||
&local_policy_server_graph_config(&graph),
|
||||
);
|
||||
graph.write_config("policy.yaml", POLICY_E2E_YAML);
|
||||
graph.write_config("policy.tests.yaml", POLICY_E2E_TESTS_YAML);
|
||||
|
||||
let validate = output_success(
|
||||
cli()
|
||||
.arg("policy")
|
||||
.arg("validate")
|
||||
.arg("--config")
|
||||
.arg(&config),
|
||||
);
|
||||
assert!(stdout_string(&validate).contains("policy valid:"));
|
||||
for config in [&config, &server_graph_config] {
|
||||
let validate = output_success(
|
||||
cli()
|
||||
.arg("policy")
|
||||
.arg("validate")
|
||||
.arg("--config")
|
||||
.arg(config),
|
||||
);
|
||||
assert!(stdout_string(&validate).contains("policy valid:"));
|
||||
|
||||
let tests = output_success(cli().arg("policy").arg("test").arg("--config").arg(&config));
|
||||
assert!(stdout_string(&tests).contains("policy tests passed: 2 cases"));
|
||||
let tests = output_success(cli().arg("policy").arg("test").arg("--config").arg(config));
|
||||
assert!(stdout_string(&tests).contains("policy tests passed: 2 cases"));
|
||||
|
||||
let explain = output_success(
|
||||
cli()
|
||||
.arg("policy")
|
||||
.arg("explain")
|
||||
.arg("--config")
|
||||
.arg(&config)
|
||||
.arg("--actor")
|
||||
.arg("act-bruno")
|
||||
.arg("--action")
|
||||
.arg("change")
|
||||
.arg("--branch")
|
||||
.arg("main"),
|
||||
);
|
||||
let explain_stdout = stdout_string(&explain);
|
||||
assert!(explain_stdout.contains("decision: deny"));
|
||||
assert!(explain_stdout.contains("branch: main"));
|
||||
let explain = output_success(
|
||||
cli()
|
||||
.arg("policy")
|
||||
.arg("explain")
|
||||
.arg("--config")
|
||||
.arg(config)
|
||||
.arg("--actor")
|
||||
.arg("act-bruno")
|
||||
.arg("--action")
|
||||
.arg("change")
|
||||
.arg("--branch")
|
||||
.arg("main"),
|
||||
);
|
||||
let explain_stdout = stdout_string(&explain);
|
||||
assert!(explain_stdout.contains("decision: deny"));
|
||||
assert!(explain_stdout.contains("branch: main"));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_cli_change_enforces_engine_layer_policy() {
|
||||
// Asserts MR-722 PR #4: when `policy.file` is configured in
|
||||
// `omnigraph.yaml`, the CLI loads PolicyEngine into Omnigraph and
|
||||
// every direct-engine write hits `enforce(action, scope, actor)` —
|
||||
// identical to what the HTTP server gets, regardless of transport.
|
||||
// Asserts MR-722 PR #4: when the selected graph has a configured
|
||||
// policy file, the CLI loads PolicyEngine into Omnigraph and every
|
||||
// direct-engine write hits `enforce(action, scope, actor)` — identical
|
||||
// to what the HTTP server gets, regardless of transport.
|
||||
//
|
||||
// Three cases, each discriminating:
|
||||
//
|
||||
|
|
@ -1135,6 +1163,32 @@ fn local_cli_change_enforces_engine_layer_policy() {
|
|||
assert_eq!(verify["rows"][0]["p.name"], "RagnorOnMain");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_cli_positional_uri_does_not_inherit_default_graph_policy() {
|
||||
let graph = SystemGraph::loaded();
|
||||
let config = graph.write_config("omnigraph-policy.yaml", &local_policy_config(&graph));
|
||||
graph.write_config("policy.yaml", POLICY_E2E_YAML);
|
||||
let mutation_file = insert_person_query(&graph, "system-local-policy-positional.gq");
|
||||
|
||||
let allowed = parse_stdout_json(&output_success(
|
||||
cli()
|
||||
.arg("--as")
|
||||
.arg("act-bruno")
|
||||
.arg("change")
|
||||
.arg("--config")
|
||||
.arg(&config)
|
||||
.arg("--uri")
|
||||
.arg(graph.path())
|
||||
.arg("--query")
|
||||
.arg(&mutation_file)
|
||||
.arg("--params")
|
||||
.arg(r#"{"name":"PositionalUriBruno","age":4}"#)
|
||||
.arg("--json"),
|
||||
));
|
||||
assert_eq!(allowed["affected_nodes"], 1);
|
||||
assert_eq!(allowed["actor_id"], "act-bruno");
|
||||
}
|
||||
|
||||
// ─── MR-722 PR A: CLI×writer matrix ───────────────────────────────────────
|
||||
//
|
||||
// The change writer is covered above by `local_cli_change_enforces_engine_layer_policy`.
|
||||
|
|
@ -1293,6 +1347,62 @@ fn local_cli_schema_apply_enforces_engine_layer_policy() {
|
|||
assert_eq!(allowed["applied"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_cli_schema_apply_rejects_stored_query_breakage_before_publish() {
|
||||
let graph = SystemGraph::loaded();
|
||||
graph.write_query(
|
||||
"stored-find-person.gq",
|
||||
"query find_person($name: String) { match { $p: Person { name: $name } } return { $p.age } }",
|
||||
);
|
||||
let config = graph.write_config(
|
||||
"omnigraph-stored-query-schema.yaml",
|
||||
&format!(
|
||||
"\
|
||||
graphs:
|
||||
local:
|
||||
uri: {}
|
||||
queries:
|
||||
find_person:
|
||||
file: ./stored-find-person.gq
|
||||
cli:
|
||||
graph: local
|
||||
branch: main
|
||||
query:
|
||||
roots:
|
||||
- .
|
||||
policy: {{}}
|
||||
",
|
||||
yaml_string(&graph.path().to_string_lossy())
|
||||
),
|
||||
);
|
||||
let renamed_schema = std::fs::read_to_string(fixture("test.pg"))
|
||||
.unwrap()
|
||||
.replace("age: I32?", "years: I32? @rename_from(\"age\")");
|
||||
let schema_path = graph.write_file("stored-query-breaks.pg", &renamed_schema);
|
||||
|
||||
let rejected = output_failure(
|
||||
cli()
|
||||
.arg("schema")
|
||||
.arg("apply")
|
||||
.arg("--config")
|
||||
.arg(&config)
|
||||
.arg("--schema")
|
||||
.arg(&schema_path)
|
||||
.arg("--json"),
|
||||
);
|
||||
let stderr = String::from_utf8_lossy(&rejected.stderr);
|
||||
assert!(
|
||||
stderr.contains("find_person") && stderr.contains("schema check"),
|
||||
"schema apply should reject the stored-query breakage before publish; stderr: {stderr}"
|
||||
);
|
||||
|
||||
let schema = stdout_string(&output_success(
|
||||
cli().arg("schema").arg("show").arg("--config").arg(&config),
|
||||
));
|
||||
assert!(schema.contains("age: I32?"));
|
||||
assert!(!schema.contains("years: I32?"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_cli_branch_create_enforces_engine_layer_policy() {
|
||||
let graph = SystemGraph::loaded();
|
||||
|
|
@ -1448,6 +1558,8 @@ project:
|
|||
graphs:
|
||||
local:
|
||||
uri: {}
|
||||
policy:
|
||||
file: ./policy.yaml
|
||||
cli:
|
||||
graph: local
|
||||
branch: main
|
||||
|
|
@ -1455,8 +1567,6 @@ cli:
|
|||
query:
|
||||
roots:
|
||||
- .
|
||||
policy:
|
||||
file: ./policy.yaml
|
||||
",
|
||||
yaml_string(&graph.path().to_string_lossy()),
|
||||
actor,
|
||||
|
|
|
|||
|
|
@ -60,10 +60,10 @@ project:
|
|||
graphs:
|
||||
local:
|
||||
uri: {}
|
||||
policy:
|
||||
file: ./policy.yaml
|
||||
server:
|
||||
graph: local
|
||||
policy:
|
||||
file: ./policy.yaml
|
||||
",
|
||||
yaml_string(&graph.path().to_string_lossy())
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[package]
|
||||
name = "omnigraph-compiler"
|
||||
version = "0.6.0"
|
||||
version = "0.6.1"
|
||||
edition = "2024"
|
||||
description = "Schema/query compiler for Omnigraph. Zero Lance dependency."
|
||||
license = "MIT"
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[package]
|
||||
name = "omnigraph-policy"
|
||||
version = "0.6.0"
|
||||
version = "0.6.1"
|
||||
edition = "2024"
|
||||
description = "Policy / authorization layer for Omnigraph — Cedar-backed PolicyEngine, PolicyChecker trait, ResourceScope enum."
|
||||
license = "MIT"
|
||||
|
|
|
|||
|
|
@ -56,6 +56,21 @@ pub enum PolicyAction {
|
|||
/// from v0.6.0; operators add and remove graphs by editing
|
||||
/// `omnigraph.yaml` and restarting.
|
||||
GraphList,
|
||||
/// Gates invoking a server-side stored query by name. Per-graph and
|
||||
/// **graph-scoped** (no branch dimension, like `Admin`): the per-branch
|
||||
/// access of the query body is enforced by the inner `Read`/`Change`
|
||||
/// gate, so branch-scoping this outer gate would be redundant (and was
|
||||
/// wrong for snapshot reads). A rule that sets `branch_scope` on
|
||||
/// `invoke_query` is rejected by `validate()`. In this release it is
|
||||
/// **coarse**: an `invoke_query` allow rule permits *any* stored query
|
||||
/// on the graph (no per-query dimension yet); a future, additive
|
||||
/// refinement adds an optional query-name scope.
|
||||
///
|
||||
/// This gate sits at the HTTP boundary. The engine `_as` writers still
|
||||
/// enforce `Read`/`Change` per the query body, so a stored *mutation*
|
||||
/// is double-gated: `invoke_query` to reach the tool, plus `change` for
|
||||
/// the write itself.
|
||||
InvokeQuery,
|
||||
}
|
||||
|
||||
impl PolicyAction {
|
||||
|
|
@ -70,6 +85,7 @@ impl PolicyAction {
|
|||
Self::BranchMerge => "branch_merge",
|
||||
Self::Admin => "admin",
|
||||
Self::GraphList => "graph_list",
|
||||
Self::InvokeQuery => "invoke_query",
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -99,7 +115,8 @@ impl PolicyAction {
|
|||
| Self::BranchCreate
|
||||
| Self::BranchDelete
|
||||
| Self::BranchMerge
|
||||
| Self::Admin => PolicyResourceKind::Graph,
|
||||
| Self::Admin
|
||||
| Self::InvokeQuery => PolicyResourceKind::Graph,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -155,6 +172,7 @@ impl FromStr for PolicyAction {
|
|||
"branch_merge" => Ok(Self::BranchMerge),
|
||||
"admin" => Ok(Self::Admin),
|
||||
"graph_list" => Ok(Self::GraphList),
|
||||
"invoke_query" => Ok(Self::InvokeQuery),
|
||||
other => bail!("unknown policy action '{other}'"),
|
||||
}
|
||||
}
|
||||
|
|
@ -806,6 +824,7 @@ namespace Omnigraph {
|
|||
action "branch_delete" appliesTo { principal: Actor, resource: Graph, context: RequestContext };
|
||||
action "branch_merge" appliesTo { principal: Actor, resource: Graph, context: RequestContext };
|
||||
action "admin" appliesTo { principal: Actor, resource: Graph, context: RequestContext };
|
||||
action "invoke_query" appliesTo { principal: Actor, resource: Graph, context: RequestContext };
|
||||
|
||||
action "graph_list" appliesTo { principal: Actor, resource: Server, context: RequestContext };
|
||||
}
|
||||
|
|
@ -1264,6 +1283,80 @@ rules:
|
|||
assert!(!deny.allowed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invoke_query_authorizes_per_graph() {
|
||||
let policy: PolicyConfig = serde_yaml::from_str(
|
||||
r#"
|
||||
version: 1
|
||||
groups:
|
||||
team: [act-alice]
|
||||
others: [act-bruno]
|
||||
rules:
|
||||
- id: team-invoke-queries
|
||||
allow:
|
||||
actors: { group: team }
|
||||
actions: [invoke_query]
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let engine = PolicyCompiler::compile(&policy, "graph").unwrap();
|
||||
|
||||
let allow = engine
|
||||
.authorize(
|
||||
"act-alice",
|
||||
&PolicyRequest {
|
||||
action: PolicyAction::InvokeQuery,
|
||||
branch: None,
|
||||
target_branch: None,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
assert!(allow.allowed);
|
||||
assert_eq!(
|
||||
allow.matched_rule_id.as_deref(),
|
||||
Some("team-invoke-queries")
|
||||
);
|
||||
|
||||
// Actor outside the group → deny.
|
||||
let deny = engine
|
||||
.authorize(
|
||||
"act-bruno",
|
||||
&PolicyRequest {
|
||||
action: PolicyAction::InvokeQuery,
|
||||
branch: None,
|
||||
target_branch: None,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
assert!(!deny.allowed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invoke_query_rejects_branch_scope() {
|
||||
// invoke_query is graph-scoped (like admin) — per-branch access is
|
||||
// enforced by the inner read/change gate — so a rule that puts a
|
||||
// `branch_scope` qualifier on it is rejected at validate().
|
||||
let policy: PolicyConfig = serde_yaml::from_str(
|
||||
r#"
|
||||
version: 1
|
||||
groups:
|
||||
team: [act-alice]
|
||||
rules:
|
||||
- id: team-invoke-any-branch
|
||||
allow:
|
||||
actors: { group: team }
|
||||
actions: [invoke_query]
|
||||
branch_scope: any
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let err = policy.validate().unwrap_err().to_string();
|
||||
assert!(
|
||||
err.contains("branch_scope") && err.contains("invoke_query"),
|
||||
"branch_scope on invoke_query must be rejected: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn server_scoped_rule_cannot_use_branch_scope() {
|
||||
let policy: PolicyConfig = serde_yaml::from_str(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[package]
|
||||
name = "omnigraph-server"
|
||||
version = "0.6.0"
|
||||
version = "0.6.1"
|
||||
edition = "2024"
|
||||
description = "HTTP server for the Omnigraph graph database."
|
||||
license = "MIT"
|
||||
|
|
@ -19,9 +19,9 @@ default = []
|
|||
aws = ["dep:aws-config", "dep:aws-sdk-secretsmanager"]
|
||||
|
||||
[dependencies]
|
||||
omnigraph = { package = "omnigraph-engine", path = "../omnigraph", version = "0.6.0" }
|
||||
omnigraph-compiler = { path = "../omnigraph-compiler", version = "0.6.0" }
|
||||
omnigraph-policy = { path = "../omnigraph-policy", version = "0.6.0" }
|
||||
omnigraph = { package = "omnigraph-engine", path = "../omnigraph", version = "0.6.1" }
|
||||
omnigraph-compiler = { path = "../omnigraph-compiler", version = "0.6.1" }
|
||||
omnigraph-policy = { path = "../omnigraph-policy", version = "0.6.1" }
|
||||
axum = { workspace = true }
|
||||
clap = { workspace = true }
|
||||
color-eyre = { workspace = true }
|
||||
|
|
|
|||
|
|
@ -1,8 +1,11 @@
|
|||
use omnigraph::db::{GraphCommit, MergeOutcome, ReadTarget, SchemaApplyResult, Snapshot};
|
||||
use omnigraph::error::{MergeConflict, MergeConflictKind};
|
||||
use omnigraph::loader::{IngestResult, LoadMode};
|
||||
use crate::queries::StoredQuery;
|
||||
use omnigraph_compiler::SchemaMigrationStep;
|
||||
use omnigraph_compiler::query::ast::Param;
|
||||
use omnigraph_compiler::result::QueryResult;
|
||||
use omnigraph_compiler::types::{PropType, ScalarType};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use utoipa::{IntoParams, ToSchema};
|
||||
|
|
@ -300,6 +303,162 @@ pub struct ChangeRequest {
|
|||
pub branch: Option<String>,
|
||||
}
|
||||
|
||||
/// Body for `POST /queries/{name}` — invokes the server-side stored query
|
||||
/// named in the path. The query source and name come from the registry,
|
||||
/// never the body; only the runtime inputs are supplied here.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize, ToSchema)]
|
||||
pub struct InvokeStoredQueryRequest {
|
||||
/// JSON object whose keys match the stored query's declared parameters.
|
||||
#[serde(default)]
|
||||
pub params: Option<Value>,
|
||||
/// Branch to run against. Defaults to `main`; for a stored mutation the
|
||||
/// write targets this branch.
|
||||
#[serde(default)]
|
||||
pub branch: Option<String>,
|
||||
/// Snapshot id to read from (read queries only — rejected for a stored
|
||||
/// mutation). Mutually exclusive with `branch`.
|
||||
#[serde(default)]
|
||||
pub snapshot: Option<String>,
|
||||
}
|
||||
|
||||
/// Response for `POST /queries/{name}`: the read envelope for a stored
|
||||
/// read, or the mutation envelope for a stored mutation. Serialized
|
||||
/// **untagged**, so the wire shape is exactly [`ReadOutput`] or
|
||||
/// [`ChangeOutput`] — classification follows the stored query, not a
|
||||
/// wrapper field.
|
||||
#[derive(Debug, Serialize, ToSchema)]
|
||||
#[serde(untagged)]
|
||||
pub enum InvokeStoredQueryResponse {
|
||||
Read(ReadOutput),
|
||||
Change(ChangeOutput),
|
||||
}
|
||||
|
||||
/// The kind of a stored-query parameter, decomposed so a client (e.g. an
|
||||
/// MCP server) can build a typed input schema with a closed `match` and
|
||||
/// never re-parse omnigraph's type spelling. `bigint`/`date`/`datetime`/
|
||||
/// `blob` are carried as JSON strings on the wire: a 64-bit integer past
|
||||
/// 2^53 loses precision as a JSON number, and Date/DateTime are ISO
|
||||
/// strings, Blob a blob-URI string.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, ToSchema)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ParamKind {
|
||||
String,
|
||||
Bool,
|
||||
Int,
|
||||
#[serde(rename = "bigint")]
|
||||
BigInt,
|
||||
Float,
|
||||
Date,
|
||||
#[serde(rename = "datetime")]
|
||||
DateTime,
|
||||
Blob,
|
||||
Vector,
|
||||
List,
|
||||
}
|
||||
|
||||
/// One declared parameter of a stored query, projected for the catalog.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
||||
pub struct ParamDescriptor {
|
||||
pub name: String,
|
||||
pub kind: ParamKind,
|
||||
/// Element kind when `kind == list` (always a scalar — the grammar
|
||||
/// forbids lists of vectors or nested lists).
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub item_kind: Option<ParamKind>,
|
||||
/// Dimension when `kind == vector`.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub vector_dim: Option<u32>,
|
||||
/// `false` → the caller must supply it; `true` → optional.
|
||||
pub nullable: bool,
|
||||
}
|
||||
|
||||
/// One entry in the stored-query catalog (`GET /queries`).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
||||
pub struct QueryCatalogEntry {
|
||||
/// Registry key / invoke path segment (`POST /queries/{name}`).
|
||||
pub name: String,
|
||||
/// MCP tool id (the `tool_name` override, else `name`).
|
||||
pub tool_name: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub instruction: Option<String>,
|
||||
/// `true` for a stored mutation → an MCP read-only hint of `false`.
|
||||
pub mutation: bool,
|
||||
pub params: Vec<ParamDescriptor>,
|
||||
}
|
||||
|
||||
/// Response for `GET /queries`: the `mcp.expose` subset of a graph's
|
||||
/// stored-query registry, each with typed parameters.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
||||
pub struct QueriesCatalogOutput {
|
||||
pub queries: Vec<QueryCatalogEntry>,
|
||||
}
|
||||
|
||||
/// Total map from a resolved scalar to its catalog kind. Exhaustive on
|
||||
/// purpose: a new `ScalarType` is a compile error here until catalogued.
|
||||
fn scalar_kind(scalar: ScalarType) -> ParamKind {
|
||||
match scalar {
|
||||
ScalarType::String => ParamKind::String,
|
||||
ScalarType::Bool => ParamKind::Bool,
|
||||
ScalarType::I32 | ScalarType::U32 => ParamKind::Int,
|
||||
ScalarType::I64 | ScalarType::U64 => ParamKind::BigInt,
|
||||
ScalarType::F32 | ScalarType::F64 => ParamKind::Float,
|
||||
ScalarType::Date => ParamKind::Date,
|
||||
ScalarType::DateTime => ParamKind::DateTime,
|
||||
ScalarType::Blob => ParamKind::Blob,
|
||||
ScalarType::Vector(_) => ParamKind::Vector,
|
||||
}
|
||||
}
|
||||
|
||||
fn param_descriptor(param: &Param) -> ParamDescriptor {
|
||||
match PropType::from_param_type_name(¶m.type_name, param.nullable) {
|
||||
Some(pt) if pt.list => ParamDescriptor {
|
||||
name: param.name.clone(),
|
||||
kind: ParamKind::List,
|
||||
item_kind: Some(scalar_kind(pt.scalar)),
|
||||
vector_dim: None,
|
||||
nullable: param.nullable,
|
||||
},
|
||||
Some(pt) => {
|
||||
let (kind, vector_dim) = match pt.scalar {
|
||||
ScalarType::Vector(dim) => (ParamKind::Vector, Some(dim)),
|
||||
other => (scalar_kind(other), None),
|
||||
};
|
||||
ParamDescriptor {
|
||||
name: param.name.clone(),
|
||||
kind,
|
||||
item_kind: None,
|
||||
vector_dim,
|
||||
nullable: param.nullable,
|
||||
}
|
||||
}
|
||||
// Unreachable for a parsed query (every declared param type is
|
||||
// grammatical); fall back to an opaque string so the field is still
|
||||
// usable rather than dropped.
|
||||
None => ParamDescriptor {
|
||||
name: param.name.clone(),
|
||||
kind: ParamKind::String,
|
||||
item_kind: None,
|
||||
vector_dim: None,
|
||||
nullable: param.nullable,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Project a loaded stored query into its catalog entry (typed params,
|
||||
/// MCP tool name, read/mutate flag, description/instruction).
|
||||
pub fn query_catalog_entry(query: &StoredQuery) -> QueryCatalogEntry {
|
||||
QueryCatalogEntry {
|
||||
name: query.name.clone(),
|
||||
tool_name: query.effective_tool_name().to_string(),
|
||||
description: query.decl.description.clone(),
|
||||
instruction: query.decl.instruction.clone(),
|
||||
mutation: query.is_mutation(),
|
||||
params: query.decl.params.iter().map(param_descriptor).collect(),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize, ToSchema)]
|
||||
pub struct SchemaApplyRequest {
|
||||
/// Project schema in `.pg` source form. The diff against the current
|
||||
|
|
|
|||
|
|
@ -9,6 +9,13 @@ use serde::{Deserialize, Serialize};
|
|||
|
||||
pub const DEFAULT_CONFIG_FILE: &str = "omnigraph.yaml";
|
||||
|
||||
pub fn graph_resource_id_for_selection(
|
||||
selected_graph: Option<&str>,
|
||||
normalized_uri: &str,
|
||||
) -> String {
|
||||
selected_graph.unwrap_or(normalized_uri).to_string()
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct ProjectConfig {
|
||||
pub name: Option<String>,
|
||||
|
|
@ -24,6 +31,14 @@ pub struct TargetConfig {
|
|||
/// graph's HTTP-layer Cedar enforcement.
|
||||
#[serde(default)]
|
||||
pub policy: PolicySettings,
|
||||
/// Per-graph stored-query registry: an inline `name -> entry`
|
||||
/// map. Mirrors the per-graph `policy` shape — each
|
||||
/// `graphs.<id>.queries` declares that graph's stored queries. Absent
|
||||
/// (or empty) = no stored queries for the graph. v1 is inline-only;
|
||||
/// an external `queries.yaml` manifest indirection is a deferred
|
||||
/// convenience.
|
||||
#[serde(default)]
|
||||
pub queries: BTreeMap<String, QueryEntry>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Serialize, Deserialize, ValueEnum)]
|
||||
|
|
@ -90,6 +105,50 @@ pub struct PolicySettings {
|
|||
pub file: Option<String>,
|
||||
}
|
||||
|
||||
/// One stored-query registry entry. The map **key** is the query's
|
||||
/// identity — it must equal the `query <name>` symbol declared inside
|
||||
/// the referenced `.gq` file (asserted when the registry loads).
|
||||
/// Renaming the key (or the symbol) is a breaking change to callers, by
|
||||
/// design.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct QueryEntry {
|
||||
/// Path to the `.gq` file (relative to the config's `base_dir`). The
|
||||
/// file may declare several queries; the registry selects the one
|
||||
/// whose symbol matches the map key.
|
||||
pub file: String,
|
||||
#[serde(default)]
|
||||
pub mcp: McpSettings,
|
||||
}
|
||||
|
||||
/// MCP exposure for a stored query. A *deployment* concern (the same
|
||||
/// `.gq` may be exposed in one graph and hidden in another), so it lives
|
||||
/// in YAML rather than in the `.gq` source. **Default `expose: true`** —
|
||||
/// declaring a query in the manifest *is* the opt-in, so it appears in the
|
||||
/// MCP tool catalog (`GET /queries`) by default; set `expose: false` to
|
||||
/// keep a query HTTP/service-callable but hidden from the agent tool list.
|
||||
/// `expose` governs catalog membership only — it is **not** an
|
||||
/// authorization gate (invocation is gated by `invoke_query`), so a hidden
|
||||
/// query is still invocable by name with the right permission.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpSettings {
|
||||
#[serde(default = "mcp_expose_default")]
|
||||
pub expose: bool,
|
||||
pub tool_name: Option<String>,
|
||||
}
|
||||
|
||||
fn mcp_expose_default() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
impl Default for McpSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
expose: mcp_expose_default(),
|
||||
tool_name: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum AliasCommand {
|
||||
|
|
@ -137,6 +196,12 @@ pub struct OmnigraphConfig {
|
|||
pub aliases: BTreeMap<String, AliasConfig>,
|
||||
#[serde(default)]
|
||||
pub policy: PolicySettings,
|
||||
/// Top-level stored-query registry, used in single-graph
|
||||
/// mode — mirrors how the top-level `policy` applies to the single
|
||||
/// graph. In multi-graph mode this is unused; each graph's
|
||||
/// `graphs.<id>.queries` applies instead.
|
||||
#[serde(default)]
|
||||
pub queries: BTreeMap<String, QueryEntry>,
|
||||
#[serde(skip)]
|
||||
base_dir: PathBuf,
|
||||
}
|
||||
|
|
@ -152,6 +217,7 @@ impl Default for OmnigraphConfig {
|
|||
query: QueryDefaults::default(),
|
||||
aliases: BTreeMap::new(),
|
||||
policy: PolicySettings::default(),
|
||||
queries: BTreeMap::new(),
|
||||
base_dir: PathBuf::new(),
|
||||
}
|
||||
}
|
||||
|
|
@ -244,6 +310,124 @@ impl OmnigraphConfig {
|
|||
.map(|path| self.resolve_config_path(path))
|
||||
}
|
||||
|
||||
/// The top-level stored-query registry entries (single-graph mode).
|
||||
pub fn query_entries(&self) -> &BTreeMap<String, QueryEntry> {
|
||||
&self.queries
|
||||
}
|
||||
|
||||
/// The per-graph stored-query registry entries for a named target
|
||||
/// (multi-graph mode). Returns `None` if the target is unknown.
|
||||
pub fn target_query_entries(
|
||||
&self,
|
||||
target_name: &str,
|
||||
) -> Option<&BTreeMap<String, QueryEntry>> {
|
||||
self.graphs.get(target_name).map(|target| &target.queries)
|
||||
}
|
||||
|
||||
/// The stored-query registry entries that apply for a graph
|
||||
/// selection — the single definition of "which `queries:` block
|
||||
/// governs graph X", shared by server boot and the CLI so the two
|
||||
/// can't drift. A named graph present in `graphs:` uses its
|
||||
/// per-graph block; everything else (no selection, or a name that is
|
||||
/// not a known graph, e.g. a bare URI) falls back to the top-level
|
||||
/// block (single-graph mode).
|
||||
pub fn query_entries_for(&self, graph: Option<&str>) -> &BTreeMap<String, QueryEntry> {
|
||||
match graph {
|
||||
Some(name) if self.graphs.contains_key(name) => &self.graphs[name].queries,
|
||||
_ => &self.queries,
|
||||
}
|
||||
}
|
||||
|
||||
/// The single CLI gate that turns a raw graph selection into a *validated*
|
||||
/// one — the fallible counterpart to the infallible
|
||||
/// [`OmnigraphConfig::query_entries_for`]. Both `queries` subcommands route
|
||||
/// their selection through here so neither can skip a check the other (or
|
||||
/// server boot) applies:
|
||||
/// * a known name passes through, but only after the same coherence check
|
||||
/// server boot enforces
|
||||
/// ([`OmnigraphConfig::ensure_top_level_blocks_honored`]) — a named graph
|
||||
/// with a populated top-level block is rejected;
|
||||
/// * an unknown name errors with the **same** message
|
||||
/// [`OmnigraphConfig::resolve_target_uri`] produces, so a command that
|
||||
/// opens no URI rejects an unknown `--target` exactly like the
|
||||
/// URI-resolving commands do;
|
||||
/// * an anonymous selection (`None`, e.g. a bare URI) stays anonymous,
|
||||
/// resolving to the top-level registry downstream (top-level honored).
|
||||
pub fn resolve_graph_selection<'a>(&self, graph: Option<&'a str>) -> Result<Option<&'a str>> {
|
||||
match graph {
|
||||
Some(name) if self.graphs.contains_key(name) => {
|
||||
self.ensure_top_level_blocks_honored(Some(name))?;
|
||||
Ok(Some(name))
|
||||
}
|
||||
Some(name) => bail!("graph '{}' not found in {}", name, DEFAULT_CONFIG_FILE),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn resolve_policy_tooling_graph_selection(&self) -> Result<Option<&str>> {
|
||||
self.resolve_graph_selection(self.cli_graph_name().or_else(|| self.server_graph_name()))
|
||||
}
|
||||
|
||||
/// The policy file that applies for a graph selection — the policy
|
||||
/// sibling of [`OmnigraphConfig::query_entries_for`], so policy and
|
||||
/// queries resolve by the same identity rule. A named graph in
|
||||
/// `graphs:` uses its per-graph `policy.file` with **no** top-level
|
||||
/// fallback (a named graph with no per-graph policy has no policy —
|
||||
/// that keeps the boot-time coherence check meaningful); anything else
|
||||
/// (no selection, or a bare URI) uses the top-level `policy.file`.
|
||||
pub fn resolve_policy_file_for(&self, graph: Option<&str>) -> Option<PathBuf> {
|
||||
match graph {
|
||||
Some(name) if self.graphs.contains_key(name) => self.resolve_target_policy_file(name),
|
||||
_ => self.resolve_policy_file(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Names of any top-level config blocks (`policy.file`, `queries:`)
|
||||
/// that are populated. Used by the boot-time coherence check: when a
|
||||
/// **named** graph is served (single-mode by name, or multi-mode),
|
||||
/// the top-level blocks are not honored, so a populated one is a
|
||||
/// configuration error rather than a silent no-op.
|
||||
pub fn populated_top_level_blocks(&self) -> Vec<&'static str> {
|
||||
let mut blocks = Vec::new();
|
||||
if self.policy.file.is_some() {
|
||||
blocks.push("policy.file");
|
||||
}
|
||||
if !self.queries.is_empty() {
|
||||
blocks.push("queries");
|
||||
}
|
||||
blocks
|
||||
}
|
||||
|
||||
/// A named graph uses its own `graphs.<name>` block, so a populated
|
||||
/// top-level block would be silently ignored — a config error. The single
|
||||
/// definition of that rule, shared by server boot and the CLI selection
|
||||
/// gate ([`OmnigraphConfig::resolve_graph_selection`]) so the two can't
|
||||
/// drift. An anonymous selection (`None`, e.g. a bare URI) legitimately
|
||||
/// honors the top-level blocks, so it is never rejected here.
|
||||
pub fn ensure_top_level_blocks_honored(&self, selected: Option<&str>) -> Result<()> {
|
||||
if let Some(name) = selected {
|
||||
let unhonored = self.populated_top_level_blocks();
|
||||
if !unhonored.is_empty() {
|
||||
bail!(
|
||||
"named graph '{name}' uses its own `graphs.{name}.…` block, but top-level {} \
|
||||
{} set and would be ignored. Move it to `graphs.{name}` (e.g. \
|
||||
`graphs.{name}.policy.file`, `graphs.{name}.queries`).",
|
||||
unhonored.join(" and "),
|
||||
if unhonored.len() == 1 { "is" } else { "are" },
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Resolve a stored-query `.gq` file path (from a registry entry),
|
||||
/// relative to the config's `base_dir`. Mirrors policy-file
|
||||
/// resolution; the registry loader calls this to turn each entry's
|
||||
/// `file:` value into an absolute path.
|
||||
pub fn resolve_query_file(&self, value: &str) -> PathBuf {
|
||||
self.resolve_config_path(value)
|
||||
}
|
||||
|
||||
/// Resolve the server-level policy file path (used by management
|
||||
/// endpoints). Returns `None` if `server.policy.file` is not set.
|
||||
pub fn resolve_server_policy_file(&self) -> Option<PathBuf> {
|
||||
|
|
@ -387,7 +571,9 @@ mod tests {
|
|||
|
||||
use tempfile::tempdir;
|
||||
|
||||
use super::{ReadOutputFormat, TableCellLayout, load_config_in};
|
||||
use super::{
|
||||
ReadOutputFormat, TableCellLayout, graph_resource_id_for_selection, load_config_in,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn load_config_reads_yaml_defaults_from_current_dir() {
|
||||
|
|
@ -451,6 +637,114 @@ policy: {}
|
|||
assert!(config.graphs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn graph_resource_id_for_selection_uses_name_or_anonymous_uri() {
|
||||
assert_eq!(
|
||||
graph_resource_id_for_selection(Some("local"), "/tmp/graph.omni"),
|
||||
"local"
|
||||
);
|
||||
assert_eq!(
|
||||
graph_resource_id_for_selection(None, "/tmp/graph.omni"),
|
||||
"/tmp/graph.omni"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_graph_selection_validates_membership_and_coherence() {
|
||||
let temp = tempdir().unwrap();
|
||||
fs::write(
|
||||
temp.path().join("omnigraph.yaml"),
|
||||
"graphs:\n local:\n uri: ./demo.omni\n",
|
||||
)
|
||||
.unwrap();
|
||||
let config = load_config_in(temp.path(), None).unwrap();
|
||||
|
||||
// A known graph passes through unchanged.
|
||||
assert_eq!(config.resolve_graph_selection(Some("local")).unwrap(), Some("local"));
|
||||
// An anonymous selection stays anonymous (→ top-level registry downstream).
|
||||
assert_eq!(config.resolve_graph_selection(None).unwrap(), None);
|
||||
// An unknown name errors, naming the graph (matching resolve_target_uri).
|
||||
let err = config.resolve_graph_selection(Some("ghost")).unwrap_err().to_string();
|
||||
assert!(
|
||||
err.contains("ghost") && err.contains("not found"),
|
||||
"unknown graph must error naming it: {err}"
|
||||
);
|
||||
|
||||
// Coherence: a named graph plus a populated top-level block is the
|
||||
// config server boot refuses, so the gate rejects it too (shared rule
|
||||
// via ensure_top_level_blocks_honored). An anonymous selection still
|
||||
// passes — top-level is honored when no graph is named.
|
||||
let temp2 = tempdir().unwrap();
|
||||
fs::write(
|
||||
temp2.path().join("omnigraph.yaml"),
|
||||
"graphs:\n local:\n uri: ./demo.omni\npolicy:\n file: ./top.yaml\n",
|
||||
)
|
||||
.unwrap();
|
||||
let incoherent = load_config_in(temp2.path(), None).unwrap();
|
||||
let err = incoherent
|
||||
.resolve_graph_selection(Some("local"))
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
assert!(
|
||||
err.contains("local") && err.contains("policy.file"),
|
||||
"named graph + populated top-level block must be rejected, naming both: {err}"
|
||||
);
|
||||
assert_eq!(
|
||||
incoherent.resolve_graph_selection(None).unwrap(),
|
||||
None,
|
||||
"anonymous selection still honors top-level"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_tooling_graph_selection_prefers_cli_then_server_and_validates() {
|
||||
let temp = tempdir().unwrap();
|
||||
fs::write(
|
||||
temp.path().join("omnigraph.yaml"),
|
||||
"graphs:\n local:\n uri: ./local.omni\n prod:\n uri: ./prod.omni\n\
|
||||
server:\n graph: local\ncli:\n graph: prod\n",
|
||||
)
|
||||
.unwrap();
|
||||
let config = load_config_in(temp.path(), None).unwrap();
|
||||
assert_eq!(
|
||||
config.resolve_policy_tooling_graph_selection().unwrap(),
|
||||
Some("prod")
|
||||
);
|
||||
|
||||
let temp = tempdir().unwrap();
|
||||
fs::write(
|
||||
temp.path().join("omnigraph.yaml"),
|
||||
"graphs:\n local:\n uri: ./local.omni\nserver:\n graph: local\n",
|
||||
)
|
||||
.unwrap();
|
||||
let config = load_config_in(temp.path(), None).unwrap();
|
||||
assert_eq!(
|
||||
config.resolve_policy_tooling_graph_selection().unwrap(),
|
||||
Some("local")
|
||||
);
|
||||
|
||||
let temp = tempdir().unwrap();
|
||||
fs::write(temp.path().join("omnigraph.yaml"), "policy: {}\n").unwrap();
|
||||
let config = load_config_in(temp.path(), None).unwrap();
|
||||
assert_eq!(config.resolve_policy_tooling_graph_selection().unwrap(), None);
|
||||
|
||||
let temp = tempdir().unwrap();
|
||||
fs::write(
|
||||
temp.path().join("omnigraph.yaml"),
|
||||
"graphs:\n local:\n uri: ./local.omni\nserver:\n graph: ghost\n",
|
||||
)
|
||||
.unwrap();
|
||||
let config = load_config_in(temp.path(), None).unwrap();
|
||||
let err = config
|
||||
.resolve_policy_tooling_graph_selection()
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
assert!(
|
||||
err.contains("ghost") && err.contains("not found"),
|
||||
"unknown server.graph must use graph-selection validation: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_query_path_searches_config_roots() {
|
||||
let temp = tempdir().unwrap();
|
||||
|
|
@ -489,6 +783,118 @@ policy: {}
|
|||
assert_eq!(resolved, config_dir.join("local.gq"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queries_block_round_trips_inline_and_per_graph() {
|
||||
let temp = tempdir().unwrap();
|
||||
fs::write(
|
||||
temp.path().join("omnigraph.yaml"),
|
||||
r#"
|
||||
graphs:
|
||||
prod:
|
||||
uri: s3://bucket/prod
|
||||
queries:
|
||||
find_user:
|
||||
file: ./queries/find_user.gq
|
||||
mcp:
|
||||
expose: true
|
||||
tool_name: lookup_user
|
||||
internal_audit:
|
||||
file: ./queries/audit.gq
|
||||
queries:
|
||||
single_mode_q:
|
||||
file: ./q.gq
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = load_config_in(temp.path(), None).unwrap();
|
||||
|
||||
// Per-graph registry (multi-graph mode).
|
||||
let prod = config.target_query_entries("prod").unwrap();
|
||||
assert_eq!(prod.len(), 2);
|
||||
let find_user = &prod["find_user"];
|
||||
assert_eq!(find_user.file, "./queries/find_user.gq");
|
||||
assert!(find_user.mcp.expose);
|
||||
assert_eq!(find_user.mcp.tool_name.as_deref(), Some("lookup_user"));
|
||||
// Default exposure is true (the manifest entry is the opt-in); tool_name absent.
|
||||
let audit = &prod["internal_audit"];
|
||||
assert!(audit.mcp.expose);
|
||||
assert!(audit.mcp.tool_name.is_none());
|
||||
|
||||
// Top-level registry (single-graph mode).
|
||||
assert_eq!(config.query_entries().len(), 1);
|
||||
|
||||
// The shared selector resolves the same blocks the server boot
|
||||
// and the CLI use: a known graph → its per-graph block; no
|
||||
// selection or an unknown name → the top-level block (the latter
|
||||
// pins the behavior of the CLI's now-deleted fallback arm).
|
||||
assert_eq!(config.query_entries_for(Some("prod")).len(), 2);
|
||||
assert_eq!(config.query_entries_for(None).len(), 1);
|
||||
assert_eq!(config.query_entries_for(Some("nonexistent")).len(), 1);
|
||||
|
||||
// Path resolution joins against base_dir, like policy files.
|
||||
assert_eq!(
|
||||
config.resolve_query_file(&find_user.file),
|
||||
temp.path().join("./queries/find_user.gq")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_policy_file_for_follows_identity() {
|
||||
let temp = tempdir().unwrap();
|
||||
fs::write(
|
||||
temp.path().join("omnigraph.yaml"),
|
||||
"policy:\n file: ./top.yaml\ngraphs:\n prod:\n uri: s3://b/prod\n \
|
||||
policy:\n file: ./prod.yaml\n bare:\n uri: s3://b/bare\n",
|
||||
)
|
||||
.unwrap();
|
||||
let config = load_config_in(temp.path(), None).unwrap();
|
||||
|
||||
// Named graph with its own policy → per-graph (not top-level).
|
||||
assert!(
|
||||
config
|
||||
.resolve_policy_file_for(Some("prod"))
|
||||
.unwrap()
|
||||
.ends_with("prod.yaml")
|
||||
);
|
||||
// Named graph with NO per-graph policy → None (no top-level fallback;
|
||||
// load-bearing for the boot coherence check).
|
||||
assert!(config.resolve_policy_file_for(Some("bare")).is_none());
|
||||
// Anonymous (bare URI) or an unknown name → top-level.
|
||||
assert!(
|
||||
config
|
||||
.resolve_policy_file_for(None)
|
||||
.unwrap()
|
||||
.ends_with("top.yaml")
|
||||
);
|
||||
assert!(
|
||||
config
|
||||
.resolve_policy_file_for(Some("nope"))
|
||||
.unwrap()
|
||||
.ends_with("top.yaml")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queries_block_absent_yields_empty_registry() {
|
||||
let temp = tempdir().unwrap();
|
||||
fs::write(
|
||||
temp.path().join("omnigraph.yaml"),
|
||||
"graphs:\n local:\n uri: ./demo.omni\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = load_config_in(temp.path(), None).unwrap();
|
||||
// Additive: no `queries:` anywhere → empty registries everywhere.
|
||||
assert!(config.query_entries().is_empty());
|
||||
assert!(
|
||||
config
|
||||
.target_query_entries("local")
|
||||
.unwrap()
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_block_accepts_non_empty_mapping() {
|
||||
let temp = tempdir().unwrap();
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ pub mod config;
|
|||
pub mod graph_id;
|
||||
pub mod identity;
|
||||
pub mod policy;
|
||||
pub mod queries;
|
||||
pub mod registry;
|
||||
pub mod workload;
|
||||
|
||||
|
|
@ -11,6 +12,8 @@ pub use graph_id::GraphId;
|
|||
pub use identity::{AuthSource, GraphKey, ResolvedActor, Scope, TenantId};
|
||||
pub use registry::{GraphHandle, GraphRegistry, InsertError, RegistryLookup, RegistrySnapshot};
|
||||
|
||||
use crate::queries::{QueryRegistry, check, format_check_breakages};
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::fs;
|
||||
use std::io;
|
||||
|
|
@ -22,7 +25,8 @@ use api::{
|
|||
BranchCreateOutput, BranchCreateRequest, BranchDeleteOutput, BranchListOutput,
|
||||
BranchMergeOutput, BranchMergeRequest, ChangeOutput, ChangeRequest, CommitListOutput,
|
||||
CommitListQuery, ErrorCode, ErrorOutput, ExportRequest, GraphInfo, GraphListResponse,
|
||||
HealthOutput, IngestOutput, IngestRequest, QueryRequest, ReadOutput, ReadRequest,
|
||||
HealthOutput, IngestOutput, IngestRequest, InvokeStoredQueryRequest,
|
||||
InvokeStoredQueryResponse, QueriesCatalogOutput, QueryRequest, ReadOutput, ReadRequest,
|
||||
SchemaApplyOutput, SchemaApplyRequest, SchemaOutput, SnapshotQuery, ingest_output,
|
||||
schema_apply_output, snapshot_payload,
|
||||
};
|
||||
|
|
@ -40,12 +44,13 @@ use color_eyre::eyre::{Result, WrapErr, bail};
|
|||
pub use config::{
|
||||
AliasCommand, AliasConfig, CliDefaults, DEFAULT_CONFIG_FILE, OmnigraphConfig, PolicySettings,
|
||||
ProjectConfig, QueryDefaults, ReadOutputFormat, ServerDefaults, TableCellLayout, TargetConfig,
|
||||
load_config,
|
||||
graph_resource_id_for_selection, load_config,
|
||||
};
|
||||
use futures::stream;
|
||||
use omnigraph::db::{Omnigraph, ReadTarget};
|
||||
use omnigraph::error::{ManifestConflictDetails, ManifestErrorKind, OmniError};
|
||||
use omnigraph::storage::normalize_root_uri;
|
||||
use omnigraph_compiler::catalog::Catalog;
|
||||
use omnigraph_compiler::json_params_to_param_map;
|
||||
use omnigraph_compiler::query::parser::parse_query;
|
||||
use omnigraph_compiler::{JsonParamMode, ParamMap};
|
||||
|
|
@ -93,6 +98,8 @@ fn hash_bearer_token(token: &str) -> BearerTokenHash {
|
|||
server_export,
|
||||
#[allow(deprecated)] server_change,
|
||||
server_mutate,
|
||||
server_list_queries,
|
||||
server_invoke_query,
|
||||
server_schema_apply,
|
||||
server_schema_get,
|
||||
server_ingest,
|
||||
|
|
@ -157,8 +164,16 @@ pub enum ServerConfigMode {
|
|||
/// set to a named target.
|
||||
Single {
|
||||
uri: String,
|
||||
/// Cedar graph resource id for the single graph. A named selection
|
||||
/// uses the graph name; an anonymous URI uses the normalized URI to
|
||||
/// preserve legacy single-graph policy identity.
|
||||
graph_id: String,
|
||||
/// Top-level `policy.file` (single-graph Cedar policy).
|
||||
policy_file: Option<PathBuf>,
|
||||
/// Top-level stored-query registry, loaded and identity-checked
|
||||
/// at settings-build time; type-checked against the schema when
|
||||
/// the engine opens.
|
||||
queries: QueryRegistry,
|
||||
},
|
||||
/// Multi-graph invocation — `--config omnigraph.yaml` with a
|
||||
/// non-empty `graphs:` map and no single-mode selector.
|
||||
|
|
@ -185,6 +200,10 @@ pub struct GraphStartupConfig {
|
|||
pub graph_id: String,
|
||||
pub uri: String,
|
||||
pub policy_file: Option<PathBuf>,
|
||||
/// Per-graph stored-query registry, loaded and identity-checked at
|
||||
/// settings-build time; type-checked against the schema when this
|
||||
/// graph's engine opens.
|
||||
pub queries: QueryRegistry,
|
||||
}
|
||||
|
||||
/// Runtime routing for the server. Single mode = legacy
|
||||
|
|
@ -285,7 +304,31 @@ impl AppState {
|
|||
) -> Self {
|
||||
let bearer_tokens = hash_bearer_tokens(bearer_tokens);
|
||||
let per_graph_policy = policy_engine.map(Arc::new);
|
||||
Self::build_single_mode(uri, db, bearer_tokens, per_graph_policy, Arc::new(workload))
|
||||
Self::build_single_mode(uri, db, bearer_tokens, per_graph_policy, Arc::new(workload), None)
|
||||
}
|
||||
|
||||
/// Like `new_single`, but attaches a pre-validated stored-query
|
||||
/// registry. Private — the production single-mode boot path
|
||||
/// (`open_single_with_queries`) is the only caller; every public
|
||||
/// `new_*` constructor builds with no stored queries.
|
||||
fn new_single_with_queries(
|
||||
uri: String,
|
||||
db: Omnigraph,
|
||||
bearer_tokens: Vec<(String, String)>,
|
||||
policy_engine: Option<PolicyEngine>,
|
||||
workload: workload::WorkloadController,
|
||||
queries: Option<Arc<QueryRegistry>>,
|
||||
) -> Self {
|
||||
let bearer_tokens = hash_bearer_tokens(bearer_tokens);
|
||||
let per_graph_policy = policy_engine.map(Arc::new);
|
||||
Self::build_single_mode(
|
||||
uri,
|
||||
db,
|
||||
bearer_tokens,
|
||||
per_graph_policy,
|
||||
Arc::new(workload),
|
||||
queries,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn new(uri: String, db: Omnigraph) -> Self {
|
||||
|
|
@ -377,6 +420,39 @@ impl AppState {
|
|||
uri: impl Into<String>,
|
||||
bearer_tokens: Vec<(String, String)>,
|
||||
policy_file: Option<&PathBuf>,
|
||||
) -> Result<Self> {
|
||||
Self::open_single_with_queries(
|
||||
uri,
|
||||
bearer_tokens,
|
||||
policy_file,
|
||||
QueryRegistry::default(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Single-mode boot with a stored-query registry: open the engine,
|
||||
/// **type-check the registry against the live schema and refuse to
|
||||
/// start on a breakage** (same posture as bad policy YAML), log
|
||||
/// non-blocking warnings, then attach the registry to the handle.
|
||||
/// With an empty registry the check is a no-op and no registry is
|
||||
/// attached — that is the path `open_with_bearer_tokens_and_policy`
|
||||
/// (no stored queries) takes.
|
||||
pub async fn open_single_with_queries(
|
||||
uri: impl Into<String>,
|
||||
bearer_tokens: Vec<(String, String)>,
|
||||
policy_file: Option<&PathBuf>,
|
||||
queries: QueryRegistry,
|
||||
) -> Result<Self> {
|
||||
Self::open_single_with_queries_for_graph_id(uri, bearer_tokens, policy_file, queries, None)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn open_single_with_queries_for_graph_id(
|
||||
uri: impl Into<String>,
|
||||
bearer_tokens: Vec<(String, String)>,
|
||||
policy_file: Option<&PathBuf>,
|
||||
queries: QueryRegistry,
|
||||
graph_id: Option<String>,
|
||||
) -> Result<Self> {
|
||||
// The "policy requires tokens" invariant is enforced once by
|
||||
// `classify_server_runtime_state` in `serve()`, before either
|
||||
|
|
@ -384,16 +460,24 @@ impl AppState {
|
|||
// time we get here, the (policy, no-tokens) combination has
|
||||
// already been rejected — no second bail needed.
|
||||
let uri = normalize_root_uri(&uri.into()).wrap_err("normalize graph URI")?;
|
||||
let graph_id = graph_id.unwrap_or_else(|| uri.clone());
|
||||
let db = Omnigraph::open(&uri).await?;
|
||||
|
||||
// Validate the registry against the live schema and resolve it to
|
||||
// an attachable handle (refuse boot on breakage).
|
||||
let registry = validate_and_attach(queries, &db.catalog(), &graph_id)?;
|
||||
|
||||
let policy_engine = match policy_file {
|
||||
Some(path) => Some(PolicyEngine::load_graph(path, &uri)?),
|
||||
Some(path) => Some(PolicyEngine::load_graph(path, &graph_id)?),
|
||||
None => None,
|
||||
};
|
||||
Ok(Self::new_with_bearer_tokens_and_policy(
|
||||
Ok(Self::new_single_with_queries(
|
||||
uri,
|
||||
db,
|
||||
bearer_tokens,
|
||||
policy_engine,
|
||||
workload::WorkloadController::from_env(),
|
||||
registry,
|
||||
))
|
||||
}
|
||||
|
||||
|
|
@ -408,6 +492,7 @@ impl AppState {
|
|||
bearer_tokens: Arc<[(BearerTokenHash, Arc<str>)]>,
|
||||
policy_engine: Option<Arc<PolicyEngine>>,
|
||||
workload: Arc<workload::WorkloadController>,
|
||||
queries: Option<Arc<QueryRegistry>>,
|
||||
) -> Self {
|
||||
// Engine-layer policy gate (MR-722). With a per-graph policy
|
||||
// installed, every `_as` writer on `Omnigraph` calls into the
|
||||
|
|
@ -436,6 +521,7 @@ impl AppState {
|
|||
uri,
|
||||
engine: Arc::new(db),
|
||||
policy: policy_engine,
|
||||
queries,
|
||||
});
|
||||
Self {
|
||||
routing: GraphRouting::Single { handle },
|
||||
|
|
@ -750,6 +836,58 @@ pub fn init_tracing() {
|
|||
let _ = tracing_subscriber::fmt().with_env_filter(filter).try_init();
|
||||
}
|
||||
|
||||
/// Log each non-blocking advisory from a registry check report.
|
||||
fn log_registry_warnings(label: &str, report: &queries::CheckReport) {
|
||||
for warning in &report.warnings {
|
||||
warn!(graph = label, query = %warning.query, "stored query: {}", warning.message);
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_registry_against_catalog(
|
||||
registry: &QueryRegistry,
|
||||
catalog: &Catalog,
|
||||
label: &str,
|
||||
) -> omnigraph::error::Result<()> {
|
||||
let report = check(registry, catalog);
|
||||
if report.has_breakages() {
|
||||
return Err(OmniError::manifest(format_check_breakages(label, &report)));
|
||||
}
|
||||
log_registry_warnings(label, &report);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Validate a loaded stored-query registry against the live schema and
|
||||
/// resolve it to an attachable handle. Refuses boot on any breakage
|
||||
/// (same posture as bad policy YAML), logs the non-blocking warnings,
|
||||
/// and collapses an empty registry to `None` (nothing attached). This is
|
||||
/// the single gate every open path funnels through, so no opener can
|
||||
/// attach a registry that has not been schema-checked. `label` names the
|
||||
/// graph in messages.
|
||||
fn validate_and_attach(
|
||||
queries: QueryRegistry,
|
||||
catalog: &Catalog,
|
||||
label: &str,
|
||||
) -> Result<Option<Arc<QueryRegistry>>> {
|
||||
validate_registry_against_catalog(&queries, catalog, label)
|
||||
.map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?;
|
||||
Ok(if queries.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(Arc::new(queries))
|
||||
})
|
||||
}
|
||||
|
||||
/// Format every load error (parse / identity failure) into a multi-line
|
||||
/// boot-abort message.
|
||||
fn format_registry_load_errors(label: &str, errors: &[queries::LoadError]) -> String {
|
||||
let joined = errors
|
||||
.iter()
|
||||
.map(|e| e.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n ");
|
||||
format!("graph '{label}': stored-query registry failed to load:\n {joined}")
|
||||
}
|
||||
|
||||
pub fn load_server_settings(
|
||||
config_path: Option<&PathBuf>,
|
||||
cli_uri: Option<String>,
|
||||
|
|
@ -799,15 +937,43 @@ pub fn load_server_settings(
|
|||
let uri = normalize_root_uri(&raw_uri).wrap_err_with(|| {
|
||||
format!("normalize single-graph URI '{raw_uri}' from server settings")
|
||||
})?;
|
||||
let policy_file = config.resolve_policy_file();
|
||||
ServerConfigMode::Single { uri, policy_file }
|
||||
// Config follows graph IDENTITY, not mode: a bare URI is anonymous
|
||||
// (top-level config); a graph chosen by name uses its per-graph
|
||||
// `graphs.<name>.{policy,queries}`. `resolve_target_uri` already
|
||||
// errored on an unknown name, so a `Some(name)` here is a known graph.
|
||||
let selected: Option<&str> = if has_cli_uri {
|
||||
None
|
||||
} else {
|
||||
cli_target.as_deref().or_else(|| config.server_graph_name())
|
||||
};
|
||||
// A named selection must not leave a populated top-level block
|
||||
// silently unused — refuse boot and point at the per-graph block. The
|
||||
// same rule the CLI selection gate enforces, shared via one helper so
|
||||
// the boot check and `omnigraph queries validate`/`list` can't drift.
|
||||
config.ensure_top_level_blocks_honored(selected)?;
|
||||
// Load + identity-check now (no engine needed); the schema
|
||||
// type-check happens when the engine opens.
|
||||
let policy_file = config.resolve_policy_file_for(selected);
|
||||
let queries = QueryRegistry::load(&config, config.query_entries_for(selected))
|
||||
.map_err(|errs| color_eyre::eyre::eyre!(format_registry_load_errors(&uri, &errs)))?;
|
||||
let graph_id = graph_resource_id_for_selection(selected, &uri);
|
||||
ServerConfigMode::Single {
|
||||
uri,
|
||||
graph_id,
|
||||
policy_file,
|
||||
queries,
|
||||
}
|
||||
} else if has_explicit_config && has_graphs_map {
|
||||
if config.resolve_policy_file().is_some() {
|
||||
// Multi mode: every graph uses its per-graph block; top-level
|
||||
// policy/queries are never honored, so a populated one is an error.
|
||||
let unhonored = config.populated_top_level_blocks();
|
||||
if !unhonored.is_empty() {
|
||||
bail!(
|
||||
"top-level `policy.file` is single-graph/CLI-local policy only; \
|
||||
in multi-graph mode move per-graph rules to \
|
||||
`graphs.<graph_id>.policy.file` and move `graph_list` rules to \
|
||||
`server.policy.file`."
|
||||
"multi-graph mode: top-level {} {} not honored — each graph uses its own \
|
||||
`graphs.<graph_id>.…` block. Move per-graph rules there (and any \
|
||||
`graph_list` policy to `server.policy.file`).",
|
||||
unhonored.join(" and "),
|
||||
if unhonored.len() == 1 { "is" } else { "are" },
|
||||
);
|
||||
}
|
||||
// Rule 4 → Multi mode. Build a startup config per graph.
|
||||
|
|
@ -823,10 +989,17 @@ pub fn load_server_settings(
|
|||
let uri = normalize_root_uri(&raw_uri).wrap_err_with(|| {
|
||||
format!("normalize URI '{raw_uri}' for graph '{name}' in omnigraph.yaml")
|
||||
})?;
|
||||
// Per-graph `queries:`, selected through the shared
|
||||
// `query_entries_for` so server and CLI resolve identically.
|
||||
// Load + identity-check now; the schema type-check happens
|
||||
// when this graph's engine opens.
|
||||
let queries = QueryRegistry::load(&config, config.query_entries_for(Some(name.as_str())))
|
||||
.map_err(|errs| color_eyre::eyre::eyre!(format_registry_load_errors(name, &errs)))?;
|
||||
graphs.push(GraphStartupConfig {
|
||||
graph_id: name.clone(),
|
||||
uri,
|
||||
policy_file: config.resolve_target_policy_file(name),
|
||||
queries,
|
||||
});
|
||||
}
|
||||
let config_path = config_path
|
||||
|
|
@ -949,6 +1122,8 @@ pub fn build_app(state: AppState) -> Router {
|
|||
server_change
|
||||
}))
|
||||
.route("/mutate", post(server_mutate))
|
||||
.route("/queries", get(server_list_queries))
|
||||
.route("/queries/{name}", post(server_invoke_query))
|
||||
.route("/schema", get(server_schema_get))
|
||||
.route("/schema/apply", post(server_schema_apply))
|
||||
.route(
|
||||
|
|
@ -1046,10 +1221,28 @@ pub async fn serve(config: ServerConfig) -> Result<()> {
|
|||
|
||||
let bind = config.bind.clone();
|
||||
let state = match config.mode {
|
||||
ServerConfigMode::Single { uri, policy_file } => {
|
||||
ServerConfigMode::Single {
|
||||
uri,
|
||||
graph_id,
|
||||
policy_file,
|
||||
queries,
|
||||
} => {
|
||||
let uri_for_log = uri.clone();
|
||||
info!(uri = %uri_for_log, bind = %bind, mode = "single", "serving omnigraph");
|
||||
AppState::open_with_bearer_tokens_and_policy(uri, tokens, policy_file.as_ref()).await?
|
||||
info!(
|
||||
uri = %uri_for_log,
|
||||
graph_id = %graph_id,
|
||||
bind = %bind,
|
||||
mode = "single",
|
||||
"serving omnigraph"
|
||||
);
|
||||
AppState::open_single_with_queries_for_graph_id(
|
||||
uri,
|
||||
tokens,
|
||||
policy_file.as_ref(),
|
||||
queries,
|
||||
Some(graph_id),
|
||||
)
|
||||
.await?
|
||||
}
|
||||
ServerConfigMode::Multi {
|
||||
graphs,
|
||||
|
|
@ -1131,6 +1324,12 @@ async fn open_single_graph(cfg: GraphStartupConfig) -> Result<Arc<GraphHandle>>
|
|||
.await
|
||||
.map_err(|err| color_eyre::eyre::eyre!("open graph '{}' at {}: {err}", graph_id, uri))?;
|
||||
|
||||
// Validate this graph's stored queries against the live schema and
|
||||
// resolve them to an attachable handle (refuse boot on breakage).
|
||||
// Done before the policy match rebinds `db`; the catalog handle is an
|
||||
// owned `Arc`, so no borrow of `db` survives into the match.
|
||||
let queries = validate_and_attach(cfg.queries, &db.catalog(), graph_id.as_str())?;
|
||||
|
||||
let (policy_arc, db) = match &cfg.policy_file {
|
||||
Some(path) => {
|
||||
let policy = PolicyEngine::load_graph(path, graph_id.as_str())?;
|
||||
|
|
@ -1146,6 +1345,7 @@ async fn open_single_graph(cfg: GraphStartupConfig) -> Result<Arc<GraphHandle>>
|
|||
uri,
|
||||
engine: Arc::new(db),
|
||||
policy: policy_arc,
|
||||
queries,
|
||||
}))
|
||||
}
|
||||
|
||||
|
|
@ -1479,7 +1679,21 @@ fn log_policy_decision(actor_id: &str, request: &PolicyRequest, decision: &Polic
|
|||
);
|
||||
}
|
||||
|
||||
/// HTTP-layer Cedar policy gate. Two sources of the policy engine:
|
||||
/// The allow/deny **decision** an authorization check produces, kept
|
||||
/// separate from the operational failures (`Err`) that can occur while
|
||||
/// computing it. [`authorize_request`] collapses `Denied` to a 403; a caller
|
||||
/// that needs to remap a denial without also remapping operational failures
|
||||
/// (the stored-query invoke handler hides a denial as a 404) matches on this
|
||||
/// directly, so a real 401 (missing bearer) or 500 (policy-evaluation error)
|
||||
/// keeps its true status instead of being masked as the denial's response.
|
||||
enum Authz {
|
||||
Allowed,
|
||||
Denied(String),
|
||||
}
|
||||
|
||||
/// HTTP-layer Cedar policy gate, returning the allow/deny [`Authz`] decision
|
||||
/// and reserving `Err` for operational failures (401 missing bearer, 500
|
||||
/// policy-evaluation error). Two sources of the policy engine:
|
||||
/// * Per-graph handler — passes `handle.policy.as_deref()` so the
|
||||
/// graph's Cedar rules govern read/change/branch_*/schema_apply.
|
||||
/// * Management handler — passes `state.server_policy.as_deref()` so
|
||||
|
|
@ -1493,11 +1707,11 @@ fn log_policy_decision(actor_id: &str, request: &PolicyRequest, decision: &Polic
|
|||
/// dropped from the type), so handlers cannot smuggle it through the
|
||||
/// request. See `actor_id_resolves_from_bearer_token_ignoring_client_supplied_headers`
|
||||
/// at `tests/server.rs`.
|
||||
fn authorize_request(
|
||||
fn authorize(
|
||||
actor: Option<&ResolvedActor>,
|
||||
policy: Option<&PolicyEngine>,
|
||||
request: PolicyRequest,
|
||||
) -> std::result::Result<(), ApiError> {
|
||||
) -> std::result::Result<Authz, ApiError> {
|
||||
let Some(engine) = policy else {
|
||||
// No PolicyEngine installed. Three runtime states can reach this:
|
||||
//
|
||||
|
|
@ -1524,21 +1738,23 @@ fn authorize_request(
|
|||
// operator's only path to enabling it is configuring an
|
||||
// explicit `server.policy.file` in omnigraph.yaml.
|
||||
if request.action.resource_kind() == PolicyResourceKind::Server {
|
||||
return Err(ApiError::forbidden(
|
||||
return Ok(Authz::Denied(
|
||||
"server-scoped actions require an explicit `server.policy.file` \
|
||||
configured in omnigraph.yaml — the management surface is closed \
|
||||
by default in every runtime state, including --unauthenticated, \
|
||||
so that server topology is never exposed without operator opt-in.",
|
||||
so that server topology is never exposed without operator opt-in."
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
if actor.is_some() && request.action != PolicyAction::Read {
|
||||
return Err(ApiError::forbidden(
|
||||
return Ok(Authz::Denied(
|
||||
"server runs in default-deny mode (bearer tokens configured but no \
|
||||
policy file). Only `read` actions are permitted; configure \
|
||||
`policy.file` in omnigraph.yaml to enable other actions.",
|
||||
`policy.file` in omnigraph.yaml to enable other actions."
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
return Ok(());
|
||||
return Ok(Authz::Allowed);
|
||||
};
|
||||
let Some(actor) = actor else {
|
||||
return Err(ApiError::unauthorized("missing bearer token"));
|
||||
|
|
@ -1560,9 +1776,26 @@ fn authorize_request(
|
|||
.map_err(|err| ApiError::internal(format!("policy: {err}")))?;
|
||||
log_policy_decision(actor_id, &request, &decision);
|
||||
if decision.allowed {
|
||||
Ok(())
|
||||
Ok(Authz::Allowed)
|
||||
} else {
|
||||
Err(ApiError::forbidden(decision.message))
|
||||
Ok(Authz::Denied(decision.message))
|
||||
}
|
||||
}
|
||||
|
||||
/// Thin wrapper over [`authorize`] for the handlers that treat any denial as a
|
||||
/// 403: a denial becomes `ApiError::forbidden`, and operational failures
|
||||
/// (401 missing bearer, 500 policy-evaluation error) propagate unchanged. The
|
||||
/// stored-query invoke handler does **not** use this — it consumes the
|
||||
/// [`Authz`] decision directly to hide a denial as a 404 while letting an
|
||||
/// operational failure keep its true status.
|
||||
fn authorize_request(
|
||||
actor: Option<&ResolvedActor>,
|
||||
policy: Option<&PolicyEngine>,
|
||||
request: PolicyRequest,
|
||||
) -> std::result::Result<(), ApiError> {
|
||||
match authorize(actor, policy, request)? {
|
||||
Authz::Allowed => Ok(()),
|
||||
Authz::Denied(message) => Err(ApiError::forbidden(message)),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -2001,6 +2234,194 @@ async fn server_mutate(
|
|||
))
|
||||
}
|
||||
|
||||
/// Path parameter for `POST /queries/{name}`.
|
||||
#[derive(Deserialize)]
|
||||
struct QueryNamePath {
|
||||
name: String,
|
||||
}
|
||||
|
||||
fn parse_optional_invoke_body(
|
||||
body: Bytes,
|
||||
) -> std::result::Result<InvokeStoredQueryRequest, ApiError> {
|
||||
if body.is_empty() {
|
||||
return Ok(InvokeStoredQueryRequest::default());
|
||||
}
|
||||
serde_json::from_slice::<Option<InvokeStoredQueryRequest>>(&body)
|
||||
.map(|request| request.unwrap_or_default())
|
||||
.map_err(|err| {
|
||||
ApiError::bad_request(format!("invalid stored-query invocation body: {err}"))
|
||||
})
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
post,
|
||||
path = "/queries/{name}",
|
||||
tag = "queries",
|
||||
operation_id = "invoke_query",
|
||||
params(("name" = String, Path, description = "Stored query name (the registry key)")),
|
||||
request_body = Option<InvokeStoredQueryRequest>,
|
||||
responses(
|
||||
(status = 200, description = "Read envelope (ReadOutput) or mutation envelope (ChangeOutput), serialized untagged", body = InvokeStoredQueryResponse),
|
||||
(status = 400, description = "Bad request (param type error; snapshot on a stored mutation)", body = ErrorOutput),
|
||||
(status = 401, description = "Unauthorized", body = ErrorOutput),
|
||||
(status = 403, description = "Forbidden (the inner `change` gate for a stored mutation)", body = ErrorOutput),
|
||||
(status = 404, description = "Unknown stored query, or `invoke_query` denied — indistinguishable to a caller without the grant", body = ErrorOutput),
|
||||
(status = 409, description = "Merge conflict", body = ErrorOutput),
|
||||
(status = 429, description = "Per-actor admission cap exceeded; honor `Retry-After` header", body = ErrorOutput),
|
||||
(status = 500, description = "Policy evaluation error (a denial is reported as 404, not 500)", body = ErrorOutput),
|
||||
),
|
||||
security(("bearer_token" = [])),
|
||||
)]
|
||||
/// Invoke a curated, server-side stored query by name.
|
||||
///
|
||||
/// The query source comes from the graph's `queries:` registry, not the
|
||||
/// request body — callers send only runtime inputs (`params`, `branch`,
|
||||
/// `snapshot`). Gated by the `invoke_query` Cedar action at the boundary;
|
||||
/// a stored *mutation* additionally passes the engine's `change` gate
|
||||
/// (double-gated). An actor **without** `invoke_query` cannot tell a denied
|
||||
/// query from a missing one — both return the same 404, so the catalog
|
||||
/// can't be probed without the grant. Once `invoke_query` is held, the
|
||||
/// inner `read`/`change` gate may surface a 403 for an existing query the
|
||||
/// actor can't run (the intended double-gate signal).
|
||||
async fn server_invoke_query(
|
||||
State(state): State<AppState>,
|
||||
Extension(handle): Extension<Arc<GraphHandle>>,
|
||||
actor: Option<Extension<ResolvedActor>>,
|
||||
Path(QueryNamePath { name }): Path<QueryNamePath>,
|
||||
body: Bytes,
|
||||
) -> std::result::Result<Json<InvokeStoredQueryResponse>, ApiError> {
|
||||
let req = parse_optional_invoke_body(body)?;
|
||||
// A caller without `invoke_query` can't tell a denial from a missing
|
||||
// query: both 404 with this exact message, so the catalog can't be
|
||||
// probed without the grant. (A caller that holds invoke_query may still
|
||||
// see the inner gate's 403 for an existing query it can't run — intended.)
|
||||
const NOT_FOUND: &str = "stored query not found";
|
||||
let actor_ref = actor.as_ref().map(|Extension(actor)| actor);
|
||||
|
||||
// Boundary gate (authentication already ran in `require_bearer_auth`).
|
||||
// A denial is hidden as 404 (deny == missing, so the catalog can't be
|
||||
// probed without the grant), but operational failures (401 missing bearer,
|
||||
// 500 policy-evaluation error) propagate with their true status via `?`
|
||||
// rather than being masked as a missing query.
|
||||
match authorize(
|
||||
actor_ref,
|
||||
handle.policy.as_deref(),
|
||||
PolicyRequest {
|
||||
action: PolicyAction::InvokeQuery,
|
||||
// Graph-scoped: no branch dimension. The per-branch/snapshot
|
||||
// access is enforced by the inner read/change gate in the
|
||||
// runner, so the outer gate must not resolve a branch (doing so
|
||||
// was wrong for snapshot reads).
|
||||
branch: None,
|
||||
target_branch: None,
|
||||
},
|
||||
)? {
|
||||
Authz::Allowed => {}
|
||||
Authz::Denied(_) => return Err(ApiError::not_found(NOT_FOUND)),
|
||||
}
|
||||
|
||||
// Resolve against the per-graph registry (same 404 on a miss).
|
||||
let stored = handle
|
||||
.queries
|
||||
.as_ref()
|
||||
.and_then(|registry| registry.lookup(&name))
|
||||
.ok_or_else(|| ApiError::not_found(NOT_FOUND))?;
|
||||
|
||||
// Detach what we need before `handle` moves into the runner — the
|
||||
// registry borrow lives inside `handle`.
|
||||
let source = Arc::clone(&stored.source);
|
||||
let query_name = stored.name.clone();
|
||||
let is_mutation = stored.is_mutation();
|
||||
|
||||
info!(
|
||||
graph = %handle.uri,
|
||||
actor = ?actor_ref.map(|a| a.actor_id.as_ref()),
|
||||
query = %query_name,
|
||||
kind = if is_mutation { "mutate" } else { "read" },
|
||||
"stored query invoked"
|
||||
);
|
||||
|
||||
if is_mutation {
|
||||
if req.snapshot.is_some() {
|
||||
return Err(ApiError::bad_request(
|
||||
"stored mutation cannot target a snapshot",
|
||||
));
|
||||
}
|
||||
let branch = req.branch.unwrap_or_else(|| "main".to_string());
|
||||
let output = run_mutate(
|
||||
state,
|
||||
handle,
|
||||
actor_ref,
|
||||
&source,
|
||||
Some(&query_name),
|
||||
req.params.as_ref(),
|
||||
branch,
|
||||
)
|
||||
.await?;
|
||||
Ok(Json(InvokeStoredQueryResponse::Change(output)))
|
||||
} else {
|
||||
let (selected, target, result) = run_query(
|
||||
handle,
|
||||
actor_ref,
|
||||
&source,
|
||||
Some(&query_name),
|
||||
req.params.as_ref(),
|
||||
req.branch,
|
||||
req.snapshot,
|
||||
true,
|
||||
)
|
||||
.await?;
|
||||
Ok(Json(InvokeStoredQueryResponse::Read(api::read_output(
|
||||
selected, &target, result,
|
||||
))))
|
||||
}
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
get,
|
||||
path = "/queries",
|
||||
tag = "queries",
|
||||
operation_id = "list_queries",
|
||||
responses(
|
||||
(status = 200, description = "Stored-query catalog (the mcp.expose subset, with typed params)", body = QueriesCatalogOutput),
|
||||
(status = 401, description = "Unauthorized", body = ErrorOutput),
|
||||
(status = 403, description = "Forbidden", body = ErrorOutput),
|
||||
),
|
||||
security(("bearer_token" = [])),
|
||||
)]
|
||||
/// List the graph's exposed stored queries as a typed tool catalog.
|
||||
///
|
||||
/// Returns the `mcp.expose == true` subset of the `queries:` registry, each
|
||||
/// with its MCP tool name, read/mutate flag, description/instruction, and
|
||||
/// typed parameters — enough for a client to register them as tools without
|
||||
/// fetching `.gq` source. Read-gated; the catalog is graph-wide (branch
|
||||
/// independent — `read` is authorized against `main`). **Not** Cedar-filtered
|
||||
/// per query yet, so it can list a query whose `invoke_query` the caller
|
||||
/// lacks (a known gap until per-query authorization lands).
|
||||
async fn server_list_queries(
|
||||
Extension(handle): Extension<Arc<GraphHandle>>,
|
||||
actor: Option<Extension<ResolvedActor>>,
|
||||
) -> std::result::Result<Json<QueriesCatalogOutput>, ApiError> {
|
||||
authorize_request(
|
||||
actor.as_ref().map(|Extension(actor)| actor),
|
||||
handle.policy.as_deref(),
|
||||
PolicyRequest {
|
||||
action: PolicyAction::Read,
|
||||
branch: Some("main".to_string()),
|
||||
target_branch: None,
|
||||
},
|
||||
)?;
|
||||
let queries = match handle.queries.as_ref() {
|
||||
Some(registry) => registry
|
||||
.iter()
|
||||
.filter(|q| q.expose)
|
||||
.map(api::query_catalog_entry)
|
||||
.collect(),
|
||||
None => Vec::new(),
|
||||
};
|
||||
Ok(Json(QueriesCatalogOutput { queries }))
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
get,
|
||||
path = "/schema",
|
||||
|
|
@ -2088,18 +2509,26 @@ async fn server_schema_apply(
|
|||
.map_err(ApiError::from_workload_reject)?;
|
||||
let result = {
|
||||
let db = &handle.engine;
|
||||
let registry = handle.queries.as_deref();
|
||||
let label = handle.key.graph_id.as_str().to_string();
|
||||
// Engine-layer policy enforcement (MR-722): pass the resolved
|
||||
// actor through so apply_schema_as can call enforce() with the
|
||||
// authoritative identity. With a policy installed in AppState,
|
||||
// engine-side enforcement re-checks the same decision the
|
||||
// HTTP-layer authorize_request just made above. PR #3 collapses
|
||||
// the redundancy.
|
||||
db.apply_schema_as(
|
||||
db.apply_schema_as_with_catalog_check(
|
||||
&request.schema_source,
|
||||
omnigraph::db::SchemaApplyOptions {
|
||||
allow_data_loss: request.allow_data_loss,
|
||||
},
|
||||
actor_id,
|
||||
|catalog| {
|
||||
if let Some(registry) = registry {
|
||||
validate_registry_against_catalog(registry, catalog, &label)?;
|
||||
}
|
||||
Ok(())
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(ApiError::from_omni)?
|
||||
|
|
@ -2658,12 +3087,133 @@ mod tests {
|
|||
use std::fs;
|
||||
use tempfile::tempdir;
|
||||
|
||||
/// `authorize` returns the allow/deny **decision** (`Authz`) and reserves
|
||||
/// `Err` for operational failures, so the invoke handler can hide a denial
|
||||
/// as 404 without also masking a 401/500. Pins each outcome.
|
||||
#[test]
|
||||
fn authorize_splits_decision_from_operational_error() {
|
||||
use super::{Authz, PolicyAction, PolicyCompiler, PolicyConfig, PolicyRequest, ResolvedActor, authorize};
|
||||
use std::sync::Arc;
|
||||
|
||||
fn req(action: PolicyAction) -> PolicyRequest {
|
||||
PolicyRequest { action, branch: None, target_branch: None }
|
||||
}
|
||||
let actor = ResolvedActor::cluster_static(Arc::from("act-alice"));
|
||||
|
||||
// --- No policy engine installed (open / default-deny modes) ---
|
||||
// A server-scoped action is denied in every no-policy state.
|
||||
assert!(matches!(
|
||||
authorize(Some(&actor), None, req(PolicyAction::GraphList)).unwrap(),
|
||||
Authz::Denied(_)
|
||||
));
|
||||
// Authenticated actor + a non-read per-graph action → default-deny.
|
||||
assert!(matches!(
|
||||
authorize(Some(&actor), None, req(PolicyAction::Change)).unwrap(),
|
||||
Authz::Denied(_)
|
||||
));
|
||||
// `read` is the one per-graph action permitted without a policy.
|
||||
assert!(matches!(
|
||||
authorize(Some(&actor), None, req(PolicyAction::Read)).unwrap(),
|
||||
Authz::Allowed
|
||||
));
|
||||
// Open mode (no actor, no policy) → allowed.
|
||||
assert!(matches!(
|
||||
authorize(None, None, req(PolicyAction::Read)).unwrap(),
|
||||
Authz::Allowed
|
||||
));
|
||||
|
||||
// --- Policy engine installed ---
|
||||
let policy: PolicyConfig = serde_yaml::from_str(
|
||||
"version: 1\n\
|
||||
groups:\n team: [act-alice]\n\
|
||||
rules:\n - id: team-read\n allow:\n actors: { group: team }\n actions: [read]\n branch_scope: any\n",
|
||||
)
|
||||
.unwrap();
|
||||
let engine = PolicyCompiler::compile(&policy, "graph").unwrap();
|
||||
|
||||
// A matched allow rule → Allowed.
|
||||
assert!(matches!(
|
||||
authorize(
|
||||
Some(&actor),
|
||||
Some(&engine),
|
||||
PolicyRequest { action: PolicyAction::Read, branch: Some("main".to_string()), target_branch: None },
|
||||
)
|
||||
.unwrap(),
|
||||
Authz::Allowed
|
||||
));
|
||||
// Known actor, no matching allow rule → Denied, carrying the decision message.
|
||||
match authorize(
|
||||
Some(&actor),
|
||||
Some(&engine),
|
||||
PolicyRequest { action: PolicyAction::Change, branch: Some("main".to_string()), target_branch: None },
|
||||
)
|
||||
.unwrap()
|
||||
{
|
||||
Authz::Denied(message) => assert!(!message.is_empty(), "a deny carries its decision message"),
|
||||
Authz::Allowed => panic!("change must be denied: only read is allowed"),
|
||||
}
|
||||
// Policy installed but no actor → operational failure (`Err`), NOT a
|
||||
// decision. This is the split that keeps a 401/500 from being masked
|
||||
// as the denial's response in the invoke handler.
|
||||
assert!(
|
||||
authorize(None, Some(&engine), req(PolicyAction::Read)).is_err(),
|
||||
"a missing actor with a policy installed is an operational error, not a deny"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hash_bearer_token_produces_32_byte_output() {
|
||||
let hash = hash_bearer_token("any-token");
|
||||
assert_eq!(hash.len(), 32);
|
||||
}
|
||||
|
||||
/// The single gate both open paths funnel through: it refuses a
|
||||
/// schema breakage (naming the graph label + query), attaches a clean
|
||||
/// registry, and collapses an empty one to `None`. Pure over its args
|
||||
/// (no engine), so it covers the multi-graph path's logic too — the
|
||||
/// only per-path difference is the `label`, asserted here.
|
||||
#[test]
|
||||
fn validate_and_attach_gates_on_schema_and_collapses_empty() {
|
||||
use crate::queries::{QueryRegistry, RegistrySpec};
|
||||
use omnigraph_compiler::catalog::build_catalog;
|
||||
use omnigraph_compiler::schema::parser::parse_schema;
|
||||
|
||||
let schema = parse_schema("node User {\nname: String\n}\n").unwrap();
|
||||
let catalog = build_catalog(&schema).unwrap();
|
||||
let spec = |name: &str, source: &str| RegistrySpec {
|
||||
name: name.to_string(),
|
||||
source: source.to_string(),
|
||||
expose: false,
|
||||
tool_name: None,
|
||||
};
|
||||
|
||||
// Empty registry → nothing attached, no error.
|
||||
let empty =
|
||||
super::validate_and_attach(QueryRegistry::default(), &catalog, "g").unwrap();
|
||||
assert!(empty.is_none());
|
||||
|
||||
// A query that type-checks → attached.
|
||||
let ok = QueryRegistry::from_specs(vec![spec(
|
||||
"find_user",
|
||||
"query find_user() { match { $u: User } return { $u.name } }",
|
||||
)])
|
||||
.unwrap();
|
||||
assert!(super::validate_and_attach(ok, &catalog, "g").unwrap().is_some());
|
||||
|
||||
// A query referencing a type the schema lacks → boot refusal that
|
||||
// names both the graph label and the offending query.
|
||||
let broken = QueryRegistry::from_specs(vec![spec(
|
||||
"ghost",
|
||||
"query ghost() { match { $w: Widget } return { $w.name } }",
|
||||
)])
|
||||
.unwrap();
|
||||
let err = super::validate_and_attach(broken, &catalog, "graph-x").unwrap_err();
|
||||
let msg = err.to_string();
|
||||
assert!(msg.contains("graph-x"), "labels the graph: {msg}");
|
||||
assert!(msg.contains("ghost"), "names the query: {msg}");
|
||||
assert!(msg.contains("schema check"), "mentions the schema check: {msg}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hash_bearer_token_is_deterministic() {
|
||||
assert_eq!(
|
||||
|
|
@ -2707,7 +3257,10 @@ server:
|
|||
|
||||
let settings = load_server_settings(Some(&config), None, None, None, false).unwrap();
|
||||
match &settings.mode {
|
||||
ServerConfigMode::Single { uri, .. } => assert_eq!(uri, "/tmp/demo.omni"),
|
||||
ServerConfigMode::Single { uri, graph_id, .. } => {
|
||||
assert_eq!(uri, "/tmp/demo.omni");
|
||||
assert_eq!(graph_id, "local");
|
||||
}
|
||||
ServerConfigMode::Multi { .. } => panic!("expected Single mode, got Multi"),
|
||||
}
|
||||
assert_eq!(settings.bind, "0.0.0.0:9090");
|
||||
|
|
@ -2739,7 +3292,10 @@ server:
|
|||
)
|
||||
.unwrap();
|
||||
match &settings.mode {
|
||||
ServerConfigMode::Single { uri, .. } => assert_eq!(uri, "/tmp/override.omni"),
|
||||
ServerConfigMode::Single { uri, graph_id, .. } => {
|
||||
assert_eq!(uri, "/tmp/override.omni");
|
||||
assert_eq!(graph_id, "/tmp/override.omni");
|
||||
}
|
||||
ServerConfigMode::Multi { .. } => panic!("expected Single mode, got Multi"),
|
||||
}
|
||||
assert_eq!(settings.bind, "0.0.0.0:9999");
|
||||
|
|
@ -2768,7 +3324,10 @@ server:
|
|||
load_server_settings(Some(&config), None, Some("dev".to_string()), None, false)
|
||||
.unwrap();
|
||||
match &settings.mode {
|
||||
ServerConfigMode::Single { uri, .. } => assert_eq!(uri, "http://127.0.0.1:8080"),
|
||||
ServerConfigMode::Single { uri, graph_id, .. } => {
|
||||
assert_eq!(uri, "http://127.0.0.1:8080");
|
||||
assert_eq!(graph_id, "dev");
|
||||
}
|
||||
ServerConfigMode::Multi { .. } => panic!("expected Single mode, got Multi"),
|
||||
}
|
||||
}
|
||||
|
|
@ -2848,6 +3407,7 @@ server:
|
|||
.to_string_lossy()
|
||||
.into_owned(),
|
||||
policy_file: None,
|
||||
queries: crate::queries::QueryRegistry::default(),
|
||||
}],
|
||||
config_path: temp.path().join("omnigraph.yaml"),
|
||||
server_policy_file: Some(policy_path),
|
||||
|
|
@ -2895,7 +3455,9 @@ server:
|
|||
.join("graph.omni")
|
||||
.to_string_lossy()
|
||||
.into_owned(),
|
||||
graph_id: "default".to_string(),
|
||||
policy_file: None,
|
||||
queries: crate::queries::QueryRegistry::default(),
|
||||
},
|
||||
bind: "127.0.0.1:0".to_string(),
|
||||
allow_unauthenticated: false,
|
||||
|
|
|
|||
688
crates/omnigraph-server/src/queries.rs
Normal file
688
crates/omnigraph-server/src/queries.rs
Normal file
|
|
@ -0,0 +1,688 @@
|
|||
//! Stored-query registry.
|
||||
//!
|
||||
//! A server-side registry of named, parameter-typed `.gq` queries that
|
||||
//! operators declare in `omnigraph.yaml` (per-graph, or top-level in
|
||||
//! single mode) and the server loads at startup. Each entry is parsed
|
||||
//! and its identity asserted here (`load`); type-checking against the
|
||||
//! live schema happens separately (a `check` pass) so the loader stays
|
||||
//! callable without an open engine (the CLI's offline `queries check`).
|
||||
//!
|
||||
//! Identity is the query **name**: the manifest key must equal the
|
||||
//! `query <name>` symbol declared in the referenced `.gq` file. The two
|
||||
//! are asserted equal at load — one name, two places that must agree.
|
||||
//! Renaming either is a breaking change to callers, by design.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::fs;
|
||||
use std::sync::Arc;
|
||||
|
||||
use omnigraph_compiler::catalog::Catalog;
|
||||
use omnigraph_compiler::query::ast::QueryDecl;
|
||||
use omnigraph_compiler::query::parser::parse_query;
|
||||
use omnigraph_compiler::query::typecheck::typecheck_query_decl;
|
||||
use omnigraph_compiler::types::{PropType, ScalarType};
|
||||
|
||||
use crate::config::{OmnigraphConfig, QueryEntry};
|
||||
|
||||
/// One loaded stored query. `source` is the full `.gq` file text — the
|
||||
/// invocation handler hands it to `run_query` / `run_mutate` verbatim,
|
||||
/// which reuse the same parse/IR/exec path as the inline routes (no
|
||||
/// parallel implementation).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StoredQuery {
|
||||
/// Identity: manifest key == `query <name>` symbol.
|
||||
pub name: String,
|
||||
/// Full `.gq` source text the query was selected from.
|
||||
pub source: Arc<str>,
|
||||
/// Parsed declaration (params, mutations, description, …).
|
||||
pub decl: QueryDecl,
|
||||
/// Whether this query is listed in the MCP tool catalog (`GET /queries`).
|
||||
/// Default `true` (the manifest entry is the opt-in); `expose: false`
|
||||
/// keeps it HTTP/service-callable but hidden from the agent tool list.
|
||||
/// Catalog membership only — not an authorization gate.
|
||||
pub expose: bool,
|
||||
/// Optional MCP tool-name override; defaults to `name`.
|
||||
pub tool_name: Option<String>,
|
||||
}
|
||||
|
||||
impl StoredQuery {
|
||||
/// `true` if the selected declaration contains insert/update/delete
|
||||
/// statements — drives read-vs-mutate routing at invocation time.
|
||||
pub fn is_mutation(&self) -> bool {
|
||||
!self.decl.mutations.is_empty()
|
||||
}
|
||||
|
||||
/// The MCP tool name this query is catalogued under: the explicit
|
||||
/// `tool_name` override, else the query `name`. The catalog key —
|
||||
/// enforced unique across exposed queries at load. Server-side
|
||||
/// consumers (the uniqueness check, the future catalog projection) read
|
||||
/// this; the CLI `queries list` resolves the same rule on its own DTO.
|
||||
pub fn effective_tool_name(&self) -> &str {
|
||||
self.tool_name.as_deref().unwrap_or(&self.name)
|
||||
}
|
||||
}
|
||||
|
||||
/// A loaded, identity-checked stored-query registry for one graph.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct QueryRegistry {
|
||||
by_name: BTreeMap<String, StoredQuery>,
|
||||
}
|
||||
|
||||
/// In-memory registry entry before file I/O. Used by [`QueryRegistry::load`]
|
||||
/// (after reading each `.gq` from disk) and directly by tests.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RegistrySpec {
|
||||
pub name: String,
|
||||
pub source: String,
|
||||
pub expose: bool,
|
||||
pub tool_name: Option<String>,
|
||||
}
|
||||
|
||||
/// A single registry load failure. Collected (not fail-fast) so a bad
|
||||
/// `omnigraph.yaml` surfaces every broken entry at once, matching the
|
||||
/// bad-policy-YAML posture.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LoadError {
|
||||
/// The offending query name, when the failure is entry-scoped.
|
||||
pub query: Option<String>,
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for LoadError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match &self.query {
|
||||
Some(name) => write!(f, "stored query '{name}': {}", self.message),
|
||||
None => write!(f, "stored query registry: {}", self.message),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl QueryRegistry {
|
||||
/// Build a registry from in-memory specs: parse each source, select
|
||||
/// the declaration whose symbol equals the manifest key, and assert
|
||||
/// they agree. Collects every failure. No schema type-checking here
|
||||
/// — that is [`check`].
|
||||
pub fn from_specs(specs: Vec<RegistrySpec>) -> Result<Self, Vec<LoadError>> {
|
||||
let mut by_name = BTreeMap::new();
|
||||
let mut errors = Vec::new();
|
||||
|
||||
for spec in specs {
|
||||
match parse_query(&spec.source) {
|
||||
Ok(file) => {
|
||||
match file.queries.into_iter().find(|q| q.name == spec.name) {
|
||||
Some(decl) => {
|
||||
by_name.insert(
|
||||
spec.name.clone(),
|
||||
StoredQuery {
|
||||
name: spec.name,
|
||||
source: Arc::from(spec.source),
|
||||
decl,
|
||||
expose: spec.expose,
|
||||
tool_name: spec.tool_name,
|
||||
},
|
||||
);
|
||||
}
|
||||
None => errors.push(LoadError {
|
||||
query: Some(spec.name.clone()),
|
||||
message: format!(
|
||||
"no `query {}` declaration found in its `.gq` file \
|
||||
(the registry key must match the query symbol)",
|
||||
spec.name
|
||||
),
|
||||
}),
|
||||
}
|
||||
}
|
||||
Err(err) => errors.push(LoadError {
|
||||
query: Some(spec.name),
|
||||
message: format!("parse error: {err}"),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// Exposed queries are catalogued under their effective tool name;
|
||||
// two claiming one name is an MCP-namespace collision. Refuse it at
|
||||
// load (collected, not fail-fast), naming the loser and the winner.
|
||||
// Iterating the `BTreeMap` makes the winner deterministic (the
|
||||
// lexicographically-first query name; config is a map, so YAML
|
||||
// declaration order isn't preserved anyway) and the error order
|
||||
// stable. Scoped to a block so these borrows of `by_name` end
|
||||
// before it is moved into `Self`.
|
||||
{
|
||||
let mut claimed: BTreeMap<&str, &str> = BTreeMap::new();
|
||||
for query in by_name.values().filter(|q| q.expose) {
|
||||
let tool = query.effective_tool_name();
|
||||
if let Some(winner) = claimed.insert(tool, &query.name) {
|
||||
errors.push(LoadError {
|
||||
query: Some(query.name.clone()),
|
||||
message: format!(
|
||||
"MCP tool name '{tool}' already claimed by exposed query '{winner}'"
|
||||
),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if errors.is_empty() {
|
||||
Ok(Self { by_name })
|
||||
} else {
|
||||
Err(errors)
|
||||
}
|
||||
}
|
||||
|
||||
/// Read each registry entry's `.gq` file from disk and build the
|
||||
/// registry. `entries` is either the top-level `queries` map (single
|
||||
/// mode) or a graph's `queries` map (multi mode); `config` resolves
|
||||
/// each entry's relative `file:` path against `base_dir`.
|
||||
pub fn load(
|
||||
config: &OmnigraphConfig,
|
||||
entries: &BTreeMap<String, QueryEntry>,
|
||||
) -> Result<Self, Vec<LoadError>> {
|
||||
let mut specs = Vec::with_capacity(entries.len());
|
||||
let mut errors = Vec::new();
|
||||
for (name, entry) in entries {
|
||||
let path = config.resolve_query_file(&entry.file);
|
||||
match fs::read_to_string(&path) {
|
||||
Ok(source) => specs.push(RegistrySpec {
|
||||
name: name.clone(),
|
||||
source,
|
||||
expose: entry.mcp.expose,
|
||||
tool_name: entry.mcp.tool_name.clone(),
|
||||
}),
|
||||
Err(err) => errors.push(LoadError {
|
||||
query: Some(name.clone()),
|
||||
message: format!("cannot read '{}': {err}", path.display()),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// Parse/identity/uniqueness-check the readable specs even when some
|
||||
// files failed to read, so every broken entry (I/O, parse, identity,
|
||||
// tool-name collision) surfaces in one pass rather than one per
|
||||
// restart. I/O errors come first (in `entries` key order), then the
|
||||
// spec errors. A non-empty `errors` always fails the load.
|
||||
match Self::from_specs(specs) {
|
||||
Ok(registry) if errors.is_empty() => Ok(registry),
|
||||
Ok(_) => Err(errors),
|
||||
Err(spec_errors) => {
|
||||
errors.extend(spec_errors);
|
||||
Err(errors)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn lookup(&self, name: &str) -> Option<&StoredQuery> {
|
||||
self.by_name.get(name)
|
||||
}
|
||||
|
||||
pub fn iter(&self) -> impl Iterator<Item = &StoredQuery> {
|
||||
self.by_name.values()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.by_name.is_empty()
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.by_name.len()
|
||||
}
|
||||
}
|
||||
|
||||
/// A stored query that fails to type-check against the live schema —
|
||||
/// e.g. it references a node/edge type or property that was renamed or
|
||||
/// removed by a migration. Breakages **block server boot** (same posture
|
||||
/// as bad policy YAML), surfacing schema drift at the deploy boundary
|
||||
/// rather than silently at invocation time.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Breakage {
|
||||
pub query: String,
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
/// A non-blocking advisory found during validation. Logged at boot;
|
||||
/// never blocks startup. Currently: an MCP-exposed query that declares a
|
||||
/// parameter an agent cannot realistically supply.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Warning {
|
||||
pub query: String,
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
/// Outcome of validating a registry against a schema. Breakages are
|
||||
/// fatal (boot refuses); warnings are advisory.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct CheckReport {
|
||||
pub breakages: Vec<Breakage>,
|
||||
pub warnings: Vec<Warning>,
|
||||
}
|
||||
|
||||
impl CheckReport {
|
||||
pub fn has_breakages(&self) -> bool {
|
||||
!self.breakages.is_empty()
|
||||
}
|
||||
|
||||
pub fn is_clean(&self) -> bool {
|
||||
self.breakages.is_empty() && self.warnings.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate a loaded registry against the live schema.
|
||||
///
|
||||
/// Pure over `(registry, catalog)` — takes an already-parsed registry and
|
||||
/// a catalog, so it is callable both at server boot (with the engine's
|
||||
/// `catalog()`) and offline from the CLI (`omnigraph queries check`),
|
||||
/// without coupling to server config or an open engine connection.
|
||||
///
|
||||
/// Every query is type-checked via the same `typecheck_query_decl` the
|
||||
/// engine runs for inline queries — no parallel implementation. Failures
|
||||
/// are **collected, not fail-fast**, so an operator sees every broken
|
||||
/// query in one pass.
|
||||
///
|
||||
/// Advisory lint (warn, never block): an `mcp.expose: true` query that
|
||||
/// declares a `Vector(N)` parameter. An LLM cannot supply a raw embedding
|
||||
/// vector; such a query should take a `String` parameter and let the
|
||||
/// engine embed it server-side at query time. Service-to-service callers
|
||||
/// may legitimately pass vectors, so this warns rather than rejects.
|
||||
pub fn check(registry: &QueryRegistry, catalog: &Catalog) -> CheckReport {
|
||||
let mut report = CheckReport::default();
|
||||
for query in registry.iter() {
|
||||
if let Err(err) = typecheck_query_decl(catalog, &query.decl) {
|
||||
report.breakages.push(Breakage {
|
||||
query: query.name.clone(),
|
||||
message: err.to_string(),
|
||||
});
|
||||
}
|
||||
if query.expose {
|
||||
for param in &query.decl.params {
|
||||
// Resolve to the structured type via the compiler's own
|
||||
// resolver rather than string-matching `Vector(` — one
|
||||
// canonical definition of "is a vector", so this lint can't
|
||||
// drift from how the parser/type system spells the type.
|
||||
let is_vector = PropType::from_param_type_name(¶m.type_name, param.nullable)
|
||||
.is_some_and(|pt| matches!(pt.scalar, ScalarType::Vector(_)));
|
||||
if is_vector {
|
||||
report.warnings.push(Warning {
|
||||
query: query.name.clone(),
|
||||
message: format!(
|
||||
"MCP-exposed query declares a `{}` parameter `${}` that agents \
|
||||
cannot supply; use a `String` parameter for server-side embedding",
|
||||
param.type_name, param.name
|
||||
),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
report
|
||||
}
|
||||
|
||||
/// Format every breakage in a registry check report into a multi-line
|
||||
/// operator-facing message, naming each offending query.
|
||||
pub fn format_check_breakages(label: &str, report: &CheckReport) -> String {
|
||||
let joined = report
|
||||
.breakages
|
||||
.iter()
|
||||
.map(|b| format!("query '{}': {}", b.query, b.message))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n ");
|
||||
format!(
|
||||
"graph '{label}': {} stored quer{} failed the schema check:\n {joined}",
|
||||
report.breakages.len(),
|
||||
if report.breakages.len() == 1 {
|
||||
"y"
|
||||
} else {
|
||||
"ies"
|
||||
}
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn spec(name: &str, source: &str, expose: bool) -> RegistrySpec {
|
||||
RegistrySpec {
|
||||
name: name.to_string(),
|
||||
source: source.to_string(),
|
||||
expose,
|
||||
tool_name: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn spec_tool(name: &str, source: &str, expose: bool, tool_name: &str) -> RegistrySpec {
|
||||
RegistrySpec {
|
||||
name: name.to_string(),
|
||||
source: source.to_string(),
|
||||
expose,
|
||||
tool_name: Some(tool_name.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn key_equal_symbol_loads() {
|
||||
let reg = QueryRegistry::from_specs(vec![spec(
|
||||
"find_user",
|
||||
"query find_user($id: String) { match { $u: User } return { $u.name } }",
|
||||
true,
|
||||
)])
|
||||
.unwrap();
|
||||
let q = reg.lookup("find_user").unwrap();
|
||||
assert_eq!(q.name, "find_user");
|
||||
assert!(q.expose);
|
||||
assert_eq!(q.decl.params.len(), 1);
|
||||
assert!(!q.is_mutation());
|
||||
// No override → the effective tool name is the query name.
|
||||
assert_eq!(q.effective_tool_name(), "find_user");
|
||||
|
||||
// An explicit override is what the catalog keys on.
|
||||
let with_tool = QueryRegistry::from_specs(vec![spec_tool(
|
||||
"find_user",
|
||||
"query find_user($id: String) { match { $u: User } return { $u.name } }",
|
||||
true,
|
||||
"lookup_user",
|
||||
)])
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
with_tool.lookup("find_user").unwrap().effective_tool_name(),
|
||||
"lookup_user"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn key_mismatch_is_an_identity_error() {
|
||||
let errors = QueryRegistry::from_specs(vec![spec(
|
||||
"find_user",
|
||||
// symbol is `lookup`, key is `find_user` — must be rejected.
|
||||
"query lookup($id: String) { match { $u: User } return { $u.name } }",
|
||||
false,
|
||||
)])
|
||||
.unwrap_err();
|
||||
assert_eq!(errors.len(), 1);
|
||||
assert_eq!(errors[0].query.as_deref(), Some("find_user"));
|
||||
assert!(errors[0].message.contains("must match the query symbol"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multi_query_file_selects_the_matching_symbol() {
|
||||
let source = "query a($x: I64) { match { $u: User } return { $u.name } }\n\
|
||||
query b($y: String) { match { $u: User } return { $u.name } }";
|
||||
let reg = QueryRegistry::from_specs(vec![spec("b", source, false)]).unwrap();
|
||||
let q = reg.lookup("b").unwrap();
|
||||
assert_eq!(q.name, "b");
|
||||
assert_eq!(q.decl.params[0].name, "y");
|
||||
assert!(reg.lookup("a").is_none(), "only the selected symbol is registered");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn duplicate_exposed_tool_name_is_a_load_error() {
|
||||
// Two MCP-exposed queries claiming one tool name is an ambiguity in
|
||||
// the catalog key space — refused at load, naming both queries and
|
||||
// the contested tool.
|
||||
let errors = QueryRegistry::from_specs(vec![
|
||||
spec_tool("a", "query a() { match { $u: User } return { $u.name } }", true, "dup"),
|
||||
spec_tool("b", "query b() { match { $u: User } return { $u.name } }", true, "dup"),
|
||||
])
|
||||
.unwrap_err();
|
||||
assert_eq!(errors.len(), 1);
|
||||
let msg = errors[0].to_string();
|
||||
assert!(msg.contains("'dup'"), "names the contested tool: {msg}");
|
||||
assert!(msg.contains("'a'"), "names the winning query: {msg}");
|
||||
assert!(msg.contains("'b'"), "names the losing query: {msg}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn duplicate_tool_name_among_unexposed_is_allowed() {
|
||||
// Unexposed queries have no MCP tool, so a shared effective tool
|
||||
// name is inert — must not error (pins the exposed-only scope).
|
||||
let reg = QueryRegistry::from_specs(vec![
|
||||
spec_tool("a", "query a() { match { $u: User } return { $u.name } }", false, "dup"),
|
||||
spec_tool("b", "query b() { match { $u: User } return { $u.name } }", false, "dup"),
|
||||
])
|
||||
.unwrap();
|
||||
assert_eq!(reg.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_error_surfaces_per_entry() {
|
||||
let errors =
|
||||
QueryRegistry::from_specs(vec![spec("broken", "query broken( {{ not valid", false)])
|
||||
.unwrap_err();
|
||||
assert_eq!(errors[0].query.as_deref(), Some("broken"));
|
||||
assert!(errors[0].message.contains("parse error"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn errors_collect_rather_than_fail_fast() {
|
||||
let errors = QueryRegistry::from_specs(vec![
|
||||
spec("good", "query good() { match { $u: User } return { $u.name } }", false),
|
||||
spec("mismatch", "query other() { match { $u: User } return { $u.name } }", false),
|
||||
spec("broken", "query broken(", false),
|
||||
])
|
||||
.unwrap_err();
|
||||
// `good` loads cleanly; only the mismatch and the parse error are
|
||||
// reported, and both surface in one pass (not fail-fast).
|
||||
assert_eq!(errors.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mutation_body_classifies_as_mutation() {
|
||||
let reg = QueryRegistry::from_specs(vec![spec(
|
||||
"add_user",
|
||||
"query add_user($name: String) { insert User { name: $name } }",
|
||||
false,
|
||||
)])
|
||||
.unwrap();
|
||||
assert!(reg.lookup("add_user").unwrap().is_mutation());
|
||||
}
|
||||
|
||||
// --- check(registry, catalog) ---
|
||||
|
||||
use omnigraph_compiler::catalog::build_catalog;
|
||||
use omnigraph_compiler::schema::parser::parse_schema;
|
||||
|
||||
fn test_catalog() -> Catalog {
|
||||
let schema = parse_schema(
|
||||
r#"
|
||||
node User {
|
||||
name: String
|
||||
age: I32?
|
||||
embedding: Vector(4)
|
||||
}
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
build_catalog(&schema).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_passes_for_valid_query() {
|
||||
let reg = QueryRegistry::from_specs(vec![spec(
|
||||
"find_user",
|
||||
"query find_user($name: String) { match { $u: User { name: $name } } return { $u.age } }",
|
||||
false,
|
||||
)])
|
||||
.unwrap();
|
||||
let report = check(®, &test_catalog());
|
||||
assert!(report.is_clean(), "unexpected: {:?}", report);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_reports_unknown_type_as_breakage() {
|
||||
let reg = QueryRegistry::from_specs(vec![spec(
|
||||
"ghost",
|
||||
// `Widget` is not in the schema.
|
||||
"query ghost() { match { $w: Widget } return { $w.name } }",
|
||||
false,
|
||||
)])
|
||||
.unwrap();
|
||||
let report = check(®, &test_catalog());
|
||||
assert!(report.has_breakages());
|
||||
assert_eq!(report.breakages[0].query, "ghost");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_reports_unknown_property_as_breakage() {
|
||||
let reg = QueryRegistry::from_specs(vec![spec(
|
||||
"bad_prop",
|
||||
// `User` exists but has no `nickname`.
|
||||
"query bad_prop() { match { $u: User } return { $u.nickname } }",
|
||||
false,
|
||||
)])
|
||||
.unwrap();
|
||||
let report = check(®, &test_catalog());
|
||||
assert!(report.has_breakages());
|
||||
assert_eq!(report.breakages[0].query, "bad_prop");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_collects_every_breakage_not_fail_fast() {
|
||||
let reg = QueryRegistry::from_specs(vec![
|
||||
spec("a", "query a() { match { $w: Widget } return { $w.x } }", false),
|
||||
spec("b", "query b() { match { $g: Gadget } return { $g.y } }", false),
|
||||
spec(
|
||||
"ok",
|
||||
"query ok() { match { $u: User } return { $u.name } }",
|
||||
false,
|
||||
),
|
||||
])
|
||||
.unwrap();
|
||||
let report = check(®, &test_catalog());
|
||||
assert_eq!(report.breakages.len(), 2, "both bad queries reported: {:?}", report);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vector_param_on_exposed_query_warns() {
|
||||
let reg = QueryRegistry::from_specs(vec![spec(
|
||||
"vec_search",
|
||||
"query vec_search($q: Vector(4)) { match { $u: User } return { $u.name } \
|
||||
order { nearest($u.embedding, $q) } limit 3 }",
|
||||
true, // mcp.expose
|
||||
)])
|
||||
.unwrap();
|
||||
let report = check(®, &test_catalog());
|
||||
assert!(!report.has_breakages(), "valid query: {:?}", report);
|
||||
assert_eq!(report.warnings.len(), 1);
|
||||
assert_eq!(report.warnings[0].query, "vec_search");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vector_param_on_unexposed_query_is_silent() {
|
||||
let reg = QueryRegistry::from_specs(vec![spec(
|
||||
"vec_search",
|
||||
"query vec_search($q: Vector(4)) { match { $u: User } return { $u.name } \
|
||||
order { nearest($u.embedding, $q) } limit 3 }",
|
||||
false, // not exposed — vector param is fine for service-to-service callers
|
||||
)])
|
||||
.unwrap();
|
||||
let report = check(®, &test_catalog());
|
||||
assert!(report.is_clean(), "unexpected: {:?}", report);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_vector_param_on_exposed_query_does_not_warn() {
|
||||
// The recommended `String` alternative on an exposed query does not
|
||||
// resolve to a Vector, so the embedding advisory stays silent. Guards
|
||||
// the structured type check against a false positive (and pins that
|
||||
// only `Vector(_)` triggers the warning).
|
||||
let reg = QueryRegistry::from_specs(vec![spec(
|
||||
"search",
|
||||
"query search($name: String) { match { $u: User { name: $name } } return { $u.name } }",
|
||||
true,
|
||||
)])
|
||||
.unwrap();
|
||||
let report = check(®, &test_catalog());
|
||||
assert!(report.is_clean(), "no breakage or warning expected: {:?}", report);
|
||||
}
|
||||
|
||||
// --- catalog projection (api::query_catalog_entry) ---
|
||||
|
||||
#[test]
|
||||
fn catalog_entry_projects_every_param_kind() {
|
||||
use crate::api::{self, ParamKind};
|
||||
let reg = QueryRegistry::from_specs(vec![spec_tool(
|
||||
"all_types",
|
||||
"query all_types($s: String, $i: I32, $big: I64, $u: U64, $f: F64, $b: Bool, \
|
||||
$d: Date, $dt: DateTime, $blob: Blob, $opt: String?, $list: [I32], $vec: Vector(4)) \
|
||||
{ match { $x: User } return { $x.name } }",
|
||||
true,
|
||||
"all",
|
||||
)])
|
||||
.unwrap();
|
||||
let entry = api::query_catalog_entry(reg.lookup("all_types").unwrap());
|
||||
assert_eq!(entry.name, "all_types");
|
||||
assert_eq!(entry.tool_name, "all");
|
||||
assert!(!entry.mutation);
|
||||
|
||||
let by: std::collections::HashMap<_, _> =
|
||||
entry.params.iter().map(|p| (p.name.as_str(), p)).collect();
|
||||
assert_eq!(by["s"].kind, ParamKind::String);
|
||||
assert_eq!(by["i"].kind, ParamKind::Int);
|
||||
assert_eq!(by["big"].kind, ParamKind::BigInt, "I64 → bigint (string on the wire)");
|
||||
assert_eq!(by["u"].kind, ParamKind::BigInt, "U64 → bigint");
|
||||
assert_eq!(by["f"].kind, ParamKind::Float);
|
||||
assert_eq!(by["b"].kind, ParamKind::Bool);
|
||||
assert_eq!(by["d"].kind, ParamKind::Date);
|
||||
assert_eq!(by["dt"].kind, ParamKind::DateTime);
|
||||
assert_eq!(by["blob"].kind, ParamKind::Blob);
|
||||
assert!(!by["s"].nullable);
|
||||
assert!(by["opt"].nullable, "String? → nullable");
|
||||
assert_eq!(by["list"].kind, ParamKind::List);
|
||||
assert_eq!(by["list"].item_kind, Some(ParamKind::Int), "[I32] → list of int");
|
||||
assert_eq!(by["vec"].kind, ParamKind::Vector);
|
||||
assert_eq!(by["vec"].vector_dim, Some(4));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn catalog_entry_flags_mutation_and_empty_params() {
|
||||
use crate::api;
|
||||
let reg = QueryRegistry::from_specs(vec![spec(
|
||||
"add_user",
|
||||
"query add_user($name: String) { insert User { name: $name } }",
|
||||
true,
|
||||
)])
|
||||
.unwrap();
|
||||
let entry = api::query_catalog_entry(reg.lookup("add_user").unwrap());
|
||||
assert!(entry.mutation, "insert body → mutation flag");
|
||||
|
||||
let reg2 = QueryRegistry::from_specs(vec![spec(
|
||||
"no_params",
|
||||
"query no_params() { match { $u: User } return { $u.name } }",
|
||||
true,
|
||||
)])
|
||||
.unwrap();
|
||||
let entry2 = api::query_catalog_entry(reg2.lookup("no_params").unwrap());
|
||||
assert!(entry2.params.is_empty(), "no declared params → empty list");
|
||||
}
|
||||
|
||||
// --- load() error collection (file I/O + parse in one pass) ---
|
||||
|
||||
#[test]
|
||||
fn load_collects_io_and_parse_errors_in_one_pass() {
|
||||
use crate::config::load_config;
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
std::fs::write(
|
||||
temp.path().join("good.gq"),
|
||||
"query good() { match { $u: User } return { $u.name } }",
|
||||
)
|
||||
.unwrap();
|
||||
std::fs::write(temp.path().join("broken.gq"), "query broken( {{ not valid").unwrap();
|
||||
// `missing.gq` is deliberately not written (an I/O failure).
|
||||
std::fs::write(
|
||||
temp.path().join("omnigraph.yaml"),
|
||||
"queries:\n good:\n file: ./good.gq\n \
|
||||
missing:\n file: ./missing.gq\n broken:\n file: ./broken.gq\n",
|
||||
)
|
||||
.unwrap();
|
||||
let config = load_config(Some(&temp.path().join("omnigraph.yaml"))).unwrap();
|
||||
|
||||
let errors = QueryRegistry::load(&config, config.query_entries()).unwrap_err();
|
||||
let joined = errors.iter().map(|e| e.to_string()).collect::<Vec<_>>().join("\n");
|
||||
// Both the missing file AND the parse error surface in one pass —
|
||||
// the I/O failure must not mask the parse failure.
|
||||
assert!(joined.contains("missing"), "I/O error must surface: {joined}");
|
||||
assert!(
|
||||
joined.contains("broken") && joined.contains("parse error"),
|
||||
"the parse error in a readable file must surface in the same pass: {joined}"
|
||||
);
|
||||
assert!(!joined.contains("'good'"), "the valid entry is not an error: {joined}");
|
||||
}
|
||||
}
|
||||
|
|
@ -29,6 +29,7 @@ use tokio::sync::Mutex;
|
|||
|
||||
use crate::identity::GraphKey;
|
||||
use crate::policy::PolicyEngine;
|
||||
use crate::queries::QueryRegistry;
|
||||
|
||||
/// Open handle for a single graph in the registry. Cheap to clone (`Arc`-wrapped
|
||||
/// engine + policy). Cluster-mode handlers extract this via
|
||||
|
|
@ -47,6 +48,11 @@ pub struct GraphHandle {
|
|||
/// `_as` writers"; the HTTP-layer `require_bearer_auth` middleware still
|
||||
/// runs regardless.
|
||||
pub policy: Option<Arc<PolicyEngine>>,
|
||||
/// Per-graph stored-query registry, loaded and validated at
|
||||
/// startup. `None` means the operator declared no stored queries for
|
||||
/// this graph — `POST /queries/{name}` then 404s. Mirrors the
|
||||
/// optional `policy` shape.
|
||||
pub queries: Option<Arc<QueryRegistry>>,
|
||||
}
|
||||
|
||||
/// Immutable snapshot of the registry's current state. Replaced atomically
|
||||
|
|
@ -245,6 +251,7 @@ fn canonicalize_handle_uri(
|
|||
uri: canonical_uri.clone(),
|
||||
engine: Arc::clone(&handle.engine),
|
||||
policy: handle.policy.clone(),
|
||||
queries: handle.queries.clone(),
|
||||
});
|
||||
Ok((canonical_uri, canonical_handle))
|
||||
}
|
||||
|
|
@ -276,6 +283,7 @@ mod tests {
|
|||
uri: graph_uri,
|
||||
engine: Arc::new(engine),
|
||||
policy: None,
|
||||
queries: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -340,12 +348,14 @@ mod tests {
|
|||
uri: shared_uri.clone(),
|
||||
engine: Arc::clone(&engine),
|
||||
policy: None,
|
||||
queries: None,
|
||||
});
|
||||
let h2 = Arc::new(GraphHandle {
|
||||
key: GraphKey::cluster(GraphId::try_from("beta").unwrap()),
|
||||
uri: shared_uri,
|
||||
engine,
|
||||
policy: None,
|
||||
queries: None,
|
||||
});
|
||||
|
||||
let registry = GraphRegistry::new();
|
||||
|
|
@ -411,12 +421,14 @@ mod tests {
|
|||
uri: shared_uri.clone(),
|
||||
engine: Arc::clone(&engine),
|
||||
policy: None,
|
||||
queries: None,
|
||||
});
|
||||
let h2 = Arc::new(GraphHandle {
|
||||
key: GraphKey::cluster(GraphId::try_from("beta").unwrap()),
|
||||
uri: shared_uri,
|
||||
engine,
|
||||
policy: None,
|
||||
queries: None,
|
||||
});
|
||||
let err = match GraphRegistry::from_handles(vec![h1, h2]) {
|
||||
Ok(_) => panic!("expected DuplicateUri, got Ok"),
|
||||
|
|
|
|||
|
|
@ -168,6 +168,8 @@ const EXPECTED_PATHS: &[&str] = &[
|
|||
"/export",
|
||||
"/change",
|
||||
"/mutate",
|
||||
"/queries",
|
||||
"/queries/{name}",
|
||||
"/schema",
|
||||
"/schema/apply",
|
||||
"/ingest",
|
||||
|
|
@ -701,6 +703,8 @@ fn protected_endpoints_reference_bearer_token_security() {
|
|||
("/read", "post"),
|
||||
("/change", "post"),
|
||||
("/schema/apply", "post"),
|
||||
("/queries", "get"),
|
||||
("/queries/{name}", "post"),
|
||||
("/ingest", "post"),
|
||||
("/export", "post"),
|
||||
("/snapshot", "get"),
|
||||
|
|
@ -913,6 +917,34 @@ fn post_endpoints_have_request_body() {
|
|||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invoke_stored_query_request_body_is_optional() {
|
||||
let doc = openapi_json();
|
||||
let request_body = &doc["paths"]["/queries/{name}"]["post"]["requestBody"];
|
||||
assert!(
|
||||
request_body.is_object(),
|
||||
"POST /queries/{{name}} should document its optional request body"
|
||||
);
|
||||
assert_eq!(
|
||||
request_body["required"].as_bool().unwrap_or(false),
|
||||
false,
|
||||
"stored-query invocation body should be optional"
|
||||
);
|
||||
let schema = &request_body["content"]["application/json"]["schema"];
|
||||
let ref_path = schema["$ref"]
|
||||
.as_str()
|
||||
.or_else(|| {
|
||||
schema["oneOf"]
|
||||
.as_array()
|
||||
.and_then(|schemas| schemas.iter().find_map(|schema| schema["$ref"].as_str()))
|
||||
})
|
||||
.unwrap();
|
||||
assert!(
|
||||
ref_path.contains("InvokeStoredQueryRequest"),
|
||||
"POST /queries/{{name}} requestBody should reference InvokeStoredQueryRequest, got {ref_path}"
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Serialization round-trip test
|
||||
// ---------------------------------------------------------------------------
|
||||
|
|
@ -1117,6 +1149,7 @@ async fn app_for_multi_mode(graph_ids: &[&str]) -> (Vec<tempfile::TempDir>, Rout
|
|||
uri: graph_uri,
|
||||
engine: Arc::new(engine),
|
||||
policy: None,
|
||||
queries: None,
|
||||
}));
|
||||
dirs.push(dir);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ use axum::body::{Body, to_bytes};
|
|||
use axum::http::header::AUTHORIZATION;
|
||||
use axum::http::{Method, Request, StatusCode};
|
||||
use lance::index::DatasetIndexExt;
|
||||
use omnigraph::db::{Omnigraph, ReadTarget, SchemaApplyOptions};
|
||||
use omnigraph::db::{Omnigraph, ReadTarget};
|
||||
use omnigraph::error::OmniError;
|
||||
use omnigraph::loader::{LoadMode, load_jsonl};
|
||||
use omnigraph_policy::{PolicyChecker, PolicyEngine};
|
||||
|
|
@ -16,6 +16,7 @@ use omnigraph_server::api::{
|
|||
BranchCreateRequest, BranchMergeRequest, ChangeRequest, ErrorOutput, ExportRequest,
|
||||
IngestRequest, QueryRequest, ReadRequest, SchemaApplyRequest, SchemaOutput,
|
||||
};
|
||||
use omnigraph_server::queries::{QueryRegistry, RegistrySpec};
|
||||
use omnigraph_server::{AppState, build_app};
|
||||
use serde_json::{Value, json};
|
||||
use serial_test::serial;
|
||||
|
|
@ -141,6 +142,469 @@ fn graph_path(root: &Path) -> PathBuf {
|
|||
root.join("server.omni")
|
||||
}
|
||||
|
||||
fn stored_query_registry(specs: &[(&str, &str, bool)]) -> QueryRegistry {
|
||||
QueryRegistry::from_specs(
|
||||
specs
|
||||
.iter()
|
||||
.map(|(name, source, expose)| RegistrySpec {
|
||||
name: name.to_string(),
|
||||
source: source.to_string(),
|
||||
expose: *expose,
|
||||
tool_name: None,
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
.expect("specs parse and key==symbol")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn server_boots_with_a_valid_stored_query_registry() {
|
||||
// A stored query that type-checks against the fixture schema
|
||||
// (`Person { name, age }`) must let the server boot.
|
||||
let temp = init_loaded_graph().await;
|
||||
let graph = graph_path(temp.path());
|
||||
let registry = stored_query_registry(&[(
|
||||
"find_person",
|
||||
"query find_person($name: String) { match { $p: Person { name: $name } } return { $p.age } }",
|
||||
false,
|
||||
)]);
|
||||
let state = AppState::open_single_with_queries(
|
||||
graph.to_string_lossy().to_string(),
|
||||
vec![],
|
||||
None,
|
||||
registry,
|
||||
)
|
||||
.await;
|
||||
assert!(state.is_ok(), "valid registry should boot: {:?}", state.err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn server_refuses_boot_on_type_broken_stored_query() {
|
||||
// A stored query referencing a type not in the schema (`Widget`)
|
||||
// must abort boot, naming the offending query.
|
||||
let temp = init_loaded_graph().await;
|
||||
let graph = graph_path(temp.path());
|
||||
let registry = stored_query_registry(&[(
|
||||
"ghost",
|
||||
"query ghost() { match { $w: Widget } return { $w.name } }",
|
||||
false,
|
||||
)]);
|
||||
let result = AppState::open_single_with_queries(
|
||||
graph.to_string_lossy().to_string(),
|
||||
vec![],
|
||||
None,
|
||||
registry,
|
||||
)
|
||||
.await;
|
||||
// `AppState` is not `Debug`, so match rather than `expect_err`.
|
||||
let err = match result {
|
||||
Ok(_) => panic!("type-broken stored query must refuse boot"),
|
||||
Err(err) => err,
|
||||
};
|
||||
let msg = err.to_string();
|
||||
assert!(msg.contains("ghost"), "error should name the broken query: {msg}");
|
||||
assert!(
|
||||
msg.contains("schema check"),
|
||||
"error should mention the schema check: {msg}"
|
||||
);
|
||||
}
|
||||
|
||||
/// Build a single-mode app with a stored-query registry plus a bearer→actor
|
||||
/// pairing and a policy, so invoke tests exercise the `invoke_query`
|
||||
/// boundary gate and the inner read/change gates together.
|
||||
async fn app_with_stored_queries(
|
||||
specs: &[(&str, &str, bool)],
|
||||
tokens: &[(&str, &str)],
|
||||
policy: &str,
|
||||
) -> (tempfile::TempDir, Router) {
|
||||
let temp = init_loaded_graph().await;
|
||||
let graph = graph_path(temp.path());
|
||||
let policy_path = temp.path().join("policy.yaml");
|
||||
fs::write(&policy_path, policy).unwrap();
|
||||
let registry = stored_query_registry(specs);
|
||||
let state = AppState::open_single_with_queries(
|
||||
graph.to_string_lossy().to_string(),
|
||||
tokens
|
||||
.iter()
|
||||
.map(|(actor, token)| ((*actor).to_string(), (*token).to_string()))
|
||||
.collect(),
|
||||
Some(&policy_path),
|
||||
registry,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
(temp, build_app(state))
|
||||
}
|
||||
|
||||
/// - `act-invoke`: invoke_query + read (stored reads, not mutations)
|
||||
/// - `act-full`: invoke_query + read + change (stored mutations)
|
||||
/// - `act-noinvoke`: read only, no invoke_query (boundary-denied)
|
||||
/// - `act-invokeonly`: invoke_query only, no read (clears the boundary, inner read denies)
|
||||
const INVOKE_POLICY_YAML: &str = r#"
|
||||
version: 1
|
||||
groups:
|
||||
invokers: ["act-invoke"]
|
||||
full: ["act-full"]
|
||||
readers: ["act-noinvoke"]
|
||||
invoke_only: ["act-invokeonly"]
|
||||
protected_branches: [main]
|
||||
rules:
|
||||
# invoke_query is graph-scoped — its own rules, no branch_scope.
|
||||
- id: invokers-can-invoke
|
||||
allow:
|
||||
actors: { group: invokers }
|
||||
actions: [invoke_query]
|
||||
- id: full-can-invoke
|
||||
allow:
|
||||
actors: { group: full }
|
||||
actions: [invoke_query]
|
||||
- id: invoke-only-can-invoke
|
||||
allow:
|
||||
actors: { group: invoke_only }
|
||||
actions: [invoke_query]
|
||||
# read / change are branch-scoped.
|
||||
- id: invokers-can-read
|
||||
allow:
|
||||
actors: { group: invokers }
|
||||
actions: [read]
|
||||
branch_scope: any
|
||||
- id: full-can-read-change
|
||||
allow:
|
||||
actors: { group: full }
|
||||
actions: [read, change]
|
||||
branch_scope: any
|
||||
- id: readers-can-read
|
||||
allow:
|
||||
actors: { group: readers }
|
||||
actions: [read]
|
||||
branch_scope: any
|
||||
"#;
|
||||
|
||||
const STORED_QUERY_SCHEMA_APPLY_POLICY_YAML: &str = r#"
|
||||
version: 1
|
||||
groups:
|
||||
admins: [act-ragnor]
|
||||
protected_branches: [main]
|
||||
rules:
|
||||
- id: admins-can-invoke
|
||||
allow:
|
||||
actors: { group: admins }
|
||||
actions: [invoke_query]
|
||||
- id: admins-can-read
|
||||
allow:
|
||||
actors: { group: admins }
|
||||
actions: [read]
|
||||
branch_scope: any
|
||||
- id: admins-can-schema-apply
|
||||
allow:
|
||||
actors: { group: admins }
|
||||
actions: [schema_apply]
|
||||
target_branch_scope: protected
|
||||
"#;
|
||||
|
||||
const FIND_PERSON_GQ: &str =
|
||||
"query find_person($name: String) { match { $p: Person { name: $name } } return { $p.age } }";
|
||||
|
||||
fn invoke_request(name: &str, token: &str, body: Value) -> Request<Body> {
|
||||
Request::builder()
|
||||
.uri(format!("/queries/{name}"))
|
||||
.method(Method::POST)
|
||||
.header("content-type", "application/json")
|
||||
.header("authorization", format!("Bearer {token}"))
|
||||
.body(Body::from(serde_json::to_vec(&body).unwrap()))
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn invoke_request_bytes(
|
||||
name: &str,
|
||||
token: &str,
|
||||
body: impl Into<Body>,
|
||||
content_type: Option<&str>,
|
||||
) -> Request<Body> {
|
||||
let mut builder = Request::builder()
|
||||
.uri(format!("/queries/{name}"))
|
||||
.method(Method::POST)
|
||||
.header("authorization", format!("Bearer {token}"));
|
||||
if let Some(content_type) = content_type {
|
||||
builder = builder.header("content-type", content_type);
|
||||
}
|
||||
builder.body(body.into()).unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn invoke_stored_read_returns_rows() {
|
||||
let (_temp, app) = app_with_stored_queries(
|
||||
&[("find_person", FIND_PERSON_GQ, false)],
|
||||
&[("act-invoke", "t-invoke")],
|
||||
INVOKE_POLICY_YAML,
|
||||
)
|
||||
.await;
|
||||
let (status, body) = json_response(
|
||||
&app,
|
||||
invoke_request("find_person", "t-invoke", json!({ "params": { "name": "Alice" } })),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK, "body: {body}");
|
||||
assert_eq!(body["query_name"], "find_person");
|
||||
assert_eq!(body["row_count"], 1, "Alice is in the fixture; body: {body}");
|
||||
assert!(body["rows"].is_array(), "read envelope shape; body: {body}");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn invoke_stored_read_accepts_absent_or_empty_body() {
|
||||
let no_param_query = "query list_people() { match { $p: Person } return { $p.name } }";
|
||||
let (_temp, app) = app_with_stored_queries(
|
||||
&[("list_people", no_param_query, false)],
|
||||
&[("act-invoke", "t-invoke")],
|
||||
INVOKE_POLICY_YAML,
|
||||
)
|
||||
.await;
|
||||
|
||||
let (status, body) = json_response(
|
||||
&app,
|
||||
invoke_request_bytes("list_people", "t-invoke", Body::empty(), None),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK, "body: {body}");
|
||||
assert_eq!(body["query_name"], "list_people");
|
||||
|
||||
let (status, body) = json_response(
|
||||
&app,
|
||||
invoke_request_bytes(
|
||||
"list_people",
|
||||
"t-invoke",
|
||||
Body::empty(),
|
||||
Some("application/json"),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK, "body: {body}");
|
||||
|
||||
let (status, body) = json_response(
|
||||
&app,
|
||||
invoke_request_bytes(
|
||||
"list_people",
|
||||
"t-invoke",
|
||||
Body::from("{}"),
|
||||
Some("application/json"),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK, "body: {body}");
|
||||
|
||||
let (status, body) = json_response(
|
||||
&app,
|
||||
invoke_request_bytes(
|
||||
"list_people",
|
||||
"t-invoke",
|
||||
Body::from("{"),
|
||||
Some("application/json"),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST, "body: {body}");
|
||||
assert!(
|
||||
body["error"]
|
||||
.as_str()
|
||||
.unwrap_or_default()
|
||||
.contains("invalid stored-query invocation body"),
|
||||
"malformed JSON should be rejected as bad request; body: {body}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn invoke_stored_mutation_double_gates_on_change() {
|
||||
let specs: &[(&str, &str, bool)] = &[(
|
||||
"add_person",
|
||||
"query add_person($name: String) { insert Person { name: $name } }",
|
||||
false,
|
||||
)];
|
||||
let (_temp, app) = app_with_stored_queries(
|
||||
specs,
|
||||
&[("act-invoke", "t-invoke"), ("act-full", "t-full")],
|
||||
INVOKE_POLICY_YAML,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Has invoke_query but NOT change → the inner change gate denies (403).
|
||||
let (status, body) = json_response(
|
||||
&app,
|
||||
invoke_request("add_person", "t-invoke", json!({ "params": { "name": "Eve" } })),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(
|
||||
status,
|
||||
StatusCode::FORBIDDEN,
|
||||
"invoke_query without change must 403; body: {body}"
|
||||
);
|
||||
|
||||
// Has invoke_query + change → applied.
|
||||
let (status, body) = json_response(
|
||||
&app,
|
||||
invoke_request("add_person", "t-full", json!({ "params": { "name": "Eve" } })),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::OK, "body: {body}");
|
||||
assert_eq!(body["affected_nodes"], 1, "body: {body}");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn invoke_stored_query_bad_param_is_400() {
|
||||
let (_temp, app) = app_with_stored_queries(
|
||||
&[("find_person", FIND_PERSON_GQ, false)],
|
||||
&[("act-invoke", "t-invoke")],
|
||||
INVOKE_POLICY_YAML,
|
||||
)
|
||||
.await;
|
||||
// `name` is declared String; pass a number.
|
||||
let (status, body) = json_response(
|
||||
&app,
|
||||
invoke_request("find_person", "t-invoke", json!({ "params": { "name": 123 } })),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST, "body: {body}");
|
||||
assert!(
|
||||
body["error"].as_str().unwrap_or_default().contains("name"),
|
||||
"400 should name the offending param; body: {body}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn invoke_unknown_query_and_denied_actor_return_identical_404() {
|
||||
let (_temp, app) = app_with_stored_queries(
|
||||
&[("find_person", FIND_PERSON_GQ, false)],
|
||||
&[("act-invoke", "t-invoke"), ("act-noinvoke", "t-noinvoke")],
|
||||
INVOKE_POLICY_YAML,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Authorized actor, unknown query name → 404.
|
||||
let (unknown_status, unknown_body) =
|
||||
json_response(&app, invoke_request("does_not_exist", "t-invoke", json!({}))).await;
|
||||
// Denied actor (no invoke_query), real query name → 404.
|
||||
let (denied_status, denied_body) = json_response(
|
||||
&app,
|
||||
invoke_request("find_person", "t-noinvoke", json!({ "params": { "name": "Alice" } })),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(unknown_status, StatusCode::NOT_FOUND);
|
||||
assert_eq!(denied_status, StatusCode::NOT_FOUND);
|
||||
assert_eq!(
|
||||
unknown_body, denied_body,
|
||||
"deny must be byte-identical to a missing query (no catalog probing)"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn invoke_query_holder_without_read_sees_403_not_404() {
|
||||
// The 404-hiding is for callers WITHOUT invoke_query. An actor that
|
||||
// HOLDS invoke_query but lacks `read` clears the boundary gate, then the
|
||||
// inner read gate denies → 403 for an EXISTING read query, vs 404 for an
|
||||
// unknown one. Existence is visible to grant-holders by design (the
|
||||
// documented double-gate); this pins that actual contract.
|
||||
let (_temp, app) = app_with_stored_queries(
|
||||
&[("find_person", FIND_PERSON_GQ, false)],
|
||||
&[("act-invokeonly", "t-invokeonly")],
|
||||
INVOKE_POLICY_YAML,
|
||||
)
|
||||
.await;
|
||||
let (exists_status, _) = json_response(
|
||||
&app,
|
||||
invoke_request("find_person", "t-invokeonly", json!({ "params": { "name": "Alice" } })),
|
||||
)
|
||||
.await;
|
||||
let (absent_status, _) =
|
||||
json_response(&app, invoke_request("does_not_exist", "t-invokeonly", json!({}))).await;
|
||||
assert_eq!(
|
||||
exists_status,
|
||||
StatusCode::FORBIDDEN,
|
||||
"an existing read query the holder can't read → inner-gate 403"
|
||||
);
|
||||
assert_eq!(absent_status, StatusCode::NOT_FOUND, "unknown query still 404s");
|
||||
}
|
||||
|
||||
fn get_request(uri: &str, token: &str) -> Request<Body> {
|
||||
Request::builder()
|
||||
.uri(uri)
|
||||
.method(Method::GET)
|
||||
.header("authorization", format!("Bearer {token}"))
|
||||
.body(Body::empty())
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn list_queries_returns_only_exposed_with_typed_params() {
|
||||
let (_temp, app) = app_with_stored_queries(
|
||||
&[
|
||||
("find_person", FIND_PERSON_GQ, true),
|
||||
(
|
||||
"add_person",
|
||||
"query add_person($name: String) { insert Person { name: $name } }",
|
||||
true,
|
||||
),
|
||||
("hidden", "query hidden() { match { $p: Person } return { $p.name } }", false),
|
||||
],
|
||||
&[("act-invoke", "t-invoke")],
|
||||
INVOKE_POLICY_YAML,
|
||||
)
|
||||
.await;
|
||||
let (status, body) = json_response(&app, get_request("/queries", "t-invoke")).await;
|
||||
assert_eq!(status, StatusCode::OK, "body: {body}");
|
||||
|
||||
let entries = body["queries"].as_array().unwrap();
|
||||
let names: Vec<&str> = entries.iter().map(|q| q["name"].as_str().unwrap()).collect();
|
||||
assert!(
|
||||
names.contains(&"find_person") && names.contains(&"add_person"),
|
||||
"exposed queries listed: {names:?}"
|
||||
);
|
||||
assert!(!names.contains(&"hidden"), "non-exposed query hidden from the catalog: {names:?}");
|
||||
|
||||
let fp = entries.iter().find(|q| q["name"] == "find_person").unwrap();
|
||||
assert_eq!(fp["mutation"], false);
|
||||
assert_eq!(fp["tool_name"], "find_person");
|
||||
assert_eq!(fp["params"][0]["name"], "name");
|
||||
assert_eq!(fp["params"][0]["kind"], "string");
|
||||
let ap = entries.iter().find(|q| q["name"] == "add_person").unwrap();
|
||||
assert_eq!(ap["mutation"], true, "stored insert → mutation");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn list_queries_is_read_gated_so_a_non_invoker_can_list() {
|
||||
// The catalog is read-gated (not invoke_query-gated), so a reader who
|
||||
// lacks invoke_query still enumerates the exposed queries — the
|
||||
// documented probe-oracle gap until per-query Cedar filtering lands.
|
||||
let (_temp, app) = app_with_stored_queries(
|
||||
&[("find_person", FIND_PERSON_GQ, true)],
|
||||
&[("act-noinvoke", "t-noinvoke")],
|
||||
INVOKE_POLICY_YAML,
|
||||
)
|
||||
.await;
|
||||
let (status, body) = json_response(&app, get_request("/queries", "t-noinvoke")).await;
|
||||
assert_eq!(status, StatusCode::OK, "read-gated catalog; body: {body}");
|
||||
let names: Vec<&str> = body["queries"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|q| q["name"].as_str().unwrap())
|
||||
.collect();
|
||||
assert!(
|
||||
names.contains(&"find_person"),
|
||||
"a reader lists the catalog despite lacking invoke_query: {names:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn list_queries_is_empty_when_no_registry() {
|
||||
let (_temp, app) = app_for_loaded_graph_with_auth("demo-token").await;
|
||||
let (status, body) = json_response(&app, get_request("/queries", "demo-token")).await;
|
||||
assert_eq!(status, StatusCode::OK, "body: {body}");
|
||||
assert!(
|
||||
body["queries"].as_array().unwrap().is_empty(),
|
||||
"no stored-query registry → empty catalog"
|
||||
);
|
||||
}
|
||||
|
||||
fn drifted_test_schema() -> String {
|
||||
fs::read_to_string(fixture("test.pg"))
|
||||
.unwrap()
|
||||
|
|
@ -423,6 +887,83 @@ async fn schema_apply_route_updates_graph_for_authorized_admin() {
|
|||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn schema_apply_route_rejects_stored_query_breakage_before_publish() {
|
||||
let (temp, app) = app_with_stored_queries(
|
||||
&[("find_person", FIND_PERSON_GQ, true)],
|
||||
&[("act-ragnor", "admin-token")],
|
||||
STORED_QUERY_SCHEMA_APPLY_POLICY_YAML,
|
||||
)
|
||||
.await;
|
||||
|
||||
let request = Request::builder()
|
||||
.method(Method::POST)
|
||||
.uri("/schema/apply")
|
||||
.header("content-type", "application/json")
|
||||
.header("authorization", "Bearer admin-token")
|
||||
.body(Body::from(
|
||||
serde_json::to_vec(&SchemaApplyRequest {
|
||||
schema_source: renamed_age_schema(),
|
||||
..Default::default()
|
||||
})
|
||||
.unwrap(),
|
||||
))
|
||||
.unwrap();
|
||||
let (status, payload) = json_response(&app, request).await;
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST, "body: {payload}");
|
||||
let message = payload["error"].as_str().unwrap_or_default();
|
||||
assert!(
|
||||
message.contains("find_person") && message.contains("schema check"),
|
||||
"registry breakage should name the stored query; body: {payload}"
|
||||
);
|
||||
|
||||
let reopened = Omnigraph::open(graph_path(temp.path()).to_str().unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
let person = &reopened.catalog().node_types["Person"];
|
||||
assert!(person.properties.contains_key("age"));
|
||||
assert!(!person.properties.contains_key("years"));
|
||||
|
||||
let (invoke_status, invoke_body) = json_response(
|
||||
&app,
|
||||
invoke_request(
|
||||
"find_person",
|
||||
"admin-token",
|
||||
json!({ "params": { "name": "Alice" } }),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(invoke_status, StatusCode::OK, "body: {invoke_body}");
|
||||
assert_eq!(invoke_body["row_count"], 1);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn schema_apply_route_noop_keeps_valid_stored_query_registry() {
|
||||
let (_temp, app) = app_with_stored_queries(
|
||||
&[("find_person", FIND_PERSON_GQ, true)],
|
||||
&[("act-ragnor", "admin-token")],
|
||||
STORED_QUERY_SCHEMA_APPLY_POLICY_YAML,
|
||||
)
|
||||
.await;
|
||||
|
||||
let request = Request::builder()
|
||||
.method(Method::POST)
|
||||
.uri("/schema/apply")
|
||||
.header("content-type", "application/json")
|
||||
.header("authorization", "Bearer admin-token")
|
||||
.body(Body::from(
|
||||
serde_json::to_vec(&SchemaApplyRequest {
|
||||
schema_source: fs::read_to_string(fixture("test.pg")).unwrap(),
|
||||
..Default::default()
|
||||
})
|
||||
.unwrap(),
|
||||
))
|
||||
.unwrap();
|
||||
let (status, payload) = json_response(&app, request).await;
|
||||
assert_eq!(status, StatusCode::OK, "body: {payload}");
|
||||
assert_eq!(payload["applied"], false);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn schema_apply_route_requires_schema_apply_policy_permission() {
|
||||
let (_temp, app) = app_for_graph_with_auth_tokens_and_policy(
|
||||
|
|
@ -4690,6 +5231,7 @@ mod multi_graph_startup {
|
|||
uri: graph_uri,
|
||||
engine: Arc::new(engine),
|
||||
policy: None,
|
||||
queries: None,
|
||||
}));
|
||||
dirs.push(dir);
|
||||
}
|
||||
|
|
@ -4985,12 +5527,14 @@ graphs:
|
|||
uri: graph_uri.clone(),
|
||||
engine: Arc::clone(&engine),
|
||||
policy: None,
|
||||
queries: None,
|
||||
});
|
||||
let beta = Arc::new(GraphHandle {
|
||||
key: GraphKey::cluster(GraphId::try_from("beta").unwrap()),
|
||||
uri: format!("file://{graph_uri}/"),
|
||||
engine,
|
||||
policy: None,
|
||||
queries: None,
|
||||
});
|
||||
|
||||
match GraphRegistry::from_handles(vec![alpha, beta]) {
|
||||
|
|
@ -5016,6 +5560,7 @@ graphs:
|
|||
uri: format!("file://{graph_uri}/"),
|
||||
engine: Arc::new(engine),
|
||||
policy: None,
|
||||
queries: None,
|
||||
});
|
||||
|
||||
let registry = GraphRegistry::from_handles(vec![handle]).unwrap();
|
||||
|
|
@ -5138,11 +5683,11 @@ graphs:
|
|||
let err = load_server_settings(Some(&config_path), None, None, None, true).unwrap_err();
|
||||
let msg = err.to_string();
|
||||
assert!(
|
||||
msg.contains("top-level `policy.file` is single-graph/CLI-local policy only"),
|
||||
"expected single-graph policy guidance, got: {msg}"
|
||||
msg.contains("top-level") && msg.contains("policy.file") && msg.contains("not honored"),
|
||||
"expected top-level-not-honored guidance, got: {msg}"
|
||||
);
|
||||
assert!(
|
||||
msg.contains("graphs.<graph_id>.policy.file"),
|
||||
msg.contains("graphs.<graph_id>"),
|
||||
"expected per-graph migration guidance, got: {msg}"
|
||||
);
|
||||
assert!(
|
||||
|
|
@ -5151,6 +5696,88 @@ graphs:
|
|||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mode_inference_multi_rejects_top_level_queries() {
|
||||
// Symmetric to the policy guard: a top-level `queries:` block in
|
||||
// multi-graph mode is not honored (each graph uses its own), so it
|
||||
// is a loud error rather than a silent no-op.
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let config_path = temp.path().join("omnigraph.yaml");
|
||||
fs::write(
|
||||
&config_path,
|
||||
"queries:\n q:\n file: ./q.gq\ngraphs:\n alpha:\n uri: /tmp/alpha.omni\n",
|
||||
)
|
||||
.unwrap();
|
||||
let err = load_server_settings(Some(&config_path), None, None, None, true).unwrap_err();
|
||||
let msg = err.to_string();
|
||||
assert!(
|
||||
msg.contains("queries") && msg.contains("not honored"),
|
||||
"top-level queries must be rejected in multi-graph mode: {msg}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_mode_named_graph_rejects_top_level_blocks() {
|
||||
// Serving a graph by name (`--target`/`server.graph`) uses its
|
||||
// per-graph block; a populated top-level block would be silently
|
||||
// shadowed, so boot refuses and names the per-graph location.
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let config_path = temp.path().join("omnigraph.yaml");
|
||||
fs::write(
|
||||
&config_path,
|
||||
"policy:\n file: ./top.yaml\ngraphs:\n prod:\n uri: /tmp/prod.omni\n",
|
||||
)
|
||||
.unwrap();
|
||||
let err =
|
||||
load_server_settings(Some(&config_path), None, Some("prod".to_string()), None, true)
|
||||
.unwrap_err();
|
||||
let msg = err.to_string();
|
||||
assert!(
|
||||
msg.contains("prod") && msg.contains("policy.file") && msg.contains("graphs.prod"),
|
||||
"named single-mode + top-level policy must refuse, naming the graph: {msg}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_mode_named_graph_uses_per_graph_policy_and_queries() {
|
||||
// The identity rule: `--target prod` attaches `graphs.prod`'s own
|
||||
// policy + queries, not the top-level ones (which are absent here).
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
fs::write(
|
||||
temp.path().join("prod.gq"),
|
||||
"query pq() { match { $u: User } return { $u.name } }",
|
||||
)
|
||||
.unwrap();
|
||||
let config_path = temp.path().join("omnigraph.yaml");
|
||||
fs::write(
|
||||
&config_path,
|
||||
"graphs:\n prod:\n uri: /tmp/prod.omni\n policy:\n file: ./prod-policy.yaml\n \
|
||||
queries:\n pq:\n file: ./prod.gq\n",
|
||||
)
|
||||
.unwrap();
|
||||
let settings =
|
||||
load_server_settings(Some(&config_path), None, Some("prod".to_string()), None, true)
|
||||
.unwrap();
|
||||
match settings.mode {
|
||||
ServerConfigMode::Single {
|
||||
graph_id,
|
||||
policy_file,
|
||||
queries,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(graph_id, "prod", "named single-mode keeps graph identity");
|
||||
assert!(
|
||||
policy_file
|
||||
.as_ref()
|
||||
.is_some_and(|p| p.ends_with("prod-policy.yaml")),
|
||||
"per-graph policy attached: {policy_file:?}"
|
||||
);
|
||||
assert!(queries.lookup("pq").is_some(), "per-graph query attached");
|
||||
}
|
||||
other => panic!("expected Single mode, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mode_inference_normalizes_multi_graph_uris() {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
|
|
@ -5383,6 +6010,7 @@ graphs:
|
|||
uri: graph_uri,
|
||||
engine: Arc::new(engine),
|
||||
policy: None,
|
||||
queries: None,
|
||||
});
|
||||
let tokens = vec![("act-andrew".to_string(), "secret-token".to_string())];
|
||||
let workload = omnigraph_server::workload::WorkloadController::from_env();
|
||||
|
|
@ -5450,6 +6078,7 @@ graphs:
|
|||
uri: graph_uri,
|
||||
engine: Arc::new(engine),
|
||||
policy: None,
|
||||
queries: None,
|
||||
}));
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[package]
|
||||
name = "omnigraph-engine"
|
||||
version = "0.6.0"
|
||||
version = "0.6.1"
|
||||
edition = "2024"
|
||||
description = "Runtime engine for the Omnigraph graph database."
|
||||
license = "MIT"
|
||||
|
|
@ -16,8 +16,8 @@ default = []
|
|||
failpoints = ["dep:fail", "fail/failpoints"]
|
||||
|
||||
[dependencies]
|
||||
omnigraph-compiler = { path = "../omnigraph-compiler", version = "0.6.0" }
|
||||
omnigraph-policy = { path = "../omnigraph-policy", version = "0.6.0" }
|
||||
omnigraph-compiler = { path = "../omnigraph-compiler", version = "0.6.1" }
|
||||
omnigraph-policy = { path = "../omnigraph-policy", version = "0.6.1" }
|
||||
lance = { workspace = true }
|
||||
lance-datafusion = { workspace = true }
|
||||
datafusion = { workspace = true }
|
||||
|
|
@ -51,7 +51,7 @@ chrono = { workspace = true }
|
|||
arc-swap = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
omnigraph-compiler = { path = "../omnigraph-compiler", version = "0.6.0" }
|
||||
omnigraph-compiler = { path = "../omnigraph-compiler", version = "0.6.1" }
|
||||
tokio = { workspace = true }
|
||||
lance-namespace-impls = { workspace = true }
|
||||
serial_test = "3"
|
||||
|
|
|
|||
|
|
@ -169,6 +169,37 @@ impl CommitGraph {
|
|||
self.refresh().await
|
||||
}
|
||||
|
||||
/// Idempotently drop the commit-graph branch `name`, tolerating an
|
||||
/// already-absent branch (see [`TableStore::force_delete_branch`] for the
|
||||
/// same semantics). Used by the best-effort reclaim in `branch_delete` and
|
||||
/// the `cleanup` orphan reconciler. `RefConflict` (referencing descendants)
|
||||
/// is still surfaced.
|
||||
pub async fn force_delete_branch(&mut self, name: &str) -> Result<()> {
|
||||
let mut ds = Dataset::open(&graph_commits_uri(&self.root_uri))
|
||||
.await
|
||||
.map_err(|e| OmniError::Lance(e.to_string()))?;
|
||||
match ds.force_delete_branch(name).await {
|
||||
Ok(()) => {}
|
||||
Err(lance::Error::RefNotFound { .. }) | Err(lance::Error::NotFound { .. }) => {}
|
||||
Err(e) => return Err(OmniError::Lance(e.to_string())),
|
||||
}
|
||||
self.refresh().await
|
||||
}
|
||||
|
||||
/// List the named branches present on the commit-graph dataset. The
|
||||
/// `cleanup` reconciler diffs this against the manifest branch set to find
|
||||
/// orphaned commit-graph branches to reclaim.
|
||||
pub async fn list_branches(&self) -> Result<Vec<String>> {
|
||||
let ds = Dataset::open(&graph_commits_uri(&self.root_uri))
|
||||
.await
|
||||
.map_err(|e| OmniError::Lance(e.to_string()))?;
|
||||
let branches = ds
|
||||
.list_branches()
|
||||
.await
|
||||
.map_err(|e| OmniError::Lance(e.to_string()))?;
|
||||
Ok(branches.into_keys().collect())
|
||||
}
|
||||
|
||||
pub async fn append_commit(
|
||||
&mut self,
|
||||
manifest_branch: Option<&str>,
|
||||
|
|
@ -345,7 +376,7 @@ impl CommitGraph {
|
|||
}
|
||||
}
|
||||
|
||||
fn graph_commits_uri(root_uri: &str) -> String {
|
||||
pub(crate) fn graph_commits_uri(root_uri: &str) -> String {
|
||||
format!("{}/{}", root_uri.trim_end_matches('/'), GRAPH_COMMITS_DIR)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -211,14 +211,47 @@ impl GraphCoordinator {
|
|||
let branch = normalize_branch_name(name)?
|
||||
.ok_or_else(|| OmniError::manifest("cannot create branch 'main'".to_string()))?;
|
||||
self.ensure_commit_graph_initialized().await?;
|
||||
|
||||
// Manifest authority flip first.
|
||||
self.manifest.create_branch(&branch).await?;
|
||||
failpoints::maybe_fail("branch_create.after_manifest_branch_create")?;
|
||||
if let Some(commit_graph) = &mut self.commit_graph {
|
||||
commit_graph.create_branch(&branch).await?;
|
||||
|
||||
// Derived commit-graph branch. If anything after the authority flip
|
||||
// fails, roll back the manifest branch so the branch never half-exists
|
||||
// (a manifest branch with no commit-graph branch breaks the next write).
|
||||
if let Err(err) = self.create_commit_graph_branch(&branch).await {
|
||||
if let Err(rollback_err) = self.manifest.delete_branch(&branch).await {
|
||||
tracing::warn!(
|
||||
target: "omnigraph::branch_create",
|
||||
branch = %branch,
|
||||
error = %rollback_err,
|
||||
"rollback of manifest branch failed after commit-graph create failure",
|
||||
);
|
||||
}
|
||||
return Err(err);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Create the derived commit-graph branch for `branch`, healing a zombie ref
|
||||
/// left by an incomplete prior delete. The manifest branch was just created
|
||||
/// fresh, so any existing commit-graph branch with this name is provably
|
||||
/// orphaned and is force-dropped before recreating.
|
||||
async fn create_commit_graph_branch(&mut self, branch: &str) -> Result<()> {
|
||||
failpoints::maybe_fail("branch_create.after_manifest_branch_create")?;
|
||||
let Some(commit_graph) = &mut self.commit_graph else {
|
||||
return Ok(());
|
||||
};
|
||||
if commit_graph
|
||||
.list_branches()
|
||||
.await?
|
||||
.iter()
|
||||
.any(|existing| existing == branch)
|
||||
{
|
||||
commit_graph.force_delete_branch(branch).await?;
|
||||
}
|
||||
commit_graph.create_branch(branch).await
|
||||
}
|
||||
|
||||
pub async fn branch_delete(&mut self, name: &str) -> Result<()> {
|
||||
let branch = normalize_branch_name(name)?
|
||||
.ok_or_else(|| OmniError::manifest("cannot delete branch 'main'".to_string()))?;
|
||||
|
|
@ -229,20 +262,43 @@ impl GraphCoordinator {
|
|||
)));
|
||||
}
|
||||
|
||||
// Manifest authority flip — the single atomic op that makes the branch
|
||||
// cease to exist. Must succeed; everything after is derived state
|
||||
// reclaimed best-effort.
|
||||
self.manifest.delete_branch(&branch).await?;
|
||||
|
||||
// Commit-graph branch is derived state. Reclaim best-effort with the
|
||||
// idempotent force variant: a failure here (or a missing dataset) is
|
||||
// reconciled by `cleanup` and must not fail the delete after the
|
||||
// authority already flipped.
|
||||
if let Err(err) = self.reclaim_commit_graph_branch(&branch).await {
|
||||
tracing::warn!(
|
||||
target: "omnigraph::branch_delete::cleanup",
|
||||
branch = %branch,
|
||||
error = %err,
|
||||
"best-effort commit-graph branch reclaim failed; cleanup will reconcile",
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Best-effort, idempotent reclaim of the commit-graph branch `branch`.
|
||||
/// Tolerates an absent commit-graph dataset (a graph that never committed).
|
||||
async fn reclaim_commit_graph_branch(&mut self, branch: &str) -> Result<()> {
|
||||
failpoints::maybe_fail("branch_delete.before_commit_graph_reclaim")?;
|
||||
if let Some(commit_graph) = &mut self.commit_graph {
|
||||
commit_graph.delete_branch(&branch).await?;
|
||||
commit_graph.force_delete_branch(branch).await
|
||||
} else if self
|
||||
.storage
|
||||
.exists(&graph_commits_uri(self.root_uri()))
|
||||
.await?
|
||||
{
|
||||
let mut commit_graph = CommitGraph::open(self.root_uri()).await?;
|
||||
commit_graph.delete_branch(&branch).await?;
|
||||
commit_graph.force_delete_branch(branch).await
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn snapshot_at_version(&self, version: u64) -> Result<Snapshot> {
|
||||
|
|
|
|||
|
|
@ -48,6 +48,22 @@ const OBJECT_TYPE_TABLE_VERSION: &str = "table_version";
|
|||
const OBJECT_TYPE_TABLE_TOMBSTONE: &str = "table_tombstone";
|
||||
const TABLE_VERSION_MANAGEMENT_KEY: &str = "table_version_management";
|
||||
|
||||
/// Apply pending internal-schema migrations against `__manifest` on the
|
||||
/// open-for-write path, independent of a publish.
|
||||
///
|
||||
/// `Omnigraph::open(ReadWrite)` calls this before the coordinator reads branch
|
||||
/// state, so branch-observing code (`branch_list`, the schema-apply
|
||||
/// blocking-branch checks) sees the post-migration graph. In particular the
|
||||
/// v2→v3 step sweeps legacy `__run__*` staging branches off `__manifest`
|
||||
/// (MR-770); running it here closes the window where those branches would
|
||||
/// otherwise block schema apply before the first publish runs the migration.
|
||||
///
|
||||
/// Idempotent: a no-op stamp read when the on-disk version already matches.
|
||||
pub(crate) async fn migrate_on_open(root_uri: &str) -> Result<()> {
|
||||
let mut dataset = open_manifest_dataset(root_uri, None).await?;
|
||||
migrations::migrate_internal_schema(&mut dataset).await
|
||||
}
|
||||
|
||||
/// Immutable point-in-time view of the database.
|
||||
///
|
||||
/// Cheap to create (no storage I/O). All reads within a query go through one
|
||||
|
|
|
|||
|
|
@ -46,7 +46,11 @@ use crate::error::{OmniError, Result};
|
|||
/// - v2 — `__manifest.object_id` carries the unenforced-PK annotation,
|
||||
/// engaging Lance's bloom-filter conflict resolver at commit time. Added
|
||||
/// alongside `expected_table_versions` OCC on `ManifestBatchPublisher::publish`.
|
||||
pub(super) const INTERNAL_MANIFEST_SCHEMA_VERSION: u32 = 2;
|
||||
/// - v3 — one-time sweep of legacy `__run__<id>` staging branches left on the
|
||||
/// `__manifest` dataset by the pre-v0.4.0 Run state machine (removed in
|
||||
/// MR-771). Once swept, the `is_internal_run_branch` defense-in-depth guard
|
||||
/// is no longer needed (MR-770).
|
||||
pub(super) const INTERNAL_MANIFEST_SCHEMA_VERSION: u32 = 3;
|
||||
|
||||
const INTERNAL_SCHEMA_VERSION_KEY: &str = "omnigraph:internal_schema_version";
|
||||
const OBJECT_ID_PK_KEY: &str = "lance-schema:unenforced-primary-key";
|
||||
|
|
@ -89,6 +93,10 @@ pub(super) async fn migrate_internal_schema(dataset: &mut Dataset) -> Result<()>
|
|||
migrate_v1_to_v2(dataset).await?;
|
||||
current = 2;
|
||||
}
|
||||
2 => {
|
||||
migrate_v2_to_v3(dataset).await?;
|
||||
current = 3;
|
||||
}
|
||||
other => {
|
||||
return Err(OmniError::manifest_internal(format!(
|
||||
"no internal-schema migration registered for v{} → v{}",
|
||||
|
|
@ -122,6 +130,51 @@ async fn migrate_v1_to_v2(dataset: &mut Dataset) -> Result<()> {
|
|||
set_stamp(dataset, 2).await
|
||||
}
|
||||
|
||||
/// v2 → v3: sweep legacy `__run__<id>` staging branches off the `__manifest`
|
||||
/// dataset, then bump the stamp.
|
||||
///
|
||||
/// The pre-v0.4.0 Run state machine (removed in MR-771) created graph-level
|
||||
/// staging branches named `__run__<ulid>` on `__manifest`. MR-771 stopped
|
||||
/// creating them but left any pre-existing ones in place; Lance's
|
||||
/// `list_branches` still enumerates them, so they leak into `branch_list()`
|
||||
/// and count as blocking branches at schema-apply time. This one-time sweep
|
||||
/// removes them so the `is_internal_run_branch` guard can retire (MR-770).
|
||||
///
|
||||
/// The `"__run__"` prefix is inlined here on purpose: this migration must keep
|
||||
/// working after the `run_registry` module (the guard) is deleted, so it does
|
||||
/// not depend on it.
|
||||
///
|
||||
/// Idempotent under both sequential retry and concurrent runners: each run
|
||||
/// re-enumerates `list_branches` fresh, and `force_delete_branch` tolerates a
|
||||
/// branch that is already gone — so a crash before the stamp bump, or a second
|
||||
/// process opening the same legacy graph at the same time, never errors out.
|
||||
async fn migrate_v2_to_v3(dataset: &mut Dataset) -> Result<()> {
|
||||
const LEGACY_RUN_BRANCH_PREFIX: &str = "__run__";
|
||||
let branches = dataset
|
||||
.list_branches()
|
||||
.await
|
||||
.map_err(|e| OmniError::Lance(e.to_string()))?;
|
||||
let run_branches: Vec<String> = branches
|
||||
.into_keys()
|
||||
.filter(|name| {
|
||||
name.trim_start_matches('/')
|
||||
.starts_with(LEGACY_RUN_BRANCH_PREFIX)
|
||||
})
|
||||
.collect();
|
||||
for name in run_branches {
|
||||
// `force_delete_branch` deletes even when the `BranchContents` is
|
||||
// already gone. Plain `delete_branch` errors "BranchContents not
|
||||
// found", which would fail a second concurrent open (or a retry that
|
||||
// raced another runner) after the first one swept the branch. Force is
|
||||
// exactly Lance's documented path for cleaning up zombie branches.
|
||||
dataset
|
||||
.force_delete_branch(&name)
|
||||
.await
|
||||
.map_err(|e| OmniError::Lance(e.to_string()))?;
|
||||
}
|
||||
set_stamp(dataset, 3).await
|
||||
}
|
||||
|
||||
async fn set_stamp(dataset: &mut Dataset, version: u32) -> Result<()> {
|
||||
dataset
|
||||
.update_schema_metadata([(INTERNAL_SCHEMA_VERSION_KEY.to_string(), version.to_string())])
|
||||
|
|
|
|||
|
|
@ -1461,6 +1461,80 @@ async fn test_publish_migrates_pre_stamp_manifest_to_current_version() {
|
|||
assert!(reopened.snapshot().entry("node:Person").is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_v2_to_v3_sweeps_legacy_run_branches_on_write_open() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().to_str().unwrap();
|
||||
let catalog = build_test_catalog();
|
||||
let mut mc = ManifestCoordinator::init(uri, &catalog).await.unwrap();
|
||||
|
||||
// Synthesize a pre-MR-770 graph: several stale `__run__` staging branches
|
||||
// left on `__manifest` (a real legacy graph accumulates one per run), plus
|
||||
// a real user branch that must survive the sweep. Multiple run branches
|
||||
// exercise the migration's delete loop on a single reused dataset handle.
|
||||
mc.create_branch("__run__01J9LEGACY").await.unwrap();
|
||||
mc.create_branch("__run__01J9SECOND").await.unwrap();
|
||||
mc.create_branch("__run__01J9THIRD").await.unwrap();
|
||||
mc.create_branch("feature").await.unwrap();
|
||||
let before = mc.list_branches().await.unwrap();
|
||||
assert_eq!(
|
||||
before.iter().filter(|b| b.starts_with("__run__")).count(),
|
||||
3,
|
||||
"precondition: three legacy run branches exist on __manifest; got {before:?}",
|
||||
);
|
||||
|
||||
// Rewind the internal-schema stamp to v2 so the next write-open runs the
|
||||
// v2 → v3 sweep arm (init stamps at the current version, which is past it).
|
||||
{
|
||||
let mut ds = open_manifest_dataset(uri, None).await.unwrap();
|
||||
ds.update_schema_metadata([(
|
||||
"omnigraph:internal_schema_version".to_string(),
|
||||
Some("2".to_string()),
|
||||
)])
|
||||
.await
|
||||
.unwrap();
|
||||
let post = open_manifest_dataset(uri, None).await.unwrap();
|
||||
assert_eq!(super::migrations::read_stamp(&post), 2, "stamp rewound to v2");
|
||||
}
|
||||
|
||||
// A no-op publish forces the open-for-write path, which runs the migration.
|
||||
let mut expected = HashMap::new();
|
||||
expected.insert("node:Person".to_string(), 1);
|
||||
GraphNamespacePublisher::new(uri, None)
|
||||
.publish(&[], &expected)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Stamp advanced to current; the legacy run branch is physically gone from
|
||||
// `__manifest` (checked via the raw, unfiltered manifest list — not the
|
||||
// guard-filtered `branch_list`), and the real branch + `main` survive.
|
||||
let post = open_manifest_dataset(uri, None).await.unwrap();
|
||||
assert_eq!(
|
||||
super::migrations::read_stamp(&post),
|
||||
super::migrations::INTERNAL_MANIFEST_SCHEMA_VERSION,
|
||||
);
|
||||
let reopened = ManifestCoordinator::open(uri).await.unwrap();
|
||||
let after = reopened.list_branches().await.unwrap();
|
||||
assert!(
|
||||
!after.iter().any(|b| b.starts_with("__run__")),
|
||||
"legacy run branch must be swept; got {after:?}",
|
||||
);
|
||||
assert!(after.iter().any(|b| b == "feature"), "user branch must survive");
|
||||
assert!(after.iter().any(|b| b == "main"), "main must survive");
|
||||
|
||||
// Idempotent: a second write-open finds the stamp at current and does not
|
||||
// re-run the sweep or error.
|
||||
GraphNamespacePublisher::new(uri, None)
|
||||
.publish(&[], &expected)
|
||||
.await
|
||||
.unwrap();
|
||||
let final_ds = open_manifest_dataset(uri, None).await.unwrap();
|
||||
assert_eq!(
|
||||
super::migrations::read_stamp(&final_ds),
|
||||
super::migrations::INTERNAL_MANIFEST_SCHEMA_VERSION,
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_publish_rejects_manifest_stamped_at_future_version() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ pub mod graph_coordinator;
|
|||
pub mod manifest;
|
||||
mod omnigraph;
|
||||
mod recovery_audit;
|
||||
mod run_registry;
|
||||
mod schema_state;
|
||||
pub(crate) mod write_queue;
|
||||
|
||||
|
|
@ -13,9 +12,8 @@ pub use manifest::{Snapshot, SubTableEntry, SubTableUpdate};
|
|||
pub(crate) use omnigraph::ensure_public_branch_ref;
|
||||
pub use omnigraph::{
|
||||
CleanupPolicyOptions, InitOptions, MergeOutcome, Omnigraph, OpenMode, SchemaApplyOptions,
|
||||
SchemaApplyResult, TableCleanupStats, TableOptimizeStats,
|
||||
SchemaApplyResult, SkipReason, TableCleanupStats, TableOptimizeStats,
|
||||
};
|
||||
pub(crate) use run_registry::is_internal_run_branch;
|
||||
|
||||
pub(crate) const SCHEMA_APPLY_LOCK_BRANCH: &str = "__schema_apply_lock__";
|
||||
|
||||
|
|
@ -69,5 +67,8 @@ pub(crate) fn is_schema_apply_lock_branch(name: &str) -> bool {
|
|||
}
|
||||
|
||||
pub(crate) fn is_internal_system_branch(name: &str) -> bool {
|
||||
is_internal_run_branch(name) || is_schema_apply_lock_branch(name)
|
||||
// Legacy `__run__*` staging branches (Run state machine, removed MR-771)
|
||||
// are swept off `__manifest` by the v2→v3 internal-schema migration, so the
|
||||
// only internal branch the engine still creates is the schema-apply lock.
|
||||
is_schema_apply_lock_branch(name)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ mod optimize;
|
|||
mod schema_apply;
|
||||
mod table_ops;
|
||||
|
||||
pub use optimize::{CleanupPolicyOptions, TableCleanupStats, TableOptimizeStats};
|
||||
pub use optimize::{CleanupPolicyOptions, SkipReason, TableCleanupStats, TableOptimizeStats};
|
||||
pub use schema_apply::SchemaApplyOptions;
|
||||
|
||||
use super::commit_graph::GraphCommit;
|
||||
|
|
@ -67,6 +67,12 @@ pub struct SchemaApplyResult {
|
|||
pub steps: Vec<SchemaMigrationStep>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SchemaApplyPreview {
|
||||
pub plan: SchemaMigrationPlan,
|
||||
pub catalog: Catalog,
|
||||
}
|
||||
|
||||
/// Top-level handle to an Omnigraph database.
|
||||
///
|
||||
/// An Omnigraph is a Lance-native graph database with git-style branching.
|
||||
|
|
@ -340,6 +346,16 @@ impl Omnigraph {
|
|||
mode: OpenMode,
|
||||
) -> Result<Self> {
|
||||
let root = normalize_root_uri(uri)?;
|
||||
// Apply pending internal-schema migrations before the coordinator reads
|
||||
// branch state, so `branch_list` and the schema-apply blocking-branch
|
||||
// checks observe the post-migration graph — notably the v2→v3 sweep of
|
||||
// legacy `__run__*` staging branches (MR-770). ReadWrite only: a
|
||||
// read-only open must not trigger object-store writes, so a read-only
|
||||
// open of an unmigrated legacy graph still lists `__run__*` until its
|
||||
// first read-write open (an accepted, documented limitation).
|
||||
if matches!(mode, OpenMode::ReadWrite) {
|
||||
crate::db::manifest::migrate_on_open(&root).await?;
|
||||
}
|
||||
// Open the coordinator first so the schema-staging recovery sweep can
|
||||
// compare its snapshot against any leftover staging files.
|
||||
let mut coordinator = GraphCoordinator::open(&root, Arc::clone(&storage)).await?;
|
||||
|
|
@ -493,6 +509,14 @@ impl Omnigraph {
|
|||
schema_apply::plan_schema(self, desired_schema_source, options).await
|
||||
}
|
||||
|
||||
pub async fn preview_schema_apply_with_options(
|
||||
&self,
|
||||
desired_schema_source: &str,
|
||||
options: SchemaApplyOptions,
|
||||
) -> Result<SchemaApplyPreview> {
|
||||
schema_apply::preview_schema_apply(self, desired_schema_source, options).await
|
||||
}
|
||||
|
||||
pub async fn apply_schema(&self, desired_schema_source: &str) -> Result<SchemaApplyResult> {
|
||||
self.apply_schema_as(desired_schema_source, SchemaApplyOptions::default(), None)
|
||||
.await
|
||||
|
|
@ -523,7 +547,28 @@ impl Omnigraph {
|
|||
options: SchemaApplyOptions,
|
||||
actor: Option<&str>,
|
||||
) -> Result<SchemaApplyResult> {
|
||||
schema_apply::apply_schema(self, desired_schema_source, options, actor).await
|
||||
self.apply_schema_as_with_catalog_check(desired_schema_source, options, actor, |_| Ok(()))
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn apply_schema_as_with_catalog_check<F>(
|
||||
&self,
|
||||
desired_schema_source: &str,
|
||||
options: SchemaApplyOptions,
|
||||
actor: Option<&str>,
|
||||
validate_catalog: F,
|
||||
) -> Result<SchemaApplyResult>
|
||||
where
|
||||
F: FnOnce(&Catalog) -> Result<()>,
|
||||
{
|
||||
schema_apply::apply_schema(
|
||||
self,
|
||||
desired_schema_source,
|
||||
options,
|
||||
actor,
|
||||
validate_catalog,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn ensure_schema_apply_idle(&self, operation: &str) -> Result<()> {
|
||||
|
|
@ -1058,11 +1103,14 @@ impl Omnigraph {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
async fn cleanup_deleted_branch_tables(
|
||||
&self,
|
||||
branch: &str,
|
||||
owned_tables: &[(String, String)],
|
||||
) -> Result<()> {
|
||||
/// Best-effort reclaim of the per-table Lance forks a just-deleted branch
|
||||
/// owned. Runs AFTER the manifest authority flip, so the branch is already
|
||||
/// gone and these forks are unreachable orphans. A failure here (transient
|
||||
/// object-store error, the `branch_delete.before_table_cleanup` failpoint)
|
||||
/// is logged and swallowed: the `cleanup` reconciler is the guaranteed
|
||||
/// backstop that converges any leftover orphan. Uses `force_delete_branch`
|
||||
/// so a partially-reclaimed retry is idempotent.
|
||||
async fn cleanup_deleted_branch_tables(&self, branch: &str, owned_tables: &[(String, String)]) {
|
||||
let mut seen_paths = HashSet::new();
|
||||
let mut cleanup_targets = owned_tables
|
||||
.iter()
|
||||
|
|
@ -1073,15 +1121,21 @@ impl Omnigraph {
|
|||
|
||||
for (table_key, table_path) in cleanup_targets {
|
||||
let dataset_uri = self.table_store.dataset_uri(&table_path);
|
||||
if let Err(err) = self.table_store.delete_branch(&dataset_uri, branch).await {
|
||||
return Err(OmniError::manifest_internal(format!(
|
||||
"branch '{}' was deleted but cleanup failed for {}: {}",
|
||||
branch, table_key, err
|
||||
)));
|
||||
let outcome = match crate::failpoints::maybe_fail("branch_delete.before_table_cleanup")
|
||||
{
|
||||
Ok(()) => self.table_store.force_delete_branch(&dataset_uri, branch).await,
|
||||
Err(injected) => Err(injected),
|
||||
};
|
||||
if let Err(err) = outcome {
|
||||
tracing::warn!(
|
||||
target: "omnigraph::branch_delete::cleanup",
|
||||
branch = %branch,
|
||||
table = %table_key,
|
||||
error = %err,
|
||||
"best-effort fork reclaim failed; cleanup will reconcile the orphan",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete_branch_storage_only(&self, branch: &str) -> Result<()> {
|
||||
|
|
@ -1105,9 +1159,12 @@ impl Omnigraph {
|
|||
.map(|entry| (entry.table_key.clone(), entry.table_path.clone()))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
// Authority flip (+ best-effort commit-graph reclaim) — must succeed.
|
||||
self.coordinator.write().await.branch_delete(branch).await?;
|
||||
// Best-effort per-table fork reclaim; cleanup reconciles any leftover.
|
||||
self.cleanup_deleted_branch_tables(branch, &owned_tables)
|
||||
.await
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_branch_name(branch: &str) -> Result<Option<String>> {
|
||||
|
|
@ -1444,12 +1501,6 @@ pub(crate) fn normalize_branch_name(branch: &str) -> Result<Option<String>> {
|
|||
}
|
||||
|
||||
pub(crate) fn ensure_public_branch_ref(branch: &str, operation: &str) -> Result<()> {
|
||||
if super::is_internal_run_branch(branch) {
|
||||
return Err(OmniError::manifest(format!(
|
||||
"{} does not allow internal run ref '{}'",
|
||||
operation, branch
|
||||
)));
|
||||
}
|
||||
if is_internal_system_branch(branch) {
|
||||
return Err(OmniError::manifest(format!(
|
||||
"{} does not allow internal system ref '{}'",
|
||||
|
|
@ -1853,7 +1904,6 @@ fn json_value_from_array(array: &dyn Array, row: usize) -> Result<serde_json::Va
|
|||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::db::is_internal_run_branch;
|
||||
use crate::db::manifest::ManifestCoordinator;
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
|
@ -2191,11 +2241,11 @@ edge WorksAt: Person -> Company
|
|||
#[tokio::test]
|
||||
async fn test_apply_schema_succeeds_after_load() {
|
||||
// Historical: schema apply used to be blocked by leftover
|
||||
// `__run__` branches. A defense-in-depth filter now skips
|
||||
// internal system branches, and run branches were made
|
||||
// ephemeral on every terminal state — so in practice no
|
||||
// `__run__` branch survives publish. The filter still guards
|
||||
// the invariant.
|
||||
// `__run__` branches. The Run state machine was removed in
|
||||
// MR-771, so a fresh graph never creates a `__run__` branch;
|
||||
// legacy ones are swept by the v2→v3 manifest migration. This
|
||||
// asserts the invariant a current graph upholds: publish leaves
|
||||
// no `__run__` branch behind, so schema apply proceeds.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().to_str().unwrap();
|
||||
let mut db = Omnigraph::init(uri, TEST_SCHEMA).await.unwrap();
|
||||
|
|
@ -2210,8 +2260,8 @@ edge WorksAt: Person -> Company
|
|||
|
||||
let all_branches = db.coordinator.read().await.all_branches().await.unwrap();
|
||||
assert!(
|
||||
!all_branches.iter().any(|b| is_internal_run_branch(b)),
|
||||
"run branch should be deleted after publish, got: {:?}",
|
||||
!all_branches.iter().any(|b| b.starts_with("__run__")),
|
||||
"no __run__ branch should exist after publish, got: {:?}",
|
||||
all_branches
|
||||
);
|
||||
|
||||
|
|
@ -2223,6 +2273,56 @@ edge WorksAt: Person -> Company
|
|||
assert!(result.applied, "schema apply should have applied");
|
||||
}
|
||||
|
||||
/// Regression (MR-770): a pre-v0.4.0 graph that still carries a stale
|
||||
/// `__run__*` branch on `__manifest` must not block schema apply. The
|
||||
/// v2→v3 sweep runs in `Omnigraph::open(ReadWrite)` — before the
|
||||
/// schema-apply blocking-branch check — so apply succeeds with no
|
||||
/// intervening publish.
|
||||
///
|
||||
/// Confirmed to fail before the open-time migration landed: the reopened
|
||||
/// graph still listed `__run__legacy`, and `apply_schema` returned
|
||||
/// "found non-main branches: __run__legacy".
|
||||
#[tokio::test]
|
||||
async fn legacy_run_branch_is_swept_on_open_and_does_not_block_schema_apply() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().to_str().unwrap();
|
||||
let mut db = Omnigraph::init(uri, TEST_SCHEMA).await.unwrap();
|
||||
|
||||
// Synthesize a legacy graph: a stale `__run__` branch on `__manifest`
|
||||
// plus the manifest stamp rewound to v2 (pre-sweep).
|
||||
db.branch_create("__run__legacy").await.unwrap();
|
||||
drop(db);
|
||||
{
|
||||
let mut ds = lance::Dataset::open(&format!("{}/__manifest", uri))
|
||||
.await
|
||||
.unwrap();
|
||||
ds.update_schema_metadata([(
|
||||
"omnigraph:internal_schema_version".to_string(),
|
||||
Some("2".to_string()),
|
||||
)])
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Reopen (ReadWrite): the open-time migration must sweep `__run__legacy`
|
||||
// before any branch-observing code runs.
|
||||
let db = Omnigraph::open(uri).await.unwrap();
|
||||
let branches = db.branch_list().await.unwrap();
|
||||
assert!(
|
||||
!branches.iter().any(|b| b.starts_with("__run__")),
|
||||
"open-time migration must sweep legacy __run__ branches; got {branches:?}",
|
||||
);
|
||||
|
||||
// Schema apply must proceed with no intervening publish — the
|
||||
// blocking-branch check no longer sees `__run__legacy`.
|
||||
let desired = TEST_SCHEMA.replace(
|
||||
" age: I32?\n}",
|
||||
" age: I32?\n nickname: String?\n}",
|
||||
);
|
||||
let result = db.apply_schema(&desired).await.unwrap();
|
||||
assert!(result.applied, "schema apply should have applied");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_apply_schema_adds_index_for_existing_property() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
|
|
|||
|
|
@ -40,6 +40,20 @@ fn maint_concurrency() -> usize {
|
|||
.unwrap_or(DEFAULT_MAINT_CONCURRENCY)
|
||||
}
|
||||
|
||||
/// Whether the installed Lance can compact a dataset that contains blob
|
||||
/// columns. `false` today: Lance `compact_files` forces
|
||||
/// `BlobHandling::AllBinary` on the read side, and the blob-v2 struct decoder
|
||||
/// mis-counts columns ("there were more fields in the schema than provided
|
||||
/// column indices"), failing even a pristine uniform-V2_2 multi-fragment blob
|
||||
/// table. Reads are unaffected (queries use descriptor handling).
|
||||
///
|
||||
/// While `false`, [`optimize_all_tables`] skips blob-bearing tables and reports
|
||||
/// [`SkipReason::BlobColumnsUnsupportedByLance`] instead of aborting the whole
|
||||
/// sweep. Flip to `true` once the upstream Lance fix ships — the
|
||||
/// `lance_surface_guards.rs::compact_files_still_fails_on_blob_columns` guard
|
||||
/// turns red on that bump and forces this flip. Tracked in `docs/dev/lance.md`.
|
||||
const LANCE_SUPPORTS_BLOB_COMPACTION: bool = false;
|
||||
|
||||
/// Retention knobs for [`cleanup_all_tables`]. At least one must be set or
|
||||
/// nothing is cleaned. If both are set, Lance applies them as AND (a manifest
|
||||
/// is kept if it satisfies either — i.e. only manifests older than BOTH the
|
||||
|
|
@ -52,8 +66,45 @@ pub struct CleanupPolicyOptions {
|
|||
pub older_than: Option<Duration>,
|
||||
}
|
||||
|
||||
/// Per-table outcome of `optimize_all_tables`.
|
||||
/// Why `optimize` did not compact a table. Typed so callers branch on the
|
||||
/// reason rather than sniffing a string. One variant today, gated by
|
||||
/// [`LANCE_SUPPORTS_BLOB_COMPACTION`].
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[non_exhaustive]
|
||||
pub enum SkipReason {
|
||||
/// The table has one or more `Blob` columns. Lance `compact_files` forces
|
||||
/// `BlobHandling::AllBinary`, which mis-decodes blob-v2 columns; see
|
||||
/// [`LANCE_SUPPORTS_BLOB_COMPACTION`] and `docs/dev/lance.md`.
|
||||
BlobColumnsUnsupportedByLance,
|
||||
}
|
||||
|
||||
impl SkipReason {
|
||||
/// Stable machine-readable token for serialized output (e.g. CLI `--json`).
|
||||
/// Once emitted this is part of the output contract — keep it stable.
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
SkipReason::BlobColumnsUnsupportedByLance => "blob_columns_unsupported_by_lance",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for SkipReason {
|
||||
/// Human-readable reason for CLI and log output.
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
let msg = match self {
|
||||
SkipReason::BlobColumnsUnsupportedByLance => {
|
||||
"blob columns — Lance compaction unsupported"
|
||||
}
|
||||
};
|
||||
f.write_str(msg)
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-table outcome of `optimize_all_tables`. This is a returned result type,
|
||||
/// not built by callers, so it is `#[non_exhaustive]`: future fields stay
|
||||
/// non-breaking and downstream code reads fields rather than constructing it.
|
||||
#[derive(Debug, Clone)]
|
||||
#[non_exhaustive]
|
||||
pub struct TableOptimizeStats {
|
||||
pub table_key: String,
|
||||
/// Number of source fragments that were rewritten by Lance.
|
||||
|
|
@ -62,14 +113,44 @@ pub struct TableOptimizeStats {
|
|||
pub fragments_added: usize,
|
||||
/// Did this table get a new Lance manifest version from the compaction?
|
||||
pub committed: bool,
|
||||
/// `Some(reason)` if this table was deliberately not compacted. When set,
|
||||
/// `fragments_removed == 0`, `fragments_added == 0`, and `!committed`.
|
||||
pub skipped: Option<SkipReason>,
|
||||
}
|
||||
|
||||
/// Per-table outcome of `cleanup_all_tables`.
|
||||
impl TableOptimizeStats {
|
||||
/// Stat for a table that Lance actually compacted.
|
||||
fn compacted(table_key: String, metrics: &CompactionMetrics, committed: bool) -> Self {
|
||||
Self {
|
||||
table_key,
|
||||
fragments_removed: metrics.fragments_removed,
|
||||
fragments_added: metrics.fragments_added,
|
||||
committed,
|
||||
skipped: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Stat for a table that was deliberately skipped (compaction not attempted).
|
||||
fn skipped(table_key: String, reason: SkipReason) -> Self {
|
||||
Self {
|
||||
table_key,
|
||||
fragments_removed: 0,
|
||||
fragments_added: 0,
|
||||
committed: false,
|
||||
skipped: Some(reason),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-table outcome of `cleanup_all_tables`. `error` is `Some` when this
|
||||
/// table's version GC failed; cleanup is fault-isolated per table, so a single
|
||||
/// table's failure is recorded here rather than aborting the whole sweep.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TableCleanupStats {
|
||||
pub table_key: String,
|
||||
pub bytes_removed: u64,
|
||||
pub old_versions_removed: u64,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// Run Lance `compact_files` on every node + edge table on `main`.
|
||||
|
|
@ -81,14 +162,21 @@ pub async fn optimize_all_tables(db: &Omnigraph) -> Result<Vec<TableOptimizeStat
|
|||
let resolved = db.resolved_branch_target(None).await?;
|
||||
let snapshot = resolved.snapshot;
|
||||
|
||||
let table_tasks: Vec<_> = all_table_keys(&db.catalog())
|
||||
.into_iter()
|
||||
.filter_map(|table_key| {
|
||||
let entry = snapshot.entry(&table_key)?;
|
||||
// Compute per-table state (path + whether it has blob columns) up front, in
|
||||
// a scope that drops the catalog handle before the async stream starts.
|
||||
let table_tasks: Vec<(String, String, bool)> = {
|
||||
let catalog = db.catalog();
|
||||
let mut tasks = Vec::new();
|
||||
for table_key in all_table_keys(&catalog) {
|
||||
let Some(entry) = snapshot.entry(&table_key) else {
|
||||
continue;
|
||||
};
|
||||
let full_path = format!("{}/{}", db.root_uri, entry.table_path);
|
||||
Some((table_key, full_path))
|
||||
})
|
||||
.collect();
|
||||
let has_blob = !blob_properties_for_table_key(&catalog, &table_key)?.is_empty();
|
||||
tasks.push((table_key, full_path, has_blob));
|
||||
}
|
||||
tasks
|
||||
};
|
||||
|
||||
if table_tasks.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
|
|
@ -98,7 +186,24 @@ pub async fn optimize_all_tables(db: &Omnigraph) -> Result<Vec<TableOptimizeStat
|
|||
let table_store = &db.table_store;
|
||||
|
||||
let stats: Vec<Result<TableOptimizeStats>> = futures::stream::iter(table_tasks.into_iter())
|
||||
.map(|(table_key, full_path)| async move {
|
||||
.map(|(table_key, full_path, has_blob)| async move {
|
||||
// Lance `compact_files` mis-decodes blob-v2 columns under the forced
|
||||
// `BlobHandling::AllBinary` read (see LANCE_SUPPORTS_BLOB_COMPACTION).
|
||||
// Skip blob-bearing tables and report it rather than aborting the
|
||||
// whole sweep — the other tables still compact.
|
||||
if has_blob && !LANCE_SUPPORTS_BLOB_COMPACTION {
|
||||
tracing::warn!(
|
||||
target: "omnigraph::optimize",
|
||||
table = %table_key,
|
||||
"skipping compaction: table has blob columns the current Lance \
|
||||
cannot rewrite (blob-v2 AllBinary decode bug); other tables \
|
||||
unaffected — rerun after the Lance fix",
|
||||
);
|
||||
return Ok(TableOptimizeStats::skipped(
|
||||
table_key,
|
||||
SkipReason::BlobColumnsUnsupportedByLance,
|
||||
));
|
||||
}
|
||||
let mut ds = table_store
|
||||
.open_dataset_head_for_write(&table_key, &full_path, None)
|
||||
.await?;
|
||||
|
|
@ -108,12 +213,11 @@ pub async fn optimize_all_tables(db: &Omnigraph) -> Result<Vec<TableOptimizeStat
|
|||
.await
|
||||
.map_err(|e| OmniError::Lance(e.to_string()))?;
|
||||
let version_after = ds.version().version;
|
||||
Ok(TableOptimizeStats {
|
||||
Ok(TableOptimizeStats::compacted(
|
||||
table_key,
|
||||
fragments_removed: metrics.fragments_removed,
|
||||
fragments_added: metrics.fragments_added,
|
||||
committed: version_after != version_before,
|
||||
})
|
||||
&metrics,
|
||||
version_after != version_before,
|
||||
))
|
||||
})
|
||||
.buffer_unordered(concurrency)
|
||||
.collect()
|
||||
|
|
@ -138,6 +242,26 @@ pub async fn cleanup_all_tables(
|
|||
db.ensure_schema_state_valid().await?;
|
||||
db.ensure_schema_apply_idle("cleanup").await?;
|
||||
|
||||
// Reclaim orphaned branch forks (from an incomplete prior `branch_delete`)
|
||||
// before version GC. Authority-derived and idempotent; the eager
|
||||
// best-effort reclaim in `branch_delete` covers the common case, this is
|
||||
// the guaranteed backstop. Logged for observability.
|
||||
let reconciled = reconcile_orphaned_branches(db).await?;
|
||||
if !reconciled.reclaimed.is_empty() {
|
||||
tracing::info!(
|
||||
count = reconciled.reclaimed.len(),
|
||||
reclaimed = ?reconciled.reclaimed,
|
||||
"cleanup reconciled orphaned branch forks"
|
||||
);
|
||||
}
|
||||
if !reconciled.failures.is_empty() {
|
||||
tracing::warn!(
|
||||
count = reconciled.failures.len(),
|
||||
failures = ?reconciled.failures,
|
||||
"cleanup could not reconcile some orphaned forks; will retry next cleanup"
|
||||
);
|
||||
}
|
||||
|
||||
let before_timestamp = options.older_than.map(|d| Utc::now() - d);
|
||||
let keep_versions = options.keep_versions;
|
||||
|
||||
|
|
@ -160,36 +284,205 @@ pub async fn cleanup_all_tables(
|
|||
let concurrency = maint_concurrency().min(table_tasks.len()).max(1);
|
||||
let table_store = &db.table_store;
|
||||
|
||||
let results: Vec<Result<TableCleanupStats>> = futures::stream::iter(table_tasks.into_iter())
|
||||
// Fault-isolated per table: a single table's GC failure is recorded on its
|
||||
// stats row (`error: Some`) and logged, never aborting the healthy tables.
|
||||
// cleanup is the convergence backstop, so it must do as much as it can and
|
||||
// converge on re-run rather than fail wholesale (invariant 13).
|
||||
let results: Vec<TableCleanupStats> = futures::stream::iter(table_tasks.into_iter())
|
||||
.map(|(table_key, full_path)| async move {
|
||||
let ds = table_store
|
||||
.open_dataset_head_for_write(&table_key, &full_path, None)
|
||||
.await?;
|
||||
let before_version = keep_versions
|
||||
.map(|n| ds.version().version.saturating_sub(n as u64))
|
||||
.filter(|v| *v > 0);
|
||||
let policy = CleanupPolicy {
|
||||
before_timestamp,
|
||||
before_version,
|
||||
delete_unverified: false,
|
||||
error_if_tagged_old_versions: false,
|
||||
clean_referenced_branches: false,
|
||||
delete_rate_limit: None,
|
||||
};
|
||||
let removed: RemovalStats = lance::dataset::cleanup::cleanup_old_versions(&ds, policy)
|
||||
.await
|
||||
.map_err(|e| OmniError::Lance(e.to_string()))?;
|
||||
Ok(TableCleanupStats {
|
||||
table_key,
|
||||
bytes_removed: removed.bytes_removed,
|
||||
old_versions_removed: removed.old_versions,
|
||||
})
|
||||
let outcome: Result<RemovalStats> = async {
|
||||
crate::failpoints::maybe_fail("cleanup.table_gc")?;
|
||||
let ds = table_store
|
||||
.open_dataset_head_for_write(&table_key, &full_path, None)
|
||||
.await?;
|
||||
let before_version = keep_versions
|
||||
.map(|n| ds.version().version.saturating_sub(n as u64))
|
||||
.filter(|v| *v > 0);
|
||||
let policy = CleanupPolicy {
|
||||
before_timestamp,
|
||||
before_version,
|
||||
delete_unverified: false,
|
||||
error_if_tagged_old_versions: false,
|
||||
clean_referenced_branches: false,
|
||||
delete_rate_limit: None,
|
||||
};
|
||||
lance::dataset::cleanup::cleanup_old_versions(&ds, policy)
|
||||
.await
|
||||
.map_err(|e| OmniError::Lance(e.to_string()))
|
||||
}
|
||||
.await;
|
||||
match outcome {
|
||||
Ok(removed) => TableCleanupStats {
|
||||
table_key,
|
||||
bytes_removed: removed.bytes_removed,
|
||||
old_versions_removed: removed.old_versions,
|
||||
error: None,
|
||||
},
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
target: "omnigraph::cleanup",
|
||||
table = %table_key,
|
||||
error = %err,
|
||||
"version GC failed for table; other tables unaffected",
|
||||
);
|
||||
TableCleanupStats {
|
||||
table_key,
|
||||
bytes_removed: 0,
|
||||
old_versions_removed: 0,
|
||||
error: Some(err.to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
.buffer_unordered(concurrency)
|
||||
.collect()
|
||||
.await;
|
||||
|
||||
results.into_iter().collect()
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Outcome of [`reconcile_orphaned_branches`]: the `(owner, branch)` pairs
|
||||
/// reclaimed and the `(owner, error)` pairs that failed, where `owner` is a
|
||||
/// table key (e.g. `node:Person`) or `"_graph_commits"`. Per-owner failures are
|
||||
/// isolated and recorded here, not propagated — the next reconcile converges.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct BranchReconcileStats {
|
||||
pub reclaimed: Vec<(String, String)>,
|
||||
pub failures: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
/// Drop every per-table and commit-graph Lance branch that the manifest no
|
||||
/// longer references.
|
||||
///
|
||||
/// Orphaned forks arise when a `branch_delete` flips the manifest authority
|
||||
/// (atomic) but a downstream best-effort reclaim does not complete. They are
|
||||
/// unreachable through any snapshot — no manifest entry can name them — yet
|
||||
/// they pin their `tree/{branch}/` storage and can block reusing the branch
|
||||
/// name. This is the guaranteed convergence backstop: it is idempotent and
|
||||
/// derived purely from the manifest authority, so it no-ops once everything is
|
||||
/// reconciled, and it would harmlessly find nothing if a future Lance atomic
|
||||
/// multi-dataset branch op prevented orphans from forming.
|
||||
///
|
||||
/// The keep-set is the full (unfiltered) manifest branch list, so system
|
||||
/// branches' forks are never reclaimed; `main`/default is not a named Lance
|
||||
/// branch and so is never a candidate. Referencing children are dropped before
|
||||
/// parents (Lance refuses to delete a referenced parent) by ordering longest
|
||||
/// branch names first.
|
||||
pub async fn reconcile_orphaned_branches(db: &Omnigraph) -> Result<BranchReconcileStats> {
|
||||
use std::collections::HashSet;
|
||||
|
||||
let keep: HashSet<String> = db
|
||||
.coordinator
|
||||
.read()
|
||||
.await
|
||||
.all_branches()
|
||||
.await?
|
||||
.into_iter()
|
||||
.collect();
|
||||
|
||||
let resolved = db.resolved_branch_target(None).await?;
|
||||
let snapshot = resolved.snapshot;
|
||||
let table_targets: Vec<(String, String)> = all_table_keys(&db.catalog())
|
||||
.into_iter()
|
||||
.filter_map(|table_key| {
|
||||
let entry = snapshot.entry(&table_key)?;
|
||||
let full_path = format!("{}/{}", db.root_uri, entry.table_path);
|
||||
Some((table_key, full_path))
|
||||
})
|
||||
.collect();
|
||||
|
||||
let mut stats = BranchReconcileStats::default();
|
||||
|
||||
// Per-table fault isolation: one table's transient failure is recorded and
|
||||
// logged, never aborting the rest of the sweep.
|
||||
for (table_key, full_path) in table_targets {
|
||||
let listed = match db.table_store.list_branches(&full_path).await {
|
||||
Ok(listed) => listed,
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
target: "omnigraph::cleanup",
|
||||
table = %table_key,
|
||||
error = %err,
|
||||
"listing branches failed during reconcile; skipping table",
|
||||
);
|
||||
stats.failures.push((table_key.clone(), err.to_string()));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
for branch in orphan_branches(listed, &keep) {
|
||||
let outcome = match crate::failpoints::maybe_fail("cleanup.reconcile_fork") {
|
||||
Ok(()) => db.table_store.force_delete_branch(&full_path, &branch).await,
|
||||
Err(injected) => Err(injected),
|
||||
};
|
||||
match outcome {
|
||||
Ok(()) => stats.reclaimed.push((table_key.clone(), branch)),
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
target: "omnigraph::cleanup",
|
||||
table = %table_key,
|
||||
branch = %branch,
|
||||
error = %err,
|
||||
"reclaiming orphaned fork failed; will retry next cleanup",
|
||||
);
|
||||
stats.failures.push((table_key.clone(), err.to_string()));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Commit-graph orphans (best-effort: the dataset may not exist on a graph
|
||||
// that has never committed; any failure is isolated and retried next time).
|
||||
if let Err(err) = reconcile_commit_graph_orphans(db, &keep, &mut stats).await {
|
||||
tracing::warn!(
|
||||
target: "omnigraph::cleanup",
|
||||
error = %err,
|
||||
"commit-graph orphan reconcile failed; will retry next cleanup",
|
||||
);
|
||||
stats.failures.push(("_graph_commits".to_string(), err.to_string()));
|
||||
}
|
||||
|
||||
Ok(stats)
|
||||
}
|
||||
|
||||
/// Commit-graph half of [`reconcile_orphaned_branches`], split out so its
|
||||
/// errors can be isolated. Returns `Ok` when the commit-graph dataset is absent.
|
||||
async fn reconcile_commit_graph_orphans(
|
||||
db: &Omnigraph,
|
||||
keep: &std::collections::HashSet<String>,
|
||||
stats: &mut BranchReconcileStats,
|
||||
) -> Result<()> {
|
||||
let commits_uri = crate::db::commit_graph::graph_commits_uri(db.root_uri());
|
||||
if !db.storage_adapter().exists(&commits_uri).await? {
|
||||
return Ok(());
|
||||
}
|
||||
let mut commit_graph = crate::db::commit_graph::CommitGraph::open(db.root_uri()).await?;
|
||||
for branch in orphan_branches(commit_graph.list_branches().await?, keep) {
|
||||
match commit_graph.force_delete_branch(&branch).await {
|
||||
Ok(()) => stats.reclaimed.push(("_graph_commits".to_string(), branch)),
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
target: "omnigraph::cleanup",
|
||||
branch = %branch,
|
||||
error = %err,
|
||||
"reclaiming orphaned commit-graph branch failed; will retry next cleanup",
|
||||
);
|
||||
stats.failures.push(("_graph_commits".to_string(), err.to_string()));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Filter `present` Lance branches down to those absent from the manifest
|
||||
/// `keep` set, ordered children-before-parents (longest name first) so Lance's
|
||||
/// referenced-parent `RefConflict` cannot block reclamation.
|
||||
fn orphan_branches(present: Vec<String>, keep: &std::collections::HashSet<String>) -> Vec<String> {
|
||||
let mut orphans: Vec<String> = present
|
||||
.into_iter()
|
||||
.filter(|branch| !keep.contains(branch))
|
||||
.collect();
|
||||
orphans.sort_by(|a, b| b.len().cmp(&a.len()).then_with(|| a.cmp(b)));
|
||||
orphans
|
||||
}
|
||||
|
||||
fn all_table_keys(catalog: &omnigraph_compiler::catalog::Catalog) -> Vec<String> {
|
||||
|
|
|
|||
|
|
@ -48,57 +48,24 @@ pub(super) async fn plan_schema(
|
|||
Ok(plan)
|
||||
}
|
||||
|
||||
pub(super) async fn apply_schema(
|
||||
db: &Omnigraph,
|
||||
desired_schema_source: &str,
|
||||
options: SchemaApplyOptions,
|
||||
actor: Option<&str>,
|
||||
) -> Result<SchemaApplyResult> {
|
||||
// Engine-layer policy gate (MR-722 chassis core).
|
||||
//
|
||||
// Fires BEFORE acquiring the schema-apply lock or doing any other
|
||||
// work. When no PolicyChecker is installed this is a no-op and
|
||||
// the apply path behaves exactly as it did before MR-722. When
|
||||
// a PolicyChecker IS installed and the actor is None, this is a
|
||||
// hard error — see Omnigraph::enforce's docstring for the
|
||||
// forget-the-actor-footgun reasoning.
|
||||
//
|
||||
// Scope is TargetBranch("main") to match the HTTP-layer convention
|
||||
// for SchemaApply: branch=None, target_branch=Some("main"). Cedar
|
||||
// policies in the wild use `target_branch_scope: protected` to
|
||||
// gate schema applies, so the engine-layer call has to set the
|
||||
// target_branch shape that activates that predicate. Wrong scope
|
||||
// here = silent policy mismatch with HTTP. See
|
||||
// `omnigraph_policy::ResourceScope::to_branch_pair` for the mapping.
|
||||
db.enforce(
|
||||
omnigraph_policy::PolicyAction::SchemaApply,
|
||||
&omnigraph_policy::ResourceScope::TargetBranch("main".to_string()),
|
||||
actor,
|
||||
)?;
|
||||
|
||||
acquire_schema_apply_lock(db).await?;
|
||||
let result = apply_schema_with_lock(db, desired_schema_source, options).await;
|
||||
let release_result = release_schema_apply_lock(db).await;
|
||||
match (result, release_result) {
|
||||
(Ok(result), Ok(())) => Ok(result),
|
||||
(Ok(_), Err(err)) => Err(err),
|
||||
(Err(err), Ok(())) => Err(err),
|
||||
(Err(err), Err(_)) => Err(err),
|
||||
}
|
||||
struct PlannedSchemaApply {
|
||||
plan: SchemaMigrationPlan,
|
||||
desired_ir: SchemaIR,
|
||||
desired_catalog: Catalog,
|
||||
}
|
||||
|
||||
pub(super) async fn apply_schema_with_lock(
|
||||
async fn plan_schema_for_apply(
|
||||
db: &Omnigraph,
|
||||
desired_schema_source: &str,
|
||||
options: SchemaApplyOptions,
|
||||
) -> Result<SchemaApplyResult> {
|
||||
) -> Result<PlannedSchemaApply> {
|
||||
db.ensure_schema_state_valid().await?;
|
||||
let branches = db.coordinator.read().await.all_branches().await?;
|
||||
// Skip `main` and internal system branches. The schema-apply lock branch
|
||||
// is excluded because it is the cluster-wide schema-apply serializer.
|
||||
// `__run__*` branches are no longer created; the filter remains as
|
||||
// defense-in-depth for legacy graphs with leftover staging branches.
|
||||
// A future production sweep will let this guard go.
|
||||
// Skip `main` and internal system branches (the schema-apply lock branch,
|
||||
// the cluster-wide schema-apply serializer). Legacy `__run__*` staging
|
||||
// branches were swept off `__manifest` by the v2→v3 migration that runs in
|
||||
// `Omnigraph::open(ReadWrite)` before this check (MR-770), so they no
|
||||
// longer appear here.
|
||||
let blocking_branches = branches
|
||||
.into_iter()
|
||||
.filter(|branch| branch != "main" && !is_internal_system_branch(branch))
|
||||
|
|
@ -123,6 +90,87 @@ pub(super) async fn apply_schema_with_lock(
|
|||
.unwrap_or_else(|| "unsupported schema migration plan".to_string());
|
||||
return Err(OmniError::manifest(message));
|
||||
}
|
||||
|
||||
let mut desired_catalog = build_catalog_from_ir(&desired_ir)?;
|
||||
fixup_blob_schemas(&mut desired_catalog);
|
||||
Ok(PlannedSchemaApply {
|
||||
plan,
|
||||
desired_ir,
|
||||
desired_catalog,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) async fn preview_schema_apply(
|
||||
db: &Omnigraph,
|
||||
desired_schema_source: &str,
|
||||
options: SchemaApplyOptions,
|
||||
) -> Result<SchemaApplyPreview> {
|
||||
let planned = plan_schema_for_apply(db, desired_schema_source, options).await?;
|
||||
Ok(SchemaApplyPreview {
|
||||
plan: planned.plan,
|
||||
catalog: planned.desired_catalog,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) async fn apply_schema<F>(
|
||||
db: &Omnigraph,
|
||||
desired_schema_source: &str,
|
||||
options: SchemaApplyOptions,
|
||||
actor: Option<&str>,
|
||||
validate_catalog: F,
|
||||
) -> Result<SchemaApplyResult>
|
||||
where
|
||||
F: FnOnce(&Catalog) -> Result<()>,
|
||||
{
|
||||
// Engine-layer policy gate (MR-722 chassis core).
|
||||
//
|
||||
// Fires BEFORE acquiring the schema-apply lock or doing any other
|
||||
// work. When no PolicyChecker is installed this is a no-op and
|
||||
// the apply path behaves exactly as it did before MR-722. When
|
||||
// a PolicyChecker IS installed and the actor is None, this is a
|
||||
// hard error — see Omnigraph::enforce's docstring for the
|
||||
// forget-the-actor-footgun reasoning.
|
||||
//
|
||||
// Scope is TargetBranch("main") to match the HTTP-layer convention
|
||||
// for SchemaApply: branch=None, target_branch=Some("main"). Cedar
|
||||
// policies in the wild use `target_branch_scope: protected` to
|
||||
// gate schema applies, so the engine-layer call has to set the
|
||||
// target_branch shape that activates that predicate. Wrong scope
|
||||
// here = silent policy mismatch with HTTP. See
|
||||
// `omnigraph_policy::ResourceScope::to_branch_pair` for the mapping.
|
||||
db.enforce(
|
||||
omnigraph_policy::PolicyAction::SchemaApply,
|
||||
&omnigraph_policy::ResourceScope::TargetBranch("main".to_string()),
|
||||
actor,
|
||||
)?;
|
||||
|
||||
acquire_schema_apply_lock(db).await?;
|
||||
let result = apply_schema_with_lock(db, desired_schema_source, options, validate_catalog).await;
|
||||
let release_result = release_schema_apply_lock(db).await;
|
||||
match (result, release_result) {
|
||||
(Ok(result), Ok(())) => Ok(result),
|
||||
(Ok(_), Err(err)) => Err(err),
|
||||
(Err(err), Ok(())) => Err(err),
|
||||
(Err(err), Err(_)) => Err(err),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn apply_schema_with_lock<F>(
|
||||
db: &Omnigraph,
|
||||
desired_schema_source: &str,
|
||||
options: SchemaApplyOptions,
|
||||
validate_catalog: F,
|
||||
) -> Result<SchemaApplyResult>
|
||||
where
|
||||
F: FnOnce(&Catalog) -> Result<()>,
|
||||
{
|
||||
let planned = plan_schema_for_apply(db, desired_schema_source, options).await?;
|
||||
validate_catalog(&planned.desired_catalog)?;
|
||||
let PlannedSchemaApply {
|
||||
plan,
|
||||
desired_ir,
|
||||
desired_catalog,
|
||||
} = planned;
|
||||
if plan.steps.is_empty() {
|
||||
return Ok(SchemaApplyResult {
|
||||
supported: true,
|
||||
|
|
@ -132,9 +180,6 @@ pub(super) async fn apply_schema_with_lock(
|
|||
});
|
||||
}
|
||||
|
||||
let mut desired_catalog = build_catalog_from_ir(&desired_ir)?;
|
||||
fixup_blob_schemas(&mut desired_catalog);
|
||||
|
||||
let snapshot = db.snapshot().await;
|
||||
let base_manifest_version = snapshot.version();
|
||||
let mut added_tables = BTreeSet::new();
|
||||
|
|
|
|||
|
|
@ -483,6 +483,22 @@ pub(super) async fn open_owned_dataset_for_branch_write(
|
|||
Ok((ds, Some(active_branch.to_string())))
|
||||
}
|
||||
source_branch => {
|
||||
crate::failpoints::maybe_fail("fork.before_classify")?;
|
||||
// Authority check before forking: re-read the live manifest. If this
|
||||
// table is already forked on active_branch, a concurrent first-write
|
||||
// won the race and our snapshot is stale — that is a retryable
|
||||
// conflict, not an orphan. (A zombie fork is never in the manifest,
|
||||
// so this only fires for a live concurrent fork.)
|
||||
let live = db.snapshot_for_branch(Some(active_branch)).await?;
|
||||
if let Some(entry) = live.entry(table_key) {
|
||||
if entry.table_branch.as_deref() == Some(active_branch) {
|
||||
return Err(OmniError::manifest_expected_version_mismatch(
|
||||
table_key,
|
||||
entry_version,
|
||||
entry.table_version,
|
||||
));
|
||||
}
|
||||
}
|
||||
fork_dataset_from_entry_state(
|
||||
db,
|
||||
table_key,
|
||||
|
|
|
|||
|
|
@ -1,16 +0,0 @@
|
|||
// The Run state machine has been removed. Mutations now write directly
|
||||
// to target tables and use the publisher's `expected_table_versions`
|
||||
// CAS for cross-table OCC; `__run__<id>` staging branches and the
|
||||
// `_graph_runs.lance` state machine no longer exist.
|
||||
//
|
||||
// What remains is the branch-name predicate, kept as a defense-in-depth
|
||||
// guard against users naming a public branch `__run__*`. A future
|
||||
// production sweep of legacy `_graph_runs.lance` rows and stale
|
||||
// `__run__*` branches will let this predicate (and this file) go too.
|
||||
|
||||
pub(crate) const INTERNAL_RUN_BRANCH_PREFIX: &str = "__run__";
|
||||
|
||||
pub(crate) fn is_internal_run_branch(name: &str) -> bool {
|
||||
name.trim_start_matches('/')
|
||||
.starts_with(INTERNAL_RUN_BRANCH_PREFIX)
|
||||
}
|
||||
|
|
@ -1087,9 +1087,9 @@ impl Omnigraph {
|
|||
target: &str,
|
||||
actor_id: Option<&str>,
|
||||
) -> Result<MergeOutcome> {
|
||||
if is_internal_run_branch(source) || is_internal_run_branch(target) {
|
||||
if is_internal_system_branch(source) || is_internal_system_branch(target) {
|
||||
return Err(OmniError::manifest(format!(
|
||||
"branch_merge does not allow internal run refs ('{}' -> '{}')",
|
||||
"branch_merge does not allow internal system refs ('{}' -> '{}')",
|
||||
source, target
|
||||
)));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ use time::format_description::well_known::Rfc3339;
|
|||
|
||||
use crate::db::commit_graph::CommitGraph;
|
||||
use crate::db::manifest::ManifestCoordinator;
|
||||
use crate::db::{MergeOutcome, Omnigraph, is_internal_run_branch};
|
||||
use crate::db::{MergeOutcome, Omnigraph, is_internal_system_branch};
|
||||
use crate::db::{ReadTarget, Snapshot};
|
||||
use crate::embedding::EmbeddingClient;
|
||||
use crate::error::{MergeConflict, MergeConflictKind, OmniError, Result};
|
||||
|
|
|
|||
|
|
@ -288,21 +288,24 @@ async fn load_jsonl_reader<R: BufRead>(
|
|||
let mut node_rows: HashMap<String, Vec<JsonValue>> = HashMap::new();
|
||||
let mut edge_rows: HashMap<String, Vec<(String, String, JsonValue)>> = HashMap::new();
|
||||
|
||||
for (line_num, line) in reader.lines().enumerate() {
|
||||
let line = line?;
|
||||
let line = line.trim();
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let value: JsonValue = serde_json::from_str(line).map_err(|e| {
|
||||
OmniError::manifest(format!("invalid JSON on line {}: {}", line_num + 1, e))
|
||||
// Parse a stream of JSON values. Accepts both compact JSONL (one object
|
||||
// per line) and pretty-printed JSON where a single object spans multiple
|
||||
// lines — serde's streaming deserializer treats any whitespace (including
|
||||
// newlines) between top-level values as a separator.
|
||||
for (idx, parsed) in serde_json::Deserializer::from_reader(reader)
|
||||
.into_iter::<JsonValue>()
|
||||
.enumerate()
|
||||
{
|
||||
let record_num = idx + 1;
|
||||
let value: JsonValue = parsed.map_err(|e| {
|
||||
OmniError::manifest(format!("invalid JSON at record {}: {}", record_num, e))
|
||||
})?;
|
||||
|
||||
if let Some(type_name) = value.get("type").and_then(|v| v.as_str()) {
|
||||
if !catalog.node_types.contains_key(type_name) {
|
||||
return Err(OmniError::manifest(format!(
|
||||
"line {}: unknown node type '{}'",
|
||||
line_num + 1,
|
||||
"record {}: unknown node type '{}'",
|
||||
record_num,
|
||||
type_name
|
||||
)));
|
||||
}
|
||||
|
|
@ -317,8 +320,8 @@ async fn load_jsonl_reader<R: BufRead>(
|
|||
} else if let Some(edge_name) = value.get("edge").and_then(|v| v.as_str()) {
|
||||
if catalog.lookup_edge_by_name(edge_name).is_none() {
|
||||
return Err(OmniError::manifest(format!(
|
||||
"line {}: unknown edge type '{}'",
|
||||
line_num + 1,
|
||||
"record {}: unknown edge type '{}'",
|
||||
record_num,
|
||||
edge_name
|
||||
)));
|
||||
}
|
||||
|
|
@ -326,14 +329,14 @@ async fn load_jsonl_reader<R: BufRead>(
|
|||
.get("from")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
OmniError::manifest(format!("line {}: edge missing 'from'", line_num + 1))
|
||||
OmniError::manifest(format!("record {}: edge missing 'from'", record_num))
|
||||
})?
|
||||
.to_string();
|
||||
let to = value
|
||||
.get("to")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
OmniError::manifest(format!("line {}: edge missing 'to'", line_num + 1))
|
||||
OmniError::manifest(format!("record {}: edge missing 'to'", record_num))
|
||||
})?
|
||||
.to_string();
|
||||
let data = value
|
||||
|
|
@ -347,8 +350,8 @@ async fn load_jsonl_reader<R: BufRead>(
|
|||
.push((from, to, data));
|
||||
} else {
|
||||
return Err(OmniError::manifest(format!(
|
||||
"line {}: expected 'type' or 'edge' field",
|
||||
line_num + 1
|
||||
"record {}: expected 'type' or 'edge' field",
|
||||
record_num
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -177,6 +177,45 @@ impl TableStore {
|
|||
.map_err(|e| OmniError::Lance(e.to_string()))
|
||||
}
|
||||
|
||||
/// List the named Lance branches present on the dataset at `dataset_uri`.
|
||||
/// The `cleanup` orphan reconciler diffs this against the manifest branch
|
||||
/// set to find orphaned per-table forks. `main`/default is not a named
|
||||
/// branch and never appears here.
|
||||
pub async fn list_branches(&self, dataset_uri: &str) -> Result<Vec<String>> {
|
||||
let ds = Dataset::open(dataset_uri)
|
||||
.await
|
||||
.map_err(|e| OmniError::Lance(e.to_string()))?;
|
||||
let branches = ds
|
||||
.list_branches()
|
||||
.await
|
||||
.map_err(|e| OmniError::Lance(e.to_string()))?;
|
||||
Ok(branches.into_keys().collect())
|
||||
}
|
||||
|
||||
/// Idempotently drop `branch` from the dataset at `dataset_uri`.
|
||||
///
|
||||
/// Unlike [`delete_branch`](Self::delete_branch), this tolerates an
|
||||
/// already-absent branch — both a missing contents ref (Lance's
|
||||
/// `force_delete_branch` handles that) and a missing `tree/{branch}/`
|
||||
/// directory (the local-store `NotFound` quirk pinned by
|
||||
/// `lance_surface_guards::force_delete_branch_semantics`). Safe to call on a
|
||||
/// possibly-orphaned or already-reclaimed fork.
|
||||
///
|
||||
/// A branch that still has referencing descendants (`RefConflict`) is NOT
|
||||
/// tolerated: that is a real ordering error and surfaces as `OmniError::Lance`.
|
||||
/// Used by the eager best-effort reclaim in `cleanup_deleted_branch_tables`
|
||||
/// and the `cleanup` orphan reconciler.
|
||||
pub async fn force_delete_branch(&self, dataset_uri: &str, branch: &str) -> Result<()> {
|
||||
let mut ds = Dataset::open(dataset_uri)
|
||||
.await
|
||||
.map_err(|e| OmniError::Lance(e.to_string()))?;
|
||||
match ds.force_delete_branch(branch).await {
|
||||
Ok(()) => Ok(()),
|
||||
Err(lance::Error::RefNotFound { .. }) | Err(lance::Error::NotFound { .. }) => Ok(()),
|
||||
Err(e) => Err(OmniError::Lance(e.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn open_dataset_at_state(
|
||||
&self,
|
||||
table_path: &str,
|
||||
|
|
@ -243,21 +282,24 @@ impl TableStore {
|
|||
.map_err(|e| OmniError::Lance(e.to_string()))?;
|
||||
self.ensure_expected_version(&source_ds, table_key, source_version)?;
|
||||
|
||||
match source_ds
|
||||
if source_ds
|
||||
.create_branch(target_branch, source_version, None)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
Ok(_) => {}
|
||||
Err(create_err) => match self
|
||||
.open_dataset_head(dataset_uri, Some(target_branch))
|
||||
.await
|
||||
{
|
||||
Ok(ds) => {
|
||||
self.ensure_expected_version(&ds, table_key, source_version)?;
|
||||
return Ok(ds);
|
||||
}
|
||||
Err(_) => return Err(OmniError::Lance(create_err.to_string())),
|
||||
},
|
||||
// The target branch ref already exists. The caller
|
||||
// (`open_owned_dataset_for_branch_write`) re-reads the live manifest
|
||||
// before forking and returns a retryable error when a concurrent
|
||||
// writer legitimately holds the fork, so reaching here means the
|
||||
// manifest does NOT reference this fork: it is an orphan from an
|
||||
// incomplete prior `branch_delete`. Surface the actionable cleanup
|
||||
// error rather than guessing from Lance branch versions.
|
||||
return Err(OmniError::manifest_conflict(format!(
|
||||
"branch '{}' has orphaned table state for '{}' from an incomplete \
|
||||
prior delete; run `omnigraph cleanup` to reclaim it before reusing \
|
||||
this branch name",
|
||||
target_branch, table_key
|
||||
)));
|
||||
}
|
||||
|
||||
let ds = self
|
||||
|
|
|
|||
|
|
@ -41,6 +41,452 @@ async fn branch_create_failpoint_triggers() {
|
|||
);
|
||||
}
|
||||
|
||||
// Branch delete flips the manifest authority first, then reclaims the per-table
|
||||
// forks best-effort. A failure during that reclaim (here, the
|
||||
// `branch_delete.before_table_cleanup` failpoint, standing in for a transient
|
||||
// object-store error) must NOT fail the call: the branch is already gone, and
|
||||
// `cleanup` reconciles the stranded fork. The branch name is reusable after.
|
||||
#[tokio::test]
|
||||
async fn branch_delete_partial_failure_converges_via_cleanup() {
|
||||
let _scenario = FailScenario::setup();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().to_str().unwrap().to_string();
|
||||
let mut main = helpers::init_and_load(&dir).await;
|
||||
|
||||
main.branch_create("feature").await.unwrap();
|
||||
let mut feature = Omnigraph::open(&uri).await.unwrap();
|
||||
helpers::mutate_branch(
|
||||
&mut feature,
|
||||
"feature",
|
||||
MUTATION_QUERIES,
|
||||
"insert_person",
|
||||
&mixed_params(&[("$name", "Eve")], &[("$age", 22)]),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
drop(feature);
|
||||
|
||||
let person_uri = node_table_uri(&uri, "Person");
|
||||
{
|
||||
let ds = lance::Dataset::open(&person_uri).await.unwrap();
|
||||
assert!(
|
||||
ds.list_branches().await.unwrap().contains_key("feature"),
|
||||
"precondition: the owned table fork exists before delete"
|
||||
);
|
||||
}
|
||||
|
||||
// Inject a failure during per-table cleanup, AFTER the manifest authority
|
||||
// flip. branch_delete must still succeed (best-effort reclaim).
|
||||
{
|
||||
let _fp = ScopedFailPoint::new("branch_delete.before_table_cleanup", "return");
|
||||
main.branch_delete("feature").await.expect(
|
||||
"branch_delete is best-effort after the manifest flip: a cleanup-step \
|
||||
failure must not fail the call",
|
||||
);
|
||||
}
|
||||
|
||||
// Authority flipped: the branch is gone.
|
||||
assert_eq!(main.branch_list().await.unwrap(), vec!["main".to_string()]);
|
||||
|
||||
// The eager reclaim failed, so the orphan is stranded until cleanup.
|
||||
{
|
||||
let ds = lance::Dataset::open(&person_uri).await.unwrap();
|
||||
assert!(
|
||||
ds.list_branches().await.unwrap().contains_key("feature"),
|
||||
"failed eager reclaim should leave the orphan for cleanup to reconcile"
|
||||
);
|
||||
}
|
||||
|
||||
// cleanup converges: the orphan is reclaimed.
|
||||
main.cleanup(omnigraph::db::CleanupPolicyOptions {
|
||||
keep_versions: Some(1),
|
||||
older_than: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
{
|
||||
let ds = lance::Dataset::open(&person_uri).await.unwrap();
|
||||
assert!(
|
||||
!ds.list_branches().await.unwrap().contains_key("feature"),
|
||||
"cleanup should reconcile the orphaned fork away"
|
||||
);
|
||||
}
|
||||
|
||||
// The name is reusable after cleanup reclaims the orphan.
|
||||
main.branch_create("feature").await.unwrap();
|
||||
let mut feature2 = Omnigraph::open(&uri).await.unwrap();
|
||||
helpers::mutate_branch(
|
||||
&mut feature2,
|
||||
"feature",
|
||||
MUTATION_QUERIES,
|
||||
"insert_person",
|
||||
&mixed_params(&[("$name", "Frank")], &[("$age", 41)]),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Reusing a branch name whose delete left an orphaned fork (before `cleanup`
|
||||
// reconciles it) must fail with a clear, actionable error pointing at
|
||||
// `cleanup`, not the opaque `ExpectedVersionMismatch` that leaks from the fork
|
||||
// path. The recreate itself succeeds; the first write to the previously-forked
|
||||
// table is where the stale orphan collides.
|
||||
#[tokio::test]
|
||||
async fn recreate_over_orphaned_fork_before_cleanup_is_actionable() {
|
||||
let _scenario = FailScenario::setup();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().to_str().unwrap().to_string();
|
||||
let mut main = helpers::init_and_load(&dir).await;
|
||||
|
||||
main.branch_create("feature").await.unwrap();
|
||||
let mut feature = Omnigraph::open(&uri).await.unwrap();
|
||||
helpers::mutate_branch(
|
||||
&mut feature,
|
||||
"feature",
|
||||
MUTATION_QUERIES,
|
||||
"insert_person",
|
||||
&mixed_params(&[("$name", "Eve")], &[("$age", 22)]),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
drop(feature);
|
||||
|
||||
// Partial delete: leaves the Person fork orphaned (cleanup not yet run).
|
||||
{
|
||||
let _fp = ScopedFailPoint::new("branch_delete.before_table_cleanup", "return");
|
||||
main.branch_delete("feature").await.unwrap();
|
||||
}
|
||||
|
||||
// Recreate the name and write to the previously-forked table WITHOUT a
|
||||
// cleanup in between.
|
||||
main.branch_create("feature").await.unwrap();
|
||||
let mut feature2 = Omnigraph::open(&uri).await.unwrap();
|
||||
let err = helpers::mutate_branch(
|
||||
&mut feature2,
|
||||
"feature",
|
||||
MUTATION_QUERIES,
|
||||
"insert_person",
|
||||
&mixed_params(&[("$name", "Frank")], &[("$age", 41)]),
|
||||
)
|
||||
.await
|
||||
.expect_err("write should collide with the stale orphaned fork");
|
||||
|
||||
let msg = err.to_string();
|
||||
assert!(
|
||||
msg.contains("cleanup")
|
||||
&& (msg.contains("orphan") || msg.contains("incomplete prior delete")),
|
||||
"expected an actionable orphaned-fork error pointing at cleanup, got: {msg}"
|
||||
);
|
||||
assert!(
|
||||
!msg.contains("expected manifest table version"),
|
||||
"should not surface the opaque ExpectedVersionMismatch, got: {msg}"
|
||||
);
|
||||
}
|
||||
|
||||
// cleanup is the guaranteed convergence backstop, so one table's transient
|
||||
// failure must not abort the whole sweep. Inject a one-shot version-GC failure
|
||||
// for a single table and assert: cleanup still succeeds, the failure is
|
||||
// surfaced per-table in the returned stats, and the independent reconcile pass
|
||||
// still reclaimed an orphan.
|
||||
#[tokio::test]
|
||||
async fn cleanup_isolates_single_table_failure() {
|
||||
let _scenario = FailScenario::setup();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().to_str().unwrap().to_string();
|
||||
let mut db = helpers::init_and_load(&dir).await;
|
||||
|
||||
// Forge an orphaned fork on the Person table (a reconcile target).
|
||||
let person_uri = node_table_uri(&uri, "Person");
|
||||
{
|
||||
let mut ds = lance::Dataset::open(&person_uri).await.unwrap();
|
||||
let base = ds.version().version;
|
||||
ds.create_branch("ghost", base, None).await.unwrap();
|
||||
}
|
||||
|
||||
// One table's version GC fails once; the sweep must isolate it.
|
||||
let _fp = ScopedFailPoint::new("cleanup.table_gc", "1*return");
|
||||
let stats = db
|
||||
.cleanup(omnigraph::db::CleanupPolicyOptions {
|
||||
keep_versions: Some(1),
|
||||
older_than: None,
|
||||
})
|
||||
.await
|
||||
.expect("a single table's GC failure must not abort cleanup");
|
||||
|
||||
let errored = stats.iter().filter(|s| s.error.is_some()).count();
|
||||
assert_eq!(
|
||||
errored, 1,
|
||||
"exactly one table's GC failure should be surfaced in stats, got {errored}"
|
||||
);
|
||||
assert!(
|
||||
stats.len() >= 4,
|
||||
"every node+edge table should still appear in the stats"
|
||||
);
|
||||
|
||||
// The reconcile pass is independent of the GC failure, so the orphan is gone.
|
||||
{
|
||||
let ds = lance::Dataset::open(&person_uri).await.unwrap();
|
||||
assert!(
|
||||
!ds.list_branches().await.unwrap().contains_key("ghost"),
|
||||
"reconcile should reclaim the orphan despite the GC failure"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Companion to the version-GC isolation test, exercising the OTHER cleanup
|
||||
// loop: a force-delete failure while reconciling one orphaned fork must be
|
||||
// isolated (logged, not propagated) so the sweep continues, and a later
|
||||
// cleanup converges. This is the loop the Devin finding was about.
|
||||
#[tokio::test]
|
||||
async fn cleanup_isolates_reconcile_failure() {
|
||||
let _scenario = FailScenario::setup();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().to_str().unwrap().to_string();
|
||||
let mut db = helpers::init_and_load(&dir).await;
|
||||
|
||||
// Forge an orphaned fork the reconcile pass will try to reclaim.
|
||||
let person_uri = node_table_uri(&uri, "Person");
|
||||
{
|
||||
let mut ds = lance::Dataset::open(&person_uri).await.unwrap();
|
||||
let base = ds.version().version;
|
||||
ds.create_branch("ghost", base, None).await.unwrap();
|
||||
}
|
||||
|
||||
// Inject a one-shot failure into the reconcile force-delete. The sweep must
|
||||
// not abort.
|
||||
{
|
||||
let _fp = ScopedFailPoint::new("cleanup.reconcile_fork", "1*return");
|
||||
db.cleanup(omnigraph::db::CleanupPolicyOptions {
|
||||
keep_versions: Some(1),
|
||||
older_than: None,
|
||||
})
|
||||
.await
|
||||
.expect("a reconcile force-delete failure must not abort cleanup");
|
||||
}
|
||||
// The blocked orphan is still present (the failure was isolated, not retried).
|
||||
{
|
||||
let ds = lance::Dataset::open(&person_uri).await.unwrap();
|
||||
assert!(
|
||||
ds.list_branches().await.unwrap().contains_key("ghost"),
|
||||
"the orphan whose reclaim was injected-to-fail should remain"
|
||||
);
|
||||
}
|
||||
// A second cleanup with no injected failure converges.
|
||||
db.cleanup(omnigraph::db::CleanupPolicyOptions {
|
||||
keep_versions: Some(1),
|
||||
older_than: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
{
|
||||
let ds = lance::Dataset::open(&person_uri).await.unwrap();
|
||||
assert!(
|
||||
!ds.list_branches().await.unwrap().contains_key("ghost"),
|
||||
"the second cleanup should reconcile the orphan"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// The cleanup reconciler must reclaim orphaned commit-graph branches, not just
|
||||
// per-table forks. A delete whose best-effort commit-graph reclaim fails leaves
|
||||
// a commit-graph orphan; the next cleanup must drop it.
|
||||
#[tokio::test]
|
||||
async fn cleanup_reclaims_orphaned_commit_graph_branch() {
|
||||
let _scenario = FailScenario::setup();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().to_str().unwrap().to_string();
|
||||
let mut db = helpers::init_and_load(&dir).await;
|
||||
|
||||
db.branch_create("feature").await.unwrap();
|
||||
// Delete, failing the commit-graph reclaim → commit-graph "feature" orphan
|
||||
// (manifest branch gone, commit-graph branch left behind).
|
||||
{
|
||||
let _fp = ScopedFailPoint::new("branch_delete.before_commit_graph_reclaim", "return");
|
||||
db.branch_delete("feature").await.unwrap();
|
||||
}
|
||||
|
||||
let commits_uri = format!("{}/_graph_commits.lance", uri.trim_end_matches('/'));
|
||||
{
|
||||
let ds = lance::Dataset::open(&commits_uri).await.unwrap();
|
||||
assert!(
|
||||
ds.list_branches().await.unwrap().contains_key("feature"),
|
||||
"precondition: the commit-graph branch should be orphaned after the failed reclaim"
|
||||
);
|
||||
}
|
||||
|
||||
db.cleanup(omnigraph::db::CleanupPolicyOptions {
|
||||
keep_versions: Some(1),
|
||||
older_than: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
{
|
||||
let ds = lance::Dataset::open(&commits_uri).await.unwrap();
|
||||
assert!(
|
||||
!ds.list_branches().await.unwrap().contains_key("feature"),
|
||||
"cleanup should reclaim the orphaned commit-graph branch"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// A branch_delete whose best-effort commit-graph reclaim fails leaves a
|
||||
// commit-graph "zombie" branch. Recreating that name must heal the zombie and
|
||||
// succeed (branch_create force-deletes a stale commit-graph ref since the
|
||||
// manifest branch is created fresh), instead of dying on the leftover ref.
|
||||
#[tokio::test]
|
||||
async fn branch_create_recreates_over_commit_graph_zombie() {
|
||||
let _scenario = FailScenario::setup();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = Omnigraph::init(dir.path().to_str().unwrap(), helpers::TEST_SCHEMA)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
db.branch_create("feature").await.unwrap();
|
||||
{
|
||||
// Fail the best-effort commit-graph reclaim → commit-graph "feature"
|
||||
// zombie survives the delete (manifest authority still flips).
|
||||
let _fp = ScopedFailPoint::new("branch_delete.before_commit_graph_reclaim", "return");
|
||||
db.branch_delete("feature").await.unwrap();
|
||||
}
|
||||
assert_eq!(db.branch_list().await.unwrap(), vec!["main".to_string()]);
|
||||
|
||||
db.branch_create("feature")
|
||||
.await
|
||||
.expect("branch_create should heal the zombie commit-graph branch and succeed");
|
||||
assert!(
|
||||
db.branch_list()
|
||||
.await
|
||||
.unwrap()
|
||||
.contains(&"feature".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
// branch_create is authority-then-derived: if the derived commit-graph branch
|
||||
// cannot be created, the manifest branch (the authority) must be rolled back so
|
||||
// the branch does not half-exist. The existing failpoint fires right after the
|
||||
// manifest create, standing in for any post-authority failure.
|
||||
#[tokio::test]
|
||||
async fn branch_create_rolls_back_manifest_on_commit_graph_failure() {
|
||||
let _scenario = FailScenario::setup();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = Omnigraph::init(dir.path().to_str().unwrap(), helpers::TEST_SCHEMA)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let err = {
|
||||
let _fp = ScopedFailPoint::new("branch_create.after_manifest_branch_create", "return");
|
||||
db.branch_create("feature").await.unwrap_err()
|
||||
};
|
||||
assert!(
|
||||
!db.branch_list()
|
||||
.await
|
||||
.unwrap()
|
||||
.contains(&"feature".to_string()),
|
||||
"branch_create must roll back the manifest branch when the derived \
|
||||
commit-graph branch fails, got error: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
// A fork collision must be classified by the manifest authority, not by Lance
|
||||
// branch versions. When a concurrent first-write legitimately wins the fork
|
||||
// race, the loser sees a version mismatch — but that is a stale snapshot, not
|
||||
// an orphan, so it must be a retryable "refresh and retry", never a misleading
|
||||
// "run cleanup".
|
||||
//
|
||||
// Ordering is made deterministic (no sleeps) via a callback at the fork point:
|
||||
// `compare_exchange` lets only the FIRST arrival (writer A) record readiness and
|
||||
// block until released; later arrivals (writer B) fall through. The test waits
|
||||
// on the readiness flag, lets B win and commit the fork, then releases A.
|
||||
static FORK_A_AT_POINT: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||
static FORK_RELEASE_A: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn fork_collision_with_live_concurrent_fork_is_retryable() {
|
||||
use std::sync::atomic::Ordering::SeqCst;
|
||||
|
||||
let _scenario = FailScenario::setup();
|
||||
FORK_A_AT_POINT.store(false, SeqCst);
|
||||
FORK_RELEASE_A.store(false, SeqCst);
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().to_str().unwrap().to_string();
|
||||
let main = helpers::init_and_load(&dir).await;
|
||||
main.branch_create("feature").await.unwrap();
|
||||
|
||||
// First arrival (A) records readiness and blocks until released; the rest
|
||||
// (B) fall through immediately. Bounded spin so a mistake can't hang forever.
|
||||
fail::cfg_callback("fork.before_classify", || {
|
||||
if FORK_A_AT_POINT
|
||||
.compare_exchange(false, true, SeqCst, SeqCst)
|
||||
.is_ok()
|
||||
{
|
||||
for _ in 0..2000 {
|
||||
if FORK_RELEASE_A.load(SeqCst) {
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(5));
|
||||
}
|
||||
}
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let uri_a = uri.clone();
|
||||
let writer_a = tokio::spawn(async move {
|
||||
let mut a = Omnigraph::open(&uri_a).await.unwrap();
|
||||
helpers::mutate_branch(
|
||||
&mut a,
|
||||
"feature",
|
||||
MUTATION_QUERIES,
|
||||
"insert_person",
|
||||
&mixed_params(&[("$name", "Eve")], &[("$age", 22)]),
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
// Wait (bounded) until A is parked at the fork point.
|
||||
for _ in 0..600 {
|
||||
if FORK_A_AT_POINT.load(SeqCst) {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
|
||||
}
|
||||
assert!(
|
||||
FORK_A_AT_POINT.load(SeqCst),
|
||||
"writer A never reached the fork point"
|
||||
);
|
||||
|
||||
// B wins the fork and commits it.
|
||||
let mut b = Omnigraph::open(&uri).await.unwrap();
|
||||
helpers::mutate_branch(
|
||||
&mut b,
|
||||
"feature",
|
||||
MUTATION_QUERIES,
|
||||
"insert_person",
|
||||
&mixed_params(&[("$name", "Frank")], &[("$age", 41)]),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Release A; it resumes, re-reads the manifest, and sees the fork is live.
|
||||
FORK_RELEASE_A.store(true, SeqCst);
|
||||
let err = writer_a
|
||||
.await
|
||||
.unwrap()
|
||||
.expect_err("A's stale-snapshot fork should be a retryable conflict");
|
||||
fail::remove("fork.before_classify");
|
||||
|
||||
let msg = err.to_string();
|
||||
assert!(
|
||||
!msg.contains("cleanup"),
|
||||
"a live concurrent fork must not be misclassified as an orphan, got: {msg}"
|
||||
);
|
||||
assert!(
|
||||
msg.contains("refresh and retry") || msg.contains("expected manifest table version"),
|
||||
"expected a retryable stale-view error, got: {msg}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn graph_publish_failpoint_triggers_before_commit_append() {
|
||||
let _scenario = FailScenario::setup();
|
||||
|
|
|
|||
|
|
@ -242,3 +242,136 @@ async fn _compile_delete_result_field_shape() -> lance::Result<()> {
|
|||
let _num_deleted: u64 = result.num_deleted_rows;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// --- Guard 9: force_delete_branch semantics --------------------------------
|
||||
//
|
||||
// The branch-delete reconciler (`db/omnigraph/optimize.rs::reconcile_orphaned_branches`)
|
||||
// and the eager best-effort reclaim in `cleanup_deleted_branch_tables` call
|
||||
// `force_delete_branch` to drop orphaned branch refs. The single-authority
|
||||
// design relies on three facts pinned here:
|
||||
// 1. plain `delete_branch` errors on a missing ref (so the design uses the
|
||||
// force variant instead);
|
||||
// 2. `force_delete_branch` removes an existing (forked) branch — the orphan
|
||||
// case, where a `tree/{branch}/` exists;
|
||||
// 3. `force_delete_branch` on a *fully-absent* branch (no tree dir) still
|
||||
// errors on the local store, because `remove_dir_all`'s NotFound is not
|
||||
// caught for Lance's native error variant. `TableStore::force_delete_branch`
|
||||
// wraps this to be fully idempotent. Pin the raw quirk so a future Lance
|
||||
// fix (which would let us simplify the wrapper) is noticed.
|
||||
|
||||
#[tokio::test]
|
||||
async fn force_delete_branch_semantics() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().join("guard9.lance");
|
||||
let uri = uri.to_str().unwrap();
|
||||
let mut ds = fresh_dataset(uri).await;
|
||||
|
||||
// (1) Plain delete of a never-created branch errors (RefNotFound).
|
||||
assert!(
|
||||
ds.delete_branch("nope").await.is_err(),
|
||||
"Dataset::delete_branch on a missing ref should error; if this is now \
|
||||
Ok, the reconciler could drop the force variant."
|
||||
);
|
||||
|
||||
// (2) force_delete_branch removes an existing (forked) branch.
|
||||
let base = ds.version().version;
|
||||
ds.create_branch("feature", base, None).await.unwrap();
|
||||
ds.force_delete_branch("feature").await.unwrap();
|
||||
assert!(
|
||||
!ds.list_branches().await.unwrap().contains_key("feature"),
|
||||
"force_delete_branch should remove an existing branch ref"
|
||||
);
|
||||
|
||||
// (3) Quirk: force_delete on a fully-absent branch errors on the local
|
||||
// store (worked around by TableStore::force_delete_branch).
|
||||
assert!(
|
||||
ds.force_delete_branch("never").await.is_err(),
|
||||
"force_delete_branch on a fully-absent branch no longer errors — \
|
||||
TableStore::force_delete_branch's NotFound tolerance can be simplified."
|
||||
);
|
||||
}
|
||||
|
||||
// --- Guard 10: blob-column compaction is still broken in this Lance --------
|
||||
//
|
||||
// `db/omnigraph/optimize.rs` skips tables with blob columns while
|
||||
// `LANCE_SUPPORTS_BLOB_COMPACTION = false`: Lance `compact_files` forces
|
||||
// `BlobHandling::AllBinary`, and the blob-v2 struct decoder mis-counts columns
|
||||
// ("more fields in the schema than provided column indices"), failing even a
|
||||
// pristine uniform-V2_2 multi-fragment blob table. Reads are unaffected (they
|
||||
// use descriptor handling).
|
||||
//
|
||||
// WHEN THIS TEST TURNS RED (compact_files no longer errors), the Lance bug is
|
||||
// fixed: flip `LANCE_SUPPORTS_BLOB_COMPACTION` to true in optimize.rs, drop the
|
||||
// blob-skip branch + the `optimize_skips_blob_table_and_reports_skip`
|
||||
// skip assertions in maintenance.rs, and re-pin docs/dev/lance.md.
|
||||
|
||||
#[tokio::test]
|
||||
async fn compact_files_still_fails_on_blob_columns() {
|
||||
use arrow_array::{LargeBinaryArray, StructArray};
|
||||
|
||||
fn blob_batch(start: i32, n: i32) -> RecordBatch {
|
||||
let ids: Vec<String> = (start..start + n).map(|i| format!("n{i}")).collect();
|
||||
let data =
|
||||
LargeBinaryArray::from_iter_values((start..start + n).map(|i| format!("blob{i}")));
|
||||
let blob_uri = StringArray::from(vec![None::<&str>; n as usize]);
|
||||
let DataType::Struct(fields) = lance::blob::blob_field("content", true).data_type().clone()
|
||||
else {
|
||||
unreachable!("blob_field is always a Struct");
|
||||
};
|
||||
let content = StructArray::new(
|
||||
fields,
|
||||
vec![Arc::new(data) as _, Arc::new(blob_uri) as _],
|
||||
None,
|
||||
);
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Utf8, false),
|
||||
lance::blob::blob_field("content", true),
|
||||
]));
|
||||
RecordBatch::try_new(
|
||||
schema,
|
||||
vec![Arc::new(StringArray::from(ids)) as _, Arc::new(content) as _],
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn write(uri: &str, batch: RecordBatch, mode: WriteMode) {
|
||||
let schema = batch.schema();
|
||||
let reader = RecordBatchIterator::new(vec![Ok(batch)], schema);
|
||||
// Blob v2 requires file version >= 2.2; without the pin the *write*
|
||||
// would fail with a different error, masking the guard's intent.
|
||||
let params = WriteParams {
|
||||
mode,
|
||||
enable_stable_row_ids: true,
|
||||
data_storage_version: Some(LanceFileVersion::V2_2),
|
||||
..Default::default()
|
||||
};
|
||||
Dataset::write(reader, uri, Some(params)).await.unwrap();
|
||||
}
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().join("guard10-blob.lance");
|
||||
let uri = uri.to_str().unwrap();
|
||||
|
||||
// Uniform V2_2, two fragments → forces compaction to actually rewrite.
|
||||
write(uri, blob_batch(0, 2), WriteMode::Create).await;
|
||||
write(uri, blob_batch(100, 2), WriteMode::Append).await;
|
||||
|
||||
let mut ds = Dataset::open(uri).await.unwrap();
|
||||
assert!(
|
||||
ds.get_fragments().len() >= 2,
|
||||
"guard needs a multi-fragment table to trigger a real compaction rewrite"
|
||||
);
|
||||
|
||||
let result = compact_files(&mut ds, CompactionOptions::default(), None).await;
|
||||
let err = result.expect_err(
|
||||
"compact_files unexpectedly SUCCEEDED on a blob table — the Lance blob-v2 \
|
||||
compaction bug is fixed. Flip LANCE_SUPPORTS_BLOB_COMPACTION to true in \
|
||||
db/omnigraph/optimize.rs, remove the blob-skip branch, and re-pin docs/dev/lance.md.",
|
||||
);
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("more fields in the schema than provided column indices"),
|
||||
"blob compaction failed with an unexpected error (Lance internals may have \
|
||||
shifted): {err}"
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,11 +7,24 @@ mod helpers;
|
|||
|
||||
use std::time::Duration;
|
||||
|
||||
use omnigraph::db::{CleanupPolicyOptions, Omnigraph};
|
||||
use lance::Dataset;
|
||||
use omnigraph::db::{CleanupPolicyOptions, Omnigraph, SkipReason};
|
||||
use omnigraph::loader::{LoadMode, load_jsonl};
|
||||
|
||||
use helpers::{TEST_DATA, TEST_SCHEMA, count_rows, init_and_load};
|
||||
|
||||
/// Filesystem URI of a node sub-table, mirroring the engine's layout
|
||||
/// (FNV-1a of the type name under `nodes/`). Matches the helper in
|
||||
/// `failpoints.rs`; used to inspect/forge Lance branches directly in tests.
|
||||
fn node_table_uri(root: &str, type_name: &str) -> String {
|
||||
let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
|
||||
for &b in type_name.as_bytes() {
|
||||
hash ^= b as u64;
|
||||
hash = hash.wrapping_mul(0x100_0000_01b3);
|
||||
}
|
||||
format!("{}/nodes/{hash:016x}", root.trim_end_matches('/'))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn optimize_on_empty_graph_returns_stats_per_table_with_no_changes() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
|
@ -59,6 +72,97 @@ async fn optimize_after_load_then_again_is_idempotent() {
|
|||
}
|
||||
}
|
||||
|
||||
// Regression: `optimize` must not crash on a graph that has a `Blob` table.
|
||||
//
|
||||
// Lance `compact_files` forces `BlobHandling::AllBinary`, which mis-decodes
|
||||
// blob-v2 columns ("more fields in the schema than provided column indices"),
|
||||
// failing even a pristine uniform-V2_2 multi-fragment blob table. `optimize`
|
||||
// must skip blob-bearing tables (and report the skip) rather than aborting the
|
||||
// whole sweep.
|
||||
//
|
||||
// Before the skip fix, `optimize()` returned that Lance error here and aborted
|
||||
// the whole sweep; it now skips the blob table (`doc.skipped == Some(..)`)
|
||||
// while the sibling non-blob `Tag` table still compacts. The skip is gated by
|
||||
// `LANCE_SUPPORTS_BLOB_COMPACTION`; the surface guard
|
||||
// `compact_files_still_fails_on_blob_columns` flags when the upstream Lance fix
|
||||
// makes the skip (and this test's blob arm) removable.
|
||||
#[tokio::test]
|
||||
async fn optimize_skips_blob_table_and_reports_skip() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().to_str().unwrap();
|
||||
// One Blob node type (`Doc`) + one plain node type (`Tag`): proves the blob
|
||||
// table is skipped while a non-blob table in the same sweep still compacts.
|
||||
let schema = "\
|
||||
node Doc {\n slug: String @key\n content: Blob\n}\n\
|
||||
node Tag {\n slug: String @key\n}\n";
|
||||
let mut db = Omnigraph::init(uri, schema).await.unwrap();
|
||||
|
||||
// Multi-fragment blob table: Overwrite creates fragment 1; each Merge of
|
||||
// new keys appends another. A >=2-fragment blob table is exactly what
|
||||
// crashes `compact_files` today (single fragment would no-op and not crash).
|
||||
load_jsonl(
|
||||
&mut db,
|
||||
"{\"type\":\"Doc\",\"data\":{\"slug\":\"d1\",\"content\":\"base64:aGVsbG8x\"}}\n{\"type\":\"Doc\",\"data\":{\"slug\":\"d2\",\"content\":\"base64:aGVsbG8y\"}}",
|
||||
LoadMode::Overwrite,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
load_jsonl(
|
||||
&mut db,
|
||||
"{\"type\":\"Doc\",\"data\":{\"slug\":\"d3\",\"content\":\"base64:aGVsbG8z\"}}",
|
||||
LoadMode::Merge,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
load_jsonl(
|
||||
&mut db,
|
||||
"{\"type\":\"Doc\",\"data\":{\"slug\":\"d4\",\"content\":\"base64:aGVsbG80\"}}",
|
||||
LoadMode::Merge,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
// Plain table, also multi-fragment so it has something to compact.
|
||||
load_jsonl(
|
||||
&mut db,
|
||||
"{\"type\":\"Tag\",\"data\":{\"slug\":\"t1\"}}\n{\"type\":\"Tag\",\"data\":{\"slug\":\"t2\"}}",
|
||||
LoadMode::Merge,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
load_jsonl(
|
||||
&mut db,
|
||||
"{\"type\":\"Tag\",\"data\":{\"slug\":\"t3\"}}",
|
||||
LoadMode::Merge,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let stats = db
|
||||
.optimize()
|
||||
.await
|
||||
.expect("optimize must not crash on a graph with a Blob table");
|
||||
|
||||
let doc = stats
|
||||
.iter()
|
||||
.find(|s| s.table_key == "node:Doc")
|
||||
.expect("Doc stat present");
|
||||
let tag = stats
|
||||
.iter()
|
||||
.find(|s| s.table_key == "node:Tag")
|
||||
.expect("Tag stat present");
|
||||
// The blob table is skipped (and reported), not compacted.
|
||||
assert_eq!(
|
||||
doc.skipped,
|
||||
Some(SkipReason::BlobColumnsUnsupportedByLance),
|
||||
"blob table must be reported as skipped",
|
||||
);
|
||||
assert!(!doc.committed, "skipped blob table is not compacted");
|
||||
assert_eq!(doc.fragments_removed, 0);
|
||||
assert_eq!(doc.fragments_added, 0);
|
||||
// The plain (non-blob) table is unaffected by the skip.
|
||||
assert_eq!(tag.skipped, None, "non-blob table must not be skipped");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cleanup_without_any_policy_option_errors() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
|
@ -158,3 +262,59 @@ async fn cleanup_then_optimize_preserves_rows_and_table_remains_writable() {
|
|||
.unwrap();
|
||||
assert_eq!(count_rows(&db, "node:Person").await, people_before);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cleanup_reconciles_orphaned_branch_forks() {
|
||||
// An incomplete prior `branch_delete` can leave a per-table Lance branch
|
||||
// that the manifest no longer references (a "zombie" fork). It is
|
||||
// unreachable through any snapshot but pins its `tree/{branch}/` storage.
|
||||
// `cleanup` must reconcile it away: drop every Lance branch absent from the
|
||||
// manifest authority, without touching `main`.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let uri = dir.path().to_str().unwrap().to_string();
|
||||
let mut db = init_and_load(&dir).await;
|
||||
|
||||
let people_before = count_rows(&db, "node:Person").await;
|
||||
assert!(people_before > 0, "fixture should seed Person rows");
|
||||
|
||||
// Forge an orphaned fork the manifest never knew about.
|
||||
let person_uri = node_table_uri(&uri, "Person");
|
||||
{
|
||||
let mut ds = Dataset::open(&person_uri).await.unwrap();
|
||||
let base = ds.version().version;
|
||||
ds.create_branch("ghost", base, None).await.unwrap();
|
||||
assert!(
|
||||
ds.list_branches().await.unwrap().contains_key("ghost"),
|
||||
"precondition: orphaned fork staged"
|
||||
);
|
||||
}
|
||||
|
||||
db.cleanup(CleanupPolicyOptions {
|
||||
keep_versions: Some(1),
|
||||
older_than: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Orphan reclaimed; main untouched.
|
||||
{
|
||||
let ds = Dataset::open(&person_uri).await.unwrap();
|
||||
assert!(
|
||||
!ds.list_branches().await.unwrap().contains_key("ghost"),
|
||||
"cleanup should reconcile the orphaned 'ghost' fork away"
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
count_rows(&db, "node:Person").await,
|
||||
people_before,
|
||||
"cleanup must not disturb main while reconciling orphans"
|
||||
);
|
||||
|
||||
// Idempotent: a second cleanup with the orphan already gone is a no-op.
|
||||
db.cleanup(CleanupPolicyOptions {
|
||||
keep_versions: Some(1),
|
||||
older_than: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -371,11 +371,10 @@ async fn cancelled_mutation_future_leaves_no_state() {
|
|||
|
||||
// Cancel-safety property: no graph-level run/staging state remains.
|
||||
//
|
||||
// Note: `branch_list()` already filters `__run__*` via
|
||||
// `is_internal_system_branch`, so a runtime "no `__run__` branches" check
|
||||
// would be vacuous. The structural property that no `__run__` branches
|
||||
// can ever be created is enforced by deletion of `begin_run` etc. in
|
||||
// (verified by the build itself — those symbols no longer exist).
|
||||
// No `__run__` branches can ever be created: the Run state machine
|
||||
// (`begin_run` etc.) was deleted in MR-771 — verified by the build itself,
|
||||
// those symbols no longer exist. Any legacy `__run__*` branch on an
|
||||
// upgraded graph is swept by the v2→v3 manifest migration.
|
||||
//
|
||||
// (1) The branch list is unchanged: cancellation/completion cannot
|
||||
// synthesize new public branches.
|
||||
|
|
@ -442,34 +441,40 @@ async fn repeated_loads_do_not_accumulate_branches() {
|
|||
assert_eq!(db.branch_list().await.unwrap(), vec!["main".to_string()]);
|
||||
}
|
||||
|
||||
/// User code must not be able to write to internal `__run__*` names.
|
||||
/// The branch-name guard predicate is kept as defense-in-depth; it
|
||||
/// will be removed once a future production sweep retires the legacy
|
||||
/// branches.
|
||||
/// After MR-770, `__run__*` is an ordinary branch name — the Run state machine
|
||||
/// and its `is_internal_run_branch` guard are gone. The surviving internal-ref
|
||||
/// guard still rejects the active `__schema_apply_lock__` branch on the public
|
||||
/// create/merge APIs.
|
||||
#[tokio::test]
|
||||
async fn public_branch_apis_reject_internal_run_refs() {
|
||||
async fn public_branch_apis_reject_internal_system_refs() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut db = init_and_load(&dir).await;
|
||||
|
||||
let create_err = db.branch_create("__run__synthetic").await.unwrap_err();
|
||||
// `__run__*` is no longer reserved — creating it now succeeds.
|
||||
db.branch_create("__run__formerly_reserved")
|
||||
.await
|
||||
.expect("__run__ prefix is a normal branch name post-MR-770");
|
||||
|
||||
// The schema-apply lock branch is still rejected on public branch APIs.
|
||||
let create_err = db.branch_create("__schema_apply_lock__").await.unwrap_err();
|
||||
let OmniError::Manifest(err) = create_err else {
|
||||
panic!("expected Manifest error");
|
||||
};
|
||||
assert!(
|
||||
err.message.contains("internal run ref"),
|
||||
err.message.contains("internal system ref"),
|
||||
"unexpected error: {}",
|
||||
err.message
|
||||
);
|
||||
|
||||
let merge_err = db
|
||||
.branch_merge("__run__synthetic", "main")
|
||||
.branch_merge("__schema_apply_lock__", "main")
|
||||
.await
|
||||
.unwrap_err();
|
||||
let OmniError::Manifest(err) = merge_err else {
|
||||
panic!("expected Manifest error");
|
||||
};
|
||||
assert!(
|
||||
err.message.contains("internal run refs"),
|
||||
err.message.contains("internal system refs"),
|
||||
"unexpected error: {}",
|
||||
err.message
|
||||
);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue