codekingpro/portable-devtools

1# dialects/sqlite/provision.py
2# Copyright (C) 2005-2024 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 exc
13from ...engine import url as sa_url
14from ...testing.provision import create_db
15from ...testing.provision import drop_db
16from ...testing.provision import follower_url_from_main
17from ...testing.provision import generate_driver_url
18from ...testing.provision import log
19from ...testing.provision import post_configure_engine
20from ...testing.provision import run_reap_dbs
21from ...testing.provision import stop_test_class_outside_fixtures
22from ...testing.provision import temp_table_keyword_args
23from ...testing.provision import upsert
24
25
26# TODO: I can't get this to build dynamically with pytest-xdist procs
27_drivernames = {
28    "pysqlite",
29    "aiosqlite",
30    "pysqlcipher",
31    "pysqlite_numeric",
32    "pysqlite_dollar",
33}
34
35
36def _format_url(url, driver, ident):
37    """given a sqlite url + desired driver + ident, make a canonical
38    URL out of it
39
40    """
41    url = sa_url.make_url(url)
42
43    if driver is None:
44        driver = url.get_driver_name()
45
46    filename = url.database
47
48    needs_enc = driver == "pysqlcipher"
49    name_token = None
50
51    if filename and filename != ":memory:":
52        assert "test_schema" not in filename
53        tokens = re.split(r"[_\.]", filename)
54
55        new_filename = f"{driver}"
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@create_db.for_db("sqlite")
145def _sqlite_create_db(cfg, eng, ident):
146    pass
147
148
149@drop_db.for_db("sqlite")
150def _sqlite_drop_db(cfg, eng, ident):
151    _drop_dbs_w_ident(eng.url.database, eng.driver, ident)
152
153
154def _drop_dbs_w_ident(databasename, driver, ident):
155    for path in os.listdir("."):
156        fname, ext = os.path.split(path)
157        if ident in fname and ext in [".db", ".db.enc"]:
158            log.info("deleting SQLite database file: %s", path)
159            os.remove(path)
160
161
162@stop_test_class_outside_fixtures.for_db("sqlite")
163def stop_test_class_outside_fixtures(config, db, cls):
164    db.dispose()
165
166
167@temp_table_keyword_args.for_db("sqlite")
168def _sqlite_temp_table_keyword_args(cfg, eng):
169    return {"prefixes": ["TEMPORARY"]}
170
171
172@run_reap_dbs.for_db("sqlite")
173def _reap_sqlite_dbs(url, idents):
174    log.info("db reaper connecting to %r", url)
175    log.info("identifiers in file: %s", ", ".join(idents))
176    url = sa_url.make_url(url)
177    for ident in idents:
178        for drivername in _drivernames:
179            _drop_dbs_w_ident(url.database, drivername, ident)
180
181
182@upsert.for_db("sqlite")
183def _upsert(
184    cfg, table, returning, *, set_lambda=None, sort_by_parameter_order=False
185):
186    from sqlalchemy.dialects.sqlite import insert
187
188    stmt = insert(table)
189
190    if set_lambda:
191        stmt = stmt.on_conflict_do_update(set_=set_lambda(stmt.excluded))
192    else:
193        stmt = stmt.on_conflict_do_nothing()
194
195    stmt = stmt.returning(
196        *returning, sort_by_parameter_order=sort_by_parameter_order
197    )
198    return stmt
199