refactoring

.env

 secret=ASLKDHSADKJFLKSJDFKJSDLFKSJ
-algorithm=HS256
\ No newline at end of file
+algorithm=HS256
+ADMIN_PASS=12345678
\ No newline at end of file

.gitignore

@@ -4,6 +4,8 @@
 .env
 .ruff_cache/
 
+.env
+
 # Packaging files:
 *.egg*
 

core/models.py

-from pydantic import BaseModel, EmailStr
+from pydantic import BaseModel, EmailStr, Field
 
 
 class User(BaseModel):
@@ -7,10 +7,6 @@ class User(BaseModel):
     password: str
 
 
-class DataCenter(BaseModel):
-    name: str
-
-
 class Token(BaseModel):
     datacenter: str
     token: str
@@ -22,3 +18,9 @@ class UserLoginSchema(BaseModel):
 
     class Config:
         schema_extra = {"example": {"username": "user", "password": "weakpassword"}}
+
+class DataCenterSchema(BaseModel):
+    name: str = Field(min_length=3, max_length=50)
+    url: str
+    active: bool = False
+    weight: float = 1.0
\ No newline at end of file

envtemplate

+secret=
+algorithm=
+ADMIN_PASS=
\ No newline at end of file

main.py

 from fastapi import FastAPI, Response, Body, Depends
+import os
 from typing import List
 from balancer.util import proxy_datacenters
-from sredis.sredis import check_user, create_puser, add_datacenter, set_token
+from sredis.sredis import check_user, create_puser, add_datacenter, set_token, create_superpuser
 import logging
 import requests
-from core.models import User, DataCenter, Token, UserLoginSchema
+from core.models import User, Token, UserLoginSchema, DataCenterSchema
 from core.auth import signJWT
 from core.bearer import get_current_user, admin_user
 from redis_om import Migrator
@@ -20,6 +21,9 @@ app = FastAPI()
 add_datacenter("https://kappa1.fured.cloud.bme.hu")
 add_datacenter("https://kappa2.fured.cloud.bme.hu")
 add_datacenter("https://kappa3.fured.cloud.bme.hu")
+create_superpuser(
+    User(username='admin', email='example@domain.com', password=os.getenv('ADMIN_PASS'))
+)
 
 
 @app.post("/user/signup", tags=["user"])
@@ -46,8 +50,8 @@ def proxy_post(server_path: str = "/", username=Depends(get_current_user)):
 
 
 @app.post("/add_datacenter/")
-def create_datacenter(dc: DataCenter = None, username=Depends(admin_user)):
-    add_datacenter(dc.name)
+def create_datacenter(dc: DataCenterSchema = None, username=Depends(admin_user)):
+    add_datacenter(dc.url)
     return Response(status_code=201)
 
 

pyproject.toml

@@ -26,7 +26,7 @@ shell = "poetry run uvicorn main:app --reload --port 6973 --host 0.0.0.0"
 help   = "Start the microservice on port 6973"
 
 [tool.poe.tasks.prod]
-shell = "poetry run uvicorn main:app --host 0.0.0.0 --port 8080 --workers 4 --log-level critical"
+shell = "poetry run uvicorn main:app --host 0.0.0.0 --port 8080 --workers 8 --log-level critical"
 help   = "Start the prod deployment"
 
 [tool.poe.tasks.lint]

sredis/models.py

@@ -6,4 +6,9 @@ class PUser(HashModel):
     username: str = Field(index=True)
     email: EmailStr
     password: str
-    admin: bool = False
+    admin: int = 0
+
+class DataCenter(HashModel):
+    name: str = Field(index=True)
+    url: str = Field(index=True)
+    active: bool = Field(index=True)
\ No newline at end of file

sredis/sredis.py

@@ -2,9 +2,10 @@ import redis
 from core.models import UserLoginSchema
 from .models import PUser
 from passlib.hash import pbkdf2_sha256
-from core.models import User
+from core.models import User, DataCenterSchema
 from core.auth import hash_pass
 from fastapi import HTTPException
+import random
 
 r = redis.Redis(host="localhost", port=6379, db=0, decode_responses=True)
 r.set("datacenters_cnt", 1)
@@ -19,7 +20,6 @@ def add_datacenter(datacenter: str):
     r.incr("datacenters_cnt")
     r.set("roundrobin_cnt", 1)
 
-
 def rr_get():
     cnt = int(r.get("datacenters_cnt"))
     rr = int(r.get("roundrobin_cnt"))
@@ -29,6 +29,10 @@ def rr_get():
         r.incr("roundrobin_cnt")
     return str(r.hget("datacenters_hash", rr))
 
+def wr_get(centers: DataCenterSchema):
+    cnt = int(r.get("datacenters_cnt"))
+    indexis = range(1, cnt+1)
+    return random.choice(indexis, weights=())[0] 
 
 def set_token(username: str, datacenter: str, token: str):
     print(f"tokens:{username}" + datacenter)
@@ -54,3 +58,14 @@ def create_puser(user: User):
         username=user.username, email=user.email, password=hash_pass(user.password)
     )
     user.save()
+
+
+def create_superpuser(user: User):
+    s = PUser.find(PUser.username == user.username).all()
+    if not s:
+        puser = PUser(
+            username=user.username, email=user.email, password=hash_pass(user.password), admin=1
+        )
+        puser.save()
+        return True
+    return False
\ No newline at end of file

stresstest/main.py

 import json
-from locust import TaskSet, task, between, HttpUser
+from locust import TaskSet, task, HttpUser
 
 class PerformanceTests(TaskSet):
     def on_start(self):
@@ -15,6 +15,4 @@ class PerformanceTests(TaskSet):
 
 
 class WebsiteUser(HttpUser):
-    tasks = [PerformanceTests]
-    min_wait = 5000
-    max_wait = 9000
\ No newline at end of file
+    tasks = [PerformanceTests]
\ No newline at end of file