Ticket #9536: 9536-db-context-manager-r10102.3.patch

trac/db/api.py

@@ -20 +20 @@
 
 from trac.config import BoolOption, IntOption, Option
 from trac.core import *
-from trac.db.pool import ConnectionPool
 from trac.util.concurrency import ThreadLocal
 from trac.util.text import unicode_passwd
 from trac.util.translation import _
 
+from .pool import ConnectionPool
+from .util import ConnectionWrapper
+
 _transaction_local = ThreadLocal(db=None)
 
 def with_transaction(env, db=None):
     """Function decorator to emulate a context manager for database
     transactions.
-    
+
     >>> def api_method(p1, p2):
     >>>     result[0] = value1
     >>>     @with_transaction(env)
@@ -38 +40 @@
     >>>         # implementation
     >>>         result[0] = value2
     >>>     return result[0]
+
+    :deprecated: use instead the new context manager:
+
+    >>> def api_method(p1, p2):
+    >>>     result = value1
+    >>>     with env.db_transaction as db:
+    >>>         # implementation
+    >>>         result = value2
+    >>>     return result
     
     In this example, the `implementation()` function is called automatically
     right after its definition, with a database connection as an argument.
@@ -51 +62 @@
     or rollback, for mutating database accesses. Its automatic handling of
     commit, rollback and nesting makes it much more robust.
     
-    This decorator will be replaced by a context manager once python 2.4
-    support is dropped.
-
     The optional `db` argument is intended for legacy code and should not
     be used in new code.
+
+    This decorator is in turn deprecated in favor of context managers 
+    now that python 2.4 support has been dropped.
     """
     def transaction_wrapper(fn):
         ldb = _transaction_local.db
@@ -72 +83 @@
         elif ldb:
             fn(ldb)
         else:
-            ldb = _transaction_local.db = env.get_db_cnx()
+            ldb = _transaction_local.db = get_write_db(env)
             try:
                 fn(ldb)
                 ldb.commit()
@@ -87 +98 @@
 
 def get_read_db(env):
     """Get a database connection for reading only."""
+    db = get_write_db(env)
+    if not db.readonly:
+        db = ConnectionWrapper(db, readonly=True)
+    return db
+
+def get_write_db(env):
+    """Get a database connection."""
     return _transaction_local.db or DatabaseManager(env).get_connection()
 
+class TransactionContextManager(object):
+    """Transactioned Database Context Manager
+
+    The outermost such context manager will either commit or rollback,
+    depending on the context being exited normally or after an exception.
+    """
+
+    db = None
+
+    def __init__(self, env):
+        self.env = env
+
+    def __enter__(self):
+        db = _transaction_local.db
+        if not db:
+            _transaction_local.db = self.db = db = get_write_db(self.env)
+        return db
+
+    def __exit__(self, et, ev, tb): 
+        if self.db is not None: 
+            _transaction_local.db = None
+            if et is None: 
+                self.db.commit()
+            else: 
+                self.db.rollback()
+
+
+class QueryContextManager(object):
+    """Database Context Manager for retrieving a readonly ConnectionWrapper"""
+
+    def __init__(self, env):
+        self.env = env
+
+    def __enter__(self):
+        db = _transaction_local.db
+        if not db:
+            db = get_read_db(self.env)
+        return db
+
+    def __exit__(self, et, ev, tb): 
+        pass
+
 
 class IDatabaseConnector(Interface):
     """Extension point interface for components that support the connection to

trac/db/util.py

@@ -91 +91 @@
     
     :since 0.12: This wrapper no longer makes cursors produced by the
     connection iterable using `IterableCursor`.
+
+    :since 0.13: added a 'readonly' flag preventing the forwarding
+    of `commit` and `rollback`.
     """
-    __slots__ = ('cnx', 'log')
+    __slots__ = ('cnx', 'log', 'readonly')
 
-    def __init__(self, cnx, log=None):
+    def __init__(self, cnx, log=None, readonly=False):
         self.cnx = cnx
         self.log = log
+        self.readonly = readonly
 
     def __getattr__(self, name):
+        if self.readonly and name in ('commit', 'rollback'):
+            raise AttributeError
         return getattr(self.cnx, name)

trac/env.py

@@ -25 +25 @@
 from trac.config import *
 from trac.core import Component, ComponentManager, implements, Interface, \
                       ExtensionPoint, TracError
-from trac.db.api import DatabaseManager, get_read_db, with_transaction
+from trac.db.api import (DatabaseManager, QueryContextManager, 
+                         TransactionContextManager, get_read_db, get_write_db,
+                         with_transaction)
 from trac.util import copytree, create_file, get_pkginfo, makedirs
 from trac.util.compat import any
 from trac.util.concurrency import threading
@@ -323 +325 @@
         Use `with_transaction` for obtaining a writable database connection
         and `get_read_db` for anything else.
         """
-        return get_read_db(self)
+        return get_write_db(self)
 
     def with_transaction(self, db=None):
-        """Decorator for transaction functions.
-
-        See `trac.db.api.with_transaction` for detailed documentation."""
+        """Decorator for transaction functions (deprecated)"""
         return with_transaction(self, db)
 
     def get_read_db(self):
@@ -337 +337 @@
         See `trac.db.api.get_read_db` for detailed documentation."""
         return get_read_db(self)
 
+    @property
+    def db_query(self):
+        return QueryContextManager(self)
+
+    @property
+    def db_transaction(self):
+        return TransactionContextManager(self)
+
     def shutdown(self, tid=None):
         """Close the environment."""
         RepositoryManager(self).shutdown(tid)
@@ -507 +515 @@
         @return: whether the upgrade was performed
         """
         upgraders = []
-        db = self.get_read_db()
+        db = self.get_write_db()
         for participant in self.setup_participants:
             if participant.environment_needs_upgrade(db):
                 upgraders.append(participant)

trac/test.py

@@ -32 +32 @@
 from trac.env import Environment
 from trac.db.api import _parse_db_str, DatabaseManager
 from trac.db.sqlite_backend import SQLiteConnection
+from trac.db.util import ConnectionWrapper
 import trac.db.postgres_backend
 import trac.db.mysql_backend
 from trac.ticket.default_workflow import load_workflow_config_snippet
@@ -289 +290 @@
     def get_read_db(self):
         return self.get_db_cnx()
     
+    def get_write_db(self):
+        return self.get_db_cnx()
+    
     def get_db_cnx(self, destroying=False):
         if self.db:
             return self.db # in-memory SQLite
@@ -306 +310 @@
                 self.reset_db() # make sure we get rid of previous garbage
         return DatabaseManager(dbenv).get_connection()
 
+    # see below, functions get_read_db_overrider, get_write_db_overrider
+
     def reset_db(self, default_data=None):
         """Remove all data from Trac tables, keeping the tables themselves.
         :param default_data: after clean-up, initialize with default data
@@ -387 +393 @@
         return self.known_users
 
 
+# As we'll create lots of Environments during testing, we want to reuse 
+# the same DatabaseManager instance (see EnvironmentStub.get_db_cnx above),
+# so we need to override the connection dispatching in trac.db.api.
+
+def get_read_db_overrider(env):
+    db = env.get_db_cnx()
+    if not db.readonly:
+        db = ConnectionWrapper(db, readonly=True)
+    return db
+
+def get_write_db_overrider(env):
+    return env.get_db_cnx()
+
+from trac.db import api as dbapi
+dbapi.get_read_db = get_read_db_overrider
+dbapi.get_write_db = get_write_db_overrider
+
+
 def locate(fn):
     """Locates a binary on the path.