Create custom querier for sourcedb

crates/pyrefly_build/src/lib.rs

@@ -27,33 +27,39 @@
 #![feature(if_let_guard)]
 
 use std::path::PathBuf;
+use std::sync::Arc;
 
 use serde::Deserialize;
 use serde::Serialize;
 
-use crate::buck::query::BxlArgs;
-
-pub mod buck;
 pub mod handle;
-pub mod map_db;
 pub mod source_db;
+pub use source_db::SourceDatabase;
+mod query;
+
+use crate::query::SourceDbQuerier;
+use crate::query::buck::BxlArgs;
+use crate::query::buck::BxlQuerier;
+use crate::query::custom::CustomQuerier;
+use crate::query::custom::CustomQueryArgs;
+use crate::source_db::query_source_db::QuerySourceDatabase;
 
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
 #[serde(rename_all = "kebab-case", tag = "type")]
 pub enum BuildSystem {
     Buck(BxlArgs),
+    Custom(CustomQueryArgs),
 }
 
 impl BuildSystem {
     pub fn get_source_db(
         &self,
         config_root: PathBuf,
     ) -> Box<dyn source_db::SourceDatabase + 'static> {
-        match &self {
-            Self::Buck(args) => Box::new(buck::bxl::BuckSourceDatabase::new(
-                config_root,
-                args.clone(),
-            )),
-        }
+        let querier: Arc<dyn SourceDbQuerier> = match &self {
+            Self::Buck(args) => Arc::new(BxlQuerier::new(args.clone())),
+            Self::Custom(args) => Arc::new(CustomQuerier::new(args.clone())),
+        };
+        Box::new(QuerySourceDatabase::new(config_root, querier))
     }
 }

crates/pyrefly_build/src/query/buck.rs

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use std::fmt::Debug;
+use std::process::Command;
+
+use serde::Deserialize;
+use serde::Serialize;
+
+use crate::query::SourceDbQuerier;
+
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, Default)]
+#[serde(rename_all = "kebab-case")]
+pub struct BxlArgs {
+    isolation_dir: Option<String>,
+    extras: Option<Vec<String>>,
+}
+
+#[derive(Debug)]
+pub struct BxlQuerier(BxlArgs);
+
+impl BxlQuerier {
+    pub fn new(args: BxlArgs) -> Self {
+        Self(args)
+    }
+}
+
+impl SourceDbQuerier for BxlQuerier {
+    fn construct_command(&self) -> Command {
+        let mut cmd = Command::new("buck2");
+        if let Some(isolation_dir) = &self.0.isolation_dir {
+            cmd.arg("--isolation-dir");
+            cmd.arg(isolation_dir);
+        }
+        cmd.arg("bxl");
+        cmd.arg("--reuse-current-config");
+        if let Some(metadata) = &self.0.extras {
+            cmd.args(metadata);
+        }
+        cmd.arg("prelude//python/sourcedb/pyrefly.bxl:main");
+        cmd
+    }
+}

crates/pyrefly_build/src/query/custom.rs

@@ -0,0 +1,52 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use std::fmt::Debug;
+use std::process::Command;
+
+use serde::Deserialize;
+use serde::Serialize;
+use vec1::Vec1;
+
+use crate::query::SourceDbQuerier;
+
+/// Args and settings for querying a custom source DB.
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, Default)]
+#[serde(rename_all = "kebab-case")]
+pub struct CustomQueryArgs {
+    /// The command to run, requires at least the command to execute.
+    /// Pyrefly will call this in the form `<command> @<argfile>`,
+    /// where `<argfile>` has the format
+    /// ```
+    /// --
+    /// <arg-flag>
+    /// <arg>
+    /// ...
+    /// ```
+    /// and `<arg-flag>` is either `--file` or `--target`, depeding on the type
+    /// of `<arg>`
+    /// and `<arg>` is an absolute path to a file or a build system's target.
+    command: Vec1<String>,
+}
+
+/// A querier allowing for a custom command when querying and constructing source DB.
+#[derive(Debug)]
+pub struct CustomQuerier(CustomQueryArgs);
+
+impl CustomQuerier {
+    pub fn new(args: CustomQueryArgs) -> Self {
+        Self(args)
+    }
+}
+
+impl SourceDbQuerier for CustomQuerier {
+    fn construct_command(&self) -> Command {
+        let mut cmd = Command::new(self.0.command.first());
+        cmd.args(self.0.command.iter().skip(1));
+        cmd
+    }
+}

crates/pyrefly_build/src/buck/query.rs → crates/pyrefly_build/src/query/mod.rs

@@ -7,8 +7,8 @@
 
 use std::ffi::OsStr;
 use std::ffi::OsString;
-use std::fmt::Debug;
-use std::io::Write;
+use std::fmt;
+use std::io::Write as _;
 use std::path::Path;
 use std::path::PathBuf;
 use std::process::Command;
@@ -18,14 +18,16 @@ use dupe::Dupe as _;
 use pyrefly_python::module_name::ModuleName;
 use pyrefly_python::sys_info::SysInfo;
 use serde::Deserialize;
-use serde::Serialize;
 use starlark_map::small_map::SmallMap;
 use starlark_map::small_set::SmallSet;
 use tempfile::NamedTempFile;
 use vec1::Vec1;
 
 use crate::source_db::Target;
 
+pub mod buck;
+pub mod custom;
+
 /// An enum representing something that has been included by the build system, and
 /// which the build system should query for when building the sourcedb.
 #[derive(Debug, PartialEq, Eq, Hash)]
@@ -40,76 +42,68 @@ impl Include {
         Self::Path(path)
     }
 
