Result.php

<?php
/**
 * Copyright 2016 Google Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

namespace Google\Cloud\Spanner;

use Generator;
use Google\ApiCore\RetrySettings;
use Google\Cloud\Core\Exception\ServiceException;
use Google\Cloud\Core\ExponentialBackoff;
use Google\Cloud\Core\TimeTrait;
use Google\Cloud\Spanner\Session\SessionCache;
use Google\Cloud\Spanner\V1\ExecuteSqlRequest\QueryMode;
use Google\Cloud\Spanner\V1\MultiplexedSessionPrecommitToken;
use Grpc;

/**
 * Represent a Cloud Spanner lookup result (either read or executeSql).
 *
 * Example:
 * ```
 * use Google\Cloud\Spanner\SpannerClient;
 *
 * $spanner = new SpannerClient(['projectId' => 'my-project']);
 *
 * $database = $spanner->connect('my-instance', 'my-database');
 * $result = $database->execute('SELECT * FROM Posts');
 * ```
 *
 * @see https://cloud.google.com/spanner/docs/reference/rpc/google.spanner.v1#google.spanner.v1.ResultSet ResultSet
 */
class Result implements \IteratorAggregate
{
    use TimeTrait;

    private const BUFFER_RESULT_LIMIT = 10;
    const RETURN_NAME_VALUE_PAIR = 'nameValuePair';
    const RETURN_ASSOCIATIVE = 'associative';
    const RETURN_ZERO_INDEXED = 'zeroIndexed';
    const MODE_NORMAL = QueryMode::NORMAL;
    const MODE_PLAN = QueryMode::PLAN;
    const MODE_PROFILE = QueryMode::PROFILE;

    private array $columns = [];
    private int $columnCount = 0;
    private array $columnNames;
    private array|null $metadata;
    private int $retries;
    private string|null $resumeToken = null;
    private TransactionalReadInterface $snapshot;
    private Transaction $transaction;
    private array|bool|null $stats;

    /**
     * @var callable
     */
    private $call;
    private Generator $generator;

    /**
     * @param Operation $operation Runs operations against Google Cloud Spanner.
     * @param SessionCache $session The session used for any operations executed.
     * @param callable $call A callable, yielding a generator filled with results.
     * @param string $transactionContext The transaction's context.
     * @param ValueMapper $mapper Maps values.
     * @param RetrySettings|null $retrySettings {
     *     Retry configuration options. Currently, only the `maxRetries` option
     *     is supported.
     *
     *     @type int $maxRetries The maximum number of retry attempts before the operation
     *           fails. Defaults to 3.
     * }
     */
    public function __construct(
        private Operation $operation,
        private SessionCache $session,
        callable $call,
        private string|null $transactionContext,
        private ValueMapper $mapper,
        RetrySettings|null $retrySettings = null
    ) {
        $this->session = $session;
        $this->call = $call;
        $this->retries = isset($retrySettings) ? $retrySettings->getMaxRetries() : 3;
        $this->createGenerator();
    }

