grpc_graphql_gateway 1.2.5

#![allow(clippy::result_large_err, clippy::const_is_empty)]
// @generated by protoc-gen-graphql-template
// Source files:
//   - live_query_example.proto
//
// This is a starter gateway. Update endpoint URLs and tweak as needed.

use async_graphql::{Name, Value as GqlValue};
use grpc_graphql_gateway::{
    CacheConfig, Gateway, GatewayBuilder, GrpcClient, PersistedQueryConfig,
    RequestCollapsingConfig, Result as GatewayResult,
};
use std::net::{SocketAddr, ToSocketAddrs};
use std::sync::Arc;
use std::time::Duration;
use tonic::{transport::Server, Request, Response, Status};
use tracing_subscriber::prelude::*;

type ServiceResult<T> = std::result::Result<T, Status>;

const DESCRIPTOR_SET: &[u8] =
    include_bytes!("../../src/generated/live_query_example_descriptor.bin");

#[allow(dead_code)]
const DEFAULT_GRPC_ADDR: &str = "0.0.0.0:50051";

fn describe(list: &[&str]) -> String {
    if list.is_empty() {
        "none".to_string()
    } else {
        list.join(", ")
    }
}

fn describe_resolvers(list: &[&str]) -> String {
    if list.is_empty() {
        "none".to_string()
    } else {
        list.join(", ")
    }
}

#[allow(dead_code)]
fn listen_addr(endpoint: &str, fallback: &str) -> GatewayResult<SocketAddr> {
    let mut addr = endpoint.trim();
    if let Some(stripped) = addr
        .strip_prefix("http://")
        .or_else(|| addr.strip_prefix("https://"))
    {
        addr = stripped;
    }
    if let Some((host, _rest)) = addr.split_once('/') {
        addr = host;
    }
    if let Ok(sock) = addr.parse() {
        return Ok(sock);
    }
    if let Ok(mut iter) = addr.to_socket_addrs() {
        if let Some(sock) = iter.next() {
            return Ok(sock);
        }
    }
    fallback
        .parse()
        .map_err(|e| grpc_graphql_gateway::Error::Other(anyhow::Error::new(e)))
}

fn describe_key_sets(keys: &[&[&str]]) -> String {
    if keys.is_empty() {
        "none".to_string()
    } else {
        keys.iter()
            .map(|set| set.join(" "))
            .collect::<Vec<_>>()
            .join(" | ")
    }
}

fn describe_entities() -> String {
    if ENTITY_CONFIGS.is_empty() {
        "none".to_string()
    } else {
        ENTITY_CONFIGS
            .iter()
            .map(|e| format!("{} (keys: {})", e.type_name, describe_key_sets(e.keys)))
            .collect::<Vec<_>>()
            .join(", ")
    }
}

const QUERIES: &[&str] = &["user", "userStatus", "users"];
const MUTATIONS: &[&str] = &["createUser", "deleteUser", "updateUser"];
const SUBSCRIPTIONS: &[&str] = &[];
const RESOLVERS: &[&str] = &["_entities"];
#[allow(dead_code)]
const FEDERATION_ENABLED: bool = true;

pub struct EntityConfigInfo {
    pub type_name: &'static str,
    pub keys: &'static [&'static [&'static str]],
    pub extend: bool,
    pub resolvable: bool,
}

pub const ENTITY_CONFIGS: &[EntityConfigInfo] = &[EntityConfigInfo {
    type_name: "live_example_User",
    keys: &[&["id"]],
    extend: false,
    resolvable: true,
}];

/// Live Query configuration for reactive subscriptions.
/// Use the @live directive on supported queries to receive automatic updates.
#[derive(Debug, Clone)]
pub struct LiveQueryConfigInfo {
    /// GraphQL operation name
    pub operation_name: &'static str,
    /// Throttle interval in milliseconds
    pub throttle_ms: u32,
    /// Invalidation triggers
    pub triggers: &'static [&'static str],
    /// Maximum connections per client
    pub max_connections: u32,
    /// TTL in seconds (0 = infinite)
    pub ttl_seconds: u32,
    /// Strategy: INVALIDATION, POLLING, or HASH_DIFF
    pub strategy: &'static str,
    /// Poll interval for POLLING strategy
    pub poll_interval_ms: u32,
    /// Entity dependencies
    pub depends_on: &'static [&'static str],
}

