A unified, robust and bug-free connection management interface for the ORM

docs/connections.rst

@@ -0,0 +1,58 @@
+.. _connections:
+
+===========
+Connections
+===========
+
+This document describes how to access the underlying connection object (:ref:`BaseDBAsyncClient<base_db_client>`) for the aliases defined
+as part of the DB config passed to the :meth:`Tortoise.init<tortoise.Tortoise.init>` call.
+
+Below is a simple code snippet which shows how the interface can be accessed:
+
+.. code-block::  python3
+
+    # connections is a singleton instance of the ConnectionHandler class and serves as the
+    # entrypoint to access all connection management APIs.
+    from tortoise import connections
+
+
+    # Assume that this is the Tortoise configuration used
+    await Tortoise.init(
+        {
+            "connections": {
+                "default": {
+                    "engine": "tortoise.backends.sqlite",
+                    "credentials": {"file_path": "example.sqlite3"},
+                }
+            },
+            "apps": {
+                "events": {"models": ["__main__"], "default_connection": "default"}
+            },
+        }
+    )
+
+    conn: BaseDBAsyncClient = connections.get("default")
+    try:
+        await conn.execute_query('SELECT * FROM "event"')
+    except OperationalError:
+        print("Expected it to fail")
+
+.. important::
+    The :ref:`tortoise.connection.ConnectionHandler<connection_handler>` class has been implemented with the singleton
+    pattern in mind and so when the ORM initializes, a singleton instance of this class
+    ``tortoise.connection.connections`` is created automatically and lives in memory up until the lifetime of the app.
+    Any attempt to modify or override its behaviour at runtime is risky and not recommended.
+
+
+Please refer to :ref:`this example<example_two_databases>` for a detailed demonstration of how this API can be used
+in practice.
+
+
+API Reference
+===========
+
+.. _connection_handler:
+
+.. automodule:: tortoise.connection
+    :members:
+    :undoc-members:

docs/databases.rst

@@ -214,6 +214,7 @@ handle complex objects.
         }
     )
 
+.. _base_db_client:
 
 Base DB client
 ==============

docs/reference.rst

@@ -16,6 +16,7 @@ Reference
    functions
    expressions
    transactions
+   connections
    exceptions
    signals
    migration

examples/manual_sql.py

@@ -1,7 +1,7 @@
 """
 This example demonstrates executing manual SQL queries
 """
-from tortoise import Tortoise, fields, run_async
+from tortoise import Tortoise, connections, fields, run_async
 from tortoise.models import Model
 from tortoise.transactions import in_transaction
 
@@ -17,7 +17,7 @@ async def run():
     await Tortoise.generate_schemas()
 
     # Need to get a connection. Unless explicitly specified, the name should be 'default'
-    conn = Tortoise.get_connection("default")
+    conn = connections.get("default")
 
     # Now we can execute queries in the normal autocommit mode
     await conn.execute_query("INSERT INTO event (name) VALUES ('Foo')")

examples/schema_create.py

@@ -1,7 +1,7 @@
 """
 This example demonstrates SQL Schema generation for each DB type supported.
 """
-from tortoise import Tortoise, fields, run_async
+from tortoise import Tortoise, connections, fields, run_async
 from tortoise.models import Model
 from tortoise.utils import get_schema_sql
 
@@ -49,19 +49,19 @@ class Meta:
 async def run():
     print("SQLite:\n")
     await Tortoise.init(db_url="sqlite://:memory:", modules={"models": ["__main__"]})
-    sql = get_schema_sql(Tortoise.get_connection("default"), safe=False)
+    sql = get_schema_sql(connections.get("default"), safe=False)
     print(sql)
 
     print("\n\nMySQL:\n")
     await Tortoise.init(db_url="mysql://root:@127.0.0.1:3306/", modules={"models": ["__main__"]})
-    sql = get_schema_sql(Tortoise.get_connection("default"), safe=False)
+    sql = get_schema_sql(connections.get("default"), safe=False)
     print(sql)
 
     print("\n\nPostgreSQL:\n")
     await Tortoise.init(
         db_url="postgres://postgres:@127.0.0.1:5432/", modules={"models": ["__main__"]}
     )
-    sql = get_schema_sql(Tortoise.get_connection("default"), safe=False)
+    sql = get_schema_sql(connections.get("default"), safe=False)
     print(sql)
 
 

examples/two_databases.py

