Allow admin to add datasets

1 files changed, 41 insertions, 2 deletions

diff --git a/app/admin/forms.py b/app/admin/forms.py
index b33b1c6..6b39e2e 100644
--- a/app/admin/forms.py
+++ b/app/admin/forms.py
@@ -1,12 +1,51 @@
 # -*- coding: utf-8 -*-
 
+import os
+
+from flask import current_app
 from flask_wtf import FlaskForm
+from flask_wtf.file import FileField, FileRequired
+
+from wtforms import StringField, SubmitField
+from wtforms.validators import DataRequired, ValidationError
+
+from werkzeug.utils import secure_filename
 
-from wtforms import StringField, PasswordField, BooleanField, SubmitField
-from wtforms.validators import DataRequired, ValidationError, Email, EqualTo
+from app.models import Dataset
+from app.admin.datasets import validate_dataset, get_name_from_dataset
 
 
 class AdminAssignTaskForm(FlaskForm):
     username = StringField("Username", validators=[DataRequired()])
     dataset = StringField("Dataset", validators=[DataRequired()])
     submit = SubmitField("Assign")
+
+
+class AdminAddDatasetForm(FlaskForm):
+    file_ = FileField("File", validators=[FileRequired()])
+    submit = SubmitField("Upload")
+
+    def validate_file_(self, field):
+        filename = secure_filename(field.data.filename)
+        if not "." in filename:
+            raise ValidationError("The file should be a '.json' file.")
+        if not filename.rsplit(".", 1)[1].lower() == "json":
+            raise ValidationError("The file should be a '.json' file.")
+
+        temp_filename = os.path.join(
+            current_app.instance_path, current_app.config["TEMP_DIR"], filename
+        )
+
+        field.data.save(temp_filename)
+        error = validate_dataset(temp_filename)
+        if not error is None:
+            os.unlink(temp_filename)
+            raise ValidationError("Error validating dataset: %s" % error)
+
+        name = get_name_from_dataset(temp_filename)
+        dataset = Dataset.query.filter_by(name=name).first()
+        if dataset is not None:
+            os.unlink(temp_filename)
+            raise ValidationError(
+                "A dataset with the name '%s' already exists." % name
+            )