pub const LIVE_QUERY_CONFIGS: &[LiveQueryConfigInfo] = &[
    LiveQueryConfigInfo {
        operation_name: "user",
        throttle_ms: 100,
        triggers: &["User.update", "User.delete"],
        max_connections: 10,
        ttl_seconds: 3600,
        strategy: "INVALIDATION",
        poll_interval_ms: 0,
        depends_on: &["User"],
    },
    LiveQueryConfigInfo {
        operation_name: "users",
        throttle_ms: 500,
        triggers: &["User.create", "User.update", "User.delete"],
        max_connections: 10,
        ttl_seconds: 0,
        strategy: "POLLING",
        poll_interval_ms: 5000,
        depends_on: &[],
    },
    LiveQueryConfigInfo {
        operation_name: "userStatus",
        throttle_ms: 50,
        triggers: &["User.status_change"],
        max_connections: 10,
        ttl_seconds: 0,
        strategy: "HASH_DIFF",
        poll_interval_ms: 0,
        depends_on: &[],
    },
];

#[allow(dead_code)]
const LIVE_QUERIES_ENABLED: bool = true;

#[allow(dead_code)]
fn describe_live_queries() -> String {
    if LIVE_QUERY_CONFIGS.is_empty() {
        "none".to_string()
    } else {
        LIVE_QUERY_CONFIGS
            .iter()
            .map(|lq| {
                format!(
                    "{}@live ({}ms, {})",
                    lq.operation_name, lq.throttle_ms, lq.strategy
                )
            })
            .collect::<Vec<_>>()
            .join(", ")
    }
}

pub mod live_example {
    include!("../../src/generated/live_example.rs");
}

pub struct ServiceConfig {
    pub name: &'static str,
    pub endpoint: &'static str,
    pub insecure: bool,
    pub queries: &'static [&'static str],
    pub mutations: &'static [&'static str],
    pub subscriptions: &'static [&'static str],
    pub resolvers: &'static [&'static str],
}

pub mod services {
    use super::ServiceConfig;

    pub const LIVE_EXAMPLE_USERSERVICE: ServiceConfig = ServiceConfig {
        name: "live_example.UserService",
        endpoint: "http://localhost:50052",
        insecure: true,
        queries: &["user", "users", "userStatus", "post"],
        mutations: &["createUser", "updateUser", "deleteUser", "createPost"],
        subscriptions: &[],
        resolvers: &["_entities"],
    };

    pub const ALL: &[ServiceConfig] = &[LIVE_EXAMPLE_USERSERVICE];
}

/// Example entity resolver stub for federation. Replace with your own logic or DataLoader.
#[derive(Clone, Default)]
pub struct LiveExampleEntityResolver;

#[async_trait::async_trait]
impl grpc_graphql_gateway::EntityResolver for LiveExampleEntityResolver {
    async fn resolve_entity(
        &self,
        _entity_config: &grpc_graphql_gateway::federation::EntityConfig,
        representation: &async_graphql::indexmap::IndexMap<Name, GqlValue>,
    ) -> grpc_graphql_gateway::Result<GqlValue> {
        let mut obj = representation.clone();
        obj.shift_remove(&Name::new("__typename"));
        Ok(GqlValue::Object(obj))
    }

    async fn batch_resolve_entities(
        &self,
        entity_config: &grpc_graphql_gateway::federation::EntityConfig,
        representations: Vec<async_graphql::indexmap::IndexMap<Name, GqlValue>>,
    ) -> grpc_graphql_gateway::Result<Vec<GqlValue>> {
        let mut results = Vec::with_capacity(representations.len());
        for repr in representations {
            results.push(self.resolve_entity(entity_config, &repr).await?);
        }
        Ok(results)
    }
}

