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.
  -}
  
  {-# LANGUAGE ApplicativeDo #-}
  module GleanCLI.Write (WriteCommand, FinishCommand) where
  
  import Control.Monad
  import qualified Data.ByteString as B
  import Data.Default
  import Data.Proxy
  import qualified Data.HashMap.Strict as HashMap
  import Data.Text (Text)
  import qualified Data.Text as Text
  import qualified Data.Text.IO as Text
  import qualified Data.Text.Encoding as Encode
  import Options.Applicative
  
  import Control.Concurrent.Stream (stream)
  import Thrift.Protocol.Compact (Compact)
  import Thrift.Protocol
  import Util.Control.Exception
  import Util.IO
  import Util.OptParse
  import Util.STM
  
  import Glean
  import Glean.LocalOrRemote (loadDbSchema)
  import qualified Glean.LocalOrRemote as LocalOrRemote
  import Glean.Database.Schema
  import Glean.Database.Write.Batch (syncWriteDatabase)
  import Glean.Types as Thrift
  import Util.Time
  import Glean.Write
  import Glean.Write.JSON ( buildJsonBatch, syncWriteJsonBatch )
  
  import GleanCLI.Common
  import GleanCLI.Finish
  import GleanCLI.Types
  import Data.Time.Clock (UTCTime)
  import Glean.Database.Meta (utcTimeToPosixEpochTime)
  import Data.ByteString (ByteString)
  import Glean.Util.ThriftService (queueTimeout)
  
  data WriteCommand
    = Write
        { writeRepo :: Repo
        , writeRepoTime :: Maybe UTCTime
        , create :: Bool
        , dependencies :: Maybe (IO Thrift.Dependencies)
        , finish :: Bool
        , properties :: [(Text,Text)]
        , writeMaxConcurrency :: Int
        , useLocalCache :: Bool
        , sendQueueSettings :: Glean.SendAndRebaseQueueSettings
        , writeFileFormat :: FileFormat
        , updateSchemaForStacked :: Bool
        , factsSource :: FactsSource
        }
  
  data FactsSource = Locations [Text] | Files [FilePath]
  
  factsSourceOpts :: Parser FactsSource
  factsSourceOpts = Files <$> fileArg <|> Locations <$> locationOpt
  
  fileArg :: Parser [FilePath]
  fileArg = many $ strArgument
    (  metavar "FILE..."
    <> help ("File(s) of facts to add to the DB. "
    <> "You can specify the format of the file with --file-format. ")
    )
  
  locationOpt :: Parser [Text]
  locationOpt = many $ strOption
    (  long "location"
    <> metavar "LOCATION"
    <> help "Location(s) of facts to be downloaded and written to the db"
    )
  
  repoTimeOpt :: Parser UTCTime
  repoTimeOpt = option readTime
    (  long "repo-hash-time"
    <> metavar "yyyy-mm-ddThh:mm:ssZ"
    <> help "Timestamp of the source data to be indexed."
    )
    where
      readTime :: ReadM UTCTime
      readTime = eitherReader $ \str ->
        case readUTC $ Text.pack str of
          Just value -> Right value
          Nothing ->
            Left "expecting UTC time e.g. 2021-01-01T12:30:00Z"
  
  dbPropertiesOpt :: Parser [(Text, Text)]
  dbPropertiesOpt = many $ option readProperty
    (  long "property"
    <> metavar "NAME=VALUE"
    <> help "Set DB's properties when creating a DB."
    )
    where
      readProperty :: ReadM (Text,Text)
      readProperty = eitherReader $ \str ->
        case break (=='=') str of
          (name, '=':value) -> Right (Text.pack name, Text.pack value)
          _other -> Left "--property: expecting NAME=VALUE"
  
  finishOpt :: Parser Bool
  finishOpt = switch
    (  long "finish"
    <> help ("Mark the DB as complete. When a DB is complete, "
    <> "it is not possible to add any more facts to it.")
    )
  
  stackedOpt :: Parser Repo
  stackedOpt = option (maybeReader Glean.parseRepo)
    (  long "stacked"
    <> metavar "DB"
    <> help ("Created DB will be stacked on top of this DB. "
    <> "For more details about its schema, see --update-schema-for-stacked.")
    )
  
  useLocalSwitchOpt :: Parser Bool
  useLocalSwitchOpt = switch
    (  long "use-local-cache"
    <> help ("Use a local cache to avoid resending duplicate facts. May improve "
    <> "write performance.")
    )
  
  incrementalOpt :: Parser Repo
  incrementalOpt = option (maybeReader Glean.parseRepo)
    (  long "incremental"
    <> metavar "DB"
    <> help "Create an incremental DB on top of this DB."
    )
  
  splitUnits :: Text -> [ByteString]
  splitUnits = map Encode.encodeUtf8 . Text.splitOn ","
  
  extractLines :: FilePath -> IO [ByteString]
  extractLines file = map Encode.encodeUtf8 . Text.lines <$> Text.readFile file
  
  includeOptString :: Parser (IO [ByteString])
  includeOptString = return . splitUnits <$> strOption
    (  long "include"
    <> metavar "unit,unit,.."
    <> help "For incremental DBs only. Include these units."
    )
  
  includeOptFile :: Parser (IO [ByteString])
  includeOptFile = extractLines <$> strOption
    (  long "include-file"
    <> metavar "FILE"
    <> help ("For incremental DBs only. Include units in FILE "
    <> "(one per line).")
    )
  
  includeOpt :: Parser (IO [ByteString])
  includeOpt = includeOptFile <|> includeOptString
  
  excludeOptString :: Parser (IO [ByteString])
  excludeOptString =  return . splitUnits <$> strOption
    (  long "exclude"
    <> metavar "unit,unit,.."
    <> help "For incremental DBs only. Exclude these units."
    )
  
  excludeOptFile :: Parser (IO [ByteString])
  excludeOptFile =  extractLines <$> strOption
    (  long "exclude-file"
    <> metavar "FILE"
    <> help ("For incremental DBs only. Exclude units in FILE "
    <> "(one per line).")
    )
  
  excludeOpt :: Parser (IO [ByteString])
  excludeOpt = excludeOptFile <|> excludeOptString
  
  updateSchemaForStackedOpt :: Parser Bool
  updateSchemaForStackedOpt = switch
    (  long "update-schema-for-stacked"
    <> help (
      "When creating a stacked DB, use the current schema instead " <>
      "of the schema from the base DB.")
    )
  
  instance Plugin WriteCommand where
    parseCommand = createCmd <|> writeCmd
      where
      createCmd =
        commandParser "create" (progDesc (
            "Create a new standalone|stacked|incremental DB. "
            <> "Please carefully read help above to understand how various "
            <> "options are related to each other.")) $ do
          writeRepo <- dbOpts
          writeRepoTime <- optional repoTimeOpt
          finish <- finishOpt
          dependencies <- optional (stackedOptions <|> updateOptions)
          properties <- dbPropertiesOpt
          writeMaxConcurrency <- maxConcurrencyOpt
          useLocalCache <- useLocalSwitchOpt
          sendQueueSettings <- Glean.sendAndRebaseQueueOptions
          writeFileFormat <- fileFormatOpt JsonFormat
          updateSchemaForStacked <- updateSchemaForStackedOpt
          factsSource <- factsSourceOpts
          return Write
            { create=True
            , ..
            }
  
      writeCmd =
        commandParser "write" (progDesc (
            "Write facts to an existing DB. "
            <> "Please carefully read help above to understand how various "
            <> "options are related to each other.")) $ do
          writeRepo <- dbOpts
          finish <- finishOpt
          writeMaxConcurrency <- maxConcurrencyOpt
          useLocalCache <- useLocalSwitchOpt
          sendQueueSettings <- Glean.sendAndRebaseQueueOptions
          writeFileFormat <- fileFormatOpt JsonFormat
          factsSource <- factsSourceOpts
          return Write
            { create=False, writeRepoTime=Nothing
            , properties=[], dependencies=Nothing
            , updateSchemaForStacked = False
            , ..
            }
  
      stackedOptions :: Parser (IO Dependencies)
      stackedOptions = f <$> stackedOpt
        where
          f Repo{..} =
            return $ Thrift.Dependencies_stacked $
              Thrift.Stacked repo_name repo_hash Nothing
  
      updateOptions :: Parser (IO Dependencies)
      updateOptions = do
        repo <- incrementalOpt
        let
          include = (,False) <$> includeOpt
          exclude = (,True) <$> excludeOpt
        ~(units, exclude) <- include <|> exclude
        return $ fmap (prune repo exclude) units
        where
          prune :: Repo -> Bool -> [ByteString] -> Dependencies
          prune repo exclude units = Thrift.Dependencies_pruned $
            Thrift.Pruned repo units exclude Nothing
  
    runCommand _ _ backend Write{..} =
      tryBracket
         (when create $ do
              putStrLn $ "Creating DB " <> showRepo writeRepo
              deps <- maybe (return Nothing) (fmap Just) dependencies
              Thrift.KickOffResponse alreadyExists <-
                Glean.kickOffDatabase backend def
                  { kickOff_repo = writeRepo
                  , kickOff_properties = HashMap.fromList properties
                  , kickOff_dependencies = deps
                  , kickOff_repo_hash_time =
                      utcTimeToPosixEpochTime <$> writeRepoTime
                  , kickOff_update_schema_for_stacked = updateSchemaForStacked
                  }
              when alreadyExists $ die 3 "DB create failure: already exists"
         )
         (\_ result -> case resultToFailure result of
            Just err -> die 3 $ "DB create failure: " ++ err
            Nothing -> when finish $ finished backend writeRepo)
         (\_ ->
            write Write{..})
      where
      write Write{..} = do
        case factsSource of
          Files writeFiles -> writeBatches writeFiles Write{..}
          Locations locations -> writeBatchDescriptors locations Write{..}
      writeBatches writeFiles Write{useLocalCache = True, ..} = do
        dbSchema <- loadDbSchema backend writeRepo
        logMessages <- newTQueueIO
        let inventory = schemaInventory dbSchema
            queueSettings = sendQueueSettings
              { sendAndRebaseQueueAllowRemoteReferences =
                  case writeFileFormat of
                    -- we expect binary batches to be self-contained.
                    BinaryFormat -> False
                    JsonFormat -> True
              }
        Glean.withSendAndRebaseQueue backend writeRepo inventory queueSettings $
          \queue ->
            stream writeMaxConcurrency (forM_ writeFiles) $ \file -> do
              batch <- case writeFileFormat of
                BinaryFormat -> do
                  r <- B.readFile file
                  case deserializeGen (Proxy :: Proxy Compact) r of
                    Left parseError -> die 3 $ "Parse error: " <> parseError
                    Right result -> return result
                JsonFormat -> do
                  (batches, schema_id) <- fileToBatches file
                  buildJsonBatch dbSchema (schemaIdToOpts schema_id) batches
              _ <- Glean.writeSendAndRebaseQueue queue batch $
                \_ -> writeTQueue logMessages $ "Wrote " <> file
              atomically (flushTQueue logMessages) >>= mapM_ putStrLn
              return ()
        atomically (flushTQueue logMessages) >>= mapM_ putStrLn
  
      writeBatches writeFiles Write{useLocalCache = False, ..}
        | LocalOrRemote.BackendEnv env <- LocalOrRemote.backendKind backend = do
          logMessages <- newTQueueIO
          case writeFileFormat of
            BinaryFormat ->
              stream writeMaxConcurrency (forM_ writeFiles) $ \file -> do
                r <- B.readFile file
                case deserializeGen (Proxy :: Proxy Compact) r of
                  Left parseError -> die 3 $ "Parse error: " <> parseError
                  Right batch -> do
                    void $ syncWriteDatabase env writeRepo batch
                    atomically $ writeTQueue logMessages $ "Wrote " <> file
                    atomically (flushTQueue logMessages) >>= mapM_ putStrLn
            JsonFormat ->
              stream writeMaxConcurrency (forM_ writeFiles) $ \file -> do
                (batches, schema_id) <- fileToBatches file
                syncWriteJsonBatch
                  env writeRepo batches (schemaIdToOpts schema_id)
                atomically $ writeTQueue logMessages $ "Wrote " <> file
                atomically (flushTQueue logMessages) >>= mapM_ putStrLn
  
      writeBatches writeFiles Write{useLocalCache = False, ..} = do
        logMessages <- newTQueueIO
        let settings = sendAndRebaseQueueSendQueueSettings sendQueueSettings
        Glean.withSendQueue backend writeRepo settings $ \queue ->
          stream writeMaxConcurrency (forM_ writeFiles) $ \file -> do
            case writeFileFormat of
              BinaryFormat -> do
                r <- B.readFile file
                case deserializeGen (Proxy :: Proxy Compact) r of
                  Left parseError -> die 3 $ "Parse error: " <> parseError
                  Right batch ->
                    atomically $ Glean.writeSendQueue queue batch $ \_ ->
                      writeTQueue logMessages $ "Wrote " <> file
              JsonFormat -> do
                (batches, _) <- fileToBatches file
                atomically $ Glean.writeSendQueueJson queue batches $ \_ ->
                  writeTQueue logMessages $ "Wrote " <> file
            atomically (flushTQueue logMessages) >>= mapM_ putStrLn
            return ()
        atomically (flushTQueue logMessages) >>= mapM_ putStrLn
      writeBatchDescriptors locations Write{..} = do
        logMessages <- newTQueueIO
        let settings = sendAndRebaseQueueSendQueueSettings sendQueueSettings
        Glean.withSendQueue backend writeRepo settings $ \queue ->
          stream writeMaxConcurrency (forM_ locations) $ \location -> do
            let batchDescriptor = Thrift.BatchDescriptor {
              Thrift.batchDescriptor_location = location,
              Thrift.batchDescriptor_format = batchFormat writeFileFormat}
            atomically
              $ Glean.writeSendQueueDescriptor queue batchDescriptor $ \_ ->
                writeTQueue logMessages $ "Wrote " <> Text.unpack location
            atomically (flushTQueue logMessages) >>= mapM_ putStrLn
            return ()
        atomically (flushTQueue logMessages) >>= mapM_ putStrLn
  
      batchFormat writeFileFormat = case writeFileFormat of
        BinaryFormat -> Thrift.BatchFormat_Binary
        JsonFormat -> Thrift.BatchFormat_JSON
      resultToFailure Right{} = Nothing
      resultToFailure (Left err) = Just (show err)
  
    withService evb cfgAPI svc c =
      -- Ensure the queue timeout is short to avoid OOMing the server
      LocalOrRemote.withBackend evb cfgAPI svc Nothing settings $ \backend ->
        runCommand evb cfgAPI backend c
      where
        settings (clientConf, serviceOpts) =
          (clientConf, serviceOpts{queueTimeout = Just 0.1})
  