WIP Support query_string macro attribute

dphm · dphm · commit 51cf91d6f39c · 2024-01-04T17:06:03.000-05:00
diff --git a/graphql_client_codegen/src/lib.rs b/graphql_client_codegen/src/lib.rs
@@ -44,29 +44,24 @@ impl std::error::Error for GeneralError {}
 
 type BoxError = Box<dyn std::error::Error + Send + Sync + 'static>;
 type CacheMap<T> = std::sync::Mutex<BTreeMap<std::path::PathBuf, T>>;
+type CachedQuery = (String, graphql_parser::query::Document<'static, String>);
 
 lazy_static! {
     static ref SCHEMA_CACHE: CacheMap<schema::Schema> = CacheMap::default();
-    static ref QUERY_CACHE: CacheMap<(String, graphql_parser::query::Document<'static, String>)> =
-        CacheMap::default();
+    static ref QUERY_CACHE: CacheMap<CachedQuery> = CacheMap::default();
 }
 
-/// Generates Rust code given a query document, a schema and options.
-pub fn generate_module_token_stream(
-    query_path: std::path::PathBuf,
-    schema_path: &std::path::Path,
-    options: GraphQLClientCodegenOptions,
-) -> Result<TokenStream, BoxError> {
+fn get_schema(schema_path: &std::path::Path) -> Result<schema::Schema, BoxError> {
     use std::collections::btree_map;
 
     let schema_extension = schema_path
         .extension()
         .and_then(std::ffi::OsStr::to_str)
         .unwrap_or("INVALID");
-    let schema_string;
+    let schema_string: String;
 
     // Check the schema cache.
-    let schema: schema::Schema = {
+    let cached_schema: schema::Schema = {
         let mut lock = SCHEMA_CACHE.lock().expect("schema cache is poisoned");
         match lock.entry(schema_path.to_path_buf()) {
             btree_map::Entry::Occupied(o) => o.get().clone(),
@@ -89,37 +84,68 @@ pub fn generate_module_token_stream(
         }
     };
 
-    // We need to qualify the query with the path to the crate it is part of
-    let (query_string, query) = {
-        let mut lock = QUERY_CACHE.lock().expect("query cache is poisoned");
-        match lock.entry(query_path) {
-            btree_map::Entry::Occupied(o) => o.get().clone(),
-            btree_map::Entry::Vacant(v) => {
-                let query_string = read_file(v.key())?;
-                let query = graphql_parser::parse_query(&query_string)
-                    .map_err(|err| GeneralError(format!("Query parser error: {}", err)))?
-                    .into_static();
-                v.insert((query_string, query)).clone()
-            }
+    Ok(cached_schema)
+}
+
+fn get_query(query_path: std::path::PathBuf) -> Result<CachedQuery, BoxError> {
+    use std::collections::btree_map;
+
+    let mut lock = QUERY_CACHE.lock().expect("query cache is poisoned");
+    let cached_query = match lock.entry(query_path) {
+        btree_map::Entry::Occupied(o) => o.get().clone(),
+        btree_map::Entry::Vacant(v) => {
+            let query_string = read_file(v.key())?;
+            let query_document = graphql_parser::parse_query(&query_string)
+                .map_err(|err| GeneralError(format!("Query parser error: {}", err)))?
+                .into_static();
+            v.insert((query_string, query_document)).clone()
         }
     };
 
-    let query = crate::query::resolve(&schema, &query)?;
+    Ok(cached_query)
+}
+
+/// Generates Rust code given a path to a query file, a path to a schema file, and options.
+pub fn generate_module_token_stream(
+    query_path: std::path::PathBuf,
+    schema_path: &std::path::Path,
+    options: GraphQLClientCodegenOptions,
+) -> Result<TokenStream, BoxError> {
+    let (query_string, _) = get_query(query_path)?;
+
+    generate_module_token_stream_from_string(query_string, schema_path, options)
+}
+
+/// Generates Rust code given a query string, a path to a schema file, and options.
+pub fn generate_module_token_stream_from_string(
+    query_string: String,
+    schema_path: &std::path::Path,
+    options: GraphQLClientCodegenOptions,
+) -> Result<TokenStream, BoxError> {
+    let schema = get_schema(schema_path)?;
+    let query_document = graphql_parser::parse_query(&query_string)
+        .map_err(|err| GeneralError(format!("Query parser error: {}", err)))?
+        .into_static();
+
+    // We need to qualify the query with the path to the crate it is part of
+    let generated_query = crate::query::resolve(&schema, &query_document)?;
 
     // Determine which operation we are generating code for. This will be used in operationName.
     let operations = options
         .operation_name
         .as_ref()
-        .and_then(|operation_name| query.select_operation(operation_name, *options.normalization()))
+        .and_then(|operation_name| {
+            generated_query.select_operation(operation_name, *options.normalization())
+        })
         .map(|op| vec![op]);
 
     let operations = match (operations, &options.mode) {
         (Some(ops), _) => ops,
-        (None, &CodegenMode::Cli) => query.operations().collect(),
+        (None, &CodegenMode::Cli) => generated_query.operations().collect(),
         (None, &CodegenMode::Derive) => {
             return Err(GeneralError(derive_operation_not_found_error(
                 options.struct_ident(),
-                &query,
+                &generated_query,
             ))
             .into());
         }
@@ -132,7 +158,7 @@ pub fn generate_module_token_stream(
         let generated = generated_module::GeneratedModule {
             query_string: query_string.as_str(),
             schema: &schema,
-            resolved_query: &query,
+            resolved_query: &generated_query,
             operation: &operation.1.name,
             options: &options,
         }
diff --git a/graphql_query_derive/src/lib.rs b/graphql_query_derive/src/lib.rs
@@ -4,7 +4,8 @@ extern crate proc_macro;
 mod attributes;
 
 use graphql_client_codegen::{
-    generate_module_token_stream, CodegenMode, GraphQLClientCodegenOptions,
+    generate_module_token_stream, generate_module_token_stream_from_string, CodegenMode,
+    GraphQLClientCodegenOptions,
 };
 use std::{
     env,
@@ -26,38 +27,44 @@ fn graphql_query_derive_inner(
 ) -> Result<proc_macro::TokenStream, syn::Error> {
     let input = TokenStream::from(input);
     let ast = syn::parse2(input)?;
-    let (query_path, schema_path) = build_query_and_schema_path(&ast)?;
-    let options = build_graphql_client_derive_options(&ast, query_path.clone())?;
-
-    generate_module_token_stream(query_path, &schema_path, options)
-        .map(Into::into)
-        .map_err(|err| {
-            syn::Error::new_spanned(
-                ast,
-                format!("Failed to generate GraphQLQuery impl: {}", err),
-            )
-        })
-}
 
-fn build_query_and_schema_path(input: &syn::DeriveInput) -> Result<(PathBuf, PathBuf), syn::Error> {
     let cargo_manifest_dir = env::var("CARGO_MANIFEST_DIR").map_err(|_err| {
         syn::Error::new_spanned(
-            input,
+            &ast,
             "Error checking that the CARGO_MANIFEST_DIR env variable is defined.",
         )
     })?;
 
-    let query_path = attributes::extract_attr(input, "query_path")?;
-    let query_path = format!("{}/{}", cargo_manifest_dir, query_path);
-    let query_path = Path::new(&query_path).to_path_buf();
-    let schema_path = attributes::extract_attr(input, "schema_path")?;
+    let schema_path = attributes::extract_attr(&ast, "schema_path")?;
     let schema_path = Path::new(&cargo_manifest_dir).join(schema_path);
-    Ok((query_path, schema_path))
+
+    let query_string = attributes::extract_attr(&ast, "query_string").ok();
+    let module_token_stream = match query_string {
+        Some(query_string) => {
+            let options = build_graphql_client_derive_options(&ast, None)?;
+            generate_module_token_stream_from_string(query_string, schema_path.as_path(), options)
+        }
+        None => {
+            let query_path = attributes::extract_attr(&ast, "query_path")?;
+            let query_path = format!("{}/{}", cargo_manifest_dir, query_path);
+            let query_path = Path::new(&query_path).to_path_buf();
+            let options = build_graphql_client_derive_options(&ast, Some(query_path.clone()))?;
+
+            generate_module_token_stream(query_path, &schema_path, options)
+        }
+    };
+
+    module_token_stream.map(Into::into).map_err(|err| {
+        syn::Error::new_spanned(
+            ast,
+            format!("Failed to generate GraphQLQuery impl: {}", err),
+        )
+    })
 }
 
 fn build_graphql_client_derive_options(
     input: &syn::DeriveInput,
-    query_path: PathBuf,
+    query_path: Option<PathBuf>,
 ) -> Result<GraphQLClientCodegenOptions, syn::Error> {
     let variables_derives = attributes::extract_attr(input, "variables_derives").ok();
     let response_derives = attributes::extract_attr(input, "response_derives").ok();
@@ -67,7 +74,11 @@ fn build_graphql_client_derive_options(
     let skip_serializing_none: bool = attributes::extract_skip_serializing_none(input);
 
     let mut options = GraphQLClientCodegenOptions::new(CodegenMode::Derive);
-    options.set_query_file(query_path);
+
+    if let Some(query_path) = query_path {
+        options.set_query_file(query_path);
+    }
+
     options.set_fragments_other_variant(fragments_other_variant);
     options.set_skip_serializing_none(skip_serializing_none);
 
