ext/pdo_pgsql: Add Pdo\Pgsql::ATTR_CHUNK_SIZE for chunked result fetching

Author: KentarouTakeda

NEWS

@@ -86,6 +86,8 @@ PHP                                                                        NEWS
     connection busy for the next one. (KentarouTakeda)
   . Fixed a lazy fetch returning a row of NULLs after another statement
     took over the connection. (KentarouTakeda)
+  . Added Pdo\Pgsql::ATTR_CHUNK_SIZE to fetch a result set in chunks of the
+    given number of rows. (KentarouTakeda)
 
 - Readline:
   . Fixed the interactive shell not waiting for the pager process to exit.

UPGRADING

@@ -389,6 +389,11 @@ PHP 8.6 UPGRADE NOTES
     outcome is reported as 'accepted', 'rejected' or 'not_sent' in the
     early_data key of the crypto stream_get_meta_data() array.
 
+- PDO_PGSQL:
+  . Added Pdo\Pgsql::ATTR_CHUNK_SIZE, the number of rows a statement fetches
+    per chunk. A value of 1 or more enters the lazy fetch mode of
+    PDO::ATTR_PREFETCH => 0. Requires libpq 17 or later.
+
 - Phar:
   . Overriding the getMTime() and getPathname() methods of SplFileInfo now
     influences the result of the phar buildFrom family of functions.

ext/pdo_pgsql/config.m4

@@ -25,6 +25,20 @@ if test "$PHP_PDO_PGSQL" != "no"; then
       or later).])],,
     [$PGSQL_LIBS])
 
