# -*- coding: utf-8 -*-

# Copyright 2020 Google LLC
#
# 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.
#

import proto  # type: ignore


from google.cloud.bigquery_storage_v1beta2.types import arrow
from google.cloud.bigquery_storage_v1beta2.types import avro
from google.cloud.bigquery_storage_v1beta2.types import protobuf
from google.cloud.bigquery_storage_v1beta2.types import stream
from google.cloud.bigquery_storage_v1beta2.types import table
from google.protobuf import timestamp_pb2 as timestamp  # type: ignore
from google.protobuf import wrappers_pb2 as wrappers  # type: ignore
from google.rpc import status_pb2 as status  # type: ignore


__protobuf__ = proto.module(
    package="google.cloud.bigquery.storage.v1beta2",
    manifest={
        "CreateReadSessionRequest",
        "ReadRowsRequest",
        "ThrottleState",
        "StreamStats",
        "ReadRowsResponse",
        "SplitReadStreamRequest",
        "SplitReadStreamResponse",
        "CreateWriteStreamRequest",
        "AppendRowsRequest",
        "AppendRowsResponse",
        "GetWriteStreamRequest",
        "BatchCommitWriteStreamsRequest",
        "BatchCommitWriteStreamsResponse",
        "FinalizeWriteStreamRequest",
        "FinalizeWriteStreamResponse",
        "FlushRowsRequest",
        "FlushRowsResponse",
        "StorageError",
    },
)


class CreateReadSessionRequest(proto.Message):
    r"""Request message for ``CreateReadSession``.

    Attributes:
        parent (str):
            Required. The request project that owns the session, in the
            form of ``projects/{project_id}``.
        read_session (google.cloud.bigquery_storage_v1beta2.types.ReadSession):
            Required. Session to be created.
        max_stream_count (int):
            Max initial number of streams. If unset or
            zero, the server will provide a value of streams
            so as to produce reasonable throughput. Must be
            non-negative. The number of streams may be lower
            than the requested number, depending on the
            amount parallelism that is reasonable for the
            table. Error will be returned if the max count
            is greater than the current system max limit of
            1,000.

            Streams must be read starting from offset 0.
    """

    parent = proto.Field(proto.STRING, number=1)

    read_session = proto.Field(proto.MESSAGE, number=2, message=stream.ReadSession,)

    max_stream_count = proto.Field(proto.INT32, number=3)


class ReadRowsRequest(proto.Message):
    r"""Request message for ``ReadRows``.

    Attributes:
        read_stream (str):
            Required. Stream to read rows from.
        offset (int):
            The offset requested must be less than the
            last row read from Read. Requesting a larger
            offset is undefined. If not specified, start
            reading from offset zero.
    """

    read_stream = proto.Field(proto.STRING, number=1)

    offset = proto.Field(proto.INT64, number=2)


class ThrottleState(proto.Message):
    r"""Information on if the current connection is being throttled.

    Attributes:
        throttle_percent (int):
            How much this connection is being throttled.
            Zero means no throttling, 100 means fully
            throttled.
    """

    throttle_percent = proto.Field(proto.INT32, number=1)


class StreamStats(proto.Message):
    r"""Estimated stream statistics for a given Stream.

    Attributes:
        progress (google.cloud.bigquery_storage_v1beta2.types.StreamStats.Progress):
            Represents the progress of the current
            stream.
    """

    class Progress(proto.Message):
        r"""

        Attributes:
            at_response_start (float):
                The fraction of rows assigned to the stream that have been
                processed by the server so far, not including the rows in
                the current response message.

                This value, along with ``at_response_end``, can be used to
                interpolate the progress made as the rows in the message are
                being processed using the following formula:
                ``at_response_start + (at_response_end - at_response_start) * rows_processed_from_response / rows_in_response``.

                Note that if a filter is provided, the ``at_response_end``
                value of the previous response may not necessarily be equal
                to the ``at_response_start`` value of the current response.
            at_response_end (float):
                Similar to ``at_response_start``, except that this value
                includes the rows in the current response.
        """

        at_response_start = proto.Field(proto.DOUBLE, number=1)

        at_response_end = proto.Field(proto.DOUBLE, number=2)

    progress = proto.Field(proto.MESSAGE, number=2, message=Progress,)


class ReadRowsResponse(proto.Message):
    r"""Response from calling ``ReadRows`` may include row data, progress
    and throttling information.

    Attributes:
        avro_rows (google.cloud.bigquery_storage_v1beta2.types.AvroRows):
            Serialized row data in AVRO format.
        arrow_record_batch (google.cloud.bigquery_storage_v1beta2.types.ArrowRecordBatch):
            Serialized row data in Arrow RecordBatch
            format.
        row_count (int):
            Number of serialized rows in the rows block.
        stats (google.cloud.bigquery_storage_v1beta2.types.StreamStats):
            Statistics for the stream.
        throttle_state (google.cloud.bigquery_storage_v1beta2.types.ThrottleState):
            Throttling state. If unset, the latest
            response still describes the current throttling
            status.
    """

    avro_rows = proto.Field(
        proto.MESSAGE, number=3, oneof="rows", message=avro.AvroRows,
    )

    arrow_record_batch = proto.Field(
        proto.MESSAGE, number=4, oneof="rows", message=arrow.ArrowRecordBatch,
    )

    row_count = proto.Field(proto.INT64, number=6)

    stats = proto.Field(proto.MESSAGE, number=2, message="StreamStats",)

    throttle_state = proto.Field(proto.MESSAGE, number=5, message="ThrottleState",)


