From: Michael Gorven <michael@gorven.net>
Date: Wed, 4 Jan 2023 17:30:42 +0000 (-0500)
Subject: [asyncpg] Extract rowcount for SELECT statements
X-Git-Tag: rel_2_0_0rc2~11^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=9c502f5788737fa65029716c73fe0f65f3dafb53;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

[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
---

diff --git a/doc/build/changelog/unreleased_14/9048.rst b/doc/build/changelog/unreleased_14/9048.rst
new file mode 100644
index 0000000000..cf0c818349
--- /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 0e3f4b5307..e005845033 100644
--- a/lib/sqlalchemy/dialects/postgresql/asyncpg.py
+++ b/lib/sqlalchemy/dialects/postgresql/asyncpg.py
@@ -483,7 +483,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 e8d9a8eb6f..78f561a352 100644
--- a/test/dialect/postgresql/test_dialect.py
+++ b/test/dialect/postgresql/test_dialect.py
@@ -1140,6 +1140,11 @@ $$ LANGUAGE plpgsql;
                 TransactionStatus.INTRANS,
             )
 
+    def test_select_rowcount(self):
+        conn = testing.db.connect()
+        cursor = conn.exec_driver_sql("SELECT 1")
+        eq_(cursor.rowcount, 1)
+
 
 class Psycopg3Test(fixtures.TestBase):
     __only_on__ = ("postgresql+psycopg",)