walrus.py

# "Walrus"
#
# A quick and dirty hack to experiment with a kind of automatic query
# plan cache.
#
# This is not real code.  It's entirely imaginary.  But you could imagine
# turning it into a wrapper, a cursor factory, a mix-in, a shirt,
# a sock, a glove, a hat.
#
# Thought: types are going to suck with this transparent approach, right?

import psycopg2

def mangle(sql):
  """Convert the placeholders from format foo to format bar."""
  # this is garbage placeholder code
  i = 1
  while sql.find("%s") != -1:
    sql = sql.replace("%s", "$%d" % (i,), 1)
    i += 1
  print sql
  return sql

class Statement:
  """A database statement which may or may not have been prepared."""
  def __init__(self, sql):
    self.name = None
    self.sql = sql
    self.execute_statement = None
    self.next = None
    self.prev = None
    self.used = 0

class StatementCache:
  """A cache for statements associated with one connection."""

  def __init__(self, min_reuses_before_prepare, max_statements):
    self.lru = None
    self.mru = None
    self.index = {}
    self.min_reuses_before_prepare = min_reuses_before_prepare
    self.max_statements = max_statements
    self.counter = 0

  def unlink(self, statement):
    """Remove a statement from the list."""
    if statement.prev:
      statement.prev.next = statement.next
    if statement.next:
      statement.next.prev = statement.prev
    if self.lru == statement:
      self.lru = statement.next
    if self.mru == statement:
      self.mru = statement.prev
    statement.prev = None
    statement.next = None

  def push_mru(self, statement):
    """Add a statement to the most-recently-used end of the list."""
    statement.prev = self.mru
    if self.mru:
      self.mru.next = statement
      statement.prev = self.mru
    if self.lru == None:
      self.lru = statement
    self.mru = statement

  def get_statement(self, cursor, sql, nargs):
    """Obtain, or recycle, a statement.  Also PREPARE it if it's been
       reused enough times."""

    if self.index.has_key(sql):
      # we have this statement already!
      statement = self.index[sql]
      print "[reusing statement for %s] %s" % (statement.sql, statement.name)
      self.unlink(statement)
      self.push_mru(statement)
    else:
      # we have to make a new statement
      if len(self.index) >= self.max_statements:
        # first, we need to get rid of the least recently used statement
        if self.lru.name:
          # we need to discard the prepared statement on the server
          print "[deallocating server-side statement %s]" % (self.lru.name,)
          cursor.execute("""DEALLOCATE %s""" % (self.lru.name,))
        print "[forgetting statement %s]" % (self.lru.sql,)
        del(self.index[self.lru.sql])
        self.unlink(self.lru)
      statement = Statement(sql)
      self.push_mru(statement)
      self.index[sql] = statement
      print "[making new statement for %s]" % (statement.sql,)

    # bump the usage count
    if statement.name == None:
      statement.used += 1
      # time to PREPARE this statement, it's been used enough
      if statement.used > self.min_reuses_before_prepare:
        statement.name = "walrus%d" % (self.counter,)
        cursor.execute("""PREPARE %s AS %s""" % (statement.name, mangle(sql)))
        statement.execute_statement = "EXECUTE walrus%s(%s)" % (self.counter, ",".join(["%s"] * nargs))
        self.counter += 1
        print "[preparing server-side statement for %s as %s]" % (statement.sql, statement.name)

    return statement

class MagicCursor:
  def __init__(self, cache, cursor):
    self.cache = cache
    self.real_cursor = cursor

  def execute(self, sql, args = []):
    statement = self.cache.get_statement(self.real_cursor, sql, len(args))
    if statement.execute_statement:
      self.real_cursor.execute(statement.execute_statement, args)
    else:
      self.real_cursor.execute(sql, args)

  def fetchone(self):
    return self.real_cursor.fetchone()

