# Copyright 2021 Google LLC All rights reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import collections

import datetime

import decimal

import math

import struct

import threading

import time

import pytest

import grpc

from google.rpc import code_pb2

from google.api_core import datetime_helpers

from google.api_core import exceptions

from google.cloud import spanner_v1

from google.cloud.spanner_admin_database_v1 import DatabaseDialect

from google.cloud._helpers import UTC

from google.cloud.spanner_v1.data_types import JsonObject

from tests import _helpers as ot_helpers

from . import _helpers

from . import _sample_data

SOME_DATE = datetime.date(2011, 1, 17)

SOME_TIME = datetime.datetime(1989, 1, 17, 17, 59, 12, 345612)

NANO_TIME = datetime_helpers.DatetimeWithNanoseconds(1995, 8, 31, nanosecond=987654321)

POS_INF = float("+inf")

NEG_INF = float("-inf")

(OTHER_NAN,) = struct.unpack("<d", b"\x01\x00\x01\x00\x00\x00\xf8\xff")

BYTES_1 = b"Ymlu"

BYTES_2 = b"Ym9vdHM="

NUMERIC_1 = decimal.Decimal("0.123456789")

NUMERIC_2 = decimal.Decimal("1234567890")

JSON_1 = JsonObject(

{

"sample_boolean": True,

"sample_int": 872163,

"sample float": 7871.298,

"sample_null": None,

"sample_string": "abcdef",

"sample_array": [23, 76, 19],

}

)

JSON_2 = JsonObject(

{"sample_object": {"name": "Anamika", "id": 2635}},

)

COUNTERS_TABLE = "counters"

COUNTERS_COLUMNS = ("name", "value")

ALL_TYPES_TABLE = "all_types"

LIVE_ALL_TYPES_COLUMNS = (

"pkey",

"int_value",

"int_array",

"bool_value",

"bool_array",

"bytes_value",

"bytes_array",

"date_value",

"date_array",

"float_value",

"float_array",

"string_value",

"string_array",

"timestamp_value",

"timestamp_array",

"numeric_value",

"numeric_array",

"json_value",

"json_array",

)

EMULATOR_ALL_TYPES_COLUMNS = LIVE_ALL_TYPES_COLUMNS[:-4]

# ToDo: Clean up generation of POSTGRES_ALL_TYPES_COLUMNS

POSTGRES_ALL_TYPES_COLUMNS = (

LIVE_ALL_TYPES_COLUMNS[:1]

+ LIVE_ALL_TYPES_COLUMNS[1:7:2]

+ LIVE_ALL_TYPES_COLUMNS[9:17:2]

)

AllTypesRowData = collections.namedtuple("AllTypesRowData", LIVE_ALL_TYPES_COLUMNS)

AllTypesRowData.__new__.__defaults__ = tuple([None for colum in LIVE_ALL_TYPES_COLUMNS])

EmulatorAllTypesRowData = collections.namedtuple(

"EmulatorAllTypesRowData", EMULATOR_ALL_TYPES_COLUMNS

)

EmulatorAllTypesRowData.__new__.__defaults__ = tuple(

[None for colum in EMULATOR_ALL_TYPES_COLUMNS]

)

PostGresAllTypesRowData = collections.namedtuple(

"PostGresAllTypesRowData", POSTGRES_ALL_TYPES_COLUMNS

)

PostGresAllTypesRowData.__new__.__defaults__ = tuple(

[None for colum in POSTGRES_ALL_TYPES_COLUMNS]

)

