hive-router 0.2.0

use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::net::SocketAddr;
use std::time::Duration;

use crate::config::primitives::absolute_path::AbsolutePath;
use crate::config::primitives::value_or_expression::ValueOrExpression;

#[derive(Debug, Deserialize, Serialize, JsonSchema, Default, Clone)]
#[serde(deny_unknown_fields)]
pub struct SubscriptionsConfig {
    /// Enables/disables subscriptions. By default, the subscriptions are disabled.
    ///
    /// You can override this setting by setting the `SUBSCRIPTIONS_ENABLED` environment variable to `true` or `false`.
    #[serde(default)]
    pub enabled: bool,
    /// The capacity of the broadcast channel used to fan out subscription events to all active listeners.
    ///
    /// Each active subscription has its own broadcast channel. This value controls how many events
    /// can be buffered in that channel before slow consumers start lagging. If a consumer falls too
    /// far behind and the buffer is full, it will skip the missed messages and continue from the
    /// latest available event.
    ///
    /// Subscription events are typically low-frequency, so the default of 32 is sufficient for most
    /// use cases. Increase this value if you expect bursts of events or have slow consumers that
    /// need more headroom to catch up.
    ///
    /// Defaults to 32.
    #[serde(default = "default_broadcast_capacity")]
    pub broadcast_capacity: usize,
    /// The capacity of the per-subscription buffer between a subgraph and the router's
    /// processing pipeline.
    ///
    /// When a subscription is established, the router reads events from the subgraph (over
    /// HTTP streaming or WebSocket) and runs each one through entity resolution before fanning
    /// it out to listeners. If that processing is slower than the rate at which the subgraph
    /// emits events, this buffer absorbs the difference so the subgraph is never throttled by
    /// the router's processing speed.
    ///
    /// When the buffer is full, the newest event is dropped (and logged) instead of slowing
    /// down or tearing down the connection to the subgraph. The subscription stays alive and
    /// the subgraph keeps emitting unaffected.
    ///
    /// A larger capacity gives the router more headroom to catch up during bursts at the cost
    /// of memory and potentially staler events under sustained backpressure. A smaller capacity
    /// keeps memory minimal and drops eagerly, which is appropriate when only the latest events
    /// matter.
    ///
    /// Defaults to 1024.
    #[serde(default = "default_subgraph_buffer_capacity")]
    pub subgraph_buffer_capacity: usize,
    /// Configuration for subgraphs using the HTTP Callback protocol.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub callback: Option<CallbackConfig>,
    /// Configuration for subgraphs using WebSocket protocol.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub websocket: Option<WebSocketConfig>,
}

/// Configuration for the HTTP Callback subscription mode.
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
#[serde(deny_unknown_fields)]
pub struct CallbackConfig {
    /// The public URL that subgraphs will use to send callback messages to this router.
    ///
    /// Your public_url must match the server address combined with the router's path.
    /// Meaning, if your server is `http://localhost:4000` and the path is `/callback`,
    /// your `public_url` should be `http://localhost:4000/callback`.
    ///
    /// Can be a static URL string or a VRL expression. Expressions are useful for
    /// service discovery in horizontally scaled deployments where the URL can be
    /// read from an environment variable:
    ///
    /// ```yaml
    /// public_url:
    ///   expression: 'env("ROUTER_PUBLIC_URL")'
    /// ```
    pub public_url: ValueOrExpression<String>,
    /// The path of the router's callback endpoint.
    /// Must be an absolute path starting with `/`. Defaults to `/callback`.
    #[serde(default = "default_callback_path")]
    pub path: AbsolutePath,
    /// The interval at which the subgraph must send heartbeat messages.
    /// If set to 0, heartbeats are disabled. Defaults to 5 seconds.
    #[serde(
        default = "default_heartbeat_interval",
        deserialize_with = "humantime_serde::deserialize",
        serialize_with = "humantime_serde::serialize"
    )]
    #[schemars(with = "String")]
    pub heartbeat_interval: Duration,
    /// The IP address and port the router will listen on for subscription callbacks.
    /// When set, the router will start a dedicated HTTP server bound to this address
    /// for receiving callback messages from subgraphs, separate from the main GraphQL server.
    /// When not set, the callback handler is registered on the main server.
    ///
    /// Example: `0.0.0.0:4001`
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub listen: Option<SocketAddr>,
    /// The list of subgraph names that use the HTTP callback protocol.
    #[serde(default)]
    pub subgraphs: HashSet<String>,
}

fn default_broadcast_capacity() -> usize {
    32
}

fn default_subgraph_buffer_capacity() -> usize {
    1024
}

fn default_callback_path() -> AbsolutePath {
    AbsolutePath::try_from("/callback").expect("default callback path is valid")
}

fn default_heartbeat_interval() -> Duration {
    Duration::from_secs(5)
}

/// Configuration for the WebSocket subscription mode.
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Default)]
#[serde(deny_unknown_fields)]
pub struct WebSocketConfig {
    /// The default configuration that will be applied to all subgraphs using
    /// WebSocket protocol, unless overridden by a specific subgraph configuration.
    ///
    /// When specified, all subgraphs (not claimed by `callback`) will use the WebSocket protocol.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub all: Option<WebSocketSubgraphConfig>,
    /// Optional per-subgraph configurations that will override the default configuration for specific subgraphs.
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    pub subgraphs: HashMap<String, WebSocketSubgraphConfig>,
}

