123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197 |
- <?php
- declare(strict_types=1);
- namespace OC\DB\QueryBuilder\Sharded;
- use OC\DB\ArrayResult;
- use OCP\DB\IResult;
- use OCP\DB\QueryBuilder\IQueryBuilder;
- class ShardQueryRunner {
- public function __construct(
- private ShardConnectionManager $shardConnectionManager,
- private ShardDefinition $shardDefinition,
- ) {
- }
-
- public function getShards(bool $allShards, array $shardKeys): ?array {
- if ($allShards) {
- return $this->shardDefinition->getAllShards();
- }
- $allConfiguredShards = $this->shardDefinition->getAllShards();
- if (count($allConfiguredShards) === 1) {
- return $allConfiguredShards;
- }
- if (empty($shardKeys)) {
- return null;
- }
- $shards = array_map(function ($shardKey) {
- return $this->shardDefinition->getShardForKey((int)$shardKey);
- }, $shardKeys);
- return array_values(array_unique($shards));
- }
-
- private function getLikelyShards(array $primaryKeys): array {
- $shards = [];
- foreach ($primaryKeys as $primaryKey) {
- $encodedShard = $primaryKey & ShardDefinition::PRIMARY_KEY_SHARD_MASK;
- if ($encodedShard < count($this->shardDefinition->shards) && !in_array($encodedShard, $shards)) {
- $shards[] = $encodedShard;
- }
- }
- return $shards;
- }
-
- public function executeQuery(
- IQueryBuilder $query,
- bool $allShards,
- array $shardKeys,
- array $primaryKeys,
- ?array $sortList = null,
- ?int $limit = null,
- ?int $offset = null,
- ): IResult {
- $shards = $this->getShards($allShards, $shardKeys);
- $results = [];
- if ($shards && count($shards) === 1) {
-
- return $query->executeQuery($this->shardConnectionManager->getConnection($this->shardDefinition, $shards[0]));
- }
-
-
- if ($limit !== null && $offset !== null) {
- $query->setMaxResults($limit + $offset);
- }
- if ($shards) {
-
- foreach ($shards as $shard) {
- $shardConnection = $this->shardConnectionManager->getConnection($this->shardDefinition, $shard);
- $subResult = $query->executeQuery($shardConnection);
- $results = array_merge($results, $subResult->fetchAll());
- $subResult->closeCursor();
- }
- } else {
-
-
-
- $likelyShards = $this->getLikelyShards($primaryKeys);
- $unlikelyShards = array_diff($this->shardDefinition->getAllShards(), $likelyShards);
- $shards = array_merge($likelyShards, $unlikelyShards);
- foreach ($shards as $shard) {
- $shardConnection = $this->shardConnectionManager->getConnection($this->shardDefinition, $shard);
- $subResult = $query->executeQuery($shardConnection);
- $rows = $subResult->fetchAll();
- $results = array_merge($results, $rows);
- $subResult->closeCursor();
- if (count($rows) >= count($primaryKeys)) {
-
- break;
- }
- }
- }
- if ($sortList) {
- usort($results, function ($a, $b) use ($sortList) {
- foreach ($sortList as $sort) {
- $valueA = $a[$sort['column']] ?? null;
- $valueB = $b[$sort['column']] ?? null;
- $cmp = $valueA <=> $valueB;
- if ($cmp === 0) {
- continue;
- }
- if ($sort['order'] === 'DESC') {
- $cmp = -$cmp;
- }
- return $cmp;
- }
- });
- }
- if ($limit !== null && $offset !== null) {
- $results = array_slice($results, $offset, $limit);
- } elseif ($limit !== null) {
- $results = array_slice($results, 0, $limit);
- } elseif ($offset !== null) {
- $results = array_slice($results, $offset);
- }
- return new ArrayResult($results);
- }
-
- public function executeStatement(IQueryBuilder $query, bool $allShards, array $shardKeys, array $primaryKeys): int {
- if ($query->getType() === \Doctrine\DBAL\Query\QueryBuilder::INSERT) {
- throw new \Exception('insert queries need special handling');
- }
- $shards = $this->getShards($allShards, $shardKeys);
- $maxCount = count($primaryKeys);
- if ($shards && count($shards) === 1) {
- return $query->executeStatement($this->shardConnectionManager->getConnection($this->shardDefinition, $shards[0]));
- } elseif ($shards) {
- $maxCount = PHP_INT_MAX;
- } else {
-
- $likelyShards = $this->getLikelyShards($primaryKeys);
- $unlikelyShards = array_diff($this->shardDefinition->getAllShards(), $likelyShards);
- $shards = array_merge($likelyShards, $unlikelyShards);
- }
- $count = 0;
- foreach ($shards as $shard) {
- $shardConnection = $this->shardConnectionManager->getConnection($this->shardDefinition, $shard);
- $count += $query->executeStatement($shardConnection);
- if ($count >= $maxCount) {
- break;
- }
- }
- return $count;
- }
- }
|