Add a helper context for optional connection args

We refactor a bunch of methods to support re-using an open connection
when available.  There's some code in the connection manager to support
nesting connections to avoid deadlocks when we call a method that can
open a connection while we already have a connection checked out - but
it might be better in the long run if we just got better at passing open
connections around whenever possible.

Add a helper method on the base db broker class to make it easier to
write methods that can optionally take an existing connection.

Change-Id: Ifd54d76ab1a5a9d82848f3cae89c3e53134aa129

swift/common/db.py

@@ -427,6 +427,14 @@ class DatabaseBroker(object):
             if self.conn:
                 self.conn.timeout = old_timeout
 
+    @contextmanager
+    def maybe_get(self, conn):
+        if conn:
+            yield conn
+        else:
+            with self.get() as conn:
+                yield conn
+
     @contextmanager
     def get(self):
         """Use with the "with" statement; returns a database connection."""

swift/container/backend.py

@@ -1653,11 +1653,8 @@ class ContainerBroker(DatabaseBroker):
             return [row for row in data]
 
         try:
-            if connection:
-                return do_query(connection)
-            else:
-                with self.get() as conn:
-                    return do_query(conn)
+            with self.maybe_get(connection) as conn:
+                return do_query(conn)
         except sqlite3.OperationalError as err:
             if ('no such table: %s' % SHARD_RANGE_TABLE) not in str(err):
                 raise

test/unit/common/test_db.py

@@ -604,6 +604,25 @@ class TestExampleBroker(unittest.TestCase):
         broker.get_info()
         self.assertEqual(1, broker.get_info()[count_key])
 
+    @with_tempdir
+    def test_maybe_get(self, tempdir):
+        broker = self.broker_class(os.path.join(tempdir, 'test.db'),
+                                   account='a', container='c')
+        broker.initialize(next(self.ts),
+                          storage_policy_index=int(self.policy))
+        qry = 'select account from %s_stat' % broker.db_type
+        with broker.maybe_get(None) as conn:
+            rows = [dict(x) for x in conn.execute(qry)]
+        self.assertEqual([{'account': 'a'}], rows)
+        self.assertEqual(conn, broker.conn)
+        with broker.get() as other_conn:
+            self.assertEqual(broker.conn, None)
+            with broker.maybe_get(other_conn) as identity_conn:
+                self.assertEqual(other_conn, identity_conn)
+                self.assertEqual(broker.conn, None)
+            self.assertEqual(broker.conn, None)
+        self.assertEqual(broker.conn, conn)
+
 
 class TestDatabaseBroker(unittest.TestCase):