# Copyright 2016 Google LLC All rights reserved.
#
# 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.


# This file is automatically generated by CrossSync. Do not edit manually.

"""Spanner read-write transaction support."""
from dataclasses import dataclass, field
import functools
from typing import Any, Optional
from google.cloud.spanner_v1._helpers import (
    AtomicCounter,
    _make_value_pb,
    _merge_query_options,
    _metadata_with_prefix,
    _metadata_with_leader_aware_routing,
    _check_rst_stream_error,
    _merge_Transaction_Options,
    _merge_client_context,
    _merge_request_options,
)
from google.cloud.spanner_v1._helpers import _retry
from google.api_core import gapic_v1
from google.api_core.exceptions import InternalServerError
from google.protobuf.struct_pb2 import Struct
from google.cloud.spanner_v1.batch import _BatchBase
from google.cloud.spanner_v1.snapshot import _SnapshotBase
from google.cloud.spanner_v1._opentelemetry_tracing import add_span_event, trace_call
from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
from google.cloud.spanner_v1.types.commit_response import CommitResponse
from google.cloud.spanner_v1.types.mutation import Mutation
from google.cloud.spanner_v1.types.result_set import ResultSet
from google.cloud.spanner_v1.types.spanner import (
    CommitRequest,
    ExecuteBatchDmlRequest,
    ExecuteBatchDmlResponse,
    ExecuteSqlRequest,
    RequestOptions,
)
from google.cloud.spanner_v1.types.transaction import TransactionOptions