LIVE_ALL_TYPES_ROWDATA = (

# all nulls

AllTypesRowData(pkey=0),

# Non-null values

AllTypesRowData(pkey=101, int_value=123),

AllTypesRowData(pkey=102, bool_value=False),

AllTypesRowData(pkey=103, bytes_value=BYTES_1),

AllTypesRowData(pkey=104, date_value=SOME_DATE),

AllTypesRowData(pkey=105, float_value=1.4142136),

AllTypesRowData(pkey=106, string_value="VALUE"),

AllTypesRowData(pkey=107, timestamp_value=SOME_TIME),

AllTypesRowData(pkey=108, timestamp_value=NANO_TIME),

AllTypesRowData(pkey=109, numeric_value=NUMERIC_1),

AllTypesRowData(pkey=110, json_value=JSON_1),

AllTypesRowData(pkey=111, json_value=[JSON_1, JSON_2]),

# empty array values

AllTypesRowData(pkey=201, int_array=[]),

AllTypesRowData(pkey=202, bool_array=[]),

AllTypesRowData(pkey=203, bytes_array=[]),

AllTypesRowData(pkey=204, date_array=[]),

AllTypesRowData(pkey=205, float_array=[]),

AllTypesRowData(pkey=206, string_array=[]),

AllTypesRowData(pkey=207, timestamp_array=[]),

AllTypesRowData(pkey=208, numeric_array=[]),

AllTypesRowData(pkey=209, json_array=[]),

# non-empty array values, including nulls

AllTypesRowData(pkey=301, int_array=[123, 456, None]),

AllTypesRowData(pkey=302, bool_array=[True, False, None]),

AllTypesRowData(pkey=303, bytes_array=[BYTES_1, BYTES_2, None]),

AllTypesRowData(pkey=304, date_array=[SOME_DATE, None]),

AllTypesRowData(pkey=305, float_array=[3.1415926, 2.71828, None]),

AllTypesRowData(pkey=306, string_array=["One", "Two", None]),

AllTypesRowData(pkey=307, timestamp_array=[SOME_TIME, NANO_TIME, None]),

AllTypesRowData(pkey=308, numeric_array=[NUMERIC_1, NUMERIC_2, None]),

AllTypesRowData(pkey=309, json_array=[JSON_1, JSON_2, None]),

)

EMULATOR_ALL_TYPES_ROWDATA = (

# all nulls

EmulatorAllTypesRowData(pkey=0),

# Non-null values

EmulatorAllTypesRowData(pkey=101, int_value=123),

EmulatorAllTypesRowData(pkey=102, bool_value=False),

EmulatorAllTypesRowData(pkey=103, bytes_value=BYTES_1),

EmulatorAllTypesRowData(pkey=104, date_value=SOME_DATE),

EmulatorAllTypesRowData(pkey=105, float_value=1.4142136),

EmulatorAllTypesRowData(pkey=106, string_value="VALUE"),

EmulatorAllTypesRowData(pkey=107, timestamp_value=SOME_TIME),

EmulatorAllTypesRowData(pkey=108, timestamp_value=NANO_TIME),

# empty array values

EmulatorAllTypesRowData(pkey=201, int_array=[]),

EmulatorAllTypesRowData(pkey=202, bool_array=[]),

EmulatorAllTypesRowData(pkey=203, bytes_array=[]),

EmulatorAllTypesRowData(pkey=204, date_array=[]),

EmulatorAllTypesRowData(pkey=205, float_array=[]),

EmulatorAllTypesRowData(pkey=206, string_array=[]),

EmulatorAllTypesRowData(pkey=207, timestamp_array=[]),

# non-empty array values, including nulls

EmulatorAllTypesRowData(pkey=301, int_array=[123, 456, None]),

EmulatorAllTypesRowData(pkey=302, bool_array=[True, False, None]),

EmulatorAllTypesRowData(pkey=303, bytes_array=[BYTES_1, BYTES_2, None]),

EmulatorAllTypesRowData(pkey=304, date_array=[SOME_DATE, None]),

EmulatorAllTypesRowData(pkey=305, float_array=[3.1415926, 2.71828, None]),

EmulatorAllTypesRowData(pkey=306, string_array=["One", "Two", None]),

EmulatorAllTypesRowData(pkey=307, timestamp_array=[SOME_TIME, NANO_TIME, None]),

)

POSTGRES_ALL_TYPES_ROWDATA = (

# all nulls

PostGresAllTypesRowData(pkey=0),

# Non-null values

PostGresAllTypesRowData(pkey=101, int_value=123),

PostGresAllTypesRowData(pkey=102, bool_value=False),

PostGresAllTypesRowData(pkey=103, bytes_value=BYTES_1),

PostGresAllTypesRowData(pkey=105, float_value=1.4142136),

PostGresAllTypesRowData(pkey=106, string_value="VALUE"),

PostGresAllTypesRowData(pkey=107, timestamp_value=SOME_TIME),

PostGresAllTypesRowData(pkey=108, timestamp_value=NANO_TIME),

PostGresAllTypesRowData(pkey=109, numeric_value=NUMERIC_1),

)

