Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
Turns out all I wanted was a Cursor Factory that spat out sequence-li…
…ke cursors
  • Loading branch information
Alex Kritikos committed Feb 17, 2012
1 parent ac32644 commit dcadd64
Showing 1 changed file with 19 additions and 60 deletions.
79 changes: 19 additions & 60 deletions __init__.py
@@ -1,73 +1,32 @@
import datetime
import collections
import psycopg2
import psycopg2.extras
import itertools
import uuid


class Cursor(object):

def __init__(self, cursor, q, params, maxlen=1000):
self.cursor = cursor
self.q = q
self.params = params
self.cls = None
self.highest = -1
# We're using the deque as a cache.
self.deque = deque(maxlen=maxlen)
self.maxlen = maxlen
def gen():
for r in cursor:
self.highest += 1
self.deque.append(r)
yield r
self._gen = gen()

def __iter__(self):
for elt in self._gen:
yield elt
class ListCursor(collections.Sequence, psycopg2.extras.NamedTupleCursor):

def __getitem__(self, index):
try:
min_idx = index.start or 0
max_idx = index.stop
step = index.step
single = False
self.scroll(index.start or 0, mode='absolute')
return self.fetchmany(index.stop - index.start)[::index.step]
except AttributeError:
# We were just handed one value. That's ok. Make it a slice anyway.
min_idx = index
max_idx = index + 1
step = 1
single = True

if min_idx <= self.highest:
self.cursor.scroll(0, mode='absolute')
self.highest = -1
# TODO remove copypasta
def gen():
for r in self.cursor:
self.highest += 1
self.deque.append(r)
yield r
self._gen = gen()

# If there's not an off-by-one error here, I'll be amazed
adj = self.highest + 1
self.scroll(index, mode='absolute')
return self.fetchone()

try:
return next(itertools.islice(self._gen, index - adj, index - adj + 1))
except TypeError:
return list(itertools.islice(self._gen, index.start - adj, index.stop - adj, index.step))
def __len__(self):
self.scroll(0, mode='absolute')
length = 0
for r in self:
length += 1
return length

def __getattr__(self, name):
return getattr(self.cursor, name)
def __repr__(self):
return "ListCursor('%s')" % self.query

# Snark
def __len__(self):
raise NotImplementedError('Have you considered "COUNT(*)"?')
def __str__(self):
return self.__repr__()

def __reversed__(self):
raise NotImplementedError('Have you considered "ORDER BY foo DESC"?')

# Database connection. Call it and you get a cursor.
class D(object):
Expand All @@ -79,8 +38,8 @@ def __call__(self, q, *params):
if q.upper().startswith('SELECT'):
# Server-side cursors are cool. Unless we know better, use them for
# everything.
cursor = self.con.cursor(str(uuid.uuid1()), cursor_factory=psycopg2.extras.NamedTupleCursor)
cursor = self.con.cursor(str(uuid.uuid1()), cursor_factory=ListCursor)
else:
cursor = self.con.cursor()
cursor.execute(q, *params)
return Cursor(cursor, q, params)
cursor.execute(q, params)
return cursor

0 comments on commit dcadd64

Please sign in to comment.