Merge improvements 2/2: Adding read of schema from JSON

Summary:
`glean merge` officially supports supplying both databases to be merged and JSON files + and inventory file to be merged.

However, if supplying raw JSONs, there's no `dbSchema` available, and the merge fails. This diff handles the raw JSON case by going off to configerator and reading the schema definitions there.

It relies on a `schema_id` field being present in the JSONs to lookup the correct schema.

Reviewed By: iamirzhan

Differential Revision: D73030289

fbshipit-source-id: 9c6ea68acfaf383d1b7f865342f2cc6f1aff5556

Author: simonhollis

glean.cabal.in

@@ -25,6 +25,7 @@ common fb-haskell
     default-extensions:
         BangPatterns
         BinaryLiterals
+        CPP
         DataKinds
         DeriveDataTypeable
         DeriveGeneric
@@ -133,6 +134,10 @@ common folly
         build-depends: fb-util < 0.2
         pkgconfig-depends: libfolly
 
+-- Facebook-specific features
+flag facebook
+     default: False
+
 common deps
     build-depends:
         fb-util,
@@ -1124,7 +1129,10 @@ executable glean
         GleanCLI.Restore
         GleanCLI.Write
     ghc-options: -main-is GleanCLI -with-rtsopts=-I0
+    default-extensions: CPP
     extra-libraries: stdc++
+    if flag(facebook)
+       ghc-options: -DGLEAN_FACEBOOK
     build-depends:
         glean:cli-types,
         glean:client-hs,

glean/db/Glean/Database/Config.hs