if _helpers.USE_EMULATOR:

ALL_TYPES_COLUMNS = EMULATOR_ALL_TYPES_COLUMNS

ALL_TYPES_ROWDATA = EMULATOR_ALL_TYPES_ROWDATA

elif _helpers.DATABASE_DIALECT:

ALL_TYPES_COLUMNS = POSTGRES_ALL_TYPES_COLUMNS

ALL_TYPES_ROWDATA = POSTGRES_ALL_TYPES_ROWDATA

else:

ALL_TYPES_COLUMNS = LIVE_ALL_TYPES_COLUMNS

ALL_TYPES_ROWDATA = LIVE_ALL_TYPES_ROWDATA

@pytest.fixture(scope="session")

def sessions_database(shared_instance, database_operation_timeout, database_dialect):

database_name = _helpers.unique_id("test_sessions", separator="_")

pool = spanner_v1.BurstyPool(labels={"testcase": "session_api"})

if database_dialect == DatabaseDialect.POSTGRESQL:

sessions_database = shared_instance.database(

database_name,

pool=pool,

database_dialect=database_dialect,

)

operation = sessions_database.create()

operation.result(database_operation_timeout)

operation = sessions_database.update_ddl(ddl_statements=_helpers.DDL_STATEMENTS)

operation.result(database_operation_timeout)

else:

sessions_database = shared_instance.database(

database_name,

ddl_statements=_helpers.DDL_STATEMENTS,

pool=pool,

)

operation = sessions_database.create()

operation.result(database_operation_timeout)

_helpers.retry_has_all_dll(sessions_database.reload)()

# Some tests expect there to be a session present in the pool.

pool.put(pool.get())

yield sessions_database

sessions_database.drop()

@pytest.fixture(scope="function")

def sessions_to_delete():

to_delete = []

yield to_delete

for session in to_delete:

session.delete()

@pytest.fixture(scope="function")

def ot_exporter():

if ot_helpers.HAS_OPENTELEMETRY_INSTALLED:

ot_helpers.use_test_ot_exporter()

ot_exporter = ot_helpers.get_test_ot_exporter()

ot_exporter.clear() # XXX?

yield ot_exporter

ot_exporter.clear()

else:

yield None

def assert_no_spans(ot_exporter):

if ot_exporter is not None:

span_list = ot_exporter.get_finished_spans()

assert len(span_list) == 0

def assert_span_attributes(

ot_exporter, name, status=ot_helpers.StatusCode.OK, attributes=None, span=None

):

if ot_exporter is not None:

if not span:

span_list = ot_exporter.get_finished_spans()

assert len(span_list) == 1

span = span_list[0]

assert span.name == name

assert span.status.status_code == status

assert dict(span.attributes) == attributes

def _make_attributes(db_instance, **kwargs):

attributes = {

"db.type": "spanner",

"db.url": "spanner.googleapis.com",

"net.host.name": "spanner.googleapis.com",

"db.instance": db_instance,

}

attributes.update(kwargs)

return attributes

class _ReadAbortTrigger(object):

"""Helper for tests provoking abort-during-read."""

KEY1 = "key1"

KEY2 = "key2"

def __init__(self):

self.provoker_started = threading.Event()

self.provoker_done = threading.Event()

self.handler_running = threading.Event()

self.handler_done = threading.Event()

def _provoke_abort_unit_of_work(self, transaction):

keyset = spanner_v1.KeySet(keys=[(self.KEY1,)])

rows = list(transaction.read(COUNTERS_TABLE, COUNTERS_COLUMNS, keyset))

assert len(rows) == 1

row = rows[0]

value = row[1]

self.provoker_started.set()

self.handler_running.wait()

transaction.update(COUNTERS_TABLE, COUNTERS_COLUMNS, [[self.KEY1, value + 1]])

def provoke_abort(self, database):

database.run_in_transaction(self._provoke_abort_unit_of_work)

