stash DbHandle in Annex state

3 files changed, 41 insertions, 10 deletions

diff --git a/Database/AssociatedFiles.hs b/Database/AssociatedFiles.hs
index 8244f15e8..d17eb8112 100644
--- a/Database/AssociatedFiles.hs
+++ b/Database/AssociatedFiles.hs
@@ -13,6 +13,7 @@
 module Database.AssociatedFiles (
 	DbHandle,
 	openDb,
+	flushDb,
 	closeDb,
 	addDb,
 	getDb,
@@ -21,6 +22,7 @@ module Database.AssociatedFiles (
 ) where
 
 import Database.Types
+import Database.AssociatedFiles.Types
 import qualified Database.Handle as H
 import Locations
 import Common hiding (delete)
@@ -33,8 +35,6 @@ import Messages
 import Database.Persist.TH
 import Database.Esqueleto hiding (Key)
 
-newtype DbHandle = DbHandle H.DbHandle
-
 share [mkPersist sqlSettings, mkMigrate "migrateAssociated"] [persistLowerCase|
 Associated
   key SKey
@@ -64,8 +64,25 @@ openDb = withExclusiveLock gitAnnexAssociatedFilesDbLock $ do
 closeDb :: DbHandle -> IO ()
 closeDb (DbHandle h) = H.closeDb h
 
-addDb :: DbHandle -> Key -> FilePath -> IO ()
-addDb (DbHandle h) k f = H.queueDb h (\_ _ -> pure True) $ do
+withDbHandle :: (H.DbHandle -> IO a) -> Annex a
+withDbHandle a = do
+	(DbHandle h) <- dbHandle
+	liftIO $ a h
+
+dbHandle :: Annex DbHandle
+dbHandle = maybe startup return =<< Annex.getState Annex.associatedfilesdbhandle
+  where
+	startup = do
+		h <- openDb
+		Annex.changeState $ \s -> s { Annex.associatedfilesdbhandle = Just h }
+		return h
+
+{- Flushes any changes made to the database. -}
+flushDb :: Annex ()
+flushDb = withDbHandle H.flushQueueDb
+
+addDb :: Key -> FilePath -> Annex ()
+addDb k f = withDbHandle $ \h -> H.queueDb h (\_ _ -> pure True) $ do
 	-- If the same file was associated with a different key before,
 	-- remove that.
 	delete $ from $ \r -> do
@@ -76,8 +93,8 @@ addDb (DbHandle h) k f = H.queueDb h (\_ _ -> pure True) $ do
 
 {- Note that the files returned used to be associated with the key, but
  - some of them may not be any longer. -}
-getDb :: DbHandle -> Key -> IO [FilePath]
-getDb (DbHandle h) = H.queryDb h . getDb' . toSKey
+getDb :: Key -> Annex [FilePath]
+getDb k = withDbHandle $ \h -> H.queryDb h $ getDb' $ toSKey k
 
 getDb' :: SKey -> SqlPersistM [FilePath]
 getDb' sk = do
@@ -86,8 +103,8 @@ getDb' sk = do
 		return (r ^. AssociatedFile)
 	return $ map unValue l
 
-removeDb :: DbHandle -> Key -> FilePath -> IO ()
-removeDb (DbHandle h) k f =  H.queueDb h (\_ _ -> pure True) $
+removeDb :: Key -> FilePath -> Annex ()
+removeDb k f =  withDbHandle $ \h -> H.queueDb h (\_ _ -> pure True) $
 	delete $ from $ \r -> do
 		where_ (r ^. AssociatedKey ==. val sk &&. r ^. AssociatedFile ==. val f)
   where
diff --git a/Database/AssociatedFiles/Types.hs b/Database/AssociatedFiles/Types.hs
new file mode 100644
index 000000000..8c32dcf22
--- /dev/null
+++ b/Database/AssociatedFiles/Types.hs
@@ -0,0 +1,14 @@
+{- Sqlite database used for tracking a key's associated files, data types.
+ -
+ - Copyright 2015 Joey Hess <id@joeyh.name>
+ -:
+ - Licensed under the GNU GPL version 3 or higher.
+ -}
+
+module Database.AssociatedFiles.Types (
+	DbHandle(..)
+) where
+
+import qualified Database.Handle as H
+
+newtype DbHandle = DbHandle H.DbHandle
diff --git a/Database/Handle.hs b/Database/Handle.hs
index 439e7c18b..6d312df68 100644
--- a/Database/Handle.hs
+++ b/Database/Handle.hs
@@ -21,7 +21,6 @@ module Database.Handle (
 
 import Utility.Exception
 import Utility.Monad
-import Messages
 
 import Database.Persist.Sqlite
 import qualified Database.Sqlite as Sqlite
@@ -35,6 +34,7 @@ import Control.Monad.Trans.Resource (runResourceT)
 import Control.Monad.Logger (runNoLoggingT)
 import Data.List
 import Data.Time.Clock
+import System.IO
 
 {- A DbHandle is a reference to a worker thread that communicates with
  - the database. It has a MVar which Jobs are submitted to. -}
@@ -79,7 +79,7 @@ type TableName = String
 workerThread :: T.Text -> TableName -> MVar Job -> IO ()
 workerThread db tablename jobs = catchNonAsync (run loop) showerr
   where
-  	showerr e = liftIO $ warningIO $
+  	showerr e = liftIO $ hPutStrLn stderr $
 		"sqlite worker thread crashed: " ++ show e
 	
 	loop = do