feat: track expected output columns in query builder

Signed-off-by: Robin Appelman <[email protected]>

lib/private/DB/QueryBuilder/ExtendedQueryBuilder.php

@@ -288,4 +288,8 @@ public function executeQuery(?IDBConnection $connection = null): IResult {
 	public function executeStatement(?IDBConnection $connection = null): int {
 		return $this->builder->executeStatement($connection);
 	}
+
+	public function getOutputColumns(): array {
+		return $this->builder->getOutputColumns();
+	}
 }

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 @@ public function select(...$selects) {
 		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 @@ public function selectAlias($select, $alias) {
 		$this->queryBuilder->addSelect(
 			$this->helper->quoteColumnName($select) . ' AS ' . $this->helper->quoteColumnName($alias)
 		);
+		$this->addOutputColumns([$alias]);
 
 		return $this;
 	}
@@ -518,6 +521,7 @@ public function selectDistinct($select) {
 		if (!is_array($select)) {
 			$select = [$select];
 		}
+		$this->addOutputColumns($select);
 
 		$quotedSelect = $this->helper->quoteColumnNames($select);
 
@@ -547,6 +551,7 @@ public function addSelect(...$selects) {
 		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 @@ public function addSelect(...$selects) {
 		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.

lib/public/DB/QueryBuilder/IQueryBuilder.php

@@ -1018,4 +1018,12 @@ public function getTableName($table);
 	 * @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;
 }