self.provoker_done.set()

def _handle_abort_unit_of_work(self, transaction):

keyset_1 = spanner_v1.KeySet(keys=[(self.KEY1,)])

rows_1 = list(transaction.read(COUNTERS_TABLE, COUNTERS_COLUMNS, keyset_1))

assert len(rows_1) == 1

row_1 = rows_1[0]

value_1 = row_1[1]

self.handler_running.set()

self.provoker_done.wait()

keyset_2 = spanner_v1.KeySet(keys=[(self.KEY2,)])

rows_2 = list(transaction.read(COUNTERS_TABLE, COUNTERS_COLUMNS, keyset_2))

assert len(rows_2) == 1

row_2 = rows_2[0]

value_2 = row_2[1]

transaction.update(

COUNTERS_TABLE, COUNTERS_COLUMNS, [[self.KEY2, value_1 + value_2]]

)

def handle_abort(self, database):

database.run_in_transaction(self._handle_abort_unit_of_work)

self.handler_done.set()

def test_session_crud(sessions_database):

session = sessions_database.session()

assert not session.exists()

session.create()

_helpers.retry_true(session.exists)()

session.delete()

_helpers.retry_false(session.exists)()

def test_batch_insert_then_read(sessions_database, ot_exporter):

db_name = sessions_database.name

sd = _sample_data

with sessions_database.batch() as batch:

batch.delete(sd.TABLE, sd.ALL)

batch.insert(sd.TABLE, sd.COLUMNS, sd.ROW_DATA)

with sessions_database.snapshot(read_timestamp=batch.committed) as snapshot:

rows = list(snapshot.read(sd.TABLE, sd.COLUMNS, sd.ALL))

sd._check_rows_data(rows)

if ot_exporter is not None:

span_list = ot_exporter.get_finished_spans()

assert len(span_list) == 4

assert_span_attributes(

ot_exporter,

"CloudSpanner.GetSession",

attributes=_make_attributes(db_name, session_found=True),

span=span_list[0],

)

assert_span_attributes(

ot_exporter,

"CloudSpanner.Commit",

attributes=_make_attributes(db_name, num_mutations=2),

span=span_list[1],

)

assert_span_attributes(

ot_exporter,

"CloudSpanner.GetSession",

attributes=_make_attributes(db_name, session_found=True),

span=span_list[2],

)

assert_span_attributes(

ot_exporter,

"CloudSpanner.ReadOnlyTransaction",

attributes=_make_attributes(db_name, columns=sd.COLUMNS, table_id=sd.TABLE),

span=span_list[3],

)

def test_batch_insert_then_read_string_array_of_string(sessions_database, not_postgres):

table = "string_plus_array_of_string"

columns = ["id", "name", "tags"]

rowdata = [

(0, None, None),

(1, "phred", ["yabba", "dabba", "do"]),

(2, "bharney", []),

(3, "wylma", ["oh", None, "phred"]),

]

sd = _sample_data

with sessions_database.batch() as batch:

batch.delete(table, sd.ALL)

batch.insert(table, columns, rowdata)

with sessions_database.snapshot(read_timestamp=batch.committed) as snapshot:

rows = list(snapshot.read(table, columns, sd.ALL))

sd._check_rows_data(rows, expected=rowdata)

def test_batch_insert_then_read_all_datatypes(sessions_database):

sd = _sample_data

with sessions_database.batch() as batch:

batch.delete(ALL_TYPES_TABLE, sd.ALL)

batch.insert(ALL_TYPES_TABLE, ALL_TYPES_COLUMNS, ALL_TYPES_ROWDATA)

with sessions_database.snapshot(read_timestamp=batch.committed) as snapshot:

rows = list(snapshot.read(ALL_TYPES_TABLE, ALL_TYPES_COLUMNS, sd.ALL))

sd._check_rows_data(rows, expected=ALL_TYPES_ROWDATA)

def test_batch_insert_or_update_then_query(sessions_database):

sd = _sample_data

with sessions_database.batch() as batch:

batch.insert_or_update(sd.TABLE, sd.COLUMNS, sd.ROW_DATA)