class SplitReadStreamRequest(proto.Message):
    r"""Request message for ``SplitReadStream``.

    Attributes:
        name (str):
            Required. Name of the stream to split.
        fraction (float):
            A value in the range (0.0, 1.0) that
            specifies the fractional point at which the
            original stream should be split. The actual
            split point is evaluated on pre-filtered rows,
            so if a filter is provided, then there is no
            guarantee that the division of the rows between
            the new child streams will be proportional to
            this fractional value. Additionally, because the
            server-side unit for assigning data is
            collections of rows, this fraction will always
            map to a data storage boundary on the server
            side.
    """

    name = proto.Field(proto.STRING, number=1)

    fraction = proto.Field(proto.DOUBLE, number=2)


class SplitReadStreamResponse(proto.Message):
    r"""

    Attributes:
        primary_stream (google.cloud.bigquery_storage_v1beta2.types.ReadStream):
            Primary stream, which contains the beginning portion of
            \|original_stream|. An empty value indicates that the
            original stream can no longer be split.
        remainder_stream (google.cloud.bigquery_storage_v1beta2.types.ReadStream):
            Remainder stream, which contains the tail of
            \|original_stream|. An empty value indicates that the
            original stream can no longer be split.
    """

    primary_stream = proto.Field(proto.MESSAGE, number=1, message=stream.ReadStream,)

    remainder_stream = proto.Field(proto.MESSAGE, number=2, message=stream.ReadStream,)


class CreateWriteStreamRequest(proto.Message):
    r"""Request message for ``CreateWriteStream``.

    Attributes:
        parent (str):
            Required. Reference to the table to which the stream
            belongs, in the format of
            ``projects/{project}/datasets/{dataset}/tables/{table}``.
        write_stream (google.cloud.bigquery_storage_v1beta2.types.WriteStream):
            Required. Stream to be created.
    """

    parent = proto.Field(proto.STRING, number=1)

    write_stream = proto.Field(proto.MESSAGE, number=2, message=stream.WriteStream,)


class AppendRowsRequest(proto.Message):
    r"""Request message for ``AppendRows``.

    Attributes:
        write_stream (str):
            Required. The stream that is the target of the append
            operation. This value must be specified for the initial
            request. If subsequent requests specify the stream name, it
            must equal to the value provided in the first request. To
            write to the \_default stream, populate this field with a
            string in the format
            ``projects/{project}/datasets/{dataset}/tables/{table}/_default``.
        offset (google.protobuf.wrappers_pb2.Int64Value):
            If present, the write is only performed if the next append
            offset is same as the provided value. If not present, the
            write is performed at the current end of stream. Specifying
            a value for this field is not allowed when calling
            AppendRows for the '_default' stream.
        proto_rows (google.cloud.bigquery_storage_v1beta2.types.AppendRowsRequest.ProtoData):
            Rows in proto format.
        trace_id (str):
            Id set by client to annotate its identity.
            Only initial request setting is respected.
    """

    class ProtoData(proto.Message):
        r"""Proto schema and data.

        Attributes:
            writer_schema (google.cloud.bigquery_storage_v1beta2.types.ProtoSchema):
                Proto schema used to serialize the data.
            rows (google.cloud.bigquery_storage_v1beta2.types.ProtoRows):
                Serialized row data in protobuf message
                format.
        """

        writer_schema = proto.Field(
            proto.MESSAGE, number=1, message=protobuf.ProtoSchema,
        )

        rows = proto.Field(proto.MESSAGE, number=2, message=protobuf.ProtoRows,)

    write_stream = proto.Field(proto.STRING, number=1)

    offset = proto.Field(proto.MESSAGE, number=2, message=wrappers.Int64Value,)

    proto_rows = proto.Field(proto.MESSAGE, number=4, oneof="rows", message=ProtoData,)

    trace_id = proto.Field(proto.STRING, number=6)


