Merge pull request #386 from kapily/add_arbitrary_pragmas

Allow passing in arbitrary sqlite pragmas

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,
     )

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