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 Schema.Lib
    ( withSchemaAndFacts
    , withSchema
    , withSchemaFile
    , angleQuery
    , mkBatch
    , mkAngleQuery
    , decodeResults
    ) where
  
  import Control.Monad.IO.Class (liftIO)
  import Control.Monad.Trans.Except
  import Data.ByteString (ByteString)
  import Data.Map (Map)
  import Control.Exception
  import Control.Monad
  import Data.Default
  import qualified Data.Map as Map
  import Data.Text (pack)
  import Data.Text.Encoding (encodeUtf8)
  import System.FilePath
  import System.IO.Temp
  
  import Util.Control.Exception
  
  import Glean.Angle.Types (latestAngleVersion, AngleVersion(..))
  import Glean.Backend.Types
  import Glean.Database.Open
  import Glean.Database.Config
  import Glean.Database.Test
  import Glean.Database.Types
  import Glean.Database.Schema.Types
  import Glean.Database.Schema
  import Glean.Types as Thrift
  import Glean.Write.JSON
  
  import qualified Glean.RTS as RTS
  import qualified Glean.RTS.Term as RTS
  import qualified Glean.RTS.Types as RTS
  
  angleQuery :: Env -> Repo -> ByteString -> IO UserQueryResults
  angleQuery env repo q = userQuery env repo $ mkAngleQuery q
  
  mkBatch :: PredicateRef -> [ByteString] -> JsonFactBatch
  mkBatch ref facts =
    JsonFactBatch
      { jsonFactBatch_predicate = ref
      , jsonFactBatch_facts = facts
      , jsonFactBatch_unit = Nothing
      }
  
  mkAngleQuery :: ByteString -> UserQuery
  mkAngleQuery q = def
    { userQuery_query = q
    , userQuery_options = Just def
      { userQueryOptions_syntax = QuerySyntax_ANGLE
      }
    }
  
  decodeResults
    :: RTS.Type
    -> (UserQueryResultsBin -> Map Id Fact) -- facts/nested
    -> Either BadQuery UserQueryResults
    -> IO (Either String [RTS.Value])
  decodeResults ty getFacts eitherRes = runExceptT $ do
    results <- case eitherRes of
      Left err -> fail $ "BadQuery: " <> show err
      Right r -> return r
    bin <- binResults results
    let keys = fmap fact_key $ Map.elems $ getFacts bin
    mapM (decodeAs ty) keys
    where
      binResults :: UserQueryResults -> ExceptT String IO UserQueryResultsBin
      binResults UserQueryResults{..} =
        case userQueryResults_results of
          UserQueryEncodedResults_bin b -> return b
          _ -> fail "wrong encoding"
  
      decodeAs :: RTS.Type -> ByteString -> ExceptT String IO RTS.Value
      decodeAs ty bs = do
        res <- liftIO $ try $ do
          print bs
          evaluate $ RTS.toValue (withUnknown $ RTS.repType ty) bs
        case res of
          Left e -> fail $ "unable to decode : " <> showException e
          Right val -> return val
        where
          showException (RTS.DecodingException e) = e
          -- we want to decode binary values that contain the unknown alternative
          withUnknown rep = case rep of
            RTS.ByteRep -> rep
            RTS.NatRep -> rep
            RTS.ArrayRep elty -> RTS.ArrayRep $ withUnknown elty
            RTS.TupleRep tys -> RTS.TupleRep $ fmap withUnknown tys
            RTS.SumRep tys ->
              let unknown = RTS.TupleRep [] in
              RTS.SumRep $ fmap withUnknown tys ++ [unknown]
            RTS.SetRep tys -> RTS.SetRep $ withUnknown tys
            RTS.StringRep -> rep
            RTS.PredicateRep _ -> rep
  
  -- | Used to test schema transformations
  -- Runs the callback on a read-only version of a db with the
  -- given schema and facts.
  withSchemaAndFacts
    :: [Setting]
    -> String                  -- ^ schema
    -> [JsonFactBatch]         -- ^ db contents
    -> ( Env -> Repo -> DbSchema -> IO a)
    -> IO a
  withSchemaAndFacts customSettings schema facts act =
    withSchemaFile latestAngleVersion schema $ \root file -> do
    let settings =
          [ setRoot root
          , setSchemaPath file
          ] ++ customSettings
  
    -- create db and write facts
    repo <- withEmptyTestDB settings $ \env repo -> do
        void $ syncWriteJsonBatch env repo facts Nothing
        completeTestDB env repo
        return repo
  
    -- get PredicateDetails
    dbSchema <- do
      schema <- either error return
        $ processOneSchema Nothing $ encodeUtf8 $ pack schema
      newDbSchema Nothing schema LatestSchema readWriteContent def
  
    -- open db for querying
    -- We need to open the db again because schema evolutions are
    -- only triggered when the db is read-only
    withTestEnv settings $ \env ->
      act env repo dbSchema
  
  withSchemaFile
    :: AngleVersion
    -> String
    -> (FilePath -> FilePath -> IO a)
    -> IO a
  withSchemaFile (AngleVersion version) str action = do
    withSystemTempDirectory "glean-dbtest" $ \root -> do
      let newSchemaFile = root </> "schema"
      appendFile newSchemaFile $ "version: " <> show version
      appendFile newSchemaFile str
      action root newSchemaFile
  
  withSchema
    :: AngleVersion
    -> String
    -> (Either SomeException () -> IO a)
    -> IO a
  withSchema version str action =
    withSchemaFile version str $ \root file -> do
      let settings =
            [ setRoot root
            , setSchemaPath file
            ]
      r <- tryAll $
        withEmptyTestDB settings $ \env repo ->
        withOpenDatabase env repo $ \_ ->
          return ()
  
      print (r :: Either SomeException ())
      action r
  