diff --git a/app/main.py b/app/main.py index 4a75aa5..5fc40ac 100644 --- a/app/main.py +++ b/app/main.py @@ -7,11 +7,12 @@ from flask import url_for from flask_sqlalchemy import SQLAlchemy from flask.ext.security import Security, SQLAlchemyUserDatastore, \ UserMixin, RoleMixin, login_required, utils +from flask.ext.security.core import current_user import uuid import hashlib from functools import wraps from flask_wtf import Form -from wtforms import DateField, IntegerField, DecimalField +from wtforms import DateField, IntegerField, DecimalField, SelectField from wtforms.validators import DataRequired, ValidationError import os @@ -27,8 +28,8 @@ app.config.from_object(__name__) roles_users = db.Table('roles_users', - db.Column('user_id', db.Integer(), db.ForeignKey('user.id')), - db.Column('role_id', db.Integer(), db.ForeignKey('role.id'))) + db.Column('user_id', db.Integer(), db.ForeignKey('user.id')), + db.Column('role_id', db.Integer(), db.ForeignKey('role.id'))) class Role(db.Model, RoleMixin): @@ -49,6 +50,9 @@ class User(db.Model, UserMixin): password = db.Column(db.String(255)) active = db.Column(db.Boolean()) confirmed_at = db.Column(db.DateTime()) + vehicles = db.relationship( + 'Vehicle' + ) roles = db.relationship( 'Role', secondary=roles_users, @@ -56,7 +60,22 @@ class User(db.Model, UserMixin): ) def __repr__(self): - return '' % self.username + return '' % self.email + + +class Vehicle(db.Model): + id = db.Column(db.Integer, primary_key=True) + owner_id = db.Column(db.Integer, db.ForeignKey('user.id')) + name = db.Column(db.String(255), unique=True) + pitstops = db.relationship( + 'Pitstop' + ) + + def __init__(self, name): + self.name = name + + def __repr__(self): + return '' % self.name class Pitstop(db.Model): @@ -64,6 +83,7 @@ class Pitstop(db.Model): date = db.Column(db.Date) odometer = db.Column(db.Integer) litres = db.Column(db.Numeric(5,2)) + vehicle_id = db.Column(db.Integer, db.ForeignKey('vehicle.id')) def __init__(self, odometer, litres, date): self.odometer = odometer @@ -99,6 +119,23 @@ def before_first_request(): user_datastore.add_role_to_user('admin@example.com', 'admin') db.session.commit() + user = User.query.filter(User.email=='someone@example.com').first() + if len(user.vehicles) == 0: + v = Vehicle('someone') + db.session.add(v) + user.vehicles.append(v) + db.session.commit() + + user = User.query.filter(User.email=='admin@example.com').first() + if len(user.vehicles) == 0: + v = Vehicle('admin_1') + db.session.add(v) + user.vehicles.append(v) + v = Vehicle('admin_2') + db.session.add(v) + user.vehicles.append(v) + db.session.commit() + @app.before_request def before_request(): @@ -127,6 +164,7 @@ def litres_check(form, field): class CreatePitstopForm(Form): + vehicle = SelectField('Vehicle', coerce=int) date = DateField('Date of Pitstop', validators=[date_check]) odometer = IntegerField('Odometer (km)', validators=[odometer_check]) litres = DecimalField('Litres (l)', places=1, validators=[litres_check]) @@ -145,9 +183,15 @@ def create_pit_stop_form(): form = CreatePitstopForm() form.set_pitstop(last_pitstop) + print(current_user.vehicles) + + form.vehicle.choices = [(g.id, g.name) for g in current_user.vehicles] if form.validate_on_submit(): + print(form.vehicle.data) + v = Vehicle.query.filter(Vehicle.id==form.vehicle.data).first() new_stop = Pitstop(form.odometer.data, form.litres.data, form.date.data) db.session.add(new_stop) + v.pitstops.append(new_stop) db.session.commit() return redirect(url_for('get_pit_stops')) diff --git a/app/templates/newPitStopForm.html b/app/templates/newPitStopForm.html index c661f10..66e1aca 100644 --- a/app/templates/newPitStopForm.html +++ b/app/templates/newPitStopForm.html @@ -3,7 +3,8 @@ {% block body %}
{{ form.hidden_tag() }} - {{ render_field_with_errors(form.date) }} + {{ render_field_with_errors(form.vehicle) }} + {{ render_field_with_errors(form.date) }} {{ render_field_with_errors(form.odometer) }} {{ render_field_with_errors(form.litres) }}