add dump in record chunk sizes

hrbonz · hrbonz · commit 66c1ff74462e · 2019-09-14T16:02:14.000+08:00
Option to dump data in a folder in fragment files of a given chunk size. This works only if the chunksize option is properly implemented in influxdb-python (see influxdata/influxdb-python#753)
diff --git a/README.rst b/README.rst
@@ -12,9 +12,25 @@ flexibility on how to load it back in the database.
 Usage
 =====
 
-Dump data::
+Dump all data from a database::
 
-    $ influxdump -u root -p -d database > data_dump.json
+    $ influxdump -u jdoe -W -d database > data_dump.json
+
+Dump data matching a pattern in chunk files of 50,000 records::
+
+    $ influxdump -u jdoe -W -d database -f _dump -c 50000 -m "node*"
+
+Load data from a dump folder::
+
+    $ influxdump -u jdoe -W -d database -f _dump
+
+
+Install
+=======
+
+.. code-block:: sh
+
+    $ pip install influxdump
 
 Packaging
 =========
diff --git a/influxdump/__init__.py b/influxdump/__init__.py
@@ -2,30 +2,36 @@
 import argparse
 import getpass
 import json
+import sys
 
-from data import dump_data, write_data, load_data
+from data import dump_data, load_file, load_folder
 from db import get_client
 
 
 __author__ = 'Stefan Berder <stefan@measureofquality.com>'
 __contact__ = 'code+influxdump@measureofquality.com'
 __version__ = "1.0.3"
 
+CHUNKSIZE = 50000
+
 
 def get_args():
     parser = argparse.ArgumentParser(description='influxDB data backup tool')
+    parser.add_argument('-c', '--chunksize',
+            help='query chunk size, default to {}'.format(CHUNKSIZE),
+            type=int, default=CHUNKSIZE)
     parser.add_argument('-d', '--database', help='database', required=True,
             type=str)
     parser.add_argument('-F', '--folder', default=None,
-            help="destination folder for fragmented dump, if this flag is not used then dump on stdoout")
+            help="destination folder for fragmented dump, if this flag is not used then dump on stdout")
     parser.add_argument('-H', '--host', help='server host',
             default="localhost", type=str)
     parser.add_argument('-i', '--input', default=None,
             help="data/metadata input file, will force action to 'load'")
     parser.add_argument('-L', '--legacy', action="store_true",
             help='influxdb legacy client (<=0.8)')
     parser.add_argument('-m', '--measurements', help='measurement pattern')
-    parser.add_argument('-n', '--dry-run', help='do mot really do anything', action="store_true")
+    parser.add_argument('-n', '--dry-run', help='do not really do anything', action="store_true")
     parser.add_argument('-p', '--port', help='server port', default=8086,
             type=int)
     parser.add_argument('-u', '--user', help='username', default='', type=str)
@@ -35,16 +41,27 @@ def get_args():
     parser.add_argument('-W', '--pwdprompt', help='password prompt',
             action="store_true")
     parser.add_argument('action', metavar="action", nargs="?", default='dump',
-            help="action, can be 'dump' or 'load', default to 'dump'",
-            choices=["load", "dump"])
+            help="""
+            action, can be 'dump' or 'load', default to 'dump'. If action is
+            'load', one input file (--input) or a folder with data to load has
+            to be provided
+            """, choices=["load", "dump"])
     args = parser.parse_args()
 
     if args.pwdprompt is True:
         pwd = getpass.getpass()
     else:
         pwd = args.password
 
+    if args.action == "load" \
+            and args.input is None and args.folder is None:
+        sys.stderr.write("Action is load, missing input file or folder\n\n")
+        parser.print_help()
+        sys.exit(1)
+
+
     return {
+        "chunksize": args.chunksize,
         "db": args.database,
         "folder": args.folder,
         "host": args.host,
@@ -62,12 +79,15 @@ def get_args():
 
 def dump(args, client):
     dump_data(client, args["measurements"], args["folder"],
-              dryrun=args["dryrun"], verbose=args["verbose"])
+              dryrun=args["dryrun"], verbose=args["verbose"],
+              chunk_size=args["chunksize"])
 
 
 def load(args, client):
