diff --git a/alembic/versions/af63c0205b19_adding_library_id_foreign_key_to_user.py b/alembic/versions/af63c0205b19_adding_library_id_foreign_key_to_user.py new file mode 100644 index 0000000..b3b2d0d --- /dev/null +++ b/alembic/versions/af63c0205b19_adding_library_id_foreign_key_to_user.py @@ -0,0 +1,71 @@ +"""Adding library_id foreign key to user +Revision ID: af63c0205b19 +Revises: dcda14f51cff +Create Date: 2025-07-09 15:00:40.189587 +""" + +# revision identifiers, used by Alembic. +revision = 'af63c0205b19' +down_revision = 'dcda14f51cff' + +from alembic import op +import sqlalchemy as sa +import json + +from sqlalchemy.dialects import postgresql + +def upgrade(): + #with app.app_context() as c: + # db.session.add(Model()) + # db.session.commit() + + # ### commands auto generated by Alembic - please adjust! ### + # Add foreign key constraint to users table + op.add_column('users', sa.Column('library_id', sa.Integer(), nullable=True)) + op.create_foreign_key(None, 'users', 'library', ['library_id'], ['id']) + + + bind = op.get_bind() + # Get all users with link_server in user_data + # Delete link_server from user_data + # But save the link_server in the library table instead to be accessible via users.library_id + + + users = bind.execute("SELECT id, user_data FROM users WHERE user_data ? 'link_server'") + + for user_id, user_data in users: + + if user_data and isinstance(user_data, dict): + link_server = user_data.get('link_server') + if link_server: + library = bind.execute( + "SELECT id FROM library WHERE libserver = %s", (link_server,) + ).fetchone() + + library_id = library[0] if library else None + + new_user_data = user_data.copy() + new_user_data.pop('link_server', None) + + bind.execute( + "UPDATE users SET user_data = %s, library_id = %s WHERE id = %s", + (new_user_data, library_id, user_id) + ) + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + + bind = op.get_bind() + + result = bind.execute("SELECT u.id, u.user_data, l.libserver FROM users u JOIN library l ON u.library_id = l.id WHERE u.library_id IS NOT NULL") + + for user_id, user_data, libserver in result: + user_data = user_data or {} + if isinstance(user_data, dict): + if libserver is not None: + user_data['link_server'] = libserver + bind.execute("UPDATE users SET user_data = %s WHERE id = %s", (json.dumps(user_data), user_id)) + + op.drop_constraint(None, 'users', type_='foreignkey') + op.drop_column('users', 'library_id') + # ### end Alembic commands ### diff --git a/vault_service/models.py b/vault_service/models.py index 2f0eed9..c19d866 100644 --- a/vault_service/models.py +++ b/vault_service/models.py @@ -5,7 +5,7 @@ Models for the users (users) of AdsWS """ -from sqlalchemy import Column, Integer, String, LargeBinary, TIMESTAMP, ForeignKey, Boolean, Text +from sqlalchemy import Column, Integer, String, LargeBinary, ForeignKey, Boolean, Text from sqlalchemy.dialects.postgresql import JSONB, ENUM, ARRAY from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.ext.mutable import MutableDict @@ -25,6 +25,7 @@ class User(Base): id = Column(Integer, primary_key=True) name = Column(String(255)) user_data = Column(MutableDict.as_mutable(JSONB)) + library_id = Column(Integer, ForeignKey('library.id'), nullable=True) created = Column(UTCDateTime, default=get_date) updated = Column(UTCDateTime, default=get_date, onupdate=get_date) diff --git a/vault_service/tests/test_user.py b/vault_service/tests/test_user.py index 61b87f7..dcf1de4 100644 --- a/vault_service/tests/test_user.py +++ b/vault_service/tests/test_user.py @@ -10,7 +10,7 @@ if project_home not in sys.path: sys.path.insert(0, project_home) -from vault_service.models import Query, User, MyADS +from vault_service.models import Query, User, MyADS, Library from vault_service.tests.base import TestCaseDatabase import adsmutils @@ -1005,5 +1005,238 @@ def test_scixplorer_referrer_updates_all_notifications(self): self.assertTrue(notification3.scix_ui) self.assertTrue(notification4.scix_ui) + def test_library_integration(self): + '''Tests library integration with user data storage''' + + # First create a library in the test database with a libserver + with self.app.session_scope() as session: + library = Library(libname='Test Library', libserver='https://resolver.example.com/sfx') + library1 = Library(libname='Test Library 2', libserver='https://another.resolver.com/sfx') + session.add(library) + session.add(library1) + session.commit() + + # Test GET request with no existing user data + r = self.client.get(url_for('user.store_data'), + headers={'Authorization': 'secret', 'X-api-uid': '10'}, + content_type='application/json') + self.assertStatus(r, 200) + self.assertEqual(r.json, {}) + + # Test POST with link_server - should create library relationship + r = self.client.post(url_for('user.store_data'), + headers={'Authorization': 'secret', 'X-api-uid': '10'}, + data=json.dumps({ + 'link_server': 'https://resolver.example.com/sfx', + 'other_data': 'test' + }), + content_type='application/json') + self.assertStatus(r, 200) + self.assertIn('link_server', r.json) + self.assertEqual(r.json['link_server'], 'https://resolver.example.com/sfx') + self.assertEqual(r.json['other_data'], 'test') + + # Verify library relationship was created + with self.app.session_scope() as session: + user = session.query(User).filter_by(id=10).first() + self.assertIsNotNone(user) + self.assertIsNotNone(user.library_id) + + library = session.query(Library).filter_by(id=user.library_id).first() + self.assertIsNotNone(library) + self.assertEqual(library.libserver, 'https://resolver.example.com/sfx') + + # Test GET after library relationship created + r = self.client.get(url_for('user.store_data'), + headers={'Authorization': 'secret', 'X-api-uid': '10'}, + content_type='application/json') + self.assertStatus(r, 200) + self.assertIn('link_server', r.json) + self.assertEqual(r.json['link_server'], 'https://resolver.example.com/sfx') + self.assertEqual(r.json['other_data'], 'test') + + # Test updating to different library + r = self.client.post(url_for('user.store_data'), + headers={'Authorization': 'secret', 'X-api-uid': '10'}, + data=json.dumps({ + 'link_server': 'https://another.resolver.com/sfx', + 'other_data': 'updated' + }), + content_type='application/json') + self.assertStatus(r, 200) + self.assertEqual(r.json['link_server'], 'https://another.resolver.com/sfx') + self.assertEqual(r.json['other_data'], 'updated') + + # Verify library relationship was updated + with self.app.session_scope() as session: + user = session.query(User).filter_by(id=10).first() + self.assertIsNotNone(user) + self.assertIsNotNone(user.library_id) + + library = session.query(Library).filter_by(id=user.library_id).first() + self.assertIsNotNone(library) + self.assertEqual(library.libserver, 'https://another.resolver.com/sfx') + + # Test clearing library (empty string) + r = self.client.post(url_for('user.store_data'), + headers={'Authorization': 'secret', 'X-api-uid': '10'}, + data=json.dumps({ + 'link_server': '', + 'other_data': 'still here' + }), + content_type='application/json') + self.assertStatus(r, 200) + self.assertNotIn('link_server', r.json) + self.assertEqual(r.json['other_data'], 'still here') + + # Verify library relationship was cleared + with self.app.session_scope() as session: + user = session.query(User).filter_by(id=10).first() + self.assertIsNotNone(user) + self.assertIsNone(user.library_id) + + # Test POST without link_server - should not affect library + r = self.client.post(url_for('user.store_data'), + headers={'Authorization': 'secret', 'X-api-uid': '10'}, + data=json.dumps({ + 'some_setting': 'value' + }), + content_type='application/json') + self.assertStatus(r, 200) + self.assertEqual(r.json['some_setting'], 'value') + self.assertEqual(r.json['other_data'], 'still here') + self.assertNotIn('link_server', r.json) + + # Verify library relationship still cleared + with self.app.session_scope() as session: + user = session.query(User).filter_by(id=10).first() + self.assertIsNotNone(user) + self.assertIsNone(user.library_id) + + def test_library_integration_edge_cases(self): + '''Tests edge cases for library integration''' + + # Test with invalid/non-existent library + r = self.client.post(url_for('user.store_data'), + headers={'Authorization': 'secret', 'X-api-uid': '11'}, + data=json.dumps({ + 'link_server': 'https://nonexistent.library.com/sfx', + 'other_data': 'test' + }), + content_type='application/json') + self.assertStatus(r, 200) + # Should still save other data even if library doesn't exist + self.assertEqual(r.json['other_data'], 'test') + self.assertNotIn('link_server', r.json) + + # Verify no library relationship created + with self.app.session_scope() as session: + user = session.query(User).filter_by(id=11).first() + self.assertIsNotNone(user) + self.assertIsNone(user.library_id) + + # Test with None value + r = self.client.post(url_for('user.store_data'), + headers={'Authorization': 'secret', 'X-api-uid': '11'}, + data=json.dumps({ + 'link_server': None, + 'other_data': 'test2' + }), + content_type='application/json') + self.assertStatus(r, 200) + self.assertEqual(r.json['other_data'], 'test2') + self.assertNotIn('link_server', r.json) + + def test_library_data_migration_logic(self): + '''Tests the migration logic for converting link_server to library_id''' + + # Create a user with old-style link_server in user_data + with self.app.session_scope() as session: + from vault_service.models import User, Library + + # Create a library first + library = Library(libname='Test Library', libserver='https://test.library.com/sfx') + library2 = Library(libname='Test Library 2', libserver='https://test.library.updated.com/sfx') + session.add(library) + session.add(library2) + session.commit() + lib_id = library.id + lib_id2 = library2.id + + # Create user with old-style data + user = User(id=12, user_data={'link_server': 'https://test.library.com/sfx', 'other': 'data'}) + session.add(user) + session.commit() + + # Verify initial state + self.assertEqual(user.user_data['link_server'], 'https://test.library.com/sfx') + self.assertIsNone(user.library_id) + + # Test GET request + r = self.client.get(url_for('user.store_data'), + headers={'Authorization': 'secret', 'X-api-uid': '12'}, + content_type='application/json') + self.assertStatus(r, 200) + self.assertEqual(r.json['link_server'], 'https://test.library.com/sfx') + self.assertEqual(r.json['other'], 'data') + + # Get request does not change the data, only the response data + with self.app.session_scope() as session: + user = session.query(User).filter_by(id=12).first() + self.assertIsNotNone(user) + self.assertEqual(user.library_id, None) + self.assertIn('link_server', user.user_data) + self.assertEqual(user.user_data['other'], 'data') + + # Simulate the Alembic migration manually + with self.app.session_scope() as session: + user = session.query(User).filter_by(id=12).first() + library = session.query(Library).filter_by(libserver='https://test.library.com/sfx').first() + + # Migration: set library_id and remove link_server from user_data + user.library_id = library.id + user_data = user.user_data.copy() + user_data.pop('link_server', None) + user.user_data = user_data + session.commit() + + # Test GET request after migration - should return new format + r = self.client.get(url_for('user.store_data'), + headers={'Authorization': 'secret', 'X-api-uid': '12'}, + content_type='application/json') + self.assertStatus(r, 200) + self.assertEqual(r.json['link_server'], 'https://test.library.com/sfx') + self.assertEqual(r.json['other'], 'data') + + # Verify migration occurred in database + with self.app.session_scope() as session: + user = session.query(User).filter_by(id=12).first() + self.assertIsNotNone(user) + self.assertEqual(user.library_id, lib_id) + self.assertNotIn('link_server', user.user_data) + self.assertEqual(user.user_data['other'], 'data') + + # Test POST request - should update library data + r = self.client.post(url_for('user.store_data'), + headers={'Authorization': 'secret', 'X-api-uid': '12'}, + data=json.dumps({ + 'new_setting': 'value', + 'link_server': 'https://test.library.updated.com/sfx' + }), + content_type='application/json') + self.assertStatus(r, 200) + self.assertEqual(r.json['link_server'], 'https://test.library.updated.com/sfx') + self.assertEqual(r.json['other'], 'data') + self.assertEqual(r.json['new_setting'], 'value') + + # Verify POST didn't affect migration - data should remain in new format + with self.app.session_scope() as session: + user = session.query(User).filter_by(id=12).first() + self.assertIsNotNone(user) + self.assertEqual(user.library_id, lib_id2) + self.assertNotIn('link_server', user.user_data) + self.assertEqual(user.user_data['other'], 'data') + self.assertEqual(user.user_data['new_setting'], 'value') + if __name__ == '__main__': unittest.main() diff --git a/vault_service/views/user.py b/vault_service/views/user.py index b0ac6f2..35370ed 100644 --- a/vault_service/views/user.py +++ b/vault_service/views/user.py @@ -9,7 +9,7 @@ from sqlalchemy import exc from sqlalchemy.orm import exc as ormexc -from ..models import Query, User, MyADS +from ..models import Query, User, MyADS, Library from .utils import check_request, cleanup_payload, make_solr_request, upsert_myads, get_keyword_query_name from flask_discoverer import advertise from dateutil import parser @@ -158,11 +158,21 @@ def store_data(): if request.method == 'GET': with current_app.session_scope() as session: - q = session.query(User).filter_by(id=user_id).first() - if not q: - return '{}', 200 # or return 404? - return json.dumps(q.user_data) or '{}', 200 + user = session.query(User).filter_by(id=user_id).first() + if not user: + return '{}', 200 + response_data = user.user_data.copy() if user.user_data else {} + if user.library_id: + library = session.query(Library).filter_by(id=user.library_id).first() + if library: + response_data['link_server'] = library.libserver + return json.dumps(response_data) or '{}', 200 elif request.method == 'POST': + # Remove link_server from payload if present + library_server = payload.pop('link_server', None) + library = None + + # limit both number of keys and length of value to keep db clean if len(max(list(payload.values()), key=len)) > current_app.config['MAX_ALLOWED_JSON_SIZE']: return json.dumps({'msg': 'You have exceeded the allowed storage limit (length of values), no data was saved'}), 400 @@ -170,39 +180,39 @@ def store_data(): return json.dumps({'msg': 'You have exceeded the allowed storage limit (number of keys), no data was saved'}), 400 with current_app.session_scope() as session: - try: - q = session.query(User).filter_by(id=user_id).with_for_update(of=User).one() + user = session.query(User).filter_by(id=user_id).with_for_update(of=User).first() + if not user: + data = payload.copy() + user = User(id=user_id, user_data=data) + session.add(user) + else: try: - data = q.user_data + data = user.user_data or {} except TypeError: data = {} - except ormexc.NoResultFound: - data = payload - u = User(id=user_id, user_data=data) - try: - session.add(u) - session.commit() - return json.dumps(data), 200 - except exc.IntegrityError: - q = session.query(User).filter_by(id=user_id).with_for_update(of=User).one() - try: - data = q.user_data - except TypeError: - data = {} - - if data is None: - data = {} - data.update(payload) - q.user_data = data + data.update(payload) + user.user_data = data + + # Handle library selection (set or clear) + if library_server is not None: + if library_server: + library = session.query(Library).filter_by(libserver=library_server).first() + user.library_id = library.id if library else None + else: + user.library_id = None session.begin_nested() try: session.commit() + # Prepare response data (do not mutate user.user_data in-place) + response_data = data.copy() + if user.library_id and library: + response_data['link_server'] = library.libserver except exc.IntegrityError: session.rollback() return json.dumps({'msg': 'We have hit a db error! The world is crumbling all around... (eh, btw, your data was not saved)'}), 500 - return json.dumps(data), 200 + return json.dumps(response_data), 200 @advertise(scopes=[], rate_limit=[1000, 3600*24])