def subset_downstream(self, table, relationships):
referencing_tables = self.__db_helper.get_redacted_table_references(
table, self.__all_tables, self.__source_conn)
if len(referencing_tables) > 0:
pk_columns = referencing_tables[0]['target_columns']
else:
return
temp_table = self.__db_helper.create_id_temp_table(
self.__destination_conn, len(pk_columns))
for r in referencing_tables:
fk_table = r['fk_table']
fk_columns = r['fk_columns']
q = 'SELECT {} FROM {} WHERE {} NOT IN (SELECT {} FROM {})'.format(
columns_joined(fk_columns),
fully_qualified_table(
mysql_db_name_hack(fk_table, self.__destination_conn)),
columns_tupled(fk_columns), columns_joined(pk_columns),
fully_qualified_table(
mysql_db_name_hack(table, self.__destination_conn)))
self.__db_helper.copy_rows(self.__destination_conn,
self.__destination_write_conn, q,
temp_table)
columns_query = columns_to_copy(table, relationships,
self.__source_conn)
cursor_name = 'table_cursor_' + str(uuid.uuid4()).replace('-', '')
cursor = self.__destination_conn.cursor(name=cursor_name,
withhold=True)
cursor_query = 'SELECT DISTINCT * FROM {}'.format(
fully_qualified_table(temp_table))
cursor.execute(cursor_query)
fetch_row_count = 100000
while True:
rows = cursor.fetchmany(fetch_row_count)
if len(rows) == 0:
break
ids = [
'(' + ','.join(['\'' + str(c) + '\'' for c in row]) + ')'
for row in rows if all([c is not None for c in row])
]
if len(ids) == 0:
break
ids_to_query = ','.join(ids)
q = 'SELECT {} FROM {} WHERE {} IN ({})'.format(
columns_query, fully_qualified_table(table),
columns_tupled(pk_columns), ids_to_query)
self.__db_helper.copy_rows(
self.__source_conn, self.__destination_conn, q,
mysql_db_name_hack(table, self.__destination_conn))
cursor.close()
def __subset_direct(self, target, relationships):
t = target['table']
columns_query = columns_to_copy(t, relationships, self.__source_conn)
if 'where' in target:
q = 'SELECT {} FROM {} WHERE {}'.format(columns_query,
fully_qualified_table(t),
target['where'])
elif 'percent' in target:
if config_reader.get_db_type() == 'postgres':
q = 'SELECT {} FROM {} WHERE random() < {}'.format(
columns_query, fully_qualified_table(t),
float(target['percent']) / 100)
else:
q = 'SELECT {} FROM {} WHERE rand() < {}'.format(
columns_query, fully_qualified_table(t),
float(target['percent']) / 100)
else:
raise ValueError(
'target table {} had no \'where\' or \'percent\' term defined, check your configuration.'
.format(t))
self.__db_helper.copy_rows(
self.__source_conn, self.__destination_conn, q,
mysql_db_name_hack(t, self.__destination_conn))
def __subset_upstream(self, target, processed_tables, relationships):
redacted_relationships = redact_relationships(relationships)
relevant_key_constraints = list(
filter(
lambda r: r['target_table'] in processed_tables and r[
'fk_table'] == target, redacted_relationships))
# this table isn't referenced by anything we've already processed, so let's leave it empty
# OR
# table was already added, this only happens if the upstream table was also a direct target
if len(relevant_key_constraints) == 0 or target in processed_tables:
return False
temp_target_name = 'subset_temp_' + table_name(target)
try:
# copy the whole table
columns_query = columns_to_copy(target, relationships,
self.__source_conn)
self.__db_helper.run_query(
'CREATE TEMPORARY TABLE {} AS SELECT * FROM {} LIMIT 0'.format(
quoter(temp_target_name),
fully_qualified_table(
mysql_db_name_hack(target, self.__destination_conn))),
self.__destination_conn)
query = 'SELECT {} FROM {}'.format(columns_query,
fully_qualified_table(target))
self.__db_helper.copy_rows(self.__source_conn,
self.__destination_conn, query,
temp_target_name)
# filter it down in the target database
table_columns = self.__db_helper.get_table_columns(
table_name(target), schema_name(target), self.__source_conn)
clauses = [
'{} IN (SELECT {} FROM {})'.format(
columns_tupled(kc['fk_columns']),
columns_joined(kc['target_columns']),
fully_qualified_table(
mysql_db_name_hack(kc['target_table'],
self.__destination_conn)))
for kc in relevant_key_constraints
]
clauses.extend(upstream_filter_match(target, table_columns))
select_query = 'SELECT * FROM {} WHERE TRUE AND {}'.format(
quoter(temp_target_name), ' AND '.join(clauses))
insert_query = 'INSERT INTO {} {}'.format(
fully_qualified_table(
mysql_db_name_hack(target, self.__destination_conn)),
select_query)
self.__db_helper.run_query(insert_query, self.__destination_conn)
self.__destination_conn.commit()
finally:
# delete temporary table
mysql_temporary = 'TEMPORARY' if config_reader.get_db_type(
) == 'mysql' else ''
self.__db_helper.run_query(
'DROP {} TABLE IF EXISTS {}'.format(mysql_temporary,
quoter(temp_target_name)),
self.__destination_conn)
return True