-    data = load_data(args["input"])
-    return write_data(client, data)
+    if args["input"] is not None:
+        load_file(client, args["input"], verbose=args["verbose"])
+    else:
+        load_folder(client, args["folder"], verbose=args["verbose"])
 
 
 def main():
diff --git a/influxdump/data.py b/influxdump/data.py
@@ -1,35 +1,41 @@
 # -*- coding: utf-8 -*-
 from datetime import datetime
 import json
+import os
 import os.path
 import sys
 
 from db import get_queries, data_to_points
 
 
-def query_data(c, queries):
+def query_data(c, queries, chunk_size):
     """Generator querying the db and sending back data for each query as
     elements.
     """
     data = []
     for q in queries:
-        res = c.query(q.get_query())
-        records = []
-        for point in c.get_points(res):
-            records.append(point)
-        yield {
-            "meta": q.get_meta(),
-            "records": records
-        }
+        res = c.query(q.get_query(),
+                chunked=True,
+                chunk_size=chunk_size)
+        counter = 0
+        for r in res:
+            records = []
+            counter += 1
+            for point in c.get_points(r):
+                records.append(point)
 
+            yield (counter, {
+                "meta": q.get_meta(),
+                "records": records
+            })
 
-def dump_data(c, pattern=None, folder=None, dryrun=False, verbose=False):
+
+def dump_data(c, pattern=None, folder=None, dryrun=False, verbose=False,
+        chunk_size=50000):
     """Get data from the database, return an `influxdb.ResultSet`
 
     :param c: an influxdb client instance
     :type c: InfluxDBClient
-    :param measurements: a list of measurements to query
-    :type measurements: list
     """
     measurements = c.get_measurements(pattern)
     if verbose is True or dryrun is True:
@@ -43,30 +49,68 @@ def dump_data(c, pattern=None, folder=None, dryrun=False, verbose=False):
         for m in measurements:
             sys.stdout.write("    {}\n".format(m))
     else:
-        for data in query_data(c, queries):
+        for (counter, data) in query_data(c, queries, chunk_size):
             if folder is None:
                 if verbose is True:
                     sys.stdout.write("> dumping {}\n".format(
                         data["meta"]["measurement"]))
                 print(json.dumps(data))
             else:
-                filename = data["meta"]["measurement"] + ".json"
-                dumpfile = os.path.join(folder, filename)
+                bundle = os.path.join(folder,
+                        data["meta"]["measurement"])
+                if not os.path.exists(bundle):
+                    os.makedirs(bundle)
+
+                fragment = "{}-{:05d}.json".format(
+                        data["meta"]["measurement"],
+                        counter)
+                dumpfile = os.path.join(bundle, fragment)
+                data["meta"]["chunk_count"] = counter
+
                 if verbose is True:
