codekingpro/portable-devtools

1# dialects/sqlite/provision.py
2# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of SQLAlchemy and is released under
6# the MIT License: https://www.opensource.org/licenses/mit-license.php
7# mypy: ignore-errors
8
9import os
10import re
11
12from ... import event
13from ... import exc
14from ...engine import url as sa_url
15from ...testing import config
16from ...testing.provision import create_db
17from ...testing.provision import drop_db
18from ...testing.provision import follower_url_from_main
19from ...testing.provision import generate_driver_url
20from ...testing.provision import log
21from ...testing.provision import post_configure_engine
22from ...testing.provision import post_configure_testing_engine
23from ...testing.provision import run_reap_dbs
24from ...testing.provision import stop_test_class_outside_fixtures
25from ...testing.provision import temp_table_keyword_args
26from ...testing.provision import upsert
27
28# TODO: I can't get this to build dynamically with pytest-xdist procs
29_drivernames = {
30    "pysqlite",
31    "aiosqlite",
32    "pysqlcipher",
33    "pysqlite_numeric",
34    "pysqlite_dollar",
35}
36
37
38def _format_url(url, driver, ident):
39    """given a sqlite url + desired driver + ident, make a canonical
40    URL out of it
41
42    """
43    url = sa_url.make_url(url)
44
45    if driver is None:
46        driver = url.get_driver_name()
47
48    filename = url.database
49
50    needs_enc = driver == "pysqlcipher"
51    name_token = None
52
53    if filename and filename != ":memory:":
54        assert "test_schema" not in filename
55        tokens = re.split(r"[_\.]", filename)
56
57        for token in tokens:
58            if token in _drivernames:
59                if driver is None:
60                    driver = token
61                continue
62            elif token in ("db", "enc"):
63                continue
64            elif name_token is None:
65                name_token = token.strip("_")
66
67        assert name_token, f"sqlite filename has no name token: {url.database}"
68
69        new_filename = f"{name_token}_{driver}"
70        if ident:
71            new_filename += f"_{ident}"
72        new_filename += ".db"
73        if needs_enc:
74            new_filename += ".enc"
75        url = url.set(database=new_filename)
76
77    if needs_enc:
78        url = url.set(password="test")
79
80    url = url.set(drivername="sqlite+%s" % (driver,))
81
82    return url
83
84
85@generate_driver_url.for_db("sqlite")
86def generate_driver_url(url, driver, query_str):
87    url = _format_url(url, driver, None)
88
89    try:
90        url.get_dialect()
91    except exc.NoSuchModuleError:
92        return None
93    else:
94        return url
95
96
97@follower_url_from_main.for_db("sqlite")
98def _sqlite_follower_url_from_main(url, ident):
99    return _format_url(url, None, ident)
100
101
102@post_configure_engine.for_db("sqlite")
103def _sqlite_post_configure_engine(url, engine, follower_ident):
104    from sqlalchemy import event
105
106    if follower_ident:
107        attach_path = f"{follower_ident}_{engine.driver}_test_schema.db"
108    else:
109        attach_path = f"{engine.driver}_test_schema.db"
110
111    @event.listens_for(engine, "connect")
112    def connect(dbapi_connection, connection_record):
113        # use file DBs in all cases, memory acts kind of strangely
114        # as an attached
115
116        # NOTE!  this has to be done *per connection*.  New sqlite connection,
117        # as we get with say, QueuePool, the attaches are gone.
118        # so schemes to delete those attached files have to be done at the
119        # filesystem level and not rely upon what attachments are in a
120        # particular SQLite connection
121        dbapi_connection.execute(
122            f'ATTACH DATABASE "{attach_path}" AS test_schema'
123        )
124
125    @event.listens_for(engine, "engine_disposed")
126    def dispose(engine):
127        """most databases should be dropped using
128        stop_test_class_outside_fixtures
129
130        however a few tests like AttachedDBTest might not get triggered on
131        that main hook
132
133        """
134
135        if os.path.exists(attach_path):
136            os.remove(attach_path)
137
138        filename = engine.url.database
139
140        if filename and filename != ":memory:" and os.path.exists(filename):
141            os.remove(filename)
142
143
144@post_configure_testing_engine.for_db("sqlite")
145def _sqlite_post_configure_testing_engine(url, engine, options, scope):
146
147    sqlite_savepoint = options.get("sqlite_savepoint", False)
148    sqlite_share_pool = options.get("sqlite_share_pool", False)
149
150    if sqlite_savepoint and engine.name == "sqlite":
151        # apply SQLite savepoint workaround
152        @event.listens_for(engine, "connect")
153        def do_connect(dbapi_connection, connection_record):
154            dbapi_connection.isolation_level = None
155
156        @event.listens_for(engine, "begin")
157        def do_begin(conn):
158            conn.exec_driver_sql("BEGIN")
159
160    if sqlite_share_pool:
161        # SingletonThreadPool, StaticPool both support "transfer"
162        # so a new pool can share the same SQLite connection
163        # (single thread only)
164        if hasattr(engine.pool, "_transfer_from"):
165            options["use_reaper"] = False
166            engine.pool._transfer_from(config.db.pool)
167
168
169@create_db.for_db("sqlite")
170def _sqlite_create_db(cfg, eng, ident):
171    pass
172
173
174@drop_db.for_db("sqlite")
175def _sqlite_drop_db(cfg, eng, ident):
176    _drop_dbs_w_ident(eng.url.database, eng.driver, ident)
177
178
179def _drop_dbs_w_ident(databasename, driver, ident):
180    for path in os.listdir("."):
181        fname, ext = os.path.split(path)
182        if ident in fname and ext in [".db", ".db.enc"]:
183            log.info("deleting SQLite database file: %s", path)
184            os.remove(path)
185
186
187@stop_test_class_outside_fixtures.for_db("sqlite")
188def stop_test_class_outside_fixtures(config, db, cls):
189    db.dispose()
190
191
192@temp_table_keyword_args.for_db("sqlite")
193def _sqlite_temp_table_keyword_args(cfg, eng):
194    return {"prefixes": ["TEMPORARY"]}
195
196
197@run_reap_dbs.for_db("sqlite")
198def _reap_sqlite_dbs(url, idents):
199    log.info("db reaper connecting to %r", url)
200    log.info("identifiers in file: %s", ", ".join(idents))
201    url = sa_url.make_url(url)
202    for ident in idents:
203        for drivername in _drivernames:
204            _drop_dbs_w_ident(url.database, drivername, ident)
205
206
207@upsert.for_db("sqlite")
208def _upsert(
209    cfg,
210    table,
211    returning,
212    *,
213    set_lambda=None,
214    sort_by_parameter_order=False,
215    index_elements=None,
216):
217    from sqlalchemy.dialects.sqlite import insert
218
219    stmt = insert(table)
220
221    if set_lambda:
222        stmt = stmt.on_conflict_do_update(set_=set_lambda(stmt.excluded))
223    else:
224        stmt = stmt.on_conflict_do_nothing()
225
226    stmt = stmt.returning(
227        *returning, sort_by_parameter_order=sort_by_parameter_order
228    )
229    return stmt
230