[asyncpg] Extract rowcount for SELECT statements

Added support to the asyncpg dialect to return the ``cursor.rowcount``
value for SELECT statements when available. While this is not a typical use
for ``cursor.rowcount``, the other PostgreSQL dialects generally provide
this value. Pull request courtesy Michael Gorven.

Fixes: #9048
Closes: #9049
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/9049
Pull-request-sha: df16160530c6001d99de059995ad5047a75fb7b0

Change-Id: I095b866779ccea7e4d50bc841fef7605e61c667f
(cherry picked from commit 9c502f5788737fa65029716c73fe0f65f3dafb53)

diff --git a/doc/build/changelog/unreleased_14/9048.rst b/doc/build/changelog/unreleased_14/9048.rst
new file mode 100644 (file)
index 0000000..cf0c818
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/9048.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: bug, postgresql
+    :tickets: 9048
+    :versions: 2.0.0
+
+    Added support to the asyncpg dialect to return the ``cursor.rowcount``
+    value for SELECT statements when available. While this is not a typical use
+    for ``cursor.rowcount``, the other PostgreSQL dialects generally provide
+    this value. Pull request courtesy Michael Gorven.

diff --git a/lib/sqlalchemy/dialects/postgresql/asyncpg.py b/lib/sqlalchemy/dialects/postgresql/asyncpg.py
index f9f7e34b54860b00d50bd6aef100e0a48844c1b2..daf26a0e509527b44c9ff9c0a357966b1109b299 100644 (file)
--- a/lib/sqlalchemy/dialects/postgresql/asyncpg.py
+++ b/lib/sqlalchemy/dialects/postgresql/asyncpg.py
@@ -443,7 +443,7 @@ class AsyncAdapt_asyncpg_cursor:
                     status = prepared_stmt.get_statusmsg()
 
                     reg = re.match(
-                        r"(?:UPDATE|DELETE|INSERT \d+) (\d+)", status
+                        r"(?:SELECT|UPDATE|DELETE|INSERT \d+) (\d+)", status
                     )
                     if reg:
                         self.rowcount = int(reg.group(1))

diff --git a/test/dialect/postgresql/test_dialect.py b/test/dialect/postgresql/test_dialect.py
index fa470a18ce189c494b8df3779ba7ed137278b591..f32915cbc3f2c23c05347023110f653909c48a4c 100644 (file)
--- a/test/dialect/postgresql/test_dialect.py
+++ b/test/dialect/postgresql/test_dialect.py
@@ -1522,6 +1522,11 @@ $$ LANGUAGE plpgsql;
         with engine.connect() as conn:
             ne_(conn.connection.status, STATUS_IN_TRANSACTION)
 
+    def test_select_rowcount(self):
+        conn = testing.db.connect()
+        cursor = conn.exec_driver_sql("SELECT 1")
+        eq_(cursor.rowcount, 1)
+
 
 class AutocommitTextTest(test_deprecations.AutocommitTextTest):
     __only_on__ = "postgresql"