1 files changed, 85 insertions, 0 deletions
diff --git a/scripts/add_problem.py b/scripts/add_problem.py
new file mode 100755
index 0000000..d8462cf
--- /dev/null
+++ b/scripts/add_problem.py
@@ -0,0 +1,85 @@
+#!/usr/bin/python3
+# coding=utf-8
+
+import db
+from .utils import filenamefy
+
+conn = db.get_connection()
+try:
+    cur = conn.cursor()
+    try:
+        new_lidentifier = None
+        # Get or add language.
+        cur.execute('select id, name, identifier from language order by id asc')
+        languages = cur.fetchall()
+        print('Languages:')
+        for lid, lname, lidentifier in languages:
+            print('  {}: {}'.format(lid, lname))
+
+        new_lid = input("Enter language ID or 'n' for new): ")
+        if new_lid == 'n':
+            new_lname = input('Enter name of the new language: ')
+            new_lidentifier = filenamefy(new_lname)
+            cur.execute('insert into language (name, identifier) values (%s, %s) returning id',
+                        (new_lname, new_lidentifier))
+            new_lid = cur.fetchone()[0]
+            print('Added new language "{}" with ID {} and identifier {}'.format(
+                new_lname, new_lid, new_lidentifier))
+        else:
+            new_lid = int(new_lid)
+            for lid, lname, lidentifier in languages:
+                print(lid, lname, lidentifier)
+                if lid == new_lid:
+                    new_lidentifier = lidentifier
+                    break
+        if new_lidentifier is None:
+            raise Exception('Language with ID {} does not exist'.format(new_lid))
+        print('Selected langauge {}'.format(new_lid))
+        print()
+
+        # Get or add problem group.
+        new_gidentifier = None
+        cur.execute('select id, name, identifier from problem_group order by id asc')
+        groups = cur.fetchall()
+        print('Problem groups:')
+        for gid, gname, gidentifier in groups:
+            print('  {}: {}'.format(gid, gname))
+        new_gid = input("Enter problem group ID or 'n' for new): ")
+        if new_gid == 'n':
+            new_gname = input('Enter name of the new problem group: ')
+            new_gidentifier = filenamefy(new_gname)
+            cur.execute('insert into problem_group (name, identifier) values (%s, %s) returning id',
+                        (new_gname, new_gidentifier))
+            new_gid = cur.fetchone()[0]
+            print('Added new problem group "{}" with ID {} and identifier {}'.format(
+                new_gname, new_gid, new_gidentifier))
+        else:
+            new_gid = int(new_gid)
+            for gid, gname, gidentifier in groups:
+                if gid == new_gid:
+                    new_gidentifier = gidentifier
+                    break
+        if new_gidentifier is None:
+            raise Exception('Group with ID {} does not exist'.format(new_gid))
+        print('Selected problem group {}'.format(new_gid))
+        print()
+
+        # Add problem.
+        new_pname = input('Enter name of the new problem: ')
+        new_pidentifier = filenamefy(new_pname)
+        cur.execute('insert into problem (language_id, problem_group_id, name, identifier, is_visible) values (%s, %s, %s, %s, %s) returning id',
+                    (new_lid, new_gid, new_pname, new_pidentifier, True))
+        new_pid = cur.fetchone()[0]
+        print('Added new problem "{}" with ID {} and identifier {}'.format(
+            new_pname, new_pid, new_pidentifier))
+        print('Data files should be placed in "{}/problems/{}/{}"'.format(
+            new_lidentifier, new_gidentifier, new_pidentifier))
+
+    finally:
+        cur.close()
+    conn.commit()
+except:
+    conn.rollback()
+    raise
+finally:
+    db.return_connection(conn)