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.Query (QueryCommand) where
  
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Char8 as B8
  import Data.Default
  import Data.Int (Int64)
  import Data.Text (Text)
  import Data.Maybe (isJust, isNothing)
  import qualified Data.Text as Text
  import qualified Data.Text.Encoding as Text
  import qualified Data.Text.IO as Text
  import Options.Applicative
  import System.IO
  
  import Thrift.Protocol.JSON (serializeJSON)
  import Util.IO
  import Util.OptParse
  
  import Glean
  import Glean.Types as Thrift
  
  import GleanCLI.Common
  import GleanCLI.Types
  import Control.Monad (when)
  
  data QueryCommand
    = Query
        { repoSpec :: Either Text Repo
        , query :: String
        , recurse :: Bool
        , queryPageOptions :: PageOptions
        , limitFacts :: Maybe Int
        , output :: Maybe FilePath
        , statsOutput :: Maybe FilePath
        , timeout :: Maybe Int64
        , omitResults :: Bool
        , debugOutput :: Maybe FilePath
        , debugBytecode :: Bool
        , debugIr :: Bool
        , profileOutput :: Bool
        }
  
  instance Plugin QueryCommand where
    parseCommand =
      commandParser "query" (progDesc "Execute an Angle query") $ do
        repoSpec <- Left <$> dbNameOpt <|> Right <$> dbSlash
        queryPageOptions <- pageOpts
        recursive <- switch $ long "recursive" <> hidden
          <> help "fetch nested facts (slower)"
        expand <- switch $ long "expand"
          <> help "recursively expand nested facts (slower)"
        limitFacts <- optional $ option auto
          ( long "limit"
          <> metavar "FACTS"
          <> help "maximum number of facts to query (default: no limit)"
          )
        output <- optional $ strOption
          ( long "output"
          <> short 'o'
          <> metavar "FILE"
          <> help "output the facts to a file"
          )
        statsOutput <- optional $ strOption
          ( long "stats"
          <> metavar "FILE"
          <> help "output stats to a file ('-' for stdout)"
          )
        profileOutput <- switch $ long "profile"
          <> help (
            "get full profiling information; " <>
            "use with --stats to include facts_searched")
        timeout <- optional $ option auto
          ( long "timeout"
          <> metavar "MILLISECONDS"
          <> help "Override the default query timeout"
          )
        omitResults <- switch $ long "omit-results"
          <> help (
            "don't print results; " <>
            "use with --stats to get a count of results")
        debugOutput <- optional $ strOption
          ( long "debug"
          <> metavar "FILE"
          <> help "output debug info to a file ('-' for stdout)"
          )
        debugBytecode <- switch $ long "debug-bytecode"
          <> help "dump the generated bytecode"
        debugIr <- switch $ long "debug-ir"
          <> help "dump the generated IR"
        query <- strArgument
          ( metavar "QUERY"
          <> help "query to execute ('@file' to read from file, '-' for stdin)"
          )
        return $
          let recurse = recursive || expand in
          Query{..}
  
    runCommand _ _ backend Query{..} = do
      when (isNothing statsOutput && profileOutput) $
        die 3 "Invalid options: --profile needs to be used together with --stats"
      query_bytes <- case query of
        "-" -> B.hGetContents stdin
        '@':path -> B.readFile path
        _ -> return $ Text.encodeUtf8 $ Text.pack query
  
      repo <- case repoSpec of
        Left name -> Glean.getLatestRepo backend name
        Right repo -> return repo
  
      let with_output f = case output of
            Just path -> withFile path WriteMode f
            Nothing -> f stdout
  
          with_stats_output f = case statsOutput of
            Just path ->
              (if path == "-" then ($ stdout) else withFile path WriteMode) $
              \out -> f $ B8.hPutStrLn out . maybe "{}" serializeJSON
            Nothing -> f $ const $ return ()
  
          with_debug_output f
            | debugBytecode || debugIr = case debugOutput of
                Just path
                  | path /= "-" -> withFile path WriteMode display
                _ -> display stdout
            | otherwise = f $ const $ return ()
            where
              display out = f $ Text.hPutStrLn out . Text.unlines
  
      with_output $ \h_out ->
        with_stats_output $ \print_stats ->
        with_debug_output $ \print_debug -> do
      let subtract_limit Nothing _ = Just Nothing
          subtract_limit (Just m) n
            | m > n = Just $ Just (m-n)
            | otherwise = Nothing
  
          loop prev_stats cont limit = do
            UserQueryResults{..} <- Glean.userQuery backend repo def
              { userQuery_query = query_bytes
              , userQuery_encodings = [UserQueryEncoding_json
                  def{ userQueryEncodingJSON_expand_results = recurse }]
              , userQuery_options = Just def
                  { userQueryOptions_max_results =
                      case (limit, fromIntegral <$> pageFacts queryPageOptions) of
                        (Just m, Just n) -> Just $ m `min` n
                        (x,y) -> x <|> y
                  , userQueryOptions_max_bytes =
                      Just $ fromIntegral $ pageBytes queryPageOptions
                  , userQueryOptions_max_time_ms = timeout
                  , userQueryOptions_continuation = cont
                  , userQueryOptions_syntax = QuerySyntax_ANGLE
                  , userQueryOptions_recursive = recurse
                  , userQueryOptions_omit_results = omitResults
                  , userQueryOptions_collect_facts_searched =
                      isJust statsOutput && profileOutput
                  , userQueryOptions_debug = case cont of
                      Nothing -> QueryDebugOptions
                        { queryDebugOptions_bytecode = debugBytecode
                        , queryDebugOptions_ir = debugIr
                        , queryDebugOptions_pred_has_facts = False
                        }
                      Just _ -> def
                  }
              }
            let stats = prev_stats <> userQueryResults_stats
            n <- case userQueryResults_results of
              UserQueryEncodedResults_json UserQueryResultsJSON{..} -> do
                mapM_ (B8.hPutStrLn h_out) userQueryResultsJSON_facts
                return $ length userQueryResultsJSON_facts
              _ -> die 1 "error: unexpected results encoding"
            print_debug userQueryResults_diagnostics
            case userQueryResults_continuation of
              Just new_cont
                | Just new_limit <- subtract_limit limit $ fromIntegral n ->
                    loop stats (Just new_cont) new_limit
              _ ->  do
                print_stats stats
                return ()
      loop Nothing Nothing $ fromIntegral <$> limitFacts
  