Skip to content

Commit e9be569

Browse files
author
Marc Hesse
committed
adding printout of total rows imported/exported
1 parent 0411ee3 commit e9be569

2 files changed

Lines changed: 18 additions & 4 deletions

File tree

‎drivers/python/rethinkdb/_export.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -359,8 +359,13 @@ def run_clients(options, db_table_set):
359359
if error_queue.empty() and not interrupt_event.is_set():
360360
print_progress(1.0)
361361

362-
# Continue past the progress output line
362+
# Continue past the progress output line and print total rows processed
363+
def plural(num, text):
364+
return "%d %s%s" % (num, text, "" if num == 1 else "s")
365+
363366
print ""
367+
print "%s exported from %s" % (plural(sum([info[0].value for info in progress_info]), "row"),
368+
plural(len(db_table_set), "table"))
364369
finally:
365370
signal.signal(signal.SIGINT, signal.SIG_DFL)
366371

‎drivers/python/rethinkdb/_import.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -329,6 +329,7 @@ def read_json_array(json_data, file_in, callback, progress_info):
329329

330330
(obj, offset) = decoder.raw_decode(json_data, idx=offset)
331331
callback(obj)
332+
progress_info[2].value += 1
332333

333334
# Read past whitespace to the next record
334335
file_offset += offset
@@ -373,6 +374,7 @@ def json_reader(task_queue, filename, db, table, primary_key, fields, progress_i
373374
json_data = read_json_array(json_data[offset + 1:], file_in, callback, progress_info)
374375
elif json_data[offset] == "{":
375376
json_data = read_json_single_object(json_data[offset:], file_in, callback)
377+
progress_info[2].value = 1
376378
else:
377379
raise RuntimeError("Error: JSON format not recognized - file does not begin with an object or array")
378380

@@ -426,6 +428,7 @@ def csv_reader(task_queue, filename, db, table, primary_key, options, progress_i
426428
if len(obj[key]) == 0:
427429
del obj[key]
428430
object_callback(obj, db, table, task_queue, object_buffers, buffer_sizes, options["fields"], exit_event)
431+
progress_info[2].value += 1
429432

430433
if len(object_buffers) > 0:
431434
task_queue.put((db, table, object_buffers))
@@ -487,7 +490,7 @@ def print_progress(ratio):
487490

488491
def update_progress(progress_info):
489492
lowest_completion = 1.0
490-
for (current, max_count) in progress_info:
493+
for (current, max_count, rows) in progress_info:
491494
curr_val = current.value
492495
max_val = max_count.value
493496
if curr_val < 0:
@@ -526,8 +529,9 @@ def spawn_import_clients(options, files_info):
526529
client_procs[-1].start()
527530

528531
for file_info in files_info:
529-
progress_info.append((multiprocessing.Value(ctypes.c_longlong, -1),
530-
multiprocessing.Value(ctypes.c_longlong, 0)))
532+
progress_info.append((multiprocessing.Value(ctypes.c_longlong, -1), # Current lines/bytes processed
533+
multiprocessing.Value(ctypes.c_longlong, 0), # Total lines/bytes to process
534+
multiprocessing.Value(ctypes.c_longlong, 0))) # Total rows processed
531535
reader_procs.append(multiprocessing.Process(target=table_reader,
532536
args=(options,
533537
file_info,
@@ -560,7 +564,12 @@ def spawn_import_clients(options, files_info):
560564
print_progress(1.0)
561565

562566
# Continue past the progress output line
567+
def plural(num, text):
568+
return "%d %s%s" % (num, text, "" if num == 1 else "s")
569+
563570
print ""
571+
print "%s imported in %s" % (plural(sum([info[2].value for info in progress_info]), "row"),
572+
plural(len(files_info), "table"))
564573
finally:
565574
signal.signal(signal.SIGINT, signal.SIG_DFL)
566575

0 commit comments

Comments
 (0)