diff --git a/aurweb/db.py b/aurweb/db.py index 9ca51de2..3f5731a9 100644 --- a/aurweb/db.py +++ b/aurweb/db.py @@ -1,8 +1,10 @@ import math import aurweb.config +import aurweb.util -engine = None # See get_engine +# See get_engine. +engine = None # ORM Session class. Session = None @@ -10,6 +12,44 @@ Session = None # Global ORM Session object. session = None +# Global introspected object memo. +introspected = dict() + + +def make_random_value(table: str, column: str): + """ Generate a unique, random value for a string column in a table. + + This can be used to generate for example, session IDs that + align with the properties of the database column with regards + to size. + + Internally, we use SQLAlchemy introspection to look at column + to decide which length to use for random string generation. + + :return: A unique string that is not in the database + """ + global introspected + + # Make sure column is converted to a string for memo interaction. + scolumn = str(column) + + # If the target column is not yet introspected, store its introspection + # object into our global `introspected` memo. + if scolumn not in introspected: + from sqlalchemy import inspect + target_column = scolumn.split('.')[-1] + col = list(filter(lambda c: c.name == target_column, + inspect(table).columns))[0] + introspected[scolumn] = col + + col = introspected.get(scolumn) + length = col.type.length + + string = aurweb.util.make_random_string(length) + while session.query(table).filter(column == string).first(): + string = aurweb.util.make_random_string(length) + return string + def query(model, *args, **kwargs): return session.query(model).filter(*args, **kwargs) diff --git a/aurweb/models/session.py b/aurweb/models/session.py new file mode 100644 index 00000000..60749303 --- /dev/null +++ b/aurweb/models/session.py @@ -0,0 +1,25 @@ +from sqlalchemy import Column, Integer +from sqlalchemy.orm import backref, mapper, relationship + +from aurweb.db import make_random_value +from aurweb.models.user import User +from aurweb.schema import Sessions + + +class Session: + UsersID = Column(Integer, nullable=True) + + def __init__(self, **kwargs): + self.UsersID = kwargs.get("UsersID") + self.SessionID = kwargs.get("SessionID") + self.LastUpdateTS = kwargs.get("LastUpdateTS") + + +mapper(Session, Sessions, primary_key=[Sessions.c.SessionID], properties={ + "User": relationship(User, backref=backref("session", + uselist=False)) +}) + + +def generate_unique_sid(): + return make_random_value(Session, Session.SessionID) diff --git a/aurweb/util.py b/aurweb/util.py new file mode 100644 index 00000000..65f18a4c --- /dev/null +++ b/aurweb/util.py @@ -0,0 +1,7 @@ +import random +import string + + +def make_random_string(length): + return ''.join(random.choices(string.ascii_lowercase + + string.digits, k=length)) diff --git a/test/test_session.py b/test/test_session.py new file mode 100644 index 00000000..560f628c --- /dev/null +++ b/test/test_session.py @@ -0,0 +1,56 @@ +""" Test our Session model. """ +from datetime import datetime +from unittest import mock + +import pytest + +from aurweb.models.account_type import AccountType +from aurweb.models.session import generate_unique_sid +from aurweb.testing import setup_test_db +from aurweb.testing.models import make_session, make_user + +user, _session = None, None + + +@pytest.fixture(autouse=True) +def setup(): + from aurweb.db import session + + global user, _session + + setup_test_db("Users", "Sessions") + + account_type = session.query(AccountType).filter( + AccountType.AccountType == "User").first() + user = make_user(Username="test", Email="test@example.org", + ResetKey="testReset", Passwd="testPassword", + AccountType=account_type) + _session = make_session(UsersID=user.ID, SessionID="testSession", + LastUpdateTS=datetime.utcnow()) + + +def test_session(): + assert _session.SessionID == "testSession" + assert _session.UsersID == user.ID + + +def test_session_user_association(): + # Make sure that the Session user attribute is correct. + assert _session.User == user + + +def test_generate_unique_sid(): + # Mock up aurweb.models.session.generate_sid by returning + # sids[i % 2] from 0 .. n. This will swap between each sid + # between each call. + sids = ["testSession", "realSession"] + i = 0 + + def mock_generate_sid(length): + nonlocal i + sid = sids[i % 2] + i += 1 + return sid + + with mock.patch("aurweb.util.make_random_string", mock_generate_sid): + assert generate_unique_sid() == "realSession"