#! /usr/bin/perl -w

use strict;
use DBI;
use Getopt::Std;
use String::Random;

use vars qw/$opt_d $opt_c/;
getopt('dc:');

#---------------------------------------------------------------------------------------------------------
# variables
#
my $RandString = new String::Random;

my $run = 1;                   # controls wether to continue running

my $mypid = $$;                # my own pid
my $pid;                       # child pid
my @pids;                      # pids of all forked children
my $sql;                       # sql command
my %session;                   # all sessions
my ($session, $batch);         # actual session and batch
my $thissession;               # child session
my $lastsession;               # parent session
my ($connect, $user, $password, $database);
my %caches;
my @insertedvalues;            # values of last insert (for foreign keys)

#---------------------------------------------------------------------------------------------------------
# catch sigint to break endless loop and kill children
#
sub catch_sigint {
	my $pid;
	foreach $pid (@pids) {
		kill 2, $pid;
	}
	$run = 0;
}
$SIG{INT} = \&catch_sigint;

#---------------------------------------------------------------------------------------------------------
# first, as the parent, read the config-file to spawn children
#
$opt_c ||= 'db-stress.cfg';
open CFG, "$opt_c" or die "config-file $opt_c not found!\n";

while (my $cfgline = <CFG>) {
	next if ($cfgline =~ /^\s*$|^#/);			                             # ignore empty lines and comments
	chomp($cfgline);
	$cfgline =~ s/#.*$//;					                                     # cut comments at end of line

	if ($cfgline =~ /^\s*connect\s*=\s*(.*)/) {                        # find connect string
    $connect = $1;
	  next;
	}

  if ($cfgline =~ /^\s*user\s*=\s*(.*)/) {                           # find user
    $user = $1;
    next;
  }

  if ($cfgline =~ /^\s*password\s*=\s*(.*)/) {                       # find password
    $password = $1;
    next;
  }

  if ($cfgline =~ /^\s*database\s*=\s*(.*)/) {                       # find database
    $database = $1;
    next;
  }

	if ($cfgline =~ /^\s*session\s*=\s*(.*)/) {		                     # find sessions
		if ($sql) {
			push @{$session{$session}->{batches}->{$batch}->{commands}}, $sql;
			$sql = '';
		}
		$session = $1;
		$session{$session}->{count} = 1;
		next;
	}
	if ($cfgline =~ /^\s*count\s*=\s*(.*)/) {		                       # find no. of children
		$session{$session}->{count} = $1;
		next;
	}
	if ($cfgline =~ /^\s*batch\s*=\s*(.*)/) {		                       # find batches
		if ($sql) {
			push @{$session{$session}->{batches}->{$batch}->{commands}}, $sql;
			$sql = '';
		}
		$batch = $1;
		$session{$session}->{batches}->{$batch}->{batch} = $1;
		$session{$session}->{batches}->{$batch}->{delay} = 1;
		next;
	}
	if ($cfgline =~ /^\s*delay\s*=\s*(.*)/) {		                       # find delay
		$session{$session}->{batches}->{$batch}->{delay} = $1;
		next;
	}
	if ($cfgline =~ /^loop$/) {				                                 # finish batch and fetch rows
		push @{$session{$session}->{batches}->{$batch}->{commands}}, $sql;
		push @{$session{$session}->{batches}->{$batch}->{commands}}, "loop";
		$sql = '';
		next;
	}
	if ($cfgline =~ /^go$/) {				                                   # finish batch
		push @{$session{$session}->{batches}->{$batch}->{commands}}, $sql;
		$sql = '';
		next;
	}
	if ($cfgline !~ /^\s*commands\s*=/) {			                         # anything else, except 'commands =' is a sql command
		$sql .= "\n" . $cfgline;
		next;
	}
}

$lastsession = $session;

#-------------------------------------------------------------------------------------------------------
# fork children
#
my $fork_this = 0;                                                   # don't fork the very first session
($thissession) = sort keys %session;                                 # remember first session

SPAWN: foreach my $key (sort keys %session) {
  print "spawning session: $key (" . $session{$key}->{count} . " children)\n";
  for (my $i = 0; $i < $session{$key}->{count}; $i++) {
    if ($fork_this && (eval {$pid = fork()})) {
      print "\t\tchild: " . ($i + 1) . " ($pid) ...\n";
			$thissession = $key;
      push @pids, $pid;
      last SPAWN;
    }
    $fork_this = 1;
  }
}

#-----------------------------------------------------------------------------------------------
# connect to database
#
my $dbh;
my $sth;

print "session: $thissession; connect($connect, $user, $password)\n" if $opt_d >= 2;
$dbh = DBI->connect($connect, $user, $password) or die "could not connect: $connect, $user, $password\n";

#---- if sybase ase, use database

$dbh->do("use $database") if ($connect =~ /sybase/i);

#---------------------------------------------------------------------------------------------------
# enter loop to executes batches of this session
#
while ($run) {
	foreach $batch (keys %{$session{$thissession}->{batches}}) {

		foreach (@{$session{$thissession}->{batches}->{$batch}->{commands}}) {
			      my $sql = $_;
                              if ($sql eq 'loop') {
				my $r = 0;
				while (my @row = $sth->fetchrow()) {
					$r++;
				}
				$sth->finish();
				next;
			}
			if ($sth) {$sth->finish();}

			$sql =~ s/=randstring\((\d+?)\)/&randstring($1)/eg;

			$sql =~ s/=randregex\((.+?)\)/$RandString->randregex($1)/eg;

			$sql =~ s/=randrow\(([^, )]+)(,\s*cache)?\)/&randrow($1, $3)/eg;

			$sql =~ s/=lastinsert\((\d+?)\)/&lastinserted($1)/eg;

			print "$sql\n" if $opt_d;
			$sth = $dbh->prepare($sql);
			$sth->execute();
			&getinserted($sql) if ($sql =~ /(^\s*update)|(^\s*insert)/ims);
		}

		sleep $session{$thissession}->{batches}->{$batch}->{delay};
	}
}