with sessions_database.snapshot(read_timestamp=batch.committed) as snapshot:

rows = list(snapshot.execute_sql(sd.SQL))

sd._check_rows_data(rows)

def test_batch_insert_w_commit_timestamp(sessions_database, not_postgres):

table = "users_history"

columns = ["id", "commit_ts", "name", "email", "deleted"]

user_id = 1234

name = "phred"

email = "phred@example.com"

row_data = [[user_id, spanner_v1.COMMIT_TIMESTAMP, name, email, False]]

sd = _sample_data

with sessions_database.batch() as batch:

batch.delete(table, sd.ALL)

batch.insert(table, columns, row_data)

with sessions_database.snapshot(read_timestamp=batch.committed) as snapshot:

rows = list(snapshot.read(table, columns, sd.ALL))

assert len(rows) == 1

r_id, commit_ts, r_name, r_email, deleted = rows[0]

assert r_id == user_id

assert commit_ts == batch.committed

assert r_name == name

assert r_email == email

assert not deleted

@_helpers.retry_mabye_aborted_txn

def test_transaction_read_and_insert_then_rollback(

sessions_database,

ot_exporter,

sessions_to_delete,

):

sd = _sample_data

db_name = sessions_database.name

session = sessions_database.session()

session.create()

sessions_to_delete.append(session)

with sessions_database.batch() as batch:

batch.delete(sd.TABLE, sd.ALL)

transaction = session.transaction()

transaction.begin()

rows = list(transaction.read(sd.TABLE, sd.COLUMNS, sd.ALL))

assert rows == []

transaction.insert(sd.TABLE, sd.COLUMNS, sd.ROW_DATA)

# Inserted rows can't be read until after commit.

rows = list(transaction.read(sd.TABLE, sd.COLUMNS, sd.ALL))

assert rows == []

transaction.rollback()

rows = list(session.read(sd.TABLE, sd.COLUMNS, sd.ALL))

assert rows == []

if ot_exporter is not None:

span_list = ot_exporter.get_finished_spans()

assert len(span_list) == 8

assert_span_attributes(

ot_exporter,

"CloudSpanner.CreateSession",

attributes=_make_attributes(db_name),

span=span_list[0],

)

assert_span_attributes(

ot_exporter,

"CloudSpanner.GetSession",

attributes=_make_attributes(db_name, session_found=True),

span=span_list[1],

)

assert_span_attributes(

ot_exporter,

"CloudSpanner.Commit",

attributes=_make_attributes(db_name, num_mutations=1),

span=span_list[2],

)

assert_span_attributes(

ot_exporter,

"CloudSpanner.BeginTransaction",

attributes=_make_attributes(db_name),

span=span_list[3],

)

assert_span_attributes(

ot_exporter,

"CloudSpanner.ReadOnlyTransaction",

attributes=_make_attributes(

db_name,

table_id=sd.TABLE,

columns=sd.COLUMNS,

),

span=span_list[4],

)

assert_span_attributes(

ot_exporter,

"CloudSpanner.ReadOnlyTransaction",

attributes=_make_attributes(

db_name,

table_id=sd.TABLE,

columns=sd.COLUMNS,

),

span=span_list[5],

)

assert_span_attributes(

ot_exporter,

"CloudSpanner.Rollback",

attributes=_make_attributes(db_name),

span=span_list[6],

)

assert_span_attributes(

ot_exporter,

"CloudSpanner.ReadOnlyTransaction",

attributes=_make_attributes(

db_name,

table_id=sd.TABLE,

columns=sd.COLUMNS,

),

span=span_list[7],

)

@_helpers.retry_mabye_conflict

def test_transaction_read_and_insert_then_exception(sessions_database):

class CustomException(Exception):

pass

sd = _sample_data

def _transaction_read_then_raise(transaction):

rows = list(transaction.read(sd.TABLE, sd.COLUMNS, sd.ALL))

assert len(rows) == 0

transaction.insert(sd.TABLE, sd.COLUMNS, sd.ROW_DATA)

raise CustomException()

with sessions_database.batch() as batch:

batch.delete(sd.TABLE, sd.ALL)

with pytest.raises(CustomException):

sessions_database.run_in_transaction(_transaction_read_then_raise)

