ext/pdo_pgsql: Retrieve the memory usage of the query result resource

Author: KentarouTakeda

ext/pdo_pgsql/pdo_pgsql.c

@@ -156,6 +156,9 @@ PHP_MINIT_FUNCTION(pdo_pgsql)
 	REGISTER_PDO_CLASS_CONST_LONG("PGSQL_TRANSACTION_INTRANS", (zend_long)PGSQL_TRANSACTION_INTRANS);
 	REGISTER_PDO_CLASS_CONST_LONG("PGSQL_TRANSACTION_INERROR", (zend_long)PGSQL_TRANSACTION_INERROR);
 	REGISTER_PDO_CLASS_CONST_LONG("PGSQL_TRANSACTION_UNKNOWN", (zend_long)PGSQL_TRANSACTION_UNKNOWN);
+#ifdef HAVE_PG_RESULT_MEMORY_SIZE
+	REGISTER_PDO_CLASS_CONST_LONG("PGSQL_ATTR_RESULT_MEMORY_SIZE", (zend_long)PRO_PGSQL_ATTR_RESULT_MEMORY_SIZE);
+#endif
 
 	PdoPgsql_ce = register_class_PdoPgsql(pdo_dbh_ce);
 	PdoPgsql_ce->create_object = pdo_dbh_new;

ext/pdo_pgsql/pgsql_statement.c

@@ -705,6 +705,26 @@ static int pdo_pgsql_stmt_cursor_closer(pdo_stmt_t *stmt)
 	return 1;
 }
 
+static int pgsql_stmt_get_attr(pdo_stmt_t *stmt, zend_long attr, zval *val)
+{
+	pdo_pgsql_stmt *S = (pdo_pgsql_stmt*)stmt->driver_data;
+
+	switch (attr) {
+#ifdef HAVE_PG_RESULT_MEMORY_SIZE
+		case PRO_PGSQL_ATTR_RESULT_MEMORY_SIZE:
+			if(stmt->executed) {
+				ZVAL_LONG(val, PQresultMemorySize(S->result));
+			} else {
+				ZVAL_NULL(val);
+			}
+			return 1;
+#endif
+
+		default:
+			return 0;
+	}
+}
+
 const struct pdo_stmt_methods pgsql_stmt_methods = {
 	pgsql_stmt_dtor,
 	pgsql_stmt_execute,
@@ -713,7 +733,7 @@ const struct pdo_stmt_methods pgsql_stmt_methods = {
 	pgsql_stmt_get_col,
 	pgsql_stmt_param_hook,
 	NULL, /* set_attr */
-	NULL, /* get_attr */
+	pgsql_stmt_get_attr,
 	pgsql_stmt_get_column_meta,
 	NULL,  /* next_rowset */
 	pdo_pgsql_stmt_cursor_closer

ext/pdo_pgsql/php_pdo_pgsql_int.h

@@ -86,6 +86,7 @@ extern const struct pdo_stmt_methods pgsql_stmt_methods;
 
 enum {
 	PDO_PGSQL_ATTR_DISABLE_PREPARES = PDO_ATTR_DRIVER_SPECIFIC,
+	PRO_PGSQL_ATTR_RESULT_MEMORY_SIZE,
 };
 
 struct pdo_pgsql_lob_self {

ext/pdo_pgsql/tests/result_memory_size.phpt

@@ -0,0 +1,43 @@
+--TEST--
+PDO PgSQL PDOStatement::getAttribute(PDO::PGSQL_ATTR_RESULT_MEMORY_SIZE)
+--EXTENSIONS--
+pdo
+pdo_pgsql
+--SKIPIF--
+<?php
+require __DIR__ . '/config.inc';
+require dirname(__DIR__, 2) . '/pdo/tests/pdo_test.inc';
+PDOTest::skip();
+if (!defined('PDO::PGSQL_ATTR_RESULT_MEMORY_SIZE')) die('skip constant PDO::PGSQL_ATTR_RESULT_MEMORY_SIZE does not exist');
+?>
+--FILE--
+<?php
+
+require_once __DIR__ . "/config.inc";
+
+/** @var Pdo */
+$db =  Pdo::connect($config['ENV']['PDOTEST_DSN']);
+
+echo 'Result set with only 1 row: ';
+$statement = $db->query('select 1');
+$result_1 = $statement->getAttribute(PDO::PGSQL_ATTR_RESULT_MEMORY_SIZE);
+var_dump($result_1);
+
+echo 'Result set with many rows: ';
+$result = $db->query('select generate_series(1, 10000)');
+$result_2 = $result->getAttribute(PDO::PGSQL_ATTR_RESULT_MEMORY_SIZE);
+var_dump($result_2);
+
+echo 'Large result sets should require more memory than small ones: ';
+var_dump($result_2 > $result_1);
+
+echo 'Statements that are not executed should not consume memory: ';
+$statement = $db->prepare('select 1');
+$result_3 = $statement->getAttribute(PDO::PGSQL_ATTR_RESULT_MEMORY_SIZE);
+var_dump($result_3);
+
+--EXPECTF--
+Result set with only 1 row: int(%d)
+Result set with many rows: int(%d)
+Large result sets should require more memory than small ones: bool(true)
+Statements that are not executed should not consume memory: NULL