Skip to content
Open
143 changes: 143 additions & 0 deletions application/database/db.py
Original file line number Diff line number Diff line change
Expand Up @@ -254,6 +254,49 @@ class StagedChangeSet(BaseModel): # type: ignore
created_at = sqla.Column(sqla.DateTime, nullable=False)


class User(BaseModel): # type: ignore
"""Persisted account identity for OIDC login (issue #586, RFC #876 TODO 1).

``google_sub`` is the immutable OIDC ``sub`` claim (the same value stored in
``session['google_id']``); email can change, so identity is anchored on the sub.
"""

__tablename__ = "users"
id = sqla.Column(sqla.String, primary_key=True, default=generate_uuid)
google_sub = sqla.Column(sqla.String, nullable=False)
email = sqla.Column(sqla.String, nullable=False)
display_name = sqla.Column(sqla.String, nullable=True)
created_at = sqla.Column(sqla.DateTime, nullable=False)
last_seen_at = sqla.Column(sqla.DateTime, nullable=False)

resource_selection = sqla.relationship(
"UserResourceSelection",
cascade="all, delete-orphan",
)

__table_args__ = (sqla.UniqueConstraint(google_sub, name="uq_users_google_sub"),)


class UserResourceSelection(BaseModel): # type: ignore
"""A single standard name a user has chosen to keep in view (issue #586)."""

__tablename__ = "user_resource_selection"
id = sqla.Column(sqla.String, primary_key=True, default=generate_uuid)
user_id = sqla.Column(
sqla.String,
sqla.ForeignKey("users.id", onupdate="CASCADE", ondelete="CASCADE"),
nullable=False,
)
standard_name = sqla.Column(sqla.String, nullable=False)
created_at = sqla.Column(sqla.DateTime, nullable=False)

__table_args__ = (
sqla.UniqueConstraint(
user_id, standard_name, name="uq_user_resource_selection"
),
)


def create_import_run(source: str, version: Optional[str] = None) -> ImportRun:
"""Create and persist an import run record. Returns the new ImportRun."""
from datetime import datetime, timezone
Expand Down Expand Up @@ -992,6 +1035,106 @@ def with_graph(self) -> "Node_collection":

return self

def upsert_user(
self, google_sub: str, email: str, display_name: Optional[str] = None
) -> User:
"""Create or refresh the User row for an OIDC ``sub`` (issue #586).

Idempotent on ``google_sub``: a repeat login updates the existing row's
profile fields and ``last_seen_at`` instead of inserting a duplicate.
"""
from datetime import datetime, timezone

now = datetime.now(timezone.utc)

def _apply_profile(u: User) -> None:
if email:
u.email = email
if display_name is not None:
u.display_name = display_name
u.last_seen_at = now

user = self.session.query(User).filter(User.google_sub == google_sub).first()
if user is not None:
_apply_profile(user)
self.session.commit()
return user

user = User(
id=generate_uuid(),
google_sub=google_sub,
email=email,
display_name=display_name,
created_at=now,
last_seen_at=now,
)
self.session.add(user)
try:
self.session.commit()
except IntegrityError:
# A concurrent login inserted the same google_sub first; recover by
# rolling back and updating the now-existing row (mirrors add_node).
self.session.rollback()
existing = (
self.session.query(User).filter(User.google_sub == google_sub).first()
)
if existing is None:
raise
_apply_profile(existing)
self.session.commit()
return existing
return user

def get_user_by_sub(self, google_sub: str) -> Optional[User]:
return self.session.query(User).filter(User.google_sub == google_sub).first()

def get_user_resource_selection(self, user_id: str) -> List[str]:
"""Return the standard names a user has selected, ordered by name."""
rows = (
self.session.query(UserResourceSelection)
.filter(UserResourceSelection.user_id == user_id)
.order_by(UserResourceSelection.standard_name)
.all()
)
return [row.standard_name for row in rows]

def set_user_resource_selection(
self, user_id: str, standard_names: List[str]
) -> List[str]:
"""Replace a user's resource selection with ``standard_names`` (deduped)."""
from datetime import datetime, timezone

now = datetime.now(timezone.utc)
deduped: List[str] = []
for name in standard_names:
if name not in deduped:
deduped.append(name)

def _replace() -> None:
self.session.query(UserResourceSelection).filter(
UserResourceSelection.user_id == user_id
).delete()
for name in deduped:
self.session.add(
UserResourceSelection(
id=generate_uuid(),
user_id=user_id,
standard_name=name,
created_at=now,
)
)
self.session.commit()

try:
_replace()
except IntegrityError:
# A concurrent PUT for the same user can collide on
# uq_user_resource_selection between our delete and our inserts;
# roll back and retry once (mirrors upsert_user).
self.session.rollback()
_replace()
return self.get_user_resource_selection(user_id)
Comment thread
coderabbitai[bot] marked this conversation as resolved.

def __get_external_links(self) -> List[Tuple[CRE, Node, str]]:
external_links: List[Tuple[CRE, Node, str]] = []

Expand Down
164 changes: 164 additions & 0 deletions application/tests/user_model_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,164 @@
"""Tests for user persistence and per-user resource selection (issue #586, RFC #876 TODO 1/2)."""

import os
import unittest

from sqlalchemy.exc import IntegrityError

from application import create_app, sqla
from application.database import db


class TestUserModel(unittest.TestCase):
def setUp(self) -> None:
# These tests exercise only the SQL layer; skip the Neo4j graph load.
os.environ["NO_LOAD_GRAPH_DB"] = "1"
self.app = create_app(mode="test")
self.app_context = self.app.app_context()
self.app_context.push()
sqla.create_all()
self.collection = db.Node_collection()

