Skip to content

Commit 0e4881d

Browse files
NeedleInAJayStackNeedleInAJayStack
committed
feat: Adds subscription support
1 parent 36c16c6 commit 0e4881d

7 files changed

Lines changed: 223 additions & 56 deletions

File tree

‎Examples/HelloWorldServer/Sources/HelloWorldServer/Resolvers.swift‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ public class Context: @unchecked Sendable {
77
// User can choose structure
88
var users: [String: User]
99
var posts: [String: Post]
10+
var onTriggerWatch: () -> Void = {}
1011

1112
init(
1213
users: [String: User],
@@ -15,6 +16,10 @@ public class Context: @unchecked Sendable {
1516
self.users = users
1617
self.posts = posts
1718
}
19+
20+
func triggerWatch() {
21+
onTriggerWatch()
22+
}
1823
}
1924
// Scalars must be represented by a Swift type of the same name, conforming to the Scalar protocol
2025
public struct EmailAddress: Scalar {
@@ -59,6 +64,7 @@ public struct EmailAddress: Scalar {
5964
struct Resolvers: ResolversProtocol {
6065
typealias Query = HelloWorldServer.Query
6166
typealias Mutation = HelloWorldServer.Mutation
67+
typealias Subscription = HelloWorldServer.Subscription
6268
}
6369
struct User: UserProtocol {
6470
// User can choose structure
@@ -149,3 +155,14 @@ struct Mutation: MutationProtocol {
149155
return user
150156
}
151157
}
158+
159+
struct Subscription: SubscriptionProtocol {
160+
// Required implementations
161+
static func watchUser(id: String, context: Context, info: GraphQLResolveInfo) async throws -> AnyAsyncSequence<(any UserProtocol)?> {
162+
return AsyncStream<(any UserProtocol)?> { continuation in
163+
context.onTriggerWatch = { [weak context] in
164+
continuation.yield(context?.users[id])
165+
}
166+
}.any()
167+
}
168+
}

‎Examples/HelloWorldServer/Sources/HelloWorldServer/schema.graphql‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,3 +118,7 @@ type Query {
118118
type Mutation {
119119
upsertUser(userInfo: UserInfo!): User!
120120
}
121+
122+
type Subscription {
123+
watchUser(id: ID!): User
124+
}

‎Examples/HelloWorldServer/Tests/HelloWorldServerTests/HelloWorldServerTests.swift‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,4 +86,59 @@ struct HelloWorldServerTests {
8686
)
8787
#expect(actual == expected)
8888
}
89+
90+
@Test func subscription() async throws {
91+
let schema = try buildGraphQLSchema(resolvers: Resolvers.self)
92+
let context = Context(
93+
users: ["1" : .init(id: "1", name: "John", email: "john@example.com", age: 18, role: .user)],
94+
posts: [:]
95+
)
96+
let stream = try await graphqlSubscribe(
97+
schema: schema,
98+
request: """
99+
subscription {
100+
watchUser(id: "1") {
101+
id
102+
name
103+
email
104+
age
105+
role
106+
}
107+
}
108+
""",
109+
context: context
110+
).get()
111+
112+
var iterator = stream.makeAsyncIterator()
113+
114+
context.triggerWatch()
115+
#expect(
116+
try await iterator.next() == GraphQLResult(
117+
data: [
118+
"watchUser": [
119+
"id": "1",
120+
"name": "John",
121+
"email": "john@example.com",
122+
"age": 18,
123+
"role": "USER"
124+
]
125+
]
126+
)
127+
)
128+
129+
context.triggerWatch()
130+
#expect(
131+
try await iterator.next() == GraphQLResult(
132+
data: [
133+
"watchUser": [
134+
"id": "1",
135+
"name": "John",
136+
"email": "john@example.com",
137+
"age": 18,
138+
"role": "USER"
139+
]
140+
]
141+
)
142+
)
143+
}
89144
}

‎Sources/GraphQLGeneratorCore/Generator/SchemaGenerator.swift‎

