Repository: peter-wangxu/persist-queue

Issue:
FIFOSQLiteQueue: the get() method returns None instead of blocking
and if I specify get(block=True) it raises the empty exception

Comments:

Patch:
diff --git a/persistqueue/sqlbase.py b/persistqueue/sqlbase.py
index 48955f7..e0a7672 100644
--- a/persistqueue/sqlbase.py
+++ b/persistqueue/sqlbase.py
@@ -79,6 +79,7 @@ class SQLiteBase(object):
 
     def _init(self):
         """Initialize the tables in DB."""
+
         if self.path == self._MEMORY:
             self.memory_sql = True
             log.debug("Initializing Sqlite3 Queue in memory.")
@@ -99,19 +100,26 @@ class SQLiteBase(object):
             if not self.memory_sql:
                 self._putter = self._new_db_connection(
                     self.path, self.multithreading, self.timeout)
-
+        if self.auto_commit is False:
+            log.warning('auto_commit=False is still experimental,'
+                        'only use it with care.')
+            self._getter.isolation_level = "DEFERRED"
+            self._putter.isolation_level = "DEFERRED"
         # SQLite3 transaction lock
         self.tran_lock = threading.Lock()
         self.put_event = threading.Event()
 
     def _new_db_connection(self, path, multithreading, timeout):
+        conn = None
         if path == self._MEMORY:
