feat: implement distributing partitioned queries over multiple shards

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

lib/composer/composer/autoload_classmap.php

@@ -247,6 +247,7 @@
     'OCP\\DB\\QueryBuilder\\IParameter' => $baseDir . '/lib/public/DB/QueryBuilder/IParameter.php',
     'OCP\\DB\\QueryBuilder\\IQueryBuilder' => $baseDir . '/lib/public/DB/QueryBuilder/IQueryBuilder.php',
     'OCP\\DB\\QueryBuilder\\IQueryFunction' => $baseDir . '/lib/public/DB/QueryBuilder/IQueryFunction.php',
+    'OCP\\DB\\QueryBuilder\\Sharded\\IShardMapper' => $baseDir . '/lib/public/DB/QueryBuilder/Sharded/IShardMapper.php',
     'OCP\\DB\\Types' => $baseDir . '/lib/public/DB/Types.php',
     'OCP\\Dashboard\\IAPIWidget' => $baseDir . '/lib/public/Dashboard/IAPIWidget.php',
     'OCP\\Dashboard\\IAPIWidgetV2' => $baseDir . '/lib/public/Dashboard/IAPIWidgetV2.php',
@@ -1418,6 +1419,15 @@
     'OC\\DB\\QueryBuilder\\QueryBuilder' => $baseDir . '/lib/private/DB/QueryBuilder/QueryBuilder.php',
     'OC\\DB\\QueryBuilder\\QueryFunction' => $baseDir . '/lib/private/DB/QueryBuilder/QueryFunction.php',
     'OC\\DB\\QueryBuilder\\QuoteHelper' => $baseDir . '/lib/private/DB/QueryBuilder/QuoteHelper.php',
+    'OC\\DB\\QueryBuilder\\Sharded\\AutoIncrementHandler' => $baseDir . '/lib/private/DB/QueryBuilder/Sharded/AutoIncrementHandler.php',
+    'OC\\DB\\QueryBuilder\\Sharded\\CrossShardMoveHelper' => $baseDir . '/lib/private/DB/QueryBuilder/Sharded/CrossShardMoveHelper.php',
+    'OC\\DB\\QueryBuilder\\Sharded\\HashShardMapper' => $baseDir . '/lib/private/DB/QueryBuilder/Sharded/HashShardMapper.php',
+    'OC\\DB\\QueryBuilder\\Sharded\\InvalidShardedQueryException' => $baseDir . '/lib/private/DB/QueryBuilder/Sharded/InvalidShardedQueryException.php',
+    'OC\\DB\\QueryBuilder\\Sharded\\RoundRobinShardMapper' => $baseDir . '/lib/private/DB/QueryBuilder/Sharded/RoundRobinShardMapper.php',
+    'OC\\DB\\QueryBuilder\\Sharded\\ShardConnectionManager' => $baseDir . '/lib/private/DB/QueryBuilder/Sharded/ShardConnectionManager.php',
+    'OC\\DB\\QueryBuilder\\Sharded\\ShardDefinition' => $baseDir . '/lib/private/DB/QueryBuilder/Sharded/ShardDefinition.php',
+    'OC\\DB\\QueryBuilder\\Sharded\\ShardQueryRunner' => $baseDir . '/lib/private/DB/QueryBuilder/Sharded/ShardQueryRunner.php',
+    'OC\\DB\\QueryBuilder\\Sharded\\ShardedQueryBuilder' => $baseDir . '/lib/private/DB/QueryBuilder/Sharded/ShardedQueryBuilder.php',
     'OC\\DB\\ResultAdapter' => $baseDir . '/lib/private/DB/ResultAdapter.php',
     'OC\\DB\\SQLiteMigrator' => $baseDir . '/lib/private/DB/SQLiteMigrator.php',
     'OC\\DB\\SQLiteSessionInit' => $baseDir . '/lib/private/DB/SQLiteSessionInit.php',

lib/composer/composer/autoload_static.php