    /**
     * Return the formatted and decoded rows. If the stream is interrupted and
     * a resume token is available, attempts will be made on your behalf to
     * resume.
     *
     * Example:
     * ```
     * $rows = $result->rows();
     * ```
     *
     * @param string $format Determines the format in which rows are returned.
     *        `Result::RETURN_NAME_VALUE_PAIR` returns items as a
     *        multi-dimensional array containing a name and a value key.
     *        Ex: `[0 => ['name' => 'column1', 'value' => 'my_value']]`.
     *        `Result::RETURN_ASSOCIATIVE` returns items as an associative array
     *        with the column name as the key. Please note with this option, if
     *        duplicate column names are present a `\RuntimeException` will be
     *        thrown. `Result::RETURN_ZERO_INDEXED` returns items as a 0 indexed
     *        array, with the key representing the column number as found by
     *        executing {@see \Google\Cloud\Spanner\Result::columns()}. Ex:
     *        `[0 => 'my_value']`. **Defaults to** `Result::RETURN_ASSOCIATIVE`.
     * @return Generator
     * @throws \InvalidArgumentException When an invalid format is provided.
     * @throws \RuntimeException When duplicate column names exist with a
     *         selected format of `Result::RETURN_ASSOCIATIVE`.
     */
    public function rows($format = self::RETURN_ASSOCIATIVE): Generator
    {
        $bufferedResults = [];
        $call = $this->call;
        $shouldRetry = false;
        $isResultsYielded = false;

        $valid = $this->createGenerator();

        while ($valid) {
            try {
                $result = $this->generator->current();
                $bufferedResults[] = $result;
                $this->setResultData($result, $format);

                $empty = false;
                if (!isset($result['values']) || $this->columnCount === 0) {
                    $empty = true;
                }

                $hasResumeToken = $this->isSetAndTrue($result, 'resumeToken');
                if ($hasResumeToken || count($bufferedResults) >= self::BUFFER_RESULT_LIMIT) {
                    $chunkedResult = null;
                    if (!$empty) {
                        list($yieldableRows, $chunkedResult) = $this->parseRowsFromBufferedResults($bufferedResults);

                        foreach ($yieldableRows as $row) {
                            $isResultsYielded = true;
                            yield $this->mapper->decodeValues($this->columns, $row, $format);
                        }
                    }

                    // Now that we've yielded all available rows, flush the buffer.
                    $bufferedResults = [];

                    // If the last item in the buffer had a chunked value let's
                    // hold on to it so we can stitch it together into a yieldable
                    // result.
                    if ($chunkedResult) {
                        $bufferedResults[] = $chunkedResult;
                    }
                }

                // retry without resume token when results have not yielded
                $shouldRetry = !$isResultsYielded || $hasResumeToken;
                $this->generator->next();
                $valid = $this->generator->valid();
            } catch (ServiceException $ex) {
                if ($shouldRetry && $ex->getCode() === Grpc\STATUS_UNAVAILABLE) {
                    $backoff = new ExponentialBackoff($this->retries, function ($ex) {
                        return $ex instanceof ServiceException &&
                            $ex->getCode() === Grpc\STATUS_UNAVAILABLE;
                    });
                    // Attempt to resume using the last stored resume token and the transaction.
                    // If we successfully resume, flush the buffer.
                    $this->generator = $backoff->execute($call, [$this->resumeToken, $this->transaction()]);
                    $bufferedResults = [];

                    continue;
                }

                throw $ex;
            }
        }

        // If there are any results remaining in the buffer, yield them.
        if ($bufferedResults) {
            list($yieldableRows, $chunkedResult) = $this->parseRowsFromBufferedResults($bufferedResults);

            foreach ($yieldableRows as $row) {
                yield $this->mapper->decodeValues($this->columns, $row, $format);
            }
        }
    }

    /**
     * Return column names.
     *
     * Will be populated once the result set is iterated upon.
     *
     * Example:
     * ```
     * $columns = $result->columns();
     * ```
     *
     * @return array|null
     */
    public function columns(): array|null
    {
        return $this->columnNames ?? null;
    }

    /**
     * Return result metadata.
     *
     * Will be populated once the result set is iterated upon.
     *
     * Example:
     * ```
     * $metadata = $result->metadata();
     * ```
     *
     * @codingStandardsIgnoreStart
     * @return array|null [ResultSetMetadata](https://cloud.google.com/spanner/docs/reference/rpc/google.spanner.v1#google.spanner.v1.ResultSetMetadata).
     * @codingStandardsIgnoreEnd
     */
    public function metadata(): array|null
    {
        return $this->metadata;
    }

    /**
     * Return the session associated with the result stream.
     *
     * Example:
     * ```
     * $session = $result->session();
     * ```
     *
     * @return SessionCache
     */
    public function session(): SessionCache
    {
        return $this->session;
    }

