SubscriptionSchema.swift

import GraphQL
import GraphQLRxSwift
import NIO
import RxSwift

// MARK: Types

struct Email: Encodable {
    let from: String
    let subject: String
    let message: String
    let unread: Bool
    let priority: Int

    init(from: String, subject: String, message: String, unread: Bool, priority: Int = 0) {
        self.from = from
        self.subject = subject
        self.message = message
        self.unread = unread
        self.priority = priority
    }
}

struct Inbox: Encodable {
    let emails: [Email]
}

struct EmailEvent: Encodable {
    let email: Email
    let inbox: Inbox
}

// MARK: Schema

let EmailType = try! GraphQLObjectType(
    name: "Email",
    fields: [
        "from": GraphQLField(
            type: GraphQLString
        ),
        "subject": GraphQLField(
            type: GraphQLString
        ),
        "message": GraphQLField(
            type: GraphQLString
        ),
        "unread": GraphQLField(
            type: GraphQLBoolean
        ),
    ]
)
let InboxType = try! GraphQLObjectType(
    name: "Inbox",
    fields: [
        "emails": GraphQLField(
            type: GraphQLList(EmailType)
        ),
        "total": GraphQLField(
            type: GraphQLInt,
            resolve: { inbox, _, _, _ in
                (inbox as! Inbox).emails.count
            }
        ),
        "unread": GraphQLField(
            type: GraphQLInt,
            resolve: { inbox, _, _, _ in
                (inbox as! Inbox).emails.filter { $0.unread }.count
            }
        ),
    ]
)
let EmailEventType = try! GraphQLObjectType(
    name: "EmailEvent",
    fields: [
        "email": GraphQLField(
            type: EmailType
        ),
        "inbox": GraphQLField(
            type: InboxType
        ),
    ]
)
let EmailQueryType = try! GraphQLObjectType(
    name: "Query",
    fields: [
        "inbox": GraphQLField(
            type: InboxType
        ),
    ]
)

// MARK: Test Helpers

let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)

class EmailDb {
    var emails: [Email]
    let publisher: PublishSubject<Any>
    let disposeBag: DisposeBag

    init() {
        emails = [
            Email(
                from: "joe@graphql.org",
                subject: "Hello",
                message: "Hello World",
                unread: false
            ),
        ]
        publisher = PublishSubject<Any>()
        disposeBag = DisposeBag()
    }

    /// Adds a new email to the database and triggers all observers
    func trigger(email: Email) {
        emails.append(email)
        publisher.onNext(email)
    }

    /// Returns the default email schema, with standard resolvers.
    func defaultSchema() -> GraphQLSchema {
        return emailSchemaWithResolvers(
            resolve: { emailAny, _, _, eventLoopGroup, _ throws -> EventLoopFuture<Any?> in
                if let email = emailAny as? Email {
                    return eventLoopGroup.next().makeSucceededFuture(EmailEvent(
                        email: email,
                        inbox: Inbox(emails: self.emails)
                    ))
                } else {
                    throw GraphQLError(message: "\(type(of: emailAny)) is not Email")
                }
            },
            subscribe: { _, args, _, eventLoopGroup, _ throws -> EventLoopFuture<Any?> in
                let priority = args["priority"].int ?? 0
                let filtered = self.publisher.filter { emailAny throws in
                    if let email = emailAny as? Email {
                        return email.priority >= priority
                    } else {
                        return true
                    }
                }
                return eventLoopGroup.next().makeSucceededFuture(filtered.toEventStream())
            }
        )
    }

    /// Generates a subscription to the database using the default schema and resolvers
    func subscription(
        query: String,
        variableValues: [String: Map] = [:]
    ) throws -> SubscriptionEventStream {
        return try createSubscription(schema: defaultSchema(), query: query, variableValues: variableValues)
    }
}

/// Generates an email schema with the specified resolve and subscribe methods
func emailSchemaWithResolvers(resolve: GraphQLFieldResolve? = nil, subscribe: GraphQLFieldResolve? = nil) -> GraphQLSchema {
    return try! GraphQLSchema(
        query: EmailQueryType,
        subscription: try! GraphQLObjectType(
            name: "Subscription",
            fields: [
                "importantEmail": GraphQLField(
                    type: EmailEventType,
                    args: [
                        "priority": GraphQLArgument(
                            type: GraphQLInt
                        ),
                    ],
                    resolve: resolve,
                    subscribe: subscribe
                ),
            ]
        )
    )
}

/// Generates a subscription from the given schema and query. It's expected that the resolver/database interactions are configured by the caller.
func createSubscription(
    schema: GraphQLSchema,
    query: String,
    variableValues: [String: Map] = [:]
) throws -> SubscriptionEventStream {
    let result = try graphqlSubscribe(
        queryStrategy: SerialFieldExecutionStrategy(),
        mutationStrategy: SerialFieldExecutionStrategy(),
        subscriptionStrategy: SerialFieldExecutionStrategy(),
        instrumentation: NoOpInstrumentation,
        schema: schema,
        request: query,
        rootValue: (),
        context: (),
        eventLoopGroup: eventLoopGroup,
        variableValues: variableValues,
        operationName: nil
    ).wait()

    if let stream = result.stream {
        return stream
    } else {
        throw result.errors.first! // We may have more than one...
    }
}