spark_graphql_client/

lib.rs

#![doc = include_str!("../README.md")]

use std::collections::HashMap;

use code_status_macros::{needs, untested};
use eyre::eyre;
use graphql_utils::compress_payload_zstd;
use reqwest::{Client, RequestBuilder, Response, StatusCode};
use serde::de::DeserializeOwned;
use serde_json::{Value, from_slice, from_value};
use url::Url;
use zstd::decode_all;

mod graphql_utils;

/// Specifies the type of payload compression to use for outgoing GraphQL requests.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub enum GraphQLPayloadCompression {
    /// Compress the request payload using the Zstandard (zstd) algorithm.
    /// This is the default if no compression type is specified.
    #[default]
    Zstd,

    /// Send the request payload without any compression.
    NeverCompress,
}

impl GraphQLPayloadCompression {
    /// Returns the corresponding `Accept-Encoding` header value for the compression type.
    fn apply_encoding_header(&self) -> &str {
        match self {
            Self::Zstd => "zstd",
            Self::NeverCompress => "",
        }
    }
}

/// A client for executing GraphQL operations (queries and mutations).
///
/// This client handles constructing the HTTP request, sending it to the specified
/// GraphQL endpoint, and processing the response.
#[derive(Clone)]
pub struct GraphqlClient<'a> {
    /// The base URL of the GraphQL endpoint.
    base_url: Url,
    /// The underlying `reqwest::Client` used for making HTTP requests.
    http_client: Client,
    /// The User-Agent string sent with each request.
    user_agent: &'a str,
    /// The compression strategy to use for request payloads.
    payload_compression: GraphQLPayloadCompression,
}

impl<'a> GraphqlClient<'a> {
    /// Creates a new `GraphqlClient` instance.
    ///
    /// # Arguments
    ///
    /// * `base_url` - The base URL of the GraphQL API endpoint (e.g., "https://api.example.com/graphql").
    /// * `user_agent` - The User-Agent string to identify this client.
    /// * `payload_compression` - An optional compression strategy. Defaults to `GraphQLPayloadCompression::Zstd` if `None`.
    ///
    /// # Panics
    ///
    /// Panics if the provided `base_url` is invalid.
    pub fn new(
        base_url: &str,
        user_agent: &'a str,
        payload_compression: Option<GraphQLPayloadCompression>,
    ) -> Self {
        let base_url = Url::parse(base_url)
            .expect("Provided invalid base URL for the GraphQL client. Please verify your URL.");

        Self {
            base_url,
            http_client: Client::new(),
            user_agent,
            payload_compression: payload_compression.unwrap_or_default(),
        }
    }

    /// Executes a GraphQL query or mutation.
    ///
    /// This method builds the request payload, applies compression if configured,
    /// sends the request, handles potential response compression, and deserializes
    /// the `data` field of the response into the specified type `T`.
    ///
    /// # Arguments
    ///
    /// * `query` - The GraphQL query or mutation string.
    ///           This should be a raw string literal (e.g., `r#"query { ... }"#`)
    ///           containing the full text of the operation. This client does not
    ///           use pre-generated query types.
    /// * `variables` - A map of variables to be included in the request.
    /// * `identity_public_key` - The public key identifying the client/user.
    /// * `session_token` - An optional session token for authentication.
    ///
    /// # Type Parameters
    ///
    /// * `T` - The type to deserialize the `data` field of the GraphQL response into.
    ///         It must implement `serde::de::DeserializeOwned`.
    ///
    /// # Returns
    ///
    /// Returns a `Result` containing the deserialized data `T` on success, or an `eyre::Report`
    /// detailing the error on failure (e.g., network issues, non-200 status code,
    /// deserialization errors, missing `data` field).
    #[needs("retry logic on network errors.")]
    pub async fn execute_graphql_request<T: DeserializeOwned>(
        &self,
        query: &str,
        variables: &HashMap<String, Value>,
        identity_public_key: &[u8],
        session_token: Option<&str>,
    ) -> eyre::Result<T> {
        let operation_name = graphql_utils::extract_operation_name(query)?;

        let payload = graphql_utils::build_graphql_request(&operation_name, query, variables);
        let payload_bytes = serde_json::to_string(&payload)?.into_bytes();
        let (payload_bytes, payload_is_compressed) = self.compress_payload(&payload_bytes)?;

        let request = self.construct_request(
            &operation_name,
            identity_public_key,
            session_token,
            payload_is_compressed,
        );

        let response = request.body(payload_bytes).send().await;
        let response = if let Ok(response) = response {
            response
        } else {
            return Err(eyre!("GraphQL request failed: {:?}", response.err()));
        };
        let status = response.status();

        if status != StatusCode::OK {
            return Err(eyre!("GraphQL request failed with status: {}", status));
        }

        let response_bytes = self.decompress_response(response).await?;
        let response_variable_map: HashMap<String, Value> = from_slice(&response_bytes)?;
        let response_data = response_variable_map
            .get("data")
            .ok_or(eyre!("No data found in the response"))?;

        let response = from_value(response_data.clone())?;
        Ok(response)
    }