    /**
     * Get the query plan and execution statistics for the query that produced
     * this result set.
     *
     * By default, statistics are returned by default only when executing a DML
     * statement. To receive statistics for a SELECT statement, set
     * `$options.queryMode` to `Result::MODE_PROFILE`, as demonstrated below.
     *
     * Setting `$options.queryMode` to `Result::MODE_PLAN` will return the query
     * plan without any results or execution statistics information.
     *
     * Example:
     * ```
     * $stats = $result->stats();
     * ```
     *
     * ```
     * // Executing a query with stats returned.
     * use Google\Cloud\Spanner\Result;
     *
     * $res = $database->execute('SELECT * FROM Posts', [
     *     'queryMode' => Result::MODE_PROFILE
     * ]);
     * ```
     *
     * @codingStandardsIgnoreStart
     * @return array|null [ResultSetStats](https://cloud.google.com/spanner/docs/reference/rpc/google.spanner.v1#google.spanner.v1.ResultSetStats).
     * @codingStandardsIgnoreEnd
     */
    public function stats(): array|bool|null
    {
        return $this->stats ?? null;
    }

    /**
     * Returns a snapshot which was begun in the read or execute, if one exists.
     *
     * Will be populated once the result set is iterated upon.
     *
     * Example:
     * ```
     * $snapshot = $result->snapshot();
     * ```
     *
     * @return TransactionalReadInterface|null
     */
    public function snapshot(): TransactionalReadInterface|null
    {
        return $this->snapshot ?? null;
    }

    /**
     * Returns a transaction which was begun in the read or execute, if one exists.
     *
     * Will be populated once the result set is iterated upon.
     *
     * Example:
     * ```
     * $transaction = $result->transaction();
     * ```
     *
     * @return Transaction|null
     */
    public function transaction(): Transaction|null
    {
        return $this->transaction ?? null;
    }

    /**
     * @access private
     * @return Generator
     */
    #[\ReturnTypeWillChange]
    public function getIterator(): Generator
    {
        return $this->rows();
    }

    /**
     * @param array $bufferedResults
     * @return array
     */
    private function parseRowsFromBufferedResults(array $bufferedResults): array
    {
        $values = [];
        $chunkedResult = null;
        $shouldMergeValues = $this->isSetAndTrue($bufferedResults[0], 'chunkedValue');

        $result = [];
        foreach ($bufferedResults as $key => $result) {
            if ($key === 0 && isset($bufferedResults[0]['values'])) {
                $values = $bufferedResults[0]['values'];
                continue;
            }

            $resultValues = $result['values'] ?? [];

            $values = $shouldMergeValues
                ? $this->mergeValues($values, $resultValues)
                : array_merge($values, $resultValues);

            $shouldMergeValues = $this->isSetAndTrue($result, 'chunkedValue')
                ? true
                : false;
        }

        $yieldableRows = $values && $this->columnCount > 0
            ? array_chunk($values, $this->columnCount)
            : [];

        if ($this->isSetAndTrue($result, 'chunkedValue')) {
            $chunkedResult = [
                'values' => array_pop($yieldableRows),
                'chunkedValue' => true
            ];
        }

        return [
            $yieldableRows,
            $chunkedResult
        ];
    }

    /**
     * @param array $result
     * @param string $format
     * @throws \RuntimeException
     */
    private function setResultData(array $result, $format): void
    {
        $this->stats = $result['stats'] ?? null;

        if ($this->isSetAndTrue($result, 'resumeToken')) {
            $this->resumeToken = $result['resumeToken'];
        }

        if (isset($result['metadata'])) {
            $this->columnNames = [];
            $this->columns = [];
            $this->columnCount = 0;
            $this->metadata = $result['metadata'];
            $this->columns = $result['metadata']['rowType']['fields'];

            foreach ($this->columns as $key => $column) {
                $this->columnNames[] = $this->isSetAndTrue($column, 'name')
                    ? $column['name']
                    : $key;
                $this->columnCount++;
            }

            if ($format === self::RETURN_ASSOCIATIVE
                && $this->columnCount !== count(array_unique($this->columnNames))
            ) {
                throw new \RuntimeException(
                    'Duplicate column names are not supported when returning' .
                    ' rows in the associative format. Please consider aliasing' .
                    ' your column names.'
                );
            }
        }

        $this->setSnapshotOrTransaction($result);
    }

