Refactoring

Author: simonmar

glean/db/Glean/Database/Config.hs

@@ -8,35 +8,41 @@
 
 {-# LANGUAGE ApplicativeDo, CPP #-}
 module Glean.Database.Config (
+  -- * DataStore
   DataStore(..),
   fileDataStore,
   tmpDataStore,
   memoryDataStore,
+
+  -- * Config, and options parser
+  options,
   Config(..),
+  DebugFlags(..),
+
+  -- * Finding and parsing the schema
   ServerConfig.SchemaLocation(..),
   showSchemaLocation,
   schemaLocation,
   schemaLocationOption,
-  DebugFlags(..),
-  options,
   processSchema,
   processSchemaCached,
   processOneSchema,
   SchemaIndex(..),
   schemaForSchemaId,
   ProcessedSchema(..),
-  schemaSourceIndexConfig,
   catSchemaFiles,
   schemaLocationToSource,
-  schemaLocationFiles,
-  schemaSourceFilesFromDir,
-  schemaSourceDir,
-  schemaSourceFile,
-  schemaSourceIndexFile,
   parseSchemaDir,
   parseSchemaIndex,
-  -- testing
+  loadSchemaIndex,
+
+  -- * Testing only
+
+  -- if you're using these somewhere other than a test,
+  -- you should probably be using something from the API above instead.
   processSchemaForTesting,
+  schemaLocationFiles,
+  schemaSourceDir,
 ) where
 
 import Control.Exception
@@ -426,6 +432,16 @@ schemaLocation cfg server_cfg = do
     other ->
       return other
 
+-- | Find and load the SchemaIndex, taking into account command line
+-- flags, the ServerConfig, and the cfgSchemaHook. This is a
+-- convenience function used by the CLI.
+loadSchemaIndex :: ConfigProvider c => Config -> c -> IO SchemaIndex
+loadSchemaIndex cfg cfgAPI = do
+  serverConfig <- ThriftSource.load cfgAPI (cfgServerConfig cfg)
+  loc <- schemaLocation cfg serverConfig
+  let (schemaSource, _) = cfgSchemaHook cfg loc
+  ThriftSource.load cfgAPI schemaSource
+
 showSchemaLocation :: ServerConfig.SchemaLocation -> String
 showSchemaLocation = \case
   ServerConfig.SchemaLocation_dir d -> "dir:" <> Text.unpack d

glean/tools/gleancli/GleanCLI.hs

@@ -50,7 +50,6 @@ import qualified Glean.Types as Thrift
 import Glean.Impl.ConfigProvider
 import Glean.Util.ConfigProvider
 import Glean.Util.ShellPrint
-import qualified Glean.Util.ThriftSource as ThriftSource
 import Glean.Shell
 
 import GleanCLI.Backup
@@ -661,11 +660,9 @@ instance Plugin WriteSerializedInventoryCommand where
         let cfg = case svc of
               Glean.Local cfg _ -> cfg
               Glean.Remote{} -> def
-        serverConfig <- ThriftSource.load cfgAPI (cfgServerConfig cfg)
-        loc <- GleanDB.schemaLocation cfg serverConfig
-        let (schemaSource, _) = GleanDB.cfgSchemaHook cfg loc
-        index <- ThriftSource.load cfgAPI schemaSource
-        dbSchema <- newDbSchema Nothing index selector readWriteContent def
+        index <- GleanDB.loadSchemaIndex cfg cfgAPI
+        dbSchema <- newDbSchema Nothing index selector
+          readWriteContent def
         return $ Inventory.serialize $ schemaInventory dbSchema
 
     inventory <- case writeSerializedInventoryFrom of

glean/tools/gleancli/GleanCLI/Merge.hs

@@ -25,6 +25,9 @@ import Util.OptParse
 import Util.Log ( logInfo )
 import Thrift.Protocol.Compact
 
+import qualified Glean.Database.Config as GleanDB
+import Glean.Database.Schema
+import Glean.Database.Schema.Types
 import Glean.LocalOrRemote (loadDbSchema)
 import qualified Glean.LocalOrRemote as Glean
 import Glean.Types
@@ -34,11 +37,6 @@ import qualified Glean.RTS.Foreign.FactSet as FactSet
 import Glean.RTS.Foreign.Define (DefineFlags(..), defineBatch)
 import qualified Glean.RTS.Foreign.Inventory as Inventory
 import Glean.RTS.Foreign.Ownership
-import qualified Glean.Database.Config as GleanDB
-import Glean.Database.Schema.Types
-import qualified Glean.Util.ThriftSource as ThriftSource
-import Glean.Database.Config (cfgServerConfig)
-import Glean.Database.Schema
 
 import GleanCLI.Types
 import GleanCLI.Common (dbOpts, fileFormatOpt, FileFormat (..))
@@ -51,7 +49,7 @@ data MergeCommand = MergeCommand
   , mergeFileSize :: Int
   , mergeOutDir :: FilePath
   , fileFormat :: FileFormat
-  , inventorySource :: Either Repo FilePath
+  , inventorySource :: Maybe (Either Repo FilePath)
   }
 
 inventoryOpt :: Parser FilePath
