Allow user/dataset selection with dropdown

2 files changed, 15 insertions, 7 deletions

diff --git a/app/admin/forms.py b/app/admin/forms.py
index 6b39e2e..6e48ee9 100644
--- a/app/admin/forms.py
+++ b/app/admin/forms.py
@@ -6,8 +6,8 @@ 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 wtforms import StringField, SubmitField, SelectField
+from wtforms.validators import DataRequired, ValidationError, InputRequired
 
 from werkzeug.utils import secure_filename
 
@@ -16,8 +16,10 @@ from app.admin.datasets import validate_dataset, get_name_from_dataset
 
 
 class AdminAssignTaskForm(FlaskForm):
-    username = StringField("Username", validators=[DataRequired()])
-    dataset = StringField("Dataset", validators=[DataRequired()])
+    username = SelectField(
+        "Username", coerce=int, validators=[InputRequired()]
+    )
+    dataset = SelectField("Dataset", coerce=int, validators=[InputRequired()])
     submit = SubmitField("Assign")
 
 
diff --git a/app/admin/routes.py b/app/admin/routes.py
index 8d40665..7617943 100644
--- a/app/admin/routes.py
+++ b/app/admin/routes.py
@@ -17,20 +17,26 @@ from app.models import User, Dataset, Task
 @bp.route("/assign", methods=("GET", "POST"))
 @admin_required
 def assign():
+    user_list = [(u.id, u.username) for u in User.query.all()]
+    dataset_list = [(d.id, d.name) for d in Dataset.query.all()]
+
     form = AdminAssignTaskForm()
+    form.username.choices = user_list
+    form.dataset.choices = dataset_list
+
     if form.validate_on_submit():
-        user = User.query.filter_by(username=form.username.data).first()
+        user = User.query.filter_by(id=form.username.data).first()
         if user is None:
             flash("User does not exist.")
             return redirect(url_for("admin.assign"))
-        dataset = Dataset.query.filter_by(name=form.dataset.data).first()
+        dataset = Dataset.query.filter_by(id=form.dataset.data).first()
         if dataset is None:
             flash("Dataset does not exist.")
             return redirect(url_for("admin.assign"))
 
         task = Task.query.filter_by(
             annotator_id=user.id, dataset_id=dataset.id
-        )
+        ).first()
         if not task is None:
             flash("Task assignment already exists.")
             return redirect(url_for("admin.assign"))