Close sync connection on allow_sync() exit; update README.md

README.md

@@ -7,59 +7,97 @@ ORM powered by **[asyncio](https://docs.python.org/3/library/asyncio.html)**.
 [![Build Status](https://travis-ci.org/05bit/peewee-async.svg)](https://travis-ci.org/05bit/peewee-async) [![PyPi Version](https://img.shields.io/pypi/v/peewee-async.svg)](https://pypi.python.org/pypi/peewee-async)
  [![Documentation Status](https://readthedocs.org/projects/peewee-async/badge/?version=latest)](http://peewee-async.readthedocs.org/en/latest/?badge=latest)
 
-
-Documentation
--------------
-
+Overview
+--------
+
+- Works on Python 3.4+
+- Has support PostgreSQL via `aiopg`
+- Has support MySQL via `aiomysql`
+- Single point high level async API
+- Drop-in replacement for sync code, sync will remain sync
+- Basic operations are supported
+- Transactions support is present, yet not heavily tested
+
+The complete documentation:  
 http://peewee-async.readthedocs.org
 
 Install
 -------
 
-Works on Python 3.3+ and PostgreSQL database.
+Install with `pip` for PostgreSQL:
+
+```
+pip install peewee-async aiopg
+```
 
-Install with `pip`:
+or for MySQL:
 
 ```
-pip install peewee-async
+pip install peewee-async aiomysql
 ```
 
 Quickstart
 ----------
 
-Create test PostgreSQL database, i.e. 'test' for running this snippet:
+Create 'test' PostgreSQL database for running this snippet:
+
+    createdb -E utf-8 test
+
+The code below is using new Python 3.5 `async` / `await` syntax, but older `yield from` will also work!
 
 ```python
 import asyncio
 import peewee
 import peewee_async
 
+# Nothing special, just define model and database:
+
 database = peewee_async.PostgresqlDatabase('test')
-loop = asyncio.get_event_loop()
 
 class TestModel(peewee.Model):
     text = peewee.CharField()
 
     class Meta:
         database = database
 
-# Create table synchronously!
+# Look, sync code is working!
+
 TestModel.create_table(True)
-# This is optional: close sync connection
+TestModel.create(text="Yo, I can do it sync!")
 database.close()
 
-@asyncio.coroutine
-def my_handler():
-    TestModel.create(text="Yo, I can do it sync!")
-    yield from peewee_async.create_object(TestModel, text="Not bad. Watch this, I'm async!")
-    all_objects = yield from peewee_async.execute(TestModel.select())
+# Create async database manager:
+
+objects = peewee_async.Manager(database)
+
+# No need for sync anymore!
+
+database.allow_sync = False
+
+async def handler():
+    await objects.create(TestModel, text="Not bad. Watch this, I'm async!")
+    all_objects = await objects.execute(TestModel.select())
     for obj in all_objects:
         print(obj.text)
 
-loop.run_until_complete(database.connect_async(loop=loop))
-loop.run_until_complete(my_handler())
+loop = asyncio.get_event_loop()
+loop.run_until_complete(handler())
+loop.close()
+
+# Clean up, can do it sync again:
+with objects.allow_sync():
+    TestModel.drop_table(True)
+
+# Expected output:"
+# Yo, I can do it sync!
+# Not bad. Watch this, I'm async!
 ```
 
+Documentation
+-------------
+
+http://peewee-async.readthedocs.org
+
 Discuss
 -------
 

peewee_async.py

@@ -292,12 +292,17 @@ def savepoint(self, sid=None):
         return savepoint(self.database, sid=sid)
 
     @contextlib.contextmanager
-    def allow_sync(self, allow=True):
-        """Allow sync queries within context.
+    def allow_sync(self):
+        """Allow sync queries within context. Close sync
+        connection on exit if connected.
         """
         old_allow = self.database.allow_sync
-        self.database.allow_sync = allow
+        self.database.allow_sync = True
         yield
+        try:
+            self.database.close()
+        except self.database.Error:
+            pass # already closed
         self.database.allow_sync = old_allow
 
     def _swap_database(self, query):
@@ -970,6 +975,10 @@ class AsyncPostgresqlMixin(AsyncDatabase):
     """Mixin for `peewee.PostgresqlDatabase` providing extra methods
     for managing async connection.
     """
+    if aiopg:
+        import psycopg2
+        Error = psycopg2.Error
+
     def init_async(self, conn_cls=AsyncPostgresqlConnection,
                    enable_json=False,
                    enable_hstore=False):
@@ -1131,6 +1140,10 @@ class MySQLDatabase(AsyncDatabase, peewee.MySQLDatabase):
     See also:
     http://peewee.readthedocs.org/en/latest/peewee/api.html#MySQLDatabase
     """
+    if aiomysql:
+        import pymysql
+        Error = pymysql.Error
+
     def init(self, database, **kwargs):
         if not aiomysql:
             raise Exception("Error, aiomysql is not installed!")

tests/__init__.py

@@ -29,43 +29,61 @@
     'postgres': {
         'database': 'test',
         'host': '127.0.0.1',
-        'port': 5432,
+        # 'port': 5432,
         'user': 'postgres',
     },
     'postgres-ext': {
         'database': 'test',
         'host': '127.0.0.1',
-        'port': 5432,
+        # 'port': 5432,
         'user': 'postgres',
     },
     'postgres-pool': {
         'database': 'test',
         'host': '127.0.0.1',
-        'port': 5432,
+        # 'port': 5432,
         'user': 'postgres',
         'max_connections': 4,
     },
     'postgres-pool-ext': {
         'database': 'test',
         'host': '127.0.0.1',
-        'port': 5432,
+        # 'port': 5432,
         'user': 'postgres',
         'max_connections': 4,
     },
     'mysql': {
         'database': 'test',
         'host': '127.0.0.1',
-        'port': 3306,
+        # 'port': 3306,
         'user': 'root',
     },
     'mysql-pool': {
         'database': 'test',
         'host': '127.0.0.1',
-        'port': 3306,
+        # 'port': 3306,
         'user': 'root',
     }
 }
 
+try:
+    import aiopg
+except ImportError:
+    print("aiopg is not installed, ignoring PostgreSQL tests")
+    for k in list(deafults.keys()):
+        if k.startswith('postgres'):
+            deafults.pop(k)
+
+
+try:
+    import aiomysql
+except ImportError:
+    print("aiomysql is not installed, ignoring MySQL tests")
+    for k in list(deafults.keys()):
+        if k.startswith('mysql'):
+            deafults.pop(k)
+
+
 db_classes = {
     'postgres': peewee_async.PostgresqlDatabase,
     'postgres-ext': peewee_asyncext.PostgresqlExtDatabase,
@@ -93,7 +111,6 @@ def load_managers(*, managers=None, loop=None, only=None):
             pass
 
         database = db_classes[k](**config[k])
-        database.allow_sync = False
         managers[k] = peewee_async.Manager(database, loop=loop)
 
 
@@ -154,24 +171,29 @@ class BaseManagerTestCase(unittest.TestCase):
     @classmethod
     @contextlib.contextmanager
     def manager(cls, objects, allow_sync=False):
-        old_allow_sync = objects.database.allow_sync
-        objects.database.allow_sync = allow_sync
         for model in cls.models:
             model._meta.database = objects.database
-        yield
-        objects.database.allow_sync = old_allow_sync
+        if allow_sync:
+            with objects.allow_sync():
+                yield
+        else:
+            yield
 
     @classmethod
     def setUpClass(cls, *args, **kwargs):
         """Configure database managers, create test tables.
         """
         cls.managers = {}
         cls.loop = asyncio.new_event_loop()
-        # cls.loop.set_debug(True)
+        cls.loop.set_debug(True)
 
-        load_managers(managers=cls.managers, loop=cls.loop, only=cls.only)
+        load_managers(managers=cls.managers,
+                      loop=cls.loop,
+                      only=cls.only)
 
         for k, objects in cls.managers.items():
+            objects.database.allow_sync = False
+
             with cls.manager(objects, allow_sync=True):
                 for model in cls.models:
                     model.create_table(True)
@@ -182,12 +204,12 @@ def tearDownClass(cls, *args, **kwargs):
         """
         for k, objects in cls.managers.items():
             cls.loop.run_until_complete(objects.close())
+        cls.loop.close()
 
         for k, objects in cls.managers.items():
             with cls.manager(objects, allow_sync=True):
                 for model in reversed(cls.models):
                     model.drop_table(fail_silently=True, cascade=True)
-
         cls.managers = None
 
     def setUp(self):