# Transaction was rolled back.

with sessions_database.snapshot() as snapshot:

rows = list(snapshot.read(sd.TABLE, sd.COLUMNS, sd.ALL))

assert rows == []

@_helpers.retry_mabye_conflict

def test_transaction_read_and_insert_or_update_then_commit(

sessions_database,

sessions_to_delete,

):

# [START spanner_test_dml_read_your_writes]

sd = _sample_data

session = sessions_database.session()

session.create()

sessions_to_delete.append(session)

with session.batch() as batch:

batch.delete(sd.TABLE, sd.ALL)

with session.transaction() as transaction:

rows = list(transaction.read(sd.TABLE, sd.COLUMNS, sd.ALL))

assert rows == []

transaction.insert_or_update(sd.TABLE, sd.COLUMNS, sd.ROW_DATA)

# Inserted rows can't be read until after commit.

rows = list(transaction.read(sd.TABLE, sd.COLUMNS, sd.ALL))

assert rows == []

rows = list(session.read(sd.TABLE, sd.COLUMNS, sd.ALL))

sd._check_rows_data(rows)

# [END spanner_test_dml_read_your_writes]

def _generate_insert_statements():

table = _sample_data.TABLE

column_list = ", ".join(_sample_data.COLUMNS)

for row in _sample_data.ROW_DATA:

row_data = "{}, '{}', '{}', '{}'".format(*row)

yield f"INSERT INTO {table} ({column_list}) VALUES ({row_data})"

@_helpers.retry_mabye_conflict

@_helpers.retry_mabye_aborted_txn

def test_transaction_execute_sql_w_dml_read_rollback(

sessions_database,

sessions_to_delete,

):

# [START spanner_test_dml_rollback_txn_not_committed]

sd = _sample_data

session = sessions_database.session()

session.create()

sessions_to_delete.append(session)

with session.batch() as batch:

batch.delete(sd.TABLE, sd.ALL)

transaction = session.transaction()

transaction.begin()

rows = list(transaction.read(sd.TABLE, sd.COLUMNS, sd.ALL))

assert rows == []

for insert_statement in _generate_insert_statements():

result = transaction.execute_sql(insert_statement)

list(result) # iterate to get stats

assert result.stats.row_count_exact == 1

# Rows inserted via DML *can* be read before commit.

during_rows = list(transaction.read(sd.TABLE, sd.COLUMNS, sd.ALL))

sd._check_rows_data(during_rows)

transaction.rollback()

rows = list(session.read(sd.TABLE, sd.COLUMNS, sd.ALL))

sd._check_rows_data(rows, [])

# [END spanner_test_dml_rollback_txn_not_committed]

@_helpers.retry_mabye_conflict

def test_transaction_execute_update_read_commit(sessions_database, sessions_to_delete):

# [START spanner_test_dml_read_your_writes]

sd = _sample_data

session = sessions_database.session()

session.create()

sessions_to_delete.append(session)

with session.batch() as batch:

batch.delete(sd.TABLE, sd.ALL)

with session.transaction() as transaction:

rows = list(transaction.read(sd.TABLE, sd.COLUMNS, sd.ALL))

assert rows == []

for insert_statement in _generate_insert_statements():

row_count = transaction.execute_update(insert_statement)

assert row_count == 1

# Rows inserted via DML *can* be read before commit.

during_rows = list(transaction.read(sd.TABLE, sd.COLUMNS, sd.ALL))

sd._check_rows_data(during_rows)

rows = list(session.read(sd.TABLE, sd.COLUMNS, sd.ALL))

sd._check_rows_data(rows)

# [END spanner_test_dml_read_your_writes]

@_helpers.retry_mabye_conflict

def test_transaction_execute_update_then_insert_commit(

sessions_database, sessions_to_delete

):

# [START spanner_test_dml_with_mutation]

# [START spanner_test_dml_update]

sd = _sample_data

session = sessions_database.session()

session.create()

sessions_to_delete.append(session)

with session.batch() as batch:

batch.delete(sd.TABLE, sd.ALL)

insert_statement = list(_generate_insert_statements())[0]

with session.transaction() as transaction:

