sqlalchemy/test/aaa_profiling/test_threading.py

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)