Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
132 changes: 132 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,95 @@ 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)
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()
return self.get_user_resource_selection(user_id)

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