Class: Ask::State::Providers::Postgres

Inherits:
Adapter
  • Object
show all
Defined in:
lib/ask/state/providers/postgres.rb

Overview

A persistent key-value store backed by PostgreSQL.

Implements the full Adapter contract: key-value storage, distributed locking, message queues, and ordered lists. Uses the pg gem for database access.

Examples:

store = Ask::State::Providers::Postgres.new(
  url: ENV["DATABASE_URL"]
)
store.set("key", { hello: "world" })
store.get("key")  # => { "hello" => "world" }

Constant Summary collapse

MIGRATIONS = 
<<~SQL
  CREATE TABLE IF NOT EXISTS state_store (
    key         TEXT PRIMARY KEY NOT NULL,
    value       TEXT NOT NULL,
    expires_at  TIMESTAMPTZ
  );

  CREATE TABLE IF NOT EXISTS locks (
    key         TEXT PRIMARY KEY NOT NULL,
    token       TEXT NOT NULL,
    expires_at  TIMESTAMPTZ NOT NULL
  );

  CREATE TABLE IF NOT EXISTS queues (
    id           BIGSERIAL PRIMARY KEY,
    queue_name   TEXT NOT NULL,
    value        TEXT NOT NULL,
    enqueued_at  TIMESTAMPTZ NOT NULL
  );
  CREATE INDEX IF NOT EXISTS idx_queues_queue_name
    ON queues (queue_name, id);

  CREATE TABLE IF NOT EXISTS lists (
    id         BIGSERIAL PRIMARY KEY,
    list_key   TEXT NOT NULL,
    value      TEXT NOT NULL
  );
  CREATE INDEX IF NOT EXISTS idx_lists_list_key
    ON lists (list_key, id);
SQL

Instance Method Summary collapse

Constructor Details

#initialize(url: ENV.fetch("DATABASE_URL", "postgres://localhost:5432/ask_state"), pool_size: 5) ⇒ Postgres

Creates a new PostgreSQL-backed state store.

Parameters:

  • url (String) (defaults to: ENV.fetch("DATABASE_URL", "postgres://localhost:5432/ask_state")) —

    Postgres connection URL

  • pool_size (Integer) (defaults to: 5) —

    connection pool size (default 5)



59
60
61
62
63
64
65
66
67
68
 
# File 'lib/ask/state/providers/postgres.rb', line 59

def initialize(url: ENV.fetch("DATABASE_URL", "postgres://localhost:5432/ask_state"),
               pool_size: 5)
  require "pg"

  @pool = ConnectionPool.new(size: pool_size) do
    ::PG.connect(url)
  end

  migrate
end

Instance Method Details

#acquire_lock(key, ttl: 10) ⇒ Object

-- distributed locking --



160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
 
# File 'lib/ask/state/providers/postgres.rb', line 160

def acquire_lock(key, ttl: 10)
  now = Time.now.utc
  expires_at_time = now + ttl
  token = SecureRandom.hex(16)

  acquired = @pool.with do |conn|
    # First clean up expired locks
    conn.exec_params("DELETE FROM locks WHERE key = $1 AND expires_at <= $2",
                   [key, now])

    result = conn.exec_params(<<~SQL, [key, expires_at_time, token, now])
      INSERT INTO locks (key, expires_at, token)
      SELECT $1, $2, $3
      WHERE NOT EXISTS (
        SELECT 1 FROM locks
        WHERE key = $1
      )
    SQL
    result.cmd_tuples > 0
  end

  acquired ? Lock.new(id: key, token: token, expires_at: expires_at_time) : nil
end

#clear ⇒ Object 



121
122
123
124
125
126
127
128
 
# File 'lib/ask/state/providers/postgres.rb', line 121

def clear
  @pool.with do |conn|
    conn.exec("DELETE FROM state_store")
    conn.exec("DELETE FROM locks")
    conn.exec("DELETE FROM queues")
    conn.exec("DELETE FROM lists")
  end
end

#close ⇒ Object 



305
306
307
 
# File 'lib/ask/state/providers/postgres.rb', line 305