@@ -25,6 +25,7 @@ module Glean.Database.Config (
   SchemaIndex(..),
   schemaForSchemaId,
   ProcessedSchema(..),
+  schemaSourceIndexConfig,
   catSchemaFiles,
   schemaLocationToSource,
   schemaLocationFiles,

glean/db/Glean/Database/Schema.hs

@@ -16,6 +16,8 @@ module Glean.Database.Schema
   , newDbSchema
   , newMergedDbSchema
   , mkDbSchemaFromSource
+  , mkDbSchema
+  , findSchemaInIndex
   , lookupPid
   , compareSchemaPredicates
   , validateNewSchema

glean/db/Glean/Database/Util.hs

@@ -0,0 +1,48 @@
+{-
+  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.Database.Util
+  (getDbSchemaFromId
+  ) where
+
+import Control.Exception (throwIO)
+import Control.Monad.IO.Class (MonadIO, liftIO)
+import Data.Default (def)
+
+import Glean.Database.Config (DebugFlags, SchemaIndex, schemaForSchemaId)
+import Glean.Database.Schema (DbSchema, DbContent,
+  readWriteContent, newDbSchema)
+import Glean.Database.Schema.Types (SchemaSelector(..))
+import Glean.Types (SchemaId(..))
+import qualified Data.Text as Text
+
+
+-- | Get a DbSchema instance from a SchemaId
+dbSchemaFromId :: MonadIO m
+               => SchemaIndex          -- ^ Schema index to search in
+               -> SchemaId             -- ^ Schema ID to look up
+               -> DbContent            -- ^ Database content information
+               -> DebugFlags           -- ^ Debug flags
+               -> m DbSchema
+dbSchemaFromId index schemaId dbContent debug = liftIO $ do
+  case schemaForSchemaId index schemaId of
+    Nothing -> throwIO $ userError $ "Schema ID not found: " ++
+      Text.unpack (unSchemaId schemaId)
+    Just _ -> do
+      -- Create a new DbSchema from the processed schema
+      newDbSchema Nothing index (SpecificSchemaId schemaId) dbContent debug
+
+-- | Get a DbSchema instance from a SchemaId using default settings
+getDbSchemaFromId :: MonadIO m
+                  => Maybe SchemaIndex  -- ^ Schema index to search in
+                  -> SchemaId           -- ^ Schema ID to look up
+                  -> m DbSchema
+getDbSchemaFromId Nothing _ =
+  liftIO $ throwIO $ userError "No schema index provided"
+getDbSchemaFromId (Just index) schemaId =
+  dbSchemaFromId index schemaId readWriteContent def

glean/tools/gleancli/GleanCLI/Merge.hs

@@ -5,8 +5,7 @@
   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 #-}
+{-# LANGUAGE CPP, ApplicativeDo #-}
 module GleanCLI.Merge (MergeCommand) where
 
 import Control.Exception
@@ -19,10 +18,11 @@ import System.Directory
 import System.FilePath
 import System.IO
 import System.Process
-
+import Data.Text(unpack)
 import Control.Concurrent.Stream
 import Util.OptParse
-import Util.Log
+import Util.Log ( logInfo )
+import Glean.Util.ThriftSource (load)
 import Thrift.Protocol.Compact
 
 import Glean.LocalOrRemote (loadDbSchema)
@@ -35,13 +35,23 @@ import Glean.RTS.Foreign.Define (DefineFlags(..), defineBatch)
 import qualified Glean.RTS.Foreign.Inventory as Inventory
 import Glean.RTS.Foreign.Ownership
 import Glean.Database.Schema.Types
+import Glean.Database.Util (getDbSchemaFromId)
 
 import GleanCLI.Types
 import GleanCLI.Common (dbOpts, fileFormatOpt, FileFormat (..))
 import Glean.Write (fileToBatches, schemaIdToOpts)
 import Glean.Write.JSON (buildJsonBatch)
+import Glean.Database.Config (schemaSourceIndexConfig)
+import Glean.DefaultConfigs(schemaConfigPath)
 import System.Directory.Extra (listFiles)
 
+#if GLEAN_FACEBOOK
+import Configerator (
+  withConfigeratorAPI,
+  defaultConfigeratorOptions,
+ )
+#endif
+
 
 data MergeCommand = MergeCommand
   { mergeFiles :: [FilePath]
@@ -88,6 +98,9 @@ instance Plugin MergeCommand where
         return (schemaInventory dbSchema, Just dbSchema)
       Right mergeInventory -> do
         inventory <- Inventory.deserialize <$> B.readFile mergeInventory
+        -- Get the schema_id from the JSON's schema ID field
+        --- this needs to be done when the file is considered later on
+        -- dbSchema <-
         return (inventory, Nothing)
     createDirectoryIfMissing True mergeOutDir
     hSetBuffering stderr LineBuffering
@@ -136,21 +149,71 @@ instance Plugin MergeCommand where
             let batch = serializeCompact $
                   batch0 { batch_owned =
                     ownershipUnits (unionOwnership ownership) }
-            hPutStrLn stderr $ "Writing " <> out <>
-              " (" <> show (B.length batch) <> ")"
+            logInfo $ "Writing " <> out <>
+              " (" <> show (B.length batch) <> " bytes)"
             B.writeFile out batch
 
       merge fileFormat inventory dbSchema files write = loop 0 0 Nothing files
         where
         read :: FilePath -> Int -> FactSet -> IO FactOwnership
         read file size factSet = do
           logInfo $ "Reading " <> file <> " (" <> show size <> " bytes)"
+          -- Merge can take an existing db or an inventory.
+          -- A DB has a dbSchema, so we can use that to build the batches
+          -- to merge
+          -- An inventory doesn't have a dbSchema, so leads to "Nothing"
+          -- This means that we need to get the schema from somewhere
+          --- this will be the first JSON file we see
+          -- read the 'schema_id' field in that file and create a
+          -- new dbSchema intance with that using configerator's
+          -- stored default schema index
+
           batch <- case fileFormat of
             JsonFormat -> do
               case dbSchema of
-                Nothing -> throwIO $ ErrorCall $
-                  "No db schema to serialize json format file. "
-                  <> "Please specify the database"
+                Nothing ->
+                  #if GLEAN_FACEBOOK
+                  withConfigeratorAPI defaultConfigeratorOptions $ \configAPI ->
+                    -- Load the schema index from the schema directory
+                    let schemaIndexResult =
+                          Right $ schemaSourceIndexConfig schemaConfigPath in do
+                    hPutStrLn stderr $
+                        "Reading current schema index from default "
+                        <> "configerator with path " <> unpack schemaConfigPath
+
+                    -- Get the schema_id from the JSON file
+                    (batches, Just schema_id) <- fileToBatches file
+                    logInfo("Schema_ID from JSON file " <> file <> " is "
+                      <> show schema_id)
+
+                    -- Try to get the DbSchema from the schema ID string
+                    case schemaIndexResult of
+                      Left err -> do
+                        throwIO $ ErrorCall err
+                      Right schemaIndex -> do
+                        logInfo $ "Looking up schema with ID: " <>
+                            show schema_id
+                        -- Materialise the schema index from its Thrift source
+                        concreteIndex <- load configAPI schemaIndex
+                        dbSchemaResult <-
+                          (Right <$>
+                              getDbSchemaFromId (Just concreteIndex) schema_id
+                          )
+                          `catch` \(e :: SomeException) ->
+                          return $ Left $ "Failed to get schema: " <> show e
+
+                        case dbSchemaResult of
+                          Left err -> do
+                            throwIO $ ErrorCall("Error getting dbSchemaResult"
+                              <> " - exiting. Error is:\n" <> err)
+                          Right dbSchema -> do
+                            logInfo $ "Successfully loaded schema with ID: "
+                              <> show schema_id
+                            buildJsonBatch dbSchema
+                              (schemaIdToOpts $ Just(schemaId dbSchema)) batches
+                  #else
+                  return $ Left $ "Failed to get schema: "
+                  #endif
                 Just schema -> do
                   (batches, schema_id_file) <- fileToBatches file
                   if Just(schemaId schema) == schema_id_file then
@@ -167,7 +230,6 @@ instance Plugin MergeCommand where
                     buildJsonBatch schema
                         (schemaIdToOpts $ getSchemaId schema) batches
 
-
             BinaryFormat -> do
               bytes <- B.readFile file
               case deserializeCompact bytes of