diff --git a/Cargo.lock b/Cargo.lock index 4d16ee2496c8..01304f86e0cf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3419,6 +3419,18 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" +[[package]] +name = "pretty" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b55c4d17d994b637e2f4daf6e5dc5d660d209d5642377d675d7a1c3ab69fa579" +dependencies = [ + "arrayvec 0.5.2", + "termcolor", + "typed-arena", + "unicode-width", +] + [[package]] name = "pretty-hex" version = "0.3.0" @@ -3769,6 +3781,7 @@ dependencies = [ "lru 0.7.8", "once_cell", "petgraph", + "pretty", "prisma-metrics", "psl", "quaint", @@ -5891,6 +5904,12 @@ dependencies = [ "static_assertions", ] +[[package]] +name = "typed-arena" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a" + [[package]] name = "typed-builder" version = "0.10.0" diff --git a/query-engine/core/Cargo.toml b/query-engine/core/Cargo.toml index d8f97ec2cec1..a6aa37e1f63a 100644 --- a/query-engine/core/Cargo.toml +++ b/query-engine/core/Cargo.toml @@ -23,6 +23,7 @@ petgraph = "0.4" query-structure = { path = "../query-structure", features = [ "default_generators", ] } +pretty = { version = "0.12", features = ["termcolor"] } prisma-metrics = { path = "../../libs/metrics", optional = true } serde.workspace = true serde_json.workspace = true diff --git a/query-engine/core/src/compiler/expression.rs b/query-engine/core/src/compiler/expression.rs index 5669e82a695b..0cb342be9881 100644 --- a/query-engine/core/src/compiler/expression.rs +++ b/query-engine/core/src/compiler/expression.rs @@ -1,4 +1,7 @@ -use itertools::Itertools; +use pretty::{ + termcolor::{Color, ColorSpec}, + DocAllocator, DocBuilder, +}; use query_structure::PrismaValue; use serde::Serialize; @@ -88,122 +91,201 @@ pub enum Expression { MapField { field: String, records: Box }, } -impl Expression { - fn display(&self, f: &mut std::fmt::Formatter<'_>, level: usize) -> std::fmt::Result { - let indent = " ".repeat(level); - - match self { - Self::Seq(exprs) => { - writeln!(f, "{indent}{{")?; - for expr in exprs { - expr.display(f, level + 1)?; - writeln!(f, ";")?; - } - write!(f, "{indent}}}")?; - } - - Self::Get { name } => { - write!(f, "{indent}get {name}")?; - } - - Self::Let { bindings, expr } => { - writeln!(f, "{indent}let")?; - for Binding { name, expr } in bindings { - writeln!(f, "{indent} {name} =")?; - expr.display(f, level + 2)?; - writeln!(f, ";")?; - } - writeln!(f, "{indent}in")?; - expr.display(f, level + 1)?; - } - - Self::GetFirstNonEmpty { names } => { - write!(f, "{indent}getFirstNonEmpty")?; - for name in names { - write!(f, " {}", name)?; - } - } - - Self::Query(query) => self.display_query("query", query, f, level)?, - - Self::Execute(query) => self.display_query("execute", query, f, level)?, - - Self::Reverse(expr) => { - writeln!(f, "{indent}reverse (")?; - expr.display(f, level + 1)?; - write!(f, "{indent})")?; - } - - Self::Sum(exprs) => self.display_function("sum", exprs, f, level)?, - - Self::Concat(exprs) => self.display_function("concat", exprs, f, level)?, - - Self::Unique(expr) => { - writeln!(f, "{indent}unique (")?; - expr.display(f, level + 1)?; - write!(f, "{indent})")?; - } - - Self::Required(expr) => { - writeln!(f, "{indent}required (")?; - expr.display(f, level + 1)?; - write!(f, "{indent})")?; - } - - Self::Join { parent, children } => { - writeln!(f, "{indent}join (")?; - parent.display(f, level + 1)?; - for nested in children { - let left = nested.on.iter().map(|(l, _)| l).cloned().join(", "); - let right = nested.on.iter().map(|(_, r)| r).cloned().join(", "); - writeln!(f, "\n{indent} with (")?; - nested.child.display(f, level + 2)?; - writeln!(f, "\n{indent} ) on left.{left} = right.{right},")?; - } - write!(f, "{indent})")?; - } - - Self::MapField { field, records } => { - writeln!(f, "{indent}mapField {field} (")?; - records.display(f, level + 1)?; - write!(f, "\n{indent})")?; - } - } +#[derive(thiserror::Error, Debug)] +pub enum PrettyPrintError { + #[error("{0}")] + IoError(#[from] std::io::Error), + #[error("{0}")] + FromUtf8Error(#[from] std::string::FromUtf8Error), +} - Ok(()) +impl Expression { + pub fn pretty_print(&self, color: bool, width: usize) -> Result { + let arena = pretty::Arena::new(); + let doc = self.to_doc(&arena); + + let mut buf = if color { + pretty::termcolor::Buffer::ansi() + } else { + pretty::termcolor::Buffer::no_color() + }; + + doc.render_colored(width, &mut buf)?; + Ok(String::from_utf8(buf.into_inner())?) } - fn display_query( - &self, - op: &str, - db_query: &DbQuery, - f: &mut std::fmt::Formatter<'_>, - level: usize, - ) -> std::fmt::Result { - let indent = " ".repeat(level); - let DbQuery { query, params } = db_query; - write!(f, "{indent}{op} (\n{indent} {query}\n{indent}) with {params:?}") - } + fn to_doc<'a, D>(&'a self, d: &'a D) -> DocBuilder<'a, D, ColorSpec> + where + D: DocAllocator<'a, ColorSpec>, + D::Doc: Clone, + { + let color_kw = || ColorSpec::new().set_fg(Some(Color::Blue)).clone(); + let color_fn = || ColorSpec::new().set_underline(true).clone(); + let color_var = || ColorSpec::new().set_bold(true).clone(); + let color_lit = || ColorSpec::new().set_italic(true).set_fg(Some(Color::Green)).clone(); + + let format_query = |tag: &'static str, db_query: &'a DbQuery| { + d.text(tag) + .annotate(color_kw()) + .append(d.softline()) + .append( + d.reflow(&db_query.query) + .align() + .enclose("«", "»") + .annotate(color_lit()), + ) + .append(d.line()) + .append(d.text("with params").annotate(color_kw())) + .append(d.space()) + .append( + d.intersperse( + db_query.params.iter().map(|param| match param { + PrismaValue::Placeholder { name, r#type } => d.text("var").annotate(color_kw()).append( + d.text(name) + .annotate(color_var()) + .append(d.space()) + .append(d.text("as").annotate(color_kw())) + .append(d.space()) + .append(match r#type { + query_structure::PlaceholderType::Array(inner) => format!("{inner:?}[]"), + _ => format!("{type:?}"), + }) + .parens(), + ), + _ => d + .text("const") + .annotate(color_kw()) + .append(d.text(format!("{param:?}")).annotate(color_lit()).parens()), + }), + d.text(",").append(d.softline()), + ) + .align() + .brackets(), + ) + .align() + }; + + let format_function = |name: &'static str, args: &'a [Expression]| { + d.text(name).annotate(color_fn()).append(d.space()).append( + d.intersperse(args.iter().map(|expr| expr.to_doc(d)), d.space()) + .parens(), + ) + }; + + let format_unary_function = |name: &'static str, arg: &'a Expression| { + d.text(name) + .annotate(color_fn()) + .append(d.space()) + .append(arg.to_doc(d).parens()) + }; - fn display_function( - &self, - name: &str, - args: &[Expression], - f: &mut std::fmt::Formatter<'_>, - level: usize, - ) -> std::fmt::Result { - let indent = " ".repeat(level); - write!(f, "{indent}{name} (")?; - for arg in args { - arg.display(f, level + 1)?; - writeln!(f, ",")?; + match self { + Expression::Seq(vec) => d.intersperse(vec.iter().map(|expr| expr.to_doc(d)), d.text(";").append(d.line())), + + Expression::Get { name } => d + .text("get") + .annotate(color_kw()) + .append(d.space()) + .append(d.text(name).annotate(color_var())), + + Expression::Let { bindings, expr } => d + .text("let") + .annotate(color_kw()) + .append(d.softline()) + .append( + d.intersperse( + bindings.iter().map(|binding| { + d.text(&binding.name) + .annotate(color_var()) + .append(d.space()) + .append("=") + .append(d.softline()) + .append(binding.expr.to_doc(d)) + }), + d.text(";").append(d.line()), + ) + .align(), + ) + .append(d.line()) + .append(d.text("in").annotate(color_kw())) + .append(d.softline()) + .append(expr.to_doc(d).align()), + + Expression::GetFirstNonEmpty { names } => d + .text("getFirstNonEmpty") + .annotate(color_fn()) + .append(d.intersperse(names.iter().map(|name| d.text(name).annotate(color_var())), d.space())), + + Expression::Query(db_query) => format_query("query", db_query), + + Expression::Execute(db_query) => format_query("execute", db_query), + + Expression::Reverse(expression) => format_unary_function("reverse", expression), + + Expression::Sum(vec) => format_function("sum", vec), + + Expression::Concat(vec) => format_function("concat", vec), + + Expression::Unique(expression) => format_unary_function("unique", expression), + + Expression::Required(expression) => format_unary_function("required", expression), + + Expression::Join { parent, children } => d + .text("join") + .annotate(color_kw()) + .append(d.space()) + .append(parent.to_doc(d).parens()) + .append(d.line()) + .append(d.text("with").annotate(color_kw())) + .append(d.space()) + .append( + d.intersperse( + children.iter().map(|join| { + join.child + .to_doc(d) + .parens() + .append(d.space()) + .append(d.text("on").annotate(color_kw())) + .append(d.space()) + .append(d.intersperse( + join.on.iter().map(|(l, r)| { + d.text("left") + .annotate(color_kw()) + .append(".") + .append(d.text(l).annotate(color_var())) + .parens() + .append(d.space()) + .append("=") + .append(d.space()) + .append( + d.text("right") + .annotate(color_kw()) + .append(".") + .append(d.text(r).annotate(color_var())) + .parens(), + ) + }), + d.text(", "), + )) + }), + d.text(",").append(d.line()), + ) + .align(), + ), + + Expression::MapField { field, records } => d + .text("mapField") + .annotate(color_fn()) + .append(d.space()) + .append(d.text(field).double_quotes().annotate(color_lit())) + .append(d.space()) + .append(records.to_doc(d).parens()), } - write!(f, ")") } } impl std::fmt::Display for Expression { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - self.display(f, 0) + self.pretty_print(false, 80).map_err(|_| std::fmt::Error)?.fmt(f) } } diff --git a/query-engine/query-engine/examples/compiler.rs b/query-engine/query-engine/examples/compiler.rs index 950fd99f255e..6292910ea3d7 100644 --- a/query-engine/query-engine/examples/compiler.rs +++ b/query-engine/query-engine/examples/compiler.rs @@ -22,6 +22,7 @@ pub fn main() -> anyhow::Result<()> { // select: { // val: true, // posts: true, + // profile: true, // } // }) let query: JsonSingleQuery = serde_json::from_value(json!({ @@ -43,6 +44,12 @@ pub fn main() -> anyhow::Result<()> { "selection": { "$scalars": true } + }, + "profile": { + "arguments": {}, + "selection": { + "$scalars": true + } } } } @@ -61,7 +68,7 @@ pub fn main() -> anyhow::Result<()> { let expr = query_core::compiler::translate(graph)?; - println!("{expr}"); + println!("{}", expr.pretty_print(true, 80)?); Ok(()) } diff --git a/query-engine/query-engine/examples/schema.prisma b/query-engine/query-engine/examples/schema.prisma index ab9cd218da49..ff51ae3ddfe0 100644 --- a/query-engine/query-engine/examples/schema.prisma +++ b/query-engine/query-engine/examples/schema.prisma @@ -8,11 +8,17 @@ datasource db { } model User { - id String @id @default(cuid()) - email String @unique - name String? - posts Post[] - val Int? + id String @id @default(cuid()) + email String @unique + name String? + posts Post[] + val Int? + profile Profile? +} + +model Profile { + userId String @id + user User @relation(fields: [userId], references: [id]) } model Post {