Merge pull request #386 from kapily/add_arbitrary_pragmas

Browse Source 
Allow passing in arbitrary sqlite pragmas
This commit is contained in:
Friedrich Lindenberg 2021-11-17 17:53:22 +01:00 committed by GitHub
commit d0491ec0e3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 16 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
dataset/__init__.py
View File

@ -21,6 +21,7 @@ def connect(
ensure_schema=True, ensure_schema=True,
row_type=row_type, row_type=row_type,
sqlite_wal_mode=True, sqlite_wal_mode=True,
sqlite_pragmas=None,
): ):
"""Opens a new connection to a database. """Opens a new connection to a database.
@ -39,6 +40,10 @@ def connect(
the `ensure_schema` argument. It can also be overridden in a lot of the the `ensure_schema` argument. It can also be overridden in a lot of the
data manipulation methods using the `ensure` flag. data manipulation methods using the `ensure` flag.
If you want to run custom SQLite pragmas on database connect, you can add them to sqlite_pragmas
You can view a full list of PRAGMAs here:
https://www.sqlite.org/pragma.html
.. _SQLAlchemy Engine URL: http://docs.sqlalchemy.org/en/latest/core/engines.html#sqlalchemy.create_engine .. _SQLAlchemy Engine URL: http://docs.sqlalchemy.org/en/latest/core/engines.html#sqlalchemy.create_engine
.. _DB connection timeout: http://docs.sqlalchemy.org/en/latest/core/pooling.html#setting-pool-recycle .. _DB connection timeout: http://docs.sqlalchemy.org/en/latest/core/pooling.html#setting-pool-recycle
""" """
@ -52,4 +57,5 @@ def connect(
ensure_schema=ensure_schema, ensure_schema=ensure_schema,
row_type=row_type, row_type=row_type,
sqlite_wal_mode=sqlite_wal_mode, sqlite_wal_mode=sqlite_wal_mode,
sqlite_pragmas=sqlite_pragmas,
) )

14
dataset/database.py
View File

@ -30,6 +30,7 @@ class Database(object):
ensure_schema=True, ensure_schema=True,
row_type=row_type, row_type=row_type,
sqlite_wal_mode=True, sqlite_wal_mode=True,
sqlite_pragmas=None,
): ):
"""Configure and connect to the database.""" """Configure and connect to the database."""
if engine_kwargs is None: if engine_kwargs is None:
@ -56,16 +57,21 @@ class Database(object):
self.engine = create_engine(url, **engine_kwargs) self.engine = create_engine(url, **engine_kwargs)
self.is_postgres = self.engine.dialect.name == "postgresql" self.is_postgres = self.engine.dialect.name == "postgresql"
self.is_sqlite = self.engine.dialect.name == "sqlite" self.is_sqlite = self.engine.dialect.name == "sqlite"
if sqlite_pragmas is None:
def _enable_sqlite_wal_mode(dbapi_con, con_record): sqlite_pragmas = []
def _run_sqlite_pragmas(dbapi_con, con_record):
# reference: # reference:
# https://stackoverflow.com/questions/9671490/how-to-set-sqlite-pragma-statements-with-sqlalchemy # https://stackoverflow.com/questions/9671490/how-to-set-sqlite-pragma-statements-with-sqlalchemy
# https://stackoverflow.com/a/7831210/1890086 # https://stackoverflow.com/a/7831210/1890086
dbapi_con.execute("PRAGMA journal_mode=WAL") for pragma in sqlite_pragmas:
dbapi_con.execute(pragma)
if self.is_sqlite and parsed_url.path != "" and sqlite_wal_mode: if self.is_sqlite and parsed_url.path != "" and sqlite_wal_mode:
# we only enable WAL mode for sqlite databases that are not in-memory # we only enable WAL mode for sqlite databases that are not in-memory
event.listen(self.engine, "connect", _enable_sqlite_wal_mode) sqlite_pragmas.append("PRAGMA journal_mode=WAL")
if sqlite_pragmas:
event.listen(self.engine, "connect", _run_sqlite_pragmas)
self.types = Types(is_postgres=self.is_postgres) self.types = Types(is_postgres=self.is_postgres)
self.url = url self.url = url