1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291
|
import random
import threading
import time
import sqlalchemy as sa
from sqlalchemy import Integer
from sqlalchemy import MetaData
from sqlalchemy import String
from sqlalchemy import testing
from sqlalchemy.orm import scoped_session
from sqlalchemy.orm import sessionmaker
from sqlalchemy.testing import eq_
from sqlalchemy.testing import fixtures
from sqlalchemy.testing.schema import Column
from sqlalchemy.testing.schema import Table
NUM_THREADS = 10
ITERATIONS = 10
class _ThreadTest:
def run_threaded(
self,
func,
*thread_args,
nthreads=NUM_THREADS,
use_barrier=False,
**thread_kwargs,
):
barrier = threading.Barrier(nthreads)
results = []
errors = []
def thread_func(*args, **kwargs):
thread_name = threading.current_thread().name
if use_barrier:
barrier.wait()
local_result = []
try:
func(local_result, thread_name, *args, **kwargs)
results.append(tuple(local_result))
except Exception as e:
# raise
errors.append((thread_name, repr(e)))
threads = [
threading.Thread(
name=f"thread-{i}",
target=thread_func,
args=thread_args,
kwargs=thread_kwargs,
)
for i in range(nthreads)
]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
return results, errors
@testing.fixture
def num_threads_engine(self, testing_engine):
return testing_engine(options=dict(pool_size=NUM_THREADS))
@testing.add_to_marker.timing_intensive
class EngineThreadSafetyTest(_ThreadTest, fixtures.TablesTest):
run_dispose_bind = "once"
__requires__ = ("multithreading_support",)
@classmethod
def define_tables(cls, metadata):
Table(
"test_table",
metadata,
Column(
"id", Integer, primary_key=True, test_needs_autoincrement=True
),
Column("thread_id", Integer),
Column("data", String(50)),
)
@testing.combinations(
(NUM_THREADS, 0),
(3, 5),
(3, 0),
(7, 0),
argnames="pool_size, max_overflow",
)
def test_engine_thread_safe(self, testing_engine, pool_size, max_overflow):
"""Test that a single Engine can be safely shared across threads."""
test_table = self.tables.test_table
engine = testing_engine(
options=dict(pool_size=pool_size, max_overflow=max_overflow)
)
def worker(results, thread_name):
for _ in range(ITERATIONS):
with engine.connect() as conn:
conn.execute(
test_table.insert(),
{"data": thread_name},
)
conn.commit()
result = conn.execute(
sa.select(test_table.c.data).where(
test_table.c.data == thread_name
)
).scalar()
results.append(result)
results, errors = self.run_threaded(worker)
eq_(errors, [])
eq_(
set(results),
{
tuple([f"thread-{i}" for j in range(ITERATIONS)])
for i in range(NUM_THREADS)
},
)
def test_metadata_thread_safe(self, num_threads_engine):
"""Test that MetaData objects are thread-safe for reads."""
metadata = sa.MetaData()
for thread_id in range(NUM_THREADS):
Table(
f"thread-{thread_id}",
metadata,
Column("id", Integer, primary_key=True),
Column("data", String(50)),
)
metadata.create_all(testing.db)
def worker(results, thread_name):
table_key = thread_name
assert table_key in metadata.tables, f"{table_key} does not exist"
with num_threads_engine.connect() as conn:
# Will raise if it cannot connect so erros will be populated
conn.execute(sa.select(metadata.tables[table_key]))
_, errors = self.run_threaded(worker)
eq_(errors, [])
@testing.add_to_marker.timing_intensive
class SessionThreadingTest(_ThreadTest, fixtures.MappedTest):
run_dispose_bind = "once"
__requires__ = ("multithreading_support",)
@classmethod
def define_tables(cls, metadata):
Table(
"users",
metadata,
Column(
"id", Integer, primary_key=True, test_needs_autoincrement=True
),
Column("name", String(50)),
Column("thread_id", String(50)),
)
@classmethod
def setup_classes(cls):
class User(cls.Comparable):
pass
def test_sessionmaker_thread_safe(self, num_threads_engine):
"""Test that sessionmaker factory is thread-safe."""
users, User = self.tables.users, self.classes.User
self.mapper_registry.map_imperatively(User, users)
# Single sessionmaker shared across threads
SessionFactory = sessionmaker(num_threads_engine)
def worker(results, thread_name):
thread_id = thread_name
for _ in range(ITERATIONS):
with SessionFactory() as session:
for i in range(3):
user = User(
name=f"user_{thread_id}_{i}", thread_id=thread_id
)
session.add(user)
session.commit()
count = (
session.query(User)
.filter_by(thread_id=thread_id)
.count()
)
results.append(count)
results, errors = self.run_threaded(worker)
eq_(errors, [])
eq_(
results,
[
tuple(range(3, 3 * ITERATIONS + 3, 3))
for _ in range(NUM_THREADS)
],
)
def test_scoped_session_thread_local(self, num_threads_engine):
"""Test that scoped_session provides thread-local sessions."""
users, User = self.tables.users, self.classes.User
self.mapper_registry.map_imperatively(User, users)
# Create scoped session
Session = scoped_session(sessionmaker(num_threads_engine))
session_ids = {}
def worker(results, thread_name):
thread_id = thread_name
session = Session()
session_ids[thread_id] = id(session)
session.close()
for _ in range(ITERATIONS):
user = User(
name=f"scoped_user_{thread_id}", thread_id=thread_id
)
Session.add(user)
Session.commit()
session2 = Session()
assert id(session2) == session_ids[thread_id]
session2.close()
count = (
Session.query(User).filter_by(thread_id=thread_id).count()
)
results.append(count)
Session.remove()
results, errors = self.run_threaded(worker)
eq_(errors, [])
unique_sessions = set(session_ids.values())
eq_(len(unique_sessions), NUM_THREADS)
eq_(
results,
[tuple(range(1, ITERATIONS + 1)) for _ in range(NUM_THREADS)],
)
@testing.add_to_marker.timing_intensive
class FromClauseConcurrencyTest(_ThreadTest, fixtures.TestBase):
"""test for issue #12302"""
@testing.variation("collection", ["c", "primary_key", "foreign_keys"])
def test_c_collection(self, collection):
dictionary_meta = MetaData()
all_indexes_table = Table(
"all_indexes",
dictionary_meta,
*[Column(f"col{i}", Integer) for i in range(50)],
)
def use_table(results, errors):
for i in range(3):
time.sleep(random.random() * 0.0001)
if collection.c:
all_indexes.c.col35
elif collection.primary_key:
all_indexes.primary_key
elif collection.foreign_keys:
all_indexes.foreign_keys
for j in range(1000):
all_indexes = all_indexes_table.alias("a_indexes")
results, errors = self.run_threaded(
use_table, use_barrier=False, nthreads=5
)
eq_(errors, [])
eq_(len(results), 5)
|