feat: add UserScopedMixin.for_user() for per-user models

asheshv · asheshv · commit 2f10925d2d76 · 2026-04-08T17:51:44.000+05:30
Mixin providing Model.for_user(**kwargs) as the standard scoped query
entry point. Applied to 11 user-scoped models. Converted all Process
query callsites. Updated BatchProcess test mocks.
diff --git a/web/pgadmin/misc/bgprocess/processes.py b/web/pgadmin/misc/bgprocess/processes.py
@@ -153,7 +153,7 @@ def __init__(self, **kwargs):
             self.manager_obj = kwargs['manager_obj']
 
     def _retrieve_process(self, _id):
-        p = Process.query.filter_by(pid=_id, user_id=current_user.id).first()
+        p = Process.for_user(pid=_id).first()
 
         if p is None:
             raise LookupError(PROCESS_NOT_FOUND)
@@ -372,19 +372,15 @@ def start(self, cb=None):
             # There is no way to find out the error message from this process
             # as standard output, and standard error were redirected to
             # devnull.
-            p = Process.query.filter_by(
-                pid=self.id, user_id=current_user.id
-            ).first()
+            p = Process.for_user(pid=self.id).first()
             p.start_time = p.end_time = get_current_time()
             if not p.exit_code:
                 p.exit_code = self.ecode
             p.process_state = PROCESS_FINISHED
             db.session.commit()
         else:
             # Update the process state to "Started"
-            p = Process.query.filter_by(
-                pid=self.id, user_id=current_user.id
-            ).first()
+            p = Process.for_user(pid=self.id).first()
             p.process_state = PROCESS_STARTED
             db.session.commit()
 
@@ -530,9 +526,7 @@ def update_cloud_details(self):
         """
         _pid = self.id
 
-        _process = Process.query.filter_by(
-            user_id=current_user.id, pid=_pid
-        ).first()
+        _process = Process.for_user(pid=_pid).first()
 
         if _process is None:
             raise LookupError(PROCESS_NOT_FOUND)
@@ -588,9 +582,7 @@ def status(self, out=0, err=0):
         out_completed = err_completed = False
         process_output = (out != -1 and err != -1)
 
-        j = Process.query.filter_by(
-            pid=self.id, user_id=current_user.id
-        ).first()
+        j = Process.for_user(pid=self.id).first()
         enc = sys.getdefaultencoding()
         if enc == 'ascii':
             enc = 'utf-8'
@@ -739,7 +731,7 @@ def _check_process_desc(p):
 
     @staticmethod
     def list():
-        processes = Process.query.filter_by(user_id=current_user.id)
+        processes = Process.for_user()
         changed = False
 
         browser_preference = Preferences.module('browser')
@@ -812,9 +804,7 @@ def acknowledge(_pid):
         And, delete the process information from the configuration, and the log
         files related to the process, if it has already been completed.
         """
-        p = Process.query.filter_by(
-            user_id=current_user.id, pid=_pid
-        ).first()
+        p = Process.for_user(pid=_pid).first()
 
         if p is None:
             raise LookupError(PROCESS_NOT_FOUND)
@@ -886,9 +876,7 @@ def set_env_variables(self, server, **kwargs):
     def stop_process(_pid):
         """
         """
-        p = Process.query.filter_by(
-            user_id=current_user.id, pid=_pid
-        ).first()
+        p = Process.for_user(pid=_pid).first()
 
         if p is None:
             raise LookupError(PROCESS_NOT_FOUND)
@@ -910,9 +898,7 @@ def stop_process(_pid):
 
     @staticmethod
     def update_server_id(_pid, _sid):
-        p = Process.query.filter_by(
-            user_id=current_user.id, pid=_pid
-        ).first()
+        p = Process.for_user(pid=_pid).first()
 
         if p is None:
             raise LookupError(PROCESS_NOT_FOUND)
diff --git a/web/pgadmin/misc/cloud/__init__.py b/web/pgadmin/misc/cloud/__init__.py
@@ -212,8 +212,9 @@ def clear_cloud_session(pid=None):
 @pga_login_required
 def update_cloud_process(sid):
     """Update Cloud Server Process"""
-    _process = Process.query.filter_by(user_id=current_user.id,
-                                       server_id=sid).first()
+    _process = Process.for_user(server_id=sid).first()
+    if _process is None:
+        return success_return()
     _process.acknowledge = None
     db.session.commit()
     return success_return()
diff --git a/web/pgadmin/model/__init__.py b/web/pgadmin/model/__init__.py
@@ -33,7 +33,7 @@
 #
 ##########################################################################
 
-SCHEMA_VERSION = 49
+SCHEMA_VERSION = 50
 
 ##########################################################################
 #
@@ -51,6 +51,60 @@
 SERVER_ID = 'server.id'
 CASCADE_STR = "all, delete-orphan"
 