if ($sth) {$sth->finish();}
if ($dbh) {$dbh->disconnect();}

#--------------------------------------------------------------------------------------
# subroutines
#--------------------------------------------------------------------------------------

#--------------------------------------------------------------------------------------
# create string of random characters
#
sub randstring() {
	my $count = shift;
	my $string = '';
	for (my $i = 0; $i < rand($count); $i++) {
		$string .= substr('abcdefghijklmnopqrstuvwxyz', rand(26), 1);
	}
	$string;
}

#---------------------------------------------------------------------------------------
# get random row of table
#
sub randrow(@) {
	my ($tablecol, $cache) = @_;
	if (!defined $caches{$tablecol} || !defined($cache)) {
		delete $caches{$tablecol} if (!defined($cache));
		my ($table, $col) = $tablecol =~ /(.+)\.(.+)/;
		my $sql = "select $col from $table";
		my $sth = $dbh->prepare($sql);
		$sth->execute();
		while (my @row = $sth->fetchrow()) {
			push @{$caches{$tablecol}}, $row[0];
		}
		$sth->finish();
	}
	
	return $caches{$tablecol}[rand(@{$caches{$tablecol}})];
}

#----------------------------------------------------------------------------------------
# get updated/inserted value 
#
sub lastinserted() {
	my $i = shift;
	$insertedvalues[$i - 1];
}

#----------------------------------------------------------------------------------------
# get inserted/updated columns
#
sub getinserted() {
	my $sql = shift;
	if ($sql =~ /^\s*insert/ims) {
		my ($s) = $sql =~ /values\((.*)\)/i;
		$s .= ",";
		@insertedvalues = $s =~ /((?:["'].*?["'])|(?:[^']*?))\s*,\s*/g;
	} else {
		my ($s) = $sql =~ /set\s+(.*)\s+(?:where.*|$)/i;
		$s = "," . $s . ",";
		@insertedvalues = $s =~ /,.*?=\s*((?:["'].*?["'])|(?:[^']*?))(?=\s*,\s*)/g;
	}
}

__END__

#-----------------------------------------------------------------------------------------------------

=head1 NAME 

db-stress - tool for generating data and putting heavy load on database servers

=head1 SYNOPSIS

B<db-stress.pl> [-d lvl] [-c cfg-file]

=head1 DESCRIPTION

The main purpose of db-stress is to put a heavy load on a database server, so you have the ability try out several tuning measurements.
Another purpose is to use it as a data generator to test replication.

You can define in a configuration file several sessions in which several batches may run against a database. Each of these batches may run in several instances. So you are able to simulate any nuber of clients you want.

db-stress is written in Perl and relies on the DBI by Tim Bunce. It has been tested with the following DBDs:

 DBD::Sybase  Linux, WinNT  Sybase SQL-Server, ASE 11.x
 DBD::ASAny   Linux         Sybase ASA 6.0
 DBD::ODBC    WinNT         Sybase ASE 11.x, Sybase ASA 6.0, 
                            MS SQL-Server 7.0

