def rebuild_partitions(athena_client, options, config): """Rebuild an Athena table's partitions Steps: - Get the list of current partitions - Destroy existing table - Re-create tables - Re-create partitions Args: athena_client (boto3.client): Instantiated CLI AthenaClient options (namedtuple): The parsed args passed from the CLI config (CLIConfig): Loaded StreamAlert CLI """ if not options.table_name: LOGGER_CLI.error('Missing command line argument --table_name') return if not options.bucket: LOGGER_CLI.error('Missing command line argument --bucket') return sa_firehose = StreamAlertFirehose( config['global']['account']['region'], config['global']['infrastructure']['firehose'], config['logs']) sanitized_table_name = sa_firehose.firehose_log_name(options.table_name) if options.type == 'data': # Get the current set of partitions partition_success, partitions = athena_client.run_athena_query( query='SHOW PARTITIONS {}'.format(sanitized_table_name), database='streamalert') if not partition_success: LOGGER_CLI.error('An error occured when loading partitions for %s', sanitized_table_name) return unique_partitions = athena_helpers.unique_values_from_query(partitions) # Drop the table LOGGER_CLI.info('Dropping table %s', sanitized_table_name) drop_success, _ = athena_client.run_athena_query( query='DROP TABLE {}'.format(sanitized_table_name), database='streamalert') if not drop_success: LOGGER_CLI.error('An error occured when dropping the %s table', sanitized_table_name) return LOGGER_CLI.info('Dropped table %s', sanitized_table_name) new_partitions_statement = athena_helpers.partition_statement( unique_partitions, options.bucket, sanitized_table_name) # Make sure our new alter table statement is within the query API limits if len(new_partitions_statement) > MAX_QUERY_LENGTH: LOGGER_CLI.error( 'Partition statement too large, writing to local file') with open('partitions_{}.txt'.format(sanitized_table_name), 'w') as partition_file: partition_file.write(new_partitions_statement) return # Re-create the table with previous partitions options.refresh_type = 'add_hive_partition' create_table(athena_client, options, config) LOGGER_CLI.info('Creating %d new partitions for %s', len(unique_partitions), sanitized_table_name) new_part_success, _ = athena_client.run_athena_query( query=new_partitions_statement, database='streamalert') if not new_part_success: LOGGER_CLI.error('Error re-creating new partitions for %s', sanitized_table_name) return LOGGER_CLI.info('Successfully rebuilt partitions for %s', sanitized_table_name) else: LOGGER_CLI.info('Refreshing alerts tables unsupported')
def create_table(athena_client, options, config): """Create a 'streamalert' Athena table Args: athena_client (boto3.client): Instantiated CLI AthenaClient options (namedtuple): The parsed args passed from the CLI config (CLIConfig): Loaded StreamAlert CLI """ sa_firehose = StreamAlertFirehose( config['global']['account']['region'], config['global']['infrastructure']['firehose'], config['logs']) if not options.bucket: LOGGER_CLI.error('Missing command line argument --bucket') return if not options.refresh_type: LOGGER_CLI.error('Missing command line argument --refresh_type') return if options.type == 'data': if not options.table_name: LOGGER_CLI.error('Missing command line argument --table_name') return # Convert special characters in schema name to underscores sanitized_table_name = sa_firehose.firehose_log_name( options.table_name) # Check that the log type is enabled via Firehose if sanitized_table_name not in sa_firehose.enabled_logs: LOGGER_CLI.error( 'Table name %s missing from configuration or ' 'is not enabled.', sanitized_table_name) return # Check if the table exists if athena_client.check_table_exists(sanitized_table_name): LOGGER_CLI.info('The \'%s\' table already exists.', sanitized_table_name) return log_info = config['logs'][options.table_name.replace('_', ':', 1)] schema = dict(log_info['schema']) sanitized_schema = StreamAlertFirehose.sanitize_keys(schema) athena_schema = handler_helpers.to_athena_schema(sanitized_schema) # Add envelope keys to Athena Schema configuration_options = log_info.get('configuration') if configuration_options: envelope_keys = configuration_options.get('envelope_keys') if envelope_keys: sanitized_envelope_key_schema = StreamAlertFirehose.sanitize_keys( envelope_keys) # Note: this key is wrapped in backticks to be Hive compliant athena_schema[ '`streamalert:envelope_keys`'] = handler_helpers.to_athena_schema( sanitized_envelope_key_schema) # Handle Schema overrides # This is useful when an Athena schema needs to differ from the normal log schema if options.schema_override: for override in options.schema_override: if '=' not in override: LOGGER_CLI.error( 'Invalid schema override [%s], use column_name=type format', override) return column_name, column_type = override.split('=') if not all([column_name, column_type]): LOGGER_CLI.error( 'Invalid schema override [%s], use column_name=type format', override) # Columns are escaped to avoid Hive issues with special characters column_name = '`{}`'.format(column_name) if column_name in athena_schema: athena_schema[column_name] = column_type LOGGER_CLI.info('Applied schema override: %s:%s', column_name, column_type) else: LOGGER_CLI.error( 'Schema override column %s not found in Athena Schema, skipping', column_name) query = _construct_create_table_statement( schema=athena_schema, table_name=sanitized_table_name, bucket=options.bucket) elif options.type == 'alerts': if athena_client.check_table_exists(options.type): LOGGER_CLI.info('The \'alerts\' table already exists.') return query = ALERTS_TABLE_STATEMENT.format(bucket=options.bucket) if query: create_table_success, _ = athena_client.run_athena_query( query=query, database='streamalert') if create_table_success: # Update the CLI config config['lambda']['athena_partition_refresh_config'] \ ['refresh_type'][options.refresh_type][options.bucket] = options.type config.write() table_name = options.type if options.type == 'alerts' else sanitized_table_name LOGGER_CLI.info('The %s table was successfully created!', table_name)
def rebuild_partitions(table, bucket, config): """Rebuild an Athena table's partitions Steps: - Get the list of current partitions - Destroy existing table - Re-create tables - Re-create partitions Args: table (str): The name of the table being rebuilt bucket (str): The s3 bucket to be used as the location for Athena data table_type (str): The type of table being refreshed Types of 'data' and 'alert' are accepted, but only 'data' is implemented config (CLIConfig): Loaded StreamAlert CLI """ athena_client = StreamAlertAthenaClient( config, results_key_prefix='stream_alert_cli') sa_firehose = StreamAlertFirehose( config['global']['account']['region'], config['global']['infrastructure']['firehose'], config['logs']) sanitized_table_name = sa_firehose.firehose_log_name(table) # Get the current set of partitions partition_success, partitions = athena_client.run_athena_query( query='SHOW PARTITIONS {}'.format(sanitized_table_name), database=athena_client.sa_database) if not partition_success: LOGGER_CLI.error('An error occurred when loading partitions for %s', sanitized_table_name) return unique_partitions = athena_helpers.unique_values_from_query(partitions) if not unique_partitions: LOGGER_CLI.info('No partitions to rebuild for %s, nothing to do', sanitized_table_name) return # Drop the table LOGGER_CLI.info('Dropping table %s', sanitized_table_name) drop_success, _ = athena_client.run_athena_query( query='DROP TABLE {}'.format(sanitized_table_name), database=athena_client.sa_database) if not drop_success: LOGGER_CLI.error('An error occurred when dropping the %s table', sanitized_table_name) return LOGGER_CLI.info('Dropped table %s', sanitized_table_name) LOGGER_CLI.info('Creating table %s', sanitized_table_name) # Re-create the table with previous partitions create_table(table, bucket, config) new_partitions_statement = athena_helpers.partition_statement( unique_partitions, bucket, sanitized_table_name) # Make sure our new alter table statement is within the query API limits if len(new_partitions_statement) > MAX_QUERY_LENGTH: LOGGER_CLI.error( 'Partition statement too large, writing to local file') with open('partitions_{}.txt'.format(sanitized_table_name), 'w') as partition_file: partition_file.write(new_partitions_statement) return LOGGER_CLI.info('Creating %d new partitions for %s', len(unique_partitions), sanitized_table_name) new_part_success, _ = athena_client.run_athena_query( query=new_partitions_statement, database=athena_client.sa_database) if not new_part_success: LOGGER_CLI.error('Error re-creating new partitions for %s', sanitized_table_name) return LOGGER_CLI.info('Successfully rebuilt partitions for %s', sanitized_table_name)
def create_table(table, bucket, config, schema_override=None): """Create a 'streamalert' Athena table Args: table (str): The name of the table being rebuilt bucket (str): The s3 bucket to be used as the location for Athena data table_type (str): The type of table being refreshed config (CLIConfig): Loaded StreamAlert CLI schema_override (set): An optional set of key=value pairs to be used for overriding the configured column_name=value_type. """ athena_client = StreamAlertAthenaClient( config, results_key_prefix='stream_alert_cli') sa_firehose = StreamAlertFirehose( config['global']['account']['region'], config['global']['infrastructure']['firehose'], config['logs']) # Convert special characters in schema name to underscores sanitized_table_name = sa_firehose.firehose_log_name(table) # Check that the log type is enabled via Firehose if sanitized_table_name != 'alerts' and sanitized_table_name not in sa_firehose.enabled_logs: LOGGER_CLI.error( 'Table name %s missing from configuration or ' 'is not enabled.', sanitized_table_name) return # Check if the table exists if athena_client.check_table_exists(sanitized_table_name, True): LOGGER_CLI.info('The \'%s\' table already exists.', sanitized_table_name) return if table == 'alerts': # get a fake alert so we can get the keys needed and their types alert = Alert('temp_rule_name', {}, {}) output = alert.output_dict() schema = record_to_schema(output) athena_schema = handler_helpers.to_athena_schema(schema) query = _construct_create_table_statement(schema=athena_schema, table_name=table, bucket=bucket) else: # all other tables are log types log_info = config['logs'][table.replace('_', ':', 1)] schema = dict(log_info['schema']) sanitized_schema = StreamAlertFirehose.sanitize_keys(schema) athena_schema = handler_helpers.to_athena_schema(sanitized_schema) # Add envelope keys to Athena Schema configuration_options = log_info.get('configuration') if configuration_options: envelope_keys = configuration_options.get('envelope_keys') if envelope_keys: sanitized_envelope_key_schema = StreamAlertFirehose.sanitize_keys( envelope_keys) # Note: this key is wrapped in backticks to be Hive compliant athena_schema[ '`streamalert:envelope_keys`'] = handler_helpers.to_athena_schema( sanitized_envelope_key_schema) # Handle Schema overrides # This is useful when an Athena schema needs to differ from the normal log schema if schema_override: for override in schema_override: column_name, column_type = override.split('=') if not all([column_name, column_type]): LOGGER_CLI.error( 'Invalid schema override [%s], use column_name=type format', override) # Columns are escaped to avoid Hive issues with special characters column_name = '`{}`'.format(column_name) if column_name in athena_schema: athena_schema[column_name] = column_type LOGGER_CLI.info('Applied schema override: %s:%s', column_name, column_type) else: LOGGER_CLI.error( 'Schema override column %s not found in Athena Schema, skipping', column_name) query = _construct_create_table_statement( schema=athena_schema, table_name=sanitized_table_name, bucket=bucket) create_table_success, _ = athena_client.run_athena_query( query=query, database=athena_client.sa_database) if not create_table_success: LOGGER_CLI.error('The %s table could not be created', sanitized_table_name) return # Update the CLI config if (table != 'alerts' and bucket not in config['lambda'] ['athena_partition_refresh_config']['buckets']): config['lambda']['athena_partition_refresh_config']['buckets'][ bucket] = 'data' config.write() LOGGER_CLI.info('The %s table was successfully created!', sanitized_table_name)
def rebuild_partitions(table, bucket, config): """Rebuild an Athena table's partitions Steps: - Get the list of current partitions - Destroy existing table - Re-create tables - Re-create partitions Args: table (str): The name of the table being rebuilt bucket (str): The s3 bucket to be used as the location for Athena data table_type (str): The type of table being refreshed Types of 'data' and 'alert' are accepted, but only 'data' is implemented config (CLIConfig): Loaded StreamAlert CLI """ athena_client = get_athena_client(config) sa_firehose = StreamAlertFirehose( config['global']['account']['region'], config['global']['infrastructure']['firehose'], config['logs']) sanitized_table_name = sa_firehose.firehose_log_name(table) # Get the current set of partitions partitions = athena_client.get_table_partitions(sanitized_table_name) if not partitions: LOGGER_CLI.info('No partitions to rebuild for %s, nothing to do', sanitized_table_name) return # Drop the table LOGGER_CLI.info('Dropping table %s', sanitized_table_name) success = athena_client.drop_table(sanitized_table_name) if not success: return LOGGER_CLI.info('Creating table %s', sanitized_table_name) # Re-create the table with previous partitions create_table(table, bucket, config) new_partitions_statement = helpers.add_partition_statement( partitions, bucket, sanitized_table_name) # Make sure our new alter table statement is within the query API limits if len(new_partitions_statement) > MAX_QUERY_LENGTH: LOGGER_CLI.error( 'Partition statement too large, writing to local file') with open('partitions_{}.txt'.format(sanitized_table_name), 'w') as partition_file: partition_file.write(new_partitions_statement) return LOGGER_CLI.info('Creating %d new partitions for %s', len(partitions), sanitized_table_name) success = athena_client.run_query(query=new_partitions_statement) if not success: LOGGER_CLI.error('Error re-creating new partitions for %s', sanitized_table_name) return LOGGER_CLI.info('Successfully rebuilt partitions for %s', sanitized_table_name)
def create_table(athena_client, options, config): """Create a 'streamalert' Athena table Args: athena_client (boto3.client): Instantiated CLI AthenaClient options (namedtuple): The parsed args passed from the CLI config (CLIConfig): Loaded StreamAlert CLI """ sa_firehose = StreamAlertFirehose( config['global']['account']['region'], config['global']['infrastructure']['firehose'], config['logs']) if not options.bucket: LOGGER_CLI.error('Missing command line argument --bucket') return if not options.refresh_type: LOGGER_CLI.error('Missing command line argument --refresh_type') return if options.type == 'data': if not options.table_name: LOGGER_CLI.error('Missing command line argument --table_name') return sanitized_table_name = sa_firehose.firehose_log_name( options.table_name) if sanitized_table_name not in sa_firehose.enabled_logs: LOGGER_CLI.error( 'Table name %s missing from configuration or ' 'is not enabled.', sanitized_table_name) return if athena_client.check_table_exists(sanitized_table_name): LOGGER_CLI.info('The \'%s\' table already exists.', sanitized_table_name) return log_info = config['logs'][options.table_name.replace('_', ':', 1)] schema = dict(log_info['schema']) schema_statement = '' sanitized_schema = StreamAlertFirehose.sanitize_keys(schema) athena_schema = {} _add_to_athena_schema(sanitized_schema, athena_schema) # Support envelope keys configuration_options = log_info.get('configuration') if configuration_options: envelope_keys = configuration_options.get('envelope_keys') if envelope_keys: sanitized_envelope_key_schema = StreamAlertFirehose.sanitize_keys( envelope_keys) # Note: this key is wrapped in backticks to be Hive compliant _add_to_athena_schema(sanitized_envelope_key_schema, athena_schema, '`streamalert:envelope_keys`') for key_name, key_type in athena_schema.iteritems(): # Account for nested structs if isinstance(key_type, dict): struct_schema = ''.join([ '{0}:{1},'.format(sub_key, sub_type) for sub_key, sub_type in key_type.iteritems() ]) nested_schema_statement = '{0} struct<{1}>, '.format( key_name, # Use the minus index to remove the last comma struct_schema[:-1]) schema_statement += nested_schema_statement else: schema_statement += '{0} {1},'.format(key_name, key_type) query = ( 'CREATE EXTERNAL TABLE {table_name} ({schema}) ' 'PARTITIONED BY (dt string) ' 'ROW FORMAT SERDE \'org.openx.data.jsonserde.JsonSerDe\' ' 'WITH SERDEPROPERTIES ( \'ignore.malformed.json\' = \'true\') ' 'LOCATION \'s3://{bucket}/{table_name}/\''.format( table_name=sanitized_table_name, # Use the minus index to remove the last comma schema=schema_statement[:-1], bucket=options.bucket)) elif options.type == 'alerts': if athena_client.check_table_exists(options.type): LOGGER_CLI.info('The \'alerts\' table already exists.') return query = ('CREATE EXTERNAL TABLE alerts (' 'log_source string,' 'log_type string,' 'outputs array<string>,' 'record string,' 'rule_description string,' 'rule_name string,' 'source_entity string,' 'source_service string)' 'PARTITIONED BY (dt string)' 'ROW FORMAT SERDE \'org.openx.data.jsonserde.JsonSerDe\'' 'LOCATION \'s3://{bucket}/alerts/\''.format( bucket=options.bucket)) if query: create_table_success, _ = athena_client.run_athena_query( query=query, database='streamalert') if create_table_success: # Update the CLI config config['lambda']['athena_partition_refresh_config'] \ ['refresh_type'][options.refresh_type][options.bucket] = options.type config.write() table_name = options.type if options.type == 'alerts' else sanitized_table_name LOGGER_CLI.info('The %s table was successfully created!', table_name)