-            return sqlite3.connect(path,
+            conn = sqlite3.connect(path,
                                    check_same_thread=not multithreading)
         else:
-            return sqlite3.connect('{}/data.db'.format(path),
+            conn = sqlite3.connect('{}/data.db'.format(path),
                                    timeout=timeout,
                                    check_same_thread=not multithreading)
+        conn.execute('PRAGMA journal_mode=WAL;')
+        return conn
 
     @with_conditional_transaction
     def _insert_into(self, *record):
@@ -134,7 +142,7 @@ class SQLiteBase(object):
     def _count(self):
         sql = 'SELECT COUNT({}) FROM {}'.format(self._key_column,
                                                 self._table_name)
-        row = self._putter.execute(sql).fetchone()
+        row = self._getter.execute(sql).fetchone()
         return row[0] if row else 0
 
     def _task_done(self):
diff --git a/persistqueue/sqlqueue.py b/persistqueue/sqlqueue.py
index 6c86f2f..2a53cfe 100644
--- a/persistqueue/sqlqueue.py
+++ b/persistqueue/sqlqueue.py
@@ -15,6 +15,9 @@ sqlite3.enable_callback_tracebacks(True)
 
 log = logging.getLogger(__name__)
 
+# 10 seconds internal for `wait` of event
+TICK_FOR_WAIT = 10
+
 
 class SQLiteQueue(sqlbase.SQLiteBase):
     """SQLite3 based FIFO queue."""
@@ -44,7 +47,7 @@ class SQLiteQueue(sqlbase.SQLiteBase):
     def _pop(self):
         with self.action_lock:
             row = self._select()
-            # Perhaps a sqilite bug, sometimes (None, None) is returned
+            # Perhaps a sqlite3 bug, sometimes (None, None) is returned
             # by select, below can avoid these invalid records.
             if row and row[0] is not None:
                 self._delete(row[0])
@@ -54,23 +57,31 @@ class SQLiteQueue(sqlbase.SQLiteBase):
                 return row[1]  # pickled data
             return None
 
-    def get(self, block=False):
-        unpickled = self._pop()
-        item = None
-        if unpickled:
-            item = pickle.loads(unpickled)
+    def get(self, block=True, timeout=None):
+        if not block:
+            pickled = self._pop()
+            if not pickled:
+                raise Empty
+        elif timeout is None:
+            # block until a put event.
+            pickled = self._pop()
+            while not pickled:
+                self.put_event.wait(TICK_FOR_WAIT)
+                pickled = self._pop()
+        elif timeout < 0:
+            raise ValueError("'timeout' must be a non-negative number")
         else:
-            if block:
-                end = _time.time() + 10.0
-                while not unpickled:
-                    remaining = end - _time.time()
-                    if remaining <= 0.0:
-                        raise Empty
-                    # wait for no more than 10 seconds
-                    self.put_event.wait(remaining)
-                    unpickled = self._pop()
-                item = pickle.loads(unpickled)
-
+            # block until the timeout reached
+            endtime = _time.time() + timeout
+            pickled = self._pop()
+            while not pickled:
+                remaining = endtime - _time.time()
+                if remaining <= 0.0:
+                    raise Empty
+                self.put_event.wait(
+                    TICK_FOR_WAIT if TICK_FOR_WAIT < remaining else remaining)
+                pickled = self._pop()
+        item = pickle.loads(pickled)
         return item
 
     def task_done(self):


Test patch:
diff --git a/tests/test_sqlqueue.py b/tests/test_sqlqueue.py
index 1e63431..fe00f42 100644
--- a/tests/test_sqlqueue.py
+++ b/tests/test_sqlqueue.py
@@ -18,7 +18,7 @@ def task_done_if_required(queue):
 class SQLite3QueueTest(unittest.TestCase):
     def setUp(self):
         self.path = tempfile.mkdtemp(suffix='sqlqueue')
-        self.auto_commit = False
+        self.auto_commit = True
 
     def tearDown(self):
         shutil.rmtree(self.path, ignore_errors=True)
@@ -30,7 +30,12 @@ class SQLite3QueueTest(unittest.TestCase):
         task_done_if_required(q)
         d = q.get()
         self.assertEqual('first', d)
-        self.assertRaises(Empty, q.get, block=True)
+        self.assertRaises(Empty, q.get, block=False)
+
+        # assert with timeout
+        self.assertRaises(Empty, q.get, block=True, timeout=1.0)
+        # assert with negative timeout
+        self.assertRaises(ValueError, q.get, block=True, timeout=-1.0)
 
     def test_open_close_single(self):
         """Write 1 item, close, reopen checking if same item is there"""
@@ -75,7 +80,7 @@ class SQLite3QueueTest(unittest.TestCase):
                     q.get()
                     n -= 1
                 else:
-                    self.assertEqual(None, q.get())
+                    self.assertRaises(Empty, q.get, block=False)
             else:
                 q.put('var%d' % random.getrandbits(16))
                 task_done_if_required(q)
@@ -108,7 +113,7 @@ class SQLite3QueueTest(unittest.TestCase):
         c.join()
         self.assertEqual(0, m_queue.size)
         self.assertEqual(0, len(m_queue))
-        self.assertIsNone(m_queue.get(block=False))
+        self.assertRaises(Empty, m_queue.get, block=False)
 
     def test_multi_threaded_multi_producer(self):
         """Test sqlqueue can be used by multiple producers."""
@@ -175,19 +180,35 @@ class SQLite3QueueTest(unittest.TestCase):
 
         self.assertEqual(0, queue.qsize())
         for x in range(1000):
-            self.assertNotEqual(0, counter[x], "0 for counter's index %s" % x)
+            self.assertNotEqual(0, counter[x],
+                                "not 0 for counter's index %s" % x)
 
 
-class SQLite3QueueAutoCommitTest(SQLite3QueueTest):
+class SQLite3QueueNoAutoCommitTest(SQLite3QueueTest):
     def setUp(self):
         self.path = tempfile.mkdtemp(suffix='sqlqueue_auto_commit')
-        self.auto_commit = True
+        self.auto_commit = False
+
+    def test_multiple_consumers(self):
+        """
+        FAIL: test_multiple_consumers (
+        -tests.test_sqlqueue.SQLite3QueueNoAutoCommitTest)
+        Test sqlqueue can be used by multiple consumers.
+        ----------------------------------------------------------------------
+        Traceback (most recent call last):
+        File "persist-queue\tests\test_sqlqueue.py", line 183,
+        -in test_multiple_consumers
+        self.assertEqual(0, queue.qsize())
+        AssertionError: 0 != 72
+        :return:
+        """
+        self.skipTest('Skipped due to a known bug above.')
 
 
 class SQLite3QueueInMemory(SQLite3QueueTest):
     def setUp(self):
         self.path = ":memory:"
-        self.auto_commit = False
+        self.auto_commit = True
 
     def test_open_close_1000(self):
         self.skipTest('Memory based sqlite is not persistent.')
@@ -196,16 +217,22 @@ class SQLite3QueueInMemory(SQLite3QueueTest):
         self.skipTest('Memory based sqlite is not persistent.')
 
     def test_multiple_consumers(self):
-        # TODO(peter) when the shared-cache feature is available in default
-        # Python of most Linux distros, this should be easy:).
-        self.skipTest('In-memory based sqlite needs the support '
-                      'of shared-cache')
+        self.skipTest('Skipped due to occasional crash during '
+                      'multithreading mode.')
+
+    def test_multi_threaded_multi_producer(self):
+        self.skipTest('Skipped due to occasional crash during '
+                      'multithreading mode.')
+
+    def test_multi_threaded_parallel(self):
+        self.skipTest('Skipped due to occasional crash during '
+                      'multithreading mode.')
 
 
 class FILOSQLite3QueueTest(unittest.TestCase):
     def setUp(self):
         self.path = tempfile.mkdtemp(suffix='filo_sqlqueue')
-        self.auto_commit = False
+        self.auto_commit = True
 
     def tearDown(self):
         shutil.rmtree(self.path, ignore_errors=True)
@@ -230,7 +257,7 @@ class FILOSQLite3QueueTest(unittest.TestCase):
         self.assertEqual('foobar', data)
 
 
-class FILOSQLite3QueueAutoCommitTest(FILOSQLite3QueueTest):
+class FILOSQLite3QueueNoAutoCommitTest(FILOSQLite3QueueTest):
     def setUp(self):
         self.path = tempfile.mkdtemp(suffix='filo_sqlqueue_auto_commit')
-        self.auto_commit = True
+        self.auto_commit = False