content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  module Glean.Database.Restore (
    restoreDatabase,
    restoreDatabaseFromSite,
    forRestoreSitesM, ifRestoreRepo, restorable,
  ) where
  
  import Control.Exception hiding(handle)
  import qualified Data.Set as Set
  import Data.Text (Text)
  
  import Util.STM
  
  import qualified Glean.Database.Backup.Backend as Backup
  import qualified Glean.Database.Backup.Locator as Backup
  import qualified Glean.Database.Catalog as Catalog
  import Glean.Database.Types
  import qualified Glean.ServerConfig.Types as ServerConfig
  import Glean.Types hiding (Database)
  import qualified Glean.Types as Thrift
  import Glean.Util.Observed as Observed
  
  
  forRestoreSitesM
    :: Env
    -> a
    -> (forall site. Backup.Site site => Text -> site -> IO a)
    -> IO [a]
  forRestoreSitesM env@Env{..} none inner = do
    ServerConfig.DatabaseRestorePolicy{..} <-
      ServerConfig.config_restore <$> Observed.get envServerConfig
    r <- atomically $ Backup.getAllSites env
    case r of
      sites@(_:_)
        | databaseRestorePolicy_enabled
          || not (Set.null databaseRestorePolicy_override) ->
        mapM (uncurry inner) sites
      _ -> return [none]
  
  ifRestoreRepo
    :: Env
    -> a
    -> Repo
    -> IO a
    -> IO a
  ifRestoreRepo Env{..} none repo inner = do
    policy <- ServerConfig.config_restore <$> Observed.get envServerConfig
    if restorable policy repo
      then inner
      else return none
  
  restorable :: ServerConfig.DatabaseRestorePolicy -> Repo -> Bool
  restorable ServerConfig.DatabaseRestorePolicy{..} repo =
    (repoName `Set.member` databaseRestorePolicy_override)
       /= databaseRestorePolicy_enabled
    where repoName = Thrift.repo_name repo
  
  restoreDatabase :: Env -> Text -> IO ()
  restoreDatabase env loc
    | Just (_, site, repo) <-
        Backup.fromRepoLocator (envBackupBackends env) loc =
          atomically =<< restoreDatabaseFromSite env site repo
    | otherwise = throwIO $
        Thrift.InvalidLocator $ "invalid locator '" <> loc <>  "'"
  
  restoreDatabaseFromSite
    :: Backup.Site site
    => Env
    -> site
    -> Repo
    -> IO (STM ())
  restoreDatabaseFromSite Env{..} site repo = do
    props <- Backup.inspect site repo
    return $ Catalog.startRestoring envCatalog repo props
  