lab_ai/apps/update_delete.py

from flask import current_app, jsonify, request
from helpers import CORS_HEADERS, db_row_to_dict, init_db_connection
from psycopg2 import IntegrityError
from pydantic_core import ValidationError
from schemas import AiUserUpdate


def main():
    """
    ```fission
    {
      "name": "ai-admin-update-delete-user",
      "fntimeout": 300,
      "http_triggers": {
        "ai-admin-update-delete-user-http": {
          "url": "/ai/admin/users/{UserID}",
          "methods": ["DELETE", "PUT"]
        }
      }
    }
    ```
    """
    try:
        if request.method == "DELETE":
            return make_delete_request()
        elif request.method == "PUT":
            return make_update_request()
        else:
            return {"error": "Method not allow"}, 405, CORS_HEADERS
    except Exception as err:
        print(f"ErrorType={type(err)}")
        return {"error": str(err)}, 500, CORS_HEADERS


def make_update_request():
    user_id = request.headers.get("X-Fission-Params-UserID")
    if not user_id:
        return jsonify({"errorCode": "MISSING_USER_ID"}), 400, CORS_HEADERS

    try:
        body = AiUserUpdate(**(request.get_json(silent=True) or {}))
    except ValidationError as e:
        return (
            jsonify({"error": "Validation failed", "details": e.errors()}),
            400,
            CORS_HEADERS,
        )
    conn = None
    try:
        conn = init_db_connection()
        with conn:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT * FROM ai_user WHERE id=%s FOR UPDATE", (user_id,))
                row = cur.fetchone()
                if not row:
                    return jsonify({"errorCode": "USER_NOT_FOUND"}), 404, CORS_HEADERS


                sets, params = [], {"id": user_id}
                if body.name is not None:
                    sets.append("name=%(name)s")
                    params["name"] = body.name

                if body.email is not None:
                    sets.append("email=%(email)s")
                    params["email"] = body.email

                if body.dob is not None:
                    sets.append("dob=%(dob)s")
                    params["dob"] = body.dob

                if body.gender is not None:
                    sets.append("gender=%(gender)s")
                    params["gender"] = body.gender

                sets.append("modified=CURRENT_TIMESTAMP")
                cur.execute(
                    f"UPDATE ai_user SET {', '.join(sets)} WHERE id=%(id)s RETURNING *",
                    params,
                )
                updated = db_row_to_dict(cur, cur.fetchone())
                return jsonify(updated), 200, CORS_HEADERS
    except IntegrityError as e:
        return (
            jsonify({"errorCode": "DUPLICATE_USER", "details": str(e)}),
            409,
            CORS_HEADERS,
        )
    finally:
        if conn:
            conn.close()




def __delete_user(cursor, id: str):
    cursor.execute("SELECT 1 FROM ai_user WHERE id = %(id)s", {"id": id})
    if not cursor.fetchone():
        return "USER_NOT_FOUND"

    cursor.execute("DELETE FROM ai_user WHERE id = %(id)s RETURNING *", {"id": id})
    row = cursor.fetchone()
    return db_row_to_dict(cursor, row)

def make_delete_request():

    user_id = request.headers.get("X-Fission-Params-UserID")
    if not user_id:
        return jsonify({"errorCode": "MISSING_USER_ID"}), 400, CORS_HEADERS

    conn = None
    try:
        conn = init_db_connection()
        with conn.cursor() as cursor:
            result = __delete_user(cursor, id=user_id)
            if result == "USER_NOT_FOUND":
                return jsonify({"errorCode": "USER_NOT_FOUND"}), 404
            conn.commit()
            return jsonify(result), 200
    except Exception as ex:
        return jsonify({"error": str(ex)}), 500
    finally:
        if conn is not None:
            conn.close()
            current_app.logger.info("Close DB connection")