Merge pull request #49832 from nextcloud/sharding-select-fixes

1 files changed, 24 insertions, 3 deletions

diff --git a/lib/private/DB/QueryBuilder/Partitioned/PartitionedQueryBuilder.php b/lib/private/DB/QueryBuilder/Partitioned/PartitionedQueryBuilder.php
index a9cc3e74c97..2942eeccdf7 100644
--- a/lib/private/DB/QueryBuilder/Partitioned/PartitionedQueryBuilder.php
+++ b/lib/private/DB/QueryBuilder/Partitioned/PartitionedQueryBuilder.php
@@ -76,6 +76,9 @@ class PartitionedQueryBuilder extends ShardedQueryBuilder {
 
 	// we need to save selects until we know all the table aliases
 	public function select(...$selects) {
+		if (count($selects) === 1 && is_array($selects[0])) {
+			$selects = $selects[0];
+		}
 		$this->selects = [];
 		$this->addSelect(...$selects);
 		return $this;
@@ -99,16 +102,33 @@ class PartitionedQueryBuilder extends ShardedQueryBuilder {
 	 *
 	 * This is mainly used to ensure that the returned rows from both sides of a partition contains the columns of the join predicate
 	 *
-	 * @param string $column
+	 * @param string|IQueryFunction $column
 	 * @return void
 	 */
 	private function ensureSelect(string|IQueryFunction $column, ?string $alias = null): void {
 		$checkColumn = $alias ?: $column;
 		if (str_contains($checkColumn, '.')) {
-			[, $checkColumn] = explode('.', $checkColumn);
+			[$table, $checkColumn] = explode('.', $checkColumn);
+			$partition = $this->getPartition($table);
+		} else {
+			$partition = null;
 		}
 		foreach ($this->selects as $select) {
-			if ($select['select'] === $checkColumn || $select['select'] === '*' || str_ends_with($select['select'], '.' . $checkColumn)) {
+			$select = $select['select'];
+			if (!is_string($select)) {
+				continue;
+			}
+
+			if (str_contains($select, '.')) {
+				[$table, $select] = explode('.', $select);
+				$selectPartition = $this->getPartition($table);
+			} else {
+				$selectPartition = null;
+			}
+			if (
+				($select === $checkColumn || $select === '*') &&
+				$selectPartition === $partition
+			) {
 				return;
 			}
 		}
@@ -284,6 +304,7 @@ class PartitionedQueryBuilder extends ShardedQueryBuilder {
 
 	/**
 	 * Split an array of predicates (WHERE query parts) by the partition they reference
+	 *
 	 * @param array $predicates
 	 * @return array<string, array>
 	 */