This class will populate a database with randomly generated data. The population includes table creation and data generation. Table names are hard coded as table_ .
| 82 | |
| 83 | |
| 84 | class DbPopulator(object): |
| 85 | '''This class will populate a database with randomly generated data. The population |
| 86 | includes table creation and data generation. Table names are hard coded as |
| 87 | table_<table number>. |
| 88 | |
| 89 | ''' |
| 90 | |
| 91 | def __init__(self, db_engine=db_connection.IMPALA): |
| 92 | self.cluster = None |
| 93 | self.db_name = None |
| 94 | self.db_engine = db_engine |
| 95 | |
| 96 | self.min_col_count = None |
| 97 | self.max_col_count = None |
| 98 | self.min_row_count = None |
| 99 | self.max_row_count = None |
| 100 | self.allowed_storage_formats = None |
| 101 | self.randomization_seed = None |
| 102 | |
| 103 | def populate_db(self, table_count, postgresql_conn=None): |
| 104 | '''Create tables with a random number of cols. |
| 105 | |
| 106 | The given db_name must have already been created. |
| 107 | ''' |
| 108 | self.cluster.hdfs.ensure_home_dir() |
| 109 | hdfs = self.cluster.hdfs.create_client() |
| 110 | |
| 111 | table_and_generators = list() |
| 112 | for table_idx in range(table_count): |
| 113 | table = self._create_random_table( |
| 114 | 'table_%s' % (table_idx + 1), |
| 115 | self.min_col_count, |
| 116 | self.max_col_count, |
| 117 | self.allowed_storage_formats) |
| 118 | self._prepare_table_storage(table, self.db_name) |
| 119 | if table.storage_format == 'TEXTFILE': |
| 120 | text_table = table |
| 121 | else: |
| 122 | text_table = deepcopy(table) |
| 123 | text_table.name += '_text' |
| 124 | text_table.storage_format = 'TEXTFILE' |
| 125 | text_table.storage_location = None |
| 126 | text_table.schema_location = None |
| 127 | self._prepare_table_storage(text_table, self.db_name) |
| 128 | table_data_generator = TextTableDataGenerator() |
| 129 | table_data_generator.randomization_seed = self.randomization_seed |
| 130 | table_data_generator.table = text_table |
| 131 | table_data_generator.row_count = randint(self.min_row_count, self.max_row_count) |
| 132 | table_and_generators.append((table, table_data_generator)) |
| 133 | |
| 134 | self._run_data_generator_mr_job([g for _, g in table_and_generators], self.db_name) |
| 135 | |
| 136 | with self.cluster.hive.cursor(db_name=self.db_name) as cursor: |
| 137 | for table, table_data_generator in table_and_generators: |
| 138 | cursor.create_table(table) |
| 139 | text_table = table_data_generator.table |
| 140 | if postgresql_conn: |
| 141 | with postgresql_conn.cursor() as postgresql_cursor: |