Skip to content

Commit f365785

Browse files
committed
feat: add UserScopedMixin.for_user() for per-user models
Introduce UserScopedMixin — a model mixin providing for_user() as the default scoped query entry point. Handles both user_id and uid column naming conventions automatically. # Instead of: Process.query.filter_by(user_id=current_user.id, pid=pid) # Use: Process.for_user(pid=pid) Applied to: Setting, ServerGroup, Server, UserPreference, DebuggerFunctionArguments, Process, QueryHistoryModel, ApplicationState, SharedServer, UserMacros, UserMFA. Converted all Process.query.filter_by(user_id=...) callsites to Process.for_user(). Updated BatchProcess test mocks to wire process_mock.for_user to process_mock.query.filter_by.
1 parent d9a22aa commit f365785

File tree

8 files changed

+91
-38
lines changed

8 files changed

+91
-38
lines changed

web/pgadmin/misc/bgprocess/processes.py

Lines changed: 9 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ def __init__(self, **kwargs):
153153
self.manager_obj = kwargs['manager_obj']
154154

155155
def _retrieve_process(self, _id):
156-
p = Process.query.filter_by(pid=_id, user_id=current_user.id).first()
156+
p = Process.for_user(pid=_id).first()
157157

158158
if p is None:
159159
raise LookupError(PROCESS_NOT_FOUND)
@@ -372,19 +372,15 @@ def start(self, cb=None):
372372
# There is no way to find out the error message from this process
373373
# as standard output, and standard error were redirected to
374374
# devnull.
375-
p = Process.query.filter_by(
376-
pid=self.id, user_id=current_user.id
377-
).first()
375+
p = Process.for_user(pid=self.id).first()
378376
p.start_time = p.end_time = get_current_time()
379377
if not p.exit_code:
380378
p.exit_code = self.ecode
381379
p.process_state = PROCESS_FINISHED
382380
db.session.commit()
383381
else:
384382
# Update the process state to "Started"
385-
p = Process.query.filter_by(
386-
pid=self.id, user_id=current_user.id
387-
).first()
383+
p = Process.for_user(pid=self.id).first()
388384
p.process_state = PROCESS_STARTED
389385
db.session.commit()
390386

@@ -530,9 +526,7 @@ def update_cloud_details(self):
530526
"""
531527
_pid = self.id
532528

533-
_process = Process.query.filter_by(
534-
user_id=current_user.id, pid=_pid
535-
).first()
529+
_process = Process.for_user(pid=_pid).first()
536530

537531
if _process is None:
538532
raise LookupError(PROCESS_NOT_FOUND)
@@ -588,9 +582,7 @@ def status(self, out=0, err=0):
588582
out_completed = err_completed = False
589583
process_output = (out != -1 and err != -1)
590584

591-
j = Process.query.filter_by(
592-
pid=self.id, user_id=current_user.id
593-
).first()
585+
j = Process.for_user(pid=self.id).first()
594586
enc = sys.getdefaultencoding()
595587
if enc == 'ascii':
596588
enc = 'utf-8'
@@ -739,7 +731,7 @@ def _check_process_desc(p):
739731

740732
@staticmethod
741733
def list():
742-
processes = Process.query.filter_by(user_id=current_user.id)
734+
processes = Process.for_user()
743735
changed = False
744736

745737
browser_preference = Preferences.module('browser')
@@ -812,9 +804,7 @@ def acknowledge(_pid):
812804
And, delete the process information from the configuration, and the log
813805
files related to the process, if it has already been completed.
814806
"""
815-
p = Process.query.filter_by(
816-
user_id=current_user.id, pid=_pid
817-
).first()
807+
p = Process.for_user(pid=_pid).first()
818808

819809
if p is None:
820810
raise LookupError(PROCESS_NOT_FOUND)
@@ -886,9 +876,7 @@ def set_env_variables(self, server, **kwargs):
886876
def stop_process(_pid):
887877
"""
888878
"""
889-
p = Process.query.filter_by(
890-
user_id=current_user.id, pid=_pid
891-
).first()
879+
p = Process.for_user(pid=_pid).first()
892880