fn default_entity_resolver() -> Arc<dyn grpc_graphql_gateway::EntityResolver> {
    Arc::new(LiveExampleEntityResolver)
}

// In-memory user store for the example
use grpc_graphql_gateway::{InvalidationEvent, SharedLiveQueryStore};
use parking_lot::RwLock;
use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};

lazy_static::lazy_static! {
    static ref USER_STORE: RwLock<HashMap<String, live_example::User>> = {
        let mut store = HashMap::new();
        // Seed with some initial users
        let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() as i64;
        store.insert("1".to_string(), live_example::User {
            id: "1".to_string(),
            name: "Alice".to_string(),
            email: "alice@example.com".to_string(),
            status: Some(live_example::UserStatus {
                user_id: "1".to_string(),
                is_online: true,
                last_active: "2024-01-01T12:00:00Z".to_string(),
                current_activity: "Browsing".to_string(),
            }),
            created_at: now,
            updated_at: now,
            posts: vec![],
        });
        store.insert("2".to_string(), live_example::User {
            id: "2".to_string(),
            name: "Bob".to_string(),
            email: "bob@example.com".to_string(),
            status: Some(live_example::UserStatus {
                user_id: "2".to_string(),
                is_online: false,
                last_active: "2024-01-01T10:00:00Z".to_string(),
                current_activity: "Offline".to_string(),
            }),
            created_at: now,
            updated_at: now,
            posts: vec![],
        });
        store.insert("3".to_string(), live_example::User {
            id: "3".to_string(),
            name: "Charlie".to_string(),
            email: "charlie@example.com".to_string(),
            status: Some(live_example::UserStatus {
                user_id: "3".to_string(),
                is_online: true,
                last_active: "2024-01-01T11:30:00Z".to_string(),
                current_activity: "Coding".to_string(),
            }),
            created_at: now,
            updated_at: now,
            posts: vec![],
        });
        RwLock::new(store)
    };

    static ref POST_STORE: RwLock<HashMap<String, live_example::Post>> = RwLock::new(HashMap::new());
    static ref NEXT_USER_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(4);
}

/// Get the live query store for triggering invalidations
/// Uses the global singleton store shared with WebSocket connections
pub fn get_live_query_store() -> SharedLiveQueryStore {
    grpc_graphql_gateway::global_live_query_store()
}

/// gRPC service implementation with live query support
#[derive(Clone)]
pub struct ServiceImpl {
    live_query_store: SharedLiveQueryStore,
}

impl Default for ServiceImpl {
    fn default() -> Self {
        Self {
            live_query_store: grpc_graphql_gateway::global_live_query_store(),
        }
    }
}

impl ServiceImpl {
    pub fn new(live_query_store: SharedLiveQueryStore) -> Self {
        Self { live_query_store }
    }

    /// Trigger a live query invalidation event
    fn invalidate(&self, type_name: &str, action: &str) {
        let event = InvalidationEvent::new(type_name, action);
        let affected = self.live_query_store.invalidate(event);
        if affected > 0 {
            tracing::info!(
                type_name = %type_name,
                action = %action,
                affected_count = affected,
                "Triggered live query invalidation"
            );
        }
    }
}

#[tonic::async_trait]
impl live_example::user_service_server::UserService for ServiceImpl {
    async fn get_user(
        &self,
        request: Request<live_example::GetUserRequest>,
    ) -> ServiceResult<Response<live_example::User>> {
        let req = request.into_inner();
        tracing::debug!(user_id = %req.id, "GetUser request");

        let store = USER_STORE.read();
        match store.get(&req.id) {
            Some(user) => {
                let mut user = user.clone();
                // Enrich with posts
                let post_store = POST_STORE.read();
                user.posts = post_store
                    .values()
                    .filter(|p| p.author_id == req.id)
                    .cloned()
                    .collect();
                Ok(Response::new(user))
            }
            None => Err(Status::not_found(format!("User {} not found", req.id))),
        }
    }