+
+class UserScopedMixin:
+    """Mixin for models that store per-user data.
+
+    Provides for_user() as the default scoped query entry point.
+    Models with a 'user_id' column or a 'uid' column are supported
+    automatically — the mixin detects which column name is used.
+
+    Usage:
+        # Instead of:
+        Process.query.filter_by(user_id=current_user.id, pid=pid)
+        # Use:
+        Process.for_user(pid=pid)
+    """
+
+    @classmethod
+    def _user_column(cls):
+        """Return the user-scoping column for this model."""
+        if hasattr(cls, 'user_id'):
+            return cls.user_id
+        if hasattr(cls, 'uid'):
+            return cls.uid
+        raise AttributeError(
+            f"{cls.__name__} has no user_id or uid column"
+        )
+
+    @classmethod
+    def _user_column_name(cls):
+        """Return the column name string ('user_id' or 'uid')."""
+        if hasattr(cls, 'user_id'):
+            return 'user_id'
+        if hasattr(cls, 'uid'):
+            return 'uid'
+        raise AttributeError(
+            f"{cls.__name__} has no user_id or uid column"
+        )
+
+    @classmethod
+    def for_user(cls, user_id=None, **kwargs):
+        """Query scoped to a specific user (defaults to current_user).
+
+        Args:
+            user_id: Explicit user ID. If None, uses current_user.id.
+            **kwargs: Additional filter_by arguments.
+
+        Returns:
+            A SQLAlchemy query filtered by the user's ID.
+        """
+        from flask_security import current_user as cu
+        uid = user_id if user_id is not None else cu.id
+        kwargs[cls._user_column_name()] = uid
+        return cls.query.filter_by(**kwargs)
+
+
 # Define models
 roles_users = db.Table(
     'roles_users',
@@ -158,15 +212,15 @@ class User(db.Model, UserMixin):
     locked = db.Column(db.Boolean(), default=False)
 
 
-class Setting(db.Model):
+class Setting(db.Model, UserScopedMixin):
     """Define a setting object"""
     __tablename__ = 'setting'
     user_id = db.Column(db.Integer, db.ForeignKey(USER_ID), primary_key=True)
     setting = db.Column(db.String(256), primary_key=True)
     value = db.Column(db.Text())
 
 
-class ServerGroup(db.Model):
+class ServerGroup(db.Model, UserScopedMixin):
     """Define a server group for the treeview"""
     __tablename__ = 'servergroup'
     id = db.Column(db.Integer, primary_key=True)
@@ -185,7 +239,7 @@ def serialize(self):
         }
 
 
-class Server(db.Model):
+class Server(db.Model, UserScopedMixin):
     """Define a registered Postgres server"""
     __tablename__ = 'server'
     id = db.Column(db.Integer, primary_key=True)
@@ -306,7 +360,7 @@ class Preferences(db.Model):
     name = db.Column(db.String(1024), nullable=False)
 
 