class AppendRowsResponse(proto.Message):
    r"""Response message for ``AppendRows``.

    Attributes:
        append_result (google.cloud.bigquery_storage_v1beta2.types.AppendRowsResponse.AppendResult):
            Result if the append is successful.
        error (google.rpc.status_pb2.Status):
            Error in case of request failed. If set, it means rows are
            not accepted into the system. Users can retry or continue
            with other requests within the same connection.
            ALREADY_EXISTS: happens when offset is specified, it means
            the entire request is already appended, it is safe to ignore
            this error. OUT_OF_RANGE: happens when offset is specified,
            it means the specified offset is beyond the end of the
            stream. INVALID_ARGUMENT: error caused by malformed request
            or data. RESOURCE_EXHAUSTED: request rejected due to
            throttling. Only happens when append without offset.
            ABORTED: request processing is aborted because of prior
            failures, request can be retried if previous failure is
            fixed. INTERNAL: server side errors that can be retried.
        updated_schema (google.cloud.bigquery_storage_v1beta2.types.TableSchema):
            If backend detects a schema update, pass it
            to user so that user can use it to input new
            type of message. It will be empty when there is
            no schema updates.
    """

    class AppendResult(proto.Message):
        r"""A success append result.

        Attributes:
            offset (google.protobuf.wrappers_pb2.Int64Value):
                The row offset at which the last append
                occurred. The offset will not be set if
                appending using default streams.
        """

        offset = proto.Field(proto.MESSAGE, number=1, message=wrappers.Int64Value,)

    append_result = proto.Field(
        proto.MESSAGE, number=1, oneof="response", message=AppendResult,
    )

    error = proto.Field(
        proto.MESSAGE, number=2, oneof="response", message=status.Status,
    )

    updated_schema = proto.Field(proto.MESSAGE, number=3, message=table.TableSchema,)


class GetWriteStreamRequest(proto.Message):
    r"""Request message for ``GetWriteStreamRequest``.

    Attributes:
        name (str):
            Required. Name of the stream to get, in the form of
            ``projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}``.
    """

    name = proto.Field(proto.STRING, number=1)


class BatchCommitWriteStreamsRequest(proto.Message):
    r"""Request message for ``BatchCommitWriteStreams``.

    Attributes:
        parent (str):
            Required. Parent table that all the streams should belong
            to, in the form of
            ``projects/{project}/datasets/{dataset}/tables/{table}``.
        write_streams (Sequence[str]):
            Required. The group of streams that will be
            committed atomically.
    """

    parent = proto.Field(proto.STRING, number=1)

    write_streams = proto.RepeatedField(proto.STRING, number=2)


class BatchCommitWriteStreamsResponse(proto.Message):
    r"""Response message for ``BatchCommitWriteStreams``.

    Attributes:
        commit_time (google.protobuf.timestamp_pb2.Timestamp):
            The time at which streams were committed in
            microseconds granularity. This field will only
            exist when there is no stream errors.
        stream_errors (Sequence[google.cloud.bigquery_storage_v1beta2.types.StorageError]):
            Stream level error if commit failed. Only
            streams with error will be in the list.
    """

    commit_time = proto.Field(proto.MESSAGE, number=1, message=timestamp.Timestamp,)

    stream_errors = proto.RepeatedField(
        proto.MESSAGE, number=2, message="StorageError",
    )


class FinalizeWriteStreamRequest(proto.Message):
    r"""Request message for invoking ``FinalizeWriteStream``.

    Attributes:
        name (str):
            Required. Name of the stream to finalize, in the form of
            ``projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}``.
    """

    name = proto.Field(proto.STRING, number=1)


class FinalizeWriteStreamResponse(proto.Message):
    r"""Response message for ``FinalizeWriteStream``.

    Attributes:
        row_count (int):
            Number of rows in the finalized stream.
    """

    row_count = proto.Field(proto.INT64, number=1)


class FlushRowsRequest(proto.Message):
    r"""Request message for ``FlushRows``.

    Attributes:
        write_stream (str):
            Required. The stream that is the target of
            the flush operation.
        offset (google.protobuf.wrappers_pb2.Int64Value):
            Ending offset of the flush operation. Rows
            before this offset(including this offset) will
            be flushed.
    """

    write_stream = proto.Field(proto.STRING, number=1)

    offset = proto.Field(proto.MESSAGE, number=2, message=wrappers.Int64Value,)


class FlushRowsResponse(proto.Message):
    r"""Respond message for ``FlushRows``.

    Attributes:
        offset (int):
            The rows before this offset (including this
            offset) are flushed.
    """

    offset = proto.Field(proto.INT64, number=1)


class StorageError(proto.Message):
    r"""Structured custom BigQuery Storage error message. The error
    can be attached as error details in the returned rpc Status.
    User can use the info to process errors in a structural way,
    rather than having to parse error messages.

    Attributes:
        code (google.cloud.bigquery_storage_v1beta2.types.StorageError.StorageErrorCode):
            BigQuery Storage specific error code.
        entity (str):
            Name of the failed entity.
        error_message (str):
            Message that describes the error.
    """

    class StorageErrorCode(proto.Enum):
        r"""Error code for ``StorageError``."""
        STORAGE_ERROR_CODE_UNSPECIFIED = 0
        TABLE_NOT_FOUND = 1
        STREAM_ALREADY_COMMITTED = 2
        STREAM_NOT_FOUND = 3
        INVALID_STREAM_TYPE = 4
        INVALID_STREAM_STATE = 5

    code = proto.Field(proto.ENUM, number=1, enum=StorageErrorCode,)

    entity = proto.Field(proto.STRING, number=2)

    error_message = proto.Field(proto.STRING, number=3)


__all__ = tuple(sorted(__protobuf__.manifest))
