%PDF-
%PDF-
Mini Shell
Mini Shell
<?php
declare(strict_types=1);
/**
* SPDX-FileCopyrightText: 2024 Robin Appelman <robin@icewind.nl>
* SPDX-License-Identifier: AGPL-3.0-or-later
*/
namespace OC\DB\QueryBuilder\Sharded;
use OC\DB\QueryBuilder\CompositeExpression;
use OC\DB\QueryBuilder\ExtendedQueryBuilder;
use OC\DB\QueryBuilder\Parameter;
use OCP\DB\IResult;
use OCP\DB\QueryBuilder\IQueryBuilder;
use OCP\IDBConnection;
/**
* A special query builder that automatically distributes queries over multiple database shards.
*
* This relies on `PartitionedQueryBuilder` to handle splitting of parts of the query that touch the sharded tables
* from the non-sharded tables. So the query build here should only either touch only sharded table or only non-sharded tables.
*
* Most of the logic in this class is concerned with extracting either the shard key (e.g. "storage") or primary key (e.g. "fileid")
* from the query. The logic for actually running the query across the shards is mostly delegated to `ShardQueryRunner`.
*/
class ShardedQueryBuilder extends ExtendedQueryBuilder {
private array $shardKeys = [];
private array $primaryKeys = [];
private ?ShardDefinition $shardDefinition = null;
/** @var bool Run the query across all shards */
private bool $allShards = false;
private ?string $insertTable = null;
private mixed $lastInsertId = null;
private ?IDBConnection $lastInsertConnection = null;
private ?int $updateShardKey = null;
private ?int $limit = null;
private ?int $offset = null;
/** @var array{column: string, order: string}[] */
private array $sortList = [];
private string $mainTable = '';
public function __construct(
IQueryBuilder $builder,
protected array $shardDefinitions,
protected ShardConnectionManager $shardConnectionManager,
protected AutoIncrementHandler $autoIncrementHandler,
) {
parent::__construct($builder);
}
public function getShardKeys(): array {
return $this->getKeyValues($this->shardKeys);
}
public function getPrimaryKeys(): array {
return $this->getKeyValues($this->primaryKeys);
}
private function getKeyValues(array $keys): array {
$values = [];
foreach ($keys as $key) {
$values = array_merge($values, $this->getKeyValue($key));
}
return array_values(array_unique($values));
}
private function getKeyValue($value): array {
if ($value instanceof Parameter) {
$value = (string)$value;
}
if (is_string($value) && str_starts_with($value, ':')) {
$param = $this->getParameter(substr($value, 1));
if (is_array($param)) {
return $param;
} else {
return [$param];
}
} elseif ($value !== null) {
return [$value];
} else {
return [];
}
}
public function where(...$predicates) {
return $this->andWhere(...$predicates);
}
public function andWhere(...$where) {
if ($where) {
foreach ($where as $predicate) {
$this->tryLoadShardKey($predicate);
}
parent::andWhere(...$where);
}
return $this;
}
private function tryLoadShardKey($predicate): void {
if (!$this->shardDefinition) {
return;
}
if ($keys = $this->tryExtractShardKeys($predicate, $this->shardDefinition->shardKey)) {
$this->shardKeys += $keys;
}
if ($keys = $this->tryExtractShardKeys($predicate, $this->shardDefinition->primaryKey)) {
$this->primaryKeys += $keys;
}
foreach ($this->shardDefinition->companionKeys as $companionKey) {
if ($keys = $this->tryExtractShardKeys($predicate, $companionKey)) {
$this->primaryKeys += $keys;
}
}
}
/**
* @param $predicate
* @param string $column
* @return string[]
*/
private function tryExtractShardKeys($predicate, string $column): array {
if ($predicate instanceof CompositeExpression) {
$values = [];
foreach ($predicate->getParts() as $part) {
$partValues = $this->tryExtractShardKeys($part, $column);
// for OR expressions, we can only rely on the predicate if all parts contain the comparison
if ($predicate->getType() === CompositeExpression::TYPE_OR && !$partValues) {
return [];
}
$values = array_merge($values, $partValues);
}
return $values;
}
$predicate = (string)$predicate;
// expect a condition in the form of 'alias1.column1 = placeholder' or 'alias1.column1 in placeholder'
if (substr_count($predicate, ' ') > 2) {
return [];
}
if (str_contains($predicate, ' = ')) {
$parts = explode(' = ', $predicate);
if ($parts[0] === "`{$column}`" || str_ends_with($parts[0], "`.`{$column}`")) {
return [$parts[1]];
} else {
return [];
}
}
if (str_contains($predicate, ' IN ')) {
$parts = explode(' IN ', $predicate);
if ($parts[0] === "`{$column}`" || str_ends_with($parts[0], "`.`{$column}`")) {
return [trim(trim($parts[1], '('), ')')];
} else {
return [];
}
}
return [];
}
public function set($key, $value) {
if ($this->shardDefinition && $key === $this->shardDefinition->shardKey) {
$updateShardKey = $value;
}
return parent::set($key, $value);
}
public function setValue($column, $value) {
if ($this->shardDefinition) {
if ($this->shardDefinition->isKey($column)) {
$this->primaryKeys[] = $value;
}
if ($column === $this->shardDefinition->shardKey) {
$this->shardKeys[] = $value;
}
}
return parent::setValue($column, $value);
}
public function values(array $values) {
foreach ($values as $column => $value) {
$this->setValue($column, $value);
}
return $this;
}
private function actOnTable(string $table): void {
$this->mainTable = $table;
foreach ($this->shardDefinitions as $shardDefinition) {
if ($shardDefinition->hasTable($table)) {
$this->shardDefinition = $shardDefinition;
}
}
}
public function from($from, $alias = null) {
if (is_string($from) && $from) {
$this->actOnTable($from);
}
return parent::from($from, $alias);
}
public function update($update = null, $alias = null) {
if (is_string($update) && $update) {
$this->actOnTable($update);
}
return parent::update($update, $alias);
}
public function insert($insert = null) {
if (is_string($insert) && $insert) {
$this->insertTable = $insert;
$this->actOnTable($insert);
}
return parent::insert($insert);
}
public function delete($delete = null, $alias = null) {
if (is_string($delete) && $delete) {
$this->actOnTable($delete);
}
return parent::delete($delete, $alias);
}
private function checkJoin(string $table): void {
if ($this->shardDefinition) {
if ($table === $this->mainTable) {
throw new InvalidShardedQueryException("Sharded query on {$this->mainTable} isn't allowed to join on itself");
}
if (!$this->shardDefinition->hasTable($table)) {
// this generally shouldn't happen as the partitioning logic should prevent this
// but the check is here just in case
throw new InvalidShardedQueryException("Sharded query on {$this->shardDefinition->table} isn't allowed to join on $table");
}
}
}
public function innerJoin($fromAlias, $join, $alias, $condition = null) {
$this->checkJoin($join);
return parent::innerJoin($fromAlias, $join, $alias, $condition);
}
public function leftJoin($fromAlias, $join, $alias, $condition = null) {
$this->checkJoin($join);
return parent::leftJoin($fromAlias, $join, $alias, $condition);
}
public function rightJoin($fromAlias, $join, $alias, $condition = null) {
if ($this->shardDefinition) {
throw new InvalidShardedQueryException("Sharded query on {$this->shardDefinition->table} isn't allowed to right join");
}
return parent::rightJoin($fromAlias, $join, $alias, $condition);
}
public function join($fromAlias, $join, $alias, $condition = null) {
return $this->innerJoin($fromAlias, $join, $alias, $condition);
}
public function setMaxResults($maxResults) {
if ($maxResults > 0) {
$this->limit = (int)$maxResults;
}
return parent::setMaxResults($maxResults);
}
public function setFirstResult($firstResult) {
if ($firstResult > 0) {
$this->offset = (int)$firstResult;
}
if ($this->shardDefinition && count($this->shardDefinition->shards) > 1) {
// we have to emulate offset
return $this;
} else {
return parent::setFirstResult($firstResult);
}
}
public function addOrderBy($sort, $order = null) {
$this->registerOrder((string)$sort, (string)$order ?? 'ASC');
return parent::addOrderBy($sort, $order);
}
public function orderBy($sort, $order = null) {
$this->sortList = [];
$this->registerOrder((string)$sort, (string)$order ?? 'ASC');
return parent::orderBy($sort, $order);
}
private function registerOrder(string $column, string $order): void {
// handle `mime + 0` and similar by just sorting on the first part of the expression
[$column] = explode(' ', $column);
$column = trim($column, '`');
$this->sortList[] = [
'column' => $column,
'order' => strtoupper($order),
];
}
public function hintShardKey(string $column, mixed $value, bool $overwrite = false): self {
if ($overwrite) {
$this->primaryKeys = [];
$this->shardKeys = [];
}
if ($this->shardDefinition?->isKey($column)) {
$this->primaryKeys[] = $value;
}
if ($column === $this->shardDefinition?->shardKey) {
$this->shardKeys[] = $value;
}
return $this;
}
public function runAcrossAllShards(): self {
$this->allShards = true;
return $this;
}
/**
* @throws InvalidShardedQueryException
*/
public function validate(): void {
if ($this->shardDefinition && $this->insertTable) {
if ($this->allShards) {
throw new InvalidShardedQueryException("Can't insert across all shards");
}
if (empty($this->getShardKeys())) {
throw new InvalidShardedQueryException("Can't insert without shard key");
}
}
if ($this->shardDefinition && !$this->allShards) {
if (empty($this->getShardKeys()) && empty($this->getPrimaryKeys())) {
throw new InvalidShardedQueryException('No shard key or primary key set for query');
}
}
if ($this->shardDefinition && $this->updateShardKey) {
$newShardKey = $this->getKeyValue($this->updateShardKey);
$oldShardKeys = $this->getShardKeys();
if (count($newShardKey) !== 1) {
throw new InvalidShardedQueryException("Can't set shard key to an array");
}
$newShardKey = current($newShardKey);
if (empty($oldShardKeys)) {
throw new InvalidShardedQueryException("Can't update without shard key");
}
$oldShards = array_values(array_unique(array_map(function ($shardKey) {
return $this->shardDefinition->getShardForKey((int)$shardKey);
}, $oldShardKeys)));
$newShard = $this->shardDefinition->getShardForKey((int)$newShardKey);
if ($oldShards === [$newShard]) {
throw new InvalidShardedQueryException('Update statement would move rows to a different shard');
}
}
}
public function executeQuery(?IDBConnection $connection = null): IResult {
$this->validate();
if ($this->shardDefinition) {
$runner = new ShardQueryRunner($this->shardConnectionManager, $this->shardDefinition);
return $runner->executeQuery($this->builder, $this->allShards, $this->getShardKeys(), $this->getPrimaryKeys(), $this->sortList, $this->limit, $this->offset);
}
return parent::executeQuery($connection);
}
public function executeStatement(?IDBConnection $connection = null): int {
$this->validate();
if ($this->shardDefinition) {
$runner = new ShardQueryRunner($this->shardConnectionManager, $this->shardDefinition);
if ($this->insertTable) {
$shards = $runner->getShards($this->allShards, $this->getShardKeys());
if (!$shards) {
throw new InvalidShardedQueryException("Can't insert without shard key");
}
$count = 0;
foreach ($shards as $shard) {
$shardConnection = $this->shardConnectionManager->getConnection($this->shardDefinition, $shard);
if (!$this->primaryKeys && $this->shardDefinition->table === $this->insertTable) {
$id = $this->autoIncrementHandler->getNextPrimaryKey($this->shardDefinition, $shard);
parent::setValue($this->shardDefinition->primaryKey, $this->createParameter('__generated_primary_key'));
$this->setParameter('__generated_primary_key', $id, self::PARAM_INT);
$this->lastInsertId = $id;
}
$count += parent::executeStatement($shardConnection);
$this->lastInsertConnection = $shardConnection;
}
return $count;
} else {
return $runner->executeStatement($this->builder, $this->allShards, $this->getShardKeys(), $this->getPrimaryKeys());
}
}
return parent::executeStatement($connection);
}
public function getLastInsertId(): int {
if ($this->lastInsertId) {
return $this->lastInsertId;
}
if ($this->lastInsertConnection) {
$table = $this->builder->prefixTableName($this->insertTable);
return $this->lastInsertConnection->lastInsertId($table);
} else {
return parent::getLastInsertId();
}
}
}
Zerion Mini Shell 1.0