#!/usr/bin/perl
#
# jobarray-script-update -- Find and update job script data for jobs
#                           in a job array.
#
# Copyright 2009 University of Tennessee
# Portions copyright 2006 Ohio Supercomputer Center
#
# License:  GNU GPL v2; see ../COPYING for details.
# Revision info:
# $HeadURL$
# $Revision$
# $Date$
#
use strict;
use POSIX;
use DBI;

sub unique
{
    # Seen at http://www.perl.com/doc/FAQs/FAQ/oldfaq-html/Q5.4.html
    my %saw;
    @saw{@_} = ();
    return sort keys %saw;  # remove sort if undesired
}

#
# Main program begins here
#

# my system name in DB
my $system=`hostname`;
chomp($system);

my $dbhost="localhost";
my $dbuser="pbsacct";
my $dbpasswd="pbsRroxor";
my $dbname="pbsacct";
my $dbtable="Jobs";

my $start_date;
my $end_date;

while ( $ARGV[0] =~ /^-.*/ )
  {
    if ( $ARGV[0] eq "-d" )
      {
        $dbname=$ARGV[1];
        shift(@ARGV);
      }
    elsif ( $ARGV[0] eq "-h" )
      {
        $system=$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);
      }
    elsif ( $ARGV[0] eq "-E" )
     {
        $end_date = $ARGV[1];
        shift(@ARGV);
     }
    elsif ( $ARGV[0] eq "-S" )
     {
        $start_date = $ARGV[1];
        shift(@ARGV);
     }

    shift(@ARGV);
  }

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

my @scriptlesssubjob = ();
# "old" jobarray format
#my $initselect = "SELECT jobid FROM Jobs WHERE system='$system' AND jobid REGEXP '^[0-9]+-[0-9]+..*\$' AND script IS NULL";
# "new" jobarray jobid format
my $initselect = "SELECT jobid FROM Jobs WHERE system='$system' AND jobid REGEXP '^[0-9]+\[[0-9]+\]..*\$' AND script IS NULL";
{
    $initselect .= " AND submit_date>='".$start_date."'";
}
if ( defined($end_date) )
{
    $initselect .= " AND submit_date<='".$end_date."'";
}
my $sth = $dbh->prepare($initselect)
    or die "Could not prepare query:  ".$dbh::errstr;
my $n = $sth->execute()
    or die "Could not execute query:  ".$sth::errstr;
while ( my @data = $sth->fetchrow_array() )
{
    push(@scriptlesssubjob,$data[0]);
}

my @jobid = ();
foreach ( @scriptlesssubjob )
{
    my $job = $_;
    # "old" jobarray format
    #$job =~ s/^([0-9]+)-[0-9]+\./$1\./;
    # "new" jobarray jobid format
    $job =~ s/^([0-9]+)\[[0-9]+\]\./$1\./;
    push(@jobid,$job);
}
@jobid = unique(@jobid);

my $selectquery = "SELECT script FROM Jobs WHERE system=? AND jobid=? AND script IS NOT NULL";
if ( defined($start_date) )
{
    $selectquery .= " AND submit_date>='".$start_date."'";
}
if ( defined($end_date) )
{
    $selectquery .= " AND submit_date<='".$end_date."'";
}
my $select = $dbh->prepare($selectquery) or die "Could not prepare query:  ".$dbh::errstr;
my $updatequery = "UPDATE Jobs SET script=? WHERE system=? AND jobid REGEXP ? AND script IS NULL";
if ( defined($start_date) )
{
    $updatequery .= " AND submit_date>='".$start_date."'";
}
if ( defined($end_date) )
{  
    $updatequery .= " AND submit_date<='".$end_date."'";
}
my $update = $dbh->prepare($updatequery) or die "Could not prepare query:  ".$dbh::errstr;
foreach ( @jobid )
{
	
    $select->execute($system,$_)
	or die "Could not execute query:  ".$select::errstr;
    my $script = ($select->fetchrow_array())[0];
    if ( $script ne "" )
    {
	my $taskptn = $_;
	# "old" jobarray format
	#$taskptn =~ s/^([0-9]+)\.(.+)/\^$1-\[0-9\]\+\.$2\$/;
	# "new" jobarray jobid format
	$taskptn =~ s/^([0-9]+)\.(.+)/\^$1\\\[\[0-9\]\+\\\]\.$2\$/;
	$update->execute($script,$system,$taskptn)
	    or die "Could not execute query:  ".$update::errstr;
    }
}