1
+ import subprocess , argparse , os
2
+
3
+ parser = argparse .ArgumentParser ()
4
+
5
+ parser .add_argument ('-n' , '--dbName' , type = str , help = 'Database Name' , required = True )
6
+ parser .add_argument ('-y' , '--databaseSystem' , type = str , help = 'Database Type' , required = False )
7
+ parser .add_argument ('-b' , '--benchmark' , type = str , help = "Benchmark Name" , required = True )
8
+ parser .add_argument ('-t' , '--tableCount' , type = str , help = 'Number of Tables' , required = True )
9
+ parser .add_argument ('-r' , '--recordCount' , type = str , help = 'Number of Records' , required = False )
10
+ parser .add_argument ('-u' , '--warehouses' , type = str , help = 'Warehouse Count' , required = False )
11
+ parser .add_argument ('-e' , '--threadCount' , type = str , help = "Number of Threads" , required = True )
12
+ parser .add_argument ('--password' , type = str , help = "PostgreSQL Password" , required = False )
13
+ parser .add_argument ('--host' , type = str , help = "Database Server Host IP Address" , required = False )
14
+
15
+
16
+ args = parser .parse_args ()
17
+ dbName = args .dbName
18
+ databaseSystem = args .databaseSystem
19
+ benchmark = args .benchmark
20
+ warehouses = args .warehouses
21
+ tableCount = args .tableCount
22
+ recordCount = args .recordCount
23
+ threadCount = args .threadCount
24
+ password = args .password
25
+ host = args .host
26
+
27
+ def add_host_if_needed (command_base , host , benchmark ):
28
+ if host and benchmark != "TPCC" :
29
+ if "sysbench" not in command_base :
30
+ return command_base .replace ('-u' , f'-h { host } -u' )
31
+ else :
32
+ return f"{ command_base } --mysql-host={ host } "
33
+ else :
34
+ return command_base
35
+
36
+ if databaseSystem == "MySQL" :
37
+ if benchmark == "TPCC" :
38
+ subprocess .run (f'sudo src/sysbench tpcc --tables={ tableCount } --scale={ warehouses } --threads={ threadCount } --mysql-db={ dbName } --use_fk=0 prepare' , shell = True , check = True )
39
+ if int (warehouses ) == 1 :
40
+ for i in range (1 ,int (tableCount )+ 1 ):
41
+ table = str (i )
42
+ # drop idxs
43
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "DROP INDEX idx_customer{ i } ON customer{ i } ;"' , shell = True , check = True )
44
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "DROP INDEX idx_orders{ i } ON orders{ i } ;"' , shell = True , check = True )
45
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "DROP INDEX fkey_stock_2{ i } ON stock{ i } ;"' , shell = True , check = True )
46
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "DROP INDEX fkey_order_line_2{ i } ON order_line{ i } ;"' , shell = True , check = True )
47
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "DROP INDEX fkey_history_1{ i } ON history{ i } ;"' , shell = True , check = True )
48
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "DROP INDEX fkey_history_2{ i } ON history{ i } ;"' , shell = True , check = True )
49
+
50
+ # truncate, to make distributing faster
51
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "TRUNCATE TABLE warehouse{ i } ;"' , shell = True , check = True )
52
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "TRUNCATE TABLE district{ i } ;"' , shell = True , check = True )
53
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "TRUNCATE TABLE customer{ i } ;"' , shell = True , check = True )
54
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "TRUNCATE TABLE history{ i } ;"' , shell = True , check = True )
55
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "TRUNCATE TABLE orders{ i } ;"' , shell = True , check = True )
56
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "TRUNCATE TABLE new_orders{ i } ;"' , shell = True , check = True )
57
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "TRUNCATE TABLE order_line{ i } ;"' , shell = True , check = True )
58
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "TRUNCATE TABLE stock{ i } ;"' , shell = True , check = True )
59
+ subprocess .run (f'sudo mysql -u { dbName } { dbName } -e "TRUNCATE TABLE item{ i } ;"' , shell = True , check = True )
60
+ else :
61
+ command_base = f'sudo src/sysbench oltp_common --tables={ tableCount } --table-size={ recordCount } --threads={ threadCount } --mysql-db={ dbName } prepare'
62
+ command = add_host_if_needed (command_base , host , benchmark )
63
+ subprocess .run (command , shell = True , check = True )
64
+ if int (recordCount ) == 1 :
65
+ for i in range (1 , int (tableCount ) + 1 ):
66
+ drop_index_command = f'sudo mysql -u { dbName } { dbName } -e "DROP INDEX k_{ i } ON sbtest{ i } ;"'
67
+ drop_index_command = add_host_if_needed (drop_index_command , host , benchmark )
68
+ subprocess .run (drop_index_command , shell = True , check = True )
69
+ elif databaseSystem == "PostgreSQL" :
70
+ os .environ ['PGPASSWORD' ] = password
71
+ if benchmark == "TPCC" :
72
+ subprocess .run (f'sudo src/sysbench tpcc --db-driver=pgsql --tables={ tableCount } --scale={ warehouses } --threads={ threadCount } --pgsql-user=postgres --pgsql-password={ password } --pgsql-db={ dbName } --use_fk=0 prepare' , shell = True , check = True )
73
+ if int (warehouses ) == 1 :
74
+ for i in range (1 ,int (tableCount )+ 1 ):
75
+ table = str (i )
76
+ # drop idxs
77
+ subprocess .run (f'psql -U postgres -d { dbName } -c "DROP INDEX IF EXISTS idx_customer{ i } ;"' , shell = True , check = True )
78
+ subprocess .run (f'psql -U postgres -d { dbName } -c "DROP INDEX IF EXISTS idx_orders{ i } ;"' , shell = True , check = True )
79
+ subprocess .run (f'psql -U postgres -d { dbName } -c "DROP INDEX IF EXISTS fkey_stock_2{ i } ;"' , shell = True , check = True )
80
+ subprocess .run (f'psql -U postgres -d { dbName } -c "DROP INDEX IF EXISTS fkey_order_line_2{ i } ;"' , shell = True , check = True )
81
+ subprocess .run (f'psql -U postgres -d { dbName } -c "DROP INDEX IF EXISTS fkey_history_1{ i } ;"' , shell = True , check = True )
82
+ subprocess .run (f'psql -U postgres -d { dbName } -c "DROP INDEX IF EXISTS fkey_history_2{ i } ;"' , shell = True , check = True )
83
+
84
+ # truncate, to make distributing faster
85
+ subprocess .run (f'psql -U postgres -d { dbName } -c "TRUNCATE TABLE warehouse{ i } ;"' , shell = True , check = True )
86
+ subprocess .run (f'psql -U postgres -d { dbName } -c "TRUNCATE TABLE district{ i } ;"' , shell = True , check = True )
87
+ subprocess .run (f'psql -U postgres -d { dbName } -c "TRUNCATE TABLE customer{ i } ;"' , shell = True , check = True )
88
+ subprocess .run (f'psql -U postgres -d { dbName } -c "TRUNCATE TABLE history{ i } ;"' , shell = True , check = True )
89
+ subprocess .run (f'psql -U postgres -d { dbName } -c "TRUNCATE TABLE orders{ i } ;"' , shell = True , check = True )
90
+ subprocess .run (f'psql -U postgres -d { dbName } -c "TRUNCATE TABLE new_orders{ i } ;"' , shell = True , check = True )
91
+ subprocess .run (f'psql -U postgres -d { dbName } -c "TRUNCATE TABLE order_line{ i } ;"' , shell = True , check = True )
92
+ subprocess .run (f'psql -U postgres -d { dbName } -c "TRUNCATE TABLE stock{ i } ;"' , shell = True , check = True )
93
+ subprocess .run (f'psql -U postgres -d { dbName } -c "TRUNCATE TABLE item{ i } ;"' , shell = True , check = True )
94
+ else :
95
+ subprocess .run (f'sudo src/sysbench oltp_common --db-driver=pgsql --tables={ tableCount } --table-size={ recordCount } --threads={ threadCount } --pgsql-user=postgres --pgsql-password={ password } --pgsql-db={ dbName } prepare' , shell = True , check = True )
96
+ if int (recordCount ) == 1 :
97
+ for i in range (1 ,int (tableCount )+ 1 ):
98
+ table = str (i )
99
+ subprocess .run (f'psql -U postgres -d { dbName } -c "DROP INDEX IF EXISTS k_{ i } ;"' , shell = True , check = True )
100
+ else :
101
+ parser .error ("You are running on a database system type that has not been onboarded to Virtual Client. Available options are: MySQL, PostgreSQL" )
0 commit comments