feat: track expected output columns in query builder

Signed-off-by: Robin Appelman <robin@icewind.nl>

3 files changed, 41 insertions, 0 deletions

diff --git a/lib/private/DB/QueryBuilder/ExtendedQueryBuilder.php b/lib/private/DB/QueryBuilder/ExtendedQueryBuilder.php
index bde6523567f..f96ed76f6bd 100644
--- a/lib/private/DB/QueryBuilder/ExtendedQueryBuilder.php
+++ b/lib/private/DB/QueryBuilder/ExtendedQueryBuilder.php
@@ -288,4 +288,8 @@ abstract class ExtendedQueryBuilder implements IQueryBuilder {
 	public function executeStatement(?IDBConnection $connection = null): int {
 		return $this->builder->executeStatement($connection);
 	}
+
+	public function getOutputColumns(): array {
+		return $this->builder->getOutputColumns();
+	}
 }
diff --git a/lib/private/DB/QueryBuilder/QueryBuilder.php b/lib/private/DB/QueryBuilder/QueryBuilder.php
index 98280d610b1..7f26cec3403 100644
--- a/lib/private/DB/QueryBuilder/QueryBuilder.php
+++ b/lib/private/DB/QueryBuilder/QueryBuilder.php
@@ -49,6 +49,7 @@ class QueryBuilder implements IQueryBuilder {
 
 	/** @var string */
 	protected $lastInsertedTable;
+	private array $selectedColumns = [];
 
 	/**
 	 * Initializes a new QueryBuilder.
@@ -470,6 +471,7 @@ class QueryBuilder implements IQueryBuilder {
 		if (count($selects) === 1 && is_array($selects[0])) {
 			$selects = $selects[0];
 		}
+		$this->addOutputColumns($selects);
 
 		$this->queryBuilder->select(
 			$this->helper->quoteColumnNames($selects)
@@ -497,6 +499,7 @@ class QueryBuilder implements IQueryBuilder {
 		$this->queryBuilder->addSelect(
 			$this->helper->quoteColumnName($select) . ' AS ' . $this->helper->quoteColumnName($alias)
 		);
+		$this->addOutputColumns([$alias]);
 
 		return $this;
 	}
@@ -518,6 +521,7 @@ class QueryBuilder implements IQueryBuilder {
 		if (!is_array($select)) {
 			$select = [$select];
 		}
+		$this->addOutputColumns($select);
 
 		$quotedSelect = $this->helper->quoteColumnNames($select);
 
@@ -547,6 +551,7 @@ class QueryBuilder implements IQueryBuilder {
 		if (count($selects) === 1 && is_array($selects[0])) {
 			$selects = $selects[0];
 		}
+		$this->addOutputColumns($selects);
 
 		$this->queryBuilder->addSelect(
 			$this->helper->quoteColumnNames($selects)
@@ -555,6 +560,30 @@ class QueryBuilder implements IQueryBuilder {
 		return $this;
 	}
 
+	private function addOutputColumns(array $columns) {
+		foreach ($columns as $column) {
+			if (is_array($column)) {
+				$this->addOutputColumns($column);
+			} elseif (is_string($column) && !str_contains($column, '*')) {
+				if (str_contains($column, '.')) {
+					[, $column] = explode('.', $column);
+				}
+				$this->selectedColumns[] = $column;
+			}
+		}
+	}
+
+	public function getOutputColumns(): array {
+		return array_unique(array_map(function (string $column) {
+			if (str_contains($column, '.')) {
+				[, $column] = explode('.', $column);
+				return $column;
+			} else {
+				return $column;
+			}
+		}, $this->selectedColumns));
+	}
+
 	/**
 	 * Turns the query being built into a bulk delete query that ranges over
 	 * a certain table.
diff --git a/lib/public/DB/QueryBuilder/IQueryBuilder.php b/lib/public/DB/QueryBuilder/IQueryBuilder.php
index f6a2f8bd9dc..e14514ac254 100644
--- a/lib/public/DB/QueryBuilder/IQueryBuilder.php
+++ b/lib/public/DB/QueryBuilder/IQueryBuilder.php
@@ -1018,4 +1018,12 @@ interface IQueryBuilder {
 	 * @since 9.0.0
 	 */
 	public function getColumnName($column, $tableAlias = '');
+
+	/**
+	 * Get a list of column names that are expected in the query output
+	 *
+	 * @return array
+	 * @since 30.0.0
+	 */
+	public function getOutputColumns(): array;
 }