diff --git a/dataset/__init__.py b/dataset/__init__.py index e4266e2..8821179 100644 --- a/dataset/__init__.py +++ b/dataset/__init__.py @@ -21,6 +21,7 @@ def connect( ensure_schema=True, row_type=row_type, sqlite_wal_mode=True, + sqlite_pragmas=None, ): """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 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 .. _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, row_type=row_type, sqlite_wal_mode=sqlite_wal_mode, + sqlite_pragmas=sqlite_pragmas, ) diff --git a/dataset/database.py b/dataset/database.py index 22d1a49..42c12f4 100644 --- a/dataset/database.py +++ b/dataset/database.py @@ -30,6 +30,7 @@ class Database(object): ensure_schema=True, row_type=row_type, sqlite_wal_mode=True, + sqlite_pragmas=None, ): """Configure and connect to the database.""" if engine_kwargs is None: @@ -56,16 +57,21 @@ class Database(object): self.engine = create_engine(url, **engine_kwargs) self.is_postgres = self.engine.dialect.name == "postgresql" self.is_sqlite = self.engine.dialect.name == "sqlite" - - def _enable_sqlite_wal_mode(dbapi_con, con_record): + if sqlite_pragmas is None: + sqlite_pragmas = [] + def _run_sqlite_pragmas(dbapi_con, con_record): # reference: # https://stackoverflow.com/questions/9671490/how-to-set-sqlite-pragma-statements-with-sqlalchemy # 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: # 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.url = url