#!/usr/bin/perl
#
# job-db-update -- Insert job records into a MySQL DB
# Copyright 2006, 2007 Ohio Supercomputer Center
# Copyright 2008, 2009, 2010 University of Tennessee
#
# License:  GNU GPL v2; see ../COPYING for details.
# Revision info:
# $HeadURL$
# $Revision$
# $Date$
#
# Usage:  job-db-update [-h system] [-s dbhost] [-d dbname] [-t dbtable]
#                       [-u dbuser] [-p dbpasswd] acctfile [...acctfiles...]
#
# Can be run with no acct files listed to check DB connectivity

use POSIX;
use DBI;

sub resourcelist
  {
    local($resources) = @_;
    
    @resourcelist=split(/ /,$resources);
    foreach $field ( @resourcelist )
      {
        ($key,$value) = ($field =~ /^([^=]*)=(.*)/);
        $resource{$key}=$value;
      }
    %resource;
  }

sub max
  {
    local(@item) = @_;
    local($nitems,$maxval,$i);
    
    $nitems=$#item;
    $maxval=$item[0];
    if ( $nitems > 0 )
      {
	for ( $i=0 ; $i <= $nitems ; $i++ )
	  {
	    if ( $item[$i] > $maxval )
	      {
		$maxval=$item[$i];
	      }
	  }
      }
    $maxval;
  }

sub keyval_exists
  {
    local($key,$val,$dbtable,$dbh) = @_;
    local($sth,$n);
    
    $sth=$dbh->prepare("SELECT ".$key." FROM ".$dbtable." WHERE ".$key." = \'".$val."\';")
      or die "Could not prepare query:  ".$dbh::errstr;
    $n=$sth->execute
      or die "Could not execute query:  ".$sth::errstr;
    
    $n;
  }

#
# Main program starts here
#
$dbhost="localhost";
$dbuser="pbsacct";
$dbpasswd="pbsRroxor";
$dbname="pbsacct";
$dbtable="Jobs";
$hostname=`hostname`;

while ( $ARGV[0] =~ /^-.*/ )
  {
    if ( $ARGV[0] eq "-d" )
      {
	$dbname=$ARGV[1];
	shift(@ARGV);
      }
    elsif ( $ARGV[0] eq "-h" )
      {
	$hostname=$ARGV[1];
	shift(@ARGV);
      }
    elsif ( $ARGV[0] eq "-p" )
      {
	$dbpasswd=$ARGV[1];
	shift(@ARGV);
      }
    elsif ( $ARGV[0] eq "-s" )
      {
	$dbhost=$ARGV[1];
	shift(@ARGV);
      }
    elsif ( $ARGV[0] eq "-t" )
      {
	$dbtable=$ARGV[1];
	shift(@ARGV);
      }
    elsif ( $ARGV[0] eq "-u" )
      {
	$dbuser=$ARGV[1];
	shift(@ARGV);
      }

    shift(@ARGV);
   }

# connect to DB
$dsn = "DBI:mysql:database=$dbname;host=$dbhost;";
$dbh = DBI->connect($dsn,$dbuser,$dbpasswd) ||
  die "$0:  unable to connect to DB!\n";

