Implemented import API endpoint

aw_server/api.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Any, Optional
+from typing import Dict, List, Any, Optional, Union
 from datetime import datetime
 from socket import gethostname
 import functools
@@ -33,7 +33,7 @@ class ServerAPI:
     def __init__(self, db, testing) -> None:
         self.db = db
         self.testing = testing
-        self.last_event = {} #type: dict
+        self.last_event = {}  # type: dict
 
     def get_info(self) -> Dict[str, Dict]:
         """Get server info"""
@@ -73,19 +73,49 @@ def export_bucket(self, bucket_id: str) -> Dict[str, Any]:
             del event["id"]
         return bucket
 
-    def create_bucket(self, bucket_id: str, event_type: str, client: str, hostname: str) -> bool:
+    def export_all(self) -> Dict[str, Any]:
+        """Exports all buckets and their events to a format consistent across versions"""
+        buckets = self.get_buckets()
+        exported_buckets = {}
+        for bid in buckets.keys():
+            exported_buckets[bid] = self.export_bucket(bid)
+        return exported_buckets
+
+    def import_bucket(self, bucket_data: Any):
+        bucket_id = bucket_data["id"]
+        logger.info("Importing bucket {}".format(bucket_id))
+        # TODO: Check that bucket doesn't already exist
+        self.db.create_bucket(
+            bucket_id,
+            type=bucket_data["type"],
+            client=bucket_data["client"],
+            hostname=bucket_data["hostname"],
+            created=(bucket_data["created"]
+                     if isinstance(bucket_data["created"], datetime)
+                     else iso8601.parse_date(bucket_data["created"])),
+        )
+        self.create_events(bucket_id, [Event(**e) if isinstance(e, dict) else e for e in bucket_data["events"]])
+
+    def import_all(self, buckets: Dict[str, Any]):
+        for bid, bucket in buckets.items():
+            self.import_bucket(bucket)
+
+    def create_bucket(self, bucket_id: str, event_type: str, client: str,
+                      hostname: str, created: Optional[datetime] = None)-> bool:
         """
         Create bucket.
         Returns True if successful, otherwise false if a bucket with the given ID already existed.
         """
+        if created is None:
+            created = datetime.now()
         if bucket_id in self.db.buckets():
             return False
         self.db.create_bucket(
             bucket_id,
             type=event_type,
             client=client,
             hostname=hostname,
-            created=datetime.now()
+            created=created
         )
         return True
 
@@ -207,10 +237,3 @@ def get_log(self):
             for line in log_file.readlines()[::-1]:
                 payload.append(json.loads(line))
         return payload, 200
-
-    def export_all(self) -> Dict[str, dict]:
-        """Exports all buckets and their events to a format consistent across versions"""
-        buckets = self.get_buckets()
-        for bid in buckets.keys():
-            buckets[bid] = self.export_bucket(bid)
-        return buckets

aw_server/rest.py

@@ -2,7 +2,7 @@
 import traceback
 import json
 
-from flask import request, Blueprint, jsonify
+from flask import request, Blueprint, jsonify, make_response
 from flask_restplus import Api, Resource, fields
 import iso8601
 from datetime import datetime, timedelta
@@ -26,6 +26,7 @@
 blueprint = Blueprint('api', __name__, url_prefix='/api')
 api = Api(blueprint, doc='/')
 
+
 # TODO: Clean up JSONEncoder code?
 class CustomJSONEncoder(json.JSONEncoder):
     def __init__(self, *args, **kwargs):
@@ -40,6 +41,8 @@ def default(self, obj, *args, **kwargs):
         except TypeError:
             pass
         return json.JSONEncoder.default(self, obj)
+
+
 app.json_encoder = CustomJSONEncoder
 
 app.register_blueprint(blueprint)
@@ -54,37 +57,32 @@ def format(self, value):
             return json.loads(value)
 
 
-# Loads event schema from JSONSchema in aw_core
+# Loads event and bucket schema from JSONSchema in aw_core
 event = api.schema_model('Event', schema.get_json_schema("event"))