@@ -61,6 +59,11 @@ inventoryOpt = strOption $
     help ("Inventory created with --write-serialized-inventory and which "
       <> "was used to create binary format files of facts")
 
+data SchemaData
+  = HaveSchema DbSchema
+  | HaveInventory Inventory.Inventory
+  | HaveNothing
+
 instance Plugin MergeCommand where
   parseCommand = commandParser "merge" (progDesc "Merge fact files") $ do
     mergeFiles <- many $ strArgument (
@@ -77,29 +80,40 @@ instance Plugin MergeCommand where
       long "output" <>
       metavar "DIR" <>
       help "Destination directory for the merged fact files"
-    inventorySource <- Left <$> dbOpts <|> Right <$> inventoryOpt
+    inventorySource <- optional $ Left <$> dbOpts <|> Right <$> inventoryOpt
     fileFormat <- fileFormatOpt BinaryFormat
     return MergeCommand{..}
 
   withService evb cfgAPI svc MergeCommand{..} = do
-    (inventory, dbSchema) <- case inventorySource of
-      Left repo -> do
+    -- * --file-format=binary requires an inventory, from either --db or
+    --   --inventory.
+    -- * --file-format=json requires a dbSchema, from either --db or
+    --   schema_id in the JSON file. It's pointless to use --inventory with
+    --   --file-format=json, because the inventory isn't used.
+    schemaData <- case inventorySource of
+      Just (Left repo) -> do
         dbSchema <- Glean.withBackendWithDefaultOptions
           evb cfgAPI svc Nothing $ \backend -> do
             loadDbSchema backend repo
-        logInfo("db's schema ID is: "  <> show(schemaId dbSchema))
-        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)
+        logInfo("db's schema ID is: "  <> show (schemaId dbSchema))
+        return (HaveSchema dbSchema)
+      Just (Right mergeInventory) -> do
+        case fileFormat of
+          JsonFormat -> do
+            hPutStrLn stderr $
+              "Warning: --inventory is ignored with --file-format=json"
+            return HaveNothing
+          _ -> do
+          inventory <- Inventory.deserialize <$> B.readFile mergeInventory
+          return (HaveInventory inventory)
+      Nothing ->
+        return HaveNothing
     createDirectoryIfMissing True mergeOutDir
     hSetBuffering stderr LineBuffering
     outputs <- newIORef []
     expandedMergeFiles <- mapM expandFile mergeFiles
-    stream 1 (merge fileFormat inventory dbSchema $ concat expandedMergeFiles)
+    stream 1
+      (merge fileFormat schemaData (concat expandedMergeFiles))
       (writeToFile outputs)
       -- stream overlaps writing with reading
     files <- readIORef outputs
@@ -146,84 +160,87 @@ instance Plugin MergeCommand where
               " (" <> show (B.length batch) <> " bytes)"
             B.writeFile out batch
 
-      merge fileFormat inventory dbSchema files write = loop 0 0 Nothing files
+      merge fileFormat schemaData files write =
+        loop 0 0 Nothing schemaData files
         where
-        read :: FilePath -> Int -> FactSet -> IO FactOwnership
-        read file size factSet = do
+        read
+          :: FilePath
+          -> Int
+          -> FactSet
+          -> SchemaData
+          -> IO (FactOwnership, SchemaData)
+        read file size factSet schemaData = 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
+
+          let
+            define batch inventory newSchemaData = do
+              subst <- defineBatch factSet inventory batch
+                DefineFlags {
+                  trustRefs = True,
+                  ignoreRedef = True }
+              own <- substOwnership subst $ FactOwnership (batch_owned batch)
+              return (own, newSchemaData)
+
+          case fileFormat of
             JsonFormat -> do
               (batches, schema_id_file) <- fileToBatches file