class Transaction(_SnapshotBase, _BatchBase):
    """Implement read-write transaction semantics for a session.

    :type session: :class:`~google.cloud.spanner_v1.session.Session`
    :param session: the session used to perform the commit

    :raises ValueError: if session has an existing transaction
    """

    exclude_txn_from_change_streams: bool = False
    isolation_level: TransactionOptions.IsolationLevel = (
        TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED
    )
    read_lock_mode: TransactionOptions.ReadWrite.ReadLockMode = (
        TransactionOptions.ReadWrite.ReadLockMode.READ_LOCK_MODE_UNSPECIFIED
    )
    _multi_use: bool = True
    _read_only: bool = False

    def __init__(self, session, client_context=None):
        super(Transaction, self).__init__(session, client_context=client_context)
        self.rolled_back: bool = False
        self._multiplexed_session_previous_transaction_id: Optional[bytes] = None

    def _build_transaction_options_pb(self) -> TransactionOptions:
        """Builds and returns transaction options for this transaction.

        :rtype: :class:`~.transaction_pb2.TransactionOptions`
        :returns: transaction options for this transaction."""
        default_transaction_options = (
            self._session._database.default_transaction_options.default_read_write_transaction_options
        )
        merge_transaction_options = TransactionOptions(
            read_write=TransactionOptions.ReadWrite(
                multiplexed_session_previous_transaction_id=self._multiplexed_session_previous_transaction_id,
                read_lock_mode=self.read_lock_mode,
            ),
            exclude_txn_from_change_streams=self.exclude_txn_from_change_streams,
            isolation_level=self.isolation_level,
        )
        return _merge_Transaction_Options(
            defaultTransactionOptions=default_transaction_options,
            mergeTransactionOptions=merge_transaction_options,
        )

    def _execute_request(
        self, method, request, metadata, trace_name=None, attributes=None
    ):
        """Helper method to execute request after fetching transaction selector.

        :type method: callable
        :param method: function returning iterator

        :type request: proto
        :param request: request proto to call the method with

        :raises: ValueError: if the transaction is not ready to update."""
        if self.committed is not None:
            raise ValueError("Transaction already committed.")
        if self.rolled_back:
            raise ValueError("Transaction already rolled back.")
        session = self._session
        transaction = self._build_transaction_selector_pb()
        request.transaction = transaction
        with trace_call(
            trace_name,
            session,
            attributes,
            observability_options=getattr(
                session._database, "observability_options", None
            ),
            metadata=metadata,
        ), MetricsCapture(self._resource_info):
            method = functools.partial(method, request=request)
            response = _retry(
                method,
                allowed_exceptions={InternalServerError: _check_rst_stream_error},
            )
        return response

    def rollback(self) -> None:
        """Roll back a transaction on the database.

        :raises: ValueError: if the transaction is not ready to roll back."""
        if self.committed is not None:
            raise ValueError("Transaction already committed.")
        if self.rolled_back:
            raise ValueError("Transaction already rolled back.")
        if self._transaction_id is not None:
            session = self._session
            database = session._database
            api = database.spanner_api
            metadata = _metadata_with_prefix(database.name)
            if database._route_to_leader_enabled:
                metadata.append(
                    _metadata_with_leader_aware_routing(
                        database._route_to_leader_enabled
                    )
                )
            observability_options = getattr(database, "observability_options", None)
            with trace_call(
                f"CloudSpanner.{type(self).__name__}.rollback",
                session,
                observability_options=observability_options,
                metadata=metadata,
            ) as span, MetricsCapture(self._resource_info):
                attempt = AtomicCounter(0)
                nth_request = database._next_nth_request

                def wrapped_method(*args, **kwargs):
                    attempt.increment()
                    call_metadata, error_augmenter = database.with_error_augmentation(
                        nth_request, attempt.value, metadata, span
                    )
                    rollback_method = functools.partial(
                        api.rollback,
                        session=session.name,
                        transaction_id=self._transaction_id,
                        metadata=call_metadata,
                    )
                    with error_augmenter:
                        return rollback_method(*args, **kwargs)

                _retry(
                    wrapped_method,
                    allowed_exceptions={InternalServerError: _check_rst_stream_error},
                )
        self.rolled_back = True

    def _reset_and_begin(self):
        """This function can be used to reset the transaction and execute an explicit BeginTransaction RPC if the first statement in the transaction failed, and that statement included an inlined BeginTransaction option."""
        self._read_request_count = 0
        self._execute_sql_request_count = 0
        self.begin()

    def commit(
        self, return_commit_stats=False, request_options=None, max_commit_delay=None
    ):
        """Commit mutations to the database.

        :type return_commit_stats: bool
        :param return_commit_stats:
          If true, the response will return commit stats which can be accessed though commit_stats.

        :type request_options:
            :class:`google.cloud.spanner_v1.types.RequestOptions`
        :param request_options:
                (Optional) Common options for this request.
                If a dict is provided, it must be of the same form as the protobuf
                message :class:`~google.cloud.spanner_v1.types.RequestOptions`.

        :type max_commit_delay: :class:`datetime.timedelta`
        :param max_commit_delay:
                (Optional) The amount of latency this request is willing to incur
                in order to improve throughput.
                :class:`~google.cloud.spanner_v1.types.MaxCommitDelay`.

        :rtype: datetime
        :returns: timestamp of the committed changes.

        :raises: ValueError: if the transaction is not ready to commit."""
        mutations = self._mutations
        num_mutations = len(mutations)
        session = self._session
        database = session._database
        api = database.spanner_api
        metadata = _metadata_with_prefix(database.name)
        if database._route_to_leader_enabled:
            metadata.append(
                _metadata_with_leader_aware_routing(database._route_to_leader_enabled)
            )
        with trace_call(
            name=f"CloudSpanner.{type(self).__name__}.commit",
            session=session,
            extra_attributes={"num_mutations": num_mutations},
            observability_options=getattr(database, "observability_options", None),
            metadata=metadata,
        ) as span, MetricsCapture(self._resource_info):
            if self.committed is not None:
                raise ValueError("Transaction already committed.")
            if self.rolled_back:
                raise ValueError("Transaction already rolled back.")
            if self._transaction_id is None:
                if num_mutations > 0:
                    self._begin_mutations_only_transaction()
                else:
                    raise ValueError("Transaction has not begun.")
            client_context = _merge_client_context(
                database._instance._client._client_context, self._client_context
            )
            request_options = _merge_request_options(request_options, client_context)
            if request_options is None:
                request_options = RequestOptions()
            elif type(request_options) is dict:
                request_options = RequestOptions(request_options)
            if self.transaction_tag is not None:
                request_options.transaction_tag = self.transaction_tag
            request_options.request_tag = None
            common_commit_request_args = {
                "session": session.name,
                "transaction_id": self._transaction_id,
                "return_commit_stats": return_commit_stats,
                "max_commit_delay": max_commit_delay,
                "request_options": request_options,
            }
            add_span_event(span, "Starting Commit")
            attempt = AtomicCounter(0)
            nth_request = database._next_nth_request

            def wrapped_method(*args, **kwargs):
                attempt.increment()
                commit_request_args = {
                    "mutations": mutations,
                    **common_commit_request_args,
                }
                is_multiplexed = getattr(self._session, "is_multiplexed", False)
                if is_multiplexed and self._precommit_token is not None:
                    commit_request_args["precommit_token"] = self._precommit_token
                call_metadata, error_augmenter = database.with_error_augmentation(
                    nth_request, attempt.value, metadata, span
                )
                commit_method = functools.partial(
                    api.commit,
                    request=CommitRequest(**commit_request_args),
                    metadata=call_metadata,
                )
                with error_augmenter:
                    return commit_method(*args, **kwargs)

            commit_retry_event_name = "Transaction Commit Attempt Failed. Retrying"

            def before_next_retry(nth_retry, delay_in_seconds):
                add_span_event(
                    span=span,
                    event_name=commit_retry_event_name,
                    event_attributes={
                        "attempt": nth_retry,
                        "sleep_seconds": delay_in_seconds,
                    },
                )

            commit_response_pb: CommitResponse = _retry(
                wrapped_method,
                allowed_exceptions={InternalServerError: _check_rst_stream_error},
                before_next_retry=before_next_retry,
            )
            if commit_response_pb._pb.HasField("precommit_token"):
                add_span_event(span, commit_retry_event_name)
                nth_request = database._next_nth_request
                call_metadata, error_augmenter = database.with_error_augmentation(
                    nth_request, 1, metadata, span
                )
                with error_augmenter:
                    commit_response_pb = api.commit(
                        request=CommitRequest(
                            precommit_token=commit_response_pb.precommit_token,
                            **common_commit_request_args,
                        ),
                        metadata=call_metadata,
                    )
            add_span_event(span, "Commit Done")
        self.committed = commit_response_pb.commit_timestamp
        if return_commit_stats:
            self.commit_stats = commit_response_pb.commit_stats
        return self.committed

    @staticmethod
    def _make_params_pb(params, param_types):
        """Helper for :meth:`execute_update`.

        :type params: dict, {str -> column value}
        :param params: values for parameter replacement.  Keys must match
                       the names used in ``dml``.

        :type param_types: dict[str -> Union[dict, .types.Type]]
        :param param_types:
            (Optional) maps explicit types for one or more param values;
            required if parameters are passed.

        :rtype: Union[None, :class:`Struct`]
        :returns: a struct message for the passed params, or None
        :raises ValueError:
            If ``param_types`` is None but ``params`` is not None.
        :raises ValueError:
            If ``params`` is None but ``param_types`` is not None."""
        if params:
            return Struct(
                fields={key: _make_value_pb(value) for key, value in params.items()}
            )
        return {}

    def execute_update(
        self,
        dml,
        params=None,
        param_types=None,
        query_mode=None,
        query_options=None,
        request_options=None,
        last_statement=False,
        *,
        retry=gapic_v1.method.DEFAULT,
        timeout=gapic_v1.method.DEFAULT,
    ):
        """Perform an ``ExecuteSql`` API request with DML.

        :type dml: str
        :param dml: SQL DML statement

        :type params: dict, {str -> column value}
        :param params: values for parameter replacement.  Keys must match
                       the names used in ``dml``.

        :type param_types: dict[str -> Union[dict, .types.Type]]
        :param param_types:
            (Optional) maps explicit types for one or more param values;
            required if parameters are passed.

        :type query_mode:
            :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryMode`
        :param query_mode: Mode governing return of results / query plan.
            See:
            `QueryMode <https://cloud.google.com/spanner/reference/rpc/google.spanner.v1#google.spanner.v1.ExecuteSqlRequest.QueryMode>`_.

        :type query_options:
            :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryOptions`
            or :class:`dict`
        :param query_options: (Optional) Options that are provided for query plan stability.

        :type request_options:
            :class:`google.cloud.spanner_v1.types.RequestOptions`
        :param request_options:
                (Optional) Common options for this request.
                If a dict is provided, it must be of the same form as the protobuf
                message :class:`~google.cloud.spanner_v1.types.RequestOptions`.

        :type last_statement: bool
        :param last_statement:
                If set to true, this option marks the end of the transaction. The
                transaction should be committed or aborted after this statement
                executes, and attempts to execute any other requests against this
                transaction (including reads and queries) will be rejected. Mixing
                mutations with statements that are marked as the last statement is
                not allowed.
                For DML statements, setting this option may cause some error
                reporting to be deferred until commit time (e.g. validation of
                unique constraints). Given this, successful execution of a DML
                statement should not be assumed until the transaction commits.

        :type retry: :class:`~google.api_core.retry.Retry`
        :param retry: (Optional) The retry settings for this request.

        :type timeout: float
        :param timeout: (Optional) The timeout for this request.

        :rtype: int
        :returns: Count of rows affected by the DML statement."""
        session = self._session
        database = session._database
        api = database.spanner_api
        params_pb = self._make_params_pb(params, param_types)
        metadata = _metadata_with_prefix(database.name)
        if database._route_to_leader_enabled:
            metadata.append(
                _metadata_with_leader_aware_routing(database._route_to_leader_enabled)
            )
        seqno, self._execute_sql_request_count = (
            self._execute_sql_request_count,
            self._execute_sql_request_count + 1,
        )
        default_query_options = database._instance._client._query_options
        query_options = _merge_query_options(default_query_options, query_options)
        client_context = _merge_client_context(
            database._instance._client._client_context, self._client_context
        )
        request_options = _merge_request_options(request_options, client_context)
        if request_options is None:
            request_options = RequestOptions()
        elif type(request_options) is dict:
            request_options = RequestOptions(request_options)
        request_options.transaction_tag = self.transaction_tag
        trace_attributes = {"db.statement": dml, "request_options": request_options}
        is_inline_begin = False
        if self._transaction_id is None:
            is_inline_begin = True
            self._lock.acquire()
        execute_sql_request = ExecuteSqlRequest(
            session=session.name,
            transaction=self._build_transaction_selector_pb(),
            sql=dml,
            params=params_pb,
            param_types=param_types,
            query_mode=query_mode,
            query_options=query_options,
            seqno=seqno,
            request_options=request_options,
            last_statement=last_statement,
        )
        nth_request = database._next_nth_request
        attempt = AtomicCounter(0)

        def wrapped_method(*args, **kwargs):
            attempt.increment()
            call_metadata, error_augmenter = database.with_error_augmentation(
                nth_request, attempt.value, metadata
            )
            execute_sql_method = functools.partial(
                api.execute_sql,
                request=execute_sql_request,
                metadata=call_metadata,
                retry=retry,
                timeout=timeout,
            )
            with error_augmenter:
                return execute_sql_method(*args, **kwargs)

        result_set_pb: ResultSet = self._execute_request(
            wrapped_method,
            execute_sql_request,
            metadata,
            f"CloudSpanner.{type(self).__name__}.execute_update",
            trace_attributes,
        )
        self._update_for_result_set_pb(result_set_pb)
        if is_inline_begin:
            self._lock.release()
        if result_set_pb._pb.HasField("precommit_token"):
            self._update_for_precommit_token_pb(result_set_pb.precommit_token)
        return result_set_pb.stats.row_count_exact

    def batch_update(
        self,
        statements,
        request_options=None,
        last_statement=False,
        *,
        retry=gapic_v1.method.DEFAULT,
        timeout=gapic_v1.method.DEFAULT,
    ):
        """Perform a batch of DML statements via an ``ExecuteBatchDml`` request.

        :type statements:
            Sequence[Union[ str, Tuple[str, Dict[str, Any], Dict[str, Union[dict, .types.Type]]]]]

        :param statements:
            List of DML statements, with optional params / param types.
            If passed, 'params' is a dict mapping names to the values
            for parameter replacement.  Keys must match the names used in the
            corresponding DML statement.  If 'params' is passed, 'param_types'
            must also be passed, as a dict mapping names to the type of
            value passed in 'params'.

        :type request_options:
            :class:`google.cloud.spanner_v1.types.RequestOptions`
        :param request_options:
                (Optional) Common options for this request.
                If a dict is provided, it must be of the same form as the protobuf
                message :class:`~google.cloud.spanner_v1.types.RequestOptions`.

        :type last_statement: bool
        :param last_statement:
                If set to true, this option marks the end of the transaction. The
                transaction should be committed or aborted after this statement
                executes, and attempts to execute any other requests against this
                transaction (including reads and queries) will be rejected. Mixing
                mutations with statements that are marked as the last statement is
                not allowed.
                For DML statements, setting this option may cause some error
                reporting to be deferred until commit time (e.g. validation of
                unique constraints). Given this, successful execution of a DML
                statement should not be assumed until the transaction commits.

        :type retry: :class:`~google.api_core.retry.Retry`
        :param retry: (Optional) The retry settings for this request.

        :type timeout: float
        :param timeout: (Optional) The timeout for this request.

        :rtype:
            Tuple(status, Sequence[int])
        :returns:
            Status code, plus counts of rows affected by each completed DML
            statement.  Note that if the status code is not ``OK``, the
            statement triggering the error will not have an entry in the
            list, nor will any statements following that one."""
        session = self._session
        database = session._database
        api = database.spanner_api
        parsed = []
        for statement in statements:
            if isinstance(statement, str):
                parsed.append(ExecuteBatchDmlRequest.Statement(sql=statement))
            else:
                dml, params, param_types = statement
                params_pb = self._make_params_pb(params, param_types)
                parsed.append(
                    ExecuteBatchDmlRequest.Statement(
                        sql=dml, params=params_pb, param_types=param_types
                    )
                )
        metadata = _metadata_with_prefix(database.name)
        if database._route_to_leader_enabled:
            metadata.append(
                _metadata_with_leader_aware_routing(database._route_to_leader_enabled)
            )
        seqno, self._execute_sql_request_count = (
            self._execute_sql_request_count,
            self._execute_sql_request_count + 1,
        )
        client_context = _merge_client_context(
            database._instance._client._client_context, self._client_context
        )
        request_options = _merge_request_options(request_options, client_context)
        if request_options is None:
            request_options = RequestOptions()
        elif type(request_options) is dict:
            request_options = RequestOptions(request_options)
        request_options.transaction_tag = self.transaction_tag
        trace_attributes = {
            "db.statement": ";".join([statement.sql for statement in parsed]),
            "request_options": request_options,
        }
        is_inline_begin = False
        if self._transaction_id is None:
            is_inline_begin = True
            self._lock.acquire()
        execute_batch_dml_request = ExecuteBatchDmlRequest(
            session=session.name,
            transaction=self._build_transaction_selector_pb(),
            statements=parsed,
            seqno=seqno,
            request_options=request_options,
            last_statements=last_statement,
        )
        nth_request = database._next_nth_request
        attempt = AtomicCounter(0)

        def wrapped_method(*args, **kwargs):
            attempt.increment()
            call_metadata, error_augmenter = database.with_error_augmentation(
                nth_request, attempt.value, metadata
            )
            execute_batch_dml_method = functools.partial(
                api.execute_batch_dml,
                request=execute_batch_dml_request,
                metadata=call_metadata,
                retry=retry,
                timeout=timeout,
            )
            with error_augmenter:
                return execute_batch_dml_method(*args, **kwargs)

        response_pb: ExecuteBatchDmlResponse = self._execute_request(
            wrapped_method,
            execute_batch_dml_request,
            metadata,
            "CloudSpanner.DMLTransaction",
            trace_attributes,
        )
        self._update_for_execute_batch_dml_response_pb(response_pb)
        if is_inline_begin:
            self._lock.release()
        if (
            len(response_pb.result_sets) > 0
            and response_pb.result_sets[0].precommit_token
        ):
            self._update_for_precommit_token_pb(
                response_pb.result_sets[0].precommit_token
            )
        row_counts = [
            result_set.stats.row_count_exact for result_set in response_pb.result_sets
        ]
        return (response_pb.status, row_counts)

    def _begin_transaction(self, mutation: Mutation = None) -> bytes:
        """Begins a transaction on the database.

        :type mutation: :class:`~google.cloud.spanner_v1.mutation.Mutation`
        :param mutation: (Optional) Mutation to include in the begin transaction
            request. Required for mutation-only transactions with multiplexed sessions.

        :rtype: bytes
        :returns: identifier for the transaction.

        :raises ValueError: if the transaction has already begun or is single-use."""
        if self.committed is not None:
            raise ValueError("Transaction is already committed")
        if self.rolled_back:
            raise ValueError("Transaction is already rolled back")
        return super(Transaction, self)._begin_transaction(
            mutation=mutation, transaction_tag=self.transaction_tag
        )

    def _begin_mutations_only_transaction(self) -> None:
        """Begins a mutations-only transaction on the database."""
        mutation = self._get_mutation_for_begin_mutations_only_transaction()
        self._begin_transaction(mutation=mutation)

    def _get_mutation_for_begin_mutations_only_transaction(self) -> Optional[Mutation]:
        """Returns a mutation to use for beginning a mutations-only transaction.
        Returns None if a mutation does not need to be included.

        :rtype: :class:`~google.cloud.spanner_v1.types.Mutation`
        :returns: A mutation to use for beginning a mutations-only transaction."""
        if not self._session.is_multiplexed:
            return None
        mutations: list[Mutation] = self._mutations
        insert_mutation: Mutation = None
        max_insert_values: int = -1
        for mut in mutations:
            if mut.insert:
                num_values = len(mut.insert.values)
                if num_values > max_insert_values:
                    insert_mutation = mut
                    max_insert_values = num_values
            else:
                return mut
        return insert_mutation

    def _update_for_execute_batch_dml_response_pb(
        self, response_pb: ExecuteBatchDmlResponse
    ) -> None:
        """Update the transaction for the given execute batch DML response.

        :type response_pb: :class:`~google.cloud.spanner_v1.types.ExecuteBatchDmlResponse`
        :param response_pb: The execute batch DML response to update the transaction with.
        """
        if len(response_pb.result_sets) > 0:
            self._update_for_result_set_pb(response_pb.result_sets[0])

    def __enter__(self):
        """Begin ``with`` block."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """End ``with`` block."""
        if exc_type is None:
            self.commit()
        else:
            self.rollback()


@dataclass
class BatchTransactionId:
    transaction_id: str
    session_id: str
    read_timestamp: Any


@dataclass
class DefaultTransactionOptions:
    isolation_level: str = TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED
    read_lock_mode: str = (
        TransactionOptions.ReadWrite.ReadLockMode.READ_LOCK_MODE_UNSPECIFIED
    )
    _defaultReadWriteTransactionOptions: Optional[TransactionOptions] = field(
        init=False, repr=False
    )

    def __post_init__(self):
        """Initialize _defaultReadWriteTransactionOptions automatically"""
        self._defaultReadWriteTransactionOptions = TransactionOptions(
            read_write=TransactionOptions.ReadWrite(read_lock_mode=self.read_lock_mode),
            isolation_level=self.isolation_level,
        )

    @property
    def default_read_write_transaction_options(self) -> TransactionOptions:
        """Public accessor for _defaultReadWriteTransactionOptions"""
        return self._defaultReadWriteTransactionOptions
