From 259c9ee1abb6fccbbea5f9a9f7639a407f738d6b Mon Sep 17 00:00:00 2001
From: Ethan Dalool <git@voussoir.net>
Date: Thu, 16 Nov 2017 18:02:05 -0800
Subject: [PATCH] Add PhotoDB init arg `ephemeral`. Uses :memory: sql and
 tempdir.

Primarily necessitated by unit testing. Running through the DB_INIT
is quite slow on disk, so this argument causes the sql to be done
on an in-memory database and all the other files are put into a
TemporaryDirectory.

Eventually I would like to have the other files be in-memory too
but that may be overcomplicated and underuseful.
---
 etiquette/photodb.py | 38 +++++++++++++++++++++++++++++++++-----
 1 file changed, 33 insertions(+), 5 deletions(-)

diff --git a/etiquette/photodb.py b/etiquette/photodb.py
index 398d8f9..631d87f 100644
--- a/etiquette/photodb.py
+++ b/etiquette/photodb.py
@@ -6,6 +6,7 @@ import os
 import random
 import sqlite3
 import string
+import tempfile
 import time
 
 from . import constants
@@ -1092,11 +1093,22 @@ class PhotoDB(PDBAlbumMixin, PDBBookmarkMixin, PDBPhotoMixin, PDBTagMixin, PDBUs
     def __init__(
             self,
             data_directory=None,
+            ephemeral=False,
         ):
         super().__init__()
 
+        self.ephemeral = ephemeral
+
         if data_directory is None:
-            data_directory = constants.DEFAULT_DATADIR
+            if self.ephemeral:
+                # In addition to the data_dir as a pathclass object, keep the
+                # TempDir object so we can use the cleanup method later.
+                self.ephemeral_directory = tempfile.TemporaryDirectory(prefix='etiquette_ephem_')
+                data_directory = self.ephemeral_directory.name
+            else:
+                data_directory = constants.DEFAULT_DATADIR
+        elif self.ephemeral:
+            raise exceptions.NotExclusive(['data_directory', 'ephemeral'])
 
         # DATA DIR PREP
         data_directory = helpers.normalize_filepath(data_directory, allowed=':/\\')
@@ -1104,9 +1116,14 @@ class PhotoDB(PDBAlbumMixin, PDBBookmarkMixin, PDBPhotoMixin, PDBTagMixin, PDBUs
         os.makedirs(self.data_directory.absolute_path, exist_ok=True)
 
         # DATABASE
-        self.database_file = self.data_directory.with_child(constants.DEFAULT_DBNAME)
-        existing_database = self.database_file.exists
-        self.sql = sqlite3.connect(self.database_file.absolute_path)
+        if self.ephemeral:
+            self.sql = sqlite3.connect(':memory:')
+            existing_database = False
+        else:
+            self.database_file = self.data_directory.with_child(constants.DEFAULT_DBNAME)
+            existing_database = self.database_file.exists
+            self.sql = sqlite3.connect(self.database_file.absolute_path)
+
         self.cur = self.sql.cursor()
 
         if existing_database:
@@ -1159,12 +1176,23 @@ class PhotoDB(PDBAlbumMixin, PDBBookmarkMixin, PDBPhotoMixin, PDBTagMixin, PDBUs
             'user': self._user_cache,
         }
 
+    def __del__(self):
+        self.close()
+
     def __repr__(self):
-        return 'PhotoDB(data_directory={datadir})'.format(datadir=repr(self.data_directory))
+        if self.ephemeral:
+            return 'PhotoDB(ephemeral=True)'
+        else:
+            return 'PhotoDB(data_directory={datadir})'.format(datadir=repr(self.data_directory))
 
     def _uncache(self):
         self._cached_frozen_children = None
 
+    def close(self):
+        self.sql.close()
+        if self.ephemeral:
+            self.ephemeral_directory.cleanup()
+
     def commit(self):
         while self.on_commit_queue:
             task = self.on_commit_queue.pop()