    /// Compresses the request payload based on the client's configuration.
    ///
    /// # Arguments
    ///
    /// * `payload_bytes` - The raw bytes of the request payload.
    ///
    /// # Returns
    ///
    /// A tuple containing:
    ///  - The potentially compressed payload bytes.
    ///  - A boolean indicating whether compression was applied.
    /// Returns an error if compression fails.
    #[untested]
    fn compress_payload(&self, payload_bytes: &[u8]) -> eyre::Result<(Vec<u8>, bool)> {
        let (payload_bytes, payload_is_compressed) = match self.payload_compression {
            GraphQLPayloadCompression::Zstd => compress_payload_zstd(&payload_bytes)?,
            GraphQLPayloadCompression::NeverCompress => (payload_bytes.to_vec(), false),
        };

        Ok((payload_bytes, payload_is_compressed))
    }

    /// Decompresses the response body if it's compressed (currently supports zstd).
    ///
    /// Checks the `Content-Encoding` header of the response.
    ///
    /// # Arguments
    ///
    /// * `graphql_response` - The `reqwest::Response` received from the server.
    ///
    /// # Returns
    ///
    /// Returns the decompressed response body as bytes. Returns an error if decompression fails
    /// or an unsupported compression algorithm is used.
    #[untested]
    #[needs("a formatted matching for response encoding algorithms used.")]
    async fn decompress_response(&self, graphql_response: Response) -> eyre::Result<Vec<u8>> {
        // Clone the header value to avoid borrowing graphql_response
        let encoding = graphql_response
            .headers()
            .get("Content-Encoding")
            .and_then(|h| h.to_str().ok())
            .map(|s| s.to_string()); // Clone the relevant part

        let body_bytes = graphql_response.bytes().await?;

        if let Some(encoding) = encoding {
            if encoding == "zstd" {
                Ok(decode_all(&*body_bytes)?)
            } else {
                Err(eyre!("Unsupported compression algorithm: {}", encoding))
            }
        } else {
            Ok(body_bytes.to_vec())
        }
    }

    /// Constructs the basic `reqwest::RequestBuilder` with necessary headers.
    ///
    /// Sets headers like `Content-Type`, `User-Agent`, `X-GraphQL-Operation`,
    /// `Spark-Identity-Public-Key`, `Accept-Encoding` (if compression is used),
    /// and `Spark-Session-Token` (if provided).
    ///
    /// # Arguments
    ///
    /// * `operation_name` - The extracted name of the GraphQL operation.
    /// * `identity_public_key` - The public key identifying the client/user.
    /// * `session_token` - An optional session token for authentication.
    /// * `is_compressed` - Indicates if the payload being sent is compressed (affects `Accept-Encoding` header).
    ///
    /// # Returns
    ///
    /// A `reqwest::RequestBuilder` ready to have the body added and be sent.
    #[untested]
    #[needs("refactoring")]
    fn construct_request(
        &self,
        operation_name: &str,
        identity_public_key: &[u8],
        session_token: Option<&str>,
        is_compressed: bool,
    ) -> RequestBuilder {
        // Build the HTTP request
        let mut request = self
            .http_client
            .post(self.base_url.clone())
            .header("Content-Type", "application/json")
            .header("User-Agent", self.user_agent)
            .header("X-GraphQL-Operation", operation_name)
            .header(
                "Spark-Identity-Public-Key",
                hex::encode(identity_public_key),
            );

        if is_compressed {
            let applied_encoding = self.payload_compression.apply_encoding_header();
            request = request.header("Accept-Encoding", applied_encoding);
        }

        if let Some(session_token) = session_token {
            request = request.header("Spark-Session-Token", session_token);
        }

        request
    }
}

#[cfg(test)]
mod tests {
    use serde::Deserialize;
    use tokio;

    use super::*;

    // Define structs mirroring the expected GraphQL response structure
    #[derive(Deserialize, Debug, PartialEq)]
    struct Character {
        name: String,
    }

    #[derive(Deserialize, Debug, PartialEq)]
    struct Characters {
        results: Vec<Character>,
    }

    #[derive(Deserialize, Debug, PartialEq)]
    struct CharacterData {
        characters: Characters,
    }

    #[tokio::test]
    async fn test_rick_and_morty_graphql_request() {
        let client = GraphqlClient::new(
            "https://rickandmortyapi.com/graphql",
            "test-agent/0.1.0",
            Some(GraphQLPayloadCompression::NeverCompress), // API might not support compression
        );

        // Simple query to get the name of the first character
        let query = r#"
            query GetCharacterName {
                characters(page: 1) {
                    results {
                        name
                    }
                }
            }
        "#;

        let variables = HashMap::new(); // No variables needed for this query
        let dummy_identity_public_key = b""; // Dummy key for testing
        let session_token = None;

        let result = client
            .execute_graphql_request::<CharacterData>(
                query,
                &variables,
                dummy_identity_public_key,
                session_token,
            )
            .await;

        assert!(result.is_ok(), "Request failed: {:?}", result.err());

        let data = result.unwrap();
        // We expect at least one character result
        assert!(!data.characters.results.is_empty(), "No characters found");
        // We can check the first character's name if we know it,
        // but simply checking for non-empty results is safer for a public API.
        // Example: assert_eq!(data.characters.results[0].name, "Rick Sanchez");
        println!("First character name: {}", data.characters.results[0].name);
    }
}