def close
  @pool&.shutdown { |c| c.close }
end

#delete(key) ⇒ Object 



94
95
96
97
98
99
100
101
 
# File 'lib/ask/state/providers/postgres.rb', line 94

def delete(key)
  @pool.with do |conn|
    conn.exec_params("DELETE FROM state_store WHERE key = $1", [key])
    # delete removes everything under the key, including ordered
    # lists (consumers store event feeds as lists).
    conn.exec_params("DELETE FROM lists WHERE list_key = $1", [key])
  end
end

#dequeue(queue) ⇒ Object 



209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
 
# File 'lib/ask/state/providers/postgres.rb', line 209

def dequeue(queue)
  @pool.with do |conn|
    result = conn.exec_params(<<~SQL, [queue])
      DELETE FROM queues
      WHERE id = (
        SELECT id FROM queues
        WHERE queue_name = $1
        ORDER BY id ASC
        LIMIT 1
      )
      RETURNING value, enqueued_at
    SQL
    return nil if result.ntuples == 0

    QueueEntry.new(
      id: SecureRandom.uuid,
      value: JSON.parse(result[0]["value"]),
      enqueued_at: Time.parse(result[0]["enqueued_at"])
    )
  end
end

#enqueue(queue, value) ⇒ Object

-- message queues --



196
197
198
199
200
201
202
203
204
205
206
207
 
# File 'lib/ask/state/providers/postgres.rb', line 196

def enqueue(queue, value)
  id = SecureRandom.uuid

  @pool.with do |conn|
    conn.exec_params(<<~SQL, [queue, JSON.generate(value), Time.now.utc])
      INSERT INTO queues (queue_name, value, enqueued_at)
      VALUES ($1, $2, $3)
    SQL
  end

  QueueEntry.new(id: id, value: value, enqueued_at: Time.now)
end

#exists?(key) ⇒ Boolean

Returns:

  • (Boolean) 


130
131
132
133
134
135
136
137
138
 
# File 'lib/ask/state/providers/postgres.rb', line 130

def exists?(key)
  @pool.with do |conn|
    result = conn.exec_params(<<~SQL, [key, Time.now.utc])
      SELECT 1 FROM state_store
      WHERE key = $1 AND (expires_at IS NULL OR expires_at > $2)
    SQL
    result.ntuples > 0
  end
end

#get(key) ⇒ Object

-- key-value --



72
73
74
75
76
77
78
79
80
 
# File 'lib/ask/state/providers/postgres.rb', line 72

def get(key)
  @pool.with do |conn|
    row = conn.exec_params(<<~SQL, [key, Time.now.utc])
      SELECT value FROM state_store
      WHERE key = $1 AND (expires_at IS NULL OR expires_at > $2)
    SQL
    row.ntuples > 0 ? JSON.parse(row[0]["value"]) : nil
  end
end

#keys(pattern: nil) ⇒ Object 



140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
 
# File 'lib/ask/state/providers/postgres.rb', line 140

def keys(pattern: nil)
  @pool.with do |conn|
    sql, params = if pattern
      like = self.class.glob_to_like(pattern)
      [<<~SQL, [like, Time.now.utc]]
        SELECT key FROM state_store
        WHERE key LIKE $1 AND (expires_at IS NULL OR expires_at > $2)
      SQL
    else
      [<<~SQL, [Time.now.utc]]
        SELECT key FROM state_store
        WHERE (expires_at IS NULL OR expires_at > $1)
      SQL
    end
    conn.exec_params(sql, params).map { |r| r["key"] }
  end
end

#list_append(key, value, max_length: nil) ⇒ Object

-- ordered lists --



242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
 
# File 'lib/ask/state/providers/postgres.rb', line 242

