Merge pull request #2 from MissterHao/feature/backends

Fix problems while saving session data into dynamodb

README.md

@@ -40,3 +40,15 @@
 ## Installation
 
 ## Example
+
+
+```python
+INSTALLED_APPS = [
+    ...
+    "dysession", # add dysession to installed apps
+    # 'django.contrib.sessions', # remove this default session
+    ...
+]
+
+SESSION_ENGINE = "dysession.backends.db"
+```

dysession/aws/dynamodb.py

@@ -1,13 +1,17 @@
 from datetime import datetime
-from typing import Any, Dict, Literal, Optional, Union
+import logging
+from typing import Any, Callable, Dict, Literal, Optional, Union
 
 import boto3
 from botocore import client as botoClitent
 from django.utils import timezone
 
 from dysession.aws.error import DynamodbItemNotFound, DynamodbTableNotFound
-from dysession.backends.error import (SessionKeyDoesNotExist,
-                                      SessionKeyDuplicated)
+from dysession.backends.error import (
+    SessionExpired,
+    SessionKeyDoesNotExist,
+    SessionKeyDuplicated,
+)
 from dysession.backends.model import SessionDataModel
 
 from ..settings import get_config
@@ -78,33 +82,36 @@ def key_exists(session_key: str, table_name: Optional[str] = None, client=None)
         },
         ProjectionExpression=f"{pk}",
     )
-
     return "Item" in response
 
 
-def get_item(session_key: str, table_name: Optional[str] = None, client=None) -> bool:
-
-    if client is None:
-        client = boto3.client("dynamodb", region_name=get_config()["DYNAMODB_REGION"])
+def get_item(session_key: str, table_name: Optional[str] = None) -> SessionDataModel:
 
     if table_name is None:
         table_name = get_config()["DYNAMODB_TABLENAME"]
 
     assert type(session_key) is str, "session_key should be string type"
 
+    logging.info("Get Item from DynamoDB")
+
     pk = get_config()["PARTITION_KEY_NAME"]
 
