From bba801f2d2d4bfc9e47863f92208a6450ae8c10c Mon Sep 17 00:00:00 2001
From: hasslesstech <hasslesstech@tutanota.com>
Date: Sat, 28 Dec 2024 18:50:46 +0200
Subject: [PATCH] add balance entity

Add respecting schema, model and adapt all endpoints to utilize it.
Add two new endpoints for increasing and checking the balance.
---
 app/__init__.py | 107 +++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 102 insertions(+), 5 deletions(-)

diff --git a/app/__init__.py b/app/__init__.py
index 1c9cb05..e98637b 100644
--- a/app/__init__.py
+++ b/app/__init__.py
@@ -15,6 +15,7 @@ class UserModel(db.Model):
     __tablename__ = "user"
     uuid = db.Column(db.String(32), unique=True, primary_key=True, nullable=False)
     name = db.Column(db.String(64), nullable=False)
+    bal_uuid = db.Column(db.String(32), db.ForeignKey('balance.uuid'))
 
 class CategoryModel(db.Model):
     __tablename__ = "category"
@@ -29,9 +30,16 @@ class RecordModel(db.Model):
     date = db.Column(db.Date)
     amount = db.Column(db.Integer)
 
+class BalanceModel(db.Model):
+    __tablename__ = "balance"
+    uuid = db.Column(db.String(32), primary_key=True, nullable=False)
+    value = db.Column(db.Integer, nullable=False)
+
+
 class UserSchema(Schema):
     uuid = fields.Str()
     name = fields.Str()
+    bal_uuid = fields.Str()
 
 class CategorySchema(Schema):
     uuid = fields.Str()
@@ -44,6 +52,10 @@ class RecordSchema(Schema):
     date = fields.Date()
     amount = fields.Integer()
 
+class BalanceSchema(Schema):
+    uuid = fields.Str()
+    value = fields.Integer()
+
 user_schema = UserSchema()
 users_schema = UserSchema(many = True)
 
@@ -53,6 +65,8 @@ categories_schema = CategorySchema(many = True)
 record_schema = RecordSchema()
 records_schema = RecordSchema(many = True)
 
+balance_schema = BalanceSchema()
+
 # "migration"
 with app.app_context():
     db.create_all()
@@ -69,6 +83,7 @@ def ep_reset():
     db.session.query(RecordModel).delete()
     db.session.query(CategoryModel).delete()
     db.session.query(UserModel).delete()
+    db.session.query(BalanceModel).delete()
     db.session.commit()
     return {}, 200
 
@@ -96,7 +111,9 @@ def ep_user_post():
     if 'uuid' in body:
         return {}, 403
 
+    b = BalanceModel(uuid=uuid.uuid4().hex, value=0)
     body.update({'uuid': uuid.uuid4().hex})
+    body.update({'bal_uuid': b.uuid})
 
     try:
         _ = user_schema.load(body)
@@ -106,6 +123,7 @@ def ep_user_post():
     u = UserModel(**body)
 
     try:
+        db.session.add(b)
         db.session.add(u)
         db.session.commit()
     except Exception as e:
@@ -124,8 +142,13 @@ def ep_user_delete(user_id):
     if len(result) == 0:
         return {}, 404
 
-    db.session.query(UserModel).filter(UserModel.uuid == user_id).delete()
-    db.session.commit()
+    try:
+        db.session.query(UserModel).filter(UserModel.uuid == user_id).delete()
+        db.session.query(BalanceModel).filter(BalanceModel.uuid == result[0].bal_uuid).delete()
+        db.session.commit()
+    except Exception as e:
+        db.session.rollback()
+        return {}, 403
 
     return user_schema.dumps(result[0]), 200
 
@@ -188,8 +211,11 @@ def ep_category_delete():
     if len(result) == 0:
         return {}, 404
 
-    db.session.query(CategoryModel).filter(CategoryModel.uuid == cat_id).delete()
-    db.session.commit()
+    try:
+        db.session.query(CategoryModel).filter(CategoryModel.uuid == cat_id).delete()
+        db.session.commit()
+    except Exception as e:
+        return {}, 403
 
     return category_schema.dumps(result[0]), 200
 
@@ -266,12 +292,25 @@ def ep_record_post():
 
     try:
         _ = record_schema.load(body)
-
     except Exception as e:
         return {}, 403
 
     r = RecordModel(**body)
 
+    b_id = db.session \
+            .query(UserModel) \
+            .filter(UserModel.uuid == body['user_uuid']) \
+            .all()[0] \
+            .bal_uuid
+
+    v = db.session \
+            .query(BalanceModel) \
+            .filter(BalanceModel.uuid == b_id) \
+            .all()[0] \
+            .value
+
+    BalanceModel.metadata.tables.get("balance").update().where(BalanceModel.metadata.tables.get("balance").c.uuid == b_id).values(value = v-body['amount'])
+
     try:
         db.session.add(r)
         db.session.commit()
@@ -279,4 +318,62 @@ def ep_record_post():
         db.session.rollback()
         return {}, 403
 
+
     return jsonify(record_schema.load(body)), 200
+
+@app.route("/balance_up", methods = ["POST"])
+def ep_balance_up():
+    body = request.json
+
+    if 'user_id' in body:
+        body.update({'user_uuid': body['user_id']})
+        del body['user_id']
+
+    if 'user_uuid' not in body:
+        return {}, 403
+
+    try:
+        b_id = db.session \
+                .query(UserModel) \
+                .filter(UserModel.uuid == body['user_uuid']) \
+                .all()[0] \
+                .bal_uuid
+
+        v = db.session \
+                .query(BalanceModel) \
+                .filter(BalanceModel.uuid == b_id) \
+                .all()[0] \
+                .value
+
+        BalanceModel.metadata.tables.get("balance").update().where(BalanceModel.metadata.tables.get("balance").c.uuid == b_id).values(value = v + body['amount'])
+    except Exception as e:
+        return {}, 403
+
+    return {}, 200
+
+@app.route("/balance", methods = ["GET"])
+def ep_balance_get():
+    body = request.json
+
+    if 'user_id' in body:
+        body.update({'user_uuid': body['user_id']})
+        del body['user_id']
+
+    if 'user_uuid' not in body:
+        return {}, 403
+
+    try:
+        b_id = db.session \
+                .query(UserModel) \
+                .filter(UserModel.uuid == body['user_uuid']) \
+                .all()[0] \
+                .bal_uuid
+
+        result = db.session.query(BalanceModel).filter(BalanceModel.uuid == b_id).all()
+    except Exception as e:
+        return {}, 403
+
+    if len(result) == 1:
+        return user_schema.dumps(result[0]), 200
+    else:
+        return {}, 404