@@ -8,7 +8,7 @@
 Key notes of this example is using db_route for Tortoise init
 and explicitly declaring model apps in class Meta
 """
-from tortoise import Tortoise, fields, run_async
+from tortoise import Tortoise, connections, fields, run_async
 from tortoise.exceptions import OperationalError
 from tortoise.models import Model
 
@@ -73,8 +73,8 @@ async def run():
         }
     )
     await Tortoise.generate_schemas()
-    client = Tortoise.get_connection("first")
-    second_client = Tortoise.get_connection("second")
+    client = connections.get("first")
+    second_client = connections.get("second")
 
     tournament = await Tournament.create(name="Tournament")
     await Event(name="Event", tournament_id=tournament.id).save()

tests/backends/test_capabilities.py

@@ -1,4 +1,4 @@
-from tortoise import Tortoise
+from tortoise import connections
 from tortoise.contrib import test
 
 
@@ -7,7 +7,7 @@ class TestCapabilities(test.TestCase):
 
     async def asyncSetUp(self) -> None:
         await super(TestCapabilities, self).asyncSetUp()
-        self.db = Tortoise.get_connection("models")
+        self.db = connections.get("models")
         self.caps = self.db.capabilities
 
     def test_str(self):

tests/backends/test_connection_params.py

@@ -2,20 +2,22 @@
 
 import asyncpg
 
-from tortoise import Tortoise
+from tortoise import connections
 from tortoise.contrib import test
 
 
-class TestConnectionParams(test.TestCase):
+class TestConnectionParams(test.SimpleTestCase):
     async def asyncSetUp(self) -> None:
-        pass
+        await super().asyncSetUp()
 
     async def asyncTearDown(self) -> None:
-        pass
+        await super().asyncTearDown()
 
     async def test_mysql_connection_params(self):
-        with patch("asyncmy.create_pool", new=AsyncMock()) as mysql_connect:
-            await Tortoise._init_connections(
+        with patch(
+            "tortoise.backends.mysql.client.mysql.create_pool", new=AsyncMock()
+        ) as mysql_connect:
+            await connections._init(
                 {
                     "models": {
                         "engine": "tortoise.backends.mysql",
@@ -32,6 +34,7 @@ async def test_mysql_connection_params(self):
                 },
                 False,
             )
+            await connections.get("models").create_connection(with_db=True)
 
             mysql_connect.assert_awaited_once_with(  # nosec
                 autocommit=True,
@@ -49,8 +52,10 @@ async def test_mysql_connection_params(self):
 
     async def test_asyncpg_connection_params(self):
         try:
-            with patch("asyncpg.create_pool", new=AsyncMock()) as asyncpg_connect:
-                await Tortoise._init_connections(
+            with patch(
+                "tortoise.backends.asyncpg.client.asyncpg.create_pool", new=AsyncMock()
+            ) as asyncpg_connect:
+                await connections._init(
                     {
                         "models": {
                             "engine": "tortoise.backends.asyncpg",
@@ -67,6 +72,7 @@ async def test_asyncpg_connection_params(self):
                     },
                     False,
                 )
+                await connections.get("models").create_connection(with_db=True)
 
                 asyncpg_connect.assert_awaited_once_with(  # nosec
                     None,
@@ -88,8 +94,12 @@ async def test_asyncpg_connection_params(self):
 
     async def test_psycopg_connection_params(self):
         try:
-            with patch("psycopg_pool.AsyncConnectionPool.open", new=AsyncMock()) as psycopg_connect:
-                await Tortoise._init_connections(
+            with patch(
+                "tortoise.backends.psycopg.client.PsycopgClient.create_pool", new=AsyncMock()
+            ) as patched_create_pool:
+                mocked_pool = AsyncMock()
+                patched_create_pool.return_value = mocked_pool
+                await connections._init(
                     {
                         "models": {
                             "engine": "tortoise.backends.psycopg",
@@ -106,8 +116,9 @@ async def test_psycopg_connection_params(self):
                     },
                     False,
                 )
-
-                psycopg_connect.assert_awaited_once_with(  # nosec
+                await connections.get("models").create_connection(with_db=True)
+                patched_create_pool.assert_awaited_once()
+                mocked_pool.open.assert_awaited_once_with(  # nosec
                     wait=True,
                     timeout=1,
                 )

tests/backends/test_mysql.py

@@ -18,16 +18,17 @@ async def asyncSetUp(self):
     async def asyncTearDown(self) -> None:
         if Tortoise._inited:
             await Tortoise._drop_databases()
+        await super().asyncTearDown()
 
     async def test_bad_charset(self):
         self.db_config["connections"]["models"]["credentials"]["charset"] = "terrible"
         with self.assertRaisesRegex(ConnectionError, "Unknown charset"):
-            await Tortoise.init(self.db_config)
+            await Tortoise.init(self.db_config, _create_db=True)
 
     async def test_ssl_true(self):
         self.db_config["connections"]["models"]["credentials"]["ssl"] = True
         try:
-            await Tortoise.init(self.db_config)
+            await Tortoise.init(self.db_config, _create_db=True)
         except (ConnectionError, ssl.SSLError):
             pass
         else:

tests/backends/test_postgres.py

@@ -4,7 +4,7 @@
 import ssl
 
 from tests.testmodels import Tournament
-from tortoise import Tortoise
+from tortoise import Tortoise, connections
 from tortoise.contrib import test
 from tortoise.exceptions import OperationalError
 
@@ -29,6 +29,7 @@ def is_asyncpg(self) -> bool:
     async def asyncTearDown(self) -> None:
         if Tortoise._inited:
             await Tortoise._drop_databases()
+        await super().asyncTearDown()
 
     async def test_schema(self):
         if self.is_asyncpg:
@@ -42,20 +43,20 @@ async def test_schema(self):
         with self.assertRaises(InvalidSchemaNameError):
             await Tortoise.generate_schemas()
 
-        conn = Tortoise.get_connection("models")
+        conn = connections.get("models")
         await conn.execute_script("CREATE SCHEMA mytestschema;")
         await Tortoise.generate_schemas()
 
         tournament = await Tournament.create(name="Test")
-        await Tortoise.close_connections()
+        await connections.close_all()
 
         del self.db_config["connections"]["models"]["credentials"]["schema"]
         await Tortoise.init(self.db_config)
 
         with self.assertRaises(OperationalError):
             await Tournament.filter(name="Test").first()
 
-        conn = Tortoise.get_connection("models")
+        conn = connections.get("models")
         _, res = await conn.execute_query(
             "SELECT id, name FROM mytestschema.tournament WHERE name='Test' LIMIT 1"
         )
@@ -67,7 +68,7 @@ async def test_schema(self):
     async def test_ssl_true(self):
         self.db_config["connections"]["models"]["credentials"]["ssl"] = True
         try:
-            await Tortoise.init(self.db_config)
+            await Tortoise.init(self.db_config, _create_db=True)
         except (ConnectionError, ssl.SSLError):
             pass
         else:
@@ -91,7 +92,7 @@ async def test_application_name(self):
         ] = "mytest_application"
         await Tortoise.init(self.db_config, _create_db=True)
 
-        conn = Tortoise.get_connection("models")
+        conn = connections.get("models")
         _, res = await conn.execute_query(
             "SELECT application_name FROM pg_stat_activity WHERE pid = pg_backend_pid()"
         )

tests/backends/test_reconnect.py

@@ -1,5 +1,5 @@
 from tests.testmodels import Tournament
-from tortoise import Tortoise
+from tortoise import connections
 from tortoise.contrib import test
 from tortoise.transactions import in_transaction
 
@@ -9,11 +9,11 @@ class TestReconnect(test.IsolatedTestCase):
     async def test_reconnect(self):
         await Tournament.create(name="1")
 
-        await Tortoise._connections["models"]._expire_connections()
+        await connections.get("models")._expire_connections()
 
         await Tournament.create(name="2")
 
-        await Tortoise._connections["models"]._expire_connections()
+        await connections.get("models")._expire_connections()
 
         await Tournament.create(name="3")
 
@@ -26,12 +26,12 @@ async def test_reconnect_transaction_start(self):
         async with in_transaction():
             await Tournament.create(name="1")
 
-        await Tortoise._connections["models"]._expire_connections()
+        await connections.get("models")._expire_connections()
 
         async with in_transaction():
             await Tournament.create(name="2")
 
-        await Tortoise._connections["models"]._expire_connections()
+        await connections.get("models")._expire_connections()
 
         async with in_transaction():
             self.assertEqual([f"{a.id}:{a.name}" for a in await Tournament.all()], ["1:1", "2:2"])

tests/contrib/test_functions.py

@@ -1,5 +1,5 @@
 from tests.testmodels import IntFields
-from tortoise import Tortoise
+from tortoise import connections
 from tortoise.contrib import test
 from tortoise.contrib.mysql.functions import Rand
 from tortoise.contrib.postgres.functions import Random as PostgresRandom
@@ -10,7 +10,7 @@ class TestFunction(test.TestCase):
     async def asyncSetUp(self):
         await super().asyncSetUp()
         self.intfields = [await IntFields.create(intnum=val) for val in range(10)]
-        self.db = Tortoise.get_connection("models")
+        self.db = connections.get("models")
 
     @test.requireCapability(dialect="mysql")
     async def test_mysql_func_rand(self):