[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 (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 0e3f4b5307bb5e1c5c8ac264932c59e29b24f255..e0058450331a6e22a0e7926bd8fdcfe8498254ad 100644 (file)
--- 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 e8d9a8eb6fd4a54f3caa5ef31b891ae8e0be231a..78f561a35256b5d77f1a4d5d2f534431e96b4952 100644 (file)
--- 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",)