Merge branch 'aarthy/auth-plugin' into 'enterprise'

SSO auth support + refactor Users page

See merge request dkinternal/testgen/dataops-testgen!297

Author: ci bot

testgen/commands/run_execute_tests.py

@@ -123,7 +123,7 @@ def run_execution_steps_in_background(project_code, test_suite):
         empty_cache()
         background_thread = threading.Thread(
             target=run_execution_steps,
-            args=(project_code, test_suite, session.username),
+            args=(project_code, test_suite, session.auth.user_display),
         )
         background_thread.start()
     else:

testgen/commands/run_profiling_bridge.py

@@ -211,7 +211,7 @@ def run_profiling_in_background(table_group_id):
         empty_cache()
         background_thread = threading.Thread(
             target=run_profiling_queries,
-            args=(table_group_id, session.username),
+            args=(table_group_id, session.auth.user_display),
         )
         background_thread.start()
     else:

testgen/common/mixpanel_service.py

@@ -57,7 +57,7 @@ def send_event(self, event_name, include_usage=False, **properties):
         properties.setdefault("instance_id", self.instance_id)
         properties.setdefault("edition", settings.DOCKER_HUB_REPOSITORY)
         properties.setdefault("version", settings.VERSION)
-        properties.setdefault("username", session.username)
+        properties.setdefault("username", session.auth.user_display)
         properties.setdefault("distinct_id", self.get_distinct_id(properties["username"]))
         if include_usage:
             properties.update(self.get_usage())

testgen/common/models/entity.py

@@ -7,13 +7,14 @@
 from sqlalchemy import delete, select
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.orm import InstrumentedAttribute
-from sqlalchemy.sql.elements import BinaryExpression
+from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
 
 from testgen.common.models import Base, get_current_session
 from testgen.utils import is_uuid4, make_json_safe
 
 ENTITY_HASH_FUNCS = {
     BinaryExpression: lambda x: str(x.compile(compile_kwargs={"literal_binds": True})),
+    BooleanClauseList: lambda x: str(x.compile(compile_kwargs={"literal_binds": True})),
     tuple: lambda x: [str(y) for y in x],
 }
 

testgen/common/models/user.py

@@ -1,9 +1,11 @@
+from datetime import UTC, datetime
 from typing import Literal
 from uuid import UUID, uuid4
 
-from sqlalchemy import Column, String, asc
+from sqlalchemy import Column, String, asc, func, update
 from sqlalchemy.dialects import postgresql
 
+from testgen.common.models import get_current_session
 from testgen.common.models.custom_types import NullIfEmptyString
 from testgen.common.models.entity import Entity
 
@@ -19,5 +21,25 @@ class User(Entity):
     name: str = Column(NullIfEmptyString)
     password: str = Column(String)
     role: RoleType = Column(String)
+    latest_login: datetime = Column(postgresql.TIMESTAMP)
+
+    _get_by = "username"
+    _default_order_by = (asc(func.lower(username)),)
+
+    def save(self, update_latest_login: bool = False) -> None:
+        if self.id and not update_latest_login:
+            values = {
+                column.key: getattr(self, column.key, None)
+                for column in self.__table__.columns
+                if column != User.latest_login
+            }
+            query = update(User).where(User.id == self.id).values(**values)
+            db_session = get_current_session()
+            db_session.execute(query)
+            db_session.commit()
+            User.clear_cache()
+        else:
+            if update_latest_login:
+                self.latest_login = datetime.now(UTC)
+            super().save()
 
-    _default_order_by = (asc(username),)

testgen/template/dbsetup/030_initialize_new_schema_structure.sql

@@ -622,22 +622,15 @@ CREATE TABLE functional_test_results
 );
 
 CREATE TABLE auth_users (
-	id UUID 		DEFAULT gen_random_uuid()
+	id UUID 		 DEFAULT gen_random_uuid()
     CONSTRAINT pk_au_id
          PRIMARY KEY,
-	username 	VARCHAR(20),
-	email 		VARCHAR(120),
-	name 			VARCHAR(120),
-	password 	VARCHAR(120),
-	role        VARCHAR(20)
-);
-
-ALTER TABLE auth_users
-ADD CONSTRAINT username_check
-CHECK (
-    LENGTH(username) >= 4 AND       -- Minimum length of 4 characters
-    LENGTH(username) <= 20 AND      -- Maximum length of 20 characters
-    username ~ '^[a-zA-Z0-9_]+$'   -- Only alphanumeric characters and underscores allowed
+	username 	 VARCHAR(256),
+	email 		 VARCHAR(256),
+	name 			 VARCHAR(256),
+	password 	 VARCHAR(120),
+	role         VARCHAR(20),
+   latest_login TIMESTAMP
 );
 
 ALTER TABLE auth_users

