diff --git a/invokeai/app/services/shared/sqlite/sqlite_util.py b/invokeai/app/services/shared/sqlite/sqlite_util.py index fb8ca9fca38..645509f1dde 100644 --- a/invokeai/app/services/shared/sqlite/sqlite_util.py +++ b/invokeai/app/services/shared/sqlite/sqlite_util.py @@ -30,8 +30,6 @@ from invokeai.app.services.shared.sqlite_migrator.migrations.migration_25 import build_migration_25 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_26 import build_migration_26 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_27 import build_migration_27 -from invokeai.app.services.shared.sqlite_migrator.migrations.migration_28 import build_migration_28 -from invokeai.app.services.shared.sqlite_migrator.migrations.migration_29 import build_migration_29 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator @@ -79,8 +77,6 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto migrator.register_migration(build_migration_25(app_config=config, logger=logger)) migrator.register_migration(build_migration_26(app_config=config, logger=logger)) migrator.register_migration(build_migration_27()) - migrator.register_migration(build_migration_28()) - migrator.register_migration(build_migration_29()) migrator.run_migrations() return db diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_27.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_27.py index 533ec5b500a..b80ea073ef8 100644 --- a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_27.py +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_27.py @@ -1,4 +1,4 @@ -"""Migration 27: Add multi-user support. +"""Migration 27: Add multi-user support, per-user client state, and app settings. This migration adds the database schema for multi-user support, including: - users table for user accounts @@ -6,15 +6,19 @@ - user_invitations table for invitation system - shared_boards table for board sharing - Adding user_id columns to existing tables for data ownership +- Restructuring client_state table to support per-user storage +- app_settings table for storing JWT secret and other app-level settings """ +import json +import secrets import sqlite3 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration class Migration27Callback: - """Migration to add multi-user support.""" + """Migration to add multi-user support, per-user client state, and app settings.""" def __call__(self, cursor: sqlite3.Cursor) -> None: self._create_users_table(cursor) @@ -27,6 +31,9 @@ def __call__(self, cursor: sqlite3.Cursor) -> None: self._update_session_queue_table(cursor) self._update_style_presets_table(cursor) self._create_system_user(cursor) + self._update_client_state_table(cursor) + self._create_app_settings_table(cursor) + self._generate_jwt_secret(cursor) def _create_users_table(self, cursor: sqlite3.Cursor) -> None: """Create users table.""" @@ -209,11 +216,148 @@ def _create_system_user(self, cursor: sqlite3.Cursor) -> None: VALUES ('system', 'system@system.invokeai', 'System', '', FALSE, TRUE); """) + def _update_client_state_table(self, cursor: sqlite3.Cursor) -> None: + """Restructure client_state table to support per-user storage.""" + # Check if client_state table exists + cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='client_state';") + if cursor.fetchone() is None: + # Table doesn't exist, create it with the new schema + cursor.execute( + """ + CREATE TABLE client_state ( + user_id TEXT NOT NULL, + key TEXT NOT NULL, + value TEXT NOT NULL, + updated_at DATETIME NOT NULL DEFAULT (CURRENT_TIMESTAMP), + PRIMARY KEY (user_id, key), + FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE + ); + """ + ) + cursor.execute("CREATE INDEX IF NOT EXISTS idx_client_state_user_id ON client_state(user_id);") + cursor.execute( + """ + CREATE TRIGGER tg_client_state_updated_at + AFTER UPDATE ON client_state + FOR EACH ROW + BEGIN + UPDATE client_state + SET updated_at = CURRENT_TIMESTAMP + WHERE user_id = OLD.user_id AND key = OLD.key; + END; + """ + ) + return + + # Table exists with old schema - migrate it + # Get existing data if the data column is present (it may be absent if an older + # version of migration 21 was deployed without the column) + cursor.execute("PRAGMA table_info(client_state);") + columns = [row[1] for row in cursor.fetchall()] + existing_data = {} + if "data" in columns: + cursor.execute("SELECT data FROM client_state WHERE id = 1;") + row = cursor.fetchone() + if row is not None: + try: + existing_data = json.loads(row[0]) + except (json.JSONDecodeError, TypeError): + # If data is corrupt, just start fresh + pass + + # Drop the old table + cursor.execute("DROP TABLE IF EXISTS client_state;") + + # Create new table with per-user schema + cursor.execute( + """ + CREATE TABLE client_state ( + user_id TEXT NOT NULL, + key TEXT NOT NULL, + value TEXT NOT NULL, + updated_at DATETIME NOT NULL DEFAULT (CURRENT_TIMESTAMP), + PRIMARY KEY (user_id, key), + FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE + ); + """ + ) + + cursor.execute("CREATE INDEX IF NOT EXISTS idx_client_state_user_id ON client_state(user_id);") + + cursor.execute( + """ + CREATE TRIGGER tg_client_state_updated_at + AFTER UPDATE ON client_state + FOR EACH ROW + BEGIN + UPDATE client_state + SET updated_at = CURRENT_TIMESTAMP + WHERE user_id = OLD.user_id AND key = OLD.key; + END; + """ + ) + + # Migrate existing data to 'system' user + for key, value in existing_data.items(): + cursor.execute( + """ + INSERT INTO client_state (user_id, key, value) + VALUES ('system', ?, ?); + """, + (key, value), + ) + + def _create_app_settings_table(self, cursor: sqlite3.Cursor) -> None: + """Create app_settings table for storing application-level configuration.""" + cursor.execute( + """ + CREATE TABLE IF NOT EXISTS app_settings ( + key TEXT NOT NULL PRIMARY KEY, + value TEXT NOT NULL, + created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')), + updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')) + ); + """ + ) + + cursor.execute( + """ + CREATE TRIGGER IF NOT EXISTS tg_app_settings_updated_at + AFTER UPDATE ON app_settings + FOR EACH ROW + BEGIN + UPDATE app_settings SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW') + WHERE key = OLD.key; + END; + """ + ) + + def _generate_jwt_secret(self, cursor: sqlite3.Cursor) -> None: + """Generate and store a cryptographically secure JWT secret key. + + The secret is a 64-character hexadecimal string (256 bits of entropy), + which is suitable for HS256 JWT signing. + """ + # Check if JWT secret already exists + cursor.execute("SELECT value FROM app_settings WHERE key = 'jwt_secret';") + existing_secret = cursor.fetchone() + + if existing_secret is None: + # Generate a new cryptographically secure secret (256 bits) + jwt_secret = secrets.token_hex(32) # 32 bytes = 256 bits = 64 hex characters + + # Store in database + cursor.execute( + "INSERT INTO app_settings (key, value) VALUES ('jwt_secret', ?);", + (jwt_secret,), + ) + def build_migration_27() -> Migration: - """Builds the migration object for migrating from version 25 to version 27. + """Builds the migration object for migrating from version 26 to version 27. - This migration adds multi-user support to the database schema. + This migration adds multi-user support, per-user client state, and app settings + (including a JWT secret) to the database schema. """ return Migration( from_version=26, diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_28.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_28.py deleted file mode 100644 index 103332b51b9..00000000000 --- a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_28.py +++ /dev/null @@ -1,123 +0,0 @@ -"""Migration 28: Add user_id to client_state table for multi-user support. - -This migration updates the client_state table to support per-user state isolation: -- Drops the single-row constraint (CHECK(id = 1)) -- Adds user_id column -- Creates unique constraint on (user_id, key) pairs -- Migrates existing data to 'system' user -""" - -import json -import sqlite3 - -from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration - - -class Migration28Callback: - """Migration to add per-user client state support.""" - - def __call__(self, cursor: sqlite3.Cursor) -> None: - self._update_client_state_table(cursor) - - def _update_client_state_table(self, cursor: sqlite3.Cursor) -> None: - """Restructure client_state table to support per-user storage.""" - # Check if client_state table exists - cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='client_state';") - if cursor.fetchone() is None: - # Table doesn't exist, create it with the new schema - cursor.execute( - """ - CREATE TABLE client_state ( - user_id TEXT NOT NULL, - key TEXT NOT NULL, - value TEXT NOT NULL, - updated_at DATETIME NOT NULL DEFAULT (CURRENT_TIMESTAMP), - PRIMARY KEY (user_id, key), - FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE - ); - """ - ) - cursor.execute("CREATE INDEX IF NOT EXISTS idx_client_state_user_id ON client_state(user_id);") - cursor.execute( - """ - CREATE TRIGGER tg_client_state_updated_at - AFTER UPDATE ON client_state - FOR EACH ROW - BEGIN - UPDATE client_state - SET updated_at = CURRENT_TIMESTAMP - WHERE user_id = OLD.user_id AND key = OLD.key; - END; - """ - ) - return - - # Table exists with old schema - migrate it - # Get existing data if the data column is present (it may be absent if an older - # version of migration 21 was deployed without the column) - cursor.execute("PRAGMA table_info(client_state);") - columns = [row[1] for row in cursor.fetchall()] - existing_data = {} - if "data" in columns: - cursor.execute("SELECT data FROM client_state WHERE id = 1;") - row = cursor.fetchone() - if row is not None: - try: - existing_data = json.loads(row[0]) - except (json.JSONDecodeError, TypeError): - # If data is corrupt, just start fresh - pass - - # Drop the old table - cursor.execute("DROP TABLE IF EXISTS client_state;") - - # Create new table with per-user schema - cursor.execute( - """ - CREATE TABLE client_state ( - user_id TEXT NOT NULL, - key TEXT NOT NULL, - value TEXT NOT NULL, - updated_at DATETIME NOT NULL DEFAULT (CURRENT_TIMESTAMP), - PRIMARY KEY (user_id, key), - FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE - ); - """ - ) - - cursor.execute("CREATE INDEX IF NOT EXISTS idx_client_state_user_id ON client_state(user_id);") - - cursor.execute( - """ - CREATE TRIGGER tg_client_state_updated_at - AFTER UPDATE ON client_state - FOR EACH ROW - BEGIN - UPDATE client_state - SET updated_at = CURRENT_TIMESTAMP - WHERE user_id = OLD.user_id AND key = OLD.key; - END; - """ - ) - - # Migrate existing data to 'system' user - for key, value in existing_data.items(): - cursor.execute( - """ - INSERT INTO client_state (user_id, key, value) - VALUES ('system', ?, ?); - """, - (key, value), - ) - - -def build_migration_28() -> Migration: - """Builds the migration object for migrating from version 27 to version 28. - - This migration adds per-user client state support to prevent data leakage between users. - """ - return Migration( - from_version=27, - to_version=28, - callback=Migration28Callback(), - ) diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_29.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_29.py deleted file mode 100644 index 60be4b074aa..00000000000 --- a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_29.py +++ /dev/null @@ -1,77 +0,0 @@ -"""Migration 29: Add app_settings table for storing JWT secret and other app-level settings. - -This migration adds the app_settings table to securely store application-level configuration: -- Creates app_settings table with key-value storage -- Generates a random cryptographically secure JWT secret key -- Stores the JWT secret in the database for token signing/verification -""" - -import secrets -import sqlite3 - -from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration - - -class Migration29Callback: - """Migration to add app_settings table and JWT secret.""" - - def __call__(self, cursor: sqlite3.Cursor) -> None: - self._create_app_settings_table(cursor) - self._generate_jwt_secret(cursor) - - def _create_app_settings_table(self, cursor: sqlite3.Cursor) -> None: - """Create app_settings table for storing application-level configuration.""" - cursor.execute( - """ - CREATE TABLE IF NOT EXISTS app_settings ( - key TEXT NOT NULL PRIMARY KEY, - value TEXT NOT NULL, - created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')), - updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')) - ); - """ - ) - - cursor.execute( - """ - CREATE TRIGGER IF NOT EXISTS tg_app_settings_updated_at - AFTER UPDATE ON app_settings - FOR EACH ROW - BEGIN - UPDATE app_settings SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW') - WHERE key = OLD.key; - END; - """ - ) - - def _generate_jwt_secret(self, cursor: sqlite3.Cursor) -> None: - """Generate and store a cryptographically secure JWT secret key. - - The secret is a 64-character hexadecimal string (256 bits of entropy), - which is suitable for HS256 JWT signing. - """ - # Check if JWT secret already exists - cursor.execute("SELECT value FROM app_settings WHERE key = 'jwt_secret';") - existing_secret = cursor.fetchone() - - if existing_secret is None: - # Generate a new cryptographically secure secret (256 bits) - jwt_secret = secrets.token_hex(32) # 32 bytes = 256 bits = 64 hex characters - - # Store in database - cursor.execute( - "INSERT INTO app_settings (key, value) VALUES ('jwt_secret', ?);", - (jwt_secret,), - ) - - -def build_migration_29() -> Migration: - """Builds the migration object for migrating from version 28 to version 29. - - This migration adds the app_settings table and generates a JWT secret for token signing. - """ - return Migration( - from_version=28, - to_version=29, - callback=Migration29Callback(), - ) diff --git a/tests/test_sqlite_migrator.py b/tests/test_sqlite_migrator.py index fb2b26a2f63..e03224b5a9a 100644 --- a/tests/test_sqlite_migrator.py +++ b/tests/test_sqlite_migrator.py @@ -357,16 +357,30 @@ def test_migration_27_creates_users_table(logger: Logger) -> None: assert "user_id" in columns assert "is_public" in columns + # Verify client_state table has the new per-user schema + cursor.execute("PRAGMA table_info(client_state);") + columns = [row[1] for row in cursor.fetchall()] + assert "user_id" in columns + assert "key" in columns + assert "value" in columns + + # Verify app_settings table exists and contains a JWT secret + cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='app_settings';") + assert cursor.fetchone() is not None + cursor.execute("SELECT value FROM app_settings WHERE key = 'jwt_secret';") + jwt_row = cursor.fetchone() + assert jwt_row is not None + assert len(jwt_row[0]) == 64 # 32 bytes = 64 hex characters + db._conn.close() -def test_migration_28_with_existing_data_column(logger: Logger) -> None: - """Test that migration 28 correctly migrates existing data from the old schema with data column.""" +def test_migration_27_with_existing_client_state_data(logger: Logger) -> None: + """Test that migration 27 correctly migrates existing data from the old client_state schema.""" import json from invokeai.app.services.shared.sqlite_migrator.migrations.migration_21 import Migration21Callback from invokeai.app.services.shared.sqlite_migrator.migrations.migration_27 import Migration27Callback - from invokeai.app.services.shared.sqlite_migrator.migrations.migration_28 import Migration28Callback db = SqliteDatabase(db_path=None, logger=logger, verbose=False) cursor = db._conn.cursor() @@ -387,14 +401,12 @@ def test_migration_28_with_existing_data_column(logger: Logger) -> None: cursor.execute("CREATE TABLE IF NOT EXISTS session_queue (item_id INTEGER PRIMARY KEY);") cursor.execute("CREATE TABLE IF NOT EXISTS style_presets (id TEXT PRIMARY KEY);") db._conn.commit() - Migration27Callback()(cursor) - db._conn.commit() - # Run migration 28 - Migration28Callback()(cursor) + # Run migration 27 + Migration27Callback()(cursor) db._conn.commit() - # Verify new schema + # Verify new client_state schema cursor.execute("PRAGMA table_info(client_state);") columns = [row[1] for row in cursor.fetchall()] assert "user_id" in columns @@ -413,10 +425,9 @@ def test_migration_28_with_existing_data_column(logger: Logger) -> None: db._conn.close() -def test_migration_28_without_data_column(logger: Logger) -> None: - """Test that migration 28 handles old client_state table without the data column.""" +def test_migration_27_without_client_state_data_column(logger: Logger) -> None: + """Test that migration 27 handles old client_state table without the data column.""" from invokeai.app.services.shared.sqlite_migrator.migrations.migration_27 import Migration27Callback - from invokeai.app.services.shared.sqlite_migrator.migrations.migration_28 import Migration28Callback db = SqliteDatabase(db_path=None, logger=logger, verbose=False) cursor = db._conn.cursor() @@ -439,14 +450,12 @@ def test_migration_28_without_data_column(logger: Logger) -> None: cursor.execute("CREATE TABLE IF NOT EXISTS session_queue (item_id INTEGER PRIMARY KEY);") cursor.execute("CREATE TABLE IF NOT EXISTS style_presets (id TEXT PRIMARY KEY);") db._conn.commit() - Migration27Callback()(cursor) - db._conn.commit() - # Run migration 28 - should not raise even without data column - Migration28Callback()(cursor) + # Run migration 27 - should not raise even without data column + Migration27Callback()(cursor) db._conn.commit() - # Verify new schema + # Verify new client_state schema cursor.execute("PRAGMA table_info(client_state);") columns = [row[1] for row in cursor.fetchall()] assert "user_id" in columns