diff --git a/docs/bigquery_storage_v1beta2/big_query_read.rst b/docs/bigquery_storage_v1beta2/big_query_read.rst new file mode 100644 index 00000000..10b3714b --- /dev/null +++ b/docs/bigquery_storage_v1beta2/big_query_read.rst @@ -0,0 +1,6 @@ +BigQueryRead +------------------------------ + +.. automodule:: google.cloud.bigquery_storage_v1beta2.services.big_query_read + :members: + :inherited-members: diff --git a/docs/bigquery_storage_v1beta2/big_query_write.rst b/docs/bigquery_storage_v1beta2/big_query_write.rst new file mode 100644 index 00000000..c685994d --- /dev/null +++ b/docs/bigquery_storage_v1beta2/big_query_write.rst @@ -0,0 +1,6 @@ +BigQueryWrite +------------------------------- + +.. automodule:: google.cloud.bigquery_storage_v1beta2.services.big_query_write + :members: + :inherited-members: diff --git a/docs/bigquery_storage_v1beta2/services.rst b/docs/bigquery_storage_v1beta2/services.rst new file mode 100644 index 00000000..67dae0ab --- /dev/null +++ b/docs/bigquery_storage_v1beta2/services.rst @@ -0,0 +1,7 @@ +Services for Google Cloud Bigquery Storage v1beta2 API +====================================================== +.. toctree:: + :maxdepth: 2 + + big_query_read + big_query_write diff --git a/docs/bigquery_storage_v1beta2/types.rst b/docs/bigquery_storage_v1beta2/types.rst new file mode 100644 index 00000000..995806da --- /dev/null +++ b/docs/bigquery_storage_v1beta2/types.rst @@ -0,0 +1,7 @@ +Types for Google Cloud Bigquery Storage v1beta2 API +=================================================== + +.. automodule:: google.cloud.bigquery_storage_v1beta2.types + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/index.rst b/docs/index.rst index 618e74d7..6892b30c 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -21,6 +21,8 @@ API Reference bigquery_storage_v1/library bigquery_storage_v1/services bigquery_storage_v1/types + bigquery_storage_v1beta2/services + bigquery_storage_v1beta2/types Migration Guide diff --git a/google/cloud/bigquery_storage_v1beta2/proto/arrow.proto b/google/cloud/bigquery_storage_v1beta2/proto/arrow.proto new file mode 100644 index 00000000..bc2e4eb1 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/proto/arrow.proto @@ -0,0 +1,57 @@ +// 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. + +syntax = "proto3"; + +package google.cloud.bigquery.storage.v1beta2; + +option go_package = "google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1beta2;storage"; +option java_multiple_files = true; +option java_outer_classname = "ArrowProto"; +option java_package = "com.google.cloud.bigquery.storage.v1beta2"; + +// Arrow schema as specified in +// https://arrow.apache.org/docs/python/api/datatypes.html +// and serialized to bytes using IPC: +// https://arrow.apache.org/docs/format/Columnar.html#serialization-and-interprocess-communication-ipc +// +// See code samples on how this message can be deserialized. +message ArrowSchema { + // IPC serialized Arrow schema. + bytes serialized_schema = 1; +} + +// Arrow RecordBatch. +message ArrowRecordBatch { + // IPC-serialized Arrow RecordBatch. + bytes serialized_record_batch = 1; +} + +// Contains options specific to Arrow Serialization. +message ArrowSerializationOptions { + // The IPC format to use when serializing Arrow streams. + enum Format { + // If unspecied the IPC format as of 0.15 release will be used. + FORMAT_UNSPECIFIED = 0; + + // Use the legacy IPC message format as of Apache Arrow Release 0.14. + ARROW_0_14 = 1; + + // Use the message format as of Apache Arrow Release 0.15. + ARROW_0_15 = 2; + } + + // The Arrow IPC format to use. + Format format = 1; +} diff --git a/google/cloud/bigquery_storage_v1beta2/proto/avro.proto b/google/cloud/bigquery_storage_v1beta2/proto/avro.proto new file mode 100644 index 00000000..109ec86a --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/proto/avro.proto @@ -0,0 +1,35 @@ +// 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. + +syntax = "proto3"; + +package google.cloud.bigquery.storage.v1beta2; + +option go_package = "google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1beta2;storage"; +option java_multiple_files = true; +option java_outer_classname = "AvroProto"; +option java_package = "com.google.cloud.bigquery.storage.v1beta2"; + +// Avro schema. +message AvroSchema { + // Json serialized schema, as described at + // https://avro.apache.org/docs/1.8.1/spec.html. + string schema = 1; +} + +// Avro rows. +message AvroRows { + // Binary serialized rows in a block. + bytes serialized_binary_rows = 1; +} diff --git a/google/cloud/bigquery_storage_v1beta2/proto/protobuf.proto b/google/cloud/bigquery_storage_v1beta2/proto/protobuf.proto new file mode 100644 index 00000000..741e7d11 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/proto/protobuf.proto @@ -0,0 +1,41 @@ +// 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. + +syntax = "proto3"; + +package google.cloud.bigquery.storage.v1beta2; + +import "google/protobuf/descriptor.proto"; + +option go_package = "google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1beta2;storage"; +option java_multiple_files = true; +option java_outer_classname = "ProtoBufProto"; +option java_package = "com.google.cloud.bigquery.storage.v1beta2"; + +// Protobuf schema is an API presentation the proto buffer schema. +message ProtoSchema { + // Descriptor for input message. The descriptor has to be self contained, + // including all the nested types, excepted for proto buffer well known types + // (https://developers.google.com/protocol-buffers/docs/reference/google.protobuf). + google.protobuf.DescriptorProto proto_descriptor = 1; +} + +// Protobuf rows. +message ProtoRows { + // A sequence of rows serialized as a Protocol Buffer. + // + // See https://developers.google.com/protocol-buffers/docs/overview for more + // information on deserializing this field. + repeated bytes serialized_rows = 1; +} diff --git a/google/cloud/bigquery_storage_v1beta2/proto/storage.proto b/google/cloud/bigquery_storage_v1beta2/proto/storage.proto new file mode 100644 index 00000000..5538e29f --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/proto/storage.proto @@ -0,0 +1,538 @@ +// 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. + +syntax = "proto3"; + +package google.cloud.bigquery.storage.v1beta2; + +import "google/api/annotations.proto"; +import "google/api/client.proto"; +import "google/api/field_behavior.proto"; +import "google/api/resource.proto"; +import "google/cloud/bigquery/storage/v1beta2/arrow.proto"; +import "google/cloud/bigquery/storage/v1beta2/avro.proto"; +import "google/cloud/bigquery/storage/v1beta2/protobuf.proto"; +import "google/cloud/bigquery/storage/v1beta2/stream.proto"; +import "google/cloud/bigquery/storage/v1beta2/table.proto"; +import "google/protobuf/timestamp.proto"; +import "google/protobuf/wrappers.proto"; +import "google/rpc/status.proto"; + +option go_package = "google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1beta2;storage"; +option java_multiple_files = true; +option java_outer_classname = "StorageProto"; +option java_package = "com.google.cloud.bigquery.storage.v1beta2"; + +// BigQuery Read API. +// +// The Read API can be used to read data from BigQuery. +// +// New code should use the v1 Read API going forward, if they don't use Write +// API at the same time. +service BigQueryRead { + option (google.api.default_host) = "bigquerystorage.googleapis.com"; + option (google.api.oauth_scopes) = + "https://www.googleapis.com/auth/bigquery," + "https://www.googleapis.com/auth/bigquery.readonly," + "https://www.googleapis.com/auth/cloud-platform"; + + // Creates a new read session. A read session divides the contents of a + // BigQuery table into one or more streams, which can then be used to read + // data from the table. The read session also specifies properties of the + // data to be read, such as a list of columns or a push-down filter describing + // the rows to be returned. + // + // A particular row can be read by at most one stream. When the caller has + // reached the end of each stream in the session, then all the data in the + // table has been read. + // + // Data is assigned to each stream such that roughly the same number of + // rows can be read from each stream. Because the server-side unit for + // assigning data is collections of rows, the API does not guarantee that + // each stream will return the same number or rows. Additionally, the + // limits are enforced based on the number of pre-filtered rows, so some + // filters can lead to lopsided assignments. + // + // Read sessions automatically expire 24 hours after they are created and do + // not require manual clean-up by the caller. + rpc CreateReadSession(CreateReadSessionRequest) returns (ReadSession) { + option (google.api.http) = { + post: "/v1beta2/{read_session.table=projects/*/datasets/*/tables/*}" + body: "*" + }; + option (google.api.method_signature) = + "parent,read_session,max_stream_count"; + } + + // Reads rows from the stream in the format prescribed by the ReadSession. + // Each response contains one or more table rows, up to a maximum of 100 MiB + // per response; read requests which attempt to read individual rows larger + // than 100 MiB will fail. + // + // Each request also returns a set of stream statistics reflecting the current + // state of the stream. + rpc ReadRows(ReadRowsRequest) returns (stream ReadRowsResponse) { + option (google.api.http) = { + get: "/v1beta2/{read_stream=projects/*/locations/*/sessions/*/streams/*}" + }; + option (google.api.method_signature) = "read_stream,offset"; + } + + // Splits a given `ReadStream` into two `ReadStream` objects. These + // `ReadStream` objects are referred to as the primary and the residual + // streams of the split. The original `ReadStream` can still be read from in + // the same manner as before. Both of the returned `ReadStream` objects can + // also be read from, and the rows returned by both child streams will be + // the same as the rows read from the original stream. + // + // Moreover, the two child streams will be allocated back-to-back in the + // original `ReadStream`. Concretely, it is guaranteed that for streams + // original, primary, and residual, that original[0-j] = primary[0-j] and + // original[j-n] = residual[0-m] once the streams have been read to + // completion. + rpc SplitReadStream(SplitReadStreamRequest) + returns (SplitReadStreamResponse) { + option (google.api.http) = { + get: "/v1beta2/{name=projects/*/locations/*/sessions/*/streams/*}" + }; + } +} + +// BigQuery Write API. +// +// The Write API can be used to write data to BigQuery. +service BigQueryWrite { + option (google.api.default_host) = "bigquerystorage.googleapis.com"; + option (google.api.oauth_scopes) = + "https://www.googleapis.com/auth/bigquery," + "https://www.googleapis.com/auth/bigquery.insertdata," + "https://www.googleapis.com/auth/cloud-platform"; + + // Creates a write stream to the given table. + // Additionally, every table has a special COMMITTED stream named '_default' + // to which data can be written. This stream doesn't need to be created using + // CreateWriteStream. It is a stream that can be used simultaneously by any + // number of clients. Data written to this stream is considered committed as + // soon as an acknowledgement is received. + rpc CreateWriteStream(CreateWriteStreamRequest) returns (WriteStream) { + option (google.api.http) = { + post: "/v1beta2/{parent=projects/*/datasets/*/tables/*}" + body: "write_stream" + }; + option (google.api.method_signature) = "parent,write_stream"; + } + + // Appends data to the given stream. + // + // If `offset` is specified, the `offset` is checked against the end of + // stream. The server returns `OUT_OF_RANGE` in `AppendRowsResponse` if an + // attempt is made to append to an offset beyond the current end of the stream + // or `ALREADY_EXISTS` if user provids an `offset` that has already been + // written to. User can retry with adjusted offset within the same RPC + // stream. If `offset` is not specified, append happens at the end of the + // stream. + // + // The response contains the offset at which the append happened. Responses + // are received in the same order in which requests are sent. There will be + // one response for each successful request. If the `offset` is not set in + // response, it means append didn't happen due to some errors. If one request + // fails, all the subsequent requests will also fail until a success request + // is made again. + // + // If the stream is of `PENDING` type, data will only be available for read + // operations after the stream is committed. + rpc AppendRows(stream AppendRowsRequest) returns (stream AppendRowsResponse) { + option (google.api.http) = { + post: "/v1beta2/{write_stream=projects/*/datasets/*/tables/*/streams/*}" + body: "*" + }; + option (google.api.method_signature) = "write_stream"; + } + + // Gets a write stream. + rpc GetWriteStream(GetWriteStreamRequest) returns (WriteStream) { + option (google.api.http) = { + post: "/v1beta2/{name=projects/*/datasets/*/tables/*/streams/*}" + body: "*" + }; + option (google.api.method_signature) = "name"; + } + + // Finalize a write stream so that no new data can be appended to the + // stream. Finalize is not supported on the '_default' stream. + rpc FinalizeWriteStream(FinalizeWriteStreamRequest) + returns (FinalizeWriteStreamResponse) { + option (google.api.http) = { + post: "/v1beta2/{name=projects/*/datasets/*/tables/*/streams/*}" + body: "*" + }; + option (google.api.method_signature) = "name"; + } + + // Atomically commits a group of `PENDING` streams that belong to the same + // `parent` table. + // Streams must be finalized before commit and cannot be committed multiple + // times. Once a stream is committed, data in the stream becomes available + // for read operations. + rpc BatchCommitWriteStreams(BatchCommitWriteStreamsRequest) + returns (BatchCommitWriteStreamsResponse) { + option (google.api.http) = { + get: "/v1beta2/{parent=projects/*/datasets/*/tables/*}" + }; + option (google.api.method_signature) = "parent"; + } + + // Flushes rows to a BUFFERED stream. + // If users are appending rows to BUFFERED stream, flush operation is + // required in order for the rows to become available for reading. A + // Flush operation flushes up to any previously flushed offset in a BUFFERED + // stream, to the offset specified in the request. + // Flush is not supported on the _default stream, since it is not BUFFERED. + rpc FlushRows(FlushRowsRequest) returns (FlushRowsResponse) { + option (google.api.http) = { + post: "/v1beta2/{write_stream=projects/*/datasets/*/tables/*/streams/*}" + body: "*" + }; + option (google.api.method_signature) = "write_stream"; + } +} + +// Request message for `CreateReadSession`. +message CreateReadSessionRequest { + // Required. The request project that owns the session, in the form of + // `projects/{project_id}`. + string parent = 1 [ + (google.api.field_behavior) = REQUIRED, + (google.api.resource_reference) = { + type: "cloudresourcemanager.googleapis.com/Project" + } + ]; + + // Required. Session to be created. + ReadSession read_session = 2 [(google.api.field_behavior) = REQUIRED]; + + // 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. + int32 max_stream_count = 3; +} + +// Request message for `ReadRows`. +message ReadRowsRequest { + // Required. Stream to read rows from. + string read_stream = 1 [ + (google.api.field_behavior) = REQUIRED, + (google.api.resource_reference) = { + type: "bigquerystorage.googleapis.com/ReadStream" + } + ]; + + // 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. + int64 offset = 2; +} + +// Information on if the current connection is being throttled. +message ThrottleState { + // How much this connection is being throttled. Zero means no throttling, + // 100 means fully throttled. + int32 throttle_percent = 1; +} + +// Estimated stream statistics for a given Stream. +message StreamStats { + message Progress { + // 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. + double at_response_start = 1; + + // Similar to `at_response_start`, except that this value includes the + // rows in the current response. + double at_response_end = 2; + } + + // Represents the progress of the current stream. + Progress progress = 2; +} + +// Response from calling `ReadRows` may include row data, progress and +// throttling information. +message ReadRowsResponse { + // Row data is returned in format specified during session creation. + oneof rows { + // Serialized row data in AVRO format. + AvroRows avro_rows = 3; + + // Serialized row data in Arrow RecordBatch format. + ArrowRecordBatch arrow_record_batch = 4; + } + + // Number of serialized rows in the rows block. + int64 row_count = 6; + + // Statistics for the stream. + StreamStats stats = 2; + + // Throttling state. If unset, the latest response still describes + // the current throttling status. + ThrottleState throttle_state = 5; +} + +// Request message for `SplitReadStream`. +message SplitReadStreamRequest { + // Required. Name of the stream to split. + string name = 1 [ + (google.api.field_behavior) = REQUIRED, + (google.api.resource_reference) = { + type: "bigquerystorage.googleapis.com/ReadStream" + } + ]; + + // 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. + double fraction = 2; +} + +message SplitReadStreamResponse { + // Primary stream, which contains the beginning portion of + // |original_stream|. An empty value indicates that the original stream can no + // longer be split. + ReadStream primary_stream = 1; + + // Remainder stream, which contains the tail of |original_stream|. An empty + // value indicates that the original stream can no longer be split. + ReadStream remainder_stream = 2; +} + +// Request message for `CreateWriteStream`. +message CreateWriteStreamRequest { + // Required. Reference to the table to which the stream belongs, in the format + // of `projects/{project}/datasets/{dataset}/tables/{table}`. + string parent = 1 [ + (google.api.field_behavior) = REQUIRED, + (google.api.resource_reference) = { type: "bigquery.googleapis.com/Table" } + ]; + + // Required. Stream to be created. + WriteStream write_stream = 2 [(google.api.field_behavior) = REQUIRED]; +} + +// Request message for `AppendRows`. +message AppendRowsRequest { + // Proto schema and data. + message ProtoData { + // Proto schema used to serialize the data. + ProtoSchema writer_schema = 1; + + // Serialized row data in protobuf message format. + ProtoRows rows = 2; + } + + // 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`. + string write_stream = 1 [ + (google.api.field_behavior) = REQUIRED, + (google.api.resource_reference) = { + type: "bigquerystorage.googleapis.com/WriteStream" + } + ]; + + // 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. + google.protobuf.Int64Value offset = 2; + + // Input rows. The `writer_schema` field must be specified at the initial + // request and currently, it will be ignored if specified in following + // requests. Following requests must have data in the same format as the + // initial request. + oneof rows { + // Rows in proto format. + ProtoData proto_rows = 4; + } + + // Id set by client to annotate its identity. Only initial request setting is + // respected. + string trace_id = 6; +} + +// Response message for `AppendRows`. +message AppendRowsResponse { + // A success append result. + message AppendResult { + // The row offset at which the last append occurred. The offset will not be + // set if appending using default streams. + google.protobuf.Int64Value offset = 1; + } + + oneof response { + // Result if the append is successful. + AppendResult append_result = 1; + + // 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. + google.rpc.Status error = 2; + } + + // 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. + TableSchema updated_schema = 3; +} + +// Request message for `GetWriteStreamRequest`. +message GetWriteStreamRequest { + // Required. Name of the stream to get, in the form of + // `projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}`. + string name = 1 [ + (google.api.field_behavior) = REQUIRED, + (google.api.resource_reference) = { + type: "bigquerystorage.googleapis.com/WriteStream" + } + ]; +} + +// Request message for `BatchCommitWriteStreams`. +message BatchCommitWriteStreamsRequest { + // Required. Parent table that all the streams should belong to, in the form + // of `projects/{project}/datasets/{dataset}/tables/{table}`. + string parent = 1 [(google.api.field_behavior) = REQUIRED]; + + // Required. The group of streams that will be committed atomically. + repeated string write_streams = 2 [(google.api.field_behavior) = REQUIRED]; +} + +// Response message for `BatchCommitWriteStreams`. +message BatchCommitWriteStreamsResponse { + // The time at which streams were committed in microseconds granularity. + // This field will only exist when there is no stream errors. + google.protobuf.Timestamp commit_time = 1; + + // Stream level error if commit failed. Only streams with error will be in + // the list. + repeated StorageError stream_errors = 2; +} + +// Request message for invoking `FinalizeWriteStream`. +message FinalizeWriteStreamRequest { + // Required. Name of the stream to finalize, in the form of + // `projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}`. + string name = 1 [ + (google.api.field_behavior) = REQUIRED, + (google.api.resource_reference) = { + type: "bigquerystorage.googleapis.com/WriteStream" + } + ]; +} + +// Response message for `FinalizeWriteStream`. +message FinalizeWriteStreamResponse { + // Number of rows in the finalized stream. + int64 row_count = 1; +} + +// Request message for `FlushRows`. +message FlushRowsRequest { + // Required. The stream that is the target of the flush operation. + string write_stream = 1 [ + (google.api.field_behavior) = REQUIRED, + (google.api.resource_reference) = { + type: "bigquerystorage.googleapis.com/WriteStream" + } + ]; + + // Ending offset of the flush operation. Rows before this offset(including + // this offset) will be flushed. + google.protobuf.Int64Value offset = 2; +} + +// Respond message for `FlushRows`. +message FlushRowsResponse { + // The rows before this offset (including this offset) are flushed. + int64 offset = 1; +} + +// 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. +message StorageError { + // Error code for `StorageError`. + enum StorageErrorCode { + // Default error. + STORAGE_ERROR_CODE_UNSPECIFIED = 0; + + // Table is not found in the system. + TABLE_NOT_FOUND = 1; + + // Stream is already committed. + STREAM_ALREADY_COMMITTED = 2; + + // Stream is not found. + STREAM_NOT_FOUND = 3; + + // Invalid Stream type. + // For example, you try to commit a stream that is not pending. + INVALID_STREAM_TYPE = 4; + + // Invalid Stream state. + // For example, you try to commit a stream that is not fianlized or is + // garbaged. + INVALID_STREAM_STATE = 5; + } + + // BigQuery Storage specific error code. + StorageErrorCode code = 1; + + // Name of the failed entity. + string entity = 2; + + // Message that describes the error. + string error_message = 3; +} diff --git a/google/cloud/bigquery_storage_v1beta2/proto/stream.proto b/google/cloud/bigquery_storage_v1beta2/proto/stream.proto new file mode 100644 index 00000000..2b0a58c9 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/proto/stream.proto @@ -0,0 +1,189 @@ +// 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. + +syntax = "proto3"; + +package google.cloud.bigquery.storage.v1beta2; + +import "google/api/field_behavior.proto"; +import "google/api/resource.proto"; +import "google/cloud/bigquery/storage/v1beta2/arrow.proto"; +import "google/cloud/bigquery/storage/v1beta2/avro.proto"; +import "google/cloud/bigquery/storage/v1beta2/table.proto"; +import "google/protobuf/timestamp.proto"; + +option go_package = "google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1beta2;storage"; +option java_multiple_files = true; +option java_outer_classname = "StreamProto"; +option java_package = "com.google.cloud.bigquery.storage.v1beta2"; +option (google.api.resource_definition) = { + type: "bigquery.googleapis.com/Table" + pattern: "projects/{project}/datasets/{dataset}/tables/{table}" +}; + +// Data format for input or output data. +enum DataFormat { + DATA_FORMAT_UNSPECIFIED = 0; + + // Avro is a standard open source row based file format. + // See https://avro.apache.org/ for more details. + AVRO = 1; + + // Arrow is a standard open source column-based message format. + // See https://arrow.apache.org/ for more details. + ARROW = 2; +} + +// Information about the ReadSession. +message ReadSession { + option (google.api.resource) = { + type: "bigquerystorage.googleapis.com/ReadSession" + pattern: "projects/{project}/locations/{location}/sessions/{session}" + }; + + // Additional attributes when reading a table. + message TableModifiers { + // The snapshot time of the table. If not set, interpreted as now. + google.protobuf.Timestamp snapshot_time = 1; + } + + // Options dictating how we read a table. + message TableReadOptions { + // Names of the fields in the table that should be read. If empty, all + // fields will be read. If the specified field is a nested field, all + // the sub-fields in the field will be selected. The output field order is + // unrelated to the order of fields in selected_fields. + repeated string selected_fields = 1; + + // SQL text filtering statement, similar to a WHERE clause in a query. + // Aggregates are not supported. + // + // Examples: "int_field > 5" + // "date_field = CAST('2014-9-27' as DATE)" + // "nullable_field is not NULL" + // "st_equals(geo_field, st_geofromtext("POINT(2, 2)"))" + // "numeric_field BETWEEN 1.0 AND 5.0" + string row_restriction = 2; + + // Optional. Options specific to the Apache Arrow output format. + ArrowSerializationOptions arrow_serialization_options = 3 [(google.api.field_behavior) = OPTIONAL]; + } + + // Output only. Unique identifier for the session, in the form + // `projects/{project_id}/locations/{location}/sessions/{session_id}`. + string name = 1 [(google.api.field_behavior) = OUTPUT_ONLY]; + + // Output only. Time at which the session becomes invalid. After this time, subsequent + // requests to read this Session will return errors. The expire_time is + // automatically assigned and currently cannot be specified or updated. + google.protobuf.Timestamp expire_time = 2 [(google.api.field_behavior) = OUTPUT_ONLY]; + + // Immutable. Data format of the output data. + DataFormat data_format = 3 [(google.api.field_behavior) = IMMUTABLE]; + + // The schema for the read. If read_options.selected_fields is set, the + // schema may be different from the table schema as it will only contain + // the selected fields. + oneof schema { + // Output only. Avro schema. + AvroSchema avro_schema = 4 [(google.api.field_behavior) = OUTPUT_ONLY]; + + // Output only. Arrow schema. + ArrowSchema arrow_schema = 5 [(google.api.field_behavior) = OUTPUT_ONLY]; + } + + // Immutable. Table that this ReadSession is reading from, in the form + // `projects/{project_id}/datasets/{dataset_id}/tables/{table_id} + string table = 6 [ + (google.api.field_behavior) = IMMUTABLE, + (google.api.resource_reference) = { + type: "bigquery.googleapis.com/Table" + } + ]; + + // Optional. Any modifiers which are applied when reading from the specified table. + TableModifiers table_modifiers = 7 [(google.api.field_behavior) = OPTIONAL]; + + // Optional. Read options for this session (e.g. column selection, filters). + TableReadOptions read_options = 8 [(google.api.field_behavior) = OPTIONAL]; + + // Output only. A list of streams created with the session. + // + // At least one stream is created with the session. In the future, larger + // request_stream_count values *may* result in this list being unpopulated, + // in that case, the user will need to use a List method to get the streams + // instead, which is not yet available. + repeated ReadStream streams = 10 [(google.api.field_behavior) = OUTPUT_ONLY]; +} + +// Information about a single stream that gets data out of the storage system. +// Most of the information about `ReadStream` instances is aggregated, making +// `ReadStream` lightweight. +message ReadStream { + option (google.api.resource) = { + type: "bigquerystorage.googleapis.com/ReadStream" + pattern: "projects/{project}/locations/{location}/sessions/{session}/streams/{stream}" + }; + + // Output only. Name of the stream, in the form + // `projects/{project_id}/locations/{location}/sessions/{session_id}/streams/{stream_id}`. + string name = 1 [(google.api.field_behavior) = OUTPUT_ONLY]; +} + +// Information about a single stream that gets data inside the storage system. +message WriteStream { + option (google.api.resource) = { + type: "bigquerystorage.googleapis.com/WriteStream" + pattern: "projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}" + }; + + // Type enum of the stream. + enum Type { + // Unknown type. + TYPE_UNSPECIFIED = 0; + + // Data will commit automatically and appear as soon as the write is + // acknowledged. + COMMITTED = 1; + + // Data is invisible until the stream is committed. + PENDING = 2; + + // Data is only visible up to the offset to which it was flushed. + BUFFERED = 3; + } + + // Output only. Name of the stream, in the form + // `projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}`. + string name = 1 [(google.api.field_behavior) = OUTPUT_ONLY]; + + // Immutable. Type of the stream. + Type type = 2 [(google.api.field_behavior) = IMMUTABLE]; + + // Output only. Create time of the stream. For the _default stream, this is the + // creation_time of the table. + google.protobuf.Timestamp create_time = 3 [(google.api.field_behavior) = OUTPUT_ONLY]; + + // Output only. Commit time of the stream. + // If a stream is of `COMMITTED` type, then it will have a commit_time same as + // `create_time`. If the stream is of `PENDING` type, commit_time being empty + // means it is not committed. + google.protobuf.Timestamp commit_time = 4 [(google.api.field_behavior) = OUTPUT_ONLY]; + + // Output only. The schema of the destination table. It is only returned in + // `CreateWriteStream` response. Caller should generate data that's + // compatible with this schema to send in initial `AppendRowsRequest`. + // The table schema could go out of date during the life time of the stream. + TableSchema table_schema = 5 [(google.api.field_behavior) = OUTPUT_ONLY]; +} diff --git a/google/cloud/bigquery_storage_v1beta2/proto/table.proto b/google/cloud/bigquery_storage_v1beta2/proto/table.proto new file mode 100644 index 00000000..fd8a0a75 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/proto/table.proto @@ -0,0 +1,102 @@ +// 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. + +syntax = "proto3"; + +package google.cloud.bigquery.storage.v1beta2; + +import "google/api/field_behavior.proto"; + +option go_package = "google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1beta2;storage"; +option java_multiple_files = true; +option java_outer_classname = "TableProto"; +option java_package = "com.google.cloud.bigquery.storage.v1beta2"; + +// Schema of a table +message TableSchema { + // Describes the fields in a table. + repeated TableFieldSchema fields = 1; +} + +// A field in TableSchema +message TableFieldSchema { + enum Type { + // Illegal value + TYPE_UNSPECIFIED = 0; + + // 64K, UTF8 + STRING = 1; + + // 64-bit signed + INT64 = 2; + + // 64-bit IEEE floating point + DOUBLE = 3; + + // Aggregate type + STRUCT = 4; + + // 64K, Binary + BYTES = 5; + + // 2-valued + BOOL = 6; + + // 64-bit signed usec since UTC epoch + TIMESTAMP = 7; + + // Civil date - Year, Month, Day + DATE = 8; + + // Civil time - Hour, Minute, Second, Microseconds + TIME = 9; + + // Combination of civil date and civil time + DATETIME = 10; + + // Geography object + GEOGRAPHY = 11; + + // Numeric value + NUMERIC = 12; + } + + enum Mode { + // Illegal value + MODE_UNSPECIFIED = 0; + + NULLABLE = 1; + + REQUIRED = 2; + + REPEATED = 3; + } + + // Required. The field name. The name must contain only letters (a-z, A-Z), + // numbers (0-9), or underscores (_), and must start with a letter or + // underscore. The maximum length is 128 characters. + string name = 1 [(google.api.field_behavior) = REQUIRED]; + + // Required. The field data type. + Type type = 2 [(google.api.field_behavior) = REQUIRED]; + + // Optional. The field mode. The default value is NULLABLE. + Mode mode = 3 [(google.api.field_behavior) = OPTIONAL]; + + // Optional. Describes the nested schema fields if the type property is set to STRUCT. + repeated TableFieldSchema fields = 4 [(google.api.field_behavior) = OPTIONAL]; + + // Optional. The field description. The maximum length is 1,024 characters. + string description = 6 [(google.api.field_behavior) = OPTIONAL]; +} diff --git a/google/cloud/bigquery_storage_v1beta2/py.typed b/google/cloud/bigquery_storage_v1beta2/py.typed new file mode 100644 index 00000000..e71b4749 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/py.typed @@ -0,0 +1,2 @@ +# Marker file for PEP 561. +# The google-cloud-bigquery-storage package uses inline types. diff --git a/google/cloud/bigquery_storage_v1beta2/services/__init__.py b/google/cloud/bigquery_storage_v1beta2/services/__init__.py new file mode 100644 index 00000000..42ffdf2b --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/__init__.py @@ -0,0 +1,16 @@ +# -*- 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. +# diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_read/__init__.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/__init__.py new file mode 100644 index 00000000..2105a1a6 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/__init__.py @@ -0,0 +1,24 @@ +# -*- 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. +# + +from .client import BigQueryReadClient +from .async_client import BigQueryReadAsyncClient + +__all__ = ( + "BigQueryReadClient", + "BigQueryReadAsyncClient", +) diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_read/async_client.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/async_client.py new file mode 100644 index 00000000..69b6ebe2 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/async_client.py @@ -0,0 +1,462 @@ +# -*- 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. +# + +from collections import OrderedDict +import functools +import re +from typing import Dict, AsyncIterable, Awaitable, Sequence, Tuple, Type, Union +import pkg_resources + +import google.api_core.client_options as ClientOptions # type: ignore +from google.api_core import exceptions # type: ignore +from google.api_core import gapic_v1 # type: ignore +from google.api_core import retry as retries # type: ignore +from google.auth import credentials # type: ignore +from google.oauth2 import service_account # 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 storage +from google.cloud.bigquery_storage_v1beta2.types import stream +from google.protobuf import timestamp_pb2 as timestamp # type: ignore + +from .transports.base import BigQueryReadTransport, DEFAULT_CLIENT_INFO +from .transports.grpc_asyncio import BigQueryReadGrpcAsyncIOTransport +from .client import BigQueryReadClient + + +class BigQueryReadAsyncClient: + """BigQuery Read API. + The Read API can be used to read data from BigQuery. + New code should use the v1 Read API going forward, if they don't + use Write API at the same time. + """ + + _client: BigQueryReadClient + + DEFAULT_ENDPOINT = BigQueryReadClient.DEFAULT_ENDPOINT + DEFAULT_MTLS_ENDPOINT = BigQueryReadClient.DEFAULT_MTLS_ENDPOINT + + read_session_path = staticmethod(BigQueryReadClient.read_session_path) + parse_read_session_path = staticmethod(BigQueryReadClient.parse_read_session_path) + read_stream_path = staticmethod(BigQueryReadClient.read_stream_path) + parse_read_stream_path = staticmethod(BigQueryReadClient.parse_read_stream_path) + table_path = staticmethod(BigQueryReadClient.table_path) + parse_table_path = staticmethod(BigQueryReadClient.parse_table_path) + + common_billing_account_path = staticmethod( + BigQueryReadClient.common_billing_account_path + ) + parse_common_billing_account_path = staticmethod( + BigQueryReadClient.parse_common_billing_account_path + ) + + common_folder_path = staticmethod(BigQueryReadClient.common_folder_path) + parse_common_folder_path = staticmethod(BigQueryReadClient.parse_common_folder_path) + + common_organization_path = staticmethod(BigQueryReadClient.common_organization_path) + parse_common_organization_path = staticmethod( + BigQueryReadClient.parse_common_organization_path + ) + + common_project_path = staticmethod(BigQueryReadClient.common_project_path) + parse_common_project_path = staticmethod( + BigQueryReadClient.parse_common_project_path + ) + + common_location_path = staticmethod(BigQueryReadClient.common_location_path) + parse_common_location_path = staticmethod( + BigQueryReadClient.parse_common_location_path + ) + + from_service_account_info = BigQueryReadClient.from_service_account_info + from_service_account_file = BigQueryReadClient.from_service_account_file + from_service_account_json = from_service_account_file + + @property + def transport(self) -> BigQueryReadTransport: + """Return the transport used by the client instance. + + Returns: + BigQueryReadTransport: The transport used by the client instance. + """ + return self._client.transport + + get_transport_class = functools.partial( + type(BigQueryReadClient).get_transport_class, type(BigQueryReadClient) + ) + + def __init__( + self, + *, + credentials: credentials.Credentials = None, + transport: Union[str, BigQueryReadTransport] = "grpc_asyncio", + client_options: ClientOptions = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + ) -> None: + """Instantiate the big query read client. + + Args: + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + transport (Union[str, ~.BigQueryReadTransport]): The + transport to use. If set to None, a transport is chosen + automatically. + client_options (ClientOptions): Custom options for the client. It + won't take effect if a ``transport`` instance is provided. + (1) The ``api_endpoint`` property can be used to override the + default endpoint provided by the client. GOOGLE_API_USE_MTLS_ENDPOINT + environment variable can also be used to override the endpoint: + "always" (always use the default mTLS endpoint), "never" (always + use the default regular endpoint) and "auto" (auto switch to the + default mTLS endpoint if client certificate is present, this is + the default value). However, the ``api_endpoint`` property takes + precedence if provided. + (2) If GOOGLE_API_USE_CLIENT_CERTIFICATE environment variable + is "true", then the ``client_cert_source`` property can be used + to provide client certificate for mutual TLS transport. If + not provided, the default SSL client certificate will be used if + present. If GOOGLE_API_USE_CLIENT_CERTIFICATE is "false" or not + set, no client certificate will be used. + + Raises: + google.auth.exceptions.MutualTlsChannelError: If mutual TLS transport + creation failed for any reason. + """ + + self._client = BigQueryReadClient( + credentials=credentials, + transport=transport, + client_options=client_options, + client_info=client_info, + ) + + async def create_read_session( + self, + request: storage.CreateReadSessionRequest = None, + *, + parent: str = None, + read_session: stream.ReadSession = None, + max_stream_count: int = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> stream.ReadSession: + r"""Creates a new read session. A read session divides + the contents of a BigQuery table into one or more + streams, which can then be used to read data from the + table. The read session also specifies properties of the + data to be read, such as a list of columns or a push- + down filter describing the rows to be returned. + + A particular row can be read by at most one stream. When + the caller has reached the end of each stream in the + session, then all the data in the table has been read. + + Data is assigned to each stream such that roughly the + same number of rows can be read from each stream. + Because the server-side unit for assigning data is + collections of rows, the API does not guarantee that + each stream will return the same number or rows. + Additionally, the limits are enforced based on the + number of pre-filtered rows, so some filters can lead to + lopsided assignments. + + Read sessions automatically expire 24 hours after they + are created and do not require manual clean-up by the + caller. + + Args: + request (:class:`google.cloud.bigquery_storage_v1beta2.types.CreateReadSessionRequest`): + The request object. Request message for + `CreateReadSession`. + parent (:class:`str`): + Required. The request project that owns the session, in + the form of ``projects/{project_id}``. + + This corresponds to the ``parent`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + read_session (:class:`google.cloud.bigquery_storage_v1beta2.types.ReadSession`): + Required. Session to be created. + This corresponds to the ``read_session`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + max_stream_count (:class:`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. + + This corresponds to the ``max_stream_count`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.ReadSession: + Information about the ReadSession. + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([parent, read_session, max_stream_count]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + request = storage.CreateReadSessionRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if parent is not None: + request.parent = parent + if read_session is not None: + request.read_session = read_session + if max_stream_count is not None: + request.max_stream_count = max_stream_count + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.create_read_session, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata( + (("read_session.table", request.read_session.table),) + ), + ) + + # Send the request. + response = await rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + def read_rows( + self, + request: storage.ReadRowsRequest = None, + *, + read_stream: str = None, + offset: int = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> Awaitable[AsyncIterable[storage.ReadRowsResponse]]: + r"""Reads rows from the stream in the format prescribed + by the ReadSession. Each response contains one or more + table rows, up to a maximum of 100 MiB per response; + read requests which attempt to read individual rows + larger than 100 MiB will fail. + + Each request also returns a set of stream statistics + reflecting the current state of the stream. + + Args: + request (:class:`google.cloud.bigquery_storage_v1beta2.types.ReadRowsRequest`): + The request object. Request message for `ReadRows`. + read_stream (:class:`str`): + Required. Stream to read rows from. + This corresponds to the ``read_stream`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + offset (:class:`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. + + This corresponds to the ``offset`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + AsyncIterable[google.cloud.bigquery_storage_v1beta2.types.ReadRowsResponse]: + Response from calling ReadRows may include row data, progress and + throttling information. + + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([read_stream, offset]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + request = storage.ReadRowsRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if read_stream is not None: + request.read_stream = read_stream + if offset is not None: + request.offset = offset + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.read_rows, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type(exceptions.ServiceUnavailable,), + ), + default_timeout=86400.0, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata( + (("read_stream", request.read_stream),) + ), + ) + + # Send the request. + response = rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + async def split_read_stream( + self, + request: storage.SplitReadStreamRequest = None, + *, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> storage.SplitReadStreamResponse: + r"""Splits a given ``ReadStream`` into two ``ReadStream`` objects. + These ``ReadStream`` objects are referred to as the primary and + the residual streams of the split. The original ``ReadStream`` + can still be read from in the same manner as before. Both of the + returned ``ReadStream`` objects can also be read from, and the + rows returned by both child streams will be the same as the rows + read from the original stream. + + Moreover, the two child streams will be allocated back-to-back + in the original ``ReadStream``. Concretely, it is guaranteed + that for streams original, primary, and residual, that + original[0-j] = primary[0-j] and original[j-n] = residual[0-m] + once the streams have been read to completion. + + Args: + request (:class:`google.cloud.bigquery_storage_v1beta2.types.SplitReadStreamRequest`): + The request object. Request message for + `SplitReadStream`. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.SplitReadStreamResponse: + + """ + # Create or coerce a protobuf request object. + + request = storage.SplitReadStreamRequest(request) + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.split_read_stream, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Send the request. + response = await rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + +try: + DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( + gapic_version=pkg_resources.get_distribution( + "google-cloud-bigquery-storage", + ).version, + ) +except pkg_resources.DistributionNotFound: + DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo() + + +__all__ = ("BigQueryReadAsyncClient",) diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_read/client.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/client.py new file mode 100644 index 00000000..110dfd5f --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/client.py @@ -0,0 +1,681 @@ +# -*- 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. +# + +from collections import OrderedDict +from distutils import util +import os +import re +from typing import Callable, Dict, Optional, Iterable, Sequence, Tuple, Type, Union +import pkg_resources + +from google.api_core import client_options as client_options_lib # type: ignore +from google.api_core import exceptions # type: ignore +from google.api_core import gapic_v1 # type: ignore +from google.api_core import retry as retries # type: ignore +from google.auth import credentials # type: ignore +from google.auth.transport import mtls # type: ignore +from google.auth.transport.grpc import SslCredentials # type: ignore +from google.auth.exceptions import MutualTLSChannelError # type: ignore +from google.oauth2 import service_account # 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 storage +from google.cloud.bigquery_storage_v1beta2.types import stream +from google.protobuf import timestamp_pb2 as timestamp # type: ignore + +from .transports.base import BigQueryReadTransport, DEFAULT_CLIENT_INFO +from .transports.grpc import BigQueryReadGrpcTransport +from .transports.grpc_asyncio import BigQueryReadGrpcAsyncIOTransport + + +class BigQueryReadClientMeta(type): + """Metaclass for the BigQueryRead client. + + This provides class-level methods for building and retrieving + support objects (e.g. transport) without polluting the client instance + objects. + """ + + _transport_registry = OrderedDict() # type: Dict[str, Type[BigQueryReadTransport]] + _transport_registry["grpc"] = BigQueryReadGrpcTransport + _transport_registry["grpc_asyncio"] = BigQueryReadGrpcAsyncIOTransport + + def get_transport_class(cls, label: str = None,) -> Type[BigQueryReadTransport]: + """Return an appropriate transport class. + + Args: + label: The name of the desired transport. If none is + provided, then the first transport in the registry is used. + + Returns: + The transport class to use. + """ + # If a specific transport is requested, return that one. + if label: + return cls._transport_registry[label] + + # No transport is requested; return the default (that is, the first one + # in the dictionary). + return next(iter(cls._transport_registry.values())) + + +class BigQueryReadClient(metaclass=BigQueryReadClientMeta): + """BigQuery Read API. + The Read API can be used to read data from BigQuery. + New code should use the v1 Read API going forward, if they don't + use Write API at the same time. + """ + + @staticmethod + def _get_default_mtls_endpoint(api_endpoint): + """Convert api endpoint to mTLS endpoint. + Convert "*.sandbox.googleapis.com" and "*.googleapis.com" to + "*.mtls.sandbox.googleapis.com" and "*.mtls.googleapis.com" respectively. + Args: + api_endpoint (Optional[str]): the api endpoint to convert. + Returns: + str: converted mTLS api endpoint. + """ + if not api_endpoint: + return api_endpoint + + mtls_endpoint_re = re.compile( + r"(?P[^.]+)(?P\.mtls)?(?P\.sandbox)?(?P\.googleapis\.com)?" + ) + + m = mtls_endpoint_re.match(api_endpoint) + name, mtls, sandbox, googledomain = m.groups() + if mtls or not googledomain: + return api_endpoint + + if sandbox: + return api_endpoint.replace( + "sandbox.googleapis.com", "mtls.sandbox.googleapis.com" + ) + + return api_endpoint.replace(".googleapis.com", ".mtls.googleapis.com") + + DEFAULT_ENDPOINT = "bigquerystorage.googleapis.com" + DEFAULT_MTLS_ENDPOINT = _get_default_mtls_endpoint.__func__( # type: ignore + DEFAULT_ENDPOINT + ) + + @classmethod + def from_service_account_info(cls, info: dict, *args, **kwargs): + """Creates an instance of this client using the provided credentials info. + + Args: + info (dict): The service account private key info. + args: Additional arguments to pass to the constructor. + kwargs: Additional arguments to pass to the constructor. + + Returns: + BigQueryReadClient: The constructed client. + """ + credentials = service_account.Credentials.from_service_account_info(info) + kwargs["credentials"] = credentials + return cls(*args, **kwargs) + + @classmethod + def from_service_account_file(cls, filename: str, *args, **kwargs): + """Creates an instance of this client using the provided credentials + file. + + Args: + filename (str): The path to the service account private key json + file. + args: Additional arguments to pass to the constructor. + kwargs: Additional arguments to pass to the constructor. + + Returns: + BigQueryReadClient: The constructed client. + """ + credentials = service_account.Credentials.from_service_account_file(filename) + kwargs["credentials"] = credentials + return cls(*args, **kwargs) + + from_service_account_json = from_service_account_file + + @property + def transport(self) -> BigQueryReadTransport: + """Return the transport used by the client instance. + + Returns: + BigQueryReadTransport: The transport used by the client instance. + """ + return self._transport + + @staticmethod + def read_session_path(project: str, location: str, session: str,) -> str: + """Return a fully-qualified read_session string.""" + return "projects/{project}/locations/{location}/sessions/{session}".format( + project=project, location=location, session=session, + ) + + @staticmethod + def parse_read_session_path(path: str) -> Dict[str, str]: + """Parse a read_session path into its component segments.""" + m = re.match( + r"^projects/(?P.+?)/locations/(?P.+?)/sessions/(?P.+?)$", + path, + ) + return m.groupdict() if m else {} + + @staticmethod + def read_stream_path( + project: str, location: str, session: str, stream: str, + ) -> str: + """Return a fully-qualified read_stream string.""" + return "projects/{project}/locations/{location}/sessions/{session}/streams/{stream}".format( + project=project, location=location, session=session, stream=stream, + ) + + @staticmethod + def parse_read_stream_path(path: str) -> Dict[str, str]: + """Parse a read_stream path into its component segments.""" + m = re.match( + r"^projects/(?P.+?)/locations/(?P.+?)/sessions/(?P.+?)/streams/(?P.+?)$", + path, + ) + return m.groupdict() if m else {} + + @staticmethod + def table_path(project: str, dataset: str, table: str,) -> str: + """Return a fully-qualified table string.""" + return "projects/{project}/datasets/{dataset}/tables/{table}".format( + project=project, dataset=dataset, table=table, + ) + + @staticmethod + def parse_table_path(path: str) -> Dict[str, str]: + """Parse a table path into its component segments.""" + m = re.match( + r"^projects/(?P.+?)/datasets/(?P.+?)/tables/(?P.+?)$", + path, + ) + return m.groupdict() if m else {} + + @staticmethod + def common_billing_account_path(billing_account: str,) -> str: + """Return a fully-qualified billing_account string.""" + return "billingAccounts/{billing_account}".format( + billing_account=billing_account, + ) + + @staticmethod + def parse_common_billing_account_path(path: str) -> Dict[str, str]: + """Parse a billing_account path into its component segments.""" + m = re.match(r"^billingAccounts/(?P.+?)$", path) + return m.groupdict() if m else {} + + @staticmethod + def common_folder_path(folder: str,) -> str: + """Return a fully-qualified folder string.""" + return "folders/{folder}".format(folder=folder,) + + @staticmethod + def parse_common_folder_path(path: str) -> Dict[str, str]: + """Parse a folder path into its component segments.""" + m = re.match(r"^folders/(?P.+?)$", path) + return m.groupdict() if m else {} + + @staticmethod + def common_organization_path(organization: str,) -> str: + """Return a fully-qualified organization string.""" + return "organizations/{organization}".format(organization=organization,) + + @staticmethod + def parse_common_organization_path(path: str) -> Dict[str, str]: + """Parse a organization path into its component segments.""" + m = re.match(r"^organizations/(?P.+?)$", path) + return m.groupdict() if m else {} + + @staticmethod + def common_project_path(project: str,) -> str: + """Return a fully-qualified project string.""" + return "projects/{project}".format(project=project,) + + @staticmethod + def parse_common_project_path(path: str) -> Dict[str, str]: + """Parse a project path into its component segments.""" + m = re.match(r"^projects/(?P.+?)$", path) + return m.groupdict() if m else {} + + @staticmethod + def common_location_path(project: str, location: str,) -> str: + """Return a fully-qualified location string.""" + return "projects/{project}/locations/{location}".format( + project=project, location=location, + ) + + @staticmethod + def parse_common_location_path(path: str) -> Dict[str, str]: + """Parse a location path into its component segments.""" + m = re.match(r"^projects/(?P.+?)/locations/(?P.+?)$", path) + return m.groupdict() if m else {} + + def __init__( + self, + *, + credentials: Optional[credentials.Credentials] = None, + transport: Union[str, BigQueryReadTransport, None] = None, + client_options: Optional[client_options_lib.ClientOptions] = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + ) -> None: + """Instantiate the big query read client. + + Args: + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + transport (Union[str, BigQueryReadTransport]): The + transport to use. If set to None, a transport is chosen + automatically. + client_options (google.api_core.client_options.ClientOptions): Custom options for the + client. It won't take effect if a ``transport`` instance is provided. + (1) The ``api_endpoint`` property can be used to override the + default endpoint provided by the client. GOOGLE_API_USE_MTLS_ENDPOINT + environment variable can also be used to override the endpoint: + "always" (always use the default mTLS endpoint), "never" (always + use the default regular endpoint) and "auto" (auto switch to the + default mTLS endpoint if client certificate is present, this is + the default value). However, the ``api_endpoint`` property takes + precedence if provided. + (2) If GOOGLE_API_USE_CLIENT_CERTIFICATE environment variable + is "true", then the ``client_cert_source`` property can be used + to provide client certificate for mutual TLS transport. If + not provided, the default SSL client certificate will be used if + present. If GOOGLE_API_USE_CLIENT_CERTIFICATE is "false" or not + set, no client certificate will be used. + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + + Raises: + google.auth.exceptions.MutualTLSChannelError: If mutual TLS transport + creation failed for any reason. + """ + if isinstance(client_options, dict): + client_options = client_options_lib.from_dict(client_options) + if client_options is None: + client_options = client_options_lib.ClientOptions() + + # Create SSL credentials for mutual TLS if needed. + use_client_cert = bool( + util.strtobool(os.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false")) + ) + + ssl_credentials = None + is_mtls = False + if use_client_cert: + if client_options.client_cert_source: + import grpc # type: ignore + + cert, key = client_options.client_cert_source() + ssl_credentials = grpc.ssl_channel_credentials( + certificate_chain=cert, private_key=key + ) + is_mtls = True + else: + creds = SslCredentials() + is_mtls = creds.is_mtls + ssl_credentials = creds.ssl_credentials if is_mtls else None + + # Figure out which api endpoint to use. + if client_options.api_endpoint is not None: + api_endpoint = client_options.api_endpoint + else: + use_mtls_env = os.getenv("GOOGLE_API_USE_MTLS_ENDPOINT", "auto") + if use_mtls_env == "never": + api_endpoint = self.DEFAULT_ENDPOINT + elif use_mtls_env == "always": + api_endpoint = self.DEFAULT_MTLS_ENDPOINT + elif use_mtls_env == "auto": + api_endpoint = ( + self.DEFAULT_MTLS_ENDPOINT if is_mtls else self.DEFAULT_ENDPOINT + ) + else: + raise MutualTLSChannelError( + "Unsupported GOOGLE_API_USE_MTLS_ENDPOINT value. Accepted values: never, auto, always" + ) + + # Save or instantiate the transport. + # Ordinarily, we provide the transport, but allowing a custom transport + # instance provides an extensibility point for unusual situations. + if isinstance(transport, BigQueryReadTransport): + # transport is a BigQueryReadTransport instance. + if credentials or client_options.credentials_file: + raise ValueError( + "When providing a transport instance, " + "provide its credentials directly." + ) + if client_options.scopes: + raise ValueError( + "When providing a transport instance, " + "provide its scopes directly." + ) + self._transport = transport + else: + Transport = type(self).get_transport_class(transport) + self._transport = Transport( + credentials=credentials, + credentials_file=client_options.credentials_file, + host=api_endpoint, + scopes=client_options.scopes, + ssl_channel_credentials=ssl_credentials, + quota_project_id=client_options.quota_project_id, + client_info=client_info, + ) + + def create_read_session( + self, + request: storage.CreateReadSessionRequest = None, + *, + parent: str = None, + read_session: stream.ReadSession = None, + max_stream_count: int = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> stream.ReadSession: + r"""Creates a new read session. A read session divides + the contents of a BigQuery table into one or more + streams, which can then be used to read data from the + table. The read session also specifies properties of the + data to be read, such as a list of columns or a push- + down filter describing the rows to be returned. + + A particular row can be read by at most one stream. When + the caller has reached the end of each stream in the + session, then all the data in the table has been read. + + Data is assigned to each stream such that roughly the + same number of rows can be read from each stream. + Because the server-side unit for assigning data is + collections of rows, the API does not guarantee that + each stream will return the same number or rows. + Additionally, the limits are enforced based on the + number of pre-filtered rows, so some filters can lead to + lopsided assignments. + + Read sessions automatically expire 24 hours after they + are created and do not require manual clean-up by the + caller. + + Args: + request (google.cloud.bigquery_storage_v1beta2.types.CreateReadSessionRequest): + The request object. Request message for + `CreateReadSession`. + parent (str): + Required. The request project that owns the session, in + the form of ``projects/{project_id}``. + + This corresponds to the ``parent`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + read_session (google.cloud.bigquery_storage_v1beta2.types.ReadSession): + Required. Session to be created. + This corresponds to the ``read_session`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + 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. + + This corresponds to the ``max_stream_count`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.ReadSession: + Information about the ReadSession. + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([parent, read_session, max_stream_count]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + # Minor optimization to avoid making a copy if the user passes + # in a storage.CreateReadSessionRequest. + # There's no risk of modifying the input as we've already verified + # there are no flattened fields. + if not isinstance(request, storage.CreateReadSessionRequest): + request = storage.CreateReadSessionRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if parent is not None: + request.parent = parent + if read_session is not None: + request.read_session = read_session + if max_stream_count is not None: + request.max_stream_count = max_stream_count + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[self._transport.create_read_session] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata( + (("read_session.table", request.read_session.table),) + ), + ) + + # Send the request. + response = rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + def read_rows( + self, + request: storage.ReadRowsRequest = None, + *, + read_stream: str = None, + offset: int = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> Iterable[storage.ReadRowsResponse]: + r"""Reads rows from the stream in the format prescribed + by the ReadSession. Each response contains one or more + table rows, up to a maximum of 100 MiB per response; + read requests which attempt to read individual rows + larger than 100 MiB will fail. + + Each request also returns a set of stream statistics + reflecting the current state of the stream. + + Args: + request (google.cloud.bigquery_storage_v1beta2.types.ReadRowsRequest): + The request object. Request message for `ReadRows`. + read_stream (str): + Required. Stream to read rows from. + This corresponds to the ``read_stream`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + 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. + + This corresponds to the ``offset`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + Iterable[google.cloud.bigquery_storage_v1beta2.types.ReadRowsResponse]: + Response from calling ReadRows may include row data, progress and + throttling information. + + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([read_stream, offset]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + # Minor optimization to avoid making a copy if the user passes + # in a storage.ReadRowsRequest. + # There's no risk of modifying the input as we've already verified + # there are no flattened fields. + if not isinstance(request, storage.ReadRowsRequest): + request = storage.ReadRowsRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if read_stream is not None: + request.read_stream = read_stream + if offset is not None: + request.offset = offset + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[self._transport.read_rows] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata( + (("read_stream", request.read_stream),) + ), + ) + + # Send the request. + response = rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + def split_read_stream( + self, + request: storage.SplitReadStreamRequest = None, + *, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> storage.SplitReadStreamResponse: + r"""Splits a given ``ReadStream`` into two ``ReadStream`` objects. + These ``ReadStream`` objects are referred to as the primary and + the residual streams of the split. The original ``ReadStream`` + can still be read from in the same manner as before. Both of the + returned ``ReadStream`` objects can also be read from, and the + rows returned by both child streams will be the same as the rows + read from the original stream. + + Moreover, the two child streams will be allocated back-to-back + in the original ``ReadStream``. Concretely, it is guaranteed + that for streams original, primary, and residual, that + original[0-j] = primary[0-j] and original[j-n] = residual[0-m] + once the streams have been read to completion. + + Args: + request (google.cloud.bigquery_storage_v1beta2.types.SplitReadStreamRequest): + The request object. Request message for + `SplitReadStream`. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.SplitReadStreamResponse: + + """ + # Create or coerce a protobuf request object. + + # Minor optimization to avoid making a copy if the user passes + # in a storage.SplitReadStreamRequest. + # There's no risk of modifying the input as we've already verified + # there are no flattened fields. + if not isinstance(request, storage.SplitReadStreamRequest): + request = storage.SplitReadStreamRequest(request) + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[self._transport.split_read_stream] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Send the request. + response = rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + +try: + DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( + gapic_version=pkg_resources.get_distribution( + "google-cloud-bigquery-storage", + ).version, + ) +except pkg_resources.DistributionNotFound: + DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo() + + +__all__ = ("BigQueryReadClient",) diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_read/transports/__init__.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/transports/__init__.py new file mode 100644 index 00000000..87e56323 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/transports/__init__.py @@ -0,0 +1,35 @@ +# -*- 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. +# + +from collections import OrderedDict +from typing import Dict, Type + +from .base import BigQueryReadTransport +from .grpc import BigQueryReadGrpcTransport +from .grpc_asyncio import BigQueryReadGrpcAsyncIOTransport + + +# Compile a registry of transports. +_transport_registry = OrderedDict() # type: Dict[str, Type[BigQueryReadTransport]] +_transport_registry["grpc"] = BigQueryReadGrpcTransport +_transport_registry["grpc_asyncio"] = BigQueryReadGrpcAsyncIOTransport + +__all__ = ( + "BigQueryReadTransport", + "BigQueryReadGrpcTransport", + "BigQueryReadGrpcAsyncIOTransport", +) diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_read/transports/base.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/transports/base.py new file mode 100644 index 00000000..72f43ab7 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/transports/base.py @@ -0,0 +1,186 @@ +# -*- 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 abc +import typing +import pkg_resources + +from google import auth # type: ignore +from google.api_core import exceptions # type: ignore +from google.api_core import gapic_v1 # type: ignore +from google.api_core import retry as retries # type: ignore +from google.auth import credentials # type: ignore + +from google.cloud.bigquery_storage_v1beta2.types import storage +from google.cloud.bigquery_storage_v1beta2.types import stream + + +try: + DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( + gapic_version=pkg_resources.get_distribution( + "google-cloud-bigquery-storage", + ).version, + ) +except pkg_resources.DistributionNotFound: + DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo() + + +class BigQueryReadTransport(abc.ABC): + """Abstract transport class for BigQueryRead.""" + + AUTH_SCOPES = ( + "https://www.googleapis.com/auth/bigquery", + "https://www.googleapis.com/auth/bigquery.readonly", + "https://www.googleapis.com/auth/cloud-platform", + ) + + def __init__( + self, + *, + host: str = "bigquerystorage.googleapis.com", + credentials: credentials.Credentials = None, + credentials_file: typing.Optional[str] = None, + scopes: typing.Optional[typing.Sequence[str]] = AUTH_SCOPES, + quota_project_id: typing.Optional[str] = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + **kwargs, + ) -> None: + """Instantiate the transport. + + Args: + host (Optional[str]): The hostname to connect to. + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is mutually exclusive with credentials. + scope (Optional[Sequence[str]]): A list of scopes. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + """ + # Save the hostname. Default to port 443 (HTTPS) if none is specified. + if ":" not in host: + host += ":443" + self._host = host + + # If no credentials are provided, then determine the appropriate + # defaults. + if credentials and credentials_file: + raise exceptions.DuplicateCredentialArgs( + "'credentials_file' and 'credentials' are mutually exclusive" + ) + + if credentials_file is not None: + credentials, _ = auth.load_credentials_from_file( + credentials_file, scopes=scopes, quota_project_id=quota_project_id + ) + + elif credentials is None: + credentials, _ = auth.default( + scopes=scopes, quota_project_id=quota_project_id + ) + + # Save the credentials. + self._credentials = credentials + + # Lifted into its own function so it can be stubbed out during tests. + self._prep_wrapped_messages(client_info) + + def _prep_wrapped_messages(self, client_info): + # Precompute the wrapped methods. + self._wrapped_methods = { + self.create_read_session: gapic_v1.method.wrap_method( + self.create_read_session, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=client_info, + ), + self.read_rows: gapic_v1.method.wrap_method( + self.read_rows, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type(exceptions.ServiceUnavailable,), + ), + default_timeout=86400.0, + client_info=client_info, + ), + self.split_read_stream: gapic_v1.method.wrap_method( + self.split_read_stream, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=client_info, + ), + } + + @property + def create_read_session( + self, + ) -> typing.Callable[ + [storage.CreateReadSessionRequest], + typing.Union[stream.ReadSession, typing.Awaitable[stream.ReadSession]], + ]: + raise NotImplementedError() + + @property + def read_rows( + self, + ) -> typing.Callable[ + [storage.ReadRowsRequest], + typing.Union[ + storage.ReadRowsResponse, typing.Awaitable[storage.ReadRowsResponse] + ], + ]: + raise NotImplementedError() + + @property + def split_read_stream( + self, + ) -> typing.Callable[ + [storage.SplitReadStreamRequest], + typing.Union[ + storage.SplitReadStreamResponse, + typing.Awaitable[storage.SplitReadStreamResponse], + ], + ]: + raise NotImplementedError() + + +__all__ = ("BigQueryReadTransport",) diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_read/transports/grpc.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/transports/grpc.py new file mode 100644 index 00000000..50d3c3fd --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/transports/grpc.py @@ -0,0 +1,364 @@ +# -*- 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 warnings +from typing import Callable, Dict, Optional, Sequence, Tuple + +from google.api_core import grpc_helpers # type: ignore +from google.api_core import gapic_v1 # type: ignore +from google import auth # type: ignore +from google.auth import credentials # type: ignore +from google.auth.transport.grpc import SslCredentials # type: ignore + +import grpc # type: ignore + +from google.cloud.bigquery_storage_v1beta2.types import storage +from google.cloud.bigquery_storage_v1beta2.types import stream + +from .base import BigQueryReadTransport, DEFAULT_CLIENT_INFO + + +class BigQueryReadGrpcTransport(BigQueryReadTransport): + """gRPC backend transport for BigQueryRead. + + BigQuery Read API. + The Read API can be used to read data from BigQuery. + New code should use the v1 Read API going forward, if they don't + use Write API at the same time. + + This class defines the same methods as the primary client, so the + primary client can load the underlying transport implementation + and call it. + + It sends protocol buffers over the wire using gRPC (which is built on + top of HTTP/2); the ``grpcio`` package must be installed. + """ + + _stubs: Dict[str, Callable] + + def __init__( + self, + *, + host: str = "bigquerystorage.googleapis.com", + credentials: credentials.Credentials = None, + credentials_file: str = None, + scopes: Sequence[str] = None, + channel: grpc.Channel = None, + api_mtls_endpoint: str = None, + client_cert_source: Callable[[], Tuple[bytes, bytes]] = None, + ssl_channel_credentials: grpc.ChannelCredentials = None, + quota_project_id: Optional[str] = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + ) -> None: + """Instantiate the transport. + + Args: + host (Optional[str]): The hostname to connect to. + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + This argument is ignored if ``channel`` is provided. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is ignored if ``channel`` is provided. + scopes (Optional(Sequence[str])): A list of scopes. This argument is + ignored if ``channel`` is provided. + channel (Optional[grpc.Channel]): A ``Channel`` instance through + which to make calls. + api_mtls_endpoint (Optional[str]): Deprecated. The mutual TLS endpoint. + If provided, it overrides the ``host`` argument and tries to create + a mutual TLS channel with client SSL credentials from + ``client_cert_source`` or applicatin default SSL credentials. + client_cert_source (Optional[Callable[[], Tuple[bytes, bytes]]]): + Deprecated. A callback to provide client SSL certificate bytes and + private key bytes, both in PEM format. It is ignored if + ``api_mtls_endpoint`` is None. + ssl_channel_credentials (grpc.ChannelCredentials): SSL credentials + for grpc channel. It is ignored if ``channel`` is provided. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + + Raises: + google.auth.exceptions.MutualTLSChannelError: If mutual TLS transport + creation failed for any reason. + google.api_core.exceptions.DuplicateCredentialArgs: If both ``credentials`` + and ``credentials_file`` are passed. + """ + self._ssl_channel_credentials = ssl_channel_credentials + + if channel: + # Sanity check: Ensure that channel and credentials are not both + # provided. + credentials = False + + # If a channel was explicitly provided, set it. + self._grpc_channel = channel + self._ssl_channel_credentials = None + elif api_mtls_endpoint: + warnings.warn( + "api_mtls_endpoint and client_cert_source are deprecated", + DeprecationWarning, + ) + + host = ( + api_mtls_endpoint + if ":" in api_mtls_endpoint + else api_mtls_endpoint + ":443" + ) + + if credentials is None: + credentials, _ = auth.default( + scopes=self.AUTH_SCOPES, quota_project_id=quota_project_id + ) + + # Create SSL credentials with client_cert_source or application + # default SSL credentials. + if client_cert_source: + cert, key = client_cert_source() + ssl_credentials = grpc.ssl_channel_credentials( + certificate_chain=cert, private_key=key + ) + else: + ssl_credentials = SslCredentials().ssl_credentials + + # create a new channel. The provided one is ignored. + self._grpc_channel = type(self).create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + ssl_credentials=ssl_credentials, + scopes=scopes or self.AUTH_SCOPES, + quota_project_id=quota_project_id, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + self._ssl_channel_credentials = ssl_credentials + else: + host = host if ":" in host else host + ":443" + + if credentials is None: + credentials, _ = auth.default( + scopes=self.AUTH_SCOPES, quota_project_id=quota_project_id + ) + + # create a new channel. The provided one is ignored. + self._grpc_channel = type(self).create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + ssl_credentials=ssl_channel_credentials, + scopes=scopes or self.AUTH_SCOPES, + quota_project_id=quota_project_id, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + + self._stubs = {} # type: Dict[str, Callable] + + # Run the base constructor. + super().__init__( + host=host, + credentials=credentials, + credentials_file=credentials_file, + scopes=scopes or self.AUTH_SCOPES, + quota_project_id=quota_project_id, + client_info=client_info, + ) + + @classmethod + def create_channel( + cls, + host: str = "bigquerystorage.googleapis.com", + credentials: credentials.Credentials = None, + credentials_file: str = None, + scopes: Optional[Sequence[str]] = None, + quota_project_id: Optional[str] = None, + **kwargs, + ) -> grpc.Channel: + """Create and return a gRPC channel object. + Args: + address (Optional[str]): The host for the channel to use. + credentials (Optional[~.Credentials]): The + authorization credentials to attach to requests. These + credentials identify this application to the service. If + none are specified, the client will attempt to ascertain + the credentials from the environment. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is mutually exclusive with credentials. + scopes (Optional[Sequence[str]]): A optional list of scopes needed for this + service. These are only used when credentials are not specified and + are passed to :func:`google.auth.default`. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + kwargs (Optional[dict]): Keyword arguments, which are passed to the + channel creation. + Returns: + grpc.Channel: A gRPC channel object. + + Raises: + google.api_core.exceptions.DuplicateCredentialArgs: If both ``credentials`` + and ``credentials_file`` are passed. + """ + scopes = scopes or cls.AUTH_SCOPES + return grpc_helpers.create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + scopes=scopes, + quota_project_id=quota_project_id, + **kwargs, + ) + + @property + def grpc_channel(self) -> grpc.Channel: + """Return the channel designed to connect to this service. + """ + return self._grpc_channel + + @property + def create_read_session( + self, + ) -> Callable[[storage.CreateReadSessionRequest], stream.ReadSession]: + r"""Return a callable for the create read session method over gRPC. + + Creates a new read session. A read session divides + the contents of a BigQuery table into one or more + streams, which can then be used to read data from the + table. The read session also specifies properties of the + data to be read, such as a list of columns or a push- + down filter describing the rows to be returned. + + A particular row can be read by at most one stream. When + the caller has reached the end of each stream in the + session, then all the data in the table has been read. + + Data is assigned to each stream such that roughly the + same number of rows can be read from each stream. + Because the server-side unit for assigning data is + collections of rows, the API does not guarantee that + each stream will return the same number or rows. + Additionally, the limits are enforced based on the + number of pre-filtered rows, so some filters can lead to + lopsided assignments. + + Read sessions automatically expire 24 hours after they + are created and do not require manual clean-up by the + caller. + + Returns: + Callable[[~.CreateReadSessionRequest], + ~.ReadSession]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "create_read_session" not in self._stubs: + self._stubs["create_read_session"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryRead/CreateReadSession", + request_serializer=storage.CreateReadSessionRequest.serialize, + response_deserializer=stream.ReadSession.deserialize, + ) + return self._stubs["create_read_session"] + + @property + def read_rows( + self, + ) -> Callable[[storage.ReadRowsRequest], storage.ReadRowsResponse]: + r"""Return a callable for the read rows method over gRPC. + + Reads rows from the stream in the format prescribed + by the ReadSession. Each response contains one or more + table rows, up to a maximum of 100 MiB per response; + read requests which attempt to read individual rows + larger than 100 MiB will fail. + + Each request also returns a set of stream statistics + reflecting the current state of the stream. + + Returns: + Callable[[~.ReadRowsRequest], + ~.ReadRowsResponse]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "read_rows" not in self._stubs: + self._stubs["read_rows"] = self.grpc_channel.unary_stream( + "/google.cloud.bigquery.storage.v1beta2.BigQueryRead/ReadRows", + request_serializer=storage.ReadRowsRequest.serialize, + response_deserializer=storage.ReadRowsResponse.deserialize, + ) + return self._stubs["read_rows"] + + @property + def split_read_stream( + self, + ) -> Callable[[storage.SplitReadStreamRequest], storage.SplitReadStreamResponse]: + r"""Return a callable for the split read stream method over gRPC. + + Splits a given ``ReadStream`` into two ``ReadStream`` objects. + These ``ReadStream`` objects are referred to as the primary and + the residual streams of the split. The original ``ReadStream`` + can still be read from in the same manner as before. Both of the + returned ``ReadStream`` objects can also be read from, and the + rows returned by both child streams will be the same as the rows + read from the original stream. + + Moreover, the two child streams will be allocated back-to-back + in the original ``ReadStream``. Concretely, it is guaranteed + that for streams original, primary, and residual, that + original[0-j] = primary[0-j] and original[j-n] = residual[0-m] + once the streams have been read to completion. + + Returns: + Callable[[~.SplitReadStreamRequest], + ~.SplitReadStreamResponse]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "split_read_stream" not in self._stubs: + self._stubs["split_read_stream"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryRead/SplitReadStream", + request_serializer=storage.SplitReadStreamRequest.serialize, + response_deserializer=storage.SplitReadStreamResponse.deserialize, + ) + return self._stubs["split_read_stream"] + + +__all__ = ("BigQueryReadGrpcTransport",) diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_read/transports/grpc_asyncio.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/transports/grpc_asyncio.py new file mode 100644 index 00000000..9d2a7aa5 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_read/transports/grpc_asyncio.py @@ -0,0 +1,370 @@ +# -*- 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 warnings +from typing import Awaitable, Callable, Dict, Optional, Sequence, Tuple + +from google.api_core import gapic_v1 # type: ignore +from google.api_core import grpc_helpers_async # type: ignore +from google import auth # type: ignore +from google.auth import credentials # type: ignore +from google.auth.transport.grpc import SslCredentials # type: ignore + +import grpc # type: ignore +from grpc.experimental import aio # type: ignore + +from google.cloud.bigquery_storage_v1beta2.types import storage +from google.cloud.bigquery_storage_v1beta2.types import stream + +from .base import BigQueryReadTransport, DEFAULT_CLIENT_INFO +from .grpc import BigQueryReadGrpcTransport + + +class BigQueryReadGrpcAsyncIOTransport(BigQueryReadTransport): + """gRPC AsyncIO backend transport for BigQueryRead. + + BigQuery Read API. + The Read API can be used to read data from BigQuery. + New code should use the v1 Read API going forward, if they don't + use Write API at the same time. + + This class defines the same methods as the primary client, so the + primary client can load the underlying transport implementation + and call it. + + It sends protocol buffers over the wire using gRPC (which is built on + top of HTTP/2); the ``grpcio`` package must be installed. + """ + + _grpc_channel: aio.Channel + _stubs: Dict[str, Callable] = {} + + @classmethod + def create_channel( + cls, + host: str = "bigquerystorage.googleapis.com", + credentials: credentials.Credentials = None, + credentials_file: Optional[str] = None, + scopes: Optional[Sequence[str]] = None, + quota_project_id: Optional[str] = None, + **kwargs, + ) -> aio.Channel: + """Create and return a gRPC AsyncIO channel object. + Args: + address (Optional[str]): The host for the channel to use. + credentials (Optional[~.Credentials]): The + authorization credentials to attach to requests. These + credentials identify this application to the service. If + none are specified, the client will attempt to ascertain + the credentials from the environment. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is ignored if ``channel`` is provided. + scopes (Optional[Sequence[str]]): A optional list of scopes needed for this + service. These are only used when credentials are not specified and + are passed to :func:`google.auth.default`. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + kwargs (Optional[dict]): Keyword arguments, which are passed to the + channel creation. + Returns: + aio.Channel: A gRPC AsyncIO channel object. + """ + scopes = scopes or cls.AUTH_SCOPES + return grpc_helpers_async.create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + scopes=scopes, + quota_project_id=quota_project_id, + **kwargs, + ) + + def __init__( + self, + *, + host: str = "bigquerystorage.googleapis.com", + credentials: credentials.Credentials = None, + credentials_file: Optional[str] = None, + scopes: Optional[Sequence[str]] = None, + channel: aio.Channel = None, + api_mtls_endpoint: str = None, + client_cert_source: Callable[[], Tuple[bytes, bytes]] = None, + ssl_channel_credentials: grpc.ChannelCredentials = None, + quota_project_id=None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + ) -> None: + """Instantiate the transport. + + Args: + host (Optional[str]): The hostname to connect to. + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + This argument is ignored if ``channel`` is provided. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is ignored if ``channel`` is provided. + scopes (Optional[Sequence[str]]): A optional list of scopes needed for this + service. These are only used when credentials are not specified and + are passed to :func:`google.auth.default`. + channel (Optional[aio.Channel]): A ``Channel`` instance through + which to make calls. + api_mtls_endpoint (Optional[str]): Deprecated. The mutual TLS endpoint. + If provided, it overrides the ``host`` argument and tries to create + a mutual TLS channel with client SSL credentials from + ``client_cert_source`` or applicatin default SSL credentials. + client_cert_source (Optional[Callable[[], Tuple[bytes, bytes]]]): + Deprecated. A callback to provide client SSL certificate bytes and + private key bytes, both in PEM format. It is ignored if + ``api_mtls_endpoint`` is None. + ssl_channel_credentials (grpc.ChannelCredentials): SSL credentials + for grpc channel. It is ignored if ``channel`` is provided. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + + Raises: + google.auth.exceptions.MutualTlsChannelError: If mutual TLS transport + creation failed for any reason. + google.api_core.exceptions.DuplicateCredentialArgs: If both ``credentials`` + and ``credentials_file`` are passed. + """ + self._ssl_channel_credentials = ssl_channel_credentials + + if channel: + # Sanity check: Ensure that channel and credentials are not both + # provided. + credentials = False + + # If a channel was explicitly provided, set it. + self._grpc_channel = channel + self._ssl_channel_credentials = None + elif api_mtls_endpoint: + warnings.warn( + "api_mtls_endpoint and client_cert_source are deprecated", + DeprecationWarning, + ) + + host = ( + api_mtls_endpoint + if ":" in api_mtls_endpoint + else api_mtls_endpoint + ":443" + ) + + if credentials is None: + credentials, _ = auth.default( + scopes=self.AUTH_SCOPES, quota_project_id=quota_project_id + ) + + # Create SSL credentials with client_cert_source or application + # default SSL credentials. + if client_cert_source: + cert, key = client_cert_source() + ssl_credentials = grpc.ssl_channel_credentials( + certificate_chain=cert, private_key=key + ) + else: + ssl_credentials = SslCredentials().ssl_credentials + + # create a new channel. The provided one is ignored. + self._grpc_channel = type(self).create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + ssl_credentials=ssl_credentials, + scopes=scopes or self.AUTH_SCOPES, + quota_project_id=quota_project_id, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + self._ssl_channel_credentials = ssl_credentials + else: + host = host if ":" in host else host + ":443" + + if credentials is None: + credentials, _ = auth.default( + scopes=self.AUTH_SCOPES, quota_project_id=quota_project_id + ) + + # create a new channel. The provided one is ignored. + self._grpc_channel = type(self).create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + ssl_credentials=ssl_channel_credentials, + scopes=scopes or self.AUTH_SCOPES, + quota_project_id=quota_project_id, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + + # Run the base constructor. + super().__init__( + host=host, + credentials=credentials, + credentials_file=credentials_file, + scopes=scopes or self.AUTH_SCOPES, + quota_project_id=quota_project_id, + client_info=client_info, + ) + + self._stubs = {} + + @property + def grpc_channel(self) -> aio.Channel: + """Create the channel designed to connect to this service. + + This property caches on the instance; repeated calls return + the same channel. + """ + # Return the channel from cache. + return self._grpc_channel + + @property + def create_read_session( + self, + ) -> Callable[[storage.CreateReadSessionRequest], Awaitable[stream.ReadSession]]: + r"""Return a callable for the create read session method over gRPC. + + Creates a new read session. A read session divides + the contents of a BigQuery table into one or more + streams, which can then be used to read data from the + table. The read session also specifies properties of the + data to be read, such as a list of columns or a push- + down filter describing the rows to be returned. + + A particular row can be read by at most one stream. When + the caller has reached the end of each stream in the + session, then all the data in the table has been read. + + Data is assigned to each stream such that roughly the + same number of rows can be read from each stream. + Because the server-side unit for assigning data is + collections of rows, the API does not guarantee that + each stream will return the same number or rows. + Additionally, the limits are enforced based on the + number of pre-filtered rows, so some filters can lead to + lopsided assignments. + + Read sessions automatically expire 24 hours after they + are created and do not require manual clean-up by the + caller. + + Returns: + Callable[[~.CreateReadSessionRequest], + Awaitable[~.ReadSession]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "create_read_session" not in self._stubs: + self._stubs["create_read_session"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryRead/CreateReadSession", + request_serializer=storage.CreateReadSessionRequest.serialize, + response_deserializer=stream.ReadSession.deserialize, + ) + return self._stubs["create_read_session"] + + @property + def read_rows( + self, + ) -> Callable[[storage.ReadRowsRequest], Awaitable[storage.ReadRowsResponse]]: + r"""Return a callable for the read rows method over gRPC. + + Reads rows from the stream in the format prescribed + by the ReadSession. Each response contains one or more + table rows, up to a maximum of 100 MiB per response; + read requests which attempt to read individual rows + larger than 100 MiB will fail. + + Each request also returns a set of stream statistics + reflecting the current state of the stream. + + Returns: + Callable[[~.ReadRowsRequest], + Awaitable[~.ReadRowsResponse]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "read_rows" not in self._stubs: + self._stubs["read_rows"] = self.grpc_channel.unary_stream( + "/google.cloud.bigquery.storage.v1beta2.BigQueryRead/ReadRows", + request_serializer=storage.ReadRowsRequest.serialize, + response_deserializer=storage.ReadRowsResponse.deserialize, + ) + return self._stubs["read_rows"] + + @property + def split_read_stream( + self, + ) -> Callable[ + [storage.SplitReadStreamRequest], Awaitable[storage.SplitReadStreamResponse] + ]: + r"""Return a callable for the split read stream method over gRPC. + + Splits a given ``ReadStream`` into two ``ReadStream`` objects. + These ``ReadStream`` objects are referred to as the primary and + the residual streams of the split. The original ``ReadStream`` + can still be read from in the same manner as before. Both of the + returned ``ReadStream`` objects can also be read from, and the + rows returned by both child streams will be the same as the rows + read from the original stream. + + Moreover, the two child streams will be allocated back-to-back + in the original ``ReadStream``. Concretely, it is guaranteed + that for streams original, primary, and residual, that + original[0-j] = primary[0-j] and original[j-n] = residual[0-m] + once the streams have been read to completion. + + Returns: + Callable[[~.SplitReadStreamRequest], + Awaitable[~.SplitReadStreamResponse]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "split_read_stream" not in self._stubs: + self._stubs["split_read_stream"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryRead/SplitReadStream", + request_serializer=storage.SplitReadStreamRequest.serialize, + response_deserializer=storage.SplitReadStreamResponse.deserialize, + ) + return self._stubs["split_read_stream"] + + +__all__ = ("BigQueryReadGrpcAsyncIOTransport",) diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_write/__init__.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/__init__.py new file mode 100644 index 00000000..89e849a9 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/__init__.py @@ -0,0 +1,24 @@ +# -*- 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. +# + +from .client import BigQueryWriteClient +from .async_client import BigQueryWriteAsyncClient + +__all__ = ( + "BigQueryWriteClient", + "BigQueryWriteAsyncClient", +) diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_write/async_client.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/async_client.py new file mode 100644 index 00000000..378bd9ef --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/async_client.py @@ -0,0 +1,671 @@ +# -*- 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. +# + +from collections import OrderedDict +import functools +import re +from typing import ( + Dict, + AsyncIterable, + Awaitable, + AsyncIterator, + Sequence, + Tuple, + Type, + Union, +) +import pkg_resources + +import google.api_core.client_options as ClientOptions # type: ignore +from google.api_core import exceptions # type: ignore +from google.api_core import gapic_v1 # type: ignore +from google.api_core import retry as retries # type: ignore +from google.auth import credentials # type: ignore +from google.oauth2 import service_account # type: ignore + +from google.cloud.bigquery_storage_v1beta2.types import storage +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.rpc import status_pb2 as status # type: ignore + +from .transports.base import BigQueryWriteTransport, DEFAULT_CLIENT_INFO +from .transports.grpc_asyncio import BigQueryWriteGrpcAsyncIOTransport +from .client import BigQueryWriteClient + + +class BigQueryWriteAsyncClient: + """BigQuery Write API. + The Write API can be used to write data to BigQuery. + """ + + _client: BigQueryWriteClient + + DEFAULT_ENDPOINT = BigQueryWriteClient.DEFAULT_ENDPOINT + DEFAULT_MTLS_ENDPOINT = BigQueryWriteClient.DEFAULT_MTLS_ENDPOINT + + table_path = staticmethod(BigQueryWriteClient.table_path) + parse_table_path = staticmethod(BigQueryWriteClient.parse_table_path) + write_stream_path = staticmethod(BigQueryWriteClient.write_stream_path) + parse_write_stream_path = staticmethod(BigQueryWriteClient.parse_write_stream_path) + + common_billing_account_path = staticmethod( + BigQueryWriteClient.common_billing_account_path + ) + parse_common_billing_account_path = staticmethod( + BigQueryWriteClient.parse_common_billing_account_path + ) + + common_folder_path = staticmethod(BigQueryWriteClient.common_folder_path) + parse_common_folder_path = staticmethod( + BigQueryWriteClient.parse_common_folder_path + ) + + common_organization_path = staticmethod( + BigQueryWriteClient.common_organization_path + ) + parse_common_organization_path = staticmethod( + BigQueryWriteClient.parse_common_organization_path + ) + + common_project_path = staticmethod(BigQueryWriteClient.common_project_path) + parse_common_project_path = staticmethod( + BigQueryWriteClient.parse_common_project_path + ) + + common_location_path = staticmethod(BigQueryWriteClient.common_location_path) + parse_common_location_path = staticmethod( + BigQueryWriteClient.parse_common_location_path + ) + + from_service_account_info = BigQueryWriteClient.from_service_account_info + from_service_account_file = BigQueryWriteClient.from_service_account_file + from_service_account_json = from_service_account_file + + @property + def transport(self) -> BigQueryWriteTransport: + """Return the transport used by the client instance. + + Returns: + BigQueryWriteTransport: The transport used by the client instance. + """ + return self._client.transport + + get_transport_class = functools.partial( + type(BigQueryWriteClient).get_transport_class, type(BigQueryWriteClient) + ) + + def __init__( + self, + *, + credentials: credentials.Credentials = None, + transport: Union[str, BigQueryWriteTransport] = "grpc_asyncio", + client_options: ClientOptions = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + ) -> None: + """Instantiate the big query write client. + + Args: + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + transport (Union[str, ~.BigQueryWriteTransport]): The + transport to use. If set to None, a transport is chosen + automatically. + client_options (ClientOptions): Custom options for the client. It + won't take effect if a ``transport`` instance is provided. + (1) The ``api_endpoint`` property can be used to override the + default endpoint provided by the client. GOOGLE_API_USE_MTLS_ENDPOINT + environment variable can also be used to override the endpoint: + "always" (always use the default mTLS endpoint), "never" (always + use the default regular endpoint) and "auto" (auto switch to the + default mTLS endpoint if client certificate is present, this is + the default value). However, the ``api_endpoint`` property takes + precedence if provided. + (2) If GOOGLE_API_USE_CLIENT_CERTIFICATE environment variable + is "true", then the ``client_cert_source`` property can be used + to provide client certificate for mutual TLS transport. If + not provided, the default SSL client certificate will be used if + present. If GOOGLE_API_USE_CLIENT_CERTIFICATE is "false" or not + set, no client certificate will be used. + + Raises: + google.auth.exceptions.MutualTlsChannelError: If mutual TLS transport + creation failed for any reason. + """ + + self._client = BigQueryWriteClient( + credentials=credentials, + transport=transport, + client_options=client_options, + client_info=client_info, + ) + + async def create_write_stream( + self, + request: storage.CreateWriteStreamRequest = None, + *, + parent: str = None, + write_stream: stream.WriteStream = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> stream.WriteStream: + r"""Creates a write stream to the given table. Additionally, every + table has a special COMMITTED stream named '_default' to which + data can be written. This stream doesn't need to be created + using CreateWriteStream. It is a stream that can be used + simultaneously by any number of clients. Data written to this + stream is considered committed as soon as an acknowledgement is + received. + + Args: + request (:class:`google.cloud.bigquery_storage_v1beta2.types.CreateWriteStreamRequest`): + The request object. Request message for + `CreateWriteStream`. + parent (:class:`str`): + Required. Reference to the table to which the stream + belongs, in the format of + ``projects/{project}/datasets/{dataset}/tables/{table}``. + + This corresponds to the ``parent`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + write_stream (:class:`google.cloud.bigquery_storage_v1beta2.types.WriteStream`): + Required. Stream to be created. + This corresponds to the ``write_stream`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.WriteStream: + Information about a single stream + that gets data inside the storage + system. + + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([parent, write_stream]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + request = storage.CreateWriteStreamRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if parent is not None: + request.parent = parent + if write_stream is not None: + request.write_stream = write_stream + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.create_write_stream, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, + exceptions.ResourceExhausted, + exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("parent", request.parent),)), + ) + + # Send the request. + response = await rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + def append_rows( + self, + requests: AsyncIterator[storage.AppendRowsRequest] = None, + *, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> Awaitable[AsyncIterable[storage.AppendRowsResponse]]: + r"""Appends data to the given stream. + + If ``offset`` is specified, the ``offset`` is checked against + the end of stream. The server returns ``OUT_OF_RANGE`` in + ``AppendRowsResponse`` if an attempt is made to append to an + offset beyond the current end of the stream or + ``ALREADY_EXISTS`` if user provids an ``offset`` that has + already been written to. User can retry with adjusted offset + within the same RPC stream. If ``offset`` is not specified, + append happens at the end of the stream. + + The response contains the offset at which the append happened. + Responses are received in the same order in which requests are + sent. There will be one response for each successful request. If + the ``offset`` is not set in response, it means append didn't + happen due to some errors. If one request fails, all the + subsequent requests will also fail until a success request is + made again. + + If the stream is of ``PENDING`` type, data will only be + available for read operations after the stream is committed. + + Args: + requests (AsyncIterator[`google.cloud.bigquery_storage_v1beta2.types.AppendRowsRequest`]): + The request object AsyncIterator. Request message for `AppendRows`. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + AsyncIterable[google.cloud.bigquery_storage_v1beta2.types.AppendRowsResponse]: + Response message for AppendRows. + """ + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.append_rows, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.ResourceExhausted, exceptions.ServiceUnavailable, + ), + ), + default_timeout=86400.0, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + (gapic_v1.routing_header.to_grpc_metadata(()),) + + # Send the request. + response = rpc(requests, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + async def get_write_stream( + self, + request: storage.GetWriteStreamRequest = None, + *, + name: str = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> stream.WriteStream: + r"""Gets a write stream. + + Args: + request (:class:`google.cloud.bigquery_storage_v1beta2.types.GetWriteStreamRequest`): + The request object. Request message for + `GetWriteStreamRequest`. + name (:class:`str`): + Required. Name of the stream to get, in the form of + ``projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}``. + + This corresponds to the ``name`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.WriteStream: + Information about a single stream + that gets data inside the storage + system. + + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([name]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + request = storage.GetWriteStreamRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if name is not None: + request.name = name + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.get_write_stream, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Send the request. + response = await rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + async def finalize_write_stream( + self, + request: storage.FinalizeWriteStreamRequest = None, + *, + name: str = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> storage.FinalizeWriteStreamResponse: + r"""Finalize a write stream so that no new data can be appended to + the stream. Finalize is not supported on the '_default' stream. + + Args: + request (:class:`google.cloud.bigquery_storage_v1beta2.types.FinalizeWriteStreamRequest`): + The request object. Request message for invoking + `FinalizeWriteStream`. + name (:class:`str`): + Required. Name of the stream to finalize, in the form of + ``projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}``. + + This corresponds to the ``name`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.FinalizeWriteStreamResponse: + Response message for FinalizeWriteStream. + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([name]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + request = storage.FinalizeWriteStreamRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if name is not None: + request.name = name + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.finalize_write_stream, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Send the request. + response = await rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + async def batch_commit_write_streams( + self, + request: storage.BatchCommitWriteStreamsRequest = None, + *, + parent: str = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> storage.BatchCommitWriteStreamsResponse: + r"""Atomically commits a group of ``PENDING`` streams that belong to + the same ``parent`` table. Streams must be finalized before + commit and cannot be committed multiple times. Once a stream is + committed, data in the stream becomes available for read + operations. + + Args: + request (:class:`google.cloud.bigquery_storage_v1beta2.types.BatchCommitWriteStreamsRequest`): + The request object. Request message for + `BatchCommitWriteStreams`. + parent (:class:`str`): + Required. Parent table that all the streams should + belong to, in the form of + ``projects/{project}/datasets/{dataset}/tables/{table}``. + + This corresponds to the ``parent`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.BatchCommitWriteStreamsResponse: + Response message for BatchCommitWriteStreams. + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([parent]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + request = storage.BatchCommitWriteStreamsRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if parent is not None: + request.parent = parent + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.batch_commit_write_streams, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("parent", request.parent),)), + ) + + # Send the request. + response = await rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + async def flush_rows( + self, + request: storage.FlushRowsRequest = None, + *, + write_stream: str = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> storage.FlushRowsResponse: + r"""Flushes rows to a BUFFERED stream. If users are appending rows + to BUFFERED stream, flush operation is required in order for the + rows to become available for reading. A Flush operation flushes + up to any previously flushed offset in a BUFFERED stream, to the + offset specified in the request. Flush is not supported on the + \_default stream, since it is not BUFFERED. + + Args: + request (:class:`google.cloud.bigquery_storage_v1beta2.types.FlushRowsRequest`): + The request object. Request message for `FlushRows`. + write_stream (:class:`str`): + Required. The stream that is the + target of the flush operation. + + This corresponds to the ``write_stream`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.FlushRowsResponse: + Respond message for FlushRows. + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([write_stream]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + request = storage.FlushRowsRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if write_stream is not None: + request.write_stream = write_stream + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.flush_rows, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata( + (("write_stream", request.write_stream),) + ), + ) + + # Send the request. + response = await rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + +try: + DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( + gapic_version=pkg_resources.get_distribution( + "google-cloud-bigquery-storage", + ).version, + ) +except pkg_resources.DistributionNotFound: + DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo() + + +__all__ = ("BigQueryWriteAsyncClient",) diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_write/client.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/client.py new file mode 100644 index 00000000..73cc2c15 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/client.py @@ -0,0 +1,843 @@ +# -*- 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. +# + +from collections import OrderedDict +from distutils import util +import os +import re +from typing import ( + Callable, + Dict, + Optional, + Iterable, + Iterator, + Sequence, + Tuple, + Type, + Union, +) +import pkg_resources + +from google.api_core import client_options as client_options_lib # type: ignore +from google.api_core import exceptions # type: ignore +from google.api_core import gapic_v1 # type: ignore +from google.api_core import retry as retries # type: ignore +from google.auth import credentials # type: ignore +from google.auth.transport import mtls # type: ignore +from google.auth.transport.grpc import SslCredentials # type: ignore +from google.auth.exceptions import MutualTLSChannelError # type: ignore +from google.oauth2 import service_account # type: ignore + +from google.cloud.bigquery_storage_v1beta2.types import storage +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.rpc import status_pb2 as status # type: ignore + +from .transports.base import BigQueryWriteTransport, DEFAULT_CLIENT_INFO +from .transports.grpc import BigQueryWriteGrpcTransport +from .transports.grpc_asyncio import BigQueryWriteGrpcAsyncIOTransport + + +class BigQueryWriteClientMeta(type): + """Metaclass for the BigQueryWrite client. + + This provides class-level methods for building and retrieving + support objects (e.g. transport) without polluting the client instance + objects. + """ + + _transport_registry = OrderedDict() # type: Dict[str, Type[BigQueryWriteTransport]] + _transport_registry["grpc"] = BigQueryWriteGrpcTransport + _transport_registry["grpc_asyncio"] = BigQueryWriteGrpcAsyncIOTransport + + def get_transport_class(cls, label: str = None,) -> Type[BigQueryWriteTransport]: + """Return an appropriate transport class. + + Args: + label: The name of the desired transport. If none is + provided, then the first transport in the registry is used. + + Returns: + The transport class to use. + """ + # If a specific transport is requested, return that one. + if label: + return cls._transport_registry[label] + + # No transport is requested; return the default (that is, the first one + # in the dictionary). + return next(iter(cls._transport_registry.values())) + + +class BigQueryWriteClient(metaclass=BigQueryWriteClientMeta): + """BigQuery Write API. + The Write API can be used to write data to BigQuery. + """ + + @staticmethod + def _get_default_mtls_endpoint(api_endpoint): + """Convert api endpoint to mTLS endpoint. + Convert "*.sandbox.googleapis.com" and "*.googleapis.com" to + "*.mtls.sandbox.googleapis.com" and "*.mtls.googleapis.com" respectively. + Args: + api_endpoint (Optional[str]): the api endpoint to convert. + Returns: + str: converted mTLS api endpoint. + """ + if not api_endpoint: + return api_endpoint + + mtls_endpoint_re = re.compile( + r"(?P[^.]+)(?P\.mtls)?(?P\.sandbox)?(?P\.googleapis\.com)?" + ) + + m = mtls_endpoint_re.match(api_endpoint) + name, mtls, sandbox, googledomain = m.groups() + if mtls or not googledomain: + return api_endpoint + + if sandbox: + return api_endpoint.replace( + "sandbox.googleapis.com", "mtls.sandbox.googleapis.com" + ) + + return api_endpoint.replace(".googleapis.com", ".mtls.googleapis.com") + + DEFAULT_ENDPOINT = "bigquerystorage.googleapis.com" + DEFAULT_MTLS_ENDPOINT = _get_default_mtls_endpoint.__func__( # type: ignore + DEFAULT_ENDPOINT + ) + + @classmethod + def from_service_account_info(cls, info: dict, *args, **kwargs): + """Creates an instance of this client using the provided credentials info. + + Args: + info (dict): The service account private key info. + args: Additional arguments to pass to the constructor. + kwargs: Additional arguments to pass to the constructor. + + Returns: + BigQueryWriteClient: The constructed client. + """ + credentials = service_account.Credentials.from_service_account_info(info) + kwargs["credentials"] = credentials + return cls(*args, **kwargs) + + @classmethod + def from_service_account_file(cls, filename: str, *args, **kwargs): + """Creates an instance of this client using the provided credentials + file. + + Args: + filename (str): The path to the service account private key json + file. + args: Additional arguments to pass to the constructor. + kwargs: Additional arguments to pass to the constructor. + + Returns: + BigQueryWriteClient: The constructed client. + """ + credentials = service_account.Credentials.from_service_account_file(filename) + kwargs["credentials"] = credentials + return cls(*args, **kwargs) + + from_service_account_json = from_service_account_file + + @property + def transport(self) -> BigQueryWriteTransport: + """Return the transport used by the client instance. + + Returns: + BigQueryWriteTransport: The transport used by the client instance. + """ + return self._transport + + @staticmethod + def table_path(project: str, dataset: str, table: str,) -> str: + """Return a fully-qualified table string.""" + return "projects/{project}/datasets/{dataset}/tables/{table}".format( + project=project, dataset=dataset, table=table, + ) + + @staticmethod + def parse_table_path(path: str) -> Dict[str, str]: + """Parse a table path into its component segments.""" + m = re.match( + r"^projects/(?P.+?)/datasets/(?P.+?)/tables/(?P
.+?)$", + path, + ) + return m.groupdict() if m else {} + + @staticmethod + def write_stream_path(project: str, dataset: str, table: str, stream: str,) -> str: + """Return a fully-qualified write_stream string.""" + return "projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}".format( + project=project, dataset=dataset, table=table, stream=stream, + ) + + @staticmethod + def parse_write_stream_path(path: str) -> Dict[str, str]: + """Parse a write_stream path into its component segments.""" + m = re.match( + r"^projects/(?P.+?)/datasets/(?P.+?)/tables/(?P
.+?)/streams/(?P.+?)$", + path, + ) + return m.groupdict() if m else {} + + @staticmethod + def common_billing_account_path(billing_account: str,) -> str: + """Return a fully-qualified billing_account string.""" + return "billingAccounts/{billing_account}".format( + billing_account=billing_account, + ) + + @staticmethod + def parse_common_billing_account_path(path: str) -> Dict[str, str]: + """Parse a billing_account path into its component segments.""" + m = re.match(r"^billingAccounts/(?P.+?)$", path) + return m.groupdict() if m else {} + + @staticmethod + def common_folder_path(folder: str,) -> str: + """Return a fully-qualified folder string.""" + return "folders/{folder}".format(folder=folder,) + + @staticmethod + def parse_common_folder_path(path: str) -> Dict[str, str]: + """Parse a folder path into its component segments.""" + m = re.match(r"^folders/(?P.+?)$", path) + return m.groupdict() if m else {} + + @staticmethod + def common_organization_path(organization: str,) -> str: + """Return a fully-qualified organization string.""" + return "organizations/{organization}".format(organization=organization,) + + @staticmethod + def parse_common_organization_path(path: str) -> Dict[str, str]: + """Parse a organization path into its component segments.""" + m = re.match(r"^organizations/(?P.+?)$", path) + return m.groupdict() if m else {} + + @staticmethod + def common_project_path(project: str,) -> str: + """Return a fully-qualified project string.""" + return "projects/{project}".format(project=project,) + + @staticmethod + def parse_common_project_path(path: str) -> Dict[str, str]: + """Parse a project path into its component segments.""" + m = re.match(r"^projects/(?P.+?)$", path) + return m.groupdict() if m else {} + + @staticmethod + def common_location_path(project: str, location: str,) -> str: + """Return a fully-qualified location string.""" + return "projects/{project}/locations/{location}".format( + project=project, location=location, + ) + + @staticmethod + def parse_common_location_path(path: str) -> Dict[str, str]: + """Parse a location path into its component segments.""" + m = re.match(r"^projects/(?P.+?)/locations/(?P.+?)$", path) + return m.groupdict() if m else {} + + def __init__( + self, + *, + credentials: Optional[credentials.Credentials] = None, + transport: Union[str, BigQueryWriteTransport, None] = None, + client_options: Optional[client_options_lib.ClientOptions] = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + ) -> None: + """Instantiate the big query write client. + + Args: + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + transport (Union[str, BigQueryWriteTransport]): The + transport to use. If set to None, a transport is chosen + automatically. + client_options (google.api_core.client_options.ClientOptions): Custom options for the + client. It won't take effect if a ``transport`` instance is provided. + (1) The ``api_endpoint`` property can be used to override the + default endpoint provided by the client. GOOGLE_API_USE_MTLS_ENDPOINT + environment variable can also be used to override the endpoint: + "always" (always use the default mTLS endpoint), "never" (always + use the default regular endpoint) and "auto" (auto switch to the + default mTLS endpoint if client certificate is present, this is + the default value). However, the ``api_endpoint`` property takes + precedence if provided. + (2) If GOOGLE_API_USE_CLIENT_CERTIFICATE environment variable + is "true", then the ``client_cert_source`` property can be used + to provide client certificate for mutual TLS transport. If + not provided, the default SSL client certificate will be used if + present. If GOOGLE_API_USE_CLIENT_CERTIFICATE is "false" or not + set, no client certificate will be used. + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + + Raises: + google.auth.exceptions.MutualTLSChannelError: If mutual TLS transport + creation failed for any reason. + """ + if isinstance(client_options, dict): + client_options = client_options_lib.from_dict(client_options) + if client_options is None: + client_options = client_options_lib.ClientOptions() + + # Create SSL credentials for mutual TLS if needed. + use_client_cert = bool( + util.strtobool(os.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false")) + ) + + ssl_credentials = None + is_mtls = False + if use_client_cert: + if client_options.client_cert_source: + import grpc # type: ignore + + cert, key = client_options.client_cert_source() + ssl_credentials = grpc.ssl_channel_credentials( + certificate_chain=cert, private_key=key + ) + is_mtls = True + else: + creds = SslCredentials() + is_mtls = creds.is_mtls + ssl_credentials = creds.ssl_credentials if is_mtls else None + + # Figure out which api endpoint to use. + if client_options.api_endpoint is not None: + api_endpoint = client_options.api_endpoint + else: + use_mtls_env = os.getenv("GOOGLE_API_USE_MTLS_ENDPOINT", "auto") + if use_mtls_env == "never": + api_endpoint = self.DEFAULT_ENDPOINT + elif use_mtls_env == "always": + api_endpoint = self.DEFAULT_MTLS_ENDPOINT + elif use_mtls_env == "auto": + api_endpoint = ( + self.DEFAULT_MTLS_ENDPOINT if is_mtls else self.DEFAULT_ENDPOINT + ) + else: + raise MutualTLSChannelError( + "Unsupported GOOGLE_API_USE_MTLS_ENDPOINT value. Accepted values: never, auto, always" + ) + + # Save or instantiate the transport. + # Ordinarily, we provide the transport, but allowing a custom transport + # instance provides an extensibility point for unusual situations. + if isinstance(transport, BigQueryWriteTransport): + # transport is a BigQueryWriteTransport instance. + if credentials or client_options.credentials_file: + raise ValueError( + "When providing a transport instance, " + "provide its credentials directly." + ) + if client_options.scopes: + raise ValueError( + "When providing a transport instance, " + "provide its scopes directly." + ) + self._transport = transport + else: + Transport = type(self).get_transport_class(transport) + self._transport = Transport( + credentials=credentials, + credentials_file=client_options.credentials_file, + host=api_endpoint, + scopes=client_options.scopes, + ssl_channel_credentials=ssl_credentials, + quota_project_id=client_options.quota_project_id, + client_info=client_info, + ) + + def create_write_stream( + self, + request: storage.CreateWriteStreamRequest = None, + *, + parent: str = None, + write_stream: stream.WriteStream = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> stream.WriteStream: + r"""Creates a write stream to the given table. Additionally, every + table has a special COMMITTED stream named '_default' to which + data can be written. This stream doesn't need to be created + using CreateWriteStream. It is a stream that can be used + simultaneously by any number of clients. Data written to this + stream is considered committed as soon as an acknowledgement is + received. + + Args: + request (google.cloud.bigquery_storage_v1beta2.types.CreateWriteStreamRequest): + The request object. Request message for + `CreateWriteStream`. + parent (str): + Required. Reference to the table to which the stream + belongs, in the format of + ``projects/{project}/datasets/{dataset}/tables/{table}``. + + This corresponds to the ``parent`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + write_stream (google.cloud.bigquery_storage_v1beta2.types.WriteStream): + Required. Stream to be created. + This corresponds to the ``write_stream`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.WriteStream: + Information about a single stream + that gets data inside the storage + system. + + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([parent, write_stream]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + # Minor optimization to avoid making a copy if the user passes + # in a storage.CreateWriteStreamRequest. + # There's no risk of modifying the input as we've already verified + # there are no flattened fields. + if not isinstance(request, storage.CreateWriteStreamRequest): + request = storage.CreateWriteStreamRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if parent is not None: + request.parent = parent + if write_stream is not None: + request.write_stream = write_stream + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[self._transport.create_write_stream] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("parent", request.parent),)), + ) + + # Send the request. + response = rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + def append_rows( + self, + requests: Iterator[storage.AppendRowsRequest] = None, + *, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> Iterable[storage.AppendRowsResponse]: + r"""Appends data to the given stream. + + If ``offset`` is specified, the ``offset`` is checked against + the end of stream. The server returns ``OUT_OF_RANGE`` in + ``AppendRowsResponse`` if an attempt is made to append to an + offset beyond the current end of the stream or + ``ALREADY_EXISTS`` if user provids an ``offset`` that has + already been written to. User can retry with adjusted offset + within the same RPC stream. If ``offset`` is not specified, + append happens at the end of the stream. + + The response contains the offset at which the append happened. + Responses are received in the same order in which requests are + sent. There will be one response for each successful request. If + the ``offset`` is not set in response, it means append didn't + happen due to some errors. If one request fails, all the + subsequent requests will also fail until a success request is + made again. + + If the stream is of ``PENDING`` type, data will only be + available for read operations after the stream is committed. + + Args: + requests (Iterator[google.cloud.bigquery_storage_v1beta2.types.AppendRowsRequest]): + The request object iterator. Request message for `AppendRows`. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + Iterable[google.cloud.bigquery_storage_v1beta2.types.AppendRowsResponse]: + Response message for AppendRows. + """ + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[self._transport.append_rows] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + (gapic_v1.routing_header.to_grpc_metadata(()),) + + # Send the request. + response = rpc(requests, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + def get_write_stream( + self, + request: storage.GetWriteStreamRequest = None, + *, + name: str = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> stream.WriteStream: + r"""Gets a write stream. + + Args: + request (google.cloud.bigquery_storage_v1beta2.types.GetWriteStreamRequest): + The request object. Request message for + `GetWriteStreamRequest`. + name (str): + Required. Name of the stream to get, in the form of + ``projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}``. + + This corresponds to the ``name`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.WriteStream: + Information about a single stream + that gets data inside the storage + system. + + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([name]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + # Minor optimization to avoid making a copy if the user passes + # in a storage.GetWriteStreamRequest. + # There's no risk of modifying the input as we've already verified + # there are no flattened fields. + if not isinstance(request, storage.GetWriteStreamRequest): + request = storage.GetWriteStreamRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if name is not None: + request.name = name + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[self._transport.get_write_stream] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Send the request. + response = rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + def finalize_write_stream( + self, + request: storage.FinalizeWriteStreamRequest = None, + *, + name: str = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> storage.FinalizeWriteStreamResponse: + r"""Finalize a write stream so that no new data can be appended to + the stream. Finalize is not supported on the '_default' stream. + + Args: + request (google.cloud.bigquery_storage_v1beta2.types.FinalizeWriteStreamRequest): + The request object. Request message for invoking + `FinalizeWriteStream`. + name (str): + Required. Name of the stream to finalize, in the form of + ``projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}``. + + This corresponds to the ``name`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.FinalizeWriteStreamResponse: + Response message for FinalizeWriteStream. + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([name]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + # Minor optimization to avoid making a copy if the user passes + # in a storage.FinalizeWriteStreamRequest. + # There's no risk of modifying the input as we've already verified + # there are no flattened fields. + if not isinstance(request, storage.FinalizeWriteStreamRequest): + request = storage.FinalizeWriteStreamRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if name is not None: + request.name = name + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[self._transport.finalize_write_stream] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Send the request. + response = rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + def batch_commit_write_streams( + self, + request: storage.BatchCommitWriteStreamsRequest = None, + *, + parent: str = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> storage.BatchCommitWriteStreamsResponse: + r"""Atomically commits a group of ``PENDING`` streams that belong to + the same ``parent`` table. Streams must be finalized before + commit and cannot be committed multiple times. Once a stream is + committed, data in the stream becomes available for read + operations. + + Args: + request (google.cloud.bigquery_storage_v1beta2.types.BatchCommitWriteStreamsRequest): + The request object. Request message for + `BatchCommitWriteStreams`. + parent (str): + Required. Parent table that all the streams should + belong to, in the form of + ``projects/{project}/datasets/{dataset}/tables/{table}``. + + This corresponds to the ``parent`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.BatchCommitWriteStreamsResponse: + Response message for BatchCommitWriteStreams. + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([parent]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + # Minor optimization to avoid making a copy if the user passes + # in a storage.BatchCommitWriteStreamsRequest. + # There's no risk of modifying the input as we've already verified + # there are no flattened fields. + if not isinstance(request, storage.BatchCommitWriteStreamsRequest): + request = storage.BatchCommitWriteStreamsRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if parent is not None: + request.parent = parent + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[ + self._transport.batch_commit_write_streams + ] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("parent", request.parent),)), + ) + + # Send the request. + response = rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + def flush_rows( + self, + request: storage.FlushRowsRequest = None, + *, + write_stream: str = None, + retry: retries.Retry = gapic_v1.method.DEFAULT, + timeout: float = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> storage.FlushRowsResponse: + r"""Flushes rows to a BUFFERED stream. If users are appending rows + to BUFFERED stream, flush operation is required in order for the + rows to become available for reading. A Flush operation flushes + up to any previously flushed offset in a BUFFERED stream, to the + offset specified in the request. Flush is not supported on the + \_default stream, since it is not BUFFERED. + + Args: + request (google.cloud.bigquery_storage_v1beta2.types.FlushRowsRequest): + The request object. Request message for `FlushRows`. + write_stream (str): + Required. The stream that is the + target of the flush operation. + + This corresponds to the ``write_stream`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.bigquery_storage_v1beta2.types.FlushRowsResponse: + Respond message for FlushRows. + """ + # Create or coerce a protobuf request object. + # Sanity check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([write_stream]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + # Minor optimization to avoid making a copy if the user passes + # in a storage.FlushRowsRequest. + # There's no risk of modifying the input as we've already verified + # there are no flattened fields. + if not isinstance(request, storage.FlushRowsRequest): + request = storage.FlushRowsRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + + if write_stream is not None: + request.write_stream = write_stream + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[self._transport.flush_rows] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata( + (("write_stream", request.write_stream),) + ), + ) + + # Send the request. + response = rpc(request, retry=retry, timeout=timeout, metadata=metadata,) + + # Done; return the response. + return response + + +try: + DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( + gapic_version=pkg_resources.get_distribution( + "google-cloud-bigquery-storage", + ).version, + ) +except pkg_resources.DistributionNotFound: + DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo() + + +__all__ = ("BigQueryWriteClient",) diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_write/transports/__init__.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/transports/__init__.py new file mode 100644 index 00000000..33e48de0 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/transports/__init__.py @@ -0,0 +1,35 @@ +# -*- 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. +# + +from collections import OrderedDict +from typing import Dict, Type + +from .base import BigQueryWriteTransport +from .grpc import BigQueryWriteGrpcTransport +from .grpc_asyncio import BigQueryWriteGrpcAsyncIOTransport + + +# Compile a registry of transports. +_transport_registry = OrderedDict() # type: Dict[str, Type[BigQueryWriteTransport]] +_transport_registry["grpc"] = BigQueryWriteGrpcTransport +_transport_registry["grpc_asyncio"] = BigQueryWriteGrpcAsyncIOTransport + +__all__ = ( + "BigQueryWriteTransport", + "BigQueryWriteGrpcTransport", + "BigQueryWriteGrpcAsyncIOTransport", +) diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_write/transports/base.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/transports/base.py new file mode 100644 index 00000000..6b9b6a18 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/transports/base.py @@ -0,0 +1,261 @@ +# -*- 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 abc +import typing +import pkg_resources + +from google import auth # type: ignore +from google.api_core import exceptions # type: ignore +from google.api_core import gapic_v1 # type: ignore +from google.api_core import retry as retries # type: ignore +from google.auth import credentials # type: ignore + +from google.cloud.bigquery_storage_v1beta2.types import storage +from google.cloud.bigquery_storage_v1beta2.types import stream + + +try: + DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( + gapic_version=pkg_resources.get_distribution( + "google-cloud-bigquery-storage", + ).version, + ) +except pkg_resources.DistributionNotFound: + DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo() + + +class BigQueryWriteTransport(abc.ABC): + """Abstract transport class for BigQueryWrite.""" + + AUTH_SCOPES = ( + "https://www.googleapis.com/auth/bigquery", + "https://www.googleapis.com/auth/bigquery.insertdata", + "https://www.googleapis.com/auth/cloud-platform", + ) + + def __init__( + self, + *, + host: str = "bigquerystorage.googleapis.com", + credentials: credentials.Credentials = None, + credentials_file: typing.Optional[str] = None, + scopes: typing.Optional[typing.Sequence[str]] = AUTH_SCOPES, + quota_project_id: typing.Optional[str] = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + **kwargs, + ) -> None: + """Instantiate the transport. + + Args: + host (Optional[str]): The hostname to connect to. + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is mutually exclusive with credentials. + scope (Optional[Sequence[str]]): A list of scopes. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + """ + # Save the hostname. Default to port 443 (HTTPS) if none is specified. + if ":" not in host: + host += ":443" + self._host = host + + # If no credentials are provided, then determine the appropriate + # defaults. + if credentials and credentials_file: + raise exceptions.DuplicateCredentialArgs( + "'credentials_file' and 'credentials' are mutually exclusive" + ) + + if credentials_file is not None: + credentials, _ = auth.load_credentials_from_file( + credentials_file, scopes=scopes, quota_project_id=quota_project_id + ) + + elif credentials is None: + credentials, _ = auth.default( + scopes=scopes, quota_project_id=quota_project_id + ) + + # Save the credentials. + self._credentials = credentials + + # Lifted into its own function so it can be stubbed out during tests. + self._prep_wrapped_messages(client_info) + + def _prep_wrapped_messages(self, client_info): + # Precompute the wrapped methods. + self._wrapped_methods = { + self.create_write_stream: gapic_v1.method.wrap_method( + self.create_write_stream, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, + exceptions.ResourceExhausted, + exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=client_info, + ), + self.append_rows: gapic_v1.method.wrap_method( + self.append_rows, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.ResourceExhausted, exceptions.ServiceUnavailable, + ), + ), + default_timeout=86400.0, + client_info=client_info, + ), + self.get_write_stream: gapic_v1.method.wrap_method( + self.get_write_stream, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=client_info, + ), + self.finalize_write_stream: gapic_v1.method.wrap_method( + self.finalize_write_stream, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=client_info, + ), + self.batch_commit_write_streams: gapic_v1.method.wrap_method( + self.batch_commit_write_streams, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=client_info, + ), + self.flush_rows: gapic_v1.method.wrap_method( + self.flush_rows, + default_retry=retries.Retry( + initial=0.1, + maximum=60.0, + multiplier=1.3, + predicate=retries.if_exception_type( + exceptions.DeadlineExceeded, exceptions.ServiceUnavailable, + ), + ), + default_timeout=600.0, + client_info=client_info, + ), + } + + @property + def create_write_stream( + self, + ) -> typing.Callable[ + [storage.CreateWriteStreamRequest], + typing.Union[stream.WriteStream, typing.Awaitable[stream.WriteStream]], + ]: + raise NotImplementedError() + + @property + def append_rows( + self, + ) -> typing.Callable[ + [storage.AppendRowsRequest], + typing.Union[ + storage.AppendRowsResponse, typing.Awaitable[storage.AppendRowsResponse] + ], + ]: + raise NotImplementedError() + + @property + def get_write_stream( + self, + ) -> typing.Callable[ + [storage.GetWriteStreamRequest], + typing.Union[stream.WriteStream, typing.Awaitable[stream.WriteStream]], + ]: + raise NotImplementedError() + + @property + def finalize_write_stream( + self, + ) -> typing.Callable[ + [storage.FinalizeWriteStreamRequest], + typing.Union[ + storage.FinalizeWriteStreamResponse, + typing.Awaitable[storage.FinalizeWriteStreamResponse], + ], + ]: + raise NotImplementedError() + + @property + def batch_commit_write_streams( + self, + ) -> typing.Callable[ + [storage.BatchCommitWriteStreamsRequest], + typing.Union[ + storage.BatchCommitWriteStreamsResponse, + typing.Awaitable[storage.BatchCommitWriteStreamsResponse], + ], + ]: + raise NotImplementedError() + + @property + def flush_rows( + self, + ) -> typing.Callable[ + [storage.FlushRowsRequest], + typing.Union[ + storage.FlushRowsResponse, typing.Awaitable[storage.FlushRowsResponse] + ], + ]: + raise NotImplementedError() + + +__all__ = ("BigQueryWriteTransport",) diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_write/transports/grpc.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/transports/grpc.py new file mode 100644 index 00000000..d4638404 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/transports/grpc.py @@ -0,0 +1,440 @@ +# -*- 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 warnings +from typing import Callable, Dict, Optional, Sequence, Tuple + +from google.api_core import grpc_helpers # type: ignore +from google.api_core import gapic_v1 # type: ignore +from google import auth # type: ignore +from google.auth import credentials # type: ignore +from google.auth.transport.grpc import SslCredentials # type: ignore + +import grpc # type: ignore + +from google.cloud.bigquery_storage_v1beta2.types import storage +from google.cloud.bigquery_storage_v1beta2.types import stream + +from .base import BigQueryWriteTransport, DEFAULT_CLIENT_INFO + + +class BigQueryWriteGrpcTransport(BigQueryWriteTransport): + """gRPC backend transport for BigQueryWrite. + + BigQuery Write API. + The Write API can be used to write data to BigQuery. + + This class defines the same methods as the primary client, so the + primary client can load the underlying transport implementation + and call it. + + It sends protocol buffers over the wire using gRPC (which is built on + top of HTTP/2); the ``grpcio`` package must be installed. + """ + + _stubs: Dict[str, Callable] + + def __init__( + self, + *, + host: str = "bigquerystorage.googleapis.com", + credentials: credentials.Credentials = None, + credentials_file: str = None, + scopes: Sequence[str] = None, + channel: grpc.Channel = None, + api_mtls_endpoint: str = None, + client_cert_source: Callable[[], Tuple[bytes, bytes]] = None, + ssl_channel_credentials: grpc.ChannelCredentials = None, + quota_project_id: Optional[str] = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + ) -> None: + """Instantiate the transport. + + Args: + host (Optional[str]): The hostname to connect to. + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + This argument is ignored if ``channel`` is provided. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is ignored if ``channel`` is provided. + scopes (Optional(Sequence[str])): A list of scopes. This argument is + ignored if ``channel`` is provided. + channel (Optional[grpc.Channel]): A ``Channel`` instance through + which to make calls. + api_mtls_endpoint (Optional[str]): Deprecated. The mutual TLS endpoint. + If provided, it overrides the ``host`` argument and tries to create + a mutual TLS channel with client SSL credentials from + ``client_cert_source`` or applicatin default SSL credentials. + client_cert_source (Optional[Callable[[], Tuple[bytes, bytes]]]): + Deprecated. A callback to provide client SSL certificate bytes and + private key bytes, both in PEM format. It is ignored if + ``api_mtls_endpoint`` is None. + ssl_channel_credentials (grpc.ChannelCredentials): SSL credentials + for grpc channel. It is ignored if ``channel`` is provided. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + + Raises: + google.auth.exceptions.MutualTLSChannelError: If mutual TLS transport + creation failed for any reason. + google.api_core.exceptions.DuplicateCredentialArgs: If both ``credentials`` + and ``credentials_file`` are passed. + """ + self._ssl_channel_credentials = ssl_channel_credentials + + if channel: + # Sanity check: Ensure that channel and credentials are not both + # provided. + credentials = False + + # If a channel was explicitly provided, set it. + self._grpc_channel = channel + self._ssl_channel_credentials = None + elif api_mtls_endpoint: + warnings.warn( + "api_mtls_endpoint and client_cert_source are deprecated", + DeprecationWarning, + ) + + host = ( + api_mtls_endpoint + if ":" in api_mtls_endpoint + else api_mtls_endpoint + ":443" + ) + + if credentials is None: + credentials, _ = auth.default( + scopes=self.AUTH_SCOPES, quota_project_id=quota_project_id + ) + + # Create SSL credentials with client_cert_source or application + # default SSL credentials. + if client_cert_source: + cert, key = client_cert_source() + ssl_credentials = grpc.ssl_channel_credentials( + certificate_chain=cert, private_key=key + ) + else: + ssl_credentials = SslCredentials().ssl_credentials + + # create a new channel. The provided one is ignored. + self._grpc_channel = type(self).create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + ssl_credentials=ssl_credentials, + scopes=scopes or self.AUTH_SCOPES, + quota_project_id=quota_project_id, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + self._ssl_channel_credentials = ssl_credentials + else: + host = host if ":" in host else host + ":443" + + if credentials is None: + credentials, _ = auth.default( + scopes=self.AUTH_SCOPES, quota_project_id=quota_project_id + ) + + # create a new channel. The provided one is ignored. + self._grpc_channel = type(self).create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + ssl_credentials=ssl_channel_credentials, + scopes=scopes or self.AUTH_SCOPES, + quota_project_id=quota_project_id, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + + self._stubs = {} # type: Dict[str, Callable] + + # Run the base constructor. + super().__init__( + host=host, + credentials=credentials, + credentials_file=credentials_file, + scopes=scopes or self.AUTH_SCOPES, + quota_project_id=quota_project_id, + client_info=client_info, + ) + + @classmethod + def create_channel( + cls, + host: str = "bigquerystorage.googleapis.com", + credentials: credentials.Credentials = None, + credentials_file: str = None, + scopes: Optional[Sequence[str]] = None, + quota_project_id: Optional[str] = None, + **kwargs, + ) -> grpc.Channel: + """Create and return a gRPC channel object. + Args: + address (Optional[str]): The host for the channel to use. + credentials (Optional[~.Credentials]): The + authorization credentials to attach to requests. These + credentials identify this application to the service. If + none are specified, the client will attempt to ascertain + the credentials from the environment. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is mutually exclusive with credentials. + scopes (Optional[Sequence[str]]): A optional list of scopes needed for this + service. These are only used when credentials are not specified and + are passed to :func:`google.auth.default`. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + kwargs (Optional[dict]): Keyword arguments, which are passed to the + channel creation. + Returns: + grpc.Channel: A gRPC channel object. + + Raises: + google.api_core.exceptions.DuplicateCredentialArgs: If both ``credentials`` + and ``credentials_file`` are passed. + """ + scopes = scopes or cls.AUTH_SCOPES + return grpc_helpers.create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + scopes=scopes, + quota_project_id=quota_project_id, + **kwargs, + ) + + @property + def grpc_channel(self) -> grpc.Channel: + """Return the channel designed to connect to this service. + """ + return self._grpc_channel + + @property + def create_write_stream( + self, + ) -> Callable[[storage.CreateWriteStreamRequest], stream.WriteStream]: + r"""Return a callable for the create write stream method over gRPC. + + Creates a write stream to the given table. Additionally, every + table has a special COMMITTED stream named '_default' to which + data can be written. This stream doesn't need to be created + using CreateWriteStream. It is a stream that can be used + simultaneously by any number of clients. Data written to this + stream is considered committed as soon as an acknowledgement is + received. + + Returns: + Callable[[~.CreateWriteStreamRequest], + ~.WriteStream]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "create_write_stream" not in self._stubs: + self._stubs["create_write_stream"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryWrite/CreateWriteStream", + request_serializer=storage.CreateWriteStreamRequest.serialize, + response_deserializer=stream.WriteStream.deserialize, + ) + return self._stubs["create_write_stream"] + + @property + def append_rows( + self, + ) -> Callable[[storage.AppendRowsRequest], storage.AppendRowsResponse]: + r"""Return a callable for the append rows method over gRPC. + + Appends data to the given stream. + + If ``offset`` is specified, the ``offset`` is checked against + the end of stream. The server returns ``OUT_OF_RANGE`` in + ``AppendRowsResponse`` if an attempt is made to append to an + offset beyond the current end of the stream or + ``ALREADY_EXISTS`` if user provids an ``offset`` that has + already been written to. User can retry with adjusted offset + within the same RPC stream. If ``offset`` is not specified, + append happens at the end of the stream. + + The response contains the offset at which the append happened. + Responses are received in the same order in which requests are + sent. There will be one response for each successful request. If + the ``offset`` is not set in response, it means append didn't + happen due to some errors. If one request fails, all the + subsequent requests will also fail until a success request is + made again. + + If the stream is of ``PENDING`` type, data will only be + available for read operations after the stream is committed. + + Returns: + Callable[[~.AppendRowsRequest], + ~.AppendRowsResponse]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "append_rows" not in self._stubs: + self._stubs["append_rows"] = self.grpc_channel.stream_stream( + "/google.cloud.bigquery.storage.v1beta2.BigQueryWrite/AppendRows", + request_serializer=storage.AppendRowsRequest.serialize, + response_deserializer=storage.AppendRowsResponse.deserialize, + ) + return self._stubs["append_rows"] + + @property + def get_write_stream( + self, + ) -> Callable[[storage.GetWriteStreamRequest], stream.WriteStream]: + r"""Return a callable for the get write stream method over gRPC. + + Gets a write stream. + + Returns: + Callable[[~.GetWriteStreamRequest], + ~.WriteStream]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "get_write_stream" not in self._stubs: + self._stubs["get_write_stream"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryWrite/GetWriteStream", + request_serializer=storage.GetWriteStreamRequest.serialize, + response_deserializer=stream.WriteStream.deserialize, + ) + return self._stubs["get_write_stream"] + + @property + def finalize_write_stream( + self, + ) -> Callable[ + [storage.FinalizeWriteStreamRequest], storage.FinalizeWriteStreamResponse + ]: + r"""Return a callable for the finalize write stream method over gRPC. + + Finalize a write stream so that no new data can be appended to + the stream. Finalize is not supported on the '_default' stream. + + Returns: + Callable[[~.FinalizeWriteStreamRequest], + ~.FinalizeWriteStreamResponse]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "finalize_write_stream" not in self._stubs: + self._stubs["finalize_write_stream"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryWrite/FinalizeWriteStream", + request_serializer=storage.FinalizeWriteStreamRequest.serialize, + response_deserializer=storage.FinalizeWriteStreamResponse.deserialize, + ) + return self._stubs["finalize_write_stream"] + + @property + def batch_commit_write_streams( + self, + ) -> Callable[ + [storage.BatchCommitWriteStreamsRequest], + storage.BatchCommitWriteStreamsResponse, + ]: + r"""Return a callable for the batch commit write streams method over gRPC. + + Atomically commits a group of ``PENDING`` streams that belong to + the same ``parent`` table. Streams must be finalized before + commit and cannot be committed multiple times. Once a stream is + committed, data in the stream becomes available for read + operations. + + Returns: + Callable[[~.BatchCommitWriteStreamsRequest], + ~.BatchCommitWriteStreamsResponse]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "batch_commit_write_streams" not in self._stubs: + self._stubs["batch_commit_write_streams"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryWrite/BatchCommitWriteStreams", + request_serializer=storage.BatchCommitWriteStreamsRequest.serialize, + response_deserializer=storage.BatchCommitWriteStreamsResponse.deserialize, + ) + return self._stubs["batch_commit_write_streams"] + + @property + def flush_rows( + self, + ) -> Callable[[storage.FlushRowsRequest], storage.FlushRowsResponse]: + r"""Return a callable for the flush rows method over gRPC. + + Flushes rows to a BUFFERED stream. If users are appending rows + to BUFFERED stream, flush operation is required in order for the + rows to become available for reading. A Flush operation flushes + up to any previously flushed offset in a BUFFERED stream, to the + offset specified in the request. Flush is not supported on the + \_default stream, since it is not BUFFERED. + + Returns: + Callable[[~.FlushRowsRequest], + ~.FlushRowsResponse]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "flush_rows" not in self._stubs: + self._stubs["flush_rows"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryWrite/FlushRows", + request_serializer=storage.FlushRowsRequest.serialize, + response_deserializer=storage.FlushRowsResponse.deserialize, + ) + return self._stubs["flush_rows"] + + +__all__ = ("BigQueryWriteGrpcTransport",) diff --git a/google/cloud/bigquery_storage_v1beta2/services/big_query_write/transports/grpc_asyncio.py b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/transports/grpc_asyncio.py new file mode 100644 index 00000000..c41a7daa --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/services/big_query_write/transports/grpc_asyncio.py @@ -0,0 +1,445 @@ +# -*- 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 warnings +from typing import Awaitable, Callable, Dict, Optional, Sequence, Tuple + +from google.api_core import gapic_v1 # type: ignore +from google.api_core import grpc_helpers_async # type: ignore +from google import auth # type: ignore +from google.auth import credentials # type: ignore +from google.auth.transport.grpc import SslCredentials # type: ignore + +import grpc # type: ignore +from grpc.experimental import aio # type: ignore + +from google.cloud.bigquery_storage_v1beta2.types import storage +from google.cloud.bigquery_storage_v1beta2.types import stream + +from .base import BigQueryWriteTransport, DEFAULT_CLIENT_INFO +from .grpc import BigQueryWriteGrpcTransport + + +class BigQueryWriteGrpcAsyncIOTransport(BigQueryWriteTransport): + """gRPC AsyncIO backend transport for BigQueryWrite. + + BigQuery Write API. + The Write API can be used to write data to BigQuery. + + This class defines the same methods as the primary client, so the + primary client can load the underlying transport implementation + and call it. + + It sends protocol buffers over the wire using gRPC (which is built on + top of HTTP/2); the ``grpcio`` package must be installed. + """ + + _grpc_channel: aio.Channel + _stubs: Dict[str, Callable] = {} + + @classmethod + def create_channel( + cls, + host: str = "bigquerystorage.googleapis.com", + credentials: credentials.Credentials = None, + credentials_file: Optional[str] = None, + scopes: Optional[Sequence[str]] = None, + quota_project_id: Optional[str] = None, + **kwargs, + ) -> aio.Channel: + """Create and return a gRPC AsyncIO channel object. + Args: + address (Optional[str]): The host for the channel to use. + credentials (Optional[~.Credentials]): The + authorization credentials to attach to requests. These + credentials identify this application to the service. If + none are specified, the client will attempt to ascertain + the credentials from the environment. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is ignored if ``channel`` is provided. + scopes (Optional[Sequence[str]]): A optional list of scopes needed for this + service. These are only used when credentials are not specified and + are passed to :func:`google.auth.default`. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + kwargs (Optional[dict]): Keyword arguments, which are passed to the + channel creation. + Returns: + aio.Channel: A gRPC AsyncIO channel object. + """ + scopes = scopes or cls.AUTH_SCOPES + return grpc_helpers_async.create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + scopes=scopes, + quota_project_id=quota_project_id, + **kwargs, + ) + + def __init__( + self, + *, + host: str = "bigquerystorage.googleapis.com", + credentials: credentials.Credentials = None, + credentials_file: Optional[str] = None, + scopes: Optional[Sequence[str]] = None, + channel: aio.Channel = None, + api_mtls_endpoint: str = None, + client_cert_source: Callable[[], Tuple[bytes, bytes]] = None, + ssl_channel_credentials: grpc.ChannelCredentials = None, + quota_project_id=None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + ) -> None: + """Instantiate the transport. + + Args: + host (Optional[str]): The hostname to connect to. + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + This argument is ignored if ``channel`` is provided. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is ignored if ``channel`` is provided. + scopes (Optional[Sequence[str]]): A optional list of scopes needed for this + service. These are only used when credentials are not specified and + are passed to :func:`google.auth.default`. + channel (Optional[aio.Channel]): A ``Channel`` instance through + which to make calls. + api_mtls_endpoint (Optional[str]): Deprecated. The mutual TLS endpoint. + If provided, it overrides the ``host`` argument and tries to create + a mutual TLS channel with client SSL credentials from + ``client_cert_source`` or applicatin default SSL credentials. + client_cert_source (Optional[Callable[[], Tuple[bytes, bytes]]]): + Deprecated. A callback to provide client SSL certificate bytes and + private key bytes, both in PEM format. It is ignored if + ``api_mtls_endpoint`` is None. + ssl_channel_credentials (grpc.ChannelCredentials): SSL credentials + for grpc channel. It is ignored if ``channel`` is provided. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + + Raises: + google.auth.exceptions.MutualTlsChannelError: If mutual TLS transport + creation failed for any reason. + google.api_core.exceptions.DuplicateCredentialArgs: If both ``credentials`` + and ``credentials_file`` are passed. + """ + self._ssl_channel_credentials = ssl_channel_credentials + + if channel: + # Sanity check: Ensure that channel and credentials are not both + # provided. + credentials = False + + # If a channel was explicitly provided, set it. + self._grpc_channel = channel + self._ssl_channel_credentials = None + elif api_mtls_endpoint: + warnings.warn( + "api_mtls_endpoint and client_cert_source are deprecated", + DeprecationWarning, + ) + + host = ( + api_mtls_endpoint + if ":" in api_mtls_endpoint + else api_mtls_endpoint + ":443" + ) + + if credentials is None: + credentials, _ = auth.default( + scopes=self.AUTH_SCOPES, quota_project_id=quota_project_id + ) + + # Create SSL credentials with client_cert_source or application + # default SSL credentials. + if client_cert_source: + cert, key = client_cert_source() + ssl_credentials = grpc.ssl_channel_credentials( + certificate_chain=cert, private_key=key + ) + else: + ssl_credentials = SslCredentials().ssl_credentials + + # create a new channel. The provided one is ignored. + self._grpc_channel = type(self).create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + ssl_credentials=ssl_credentials, + scopes=scopes or self.AUTH_SCOPES, + quota_project_id=quota_project_id, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + self._ssl_channel_credentials = ssl_credentials + else: + host = host if ":" in host else host + ":443" + + if credentials is None: + credentials, _ = auth.default( + scopes=self.AUTH_SCOPES, quota_project_id=quota_project_id + ) + + # create a new channel. The provided one is ignored. + self._grpc_channel = type(self).create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + ssl_credentials=ssl_channel_credentials, + scopes=scopes or self.AUTH_SCOPES, + quota_project_id=quota_project_id, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + + # Run the base constructor. + super().__init__( + host=host, + credentials=credentials, + credentials_file=credentials_file, + scopes=scopes or self.AUTH_SCOPES, + quota_project_id=quota_project_id, + client_info=client_info, + ) + + self._stubs = {} + + @property + def grpc_channel(self) -> aio.Channel: + """Create the channel designed to connect to this service. + + This property caches on the instance; repeated calls return + the same channel. + """ + # Return the channel from cache. + return self._grpc_channel + + @property + def create_write_stream( + self, + ) -> Callable[[storage.CreateWriteStreamRequest], Awaitable[stream.WriteStream]]: + r"""Return a callable for the create write stream method over gRPC. + + Creates a write stream to the given table. Additionally, every + table has a special COMMITTED stream named '_default' to which + data can be written. This stream doesn't need to be created + using CreateWriteStream. It is a stream that can be used + simultaneously by any number of clients. Data written to this + stream is considered committed as soon as an acknowledgement is + received. + + Returns: + Callable[[~.CreateWriteStreamRequest], + Awaitable[~.WriteStream]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "create_write_stream" not in self._stubs: + self._stubs["create_write_stream"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryWrite/CreateWriteStream", + request_serializer=storage.CreateWriteStreamRequest.serialize, + response_deserializer=stream.WriteStream.deserialize, + ) + return self._stubs["create_write_stream"] + + @property + def append_rows( + self, + ) -> Callable[[storage.AppendRowsRequest], Awaitable[storage.AppendRowsResponse]]: + r"""Return a callable for the append rows method over gRPC. + + Appends data to the given stream. + + If ``offset`` is specified, the ``offset`` is checked against + the end of stream. The server returns ``OUT_OF_RANGE`` in + ``AppendRowsResponse`` if an attempt is made to append to an + offset beyond the current end of the stream or + ``ALREADY_EXISTS`` if user provids an ``offset`` that has + already been written to. User can retry with adjusted offset + within the same RPC stream. If ``offset`` is not specified, + append happens at the end of the stream. + + The response contains the offset at which the append happened. + Responses are received in the same order in which requests are + sent. There will be one response for each successful request. If + the ``offset`` is not set in response, it means append didn't + happen due to some errors. If one request fails, all the + subsequent requests will also fail until a success request is + made again. + + If the stream is of ``PENDING`` type, data will only be + available for read operations after the stream is committed. + + Returns: + Callable[[~.AppendRowsRequest], + Awaitable[~.AppendRowsResponse]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "append_rows" not in self._stubs: + self._stubs["append_rows"] = self.grpc_channel.stream_stream( + "/google.cloud.bigquery.storage.v1beta2.BigQueryWrite/AppendRows", + request_serializer=storage.AppendRowsRequest.serialize, + response_deserializer=storage.AppendRowsResponse.deserialize, + ) + return self._stubs["append_rows"] + + @property + def get_write_stream( + self, + ) -> Callable[[storage.GetWriteStreamRequest], Awaitable[stream.WriteStream]]: + r"""Return a callable for the get write stream method over gRPC. + + Gets a write stream. + + Returns: + Callable[[~.GetWriteStreamRequest], + Awaitable[~.WriteStream]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "get_write_stream" not in self._stubs: + self._stubs["get_write_stream"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryWrite/GetWriteStream", + request_serializer=storage.GetWriteStreamRequest.serialize, + response_deserializer=stream.WriteStream.deserialize, + ) + return self._stubs["get_write_stream"] + + @property + def finalize_write_stream( + self, + ) -> Callable[ + [storage.FinalizeWriteStreamRequest], + Awaitable[storage.FinalizeWriteStreamResponse], + ]: + r"""Return a callable for the finalize write stream method over gRPC. + + Finalize a write stream so that no new data can be appended to + the stream. Finalize is not supported on the '_default' stream. + + Returns: + Callable[[~.FinalizeWriteStreamRequest], + Awaitable[~.FinalizeWriteStreamResponse]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "finalize_write_stream" not in self._stubs: + self._stubs["finalize_write_stream"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryWrite/FinalizeWriteStream", + request_serializer=storage.FinalizeWriteStreamRequest.serialize, + response_deserializer=storage.FinalizeWriteStreamResponse.deserialize, + ) + return self._stubs["finalize_write_stream"] + + @property + def batch_commit_write_streams( + self, + ) -> Callable[ + [storage.BatchCommitWriteStreamsRequest], + Awaitable[storage.BatchCommitWriteStreamsResponse], + ]: + r"""Return a callable for the batch commit write streams method over gRPC. + + Atomically commits a group of ``PENDING`` streams that belong to + the same ``parent`` table. Streams must be finalized before + commit and cannot be committed multiple times. Once a stream is + committed, data in the stream becomes available for read + operations. + + Returns: + Callable[[~.BatchCommitWriteStreamsRequest], + Awaitable[~.BatchCommitWriteStreamsResponse]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "batch_commit_write_streams" not in self._stubs: + self._stubs["batch_commit_write_streams"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryWrite/BatchCommitWriteStreams", + request_serializer=storage.BatchCommitWriteStreamsRequest.serialize, + response_deserializer=storage.BatchCommitWriteStreamsResponse.deserialize, + ) + return self._stubs["batch_commit_write_streams"] + + @property + def flush_rows( + self, + ) -> Callable[[storage.FlushRowsRequest], Awaitable[storage.FlushRowsResponse]]: + r"""Return a callable for the flush rows method over gRPC. + + Flushes rows to a BUFFERED stream. If users are appending rows + to BUFFERED stream, flush operation is required in order for the + rows to become available for reading. A Flush operation flushes + up to any previously flushed offset in a BUFFERED stream, to the + offset specified in the request. Flush is not supported on the + \_default stream, since it is not BUFFERED. + + Returns: + Callable[[~.FlushRowsRequest], + Awaitable[~.FlushRowsResponse]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "flush_rows" not in self._stubs: + self._stubs["flush_rows"] = self.grpc_channel.unary_unary( + "/google.cloud.bigquery.storage.v1beta2.BigQueryWrite/FlushRows", + request_serializer=storage.FlushRowsRequest.serialize, + response_deserializer=storage.FlushRowsResponse.deserialize, + ) + return self._stubs["flush_rows"] + + +__all__ = ("BigQueryWriteGrpcAsyncIOTransport",) diff --git a/google/cloud/bigquery_storage_v1beta2/types/__init__.py b/google/cloud/bigquery_storage_v1beta2/types/__init__.py new file mode 100644 index 00000000..ba1cac28 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/types/__init__.py @@ -0,0 +1,96 @@ +# -*- 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. +# + +from .arrow import ( + ArrowSchema, + ArrowRecordBatch, + ArrowSerializationOptions, +) +from .avro import ( + AvroSchema, + AvroRows, +) +from .protobuf import ( + ProtoSchema, + ProtoRows, +) +from .table import ( + TableSchema, + TableFieldSchema, +) +from .stream import ( + DataFormat, + ReadSession, + ReadStream, + WriteStream, + DataFormat, +) +from .storage import ( + CreateReadSessionRequest, + ReadRowsRequest, + ThrottleState, + StreamStats, + ReadRowsResponse, + SplitReadStreamRequest, + SplitReadStreamResponse, + CreateWriteStreamRequest, + AppendRowsRequest, + AppendRowsResponse, + GetWriteStreamRequest, + BatchCommitWriteStreamsRequest, + BatchCommitWriteStreamsResponse, + FinalizeWriteStreamRequest, + FinalizeWriteStreamResponse, + FlushRowsRequest, + FlushRowsResponse, + StorageError, +) + +__all__ = ( + "ArrowSchema", + "ArrowRecordBatch", + "ArrowSerializationOptions", + "AvroSchema", + "AvroRows", + "ProtoSchema", + "ProtoRows", + "TableSchema", + "TableFieldSchema", + "DataFormat", + "ReadSession", + "ReadStream", + "WriteStream", + "DataFormat", + "CreateReadSessionRequest", + "ReadRowsRequest", + "ThrottleState", + "StreamStats", + "ReadRowsResponse", + "SplitReadStreamRequest", + "SplitReadStreamResponse", + "CreateWriteStreamRequest", + "AppendRowsRequest", + "AppendRowsResponse", + "GetWriteStreamRequest", + "BatchCommitWriteStreamsRequest", + "BatchCommitWriteStreamsResponse", + "FinalizeWriteStreamRequest", + "FinalizeWriteStreamResponse", + "FlushRowsRequest", + "FlushRowsResponse", + "StorageError", +) diff --git a/google/cloud/bigquery_storage_v1beta2/types/arrow.py b/google/cloud/bigquery_storage_v1beta2/types/arrow.py new file mode 100644 index 00000000..0ebbbe8a --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/types/arrow.py @@ -0,0 +1,71 @@ +# -*- 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 + + +__protobuf__ = proto.module( + package="google.cloud.bigquery.storage.v1beta2", + manifest={"ArrowSchema", "ArrowRecordBatch", "ArrowSerializationOptions",}, +) + + +class ArrowSchema(proto.Message): + r"""Arrow schema as specified in + https://arrow.apache.org/docs/python/api/datatypes.html and + serialized to bytes using IPC: + https://arrow.apache.org/docs/format/Columnar.html#serialization- + and-interprocess-communication-ipc + See code samples on how this message can be deserialized. + + Attributes: + serialized_schema (bytes): + IPC serialized Arrow schema. + """ + + serialized_schema = proto.Field(proto.BYTES, number=1) + + +class ArrowRecordBatch(proto.Message): + r"""Arrow RecordBatch. + + Attributes: + serialized_record_batch (bytes): + IPC-serialized Arrow RecordBatch. + """ + + serialized_record_batch = proto.Field(proto.BYTES, number=1) + + +class ArrowSerializationOptions(proto.Message): + r"""Contains options specific to Arrow Serialization. + + Attributes: + format_ (google.cloud.bigquery_storage_v1beta2.types.ArrowSerializationOptions.Format): + The Arrow IPC format to use. + """ + + class Format(proto.Enum): + r"""The IPC format to use when serializing Arrow streams.""" + FORMAT_UNSPECIFIED = 0 + ARROW_0_14 = 1 + ARROW_0_15 = 2 + + format_ = proto.Field(proto.ENUM, number=1, enum=Format,) + + +__all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/google/cloud/bigquery_storage_v1beta2/types/avro.py b/google/cloud/bigquery_storage_v1beta2/types/avro.py new file mode 100644 index 00000000..49aa017d --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/types/avro.py @@ -0,0 +1,50 @@ +# -*- 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 + + +__protobuf__ = proto.module( + package="google.cloud.bigquery.storage.v1beta2", + manifest={"AvroSchema", "AvroRows",}, +) + + +class AvroSchema(proto.Message): + r"""Avro schema. + + Attributes: + schema (str): + Json serialized schema, as described at + https://avro.apache.org/docs/1.8.1/spec.html. + """ + + schema = proto.Field(proto.STRING, number=1) + + +class AvroRows(proto.Message): + r"""Avro rows. + + Attributes: + serialized_binary_rows (bytes): + Binary serialized rows in a block. + """ + + serialized_binary_rows = proto.Field(proto.BYTES, number=1) + + +__all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/google/cloud/bigquery_storage_v1beta2/types/protobuf.py b/google/cloud/bigquery_storage_v1beta2/types/protobuf.py new file mode 100644 index 00000000..99c0543c --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/types/protobuf.py @@ -0,0 +1,64 @@ +# -*- 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.protobuf import descriptor_pb2 as descriptor # type: ignore + + +__protobuf__ = proto.module( + package="google.cloud.bigquery.storage.v1beta2", + manifest={"ProtoSchema", "ProtoRows",}, +) + + +class ProtoSchema(proto.Message): + r"""Protobuf schema is an API presentation the proto buffer + schema. + + Attributes: + proto_descriptor (google.protobuf.descriptor_pb2.DescriptorProto): + Descriptor for input message. The descriptor + has to be self contained, including all the + nested types, excepted for proto buffer well + known types + (https://developers.google.com/protocol- + buffers/docs/reference/google.protobuf). + """ + + proto_descriptor = proto.Field( + proto.MESSAGE, number=1, message=descriptor.DescriptorProto, + ) + + +class ProtoRows(proto.Message): + r"""Protobuf rows. + + Attributes: + serialized_rows (Sequence[bytes]): + A sequence of rows serialized as a Protocol + Buffer. + See https://developers.google.com/protocol- + buffers/docs/overview for more information on + deserializing this field. + """ + + serialized_rows = proto.RepeatedField(proto.BYTES, number=1) + + +__all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/google/cloud/bigquery_storage_v1beta2/types/storage.py b/google/cloud/bigquery_storage_v1beta2/types/storage.py new file mode 100644 index 00000000..85369388 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/types/storage.py @@ -0,0 +1,487 @@ +# -*- 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)) diff --git a/google/cloud/bigquery_storage_v1beta2/types/stream.py b/google/cloud/bigquery_storage_v1beta2/types/stream.py new file mode 100644 index 00000000..77fb444c --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/types/stream.py @@ -0,0 +1,203 @@ +# -*- 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 table as gcbs_table +from google.protobuf import timestamp_pb2 as timestamp # type: ignore + + +__protobuf__ = proto.module( + package="google.cloud.bigquery.storage.v1beta2", + manifest={"DataFormat", "ReadSession", "ReadStream", "WriteStream",}, +) + + +class DataFormat(proto.Enum): + r"""Data format for input or output data.""" + DATA_FORMAT_UNSPECIFIED = 0 + AVRO = 1 + ARROW = 2 + + +class ReadSession(proto.Message): + r"""Information about the ReadSession. + + Attributes: + name (str): + Output only. Unique identifier for the session, in the form + ``projects/{project_id}/locations/{location}/sessions/{session_id}``. + expire_time (google.protobuf.timestamp_pb2.Timestamp): + Output only. Time at which the session becomes invalid. + After this time, subsequent requests to read this Session + will return errors. The expire_time is automatically + assigned and currently cannot be specified or updated. + data_format (google.cloud.bigquery_storage_v1beta2.types.DataFormat): + Immutable. Data format of the output data. + avro_schema (google.cloud.bigquery_storage_v1beta2.types.AvroSchema): + Output only. Avro schema. + arrow_schema (google.cloud.bigquery_storage_v1beta2.types.ArrowSchema): + Output only. Arrow schema. + table (str): + Immutable. Table that this ReadSession is reading from, in + the form + \`projects/{project_id}/datasets/{dataset_id}/tables/{table_id} + table_modifiers (google.cloud.bigquery_storage_v1beta2.types.ReadSession.TableModifiers): + Optional. Any modifiers which are applied + when reading from the specified table. + read_options (google.cloud.bigquery_storage_v1beta2.types.ReadSession.TableReadOptions): + Optional. Read options for this session (e.g. + column selection, filters). + streams (Sequence[google.cloud.bigquery_storage_v1beta2.types.ReadStream]): + Output only. A list of streams created with the session. + + At least one stream is created with the session. In the + future, larger request_stream_count values *may* result in + this list being unpopulated, in that case, the user will + need to use a List method to get the streams instead, which + is not yet available. + """ + + class TableModifiers(proto.Message): + r"""Additional attributes when reading a table. + + Attributes: + snapshot_time (google.protobuf.timestamp_pb2.Timestamp): + The snapshot time of the table. If not set, + interpreted as now. + """ + + snapshot_time = proto.Field( + proto.MESSAGE, number=1, message=timestamp.Timestamp, + ) + + class TableReadOptions(proto.Message): + r"""Options dictating how we read a table. + + Attributes: + selected_fields (Sequence[str]): + Names of the fields in the table that should be read. If + empty, all fields will be read. If the specified field is a + nested field, all the sub-fields in the field will be + selected. The output field order is unrelated to the order + of fields in selected_fields. + row_restriction (str): + SQL text filtering statement, similar to a WHERE clause in a + query. Aggregates are not supported. + + Examples: "int_field > 5" "date_field = CAST('2014-9-27' as + DATE)" "nullable_field is not NULL" "st_equals(geo_field, + st_geofromtext("POINT(2, 2)"))" "numeric_field BETWEEN 1.0 + AND 5.0". + arrow_serialization_options (google.cloud.bigquery_storage_v1beta2.types.ArrowSerializationOptions): + Optional. Options specific to the Apache + Arrow output format. + """ + + selected_fields = proto.RepeatedField(proto.STRING, number=1) + + row_restriction = proto.Field(proto.STRING, number=2) + + arrow_serialization_options = proto.Field( + proto.MESSAGE, number=3, message=arrow.ArrowSerializationOptions, + ) + + name = proto.Field(proto.STRING, number=1) + + expire_time = proto.Field(proto.MESSAGE, number=2, message=timestamp.Timestamp,) + + data_format = proto.Field(proto.ENUM, number=3, enum="DataFormat",) + + avro_schema = proto.Field( + proto.MESSAGE, number=4, oneof="schema", message=avro.AvroSchema, + ) + + arrow_schema = proto.Field( + proto.MESSAGE, number=5, oneof="schema", message=arrow.ArrowSchema, + ) + + table = proto.Field(proto.STRING, number=6) + + table_modifiers = proto.Field(proto.MESSAGE, number=7, message=TableModifiers,) + + read_options = proto.Field(proto.MESSAGE, number=8, message=TableReadOptions,) + + streams = proto.RepeatedField(proto.MESSAGE, number=10, message="ReadStream",) + + +class ReadStream(proto.Message): + r"""Information about a single stream that gets data out of the storage + system. Most of the information about ``ReadStream`` instances is + aggregated, making ``ReadStream`` lightweight. + + Attributes: + name (str): + Output only. Name of the stream, in the form + ``projects/{project_id}/locations/{location}/sessions/{session_id}/streams/{stream_id}``. + """ + + name = proto.Field(proto.STRING, number=1) + + +class WriteStream(proto.Message): + r"""Information about a single stream that gets data inside the + storage system. + + Attributes: + name (str): + Output only. Name of the stream, in the form + ``projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}``. + type_ (google.cloud.bigquery_storage_v1beta2.types.WriteStream.Type): + Immutable. Type of the stream. + create_time (google.protobuf.timestamp_pb2.Timestamp): + Output only. Create time of the stream. For the \_default + stream, this is the creation_time of the table. + commit_time (google.protobuf.timestamp_pb2.Timestamp): + Output only. Commit time of the stream. If a stream is of + ``COMMITTED`` type, then it will have a commit_time same as + ``create_time``. If the stream is of ``PENDING`` type, + commit_time being empty means it is not committed. + table_schema (google.cloud.bigquery_storage_v1beta2.types.TableSchema): + Output only. The schema of the destination table. It is only + returned in ``CreateWriteStream`` response. Caller should + generate data that's compatible with this schema to send in + initial ``AppendRowsRequest``. The table schema could go out + of date during the life time of the stream. + """ + + class Type(proto.Enum): + r"""Type enum of the stream.""" + TYPE_UNSPECIFIED = 0 + COMMITTED = 1 + PENDING = 2 + BUFFERED = 3 + + name = proto.Field(proto.STRING, number=1) + + type_ = proto.Field(proto.ENUM, number=2, enum=Type,) + + create_time = proto.Field(proto.MESSAGE, number=3, message=timestamp.Timestamp,) + + commit_time = proto.Field(proto.MESSAGE, number=4, message=timestamp.Timestamp,) + + table_schema = proto.Field(proto.MESSAGE, number=5, message=gcbs_table.TableSchema,) + + +__all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/google/cloud/bigquery_storage_v1beta2/types/table.py b/google/cloud/bigquery_storage_v1beta2/types/table.py new file mode 100644 index 00000000..f1a209e0 --- /dev/null +++ b/google/cloud/bigquery_storage_v1beta2/types/table.py @@ -0,0 +1,94 @@ +# -*- 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 + + +__protobuf__ = proto.module( + package="google.cloud.bigquery.storage.v1beta2", + manifest={"TableSchema", "TableFieldSchema",}, +) + + +class TableSchema(proto.Message): + r"""Schema of a table + + Attributes: + fields (Sequence[google.cloud.bigquery_storage_v1beta2.types.TableFieldSchema]): + Describes the fields in a table. + """ + + fields = proto.RepeatedField(proto.MESSAGE, number=1, message="TableFieldSchema",) + + +class TableFieldSchema(proto.Message): + r"""A field in TableSchema + + Attributes: + name (str): + Required. The field name. The name must contain only letters + (a-z, A-Z), numbers (0-9), or underscores (_), and must + start with a letter or underscore. The maximum length is 128 + characters. + type_ (google.cloud.bigquery_storage_v1beta2.types.TableFieldSchema.Type): + Required. The field data type. + mode (google.cloud.bigquery_storage_v1beta2.types.TableFieldSchema.Mode): + Optional. The field mode. The default value + is NULLABLE. + fields (Sequence[google.cloud.bigquery_storage_v1beta2.types.TableFieldSchema]): + Optional. Describes the nested schema fields + if the type property is set to STRUCT. + description (str): + Optional. The field description. The maximum + length is 1,024 characters. + """ + + class Type(proto.Enum): + r"""""" + TYPE_UNSPECIFIED = 0 + STRING = 1 + INT64 = 2 + DOUBLE = 3 + STRUCT = 4 + BYTES = 5 + BOOL = 6 + TIMESTAMP = 7 + DATE = 8 + TIME = 9 + DATETIME = 10 + GEOGRAPHY = 11 + NUMERIC = 12 + + class Mode(proto.Enum): + r"""""" + MODE_UNSPECIFIED = 0 + NULLABLE = 1 + REQUIRED = 2 + REPEATED = 3 + + name = proto.Field(proto.STRING, number=1) + + type_ = proto.Field(proto.ENUM, number=2, enum=Type,) + + mode = proto.Field(proto.ENUM, number=3, enum=Mode,) + + fields = proto.RepeatedField(proto.MESSAGE, number=4, message="TableFieldSchema",) + + description = proto.Field(proto.STRING, number=6) + + +__all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/synth.py b/synth.py index fffa438f..21fb8b3a 100644 --- a/synth.py +++ b/synth.py @@ -38,6 +38,8 @@ "docs/conf.py", "docs/index.rst", f"google/cloud/bigquery_storage_{version}/__init__.py", + # v1beta2 was first generated after the microgenerator migration. + "scripts/fixup_bigquery_storage_v1beta2_keywords.py", "README.rst", "nox*.py", "setup.py", @@ -150,16 +152,27 @@ # The DataFormat enum is not exposed in bigquery_storage_v1/types, add it there. s.replace( - "google/cloud/bigquery_storage_v1/types/__init__.py", + "google/cloud/bigquery_storage_v1*/types/__init__.py", r"from \.stream import \(", "\g<0>\n DataFormat,", ) s.replace( - "google/cloud/bigquery_storage_v1/types/__init__.py", + "google/cloud/bigquery_storage_v1*/types/__init__.py", r"""["']ReadSession["']""", '"DataFormat",\n \g<0>', ) +# The append_rows method doesn't contain keyword arguments that build request +# objects, so flattened tests are not needed and break with TypeError. +s.replace( + 'tests/unit/gapic/bigquery_storage_v1*/test_big_query_write.py', + r"(@[a-z.()\n]*\n)?(async )?" + r"def test_append_rows_flattened[_a-z]*\(\):\n" + r"( {4}.*|\n)+", + '\n', +) + + # Fix library installations in nox sessions (unit and system tests) - it's # redundant to install the library twice. s.replace( @@ -168,19 +181,6 @@ ")\n", ) -# Fix test coverage plugin paths. -s.replace( - "noxfile.py", - r'"--cov=google\.cloud\.bigquerystorage"', - ( - '"--cov=google.cloud.bigquery_storage",\n' - ' "--cov=google.cloud.bigquery_storage_v1"' - ), -) -s.replace( - "noxfile.py", r"--cov=tests\.unit", "--cov=tests/unit", -) - # TODO(busunkim): Use latest sphinx after microgenerator transition s.replace("noxfile.py", """['"]sphinx['"]""", '"sphinx<3.0.0"') diff --git a/tests/unit/gapic/bigquery_storage_v1beta2/__init__.py b/tests/unit/gapic/bigquery_storage_v1beta2/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/tests/unit/gapic/bigquery_storage_v1beta2/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/unit/gapic/bigquery_storage_v1beta2/test_big_query_read.py b/tests/unit/gapic/bigquery_storage_v1beta2/test_big_query_read.py new file mode 100644 index 00000000..c83a2593 --- /dev/null +++ b/tests/unit/gapic/bigquery_storage_v1beta2/test_big_query_read.py @@ -0,0 +1,1525 @@ +# -*- 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 os +import mock + +import grpc +from grpc.experimental import aio +import math +import pytest +from proto.marshal.rules.dates import DurationRule, TimestampRule + +from google import auth +from google.api_core import client_options +from google.api_core import exceptions +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.auth import credentials +from google.auth.exceptions import MutualTLSChannelError +from google.cloud.bigquery_storage_v1beta2.services.big_query_read import ( + BigQueryReadAsyncClient, +) +from google.cloud.bigquery_storage_v1beta2.services.big_query_read import ( + BigQueryReadClient, +) +from google.cloud.bigquery_storage_v1beta2.services.big_query_read import transports +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 storage +from google.cloud.bigquery_storage_v1beta2.types import stream +from google.oauth2 import service_account +from google.protobuf import timestamp_pb2 as timestamp # type: ignore + + +def client_cert_source_callback(): + return b"cert bytes", b"key bytes" + + +# If default endpoint is localhost, then default mtls endpoint will be the same. +# This method modifies the default endpoint so the client can produce a different +# mtls endpoint for endpoint testing purposes. +def modify_default_endpoint(client): + return ( + "foo.googleapis.com" + if ("localhost" in client.DEFAULT_ENDPOINT) + else client.DEFAULT_ENDPOINT + ) + + +def test__get_default_mtls_endpoint(): + api_endpoint = "example.googleapis.com" + api_mtls_endpoint = "example.mtls.googleapis.com" + sandbox_endpoint = "example.sandbox.googleapis.com" + sandbox_mtls_endpoint = "example.mtls.sandbox.googleapis.com" + non_googleapi = "api.example.com" + + assert BigQueryReadClient._get_default_mtls_endpoint(None) is None + assert ( + BigQueryReadClient._get_default_mtls_endpoint(api_endpoint) == api_mtls_endpoint + ) + assert ( + BigQueryReadClient._get_default_mtls_endpoint(api_mtls_endpoint) + == api_mtls_endpoint + ) + assert ( + BigQueryReadClient._get_default_mtls_endpoint(sandbox_endpoint) + == sandbox_mtls_endpoint + ) + assert ( + BigQueryReadClient._get_default_mtls_endpoint(sandbox_mtls_endpoint) + == sandbox_mtls_endpoint + ) + assert BigQueryReadClient._get_default_mtls_endpoint(non_googleapi) == non_googleapi + + +def test_big_query_read_client_from_service_account_info(): + creds = credentials.AnonymousCredentials() + with mock.patch.object( + service_account.Credentials, "from_service_account_info" + ) as factory: + factory.return_value = creds + info = {"valid": True} + client = BigQueryReadClient.from_service_account_info(info) + assert client.transport._credentials == creds + + assert client.transport._host == "bigquerystorage.googleapis.com:443" + + +@pytest.mark.parametrize("client_class", [BigQueryReadClient, BigQueryReadAsyncClient,]) +def test_big_query_read_client_from_service_account_file(client_class): + creds = credentials.AnonymousCredentials() + with mock.patch.object( + service_account.Credentials, "from_service_account_file" + ) as factory: + factory.return_value = creds + client = client_class.from_service_account_file("dummy/file/path.json") + assert client.transport._credentials == creds + + client = client_class.from_service_account_json("dummy/file/path.json") + assert client.transport._credentials == creds + + assert client.transport._host == "bigquerystorage.googleapis.com:443" + + +def test_big_query_read_client_get_transport_class(): + transport = BigQueryReadClient.get_transport_class() + available_transports = [ + transports.BigQueryReadGrpcTransport, + ] + assert transport in available_transports + + transport = BigQueryReadClient.get_transport_class("grpc") + assert transport == transports.BigQueryReadGrpcTransport + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name", + [ + (BigQueryReadClient, transports.BigQueryReadGrpcTransport, "grpc"), + ( + BigQueryReadAsyncClient, + transports.BigQueryReadGrpcAsyncIOTransport, + "grpc_asyncio", + ), + ], +) +@mock.patch.object( + BigQueryReadClient, "DEFAULT_ENDPOINT", modify_default_endpoint(BigQueryReadClient) +) +@mock.patch.object( + BigQueryReadAsyncClient, + "DEFAULT_ENDPOINT", + modify_default_endpoint(BigQueryReadAsyncClient), +) +def test_big_query_read_client_client_options( + client_class, transport_class, transport_name +): + # Check that if channel is provided we won't create a new one. + with mock.patch.object(BigQueryReadClient, "get_transport_class") as gtc: + transport = transport_class(credentials=credentials.AnonymousCredentials()) + client = client_class(transport=transport) + gtc.assert_not_called() + + # Check that if channel is provided via str we will create a new one. + with mock.patch.object(BigQueryReadClient, "get_transport_class") as gtc: + client = client_class(transport=transport_name) + gtc.assert_called() + + # Check the case api_endpoint is provided. + options = client_options.ClientOptions(api_endpoint="squid.clam.whelk") + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host="squid.clam.whelk", + scopes=None, + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + # Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS_ENDPOINT is + # "never". + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "never"}): + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class() + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client.DEFAULT_ENDPOINT, + scopes=None, + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + # Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS_ENDPOINT is + # "always". + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "always"}): + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class() + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client.DEFAULT_MTLS_ENDPOINT, + scopes=None, + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + # Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS_ENDPOINT has + # unsupported value. + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "Unsupported"}): + with pytest.raises(MutualTLSChannelError): + client = client_class() + + # Check the case GOOGLE_API_USE_CLIENT_CERTIFICATE has unsupported value. + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "Unsupported"} + ): + with pytest.raises(ValueError): + client = client_class() + + # Check the case quota_project_id is provided + options = client_options.ClientOptions(quota_project_id="octopus") + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client.DEFAULT_ENDPOINT, + scopes=None, + ssl_channel_credentials=None, + quota_project_id="octopus", + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name,use_client_cert_env", + [ + (BigQueryReadClient, transports.BigQueryReadGrpcTransport, "grpc", "true"), + ( + BigQueryReadAsyncClient, + transports.BigQueryReadGrpcAsyncIOTransport, + "grpc_asyncio", + "true", + ), + (BigQueryReadClient, transports.BigQueryReadGrpcTransport, "grpc", "false"), + ( + BigQueryReadAsyncClient, + transports.BigQueryReadGrpcAsyncIOTransport, + "grpc_asyncio", + "false", + ), + ], +) +@mock.patch.object( + BigQueryReadClient, "DEFAULT_ENDPOINT", modify_default_endpoint(BigQueryReadClient) +) +@mock.patch.object( + BigQueryReadAsyncClient, + "DEFAULT_ENDPOINT", + modify_default_endpoint(BigQueryReadAsyncClient), +) +@mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "auto"}) +def test_big_query_read_client_mtls_env_auto( + client_class, transport_class, transport_name, use_client_cert_env +): + # This tests the endpoint autoswitch behavior. Endpoint is autoswitched to the default + # mtls endpoint, if GOOGLE_API_USE_CLIENT_CERTIFICATE is "true" and client cert exists. + + # Check the case client_cert_source is provided. Whether client cert is used depends on + # GOOGLE_API_USE_CLIENT_CERTIFICATE value. + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": use_client_cert_env} + ): + options = client_options.ClientOptions( + client_cert_source=client_cert_source_callback + ) + with mock.patch.object(transport_class, "__init__") as patched: + ssl_channel_creds = mock.Mock() + with mock.patch( + "grpc.ssl_channel_credentials", return_value=ssl_channel_creds + ): + patched.return_value = None + client = client_class(client_options=options) + + if use_client_cert_env == "false": + expected_ssl_channel_creds = None + expected_host = client.DEFAULT_ENDPOINT + else: + expected_ssl_channel_creds = ssl_channel_creds + expected_host = client.DEFAULT_MTLS_ENDPOINT + + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=expected_host, + scopes=None, + ssl_channel_credentials=expected_ssl_channel_creds, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + # Check the case ADC client cert is provided. Whether client cert is used depends on + # GOOGLE_API_USE_CLIENT_CERTIFICATE value. + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": use_client_cert_env} + ): + with mock.patch.object(transport_class, "__init__") as patched: + with mock.patch( + "google.auth.transport.grpc.SslCredentials.__init__", return_value=None + ): + with mock.patch( + "google.auth.transport.grpc.SslCredentials.is_mtls", + new_callable=mock.PropertyMock, + ) as is_mtls_mock: + with mock.patch( + "google.auth.transport.grpc.SslCredentials.ssl_credentials", + new_callable=mock.PropertyMock, + ) as ssl_credentials_mock: + if use_client_cert_env == "false": + is_mtls_mock.return_value = False + ssl_credentials_mock.return_value = None + expected_host = client.DEFAULT_ENDPOINT + expected_ssl_channel_creds = None + else: + is_mtls_mock.return_value = True + ssl_credentials_mock.return_value = mock.Mock() + expected_host = client.DEFAULT_MTLS_ENDPOINT + expected_ssl_channel_creds = ( + ssl_credentials_mock.return_value + ) + + patched.return_value = None + client = client_class() + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=expected_host, + scopes=None, + ssl_channel_credentials=expected_ssl_channel_creds, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + # Check the case client_cert_source and ADC client cert are not provided. + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": use_client_cert_env} + ): + with mock.patch.object(transport_class, "__init__") as patched: + with mock.patch( + "google.auth.transport.grpc.SslCredentials.__init__", return_value=None + ): + with mock.patch( + "google.auth.transport.grpc.SslCredentials.is_mtls", + new_callable=mock.PropertyMock, + ) as is_mtls_mock: + is_mtls_mock.return_value = False + patched.return_value = None + client = client_class() + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client.DEFAULT_ENDPOINT, + scopes=None, + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name", + [ + (BigQueryReadClient, transports.BigQueryReadGrpcTransport, "grpc"), + ( + BigQueryReadAsyncClient, + transports.BigQueryReadGrpcAsyncIOTransport, + "grpc_asyncio", + ), + ], +) +def test_big_query_read_client_client_options_scopes( + client_class, transport_class, transport_name +): + # Check the case scopes are provided. + options = client_options.ClientOptions(scopes=["1", "2"],) + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client.DEFAULT_ENDPOINT, + scopes=["1", "2"], + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name", + [ + (BigQueryReadClient, transports.BigQueryReadGrpcTransport, "grpc"), + ( + BigQueryReadAsyncClient, + transports.BigQueryReadGrpcAsyncIOTransport, + "grpc_asyncio", + ), + ], +) +def test_big_query_read_client_client_options_credentials_file( + client_class, transport_class, transport_name +): + # Check the case credentials file is provided. + options = client_options.ClientOptions(credentials_file="credentials.json") + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options) + patched.assert_called_once_with( + credentials=None, + credentials_file="credentials.json", + host=client.DEFAULT_ENDPOINT, + scopes=None, + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + +def test_big_query_read_client_client_options_from_dict(): + with mock.patch( + "google.cloud.bigquery_storage_v1beta2.services.big_query_read.transports.BigQueryReadGrpcTransport.__init__" + ) as grpc_transport: + grpc_transport.return_value = None + client = BigQueryReadClient(client_options={"api_endpoint": "squid.clam.whelk"}) + grpc_transport.assert_called_once_with( + credentials=None, + credentials_file=None, + host="squid.clam.whelk", + scopes=None, + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + +def test_create_read_session( + transport: str = "grpc", request_type=storage.CreateReadSessionRequest +): + client = BigQueryReadClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.create_read_session), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = stream.ReadSession( + name="name_value", + data_format=stream.DataFormat.AVRO, + table="table_value", + avro_schema=avro.AvroSchema(schema="schema_value"), + ) + + response = client.create_read_session(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.CreateReadSessionRequest() + + # Establish that the response is the type that we expect. + + assert isinstance(response, stream.ReadSession) + + assert response.name == "name_value" + + assert response.data_format == stream.DataFormat.AVRO + + assert response.table == "table_value" + + +def test_create_read_session_from_dict(): + test_create_read_session(request_type=dict) + + +@pytest.mark.asyncio +async def test_create_read_session_async( + transport: str = "grpc_asyncio", request_type=storage.CreateReadSessionRequest +): + client = BigQueryReadAsyncClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.create_read_session), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + stream.ReadSession( + name="name_value", + data_format=stream.DataFormat.AVRO, + table="table_value", + ) + ) + + response = await client.create_read_session(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.CreateReadSessionRequest() + + # Establish that the response is the type that we expect. + assert isinstance(response, stream.ReadSession) + + assert response.name == "name_value" + + assert response.data_format == stream.DataFormat.AVRO + + assert response.table == "table_value" + + +@pytest.mark.asyncio +async def test_create_read_session_async_from_dict(): + await test_create_read_session_async(request_type=dict) + + +def test_create_read_session_field_headers(): + client = BigQueryReadClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.CreateReadSessionRequest() + request.read_session.table = "read_session.table/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.create_read_session), "__call__" + ) as call: + call.return_value = stream.ReadSession() + + client.create_read_session(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "read_session.table=read_session.table/value", + ) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_create_read_session_field_headers_async(): + client = BigQueryReadAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.CreateReadSessionRequest() + request.read_session.table = "read_session.table/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.create_read_session), "__call__" + ) as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(stream.ReadSession()) + + await client.create_read_session(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "read_session.table=read_session.table/value", + ) in kw["metadata"] + + +def test_create_read_session_flattened(): + client = BigQueryReadClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.create_read_session), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = stream.ReadSession() + + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + client.create_read_session( + parent="parent_value", + read_session=stream.ReadSession(name="name_value"), + max_stream_count=1721, + ) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0].parent == "parent_value" + + assert args[0].read_session == stream.ReadSession(name="name_value") + + assert args[0].max_stream_count == 1721 + + +def test_create_read_session_flattened_error(): + client = BigQueryReadClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.create_read_session( + storage.CreateReadSessionRequest(), + parent="parent_value", + read_session=stream.ReadSession(name="name_value"), + max_stream_count=1721, + ) + + +@pytest.mark.asyncio +async def test_create_read_session_flattened_async(): + client = BigQueryReadAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.create_read_session), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = stream.ReadSession() + + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(stream.ReadSession()) + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + response = await client.create_read_session( + parent="parent_value", + read_session=stream.ReadSession(name="name_value"), + max_stream_count=1721, + ) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0].parent == "parent_value" + + assert args[0].read_session == stream.ReadSession(name="name_value") + + assert args[0].max_stream_count == 1721 + + +@pytest.mark.asyncio +async def test_create_read_session_flattened_error_async(): + client = BigQueryReadAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + await client.create_read_session( + storage.CreateReadSessionRequest(), + parent="parent_value", + read_session=stream.ReadSession(name="name_value"), + max_stream_count=1721, + ) + + +def test_read_rows(transport: str = "grpc", request_type=storage.ReadRowsRequest): + client = BigQueryReadClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.read_rows), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = iter([storage.ReadRowsResponse()]) + + response = client.read_rows(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.ReadRowsRequest() + + # Establish that the response is the type that we expect. + for message in response: + assert isinstance(message, storage.ReadRowsResponse) + + +def test_read_rows_from_dict(): + test_read_rows(request_type=dict) + + +@pytest.mark.asyncio +async def test_read_rows_async( + transport: str = "grpc_asyncio", request_type=storage.ReadRowsRequest +): + client = BigQueryReadAsyncClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.read_rows), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = mock.Mock(aio.UnaryStreamCall, autospec=True) + call.return_value.read = mock.AsyncMock( + side_effect=[storage.ReadRowsResponse()] + ) + + response = await client.read_rows(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.ReadRowsRequest() + + # Establish that the response is the type that we expect. + message = await response.read() + assert isinstance(message, storage.ReadRowsResponse) + + +@pytest.mark.asyncio +async def test_read_rows_async_from_dict(): + await test_read_rows_async(request_type=dict) + + +def test_read_rows_field_headers(): + client = BigQueryReadClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.ReadRowsRequest() + request.read_stream = "read_stream/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.read_rows), "__call__") as call: + call.return_value = iter([storage.ReadRowsResponse()]) + + client.read_rows(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "read_stream=read_stream/value",) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_read_rows_field_headers_async(): + client = BigQueryReadAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.ReadRowsRequest() + request.read_stream = "read_stream/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.read_rows), "__call__") as call: + call.return_value = mock.Mock(aio.UnaryStreamCall, autospec=True) + call.return_value.read = mock.AsyncMock( + side_effect=[storage.ReadRowsResponse()] + ) + + await client.read_rows(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "read_stream=read_stream/value",) in kw["metadata"] + + +def test_read_rows_flattened(): + client = BigQueryReadClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.read_rows), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = iter([storage.ReadRowsResponse()]) + + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + client.read_rows( + read_stream="read_stream_value", offset=647, + ) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0].read_stream == "read_stream_value" + + assert args[0].offset == 647 + + +def test_read_rows_flattened_error(): + client = BigQueryReadClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.read_rows( + storage.ReadRowsRequest(), read_stream="read_stream_value", offset=647, + ) + + +@pytest.mark.asyncio +async def test_read_rows_flattened_async(): + client = BigQueryReadAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.read_rows), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = iter([storage.ReadRowsResponse()]) + + call.return_value = mock.Mock(aio.UnaryStreamCall, autospec=True) + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + response = await client.read_rows(read_stream="read_stream_value", offset=647,) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0].read_stream == "read_stream_value" + + assert args[0].offset == 647 + + +@pytest.mark.asyncio +async def test_read_rows_flattened_error_async(): + client = BigQueryReadAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + await client.read_rows( + storage.ReadRowsRequest(), read_stream="read_stream_value", offset=647, + ) + + +def test_split_read_stream( + transport: str = "grpc", request_type=storage.SplitReadStreamRequest +): + client = BigQueryReadClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.split_read_stream), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = storage.SplitReadStreamResponse() + + response = client.split_read_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.SplitReadStreamRequest() + + # Establish that the response is the type that we expect. + + assert isinstance(response, storage.SplitReadStreamResponse) + + +def test_split_read_stream_from_dict(): + test_split_read_stream(request_type=dict) + + +@pytest.mark.asyncio +async def test_split_read_stream_async( + transport: str = "grpc_asyncio", request_type=storage.SplitReadStreamRequest +): + client = BigQueryReadAsyncClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.split_read_stream), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + storage.SplitReadStreamResponse() + ) + + response = await client.split_read_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.SplitReadStreamRequest() + + # Establish that the response is the type that we expect. + assert isinstance(response, storage.SplitReadStreamResponse) + + +@pytest.mark.asyncio +async def test_split_read_stream_async_from_dict(): + await test_split_read_stream_async(request_type=dict) + + +def test_split_read_stream_field_headers(): + client = BigQueryReadClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.SplitReadStreamRequest() + request.name = "name/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.split_read_stream), "__call__" + ) as call: + call.return_value = storage.SplitReadStreamResponse() + + client.split_read_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "name=name/value",) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_split_read_stream_field_headers_async(): + client = BigQueryReadAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.SplitReadStreamRequest() + request.name = "name/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.split_read_stream), "__call__" + ) as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + storage.SplitReadStreamResponse() + ) + + await client.split_read_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "name=name/value",) in kw["metadata"] + + +def test_credentials_transport_error(): + # It is an error to provide credentials and a transport instance. + transport = transports.BigQueryReadGrpcTransport( + credentials=credentials.AnonymousCredentials(), + ) + with pytest.raises(ValueError): + client = BigQueryReadClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # It is an error to provide a credentials file and a transport instance. + transport = transports.BigQueryReadGrpcTransport( + credentials=credentials.AnonymousCredentials(), + ) + with pytest.raises(ValueError): + client = BigQueryReadClient( + client_options={"credentials_file": "credentials.json"}, + transport=transport, + ) + + # It is an error to provide scopes and a transport instance. + transport = transports.BigQueryReadGrpcTransport( + credentials=credentials.AnonymousCredentials(), + ) + with pytest.raises(ValueError): + client = BigQueryReadClient( + client_options={"scopes": ["1", "2"]}, transport=transport, + ) + + +def test_transport_instance(): + # A client may be instantiated with a custom transport instance. + transport = transports.BigQueryReadGrpcTransport( + credentials=credentials.AnonymousCredentials(), + ) + client = BigQueryReadClient(transport=transport) + assert client.transport is transport + + +def test_transport_get_channel(): + # A client may be instantiated with a custom transport instance. + transport = transports.BigQueryReadGrpcTransport( + credentials=credentials.AnonymousCredentials(), + ) + channel = transport.grpc_channel + assert channel + + transport = transports.BigQueryReadGrpcAsyncIOTransport( + credentials=credentials.AnonymousCredentials(), + ) + channel = transport.grpc_channel + assert channel + + +@pytest.mark.parametrize( + "transport_class", + [ + transports.BigQueryReadGrpcTransport, + transports.BigQueryReadGrpcAsyncIOTransport, + ], +) +def test_transport_adc(transport_class): + # Test default credentials are used if not provided. + with mock.patch.object(auth, "default") as adc: + adc.return_value = (credentials.AnonymousCredentials(), None) + transport_class() + adc.assert_called_once() + + +def test_transport_grpc_default(): + # A client should use the gRPC transport by default. + client = BigQueryReadClient(credentials=credentials.AnonymousCredentials(),) + assert isinstance(client.transport, transports.BigQueryReadGrpcTransport,) + + +def test_big_query_read_base_transport_error(): + # Passing both a credentials object and credentials_file should raise an error + with pytest.raises(exceptions.DuplicateCredentialArgs): + transport = transports.BigQueryReadTransport( + credentials=credentials.AnonymousCredentials(), + credentials_file="credentials.json", + ) + + +def test_big_query_read_base_transport(): + # Instantiate the base transport. + with mock.patch( + "google.cloud.bigquery_storage_v1beta2.services.big_query_read.transports.BigQueryReadTransport.__init__" + ) as Transport: + Transport.return_value = None + transport = transports.BigQueryReadTransport( + credentials=credentials.AnonymousCredentials(), + ) + + # Every method on the transport should just blindly + # raise NotImplementedError. + methods = ( + "create_read_session", + "read_rows", + "split_read_stream", + ) + for method in methods: + with pytest.raises(NotImplementedError): + getattr(transport, method)(request=object()) + + +def test_big_query_read_base_transport_with_credentials_file(): + # Instantiate the base transport with a credentials file + with mock.patch.object( + auth, "load_credentials_from_file" + ) as load_creds, mock.patch( + "google.cloud.bigquery_storage_v1beta2.services.big_query_read.transports.BigQueryReadTransport._prep_wrapped_messages" + ) as Transport: + Transport.return_value = None + load_creds.return_value = (credentials.AnonymousCredentials(), None) + transport = transports.BigQueryReadTransport( + credentials_file="credentials.json", quota_project_id="octopus", + ) + load_creds.assert_called_once_with( + "credentials.json", + scopes=( + "https://www.googleapis.com/auth/bigquery", + "https://www.googleapis.com/auth/bigquery.readonly", + "https://www.googleapis.com/auth/cloud-platform", + ), + quota_project_id="octopus", + ) + + +def test_big_query_read_base_transport_with_adc(): + # Test the default credentials are used if credentials and credentials_file are None. + with mock.patch.object(auth, "default") as adc, mock.patch( + "google.cloud.bigquery_storage_v1beta2.services.big_query_read.transports.BigQueryReadTransport._prep_wrapped_messages" + ) as Transport: + Transport.return_value = None + adc.return_value = (credentials.AnonymousCredentials(), None) + transport = transports.BigQueryReadTransport() + adc.assert_called_once() + + +def test_big_query_read_auth_adc(): + # If no credentials are provided, we should use ADC credentials. + with mock.patch.object(auth, "default") as adc: + adc.return_value = (credentials.AnonymousCredentials(), None) + BigQueryReadClient() + adc.assert_called_once_with( + scopes=( + "https://www.googleapis.com/auth/bigquery", + "https://www.googleapis.com/auth/bigquery.readonly", + "https://www.googleapis.com/auth/cloud-platform", + ), + quota_project_id=None, + ) + + +def test_big_query_read_transport_auth_adc(): + # If credentials and host are not provided, the transport class should use + # ADC credentials. + with mock.patch.object(auth, "default") as adc: + adc.return_value = (credentials.AnonymousCredentials(), None) + transports.BigQueryReadGrpcTransport( + host="squid.clam.whelk", quota_project_id="octopus" + ) + adc.assert_called_once_with( + scopes=( + "https://www.googleapis.com/auth/bigquery", + "https://www.googleapis.com/auth/bigquery.readonly", + "https://www.googleapis.com/auth/cloud-platform", + ), + quota_project_id="octopus", + ) + + +def test_big_query_read_host_no_port(): + client = BigQueryReadClient( + credentials=credentials.AnonymousCredentials(), + client_options=client_options.ClientOptions( + api_endpoint="bigquerystorage.googleapis.com" + ), + ) + assert client.transport._host == "bigquerystorage.googleapis.com:443" + + +def test_big_query_read_host_with_port(): + client = BigQueryReadClient( + credentials=credentials.AnonymousCredentials(), + client_options=client_options.ClientOptions( + api_endpoint="bigquerystorage.googleapis.com:8000" + ), + ) + assert client.transport._host == "bigquerystorage.googleapis.com:8000" + + +def test_big_query_read_grpc_transport_channel(): + channel = grpc.secure_channel("http://localhost/", grpc.local_channel_credentials()) + + # Check that channel is used if provided. + transport = transports.BigQueryReadGrpcTransport( + host="squid.clam.whelk", channel=channel, + ) + assert transport.grpc_channel == channel + assert transport._host == "squid.clam.whelk:443" + assert transport._ssl_channel_credentials == None + + +def test_big_query_read_grpc_asyncio_transport_channel(): + channel = aio.secure_channel("http://localhost/", grpc.local_channel_credentials()) + + # Check that channel is used if provided. + transport = transports.BigQueryReadGrpcAsyncIOTransport( + host="squid.clam.whelk", channel=channel, + ) + assert transport.grpc_channel == channel + assert transport._host == "squid.clam.whelk:443" + assert transport._ssl_channel_credentials == None + + +@pytest.mark.parametrize( + "transport_class", + [transports.BigQueryReadGrpcTransport, transports.BigQueryReadGrpcAsyncIOTransport], +) +def test_big_query_read_transport_channel_mtls_with_client_cert_source(transport_class): + with mock.patch( + "grpc.ssl_channel_credentials", autospec=True + ) as grpc_ssl_channel_cred: + with mock.patch.object( + transport_class, "create_channel" + ) as grpc_create_channel: + mock_ssl_cred = mock.Mock() + grpc_ssl_channel_cred.return_value = mock_ssl_cred + + mock_grpc_channel = mock.Mock() + grpc_create_channel.return_value = mock_grpc_channel + + cred = credentials.AnonymousCredentials() + with pytest.warns(DeprecationWarning): + with mock.patch.object(auth, "default") as adc: + adc.return_value = (cred, None) + transport = transport_class( + host="squid.clam.whelk", + api_mtls_endpoint="mtls.squid.clam.whelk", + client_cert_source=client_cert_source_callback, + ) + adc.assert_called_once() + + grpc_ssl_channel_cred.assert_called_once_with( + certificate_chain=b"cert bytes", private_key=b"key bytes" + ) + grpc_create_channel.assert_called_once_with( + "mtls.squid.clam.whelk:443", + credentials=cred, + credentials_file=None, + scopes=( + "https://www.googleapis.com/auth/bigquery", + "https://www.googleapis.com/auth/bigquery.readonly", + "https://www.googleapis.com/auth/cloud-platform", + ), + ssl_credentials=mock_ssl_cred, + quota_project_id=None, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + assert transport.grpc_channel == mock_grpc_channel + assert transport._ssl_channel_credentials == mock_ssl_cred + + +@pytest.mark.parametrize( + "transport_class", + [transports.BigQueryReadGrpcTransport, transports.BigQueryReadGrpcAsyncIOTransport], +) +def test_big_query_read_transport_channel_mtls_with_adc(transport_class): + mock_ssl_cred = mock.Mock() + with mock.patch.multiple( + "google.auth.transport.grpc.SslCredentials", + __init__=mock.Mock(return_value=None), + ssl_credentials=mock.PropertyMock(return_value=mock_ssl_cred), + ): + with mock.patch.object( + transport_class, "create_channel" + ) as grpc_create_channel: + mock_grpc_channel = mock.Mock() + grpc_create_channel.return_value = mock_grpc_channel + mock_cred = mock.Mock() + + with pytest.warns(DeprecationWarning): + transport = transport_class( + host="squid.clam.whelk", + credentials=mock_cred, + api_mtls_endpoint="mtls.squid.clam.whelk", + client_cert_source=None, + ) + + grpc_create_channel.assert_called_once_with( + "mtls.squid.clam.whelk:443", + credentials=mock_cred, + credentials_file=None, + scopes=( + "https://www.googleapis.com/auth/bigquery", + "https://www.googleapis.com/auth/bigquery.readonly", + "https://www.googleapis.com/auth/cloud-platform", + ), + ssl_credentials=mock_ssl_cred, + quota_project_id=None, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + assert transport.grpc_channel == mock_grpc_channel + + +def test_read_session_path(): + project = "squid" + location = "clam" + session = "whelk" + + expected = "projects/{project}/locations/{location}/sessions/{session}".format( + project=project, location=location, session=session, + ) + actual = BigQueryReadClient.read_session_path(project, location, session) + assert expected == actual + + +def test_parse_read_session_path(): + expected = { + "project": "octopus", + "location": "oyster", + "session": "nudibranch", + } + path = BigQueryReadClient.read_session_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryReadClient.parse_read_session_path(path) + assert expected == actual + + +def test_read_stream_path(): + project = "cuttlefish" + location = "mussel" + session = "winkle" + stream = "nautilus" + + expected = "projects/{project}/locations/{location}/sessions/{session}/streams/{stream}".format( + project=project, location=location, session=session, stream=stream, + ) + actual = BigQueryReadClient.read_stream_path(project, location, session, stream) + assert expected == actual + + +def test_parse_read_stream_path(): + expected = { + "project": "scallop", + "location": "abalone", + "session": "squid", + "stream": "clam", + } + path = BigQueryReadClient.read_stream_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryReadClient.parse_read_stream_path(path) + assert expected == actual + + +def test_table_path(): + project = "whelk" + dataset = "octopus" + table = "oyster" + + expected = "projects/{project}/datasets/{dataset}/tables/{table}".format( + project=project, dataset=dataset, table=table, + ) + actual = BigQueryReadClient.table_path(project, dataset, table) + assert expected == actual + + +def test_parse_table_path(): + expected = { + "project": "nudibranch", + "dataset": "cuttlefish", + "table": "mussel", + } + path = BigQueryReadClient.table_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryReadClient.parse_table_path(path) + assert expected == actual + + +def test_common_billing_account_path(): + billing_account = "winkle" + + expected = "billingAccounts/{billing_account}".format( + billing_account=billing_account, + ) + actual = BigQueryReadClient.common_billing_account_path(billing_account) + assert expected == actual + + +def test_parse_common_billing_account_path(): + expected = { + "billing_account": "nautilus", + } + path = BigQueryReadClient.common_billing_account_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryReadClient.parse_common_billing_account_path(path) + assert expected == actual + + +def test_common_folder_path(): + folder = "scallop" + + expected = "folders/{folder}".format(folder=folder,) + actual = BigQueryReadClient.common_folder_path(folder) + assert expected == actual + + +def test_parse_common_folder_path(): + expected = { + "folder": "abalone", + } + path = BigQueryReadClient.common_folder_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryReadClient.parse_common_folder_path(path) + assert expected == actual + + +def test_common_organization_path(): + organization = "squid" + + expected = "organizations/{organization}".format(organization=organization,) + actual = BigQueryReadClient.common_organization_path(organization) + assert expected == actual + + +def test_parse_common_organization_path(): + expected = { + "organization": "clam", + } + path = BigQueryReadClient.common_organization_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryReadClient.parse_common_organization_path(path) + assert expected == actual + + +def test_common_project_path(): + project = "whelk" + + expected = "projects/{project}".format(project=project,) + actual = BigQueryReadClient.common_project_path(project) + assert expected == actual + + +def test_parse_common_project_path(): + expected = { + "project": "octopus", + } + path = BigQueryReadClient.common_project_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryReadClient.parse_common_project_path(path) + assert expected == actual + + +def test_common_location_path(): + project = "oyster" + location = "nudibranch" + + expected = "projects/{project}/locations/{location}".format( + project=project, location=location, + ) + actual = BigQueryReadClient.common_location_path(project, location) + assert expected == actual + + +def test_parse_common_location_path(): + expected = { + "project": "cuttlefish", + "location": "mussel", + } + path = BigQueryReadClient.common_location_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryReadClient.parse_common_location_path(path) + assert expected == actual + + +def test_client_withDEFAULT_CLIENT_INFO(): + client_info = gapic_v1.client_info.ClientInfo() + + with mock.patch.object( + transports.BigQueryReadTransport, "_prep_wrapped_messages" + ) as prep: + client = BigQueryReadClient( + credentials=credentials.AnonymousCredentials(), client_info=client_info, + ) + prep.assert_called_once_with(client_info) + + with mock.patch.object( + transports.BigQueryReadTransport, "_prep_wrapped_messages" + ) as prep: + transport_class = BigQueryReadClient.get_transport_class() + transport = transport_class( + credentials=credentials.AnonymousCredentials(), client_info=client_info, + ) + prep.assert_called_once_with(client_info) diff --git a/tests/unit/gapic/bigquery_storage_v1beta2/test_big_query_write.py b/tests/unit/gapic/bigquery_storage_v1beta2/test_big_query_write.py new file mode 100644 index 00000000..98355b8a --- /dev/null +++ b/tests/unit/gapic/bigquery_storage_v1beta2/test_big_query_write.py @@ -0,0 +1,2043 @@ +# -*- 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 os +import mock + +import grpc +from grpc.experimental import aio +import math +import pytest +from proto.marshal.rules.dates import DurationRule, TimestampRule + +from google import auth +from google.api_core import client_options +from google.api_core import exceptions +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.auth import credentials +from google.auth.exceptions import MutualTLSChannelError +from google.cloud.bigquery_storage_v1beta2.services.big_query_write import ( + BigQueryWriteAsyncClient, +) +from google.cloud.bigquery_storage_v1beta2.services.big_query_write import ( + BigQueryWriteClient, +) +from google.cloud.bigquery_storage_v1beta2.services.big_query_write import transports +from google.cloud.bigquery_storage_v1beta2.types import protobuf +from google.cloud.bigquery_storage_v1beta2.types import storage +from google.cloud.bigquery_storage_v1beta2.types import stream +from google.cloud.bigquery_storage_v1beta2.types import table +from google.oauth2 import service_account +from google.protobuf import descriptor_pb2 as descriptor # type: ignore +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 + + +def client_cert_source_callback(): + return b"cert bytes", b"key bytes" + + +# If default endpoint is localhost, then default mtls endpoint will be the same. +# This method modifies the default endpoint so the client can produce a different +# mtls endpoint for endpoint testing purposes. +def modify_default_endpoint(client): + return ( + "foo.googleapis.com" + if ("localhost" in client.DEFAULT_ENDPOINT) + else client.DEFAULT_ENDPOINT + ) + + +def test__get_default_mtls_endpoint(): + api_endpoint = "example.googleapis.com" + api_mtls_endpoint = "example.mtls.googleapis.com" + sandbox_endpoint = "example.sandbox.googleapis.com" + sandbox_mtls_endpoint = "example.mtls.sandbox.googleapis.com" + non_googleapi = "api.example.com" + + assert BigQueryWriteClient._get_default_mtls_endpoint(None) is None + assert ( + BigQueryWriteClient._get_default_mtls_endpoint(api_endpoint) + == api_mtls_endpoint + ) + assert ( + BigQueryWriteClient._get_default_mtls_endpoint(api_mtls_endpoint) + == api_mtls_endpoint + ) + assert ( + BigQueryWriteClient._get_default_mtls_endpoint(sandbox_endpoint) + == sandbox_mtls_endpoint + ) + assert ( + BigQueryWriteClient._get_default_mtls_endpoint(sandbox_mtls_endpoint) + == sandbox_mtls_endpoint + ) + assert ( + BigQueryWriteClient._get_default_mtls_endpoint(non_googleapi) == non_googleapi + ) + + +def test_big_query_write_client_from_service_account_info(): + creds = credentials.AnonymousCredentials() + with mock.patch.object( + service_account.Credentials, "from_service_account_info" + ) as factory: + factory.return_value = creds + info = {"valid": True} + client = BigQueryWriteClient.from_service_account_info(info) + assert client.transport._credentials == creds + + assert client.transport._host == "bigquerystorage.googleapis.com:443" + + +@pytest.mark.parametrize( + "client_class", [BigQueryWriteClient, BigQueryWriteAsyncClient,] +) +def test_big_query_write_client_from_service_account_file(client_class): + creds = credentials.AnonymousCredentials() + with mock.patch.object( + service_account.Credentials, "from_service_account_file" + ) as factory: + factory.return_value = creds + client = client_class.from_service_account_file("dummy/file/path.json") + assert client.transport._credentials == creds + + client = client_class.from_service_account_json("dummy/file/path.json") + assert client.transport._credentials == creds + + assert client.transport._host == "bigquerystorage.googleapis.com:443" + + +def test_big_query_write_client_get_transport_class(): + transport = BigQueryWriteClient.get_transport_class() + available_transports = [ + transports.BigQueryWriteGrpcTransport, + ] + assert transport in available_transports + + transport = BigQueryWriteClient.get_transport_class("grpc") + assert transport == transports.BigQueryWriteGrpcTransport + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name", + [ + (BigQueryWriteClient, transports.BigQueryWriteGrpcTransport, "grpc"), + ( + BigQueryWriteAsyncClient, + transports.BigQueryWriteGrpcAsyncIOTransport, + "grpc_asyncio", + ), + ], +) +@mock.patch.object( + BigQueryWriteClient, + "DEFAULT_ENDPOINT", + modify_default_endpoint(BigQueryWriteClient), +) +@mock.patch.object( + BigQueryWriteAsyncClient, + "DEFAULT_ENDPOINT", + modify_default_endpoint(BigQueryWriteAsyncClient), +) +def test_big_query_write_client_client_options( + client_class, transport_class, transport_name +): + # Check that if channel is provided we won't create a new one. + with mock.patch.object(BigQueryWriteClient, "get_transport_class") as gtc: + transport = transport_class(credentials=credentials.AnonymousCredentials()) + client = client_class(transport=transport) + gtc.assert_not_called() + + # Check that if channel is provided via str we will create a new one. + with mock.patch.object(BigQueryWriteClient, "get_transport_class") as gtc: + client = client_class(transport=transport_name) + gtc.assert_called() + + # Check the case api_endpoint is provided. + options = client_options.ClientOptions(api_endpoint="squid.clam.whelk") + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host="squid.clam.whelk", + scopes=None, + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + # Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS_ENDPOINT is + # "never". + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "never"}): + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class() + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client.DEFAULT_ENDPOINT, + scopes=None, + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + # Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS_ENDPOINT is + # "always". + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "always"}): + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class() + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client.DEFAULT_MTLS_ENDPOINT, + scopes=None, + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + # Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS_ENDPOINT has + # unsupported value. + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "Unsupported"}): + with pytest.raises(MutualTLSChannelError): + client = client_class() + + # Check the case GOOGLE_API_USE_CLIENT_CERTIFICATE has unsupported value. + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "Unsupported"} + ): + with pytest.raises(ValueError): + client = client_class() + + # Check the case quota_project_id is provided + options = client_options.ClientOptions(quota_project_id="octopus") + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client.DEFAULT_ENDPOINT, + scopes=None, + ssl_channel_credentials=None, + quota_project_id="octopus", + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name,use_client_cert_env", + [ + (BigQueryWriteClient, transports.BigQueryWriteGrpcTransport, "grpc", "true"), + ( + BigQueryWriteAsyncClient, + transports.BigQueryWriteGrpcAsyncIOTransport, + "grpc_asyncio", + "true", + ), + (BigQueryWriteClient, transports.BigQueryWriteGrpcTransport, "grpc", "false"), + ( + BigQueryWriteAsyncClient, + transports.BigQueryWriteGrpcAsyncIOTransport, + "grpc_asyncio", + "false", + ), + ], +) +@mock.patch.object( + BigQueryWriteClient, + "DEFAULT_ENDPOINT", + modify_default_endpoint(BigQueryWriteClient), +) +@mock.patch.object( + BigQueryWriteAsyncClient, + "DEFAULT_ENDPOINT", + modify_default_endpoint(BigQueryWriteAsyncClient), +) +@mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "auto"}) +def test_big_query_write_client_mtls_env_auto( + client_class, transport_class, transport_name, use_client_cert_env +): + # This tests the endpoint autoswitch behavior. Endpoint is autoswitched to the default + # mtls endpoint, if GOOGLE_API_USE_CLIENT_CERTIFICATE is "true" and client cert exists. + + # Check the case client_cert_source is provided. Whether client cert is used depends on + # GOOGLE_API_USE_CLIENT_CERTIFICATE value. + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": use_client_cert_env} + ): + options = client_options.ClientOptions( + client_cert_source=client_cert_source_callback + ) + with mock.patch.object(transport_class, "__init__") as patched: + ssl_channel_creds = mock.Mock() + with mock.patch( + "grpc.ssl_channel_credentials", return_value=ssl_channel_creds + ): + patched.return_value = None + client = client_class(client_options=options) + + if use_client_cert_env == "false": + expected_ssl_channel_creds = None + expected_host = client.DEFAULT_ENDPOINT + else: + expected_ssl_channel_creds = ssl_channel_creds + expected_host = client.DEFAULT_MTLS_ENDPOINT + + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=expected_host, + scopes=None, + ssl_channel_credentials=expected_ssl_channel_creds, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + # Check the case ADC client cert is provided. Whether client cert is used depends on + # GOOGLE_API_USE_CLIENT_CERTIFICATE value. + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": use_client_cert_env} + ): + with mock.patch.object(transport_class, "__init__") as patched: + with mock.patch( + "google.auth.transport.grpc.SslCredentials.__init__", return_value=None + ): + with mock.patch( + "google.auth.transport.grpc.SslCredentials.is_mtls", + new_callable=mock.PropertyMock, + ) as is_mtls_mock: + with mock.patch( + "google.auth.transport.grpc.SslCredentials.ssl_credentials", + new_callable=mock.PropertyMock, + ) as ssl_credentials_mock: + if use_client_cert_env == "false": + is_mtls_mock.return_value = False + ssl_credentials_mock.return_value = None + expected_host = client.DEFAULT_ENDPOINT + expected_ssl_channel_creds = None + else: + is_mtls_mock.return_value = True + ssl_credentials_mock.return_value = mock.Mock() + expected_host = client.DEFAULT_MTLS_ENDPOINT + expected_ssl_channel_creds = ( + ssl_credentials_mock.return_value + ) + + patched.return_value = None + client = client_class() + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=expected_host, + scopes=None, + ssl_channel_credentials=expected_ssl_channel_creds, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + # Check the case client_cert_source and ADC client cert are not provided. + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": use_client_cert_env} + ): + with mock.patch.object(transport_class, "__init__") as patched: + with mock.patch( + "google.auth.transport.grpc.SslCredentials.__init__", return_value=None + ): + with mock.patch( + "google.auth.transport.grpc.SslCredentials.is_mtls", + new_callable=mock.PropertyMock, + ) as is_mtls_mock: + is_mtls_mock.return_value = False + patched.return_value = None + client = client_class() + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client.DEFAULT_ENDPOINT, + scopes=None, + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name", + [ + (BigQueryWriteClient, transports.BigQueryWriteGrpcTransport, "grpc"), + ( + BigQueryWriteAsyncClient, + transports.BigQueryWriteGrpcAsyncIOTransport, + "grpc_asyncio", + ), + ], +) +def test_big_query_write_client_client_options_scopes( + client_class, transport_class, transport_name +): + # Check the case scopes are provided. + options = client_options.ClientOptions(scopes=["1", "2"],) + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client.DEFAULT_ENDPOINT, + scopes=["1", "2"], + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name", + [ + (BigQueryWriteClient, transports.BigQueryWriteGrpcTransport, "grpc"), + ( + BigQueryWriteAsyncClient, + transports.BigQueryWriteGrpcAsyncIOTransport, + "grpc_asyncio", + ), + ], +) +def test_big_query_write_client_client_options_credentials_file( + client_class, transport_class, transport_name +): + # Check the case credentials file is provided. + options = client_options.ClientOptions(credentials_file="credentials.json") + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options) + patched.assert_called_once_with( + credentials=None, + credentials_file="credentials.json", + host=client.DEFAULT_ENDPOINT, + scopes=None, + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + +def test_big_query_write_client_client_options_from_dict(): + with mock.patch( + "google.cloud.bigquery_storage_v1beta2.services.big_query_write.transports.BigQueryWriteGrpcTransport.__init__" + ) as grpc_transport: + grpc_transport.return_value = None + client = BigQueryWriteClient( + client_options={"api_endpoint": "squid.clam.whelk"} + ) + grpc_transport.assert_called_once_with( + credentials=None, + credentials_file=None, + host="squid.clam.whelk", + scopes=None, + ssl_channel_credentials=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + ) + + +def test_create_write_stream( + transport: str = "grpc", request_type=storage.CreateWriteStreamRequest +): + client = BigQueryWriteClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.create_write_stream), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = stream.WriteStream( + name="name_value", type_=stream.WriteStream.Type.COMMITTED, + ) + + response = client.create_write_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.CreateWriteStreamRequest() + + # Establish that the response is the type that we expect. + + assert isinstance(response, stream.WriteStream) + + assert response.name == "name_value" + + assert response.type_ == stream.WriteStream.Type.COMMITTED + + +def test_create_write_stream_from_dict(): + test_create_write_stream(request_type=dict) + + +@pytest.mark.asyncio +async def test_create_write_stream_async( + transport: str = "grpc_asyncio", request_type=storage.CreateWriteStreamRequest +): + client = BigQueryWriteAsyncClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.create_write_stream), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + stream.WriteStream( + name="name_value", type_=stream.WriteStream.Type.COMMITTED, + ) + ) + + response = await client.create_write_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.CreateWriteStreamRequest() + + # Establish that the response is the type that we expect. + assert isinstance(response, stream.WriteStream) + + assert response.name == "name_value" + + assert response.type_ == stream.WriteStream.Type.COMMITTED + + +@pytest.mark.asyncio +async def test_create_write_stream_async_from_dict(): + await test_create_write_stream_async(request_type=dict) + + +def test_create_write_stream_field_headers(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.CreateWriteStreamRequest() + request.parent = "parent/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.create_write_stream), "__call__" + ) as call: + call.return_value = stream.WriteStream() + + client.create_write_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "parent=parent/value",) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_create_write_stream_field_headers_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.CreateWriteStreamRequest() + request.parent = "parent/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.create_write_stream), "__call__" + ) as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(stream.WriteStream()) + + await client.create_write_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "parent=parent/value",) in kw["metadata"] + + +def test_create_write_stream_flattened(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.create_write_stream), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = stream.WriteStream() + + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + client.create_write_stream( + parent="parent_value", write_stream=stream.WriteStream(name="name_value"), + ) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0].parent == "parent_value" + + assert args[0].write_stream == stream.WriteStream(name="name_value") + + +def test_create_write_stream_flattened_error(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.create_write_stream( + storage.CreateWriteStreamRequest(), + parent="parent_value", + write_stream=stream.WriteStream(name="name_value"), + ) + + +@pytest.mark.asyncio +async def test_create_write_stream_flattened_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.create_write_stream), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = stream.WriteStream() + + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(stream.WriteStream()) + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + response = await client.create_write_stream( + parent="parent_value", write_stream=stream.WriteStream(name="name_value"), + ) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0].parent == "parent_value" + + assert args[0].write_stream == stream.WriteStream(name="name_value") + + +@pytest.mark.asyncio +async def test_create_write_stream_flattened_error_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + await client.create_write_stream( + storage.CreateWriteStreamRequest(), + parent="parent_value", + write_stream=stream.WriteStream(name="name_value"), + ) + + +def test_append_rows(transport: str = "grpc", request_type=storage.AppendRowsRequest): + client = BigQueryWriteClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + requests = [request] + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.append_rows), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = iter([storage.AppendRowsResponse()]) + + response = client.append_rows(iter(requests)) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert next(args[0]) == request + + # Establish that the response is the type that we expect. + for message in response: + assert isinstance(message, storage.AppendRowsResponse) + + +def test_append_rows_from_dict(): + test_append_rows(request_type=dict) + + +@pytest.mark.asyncio +async def test_append_rows_async( + transport: str = "grpc_asyncio", request_type=storage.AppendRowsRequest +): + client = BigQueryWriteAsyncClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + requests = [request] + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.append_rows), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = mock.Mock(aio.StreamStreamCall, autospec=True) + call.return_value.read = mock.AsyncMock( + side_effect=[storage.AppendRowsResponse()] + ) + + response = await client.append_rows(iter(requests)) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert next(args[0]) == request + + # Establish that the response is the type that we expect. + message = await response.read() + assert isinstance(message, storage.AppendRowsResponse) + + +@pytest.mark.asyncio +async def test_append_rows_async_from_dict(): + await test_append_rows_async(request_type=dict) + + +def test_get_write_stream( + transport: str = "grpc", request_type=storage.GetWriteStreamRequest +): + client = BigQueryWriteClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.get_write_stream), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = stream.WriteStream( + name="name_value", type_=stream.WriteStream.Type.COMMITTED, + ) + + response = client.get_write_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.GetWriteStreamRequest() + + # Establish that the response is the type that we expect. + + assert isinstance(response, stream.WriteStream) + + assert response.name == "name_value" + + assert response.type_ == stream.WriteStream.Type.COMMITTED + + +def test_get_write_stream_from_dict(): + test_get_write_stream(request_type=dict) + + +@pytest.mark.asyncio +async def test_get_write_stream_async( + transport: str = "grpc_asyncio", request_type=storage.GetWriteStreamRequest +): + client = BigQueryWriteAsyncClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.get_write_stream), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + stream.WriteStream( + name="name_value", type_=stream.WriteStream.Type.COMMITTED, + ) + ) + + response = await client.get_write_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.GetWriteStreamRequest() + + # Establish that the response is the type that we expect. + assert isinstance(response, stream.WriteStream) + + assert response.name == "name_value" + + assert response.type_ == stream.WriteStream.Type.COMMITTED + + +@pytest.mark.asyncio +async def test_get_write_stream_async_from_dict(): + await test_get_write_stream_async(request_type=dict) + + +def test_get_write_stream_field_headers(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.GetWriteStreamRequest() + request.name = "name/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.get_write_stream), "__call__") as call: + call.return_value = stream.WriteStream() + + client.get_write_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "name=name/value",) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_get_write_stream_field_headers_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.GetWriteStreamRequest() + request.name = "name/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.get_write_stream), "__call__") as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(stream.WriteStream()) + + await client.get_write_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "name=name/value",) in kw["metadata"] + + +def test_get_write_stream_flattened(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.get_write_stream), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = stream.WriteStream() + + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + client.get_write_stream(name="name_value",) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0].name == "name_value" + + +def test_get_write_stream_flattened_error(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.get_write_stream( + storage.GetWriteStreamRequest(), name="name_value", + ) + + +@pytest.mark.asyncio +async def test_get_write_stream_flattened_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.get_write_stream), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = stream.WriteStream() + + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(stream.WriteStream()) + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + response = await client.get_write_stream(name="name_value",) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0].name == "name_value" + + +@pytest.mark.asyncio +async def test_get_write_stream_flattened_error_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + await client.get_write_stream( + storage.GetWriteStreamRequest(), name="name_value", + ) + + +def test_finalize_write_stream( + transport: str = "grpc", request_type=storage.FinalizeWriteStreamRequest +): + client = BigQueryWriteClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.finalize_write_stream), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = storage.FinalizeWriteStreamResponse(row_count=992,) + + response = client.finalize_write_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.FinalizeWriteStreamRequest() + + # Establish that the response is the type that we expect. + + assert isinstance(response, storage.FinalizeWriteStreamResponse) + + assert response.row_count == 992 + + +def test_finalize_write_stream_from_dict(): + test_finalize_write_stream(request_type=dict) + + +@pytest.mark.asyncio +async def test_finalize_write_stream_async( + transport: str = "grpc_asyncio", request_type=storage.FinalizeWriteStreamRequest +): + client = BigQueryWriteAsyncClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.finalize_write_stream), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + storage.FinalizeWriteStreamResponse(row_count=992,) + ) + + response = await client.finalize_write_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.FinalizeWriteStreamRequest() + + # Establish that the response is the type that we expect. + assert isinstance(response, storage.FinalizeWriteStreamResponse) + + assert response.row_count == 992 + + +@pytest.mark.asyncio +async def test_finalize_write_stream_async_from_dict(): + await test_finalize_write_stream_async(request_type=dict) + + +def test_finalize_write_stream_field_headers(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.FinalizeWriteStreamRequest() + request.name = "name/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.finalize_write_stream), "__call__" + ) as call: + call.return_value = storage.FinalizeWriteStreamResponse() + + client.finalize_write_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "name=name/value",) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_finalize_write_stream_field_headers_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.FinalizeWriteStreamRequest() + request.name = "name/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.finalize_write_stream), "__call__" + ) as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + storage.FinalizeWriteStreamResponse() + ) + + await client.finalize_write_stream(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "name=name/value",) in kw["metadata"] + + +def test_finalize_write_stream_flattened(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.finalize_write_stream), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = storage.FinalizeWriteStreamResponse() + + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + client.finalize_write_stream(name="name_value",) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0].name == "name_value" + + +def test_finalize_write_stream_flattened_error(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.finalize_write_stream( + storage.FinalizeWriteStreamRequest(), name="name_value", + ) + + +@pytest.mark.asyncio +async def test_finalize_write_stream_flattened_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.finalize_write_stream), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = storage.FinalizeWriteStreamResponse() + + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + storage.FinalizeWriteStreamResponse() + ) + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + response = await client.finalize_write_stream(name="name_value",) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0].name == "name_value" + + +@pytest.mark.asyncio +async def test_finalize_write_stream_flattened_error_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + await client.finalize_write_stream( + storage.FinalizeWriteStreamRequest(), name="name_value", + ) + + +def test_batch_commit_write_streams( + transport: str = "grpc", request_type=storage.BatchCommitWriteStreamsRequest +): + client = BigQueryWriteClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.batch_commit_write_streams), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = storage.BatchCommitWriteStreamsResponse() + + response = client.batch_commit_write_streams(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.BatchCommitWriteStreamsRequest() + + # Establish that the response is the type that we expect. + + assert isinstance(response, storage.BatchCommitWriteStreamsResponse) + + +def test_batch_commit_write_streams_from_dict(): + test_batch_commit_write_streams(request_type=dict) + + +@pytest.mark.asyncio +async def test_batch_commit_write_streams_async( + transport: str = "grpc_asyncio", request_type=storage.BatchCommitWriteStreamsRequest +): + client = BigQueryWriteAsyncClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.batch_commit_write_streams), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + storage.BatchCommitWriteStreamsResponse() + ) + + response = await client.batch_commit_write_streams(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.BatchCommitWriteStreamsRequest() + + # Establish that the response is the type that we expect. + assert isinstance(response, storage.BatchCommitWriteStreamsResponse) + + +@pytest.mark.asyncio +async def test_batch_commit_write_streams_async_from_dict(): + await test_batch_commit_write_streams_async(request_type=dict) + + +def test_batch_commit_write_streams_field_headers(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.BatchCommitWriteStreamsRequest() + request.parent = "parent/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.batch_commit_write_streams), "__call__" + ) as call: + call.return_value = storage.BatchCommitWriteStreamsResponse() + + client.batch_commit_write_streams(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "parent=parent/value",) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_batch_commit_write_streams_field_headers_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.BatchCommitWriteStreamsRequest() + request.parent = "parent/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.batch_commit_write_streams), "__call__" + ) as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + storage.BatchCommitWriteStreamsResponse() + ) + + await client.batch_commit_write_streams(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "parent=parent/value",) in kw["metadata"] + + +def test_batch_commit_write_streams_flattened(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.batch_commit_write_streams), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = storage.BatchCommitWriteStreamsResponse() + + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + client.batch_commit_write_streams(parent="parent_value",) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0].parent == "parent_value" + + +def test_batch_commit_write_streams_flattened_error(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.batch_commit_write_streams( + storage.BatchCommitWriteStreamsRequest(), parent="parent_value", + ) + + +@pytest.mark.asyncio +async def test_batch_commit_write_streams_flattened_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.batch_commit_write_streams), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = storage.BatchCommitWriteStreamsResponse() + + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + storage.BatchCommitWriteStreamsResponse() + ) + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + response = await client.batch_commit_write_streams(parent="parent_value",) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0].parent == "parent_value" + + +@pytest.mark.asyncio +async def test_batch_commit_write_streams_flattened_error_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + await client.batch_commit_write_streams( + storage.BatchCommitWriteStreamsRequest(), parent="parent_value", + ) + + +def test_flush_rows(transport: str = "grpc", request_type=storage.FlushRowsRequest): + client = BigQueryWriteClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.flush_rows), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = storage.FlushRowsResponse(offset=647,) + + response = client.flush_rows(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.FlushRowsRequest() + + # Establish that the response is the type that we expect. + + assert isinstance(response, storage.FlushRowsResponse) + + assert response.offset == 647 + + +def test_flush_rows_from_dict(): + test_flush_rows(request_type=dict) + + +@pytest.mark.asyncio +async def test_flush_rows_async( + transport: str = "grpc_asyncio", request_type=storage.FlushRowsRequest +): + client = BigQueryWriteAsyncClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.flush_rows), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + storage.FlushRowsResponse(offset=647,) + ) + + response = await client.flush_rows(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0] == storage.FlushRowsRequest() + + # Establish that the response is the type that we expect. + assert isinstance(response, storage.FlushRowsResponse) + + assert response.offset == 647 + + +@pytest.mark.asyncio +async def test_flush_rows_async_from_dict(): + await test_flush_rows_async(request_type=dict) + + +def test_flush_rows_field_headers(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.FlushRowsRequest() + request.write_stream = "write_stream/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.flush_rows), "__call__") as call: + call.return_value = storage.FlushRowsResponse() + + client.flush_rows(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "write_stream=write_stream/value",) in kw[ + "metadata" + ] + + +@pytest.mark.asyncio +async def test_flush_rows_field_headers_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = storage.FlushRowsRequest() + request.write_stream = "write_stream/value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.flush_rows), "__call__") as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + storage.FlushRowsResponse() + ) + + await client.flush_rows(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ("x-goog-request-params", "write_stream=write_stream/value",) in kw[ + "metadata" + ] + + +def test_flush_rows_flattened(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.flush_rows), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = storage.FlushRowsResponse() + + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + client.flush_rows(write_stream="write_stream_value",) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + + assert args[0].write_stream == "write_stream_value" + + +def test_flush_rows_flattened_error(): + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.flush_rows( + storage.FlushRowsRequest(), write_stream="write_stream_value", + ) + + +@pytest.mark.asyncio +async def test_flush_rows_flattened_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.flush_rows), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = storage.FlushRowsResponse() + + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + storage.FlushRowsResponse() + ) + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + response = await client.flush_rows(write_stream="write_stream_value",) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + + assert args[0].write_stream == "write_stream_value" + + +@pytest.mark.asyncio +async def test_flush_rows_flattened_error_async(): + client = BigQueryWriteAsyncClient(credentials=credentials.AnonymousCredentials(),) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + await client.flush_rows( + storage.FlushRowsRequest(), write_stream="write_stream_value", + ) + + +def test_credentials_transport_error(): + # It is an error to provide credentials and a transport instance. + transport = transports.BigQueryWriteGrpcTransport( + credentials=credentials.AnonymousCredentials(), + ) + with pytest.raises(ValueError): + client = BigQueryWriteClient( + credentials=credentials.AnonymousCredentials(), transport=transport, + ) + + # It is an error to provide a credentials file and a transport instance. + transport = transports.BigQueryWriteGrpcTransport( + credentials=credentials.AnonymousCredentials(), + ) + with pytest.raises(ValueError): + client = BigQueryWriteClient( + client_options={"credentials_file": "credentials.json"}, + transport=transport, + ) + + # It is an error to provide scopes and a transport instance. + transport = transports.BigQueryWriteGrpcTransport( + credentials=credentials.AnonymousCredentials(), + ) + with pytest.raises(ValueError): + client = BigQueryWriteClient( + client_options={"scopes": ["1", "2"]}, transport=transport, + ) + + +def test_transport_instance(): + # A client may be instantiated with a custom transport instance. + transport = transports.BigQueryWriteGrpcTransport( + credentials=credentials.AnonymousCredentials(), + ) + client = BigQueryWriteClient(transport=transport) + assert client.transport is transport + + +def test_transport_get_channel(): + # A client may be instantiated with a custom transport instance. + transport = transports.BigQueryWriteGrpcTransport( + credentials=credentials.AnonymousCredentials(), + ) + channel = transport.grpc_channel + assert channel + + transport = transports.BigQueryWriteGrpcAsyncIOTransport( + credentials=credentials.AnonymousCredentials(), + ) + channel = transport.grpc_channel + assert channel + + +@pytest.mark.parametrize( + "transport_class", + [ + transports.BigQueryWriteGrpcTransport, + transports.BigQueryWriteGrpcAsyncIOTransport, + ], +) +def test_transport_adc(transport_class): + # Test default credentials are used if not provided. + with mock.patch.object(auth, "default") as adc: + adc.return_value = (credentials.AnonymousCredentials(), None) + transport_class() + adc.assert_called_once() + + +def test_transport_grpc_default(): + # A client should use the gRPC transport by default. + client = BigQueryWriteClient(credentials=credentials.AnonymousCredentials(),) + assert isinstance(client.transport, transports.BigQueryWriteGrpcTransport,) + + +def test_big_query_write_base_transport_error(): + # Passing both a credentials object and credentials_file should raise an error + with pytest.raises(exceptions.DuplicateCredentialArgs): + transport = transports.BigQueryWriteTransport( + credentials=credentials.AnonymousCredentials(), + credentials_file="credentials.json", + ) + + +def test_big_query_write_base_transport(): + # Instantiate the base transport. + with mock.patch( + "google.cloud.bigquery_storage_v1beta2.services.big_query_write.transports.BigQueryWriteTransport.__init__" + ) as Transport: + Transport.return_value = None + transport = transports.BigQueryWriteTransport( + credentials=credentials.AnonymousCredentials(), + ) + + # Every method on the transport should just blindly + # raise NotImplementedError. + methods = ( + "create_write_stream", + "append_rows", + "get_write_stream", + "finalize_write_stream", + "batch_commit_write_streams", + "flush_rows", + ) + for method in methods: + with pytest.raises(NotImplementedError): + getattr(transport, method)(request=object()) + + +def test_big_query_write_base_transport_with_credentials_file(): + # Instantiate the base transport with a credentials file + with mock.patch.object( + auth, "load_credentials_from_file" + ) as load_creds, mock.patch( + "google.cloud.bigquery_storage_v1beta2.services.big_query_write.transports.BigQueryWriteTransport._prep_wrapped_messages" + ) as Transport: + Transport.return_value = None + load_creds.return_value = (credentials.AnonymousCredentials(), None) + transport = transports.BigQueryWriteTransport( + credentials_file="credentials.json", quota_project_id="octopus", + ) + load_creds.assert_called_once_with( + "credentials.json", + scopes=( + "https://www.googleapis.com/auth/bigquery", + "https://www.googleapis.com/auth/bigquery.insertdata", + "https://www.googleapis.com/auth/cloud-platform", + ), + quota_project_id="octopus", + ) + + +def test_big_query_write_base_transport_with_adc(): + # Test the default credentials are used if credentials and credentials_file are None. + with mock.patch.object(auth, "default") as adc, mock.patch( + "google.cloud.bigquery_storage_v1beta2.services.big_query_write.transports.BigQueryWriteTransport._prep_wrapped_messages" + ) as Transport: + Transport.return_value = None + adc.return_value = (credentials.AnonymousCredentials(), None) + transport = transports.BigQueryWriteTransport() + adc.assert_called_once() + + +def test_big_query_write_auth_adc(): + # If no credentials are provided, we should use ADC credentials. + with mock.patch.object(auth, "default") as adc: + adc.return_value = (credentials.AnonymousCredentials(), None) + BigQueryWriteClient() + adc.assert_called_once_with( + scopes=( + "https://www.googleapis.com/auth/bigquery", + "https://www.googleapis.com/auth/bigquery.insertdata", + "https://www.googleapis.com/auth/cloud-platform", + ), + quota_project_id=None, + ) + + +def test_big_query_write_transport_auth_adc(): + # If credentials and host are not provided, the transport class should use + # ADC credentials. + with mock.patch.object(auth, "default") as adc: + adc.return_value = (credentials.AnonymousCredentials(), None) + transports.BigQueryWriteGrpcTransport( + host="squid.clam.whelk", quota_project_id="octopus" + ) + adc.assert_called_once_with( + scopes=( + "https://www.googleapis.com/auth/bigquery", + "https://www.googleapis.com/auth/bigquery.insertdata", + "https://www.googleapis.com/auth/cloud-platform", + ), + quota_project_id="octopus", + ) + + +def test_big_query_write_host_no_port(): + client = BigQueryWriteClient( + credentials=credentials.AnonymousCredentials(), + client_options=client_options.ClientOptions( + api_endpoint="bigquerystorage.googleapis.com" + ), + ) + assert client.transport._host == "bigquerystorage.googleapis.com:443" + + +def test_big_query_write_host_with_port(): + client = BigQueryWriteClient( + credentials=credentials.AnonymousCredentials(), + client_options=client_options.ClientOptions( + api_endpoint="bigquerystorage.googleapis.com:8000" + ), + ) + assert client.transport._host == "bigquerystorage.googleapis.com:8000" + + +def test_big_query_write_grpc_transport_channel(): + channel = grpc.secure_channel("http://localhost/", grpc.local_channel_credentials()) + + # Check that channel is used if provided. + transport = transports.BigQueryWriteGrpcTransport( + host="squid.clam.whelk", channel=channel, + ) + assert transport.grpc_channel == channel + assert transport._host == "squid.clam.whelk:443" + assert transport._ssl_channel_credentials == None + + +def test_big_query_write_grpc_asyncio_transport_channel(): + channel = aio.secure_channel("http://localhost/", grpc.local_channel_credentials()) + + # Check that channel is used if provided. + transport = transports.BigQueryWriteGrpcAsyncIOTransport( + host="squid.clam.whelk", channel=channel, + ) + assert transport.grpc_channel == channel + assert transport._host == "squid.clam.whelk:443" + assert transport._ssl_channel_credentials == None + + +@pytest.mark.parametrize( + "transport_class", + [ + transports.BigQueryWriteGrpcTransport, + transports.BigQueryWriteGrpcAsyncIOTransport, + ], +) +def test_big_query_write_transport_channel_mtls_with_client_cert_source( + transport_class, +): + with mock.patch( + "grpc.ssl_channel_credentials", autospec=True + ) as grpc_ssl_channel_cred: + with mock.patch.object( + transport_class, "create_channel" + ) as grpc_create_channel: + mock_ssl_cred = mock.Mock() + grpc_ssl_channel_cred.return_value = mock_ssl_cred + + mock_grpc_channel = mock.Mock() + grpc_create_channel.return_value = mock_grpc_channel + + cred = credentials.AnonymousCredentials() + with pytest.warns(DeprecationWarning): + with mock.patch.object(auth, "default") as adc: + adc.return_value = (cred, None) + transport = transport_class( + host="squid.clam.whelk", + api_mtls_endpoint="mtls.squid.clam.whelk", + client_cert_source=client_cert_source_callback, + ) + adc.assert_called_once() + + grpc_ssl_channel_cred.assert_called_once_with( + certificate_chain=b"cert bytes", private_key=b"key bytes" + ) + grpc_create_channel.assert_called_once_with( + "mtls.squid.clam.whelk:443", + credentials=cred, + credentials_file=None, + scopes=( + "https://www.googleapis.com/auth/bigquery", + "https://www.googleapis.com/auth/bigquery.insertdata", + "https://www.googleapis.com/auth/cloud-platform", + ), + ssl_credentials=mock_ssl_cred, + quota_project_id=None, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + assert transport.grpc_channel == mock_grpc_channel + assert transport._ssl_channel_credentials == mock_ssl_cred + + +@pytest.mark.parametrize( + "transport_class", + [ + transports.BigQueryWriteGrpcTransport, + transports.BigQueryWriteGrpcAsyncIOTransport, + ], +) +def test_big_query_write_transport_channel_mtls_with_adc(transport_class): + mock_ssl_cred = mock.Mock() + with mock.patch.multiple( + "google.auth.transport.grpc.SslCredentials", + __init__=mock.Mock(return_value=None), + ssl_credentials=mock.PropertyMock(return_value=mock_ssl_cred), + ): + with mock.patch.object( + transport_class, "create_channel" + ) as grpc_create_channel: + mock_grpc_channel = mock.Mock() + grpc_create_channel.return_value = mock_grpc_channel + mock_cred = mock.Mock() + + with pytest.warns(DeprecationWarning): + transport = transport_class( + host="squid.clam.whelk", + credentials=mock_cred, + api_mtls_endpoint="mtls.squid.clam.whelk", + client_cert_source=None, + ) + + grpc_create_channel.assert_called_once_with( + "mtls.squid.clam.whelk:443", + credentials=mock_cred, + credentials_file=None, + scopes=( + "https://www.googleapis.com/auth/bigquery", + "https://www.googleapis.com/auth/bigquery.insertdata", + "https://www.googleapis.com/auth/cloud-platform", + ), + ssl_credentials=mock_ssl_cred, + quota_project_id=None, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + assert transport.grpc_channel == mock_grpc_channel + + +def test_table_path(): + project = "squid" + dataset = "clam" + table = "whelk" + + expected = "projects/{project}/datasets/{dataset}/tables/{table}".format( + project=project, dataset=dataset, table=table, + ) + actual = BigQueryWriteClient.table_path(project, dataset, table) + assert expected == actual + + +def test_parse_table_path(): + expected = { + "project": "octopus", + "dataset": "oyster", + "table": "nudibranch", + } + path = BigQueryWriteClient.table_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryWriteClient.parse_table_path(path) + assert expected == actual + + +def test_write_stream_path(): + project = "cuttlefish" + dataset = "mussel" + table = "winkle" + stream = "nautilus" + + expected = "projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}".format( + project=project, dataset=dataset, table=table, stream=stream, + ) + actual = BigQueryWriteClient.write_stream_path(project, dataset, table, stream) + assert expected == actual + + +def test_parse_write_stream_path(): + expected = { + "project": "scallop", + "dataset": "abalone", + "table": "squid", + "stream": "clam", + } + path = BigQueryWriteClient.write_stream_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryWriteClient.parse_write_stream_path(path) + assert expected == actual + + +def test_common_billing_account_path(): + billing_account = "whelk" + + expected = "billingAccounts/{billing_account}".format( + billing_account=billing_account, + ) + actual = BigQueryWriteClient.common_billing_account_path(billing_account) + assert expected == actual + + +def test_parse_common_billing_account_path(): + expected = { + "billing_account": "octopus", + } + path = BigQueryWriteClient.common_billing_account_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryWriteClient.parse_common_billing_account_path(path) + assert expected == actual + + +def test_common_folder_path(): + folder = "oyster" + + expected = "folders/{folder}".format(folder=folder,) + actual = BigQueryWriteClient.common_folder_path(folder) + assert expected == actual + + +def test_parse_common_folder_path(): + expected = { + "folder": "nudibranch", + } + path = BigQueryWriteClient.common_folder_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryWriteClient.parse_common_folder_path(path) + assert expected == actual + + +def test_common_organization_path(): + organization = "cuttlefish" + + expected = "organizations/{organization}".format(organization=organization,) + actual = BigQueryWriteClient.common_organization_path(organization) + assert expected == actual + + +def test_parse_common_organization_path(): + expected = { + "organization": "mussel", + } + path = BigQueryWriteClient.common_organization_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryWriteClient.parse_common_organization_path(path) + assert expected == actual + + +def test_common_project_path(): + project = "winkle" + + expected = "projects/{project}".format(project=project,) + actual = BigQueryWriteClient.common_project_path(project) + assert expected == actual + + +def test_parse_common_project_path(): + expected = { + "project": "nautilus", + } + path = BigQueryWriteClient.common_project_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryWriteClient.parse_common_project_path(path) + assert expected == actual + + +def test_common_location_path(): + project = "scallop" + location = "abalone" + + expected = "projects/{project}/locations/{location}".format( + project=project, location=location, + ) + actual = BigQueryWriteClient.common_location_path(project, location) + assert expected == actual + + +def test_parse_common_location_path(): + expected = { + "project": "squid", + "location": "clam", + } + path = BigQueryWriteClient.common_location_path(**expected) + + # Check that the path construction is reversible. + actual = BigQueryWriteClient.parse_common_location_path(path) + assert expected == actual + + +def test_client_withDEFAULT_CLIENT_INFO(): + client_info = gapic_v1.client_info.ClientInfo() + + with mock.patch.object( + transports.BigQueryWriteTransport, "_prep_wrapped_messages" + ) as prep: + client = BigQueryWriteClient( + credentials=credentials.AnonymousCredentials(), client_info=client_info, + ) + prep.assert_called_once_with(client_info) + + with mock.patch.object( + transports.BigQueryWriteTransport, "_prep_wrapped_messages" + ) as prep: + transport_class = BigQueryWriteClient.get_transport_class() + transport = transport_class( + credentials=credentials.AnonymousCredentials(), client_info=client_info, + ) + prep.assert_called_once_with(client_info)