#!/usr/bin/ruby
# upgrade-wallaby-db:  
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

require 'mrg/grid/config/shell'

module Mrg
  module Grid
    module Config
      class UpgradeWallabyDb < ::Mrg::Grid::Config::Shell::Command
        # opname returns the operation name; for "wallaby foo", it
        # would return "foo".
        def self.opname
          "upgrade-wallaby-db"
        end
      
        # description returns a short description of this command, suitable 
        # for use in the output of "wallaby help commands".
        def self.description
          "Upgrade the wallaby database"
        end
      
        def init_option_parser
          @force = false
          OptionParser.new do |opts|
            opts.banner = "Usage:  wallaby #{self.class.opname}\n#{self.class.description}"
      
            opts.on("-h", "--help", "displays this message") do
              puts @oparser
              exit
            end

            opts.on("-f", "--force", "force upgrade") do
              @force = true
            end
          end
        end
      
        def act
          update_features = {"Master"=>{"cmd"=>"REMOVE",
                                        "params"=>{"SEC_DEFAULT_INTEGRITY"=>0, "SEC_DEFAULT_ENCRYPTION"=>0}},
                            "HACentralManager"=>{"cmd"=>"ADD",
                                                 "params"=>{"CONDOR_HOST"=>0}}, 
                            "TriggerService"=>{"cmd"=>"ADD",
                                               "params"=>{"ENABLE_ABSENT_NODES_DETECTION"=>"True", "DC_DAEMON_LIST"=>">= TRIGGERD"}},
                            "EC2"=>{"cmd"=>"ADD",
                                    "params"=>{"EC2_GAHP_LOG"=>"/tmp/EC2GahpLog.$(USERNAME)", "GRIDMANAGER_MAX_SUBMITTED_JOBS_PER_RESOURCE_EC2"=>"20", "EC2_GAHP"=>"$(SBIN)/ec2_gahp", "GRIDMANAGER_MAX_SUBMITTED_JOBS_PER_RESOURCE_AMAZON"=>"20"}},
                            "ConsoleCollector"=>{"cmd"=>"ADD",
                                                 "params"=>{"COLLECTOR.PLUGINS"=>">= $(LIB)/plugins/MgmtCollectorPlugin-plugin.so"}},
                            "ConsoleExecuteNode"=>{"cmd"=>"ADD",
                                                   "params"=>{"STARTD.PLUGINS"=>">= $(LIB)/plugins/MgmtStartdPlugin-plugin.so"}},
                            "ConsoleMaster"=>{"cmd"=>"ADD",
                                              "params"=>{"MASTER.PLUGINS"=>">= $(LIB)/plugins/MgmtMasterPlugin-plugin.so"}},
                            "ConsoleNegotiator"=>{"cmd"=>"ADD",
                                                  "params"=>{"NEGOTIATOR.PLUGINS"=>">= $(LIB)/plugins/MgmtNegotiatorPlugin-plugin.so"}},
                            "ConsoleScheduler"=>{"cmd"=>"ADD",
                                                  "params"=>{"SCHEDD.PLUGINS"=>">= $(LIB)/plugins/MgmtScheddPlugin-plugin.so"}},
          }

          fobj = store.getFeature("BaseDBVersion")
          if fobj != nil:
            db_ver = (fobj.params["BaseDBVersion"].to_s rescue 0)
            temp = db_ver.split('.')
            db_major = temp[0].to_i
            db_minor = temp[1].to_i
          else
            db_major = 0
            db_minor = 0
          end

          if db_major > 1 or (db_major <= 1 and db_minor >= 13)
            puts "The database is up to date"
          else
            t = Time.now.utc
            @snap_name = "Database upgrade automatically generated snapshot at #{t} -- #{((t.tv_sec * 1000000) + t.tv_usec).to_s(16)}"

            puts "Creating pre-upgrade snapshot named #{@snap_name}"
            if store.makeSnapshot(@snap_name) == nil
               exit!(1, "Failed to create pre-upgrade snapshot.  Database upgrade aborted")
            end

            puts "Upgrading database"

            # Add new params
            puts "Adding new Parameters"
            add_param("EC2_GAHP_LOG", 0, false, "Location of the EC2 Gahp log files", "String", false, "/tmp/EC2GahpLog.$(USERNAME)")

            add_param("GRIDMANAGER_MAX_SUBMITTED_JOBS_PER_RESOURCE_EC2", 0, false, "Amazon EC2 has a hard limit of 20 concurrently running instances.  This parameter limits the number of EC2 resources", "Integer", false, "20")

            add_param("EC2_GAHP", 0, false, "The location of the EC2 Gahp binary", "String", false, "$(SBIN)/ec2_gahp")

            add_param("GRIDMANAGER_MAX_SUBMITTED_JOBS_PER_RESOURCE_AMAZON", 0, false, "Amazon EC2 has a hard limit of 20 concurrently running instances.  This parameter limits the number of amazon resources", "Integer", false, "20")

            add_param("TimeToWait", 0, false, "Expression used to indicate the time since the last job was run", "String", false, "(2 * $(HOUR))")

            add_param("ShouldHibernate", 0, false, "Expression used to determine if the machine should hibernate due to inactivity", "String", false, "( (KeyboardIdle > $(StartIdleTime)) && $(CPUIdle) && ($(StateTimer) > $(TimeToWait)) )")

            add_param("HIBERNATE", 0, false, "An expression that represents a lower power state. When this state name evaluates to a valid non-NONE state, the machine will be put into the specified low power state", "String", false, 'ifThenElse( $(ShouldHibernate), "RAM", "NONE" )')

            add_param("OFFLINE_LOG", 0, false, "The full path and file name of a file that stores machine ClassAds for every hibernating machine", "String", false, "$(SPOOL)/OfflineLog")

            add_param("OFFLINE_EXPIRE_ADS_AFTER", 0, false, "The number of seconds specifying the lifetime of the persistent machine ClassAd representing a hibernating machine", "Integer", false, "28800")

            add_param("UNHIBERNATE", 0, false, "A boolean expression that specifies when an offline machine should be woken up", "String", false, "MachineLastMatchTime =!= UNDEFINED")

            add_param("ROOSTER", 0, true, "The location of the Rooster binary", "String", false, "$(LIBEXEC)/condor_rooster")

            add_param("ROOSTER_INTERVAL", 0, false, "The number of seconds between checks for offline machines that should be woken up", "Integer", false, "300")

            add_param("ROOSTER_MAX_UNHIBERNATE", 0, false, "The maximum number of machines to wake up per cycle.  A value of 0 means unlimited", "Integer", false, "0")

            add_param("ROOSTER_UNHIBERNATE_RANK", 0, false, "A ClassAd expression specifying which machines should be woken up first in a given cycle. Higher ranked machines are woken first", "String", false, "Mips*Cpus")

            add_param("ROOSTER_UNHIBERNATE", 0, false, "A boolean expression that specifies which machines should be woken up", "String", false, "Offline && Unhibernate")

            add_param("ROOSTER_WAKEUP_CMD", 0, false, "A string representing the command line to invoke by condor_rooster in order to wake up a machine", "String", false, "\"$(BIN)/condor_power -d -i -s 255.255.255.255\"")

            add_param("HIBERNATE_CHECK_INTERVAL", 0, false, "The number of seconds specifying how often the condor_startd checks to see if the machine is ready to enter a low power state", "Integer", false, "0")

            add_param("ROOSTER_SUBNET_MASK", 0, false, "The subnet used by condor_rooster when waking up a machine", "String", true, "")

            add_param("ENABLE_ABSENT_NODES_DETECTION", 0, true, "Determines whether the condor_triggerd will look for absent nodes", "Boolean", false, "TRUE")

            add_param("QMF_BROKER_AUTH_MECH", 0, true, "The mechanism to use when authenticating with a QMF broker", "String", true, "")

            add_param("QMF_BROKER_USERNAME", 0, true, "The username to use when authenticating with a QMF broker", "String", true, "")

            add_param("QMF_BROKER_PASSWORD_FILE", 0, true, "The location of a file containing a password to use when authenticating with a QMF broker", "String", true, "")

            add_param("WALLABY_FORCE_RESTART", 0, true, "A dummy param used to force all daemons to restart", "String", false, "")

            add_param("WALLABY_FORCE_CONFIG_PULL", 0, false, "A dummy param used to force a configuration pull", "String", false, "")

            add_param("SHARED_PORT", 0, true, "The Shared Port binary", "String", false, "$(LIBEXEC)/condor_shared_port")

            add_param("USE_SHARED_PORT", 0, false, "Specifies whether a condor process should rely on the Shared Port for receiving incoming connections", "Boolean", false, "False")

            add_param("SHARED_PORT_DEBUG", 0, false, "The debugging output that the Shared Port will produce in its log", "String", false, "")

            add_param("DAEMON_SOCKET_DIR", 0, false, "Specifies the directory where Unix versions of condor daemons will create named sockets so that incoming connections can be forwarded to them by the Shared Port.  Write access to this directory grants permission to receive connections through the Shared Port", "String", false, "$(RUN)")

            add_param("QUERY_SERVER.QUERY_SERVER_LOG", 0, false, "The location of the Aviary Query Server log file", "String", false, "$(LOG)/QueryServerLog")

            add_param("QUERY_SERVER.SCHEDD_NAME", 0, false, "The name of the Scehduler that the Aviary Query Server is working with", "String", true, "")

            add_param("QUERY_SERVER.QUERY_SERVER_DEBUG", 0, false, "The debugging output the Aviary Query Server will produce in its log file", "String", false, "D_ALWAYS")

            add_param("QUERY_SERVER.HTTP_PORT", 0, true, "The port the QueryServer listens on", "Integer", false, "9091")

            add_param("QUERY_SERVER_ARGS", 0, true, "Args to append to the Aviary Query Server on daemon startup", "String", false, "")

            add_param("QUERY_SERVER", 0, true, "The location of the Aviary Query Server binary", "String", false, "$(SBIN)/aviary_query_server")

            add_param("SCHEDD.HTTP_PORT", 0, false, "Port the Aviary Schedd plugin listens on", "Integer", false, "9090")

            add_param("QUERY_SERVER.HISTORY_INTERVAL", 0, false, "The number of seconds between polls of the HISTORY file", "Integer", false, "120")

            add_param("WSFCPP_HOME", 0, false, "The root of the axis2 deployment", "String", false, "/var/lib/condor/aviary/axis2.xml")

            if db_minor <= 1 and db_minor < 5
              add_param("BaseDBVersion", 0, false, "The version of the base database", "String", false, "0")
            end

            # Add new features
            puts "Adding new Features"
            add_feature("PowerManagementNode",
                        {"ShouldHibernate"=>"( (KeyboardIdle > $(StartIdleTime)) && $(CPUIdle) && ($(StateTimer) > $(TimeToWait)) )",
                         "HIBERNATE_CHECK_INTERVAL"=>"300",
                         "HIBERNATE"=>'ifThenElse( $(ShouldHibernate), "RAM", "NONE" )',
                         "TimeToWait"=>"(2 * $(HOUR))"},
                         ["ExecuteNode"], [], ["Collector", "Negotiator", "Scheduler"])

            add_feature("PowerManagementCollector",
                        {"OFFLINE_LOG"=>"$(SPOOL)/OfflineLog",
                         "OFFLINE_EXPIRE_ADS_AFTER"=>"28800",
                         "VALID_SPOOL_FILES"=>"$(VALID_SPOOL_FILES), OfflineLog"},
                         ["Collector"], [], [])

            add_feature("PowerManagementSubnetManager",
                        {"ROOSTER_MAX_UNHIBERNATE"=>"0",
                         "ROOSTER"=>"$(LIBEXEC)/condor_rooster",
                         "UNHIBERNATE"=>"MachineLastMatchTime =!= UNDEFINED",
                         "DAEMON_LIST"=>">= ROOSTER",
                         "ROOSTER_UNHIBERNATE_RANK"=>"Mips*Cpus",
                         "ROOSTER_UNHIBERNATE"=>"Offline && Unhibernate",
                         "ROOSTER_SUBNET_MASK"=>0,
                         "ROOSTER_INTERVAL"=>"300",
                         "ROOSTER_WAKEUP_CMD"=>"\"$(BIN)/condor_power -d -i -s $(ROOSTER_SUBNET_MASK)\""},
                         [], [], ["PowerManagementNode"])

            add_feature("SharedPort",
                        {"SHARED_PORT"=>"$(LIBEXEC)/condor_shared_port",
                         "USE_SHARED_PORT"=>"True",
                         "DAEMON_LIST"=>">= SHARED_PORT",
                         "DAEMON_SOCKET_DIR"=>"$(RUN)",
                         "SHARED_PORT_DEBUG"=>""},
                         [], [], [])

            add_feature("Axis2Home", {"WSFCPP_HOME"=>"/var/lib/condor/aviary/axis2.xml"}, [], [], [])

            add_feature("AviaryScheduler",
                        {"SCHEDD.PLUGINS"=>">= $(LIB)/plugins/AviaryScheddPlugin-plugin.so",
                         "SCHEDD.HTTP_PORT"=>"9090"},
                         ["Axis2Home"], [], [])

            add_feature("QueryServer",
                        {"QUERY_SERVER.QUERY_SERVER_LOG"=>"$(LOG)/QueryServerLog",
                         "DAEMON_LIST"=>">= QUERY_SERVER",
                         "QUERY_SERVER.QUERY_SERVER_DEBUG"=>"D_ALWAYS",
                         "QUERY_SERVER.HTTP_PORT"=>"9091",
                         "QUERY_SERVER_ARGS"=>"-f",
                         "QUERY_SERVER"=>"$(SBIN)/aviary_query_server",
                         "QUERY_SERVER.HISTORY_INTERVAL"=>"120"},
                         ["Axis2Home"], ["Master", "JobQueueLocation"], [])

            if fobj == nil
              add_feature("BaseDBVersion", {"BaseDBVersion"=>"1.13"}, [], [], [""])
            else
              obj = store.getFeature("BaseDBVersion")
              obj.modifyParams("REPLACE", {"BaseDBVersion"=>"1.13"}, {})
            end

            # Add new subsystems
            puts "Adding new Subsystems"
            add_subsystem("rooster", ["ROOSTER", "ROOSTER_INTERVAL", "ROOSTER_MAX_UNHIBERNATE", "ROOSTER_UNHIBERNATE", "ROOSTER_UNHIBERNATE_RANK", "ROOSTER_WAKEUP_CMD"])
            add_subsystem("shared_port", ["SHARED_PORT", "USE_SHARED_PORT", "SHARED_PORT_DEBUG", "DAEMON_SOCKET_DIR"])

            add_subsystem("query_server", ["HISTORY", "QUERY_SERVER", "QUERY_SERVER.HISTORY_INTERVAL", "QUERY_SERVER.HTTP_PORT", "QUERY_SERVER.QUERY_SERVER_DEBUG", "QUERY_SERVER.QUERY_SERVER_LOG", "QUERY_SERVER.SCHEDD_NAME", "QUERY_SERVER_ARGS", "SPOOL"])

            # Update existing features
            puts "Updating existing Features"
            update_features.each_pair do |key, value|
              obj = store.getFeature(key)
              if obj != nil
                obj.modifyParams(value["cmd"], value["params"], {})
              else
                if @force == true
                  puts "Warning: Failed to update Feature #{key}"
                else
                  puts "Error updating Feature #{key}"
                  upgrade_failed
                end
              end
            end

            # Update existing params
            puts "Updating existing Parameters"
            obj = store.getParam("AMAZON_GAHP")
            if obj != nil
              obj.setRequiresRestart(false)
            else
              if @force == true
                puts "Warning: Failed to update Parameter AMAZON_GAHP"
              else
                puts "Error updating Parameter AMAZON_GAHP"
                upgrade_failed
              end
            end

            # Update existing subsystems
            puts "Updating existing Subsystems"
            subsys_list = ["collector", "job_server", "master", "negotiator", "schedd", "startd", "triggerd"]
            subsys_list.each do |sn|
              obj = store.getSubsys(sn)
              if obj != nil
                if sn == "triggerd"
                  obj.modifyParams("ADD", %w{QMF_BROKER_AUTH_MECH QMF_BROKER_USERNAME QMF_BROKER_PASSWORD_FILE ENABLE_ABSENT_NODES_DETECTION}, {})
                elsif sn == "master"
                  obj.modifyParams("ADD", %w{QMF_BROKER_AUTH_MECH QMF_BROKER_USERNAME QMF_BROKER_PASSWORD_FILE WALLABY_FORCE_CONFIG_PULL WALLABY_FORCE_RESTART}, {})
                elsif sn == "job_server"
                  obj.modifyParams("ADD", %w{QMF_BROKER_AUTH_MECH QMF_BROKER_USERNAME QMF_BROKER_PASSWORD_FILE SPOOL}, {})
                else
                  obj.modifyParams("ADD", %w{QMF_BROKER_AUTH_MECH QMF_BROKER_USERNAME QMF_BROKER_PASSWORD_FILE}, {})
                end
              else
                if @force == true
                  puts "Warning: Failed to update Subsystem #{sn}"
                else
                  puts "Error updating Subsystem #{sn}"
                  upgrade_failed
                end
              end
            end
            puts "Database upgraded successfully"
          end
    
          return 0
        end

        def add_param(name, level, restart, desc, kind, change, default)
          obj = store.addParam(name)
          if obj != nil
            obj.setVisibilityLevel(level)
            obj.setRequiresRestart(restart)
            obj.setDescription(desc)
            obj.setKind(kind)
            obj.setMustChange(change)
            obj.setDefault(default)
          else
            if @force == true
              puts "Warning: Failed to add Parameter #{name}"
            else
              puts "Error adding Parameter #{name}.  Reverting database"
              upgrade_failed
            end
          end
        end

        def add_feature(name, params, inc, dep, con)
          obj = store.addFeature(name)
          if obj != nil
            obj.modifyParams('REPLACE', params, {})
            obj.modifyIncludedFeatures('REPLACE', inc, {})
            obj.modifyDepends('REPLACE', dep, {})
            obj.modifyConflicts('REPLACE', con, {})
          else
            if @force == true
              puts "Warning: Failed to add Feature #{name}"
            else
              puts "Error adding Feature #{name}.  Reverting database"
              upgrade_failed
            end
          end
        end

        def add_subsystem(name, params)
          obj = store.addSubsys(name)
          if obj != nil
            obj.modifyParams('REPLACE', params, {})
          else
            if @force == true
              puts "Warning: Failed to add Subsystem #{name}"
            else
              puts "Error adding Subsystem #{name}.  Reverting database"
              upgrade_failed
            end
          end
        end

        def upgrade_failed
          store.loadSnapshot(@snap_name)
          exit!(1, "Database upgrade failed")
        end
      end
    end
  end
