Restore.hs

{-
  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.
-}

{-# LANGUAGE ApplicativeDo #-}
{-# OPTIONS_GHC -Wno-incomplete-patterns -Wno-deprecations #-}
module GleanCLI.Restore (RestoreCommand) where

import Control.Exception.Safe (catch)
import Control.Monad (forM_, forM)
import Control.Monad.Extra (firstJustM)
import Control.Concurrent
import Data.Either
import Data.List (sortOn)
import Data.Text (Text)
import Data.Maybe
import qualified Data.Text as Text
import Data.Time
import Data.Time.Clock.POSIX
import Options.Applicative

import Util.IO
import Util.OptParse
import Util.Log (logInfo)
import Util.STM (atomically)

import Glean
  ( Repo(..)
  , Database(..)
  , DatabaseStatus(..)
  )
import qualified Glean
import qualified Glean.Database.Backup.Locator as Backup
import Glean.LocalOrRemote (BackendKind(..), backendKind)
import Glean.Database.Backup.Backend (Site(inspect))
import Glean.Database.Meta (metaToThriftDatabase)

import GleanCLI.Common
import GleanCLI.Types
import GleanCLI.Utils
import Glean.Database.Exception
import Glean.Database.Repo (inRepo)
import Glean.Database.Open (depParent)

data WhatToRestore
  = RestoreLocator Text
  | RestoreDb Repo
  | RestoreDbOnDay Text Day
  | RestoreDbLatest Text

data RestoreDbSpec
  = InstanceOpt Text
  | DateOpt Day
  | LatestOpt

data RestoreCommand
  = Restore
      { what :: WhatToRestore
      , ignoreDependencies :: Bool
      }

type Locator = Text

instance Plugin RestoreCommand where
  serverConfigTransform _ = disableJanitor . disableAutoBackups

  parseCommand =
    commandParser "restore" (progDesc "Restore a database") $
      Restore <$> what <*> deps
    where
      locator = strArgument
        (  metavar "LOCATOR"
        <> help "DB location, see :list-all in glean shell."
        )
      deps = switch
        (  long "ignore-dependencies"
        <> help "Don't download database dependencies when specifying a DB name.")
      what =
        (RestoreLocator <$> locator) <|>
        (RestoreDb <$> dbSlash) <|> do
          repoName <- dbNameOpt
          spec <- InstanceOpt <$> dbInstanceOpt <|> DateOpt <$> dayOpt
            <|> latestOpt
          return $ case spec of
            InstanceOpt hash -> RestoreDb (Repo repoName hash)
            DateOpt day -> RestoreDbOnDay repoName day
            LatestOpt  -> RestoreDbLatest repoName

      parseDay = return .
        parseTimeOrError False defaultTimeLocale (iso8601DateFormat Nothing)

      dayOpt = option (eitherReader parseDay)
        (  long "date"
        <> metavar "YYYY-MM-DD"
        <> help "find a DB for this date (UTC)")

      latestOpt = flag' LatestOpt
        (long "latest"
        <> help "find latest instance")

  runCommand _ _ backend Restore{..} = do
    targets <- locatorsToRestore
    restore targets
    wait (fst <$> targets)
    where
      locatorsToRestore = case what of
        -- ignores dependencies
        RestoreLocator locator -> return [(locator, Nothing)]
        RestoreDb repo
          | BackendEnv env <- backendKind backend -> do
            logInfo $ inRepo repo "Searching"
            sites <- atomically $ Backup.getAllSites env
            db <- getDatabaseFromSite sites repo
            if ignoreDependencies then withLocator [db] [repo] else do
            allDbs <- getDependenciesOneByOne sites db
            withLocator allDbs (map database_repo allDbs)
          | otherwise -> do
            databases <- listWithBackups
            let deps = if ignoreDependencies
                  then []
                  else dependencies databases repo
            withLocator databases $ repo : deps
        RestoreDbOnDay repoName day -> do
          databases <- listWithBackups
          let
            matchingDay = listToMaybe
                [ database_repo
                | db@Database{..} <- sortOn dbTime databases
                , repo_name database_repo == repoName
                , day == utctDay (posixSecondsToUTCTime $ fromIntegral $
                    Glean.unPosixEpochTime (dbTime db))
                ]
          case matchingDay of
            Just repo -> do
              let deps = if ignoreDependencies
                    then []
                    else dependencies databases repo
              withLocator databases $ repo : deps
            Nothing -> die 1 $ unwords
              ["Cannot find backup locator for", Text.unpack repoName, "on"
              , formatTime defaultTimeLocale (iso8601DateFormat Nothing) day ]
        RestoreDbLatest repoName -> do
          databases <- listWithBackups
          let latestRepo = listToMaybe $ reverse
                [ database_repo
                | Database{..} <- sortOn dbTime databases
                , repo_name database_repo == repoName
                ]
          case latestRepo of
            Just repo -> do
              let deps = if ignoreDependencies
                    then []
                    else dependencies databases repo
              withLocator databases $ repo : deps
            Nothing -> die 1 $ "Cannot find repo for " <> Text.unpack repoName

      getDependenciesOneByOne sites db = do
        case database_dependencies db of
          Nothing -> return [db]
          Just dep -> do
            dep_db <- getDatabaseFromSite sites (depParent dep)
            (db:) <$> getDependenciesOneByOne sites dep_db

      getDatabaseFromSite sites repo = do
        result <- flip firstJustM sites $ \(_prefix, site) ->
                    restorable repo <$> inspect site repo
        maybe (dieHere repo) return result

      restorable repo meta
        = Just $
          metaToThriftDatabase Glean.DatabaseStatus_Restorable Nothing repo meta

      dieHere repo = die 1 $ "Cannot find " <> show repo

      dbTime Database{..} =
              fromMaybe database_created_since_epoch
                database_repo_hash_time

      withLocator :: [Database] -> [Repo] -> IO [(Locator, Maybe Repo)]
      withLocator databases repos = forM repos $ \repo -> do
        locator <- repoLocator databases repo
        return (locator, Just repo)

      listWithBackups =
        Glean.listDatabasesResult_databases <$>
          Glean.listDatabases backend Glean.ListDatabases
            { listDatabases_includeBackups = True
            , listDatabases_client_info = Nothing
            }

      restore targets = do
        forM_ targets $ \(locator, mrepo) -> do
          logInfo $ unwords $
            [ "Restoring" ] ++
            [ Glean.showRepo repo | Just repo <- [mrepo]] ++
            [ "from", Text.unpack locator]
          Glean.restoreDatabase backend locator
              `catch` \DBAlreadyExists -> return ()

      wait locators = do
        let noRetries =
              case backendKind backend of
                -- If we're calling a remote backend, it might not have
                -- started the restoring yet. So we'll wait a bit to give
                -- some time to start the restoring.
                BackendThrift _ -> 5
                _ -> 0
        dbs <- waitForRestoreToStart noRetries
        waitForRestoreToFinish dbs
        where
          waitForRestoreToStart (retries::Int) = do
            let
              isAvailable =  (DatabaseStatus_Available ==) . database_status
              retry err =
                if
                  | retries > 0
                    -> do
                        threadDelay 1000000
                        waitForRestoreToStart (retries-1)
                  | retries == 0
                    -> die 1 err
            localDatabases <- listLocalDBs
            (nonExistingLocs, dbs) <- partitionEithers <$>
              traverse (locatorDb localDatabases) locators
            if
              | locator : _ <- nonExistingLocs ->
                retry $ unwords
                        ["error: did not find database locator"
                        ,Text.unpack locator]
              | db : _ <- filter isAvailable dbs ->
                retry $ unwords
                        ["error: timed out restoring"
                        ,show db]
              | otherwise -> return dbs

          waitForRestoreToFinish dbs = do
            let isRestoring = (DatabaseStatus_Restoring ==) . database_status
            if any isRestoring dbs
              then do threadDelay 1000000
                      localDatabases <- listLocalDBs
                      edbs <- traverse (locatorDb localDatabases) locators
                      waitForRestoreToFinish (rights edbs)
              else forM_ dbs $ \db -> case database_status db of
                DatabaseStatus_Complete -> return ()
                DatabaseStatus_Missing -> putStrLn $ unwords
                  [ "Some of"
                  , Glean.showRepo (database_repo db)
                  , "dependencies are missing."
                  , "You may want to restore those as well"
                  ] -- TODO: List missing dependencies
                status ->
                  die 1 $ "error: unexpected database status: " <> show status

          listLocalDBs = Glean.listDatabasesResult_databases <$>
            Glean.listDatabases backend Glean.ListDatabases
              { listDatabases_includeBackups = False
              , listDatabases_client_info = Nothing
              }

          locatorDb databases locator =
            case
              [ db
              | db <- databases
              , database_location db == Just locator ]
            of
              [db] -> return (Right db)
              [] -> return (Left locator)
              _ -> die 1 $ unwords
                ["error: server has multiple DBs with the locator"
                , Text.unpack locator]

repoLocator :: [Database] -> Repo -> IO Locator
repoLocator databases repo =
  case [ locator
        | Database{..} <- databases
        , database_repo == repo
        , Just locator <- [database_location] ]
  of
    (locator:_) -> return locator
    [] -> die 1 $ "Cannot find backup locator for " <> Glean.showRepo repo

dependencies :: [Database] -> Repo -> [Repo]
dependencies databases repo = repoDeps repo
  where
    repoDeps base =
      case repoDirectDep base of
        Nothing -> []
        Just repo -> repo : repoDeps repo
    repoDirectDep repo = listToMaybe
        [ depParent dep
        | Database{..} <- databases
        , database_repo == repo
        , Just dep <- [database_dependencies]
        ]