GraphQL Subscriptions in Ruby

gem install graphql sinatra sinatra-contrib puma faye-websocket

require 'sinatra'
require 'sinatra/json'
require 'graphql'
require 'faye/websocket'
require 'json'
require 'securerandom'

# In-memory storage for books
# We use global variables (prefixed with $) here so that
# the data is accessible across all classes and routes in
# our application. In a production app, you would use a
# database instead — global variables are only suitable
# for simple examples like this because they are mutable
# from anywhere, which can make code harder to reason
# about and is not thread-safe.
$books = [
  { id: '1', title: 'The Hobbit', author: 'J.R.R. Tolkien' },
  { id: '2', title: 'Harry Potter', author: 'J.K. Rowling' }
]

# In-memory subscription storage
$subscriptions = {}

# The query declares variables with $ and their types
query = <<~GRAPHQL
  mutation($title: String!, $author: String!) {
    addBook(title: $title, author: $author) {
      id
      title
      author
    }
  }
GRAPHQL

# The values are passed separately
variables = { title: "1984", author: "George Orwell" }

# Both are sent together in the request body
request.body = JSON.generate({ query: query, variables: variables })

# Define the Book type
class BookType < GraphQL::Schema::Object
  field :id, ID, null: false
  field :title, String, null: false
  field :author, String, null: false
end

# Define the Query type
class QueryType < GraphQL::Schema::Object
  field :books, [BookType], null: false

  def books
    $books
  end
end

# Define the Mutation type
class MutationType < GraphQL::Schema::Object
  field :add_book, BookType, null: false do
    argument :title, String, required: true
    argument :author, String, required: true
  end

  def add_book(title:, author:)
    new_book = {
      id: SecureRandom.uuid,
      title: title,
      author: author
    }
    $books << new_book
    
    # Trigger the subscription — this notifies all clients
    # currently subscribed to 'book_added' that a new book
    # has been created. We'll define MySchema shortly in the
    # next section; Ruby allows forward references like this
    # because the call is inside a method body, which isn't
    # evaluated until the method is actually called at runtime.
    MySchema.subscriptions.trigger('book_added', {}, new_book)
    
    new_book
  end
end

# Define the Subscription type
class SubscriptionType < GraphQL::Schema::Object
  field :book_added, BookType, null: false

  def book_added
    object
  end
end

# Simple in-memory subscription implementation
class InMemorySubscriptions < GraphQL::Subscriptions
  def initialize(**rest)
    super
    @subscriptions = {}
    @mutex = Mutex.new
  end

  # Called when a client subscribes. Stores the subscription
  # query so it can be re-executed when an event is triggered.
  def write_subscription(query, events)
    @mutex.synchronize do
      events.each do |event|
        @subscriptions[event.topic] ||= []
        @subscriptions[event.topic] << {
          query: query,
          variables: event.arguments
        }
      end
    end
  end

  # Called when an event is triggered (e.g., a book is added).
  # Iterates over all subscribers for that event so each one
  # can receive the update.
  def each_subscription_id(event)
    @mutex.synchronize do
      subscribers = @subscriptions[event.topic] || []
      subscribers.each_with_index do |subscription, index|
        yield index.to_s
      end
    end
  end

  # Retrieves a stored subscription by its ID so the query
  # can be re-executed with the new data.
  def read_subscription(subscription_id)
    topic = @subscriptions.keys.first
    @subscriptions[topic][subscription_id.to_i] if topic
  end

  # Removes a subscription when a client unsubscribes.
  def delete_subscription(subscription_id)
    @mutex.synchronize do
      @subscriptions.each do |topic, subscribers|
        subscribers.delete_at(subscription_id.to_i)
      end
    end
  end

  # Delivers the result of a triggered event to the client.
  # In this simple implementation, we store the result for
  # later retrieval. A production system would push the
  # result directly to the client over a WebSocket.
  def deliver(subscription_id, result)
    @mutex.synchronize do
      @subscriptions[:results] ||= {}
      @subscriptions[:results][subscription_id] = result
    end
  end
end

class MySchema < GraphQL::Schema
  query QueryType
  mutation MutationType
  subscription SubscriptionType
  
  use InMemorySubscriptions
end

set :server, 'puma'
set :sockets, []

# HTTP endpoint for queries and mutations
post '/graphql' do
  request_payload = JSON.parse(request.body.read)
  query = request_payload['query']
  variables = request_payload['variables'] || {}
  
  result = MySchema.execute(
    query,
    variables: variables,
    context: {}
  )
  
  json result
end