Lines changed: 59 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -123,15 +123,23 @@ package struct SchemaGenerator {
123123
if let queryType = schema.queryType {
124124
output += """
125125
126-
\(try generateQueryTypeDefinition(for: queryType).indent(1))
126+
\(try generateRootTypeDefinition(for: queryType, rootType: .query).indent(1))
127127
"""
128128
}
129129

130130
// Generate Mutation type if it exists
131131
if let mutationType = schema.mutationType {
132132
output += """
133133
134-
\(try generateMutationTypeDefinition(for: mutationType).indent(1))
134+
\(try generateRootTypeDefinition(for: mutationType, rootType: .mutation).indent(1))
135+
"""
136+
}
137+
138+
// Generate Subscription type if it exists
139+
if let subscriptionType = schema.subscriptionType {
140+
output += """
141+
142+
\(try generateRootTypeDefinition(for: subscriptionType, rootType: .subscription).indent(1))
135143
"""
136144
}
137145

@@ -151,6 +159,13 @@ package struct SchemaGenerator {
151159
"""
152160
}
153161

162+
if schema.subscriptionType != nil {
163+
output += """
164+
,
165+
subscription: subscription
166+
"""
167+
}
168+
154169
output += """
155170
156171
)
@@ -524,55 +539,25 @@ package struct SchemaGenerator {
524539
return output
525540
}
526541

527-
private func generateQueryTypeDefinition(for type: GraphQLObjectType) throws -> String {
528-
var output = """
529-
530-
let query = try GraphQLObjectType(
531-
name: "Query",
532-
"""
533-
534-
if let description = type.description {
535-
output += """
536-
537-
description: \"\"\"
538-
\(description.indent(1, includeFirst: false))
539-
\"\"\",
540-
"""
542+
private func generateRootTypeDefinition(for type: GraphQLObjectType, rootType: RootType) throws -> String {
543+
let variableName: String
544+
let target: ResolverTarget
545+
switch rootType {
546+
case .query:
547+
variableName = "query"
548+
target = .query
549+
case .mutation:
550+
variableName = "mutation"
551+
target = .mutation
552+
case .subscription:
553+
variableName = "subscription"
554+
target = .subscription
541555
}
542556

543-
output += """
544-
545-
fields: [
546-
"""
547-
548-
// Generate fields
549-
let fields = try type.fields()
550-
for (fieldName, field) in fields {
551-
output += """
552-
553-
\(try generateFieldDefinition(
554-
fieldName: fieldName,
555-
field: field,
556-
target: .query,
557-
parentType: type
558-
).indent(2))
559-
"""
560-
}
561-
562-
output += """
563-
564-
]
565-
)
566-
"""
567-
568-
return output
569-
}
570-
571-
private func generateMutationTypeDefinition(for type: GraphQLObjectType) throws -> String {
572557
var output = """
573558
574-
let mutation = try GraphQLObjectType(
575-
name: "Mutation",
559+
let \(variableName) = try GraphQLObjectType(
560+
name: "\(type.name)",
576561
"""
577562

578563
if let description = type.description {
@@ -597,7 +582,7 @@ package struct SchemaGenerator {
597582
\(try generateFieldDefinition(
598583
fieldName: fieldName,
599584
field: field,
600-
target: .mutation,
585+
target: target,
601586
parentType: type
602587
).indent(2))
603588
"""
@@ -616,7 +601,7 @@ package struct SchemaGenerator {
616601
fieldName: String,
617602
field: GraphQLField,
618603
target: ResolverTarget,
619-
parentType: GraphQLType
604+
parentType: GraphQLNamedType
620605
) throws -> String {
621606
var output = """
622607
@@ -704,10 +689,23 @@ package struct SchemaGenerator {
704689
target: ResolverTarget,
705690
parentType: GraphQLType
706691
) throws -> String {
707-
var output = """
692+
var output = ""
693+
694+
if target == .subscription {
695+
output += """
696+
697+
resolve: { source, _, _, _ in
698+
return source
699+
},
700+
subscribe: { source, args, context, info in
701+
"""
702+
} else {
703+
output += """
704+
705+
resolve: { source, args, context, info in
706+
"""
707+
}
708708

709-
resolve: { source, args, context, info in
710-
"""
711709

712710
// Build argument list
713711
var argsList: [String] = []
@@ -754,7 +752,8 @@ package struct SchemaGenerator {
754752
let targetName = switch target {
755753
case .parent: "parent"
756754
case .query: "Resolvers.Query"
757-
case .mutation: "Resolvers.Mutation"
755+
case .mutation: "Resolvers.Mutation"
756+
case .subscription: "Resolvers.Subscription"
758757
}
759758
let functionName = nameGenerator.swiftMemberName(for: fieldName)
760759
output += """
@@ -796,9 +795,16 @@ package struct SchemaGenerator {
796795
throw GeneratorError.unsupportedType("Unknown type: \(type)")
797796
}
798797

798+
private enum RootType {
799+
case query
800+
case mutation
801+
case subscription
802+
}
803+
799804
private enum ResolverTarget {
800805
case parent
801806
case query
802807
case mutation
808+
case subscription
803809
}
804810
}

‎Sources/GraphQLGeneratorCore/Generator/TypeGenerator.swift‎

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ package struct TypeGenerator {
1212
1313
import Foundation
1414
import GraphQL
15+
import GraphQLGeneratorRuntime
16+
1517
"""
1618

1719
// Generate ResolversProtocol
@@ -32,6 +34,12 @@ package struct TypeGenerator {
3234
associatedtype Mutation: MutationProtocol
3335
"""
3436
}
37+
if schema.subscriptionType != nil {
38+
output += """
39+
40+
associatedtype Subscription: SubscriptionProtocol
41+
"""
42+
}
3543
output += """
3644
3745
}
@@ -138,7 +146,13 @@ package struct TypeGenerator {
138146
"""
139147
}
140148

141-
// TODO: Add subscription types
149+
// Generate Mutation type
150+
if let subscriptionType = schema.subscriptionType {
151+
output += """
152+
153+
\(try generateRootTypeProtocol(for: subscriptionType))
154+
"""
155+
}
142156

143157
return output
144158
}
@@ -387,7 +401,10 @@ package struct TypeGenerator {
387401
"""
388402
}
389403

390-
let returnType = try swiftTypeReference(for: field.type, nameGenerator: nameGenerator)
404+
var returnType = try swiftTypeReference(for: field.type, nameGenerator: nameGenerator)
405+
if type.name == "Subscription" {
406+
returnType = "AnyAsyncSequence<\(returnType)>"
407+
}
391408

392409
var params: [String] = []
393410

‎Sources/GraphQLGeneratorRuntime/AnyAsyncSequence.swift‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
2+
/// A type-erased AsyncSequence. This exists because we cannot qualify `AsyncSequence` opaque types with `Element`
3+
/// constraints in our SubscriptionProtocol.
4+
public struct AnyAsyncSequence<Element: Sendable>: AsyncSequence, Sendable {
5+
public typealias Element = Element
6+
public typealias AsyncIterator = AnyAsyncIterator
7+
8+
private let makeAsyncIteratorClosure: @Sendable () -> AsyncIterator
9+
10+
public init<T: AsyncSequence>(_ sequence: T) where T.Element == Element, T: Sendable {
11+
makeAsyncIteratorClosure = {
12+
AnyAsyncIterator(sequence.makeAsyncIterator())
13+
}
14+
}
15+
16+
public func makeAsyncIterator() -> AsyncIterator {
17+
AnyAsyncIterator(makeAsyncIteratorClosure())
18+
}
19+
20+
public struct AnyAsyncIterator: AsyncIteratorProtocol, @unchecked Sendable {
21+
private let nextClosure: () async throws -> Element?
22+
23+
public init<T: AsyncIteratorProtocol>(_ iterator: T) where T.Element == Element {
24+
var iterator = iterator
25+
nextClosure = { try await iterator.next() }
26+
}
27+
28+
public func next() async throws -> Element? {
29+
try await nextClosure()
30+
}
31+
}
32+
}
33+
34+
public extension AsyncSequence where Self: Sendable, Element: Sendable {
35+
/// Create a type erased version of this sequence
36+
func any() -> AnyAsyncSequence<Element> {
37+
AnyAsyncSequence(self)
38+
}
39+
}

0 commit comments

Comments
 (0)