Args for send_data.py take file parameters instead string

This commit is contained in:
2017-04-21 19:15:23 +02:00
parent 8aee169713
commit 8049e0c9ef

View File

@@ -7,15 +7,17 @@
import argparse import argparse
import requests import requests
def send_data(file_name, quiet=False): def send_data(file, quiet=False):
""" """
Send a file to ELS Send a data bulk file to ELS.
'file' should be (readable) file object.
""" """
if not quiet: if not quiet:
print("Sending '{}' data file...".format(file_name)) print("Sending '{}' data file...".format(file.name))
res = requests.post(url=ELASTICSEARCH_URL + '_bulk', res = requests.post(url=ELASTICSEARCH_URL + '_bulk',
data=open(file_name, 'rb'), data=file,
headers={'Content-Type': 'application/x-ndjson'}) headers={'Content-Type': 'application/x-ndjson'})
if res.status_code != 200: if res.status_code != 200:
print("An error occured") print("An error occured")
@@ -23,7 +25,7 @@ def send_data(file_name, quiet=False):
print(res.text) print(res.text)
else: else:
if not quiet: if not quiet:
print("File '{} sended to Elasticsearch!".format(file_name)) print("File '{} sended to Elasticsearch!".format(file.name))
def delete_index(index_name, quiet=False): def delete_index(index_name, quiet=False):
@@ -55,21 +57,22 @@ parser = argparse.ArgumentParser(
) )
parser.add_argument('-s', '--song', action='store_false', parser.add_argument('-s', '--song', action='store_false',
help='Disable send song data') help='Disable send song data')
parser.add_argument('-sf', '--song-file', default='es-music-data.json', #TODO type=argparse.FileType('r'), parser.add_argument('-sf', '--song-file', default='es-music-data.json', type=argparse.FileType('r'),
help='Song file data to send (default: es-music-data.json). Need to enable send song data.') help='Song file data to send (default: es-music-data.json). Need to enable send song data.')
# FIXME If song is disable and default file not exist ==> Fail
parser.add_argument('-al', '--album', action='store_true', parser.add_argument('-al', '--album', action='store_true',
help='Enable send album data') help='Enable send album data')
parser.add_argument('-alf', '--album-file', default='es-albums-data.json', #TODO type=argparse.FileType('r'), parser.add_argument('-alf', '--album-file', default='es-albums-data.json', type=argparse.FileType('r'),
help='Album file data to send (default: es-albums-data.json). Need to enable send album data.') help='Album file data to send (default: es-albums-data.json). Need to enable send album data.')
parser.add_argument('-ar', '--artist', action='store_true', parser.add_argument('-ar', '--artist', action='store_true',
help='Enable send music data') help='Enable send music data')
parser.add_argument('-arf', '--artist-file', default='es-artist-data.json', #TODO type=argparse.FileType('r'), parser.add_argument('-arf', '--artist-file', default='es-artist-data.json', type=argparse.FileType('r'),
help='Album file data to send (default: es-artist-data.json). Need to enable send artist data.') help='Album file data to send (default: es-artist-data.json). Need to enable send artist data.')
parser.add_argument('-A', '--ALL', action='store_true', parser.add_argument('-A', '--ALL', action='store_true',
help='Send all') help='Send all')
parser.add_argument('-D', '--DELETE', action='store_true', parser.add_argument('-D', '--DELETE', action='store_true',
help='Delete old index') help='Delete old index')
parser.add_argument('-els', '--elasticsearch_url', default=ELASTICSEARCH_URL, parser.add_argument('-els', '--elasticsearch_url', default=ELASTICSEARCH_URL, nargs='?',
help="Elasticsearch URL to send data (default: {})".format(ELASTICSEARCH_URL)) help="Elasticsearch URL to send data (default: {})".format(ELASTICSEARCH_URL))
parser.add_argument('-i', '--index_name', default=INDEX_NAME, parser.add_argument('-i', '--index_name', default=INDEX_NAME,
help="Index name in Elasticsearch for data (default: {})".format(INDEX_NAME)) help="Index name in Elasticsearch for data (default: {})".format(INDEX_NAME))
@@ -91,16 +94,16 @@ if __name__ == '__main__':
if args.song or args.ALL: if args.song or args.ALL:
song_file = args.song_file song_file = args.song_file
if not args.quiet: if not args.quiet:
print("Take file '{}' to send song data".format(song_file)) print("Take file '{}' to send song data".format(song_file.name))
send_data(song_file, args.quiet) send_data(song_file, args.quiet)
if args.artist or args.ALL: if args.artist or args.ALL:
artist_file = args.artist_file artist_file = args.artist_file
if not args.quiet: if not args.quiet:
print("Take file '{}' to send song data".format(artist_file)) print("Take file '{}' to send song data".format(artist_file.name))
send_data(artist_file, args.quiet) send_data(artist_file, args.quiet)
if args.album or args.ALL: if args.album or args.ALL:
album_file = args.album_file album_file = args.album_file
if not args.quiet: if not args.quiet:
print("Take file '{}' to send song data".format(album_file)) print("Take file '{}' to send song data".format(album_file.name))
send_data(album_file, args.quiet) send_data(album_file, args.quiet)