Commits

Lynn Rees committed 72646f2

- hstore

  • Participants
  • Parent commits 78524e5

Comments (0)

Files changed (2)

File shove/stores/hstore.py

+# -*- coding: utf-8 -*-
+'''
+PostgreSQL hstore store.
+
+shove's URI for PostgreSQL hstore stores follows the form:
+
+hstore://<host>:<port>/<db>/collection/
+'''
+
+try:
+    import psycopg2
+    from psycopg2 import extras
+except ImportError:
+    raise ImportError('requires `psycopg2` library')
+
+from stuf.six import items, values
+
+from shove.store import BaseStore
+from shove._compat import urlsplit
+
+__all__ = ['HStore']
+
+
+class HStore(BaseStore):
+
+    '''PostgreSQL hstore object storage frontend.'''
+
+    init = 'mongodb://'
+
+    def __init__(self, engine, **kw):
+        super(HStore, self).__init__(engine, **kw)
+        spliturl = urlsplit(engine)
+        try:
+            self._conn = conn = psycopg2.connect(
+                host=spliturl.hostname,
+                port=spliturl.port,
+                database=spliturl.path,
+                user=spliturl.username or '',
+                password=spliturl.password or '',
+            )
+            self._store = db = conn.cursor(
+                cursor_factory=extras.RealDictCursor
+            )
+        except psycopg2.OperationalError:
+            raise TypeError('configuration error')
+        try:
+            db.execute('CREATE EXTENSION hstore')
+            conn.commit()
+        except psycopg2.ProgrammingError:
+            conn.rollback()
+        extras.register_hstore(conn)
+        try:
+            db.execute(
+                'CREATE TABLE shove (id serial PRIMARY KEY, data hstore)'
+            )
+            conn.commit()
+            db.execute(
+                'INSERT INTO shove (data) VALUES (%s)', ['"key"=>"value"'],
+            )
+            conn.commit()
+            db.execute(
+                'UPDATE shove SET data = delete(data, %s)', ['key'],
+            )
+        except psycopg2.ProgrammingError:
+            conn.rollback()
+
+    def __getitem__(self, key):
+        try:
+            self._store.execute('SELECT data -> %s as t FROM shove', [key])
+            data = self._store.fetchone()['t']
+            if data is not None:
+                return self.loads(data)
+            raise KeyError
+        except psycopg2.ProgrammingError:
+            self._conn.rollback()
+            raise KeyError(key)
+
+    def __setitem__(self, key, value):
+        try:
+            self._store.execute(
+                'UPDATE shove SET data = data || (%s)',
+                [{key: self.dumps(value)}],
+            )
+            self._conn.commit()
+        except psycopg2.ProgrammingError:
+            self._conn.rollback()
+
+    def __delitem__(self, key):
+        try:
+            self._store.execute(
+                'UPDATE shove SET data = delete(data, %s)', [key],
+            )
+            self._conn.commit()
+        except psycopg2.ProgrammingError:
+            self._conn.rollback()
+
+    def __contains__(self, key):
+        try:
+            self._store.execute('SELECT exist(data, %s) FROM shove')
+            return self._store.fetchone()['exist']
+        except psycopg2.ProgrammingError:
+            self._conn.rollback()
+
+    def __len__(self):
+        try:
+            self._store.execute(
+            'SELECT count(k) FROM (SELECT (each(data)).key FROM shove) as k'
+            )
+            return self._store.fetchone()['count']
+        except psycopg2.ProgrammingError:
+            self._conn.rollback()
+
+    def __iter__(self):
+        try:
+            self._store.execute('SELECT akeys(data) as k FROM shove')
+            return iter(self._store.fetchall()['k'])
+        except psycopg2.ProgrammingError:
+            self._conn.rollback()
+
+    def close(self):
+        self._conn.close()
+
+    def clear(self):
+        pass
+
+    def items(self):
+        loads = self.loads
+        try:
+            self._store.execute('SELECT data as i FROM shove')
+            for k, v in items(self._store.fetchall()['i']):
+                yield k, loads(v)
+        except psycopg2.ProgrammingError:
+            self._conn.rollback()
+
+    def values(self):
+        loads = self.loads
+        try:
+            self._store.execute('SELECT svalues(data) as v FROM shove')
+            for v in values(self._store.fetchall()['v']):
+                yield loads(v)
+        except psycopg2.ProgrammingError:
+            self._conn.rollback()

File tests/test_hstore.py

+# -*- coding: utf-8 -*-
+'''shove hstore tests.'''
+
+import os
+import time
+import signal
+import shutil
+from tempfile import mkdtemp
+
+from stuf.six import unittest
+from subprocess32 import Popen  # @UnresolvedImport
+
+from tests.test_store import Store
+
+
+def setUpModule():
+    temp = mkdtemp()
+    os.chdir(temp)
+    os.environ['PGDATABASE'] = 'shove'
+    os.environ['PGDATA'] = temp
+    process = Popen(
+        ['initdb', '-A', 'trust', '-E', 'utf-8', '-D', temp],
+        stdout=open('/dev/null', 'w'),
+        stderr=open('/dev/null', 'w'),
+        shell=True,
+    )
+    process.wait()
+    unittest.TestCase.process = Popen(
+        ['postgres', '-D', temp, '-h', 'localhost', '-p', '5432', '-F'],
+        stdout=open('/dev/null', 'w'),
+        stderr=open('/dev/null', 'w'),
+        shell=True,
+    )
+    time.sleep(15.0)
+    process = Popen(
+        ['createdb', '-h', 'localhost', '-p', '5432', 'shove'],
+        stdout=open('/dev/null', 'w'),
+        stderr=open('/dev/null', 'w'),
+        shell=True,
+    )
+    process.wait()
+
+
+def tearDownModule():
+    unittest.TestCase.process.send_signal(signal.SIGQUIT)
+    unittest.TestCase.process.wait()
+    temp = os.environ['PGDATA']
+    shutil.rmtree(temp)
+    del os.environ['PGDATA']
+
+
+class TestHStoreStore(Store, unittest.TestCase):
+
+    initstring = 'hstore://localhost:5432/shove'