testgen/template/dbupgrade/0146_incremental_upgrade.sql

@@ -0,0 +1,8 @@
+SET SEARCH_PATH TO {SCHEMA_NAME};
+
+ALTER TABLE auth_users
+  ADD COLUMN latest_login TIMESTAMP,
+  ALTER COLUMN username TYPE VARCHAR(256),
+  ALTER COLUMN email TYPE VARCHAR(256),
+  ALTER COLUMN name TYPE VARCHAR(256),
+  DROP CONSTRAINT username_check;

testgen/ui/app.py

@@ -10,7 +10,7 @@
 from testgen.ui import bootstrap
 from testgen.ui.assets import get_asset_path
 from testgen.ui.components import widgets as testgen
-from testgen.ui.services import javascript_service, user_session_service
+from testgen.ui.services import javascript_service
 from testgen.ui.session import session
 
 
@@ -23,12 +23,14 @@ def render(log_level: int = logging.INFO):
         # Collapse when logging out because the sidebar takes some time to be removed from the DOM
         # Collapse for Catalog role since they only have access to one page
         initial_sidebar_state="collapsed"
-        if session.logging_out or user_session_service.user_has_catalog_role()
+        if session.auth and (session.auth.logging_out or (session.auth.is_logged_in and not session.auth.user_has_permission("view")))
         else "auto",
     )
 
     application = get_application(log_level=log_level)
     application.logger.debug("Starting Streamlit re-run")
+    if not session.auth:
+        session.auth = application.auth_class()
 
     status_ok, message = check_basic_configuration()
     if not status_ok:
@@ -41,20 +43,18 @@ def render(log_level: int = logging.INFO):
         session.page_args_pending_router and session.page_args_pending_router.get("project_code")
     ) or st.query_params.get("project_code", session.sidebar_project)
 
-    if session.authentication_status is None and not session.logging_out:
-        user_session_service.load_user_session()
+    if not session.auth.is_logged_in and not session.auth.logging_out:
+        session.auth.load_user_session()
 
     application.logo.render()
 
-    if session.authentication_status and not session.logging_in:
+    if session.auth.is_logged_in and not session.auth.logging_in:
         with st.sidebar:
             testgen.sidebar(
                 projects=Project.select_where(),
                 current_project=session.sidebar_project,
                 menu=application.menu,
                 current_page=session.current_page,
-                username=session.username,
-                role=session.auth_role,
                 version=version_service.get_version(),
                 support_email=settings.SUPPORT_EMAIL,
             )

testgen/ui/assets/scripts.js

@@ -13,15 +13,6 @@ window.addEventListener('message', async function(event) {
     }
 });
 
-function removeElements(selectors) {
-    for (const selector of selectors) {
-        const element = window.top.document.querySelector(selector);
-        if (element) {
-            element.remove();
-        }
-    }
-}
-
 async function copyToClipboard(text) {
     if (navigator.clipboard && window.isSecureContext) {
         await navigator.clipboard.writeText(text || '');
@@ -55,4 +46,5 @@ window.testgen = {
     states: {},
     components: {},
     loadedStylesheets: {},
+    changeLocation: url => window.location.href = url, 
 };

testgen/ui/assets/style.css

@@ -79,7 +79,9 @@ img.dk-logo-img {
 
 /* Sidebar */
 [data-testid="stSidebarContent"] [data-testid="stSidebarHeader"] {
-    padding: 16px 20px;
+    padding: 16px 20px 20px;
+    margin-bottom: 0;
+    height: auto;
 }
 
 [data-testid="stSidebarHeader"] .stLogo {
@@ -200,6 +202,7 @@ button[title="Show password text"] {
     display: none;
 }
 
+.element-container:has(iframe[height="0"][title="extra_streamlit_components.CookieManager.cookie_manager"]),
 .element-container:has(iframe[height="0"][title="streamlit_javascript.streamlit_javascript"]),
 .element-container:has(iframe[height="0"][title="testgen.ui.components.utils.component.testgen"]) {
     display: none !important;