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.Regression.Snapshot.Query (runQuery) where
  
  import Control.Exception
  import Control.Monad
  import Control.Monad.IO.Class (liftIO)
  import qualified Data.Aeson as Aeson
  import qualified Data.ByteString.UTF8 as UTF8
  import Data.Default
  import qualified Data.HashMap.Strict as HashMap
  import Data.List
  import Data.Maybe
  import qualified Data.Map as Map
  import Data.Ord
  import Data.Text (Text)
  import qualified Data.Text as Text
  import qualified Data.Text.Encoding as Text
  import Compat.Prettyprinter (pretty)
  import qualified Data.Yaml as Yaml
  import System.Exit
  import qualified Text.JSON as JSON
  import qualified Data.Vector as Vector
  
  import Util.JSON.Pretty ()
  
  import qualified Glean
  import qualified Glean.Types as Thrift
  import Glean.Regression.Snapshot.Transform
  
  newtype TQ = TQ { fromTQ :: [(Text, Aeson.Value)] }
  
  instance Aeson.FromJSON TQ where
    parseJSON = Aeson.withArray "transform" $ go [] . Vector.toList
      where
        go ts [] = return $ TQ $ reverse ts
        go ts (Aeson.String s : v : xs) = go ((s,v) : ts) xs
        go _ _ = fail "invalid transform"
  
  data Query = Query
    { queryText :: UTF8.ByteString
    , queryRecursive :: Bool
    , queryPerf :: Bool
      -- ^ Also gather performance results, compare against test.perf
    , queryMaxResults :: Maybe Int
    , queryTransforms :: TQ
    , queryDerive :: Bool
    }
  
  instance Aeson.FromJSON Query where
    parseJSON = Aeson.withObject "query" $ \v -> do
      queryText <- Text.encodeUtf8 <$> v Aeson..: "query"
      queryRecursive <- v Aeson..:! "recursive" Aeson..!= True
      queryPerf <- v Aeson..:! "perf" Aeson..!= False
      queryMaxResults <- v Aeson..:! "max_results"
      queryTransforms <- v Aeson..:! "transform" Aeson..!= TQ []
      queryDerive <- v Aeson..:! "derive" Aeson..!= False
      return Query{..}
  
  runQuery
    :: Glean.Backend e => e
    -> Thrift.Repo
    -> Transforms
    -> FilePath
    -> IO (
      String, -- The results as transformed, pretty-printed JSON
      Maybe String -- Query performance stats, also pretty-printed JSON
    )
  runQuery backend repo xforms qfile =
    handle (\(Glean.BadQuery b) -> throwIO $ Glean.BadQuery ("in " <> Text.pack qfile <> ": " <> b)) $ do
    r <- Yaml.decodeFileEither qfile
    case r of
      Left err -> die
        $ qfile ++ ": invalid query - " ++ Yaml.prettyPrintParseException err
      Right Query{..} -> do
        transform <-
          fmap (foldl' (.) id) $ mapM getTransform $ fromTQ queryTransforms
        let mkQuery k cont = def
              { Thrift.userQuery_query = queryText
              , Thrift.userQuery_options = Just def
                  { Thrift.userQueryOptions_no_base64_binary = True
                  , Thrift.userQueryOptions_syntax = Thrift.QuerySyntax_ANGLE
                  , Thrift.userQueryOptions_expand_results = True
                  , Thrift.userQueryOptions_recursive = queryRecursive
                  , Thrift.userQueryOptions_max_results = fromIntegral <$> k
                  , Thrift.userQueryOptions_continuation = cont
                  , Thrift.userQueryOptions_collect_facts_searched = queryPerf
                  , Thrift.userQueryOptions_store_derived_facts = queryDerive
                  }
              }
  
            collect acc perf k cont = do
              res <- liftIO $ Glean.userQuery backend repo $ mkQuery k cont
              let facts = Thrift.userQueryResults_facts res
                  remaining = subtract (length facts) <$> k
                  perf' = perf <> Thrift.userQueryResults_stats res
              if isJust cont && maybe True (>0) remaining
                then collect (facts ++ acc) perf' remaining cont
                else return (facts ++ acc, perf')
  
        (facts, perf) <- collect [] Nothing queryMaxResults Nothing
  
        let generatedTag = '@':"generated"
  
        perfString <- if queryPerf
          then do
            Thrift.SchemaInfo{..} <- Glean.getSchemaInfo backend (Just repo)
              def { Thrift.getSchemaInfo_omit_source = True }
            return $ Just $ show $ pretty $ JSON.JSObject $ JSON.toJSObject
              [ (generatedTag, JSON.JSNull)
              , ("facts_searched", JSON.JSObject $ JSON.toJSObject $
                sortBy (comparing fst)
                  [ (show (pretty pred), JSON.JSRational False (fromIntegral n))
                  | Just facts <- [Thrift.userQueryStats_facts_searched =<< perf]
                  , (pid,n) <- Map.toList facts
                  , Just pred <- [Map.lookup pid schemaInfo_predicateIds]
                  ])
              , ("full_scans", JSON.JSArray $ sort $
                  [ JSON.JSString $ JSON.toJSString $ show (pretty pred)
                  | Just scans <- [Thrift.userQueryStats_full_scans <$> perf]
                  , pred <- scans
                  ])
              ]
          else
            return Nothing
  
        resultString <- fmap
          (show
            . pretty
            . JSON.JSArray
            . (JSON.JSString (JSON.toJSString generatedTag) :)
            . transform
            . map nukeIds)
          $ forM facts $ \fact -> case JSON.decode (UTF8.toString fact) of
            JSON.Error err -> die $ "invalid fact in response: " ++ err
            JSON.Ok (value :: JSON.JSValue) -> return value
  
        return (resultString, perfString)
  
    where
      getTransform (name, arg)
        | Just (Transform f) <- HashMap.lookup name xforms =
            case Aeson.fromJSON arg of
              Aeson.Success x -> return $ f x
              Aeson.Error msg ->
                die $
                  "invalid argument to transform '" ++ Text.unpack name ++ "': "
                  ++ msg
        | otherwise = die $ "unknown transform '" ++ Text.unpack name ++ "'"
  