discovery_connect/api/

query.rs

// Copyright 2023 Ikerian AG
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use graphql_client::GraphQLQuery;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};

/// This unit is responsible for sending asynchronous GraphQL queries to the Discovery server.
///
/// Asynchronously sends a GraphQL query to a server using the provided `QueryClient` and query variables.
/// This function handles the process of serializing, posting a GraphQL query and deserializing its response.
///
/// # Arguments
/// * `qc` - An `Arc<QueryClient>` instance, used to maintain state across async calls and manage the HTTP client.
/// * `variables` - Variables required for the GraphQL query, as defined by the `GraphQLQuery` trait.
///
/// # Returns
/// An implementation of `Future`, which when awaited, yields a `Result` containing either the response data
/// on success (`T::ResponseData`), or an `reqwest::Error` on failure.
///
/// # Type Constraints
/// * `T` - Represents the GraphQL query and must implement `GraphQLQuery`.
/// * `<T as GraphQLQuery>::Variables` - Must implement `Clone` trait.
///
/// # Errors
/// Returns an `Err` of type `reqwest::Error` if the request fails or if the server
/// response status is not OK.
///
/// # Examples
/// ```
/// use std::sync::Arc;
/// use std::time::Duration;
/// use crate::discovery_connect::file::{create_file, CreateFile};
/// use crate::discovery_connect::{post_query, QueryClient};
///
/// async fn query_example() {
///     let timeout = Duration::from_secs(1);
///     let query_client = Arc::new(QueryClient::new(
///         "https://api.example.discovery.retinai.com",
///         "client_id",
///         "client_secret",
///         "user@example",
///         "password123",
///         timeout,));
///    let input = create_file::CreateFileInput {
///        uuid: None,
///        filename: "example.zip".to_string(),
///        tags: Some(vec![]),
///        remarks: Some(serde_json::Value::Object(serde_json::Map::new())),
///        overwrite: Some(serde_json::Value::Object(serde_json::Map::new())),
///        workbook_uuid: Some("123e4567-e89b-12d3-a456-426614174000".to_string()),
///    };
///    let variables = create_file::Variables { input };
///    post_query::<CreateFile>(query_client, variables).await;
/// }
/// ```
pub async fn post_query<'a, T: GraphQLQuery + 'a>(
    qc: Arc<QueryClient>,
    variables: <T as GraphQLQuery>::Variables,
) -> Result<T::ResponseData, reqwest::Error>
where
    <T as GraphQLQuery>::Variables: Clone,
{
    async fn _post_query<'a, T: GraphQLQuery + 'a>(
        qc: Arc<QueryClient>,
        url: String,
        access_token: String,
        request_body: graphql_client::QueryBody<<T as GraphQLQuery>::Variables>,
    ) -> Result<T::ResponseData, reqwest::Error> {
        let result = qc
            .client
            .post(url)
            .bearer_auth(access_token)
            .json(&request_body);

        let result = result.send().await;
        match result {
            Ok(result) => {
                if result.status() != reqwest::StatusCode::OK {
                    let e = result.error_for_status().unwrap_err();
                    eprintln!("Error: {:?}", e);
                    return Err(e);
                }
                result.text().await.map(|text| {
                    let json: serde_json::Value = serde_json::from_str(&text).unwrap_or_default();
                    let data = json.get("data");
                    let data = serde_json::to_string(&data).unwrap_or_default();
                    match serde_json::from_str(&data) {
                        Ok(r) => r,
                        Err(e) => {
                            eprintln!("Error: {:?}", e);
                            // TODO: replace with proper error handling
                            // temporary hack to get this to compile
                            panic!("Failed to parse response");
                        }
                    }
                })
            }
            Err(e) => {
                eprintln!("Error: {:?}", e);
                Err(e)
            }
        }
    }

    with_fresh_token(qc.clone(), move || {
        let url = qc.config.lock().unwrap().url.clone();
        let url = format!("{}/api", url).clone();

        let access_token: String = match qc.config.lock() {
            Ok(config) => config.access_token.clone().unwrap(),
            Err(e) => {
                panic!("Mutex poisoned: {:?}", e);
            }
        };
        let request_body = T::build_query(variables.clone());
        Box::pin(_post_query::<'a, T>(
            qc.clone(),
            url,
            access_token,
            request_body,
        ))
    })
    .await
}