    async fn list_users(
        &self,
        request: Request<live_example::ListUsersRequest>,
    ) -> ServiceResult<Response<live_example::ListUsersResponse>> {
        let req = request.into_inner();
        tracing::debug!(limit = req.limit, offset = req.offset, "ListUsers request");

        let store = USER_STORE.read();
        let post_store = POST_STORE.read();
        let mut users: Vec<live_example::User> = store
            .values()
            .map(|u| {
                let mut user = u.clone();
                user.posts = post_store
                    .values()
                    .filter(|p| p.author_id == u.id)
                    .cloned()
                    .collect();
                user
            })
            .collect();

        // Apply filter if provided
        if !req.filter.is_empty() {
            let filter_lower = req.filter.to_lowercase();
            users.retain(|u| {
                u.name.to_lowercase().contains(&filter_lower)
                    || u.email.to_lowercase().contains(&filter_lower)
            });
        }

        // Sort by id for consistent ordering
        users.sort_by(|a, b| a.id.cmp(&b.id));

        let total_count = users.len() as i32;

        // Apply pagination
        let offset = req.offset.max(0) as usize;
        let limit = if req.limit > 0 {
            req.limit as usize
        } else {
            users.len()
        };

        let users: Vec<live_example::User> = users.into_iter().skip(offset).take(limit).collect();

        Ok(Response::new(live_example::ListUsersResponse {
            users,
            total_count,
        }))
    }

    async fn get_user_status(
        &self,
        request: Request<live_example::GetUserStatusRequest>,
    ) -> ServiceResult<Response<live_example::UserStatus>> {
        let req = request.into_inner();
        tracing::debug!(user_id = %req.user_id, "GetUserStatus request");

        let store = USER_STORE.read();
        match store.get(&req.user_id) {
            Some(user) => match &user.status {
                Some(status) => Ok(Response::new(status.clone())),
                None => Err(Status::not_found(format!(
                    "Status for user {} not found",
                    req.user_id
                ))),
            },
            None => Err(Status::not_found(format!("User {} not found", req.user_id))),
        }
    }

    async fn create_user(
        &self,
        request: Request<live_example::CreateUserRequest>,
    ) -> ServiceResult<Response<live_example::User>> {
        let req = request.into_inner();
        tracing::info!(name = %req.name, email = %req.email, "CreateUser request");

        let id = NEXT_USER_ID
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst)
            .to_string();
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64;

        let user = live_example::User {
            id: id.clone(),
            name: req.name,
            email: req.email,
            status: Some(live_example::UserStatus {
                user_id: id.clone(),
                is_online: true,
                last_active: chrono::Utc::now().to_rfc3339(),
                current_activity: "Just joined".to_string(),
            }),
            created_at: now,
            updated_at: now,
            posts: vec![],
        };

        {
            let mut store = USER_STORE.write();
            store.insert(id.clone(), user.clone());
        }

        // Trigger live query invalidation for User.create
        self.invalidate("User", "create");