893881
if p is None:
894882
raise LookupError(PROCESS_NOT_FOUND)
@@ -910,9 +898,7 @@ def stop_process(_pid):
910898

911899
@staticmethod
912900
def update_server_id(_pid, _sid):
913-
p = Process.query.filter_by(
914-
user_id=current_user.id, pid=_pid
915-
).first()
901+
p = Process.for_user(pid=_pid).first()
916902

917903
if p is None:
918904
raise LookupError(PROCESS_NOT_FOUND)

web/pgadmin/misc/cloud/__init__.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -212,8 +212,9 @@ def clear_cloud_session(pid=None):
212212
@pga_login_required
213213
def update_cloud_process(sid):
214214
"""Update Cloud Server Process"""
215-
_process = Process.query.filter_by(user_id=current_user.id,
216-
server_id=sid).first()
215+
_process = Process.for_user(server_id=sid).first()
216+
if _process is None:
217+
return success_return()
217218
_process.acknowledge = None
218219
db.session.commit()
219220
return success_return()

web/pgadmin/model/__init__.py

Lines changed: 70 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
#
3434
##########################################################################
3535

36-
SCHEMA_VERSION = 49
36+
SCHEMA_VERSION = 50
3737

3838
##########################################################################
3939
#
@@ -51,6 +51,60 @@
5151
SERVER_ID = 'server.id'
5252
CASCADE_STR = "all, delete-orphan"
5353

54+
55+
class UserScopedMixin:
56+
"""Mixin for models that store per-user data.
57+
58+
Provides for_user() as the default scoped query entry point.
59+
Models with a 'user_id' column or a 'uid' column are supported
60+
automatically — the mixin detects which column name is used.
61+
62+
Usage:
63+
# Instead of:
64+
Process.query.filter_by(user_id=current_user.id, pid=pid)
65+
# Use:
66+
Process.for_user(pid=pid)
67+
"""
68+
69+
@classmethod
70+
def _user_column(cls):
71+
"""Return the user-scoping column for this model."""
72+
if hasattr(cls, 'user_id'):
73+
return cls.user_id
74+
if hasattr(cls, 'uid'):
75+
return cls.uid
76+
raise AttributeError(
77+
f"{cls.__name__} has no user_id or uid column"
78+
)
79+
80+
@classmethod
81+
def _user_column_name(cls):
82+
"""Return the column name string ('user_id' or 'uid')."""
83+
if hasattr(cls, 'user_id'):
84+
return 'user_id'
85+
if hasattr(cls, 'uid'):
86+
return 'uid'
87+
raise AttributeError(
88+
f"{cls.__name__} has no user_id or uid column"
89+
)
90+
91+
@classmethod
92+
def for_user(cls, user_id=None, **kwargs):
93+
"""Query scoped to a specific user (defaults to current_user).
94+
95+
Args:
96+
user_id: Explicit user ID. If None, uses current_user.id.
97+
**kwargs: Additional filter_by arguments.
98+
99+
Returns:
100+
A SQLAlchemy query filtered by the user's ID.
101+
"""
102+
from flask_security import current_user as cu
103+
uid = user_id if user_id is not None else cu.id
104+
kwargs[cls._user_column_name()] = uid
105+
return cls.query.filter_by(**kwargs)
106+
107+
54108
# Define models
55109
roles_users = db.Table(
56110
'roles_users',
@@ -158,15 +212,15 @@ class User(db.Model, UserMixin):
158212
locked = db.Column(db.Boolean(), default=False)
159213

160214

161-
class Setting(db.Model):
215+
class Setting(db.Model, UserScopedMixin):
162216
"""Define a setting object"""
163217
__tablename__ = 'setting'
164218
user_id = db.Column(db.Integer, db.ForeignKey(USER_ID), primary_key=True)
165219
setting = db.Column(db.String(256), primary_key=True)
166220
value = db.Column(db.Text())
167221

168222

169-
class ServerGroup(db.Model):
223+
class ServerGroup(db.Model, UserScopedMixin):
170224
"""Define a server group for the treeview"""
171225
__tablename__ = 'servergroup'
172226
id = db.Column(db.Integer, primary_key=True)
@@ -185,7 +239,7 @@ def serialize(self):
185239
}
186240

