packages feed

transfer-db-0.3.1.0: src/Main.hs

module Main where

import Prelude hiding (fail)

import SQL.CLI.ODBC (odbcImplementation)

import TransferDB(transferDB, TransferOptions(TransferOptions))
import Options (readPlan, Options(Options, option_Plan, option_Threads, option_Count, option_Drop))
import qualified Generator as G
import Database.TransferDB.DumpDB (dump, restore, DumpConfig(DumpConfig), RestoreConfig(RestoreConfig))

import Database.TransferDB.Commons (ProgramOptions (ProgramOptions),
                                    DBInfo (DBInfo, dbi_Datasource, dbi_User, dbi_Password, dbi_Schema))

import CorrectionPlan (generateCorrections, CorrectionConfig(CorrectionConfig))


import Control.Concurrent.STM (newTVar, atomically)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Fail (MonadFail, fail)
import Control.Monad.Trans.Maybe (MaybeT, runMaybeT)
import Control.Monad.Trans.Reader (ReaderT, runReaderT)

import Control.Logging (setLogLevel, LogLevel(LevelDebug, LevelInfo, LevelWarn, LevelError, LevelOther))

import qualified Data.Map.Strict as Map (empty)


import System.Clock (Clock(Monotonic), getTime)
                               
import System.Console.Program (single, interactive, showUsage)
import System.Console.Command (Action, Commands, Tree(Node), io, withNonOption, withOption, command)
import System.Console.Argument (Option, file, natural, string, option)

import System.IO (hPutStrLn, stdin, stdout, stderr, withFile, IOMode(ReadMode, WriteMode), Handle)


main :: IO ()
main = do
  setLogLevel LevelInfo
  (runMaybeT $ single commands) >>= maybe (fail "execution failed") (\ _ -> hPutStrLn stderr "execution finished") 

commands :: Commands (MaybeT IO)
commands = Node (command "transfer-db" "Copy data between 2 databases through ODBC, based on a transfer plan" . io $ do
                    liftIO $ hPutStrLn stderr "For usage type help"
                    interactive commands)
           [
             Node (command "run"      "Run a transfer plan"    . (withNonOption file) $ runTransferPlan) [],
             Node (command "makePlan" "Create a transfer plan based on info collected from the source db" makeSimplePlan)
                  [
                    Node (command "splitByRows"   "makes a batch at each <INT> records" . (withNonOption natural) $ makePlanByRows) [],
                    Node (command "splitByTables" "makes a batch at each <INT> tables"  . (withNonOption natural) $ makePlanByTables) []
                  ],
             Node (command "correctivePlan" "Generate a new plan and a sql script to correct problems with a previous run" . (withNonOption file) $ generateCorrectionsAction) [],
             Node (command "dump"     "dumps a database schema to a binary file" dumpAction) [],
             Node (command "restore"  "restores a database from a a binary dump file created with dump command" restoreAction) [],
             Node (command "help"     "Show usage information" . io $ liftIO $ showUsage commands) []
           ]

-- commands

runTransferPlan :: FilePath -> Action (MaybeT IO)
runTransferPlan planFile = withOption parallelConnections
  (\ p ->
     withOption runCount
     (\ c ->
        withOption dropCount
        (\ d -> io $ do
            plan <- liftIO $ readPlan planFile
            runReaderT transferDB (TransferOptions (Options {
                                                       option_Plan    = plan,
                                                       option_Threads = fromIntegral p,
                                                       option_Count   = fromIntegral c,
                                                       option_Drop    = fromIntegral d}) odbcImplementation))))
  
makeSimplePlan :: Action (MaybeT IO)
makeSimplePlan = withProgramOptions G.makeSimplePlan

makePlanByRows :: Integer -> Action (MaybeT IO)
makePlanByRows n = withProgramOptions $ G.makePlanByRows n

makePlanByTables :: Integer -> Action (MaybeT IO)
makePlanByTables n = withProgramOptions $ G.makePlanByTables n

dumpAction :: Action (MaybeT IO)
dumpAction = withOption parallelConnections
  (\ p -> withSourceDb
    (\ dbi -> withOption descriptionOption
      (\ description -> withNonOption file
        (\ dumpFile -> io $ do
            statsVar  <- liftIO $ atomically $ newTVar Map.empty -- create a TVar to collect statistics
            startTime <- liftIO $ getTime Monotonic              -- the start time, to compute the dump rate
            runReaderT dump (DumpConfig 
                              (dbi_Datasource dbi)
                              (dbi_User dbi)
                              (dbi_Password dbi)
                              (dbi_Schema dbi)
                              description
                              dumpFile
                              (fromIntegral p)
                              statsVar
                              startTime
                            )))))

