1 files changed, 51 insertions, 34 deletions
diff --git a/python/Dawn/dawn.py b/python/Dawn/dawn.py
index 6d28d6dc..7b465e36 100644
--- a/python/Dawn/dawn.py
+++ b/python/Dawn/dawn.py
@@ -1,15 +1,18 @@
-from flask import Flask, request, render_template
+from flask import Flask, request, render_template, redirect, url_for, send_from_directory
 from models import db, ModelForm, BondDeal, Counterparties
 from sqlalchemy.sql import text
 from sqlalchemy import create_engine
+from sqlalchemy.exc import IntegrityError
 import pandas as pd
+import os
+import datetime
 
 app = Flask(__name__)
 app.config['SQLALCHEMY_DATABASE_URI'] = 'postgresql://dawn_user@debian/dawndb'
 app.config['SECRET_KEY'] = 'papa'
 app.config['UPLOAD_FOLDER'] = 'tickets'
 
-def list_counterparties():
+def cp_choices():
     return Counterparties.query.order_by('name').\
         with_entities(Counterparties.code, Counterparties.name)
 
@@ -58,41 +61,53 @@ class CounterpartyForm(ModelForm):
     class Meta:
         model = Counterparties
 
-@app.route('/', methods=['GET', 'POST'])
-def trade_entry():
-    bond_form = BondForm()
-    bond_form.counterparty.choices = list_counterparties()
-    if bond_form.is_submitted():
-        if bond_form.validate():
-            bond = BondDeal()
-            bond_form.populate_obj(bond)
-            bond.ticket.save('test.pdf')
-            bond.ticket='test.pdf'
-            bond_form.get_session().add(bond)
-            bond_form.get_session().commit()
-            return "Success!"
-        else:
-            print(bond_form.errors)
-    return render_template("trade_entry.html", form=bond_form)
-
-@app.route('/trade/<int:tradeid>', methods=['GET', 'POST'])
-def trade_update(tradeid):
-    trade = BondDeal()
-    bond_form = BondForm(obj = BondDeal.query.get(tradeid))
-    bond_form.counterparty.choices = list_counterparties()
+@app.route('/trades/<int:tradeid>', methods=['GET', 'POST'])
+@app.route('/trades/', defaults={'tradeid': None}, methods=['GET', 'POST'])
+def trade_manage(tradeid):
+    if tradeid:
+        trade = BondDeal.query.get(tradeid)
+        bond_form = BondForm(obj = BondDeal.query.get(tradeid))
+        #we change the default to update
+        bond_form.action.data = 'UPDATE'
+    else:
+        trade = BondDeal()
+        bond_form = BondForm()
+    bond_form.cp_code.choices = cp_choices()
     if bond_form.is_submitted():
         if bond_form.validate():
+            if tradeid:
+                old_ticket_name = trade.ticket
             bond_form.populate_obj(trade)
+            session = bond_form.get_session()
+            if tradeid and trade.ticket.filename == '':
+                trade.ticket = old_ticket_name
+            else:
+                ticket_name =  "{0} {1}.pdf".format(trade.description,
+                                                  str(datetime.datetime.now()))
+                trade.ticket.save(os.path.join(app.config['UPLOAD_FOLDER'],
+                                               ticket_name))
+                trade.ticket = ticket_name
+                if not tradeid:
+                    session.add(trade)
+            try:
+                session.commit()
+            except IntegrityError:
+                print("TODO: fix this")
             return redirect(url_for('list_trades'))
         else:
-            return "something wrong!"
-    return render_template("trade_entry.html", form=bond_form)
+            return str(bond_form.errors)
+    return render_template('trade_entry.html', form=bond_form, trade_id=tradeid)
 
 @app.route('/blotter/')
 def list_trades():
     trade_list = BondDeal.query.order_by(BondDeal.trade_date)
     return render_template('blotter.html', trades=trade_list.all())
 
+@app.route('/tickets/<path:filename>')
+def download_ticket(filename):
+    return send_from_directory(app.config['UPLOAD_FOLDER'],
+                               filename, as_attachment=True)
+
 @app.route('/counterparties/')
 def list_counterparties():
     cp_list = Counterparties.query.order_by(Counterparties.name)
@@ -100,18 +115,20 @@ def list_counterparties():
 
 @app.route('/edit_cp/<cpcode>', methods=['GET', 'POST'])
 def edit_counterparty(cpcode):
-    cp_form = CounterpartyForm(obj = Counterparties.query.get(cpcode))
+    cp = Counterparties.query.get(cpcode)
+    cp_form = CounterpartyForm(obj = cp)
     if cp_form.is_submitted():
         if cp_form.validate():
-            cp = Counterparties()
             cp_form.populate_obj(cp)
-            return
-    return render_template('edit_cp.html', form=cp_form)
+            session = cp_form.get_session()
+            session.commit()
+            return redirect(url_for('list_counterparties'))
+    return render_template('edit_cp.html', form=cp_form, code=cpcode)
 
 if __name__=="__main__":
     db.init_app(app)
-    db.drop_all(app=app)
-    db.create_all(app=app)
-    add_triggers()
-    load_counterparties()
+    # db.drop_all(app=app)
+    # db.create_all(app=app)
+    # add_triggers()
+    # load_counterparties()
     app.run(debug=True)