Add a RAW SQL backend (bug 927794)

Makefile

@@ -4,7 +4,7 @@ build:
 	bin/python setup.py develop
 
 test:
-	bin/flake8 addonreg && bin/nosetests --with-coverage --cover-erase --cover-package addonreg addonreg
+	bin/flake8 addonreg && SQLURI=mysql://addonreg:addonreg@localhost/addonreg bin/nosetests --with-coverage --cover-erase --cover-package addonreg addonreg
 
 clean:
 	rm -rf bin lib local include man

addonreg/__init__.py

@@ -1,13 +1,16 @@
 VERSION = '0.1'
 
 import os
+import logging
 
 from celery.app import app_or_default
 from pyramid.config import Configurator
 from pyramid.events import NewRequest
 
 from konfig import Config
 
+logger = logging.getLogger('addonreg')
+
 
 def get_config(filename=None):
     filename = filename or os.environ.get('CONFIG', 'development.ini')

addonreg/backends/rawsql.py

@@ -1,11 +1,101 @@
+import traceback
+
+from sqlalchemy import String, Index, Boolean
+from sqlalchemy import Table, MetaData, create_engine, Column
+from sqlalchemy.exc import OperationalError, TimeoutError
+from sqlalchemy.pool import NullPool
+from sqlalchemy.sql import text as sqltext
+
+from addonreg import logger
+
+
+_GET = sqltext("""\
+SELECT addonid, sha256, registered
+FROM hashes
+WHERE addonid = :addonid
+AND sha256 = :sha256
+""")
+
+_INSERT = sqltext("""\
+INSERT INTO hashes
+(addonid, sha256, registered)
+VALUES (:addonid, :sha256, 1)
+""")
+
+metadata = MetaData()
+
+hash_table = Table(
+    'hashes', metadata,
+    Column('addonid', String(255), nullable=False, primary_key=True),
+    Column('sha256', String(64), nullable=False, primary_key=True),
+    Column('registered', Boolean(), default=True),
+    Index('hash_idx', 'sha256', 'addonid', unique=True),
+    mysql_engine='InnoDB', mysql_charset='utf8')
+
+
 class RawSQLBackend(object):
-    """A backend using RAW SQL queries to go faster."""
+    """A backend using RAW SQL queries."""
+
+    def __init__(self, config=None, sqluri=False, create_tables=False,
+                 pool_size=100, pool_recycle=60, pool_timeout=30,
+                 max_overflow=10, pool_reset_on_return='rollback', **kw):
+        self.config = config or {}
+        self.sqluri = sqluri or config['SQLURI']
+        if pool_reset_on_return.lower() in ('', 'none'):
+            pool_reset_on_return = None
+
+        if self.sqluri.startswith(('mysql', 'pymysql')):
+            self._engine = create_engine(
+                sqluri,
+                pool_size=pool_size,
+                pool_recycle=pool_recycle,
+                pool_timeout=pool_timeout,
+                pool_reset_on_return=pool_reset_on_return,
+                max_overflow=max_overflow,
+                logging_name='addonreg')
+
+        else:
+            self._engine = create_engine(sqluri, poolclass=NullPool)
+
+        self._engine.echo = kw.get('echo', False)
+        self.hashes = hash_table
+
+        self.hashes.metadata.bind = self._engine
+        if create_tables:
+            self.hashes.create(checkfirst=True)
+
+    def _get_engine(self, service=None):
+        return self._engine
+
+    def _safe_execute(self, *args, **kwds):
+        """Execute an sqlalchemy query & log + raise an exception on failure"""
+        if hasattr(args[0], 'bind'):
+            engine = args[0].bind
+        else:
+            engine = None
+
+        if engine is None:
+            engine = kwds.get('engine')
+            if engine is None:
+                engine = self._get_engine(kwds.get('service'))
+            else:
+                del kwds['engine']
 
-    def __init__(self, config):
-        self.config = config
+        try:
+            return engine.execute(*args, **kwds)
+        except (OperationalError, TimeoutError):
+            err = traceback.format_exc()
+            logger.error(err)
+            raise
 
     def hash_exists(self, addon_id, hash_):
-        pass
+        res = self._safe_execute(_GET, addonid=addon_id, sha256=hash_)
+        try:
+            item = res.fetchone()
+            return item is not None
+        finally:
+            res.close()
 
     def register_hash(self, addon_id, hash_):
-        pass
+        res = self._safe_execute(_INSERT, addonid=addon_id, sha256=hash_)
+        res.close()

addonreg/tests/test_sqlbackend.py

@@ -0,0 +1,40 @@
+import os
+from unittest2 import TestCase
+
+from addonreg.backends.rawsql import RawSQLBackend
+
+
+class TestSQLBackend(TestCase):
+
+    # By default, the tests are using SQLite in order to be faster.
+    # You can change that (if you want to run the tests against a real database
+    # for instance) by changing the SQLURI environment variable.
+    _SQLURI = os.environ.get('SQLURI', 'sqlite:////tmp/wimms')
+
+    def setUp(self):
+        super(TestSQLBackend, self).setUp()
+        self.backend = RawSQLBackend(sqluri=self._SQLURI, create_tables=True)
+        self.guid = u'{9c51bd27-6ed8-4000-a2bf-36cb95c0c947}'
+        self.sha256 = (u'31f7a65e315586ac198bd798b6629ce4903d0899476d5741a9f32'
+                       'e2e521b6a66')
+        self._sqlite = self.backend._engine.driver == 'pysqlite'
+
+    def tearDown(self):
+        if self._sqlite:
+            filename = self.backend.sqluri.split('sqlite://')[-1]
+            if os.path.exists(filename):
+                os.remove(filename)
+        else:
+            self.backend._safe_execute('drop table hashes;')
+
+    def test_read(self):
+        # Let's create a hash to test if we're able to read it back.
+        self.backend._safe_execute(
+            """INSERT INTO hashes (addonid, sha256, registered)
+               VALUES ("%s", "%s", 1)""" % (self.guid, self.sha256))
+
+        self.assertTrue(self.backend.hash_exists(self.guid, self.sha256))
+
+    def test_write(self):
+        self.backend.register_hash(self.guid, self.sha256)
+        self.assertTrue(self.backend.hash_exists(self.guid, self.sha256))

requirements/development.txt

@@ -14,3 +14,5 @@ celery==3.0.23
 redis==2.8.0
 colander==0.9.9
 konfig==0.8
+SQLAlchemy==0.8.2
+MySQL-python==1.2.4