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.Repo
    ( -- * Listing repos
      getLatestRepo
    , getLatestRepos
    , getRepo
    , getRepos
    , getSCMrevisions
    , scmRevisionsOfDatabase
    , LatestRepos(..)
    , NoDatabase(..)
    ) where
  
  import Control.Exception
  import Data.Function
  import Data.List
  import Data.List.Extra (groupSortOn)
  import Data.HashMap.Strict (HashMap)
  import qualified Data.HashMap.Strict as HashMap
  import Data.Map (Map)
  import qualified Data.Map as Map
  import Data.Maybe
  import Data.Text (Text)
  import qualified Data.Text as T
  
  import Util.Log
  
  import Glean.Repo.Text
  import Glean.Backend.Types
  import Glean.Types
  
  
  newtype NoDatabase = NoDatabase Text
    deriving Show
  instance Control.Exception.Exception NoDatabase
  
  -- | Collecting 'getLatestRepo' for every available repo in the backend
  data LatestRepos = LatestRepos{
    latestRepos :: !(Map Text Repo),  -- ^ Guaranteed available
    allLatestRepos :: !(Map Text [Database]) -- ^ Possibly unavailable
    }
  
  -- | Ask the server which databases are available, and return
  -- list of repo.
  getRepos
    :: Backend be
    => be
    -> (Repo -> Bool)  -- ^ Predicate to choose repo
    -> IO ([Database], [Database])
  getRepos backend pred = do
    let
      -- If we're connecting to a tier and using shards, then we can
      -- assume that DBs in the Restoring state are usable, because
      -- hopefully some other server will have the DB. The idea is
      -- to get a consistent view of the available DBs even if we
      -- happen to contact a server that has recently come up and is
      -- still downloading DBs.
  
      ok Database{..} =
        (database_status == DatabaseStatus_Complete
          || (usingShards backend &&
             (database_status == DatabaseStatus_Restoring
              || database_status == DatabaseStatus_Available))) &&
        pred database_repo &&
        isNothing database_expire_time
    xss <- groupSortOn (repo_name . database_repo) . filter ok <$>
      localDatabases backend
    let pickRepo xs = do
          let
            sorted = sortBy (flip compare `on` database_created_since_epoch) xs
          available <- maybeToList <$> checkRestorableAvailable backend sorted
          return (available, sorted)
    mconcat <$> mapM pickRepo xss
  
  
  -- | Ask the server which databases are available, and return the first one that
  -- matches the predicate.
  getRepo
    :: Backend be
    => be
    -> (Repo -> Bool)  -- ^ Predicate to choose repo
    -> IO (Maybe Repo)
  getRepo backend pred =
    listToMaybe . map database_repo . fst <$> getRepos backend pred
  
  -- | Ask the server which databases are available, and select the
  -- latest complete database for all repo names selected by the predicate.
  getLatestRepos
    :: Backend be
    => be
    -> (Text -> Bool)  -- ^ Predicate to choose repo names to include
    -> IO LatestRepos
  getLatestRepos backend keepName = do
    (latest, all) <- getRepos backend (keepName . repo_name)
    let
      avail = Map.fromList
        [ (repo_name r, r) | d <- latest, let r = database_repo d ]
      alldb = Map.fromListWith (++)
        [ (repo_name (database_repo d), [d]) | d <- all]
    return $ LatestRepos avail alldb
  
  -- | Ask the server which databases are available, and select the
  -- latest complete database for the given repo name.
  --
  -- Otherwise throw 'NoDatabase' exception
  getLatestRepo
    :: Backend be
    => be
    -> Text  -- ^ The name of the repo, e.g. "fbsource"
    -> IO Repo
  getLatestRepo backend name = do
    lr <- getLatestRepos backend (name ==)
    case Map.lookup name (latestRepos lr) of
      Nothing -> throwIO $ NoDatabase $ "no database found for: " <> name
      Just x -> return x
  
  -- | For each DB that is "restorable" on the host that returned the
  -- listDatabases result, check whether we can route to a host that has
  -- the DB. If not, we'll pick the next available DB in the list. This
  -- logic only kicks in when the service is a Tier and supports shards,
  -- otherwise we ignore restorable DBs when filtering above.
  --
  -- Otherwise return Nothing.
  checkRestorableAvailable
    :: Backend be
    => be
    -> [Database]
    -> IO (Maybe Database)
  checkRestorableAvailable _ [] = return Nothing
  checkRestorableAvailable backend (db@Database{..}:dbs)
    | database_status `notElem`
      [DatabaseStatus_Restoring, DatabaseStatus_Available] =
      return (Just db)
    | not (usingShards backend) =
      checkRestorableAvailable backend dbs
    | otherwise = do
      let status = case database_status of
            DatabaseStatus_Restoring -> "restoring"
            DatabaseStatus_Available -> "available elsewhere"
            _ -> error "unreachable"
      vlog 1 $ "Repo " <> showRepo database_repo <>
        " is " <> status <> ", checking for shard availability..."
      avail <- hasDatabase backend database_repo
      vlog 1 $ "Repo " <> showRepo database_repo <>
         ": " <> (if avail then "some" else "no") <> " hosts have it"
      if not avail
        then checkRestorableAvailable backend dbs
        else return (Just db)
  
  -- | Returns a mapping (SCM name -> SCM revision)
  --   of the SCM repositories indexed by the DB.
  --   Assumes that the DB includes SCM properties in the conventional format.
  getSCMrevisions :: Backend be => be -> Repo -> IO (HashMap Text Text)
  getSCMrevisions backend repo = do
    GetDatabaseResult{getDatabaseResult_database = db} <- getDatabase backend repo
    return $ scmRevisionsOfDatabase db
  
  -- | Extract a mapping (SCM name -> SCM revision) from the DB properties
  scmRevisionsOfDatabase :: Database -> HashMap Text Text
  scmRevisionsOfDatabase Database{..} =
    HashMap.fromList
      [ (scm_name, prop_value)
      | (prop_name, prop_value) <- HashMap.toList database_properties
      , Just scm_name <- [T.stripPrefix "glean.scm." prop_name]
      -- Want to filter out values that are not SCM names,
      -- but lack an exhaustive list of valid SCM names.
      , scm_name `notElem` ["repo", "revision"]
      ]
  