def tearDown(self) -> None:
sqla.session.remove()
sqla.drop_all()
self.app_context.pop()
os.environ.pop("NO_LOAD_GRAPH_DB", None)

def test_upsert_user_creates_single_row(self) -> None:
user = self.collection.upsert_user(
google_sub="sub-123", email="a@example.com", display_name="Alice"
)
self.assertIsNotNone(user.id)
self.assertEqual(user.google_sub, "sub-123")
self.assertEqual(user.email, "a@example.com")
self.assertEqual(user.display_name, "Alice")
self.assertIsNotNone(user.created_at)
self.assertIsNotNone(user.last_seen_at)
self.assertEqual(sqla.session.query(db.User).count(), 1)

def test_upsert_user_is_idempotent_and_updates_in_place(self) -> None:
first = self.collection.upsert_user(
google_sub="sub-123", email="old@example.com", display_name="Alice"
)
first_id = first.id
created_at = first.created_at

second = self.collection.upsert_user(
google_sub="sub-123", email="new@example.com", display_name="Alice B"
)

# No duplicate row, same identity, refreshed profile fields.
self.assertEqual(sqla.session.query(db.User).count(), 1)
self.assertEqual(second.id, first_id)
self.assertEqual(second.email, "new@example.com")
self.assertEqual(second.display_name, "Alice B")
self.assertEqual(second.created_at, created_at)
self.assertGreaterEqual(second.last_seen_at, created_at)

def test_upsert_user_tolerates_missing_email_and_name(self) -> None:
user = self.collection.upsert_user(
google_sub="sub-noemail", email="", display_name=None
)
self.assertEqual(user.email, "")
self.assertIsNone(user.display_name)
self.assertEqual(sqla.session.query(db.User).count(), 1)

def test_get_user_by_sub(self) -> None:
self.collection.upsert_user(
google_sub="sub-123", email="a@example.com", display_name="Alice"
)
found = self.collection.get_user_by_sub("sub-123")
self.assertIsNotNone(found)
self.assertEqual(found.google_sub, "sub-123")
self.assertIsNone(self.collection.get_user_by_sub("does-not-exist"))

def test_resource_selection_round_trips(self) -> None:
user = self.collection.upsert_user(
google_sub="sub-123", email="a@example.com", display_name="Alice"
)
# Empty by default.
self.assertEqual(self.collection.get_user_resource_selection(user.id), [])

stored = self.collection.set_user_resource_selection(
user.id, ["ASVS", "CWE", "SAMM"]
)
self.assertEqual(sorted(stored), ["ASVS", "CWE", "SAMM"])
self.assertEqual(
self.collection.get_user_resource_selection(user.id),
["ASVS", "CWE", "SAMM"],
)

def test_set_resource_selection_replaces_previous(self) -> None:
user = self.collection.upsert_user(
google_sub="sub-123", email="a@example.com", display_name="Alice"
)
self.collection.set_user_resource_selection(user.id, ["ASVS", "CWE"])
self.collection.set_user_resource_selection(user.id, ["SAMM"])
self.assertEqual(self.collection.get_user_resource_selection(user.id), ["SAMM"])
self.assertEqual(
sqla.session.query(db.UserResourceSelection)
.filter(db.UserResourceSelection.user_id == user.id)
.count(),
1,
)

def test_set_resource_selection_dedupes_input(self) -> None:
user = self.collection.upsert_user(
google_sub="sub-123", email="a@example.com", display_name="Alice"
)
self.collection.set_user_resource_selection(user.id, ["ASVS", "ASVS", "CWE"])
self.assertEqual(
self.collection.get_user_resource_selection(user.id), ["ASVS", "CWE"]
)

def test_resource_selection_is_per_user(self) -> None:
alice = self.collection.upsert_user(
google_sub="sub-a", email="a@example.com", display_name="Alice"
)
bob = self.collection.upsert_user(
google_sub="sub-b", email="b@example.com", display_name="Bob"
)
self.collection.set_user_resource_selection(alice.id, ["ASVS"])
self.collection.set_user_resource_selection(bob.id, ["CWE"])
self.assertEqual(
self.collection.get_user_resource_selection(alice.id), ["ASVS"]
)
self.assertEqual(self.collection.get_user_resource_selection(bob.id), ["CWE"])

def test_resource_selection_unique_constraint_enforced(self) -> None:
user = self.collection.upsert_user(
google_sub="sub-123", email="a@example.com", display_name="Alice"
)
from datetime import datetime, timezone

sqla.session.add(
db.UserResourceSelection(
user_id=user.id,
standard_name="ASVS",
created_at=datetime.now(timezone.utc),
)
)
sqla.session.add(
db.UserResourceSelection(
user_id=user.id,
standard_name="ASVS",
created_at=datetime.now(timezone.utc),
)
)
with self.assertRaises(IntegrityError):
sqla.session.commit()
sqla.session.rollback()

def test_deleting_user_cascades_to_selection(self) -> None:
user = self.collection.upsert_user(
google_sub="sub-123", email="a@example.com", display_name="Alice"
)
self.collection.set_user_resource_selection(user.id, ["ASVS", "CWE"])
sqla.session.delete(user)
sqla.session.commit()
self.assertEqual(sqla.session.query(db.UserResourceSelection).count(), 0)


if __name__ == "__main__":
unittest.main()
Loading
Loading