restoreAction :: Action (MaybeT IO)
restoreAction = withDestinationDb
  (\ dbi -> withNonOption file
    (\ dumpFile -> io $ runReaderT restore (RestoreConfig
                                            (dbi_Datasource dbi)
                                            (dbi_User dbi)
                                            (dbi_Password dbi)
                                            (dbi_Schema dbi)
                                            dumpFile)))

generateCorrectionsAction :: (MonadIO m, MonadFail m) => FilePath -> Action m
generateCorrectionsAction file = withOption outputScript
  (\ scriptFile -> withOption outputPlan
    (\ planFile -> withOption inputLog
      (\ log -> io $ do
          plan <- liftIO $ readPlan file
          result <- liftIO $ withInputFile log
            (\ logHandle -> withOutputFile scriptFile
              (\ scriptHandle -> withOutputFile planFile
                (\ planHandle -> runMaybeT $ runReaderT generateCorrections $ CorrectionConfig plan logHandle scriptHandle planHandle)))
          maybe ((liftIO $ hPutStrLn stderr "generateCorrectionPlan failed") >> fail "generateCorrectionPlan failed") return result
      )))
  



-- command builders


withSourceDb :: (MonadIO m) => (DBInfo -> Action m) -> Action m
withSourceDb f = withOption sourceDbOption
  (\ db -> withOption sourceUserOption
    (\ user -> withOption sourcePasswordOption
      (\ password -> withOption (sourceSchemaOption user)
        (\ schema -> f $ DBInfo db user password schema))))

withDestinationDb :: (MonadIO m) => (DBInfo -> Action m) -> Action m
withDestinationDb f = withOption destDbOption
  (\ db -> withOption destUserOption
    (\ user -> withOption destPasswordOption
      (\ password -> withOption (destSchemaOption user)
        (\ schema -> f $ DBInfo db user password schema))))

withProgramOptions :: (MonadIO m) => ReaderT ProgramOptions m () -> Action m
withProgramOptions f = withSourceDb (\ sdbi -> withDestinationDb ( \ ddbi -> io $ runReaderT f $ ProgramOptions sdbi ddbi))


-- Options

sourceDbOption :: Option String
sourceDbOption = option "d" ["sd", "source-db"] string "" "ODBC datasource name for source database"

sourceUserOption :: Option String
sourceUserOption = option "u" ["su", "source-user"] string "" "User name to connect to source database"

sourcePasswordOption :: Option String
sourcePasswordOption = option "p" ["sp", "password"] string "" "Password to connect to source database"

sourceSchemaOption :: String -> Option String
sourceSchemaOption userName = option "s" ["ss", "schema"] string userName "The schema to connect to in the source database. Defaults to the user name"

destDbOption :: Option String
destDbOption = option "D" ["dd", "dest-db"] string "<datasource>" "ODBC datasource name for dest database"

destUserOption :: Option String
destUserOption = option "U" ["du", "dest-user"] string "<username>" "User name to connect to dest database"

destPasswordOption :: Option String
destPasswordOption = option "P" ["dp", "dest-password"] string "<passoword>" "Password to connect to dest database"

destSchemaOption :: String -> Option String
destSchemaOption userName = option "S" ["ds", "dest-schema"] string userName "The schema to connect to in the dest database. Defaults to the user name"

descriptionOption :: Option String
descriptionOption = option "m" ["description"] string "" "The description of the dump file"

parallelConnections :: Option Integer
parallelConnections = option "t" ["threads"] natural 0 "The numbers of parallel threads to run"

runCount :: Option Integer
runCount = option "c" ["count", "only"] natural 0 "The number of batches to be run. If 0 or not specified, all batches will be run"

dropCount :: Option Integer
dropCount = option "x" ["drop", "skip", "ignore"] natural 0 "The number of batches in the plan to skip before starting to run the plan"

inputLog :: Option FilePath
inputLog = option "i" ["log"] file "-" "The file name of the transfer-db log of the run you need to correct. If not specified or the name is '-', then the log will be read from standard input"

outputScript :: Option FilePath
outputScript = option "s" ["script"] file "-" "The name of the sql script file that will be generated. If not specified or '-' the script will be written on standard output"

outputPlan :: Option FilePath
outputPlan = option "p" ["plan"] file "-" "The name of the plan to be generated. If not specified or '-', the plan will be written on standard output"


-- | runs an IO action by passing an input file handler to it. If the file name is
-- "-", the input file will be standard input
withInputFile :: FilePath -> (Handle -> IO a) -> IO a
withInputFile fileName f = if fileName == "-"
                           then f stdin
                           else withFile fileName ReadMode f

          
-- | runs an IO action by passing an output file handler to it. If the file name is
-- "-", the output file will be the standard output
withOutputFile :: FilePath -> (Handle -> IO a) -> IO a
withOutputFile fileName f = if fileName == "-"
                            then f stdout
                            else withFile fileName WriteMode f