From b2c6f61844113a327a3e664769f9b46be9aabdb9 Mon Sep 17 00:00:00 2001
From: =?utf8?q?Martin=20Bal=C3=A1=C5=BE?= <embeembe@gmail.com>
Date: Fri, 26 Jan 2024 13:41:52 +0100
Subject: [PATCH] =?utf8?q?parent=204ccc131669ff4d5e67b7a93cbd7c7c413a17241?=
 =?utf8?q?1=20author=20Martin=20Bal=C3=A1=C5=BE=20<embeembe@gmail.com>=201?=
 =?utf8?q?706272912=20+0100=20committer=20Daniele=20Varrazzo=20<daniele.va?=
 =?utf8?q?rrazzo@gmail.com>=201706783634=20+0000?=
MIME-Version: 1.0
Content-Type: text/plain; charset=utf8
Content-Transfer-Encoding: 8bit

feat: add `scalar_row` factory

Close #723
---
 docs/api/rows.rst       |  4 ++++
 docs/news.rst           |  1 +
 psycopg/psycopg/rows.py | 22 ++++++++++++++++++++++
 tests/test_rows.py      | 10 ++++++++++
 4 files changed, 37 insertions(+)

diff --git a/docs/api/rows.rst b/docs/api/rows.rst
index d4c438242..15dfd3cad 100644
--- a/docs/api/rows.rst
+++ b/docs/api/rows.rst
@@ -14,6 +14,10 @@ Check out :ref:`row-factory-create` for information about how to use these objec
 .. autofunction:: tuple_row
 .. autofunction:: dict_row
 .. autofunction:: namedtuple_row
+.. autofunction:: scalar_row
+
+        .. versionadded:: 3.2
+
 .. autofunction:: class_row
 
     This is not a row factory, but rather a factory of row factories.
diff --git a/docs/news.rst b/docs/news.rst
index cecffd40c..aa87171a5 100644
--- a/docs/news.rst
+++ b/docs/news.rst
@@ -19,6 +19,7 @@ Psycopg 3.2 (unreleased)
   (:ticket:`340`).
 - Add :ref:`raw-query-cursors` to execute queries using placeholders in
   PostgreSQL format (`$1`, `$2`...) (:ticket:`#560`).
+- Add `~rows.scalar_row` to return scalar values from a query (:ticket:`#723`).
 - Add `~Connection.set_autocommit()` on sync connections, and similar
   transaction control methods available on the async connections.
 - Add support for libpq functions to close prepared statements and portals
diff --git a/psycopg/psycopg/rows.py b/psycopg/psycopg/rows.py
index 1aa910a64..07e0dbcaf 100644
--- a/psycopg/psycopg/rows.py
+++ b/psycopg/psycopg/rows.py
@@ -207,6 +207,28 @@ def kwargs_row(func: Callable[..., T]) -> BaseRowFactory[T]:
     return kwargs_row_
 
 
+def scalar_row(cursor: "BaseCursor[Any, Any]") -> "RowMaker[Any]":
+    """
+    Generate a row factory returning the first column
+    as a scalar value.
+    """
+    res = cursor.pgresult
+    if not res:
+        return no_result
+
+    nfields = _get_nfields(res)
+    if nfields is None:
+        return no_result
+
+    if nfields < 1:
+        raise e.ProgrammingError("at least one column expected")
+
+    def scalar_row_(values: Sequence[Any]) -> Any:
+        return values[0]
+
+    return scalar_row_
+
+
 def no_result(values: Sequence[Any]) -> NoReturn:
     """A `RowMaker` that always fail.
 
diff --git a/tests/test_rows.py b/tests/test_rows.py
index 5165b8007..93240b5eb 100644
--- a/tests/test_rows.py
+++ b/tests/test_rows.py
@@ -102,6 +102,16 @@ def test_kwargs_row(conn):
     assert p.age == 42
 
 
+def test_scalar_row(conn):
+    cur = conn.cursor(row_factory=rows.scalar_row)
+    cur.execute("select 1")
+    assert cur.fetchone() == 1
+    cur.execute("select 1, 2")
+    assert cur.fetchone() == 1
+    with pytest.raises(psycopg.ProgrammingError):
+        cur.execute("select")
+
+
 @pytest.mark.parametrize(
     "factory",
     "tuple_row dict_row namedtuple_row class_row args_row kwargs_row".split(),
-- 
2.47.2