end

# Process args
wallaby_opts = []
cmd_opts = []
op = OptionParser.new do |opts|
  opts.on("-h", "--help", "shows this message") do
    puts op
    exit
  end

  opts.on("-H", "--host HOSTNAME", "qpid broker host (default localhost)") do |h|
    wallaby_opts << "--host" << h
  end

  opts.on("-p", "--port NUM", "qpid broker port (default 5672)") do |num|
    wallaby_opts << "--port" << num
  end

  opts.on("-U", "--user NAME", "qpid username") do |name|
    wallaby_opts << "--user" << name
  end

  opts.on("-P", "--password PASS", "qpid password") do |pass|
    wallaby_opts << "--password" << pass
  end

  opts.on("-M", "--auth-mechanism PASS", %w{ANONYMOUS PLAIN GSSAPI}, "authentication mechanism (#{%w{ANONYMOUS PLAIN GSSAPI}.join(", ")})") do |mechanism|
    wallaby_opts << "--auth-mechanism" << mechanism
  end

  opts.on("-f", "--force", "force the upgrade to proceed") do
    cmd_opts << "--force"
  end
end

args = ARGV
begin
  op.parse!(args)
rescue OptionParser::InvalidOption
  puts op
  exit
rescue OptionParser::InvalidArgument => ia
  puts ia.message
  puts op
  exit
end

::Mrg::Grid::Config::Shell::register_command(::Mrg::Grid::Config::UpgradeWallabyDb)
::Mrg::Grid::Config::Shell::main(wallaby_opts + ["upgrade-wallaby-db"] + cmd_opts + args)