-    response = client.get_item(
-        TableName=table_name,
+    resource = boto3.resource("dynamodb", region_name=get_config()["DYNAMODB_REGION"])
+    table = resource.Table(table_name)
+
+    response = table.get_item(
         Key={
-            pk: {"S": session_key},
+            pk: session_key,
         },
     )
 
     if "Item" not in response:
         raise DynamodbItemNotFound()
 
-    return response
+    model = SessionDataModel(session_key=session_key)
+    for k, v in response["Item"].items():
+        model[k] = v
+    return model
 
 
 def insert_session_item(
@@ -119,6 +126,9 @@ def insert_session_item(
     if table_name is None:
         table_name = get_config()["DYNAMODB_TABLENAME"]
 
+    if key_exists(data.session_key):
+        raise SessionKeyDuplicated
+
     resource = boto3.resource("dynamodb", region_name=get_config()["DYNAMODB_REGION"])
     table = resource.Table(table_name)
     pk = get_config()["PARTITION_KEY_NAME"]
@@ -137,28 +147,54 @@ def insert_session_item(
 
 
 class DynamoDB:
-    def __init__(self, client) -> None:
+    def __init__(self, client=None) -> None:
         self.client = client
 
     def get(
-        self, session_key: Optional[str] = None, ttl: Optional[datetime] = None
+        self,
+        session_key: Optional[str] = None,
+        table_name: Optional[str] = None,
+        expired_time_fn: Callable[[], datetime] = datetime.now,
     ) -> Dict[str, Any]:
         """Return session data if dynamodb partision key is matched with inputed session_key"""
         if session_key is None:
             raise ValueError("session_key should be str type")
 
+        if table_name is None:
+            table_name = get_config()["DYNAMODB_TABLENAME"]
 
-        # if not found then raise
-        # raise SessionKeyDoesNotExist
-        # if key is expired
-        # raise SessionExpired
-
-    def set(self, session_key: Optional[str] = None, session_data=None) -> None:
-        return
-        # Partision key duplicated
-        raise SessionKeyDuplicated
+        now = expired_time_fn()
 
-    def exists(self, session_key: Optional[str] = None) -> bool:
-        return False
+        try:
+            model = get_item(session_key=session_key, table_name=table_name)
+            if get_config()["TTL_ATTRIBUTE_NAME"] in model:
+                time = model[get_config()["TTL_ATTRIBUTE_NAME"]]
+                if time < int(now.timestamp()):
+                    raise SessionExpired
         # if not found then raise
-        raise SessionKeyDoesNotExist
+        except DynamodbItemNotFound:
+            raise SessionKeyDoesNotExist
+        # if key is expired
+        except SessionExpired:
+            raise SessionExpired
+
+        return model
+
+    def set(
+        self,
+        data: SessionDataModel,
+        table_name: Optional[str] = None,
+        return_consumed_capacity: Literal["INDEXES", "TOTAL", "NONE"] = "TOTAL",
+        ignore_duplicated: bool = True,
+    ) -> None:
+        try:
+            insert_session_item(data, table_name, return_consumed_capacity)
+        except SessionKeyDuplicated:
+            if not ignore_duplicated:
+                raise SessionKeyDuplicated
+
+    def exists(self, session_key: str) -> bool:
+        if type(session_key) is not str:
+            raise TypeError(f"session_key should be type of str instead of {type(session_key)}.")
+
+        return key_exists(session_key=session_key)

dysession/backends/db.py

@@ -14,19 +14,22 @@
     SessionKeyDuplicated,
 )
 from dysession.backends.model import SessionDataModel
+from dysession.settings import get_config
 
 
 class SessionStore(SessionBase):
     """Implement DynamoDB session store"""
 
     def __init__(self, session_key: Optional[str], **kwargs: Any) -> None:
         super().__init__(session_key, **kwargs)
-        # self.client = boto3.client("dynamodb")
-        self.db = DynamoDB(client=boto3.client("dynamodb"))
+        self.db = DynamoDB(
+            client=boto3.client("dynamodb", region_name=get_config()["DYNAMODB_REGION"])
+        )
+        self._get_session()
 
     def _get_session_from_ddb(self) -> SessionDataModel:
         try:
-            return self.db.get(session_key=self.session_key, ttl=timezone.now())
+            return self.db.get(session_key=self.session_key)
         except (SessionKeyDoesNotExist, SessionExpired, SuspiciousOperation) as e:
             if isinstance(e, SuspiciousOperation):
                 logger = logging.getLogger(f"django.security.{e.__class__.__name__}")
@@ -40,10 +43,12 @@ def _get_session(self, no_load=False) -> SessionDataModel:
         """
         self.accessed = True
         try:
-            return self._session_cache
+            if isinstance(self._session_cache, SessionDataModel):
+                return self._session_cache
+            raise AttributeError
         except AttributeError:
             if self.session_key is None or no_load:
-                self._session_cache = SessionDataModel()
+                self._session_cache = SessionDataModel(self.session_key)
             else:
                 self._session_cache = self.load()
         return self._session_cache
@@ -63,7 +68,13 @@ def clear(self):
         super().clear()
         self._session_cache = SessionDataModel()
 
-    # ====== Methods that subclass must implement
+    def items(self):
+        return self._session.items()
+
+    def __str__(self):
+        return str(self._get_session())
+
+    # Methods that subclass must implement
     def exists(self, session_key: str) -> bool:
         """
         Return True if the given session_key already exists.
@@ -87,22 +98,24 @@ def create(self) -> None:
             self.modified = True
             return
 
-    def save(self, must_create: bool = ...) -> None:
+    def save(self, must_create: bool = False) -> None:
         """
         Save the session data. If 'must_create' is True, create a new session
         object (or raise CreateError). Otherwise, only update an existing
         object and don't create one (raise UpdateError if needed).
         """
         try:
-            self.db.set(
-                session_key=self._session_key,
-                session_data=self._get_session(must_create),
-            )
+            if self._session_key is None:
+                return self.create()
+
+            data = self._get_session(no_load=must_create)
+            data.session_key = self._session_key
+            self.db.set(data=data)
         except SessionKeyDuplicated:
             if must_create:
                 raise SessionKeyDuplicated
 
-    def delete(self, request, *args, **kwargs):
+    def delete(self, session_key=None):
         """
         Delete the session data under this key. If the key is None, use the
         current session key value.

dysession/backends/model.py

@@ -1,21 +1,35 @@
+import json
 from typing import Any, Optional
 
 
 class SessionDataModel:
+
+    NOTFOUND_ALLOW_LIST = ["_auth_user_id", "_auth_user_backend", "_auth_user_hash"]
+
     def __init__(self, session_key: Optional[str] = None) -> None:
 
         if type(session_key) is not str and session_key is not None:
             raise TypeError("session_key should be type str or None")
 
         self.session_key = session_key
-        self.__variables_names = set()
+        self.__variables_names = set(["session_key"])
 
     def __getitem__(self, key) -> Any:
-        return getattr(self, key)
+        # Set SESSION_EXPIRE_AT_BROWSER_CLOSE to False
+        # https://docs.djangoproject.com/en/4.1/topics/http/sessions/#browser-length-sessions-vs-persistent-sessions
+        if key == "_session_expiry":
+            return False
+
+        try:
+            return getattr(self, key)
+        except AttributeError:
+            if key in self.NOTFOUND_ALLOW_LIST:
+                raise KeyError
+            raise
 
     def __setitem__(self, key, value):
-        if key == "session_key":
-            raise ValueError()
+        # if key == "session_key":
+        #     raise ValueError()
 
         setattr(self, key, value)
         self.__variables_names.add(key)
@@ -28,7 +42,7 @@ def __iter__(self):
         return iter(self.__variables_names)
 
     def __is_empty(self):
-        return len(self.__variables_names) == 0
+        return "session_key" in self.__variables_names and len(self.__variables_names) == 1
 
     is_empty = property(__is_empty)
 
@@ -49,3 +63,14 @@ def pop(self, key, default=...):
             if default is Ellipsis:
                 raise
             return default
+
+    def items(self):
+        for key in self.__variables_names:
+            yield (key, self[key])
+
+    def __str__(self) -> str:
+        data = {}
+        for key in self.__variables_names:
+            data[key] = getattr(self, key)
+
+        return json.dumps(data)

dysession/management/commands/dysession_clear.py

@@ -23,9 +23,6 @@ def add_arguments(self, parser: CommandParser) -> None:
     def handle(self, *args: Any, **options: Any) -> Optional[str]:
         userids = options.get("uid", None)
         if userids:
-            print(f"Ready to clear {userids} session data.")
-            return
-
-        print("Clearing whole session data")
-        return 
+            ...
+        return
 

dysession/settings.py

@@ -17,6 +17,19 @@
 
 @lru_cache
 def get_config() -> Dict[str, Union[str, int]]:
+    """Return cached django-dysession config in dictionary type
+
+    Contain Items:
+        * DYNAMODB_TABLENAME
+        * PARTITION_KEY_NAME
+        * SORT_KEY_NAME
+        * TTL_ATTRIBUTE_NAME
+        * CACHE_PERIOD
+        * DYNAMODB_REGION
+
+    Returns:
+        Dict[str, Union[str, int]]
+    """
     config = DEFAULT_CONFIG.copy()
     custom_config = getattr(settings, "DYSESSION_CONFIG", {})
     config.update(custom_config)