Add possibility to manage own vote

7 files changed, 134 insertions, 21 deletions

diff --git a/app.py b/app.py
index bfb075f..f8422a9 100644
--- a/app.py
+++ b/app.py
@@ -111,6 +111,13 @@ def poll(id):
   if not validate_uuid(id):
     return error_page("Invalid poll ID", 400)
 
+  prefill_voter_name = request.args.get("prefill_voter_name")
+
+  voter_codes: list[str] = []
+  for k, _ in request.cookies.items():
+    if k.startswith("diddle_voter_code_"):
+      voter_codes.append(k.replace("diddle_voter_code_", ""))
+
   poll = db.get_poll(id)
   if poll is None:
     return error_page("Poll not found", 404)
@@ -127,11 +134,15 @@ def poll(id):
 
   voter_names_set: set[str] = set()
   selections: dict[VoterNameChoiceIdPair, int] = {}
+  managed_voter_names: dict[str, str] = {}
   for choice in poll.choices:
     for vote in choice.votes:
       selections[(vote.voter_name, choice.id)] = vote.value
       voter_names_set.add(vote.voter_name)
 
+      if vote.manage_code in voter_codes:
+        managed_voter_names[vote.voter_name] = vote.manage_code
+
   voter_names = list(voter_names_set)
   voter_names.sort()
 
@@ -140,7 +151,9 @@ def poll(id):
                     poll=poll,
                     selections=selections,
                     choices=poll.choices,
+                    prefill_voter_name=prefill_voter_name,
                     voter_names=voter_names,
+                    managed_voter_names=managed_voter_names,
                     now=datetime.datetime.now(),
                     display_mode=display_mode))
 
@@ -175,11 +188,36 @@ def vote_poll(id):
       choice_id = k.replace("choice_", "")
       selections[choice_id] = 1
 
-  db.vote_poll(id, voter_name, selections)
+  manage_code = db.vote_poll(id, voter_name, selections)
 
   email_client.send_participation_email_if_enabled(poll_id=id, voter_name=voter_name)
 
-  return redirect(f"/poll/{id}")
+  response = make_response(
+    redirect(f"/poll/{id}")
+  )
+  response.set_cookie(f"diddle_voter_code_{manage_code}", "1",
+                      samesite="Strict", secure=False)
+  return response
+
+@app.post("/poll/<id>/delete_voter")
+def delete_voter(id):
+  if not validate_uuid(id):
+    return error_page("Invalid poll ID", 400)
+
+  voter_manage_code = request.form["voter_code"]
+
+  voter_name = db.get_voter_name_by_manage_code(voter_manage_code)
+  if voter_name is None:
+    return error_page("Voter not found", 404)
+
+  db.delete_voter(voter_manage_code)
+
+  resp = make_response(
+    redirect(f"/poll/{id}?prefill_voter_name={voter_name}")
+  )
+  resp.set_cookie(f"diddle_voter_code_{voter_manage_code}", "", expires=0,
+                  samesite="Strict", secure=False)
+  return resp
 
 @app.post("/manage/<code>/update_info")
 def update_poll_info(code):
diff --git a/db.py b/db.py
index d91692a..5effa5a 100644
--- a/db.py
+++ b/db.py
@@ -3,6 +3,7 @@ from typing import Literal
 from dataclasses import dataclass
 import datetime
 import psycopg2
+import uuid
 
 BASE_URL = os.environ["BASE_URL"]
 
@@ -55,6 +56,7 @@ class Vote:
   choice_id: str
   voter_name: str
   value: int # 0 or 1
+  manage_code: str
 
 @dataclass
 class Choice:
@@ -138,6 +140,7 @@ def tuple_to_vote(vote_t: tuple) -> Vote:
     choice_id=vote_t[2],
     voter_name=vote_t[3],
     value=vote_t[4],
+    manage_code=vote_t[5]
   )
 
 def get_poll(id: str):
@@ -205,16 +208,18 @@ def create_poll(title: str,
       conn.rollback()
       raise e
 
-def vote_poll(poll_id: str, voter_name: str, selections: dict[str, int]) -> None:
+def vote_poll(poll_id: str, voter_name: str, selections: dict[str, int]) -> str:
+  """Returns the manage code of the vote."""
   with db.cursor() as (conn, cur):
     try:
+      manage_code = str(uuid.uuid4())
       for choice_id in selections:
         value = selections[choice_id]
-        cur.execute("INSERT INTO votes (poll_id, voter_name, choice_id, value)"
-                    "VALUES (%s, %s, %s, %s)",
-                    (poll_id, voter_name, choice_id, value))
-
+        cur.execute("INSERT INTO votes (poll_id, voter_name, choice_id, value, manage_code)"
+                    "VALUES (%s, %s, %s, %s, %s)",
+                    (poll_id, voter_name, choice_id, value, manage_code))
       conn.commit()
+      return manage_code
     except Exception as e:
       conn.rollback()
       raise e
@@ -310,6 +315,26 @@ def delete_poll(code: str) -> None:
       conn.rollback()
       raise e
 
+def get_voter_name_by_manage_code(voter_manage_code: str) -> str | None:
+  with db.cursor() as (conn, cur):
+    try:
+      cur.execute("SELECT voter_name FROM votes WHERE manage_code = %s", (voter_manage_code,))
+      voter_name = cur.fetchone()
+      conn.commit()
+      return voter_name[0] if voter_name else None
+    except Exception as e:
+      conn.rollback()
+      raise e
+
+def delete_voter(voter_manage_code: str) -> None:
+  with db.cursor() as (conn, cur):
+    try:
+      cur.execute("DELETE FROM votes WHERE manage_code = %s", (voter_manage_code,))
+      conn.commit()
+    except Exception as e:
+      conn.rollback()
+      raise e
+
 ### Migrations
 
 def ensure_migration_table_exists() -> None:
diff --git a/migrations/0004_vote_manage_code.sql b/migrations/0004_vote_manage_code.sql
new file mode 100644
index 0000000..69ce38f
--- /dev/null
+++ b/migrations/0004_vote_manage_code.sql
@@ -0,0 +1 @@
+ALTER TABLE votes ADD COLUMN manage_code uuid DEFAULT uuid_generate_v4();
diff --git a/static/styles.css b/static/styles.css
index fae665b..73feb61 100644
--- a/static/styles.css
+++ b/static/styles.css
@@ -122,6 +122,10 @@ table {
   margin: 0;
 }
 
+.vote-table input[name="voter_name"] {
+  max-width: 150px;
+}
+
 .danger-zone {
   margin-top: 40px;
 }
@@ -183,4 +187,10 @@ button#share-link-copy {
   color: black;
   font-size: 25px;
   line-height: 16px;
-}
\ No newline at end of file
+}
+
+input[type="submit"].delete-voter-btn {
+  margin-left: 5px;
+  background: none;
+  padding: 0;
+}
diff --git a/templates/poll.html.j2 b/templates/poll.html.j2
index e0ede1f..1661b79 100644
--- a/templates/poll.html.j2
+++ b/templates/poll.html.j2
@@ -21,14 +21,23 @@
   </form>
 </div>
 