/// WebSocket configuration for a specific subgraph or the default for all subgraphs.
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Default)]
#[serde(deny_unknown_fields)]
pub struct WebSocketSubgraphConfig {
    /// Determines the URL path to use for the subscription endpoint:
    ///
    /// - For WebSocket connections, the URL will be `ws://<subgraph-url><path>`.
    /// - If `path` is not set, the default subgraph URL is used, with the scheme adjusted to `ws`
    ///   for WebSocket connections where applicable.
    ///
    /// Note to always provide the absolute path starting with a `/`, e.g., `/ws`.
    ///
    /// For example, if the subgraph URL is `http://example.com/graphql` and the path is set to `/ws`,
    /// the resulting WebSocket URL will be `ws://example.com/ws`.
    #[serde(default)]
    pub path: Option<AbsolutePath>,
}

/// Subscription settings used by subgraph executors for one supergraph.
#[derive(Clone)]
pub struct SupergraphSubscriptionsConfig {
    pub callback_subgraphs: HashSet<String>,
    pub websocket: Option<WebSocketConfig>,
    pub subgraph_buffer_capacity: usize,
}

impl Default for SupergraphSubscriptionsConfig {
    fn default() -> Self {
        Self {
            callback_subgraphs: HashSet::new(),
            websocket: None,
            subgraph_buffer_capacity: default_subgraph_buffer_capacity(),
        }
    }
}

impl From<&SubscriptionsConfig> for SupergraphSubscriptionsConfig {
    fn from(config: &SubscriptionsConfig) -> Self {
        Self {
            callback_subgraphs: config
                .callback
                .as_ref()
                .map(|callback| callback.subgraphs.clone())
                .unwrap_or_default(),
            websocket: config.websocket.clone(),
            subgraph_buffer_capacity: config.subgraph_buffer_capacity,
        }
    }
}

impl SupergraphSubscriptionsConfig {
    /// Returns the subscription protocol for the given subgraph.
    pub fn get_protocol_for_subgraph(&self, subgraph_name: &str) -> SubscriptionProtocol {
        if self.callback_subgraphs.contains(subgraph_name) {
            return SubscriptionProtocol::HTTPCallback;
        }
        if let Some(ref websocket) = self.websocket {
            if websocket.all.is_some() || websocket.subgraphs.contains_key(subgraph_name) {
                return SubscriptionProtocol::WebSocket;
            }
        }
        SubscriptionProtocol::HTTP
    }

    /// Returns the WebSocket path for the given subgraph, if configured.
    pub fn get_websocket_path(&self, subgraph_name: &str) -> Option<&str> {
        self.websocket.as_ref().and_then(|ws| {
            ws.subgraphs
                .get(subgraph_name)
                .and_then(|s| s.path.as_ref().map(|p| p.as_str()))
                .or_else(|| {
                    ws.all
                        .as_ref()
                        .and_then(|a| a.path.as_ref().map(|p| p.as_str()))
                })
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn callback_path_must_be_absolute() {
        let err = serde_json::from_str::<CallbackConfig>(
            r#"{"public_url": "http://localhost:4000/callback", "path": "callback"}"#,
        )
        .unwrap_err();
        assert!(
            err.to_string()
                .contains("path must be absolute (start with /)"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn callback_path_absolute_is_accepted() {
        let config = serde_json::from_str::<CallbackConfig>(
            r#"{"public_url": "http://localhost:4000/callback", "path": "/callback"}"#,
        )
        .unwrap();
        assert_eq!(config.path.as_str(), "/callback");
    }

    #[test]
    fn callback_path_defaults_to_absolute() {
        let config = serde_json::from_str::<CallbackConfig>(
            r#"{"public_url": "http://localhost:4000/callback"}"#,
        )
        .unwrap();
        assert_eq!(config.path.as_str(), "/callback");
    }
}

/// The selected protocol for the subscriptions towards subgraphs.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum SubscriptionProtocol {
    /// Uses any HTTP streaming protocol that the subgraph accepts. Supported protocols are:
    /// - Server-Sent Events (SSE). Respecting only the "distinct connection mode" of the GraphQL over SSE specification. See: https://github.com/graphql/graphql-over-http/blob/main/rfcs/GraphQLOverSSE.md#distinct-connections-mode.
    /// - Apollo Multipart HTTP. Implements the Apollo's Multipart HTTP specification. See: https://www.apollographql.com/docs/graphos/routing/operations/subscriptions/multipart-protocol.
    /// - GraphQL Incremental Delivery. Implements the official GraphQL Incremental Delivery specification. See: https://github.com/graphql/graphql-over-http/blob/main/rfcs/IncrementalDelivery.md.
    #[default]
    HTTP,
    /// Uses GraphQL over WebSocket (graphql-transport-ws subprotocol).
    WebSocket,
    /// Uses the HTTP Callback protocol for subscriptions.
    /// See: https://www.apollographql.com/docs/graphos/routing/operations/subscriptions/callback-protocol
    HTTPCallback,
}