Skip to content

Commit 237ae41

Browse files
Ilya Gurovlarkee
Ilya Gurov
and
larkee
authored
fix(db_api): move connection validation into a separate method (#543)
Co-authored-by: larkee <31196561+larkee@users.noreply.github.com>
1 parent 23b1600 commit 237ae41

File tree

5 files changed

+116
-79
lines changed

5 files changed

+116
-79
lines changed

‎google/cloud/spanner_dbapi/connection.py

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
from google.cloud.spanner_dbapi.checksum import _compare_checksums
2929
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
3030
from google.cloud.spanner_dbapi.cursor import Cursor
31-
from google.cloud.spanner_dbapi.exceptions import InterfaceError
31+
from google.cloud.spanner_dbapi.exceptions import InterfaceError, OperationalError
3232
from google.cloud.spanner_dbapi.version import DEFAULT_USER_AGENT
3333
from google.cloud.spanner_dbapi.version import PY_VERSION
3434

@@ -349,6 +349,30 @@ def run_statement(self, statement, retried=False):
349349
ResultsChecksum() if retried else statement.checksum,
350350
)
351351

352+
def validate(self):
353+
"""
354+
Execute a minimal request to check if the connection
355+
is valid and the related database is reachable.
356+
357+
Raise an exception in case if the connection is closed,
358+
invalid, target database is not found, or the request result
359+
is incorrect.
360+
361+
:raises: :class:`InterfaceError`: if this connection is closed.
362+
:raises: :class:`OperationalError`: if the request result is incorrect.
363+
:raises: :class:`google.cloud.exceptions.NotFound`: if the linked instance
364+
or database doesn't exist.
365+
"""
366+
self._raise_if_closed()
367+
368+
with self.database.snapshot() as snapshot:
369+
result = list(snapshot.execute_sql("SELECT 1"))
370+
if result != [[1]]:
371+
raise OperationalError(
372+
"The checking query (SELECT 1) returned an unexpected result: %s. "
373+
"Expected: [[1]]" % result
374+
)
375+
352376
def __enter__(self):
353377
return self
354378

