def test_empty(self): from google.cloud.spanner_v1 import Type from google.cloud.spanner_v1 import StructType from google.cloud.spanner_v1 import TypeCode struct_type_pb = StructType(fields=[ StructType.Field(name="name", type_=Type(code=TypeCode.STRING)), StructType.Field(name="age", type_=Type(code=TypeCode.INT64)), ]) self.assertEqual(self._callFUT(rows=[], row_type=struct_type_pb), [])
def test_w_struct(self): from google.cloud.spanner_v1 import Type from google.cloud.spanner_v1 import StructType from google.cloud.spanner_v1 import TypeCode struct_type_pb = StructType(fields=[ StructType.Field(name="name", type_=Type(code=TypeCode.STRING)), StructType.Field(name="age", type_=Type(code=TypeCode.INT64)), ]) field_type = Type(code=TypeCode.STRUCT, struct_type=struct_type_pb) values = [u"phred", "32"] expected_values = [u"phred", 32] self.assertEqual(self._callFUT(values, field_type), expected_values)
def test_non_empty(self): from google.cloud.spanner_v1 import Type from google.cloud.spanner_v1 import StructType from google.cloud.spanner_v1 import TypeCode from google.cloud.spanner_v1._helpers import _make_list_value_pbs VALUES = [[u"phred", 32], [u"bharney", 31]] struct_type_pb = StructType(fields=[ StructType.Field(name="name", type_=Type(code=TypeCode.STRING)), StructType.Field(name="age", type_=Type(code=TypeCode.INT64)), ]) values_pbs = _make_list_value_pbs(VALUES) self.assertEqual( self._callFUT(rows=values_pbs, row_type=struct_type_pb), VALUES)
def test_w_struct(self): from google.protobuf.struct_pb2 import Value from google.cloud.spanner_v1 import Type from google.cloud.spanner_v1 import StructType from google.cloud.spanner_v1 import TypeCode from google.cloud.spanner_v1._helpers import _make_list_value_pb VALUES = [u"phred", 32] struct_type_pb = StructType(fields=[ StructType.Field(name="name", type_=Type(code=TypeCode.STRING)), StructType.Field(name="age", type_=Type(code=TypeCode.INT64)), ]) field_type = Type(code=TypeCode.STRUCT, struct_type=struct_type_pb) value_pb = Value(list_value=_make_list_value_pb(VALUES)) self.assertEqual(self._callFUT(value_pb, field_type), VALUES)
def test__merge_chunk_array_of_array_of_string(self): from google.cloud.spanner_v1 import StructType from google.cloud.spanner_v1 import Type from google.cloud.spanner_v1 import TypeCode subarray_type = Type(code=TypeCode.ARRAY, array_element_type=Type(code=TypeCode.STRING)) array_type = Type(code=TypeCode.ARRAY, array_element_type=subarray_type) iterator = _MockCancellableIterator() streamed = self._make_one(iterator) FIELDS = [StructType.Field(name="lolos", type_=array_type)] streamed._metadata = self._make_result_set_metadata(FIELDS) streamed._pending_chunk = [ [u"A", u"B"], [u"C"], ] chunk = [ [u"D"], [u"E", u"F"], ] merged = streamed._merge_chunk(chunk) expected = [ [u"A", u"B"], [u"CD"], [u"E", u"F"], ] self.assertEqual(merged, expected) self.assertIsNone(streamed._pending_chunk)
def test__merge_chunk_array_of_array_of_int(self): from google.cloud.spanner_v1 import StructType from google.cloud.spanner_v1 import Type from google.cloud.spanner_v1 import TypeCode subarray_type = Type(code=TypeCode.ARRAY, array_element_type=Type(code=TypeCode.INT64)) array_type = Type(code=TypeCode.ARRAY, array_element_type=subarray_type) iterator = _MockCancellableIterator() streamed = self._make_one(iterator) FIELDS = [StructType.Field(name="loloi", type_=array_type)] streamed._metadata = self._make_result_set_metadata(FIELDS) streamed._pending_chunk = [[0, 1], [2]] chunk = [[3], [4, 5]] merged = streamed._merge_chunk(chunk) expected = [ [0, 1], [23], [4, 5], ] self.assertEqual(merged, expected) self.assertIsNone(streamed._pending_chunk)
def _make_array_field(name, element_type_code=None, element_type=None): from google.cloud.spanner_v1 import StructType from google.cloud.spanner_v1 import Type from google.cloud.spanner_v1 import TypeCode if element_type is None: element_type = Type(code=element_type_code) array_type = Type(code=TypeCode.ARRAY, array_element_type=element_type) return StructType.Field(name=name, type_=array_type)
def test_it(self): from google.cloud.spanner_v1 import Type from google.cloud.spanner_v1 import TypeCode from google.cloud.spanner_v1 import StructType from google.cloud.spanner_v1 import param_types struct_type = StructType(fields=[ StructType.Field(name="name", type_=Type(code=TypeCode.STRING)), StructType.Field(name="count", type_=Type(code=TypeCode.INT64)), ]) expected = Type(code=TypeCode.STRUCT, struct_type=struct_type) found = param_types.Struct([ param_types.StructField("name", param_types.STRING), param_types.StructField("count", param_types.INT64), ]) self.assertEqual(found, expected)
def _make_struct_type(struct_type_fields): from google.cloud.spanner_v1 import StructType from google.cloud.spanner_v1 import Type from google.cloud.spanner_v1 import TypeCode fields = [ StructType.Field(name=key, type_=Type(code=value)) for key, value in struct_type_fields ] struct_type = StructType(fields=fields) return Type(code=TypeCode.STRUCT, struct_type=struct_type)
def StructField(name, field_type): # pylint: disable=invalid-name """Construct a field description protobuf. :type name: str :param name: the name of the field :type field_type: :class:`google.cloud.spanner_v1.Type` :param field_type: the type of the field :rtype: :class:`google.cloud.spanner_v1.StructType.Field` :returns: the appropriate struct-field-type protobuf """ return StructType.Field(name=name, type_=field_type)
def _make_scalar_field(name, type_): from google.cloud.spanner_v1 import StructType from google.cloud.spanner_v1 import Type return StructType.Field(name=name, type_=Type(code=type_))
def _execute_sql_helper( self, multi_use, first=True, count=0, partition=None, sql_count=0, query_options=None, timeout=google.api_core.gapic_v1.method.DEFAULT, retry=google.api_core.gapic_v1.method.DEFAULT, ): from google.protobuf.struct_pb2 import Struct from google.cloud.spanner_v1 import ( PartialResultSet, ResultSetMetadata, ResultSetStats, ) from google.cloud.spanner_v1 import ( TransactionSelector, TransactionOptions, ) from google.cloud.spanner_v1 import ExecuteSqlRequest from google.cloud.spanner_v1 import Type, StructType from google.cloud.spanner_v1 import TypeCode from google.cloud.spanner_v1._helpers import ( _make_value_pb, _merge_query_options, ) VALUES = [[u"bharney", u"rhubbyl", 31], [u"phred", u"phlyntstone", 32]] MODE = 2 # PROFILE struct_type_pb = StructType( fields=[ StructType.Field(name="first_name", type_=Type(code=TypeCode.STRING)), StructType.Field(name="last_name", type_=Type(code=TypeCode.STRING)), StructType.Field(name="age", type_=Type(code=TypeCode.INT64)), ] ) metadata_pb = ResultSetMetadata(row_type=struct_type_pb) stats_pb = ResultSetStats( query_stats=Struct(fields={"rows_returned": _make_value_pb(2)}) ) result_sets = [ PartialResultSet(metadata=metadata_pb), PartialResultSet(stats=stats_pb), ] for i in range(len(result_sets)): result_sets[i].values.extend(VALUES[i]) iterator = _MockIterator(*result_sets) database = _Database() api = database.spanner_api = self._make_spanner_api() api.execute_streaming_sql.return_value = iterator session = _Session(database) derived = self._makeDerived(session) derived._multi_use = multi_use derived._read_request_count = count derived._execute_sql_count = sql_count if not first: derived._transaction_id = TXN_ID result_set = derived.execute_sql( SQL_QUERY_WITH_PARAM, PARAMS, PARAM_TYPES, query_mode=MODE, query_options=query_options, partition=partition, retry=retry, timeout=timeout, ) self.assertEqual(derived._read_request_count, count + 1) if multi_use: self.assertIs(result_set._source, derived) else: self.assertIsNone(result_set._source) self.assertEqual(list(result_set), VALUES) self.assertEqual(result_set.metadata, metadata_pb) self.assertEqual(result_set.stats, stats_pb) txn_options = TransactionOptions( read_only=TransactionOptions.ReadOnly(strong=True) ) if multi_use: if first: expected_transaction = TransactionSelector(begin=txn_options) else: expected_transaction = TransactionSelector(id=TXN_ID) else: expected_transaction = TransactionSelector(single_use=txn_options) expected_params = Struct( fields={key: _make_value_pb(value) for (key, value) in PARAMS.items()} ) expected_query_options = database._instance._client._query_options if query_options: expected_query_options = _merge_query_options( expected_query_options, query_options ) expected_request = ExecuteSqlRequest( session=self.SESSION_NAME, sql=SQL_QUERY_WITH_PARAM, transaction=expected_transaction, params=expected_params, param_types=PARAM_TYPES, query_mode=MODE, query_options=expected_query_options, partition_token=partition, seqno=sql_count, ) api.execute_streaming_sql.assert_called_once_with( request=expected_request, metadata=[("google-cloud-resource-prefix", database.name)], timeout=timeout, retry=retry, ) self.assertEqual(derived._execute_sql_count, sql_count + 1) self.assertSpanAttributes( "CloudSpanner.ReadWriteTransaction", status=StatusCanonicalCode.OK, attributes=dict(BASE_ATTRIBUTES, **{"db.statement": SQL_QUERY_WITH_PARAM}), )
def _read_helper(self, multi_use, first=True, count=0, partition=None): from google.protobuf.struct_pb2 import Struct from google.cloud.spanner_v1 import ( PartialResultSet, ResultSetMetadata, ResultSetStats, ) from google.cloud.spanner_v1 import ( TransactionSelector, TransactionOptions, ) from google.cloud.spanner_v1 import ReadRequest from google.cloud.spanner_v1 import Type, StructType from google.cloud.spanner_v1 import TypeCode from google.cloud.spanner_v1.keyset import KeySet from google.cloud.spanner_v1._helpers import _make_value_pb VALUES = [[u"bharney", 31], [u"phred", 32]] struct_type_pb = StructType( fields=[ StructType.Field(name="name", type_=Type(code=TypeCode.STRING)), StructType.Field(name="age", type_=Type(code=TypeCode.INT64)), ] ) metadata_pb = ResultSetMetadata(row_type=struct_type_pb) stats_pb = ResultSetStats( query_stats=Struct(fields={"rows_returned": _make_value_pb(2)}) ) result_sets = [ PartialResultSet(metadata=metadata_pb), PartialResultSet(stats=stats_pb), ] for i in range(len(result_sets)): result_sets[i].values.extend(VALUES[i]) KEYS = [["*****@*****.**"], ["*****@*****.**"]] keyset = KeySet(keys=KEYS) INDEX = "email-address-index" LIMIT = 20 database = _Database() api = database.spanner_api = self._make_spanner_api() api.streaming_read.return_value = _MockIterator(*result_sets) session = _Session(database) derived = self._makeDerived(session) derived._multi_use = multi_use derived._read_request_count = count if not first: derived._transaction_id = TXN_ID if partition is not None: # 'limit' and 'partition' incompatible result_set = derived.read( TABLE_NAME, COLUMNS, keyset, index=INDEX, partition=partition ) else: result_set = derived.read( TABLE_NAME, COLUMNS, keyset, index=INDEX, limit=LIMIT ) self.assertEqual(derived._read_request_count, count + 1) if multi_use: self.assertIs(result_set._source, derived) else: self.assertIsNone(result_set._source) self.assertEqual(list(result_set), VALUES) self.assertEqual(result_set.metadata, metadata_pb) self.assertEqual(result_set.stats, stats_pb) txn_options = TransactionOptions( read_only=TransactionOptions.ReadOnly(strong=True) ) if multi_use: if first: expected_transaction = TransactionSelector(begin=txn_options) else: expected_transaction = TransactionSelector(id=TXN_ID) else: expected_transaction = TransactionSelector(single_use=txn_options) if partition is not None: expected_limit = 0 else: expected_limit = LIMIT expected_request = ReadRequest( session=self.SESSION_NAME, table=TABLE_NAME, columns=COLUMNS, key_set=keyset._to_pb(), transaction=expected_transaction, index=INDEX, limit=expected_limit, partition_token=partition, ) api.streaming_read.assert_called_once_with( request=expected_request, metadata=[("google-cloud-resource-prefix", database.name)], ) self.assertSpanAttributes( "CloudSpanner.ReadOnlyTransaction", attributes=dict( BASE_ATTRIBUTES, table_id=TABLE_NAME, columns=tuple(COLUMNS) ), )