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 d4c438242541908d7ef608f083a29a1af3ba8478..15dfd3cadcea60845a1aea0c846c5fd3cfd27e67 100644 (file)
--- 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 cecffd40cdad160673d0dcb1b7326d7ddf1a240e..aa87171a5198fe86525600a93c66be173809386e 100644 (file)
--- 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 1aa910a6403794dae0eaaabb1fd656ae1ee40861..07e0dbcaf294c9361a3c4e5c245c7e26691ceee5 100644 (file)
--- 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 5165b80074b0f8f1b658e840a2ebeba5e4a48e60..93240b5eba92c6c8ab830befbf226a9cf3dd6661 100644 (file)
--- 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(),