+  old_CFLAGS=$CFLAGS
+  CFLAGS="$CFLAGS $PGSQL_CFLAGS"
+
+  AC_CHECK_DECL([PGRES_TUPLES_CHUNK],
+    PHP_CHECK_LIBRARY([pq], [PQsetChunkedRowsMode],
+      [AC_DEFINE([HAVE_PG_SET_CHUNKED_ROWS_SIZE], [1],
+        [Define to 1 if libpq has the 'PQsetChunkedRowsMode' function (PostgreSQL
+        17 or later).])],,
+      [$PGSQL_LIBS]),,
+      [#include <libpq-fe.h>]
+  )
+
+  CFLAGS=$old_CFLAGS
+
   PHP_CHECK_PDO_INCLUDES
 
   PHP_NEW_EXTENSION([pdo_pgsql],

ext/pdo_pgsql/pdo_pgsql.stub.php

@@ -18,6 +18,11 @@ class Pgsql extends \PDO
     public const int ATTR_RESULT_MEMORY_SIZE = UNKNOWN;
 #endif
 
+#ifdef HAVE_PG_SET_CHUNKED_ROWS_SIZE
+    /** @cvalue PDO_PGSQL_ATTR_CHUNK_SIZE */
+    public const int ATTR_CHUNK_SIZE = UNKNOWN;
+#endif
+
     /** @cvalue PGSQL_TRANSACTION_IDLE */
     #[\Deprecated(since: "8.5", message: "as it has no effect")]
     public const int TRANSACTION_IDLE = UNKNOWN;

ext/pdo_pgsql/pdo_pgsql_arginfo.h

@@ -266,6 +266,18 @@ static void pgsql_handle_closer(pdo_dbh_t *dbh) /* {{{ */
 }
 /* }}} */
 
+#ifdef HAVE_PG_SET_CHUNKED_ROWS_SIZE
+static bool pdo_pgsql_check_chunk_size(zend_long size)
+{
+	if (size < 0 || ZEND_LONG_EXCEEDS_INT(size)) {
+		zend_value_error("Pdo\\Pgsql::ATTR_CHUNK_SIZE must be between 0 and %d", INT_MAX);
+		return false;
+	}
+
+	return true;
+}
+#endif
+
 static bool pgsql_handle_preparer(pdo_dbh_t *dbh, zend_string *sql, pdo_stmt_t *stmt, zval *driver_options)
 {
 	pdo_pgsql_db_handle *H = (pdo_pgsql_db_handle *)dbh->driver_data;
@@ -318,6 +330,38 @@ static bool pgsql_handle_preparer(pdo_dbh_t *dbh, zend_string *sql, pdo_stmt_t *
 		: H->default_fetching_laziness
 	;
 
+#ifdef HAVE_PG_SET_CHUNKED_ROWS_SIZE
+	bool chunk_size_given = driver_options
+		&& (val = zend_hash_index_find(Z_ARRVAL_P(driver_options), PDO_PGSQL_ATTR_CHUNK_SIZE));
+
+	if (chunk_size_given) {
+		if (!pdo_get_long_param(&lval, val)) {
+			return false;
+		}
+		S->chunk_size = lval;
+	} else {
+		S->chunk_size = H->default_chunk_size;
+	}
+
+	if (!pdo_pgsql_check_chunk_size(S->chunk_size)) {
+		return false;
+	}
+
+	if (S->chunk_size >= 1 && scrollable) {
+		if (chunk_size_given) {
+			zend_value_error("Pdo\\Pgsql::ATTR_CHUNK_SIZE cannot be combined with "
+				"PDO::ATTR_CURSOR set to PDO::CURSOR_SCROLL");
+			return false;
+		}
+
+		S->chunk_size = 0;
+	}
+
+	if (S->chunk_size >= 1) {
+		S->is_unbuffered = true;
+	}
+#endif
+
 	ret = pdo_parse_params(stmt, sql, &nsql);
 
 	if (ret == -1) {
@@ -473,6 +517,12 @@ static int pdo_pgsql_get_attribute(pdo_dbh_t *dbh, zend_long attr, zval *return_
 			ZVAL_BOOL(return_value, H->disable_prepares);
 			break;
 
+#ifdef HAVE_PG_SET_CHUNKED_ROWS_SIZE
+		case PDO_PGSQL_ATTR_CHUNK_SIZE:
+			ZVAL_LONG(return_value, H->default_chunk_size);
+			break;
+#endif
+
 		case PDO_ATTR_CLIENT_VERSION: {
 			char buf[16];
 			pdo_libpq_version(buf, sizeof(buf));
@@ -1377,6 +1427,20 @@ static bool pdo_pgsql_set_attr(pdo_dbh_t *dbh, zend_long attr, zval *val)
 			}
 			H->default_fetching_laziness = !bval;
 			return true;
+#ifdef HAVE_PG_SET_CHUNKED_ROWS_SIZE
+		case PDO_PGSQL_ATTR_CHUNK_SIZE: {
+			zend_long lval;
+
+			if (!pdo_get_long_param(&lval, val)) {
+				return false;
+			}
+			if (!pdo_pgsql_check_chunk_size(lval)) {
+				return false;
+			}
+			H->default_chunk_size = lval;
+			return true;
+		}
+#endif
 		default:
 			return false;
 	}

ext/pdo_pgsql/pgsql_driver.c

@@ -58,6 +58,21 @@
 #define FIN_CLOSE   0x2
 #define FIN_ABORT   0x4
 
+static bool pgsql_result_status_ok(ExecStatusType status)
+{
+	switch (status) {
+		case PGRES_COMMAND_OK:
+		case PGRES_TUPLES_OK:
+		case PGRES_SINGLE_TUPLE:
+#ifdef HAVE_PG_SET_CHUNKED_ROWS_SIZE
+		case PGRES_TUPLES_CHUNK:
+#endif
+			return true;
+		default:
+			return false;
+	}
+}
+
 
 
 static void pgsql_stmt_finish(pdo_pgsql_stmt *S, int fin_mode)
@@ -354,16 +369,24 @@ static int pgsql_stmt_execute(pdo_stmt_t *stmt)
 			return 0;
 		}
 		S->is_running_unbuffered = true;
+		/* no matter if they return 0: PQ then transparently fallbacks to full result fetching */
+#ifdef HAVE_PG_SET_CHUNKED_ROWS_SIZE
+		if (S->chunk_size >= 1) {
+			(void)PQsetChunkedRowsMode(H->server, (int)S->chunk_size);
+		} else {
+			(void)PQsetSingleRowMode(H->server);
+		}
+#else
 		(void)PQsetSingleRowMode(H->server);
-		/* no matter if it returns 0: PQ then transparently fallbacks to full result fetching */
+#endif
 
 		/* try a first fetch to at least have column names and so on */
 		S->result = PQgetResult(S->H->server);
 	}
 
 	status = PQresultStatus(S->result);
 
-	if (status != PGRES_COMMAND_OK && status != PGRES_TUPLES_OK && status != PGRES_SINGLE_TUPLE) {
+	if (!pgsql_result_status_ok(status)) {
 		pdo_pgsql_error_stmt(stmt, status, pdo_pgsql_sqlstate(S->result));
 		return 0;
 	}
@@ -607,7 +630,7 @@ static int pgsql_stmt_fetch(pdo_stmt_t *stmt,
 			}
 			status = PQresultStatus(S->result);
 
-			if (status != PGRES_COMMAND_OK && status != PGRES_TUPLES_OK && status != PGRES_SINGLE_TUPLE) {
+			if (!pgsql_result_status_ok(status)) {
 				pdo_pgsql_error_stmt(stmt, status, pdo_pgsql_sqlstate(S->result));
 				return 0;
 			}
@@ -884,6 +907,12 @@ static int pgsql_stmt_get_attr(pdo_stmt_t *stmt, zend_long attr, zval *val)
 			return 1;
 #endif
 
+#ifdef HAVE_PG_SET_CHUNKED_ROWS_SIZE
+		case PDO_PGSQL_ATTR_CHUNK_SIZE:
+			ZVAL_LONG(val, S->chunk_size);
+			return 1;
+#endif
+
 		default:
 			(void)S;
 			return 0;

ext/pdo_pgsql/pgsql_statement.c

@@ -47,6 +47,7 @@ typedef struct {
 	HashTable       *lob_streams;
 	zend_fcall_info_cache *notice_callback;
 	bool		default_fetching_laziness;
+	zend_long	default_chunk_size;
 	pdo_pgsql_stmt  *running_stmt;
 } pdo_pgsql_db_handle;
 
@@ -66,6 +67,7 @@ struct pdo_pgsql_stmt {
 	int *param_formats;
 	Oid *param_types;
 	int                     current_row;
+	zend_long chunk_size;
 	bool is_prepared;
 	bool is_unbuffered;
 	bool is_running_unbuffered;
@@ -93,6 +95,7 @@ extern const struct pdo_stmt_methods pgsql_stmt_methods;
 enum {
 	PDO_PGSQL_ATTR_DISABLE_PREPARES = PDO_ATTR_DRIVER_SPECIFIC,
 	PDO_PGSQL_ATTR_RESULT_MEMORY_SIZE,
+	PDO_PGSQL_ATTR_CHUNK_SIZE,
 };
 
 struct pdo_pgsql_lob_self {

ext/pdo_pgsql/php_pdo_pgsql_int.h

@@ -0,0 +1,63 @@
+--TEST--
+PDO PgSQL Pdo\Pgsql::ATTR_CHUNK_SIZE splits a result set into chunks of rows
+--EXTENSIONS--
+pdo_pgsql
+--SKIPIF--
+<?php
+require __DIR__ . '/config.inc';
+require __DIR__ . '/../../../ext/pdo/tests/pdo_test.inc';
+PDOTest::skip();
+if (!defined('Pdo\Pgsql::ATTR_CHUNK_SIZE')) die('skip libpq >= 17 required');
+?>
+--FILE--
+<?php
+
+require __DIR__ . '/../../../ext/pdo/tests/pdo_test.inc';
+$pdo = PDOTest::test_factory(__DIR__ . '/common.phpt');
+$pdo->setAttribute(PDO::ATTR_ERRMODE, PDO::ERRMODE_EXCEPTION);
+
+// rowCount() reports the size of the chunk being consumed while fetching unbuffered
+function run(PDO $pdo, string $label, array $options): void
+{
+    $stmt = $pdo->prepare("SELECT * FROM generate_series(1, 10)", $options);
+    $stmt->execute();
+
+    $values = [];
+    $sizes = [];
+
+    while (($row = $stmt->fetch(PDO::FETCH_NUM))) {
+        $values[] = $row[0];
+        $sizes[] = $stmt->rowCount();
+    }
+
+    printf("%s\n  values=%s\n  chunk sizes=%s\n",
+        $label, implode(',', $values), implode(',', $sizes));
+}
+
+run($pdo, 'buffered (default)', []);
+run($pdo, 'ATTR_PREFETCH => 0', [PDO::ATTR_PREFETCH => 0]);
+run($pdo, 'ATTR_CHUNK_SIZE => 1', [Pdo\Pgsql::ATTR_CHUNK_SIZE => 1]);
+run($pdo, 'ATTR_CHUNK_SIZE => 4 (last chunk is partial)', [Pdo\Pgsql::ATTR_CHUNK_SIZE => 4]);
+run($pdo, 'ATTR_CHUNK_SIZE => 5 (divides evenly)', [Pdo\Pgsql::ATTR_CHUNK_SIZE => 5]);
+run($pdo, 'ATTR_CHUNK_SIZE => 99 (larger than the result)', [Pdo\Pgsql::ATTR_CHUNK_SIZE => 99]);
+
+?>
+--EXPECT--
+buffered (default)
+  values=1,2,3,4,5,6,7,8,9,10
+  chunk sizes=10,10,10,10,10,10,10,10,10,10
+ATTR_PREFETCH => 0
+  values=1,2,3,4,5,6,7,8,9,10
+  chunk sizes=1,1,1,1,1,1,1,1,1,1
+ATTR_CHUNK_SIZE => 1
+  values=1,2,3,4,5,6,7,8,9,10
+  chunk sizes=1,1,1,1,1,1,1,1,1,1
+ATTR_CHUNK_SIZE => 4 (last chunk is partial)
+  values=1,2,3,4,5,6,7,8,9,10
+  chunk sizes=4,4,4,4,4,4,4,4,2,2
+ATTR_CHUNK_SIZE => 5 (divides evenly)
+  values=1,2,3,4,5,6,7,8,9,10
+  chunk sizes=5,5,5,5,5,5,5,5,5,5
+ATTR_CHUNK_SIZE => 99 (larger than the result)
+  values=1,2,3,4,5,6,7,8,9,10
+  chunk sizes=10,10,10,10,10,10,10,10,10,10