Move trustfall related parts to its own module

Signed-off-by: Marcel Müller <neikos@neikos.email>
This commit is contained in:
Marcel Müller 2025-02-08 10:08:58 +01:00
parent 8a453a44f9
commit 3112e78bb4
2 changed files with 213 additions and 206 deletions

View file

@ -9,25 +9,16 @@ use clap::Parser;
use clap::Subcommand;
use clap::ValueHint;
use human_panic::Metadata;
use kdl::KdlValue;
use miette::IntoDiagnostic;
use parsing::Definition;
use parsing::Record;
use tracing::debug;
use tracing::info;
use tracing::trace;
use tracing_subscriber::EnvFilter;
use trustfall::execute_query;
use trustfall::provider::field_property;
use trustfall::provider::resolve_coercion_with;
use trustfall::provider::resolve_neighbors_with;
use trustfall::provider::resolve_property_with;
use trustfall::provider::Adapter;
use trustfall::FieldValue;
use trustfall::Schema;
mod config;
mod parsing;
mod trustfall_plaixt;
#[derive(Debug, Parser)]
struct Args {
@ -73,7 +64,7 @@ async fn main() -> miette::Result<()> {
let records = parsing::load_records(root_folder, &definitions).await?;
let schema = to_schema(&definitions);
let schema = trustfall_plaixt::to_schema(&definitions);
match args.mode {
ArgMode::Query => {
@ -84,7 +75,7 @@ async fn main() -> miette::Result<()> {
let result = execute_query(
&schema,
Arc::new(PlaixtAdapter {
Arc::new(trustfall_plaixt::PlaixtAdapter {
records: records.clone(),
}),
&query,
@ -112,197 +103,3 @@ fn print_records(records: &[Record]) {
println!("}}")
}
}
fn to_schema(definitions: &BTreeMap<String, Vec<Definition>>) -> Schema {
let custom_schemas = definitions
.iter()
.map(|(name, def)| {
let fields = def
.last()
.unwrap()
.fields
.iter()
.map(|(name, def)| format!("{name}: {}!", def.trustfall_kind()))
.collect::<Vec<_>>()
.join("\n");
let field_type = format!("{name}Fields");
format!(
r#"
type {field_type} {{
{fields}
}}
type {name} implements Record {{
at: String!
kind: String!
fields: {field_type}!
}}
"#
)
})
.collect::<Vec<_>>()
.join("");
let schema = format!(
r#"schema {{
query: RootSchemaQuery
}}
{}
type RootSchemaQuery {{
RecordsAll: [Record!]!
}}
interface Record {{
at: String!,
kind: String!,
}}
{}
"#,
Schema::ALL_DIRECTIVE_DEFINITIONS,
custom_schemas
);
trace!(%schema, "Using schema");
Schema::parse(schema).unwrap()
}
struct PlaixtAdapter {
records: Vec<Record>,
}
#[derive(Clone, Debug)]
enum PlaixtVertex {
Record(Record),
Fields {
name: String,
values: BTreeMap<String, KdlValue>,
},
}
impl PlaixtVertex {
fn as_fields(&self) -> Option<&BTreeMap<String, KdlValue>> {
if let Self::Fields { values, .. } = self {
Some(values)
} else {
None
}
}
fn as_record(&self) -> Option<&Record> {
if let Self::Record(v) = self {
Some(v)
} else {
None
}
}
fn typename(&self) -> String {
match self {
PlaixtVertex::Record { .. } => "Record".to_string(),
PlaixtVertex::Fields { name, .. } => name.clone(),
}
}
}
impl<'a> Adapter<'a> for PlaixtAdapter {
type Vertex = PlaixtVertex;
fn resolve_starting_vertices(
&self,
edge_name: &Arc<str>,
_parameters: &trustfall::provider::EdgeParameters,
_resolve_info: &trustfall::provider::ResolveInfo,
) -> trustfall::provider::VertexIterator<'a, Self::Vertex> {
match edge_name.as_ref() {
"RecordsAll" => Box::new(self.records.clone().into_iter().map(PlaixtVertex::Record)),
_ => unreachable!(),
}
}
fn resolve_property<V: trustfall::provider::AsVertex<Self::Vertex> + 'a>(
&self,
contexts: trustfall::provider::ContextIterator<'a, V>,
type_name: &Arc<str>,
property_name: &Arc<str>,
_resolve_info: &trustfall::provider::ResolveInfo,
) -> trustfall::provider::ContextOutcomeIterator<'a, V, trustfall::FieldValue> {
match (type_name.as_ref(), property_name.as_ref()) {
(_, "__typename") => Box::new(contexts.map(|ctx| {
let value = match ctx.active_vertex() {
Some(_record) => _record.typename().into(),
None => FieldValue::Null,
};
(ctx, value)
})),
(_, "at") => resolve_property_with(
contexts,
field_property!(as_record, at, { at.to_string().into() }),
),
(_, "kind") => resolve_property_with(contexts, field_property!(as_record, kind)),
(name, field) => {
debug!(?name, ?field, "Asking for properties");
let field = field.to_string();
resolve_property_with(contexts, move |vertex| {
trace!(?vertex, ?field, "Getting property");
let fields = vertex.as_fields().unwrap();
match fields.get(&field).unwrap().clone() {
KdlValue::Bool(b) => FieldValue::Boolean(b),
KdlValue::Float(f) => FieldValue::Float64(f),
KdlValue::Null => FieldValue::Null,
KdlValue::Integer(i) => FieldValue::Int64(i.try_into().unwrap()),
KdlValue::String(s) => FieldValue::String(s.into()),
}
})
}
}
}
fn resolve_neighbors<V: trustfall::provider::AsVertex<Self::Vertex> + 'a>(
&self,
contexts: trustfall::provider::ContextIterator<'a, V>,
_type_name: &Arc<str>,
edge_name: &Arc<str>,
_parameters: &trustfall::provider::EdgeParameters,
_resolve_info: &trustfall::provider::ResolveEdgeInfo,
) -> trustfall::provider::ContextOutcomeIterator<
'a,
V,
trustfall::provider::VertexIterator<'a, Self::Vertex>,
> {
match edge_name.as_ref() {
"fields" => resolve_neighbors_with(contexts, |c| {
Box::new(
c.as_record()
.map(|r| PlaixtVertex::Fields {
name: format!("{}Fields", r.kind),
values: r.fields.clone(),
})
.into_iter(),
)
}),
_ => unreachable!(),
}
}
fn resolve_coercion<V: trustfall::provider::AsVertex<Self::Vertex> + 'a>(
&self,
contexts: trustfall::provider::ContextIterator<'a, V>,
type_name: &Arc<str>,
coerce_to_type: &Arc<str>,
_resolve_info: &trustfall::provider::ResolveInfo,
) -> trustfall::provider::ContextOutcomeIterator<'a, V, bool> {
debug!("Asking to coerce {type_name} into {coerce_to_type}");
let coerce_to_type = coerce_to_type.clone();
resolve_coercion_with(contexts, move |node| {
node.as_record()
.map(|r| r.kind == *coerce_to_type)
.unwrap_or(false)
})
}
}

