From 68db21f829117c1ace9ef2738e331da852eac310 Mon Sep 17 00:00:00 2001
From: Robin Appelman <robin@icewind.nl>
Date: Thu, 12 Dec 2024 17:39:04 +0100
Subject: [PATCH] fix: improve logic for ensuring join columns are selected for
 partitioned queries

Signed-off-by: Robin Appelman <robin@icewind.nl>
---
 .../Partitioned/PartitionedQueryBuilder.php   | 24 ++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/lib/private/DB/QueryBuilder/Partitioned/PartitionedQueryBuilder.php b/lib/private/DB/QueryBuilder/Partitioned/PartitionedQueryBuilder.php
index 794ac7484d70a..ba1e5a80c0f14 100644
--- a/lib/private/DB/QueryBuilder/Partitioned/PartitionedQueryBuilder.php
+++ b/lib/private/DB/QueryBuilder/Partitioned/PartitionedQueryBuilder.php
@@ -102,16 +102,33 @@ public function selectAlias($select, $alias) {
 	 *
 	 * 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;
 			}
 		}
@@ -287,6 +304,7 @@ private function flattenPredicates(array $predicates): array {
 
 	/**
 	 * Split an array of predicates (WHERE query parts) by the partition they reference
+	 *
 	 * @param array $predicates
 	 * @return array<string, array>
 	 */