        tracing::info!(user_id = %id, "Created new user");
        Ok(Response::new(user))
    }

    async fn update_user(
        &self,
        request: Request<live_example::UpdateUserRequest>,
    ) -> ServiceResult<Response<live_example::User>> {
        let req = request.into_inner();
        tracing::info!(user_id = %req.id, "UpdateUser request");

        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64;

        let user = {
            let mut store = USER_STORE.write();
            match store.get_mut(&req.id) {
                Some(user) => {
                    if !req.name.is_empty() {
                        user.name = req.name;
                    }
                    if !req.email.is_empty() {
                        user.email = req.email;
                    }
                    user.updated_at = now;
                    user.clone()
                }
                None => return Err(Status::not_found(format!("User {} not found", req.id))),
            }
        };

        // Trigger live query invalidation for User.update
        self.invalidate("User", "update");

        tracing::info!(user_id = %req.id, "Updated user");
        Ok(Response::new(user))
    }

    async fn delete_user(
        &self,
        request: Request<live_example::DeleteUserRequest>,
    ) -> ServiceResult<Response<live_example::DeleteUserResponse>> {
        let req = request.into_inner();
        tracing::info!(user_id = %req.id, "DeleteUser request");

        let removed = {
            let mut store = USER_STORE.write();
            store.remove(&req.id)
        };

        match removed {
            Some(_) => {
                // Trigger live query invalidation for User.delete
                self.invalidate("User", "delete");

                tracing::info!(user_id = %req.id, "Deleted user");
                Ok(Response::new(live_example::DeleteUserResponse {
                    success: true,
                    message: format!("User {} deleted successfully", req.id),
                }))
            }
            None => Err(Status::not_found(format!("User {} not found", req.id))),
        }
    }

    async fn get_post(
        &self,
        request: Request<live_example::GetPostRequest>,
    ) -> ServiceResult<Response<live_example::Post>> {
        let req = request.into_inner();
        tracing::debug!(post_id = %req.id, "GetPost request");

        let store = POST_STORE.read();
        match store.get(&req.id) {
            Some(post) => Ok(Response::new(post.clone())),
            None => Err(Status::not_found(format!("Post {} not found", req.id))),
        }
    }

    async fn create_post(
        &self,
        request: Request<live_example::CreatePostRequest>,
    ) -> ServiceResult<Response<live_example::Post>> {
        let req = request.into_inner();
        tracing::info!(title = %req.title, author_id = %req.author_id, "CreatePost request");

        let id = uuid::Uuid::new_v4().to_string();

        let post = live_example::Post {
            id: id.clone(),
            title: req.title,
            content: req.content,
            author_id: req.author_id.clone(),
        };

        {
            let mut store = POST_STORE.write();
            store.insert(id.clone(), post.clone());
        }

        // Trigger live query invalidation for Post.create and User.update (since user.posts changed)
        self.invalidate("Post", "create");
        self.invalidate("User", "update"); // Invalidate user to refresh their post list

        tracing::info!(post_id = %id, "Created new post");
        Ok(Response::new(post))
    }
}

pub async fn run_services() -> GatewayResult<()> {
    let addr: SocketAddr = "127.0.0.1:50052"
        .parse()
        .map_err(|e| grpc_graphql_gateway::Error::Other(anyhow::Error::new(e)))?;

    let service = ServiceImpl::default();
    tracing::info!(
        "gRPC service live_example.UserService listening on {}",
        addr
    );

    // Run the server directly (not in a spawned task)
    Server::builder()
        .add_service(live_example::user_service_server::UserServiceServer::new(
            service,
        ))
        .serve(addr)
        .await
        .map_err(|e| {
            tracing::error!(error = %e, "gRPC server failed to start");
            grpc_graphql_gateway::Error::Other(anyhow::Error::new(e))
        })?;

    Ok(())
}

pub fn gateway_builder() -> GatewayResult<GatewayBuilder> {
    // The descriptor set is produced by your build.rs using tonic-build.
    let mut builder = Gateway::builder().with_descriptor_set_bytes(DESCRIPTOR_SET);

    if FEDERATION_ENABLED {
        tracing::info!(
            "Federation enabled (entities: {entities})",
            entities = describe_entities()
        );
        builder = builder
            .enable_federation()
            .with_entity_resolver(default_entity_resolver());
        // For subgraphs, restrict the schema to the services owned by this process:
        // builder = builder.with_services(["your.package.Service"]);
    }

    // Add gRPC backends for each service discovered in your protos.
    for svc in services::ALL {
        tracing::info!(
            "{svc} -> {endpoint} (queries: {queries}; mutations: {mutations}; subscriptions: {subscriptions}; resolvers: {resolvers})",
            svc = svc.name,
            endpoint = svc.endpoint,
            queries = describe(svc.queries),
            mutations = describe(svc.mutations),
            subscriptions = describe(svc.subscriptions),
            resolvers = describe_resolvers(svc.resolvers),
        );
        let client = GrpcClient::builder(svc.endpoint)
            .insecure(svc.insecure)
            .lazy(true)
            .connect_lazy()?;
        builder = builder.add_grpc_client(svc.name, client);
    }

    // Update the endpoints above to point at your actual services.

    Ok(builder)
}