Following modules are also used:
 Getopt::Std
 String::Random

 available at http://www.cpan.org

=head1 OPTIONS

 -d n		switch on debugging of the sql-batches.
		level 1: print executed sql-commands
		level 2: print connection info

 -c file	use file instead of db-stress.cfg

=head1 CONFIGURATION FILE

Comments start with #, the rest of the line is ignored.
Blank lines are ignored.

The configuration file is separated into several sections: 
The global section at the top of the file contains connection informations. 
The remaining sections contain each a session.

Each entry in the configuration file begins with a keyword and an equal sign.

=head2 Global Section

Keywords for the global section are:

 connect	the complete connection string
		example: connect = dbi:Sybase:server=SYBASE

 user		the login
		example: user = sa

 password	the password for the login
		example: password =

 database	the database to use
		example: database = pubs2

=head2	Session Sections

Keywords for session sections:

 session	the name of the session
		example: session = 1

 count		the number of children to be forked for this
		session, default is 1
		example: count = 8

Each session can contain several batches of sql commands which are processed sequentially.
Keywords for batches:

 batch		the name of the batch
		example: batch = 1

 delay		the number of seconds execution is delayed after
		processing this batch, default is 1
		example: delay = 0

 commands	tag marking the begin of the sql commands, must
		be on a line of its own
		example: commands =

Each sql batch must be concluded by either a 'go' or a 'loop' on a line of its own.

Apart from sql the following commands and functions are available:

 loop		fetches the rows of a result set after a select
		used to generate network traffic

in a sql command:

 =randstring()	generates a string of a certain maximal length
		example: insert into xxx values('=randstring(20)')

 =randregex()	generates a string that matches a regular expression
		example: insert into xxx values('=randregex([A-Z]\d\d)')

 =randrow()	gets the value of a certain column of a random
		row from a certain table, used for inserting
		row into tables with RI. it is possible to cache
		the values in the table for faster access.
		examples:
		insert into xxx values(=randrow(tbl1.col1), 'blabla')
		insert into xxx values('=randrow(tbl1.col1)', 'blabla')
		insert into xxx values(=randrow(tbl1.col1, cache), 'blabla')

 =lastinsert()	fetches the value in a given column of the last insert or update
		if you want to use values the preceding insert/update must reside
		in a batch of its own
		examples:
		insert into x values("1", 2)
		go
		insert into y values(=lastinsert(1), lastinsert(2))
		update x set f1 = =lastinsert(1) where f1 = "9"

=head1 NOTES

Example configuration files are provided for Sybase Adaptive Server Enterprise, Adaptive Server Anywhere and ODBC. Please note the special connection syntax for ASA.
For connection information read the documentation of the DBD.

=head1 FILES

 db-stress.pl		main program
 db-stress.cfg		default configuration file
 db-stress.pod		documentation
 db-stress.1		man-page
 sybase_ase.cfg		example configuration file for
			Sybase Adaptive Server Enterprise
 sybase_asa.cfg		example configuration file for
			Sybase Adaptive Server Anywhere
 odbc_ase.cfg		example configuration file for
			Sybase ASE ODBC
 odbc_asa.cfg		example configuration file for
			Sybase ASA ODBC

=head1 TODO

 - test and support for more RDBMS
 - better error handling
 - install script

=head1 LIMITS

On systems not supporting fork(3) only the first session will be processed in only one process.

=head1 AUTHOR

-- Bernd Dulfer (bdulfer@sybase.com)

=head1 SEE ALSO

perl(1), DBI(3), String::Random(3)

=cut

#-------------------------------------------------------------------------------------------------------------------
#
# sample configuration files
#
#-------------------------------------------------------------------------------------------------------------------

# format of configuration file is described in db-stress(1)
# this is for sybase adaptive server enterprise with demo database installed
# name: sybase_ase.cfg

connect = dbi:Sybase:server=SYBASE
user = sa
password =
database = pubs2

# session is a bunch of jobs all doing the same batches
session = 1

# the count of jobs working for this session
count = 5

# a batch is a sequence of sql and db-stress commands
batch = 1
delay = 10

