Skip to content

Commit c8e7a1c

Browse files
currantwcurrantw
committed
Update SessionOptions to support GOOGLE_CLOUD_SPANNER_FORCE_DISABLE_MULTIPLEXED_SESSIONS and add unit tests.
Signed-off-by: Taylor Curran <[email protected]>
1 parent 234135d commit c8e7a1c

File tree

3 files changed

+233
-38
lines changed

3 files changed

+233
-38
lines changed

‎google/cloud/spanner_v1/session_options.py‎

Lines changed: 51 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,9 @@
1111
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
14+
import logging
1415
import os
1516
from enum import Enum
16-
from logging import Logger
1717

1818

1919
class TransactionType(Enum):
@@ -26,7 +26,7 @@ class TransactionType(Enum):
2626

2727
class SessionOptions(object):
2828
"""Represents the session options for the Cloud Spanner Python client.
29-
We can use ::class::`SessionOptions` to determine whether multiplexed sessions
29+
We can use :class:`SessionOptions` to determine whether multiplexed sessions
3030
should be used for a specific transaction type with :meth:`use_multiplexed`. The use
3131
of multiplexed session can be disabled for a specific transaction type or for all
3232
transaction types with :meth:`disable_multiplexed`.
@@ -40,6 +40,9 @@ class SessionOptions(object):
4040
ENV_VAR_ENABLE_MULTIPLEXED_FOR_READ_WRITE = (
4141
"GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW"
4242
)
43+
ENV_VAR_FORCE_DISABLE_MULTIPLEXED = (
44+
"GOOGLE_CLOUD_SPANNER_FORCE_DISABLE_MULTIPLEXED_SESSIONS"
45+
)
4346

4447
def __init__(self):
4548
# Internal overrides to disable the use of multiplexed
@@ -52,76 +55,86 @@ def __init__(self):
5255

5356
def use_multiplexed(self, transaction_type: TransactionType) -> bool:
5457
"""Returns whether to use multiplexed sessions for the given transaction type.
58+
5559
Multiplexed sessions are enabled for read-only transactions if:
56-
* ENV_VAR_ENABLE_MULTIPLEXED is set to true; and
57-
* multiplexed sessions have not been disabled for read-only transactions.
60+
* ENV_VAR_ENABLE_MULTIPLEXED is set to true;
61+
* ENV_VAR_FORCE_DISABLE_MULTIPLEXED is not set to true; and
62+
* multiplexed sessions have not been disabled for read-only transactions.
63+
5864
Multiplexed sessions are enabled for partitioned transactions if:
59-
* ENV_VAR_ENABLE_MULTIPLEXED is set to true;
60-
* ENV_VAR_ENABLE_MULTIPLEXED_FOR_PARTITIONED is set to true; and
61-
* multiplexed sessions have not been disabled for partitioned transactions.
62-
Multiplexed sessions are **currently disabled** for read / write.
65+
* ENV_VAR_ENABLE_MULTIPLEXED is set to true;
66+
* ENV_VAR_ENABLE_MULTIPLEXED_FOR_PARTITIONED is set to true;
67+
* ENV_VAR_FORCE_DISABLE_MULTIPLEXED is not set to true; and
68+
* multiplexed sessions have not been disabled for partitioned transactions.
69+
70+
Multiplexed sessions are enabled for read/write transactions if:
71+
* ENV_VAR_ENABLE_MULTIPLEXED is set to true;
72+
* ENV_VAR_ENABLE_MULTIPLEXED_FOR_READ_WRITE is set to true;
73+
* ENV_VAR_FORCE_DISABLE_MULTIPLEXED is not set to true; and
74+
* multiplexed sessions have not been disabled for read/write transactions.
75+
6376
:type transaction_type: :class:`TransactionType`
64-
:param transaction_type: the type of transaction to check whether
65-
multiplexed sessions should be used.
77+
:param transaction_type: the type of transaction
6678
"""
6779

6880
if transaction_type is TransactionType.READ_ONLY:
69-
return self._is_multiplexed_enabled[transaction_type] and self._getenv(
70-
self.ENV_VAR_ENABLE_MULTIPLEXED
81+
return (
82+
self._getenv(self.ENV_VAR_ENABLE_MULTIPLEXED)
83+
and not self._getenv(self.ENV_VAR_FORCE_DISABLE_MULTIPLEXED)
84+
and self._is_multiplexed_enabled[transaction_type]
7185
)
7286

7387
elif transaction_type is TransactionType.PARTITIONED:
7488
return (
75-
self._is_multiplexed_enabled[transaction_type]
76-
and self._getenv(self.ENV_VAR_ENABLE_MULTIPLEXED)
89+
self._getenv(self.ENV_VAR_ENABLE_MULTIPLEXED)
7790
and self._getenv(self.ENV_VAR_ENABLE_MULTIPLEXED_FOR_PARTITIONED)
91+
and not self._getenv(self.ENV_VAR_FORCE_DISABLE_MULTIPLEXED)
92+
and self._is_multiplexed_enabled[transaction_type]
7893
)
7994

8095
elif transaction_type is TransactionType.READ_WRITE:
81-
return False
96+
return (
97+
self._getenv(self.ENV_VAR_ENABLE_MULTIPLEXED)
98+
and self._getenv(self.ENV_VAR_ENABLE_MULTIPLEXED_FOR_READ_WRITE)
99+
and not self._getenv(self.ENV_VAR_FORCE_DISABLE_MULTIPLEXED)
100+
and self._is_multiplexed_enabled[transaction_type]
101+
)
82102

83103
raise ValueError(f"Transaction type {transaction_type} is not supported.")
84104

85105
def disable_multiplexed(
86-
self, logger: Logger = None, transaction_type: TransactionType = None
106+
self, logger: logging.Logger = None, transaction_type: TransactionType = None
87107
) -> None:
88108
"""Disables the use of multiplexed sessions for the given transaction type.
89109
If no transaction type is specified, disables the use of multiplexed sessions
90110
for all transaction types.
111+
91112
:type logger: :class:`Logger`
92-
:param logger: logger to use for logging the disabling the use of multiplexed
93-
sessions.
113+
:param logger: logger for logging disabling the use of multiplexed sessions.
114+
94115
:type transaction_type: :class:`TransactionType`
95116
:param transaction_type: (Optional) the type of transaction for which to disable
96117
the use of multiplexed sessions.
97118
"""
98119

99-
disable_multiplexed_log_msg_fstring = (
100-
"Disabling multiplexed sessions for {transaction_type_value} transactions"
101-
)
102-
import logging
120+
if transaction_type and transaction_type not in self._is_multiplexed_enabled:
121+
raise ValueError(f"Transaction type '{transaction_type}' is not supported.")
103122

104-
if logger is None:
105-
logger = logging.getLogger(__name__)
123+
logger = logger or logging.getLogger(__name__)
106124

107-
if transaction_type is None:
108-
logger.warning(
109-
disable_multiplexed_log_msg_fstring.format(transaction_type_value="all")
110-
)
111-
for transaction_type in TransactionType:
112-
self._is_multiplexed_enabled[transaction_type] = False
113-
return
125+
transaction_types_to_disable = (
126+
[transaction_type]
127+
if transaction_type is not None
128+
else list(TransactionType)
129+
)
114130

115-
elif transaction_type in self._is_multiplexed_enabled.keys():
131+
for transaction_type_to_disable in transaction_types_to_disable:
116132
logger.warning(
117-
disable_multiplexed_log_msg_fstring.format(
118-
transaction_type_value=transaction_type.value
119-
)
133+
f"Disabling multiplexed sessions for {transaction_type_to_disable.value} transactions"
120134
)
121-
self._is_multiplexed_enabled[transaction_type] = False
122-
return
135+
self._is_multiplexed_enabled[transaction_type_to_disable] = False
123136

124-
raise ValueError(f"Transaction type '{transaction_type}' is not supported.")
137+
return
125138

126139
@staticmethod
127140
def _getenv(name: str) -> bool:

‎tests/_builders.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
# Copyright 2025 Google LLC All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from mock import create_autospec
16+
17+
18+
def build_logger():
19+
"""Builds and returns a logger for testing."""
20+
from logging import Logger
21+
22+
return create_autospec(Logger, instance=True)

‎tests/unit/test_session_options.py‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
# Copyright 2025 Google LLC All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
from logging import Logger
15+
from os import environ
16+
from unittest import TestCase
17+
18+
from google.cloud.spanner_v1.session_options import SessionOptions, TransactionType
19+
from tests._builders import build_logger
20+
21+
22+
class TestSessionOptions(TestCase):
23+
@classmethod
24+
def setUpClass(cls):
25+
# Save the original environment variables.
26+
cls._original_env = dict(environ)
27+
28+
@classmethod
29+
def tearDownClass(cls):
30+
# Restore environment variables.
31+
environ.clear()
32+
environ.update(cls._original_env)
33+
34+
def setUp(self):
35+
self.logger: Logger = build_logger()
36+
37+
def test_use_multiplexed_for_read_only(self):
38+
session_options = SessionOptions()
39+
transaction_type = TransactionType.READ_ONLY
40+
41+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED] = "false"
42+
self.assertFalse(session_options.use_multiplexed(transaction_type))
43+
44+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED] = "true"
45+
environ[SessionOptions.ENV_VAR_FORCE_DISABLE_MULTIPLEXED] = "true"
46+
self.assertFalse(session_options.use_multiplexed(transaction_type))
47+
48+
environ[SessionOptions.ENV_VAR_FORCE_DISABLE_MULTIPLEXED] = "false"
49+
self.assertTrue(session_options.use_multiplexed(transaction_type))
50+
51+
session_options.disable_multiplexed(self.logger, transaction_type)
52+
self.assertFalse(session_options.use_multiplexed(transaction_type))
53+
54+
self.logger.warning.assert_called_once_with(
55+
"Disabling multiplexed sessions for read-only transactions"
56+
)
57+
58+
def test_use_multiplexed_for_partitioned(self):
59+
session_options = SessionOptions()
60+
transaction_type = TransactionType.PARTITIONED
61+
62+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED] = "false"
63+
self.assertFalse(session_options.use_multiplexed(transaction_type))
64+
65+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED] = "true"
66+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED_FOR_PARTITIONED] = "false"
67+
self.assertFalse(session_options.use_multiplexed(transaction_type))
68+
69+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED_FOR_PARTITIONED] = "true"
70+
environ[SessionOptions.ENV_VAR_FORCE_DISABLE_MULTIPLEXED] = "true"
71+
self.assertFalse(session_options.use_multiplexed(transaction_type))
72+
73+
environ[SessionOptions.ENV_VAR_FORCE_DISABLE_MULTIPLEXED] = "false"
74+
self.assertTrue(session_options.use_multiplexed(transaction_type))
75+
76+
session_options.disable_multiplexed(self.logger, transaction_type)
77+
self.assertFalse(session_options.use_multiplexed(transaction_type))
78+
79+
self.logger.warning.assert_called_once_with(
80+
"Disabling multiplexed sessions for partitioned transactions"
81+
)
82+
83+
def test_use_multiplexed_for_read_write(self):
84+
session_options = SessionOptions()
85+
transaction_type = TransactionType.READ_WRITE
86+
87+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED] = "false"
88+
self.assertFalse(session_options.use_multiplexed(transaction_type))
89+
90+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED] = "true"
91+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED_FOR_READ_WRITE] = "false"
92+
self.assertFalse(session_options.use_multiplexed(transaction_type))
93+
94+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED_FOR_READ_WRITE] = "true"
95+
environ[SessionOptions.ENV_VAR_FORCE_DISABLE_MULTIPLEXED] = "true"
96+
self.assertFalse(session_options.use_multiplexed(transaction_type))
97+
98+
environ[SessionOptions.ENV_VAR_FORCE_DISABLE_MULTIPLEXED] = "false"
99+
self.assertTrue(session_options.use_multiplexed(transaction_type))
100+
101+
session_options.disable_multiplexed(self.logger, transaction_type)
102+
self.assertFalse(session_options.use_multiplexed(transaction_type))
103+
104+
self.logger.warning.assert_called_once_with(
105+
"Disabling multiplexed sessions for read/write transactions"
106+
)
107+
108+
def test_disable_multiplexed_all(self):
109+
session_options = SessionOptions()
110+
111+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED] = "true"
112+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED_FOR_PARTITIONED] = "true"
113+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED_FOR_READ_WRITE] = "true"
114+
environ[SessionOptions.ENV_VAR_FORCE_DISABLE_MULTIPLEXED] = "false"
115+
116+
session_options.disable_multiplexed(self.logger)
117+
118+
self.assertFalse(session_options.use_multiplexed(TransactionType.READ_ONLY))
119+
self.assertFalse(session_options.use_multiplexed(TransactionType.PARTITIONED))
120+
self.assertFalse(session_options.use_multiplexed(TransactionType.READ_WRITE))
121+
122+
warning = self.logger.warning
123+
self.assertEqual(warning.call_count, 3)
124+
warning.assert_any_call(
125+
"Disabling multiplexed sessions for read-only transactions"
126+
)
127+
warning.assert_any_call(
128+
"Disabling multiplexed sessions for partitioned transactions"
129+
)
130+
warning.assert_any_call(
131+
"Disabling multiplexed sessions for read/write transactions"
132+
)
133+
134+
def test_unsupported_transaction_type(self):
135+
session_options = SessionOptions()
136+
unsupported_type = "UNSUPPORTED_TRANSACTION_TYPE"
137+
138+
with self.assertRaises(ValueError):
139+
session_options.use_multiplexed(unsupported_type)
140+
141+
with self.assertRaises(ValueError):
142+
session_options.disable_multiplexed(self.logger, unsupported_type)
143+
144+
def test_env_var_values(self):
145+
session_options = SessionOptions()
146+
147+
environ[SessionOptions.ENV_VAR_FORCE_DISABLE_MULTIPLEXED] = "false"
148+
149+
true_values = ["1", " 1", " 1", "true", "True", "TRUE", " true "]
150+
for value in true_values:
151+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED] = value
152+
self.assertTrue(session_options.use_multiplexed(TransactionType.READ_ONLY))
153+
154+
false_values = ["", "0", "false", "False", "FALSE", " false "]
155+
for value in false_values:
156+
environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED] = value
157+
self.assertFalse(session_options.use_multiplexed(TransactionType.READ_ONLY))
158+
159+
del environ[SessionOptions.ENV_VAR_ENABLE_MULTIPLEXED]
160+
self.assertFalse(session_options.use_multiplexed(TransactionType.READ_ONLY))

0 commit comments

Comments
 (0)