From: Martin Baláž <embeembe@gmail.com>
Date: Fri, 26 Jan 2024 12:41:52 +0000 (+0100)
Subject: parent 4ccc131669ff4d5e67b7a93cbd7c7c413a172411
X-Git-Tag: 3.2.0~83^2
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=refs%2Fpull%2F724%2Fhead;p=thirdparty%2Fpsycopg.git

parent 4ccc131669ff4d5e67b7a93cbd7c7c413a172411
author Martin Baláž <embeembe@gmail.com> 1706272912 +0100
committer Daniele Varrazzo <daniele.varrazzo@gmail.com> 1706783634 +0000

feat: add `scalar_row` factory

Close #723
---

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(),