-<form action="/poll/{{ poll.id }}/vote" method="post">
-  {% if display_mode == "table" %}
-  {% include "poll_vote_table.html.j2" %}
-  {% else %}
-  {% include "poll_vote_list.html.j2" %}
-  {% endif %}
-</form>
+
+{% if display_mode == "table" %}
+{% include "poll_vote_table.html.j2" %}
+{% else %}
+{% include "poll_vote_list.html.j2" %}
 {% endif %}
 
+<script>
+  const query = new URLSearchParams(window.location.search);
+  const prefill_voter_name = query.get("prefill_voter_name");
+  if (prefill_voter_name !== null) {
+    document.querySelector("input[name=voter_name]").focus()
+    document.querySelector("input[name=voter_name]").select();
+    query.delete("prefill_voter_name");
+    history.replaceState(null, "", window.location.pathname + (query.toString() ? "?" + query.toString() : ""));
+  }
+  </script>
 
+{% endif %}
 {% endblock %}
diff --git a/templates/poll_vote_list.html.j2 b/templates/poll_vote_list.html.j2
index 9d3a15c..ef09af4 100644
--- a/templates/poll_vote_list.html.j2
+++ b/templates/poll_vote_list.html.j2
@@ -23,8 +23,25 @@
 </div>
 
 <div>
-  <label for="voter_name" hidden>Your name</label>
-  <input type="text" name="voter_name" placeholder="Your name" required>
-  <input class="green" type="submit" value="Submit">
+  <form action="/poll/{{ poll.id }}/vote" method="post">
+    <label for="voter_name" hidden>Your name</label>
+    <input type="text" name="voter_name" placeholder="Your name" required
+           {% if prefill_voter_name %}value="{{ prefill_voter_name }}"{% endif %}>
+    <input class="green" type="submit" value="Submit">
+  </form>
 </div>
 
+{% if managed_voter_names | length != 0 %}
+<br>
+<br>
+<h3>Remove submission(s)</h3>
+{% endif %}
+<div>
+  {% for voter_name in managed_voter_names %}
+  <form action="/poll/{{ poll.id }}/delete_voter" method="post">
+    <input type="hidden" name="voter_code" value="{{ managed_voter_names[voter_name] }}">
+    {{ voter_name }}
+    <input class="delete-voter-btn" type="submit" value="❌">
+  </form>
+  {% endfor %}
+</div>
\ No newline at end of file
diff --git a/templates/poll_vote_table.html.j2 b/templates/poll_vote_table.html.j2
index a7aa434..440f2b7 100644
--- a/templates/poll_vote_table.html.j2
+++ b/templates/poll_vote_table.html.j2
@@ -12,7 +12,17 @@
   </tr>
   {% for voter_name in voter_names %}
   <tr>
-    <td>{{ voter_name }}</td>
+    <td>
+      {% if voter_name in managed_voter_names %}
+      <form action="/poll/{{ poll.id }}/delete_voter" method="post">
+        <input type="hidden" name="voter_code" value="{{ managed_voter_names[voter_name] }}">
+        {{ voter_name }}
+        <input class="delete-voter-btn" type="submit" value="❌">
+      </form>
+      {% else %}
+      {{ voter_name }}
+      {% endif %}
+    </td>
     {% for choice in choices %}
     {% if selections[(voter_name, choice.id)] == 1 %}
     <td>
@@ -32,10 +42,13 @@
   {% endfor %}
 
   <!-- Add a row for the current user -->
+  <form action="/poll/{{ poll.id }}/vote" method="post">
   <tr>
     <td>
       <label for="voter_name" hidden>Your name</label>
-      <input type="text" name="voter_name" placeholder="Your name" required>
+      <input type="text" name="voter_name" placeholder="Your name" required
+             {% if prefill_voter_name %}value="{{ prefill_voter_name }}"{% endif %}>
+      <input class="green" type="submit" value="Submit">
     </td>
     {% for choice in choices %}
     <td>
@@ -43,7 +56,7 @@
     </td>
     {% endfor %}
   </tr>
+  </form>
   <!-- End of the row for the current user -->
 </table>
 
-<input class="green" type="submit" value="Submit">