@@ -280,6 +280,7 @@ class ComposerStaticInit749170dad3f5e7f9ca158f5a9f04f6a2
         'OCP\\DB\\QueryBuilder\\IParameter' => __DIR__ . '/../../..' . '/lib/public/DB/QueryBuilder/IParameter.php',
         'OCP\\DB\\QueryBuilder\\IQueryBuilder' => __DIR__ . '/../../..' . '/lib/public/DB/QueryBuilder/IQueryBuilder.php',
         'OCP\\DB\\QueryBuilder\\IQueryFunction' => __DIR__ . '/../../..' . '/lib/public/DB/QueryBuilder/IQueryFunction.php',
+        'OCP\\DB\\QueryBuilder\\Sharded\\IShardMapper' => __DIR__ . '/../../..' . '/lib/public/DB/QueryBuilder/Sharded/IShardMapper.php',
         'OCP\\DB\\Types' => __DIR__ . '/../../..' . '/lib/public/DB/Types.php',
         'OCP\\Dashboard\\IAPIWidget' => __DIR__ . '/../../..' . '/lib/public/Dashboard/IAPIWidget.php',
         'OCP\\Dashboard\\IAPIWidgetV2' => __DIR__ . '/../../..' . '/lib/public/Dashboard/IAPIWidgetV2.php',
@@ -1451,6 +1452,15 @@ class ComposerStaticInit749170dad3f5e7f9ca158f5a9f04f6a2
         'OC\\DB\\QueryBuilder\\QueryBuilder' => __DIR__ . '/../../..' . '/lib/private/DB/QueryBuilder/QueryBuilder.php',
         'OC\\DB\\QueryBuilder\\QueryFunction' => __DIR__ . '/../../..' . '/lib/private/DB/QueryBuilder/QueryFunction.php',
         'OC\\DB\\QueryBuilder\\QuoteHelper' => __DIR__ . '/../../..' . '/lib/private/DB/QueryBuilder/QuoteHelper.php',
+        'OC\\DB\\QueryBuilder\\Sharded\\AutoIncrementHandler' => __DIR__ . '/../../..' . '/lib/private/DB/QueryBuilder/Sharded/AutoIncrementHandler.php',
+        'OC\\DB\\QueryBuilder\\Sharded\\CrossShardMoveHelper' => __DIR__ . '/../../..' . '/lib/private/DB/QueryBuilder/Sharded/CrossShardMoveHelper.php',
+        'OC\\DB\\QueryBuilder\\Sharded\\HashShardMapper' => __DIR__ . '/../../..' . '/lib/private/DB/QueryBuilder/Sharded/HashShardMapper.php',
+        'OC\\DB\\QueryBuilder\\Sharded\\InvalidShardedQueryException' => __DIR__ . '/../../..' . '/lib/private/DB/QueryBuilder/Sharded/InvalidShardedQueryException.php',
+        'OC\\DB\\QueryBuilder\\Sharded\\RoundRobinShardMapper' => __DIR__ . '/../../..' . '/lib/private/DB/QueryBuilder/Sharded/RoundRobinShardMapper.php',
+        'OC\\DB\\QueryBuilder\\Sharded\\ShardConnectionManager' => __DIR__ . '/../../..' . '/lib/private/DB/QueryBuilder/Sharded/ShardConnectionManager.php',
+        'OC\\DB\\QueryBuilder\\Sharded\\ShardDefinition' => __DIR__ . '/../../..' . '/lib/private/DB/QueryBuilder/Sharded/ShardDefinition.php',
+        'OC\\DB\\QueryBuilder\\Sharded\\ShardQueryRunner' => __DIR__ . '/../../..' . '/lib/private/DB/QueryBuilder/Sharded/ShardQueryRunner.php',
+        'OC\\DB\\QueryBuilder\\Sharded\\ShardedQueryBuilder' => __DIR__ . '/../../..' . '/lib/private/DB/QueryBuilder/Sharded/ShardedQueryBuilder.php',
         'OC\\DB\\ResultAdapter' => __DIR__ . '/../../..' . '/lib/private/DB/ResultAdapter.php',
         'OC\\DB\\SQLiteMigrator' => __DIR__ . '/../../..' . '/lib/private/DB/SQLiteMigrator.php',
         'OC\\DB\\SQLiteSessionInit' => __DIR__ . '/../../..' . '/lib/private/DB/SQLiteSessionInit.php',