# process PBS acct records
while ( defined($ARGV[0]) )
  {
    open(ACCTFILE,"<$ARGV[0]");
    
    while ( <ACCTFILE> )
      {
	chop;
	($etime,$type,$job,$resources)=split(/;/);
	undef %resource;
	%resource=&resourcelist($resources);
	
	if ( $type eq "S" || $type eq "E" )
	  {
	    ($jobid,$junk)=split(/\./,$job);
	    $pstring = $resource{"Resource_List.nodes"};
	    @node_plus_ppn = split(/\+/, $pstring);
	    $processors = 0;
	    foreach $node_plus_ppn (@node_plus_ppn)
	      {
		($nodes,$ppnstring) = split(/\:/, $node_plus_ppn);
		$processors_per = substr($ppnstring,4);
		if ( $nodes < 1 ) { $nodes = 1; }
		if ( $processors_per < 1 ) { $processors_per = 1; }
		$processors += $nodes * $processors_per;
	      }
	    $ncpus=&max($resource{"Resource_List.ncpus"},
			$resource{"resources_used.mppssp"}+4*$resource{"resources_used.mppe"});
            $ncpus=&max($ncpus,$resource{"Resource_List.size"});
	    $nprocs=&max($processors,$ncpus);
	    $mem=$resource{"resources_used.mem"};
	    $mem=~s/kb$//;
	    $vmem=$resource{"resources_used.vmem"};
	    $vmem=~s/kb$//;

	    # construct query string(s)
	    $update = "";
	    $insert = "";
	    if ( defined($resource{"start"}) &&
		 $resource{"start"}>0 )
	      {
		$uhead = "UPDATE $dbtable SET jobid = \'".$job."\'";
		$utail = " WHERE jobid = \'".$job."\';";
		$ihead = "INSERT INTO $dbtable \(jobid,system";
		$itail = " VALUES \(\'$job\',\'$hostname\'";
		if ( defined($resource{"user"}) )
		  {
		    $uhead.=",username = \'".$resource{"user"}."\'";
		    $ihead.=",username";
		    $itail.=",\'".$resource{"user"}."\'";
		  }
		if ( defined($resource{"group"}) )
		  {
		    $uhead.=",groupname = \'".$resource{"group"}."\'";
		    $ihead.=",groupname";
		    $itail.=",\'".$resource{"group"}."\'";
		  }
		if ( defined($resource{"owner"}) )
		  {
		    $submithost = $resource{"owner"};
		    $submithost =~ s/^.+\@//;
		    $uhead.=",submithost = \'".$submithost."\'";
		    $ihead.=",submithost";
		    $itail.=",\'".$submithost."\'";
		  }
		if ( defined($resource{"jobname"}) )
		  {
		    $uhead.=",jobname = \'".$resource{"jobname"}."\'";
		    $ihead.=",jobname";
		    $itail.=",\'".$resource{"jobname"}."\'";
		  }
		if ( $nprocs>0 )
		  {
		    $uhead.=",nproc = \'".$nprocs."\'";
		    $ihead.=",nproc";
		    $itail.=",\'$nprocs\'";
		  }
		if ( defined($resource{"Resource_List.nodes"}) )
		  {
		    $uhead.=",nodes = \'".$resource{"Resource_List.nodes"}."\'";
		    $ihead.=",nodes";
		    $itail.=",\'".$resource{"Resource_List.nodes"}."\'";
		  }
                if ( defined($resource{"Resource_List.feature"}) )
                  {
                    $uhead.=",feature = \'".$resource{"Resource_List.feature"}."\'";
                    $ihead.=",feature";
                    $itail.=",\'".$resource{"Resource_List.feature"}."\'";
                  }
		if ( defined($resource{"Resource_List.gattr"}) )
                  {
                    $uhead.=",gattr = \'".$resource{"Resource_List.gattr"}."\'";
                    $ihead.=",gattr";
                    $itail.=",\'".$resource{"Resource_List.gattr"}."\'";
                  }
		if ( defined($resource{"Resource_List.gres"}) )
                  {
                    $uhead.=",gres = \'".$resource{"Resource_List.gres"}."\'";
                    $ihead.=",gres";
                    $itail.=",\'".$resource{"Resource_List.gres"}."\'";
                  }
		if ( defined($resource{"queue"}) )
		  {
		    $uhead.=",queue = \'".$resource{"queue"}."\'";
		    $ihead.=",queue";
		    $itail.=",\'".$resource{"queue"}."\'";
		  }
                if ( defined($resource{"Resource_List.qos"}) )
                  {
                    $uhead.=",qos = \'".$resource{"Resource_List.qos"}."\'";
                    $ihead.=",qos";
                    $itail.=",\'".$resource{"Resource_List.qos"}."\'";
                  }
		if ( defined($resource{"ctime"}) )
		  {
		    $uhead.=",submit_ts = \'".$resource{"ctime"}."\'";
		    $uhead.=",submit_date = DATE(FROM_UNIXTIME(\'".$resource{"ctime"}."\'))";
		    $ihead.=",submit_ts";
		    $itail.=",\'".$resource{"ctime"}."\'";
		    $ihead.=",submit_date";
		    $itail.=",DATE(FROM_UNIXTIME(\'".$resource{"ctime"}."\'))";
		  }
		if ( defined($resource{"start"}) )
		  {
		    $uhead.=",start_ts = \'".$resource{"start"}."\'";
		    $uhead.=",start_date = DATE(FROM_UNIXTIME(\'".$resource{"start"}."\'))";
		    $ihead.=",start_ts";
		    $itail.=",\'".$resource{"start"}."\'";
		    $ihead.=",start_date";
		    $itail.=",DATE(FROM_UNIXTIME(\'".$resource{"start"}."\'))";
		  }
		if ( defined($resource{"end"}) )
		  {
		    $uhead.=",end_ts = \'".$resource{"end"}."\'";
		    $uhead.=",end_date = DATE(FROM_UNIXTIME(\'".$resource{"end"}."\'))";
		    $ihead.=",end_ts";
		    $itail.=",\'".$resource{"end"}."\'";
		    $ihead.=",end_date";
		    $itail.=",DATE(FROM_UNIXTIME(\'".$resource{"end"}."\'))";
		  }
		if ( defined($resource{"Resource_List.cput"}) )
		  {
		    $uhead.=",cput_req = \'".$resource{"Resource_List.cput"}."\'";
		    $ihead.=",cput_req";
		    $itail.=",\'".$resource{"Resource_List.cput"}."\'";
		  }
		if ( defined($resource{"resources_used.cput"}) )
		  {
		    $uhead.=",cput = \'".$resource{"resources_used.cput"}."\'";
		    $ihead.=",cput";
		    $itail.=",\'".$resource{"resources_used.cput"}."\'";
		  }
		if ( defined($resource{"Resource_List.walltime"}) )
		  {
		    $uhead.=",walltime_req = \'".$resource{"Resource_List.walltime"}."\'";
		    $ihead.=",walltime_req";
		    $itail.=",\'".$resource{"Resource_List.walltime"}."\'";
		  }
		if ( defined($resource{"resources_used.walltime"}) )
		  {
		    $ihead.=",walltime";
		    if ( $resource{"resources_used.walltime"} =~ /^\d+$/ ) 
		      {
                        $uhead.=",walltime = SEC_TO_TIME(".$resource{"resources_used.walltime"}.")";
		        $itail.=",SEC_TO_TIME(".$resource{"resources_used.walltime"}.")";
                      }
		    else
		      {
		        $uhead.=",walltime = \'".$resource{"resources_used.walltime"}."\'";
		        $itail.=",\'".$resource{"resources_used.walltime"}."\'";
		      }
		  }
		if ( defined($resource{"Resource_List.mem"}) )
		  {
		    $uhead.=",mem_req = \'".$resource{"Resource_List.mem"}."\'";
		    $ihead.=",mem_req";
		    $itail.=",\'".$resource{"Resource_List.mem"}."\'";
		  }
		if ( defined($resource{"resources_used.mem"}) )
		  {
		    $uhead.=",mem_kb = \'".$resource{"resources_used.mem"}."\'";
		    $ihead.=",mem_kb";
		    $itail.=",\'$mem\'";
		  }
		if ( defined($resource{"Resource_List.vmem"}) )
		  {
		    $uhead.=",vmem_req = \'".$resource{"Resource_List.vmem"}."\'";
		    $ihead.=",vmem_req";
		    $itail.=",\'".$resource{"Resource_List.vmem"}."\'";
		  }
		if ( defined($resource{"resources_used.vmem"}) )
		  {
		    $uhead.=",vmem_kb = \'".$resource{"resources_used.vmem"}."\'";
		    $ihead.=",vmem_kb";
		    $itail.=",\'$vmem\'";
		  }
		if ( defined($resource{"resources_used.mppe"}) )
		  {
		    $uhead.=",mppe = \'".$resource{"resources_used.mppe"}."\'";
		    $ihead.=",mppe";
		    $itail.=",\'".$resource{"resources_used.mppe"}."\'";
		  }		
		if ( defined($resource{"resources_used.mppssp"}) )
		  {
		    $uhead.=",mppssp = \'".$resource{"resources_used.mppssp"}."\'";
		    $ihead.=",mppssp";
		    $itail.=",\'".$resource{"resources_used.mppssp"}."\'";
		  }		
		if ( defined($resource{"exec_host"}) )
		  {
		    $uhead.=",hostlist = \'".$resource{"exec_host"}."\'";
		    $ihead.=",hostlist";
		    $itail.=",\'".$resource{"exec_host"}."\'";
		  }
		if ( defined($resource{"Exit_status"}) )
		  {
		    $uhead.=",exit_status = \'".$resource{"Exit_status"}."\'";
		    $ihead.=",exit_status";
		    $itail.=",\'".$resource{"Exit_status"}."\'";
		  }
		if ( defined($resource{"Resource_List.software"}) )
		  {
		    $uhead.=",software = \'".$resource{"Resource_List.software"}."\'";
		    $ihead.=",software";
		    $itail.=",\'".$resource{"Resource_List.software"}."\'";
		  }
                if ( defined($resource{"resources_used.energy_used"}) )
                  {
                    $uhead.=",energy = \'".$resource{"resources_used.energy_used"}."\'";
                    $ihead.=",energy";
                    $itail.=",\'".$resource{"resources_used.energy_used"}."\'";
                  }
		if ( defined($resource{"account"}) )
		  {
		    $uhead.=",account = \'".$resource{"account"}."\'";
		    $ihead.=",account";
		    $itail.=",\'".$resource{"account"}."\'";
		  }
		$ihead.="\)";
		$itail.="\);";
		
		$update = $uhead.$utail;
		$insert = $ihead.$itail;
	      }
	    if ( &keyval_exists("jobid",$job,$dbtable,$dbh)>0 )
	      {
		# This jobid already exists in the table,
		# so do an UPDATE
		$query = $update;
	      }
	    else
	      {
		# This jobid does *NOT* already exist in the table,
		# so do an INSERT
		$query = $insert;
	      }
            #print "$query\n";
	    if ( $query ne "" )
              {
	    	if ( ! $dbh->do($query) )
		  { 
	    	    if ( $query eq $insert && $update ne "" )
	    	      {
	    		# INSERT failed, which can happen due to a race
	    		# condition with job script capture; try doing
	    		# the UPDATE before squawking about it
	    		$dbh->do($update) || warn "$0:  query failed -- ".$dbh::errstr."\n".$update;
	    	      }
	    	    else
	    	      {
	    	        warn "$0:  query failed -- ".$dbh::errstr."\n".$query;
		      }
		  }
	      }
	  }
      }
    close(ACCTFILE);
    shift(@ARGV);
  }
$dbh->disconnect();