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_1_4_47~26
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=7b2cb47ba16f48462b0e6b2c8f7c24143e9f7e26;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
(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
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 f9f7e34b54..daf26a0e50 100644
--- 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 fa470a18ce..f32915cbc3 100644
--- 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"