-              case dbSchema of
-                Nothing -> do
+              dbSchema <- case schemaData of
+                HaveSchema dbSchema
+                  | Just schema_id <- schema_id_file,
+                    schema_id /= schemaId dbSchema ->
+                    throwIO $ ErrorCall $
+                      "Schema ID mismatch:\ndb: "
+                      <> show (schemaId dbSchema) <> "\nvs\nFile: "
+                      <> file <> " has " <> show schema_id
+                  | otherwise -> return dbSchema
+                _otherwise -> do
+                  -- the first time (only), we load the schema we'll
+                  -- be using to parse the JSON files.
                   schema_id <- case schema_id_file of
-                    Nothing -> throwIO $ ErrorCall "missing schema ID"
+                    Nothing -> throwIO $ ErrorCall $
+                      file <> ": missing schema_id"
                     Just id -> return id
                   let cfg = case svc of
                         Glean.Local cfg _ -> cfg
                         Glean.Remote{} -> def
-                  serverConfig <- ThriftSource.load cfgAPI (cfgServerConfig cfg)
-                  loc <- GleanDB.schemaLocation cfg serverConfig
-                  let (schemaSource, _) = GleanDB.cfgSchemaHook cfg loc
-                  index <- ThriftSource.load cfgAPI schemaSource
-                  dbSchema <- newDbSchema Nothing index
-                    (SpecificSchemaId schema_id) readWriteContent def
-                  buildJsonBatch dbSchema
-                    (schemaIdToOpts $ Just (schemaId dbSchema)) batches
-                Just schema -> do
-                  if Just(schemaId schema) == schema_id_file then
-                    logInfo(
-                      "Schema matches with db schema. Merging data from "
-                      <>  file
-                      )
-                  else
-                    throwIO $ ErrorCall $
-                        "ERROR - ABORTING MERGE\nSchema ID mismatch:\ndb: "
-                        <> show (schemaId schema) <> "\nvs\nFile: "
-                        <> file <> " has " <> show schema_id_file
-                  let getSchemaId theschema = Just(schemaId theschema) in
-                    buildJsonBatch schema
-                        (schemaIdToOpts $ getSchemaId schema) batches
+                  index <- GleanDB.loadSchemaIndex cfg cfgAPI
+                  newDbSchema Nothing index (SpecificSchemaId schema_id)
+                    readWriteContent def
+              batch <- buildJsonBatch dbSchema
+                (schemaIdToOpts $ Just (schemaId dbSchema)) batches
+              define batch (schemaInventory dbSchema) (HaveSchema dbSchema)
 
             BinaryFormat -> do
+              inventory <- case schemaData of
+                HaveSchema dbSchema -> return (schemaInventory dbSchema)
+                HaveInventory inventory -> return inventory
+                HaveNothing -> throwIO $ ErrorCall $
+                  "--file-format=binary requires either --inventory, " <>
+                  "--db, or --db-name and --db-instance"
               bytes <- B.readFile file
               case deserializeCompact bytes of
                 Left err -> throwIO $ ErrorCall $
                   "failed to deserialize " <> file <> ": " <> err
-                Right batch -> return batch
-          subst <- defineBatch factSet inventory batch
-            DefineFlags {
-              trustRefs = True,
-              ignoreRedef = True }
-          substOwnership subst $ FactOwnership (batch_owned batch)
-
-        loop !_ _ Nothing [] = return ()
-        loop !n _ (Just set) [] = write (n, Right set)
-        loop !n currentSize acc (f : files) = do
+                Right batch ->
+                  define batch inventory (HaveInventory inventory)
+
+        loop !_ _ Nothing _ [] = return ()
+        loop !n _ (Just set) _ [] = write (n, Right set)
+        loop !n currentSize acc schema (f : files) = do
           size <- fromIntegral <$> getFileSize f
           if size > mergeFileSize && fileFormat == BinaryFormat
             -- just copy huge binary files
             then do
               write (n, Left f)
-              loop (n+1) currentSize acc files
+              loop (n+1) currentSize acc schema files
             else do
               (factSet, ownership) <- case acc of
                 Nothing -> (,[]) <$> FactSet.new lowestFid
                 Just facts -> return facts
-              owners <- read f size factSet
+              (owners, schema) <- read f size factSet schema
               newSize <- factSetSize factSet
               let facts = (factSet, owners : ownership)
               (n, acc) <- if newSize > mergeFileSize
                 then do write (n, Right facts); return (n+1, Nothing)
                 else return (n, Just facts)
-              loop n newSize acc files
+              loop n newSize acc schema files