/// Macro to define types for the Disco GraphQL schema.
///
/// This macro simplifies the creation of types associated with GraphQL queries by automatically
/// deriving necessary traits and setting up schema configurations. It is intended to be used
/// in place of the `#[graphql]` attribute macro directly, providing a more concise syntax and
/// ensuring consistent configuration across different query types.
///
/// The macro generates a struct with the given identifier that represents a GraphQL query or mutation,
/// along with associated traits like `Debug`, `Default`, `PartialEq`, and `Serialize`. It also
/// handles custom scalar types, particularly mapping the GraphQL "JSON" type to the Rust equivalent.
///
/// # Usage
///
/// `disco_api!` macro is used by passing the name of the struct to be generated and the GraphQL
/// operation name. The macro expands to a struct definition with the `#[derive(GraphQLQuery)]`
/// attribute, along with necessary configurations.
///
/// # Example
///
/// ```
/// # use graphql_client::GraphQLQuery;
/// # use serde_json::Value as JSON;
/// # use discovery_connect::disco_api;
/// // TODO: the following line fails in doc tests but works in regular code:
/// // disco_api!(CreateFile, "createFile", "src/api/graphql/CreateFile.graphql");
/// ```
///
/// This example creates a `CreateFile` struct, which is tied to the "createFile" GraphQL operation
/// defined in file "src/api/graphql/CreateFile.graphql".
/// The struct will automatically derive `Debug`, `Default`, `PartialEq`, and `Serialize` traits, and
/// will use custom scalar configuration for handling JSON types.
///
/// # Parameters
///
/// * `$name`: The struct name to be generated.
/// * `$operation_name`: The name of the GraphQL operation this struct represents.
#[macro_export]
macro_rules! disco_api {
    ($name:ident, $operation_name:expr, $query_path:expr) => {
        #[derive(GraphQLQuery)]
        #[graphql(
                    query_path = $query_path,
                    schema_path = "src/api/graphql/schema.graphql",
                    response_derives = "Clone,Debug,Default,PartialEq,Deserialize,Serialize",
                    variables_derives = "Clone",
                    operation_name = $operation_name,
                    custom_scalars(
                        graphql_type = "JSON",
                        rust_type = "serde_json::Value",
                        serialize_with = "graphql_client::serialization::AsJson",
                        deserialize_with = "graphql_client::deprecation::FromJson"
                    )
                )]
        pub struct $name;
    };
}

/// Represents the authentication configuration for a service.
///
/// This structure holds various credentials and tokens necessary for authenticating
/// with a specific service. It includes the service's URL, user credentials, and tokens
/// used for accessing protected resources.
///
/// # Fields
///
/// * `url` - A `String` representing the URL of the service to which queries will be sent.
/// * `email` - A `String` representing the user's email address used for authentication.
/// * `region` - A `Region` enum representing the region of the service.
/// * `password` - A `String` representing the user's password used for authentication.
/// * `access_token` - An `Option<String>` that holds the access token if available.
///   This token is used to authenticate queries. It can be `None` if the token has not been acquired yet.
/// * `refresh_token` - An `Option<String>` that holds the refresh token if available.
///   This token is used to obtain a new access token when the current one expires.
///   It can be `None` if the token has not been acquired yet.
///
/// # Examples
///
/// ```
/// use discovery_connect::{QueryConfig};
///
/// let query_config = QueryConfig {
///     client_id: "client_id".to_string(),
///     client_secret: "client_secret".to_string(),
///     url: "https://api.europe.discovery.retinai.com".to_string(),
///     email: "user@example.com".to_string(),
///     password: "password123".to_string(),
///     access_token: None,
///     refresh_token: None,
/// };
/// ```
#[derive(Clone, Debug)]
pub struct QueryConfig {
    pub url: String,
    pub client_id: String,
    pub client_secret: String,
    pub email: String,
    pub password: String,
    pub access_token: Option<String>,
    pub refresh_token: Option<String>,
}

/// Encapsulates a client for performing authenticated queries to a service.
///
/// # Fields
///
/// * `client` - An instance of `reqwest::Client` used for making HTTP requests.
/// * `config` - A `Mutex<QueryConfig>` that holds the authentication configuration.
///   The `Mutex` ensures that the configuration can be safely accessed and modified
///   across multiple threads.
///
/// # Examples
///
/// ```
/// use discovery_connect::{QueryClient, QueryConfig};
/// use std::sync::Mutex;
///
/// let query_client = QueryClient {
///     client: reqwest::Client::new(),
///     config: Mutex::new(QueryConfig {
///         client_id: "client_id".to_string(),
///         client_secret: "client_secret".to_string(),
///         url: "https://api.europe.discovery.retinai.com".to_string(),
///         email: "user@example.com".to_string(),
///         password: "password123".to_string(),
///         access_token: None,
///         refresh_token: None,
///     }),
/// };
/// ```
#[derive(Debug)]
pub struct QueryClient {
    pub config: Mutex<QueryConfig>,
    pub client: reqwest::Client,
}

impl QueryClient {
    pub fn new(
        url: &str,
        client_id: &str,
        client_secret: &str,
        email: &str,
        password: &str,
        timeout: std::time::Duration,
    ) -> QueryClient {
        let client = reqwest::Client::builder()
            .timeout(timeout)
            .build()
            .unwrap_or_default();

        QueryClient {
            client,
            config: Mutex::new(QueryConfig {
                url: url.to_string(),
                client_id: client_id.to_string(),
                client_secret: client_secret.to_string(),
                email: email.to_string(),
                password: password.to_string(),
                access_token: None,
                refresh_token: None,
            }),
        }
    }
}