rows = list(transaction.read(sd.TABLE, sd.COLUMNS, sd.ALL))

assert rows == []

row_count = transaction.execute_update(insert_statement)

assert row_count == 1

transaction.insert(sd.TABLE, sd.COLUMNS, sd.ROW_DATA[1:])

rows = list(session.read(sd.TABLE, sd.COLUMNS, sd.ALL))

sd._check_rows_data(rows)

# [END spanner_test_dml_update]

# [END spanner_test_dml_with_mutation]

def test_transaction_batch_update_success(

sessions_database, sessions_to_delete, database_dialect

):

# [START spanner_test_dml_with_mutation]

# [START spanner_test_dml_update]

sd = _sample_data

param_types = spanner_v1.param_types

session = sessions_database.session()

session.create()

sessions_to_delete.append(session)

with session.batch() as batch:

batch.delete(sd.TABLE, sd.ALL)

keys = (

["p1", "p2"]

if database_dialect == DatabaseDialect.POSTGRESQL

else ["contact_id", "email"]

)

placeholders = (

["$1", "$2"]

if database_dialect == DatabaseDialect.POSTGRESQL

else [f"@{key}" for key in keys]

)

insert_statement = list(_generate_insert_statements())[0]

update_statement = (

f"UPDATE contacts SET email = {placeholders[1]} WHERE contact_id = {placeholders[0]};",

{keys[0]: 1, keys[1]: "phreddy@example.com"},

{keys[0]: param_types.INT64, keys[1]: param_types.STRING},

)

delete_statement = (

f"DELETE FROM contacts WHERE contact_id = {placeholders[0]};",

{keys[0]: 1},

{keys[0]: param_types.INT64},

)

def unit_of_work(transaction):

rows = list(transaction.read(sd.TABLE, sd.COLUMNS, sd.ALL))

assert rows == []

status, row_counts = transaction.batch_update(

[insert_statement, update_statement, delete_statement]

)

_check_batch_status(status.code)

assert len(row_counts) == 3

for row_count in row_counts:

assert row_count == 1

session.run_in_transaction(unit_of_work)

rows = list(session.read(sd.TABLE, sd.COLUMNS, sd.ALL))

sd._check_rows_data(rows, [])

# [END spanner_test_dml_with_mutation]

# [END spanner_test_dml_update]

def test_transaction_batch_update_and_execute_dml(

sessions_database, sessions_to_delete, database_dialect

):

sd = _sample_data

param_types = spanner_v1.param_types

session = sessions_database.session()

session.create()

sessions_to_delete.append(session)

with session.batch() as batch:

batch.delete(sd.TABLE, sd.ALL)

keys = (

["p1", "p2"]

if database_dialect == DatabaseDialect.POSTGRESQL

else ["contact_id", "email"]

)

placeholders = (

["$1", "$2"]

if database_dialect == DatabaseDialect.POSTGRESQL

else [f"@{key}" for key in keys]

)

insert_statements = list(_generate_insert_statements())

update_statements = [

(

f"UPDATE contacts SET email = {placeholders[1]} WHERE contact_id = {placeholders[0]};",

{keys[0]: 1, keys[1]: "phreddy@example.com"},

{keys[0]: param_types.INT64, keys[1]: param_types.STRING},

)

]

delete_statement = "DELETE FROM contacts WHERE TRUE;"

def unit_of_work(transaction):

rows = list(transaction.read(sd.TABLE, sd.COLUMNS, sd.ALL))

assert rows == []

status, row_counts = transaction.batch_update(

insert_statements + update_statements

)

_check_batch_status(status.code)

assert len(row_counts) == len(insert_statements) + 1

for row_count in row_counts:

assert row_count == 1

row_count = transaction.execute_update(delete_statement)

assert row_count == len(insert_statements)

session.run_in_transaction(unit_of_work)

rows = list(session.read(sd.TABLE, sd.COLUMNS, sd.ALL))

sd._check_rows_data(rows, [])

def test_transaction_batch_update_w_syntax_error(

sessions_database, sessions_to_delete, database_dialect

):

from google.rpc import code_pb2

sd = _sample_data

param_types = spanner_v1.param_types

session = sessions_database.session()