if __name__ == "__main__":
  conn = psycopg2.connect("dbname=macdice")
  real_cursor = conn.cursor()

  cache = StatementCache(3, 5);
  magic_cursor = MagicCursor(cache, real_cursor)
  magic_cursor.execute("SELECT %s || %s", ("hello ", "world"))
  magic_cursor.execute("SELECT %s || %s", ("hello ", "world"))
  magic_cursor.execute("SELECT %s || %s", ("hello ", "world"))
  magic_cursor.execute("SELECT %s || %s", ("hello ", "world"))
  magic_cursor.execute("SELECT %s || %s", ("hello ", "world"))
  magic_cursor.execute("SELECT %s || %s", ("hello ", "world"))
  magic_cursor.execute("SELECT %s || %s", ("hello ", "world"))
  magic_cursor.execute("SELECT %s || %s", ("hello ", "world"))
  magic_cursor.execute("SELECT 1", ("hello ", "world"))
  magic_cursor.execute("SELECT 2", ("hello ", "world"))
  magic_cursor.execute("SELECT 3", ("hello ", "world"))
  magic_cursor.execute("SELECT 4", ("hello ", "world"))
  magic_cursor.execute("SELECT 5", ("hello ", "world"))
  magic_cursor.execute("SELECT 5", ("hello ", "world"))
  magic_cursor.execute("SELECT 5", ("hello ", "world"))
  magic_cursor.execute("SELECT 5", ("hello ", "world"))
  magic_cursor.execute("SELECT 5", ("hello ", "world"))
  magic_cursor.execute("SELECT 5", ("hello ", "world"))
  magic_cursor.execute("SELECT 5", ("hello ", "world"))
  magic_cursor.execute("SELECT 5", ("hello ", "world"))
  magic_cursor.execute("SELECT 6", ("hello ", "world"))
  magic_cursor.execute("SELECT 6", ("hello ", "world"))
  magic_cursor.execute("SELECT 6", ("hello ", "world"))
  magic_cursor.execute("SELECT 6", ("hello ", "world"))
  magic_cursor.execute("SELECT 6", ("hello ", "world"))
  magic_cursor.execute("SELECT 6", ("hello ", "world"))
  magic_cursor.execute("SELECT 6", ("hello ", "world"))
  magic_cursor.execute("SELECT 6", ("hello ", "world"))
  magic_cursor.execute("SELECT 4", ("hello ", "world"))
  magic_cursor.execute("SELECT 4", ("hello ", "world"))
  magic_cursor.execute("SELECT 4", ("hello ", "world"))
  magic_cursor.execute("SELECT 4", ("hello ", "world"))
  magic_cursor.execute("SELECT 4", ("hello ", "world"))
  magic_cursor.execute("SELECT 4", ("hello ", "world"))
  magic_cursor.execute("SELECT 4", ("hello ", "world"))
  magic_cursor.execute("SELECT 4", ("hello ", "world"))
  magic_cursor.execute("SELECT 4", ("hello ", "world"))
  magic_cursor.execute("SELECT 4", ("hello ", "world"))
  magic_cursor.execute("SELECT 4", ("hello ", "world"))
  magic_cursor.execute("SELECT 4", ("hello ", "world"))
  magic_cursor.execute("SELECT 3", ("hello ", "world"))
  magic_cursor.execute("SELECT 3", ("hello ", "world"))
  magic_cursor.execute("SELECT 3", ("hello ", "world"))
  magic_cursor.execute("SELECT 3", ("hello ", "world"))
  magic_cursor.execute("SELECT 3", ("hello ", "world"))
  magic_cursor.execute("SELECT 3", ("hello ", "world"))
  magic_cursor.execute("SELECT 3", ("hello ", "world"))
  magic_cursor.execute("SELECT 3", ("hello ", "world"))
  magic_cursor.execute("SELECT 3", ("hello ", "world"))
  magic_cursor.execute("SELECT 2", ("hello ", "world"))
  magic_cursor.execute("SELECT 2", ("hello ", "world"))
  magic_cursor.execute("SELECT 2", ("hello ", "world"))
  magic_cursor.execute("SELECT 1", ("hello ", "world"))