/// Executes a provided asynchronous function, handling token authentication.
///
/// This function takes a configuration object and an asynchronous function `func`.
/// It attempts to execute `func`. If `func` fails due to an unauthorized error,
/// `with_fresh_token` tries to refresh the access token using the refresh token in the configuration.
/// If refreshing the token also fails with an unauthorized error, it attempts to log in again
/// with the credentials in the configuration to obtain new tokens.
///
/// # Arguments
///
/// * `config` - A shared, thread-safe reference (`Arc`) to an `QueryConfig` object
///   containing authentication details like URL, email, password, and tokens.
/// * `func` - The asynchronous function to be executed. This function should return a `Future`
///   that resolves to a `Result<R, reqwest::Error>`. `R` is the expected successful return type of the function.
///
/// # Returns
///
/// This function returns a `Future` that resolves to a `Result<R, reqwest::Error>`.
/// `R` is the return type of the provided asynchronous function `func`.
/// If `func` executes successfully, its result is returned. If `func` fails due to an unauthorized error,
/// a token refresh or re-login is attempted, and `func` is retried.
/// If the token refresh or re-login fails, or if `func` fails due to any other error,
/// the error is propagated as the return value.
///
/// # Examples
///
/// ```
/// use discovery_connect::{QueryConfig, QueryClient, with_fresh_token};
/// use std::sync::Arc;
/// use std::sync::Mutex;
///
/// async fn example() {
///     let config = Arc::new(QueryClient {
///         client: reqwest::Client::new(),
///         config: Mutex::new(QueryConfig {
///             client_id: "client_id".to_string(),
///             client_secret: "client_secret".to_string(),
///             url: "https://api.europe.discovery.retinai.com".to_string(),
///             email: "user@example".to_string(),
///             password: "password123".to_string(),
///             access_token: None,
///             refresh_token: None,
///         })
///     });
///     let result = with_fresh_token(config, move || {
///         // do something that requires authentication here
///         // ...
///         Box::pin(async move { Ok(()) })
///     }).await;
/// }
/// ```
pub async fn with_fresh_token<'a, F, R>(qc: Arc<QueryClient>, func: F) -> Result<R, reqwest::Error>
where
    F: Fn() -> Pin<Box<dyn Future<Output = Result<R, reqwest::Error>> + 'a>> + 'a,
    R: 'a,
{
    {
        // login if there is no refresh token
        let config = qc.config.lock().unwrap().clone();
        let access_token = &config.access_token;
        let refresh_token = &config.refresh_token;
        if access_token.is_none() || refresh_token.is_none() {
            match login(qc.clone(), config).await {
                Ok(config) => {
                    let mut qc_config = qc.config.lock().unwrap();
                    qc_config.access_token = config.access_token;
                    qc_config.refresh_token = config.refresh_token;
                }
                Err(e) => {
                    eprintln!("Error: {:?}", e);
                    return Err(e);
                }
            }
        }
    }

    // execute the function and check for authorization error
    match func().await {
        Ok(result) => Ok(result),
        Err(e) => {
            if e.status() == Some(reqwest::StatusCode::UNAUTHORIZED) {
                {
                    // Try to refresh the access token and retry
                    let config = qc.config.lock().unwrap().clone();
                    let refresh_token = &config.refresh_token;
                    match super::auth::access_token(
                        &qc.client,
                        &config.url,
                        refresh_token.as_ref().unwrap(),
                    )
                    .await
                    {
                        Ok(auth) => {
                            let mut qc_config = qc.config.lock().unwrap();
                            qc_config.access_token = Some(auth.access_token.clone());
                            qc_config.refresh_token = Some(auth.refresh_token.clone());
                        }
                        Err(e) => {
                            eprintln!("Error: {:?}", e);
                            // failed to refresh tokens, clear to force a fresh login
                            let mut qc_config = qc.config.lock().unwrap();
                            qc_config.access_token = None;
                            qc_config.refresh_token = None;
                            return Err(e);
                        }
                    };
                }

                // tokens refreshed, retry the function
                func().await
            } else {
                eprintln!("Error: {:?}", e);
                Err(e)
            }
        }
    }
}

/// Performs a basic password login and returns the authentication response.
///
/// # Arguments
///
/// * `qc` - A shared, thread-safe reference (`Arc`) to an `QueryClient` object
///   containing authentication details like URL, email, password, and tokens.
/// * `config` - A shared, thread-safe reference (`MutexGuard`) to an `QueryConfig` object
///   containing authentication details like URL, email, password, and tokens.
///
/// # Returns
///
/// If the login is successful, the `AuthResponse` is returned.
/// If the login fails, an `reqwest::Error` is returned.
pub async fn login(
    qc: Arc<QueryClient>,
    config: QueryConfig,
) -> Result<QueryConfig, reqwest::Error> {
    match super::auth::login(
        &qc.client,
        &config.url,
        &config.client_id,
        &config.client_secret,
        &config.email,
        &config.password,
        None,
    )
    .await
    {
        Ok(auth) => Ok(QueryConfig {
            access_token: Some(auth.access_token.clone()),
            refresh_token: Some(auth.refresh_token.clone()),
            ..config
        }),
        Err(e) => Err(e),
    }
}