@@ -399,9 +423,6 @@ def connect(
399423
:rtype: :class:`google.cloud.spanner_dbapi.connection.Connection`
400424
:returns: Connection object associated with the given Google Cloud Spanner
401425
resource.
402-
403-
:raises: :class:`ValueError` in case of given instance/database
404-
doesn't exist.
405426
"""
406427

407428
client_info = ClientInfo(
@@ -418,14 +439,7 @@ def connect(
418439
)
419440

420441
instance = client.instance(instance_id)
421-
if not instance.exists():
422-
raise ValueError("instance '%s' does not exist." % instance_id)
423-
424-
database = instance.database(database_id, pool=pool)
425-
if not database.exists():
426-
raise ValueError("database '%s' does not exist." % database_id)
427-
428-
conn = Connection(instance, database)
442+
conn = Connection(instance, instance.database(database_id, pool=pool))
429443
if pool is not None:
430444
conn._own_pool = False
431445

‎tests/system/test_dbapi.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -350,3 +350,10 @@ def test_DDL_commit(shared_instance, dbapi_database):
350350

351351
cur.execute("DROP TABLE Singers")
352352
conn.commit()
353+
354+
355+
def test_ping(shared_instance, dbapi_database):
356+
"""Check connection validation method."""
357+
conn = Connection(shared_instance, dbapi_database)
358+
conn.validate()
359+
conn.close()

‎tests/unit/spanner_dbapi/test_connect.py

Lines changed: 0 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -88,31 +88,6 @@ def test_w_explicit(self, mock_client):
8888
self.assertIs(connection.database, database)
8989
instance.database.assert_called_once_with(DATABASE, pool=pool)
9090

91-
def test_w_instance_not_found(self, mock_client):
92-
from google.cloud.spanner_dbapi import connect
93-
94-
client = mock_client.return_value
95-
instance = client.instance.return_value
96-
instance.exists.return_value = False
97-
98-
with self.assertRaises(ValueError):
99-
connect(INSTANCE, DATABASE)
100-
101-
instance.exists.assert_called_once_with()
102-
103-
def test_w_database_not_found(self, mock_client):
104-
from google.cloud.spanner_dbapi import connect
105-
106-
client = mock_client.return_value
107-
instance = client.instance.return_value
108-
database = instance.database.return_value
109-
database.exists.return_value = False
110-
111-
with self.assertRaises(ValueError):
112-
connect(INSTANCE, DATABASE)
113-
114-
database.exists.assert_called_once_with()
115-
11691
def test_w_credential_file_path(self, mock_client):
11792
from google.cloud.spanner_dbapi import connect
11893
from google.cloud.spanner_dbapi import Connection

‎tests/unit/spanner_dbapi/test_connection.py

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -624,3 +624,80 @@ def test_retry_transaction_w_empty_response(self):
624624
compare_mock.assert_called_with(checksum, retried_checkum)
625625

626626
run_mock.assert_called_with(statement, retried=True)
627+
628+
def test_validate_ok(self):
629+
def exit_func(self, exc_type, exc_value, traceback):
630+
pass
631+
632+
connection = self._make_connection()
633+
634+
# mock snapshot context manager
635+
snapshot_obj = mock.Mock()
636+
snapshot_obj.execute_sql = mock.Mock(return_value=[[1]])
637+
638+
snapshot_ctx = mock.Mock()
639+
snapshot_ctx.__enter__ = mock.Mock(return_value=snapshot_obj)
640+
snapshot_ctx.__exit__ = exit_func
641+
snapshot_method = mock.Mock(return_value=snapshot_ctx)
642+
643+
connection.database.snapshot = snapshot_method
644+
645+
connection.validate()
646+
snapshot_obj.execute_sql.assert_called_once_with("SELECT 1")
647+
648+
def test_validate_fail(self):
649+
from google.cloud.spanner_dbapi.exceptions import OperationalError
650+
651+
def exit_func(self, exc_type, exc_value, traceback):
652+
pass
653+
654+
connection = self._make_connection()
655+
656+
# mock snapshot context manager
657+
snapshot_obj = mock.Mock()
658+
snapshot_obj.execute_sql = mock.Mock(return_value=[[3]])
659+
660+
snapshot_ctx = mock.Mock()
661+
snapshot_ctx.__enter__ = mock.Mock(return_value=snapshot_obj)
662+
snapshot_ctx.__exit__ = exit_func
663+
snapshot_method = mock.Mock(return_value=snapshot_ctx)
664+
665+
connection.database.snapshot = snapshot_method
666+
667+
with self.assertRaises(OperationalError):
668+
connection.validate()
669+
670+
snapshot_obj.execute_sql.assert_called_once_with("SELECT 1")
671+
672+
def test_validate_error(self):
673+
from google.cloud.exceptions import NotFound
674+
675+
def exit_func(self, exc_type, exc_value, traceback):
676+
pass
677+
678+
connection = self._make_connection()
679+
680+
# mock snapshot context manager
681+
snapshot_obj = mock.Mock()
682+
snapshot_obj.execute_sql = mock.Mock(side_effect=NotFound("Not found"))
683+
684+
snapshot_ctx = mock.Mock()
685+
snapshot_ctx.__enter__ = mock.Mock(return_value=snapshot_obj)
686+
snapshot_ctx.__exit__ = exit_func
687+
snapshot_method = mock.Mock(return_value=snapshot_ctx)
688+
689+
connection.database.snapshot = snapshot_method
690+
691+
with self.assertRaises(NotFound):
692+
connection.validate()
693+
694+
snapshot_obj.execute_sql.assert_called_once_with("SELECT 1")
695+
696+
def test_validate_closed(self):
697+
from google.cloud.spanner_dbapi.exceptions import InterfaceError
698+
699+
connection = self._make_connection()
700+
connection.close()
701+
702+
with self.assertRaises(InterfaceError):
703+
connection.validate()

‎tests/unit/spanner_dbapi/test_cursor.py

Lines changed: 6 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -332,13 +332,7 @@ def test_executemany_delete_batch_autocommit(self):
332332

333333
sql = "DELETE FROM table WHERE col1 = %s"
334334

335-
with mock.patch(
336-
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True
337-
):
338-
with mock.patch(
339-
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
340-
):
341-
connection = connect("test-instance", "test-database")
335+
connection = connect("test-instance", "test-database")
342336

343337
connection.autocommit = True
344338
transaction = self._transaction_mock()
@@ -369,13 +363,7 @@ def test_executemany_update_batch_autocommit(self):
369363

370364
sql = "UPDATE table SET col1 = %s WHERE col2 = %s"
371365

372-
with mock.patch(
373-
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True
374-
):
375-
with mock.patch(
376-
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
377-
):
378-
connection = connect("test-instance", "test-database")
366+
connection = connect("test-instance", "test-database")
379367

380368
connection.autocommit = True
381369
transaction = self._transaction_mock()
@@ -418,13 +406,7 @@ def test_executemany_insert_batch_non_autocommit(self):
418406

419407
sql = """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (%s, %s, %s, %s)"""
420408

421-
with mock.patch(
422-
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True
423-
):
424-
with mock.patch(
425-
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
426-
):
427-
connection = connect("test-instance", "test-database")
409+
connection = connect("test-instance", "test-database")
428410

429411
transaction = self._transaction_mock()
430412

@@ -461,13 +443,7 @@ def test_executemany_insert_batch_autocommit(self):
461443

462444
sql = """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (%s, %s, %s, %s)"""
463445

464-
with mock.patch(
465-
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True
466-
):
467-
with mock.patch(
468-
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
469-
):
470-
connection = connect("test-instance", "test-database")
446+
connection = connect("test-instance", "test-database")
471447

472448
connection.autocommit = True
473449

@@ -510,13 +486,7 @@ def test_executemany_insert_batch_failed(self):
510486
sql = """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (%s, %s, %s, %s)"""
511487
err_details = "Details here"
512488

513-
with mock.patch(
514-
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True
515-
):
516-
with mock.patch(
517-
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
518-
):
519-
connection = connect("test-instance", "test-database")
489+
connection = connect("test-instance", "test-database")
520490

521491
connection.autocommit = True
522492
cursor = connection.cursor()
@@ -546,13 +516,7 @@ def test_executemany_insert_batch_aborted(self):
546516
sql = """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (%s, %s, %s, %s)"""
547517
err_details = "Aborted details here"
548518

549-
with mock.patch(
550-
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True
551-
):
552-
with mock.patch(
553-
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
554-
):
555-
connection = connect("test-instance", "test-database")
519+
connection = connect("test-instance", "test-database")
556520

557521
transaction1 = mock.Mock(committed=False, rolled_back=False)
558522
transaction1.batch_update = mock.Mock(

0 commit comments

Comments
 (0)