-class UserPreference(db.Model):
+class UserPreference(db.Model, UserScopedMixin):
     """Define the preference for a particular user."""
     __tablename__ = 'user_preferences'
     pid = db.Column(
@@ -318,9 +372,13 @@ class UserPreference(db.Model):
     value = db.Column(db.String(1024), nullable=False)
 
 
-class DebuggerFunctionArguments(db.Model):
+class DebuggerFunctionArguments(db.Model, UserScopedMixin):
     """Define the debugger input function arguments."""
     __tablename__ = 'debugger_function_arguments'
+    user_id = db.Column(
+        db.Integer, db.ForeignKey(USER_ID),
+        nullable=False, primary_key=True
+    )
     server_id = db.Column(db.Integer(), nullable=False, primary_key=True)
     database_id = db.Column(db.Integer(), nullable=False, primary_key=True)
     schema_id = db.Column(db.Integer(), nullable=False, primary_key=True)
@@ -349,7 +407,7 @@ class DebuggerFunctionArguments(db.Model):
     value = db.Column(db.String(), nullable=True)
 
 
-class Process(db.Model):
+class Process(db.Model, UserScopedMixin):
     """Define the Process table."""
     __tablename__ = 'process'
     pid = db.Column(db.String(), nullable=False, primary_key=True)
@@ -382,7 +440,7 @@ class Keys(db.Model):
     value = db.Column(db.String(), nullable=False)
 
 
-class QueryHistoryModel(db.Model):
+class QueryHistoryModel(db.Model, UserScopedMixin):
     """Define the history SQL table."""
     __tablename__ = 'query_history'
     srno = db.Column(db.Integer(), nullable=False, primary_key=True)
@@ -397,7 +455,7 @@ class QueryHistoryModel(db.Model):
     last_updated_flag = db.Column(db.String(), nullable=False)
 
 
-class ApplicationState(db.Model):
+class ApplicationState(db.Model, UserScopedMixin):
     """Define the application state SQL table."""
     __tablename__ = 'application_state'
     uid = db.Column(db.Integer(), db.ForeignKey(USER_ID), nullable=False,
@@ -422,7 +480,7 @@ class Database(db.Model):
     )
 
 
-class SharedServer(db.Model):
+class SharedServer(db.Model, UserScopedMixin):
     """Define a shared Postgres server"""
 
     __tablename__ = 'sharedserver'
@@ -510,7 +568,7 @@ class Macros(db.Model):
     key_code = db.Column(db.Integer, nullable=False)
 
 
-class UserMacros(db.Model):
+class UserMacros(db.Model, UserScopedMixin):
     """Define the macro for a particular user."""
     __tablename__ = 'user_macros'
     id = db.Column(db.Integer, primary_key=True, autoincrement=True)
@@ -524,7 +582,7 @@ class UserMacros(db.Model):
     sql = db.Column(db.Text(), nullable=False)
 
 
-class UserMFA(db.Model):
+class UserMFA(db.Model, UserScopedMixin):
     """Stores the options for the MFA for a particular user."""
     __tablename__ = 'user_mfa'
     user_id = db.Column(db.Integer, db.ForeignKey(USER_ID), primary_key=True)
diff --git a/web/pgadmin/tools/backup/tests/test_batch_process.py b/web/pgadmin/tools/backup/tests/test_batch_process.py
@@ -195,6 +195,7 @@ def __init__(self, desc, args, cmd):
                 self.utility_pid = 123
                 self.server_id = None
 
+        process_mock.for_user = process_mock.query.filter_by
         mock_result = process_mock.query.filter_by.return_value
         mock_result.first.return_value = TestMockProcess(
             backup_obj, self.class_params['args'], self.class_params['cmd'])
@@ -239,6 +240,7 @@ def __init__(self, desc, args, cmd):
                 self.utility_pid = 123
                 self.server_id = None
 
+        process_mock.for_user = process_mock.query.filter_by
         process_mock.query.filter_by.return_value = [
             TestMockProcess(backup_obj,
                             self.class_params['args'],
diff --git a/web/pgadmin/tools/import_export/tests/test_batch_process.py b/web/pgadmin/tools/import_export/tests/test_batch_process.py
@@ -204,6 +204,7 @@ def __init__(self, desc, args, cmd):
                 self.utility_pid = 123
                 self.server_id = None
 
+        process_mock.for_user = process_mock.query.filter_by
         mock_result = process_mock.query.filter_by.return_value
         mock_result.first.return_value = TestMockProcess(
             import_export_obj, self.class_params['args'],
@@ -250,6 +251,7 @@ def __init__(self, desc, args, cmd):
                 self.utility_pid = 123
                 self.server_id = None
 
+        process_mock.for_user = process_mock.query.filter_by
         process_mock.query.filter_by.return_value = [
             TestMockProcess(import_export_obj,
                             self.class_params['args'],
diff --git a/web/pgadmin/tools/maintenance/tests/test_batch_process_maintenance.py b/web/pgadmin/tools/maintenance/tests/test_batch_process_maintenance.py
@@ -137,6 +137,7 @@ def __init__(self, desc, args, cmd):
                 self.utility_pid = 123
                 self.server_id = None
 
+        process_mock.for_user = process_mock.query.filter_by
         mock_result = process_mock.query.filter_by.return_value
         mock_result.first.return_value = TestMockProcess(
             maintenance_obj, self.class_params['args'],
@@ -177,6 +178,7 @@ def __init__(self, desc, args, cmd):
                 self.utility_pid = 123
                 self.server_id = None
 
+        process_mock.for_user = process_mock.query.filter_by
         process_mock.query.filter_by.return_value = [
             TestMockProcess(maintenance_obj,
                             self.class_params['args'],
diff --git a/web/pgadmin/tools/restore/tests/test_batch_process.py b/web/pgadmin/tools/restore/tests/test_batch_process.py
@@ -134,6 +134,7 @@ def __init__(self, desc, args, cmd):
                 self.utility_pid = 123
                 self.server_id = None
 
+        process_mock.for_user = process_mock.query.filter_by
         mock_result = process_mock.query.filter_by.return_value
         mock_result.first.return_value = TestMockProcess(
             restore_obj, self.class_params['args'],
@@ -174,6 +175,7 @@ def __init__(self, desc, args, cmd):
                 self.utility_pid = 123
                 self.server_id = None
 
+        process_mock.for_user = process_mock.query.filter_by
         process_mock.query.filter_by.return_value = [
             TestMockProcess(restore_obj,
                             self.class_params['args'],
diff --git a/web/pgadmin/tools/user_management/__init__.py b/web/pgadmin/tools/user_management/__init__.py
@@ -759,7 +759,7 @@ def delete_user(uid):
 
         ServerGroup.query.filter_by(user_id=uid).delete()
 
-        Process.query.filter_by(user_id=uid).delete()
+        Process.for_user(user_id=uid).delete()
         # Delete Shared servers for current user.
         SharedServer.query.filter_by(user_id=uid).delete()
 
