Refactor main function using a loop instead of multiple if statements

This commit is contained in:
Ahrimdon 2024-03-03 02:32:51 -05:00
parent b91eea00a9
commit b22d2f3b00

View File

@ -83,68 +83,30 @@ def main():
new_conn = sqlite3.connect(output_db_path) new_conn = sqlite3.connect(output_db_path)
new_cur = new_conn.cursor() new_cur = new_conn.cursor()
# Process each table # Define a mapping of table names to their processing functions
if 'AuditLog' in config.get('tables', {}): table_functions = {
audit_log_limit = config['tables']['AuditLog'] 'AuditLog': audit_log_table,
audit_log_table(existing_cur, new_cur, limit=audit_log_limit) 'ClientConnectionHistory': client_connection_history_table,
else: 'ClientMessages': client_messages_table,
audit_log_table(existing_cur, new_cur) 'Clients': clients_table,
'IPAddresses': ip_address_table,
'InboxMessagesModified': inbox_messages_modified_table,
'Maps': maps_table,
'Metadata': metadata_table,
'Penalties': penalties_table,
'PenaltyIdentifiers': penalty_identifiers_table,
'Servers': servers_table,
}
if 'ClientConnectionHistory' in config.get('tables', {}): # Process each table according to the configuration
client_connection_history_limit = config['tables']['ClientConnectionHistory'] for table_name, func in table_functions.items():
client_connection_history_table(existing_cur, new_cur, limit=client_connection_history_limit) limit = config.get('tables', {}).get(table_name)
else: if limit is not None:
client_connection_history_table(existing_cur, new_cur) print(f"Processing {table_name} with limit: {limit}")
func(existing_cur, new_cur, limit=limit) # Assuming each function can accept a limit parameter
if 'ClientMessages' in config.get('tables', {}): else:
client_messages_limit = config['tables']['ClientMessages'] print(f"Processing {table_name} without limit")
client_messages_table(existing_cur, new_cur, limit=client_messages_limit) func(existing_cur, new_cur)
else:
client_messages_table(existing_cur, new_cur)
if 'Clients' in config.get('tables', {}):
clients_limit = config['tables']['Clients']
clients_table(existing_cur, new_cur, limit=clients_limit)
else:
clients_table(existing_cur, new_cur)
ip_address_table(existing_cur, new_cur)
if 'InboxMessagesModified' in config.get('tables', {}):
inbox_messages_modified_limit = config['tables']['InboxMessagesModified']
inbox_messages_modified_table(existing_cur, new_cur, limit=inbox_messages_modified_limit)
else:
inbox_messages_modified_table(existing_cur, new_cur)
if 'Maps' in config.get('tables', {}):
maps_limit = config['tables']['Maps']
maps_table(existing_cur, new_cur, limit=maps_limit)
else:
maps_table(existing_cur, new_cur)
if 'Metadata' in config.get('tables', {}):
metadata_limit = config['tables']['Metadata']
metadata_table(existing_cur, new_cur, limit=metadata_limit)
else:
metadata_table(existing_cur, new_cur)
if 'Penalties' in config.get('tables', {}):
penalties_limit = config['tables']['Penalties']
penalties_table(existing_cur, new_cur, limit=penalties_limit)
else:
penalties_table(existing_cur, new_cur)
if 'PenaltyIdentifiers' in config.get('tables', {}):
penalty_identifiers_limit = config['tables']['PenaltyIdentifiers']
penalty_identifiers_table(existing_cur, new_cur, limit=penalty_identifiers_limit)
else:
penalty_identifiers_table(existing_cur, new_cur)
if 'Servers' in config.get('tables', {}):
servers_limit = config['tables']['Servers']
servers_table(existing_cur, new_cur, limit=servers_limit)
else:
servers_table(existing_cur, new_cur)
# Commit and close # Commit and close
new_conn.commit() new_conn.commit()