View file

@ -0,0 +1,210 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use kdl::KdlValue;
use tracing::debug;
use tracing::trace;
use trustfall::provider::field_property;
use trustfall::provider::resolve_coercion_with;
use trustfall::provider::resolve_neighbors_with;
use trustfall::provider::resolve_property_with;
use trustfall::provider::Adapter;
use trustfall::FieldValue;
use trustfall::Schema;
use crate::parsing::Definition;
use crate::parsing::Record;
pub(crate) fn to_schema(definitions: &BTreeMap<String, Vec<Definition>>) -> Schema {
let custom_schemas = definitions
.iter()
.map(|(name, def)| {
let fields = def
.last()
.unwrap()
.fields
.iter()
.map(|(name, def)| format!("{name}: {}!", def.trustfall_kind()))
.collect::<Vec<_>>()
.join("\n");
let field_type = format!("{name}Fields");
format!(
r#"
type {field_type} {{
{fields}
}}
type {name} implements Record {{
at: String!
kind: String!
fields: {field_type}!
}}
"#
)
})
.collect::<Vec<_>>()
.join("");
let schema = format!(
r#"schema {{
query: RootSchemaQuery
}}
{}
type RootSchemaQuery {{
RecordsAll: [Record!]!
}}
interface Record {{
at: String!,
kind: String!,
}}
{}
"#,
Schema::ALL_DIRECTIVE_DEFINITIONS,
custom_schemas
);
trace!(%schema, "Using schema");
Schema::parse(schema).unwrap()
}
pub(crate) struct PlaixtAdapter {
pub(crate) records: Vec<Record>,
}
#[derive(Clone, Debug)]
pub(crate) enum PlaixtVertex {
Record(Record),
Fields {
name: String,
values: BTreeMap<String, KdlValue>,
},
}
impl PlaixtVertex {
pub(crate) fn as_fields(&self) -> Option<&BTreeMap<String, KdlValue>> {
if let Self::Fields { values, .. } = self {
Some(values)
} else {
None
}
}
pub(crate) fn as_record(&self) -> Option<&Record> {
if let Self::Record(v) = self {
Some(v)
} else {
None
}
}
pub(crate) fn typename(&self) -> String {
match self {
PlaixtVertex::Record { .. } => "Record".to_string(),
PlaixtVertex::Fields { name, .. } => name.clone(),
}
}
}
impl<'a> Adapter<'a> for PlaixtAdapter {
type Vertex = PlaixtVertex;
fn resolve_starting_vertices(
&self,
edge_name: &Arc<str>,
_parameters: &trustfall::provider::EdgeParameters,
_resolve_info: &trustfall::provider::ResolveInfo,
) -> trustfall::provider::VertexIterator<'a, Self::Vertex> {
match edge_name.as_ref() {
"RecordsAll" => Box::new(self.records.clone().into_iter().map(PlaixtVertex::Record)),
_ => unreachable!(),
}
}
fn resolve_property<V: trustfall::provider::AsVertex<Self::Vertex> + 'a>(
&self,
contexts: trustfall::provider::ContextIterator<'a, V>,
type_name: &Arc<str>,
property_name: &Arc<str>,
_resolve_info: &trustfall::provider::ResolveInfo,
) -> trustfall::provider::ContextOutcomeIterator<'a, V, trustfall::FieldValue> {
match (type_name.as_ref(), property_name.as_ref()) {
(_, "__typename") => Box::new(contexts.map(|ctx| {
let value = match ctx.active_vertex() {
Some(_record) => _record.typename().into(),
None => FieldValue::Null,
};
(ctx, value)
})),
(_, "at") => resolve_property_with(
contexts,
field_property!(as_record, at, { at.to_string().into() }),
),
(_, "kind") => resolve_property_with(contexts, field_property!(as_record, kind)),
(name, field) => {
debug!(?name, ?field, "Asking for properties");
let field = field.to_string();
resolve_property_with(contexts, move |vertex| {
trace!(?vertex, ?field, "Getting property");
let fields = vertex.as_fields().unwrap();
match fields.get(&field).unwrap().clone() {
KdlValue::Bool(b) => FieldValue::Boolean(b),
KdlValue::Float(f) => FieldValue::Float64(f),
KdlValue::Null => FieldValue::Null,
KdlValue::Integer(i) => FieldValue::Int64(i.try_into().unwrap()),
KdlValue::String(s) => FieldValue::String(s.into()),
}
})
}
}
}
fn resolve_neighbors<V: trustfall::provider::AsVertex<Self::Vertex> + 'a>(
&self,
contexts: trustfall::provider::ContextIterator<'a, V>,
_type_name: &Arc<str>,
edge_name: &Arc<str>,
_parameters: &trustfall::provider::EdgeParameters,
_resolve_info: &trustfall::provider::ResolveEdgeInfo,
) -> trustfall::provider::ContextOutcomeIterator<
'a,
V,
trustfall::provider::VertexIterator<'a, Self::Vertex>,
> {
match edge_name.as_ref() {
"fields" => resolve_neighbors_with(contexts, |c| {
Box::new(
c.as_record()
.map(|r| PlaixtVertex::Fields {
name: format!("{}Fields", r.kind),
values: r.fields.clone(),
})
.into_iter(),
)
}),
_ => unreachable!(),
}
}
fn resolve_coercion<V: trustfall::provider::AsVertex<Self::Vertex> + 'a>(
&self,
contexts: trustfall::provider::ContextIterator<'a, V>,
type_name: &Arc<str>,
coerce_to_type: &Arc<str>,
_resolve_info: &trustfall::provider::ResolveInfo,
) -> trustfall::provider::ContextOutcomeIterator<'a, V, bool> {
debug!("Asking to coerce {type_name} into {coerce_to_type}");
let coerce_to_type = coerce_to_type.clone();
resolve_coercion_with(contexts, move |node| {
node.as_record()
.map(|r| r.kind == *coerce_to_type)
.unwrap_or(false)
})
}
}