From 5e2f824f5b0a88b58bd16cdc7905d5c96aea3c95 Mon Sep 17 00:00:00 2001
From: knst-kotov <knst.kotov@gmail.com>
Date: Wed, 24 Nov 2021 15:56:09 +0300
Subject: [PATCH] add UserStorage

---
 auth/authenticate.py |  7 ++++---
 orm/__init__.py      |  3 ++-
 orm/user.py          | 48 +++++++++++++++++++++++++++++++++++---------
 resolvers/auth.py    | 12 ++++++-----
 resolvers/profile.py |  7 +------
 5 files changed, 53 insertions(+), 24 deletions(-)

diff --git a/auth/authenticate.py b/auth/authenticate.py
index ad07a39e..3e5a43a6 100644
--- a/auth/authenticate.py
+++ b/auth/authenticate.py
@@ -10,7 +10,7 @@ from auth.credentials import AuthCredentials, AuthUser
 from auth.token import Token
 from auth.authorize import Authorize
 from exceptions import InvalidToken, OperationNotAllowed
-from orm import User
+from orm import User, UserStorage
 from orm.base import local_session
 from redis import redis
 from settings import JWT_AUTH_HEADER, EMAIL_TOKEN_LIFE_SPAN
@@ -70,8 +70,9 @@ class JWTAuthenticate(AuthenticationBackend):
 		if not payload.device in ("pc", "mobile"):
 			return AuthCredentials(scopes=[]), AuthUser(user_id=None)
 
-		scopes = User.get_permission(user_id=payload.user_id)
-		return AuthCredentials(user_id=payload.user_id, scopes=scopes, logged_in=True), AuthUser(user_id=payload.user_id)
+		user = await UserStorage.get_user(payload.user_id)
+		scopes = user.get_permission()
+		return AuthCredentials(user_id=payload.user_id, scopes=scopes, logged_in=True), user
 
 class EmailAuthenticate:
 	@staticmethod
diff --git a/orm/__init__.py b/orm/__init__.py
index a18511a9..3f24bccd 100644
--- a/orm/__init__.py
+++ b/orm/__init__.py
@@ -1,6 +1,6 @@
 from orm.rbac import Operation, Resource, Permission, Role
 from orm.community import Community
-from orm.user import User, UserRating, UserRole
+from orm.user import User, UserRating, UserRole, UserStorage
 from orm.message import Message
 from orm.topic import Topic, TopicSubscription
 from orm.notification import Notification
@@ -18,3 +18,4 @@ Resource.init_table()
 with local_session() as session:
 	ShoutRatingStorage.init(session)
 	ShoutViewStorage.init(session)
+	UserStorage.init(session)
diff --git a/orm/user.py b/orm/user.py
index 73b7af93..2d00b648 100644
--- a/orm/user.py
+++ b/orm/user.py
@@ -2,13 +2,15 @@ from typing import List
 from datetime import datetime
 
 from sqlalchemy import Table, Column, Integer, String, ForeignKey, Boolean, DateTime, JSON as JSONType
-from sqlalchemy.orm import relationship
+from sqlalchemy.orm import relationship, selectinload
 
 from orm import Permission
 from orm.base import Base, local_session
 from orm.rbac import Role
 from orm.topic import Topic
 
+import asyncio
+
 class UserNotifications(Base):
 	__tablename__ = 'user_notifications'
 
@@ -62,17 +64,45 @@ class User(Base):
 	old_id: str = Column(String, nullable = True)
 
 	@classmethod
-	def get_permission(cls, user_id):
+	def get_permission(self):
 		scope = {}
-		with local_session() as session:
-			user = session.query(User).filter(User.id == user_id).first()
-			for role in user.roles:
-				for p in role.permissions:
-					if not p.resource_id in scope:
-						scope[p.resource_id] = set()
-					scope[p.resource_id].add(p.operation_id)
+		#TODO implement RoleStorage
+		#for role in self.roles:
+			# for p in role.permissions:
+				# if not p.resource_id in scope:
+					# scope[p.resource_id] = set()
+				# scope[p.resource_id].add(p.operation_id)
 		return scope
 
+class UserStorage:
+	users = {}
+	lock = asyncio.Lock()
+
+	@staticmethod
+	def init(session):
+		self = UserStorage
+		users = session.query(User).\
+			options(selectinload(User.roles)).all()
+		self.users = dict([(user.id, user) for user in users])
+
+	@staticmethod
+	async def get_user(id):
+		self = UserStorage
+		async with self.lock:
+			return self.users.get(id)
+
+	@staticmethod
+	async def add_user(user):
+		self = UserStorage
+		async with self.lock:
+			self.users[id] = user
+
+	@staticmethod
+	async def del_user(user):
+		self = UserStorage
+		async with self.lock:
+			del self.users[id]
+
 
 if __name__ == "__main__":
 	print(User.get_permission(user_id=1))
diff --git a/resolvers/auth.py b/resolvers/auth.py
index de064704..56184b86 100644
--- a/resolvers/auth.py
+++ b/resolvers/auth.py
@@ -40,12 +40,14 @@ async def register(*_, email: str, password: str = ""):
 	if not password:
 		user = User.create(**user_dict)
 		await send_confirm_email(user)
+		UserStorage.add_user(user)
 		return { "user": user }
-	else:
-		user_dict["password"] = Password.encode(password)
-		user = User.create(**user_dict)
-		token = await Authorize.authorize(user)
-		return {"user": user, "token": token }
+
+	user_dict["password"] = Password.encode(password)
+	user = User.create(**user_dict)
+	token = await Authorize.authorize(user)
+	UserStorage.add_user(user)
+	return {"user": user, "token": token }
 
 
 @query.field("signIn")
diff --git a/resolvers/profile.py b/resolvers/profile.py
index f73e3c2b..26091326 100644
--- a/resolvers/profile.py
+++ b/resolvers/profile.py
@@ -22,12 +22,7 @@ async def get_user_by_slug(_, info, slug):
 @query.field("getCurrentUser")
 @login_required
 async def get_current_user(_, info):
-	auth = info.context["request"].auth
-	user_id = auth.user_id
-
-	with local_session() as session:
-		user = session.query(User).filter(User.id == user_id).first()
-
+	user = info.context["request"].user
 	return { "user": user }
 
 @query.field("userRoles")