From 603e3868279cdf327aac50312ad1aa2d5c85cd5c Mon Sep 17 00:00:00 2001
From: Elvis Pranskevichus <elvis@edgedb.com>
Date: Mon, 15 Nov 2021 22:12:48 -0800
Subject: [PATCH] Add Pool methods to determine its min, max, current and idle
 size (#849)

The new `Pool.get_size()`, `Pool.get_min_size()`, `Pool.get_max_size()`,
`Pool.get_idle_size()` methods are added to get the size of the current,
minimum, maximum and idle connection set size at any given moment.
---
 asyncpg/pool.py    | 34 ++++++++++++++++++++++++++++++++++
 tests/test_pool.py | 21 +++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/asyncpg/pool.py b/asyncpg/pool.py
index c868097c..69daa77d 100644
--- a/asyncpg/pool.py
+++ b/asyncpg/pool.py
@@ -118,6 +118,12 @@ def __init__(self, pool, *, max_queries, setup, max_inactive_time):
         self._timeout = None
         self._generation = None
 
+    def is_connected(self):
+        return self._con is not None and not self._con.is_closed()
+
+    def is_idle(self):
+        return not self._in_use
+
     async def connect(self):
         if self._con is not None:
             raise exceptions.InternalClientError(
@@ -444,6 +450,34 @@ async def _initialize(self):
 
                 await asyncio.gather(*connect_tasks)
 
+    def get_size(self):
+        """Return the current number of connections in this pool.
+
+        .. versionadded:: 0.25.0
+        """
+        return sum(h.is_connected() for h in self._holders)
+
+    def get_min_size(self):
+        """Return the minimum number of connections in this pool.
+
+        .. versionadded:: 0.25.0
+        """
+        return self._minsize
+
+    def get_max_size(self):
+        """Return the maximum allowed number of connections in this pool.
+
+        .. versionadded:: 0.25.0
+        """
+        return self._maxsize
+
+    def get_idle_size(self):
+        """Return the current number of idle connections in this pool.
+
+        .. versionadded:: 0.25.0
+        """
+        return sum(h.is_connected() and h.is_idle() for h in self._holders)
+
     def set_connect_args(self, dsn=None, **connect_kwargs):
         r"""Set the new connection arguments for this pool.
 
diff --git a/tests/test_pool.py b/tests/test_pool.py
index 598baef7..e2c99efc 100644
--- a/tests/test_pool.py
+++ b/tests/test_pool.py
@@ -720,6 +720,27 @@ async def test_pool_handles_inactive_connection_errors(self):
 
         await pool.close()
 
+    async def test_pool_size_and_capacity(self):
+        async with self.create_pool(
+            database='postgres',
+            min_size=2,
+            max_size=3,
+        ) as pool:
+            self.assertEqual(pool.get_min_size(), 2)
+            self.assertEqual(pool.get_max_size(), 3)
+            self.assertEqual(pool.get_size(), 2)
+            self.assertEqual(pool.get_idle_size(), 2)
+
+            async with pool.acquire():
+                self.assertEqual(pool.get_idle_size(), 1)
+
+                async with pool.acquire():
+                    self.assertEqual(pool.get_idle_size(), 0)
+
+                    async with pool.acquire():
+                        self.assertEqual(pool.get_size(), 3)
+                        self.assertEqual(pool.get_idle_size(), 0)
+
     @unittest.skipIf(sys.version_info[:2] < (3, 6), 'no asyncgen support')
     async def test_pool_handles_transaction_exit_in_asyncgen_1(self):
         pool = await self.create_pool(database='postgres',