lib/private/DB/Connection.php

@@ -23,12 +23,19 @@
 use Doctrine\DBAL\Result;
 use Doctrine\DBAL\Schema\Schema;
 use Doctrine\DBAL\Statement;
-use OC\DB\QueryBuilder\Partitioned\PartitionSplit;
 use OC\DB\QueryBuilder\Partitioned\PartitionedQueryBuilder;
+use OC\DB\QueryBuilder\Partitioned\PartitionSplit;
 use OC\DB\QueryBuilder\QueryBuilder;
+use OC\DB\QueryBuilder\Sharded\AutoIncrementHandler;
+use OC\DB\QueryBuilder\Sharded\CrossShardMoveHelper;
+use OC\DB\QueryBuilder\Sharded\RoundRobinShardMapper;
+use OC\DB\QueryBuilder\Sharded\ShardConnectionManager;
+use OC\DB\QueryBuilder\Sharded\ShardDefinition;
 use OC\SystemConfig;
 use OCP\DB\QueryBuilder\IQueryBuilder;
+use OCP\DB\QueryBuilder\Sharded\IShardMapper;
 use OCP\Diagnostics\IEventLogger;
+use OCP\ICacheFactory;
 use OCP\IDBConnection;
 use OCP\IRequestId;
 use OCP\PreConditionNotMetException;
@@ -79,6 +86,10 @@ class Connection extends PrimaryReadReplicaConnection {
 
 	/** @var array<string, list<string>> */
 	protected array $partitions;
+	/** @var ShardDefinition[] */
+	protected array $shards = [];
+	protected ShardConnectionManager $shardConnectionManager;
+	protected AutoIncrementHandler $autoIncrementHandler;
 
 	/**
 	 * Initializes a new instance of the Connection class.
@@ -104,6 +115,13 @@ public function __construct(
 		$this->adapter = new $params['adapter']($this);
 		$this->tablePrefix = $params['tablePrefix'];
 
+		/** @psalm-suppress InvalidArrayOffset */
+		$this->shardConnectionManager = $this->params['shard_connection_manager'] ?? Server::get(ShardConnectionManager::class);
+		/** @psalm-suppress InvalidArrayOffset */
+		$this->autoIncrementHandler = $this->params['auto_increment_handler'] ?? new AutoIncrementHandler(
+			Server::get(ICacheFactory::class),
+			$this->shardConnectionManager,
+		);
 		$this->systemConfig = \OC::$server->getSystemConfig();
 		$this->clock = Server::get(ClockInterface::class);
 		$this->logger = Server::get(LoggerInterface::class);
@@ -122,11 +140,44 @@ public function __construct(
 			$this->_config->setSQLLogger($debugStack);
 		}
 
-		$this->partitions = $this->systemConfig->getValue('db.partitions', []);
+		// todo: only allow specific, pre-defined shard configurations, the current config exists for easy testing setup
+		$this->shards = array_map(function (array $config) {
+			$shardMapperClass = $config['mapper'] ?? RoundRobinShardMapper::class;
+			$shardMapper = Server::get($shardMapperClass);
+			if (!$shardMapper instanceof IShardMapper) {
+				throw new \Exception("Invalid shard mapper: $shardMapperClass");
+			}
+			return new ShardDefinition(
+				$config['table'],
+				$config['primary_key'],
+				$config['companion_keys'],
+				$config['shard_key'],
+				$shardMapper,
+				$config['companion_tables'],
+				$config['shards']
+			);
+		}, $this->params['sharding']);
+		$this->partitions = array_map(function (ShardDefinition $shard) {
+			return array_merge([$shard->table], $shard->companionTables);
+		}, $this->shards);
 
 		$this->setNestTransactionsWithSavepoints(true);
 	}
 
