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.Dump
    ( dump
    ) where
  
  import Control.Exception
  import Control.Monad
  import Data.Int (Int64)
  import Text.Printf
  
  import Util.Control.Exception (catchAll)
  import Util.Log
  
  import Glean.Backend.Types
  import Glean.Backend.Local
  import Glean.Database.Schema
  import Glean.Query.JSON
  import Glean.RTS.Types
  import Glean.Types
  
  dump :: Backend b => b -> Repo -> (JsonFactBatch -> IO ()) -> IO ()
  dump backend repo withBatch = doDump =<< loadDbSchema backend repo
    where
      doDump dbSchema = do
        FactIdRange{..} <- factIdRange backend repo
        go iNVALID_ID [] 0 factIdRange_start factIdRange_finish
        where
  
        doneBatch :: Id -> [(Id,Fact)] -> IO ()
        doneBatch _ [] = return ()
        doneBatch pid facts = do
          vlog 1 (printf "%d facts for predicate %d\n" (length facts) pid)
          case lookupPid (Pid pid) dbSchema of
            Nothing -> throwIO $ ErrorCall $ "unknown predicate Id: " <> show pid
            Just details@PredicateDetails{} -> do
              jsonFacts <- forM (reverse facts) $ \(fid, Fact{..}) ->
                  factToJSON
                    True
                    mempty
                    details
                    (Fid fid)
                    fact_key
                    fact_value
                  `catchAll` \exc -> throwIO $ ErrorCall $ show exc
              withBatch JsonFactBatch
                { jsonFactBatch_predicate = predicateRef details
                , jsonFactBatch_facts = jsonFacts
                , jsonFactBatch_unit = Nothing -- TODO?
                }
  
        -- limit the size of batches we give back, so that the writer
        -- can run in constant space.
        maxBatchSize :: Id
        maxBatchSize = 10000
  
        go :: Id -> [(Id,Fact)] -> Int64 -> Id -> Id ->  IO ()
        go !currentPid facts !batchSize !nextId !finalId
          | nextId < finalId = do
              r <- queryFact backend repo nextId
              case r of
                Nothing -> go currentPid facts batchSize (nextId+1) finalId
                Just fact@Fact{..} ->
                  if fact_type == currentPid && batchSize < maxBatchSize
                    then go
                      currentPid
                      ((nextId,fact):facts)
                      (batchSize + 1)
                      (nextId+1)
                      finalId
                    else do
                      doneBatch currentPid facts
                      go fact_type [(nextId,fact)] 1 (nextId+1) finalId
          | otherwise = doneBatch currentPid facts
  