def list_append(key, value, max_length: nil)
  @pool.with do |conn|
    conn.exec_params(<<~SQL, [key, JSON.generate(value)])
      INSERT INTO lists (list_key, value)
      VALUES ($1, $2)
    SQL

    return unless max_length

    conn.exec_params(<<~SQL, [key, key, max_length])
      DELETE FROM lists WHERE id <= (
        SELECT COALESCE(MIN(id), 0) FROM (
          SELECT id FROM lists
          WHERE list_key = $1
          ORDER BY id DESC
          LIMIT $3
        ) sub
      )
    SQL
  end
end

#list_range(key, start = 0, stop = -1)) ⇒ Object 



264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
 
# File 'lib/ask/state/providers/postgres.rb', line 264

def list_range(key, start = 0, stop = -1)
  @pool.with do |conn|
    rows = if stop == -1
      conn.exec_params(<<~SQL, [key, start])
        SELECT value FROM lists
        WHERE list_key = $1
        ORDER BY id ASC
        OFFSET $2
      SQL
    else
      limit = stop - start + 1
      conn.exec_params(<<~SQL, [key, limit, start])
        SELECT value FROM lists
        WHERE list_key = $1
        ORDER BY id ASC
        LIMIT $2 OFFSET $3
      SQL
    end
    rows.map { |r| JSON.parse(r["value"]) }
  end
end

#list_remove(key, value) ⇒ Object 



286
287
288
289
290
291
292
293
294
295
 
# File 'lib/ask/state/providers/postgres.rb', line 286

def list_remove(key, value)
  serialized = JSON.generate(value)
  @pool.with do |conn|
    result = conn.exec_params(
      "DELETE FROM lists WHERE list_key = $1 AND value = $2",
      [key, serialized]
    )
    result.cmd_tuples
  end
end

#queue_depth(queue) ⇒ Object 



231
232
233
234
235
236
237
238
 
# File 'lib/ask/state/providers/postgres.rb', line 231

def queue_depth(queue)
  @pool.with do |conn|
    result = conn.exec_params(
      "SELECT COUNT(*) AS cnt FROM queues WHERE queue_name = $1", [queue]
    )
    result[0]["cnt"].to_i
  end
end

#release_lock(key, lock) ⇒ Object 



184
185
186
187
188
189
190
191
192
 
# File 'lib/ask/state/providers/postgres.rb', line 184

def release_lock(key, lock)
  @pool.with do |conn|
    result = conn.exec_params(
      "DELETE FROM locks WHERE key = $1 AND token = $2",
      [key, lock.token]
    )
    result.cmd_tuples > 0
  end
end

#set(key, value, ttl: nil) ⇒ Object 



82
83
84
85
86
87
88
89
90
91
92
 
# File 'lib/ask/state/providers/postgres.rb', line 82

def set(key, value, ttl: nil)
  @pool.with do |conn|
    conn.exec_params(<<~SQL, [key, JSON.generate(value), ttl ? Time.now.utc + ttl : nil])
      INSERT INTO state_store (key, value, expires_at)
      VALUES ($1, $2, $3)
      ON CONFLICT (key) DO UPDATE SET
        value = EXCLUDED.value,
        expires_at = EXCLUDED.expires_at
    SQL
  end
end

#set_if_not_exists(key, value, ttl: nil) ⇒ Object 



103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
 
# File 'lib/ask/state/providers/postgres.rb', line 103

def set_if_not_exists(key, value, ttl: nil)
  @pool.with do |conn|
    now = Time.now.utc
    expires = ttl ? now + ttl : nil

    result = conn.exec_params(<<~SQL, [key, JSON.generate(value), expires, now])
      INSERT INTO state_store (key, value, expires_at)
      SELECT $1, $2, $3
      WHERE NOT EXISTS (
        SELECT 1 FROM state_store
        WHERE key = $1 AND (expires_at IS NULL OR expires_at > $4)
      )
      ON CONFLICT (key) DO NOTHING
    SQL
    result.cmd_tuples > 0
  end
end

#setup! ⇒ Object

Idempotent setup — creates tables if they don't exist. Called automatically on initialize. Safe to call multiple times.



301
302
303
 
# File 'lib/ask/state/providers/postgres.rb', line 301

def setup!
  @pool.with { |conn| conn.exec(MIGRATIONS) }
end