+	/**
+	 * @return IDBConnection[]
+	 */
+	public function getShardConnections(): array {
+		$connections = [];
+		foreach ($this->shards as $shardDefinition) {
+			foreach ($shardDefinition->getAllShards() as $shard) {
+				/** @var ConnectionAdapter $connection */
+				$connections[] = $this->shardConnectionManager->getConnection($shardDefinition, $shard);
+			}
+		}
+		return $connections;
+	}
+
 	/**
 	 * @throws Exception
 	 */
@@ -175,13 +226,19 @@ public function getStats(): array {
 	 */
 	public function getQueryBuilder(): IQueryBuilder {
 		$this->queriesBuilt++;
+
 		$builder = new QueryBuilder(
 			new ConnectionAdapter($this),
 			$this->systemConfig,
 			$this->logger
 		);
 		if (count($this->partitions) > 0) {
-			$builder = new PartitionedQueryBuilder($builder);
+			$builder = new PartitionedQueryBuilder(
+				$builder,
+				$this->shards,
+				$this->shardConnectionManager,
+				$this->autoIncrementHandler,
+			);
 			foreach ($this->partitions as $name => $tables) {
 				$partition = new PartitionSplit($name, $tables);
 				$builder->addPartition($partition);
@@ -703,6 +760,9 @@ public function migrateToSchema(Schema $toSchema, bool $dryRun = false) {
 			return $migrator->generateChangeScript($toSchema);
 		} else {
 			$migrator->migrate($toSchema);
+			foreach ($this->getShardConnections() as $shardConnection) {
+				$shardConnection->migrateToSchema($toSchema);
+			}
 		}
 	}
 
@@ -819,4 +879,12 @@ public function logDatabaseException(\Exception $exception): void {
 			}
 		}
 	}
+
+	public function getShardDefinition(string $name): ?ShardDefinition {
+		return $this->shards[$name] ?? null;
+	}
+
+	public function getCrossShardMoveHelper(): CrossShardMoveHelper {
+		return new CrossShardMoveHelper($this->shardConnectionManager);
+	}
 }

lib/private/DB/ConnectionAdapter.php

@@ -12,6 +12,8 @@
 use Doctrine\DBAL\Platforms\AbstractPlatform;
 use Doctrine\DBAL\Schema\Schema;
 use OC\DB\Exceptions\DbalException;
+use OC\DB\QueryBuilder\Sharded\CrossShardMoveHelper;
+use OC\DB\QueryBuilder\Sharded\ShardDefinition;
 use OCP\DB\IPreparedStatement;
 use OCP\DB\IResult;
 use OCP\DB\QueryBuilder\IQueryBuilder;
@@ -244,4 +246,12 @@ public function getServerVersion(): string {
 	public function logDatabaseException(\Exception $exception) {
 		$this->inner->logDatabaseException($exception);
 	}
+
+	public function getShardDefinition(string $name): ?ShardDefinition {
+		return $this->inner->getShardDefinition($name);
+	}
+
+	public function getCrossShardMoveHelper(): CrossShardMoveHelper {
+		return $this->inner->getCrossShardMoveHelper();
+	}
 }

lib/private/DB/ConnectionFactory.php

@@ -11,7 +11,11 @@
 use Doctrine\DBAL\Configuration;
 use Doctrine\DBAL\DriverManager;
 use Doctrine\DBAL\Event\Listeners\OracleSessionInit;
+use OC\DB\QueryBuilder\Sharded\AutoIncrementHandler;
+use OC\DB\QueryBuilder\Sharded\ShardConnectionManager;
 use OC\SystemConfig;