+bucket = api.schema_model('Bucket', schema.get_json_schema("bucket"))
+buckets_export = api.schema_model('Export', schema.get_json_schema("export"))
 
 # TODO: Construct all the models from JSONSchema?
 #       A downside to contructing from JSONSchema: flask-restplus does not have marshalling support
+
 info = api.model('Info', {
     'hostname': fields.String(),
     'version': fields.String(),
     'testing': fields.Boolean(),
 })
 
-bucket = api.model('Bucket', {
-    'id': fields.String(required=True, description='The buckets unique id'),
-    'name': fields.String(required=False, description='The buckets readable and renameable name'),
-    'type': fields.String(required=True, description='The buckets event type'),
-    'client': fields.String(required=True, description='The name of the watcher client'),
-    'hostname': fields.String(required=True, description='The hostname of the client that the bucket belongs to'),
-    'created': fields.DateTime(required=True, description='The creation datetime of the bucket'),
-})
-
 create_bucket = api.model('CreateBucket', {
     'client': fields.String(required=True),
     'type': fields.String(required=True),
     'hostname': fields.String(required=True),
 })
 
 query = api.model('Query', {
-    'timeperiods':  fields.List(fields.String, required=True, description='List of periods to query'),
+    'timeperiods': fields.List(fields.String, required=True, description='List of periods to query'),
     'query': fields.List(fields.String, required=True, description='String list of query statements'),
 })
 
+
 def copy_doc(api_method):
     """Decorator that copies another functions docstring to the decorated function.
        Used to copy the docstrings in ServerAPI over to the flask-restplus Resources.
@@ -117,7 +115,7 @@ def get(self) -> Dict[str, Dict]:
 
 @api.route("/0/buckets/<string:bucket_id>")
 class BucketResource(Resource):
-    @api.marshal_with(bucket)
+    @api.doc(model=bucket)
     @copy_doc(ServerAPI.get_bucket_metadata)
     def get(self, bucket_id):
         return app.api.get_bucket_metadata(bucket_id)
@@ -254,21 +252,53 @@ def post(self):
             return {"type": type(qe).__name__, "message": str(qe)}, 400
 
 
-# EXPORTING
+# EXPORT AND IMPORT
+
 
 @api.route("/0/export")
 class ExportAllResource(Resource):
+    @api.doc(model=buckets_export)
     @copy_doc(ServerAPI.export_all)
     def get(self):
-        return app.api.export_all(), 200
+        buckets_export = app.api.export_all()
+        payload = {"buckets": buckets_export}
+        response = make_response(json.dumps(payload))
+        filename = "aw-buckets-export.json"
+        response.headers["Content-Disposition"] = "attachment; filename={}".format(filename)
+        return response
 
 
 # TODO: Perhaps we don't need this, could be done with a query argument to /0/export instead
 @api.route("/0/buckets/<string:bucket_id>/export")
 class BucketExportResource(Resource):
+    @api.doc(model=buckets_export)
     @copy_doc(ServerAPI.export_bucket)
     def get(self, bucket_id):
-        return app.api.export_bucket(bucket_id)
+        bucket_export = app.api.export_bucket(bucket_id)
+        payload = {"buckets": {bucket_export["id"]: bucket_export}}
+        response = make_response(json.dumps(payload))
+        filename = "aw-bucket-export_{}.json".format(bucket_export["id"])
+        response.headers["Content-Disposition"] = "attachment; filename={}".format(filename)
+        return response
+
+
+@api.route("/0/import")
+class ImportAllResource(Resource):
+    @api.expect(buckets_export)
+    @copy_doc(ServerAPI.import_all)
+    def post(self):
+        # If import comes from a form in th web-ui
+        if len(request.files) > 0:
+            # web-ui form only allows one file, but technically it's possible to
+            # upload multiple files at the same time
+            for filename, f in request.files.items():
+                buckets = json.loads(f.stream.read())["buckets"]
+                app.api.import_all(buckets)
+        # Normal import from body
+        else:
+            buckets = request.get_json()["buckets"]
+            app.api.import_all(buckets)
+        return None, 200
 
 
 # LOGGING