187241

188-
class Server(db.Model):
242+
class Server(db.Model, UserScopedMixin):
189243
"""Define a registered Postgres server"""
190244
__tablename__ = 'server'
191245
id = db.Column(db.Integer, primary_key=True)
@@ -306,7 +360,7 @@ class Preferences(db.Model):
306360
name = db.Column(db.String(1024), nullable=False)
307361

308362

309-
class UserPreference(db.Model):
363+
class UserPreference(db.Model, UserScopedMixin):
310364
"""Define the preference for a particular user."""
311365
__tablename__ = 'user_preferences'
312366
pid = db.Column(
@@ -318,9 +372,13 @@ class UserPreference(db.Model):
318372
value = db.Column(db.String(1024), nullable=False)
319373

320374

321-
class DebuggerFunctionArguments(db.Model):
375+
class DebuggerFunctionArguments(db.Model, UserScopedMixin):
322376
"""Define the debugger input function arguments."""
323377
__tablename__ = 'debugger_function_arguments'
378+
user_id = db.Column(
379+
db.Integer, db.ForeignKey(USER_ID),
380+
nullable=False, primary_key=True
381+
)
324382
server_id = db.Column(db.Integer(), nullable=False, primary_key=True)
325383
database_id = db.Column(db.Integer(), nullable=False, primary_key=True)
326384
schema_id = db.Column(db.Integer(), nullable=False, primary_key=True)
@@ -349,7 +407,7 @@ class DebuggerFunctionArguments(db.Model):
349407
value = db.Column(db.String(), nullable=True)
350408

351409

352-
class Process(db.Model):
410+
class Process(db.Model, UserScopedMixin):
353411
"""Define the Process table."""
354412
__tablename__ = 'process'
355413
pid = db.Column(db.String(), nullable=False, primary_key=True)
@@ -382,7 +440,7 @@ class Keys(db.Model):
382440
value = db.Column(db.String(), nullable=False)
383441

384442

385-
class QueryHistoryModel(db.Model):
443+
class QueryHistoryModel(db.Model, UserScopedMixin):
386444
"""Define the history SQL table."""
387445
__tablename__ = 'query_history'
388446
srno = db.Column(db.Integer(), nullable=False, primary_key=True)
@@ -397,7 +455,7 @@ class QueryHistoryModel(db.Model):
397455
last_updated_flag = db.Column(db.String(), nullable=False)
398456

399457

400-
class ApplicationState(db.Model):
458+
class ApplicationState(db.Model, UserScopedMixin):
401459
"""Define the application state SQL table."""
402460
__tablename__ = 'application_state'
403461
uid = db.Column(db.Integer(), db.ForeignKey(USER_ID), nullable=False,
@@ -422,7 +480,7 @@ class Database(db.Model):
422480
)
423481

424482

425-
class SharedServer(db.Model):
483+
class SharedServer(db.Model, UserScopedMixin):
426484
"""Define a shared Postgres server"""
427485

428486
__tablename__ = 'sharedserver'
@@ -510,7 +568,7 @@ class Macros(db.Model):
510568
key_code = db.Column(db.Integer, nullable=False)
511569

512570

513-
class UserMacros(db.Model):
571+
class UserMacros(db.Model, UserScopedMixin):
514572
"""Define the macro for a particular user."""
515573
__tablename__ = 'user_macros'
516574
id = db.Column(db.Integer, primary_key=True, autoincrement=True)
@@ -524,7 +582,7 @@ class UserMacros(db.Model):
524582
sql = db.Column(db.Text(), nullable=False)
525583

526584

527-
class UserMFA(db.Model):
585+
class UserMFA(db.Model, UserScopedMixin):
528586
"""Stores the options for the MFA for a particular user."""
529587
__tablename__ = 'user_mfa'
530588
user_id = db.Column(db.Integer, db.ForeignKey(USER_ID), primary_key=True)

web/pgadmin/tools/backup/tests/test_batch_process.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,7 @@ def __init__(self, desc, args, cmd):
195195
self.utility_pid = 123
196196
self.server_id = None
197197

198+
process_mock.for_user = process_mock.query.filter_by
198199
mock_result = process_mock.query.filter_by.return_value
199200
mock_result.first.return_value = TestMockProcess(
200201
backup_obj, self.class_params['args'], self.class_params['cmd'])
@@ -239,6 +240,7 @@ def __init__(self, desc, args, cmd):
239240
self.utility_pid = 123
240241
self.server_id = None
241242

243+
process_mock.for_user = process_mock.query.filter_by
242244
process_mock.query.filter_by.return_value = [
243245
TestMockProcess(backup_obj,
244246
self.class_params['args'],

web/pgadmin/tools/import_export/tests/test_batch_process.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,7 @@ def __init__(self, desc, args, cmd):
204204
self.utility_pid = 123
205205
self.server_id = None
206206

207+
process_mock.for_user = process_mock.query.filter_by
207208
mock_result = process_mock.query.filter_by.return_value
208209
mock_result.first.return_value = TestMockProcess(
209210
import_export_obj, self.class_params['args'],
@@ -250,6 +251,7 @@ def __init__(self, desc, args, cmd):
250251
self.utility_pid = 123
251252
self.server_id = None
252253

254+
process_mock.for_user = process_mock.query.filter_by
253255
process_mock.query.filter_by.return_value = [
254256
TestMockProcess(import_export_obj,
255257
self.class_params['args'],

web/pgadmin/tools/maintenance/tests/test_batch_process_maintenance.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@ def __init__(self, desc, args, cmd):
137137
self.utility_pid = 123
138138
self.server_id = None
139139

140+
process_mock.for_user = process_mock.query.filter_by
140141
mock_result = process_mock.query.filter_by.return_value
141142
mock_result.first.return_value = TestMockProcess(
142143
maintenance_obj, self.class_params['args'],
@@ -177,6 +178,7 @@ def __init__(self, desc, args, cmd):
177178
self.utility_pid = 123
178179
self.server_id = None
179180

181+
process_mock.for_user = process_mock.query.filter_by
180182
process_mock.query.filter_by.return_value = [
181183
TestMockProcess(maintenance_obj,
182184
self.class_params['args'],

web/pgadmin/tools/restore/tests/test_batch_process.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@ def __init__(self, desc, args, cmd):
134134
self.utility_pid = 123
135135
self.server_id = None
136136

137+
process_mock.for_user = process_mock.query.filter_by
137138
mock_result = process_mock.query.filter_by.return_value
138139
mock_result.first.return_value = TestMockProcess(
139140
restore_obj, self.class_params['args'],
@@ -174,6 +175,7 @@ def __init__(self, desc, args, cmd):
174175
self.utility_pid = 123
175176
self.server_id = None
176177

178+
process_mock.for_user = process_mock.query.filter_by
177179
process_mock.query.filter_by.return_value = [
178180
TestMockProcess(restore_obj,
179181
self.class_params['args'],

web/pgadmin/tools/user_management/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -759,7 +759,7 @@ def delete_user(uid):
759759

760760
ServerGroup.query.filter_by(user_id=uid).delete()
761761

762-
Process.query.filter_by(user_id=uid).delete()
762+
Process.for_user(user_id=uid).delete()
763763
# Delete Shared servers for current user.
764764
SharedServer.query.filter_by(user_id=uid).delete()
765765

0 commit comments

Comments
 (0)