+use OCP\ICacheFactory;
+use OCP\Server;
 
 /**
  * Takes care of creating and configuring Doctrine connections.
@@ -54,9 +58,12 @@ class ConnectionFactory {
 		],
 	];
 
+	private ShardConnectionManager $shardConnectionManager;
+	private ICacheFactory $cacheFactory;
 
 	public function __construct(
-		private SystemConfig $config
+		private SystemConfig $config,
+		?ICacheFactory $cacheFactory = null,
 	) {
 		if ($this->config->getValue('mysql.utf8mb4', false)) {
 			$this->defaultConnectionParams['mysql']['charset'] = 'utf8mb4';
@@ -65,6 +72,8 @@ public function __construct(
 		if ($collationOverride) {
 			$this->defaultConnectionParams['mysql']['collation'] = $collationOverride;
 		}
+		$this->shardConnectionManager = new ShardConnectionManager($this->config, $this);
+		$this->cacheFactory = $cacheFactory ?? Server::get(ICacheFactory::class);
 	}
 
 	/**
@@ -214,6 +223,14 @@ public function createConnectionParams(string $configPrefix = '', array $additio
 		if ($this->config->getValue('dbpersistent', false)) {
 			$connectionParams['persistent'] = true;
 		}
+
+		$connectionParams['sharding'] = $this->config->getValue('dbsharding', []);
+		$connectionParams['shard_connection_manager'] = $this->shardConnectionManager;
+		$connectionParams['auto_increment_handler'] = new AutoIncrementHandler(
+			$this->cacheFactory,
+			$this->shardConnectionManager,
+		);
+
 		$connectionParams = array_merge($connectionParams, $additionalConnectionParams);
 
 		$replica = $this->config->getValue($configPrefix . 'dbreplica', $this->config->getValue('dbreplica', [])) ?: [$connectionParams];

lib/private/DB/QueryBuilder/ExtendedQueryBuilder.php

@@ -289,7 +289,21 @@ public function executeStatement(?IDBConnection $connection = null): int {
 		return $this->builder->executeStatement($connection);
 	}
 
+	public function hintShardKey(string $column, mixed $value) {
+		$this->builder->hintShardKey($column, $value);
+		return $this;
+	}
+
+	public function runAcrossAllShards() {
+		$this->builder->runAcrossAllShards();
+		return $this;
+	}
+
 	public function getOutputColumns(): array {
 		return $this->builder->getOutputColumns();
 	}
+
+	public function prefixTableName(string $table): string {
+		return $this->builder->prefixTableName($table);
+	}
 }

lib/private/DB/QueryBuilder/Partitioned/JoinCondition.php

@@ -44,7 +44,7 @@ public static function merge(array $conditions): JoinCondition {
 		$fromConditions = [];
 		$toConditions = [];
 		foreach ($conditions as $condition) {
-			if (($condition->fromColumn && $fromColumn) ||($condition->toColumn && $toColumn)) {
+			if (($condition->fromColumn && $fromColumn) || ($condition->toColumn && $toColumn)) {
 				throw new InvalidPartitionedQueryException("Can't join from {$condition->fromColumn} to {$condition->toColumn} as it already join froms {$fromColumn} to {$toColumn}");
 			}
 			if ($condition->fromColumn) {

lib/private/DB/QueryBuilder/Partitioned/PartitionQuery.php

@@ -14,12 +14,12 @@
  * A sub-query from a partitioned join
  */
 class PartitionQuery {
-	const JOIN_MODE_INNER = 'inner';
-	const JOIN_MODE_LEFT = 'left';
+	public const JOIN_MODE_INNER = 'inner';
+	public const JOIN_MODE_LEFT = 'left';
 	// left-join where the left side IS NULL
-	const JOIN_MODE_LEFT_NULL = 'left_null';
+	public const JOIN_MODE_LEFT_NULL = 'left_null';
 
-	const JOIN_MODE_RIGHT = 'right';
+	public const JOIN_MODE_RIGHT = 'right';
 
 	public function __construct(
 		public IQueryBuilder $query,

lib/private/DB/QueryBuilder/Partitioned/PartitionedQueryBuilder.php

@@ -8,18 +8,15 @@
 
 namespace OC\DB\QueryBuilder\Partitioned;
 
-use OC\DB\ConnectionAdapter;
 use OC\DB\QueryBuilder\CompositeExpression;
-use OC\DB\QueryBuilder\ExtendedQueryBuilder;
 use OC\DB\QueryBuilder\QuoteHelper;
+use OC\DB\QueryBuilder\Sharded\AutoIncrementHandler;
 use OC\DB\QueryBuilder\Sharded\ShardConnectionManager;
 use OC\DB\QueryBuilder\Sharded\ShardedQueryBuilder;
-use OC\SystemConfig;
 use OCP\DB\IResult;
 use OCP\DB\QueryBuilder\IQueryBuilder;
 use OCP\DB\QueryBuilder\IQueryFunction;
 use OCP\IDBConnection;
-use Psr\Log\LoggerInterface;
 
 /**
  * A special query builder that automatically splits queries that span across multiple database partitions[1].
@@ -38,7 +35,7 @@
  *
  * [1]: A set of tables which can't be queried together with the rest of the tables, such as when sharding is used.
  */
-class PartitionedQueryBuilder extends ExtendedQueryBuilder {
+class PartitionedQueryBuilder extends ShardedQueryBuilder {
 	/** @var array<string, PartitionQuery> $splitQueries */
 	private array $splitQueries = [];
 	/** @var list<PartitionSplit> */
@@ -53,14 +50,28 @@ class PartitionedQueryBuilder extends ExtendedQueryBuilder {
 	private ?int $offset = null;
 
 	public function __construct(
-		IQueryBuilder                  $builder,
+		IQueryBuilder          $builder,
+		array                  $shardDefinitions,
+		ShardConnectionManager $shardConnectionManager,
+		AutoIncrementHandler   $autoIncrementHandler,
 	) {
-		parent::__construct($builder);
+		parent::__construct($builder, $shardDefinitions, $shardConnectionManager, $autoIncrementHandler);
 		$this->quoteHelper = new QuoteHelper();
 	}
 
 	private function newQuery(): IQueryBuilder {
-		return $this->builder->getConnection()->getQueryBuilder();
+		// get a fresh, non-partitioning query builder
+		$builder = $this->builder->getConnection()->getQueryBuilder();
+		if ($builder instanceof PartitionedQueryBuilder) {
+			$builder = $builder->builder;
+		}
+
+		return new ShardedQueryBuilder(
+			$builder,
+			$this->shardDefinitions,
+			$this->shardConnectionManager,
+			$this->autoIncrementHandler,
+		);
 	}
 
 	// we need to save selects until we know all the table aliases
@@ -70,8 +81,8 @@ public function select(...$selects) {
 		return $this;
 	}
 
-	public function addSelect(...$selects) {
-		$selects = array_map(function($select) {
+	public function addSelect(...$select) {
+		$select = array_map(function ($select) {
 			return ['select' => $select, 'alias' => null];
 		}, $select);
 		$this->selects = array_merge($this->selects, $select);
@@ -281,7 +292,7 @@ private function splitPredicatesByParts(array $predicates): array {
 
 		$partitionPredicates = [];
 		foreach ($predicates as $predicate) {
-			$partition = $this->getPartitionForPredicate((string) $predicate);
+			$partition = $this->getPartitionForPredicate((string)$predicate);
 			if ($this->mainPartition === $partition) {
 				$partitionPredicates[''][] = $predicate;
 			} elseif ($partition) {

lib/private/DB/QueryBuilder/QueryBuilder.php

@@ -1329,7 +1329,7 @@ public function getTableName($table) {
 	 * @param string $table
 	 * @return string
 	 */
-	public function prefixTableName($table) {
+	public function prefixTableName(string $table): string {
 		if ($this->automaticTablePrefix === false || str_starts_with($table, '*PREFIX*')) {
 			return $table;
 		}