Fix to support replacing a database, closes #2465

This commit is contained in:
Simon Willison 2025-02-06 10:32:47 -08:00
commit f95ac19e71
2 changed files with 36 additions and 2 deletions

View file

@ -32,6 +32,7 @@ AttachedDatabase = namedtuple("AttachedDatabase", ("seq", "name", "file"))
class Database: class Database:
# For table counts stop at this many rows: # For table counts stop at this many rows:
count_limit = 10000 count_limit = 10000
_thread_local_id_counter = 1
def __init__( def __init__(
self, self,
@ -43,6 +44,8 @@ class Database:
mode=None, mode=None,
): ):
self.name = None self.name = None
self._thread_local_id = f"x{self._thread_local_id_counter}"
Database._thread_local_id_counter += 1
self.route = None self.route = None
self.ds = ds self.ds = ds
self.path = path self.path = path
@ -278,11 +281,11 @@ class Database:
# threaded mode # threaded mode
def in_thread(): def in_thread():
conn = getattr(connections, self.name, None) conn = getattr(connections, self._thread_local_id, None)
if not conn: if not conn:
conn = self.connect() conn = self.connect()
self.ds._prepare_connection(conn, self.name) self.ds._prepare_connection(conn, self.name)
setattr(connections, self.name, conn) setattr(connections, self._thread_local_id, conn)
return fn(conn) return fn(conn)
return await asyncio.get_event_loop().run_in_executor( return await asyncio.get_event_loop().run_in_executor(

View file

@ -721,3 +721,34 @@ async def test_hidden_tables(app_client):
"r_parent", "r_parent",
"r_rowid", "r_rowid",
] ]
@pytest.mark.asyncio
async def test_replace_database(tmpdir):
path1 = str(tmpdir / "data1.db")
(tmpdir / "two").mkdir()
path2 = str(tmpdir / "two" / "data1.db")
sqlite3.connect(path1).executescript(
"""
create table t (id integer primary key);
insert into t (id) values (1);
insert into t (id) values (2);
"""
)
sqlite3.connect(path2).executescript(
"""
create table t (id integer primary key);
insert into t (id) values (1);
"""
)
datasette = Datasette([path1])
db = datasette.get_database("data1")
count = (await db.execute("select count(*) from t")).first()[0]
assert count == 2
# Now replace that database
datasette.get_database("data1").close()
datasette.remove_database("data1")
datasette.add_database(Database(datasette, path2), "data1")
db2 = datasette.get_database("data1")
count = (await db2.execute("select count(*) from t")).first()[0]
assert count == 1