    /**
     * Merge result set values together.
     *
     * @param array $set1
     * @param array $set2
     * @return array
     */
    private function mergeValues(array $set1, array $set2): array
    {
        // `$set2` may be empty if an array value is chunked at the end of the
        // list. Handling it normally results in an additional `null` value
        // being pushed onto the list of values. Since this method is only
        // called in cases where two chunks must be merged, we can safely
        // short-circuit the operation of the second chunk is empty.
        if (empty($set2)) {
            return $set1;
        }

        $lastItemSet1 = array_pop($set1);
        $firstItemSet2 = array_shift($set2);
        $item = $firstItemSet2;

        if (is_string($lastItemSet1) && is_string($firstItemSet2)) {
            $item = $lastItemSet1 . $firstItemSet2;
        } elseif (is_array($lastItemSet1)) {
            $item = $this->mergeValues($lastItemSet1, $firstItemSet2);
        } else {
            array_push($set1, $lastItemSet1);
        }

        array_push($set1, $item);
        return array_merge($set1, $set2);
    }

    /**
     * @param array $arr
     * @param string $key
     * @return bool
     */
    private function isSetAndTrue($arr, $key): bool
    {
        return isset($arr[$key]) && $arr[$key];
    }

    /**
     * Create the ResultSet generator and use it to set the transaction.
     *
     * @return bool Whether or not the created generator is valid.
     */
    private function createGenerator(): bool
    {
        if (isset($this->generator)) {
            return $this->generator->valid();
        }

        $call = $this->call;
        $generator = null;

        $backoff = new ExponentialBackoff($this->retries, function ($ex) {
            return $ex instanceof ServiceException &&
                $ex->getCode() === Grpc\STATUS_UNAVAILABLE;
        });

        $valid = $backoff->execute(function () use ($call, &$generator) {
            $generator = $call();
            return $generator->valid();
        });
        if ($valid) {
            // Multiple calls to the current method yields the same value.
            $result = $generator->current();
            $this->setSnapshotOrTransaction($result);
            $this->generator = $generator;
            return true;
        }
        return false;
    }

    /**
     * Set the Result's snapshot or transaction.
     *
     * @param array $result The streaming call response from the server.
     */
    private function setSnapshotOrTransaction(array $result): void
    {
        if (!empty($result['metadata']['transaction']['id'])) {
            $res = $result['metadata']['transaction'];
            if ($this->transactionContext === Database::CONTEXT_READ) {
                if (isset($res['readTimestamp'])) {
                    if (!($res['readTimestamp'] instanceof Timestamp)) {
                        $time = $this->parseTimeString($res['readTimestamp']);
                        $res['readTimestamp'] = new Timestamp($time[0], $time[1]);
                    }
                }
                $this->snapshot = $this->snapshot ?? new Snapshot(
                    $this->operation,
                    $this->session,
                    $res
                );
            } else {
                $this->transaction = $this->transaction ?? new Transaction(
                    $this->operation,
                    $this->session,
                    $res['id'],
                    [],
                    $this->mapper
                );
                if (isset($result['precommitToken']['precommitToken'])) {
                    // @TODO: Can we move this logic to the serializer or value mapper?
                    $this->transaction->setPrecommitToken(
                        (new MultiplexedSessionPrecommitToken())
                            ->setPrecommitToken(base64_decode($result['precommitToken']['precommitToken']))
                            ->setSeqNum($result['precommitToken']['seqNum'] ?? 0)
                    );
                }
            }
        }
    }
}