From 6710627f55e7d09b14e1ef1bfdd50f522d893baa Mon Sep 17 00:00:00 2001 From: skypank Date: Thu, 9 Jul 2026 20:18:11 +0530 Subject: [PATCH 1/2] feat(db): add User model + resource-selection storage + login upsert (#586) Persist accounts on OIDC login and store a per-user resource (standard) selection. Additive to the existing session auth: the callback upserts the User row and sets session['user_id'] only when CRE_ENABLE_LOGIN is on; the existing session keys and 401 behaviour are unchanged. - User(google_sub UNIQUE, email, display_name, created_at, last_seen_at) - UserResourceSelection(user_id FK CASCADE, standard_name, UNIQUE(user_id, name)) - Node_collection: upsert_user / get_user_by_sub / get_user_resource_selection / set_user_resource_selection - Alembic: dedicated merge of the two open heads + create tables (up/down verified on Postgres; guardrail green) --- application/database/db.py | 113 ++++++++++++ application/tests/user_model_test.py | 164 ++++++++++++++++++ application/tests/web_main_test.py | 62 +++++++ application/web/web_main.py | 13 ++ ...d4e5f6_add_users_and_resource_selection.py | 50 ++++++ .../f0e1d2c3b4a5_merge_heads_before_users.py | 27 +++ 6 files changed, 429 insertions(+) create mode 100644 application/tests/user_model_test.py create mode 100644 migrations/versions/a1b2c3d4e5f6_add_users_and_resource_selection.py create mode 100644 migrations/versions/f0e1d2c3b4a5_merge_heads_before_users.py diff --git a/application/database/db.py b/application/database/db.py index 6600c3920..760e874ef 100644 --- a/application/database/db.py +++ b/application/database/db.py @@ -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, unique=True) + 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 @@ -992,6 +1035,76 @@ 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) + user = self.session.query(User).filter(User.google_sub == google_sub).first() + if user is None: + 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) + else: + if email: + user.email = email + if display_name is not None: + user.display_name = display_name + user.last_seen_at = now + self.session.commit() + 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]] = [] diff --git a/application/tests/user_model_test.py b/application/tests/user_model_test.py new file mode 100644 index 000000000..bf241de03 --- /dev/null +++ b/application/tests/user_model_test.py @@ -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() diff --git a/application/tests/web_main_test.py b/application/tests/web_main_test.py index 0fab4fc74..5f3c2225f 100644 --- a/application/tests/web_main_test.py +++ b/application/tests/web_main_test.py @@ -1151,6 +1151,68 @@ def test_gap_analysis_weak_links_response(self, db_mock) -> None: self.assertEqual(200, response.status_code) self.assertEqual(expected, json.loads(response.data)) + @patch("application.web.web_main.id_token") + @patch("application.web.web_main.CREFlow") + def test_callback_upserts_user_and_sets_session( + self, cre_flow_mock, id_token_mock + ) -> None: + id_token_mock.verify_oauth2_token.return_value = { + "sub": "sub-xyz", + "name": "Test User", + "email": "test@example.com", + } + flow_instance = cre_flow_mock.instance.return_value + flow_instance.flow.credentials._id_token = "tok" + self.app.secret_key = "test-secret" + with patch.dict( + os.environ, + { + "CRE_ENABLE_LOGIN": "1", + "LOGIN_ALLOWED_DOMAINS": "*", + "NO_LOAD_GRAPH_DB": "1", + "INSECURE_REQUESTS": "1", + }, + ): + with self.app.test_client() as client: + with client.session_transaction() as sess: + sess["state"] = "xyz" + client.get("/rest/v1/callback?state=xyz") + with client.session_transaction() as sess: + self.assertIn("user_id", sess) + + users = sqla.session.query(db.User).all() + self.assertEqual(len(users), 1) + self.assertEqual(users[0].google_sub, "sub-xyz") + self.assertEqual(users[0].email, "test@example.com") + self.assertEqual(users[0].display_name, "Test User") + + @patch("application.web.web_main.id_token") + @patch("application.web.web_main.CREFlow") + def test_callback_does_not_persist_when_login_flag_off( + self, cre_flow_mock, id_token_mock + ) -> None: + id_token_mock.verify_oauth2_token.return_value = { + "sub": "sub-xyz", + "name": "Test User", + "email": "test@example.com", + } + flow_instance = cre_flow_mock.instance.return_value + flow_instance.flow.credentials._id_token = "tok" + self.app.secret_key = "test-secret" + env = { + "LOGIN_ALLOWED_DOMAINS": "*", + "NO_LOAD_GRAPH_DB": "1", + "INSECURE_REQUESTS": "1", + } + with patch.dict(os.environ, env): + os.environ.pop("CRE_ENABLE_LOGIN", None) + with self.app.test_client() as client: + with client.session_transaction() as sess: + sess["state"] = "xyz" + client.get("/rest/v1/callback?state=xyz") + + self.assertEqual(sqla.session.query(db.User).count(), 0) + def test_deeplink(self) -> None: self.maxDiff = None collection = db.Node_collection().with_graph() diff --git a/application/web/web_main.py b/application/web/web_main.py index 012438476..fddd46773 100644 --- a/application/web/web_main.py +++ b/application/web/web_main.py @@ -1132,6 +1132,19 @@ def callback(): 401, description=f"You need an account with one of the following providers to access this functionality {allowed_domains}", ) + + # Persist the account when login is enabled; the session keeps working + # unchanged if this no-ops (flag off) or fails. + if is_login_enabled(): + try: + user = db.Node_collection().upsert_user( + google_sub=id_info.get("sub"), + email=id_info.get("email") or "", + display_name=id_info.get("name"), + ) + session["user_id"] = user.id + except Exception as e: + logger.error(f"failed to persist user on login: {e}") return redirect("/chatbot") diff --git a/migrations/versions/a1b2c3d4e5f6_add_users_and_resource_selection.py b/migrations/versions/a1b2c3d4e5f6_add_users_and_resource_selection.py new file mode 100644 index 000000000..e333e6890 --- /dev/null +++ b/migrations/versions/a1b2c3d4e5f6_add_users_and_resource_selection.py @@ -0,0 +1,50 @@ +"""add users and user_resource_selection tables (issue #586) + +Revision ID: a1b2c3d4e5f6 +Revises: f0e1d2c3b4a5 +Create Date: 2026-07-08 + +""" + +from alembic import op +import sqlalchemy as sa + + +revision = "a1b2c3d4e5f6" +down_revision = "f0e1d2c3b4a5" +branch_labels = None +depends_on = None + + +def upgrade(): + op.create_table( + "users", + sa.Column("id", sa.String(), primary_key=True), + sa.Column("google_sub", sa.String(), nullable=False), + sa.Column("email", sa.String(), nullable=False), + sa.Column("display_name", sa.String(), nullable=True), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.Column("last_seen_at", sa.DateTime(), nullable=False), + sa.UniqueConstraint("google_sub", name="uq_users_google_sub"), + ) + op.create_table( + "user_resource_selection", + sa.Column("id", sa.String(), primary_key=True), + sa.Column("user_id", sa.String(), nullable=False), + sa.Column("standard_name", sa.String(), nullable=False), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["user_id"], + ["users.id"], + onupdate="CASCADE", + ondelete="CASCADE", + ), + sa.UniqueConstraint( + "user_id", "standard_name", name="uq_user_resource_selection" + ), + ) + + +def downgrade(): + op.drop_table("user_resource_selection") + op.drop_table("users") diff --git a/migrations/versions/f0e1d2c3b4a5_merge_heads_before_users.py b/migrations/versions/f0e1d2c3b4a5_merge_heads_before_users.py new file mode 100644 index 000000000..3f33cf2c3 --- /dev/null +++ b/migrations/versions/f0e1d2c3b4a5_merge_heads_before_users.py @@ -0,0 +1,27 @@ +"""merge embeddings heads (ab12cd34ef56, 9b1c2d3e4f50) before users + +No schema changes; collapses the two open heads into a single lineage so the +users migration has a single parent and ``flask db downgrade`` is unambiguous. + +Revision ID: f0e1d2c3b4a5 +Revises: ab12cd34ef56, 9b1c2d3e4f50 +Create Date: 2026-07-08 + +""" + +from alembic import op +import sqlalchemy as sa + + +revision = "f0e1d2c3b4a5" +down_revision = ("ab12cd34ef56", "9b1c2d3e4f50") +branch_labels = None +depends_on = None + + +def upgrade(): + pass + + +def downgrade(): + pass From 2c0cd6e532eaa0c8b4f1efaf07e381a73ef499de Mon Sep 17 00:00:00 2001 From: skypank Date: Thu, 9 Jul 2026 21:00:28 +0530 Subject: [PATCH 2/2] fix(db): concurrency-safe upsert_user + drop redundant unique constraint (#586) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address PR review: - upsert_user now mirrors add_node: catch IntegrityError, rollback, re-query by google_sub, update profile — safe under concurrent logins. - Drop column-level unique=True on User.google_sub; keep only the named constraint so create_all matches the Alembic migration (no schema drift). - Login callback skips persistence when the OIDC 'sub' claim is missing. - Callback test asserts session['user_id'] equals the persisted User.id. --- application/database/db.py | 53 ++++++++++++++++++++---------- application/tests/web_main_test.py | 4 +++ application/web/web_main.py | 22 ++++++++----- 3 files changed, 54 insertions(+), 25 deletions(-) diff --git a/application/database/db.py b/application/database/db.py index 56d3ea336..054dadf4e 100644 --- a/application/database/db.py +++ b/application/database/db.py @@ -263,7 +263,7 @@ class User(BaseModel): # type: ignore __tablename__ = "users" id = sqla.Column(sqla.String, primary_key=True, default=generate_uuid) - google_sub = sqla.Column(sqla.String, nullable=False, unique=True) + 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) @@ -1046,24 +1046,43 @@ def upsert_user( from datetime import datetime, timezone now = datetime.now(timezone.utc) - user = self.session.query(User).filter(User.google_sub == google_sub).first() - if user is None: - 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) - else: + + def _apply_profile(u: User) -> None: if email: - user.email = email + u.email = email if display_name is not None: - user.display_name = display_name - user.last_seen_at = now - self.session.commit() + 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]: diff --git a/application/tests/web_main_test.py b/application/tests/web_main_test.py index 5f3c2225f..6b33c26d7 100644 --- a/application/tests/web_main_test.py +++ b/application/tests/web_main_test.py @@ -1173,18 +1173,22 @@ def test_callback_upserts_user_and_sets_session( "INSECURE_REQUESTS": "1", }, ): + session_user_id = None with self.app.test_client() as client: with client.session_transaction() as sess: sess["state"] = "xyz" client.get("/rest/v1/callback?state=xyz") with client.session_transaction() as sess: self.assertIn("user_id", sess) + session_user_id = sess["user_id"] users = sqla.session.query(db.User).all() self.assertEqual(len(users), 1) self.assertEqual(users[0].google_sub, "sub-xyz") self.assertEqual(users[0].email, "test@example.com") self.assertEqual(users[0].display_name, "Test User") + # Session id must match the persisted row (guards PR2's /user wiring). + self.assertEqual(session_user_id, users[0].id) @patch("application.web.web_main.id_token") @patch("application.web.web_main.CREFlow") diff --git a/application/web/web_main.py b/application/web/web_main.py index 546926532..e1c737538 100644 --- a/application/web/web_main.py +++ b/application/web/web_main.py @@ -1165,15 +1165,21 @@ def callback(): # Persist the account when login is enabled; the session keeps working # unchanged if this no-ops (flag off) or fails. if is_login_enabled(): - try: - user = db.Node_collection().upsert_user( - google_sub=id_info.get("sub"), - email=id_info.get("email") or "", - display_name=id_info.get("name"), + google_sub = id_info.get("sub") + if not google_sub: + logger.error( + "OIDC callback returned no 'sub' claim; skipping user persistence" ) - session["user_id"] = user.id - except Exception as e: - logger.error(f"failed to persist user on login: {e}") + else: + try: + user = db.Node_collection().upsert_user( + google_sub=google_sub, + email=id_info.get("email") or "", + display_name=id_info.get("name"), + ) + session["user_id"] = user.id + except Exception as e: + logger.error(f"failed to persist user on login: {e}") return redirect("/chatbot")