8000 feat: Multiplexed sessions - Update `SessionOptions` to support `GOOG… · googleapis/python-spanner@57902b4 · GitHub
[go: up one dir, main page]
Skip to content

Commit 57902b4

Browse files
currantwcurrantw
committed
feat: Multiplexed sessions - Update SessionOptions to support GOOGLE_CLOUD_SPANNER_FORCE_DISABLE_MULTIPLEXED_SESSIONS and add unit tests.
Signed-off-by: Taylor Curran <taylor.curran@improving.com>
1 parent 234135d commit 57902b4

File tree

3 files changed

+232
-36
lines changed

3 files changed

+232
-36
lines changed

‎google/cloud/spanner_v1/session_options.py

Lines changed: 50 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
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
1617
from logging import Logger
@@ -26,7 +27,7 @@ class TransactionType(Enum):
2627

2728
class SessionOptions(object):
2829
"""Represents the session options for the Cloud Spanner Python client.
29-
We can use ::class::`SessionOptions` to determine whether multiplexed sessions
30+
We can use :class:`SessionOptions` to determine whether multiplexed sessions
3031
should be used for a specific transaction type with :meth:`use_multiplexed`. The use
3132
of multiplexed session can be disabled for a specific transaction type or for all
3233
transaction types with :meth:`disable_multiplexed`.
@@ -40,6 +41,9 @@ class SessionOptions(object):
4041
ENV_VAR_ENABLE_MULTIPLEXED_FOR_READ_WRITE = (
4142
"GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW"
4243
)
44+
ENV_VAR_FORCE_DISABLE_MULTIPLEXED = (
45+
"GOOGLE_CLOUD_SPANNER_FORCE_DISABLE_MULTIPLEXED_SESSIONS"
46+
)
4347

4448
def __init__(self):
4549
# Internal overrides to disable the use of multiplexed
@@ -52,33 +56,50 @@ def __init__(self):
5256

5357
def use_multiplexed(self, transaction_type: TransactionType) -> bool:
5458
"""Returns whether to use multiplexed sessions for the given transaction type.
59+
5560
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.
61+
* ENV_VAR_ENABLE_MULTIPLEXED is set to true;
62+
* ENV_VAR_FORCE_DISABLE_MULTIPLEXED is not set to true; and
63+
* multiplexed sessions have not been disabled for read-only transactions.
64+
5865
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.
66+
* ENV_VAR_ENABLE_MULTIPLEXED is set to true;
67+
* ENV_VAR_ENABLE_MULTIPLEXED_FOR_PARTITIONED is set to true;
68+
* ENV_VAR_FORCE_DISABLE_MULTIPLEXED is not set to true; and
69+
* multiplexed sessions have not been disabled for partitioned transactions.
70+
71+
Multiplexed sessions are enabled for read/write transactions if:
72+
* ENV_VAR_ENABLE_MULTIPLEXED is set to true;
73+
* ENV_VAR_ENABLE_MULTIPLEXED_FOR_READ_WRITE is set to true;
74+
* ENV_VAR_FORCE_DISABLE_MULTIPLEXED is not set to true; and
75+
* multiplexed sessions have not been disabled for read/write transactions.
76+
6377
:type transaction_type: :class:`TransactionType`
64-
:param transaction_type: the type of transaction to check whether
65-
multiplexed sessions should be used.
78+
:param transaction_type: the type of transaction
6679
"""
6780

6881
if transaction_type is TransactionType.READ_ONLY:
69-
return self._is_multiplexed_enabled[transaction_type] and self._getenv(
70-
self.ENV_VAR_ENABLE_MULTIPLEXED
82+
return (
83+
self._getenv(self.ENV_VAR_ENABLE_MULTIPLEXED)
84+
and not self._getenv(self.ENV_VAR_FORCE_DISABLE_MULTIPLEXED)
85+
and self._is_multiplexed_enabled[transaction_type]
7186
)
7287

7388
elif transaction_type is TransactionType.PARTITIONED:
7489
return (
75-
self._is_multiplexed_enabled[transaction_type]
76-
and self._getenv(self.ENV_VAR_ENABLE_MULTIPLEXED)
90+
self._getenv(self.ENV_VAR_ENABLE_MULTIPLEXED)
7791
and self._getenv(self.ENV_VAR_ENABLE_MULTIPLEXED_FOR_PARTITIONED)
92+
and not self._getenv(self.ENV_VAR_FORCE_DISABLE_MULTIPLEXED)
93+
and self._is_multiplexed_enabled[transaction_type]
7894
)
7995

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

83104
raise ValueError(f"Transaction type {transaction_type} is not supported.")
84105

@@ -88,40 +109,33 @@ def disable_multiplexed(
88109
"""Disables the use of multiplexed sessions for the given transaction type.
89110
If no transaction type is specified, disables the use of multiplexed sessions
90111
for all transaction types.
112+
91113
:type logger: :class:`Logger`
92-
:param logger: logger to use for logging the disabling the use of multiplexed
93-
sessions.
114+
:param logger: logger for logging disabling the use of multiplexed sessions.
115+
94116
:type transaction_type: :class:`TransactionType`
95117
:param transaction_type: (Optional) the type of transaction for which to disable
96118
the use of multiplexed sessions.
97119
"""
98120

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

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

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
126+
transaction_types_to_disable = (
127+
[transaction_type]
128+
if transaction_type is not None
129+
else list(TransactionType)
130+
)
114131

115-
elif transaction_type in self._is_multiplexed_enabled.keys():
132+
for transaction_type_to_disable in transaction_types_to_disable:
116133
logger.warning(
117-
disable_multiplexed_log_msg_fstring.format(
118-
transaction_type_value=transaction_type.value
119-
)
134+
f"Disabling multiplexed sessions for {transaction_type_to_disable.value} transactions"
120135
)
121-
self._is_multiplexed_enabled[transaction_type] = False
122-
return
136+
self._is_multiplexed_enabled[transaction_type_to_disable] = False
123137

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

126140
@staticmethod
127141
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)
0