# WebSocket endpoint for subscriptions
get '/graphql' do
  if Faye::WebSocket.websocket?(request.env)
    ws = Faye::WebSocket.new(request.env)
    
    ws.on :open do |event|
      settings.sockets << ws
    end
    
    ws.on :message do |event|
      data = JSON.parse(event.data)
      
      if data['type'] == 'start'
        # Client wants to start a subscription (or run a query).
        query = data['payload']['query']
        variables = data['payload']['variables'] || {}
        
        result = MySchema.execute(
          query,
          variables: variables,
          context: { socket: ws },
          # operation_name identifies which operation to run when
          # a query string contains multiple named operations. For
          # single-operation queries (like ours), this can be nil.
          operation_name: data['payload']['operationName']
        )
        
        # result.subscription? returns true if the executed query
        # was a subscription. In that case, the client is now
        # registered and will receive future updates via the
        # deliver method. We still send an initial acknowledgment.
        if result.subscription?
          ws.send(JSON.generate({
            type: 'data',
            id: data['id'],
            payload: { data: result.to_h['data'] }
          }))
        else
          ws.send(JSON.generate({
            type: 'data',
            id: data['id'],
            payload: result.to_h
          }))
        end
      elsif data['type'] == 'stop'
        # Client wants to unsubscribe
      end
    end
    
    ws.on :close do |event|
      settings.sockets.delete(ws)
      ws = nil
    end
    
    ws.rack_response
  else
    status 400
    body 'WebSocket connection required'
  end
end

gem install faye-websocket eventmachine

require 'faye/websocket'
require 'eventmachine'
require 'json'
require 'net/http'
require 'uri'

# Define the GraphQL endpoint
GRAPHQL_ENDPOINT = 'http://localhost:4000/graphql'
WEBSOCKET_ENDPOINT = 'ws://localhost:4000/graphql'

GET_BOOKS_QUERY = <<~GRAPHQL
  query {
    books {
      id
      title
      author
    }
  }
GRAPHQL

ADD_BOOK_MUTATION = <<~GRAPHQL
  mutation($title: String!, $author: String!) {
    addBook(title: $title, author: $author) {
      id
      title
      author
    }
  }
GRAPHQL

BOOK_ADDED_SUBSCRIPTION = <<~GRAPHQL
  subscription {
    bookAdded {
      id
      title
      author
    }
  }
GRAPHQL

def fetch_graphql(query, variables = {})
  uri = URI.parse(GRAPHQL_ENDPOINT)
  http = Net::HTTP.new(uri.host, uri.port)
  
  request = Net::HTTP::Post.new(uri.path)
  request['Content-Type'] = 'application/json'
  request.body = JSON.generate({
    query: query,
    variables: variables
  })
  
  response = http.request(request)
  result = JSON.parse(response.body)
  
  if result['errors']
    raise "GraphQL error: #{result['errors'].map { |e| e['message'] }.join(', ')}"
  end
  
  result['data']
end

def setup_subscription
  # EM.run starts the EventMachine event loop. All WebSocket
  # communication happens inside this block. The loop keeps
  # running until we call EM.stop.
  EM.run do
    ws = Faye::WebSocket::Client.new(WEBSOCKET_ENDPOINT)
    subscription_id = '1'
    
    ws.on :open do |event|
      puts 'WebSocket connection opened'
      
      # Send a 'start' message to subscribe
      ws.send(JSON.generate({
        id: subscription_id,
        type: 'start',
        payload: {
          query: BOOK_ADDED_SUBSCRIPTION,
          variables: {}
        }
      }))
    end
    
    ws.on :message do |event|
      data = JSON.parse(event.data)
      
      # The server sends 'data' messages with subscription updates
      if data['type'] == 'data' && data['payload']['data']
        book = data['payload']['data']['bookAdded']
        puts "Book added: #{book}"
        
        # Send a 'stop' message to unsubscribe
        ws.send(JSON.generate({
          id: subscription_id,
          type: 'stop'
        }))
        
        sleep 1
        ws.close
        EM.stop
      end
    end
    
    ws.on :error do |event|
      puts "WebSocket error: #{event.message}"
      EM.stop
    end
    
    ws.on :close do |event|
      puts 'WebSocket connection closed'
      ws = nil
    end
  end
end

# Fetch books
begin
  books_data = fetch_graphql(GET_BOOKS_QUERY)
  puts "Books: #{books_data['books']}"
rescue => e
  puts "Error fetching books: #{e.message}"
end

# Add a new book
begin
  add_book_data = fetch_graphql(ADD_BOOK_MUTATION, { title: '1984', author: 'George Orwell' })
  puts "Added book: #{add_book_data['addBook']}"
rescue => e
  puts "Error adding book: #{e.message}"
end

# Set up subscription to listen for new books
setup_subscription