X-Git-Url: http://gitweb.michael.orlitzky.com/?a=blobdiff_plain;f=src%2FMain.hs;h=ed28f919b7e349e782b8e62e4cacc7f2ecfa120c;hb=000318e2a4b56772a9ef219a13e960acea6453b6;hp=2ca3a83816843db631ed6a5eef56366384b4c242;hpb=dd4d3adf452823702f3376d1869b07e39f03bff2;p=dead%2Fhtsn-import.git diff --git a/src/Main.hs b/src/Main.hs index 2ca3a83..ed28f91 100644 --- a/src/Main.hs +++ b/src/Main.hs @@ -1,60 +1,235 @@ -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE OverloadedStrings #-} - +{-# LANGUAGE DoAndIfThenElse #-} +{-# LANGUAGE NoMonomorphismRestriction #-} module Main where ---import Control.Monad.IO.Class ( liftIO ) -import Data.Maybe ( listToMaybe ) ---import Database.Groundhog.TH ---import Database.Groundhog.Sqlite -import Text.Show.Pretty ( ppShow ) +import Control.Arrow ( (&&&), (>>^), arr, returnA ) +import Control.Concurrent ( threadDelay ) +import Control.Exception ( SomeException, catch ) +import Control.Monad ( when ) +import Database.Groundhog.Generic ( runDbConn ) +import Database.Groundhog.Sqlite ( + withSqliteConn ) +import Database.Groundhog.Postgresql ( + withPostgresqlConn ) +import Data.Monoid ( (<>) ) +import Network.Services.TSN.Logging ( init_logging ) +import System.Console.CmdArgs ( def ) +import System.Directory ( removeFile ) +import System.Exit ( exitWith, ExitCode (ExitFailure) ) +import System.IO.Error ( catchIOError ) import Text.XML.HXT.Core ( - SysConfigList, - XmlPickler, - no, - runX, - withPreserveComment, - withRemoveWS, - withValidate, - xpickle, - xunpickleDocument, - yes ) - -import qualified TSN.Injuries as Injuries ( Message ) -import qualified TSN.InjuriesDetail as InjuriesDetail ( Message ) - -parse_opts :: SysConfigList -parse_opts = - [ withPreserveComment no, - withRemoveWS yes, - withValidate no ] - -parse_file :: XmlPickler a => FilePath -> IO (Maybe a) -parse_file path = - fmap listToMaybe $ - runX ( xunpickleDocument xpickle parse_opts path ) - --- main_sql :: IO () --- main_sql = --- withSqliteConn "foo.sqlite3" $ runDbConn $ do --- runMigration defaultMigrationLogger $ do --- migrate (undefined :: Injuries.Message) --- migrate (undefined :: Injuries.Listing) - --- msg :: Maybe Injuries.Message <- liftIO $ parse_file --- "test/xml/injuriesxml.xml" --- case msg of --- Nothing -> return () --- Just m -> do --- msg_id <- insert m --- return () + ArrowXml, + IOStateArrow, + XmlTree, + (>>>), + (/>), + getAttrl, + getText, + hasName, + readDocument, + runX ) + +import Backend ( Backend(..) ) +import CommandLine ( get_args ) +import Configuration ( Configuration(..), merge_optional ) +import ConnectionString ( ConnectionString(..) ) +import ExitCodes ( exit_no_xml_files ) +import qualified OptionalConfiguration as OC ( + OptionalConfiguration ( xml_files ), + from_rc ) +import Network.Services.TSN.Report ( + report_info, + report_error ) +import TSN.DbImport ( DbImport(..), ImportResult(..) ) +import qualified TSN.XML.Heartbeat as Heartbeat ( verify ) +import qualified TSN.XML.Injuries as Injuries ( Listing ) +import qualified TSN.XML.InjuriesDetail as InjuriesDetail ( PlayerListing ) +import qualified TSN.XML.News as News ( News ) +import qualified TSN.XML.Odds as Odds ( Odds ) +import Xml ( DtdName(..), parse_opts ) + + +-- | This is where most of the work happens. This function is called +-- on every file that we would like to import. It determines which +-- importer to use based on the DTD, processes the file, and then +-- returns whether or not any records were imported. If the file was +-- processed, the number of records imported is returned (wrapped in +-- a Just). Otherwise, if the file was not processed, 'Nothing' is +-- returned. +-- +-- Since we are already in arrow world with HXT, the +-- 'import_with_dtd' function is lifted to an 'Arrow' as well with +-- 'arr'. This prevents us from having to do a bunch of unwrapping +-- and rewrapping with the associated error checking. +-- +import_file :: Configuration -- ^ A configuration object needed for the + -- 'backend' and 'connection_string'. + + -> FilePath -- ^ The path of the XML file to import. + + -> IO (Maybe Int) -- ^ If we processed the file, Just the number + -- of records imported. Otherwise, Nothing. +import_file cfg path = do + results <- parse_and_import `catch` exception_handler + case results of + [] -> do + -- One of the arrows returned "nothing." + report_error $ "Unable to determine DTD for file " ++ path ++ "." + return Nothing + (ImportFailed errmsg:_) -> do + report_error errmsg + return Nothing + (ImportSkipped infomsg:_) -> do + -- We processed the message but didn't import anything. Return + -- "success" so that the XML file is deleted. + report_info infomsg + return $ Just 0 + (ImportSucceeded count:_) -> do + report_info $ "Successfully imported " ++ (show count) ++ + " records from " ++ path ++ "." + return $ Just count + (ImportUnsupported infomsg:_) -> do + -- For now we return "success" for these too, since we know we don't + -- support a bunch of DTDs and we want them to get deleted. + report_info infomsg + return $ Just 0 + where + -- | This will catch *any* exception, even the ones thrown by + -- Haskell's 'error' (which should never occur under normal + -- circumstances). + exception_handler :: SomeException -> IO [ImportResult] + exception_handler e = do + report_error (show e) + let errdesc = "Failed to import file " ++ path ++ "." + -- Return a nonempty list so we don't claim incorrectly that + -- we couldn't parse the DTD. + return [ImportFailed errdesc] + + -- | An arrow that reads a document into an 'XmlTree'. + readA :: IOStateArrow s a XmlTree + readA = readDocument parse_opts path + + -- | An arrow which parses the doctype "SYSTEM" of an 'XmlTree'. + -- We use these to determine the parser to use. + dtdnameA :: ArrowXml a => a XmlTree DtdName + dtdnameA = getAttrl >>> hasName "doctype-SYSTEM" /> getText >>^ DtdName + + -- | Combine the arrows above as well as the function below + -- (arrowized with 'arr') into an IO action that does everything + -- (parses and then runs the import on what was parsed). + -- + -- The result of runX has type IO [IO ImportResult]. We thus use + -- bind (>>=) and sequence to combine all of the IOs into one + -- big one outside of the list. + parse_and_import :: IO [ImportResult] + parse_and_import = + runX (readA >>> (dtdnameA &&& returnA) >>> (arr import_with_dtd)) + >>= + sequence + + -- | Takes a ('DtdName', 'XmlTree') pair and uses the 'DtdName' + -- to determine which function to call on the 'XmlTree'. + import_with_dtd :: (DtdName, XmlTree) -> IO ImportResult + import_with_dtd (DtdName dtd,xml) + -- We special-case the heartbeat so it doesn't have to run in + -- the database monad. + | dtd == "Heartbeat.dtd" = Heartbeat.verify xml + | otherwise = + -- We need NoMonomorphismRestriction here. + if backend cfg == Postgres + then withPostgresqlConn cs $ runDbConn $ importer xml + else withSqliteConn cs $ runDbConn $ importer xml + where + -- | Pull the real connection String out of the configuration. + cs :: String + cs = get_connection_string $ connection_string cfg + + importer + | dtd == "injuriesxml.dtd" = + dbimport (undefined :: Injuries.Listing) + + | dtd == "Injuries_Detail_XML.dtd" = + dbimport (undefined :: InjuriesDetail.PlayerListing) + | dtd == "newsxml.dtd" = + dbimport (undefined :: News.News) + + | dtd == "Odds_XML.dtd" = undefined + + | otherwise = \_ -> do -- Dummy arg simplifies the other cases. + let infomsg = + "Unrecognized DTD in " ++ path ++ ": " ++ dtd ++ "." + return $ ImportUnsupported infomsg + + +-- | Entry point of the program. It twiddles some knobs for +-- configuration options and then calls 'import_file' on each XML file +-- given on the command-line. +-- +-- Any file successfully processed is then removed, and we're done. +-- main :: IO () main = do - msg1 :: Maybe Injuries.Message <- parse_file "test/xml/injuriesxml.xml" - putStr $ ppShow msg1 + rc_cfg <- OC.from_rc + cmd_cfg <- get_args + + -- Merge the config file options with the command-line ones, + -- prefering the command-line ones. + let opt_config = rc_cfg <> cmd_cfg + + -- Update a default config with any options that have been set in + -- either the config file or on the command-line. We initialize + -- logging before the missing parameter checks below so that we can + -- log the errors. + let cfg = (def :: Configuration) `merge_optional` opt_config + init_logging (log_level cfg) (log_file cfg) (syslog cfg) + + -- Check the optional config for missing required options. + when (null $ OC.xml_files opt_config) $ do + report_error "No XML files given." + exitWith (ExitFailure exit_no_xml_files) + + -- We don't do this in parallel (for now?) to keep the error + -- messages nice and linear. + results <- mapM (import_file cfg) (OC.xml_files opt_config) + + -- Zip the results with the files list to find out which ones can be + -- deleted. + let result_pairs = zip (OC.xml_files opt_config) results + let victims = [ (p,c) | (p, Just c) <- result_pairs ] + let imported_count = sum $ map snd victims + report_info $ "Imported " ++ (show imported_count) ++ " records total." + mapM_ ((kill True) . fst) victims + + where + -- | Wrap these two actions into one function so that we don't + -- report that the file was removed if the exception handler is + -- run. + remove_and_report path = do + removeFile path + report_info $ "Removed processed file " ++ path ++ "." + + -- | Try to remove @path@ and potentially try again. + kill try_again path = + (remove_and_report path) `catchIOError` exception_handler + where + -- | A wrapper around threadDelay which takes seconds instead of + -- microseconds as its argument. + thread_sleep :: Int -> IO () + thread_sleep seconds = do + let microseconds = seconds * (10 ^ (6 :: Int)) + threadDelay microseconds - msg2 :: Maybe InjuriesDetail.Message <- parse_file - "test/xml/Injuries_Detail_XML.xml" - putStr $ ppShow msg2 + -- | If we can't remove the file, report that, and try once + -- more after waiting a few seconds. + exception_handler :: IOError -> IO () + exception_handler e = do + report_error (show e) + report_error $ "Failed to remove imported file " ++ path ++ "." + if try_again then do + report_info "Waiting 5 seconds to attempt removal again..." + thread_sleep 5 + kill False path + else + report_info $ "Giving up on " ++ path ++ "."