OmniSciDB  a5dc49c757
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
bench_system_tables.py
Go to the documentation of this file.
1 import argparse
2 import sys
3 import pymapd
4 
5 from heavy.thrift.ttypes import TDashboard
6 
7 def getOptions(args=None):
8  parser = argparse.ArgumentParser(description='Basic benchmark for system tables')
9  parser.add_argument('--host', help='HEAVY.AI server address', default='localhost')
10  parser.add_argument('--port', help='HEAVY.AI server port', default='6273')
11  parser.add_argument('--user', help='HEAVY.AI user name', default='admin')
12  parser.add_argument('--password', help='HEAVY.AI password', default='HyperInteractive')
13  parser.add_argument('--database_count', help='Number of databases to create', default=1)
14  parser.add_argument('--table_count', help='Number of tables to create', default=100)
15  parser.add_argument('--dashboard_count', help='Number of dashboards to create', default=100)
16  parser.add_argument('--user_count', help='Number of users to create', default=10)
17  parser.add_argument('--role_count', help='Number of roles to create', default=5)
18  parser.add_argument('--skip_object_creation', help='Skip creation of database objects', default=False)
19  parser.add_argument('--skip_object_deletion', help='Skip deletion of database objects', default=False)
20  parser.add_argument('--tag', help='Tag for test run')
21  return parser.parse_args(args)
22 
23 class HeavyAICon:
24  def __init__(self, user, password, db_name, host):
25  self.con = pymapd.connect(user=user, password=password, dbname=db_name, host=host)
26  self.cursor = self.con.cursor()
27 
28  def query(self, sql):
29  return self.cursor.execute(sql)
30 
31  def create_dashboard(self, dashboard_name):
32  dashboard = TDashboard(dashboard_name = dashboard_name)
33  return self.con.create_dashboard(dashboard)
34 
35 def create_database(heavyai_con, db_id):
36  heavyai_con.query(f"CREATE DATABASE test_db_{db_id}")
37 
38 def create_insert_and_select_from_table(heavyai_con, table_id):
39  heavyai_con.query(f"CREATE TABLE test_table_{table_id} (a INTEGER, b TEXT)")
40  for i in range(10):
41  heavyai_con.query(f"INSERT INTO test_table_{table_id} VALUES ({i}, 'abc_{i}')")
42  heavyai_con.query(f"SELECT AVG(a) FROM test_table_{table_id}")
43 
44 def create_dashboard(heavyai_con, dashboard_id):
45  heavyai_con.create_dashboard(f"test_dashboard_{dashboard_id}")
46 
47 def create_user(heavyai_con, user_id):
48  heavyai_con.query(f"CREATE USER test_user_{user_id} (password = 'test_pass')")
49 
50 def create_role(heavyai_con, role_id):
51  heavyai_con.query(f"CREATE ROLE test_role_{role_id}")
52 
53 def assign_role(heavyai_con, user_id, role_id):
54  heavyai_con.query(f"GRANT test_role_{role_id} TO test_user_{user_id}")
55 
56 def grant_role_table_select(heavyai_con, role_id, db_id):
57  heavyai_con.query(f"GRANT SELECT ON DATABASE test_db_{db_id} TO test_role_{role_id}")
58 
59 def grant_user_table_select(heavyai_con, user_id, db_id):
60  heavyai_con.query(f"GRANT SELECT ON DATABASE test_db_{db_id} TO test_user_{user_id}")
61 
62 def drop_database(heavyai_con, db_id):
63  heavyai_con.query(f"DROP DATABASE test_db_{db_id}")
64 
65 def drop_user(heavyai_con, user_id):
66  heavyai_con.query(f"DROP USER test_user_{user_id}")
67 
68 def drop_role(heavyai_con, role_id):
69  heavyai_con.query(f"DROP ROLE test_role_{role_id}")
70 
71 def query_and_time_system_table(heavyai_con, table_name):
72  query = f"SELECT COUNT(*) FROM {table_name}"
73  result = heavyai_con.query(query)
74  print(f"Query: {query}, Execution time: {result._result.execution_time_ms}ms")
75  query = f"SELECT * FROM {table_name} LIMIT 10"
76  result = heavyai_con.query(query)
77  print(f"Query: {query}, Execution time: {result._result.execution_time_ms}ms")
78 
79 def get_connection(options, db_name):
80  return HeavyAICon(options.user, options.password, db_name, options.host)
81 
82 def main(argv):
83  options = getOptions(argv)
84  default_db = "heavyai"
85  heavyai_con = get_connection(options, default_db)
86 
87  if not options.skip_object_creation:
88  print("Creating database objects")
89  for db_id in range(options.database_count):
90  create_database(heavyai_con, db_id)
91  db_name = f"test_db_{db_id}"
92  heavyai_con = get_connection(options, db_name)
93  for table_id in range(options.table_count):
94  create_insert_and_select_from_table(heavyai_con, table_id)
95  print(f"{options.table_count} tables created for {db_name}")
96  for dashboard_id in range(options.dashboard_count):
97  create_dashboard(heavyai_con, dashboard_id)
98  print(f"{options.dashboard_count} dashboards created for {db_name}")
99  print(f"{options.database_count} databases created")
100 
101  heavyai_con = get_connection(options, default_db)
102  for user_id in range(options.user_count):
103  create_user(heavyai_con, user_id)
104  print(f"{options.user_count} users created")
105 
106  for role_id in range(options.role_count):
107  create_role(heavyai_con, role_id)
108  print(f"{options.role_count} roles created")
109 
110  half_roles = int(options.role_count / 2)
111  for user_id in range(options.user_count):
112  for role_id in range(half_roles):
113  assign_role(heavyai_con, user_id, role_id)
114 
115  if options.database_count > 0:
116  db_id = 0
117  for role_id in range(half_roles):
118  grant_role_table_select(heavyai_con, role_id + half_roles, db_id)
119 
120  half_users = int(options.user_count / 2)
121  for user_id in range(half_users):
122  grant_user_table_select(heavyai_con, user_id + half_users, db_id)
123 
124  system_tables = ["tables",
125  "dashboards",
126  "databases",
127  "users",
128  "permissions",
129  "role_assignments",
130  "roles",
131  "storage_details",
132  "memory_details",
133  "memory_summary"]
134  heavyai_con = get_connection(options, "information_schema")
135  print("Executing system table queries")
136  for table_name in system_tables:
137  query_and_time_system_table(heavyai_con, table_name)
138 
139  if not options.skip_object_deletion:
140  heavyai_con = get_connection(options, default_db)
141  print("Dropping databases")
142  for db_id in range(options.database_count):
143  drop_database(heavyai_con, db_id)
144 
145  print("Dropping users")
146  for user_id in range(options.user_count):
147  drop_user(heavyai_con, user_id)
148 
149  print("Dropping roles")
150  for role_id in range(options.role_count):
151  drop_role(heavyai_con, role_id)
152 
153 if __name__ == "__main__":
154  main(sys.argv[1:])