-    fn to_bxl_args(&self) -> impl Iterator<Item = &OsStr> {
+    fn to_cli_arg(&self) -> impl Iterator<Item = &OsStr> {
         match self {
             Include::Target(target) => [OsStr::new("--target"), target.to_os_str()].into_iter(),
             Include::Path(path) => [OsStr::new("--file"), path.as_os_str()].into_iter(),
         }
     }
 }
 
-#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, Default)]
-#[serde(rename_all = "kebab-case")]
-pub struct BxlArgs {
-    isolation_dir: Option<String>,
-    extras: Option<Vec<String>>,
-}
+pub trait SourceDbQuerier: Send + Sync + fmt::Debug {
+    fn query_source_db(
+        &self,
+        files: &SmallSet<Include>,
+        cwd: &Path,
+    ) -> anyhow::Result<TargetManifestDatabase> {
+        if files.is_empty() {
+            return Ok(TargetManifestDatabase {
+                db: SmallMap::new(),
+                root: cwd.to_path_buf(),
+            });
+        }
 
-pub(crate) fn query_source_db<'a>(
-    files: impl Iterator<Item = &'a Include>,
-    cwd: &Path,
-    bxl_args: &BxlArgs,
-) -> anyhow::Result<TargetManifestDatabase> {
-    let mut files = files.peekable();
-    if files.peek().is_none() {
-        return Ok(TargetManifestDatabase {
-            db: SmallMap::new(),
-            root: cwd.to_path_buf(),
+        let mut argfile =
+            NamedTempFile::with_prefix("pyrefly_build_query_").with_context(|| {
+                "Failed to create temporary argfile for querying source DB".to_owned()
+            })?;
+        let mut argfile_args = OsString::from("--");
+        files.iter().flat_map(Include::to_cli_arg).for_each(|arg| {
+            argfile_args.push("\n");
+            argfile_args.push(arg);
         });
-    }
 
-    let mut argfile = NamedTempFile::with_prefix("pyrefly_buck_query_")
-        .with_context(|| "Failed to create temporary argfile for querying Buck".to_owned())?;
-    let mut argfile_args = OsString::from("--");
-    files.flat_map(Include::to_bxl_args).for_each(|arg| {
-        argfile_args.push("\n");
-        argfile_args.push(arg);
-    });
-
-    argfile
-        .as_file_mut()
-        .write_all(argfile_args.as_encoded_bytes())
-        .with_context(|| "Could not write to argfile when querying Buck".to_owned())?;
-
-    let mut cmd = Command::new("buck2");
-    if let Some(isolation_dir) = &bxl_args.isolation_dir {
-        cmd.arg("--isolation-dir");
-        cmd.arg(isolation_dir);
-    }
-    cmd.arg("bxl");
-    cmd.arg("--reuse-current-config");
-    if let Some(metadata) = &bxl_args.extras {
-        cmd.args(metadata);
-    }
-    cmd.arg("prelude//python/sourcedb/pyrefly.bxl:main");
-    cmd.arg(format!("@{}", argfile.path().display()));
-    cmd.current_dir(cwd);
-
-    let result = cmd.output()?;
-    if !result.status.success() {
-        let stdout = String::from_utf8(result.stdout)
-            .unwrap_or_else(|_| "<Failed to parse stdout from Buck source db query>".to_owned());
-        let stderr = String::from_utf8(result.stderr)
-            .unwrap_or_else(|_| "<Failed to parse stderr from Buck source db query>".to_owned());
-
-        return Err(anyhow::anyhow!(
-            "Buck source db query failed...\nSTDOUT: {stdout}\nSTDERR: {stderr}"
-        ));
+        argfile
+            .as_file_mut()
+            .write_all(argfile_args.as_encoded_bytes())
+            .with_context(|| "Could not write to argfile when querying source DB".to_owned())?;
+
+        let mut cmd = self.construct_command();
+        cmd.arg(format!("@{}", argfile.path().display()));
+        cmd.current_dir(cwd);
+
+        let result = cmd.output()?;
+        if !result.status.success() {
+            let stdout = String::from_utf8(result.stdout)
+                .unwrap_or_else(|_| "<Failed to parse stdout from source DB query>".to_owned());
+            let stderr = String::from_utf8(result.stderr).unwrap_or_else(|_| {
+                "<Failed to parse stderr from Buck source DB query>".to_owned()
+            });
+
+            return Err(anyhow::anyhow!(
+                "Source DB query failed...\nSTDOUT: {stdout}\nSTDERR: {stderr}"
+            ));
+        }
+
+        serde_json::from_slice(&result.stdout).with_context(|| {
+            format!(
+                "Failed to construct valid `TargetManifestDatabase` from querier result. Command run: `{}`",
+                cmd.get_program().display(),
+            )
+        })
     }
 
-    serde_json::from_slice(&result.stdout).with_context(|| {
-        "Failed to construct valid `TargetManifestDatabase` from BXL query result".to_owned()
-    })
+    fn construct_command(&self) -> Command;
 }
 
 #[derive(Debug, PartialEq, Eq, Deserialize, Clone)]

crates/pyrefly_build/src/source_db.rs → crates/pyrefly_build/src/source_db/mod.rs

@@ -24,6 +24,10 @@ use static_interner::Interner;
 
 use crate::handle::Handle;
 
+pub mod buck_check;
+pub mod map_db;
+pub(crate) mod query_source_db;
+
 // We're interning `Target`s, since they'll be duplicated all over the place,
 // and it would be nice to have something that implements `Copy`.
 // We choose Interning over `Arc`, since we want to make sure all `Target`s