-                    sys.stdout.write("> dumping {} to {} ({} records) [{}]\n".format(
-                        data["meta"]["measurement"], filename,
+                    sys.stdout.write(
+                        "> dumping {} (chunk {:05d}) to {} ({} records) [{}]\n".format(
+                        data["meta"]["measurement"], counter, dumpfile,
                         len(data["records"]), datetime.now().isoformat()))
+
                 with open(dumpfile, "w") as fd:
                     json.dump(data, fd)
 
 
 def write_data(c, data):
-    for chunk in data:
-        points = data_to_points(chunk["meta"]["measurement"],
-                                chunk["records"])
-        c.write_points(points, batch_size=10000)
+    #for chunk in data:
+    points = data_to_points(data["meta"]["measurement"],
+                            data["records"])
+    c.write_points(points, batch_size=10000)
 
 
-def load_data(datafile):
+def load_file(c, datafile, verbose=False):
     with open(datafile, 'r') as fh:
-        return json.load(fh)
+        data = json.load(fh)
+
+        if verbose is True:
+            sys.stdout.write(
+                "> loading {} in {} ({} records) [{}]\n".format(
+                datafile, data["meta"]["measurement"],
+                len(data["records"]), datetime.now().isoformat()))
+
+        write_data(c, data)
+
+
+def load_folder(c, folder, verbose=False):
+    for (dirpath, dirnames, filenames) in os.walk(folder):
+        filenames.sort()
+        for filename in filenames:
+            if filename.endswith('.json'):
+                datafile = os.path.join(dirpath, filename)
+
+                with open(datafile, 'r') as fh:
+                    data = json.load(fh)
+                    if verbose is True:
+                        sys.stdout.write(
+                            "> loading {} in {} ({} records) [{}]\n".format(
+                            datafile, data["meta"]["measurement"],
+                            len(data["records"]), datetime.now().isoformat()))
+
+                    write_data(c, data)
+                    del data
diff --git a/influxdump/db.py b/influxdump/db.py
@@ -45,6 +45,9 @@ def get_measurements(self, pattern=None):
 
         return measurements
 
+    def write_points(self, *args, **kwargs):
+        return self._client.write_points(*args, **kwargs)
+
 
 class InfluxDB08Client(InfluxDBClient):
     def __init__(self, host, port, user, pwd, db):
diff --git a/requirements-dev.in b/requirements-dev.in
@@ -1,4 +1,4 @@
 -r requirements.in
-pip-tools==1.9.0
-bumpr==0.3.6
-twine==1.8.1
+bumpr==0.3.7
+pip-tools==4.1.0
+twine==1.13.0
diff --git a/requirements-dev.txt b/requirements-dev.txt
@@ -2,19 +2,31 @@
 # This file is autogenerated by pip-compile
 # To update, run:
 #
-#    pip-compile --output-file requirements-dev.txt requirements-dev.in
+#    pip-compile requirements-dev.in
 #
-args==0.1.0               # via clint
-bumpr==0.3.6
+--index-url https://pypi.priv.measureofquality.com/gams/prod/+simple/
+
+bleach==3.1.0             # via readme-renderer
+bumpr==0.3.7
+certifi==2019.9.11        # via requests
+chardet==3.0.4            # via requests
 click==6.6                # via pip-tools
-clint==0.5.1              # via twine
-first==2.0.1              # via pip-tools
-influxdb==4.0.0
-pip-tools==1.9.0
-pkginfo==1.4.1            # via twine
+docutils==0.15.2          # via readme-renderer
+idna==2.8                 # via requests
+influxdb==5.2.3
+pip-tools==4.1.0
+pkginfo==1.5.0.1          # via twine
+pygments==2.4.2           # via readme-renderer
 python-dateutil==2.6.0    # via influxdb
 pytz==2016.10             # via influxdb
-requests-toolbelt==0.7.1  # via twine
-requests==2.12.3          # via influxdb, requests-toolbelt, twine
-six==1.10.0               # via influxdb, pip-tools, python-dateutil
-twine==1.8.1
+readme-renderer==24.0     # via twine
+requests-toolbelt==0.9.1  # via twine
+requests==2.22.0          # via influxdb, requests-toolbelt, twine
+six==1.10.0               # via bleach, influxdb, pip-tools, python-dateutil, readme-renderer
+tqdm==4.35.0              # via twine
+twine==1.13.0
+urllib3==1.25.3           # via requests
+webencodings==0.5.1       # via bleach
+
+# The following packages are considered to be unsafe in a requirements file:
+# setuptools==41.2.0        # via twine
diff --git a/requirements.in b/requirements.in
@@ -1 +1 @@
-influxdb==4.0.0
+influxdb==5.2.3
diff --git a/requirements.txt b/requirements.txt
@@ -2,10 +2,16 @@
 # This file is autogenerated by pip-compile
 # To update, run:
 #
-#    pip-compile --output-file requirements.txt requirements.in
+#    pip-compile requirements.in
 #
-influxdb==4.0.0
+--index-url https://pypi.priv.measureofquality.com/gams/prod/+simple/
+
+certifi==2019.9.11        # via requests
+chardet==3.0.4            # via requests
+idna==2.8                 # via requests
+influxdb==5.2.3
 python-dateutil==2.6.0    # via influxdb
 pytz==2016.10             # via influxdb
-requests==2.12.3          # via influxdb
+requests==2.22.0          # via influxdb
 six==1.10.0               # via influxdb, python-dateutil
+urllib3==1.25.3           # via requests