pub fn gateway_builder_for_service(svc: &ServiceConfig) -> GatewayResult<GatewayBuilder> {
    let mut builder = Gateway::builder()
        .with_descriptor_set_bytes(DESCRIPTOR_SET)
        .with_persisted_queries(PersistedQueryConfig::default())
        .with_response_cache(CacheConfig {
            max_size: 1000,
            default_ttl: Duration::from_secs(60),
            ..Default::default()
        })
        .with_request_collapsing(RequestCollapsingConfig::default());

    tracing::info!(
        "{svc} -> {endpoint} (queries: {queries}; mutations: {mutations}; subscriptions: {subscriptions}; resolvers: {resolvers})",
        svc = svc.name,
        endpoint = svc.endpoint,
        queries = describe(svc.queries),
        mutations = describe(svc.mutations),
        subscriptions = describe(svc.subscriptions),
        resolvers = describe_resolvers(svc.resolvers),
    );

    if FEDERATION_ENABLED {
        builder = builder
            .enable_federation()
            .with_entity_resolver(default_entity_resolver())
            .with_services([svc.name]);
    } else {
        builder = builder.with_services([svc.name]);
    }

    let client = GrpcClient::builder(svc.endpoint)
        .insecure(svc.insecure)
        .lazy(true)
        .connect_lazy()?;
    builder = builder.add_grpc_client(svc.name, client);

    Ok(builder)
}

pub fn gateway_builder_for(name: &str) -> GatewayResult<Option<GatewayBuilder>> {
    for svc in services::ALL {
        if svc.name == name {
            return gateway_builder_for_service(svc).map(Some);
        }
    }
    Ok(None)
}

pub async fn run_live_example_userservice_gateway() -> GatewayResult<()> {
    gateway_builder_for_service(&services::LIVE_EXAMPLE_USERSERVICE)?
        .serve("0.0.0.0:9000")
        .await
}

#[tokio::main]
async fn main() -> GatewayResult<()> {
    // Basic logging; adjust as desired.
    tracing_subscriber::registry()
        .with(tracing_subscriber::fmt::layer())
        .init();

    tracing::info!(
        "GraphQL operations -> queries: {queries}; mutations: {mutations}; subscriptions: {subscriptions}; resolvers: {resolvers}",
        queries = describe(QUERIES),
        mutations = describe(MUTATIONS),
        subscriptions = describe(SUBSCRIPTIONS),
        resolvers = describe_resolvers(RESOLVERS),
    );

    if FEDERATION_ENABLED {
        tracing::info!(
            "Federation entities -> {entities}",
            entities = describe_entities()
        );
    }

    // NOTE: Resolver entries are listed above; the runtime currently warns that they are not implemented.
    // Spawn stub gRPC services and per-service gateways in this process:
    tokio::spawn(async {
        if let Err(e) = run_services().await {
            tracing::error!(error = %e, "gRPC services exited with error");
        }
    });

    tokio::spawn(async {
        if let Err(e) = run_live_example_userservice_gateway().await {
            tracing::error!(error = %e, "GraphQL gateway for live_example.UserService exited with error");
        }
    });

    // Keep running until interrupted so the spawned servers stay alive.
    tokio::signal::ctrl_c()
        .await
        .map_err(|e| grpc_graphql_gateway::Error::Other(anyhow::Error::new(e)))?;

    Ok(())
}