session.create()

sessions_to_delete.append(session)

with session.batch() as batch:

batch.delete(sd.TABLE, sd.ALL)

keys = (

["p1", "p2"]

if database_dialect == DatabaseDialect.POSTGRESQL

else ["contact_id", "email"]

)

placeholders = (

["$1", "$2"]

if database_dialect == DatabaseDialect.POSTGRESQL

else [f"@{key}" for key in keys]

)

insert_statement = list(_generate_insert_statements())[0]

update_statement = (

f"UPDTAE contacts SET email = {placeholders[1]} WHERE contact_id = {placeholders[0]};",

{keys[0]: 1, keys[1]: "phreddy@example.com"},

{keys[0]: param_types.INT64, keys[1]: param_types.STRING},

)

delete_statement = (

f"DELETE FROM contacts WHERE contact_id = {placeholders[0]};",

{keys[0]: 1},

{keys[0]: param_types.INT64},

)

def unit_of_work(transaction):

rows = list(transaction.read(sd.TABLE, sd.COLUMNS, sd.ALL))

assert rows == []

status, row_counts = transaction.batch_update(

[insert_statement, update_statement, delete_statement]

)

_check_batch_status(status.code, code_pb2.INVALID_ARGUMENT)

assert len(row_counts) == 1

assert row_counts[0] == 1

session.run_in_transaction(unit_of_work)

def test_transaction_batch_update_wo_statements(sessions_database, sessions_to_delete):

session = sessions_database.session()

session.create()

sessions_to_delete.append(session)

with session.transaction() as transaction:

with pytest.raises(exceptions.InvalidArgument):

transaction.batch_update([])

@pytest.mark.skipif(

not ot_helpers.HAS_OPENTELEMETRY_INSTALLED,

reason="trace requires OpenTelemetry",

)

def test_transaction_batch_update_w_parent_span(

sessions_database, sessions_to_delete, ot_exporter, database_dialect

):

from opentelemetry import trace

sd = _sample_data

param_types = spanner_v1.param_types

tracer = trace.get_tracer(__name__)

session = sessions_database.session()

session.create()

sessions_to_delete.append(session)

with session.batch() as batch:

batch.delete(sd.TABLE, sd.ALL)

keys = (

["p1", "p2"]

if database_dialect == DatabaseDialect.POSTGRESQL

else ["contact_id", "email"]

)

placeholders = (

["$1", "$2"]

if database_dialect == DatabaseDialect.POSTGRESQL

else [f"@{key}" for key in keys]

)

insert_statement = list(_generate_insert_statements())[0]

update_statement = (

f"UPDATE contacts SET email = {placeholders[1]} WHERE contact_id = {placeholders[0]};",

{keys[0]: 1, keys[1]: "phreddy@example.com"},

{keys[0]: param_types.INT64, keys[1]: param_types.STRING},

)

delete_statement = (

f"DELETE FROM contacts WHERE contact_id = {placeholders[0]};",

{keys[0]: 1},

{keys[0]: param_types.INT64},

)

def unit_of_work(transaction):

status, row_counts = transaction.batch_update(

[insert_statement, update_statement, delete_statement]

)

_check_batch_status(status.code)

assert len(row_counts) == 3

for row_count in row_counts:

assert row_count == 1

with tracer.start_as_current_span("Test Span"):

session.run_in_transaction(unit_of_work)

span_list = ot_exporter.get_finished_spans()

assert len(span_list) == 6

expected_span_names = [

"CloudSpanner.CreateSession",

"CloudSpanner.Commit",

"CloudSpanner.BeginTransaction",

"CloudSpanner.DMLTransaction",

"CloudSpanner.Commit",

"Test Span",

]

assert [span.name for span in span_list] == expected_span_names

for span in span_list[2:-1]:

assert span.context.trace_id == span_list[-1].context.trace_id

assert span.parent.span_id == span_list[-1].context.span_id

def test_execute_partitioned_dml(sessions_database, database_dialect):

# [START spanner_test_dml_partioned_dml_update]

sd = _sample_data

param_types = spanner_v1.param_types

delete_statement = f"DELETE FROM {sd.TABLE} WHERE true"