From 58b338ec2999d36441dd9d3d96f301533fac044f Mon Sep 17 00:00:00 2001
From: Nikita Lapkov <5737185+laplab@users.noreply.github.com>
Date: Mon, 26 Feb 2024 14:36:54 +0000
Subject: [PATCH] feat: store GraphViz query graphs in a separate folder
 (#4720)

---
 .gitignore                                    |  3 ++
 .../core/src/query_graph_builder/builder.rs   | 46 ++++++++++++++-----
 2 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/.gitignore b/.gitignore
index 739ad7005546..d401ff68f180 100644
--- a/.gitignore
+++ b/.gitignore
@@ -56,3 +56,6 @@ package-lock.json
 
 # Human readable Wasm output
 *.wat
+
+# Folder with GraphViz representation of query graphs.
+.query_graphs
diff --git a/query-engine/core/src/query_graph_builder/builder.rs b/query-engine/core/src/query_graph_builder/builder.rs
index 1152c4974a45..239fecc55530 100644
--- a/query-engine/core/src/query_graph_builder/builder.rs
+++ b/query-engine/core/src/query_graph_builder/builder.rs
@@ -1,12 +1,30 @@
+use std::{
+    fmt,
+    sync::atomic::{AtomicU32, Ordering},
+};
+
+use chrono::Local;
 use once_cell::sync::Lazy;
-use std::{fmt, fs::File, io::Write};
 
 use super::*;
 use crate::{query_document::*, query_graph::*, schema::*, IrSerializer};
 
-pub static PRISMA_RENDER_DOT_FILE: Lazy<bool> = Lazy::new(|| match std::env::var("PRISMA_RENDER_DOT_FILE") {
-    Ok(enabled) => enabled == *("true") || enabled == *("1"),
-    Err(_) => false,
+static PRISMA_DOT_PATH: Lazy<Option<String>> = Lazy::new(|| {
+    // Query graphs are saved only if `PRISMA_RENDER_DOT_FILE` env variable is set.
+    if !matches!(std::env::var("PRISMA_RENDER_DOT_FILE").as_deref(), Ok("true") | Ok("1")) {
+        return None;
+    }
+    // If `WORKSPACE_ROOT` env variable is defined, we save query graphs there. This ensures that
+    // there is a single central place to store query graphs, no matter which target is running.
+    // If this env variable is not defined, we save query graphs in the current directory.
+    let base_path = std::env::var("WORKSPACE_ROOT")
+        .ok()
+        .filter(|path| !path.is_empty())
+        .unwrap_or(".".into());
+    let time = Local::now().format("%Y-%m-%d %H:%M:%S");
+    let path = format!("{base_path}/.query_graphs/{time}");
+    std::fs::create_dir_all(&path).expect("Could not create directory to store query graphs");
+    Some(path)
 });
 
 pub struct QueryGraphBuilder<'a> {
@@ -57,6 +75,18 @@ impl<'a> QueryGraphBuilder<'a> {
 
         if field_pair.schema_field.query_info().is_some() {
             let graph = self.dispatch_build(field_pair)?;
+
+            // Used to debug generated graph.
+            if let Some(path) = &*PRISMA_DOT_PATH {
+                static COUNTER: AtomicU32 = AtomicU32::new(1);
+                let current = COUNTER.fetch_add(1, Ordering::Relaxed);
+                std::fs::write(
+                    format!("{}/{}_{}.graphviz", path, current, serializer.key),
+                    graph.to_graphviz(),
+                )
+                .unwrap();
+            }
+
             Ok((graph, serializer))
         } else {
             Err(QueryGraphBuilderError::SchemaError(format!(
@@ -101,14 +131,6 @@ impl<'a> QueryGraphBuilder<'a> {
         graph.finalize(self.query_schema.capabilities())?;
         trace!("{}", graph);
 
-        // Used to debug generated graph.
-        if *PRISMA_RENDER_DOT_FILE {
-            let mut f = File::create("graph.dot").unwrap();
-            let output = graph.to_graphviz();
-
-            f.write_all(output.as_bytes()).unwrap();
-        }
-
         Ok(graph)
     }