transactions wip

Author: dantownsend

aragorm/engine/__init__.py

@@ -0,0 +1,5 @@
+from .base import Engine
+from .postgres import PostgresEngine
+
+
+__all__ = ['Engine', 'PostgresEngine']

aragorm/engine/base.py

@@ -0,0 +1,2 @@
+class Engine():
+    pass

aragorm/query/engine.py aragorm/engine/postgres.pyaragorm/query/engine.py renamed to aragorm/engine/postgres.py

@@ -1,11 +1,12 @@
-import asyncio
 import typing as t
 
 import asyncpg
 from asyncpg.pool import Pool
 
+from .base import Engine
 
-class Engine():
+
+class PostgresEngine(Engine):
     """
     Currently when using run ...  it sets up a connection each time.
 
@@ -14,12 +15,10 @@ class Engine():
     Needs to be a singleton that's shared by all the tables.
     """
 
-    pool: t.Optional[Pool] = None
-
     def __init__(self, config: t.Dict[str, t.Any]) -> None:
         self.config = config
 
-    async def get_pool(self):
+    async def get_pool(self) -> Pool:
         if not self.pool:
             self.pool = await asyncpg.create_pool(
                 **self.config
@@ -29,12 +28,12 @@ async def get_pool(self):
     async def run(self, query: str):
         pool = await self.get_pool()
 
-        async with pool.acquire() as connection:
-            response = connection.fetch(query)
+        connection = await pool.acquire()
+        try:
+            response = await connection.fetch(query)
+        except Exception:
+            pass
+        finally:
+            await pool.release(connection)
 
         return response
-
-    def run_sync(self, query: str):
-        return asyncio.run(
-            self.run(query)
-        )

aragorm/migration.py

@@ -11,6 +11,7 @@
 import click
 import asyncpg
 
+from aragorm.engine import PostgresEngine
 from aragorm.migrations.template import TEMPLATE
 from aragorm.migrations.table import Migration
 
@@ -129,7 +130,7 @@ def forwards():
     print('Running migrations ...')
     sys.path.insert(0, os.getcwd())
 
-    Migration.Meta.db = _get_config()
+    Migration.Meta.db = PostgresEngine(_get_config())
 
     _create_migration_table()
 
@@ -172,7 +173,7 @@ def backwards(migration_name: str):
     _get_config()  # Just required for path manipulation - needs changing
     _get_migration_modules()
 
-    Migration.Meta.db = _get_config()
+    Migration.Meta.db = PostgresEngine(_get_config())
 
     _create_migration_table()
 

aragorm/query/base.py

@@ -17,20 +17,14 @@ def __init__(self, table: 'Table', base: str = '', *args,
         self.table = table
         super().__init__()
 
-    async def run(self, as_dict=True, credentials=None):
-        """
-        Should use an engine.
-        """
-        if not credentials:
-            credentials = getattr(self.table.Meta, 'db', None)
-        if not credentials:
+    async def run(self, as_dict=True):
+        engine = getattr(self.table.Meta, 'db', None)
+        if not engine:
             raise ValueError(
                 f'Table {self.table.Meta.tablename} has no db defined in Meta'
             )
 
-        conn = await asyncpg.connect(**credentials)
-        results = await conn.fetch(self.__str__())
-        await conn.close()
+        results = await engine.run(self.__str__())
 
         if results:
             keys = results[0].keys()

aragorm/table.py

@@ -2,6 +2,7 @@
 import typing as t
 
 from .alter import Alter
+from .engine import Engine
 from .columns import Column, PrimaryKey, ForeignKey
 from .query import (
     Create,
@@ -18,17 +19,6 @@
 from .utils import _camel_to_snake
 
 
-class Database(object):
-
-    async def run(*queries):
-        """
-        Use asyncio.gather here ...
-
-        I think I'll call it Engine instead ...
-        """
-        pass
-
-
 class TableMeta(type):
 
     def __new__(cls, name, bases, namespace, **kwds):
@@ -67,7 +57,7 @@ class Table(metaclass=TableMeta):
     class Meta:
         tablename = None
         columns: t.List[Column] = []
-        db: t.Optional[t.Dict[t.Any, t.Any]] = None
+        db: t.Optional[Engine] = None
 
     def __init__(self, **kwargs):
         """

tests/base.py

@@ -10,7 +10,7 @@ class DBTestCase(TestCase):
 
     def run_sync(self, query):
         async def _run():
-            connection = await asyncpg.connect(**Pokemon.Meta.db)
+            connection = await asyncpg.connect(**Pokemon.Meta.db.config)
             await connection.execute(query)
             await connection.close()
 

tests/example_project/tables.py

@@ -1,13 +1,14 @@
 from aragorm import table
 from aragorm import columns
+from aragorm.engine import PostgresEngine
 
 
-DB = {
+DB = PostgresEngine({
     'host': 'localhost',
     'database': 'aragorm',
     'user': 'aragorm',
     'password': 'aragorm'
-}
+})
 
 
 ###############################################################################

tests/test_migration.py

@@ -1,7 +1,5 @@
-import asyncio
 from unittest import TestCase
 
-import asyncpg
 from aragorm.migrations.table import Migration
 
 from .example_project.tables import DB