commands =
insert into authors (au_id, au_lname, au_fname, address, city, state, country, postalcode) values('=randstring(4)', '=randstring(40)', '=randstring(20)', '=randstring(40)', '=randstring(20)', '=randstring(2)', '=randstring(12)', '=randstring(10)')
insert into titles values('=randregex([A-Z][A-Z])xxxx', '=randstring(80)', '=randrow(titles.type, cache)', '=randrow(publishers.pub_id, cache)', rand(1234), rand(2345), rand(4321), '=randstring(200)', getdate(), 1)
go

select * from titles
loop

session = 2
count = 20
batch = 1
delay = 5
commands =
begin tran
go
insert into sales values('=randrow(stores.stor_id, cache)', '=randstring(20)', getdate())
go
insert into salesdetail values(=lastinsert(1), =lastinsert(2), '=randrow(titles.title_id, cache)', rand(100), 0)
insert into salesdetail values(=lastinsert(1), =lastinsert(2), '=randrow(titles.title_id)', rand(100), 0)
insert into salesdetail values(=lastinsert(1), =lastinsert(2), '=randrow(titles.title_id)', rand(100), 0)
commit
go

#--------------------------------------------------------------------------------------------------------------------

# format of configuration file is described in db-stress(1)
# for sybase adaptive server anywhere with a test database
# name: sybase_asa.cfg

connect = DBI:ASAny:asatest
user = UID=dba;PWD=sql;ENG=asatest
password = na
database = 

# session is a bunch of jobs all doing the same batches
session = 1

# the count of jobs working for this session
count = 1

# a batch is a sequence of sql and db-stress commands
batch = 1
delay = 3				# no. of seconds the job will wait after this batch is run, default is 1

commands =				# mark for starting of commands
insert into tbl1 values(rand(999999), '=randstring(30)')
insert into tbl1 values(rand(999999), '=randstring(30)')
insert into tbl1 values(rand(999999), '=randstring(30)')
go
select * from tbl1
loop

#-------------------------------------------------------------------------------------------------------------------

# format of configuration file is described in db-stress(1)
# this is for sybase adaptive server enterprise with demo database installed
# you have to create the DSN on your own
# name: odbc_ase.cfg

connect = dbi:ODBC:ASE1
user = sa
password =
database = pubs2

# session is a bunch of jobs all doing the same batches
session = 1

# the count of jobs working for this session
count = 5

# a batch is a sequence of sql and db-stress commands
batch = 1
delay = 5

commands =
insert into authors (au_id, au_lname, au_fname, address, city, state, country, postalcode) values('=randstring(4)', '=randstring(40)', '=randstring(20)', '=randstring(40)', '=randstring(20)', '=randstring(2)', '=randstring(12)', '=randstring(10)')
insert into titles values('=randregex([A-Z][A-Z])xxxx', '=randstring(80)', '=randrow(titles.type, cache)', '=randrow(publishers.pub_id, cache)', rand(1234), rand(2345), rand(4321), '=randstring(200)', getdate(), 1)
go

select * from titles
loop

#------------------------------------------------------------------------------------------------------------------------

# format of configuration file is described in db-stress(1)
# this is for sybase adaptive server anywhere with demo database installed
# you have to create the DSN on your own
# name: odbc_asa.cfg

connect = dbi:ODBC:ASA1
user = dba
password = sql
database = asademo

# session is a bunch of jobs all doing the same batches
session = 1

# the count of jobs working for this session
count = 5

# a batch is a sequence of sql and db-stress commands
batch = 1
delay = 5

commands =
insert into customer values(=randregex(\d\d\d\d), '=randstring(15)', '=randstring(20)', '=randstring(35)', '=randstring(20)', '=randstring(2)', '=randstring(10)', '=randstring(12)', '=randstring(35)')
go

select * from customer
loop

#--------------------------------------------------------------------------------------------------------------------------

# format of configuration file is described in db-stress(1)
# this is for ms sql-server with demo database installed
# you have to create the DSN on your own
# name: odbc_mssql.cfg

connect = dbi:ODBC:stress
user = sa
password =
database = pubs2

# session is a bunch of jobs all doing the same batches
session = 1

# the count of jobs working for this session
count = 1

# a batch is a sequence of sql and db-stress commands
batch = 1
delay = 5

commands =
insert into authors (au_id, au_lname, au_fname, address, city, state, contract) values('=randregex([0-9][0-9][0-9]-[0-9][0-9]-[0-9][0-9][0-9][0-9])', '=randstring(40)', '=randstring(20)', '=randstring(40)', '=randstring(20)', '=randstring(2)', 0)
go

select * from titles
loop

#------------------------------------------------------------------------------------------------------------------------