diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index c39e9a3da3b..253744fabba 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -155,3 +155,96 @@ jobs: # one worker per vCPU is the right level where the runner has the headroom for it. The # matrix carries the count, because windows needs a lower one. run: uv run --no-sync pytest -n ${{ matrix.pytest-workers }} --max-worker-restart=0 + + # Decides whether the external-database job runs. That job's service containers start before any of its + # steps, so the decision has to come from a job of its own: otherwise every pull request, frontend-only ones + # included, would pull both server images and wait for them to become healthy. + external-database-changes: + name: 'external database: change detection' + runs-on: ubuntu-24.04 + timeout-minutes: 10 + outputs: + run: ${{ steps.changed-files.outputs.python_any_changed == 'true' || inputs.always_run == true }} + steps: + - name: checkout + if: ${{ inputs.always_run != true }} + uses: nschloe/action-cached-lfs-checkout@385a8ecc719e50b8c71af6ab01a624b486b7c3bc + + - name: check for changed python files + if: ${{ inputs.always_run != true }} + id: changed-files + # Pinned to the _hash_ for v45.0.9, see the matrix job above. + uses: tj-actions/changed-files@9426d40962ed5378910ee2e21d5f8c6fcbf2dd96 + with: + files_yaml: | + python: + - 'pyproject.toml' + - 'uv.lock' + - '.github/workflows/python-tests.yml' + - 'invokeai/**/*.py' + # The library's tests store the bundled workflows and style presets. + - 'invokeai/app/services/workflow_records/default_workflows/**' + - 'invokeai/app/services/style_preset_records/default_style_presets.json' + - 'tests/**' + + # The database layer against MySQL and MariaDB, at the oldest release of each that InvokeAI supports. + # Only the tests on the database fixtures (marked `uses_database`) target the server; the matrix above + # runs the whole suite, those tests included, on SQLite. + external-database: + needs: external-database-changes + if: ${{ needs.external-database-changes.outputs.run == 'true' }} + strategy: + fail-fast: false + matrix: + include: + - server: mysql + image: 'mysql:8.4.11' + health-cmd: 'mysqladmin ping -h 127.0.0.1 --silent' + url: 'mysql+pymysql://root:invokeai@127.0.0.1:3306' + - server: mariadb + image: 'mariadb:10.11.19' + health-cmd: 'mariadb-admin ping -h 127.0.0.1 --silent' + url: 'mariadb+pymysql://root:invokeai@127.0.0.1:3306' + name: 'py3.12: linux-${{ matrix.server }}' + runs-on: ubuntu-24.04 + timeout-minutes: 30 + services: + database: + image: ${{ matrix.image }} + env: + MYSQL_ROOT_PASSWORD: invokeai + MARIADB_ROOT_PASSWORD: invokeai + ports: + - '3306:3306' + options: >- + --health-cmd="${{ matrix.health-cmd }}" + --health-interval=5s + --health-timeout=5s + --health-retries=30 + env: + OMP_NUM_THREADS: '1' + MKL_NUM_THREADS: '1' + INVOKEAI_TEST_DB_URL: ${{ matrix.url }} + + steps: + - name: checkout + uses: nschloe/action-cached-lfs-checkout@385a8ecc719e50b8c71af6ab01a624b486b7c3bc + + - name: setup uv + uses: astral-sh/setup-uv@v8.1.0 + with: + version: '0.11.28' # When updating this, make sure the version in launcher's package.json gets updated too. + enable-cache: true + python-version: '3.12' + + - name: install dependencies + env: + UV_INDEX: 'https://download.pytorch.org/whl/cpu' + run: uv sync --no-progress --locked --extra test --extra mysql + + - name: run pytest + run: uv run --no-sync pytest -n logical --max-worker-restart=0 -m "uses_database and not sqlite_only and not slow" + + - name: server log + if: ${{ failure() }} + run: docker logs "${{ job.services.database.id }}" diff --git a/docker/Dockerfile b/docker/Dockerfile index 6be57bd8089..a323af9a7ed 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -84,7 +84,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # this is just to get the package manager to recognize that the project exists, without making changes to the docker layer --mount=type=bind,source=invokeai/version,target=invokeai/version \ ulimit -n 30000 && \ - uv sync --extra $GPU_DRIVER --frozen + uv sync --extra $GPU_DRIVER --extra mysql --frozen # Link amdgpu.ids for ROCm builds # contributed by https://github.com/Rubonnek @@ -114,5 +114,5 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,source=pyproject.toml,target=pyproject.toml \ --mount=type=bind,source=uv.lock,target=uv.lock \ ulimit -n 30000 && \ - uv pip install -e .[$GPU_DRIVER] + uv pip install -e .[$GPU_DRIVER,mysql] diff --git a/docs/src/content/docs/Users Guide/03.Models/01.introduction.mdx b/docs/src/content/docs/Users Guide/03.Models/01.introduction.mdx index 973b1e00c19..71708b62470 100644 --- a/docs/src/content/docs/Users Guide/03.Models/01.introduction.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/01.introduction.mdx @@ -235,16 +235,18 @@ Both actions handle partial failures: if some models succeed and others fail, th ## Finding orphaned models -If a model file is deleted or moved outside the Model Manager, its -database entry sticks around. To find these orphaned entries: +Model files can stay in the models directory after their database entry is gone, for example when the +database was reset, or files were copied into the directory by hand. To find these orphaned models: 1. Open the Model Manager. 2. Click the "..." menu at the top right of the installed model list, and pick **"Clean up orphaned models..."** -3. If there are orphaned models, a dialogue box appears that lists them and offers to remove their database entries. +3. If there are orphaned models, a dialogue box appears that lists their folders and offers to delete them. -Orphaned models are automatically excluded from selection dropdowns -(main model, LoRA, VAE, etc.), so you cannot accidentally pick one for -generation. +On an install without the web UI at hand, an administrator can do the same through the API: +`GET /api/v2/models/sync/orphaned` lists the folders and `DELETE /api/v2/models/sync/orphaned` deletes those given. + +The opposite case, a database entry whose files are missing from disk, is marked in the Model Manager's +model list, and its filter can show only those models. ## Exporting and Importing Model Settings diff --git a/docs/src/content/docs/configuration/01.invokeai-yaml.mdx b/docs/src/content/docs/configuration/01.invokeai-yaml.mdx index 8cef7ee7149..28f3514631a 100644 --- a/docs/src/content/docs/configuration/01.invokeai-yaml.mdx +++ b/docs/src/content/docs/configuration/01.invokeai-yaml.mdx @@ -249,6 +249,8 @@ If you serve InvokeAI through nginx, Caddy, or similar, set `http_compression_le #### Database Durability +This section concerns the default SQLite database. A [MySQL or MariaDB database](/configuration/external-database/), set with `db_url`, has its durability set on the server. + Every write to InvokeAI's SQLite database is committed to disk before the request returns. How hard that commit is pushed to the physical disk is controlled by `db_synchronous`: ```yaml @@ -268,7 +270,7 @@ Running the database on a network share is not a supported configuration -- SQLi The safeguard exists anyway, because WAL needs shared memory and a few setups cannot provide it: a database on a network share, some container volume drivers, and SQLite builds compiled without WAL all leave a rollback journal in place instead — and on a rollback journal, `normal` _can_ corrupt the database on power loss. InvokeAI reads back the journal mode it actually got, so in that case it ignores `normal`, uses `full`, and logs a warning naming the mode. An in-memory database (`use_memory_db`) has no durability to trade and is unaffected. -`db_synchronous` applies to the application's own database connection. Gallery maintenance is an administrator action in Settings and uses that same application database connection, so its database work follows this setting. The image import script, `invokeai/frontend/install/import_images.py`, still opens the database independently and leaves SQLite's own default in place; the measurements below describe the running app, not that script. +`db_synchronous` applies to the application's own database connection, which gallery maintenance (an administrator action in Settings) uses too, and to the image import script, `invokeai/frontend/install/import_images.py`, which reads the same `invokeai.yaml`. The measurements below describe the running app. Measured on a copy of a real library, 300 single-row inserts each committed on its own: diff --git a/docs/src/content/docs/configuration/02.docker.mdx b/docs/src/content/docs/configuration/02.docker.mdx index 9ab015bf304..64614a61fb5 100644 --- a/docs/src/content/docs/configuration/02.docker.mdx +++ b/docs/src/content/docs/configuration/02.docker.mdx @@ -60,6 +60,8 @@ Any environment variables supported by InvokeAI can be set here - please see the At the very least, you might want to set the `INVOKEAI_ROOT` environment variable to point to the location where you wish to store your InvokeAI models, configuration, and outputs. +The image includes the driver for a [MySQL or MariaDB database](/configuration/external-database/): set `INVOKEAI_DB_URL` to use one instead of the SQLite file in the root directory. + | Environment Variable | Default value | Description | | --- | --- | --- | | `INVOKEAI_ROOT` | `~/invokeai` | **Required** - the location of your InvokeAI root directory. It will be created if it does not exist. | diff --git a/docs/src/content/docs/configuration/03.external-database.mdx b/docs/src/content/docs/configuration/03.external-database.mdx new file mode 100644 index 00000000000..f1a9cc4f348 --- /dev/null +++ b/docs/src/content/docs/configuration/03.external-database.mdx @@ -0,0 +1,110 @@ +--- +title: MySQL and MariaDB +sidebar: + order: 3 +--- + +InvokeAI keeps its boards, image records, workflows, projects, queue and accounts in a database. By default that is a SQLite file in `db_dir`, which needs no setup. On a server that hosts InvokeAI for several people, the database can be a MySQL or MariaDB database instead, run and backed up alongside the server's other databases. + +## Requirements + +- MySQL 8.4 or newer, or MariaDB 10.11 or newer, with the InnoDB storage engine. +- `max_allowed_packet` of at least 64 MiB on the server. MariaDB's default is 16 MiB; InvokeAI refuses to start below 64 MiB, because a large project is stored in one statement. Set it in the server's configuration and restart the server: + + ```ini + [mysqld] + max_allowed_packet=64M + ``` + +- The `mysql` extra of the `invokeai` package, which installs the PyMySQL driver: `pip install "invokeai[mysql]"`, or `uv sync --extra mysql` in a source checkout. The Docker image includes it. + +## Creating the database + +Create an empty database and an account for InvokeAI. InvokeAI creates its tables, each with the character set and collations it needs, when it first starts. + +MySQL: + +```sql +CREATE DATABASE invokeai CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_bin; +CREATE USER 'invokeai'@'%' IDENTIFIED BY 'a-strong-password'; +GRANT ALL PRIVILEGES ON invokeai.* TO 'invokeai'@'%'; +``` + +MariaDB: + +```sql +CREATE DATABASE invokeai CHARACTER SET utf8mb4 COLLATE utf8mb4_nopad_bin; +CREATE USER 'invokeai'@'%' IDENTIFIED BY 'a-strong-password'; +GRANT ALL PRIVILEGES ON invokeai.* TO 'invokeai'@'%'; +``` + +## Configuring InvokeAI + +Set `db_url` in `invokeai.yaml`, or the `INVOKEAI_DB_URL` environment variable: + +```yaml +db_url: mariadb+pymysql://invokeai:a-strong-password@db.example:3306/invokeai +``` + +Use `mysql+pymysql://` for MySQL and `mariadb+pymysql://` for MariaDB. Characters in the password that have a meaning in a URL (`@`, `:`, `/`, `%`, `#`) must be percent-encoded, e.g. `@` as `%40`. InvokeAI reads `db_url` when it starts, and shows it without the password in its settings. `db_dir` and `db_synchronous` apply to a SQLite database only. + +On its first start InvokeAI creates the tables. On later starts it updates them to its version, as it does a SQLite database. It makes no backup of a server database before an update: back the database up yourself, e.g. with `mysqldump` or `mariadb-dump`, before upgrading InvokeAI. + +## Moving an existing install + +`invoke-db-copy` copies an install's SQLite database into a new, empty server database. It copies the database only: images, models and the other files stay where they are, in the root directory. The SQLite database is not changed. + +1. Stop InvokeAI. What it writes during the copy is not copied. +2. Set `db_url` in `invokeai.yaml` (or `INVOKEAI_DB_URL`) to the new database. Do not start InvokeAI yet: it would start the new database empty. +3. Check the copy, which copies nothing yet: + + ```bash + invoke-db-copy --check + ``` + +4. Copy: + + ```bash + invoke-db-copy + ``` + +5. Start InvokeAI. + +`invoke-db-copy` takes the install it runs in, or the one `--root` names, and copies into the database its `db_url` names. `--target URL` names another one; a password on the command line is visible to the other users of the computer and stays in the shell's history. + +It copies from a snapshot of the SQLite database, which it makes in the database's directory and deletes afterwards: that needs free space as large as the database. It brings the snapshot to no newer version: if InvokeAI has not started since an upgrade, start it once first. The target must be empty, for the check too, and no InvokeAI process may use it. + +It stops before copying anything when the server would refuse a row: + +- **Rows whose foreign keys name a row that is gone**, which a SQLite database can hold and a server refuses. `--orphans skip` copies them without the reference where the database would clear it anyway (a board whose cover image is gone keeps no cover), and leaves out the others, together with what the database deletes with them. The report lists what it left out, per table. +- **Values longer than the server's column holds**, and rows larger than half of `max_allowed_packet`. + +Two kinds of value are copied in the form a server stores: + +- **Vocabulary terms of the image index that differ only in the case of a letter outside A-Z** (`Äpfel` and `äpfel`), which the server compares as equal: it keeps one of them. +- **A model's file size stored with a fraction**, which is written as a whole number of bytes. + +After copying, `invoke-db-copy` compares every table of the copy with its source, row by row, and only then writes the records that make it a database InvokeAI takes. If the copy fails or does not match, InvokeAI refuses the target: drop the database, create it again empty, and run `invoke-db-copy` again. + +## Going back to SQLite + +`invoke-db-copy --to-sqlite FILE` copies the server database back into a new SQLite file. Removing `db_url` alone is not a way back: InvokeAI then opens the SQLite database it used before the move, as it was then. The boards, image records, workflows and accounts made since are not in it, and the gallery no longer shows the images made since, although their files are in the outputs folder. + +1. Stop InvokeAI. `invoke-db-copy` refuses to copy a database an InvokeAI process uses. +2. Copy into a new file: + + ```bash + invoke-db-copy --to-sqlite invokeai-from-server.db + ``` + +3. In the install's database directory (`db_dir`, by default `databases` in the root), rename `invokeai.db`, which holds the state before the move, and put the new file in its place, named `invokeai.db`. +4. Remove `db_url` from `invokeai.yaml`, and `INVOKEAI_DB_URL` from the environment, and start InvokeAI. + +The source is the database `db_url` names, unless `--source URL` names another; it is not changed. The new file appears only once every table matches its source, so a copy that fails leaves no file behind, and `invoke-db-copy` never overwrites an existing file. Ids continue as they did on the server: a queue item gets no id the server already gave out. + +## Limits + +- **One InvokeAI process per database.** At startup a process cancels the queue items left running and updates its bundled workflows and style presets, which a second process would undo. The first process holds a lock on the database; a second one refuses to start. If a process ends without closing its connection, because its computer crashed or lost the network, the server releases its lock within ten minutes; the refusal names the query that finds the connection to end sooner with `KILL`. +- **One database server.** Setups with several writing servers (Galera, group replication with several primaries) are not supported. +- **Busy database.** When a request waits too long for a lock held by another, finds no free connection, or cannot reach the server, it fails with HTTP 503 and `Retry-After: 1`, and can be repeated. If this happens often, the server is overloaded or another program holds locks in InvokeAI's database. +- **Command-line tools.** The account commands (`invoke-useradd`, `invoke-userdel`, `invoke-usermod`, `invoke-userlist`) and the maintenance scripts use the database `invokeai.yaml` names. They refuse a database InvokeAI has not updated to their version: start InvokeAI once after an upgrade before using them. diff --git a/docs/src/content/docs/configuration/Multi-User Mode/admin-guide.mdx b/docs/src/content/docs/configuration/Multi-User Mode/admin-guide.mdx index 1a2e518f4b9..c9894b15405 100644 --- a/docs/src/content/docs/configuration/Multi-User Mode/admin-guide.mdx +++ b/docs/src/content/docs/configuration/Multi-User Mode/admin-guide.mdx @@ -381,7 +381,7 @@ Then cut and paste the printed secret into this Sqlite3 command: sqlite3 INVOKE_ROOT/databases/invokeai.db 'update app_settings set value="THE_SECRET" where key="jwt_secret"' ``` -(replace INVOKE_ROOT with your InvokeAI root directory and THE_SECRET with the new secret). +(replace INVOKE_ROOT with your InvokeAI root directory and THE_SECRET with the new secret). With a [MySQL or MariaDB database](/configuration/external-database/), run the same `UPDATE` in the server's client instead, with the value in single quotes. After this, restart the server. All logged in users will be logged out and will need to provide their usernames and passwords again. @@ -451,7 +451,7 @@ It is best to restrict access to trusted networks and remote IP addresses, or us **Backup and Recovery:** It is always a good idea to periodically backup your InvokeAI database and images, but especially -so if the server is publicly accessible to the Internet. +so if the server is publicly accessible to the Internet. The commands below back up a SQLite database; back up a [MySQL or MariaDB database](/configuration/external-database/) with the server's own tools, such as `mysqldump` or `mariadb-dump`. **Manual Backup:** diff --git a/docs/src/content/docs/development/Architecture/model-manager.mdx b/docs/src/content/docs/development/Architecture/model-manager.mdx index fab6f1f4fcc..7472dac054b 100644 --- a/docs/src/content/docs/development/Architecture/model-manager.mdx +++ b/docs/src/content/docs/development/Architecture/model-manager.mdx @@ -110,9 +110,9 @@ The config hierarchy has a major limitation in its handling of the base model ty ## Reading and Writing Model Configuration Records -The `ModelRecordService` provides the ability to retrieve model configuration records from SQL or YAML databases, update them, and write them back. +The model record service keeps each installed model's configuration record in the `models` table of the InvokeAI database, as the config's JSON; the columns models are searched by (`name`, `base`, `type`, `format`, `path`, `hash`, ...) are generated from that document. `ModelRecordServiceSQL` implements `ModelRecordServiceBase` on the [database layer](/development/guides/database-layer/). -A application-wide `ModelRecordService` is created during API initialization and can be retrieved within an invocation from the `InvocationContext` object: +An application-wide record service is created during API initialization and can be retrieved within an invocation from the `InvocationContext` object: ```py store = context.services.model_manager.store @@ -120,68 +120,29 @@ store = context.services.model_manager.store or from elsewhere in the code by accessing `ApiDependencies.invoker.services.model_manager.store`. -### Creating a `ModelRecordService` +### Creating a `ModelRecordServiceSQL` -To create a new `ModelRecordService` database or open an existing one, you can directly create either a `ModelRecordServiceSQL` or a `ModelRecordServiceFile` object: +The service takes the application's database and a logger. To open the records of an existing, migrated database outside the server: ```py -from invokeai.app.services.model_records import ModelRecordServiceSQL, ModelRecordServiceFile - -store = ModelRecordServiceSQL.from_connection(connection, lock) -store = ModelRecordServiceSQL.from_db_file('/path/to/sqlite_database.db') -store = ModelRecordServiceFile.from_db_file('/path/to/database.yaml') -``` - -The `from_connection()` form is only available from the `ModelRecordServiceSQL` class, and is used to manage records in a previously-opened SQLITE3 database using a `sqlite3.connection` object and a `threading.lock` object. It is intended for the specific use case of storing the record information in the main InvokeAI database, usually `databases/invokeai.db`. - -The `from_db_file()` methods can be used to open new connections to the named database files. If the file doesn't exist, it will be created and initialized. - -As a convenience, `ModelRecordServiceBase` offers two methods, `from_db_file` and `open`, which will return either a SQL or File implementation depending on the context. The former looks at the file extension to determine whether to open the file as a SQL database (".db") or as a file database (".yaml"). If the file exists, but is either the wrong type or does not contain the expected schema metainformation, then an appropriate `AssertionError` will be raised: - -```py -store = ModelRecordServiceBase.from_db_file('/path/to/a/file.{yaml,db}') -``` - -The `ModelRecordServiceBase.open()` method is specifically designed for use in the InvokeAI web server. Its signature is: - -```py -def open( - cls, - config: InvokeAIAppConfig, - conn: Optional[sqlite3.Connection] = None, - lock: Optional[threading.Lock] = None - ) -> Union[ModelRecordServiceSQL, ModelRecordServiceFile]: -``` - -The way it works is as follows: - -1. Retrieve the value of the `model_config_db` option from the user's `invokeai.yaml` config file. -2. If `model_config_db` is `auto` (the default), then: - - Use the values of `conn` and `lock` to return a `ModelRecordServiceSQL` object opened on the passed connection and lock. - - Open up a new connection to `databases/invokeai.db` if `conn` and/or `lock` are missing (see note below). -3. If `model_config_db` is a Path, then use `from_db_file` to return the appropriate type of ModelRecordService. -4. If `model_config_db` is None, then retrieve the legacy `conf_path` option from `invokeai.yaml` and use the Path indicated there. This will default to `configs/models.yaml`. - -So a typical startup pattern would be: - -```py -import sqlite3 -from invokeai.app.services.thread import lock -from invokeai.app.services.model_records import ModelRecordServiceBase -from invokeai.app.services.config import InvokeAIAppConfig +from invokeai.app.services.config import get_config +from invokeai.app.services.model_records import ModelRecordServiceSQL +from invokeai.app.services.shared.database.database import Database +from invokeai.backend.util.logging import InvokeAILogger -config = InvokeAIAppConfig.get_config() -db_conn = sqlite3.connect(config.db_path.as_posix(), check_same_thread=False) -store = ModelRecordServiceBase.open(config, db_conn, lock) +config = get_config() +logger = InvokeAILogger.get_logger(config=config) +database = Database.open_sqlite(config.db_path, logger) +store = ModelRecordServiceSQL(database, logger) ``` -### Fetching a Model's Configuration from `ModelRecordServiceBase` +### Fetching a Model's Configuration Configurations can be retrieved in several ways. #### get_model(key) -> AnyModelConfig -The basic functionality is to call the record store object's `get_model()` method with the desired model's unique key. It returns the appropriate subclass of ModelConfigBase: +The basic functionality is to call the record store object's `get_model()` method with the desired model's unique key. It returns the appropriate subclass of `ModelConfigBase`: ```py model_conf = store.get_model('f13dd932c0c35c22dcb8d6cda4203764') @@ -193,62 +154,55 @@ print(model_conf.path) If the key is unrecognized, this call raises an `UnknownModelException`. -#### exists(key) -> AnyModelConfig +#### exists(key) -> bool Returns True if a model with the given key exists in the database. -#### search_by_path(path) -> AnyModelConfig +#### search_by_path(path) -> List[AnyModelConfig] -Returns the configuration of the model whose path is `path`. The path is matched using a simple string comparison and won't correctly match models referred to by different paths (e.g. using symbolic links). +Returns the configuration of the model whose path is `path`, in a list of at most one since paths are unique. The path is matched using a simple string comparison and won't correctly match models referred to by different paths (e.g. using symbolic links). -#### search_by_name(name, base, type) -> List[AnyModelConfig] +#### search_by_hash(hash) -> List[AnyModelConfig] -This method searches for models that match some combination of `name`, `BaseType` and `ModelType`. Calling without any arguments will return all the models in the database. +Returns the configurations of the models whose file has the given hash. -#### all_models() -> List[AnyModelConfig] +#### search_by_attr(model_name, base_model, model_type, model_format, order_by, direction) -> List[AnyModelConfig] -Return all the model configs in the database. Exactly equivalent to calling `search_by_name()` with no arguments. +Returns the models that match every filter given among `model_name`, `base_model`, `model_type` and `model_format`, ordered by a `ModelRecordOrderBy` (type, base, name and format by default; names and bases ignore case) in `direction`. Models that tie are ordered by key, in the same direction. Calling it without filters returns all the models in the database. -#### search_by_tag(tags) -> List[AnyModelConfig] +A stored config that no longer validates, for instance one written by a newer version of the application, is skipped with a warning in the log. -`tags` is a list of strings. This method returns a list of model configs that contain all of the given tags. Examples: +#### all_models() -> List[AnyModelConfig] -```py -# find all models that are marked as both SFW and as generating -# background scenery -configs = store.search_by_tag(['sfw', 'scenery']) -``` +Return all the model configs in the database. Exactly equivalent to calling `search_by_attr()` with no arguments. -Note that only tags are not searchable in this way. Other fields can be searched using a filter: +#### model_info_by_name(model_name, base_model, model_type) -> AnyModelConfig -```py -commercializable_models = [x for x in store.all_models() \ - if x.license.contains('allowCommercialUse=Sell')] -``` +Looks for the one model identified by name, base model and model type and returns it. It raises a `DuplicateModelException` if more than one model shares the same type, base and name, which is possible (one at `/foo/my_model`, another at `/bar/my_model`), and an `UnknownModelException` if none does. Prefer `search_by_attr()`, which can return several models, and pass the key of the one you want to `get_model()`. -#### version() -> str +#### get_model_paths() -> List[str] -Returns the version of the database, currently at `3.2` +Returns the path of every record, including records whose config no longer validates. The orphaned-model scan uses it, so that it never reports the files of a registered model. -#### model_info_by_name(name, base_model, model_type) -> ModelConfigBase +### Writing model configs to the database -This method exists to ease the transition from the previous version of the model manager, in which `get_model()` took the three arguments shown above. This looks for a unique model identified by name, base model and model type and returns it. +Several methods allow you to create and update stored model config records. -The method will generate a `DuplicateModelException` if there are more than one models that share the same type, base and name. While unlikely, it is certainly possible to have a situation in which the user had added two models with the same name, base and type, one located at path `/foo/my_model` and the other at `/bar/my_model`. It is strongly recommended to search for models using `search_by_name()`, which can return multiple results, and then to select the desired model and pass its key to `get_model()`. +#### add_model(config) -> AnyModelConfig -### Writing model configs to the database +Adds the configuration record; the config's `key` becomes the record's key. It returns the stored config, and raises a `DuplicateModelException` if a model with the same path or key is already installed. -Several methods allow you to create and update stored model config records. +#### update_model(key, changes, allow_class_change=False) -> AnyModelConfig -#### add_model(key, config) -> AnyModelConfig +Applies a `ModelRecordChanges` to the record and returns the updated config. The changes are validated against the model's config class; with `allow_class_change`, a change may turn the record into another class (a LoRA into a Main model, say), which must validate as well. The record is locked while the update is applied, so concurrent updates of one model apply one after the other. Raises `UnknownModelException` for an unknown key, or pydantic's `ValidationError` for invalid changes. -Given a key and a configuration, this will add the model's configuration record to the database. `config` can either be a subclass of `ModelConfigBase` (i.e. any class listed in `AnyModelConfig`), or a `dict` of key/value pairs. In the latter case, the correct configuration class will be picked by Pydantic's discriminated union mechanism. +#### replace_model(key, new_config) -> AnyModelConfig -If successful, the method will return the appropriate subclass of `ModelConfigBase`. It will raise a `DuplicateModelException` if a model with the same key is already in the database, or an `InvalidModelConfigException` if a dict was passed and Pydantic experienced a parse or validation error. +Replaces the whole record, for instance after a model was identified anew. `new_config.key` must equal `key`. Raises `UnknownModelException` for an unknown key. -### update_model(key, config) -> AnyModelConfig +#### del_model(key) -Given a key and a configuration, this will update the model configuration record in the database. `config` can be either a instance of `ModelConfigBase`, or a sparse `dict` containing the fields to be updated. This will return an `AnyModelConfig` on success, or raise `InvalidModelConfigException` or `UnknownModelException` exceptions on failure. +Deletes the record, and the model's relationships with other models. Raises `UnknownModelException` for an unknown key. --- @@ -284,14 +238,14 @@ from invokeai.app.services.config import get_config from invokeai.app.services.model_records import ModelRecordServiceSQL from invokeai.app.services.model_install import ModelInstallService from invokeai.app.services.download import DownloadQueueService -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database from invokeai.backend.util.logging import InvokeAILogger config = get_config() logger = InvokeAILogger.get_logger(config=config) -db = SqliteDatabase(config.db_path, logger) -record_store = ModelRecordServiceSQL(db, logger) +database = Database.open_sqlite(config.db_path, logger) +record_store = ModelRecordServiceSQL(database, logger) queue = DownloadQueueService() queue.start() diff --git a/docs/src/content/docs/development/Guides/database-layer.mdx b/docs/src/content/docs/development/Guides/database-layer.mdx new file mode 100644 index 00000000000..2e2115bc397 --- /dev/null +++ b/docs/src/content/docs/development/Guides/database-layer.mdx @@ -0,0 +1,186 @@ +--- +title: Database Layer +lastUpdated: 2026-10-04 +--- + +Every SQL statement of the application belongs in `invokeai/app/services/shared/database/`. Services reach the database through `Database.queries`, never through a connection, a cursor or SQL text of their own. `tests/app/test_database_access_guard.py` fails when code elsewhere imports a database driver or SQLAlchemy, executes a statement or carries SQL. Code that keeps a SQLite file of its own, as a scratch store in a temporary directory, is not the application's database: it is listed in the test's `NOT_THE_APPLICATION_DATABASE`, with the reason. + +Queries are written with SQLAlchemy Core, which compiles each statement for the backend in use. SQLite is the default backend. The layer and its tests also run on MySQL 8.4 and MariaDB 10.11; configuring InvokeAI itself to use one of them is not available yet. + +## Query modules + +Each domain gets a query module as it is ported: a `QueryModule` subclass in `queries/`, attached to `Queries` as a cached property, whose methods take the connection right after `self` and are decorated with `@read` or `@write`. The decorators supply the connection, so callers never see one. + +```py +from sqlalchemy import Connection, bindparam, select + +from invokeai.app.services.shared.database.queries.base import QueryModule, read +from invokeai.app.services.shared.database.schema.app_settings import app_settings + +_GET = select(app_settings.c.value).where(app_settings.c.key == bindparam("key")) + + +class AppSettingQueries(QueryModule): + @read + def get(self, conn: Connection, key: str) -> str | None: + return conn.execute(_GET, {"key": key}).scalar_one_or_none() +``` + +Build each statement once, at module level, with `bindparam()` for its values. SQLAlchemy compiles a statement once, but a statement object built per call costs several times what executing a prebuilt one does: on SQLite, about 60 µs against 18 µs for a point read. Build per call only what varies in shape, such as an update of the fields a caller gives, and keep it off hot paths; a statement that takes one of a few shapes comes from a builder cached with `functools.cache`, keyed by values whose equality means the same statement (the API's enums are also strings, so compare them with `==`, never `is`, and normalize flags with `bool()`, since `1 == True`). SQLAlchemy keeps every statement it compiled, so a client must not choose how many shapes there are: normalize what it sends (a list of categories becomes a sorted set), round a length it chooses up to one of a few sizes with the spare values bound to `NULL`, and build anything larger for its call alone, executed with `execution_options={"compiled_cache": None}`. In an `UPDATE`, a bound parameter cannot share its name with a column, since those names are reserved for the values it sets. Read result rows by position, by unpacking them: looking a value up by name costs about ten times as much, which adds up over a page of rows. Fetch a result whole with `all()` (or `scalars().all()`) even to loop over it: iterating the result itself fetches row by row, about 10 % slower over thousands of rows on SQLite. A fixed list of values in an `IN` is bound value by value (`in_([literal(value) for value in values])`): a plain list becomes an expanding parameter, whose SQL is rendered anew on every execution; that suits a list the caller passes, as `bindparam(name, expanding=True)`. + +A query method returns plain values, or its rows read completely (`first()`, `all()`), which `@mapped(...)` above `@read` turns into DTOs once the call's own transaction has ended. On SQLite a transaction holds the process-wide lock, and validating a DTO costs more than reading its row. A query method does not call other query modules and has no effects outside the database, because a call that loses a race is run again from the start. It passes execution options per statement, never to the connection: on SQLite every transaction shares one connection object. + +What backends spell differently comes from `dialect.py`: + +- `upsert()` builds an `INSERT` that updates the existing row instead when one with the same primary key exists. Every unique key of the table must be on the primary key's columns, since MySQL and MariaDB would update a row any unique key matched. A column's `onupdate` default does not apply to the update, so an upsert sets `updated_at` itself. +- `insert_ignore()` builds an `INSERT` that leaves an existing row with the same primary key as it is. It skips only that conflict: a NULL, a failed CHECK and a missing foreign key still raise, where SQLite's `INSERT OR IGNORE` would skip the first two and MySQL's `INSERT IGNORE` all three. Its row count does not tell whether a row was inserted: a skipped one counts 0 on SQLite and 1 on a server. The same rule on unique keys as for `upsert()` applies. +- `CaseInsensitiveLike(column, pattern)` matches as SQLite's `LIKE` always has, ignoring case, with `\` escaping `%` and `_`. On SQLite it is that `LIKE`, on a server `LOWER()` on both sides, which also folds letters beyond ASCII. `like_prefix(text)` and `like_contains(text)` build the patterns of the values that start with, or contain, a text taken as it is. +- `CaseInsensitiveOrder(column)` is an ORDER BY key that ignores case: `COLLATE NOCASE` on SQLite, which folds ASCII letters as SQLite's `LOWER()` does but without a function call per row, and `LOWER()` on a server. +- `fixed_limit(n)` is a LIMIT written into the statement, for a count that never changes, such as the 1 of a lookup. SQLite runs a statement whose LIMIT is a bound parameter 10–20 µs slower on every execution, with the same plan. A count the caller chooses, such as a page size, stays a bound parameter. +- `InBoundSet(column, bindparam(name))` is `column IN` a set of strings, or of integers for an integer column, bound as one parameter, which `bound_set(values)` builds: a JSON array, read with `json_each()` on SQLite and `JSON_TABLE()` on a server, strings in the tables' binary collation. A set of any size is one statement, where an `IN` list binds every value and is a statement per length. For sets held in process memory, such as the media names of active queue items, and for long lists a caller passes, such as the thousands of items whose embeddings the image map reads: there it is faster than `IN` lists of at most `IN_CHUNK` values, while at a hundred values it costs about 30 µs more on SQLite. SQLite cannot tell how many values a bound set holds, so beside another indexed condition it may prefer that condition's index: for a filter of a few values that should find their rows through its own index, such as the batch ids a cancel names, use an `IN` list. Keep it out of `UPDATE` and `DELETE`: MariaDB before 11.1 runs one with an `IN` subquery as a scan of the table that reads the set again for every row, so change rows by id in `IN` lists of at most `IN_CHUNK` values. Strings are at most 255 characters long. +- `KeysetAfter(first, second, first_value, second_value)` is the condition of the rows after a keyset position, `(first, second) > (first_value, second_value)`: a row value on SQLite, whose planner starts an index range at it, and spelled out on a server, since MariaDB does not start a range at a row value. For paging by a keyset, such as the intermediates cleanup's (creation time, name). +- `sql_true()` is TRUE as SQLite's partial indexes spell their predicates (`WHERE is_intermediate = TRUE`). SQLAlchemy renders `true()` as 1 there, and SQLite uses a partial index only for a query term that matches its predicate as written, so `is_intermediate = 1` leaves the index unused. +- `JsonString(column, path)` is the string at a member path of the JSON document in a text column, and NULL when the document is not valid JSON or the member is missing or holds anything else. SQLite's `json_extract()` alone would give 1 for a member holding true, where a server gives 'true', and raise on a document that is not JSON. +- `OrderedJoin(left, right, onclause)` is an inner join, and only that, whose left side SQLite plans as the outer loop: it renders as `CROSS JOIN ... ON` there, which SQLite's planner never reorders, and as a plain `JOIN` elsewhere. For a join that must stay proportional to its left side, such as a board's membership rather than every image an index offers first. + +## Schema + +The tables are SQLAlchemy `Table`s in `schema/`, one module per domain, all in `schema.metadata`. They are the one definition of the schema: query modules build their statements from them, and a server database is created from them. + +The metadata is the schema the SQLite migrations build. `tests/app/services/shared/database/test_schema_parity.py` migrates one database, creates another from the metadata and compares them column by column: declared type, NOT NULL, default, collation and generation expression. It compares every constraint and index the same way. A schema change is therefore a migration together with the same change to the metadata. + +Column types come from `types.py`. Each one renders the migrations' declared type on SQLite and its equivalent on a server: + +| Type | SQLite | MySQL / MariaDB | +|---|---|---| +| `Key(n)` | TEXT | VARCHAR(n) | +| `NoCaseKey(n)` | TEXT COLLATE NOCASE | VARCHAR(n), case-insensitive (`utf8mb4_0900_as_ci` / `utf8mb4_uca1400_nopad_as_ci`) | +| `LongText` | TEXT | LONGTEXT | +| `Timestamp` | DATETIME, or TEXT where a migration declared that | VARCHAR(32) | +| `BigInt` | INTEGER | BIGINT | +| `Real` | REAL | DOUBLE | +| `Blob` | BLOB | LONGBLOB | + +A server indexes the whole value of a primary key, a unique constraint or a foreign key, and one index holds at most 768 characters there. Such text, and text in indexes on identifiers, is a `Key`; a server refuses a longer value. All other text is `LongText`. An index on free text, such as a name, covers only a prefix of it on a server. `test_schema.py` holds the columns to this rule. + +A timestamp is UTC text that sorts and compares as text. Ported queries write it as `YYYY-MM-DD HH:MM:SS.fff` with `now_text()`. Rows also hold whole seconds and ISO 8601 text, at most 32 characters, written before the layer and its last migration. Columns made with `inserted_at()` and `updated_at()` get `now_text()` on insert; `updated_at()` columns also get it in an `update()` statement. An upsert must set `updated_at` in its own update values, because SQLAlchemy does not apply it to `ON CONFLICT DO UPDATE` or `ON DUPLICATE KEY UPDATE`. The SQL defaults of these columns exist only on SQLite, because the migrated schema has them. + +No database has triggers: queries set `updated_at` and the queue's `started_at`, `completed_at` and `session_revision` themselves. (The SQLite migrations created triggers for these; `2026_10_07_drop_sqlite_triggers` drops them.) + +Some DDL differs between backends: + +- A partial index is partial only on SQLite, because a server has none. A server's variant of the index leads with the condition's columns instead, so a query with the same condition still reads only the matching rows. The fonts' partial unique index is a plain unique index there, which the table's CHECK makes the same constraint. +- An index that a key or a longer index already covers exists only on SQLite, where a migration created it. +- MariaDB does not allow NOT NULL on a generated column. A CHECK constraint enforces it there instead, so a missing value raises `CheckViolation` there rather than `NotNullViolation`. +- Server tables are InnoDB, use the DYNAMIC row format and compare text byte for byte, whatever the server's defaults. +- `orphaned_projects_2026_08_06`, which holds projects a migration could not keep, always exists on a server. On SQLite it exists only where that migration had such projects. +- A MariaDB server must be reached through a `mariadb+pymysql://` URL. The dialect the URL names decides the collations and column definitions of the tables, and a `mysql+` URL to MariaDB is refused. + +## Migrations + +Migrations live in `shared/sqlite_migrator/migrations/`. The `Migrator` runs the ones a database has not had, in dependency order, and records each in `applied_migrations`. How to write one: [Database Migrations](/development/guides/sqlite-migrations/). + +- The migrations up to `PORTABLE_CUTOVER` (in `migration_loader.py`) run on SQLite only and get a cursor. They stay as written: they are the upgrade path of existing SQLite databases. `test_portable_migrations.py` fails when another one appears. +- Every later migration is a `PortableMigration`: it changes the schema with Alembic's operations (`context.op`), and creates tables with `context.create_table()`, which applies the schema's rules for each backend as the metadata does. The same commit makes the same change to the schema metadata. +- A portable migration is idempotent. On SQLite it runs in a transaction that a failure rolls back, DDL included. On MySQL and MariaDB every DDL statement commits as it runs, so a migration that failed halfway runs again from the start and finds part of its work done; it checks what is there before it changes it (`inspect(context.conn)`). Its id is recorded once it succeeds. +- On SQLite a portable migration runs with foreign keys off, so that rebuilding a table (Alembic's batch mode) does not cascade into the tables that reference it. The foreign keys are checked before it commits: a migration that deletes rows others reference deletes those itself. + +A MySQL or MariaDB database starts at the newest schema. On an empty database the migrator creates the tables from the metadata, and copies in the rows the migrations seed: the system account, the token secret, the system prompts and the lock rows. Those come from a reference database that it builds in memory with the migration chain, in a temporary root and with no settings from the environment or a config file, so the migrations' clean-ups of old files touch nothing real. The migrator's own records go in last. An interrupted creation therefore leaves a database the next start refuses, rather than one that looks complete; start again from an empty database. + +Only one process at a time serves a server database: at startup it takes the database's instance lock (`InstanceLock` in `session_lock.py`, a `GET_LOCK` on a connection of its own) before it migrates, and a second process refuses to start. A thread checks the lock every minute and takes it again should the server end its session. The migrator takes a lock of its own as well, verified before each migration, for whatever runs it without the instance lock. It is per server, so Galera clusters and other setups with several writing servers are not supported. The migrator makes no backups of a server database; they are the operator's. `test_portable_migrations.py` takes a server database as it was at the cutover (`server_schema_at_cutover/`) through every portable migration, and compares the result with a database created from the metadata. + +`copy.py` copies rows between databases, in foreign key order. The target computes generated columns itself. An AUTOINCREMENT table's ids continue after the highest one the source ever issued, which a SQLite source records in `sqlite_sequence` and a server in its tables' `AUTO_INCREMENT`, so an id deleted before the copy is not issued again after it. (SQLite's rowid tables reuse the highest id after a delete, with or without a copy.) `invoke-db-copy` uses it in both directions: an install's SQLite file into a new server database, and with `--to-sqlite` a server database back into a new SQLite file. + +## Transactions + +A call on `db.queries` runs in a transaction of its own: + +```py +user = db.queries.users.get(user_id) +``` + +Work that must succeed or fail as a whole runs in one transaction, which commits when the block exits and rolls back when it raises: + +```py +with db.queries.transaction() as q: + q.client_state.set(user_id, key, value) + q.media_references.replace(owner_kind="client_state", user_id=user_id, owner_id=key, references=references) +``` + +- A thread has at most one transaction open per database. Opening another one, including a call on `db.queries` inside a transaction, raises `NestedTransactionError`. Pass `q` to helpers instead. +- `transaction(read_only=True)` reads one consistent snapshot and refuses writes. +- A single call that loses a race (`ConflictError`: a deadlock) is run again from the start, three attempts in all. A `transaction()` block is not retried. Work that has no effects outside the database runs as `db.queries.run(work)` instead, which passes `work` the queries of one transaction and retries it the same way: on a server, transactions that write the same rows deadlock now and then. A lock held too long elsewhere (`LockTimeoutError`) is never retried, because the call has already waited for it. + +Errors are raised as backend-neutral types from `invokeai.app.services.shared.database.errors`, so code that catches them works on every backend. They are classified by the driver's error code; the one exception is a code SQLite uses for two causes, which its fixed message tells apart. The types: + +- `UniqueViolation`, `ForeignKeyViolation`, `CheckViolation` and `NotNullViolation`, all `IntegrityViolation`s. Inside a transaction they can be caught where they happen, to raise a domain error instead. +- `ConflictError` and `LockTimeoutError`, both `TransientDatabaseError`s. The transaction is rolled back and may be run again. + +A transaction does not continue past a failed statement, because backends disagree on what such a statement leaves behind. Further calls on it raise `TransactionFailedError`, and if the block is left without raising, it rolls back and raises `TransactionFailedError` instead of committing. + +## Backends + +| | SQLite | MySQL 8.4 / MariaDB 10.11 | +|---|---|---| +| Connections | One, behind a process-wide lock held for each transaction | A pool | +| Readers | `BEGIN`, one snapshot | `REPEATABLE READ`, one snapshot | +| Writers | `BEGIN IMMEDIATE`: the write lock is taken up front, so a transaction cannot fail to upgrade when another process wrote meanwhile | `READ COMMITTED`: no gap locks, and guards see current rows | +| Session | Foreign keys on, WAL, 5 s busy timeout | Strict SQL mode with `ONLY_FULL_GROUP_BY`, UTC, 10 s row-lock timeout | +| Text comparison | Byte for byte (`BINARY`) | Byte for byte, trailing spaces significant: tables and connections use `utf8mb4_0900_bin` (MySQL) / `utf8mb4_nopad_bin` (MariaDB), whatever the server's defaults | + +On SQLite the lock serialises all database work, so correctness can lean on it; on MySQL and MariaDB it cannot. Write invariants such as "only if still pending" into the statement itself (a conditional `UPDATE` whose row count is checked), not into a read followed by a write. + +An invariant that spans rows, such as "at least one active administrator remains", cannot be one statement. It gets a named lock: a member of `DatabaseLock` (`queries/locks.py`) with a row in `db_locks`, which a migration adds. Every transaction that changes what the check reads takes the lock first: + +```py +with db.queries.transaction() as q: + q.locks.acquire(DatabaseLock.ADMIN_ACCOUNTS) # SELECT ... FOR UPDATE on a server + if q.users.count_active_admins() <= 1: # deleting an active administrator + raise LastAdministratorError(LAST_ADMIN_DETAIL) + q.users.delete(admin_id) +``` + +A second transaction waits at `acquire()` until the first ends, and then reads what it committed. `acquire()` must be the first call of its transaction, taking all the transaction's locks at once, in name order: so they guard everything the transaction reads, and two transactions cannot deadlock on them. On SQLite the transaction already excludes every other, but the lock rows are read there too: a lock whose row is missing raises on every backend instead of locking nothing. + +`acquire(..., shared=True)` takes a lock that other shared holders take as well, and that waits only for, and holds off, an exclusive holder. `MEDIA_PROTECTION` works this way: every write that makes media protected from the intermediates cleanup (a project, workflow or client-state save with its media references, a browser lease, a cached output a session reuses) shares it, and the cleanup's final check and delete take it exclusively, so no such write lands between the check and the delete. The session queue makes media active too, by enqueueing, retrying or rewriting an item's session, and shares it for each of these. `acquire(*exclusive, also_shared=(...))` takes exclusive and shared locks in one call, still in name order: an enqueue takes `SESSION_QUEUE_ADMISSION` exclusively, so that two enqueues neither both fit into the queue's last free places nor both write one enqueue receipt, and `MEDIA_PROTECTION` shared. + +The queue's status changes need no lock: each is one conditional `UPDATE` (a claim applies only to a pending item, any other change only to an unfinished one), and its row count tells whether this call made it. Several workers may choose the same item to run; the claim gives it to one, and the others choose again. A bulk cancel or delete reads the ids it applies to with `SELECT ... FOR UPDATE` and changes exactly those, so its events name the rows it changed; one that spares the running workflow-call chains reads them after taking its row locks, so no item starts running, and no running item enqueues children, between that read and the change. A session save shares `MEDIA_PROTECTION` too, so on a server it waits while the intermediates cleanup checks and deletes. + +Image moves take no database lock. Their jobs are created, worked and committed under the move service's process lock (`image_mutation_lock`), which image deletes take too, and one InvokeAI process serves a database (a server database's instance lock refuses a second); so "at most one unfinished move job" is a read followed by an insert. A job's commit repoints each moved image only while its record still names the old subfolder, and checks every moved item in the same transaction, which rolls back if one is not where the move put it. + +The font catalog takes none either, for the same reason. An upload checks for the same content and for the quota, writes its file and inserts its row under the font service's managed-storage lock, which the sweep for orphaned files takes too; the file is written outside any transaction. Directory scans run one at a time under a lock of their own, so two cannot insert the same new font. + +A decision about one row needs no named lock: the row is the lock. The transaction locks the row before it reads what it decides on, with a query method marked `@locking` that selects it `FOR UPDATE` (`q.boards.lock(board_id)`, `q.projects.lock(user_id, project_id)`). Every transaction that changes what the decision reads writes or locks the same row, so it waits until the decision has committed, and the decision sees what such a transaction committed before. A project save works this way, so its revision and compatibility checks apply to the row it writes, and so does a project claiming a board: a board made public, deleted or claimed by another project meanwhile is either seen by the claim or finds the board claimed. Row locks, too, are held only inside a `transaction()`; called on its own, a lock method raises. + +Rows are locked in one order, a project's before its board's, so two transactions that lock both cannot deadlock on them. A single statement that locks a board and then reads its project, like the boards service's guarded update and delete, cannot keep that order; the rare deadlock it causes with a project write is retried by `run()`. + +## Testing + +Tests that need a database use a fixture from `tests/fixtures/database.py`, and are marked `uses_database` for it: + +- `database`: a database at the newest schema, holding the rows a new install starts with (the system account, the JWT secret, the lock rows). For tests of ported services. +- `empty_database`: a database without tables, for tests of the layer itself and of migrations. + +On a server the ids a table generates go on counting from test to test: assert on the ids a test got back, not on literal ones. + +`capture_statements(database)` collects the SQL and parameters a block of code sends, to assert on how many statements it takes or on their shape. `explain_query_plan(database, sql, parameters)` gives SQLite's plan for one of them, for tests (marked `sqlite_only`) that a query keeps using its index. + +Both are in-memory SQLite databases by default. When `INVOKEAI_TEST_DB_URL` names a MySQL or MariaDB server, they are instead schemas on that server of their own per test run and worker; the `database` schema is migrated once per run and reset to those rows before each test. Mark a test that depends on SQLite internals (query plans, PRAGMAs, the single connection) `sqlite_only`. + +To run the tests against a server, start one and point the variable at it. The account needs the privilege to create and drop databases: + +```sh +docker run -d --name invokeai-test-mysql -e MYSQL_ROOT_PASSWORD=invokeai -p 127.0.0.1:33084:3306 mysql:8.4.11 +docker run -d --name invokeai-test-mariadb -e MARIADB_ROOT_PASSWORD=invokeai -p 127.0.0.1:33011:3306 mariadb:10.11.19 + +# Your usual sync with the `mysql` extra added; keep your accelerator extra (here `cuda`). +uv sync --extra cuda --extra test --extra mysql +INVOKEAI_TEST_DB_URL=mysql+pymysql://root:invokeai@127.0.0.1:33084 \ + uv run --no-sync pytest -m "uses_database and not sqlite_only and not slow" +INVOKEAI_TEST_DB_URL=mariadb+pymysql://root:invokeai@127.0.0.1:33011 \ + uv run --no-sync pytest -m "uses_database and not sqlite_only and not slow" +``` + +CI runs the same selection against both servers in the `external-database` job of the python tests workflow. + +`scripts/benchmark_database.py` times representative operations through the services on a seeded SQLite file. Run it for the base and the head of a change and compare the results with `--compare`. diff --git a/docs/src/content/docs/development/Guides/sqlite-migrations.mdx b/docs/src/content/docs/development/Guides/sqlite-migrations.mdx index edc300d630a..4e5c671cc90 100644 --- a/docs/src/content/docs/development/Guides/sqlite-migrations.mdx +++ b/docs/src/content/docs/development/Guides/sqlite-migrations.mdx @@ -1,96 +1,85 @@ --- -title: SQLite Database Migrations -lastUpdated: 2026-06-30 +title: Database Migrations +lastUpdated: 2026-10-04 --- -InvokeAI uses a custom SQLite migrator for the main application database. Migrations live in `invokeai/app/services/shared/sqlite_migrator/migrations/` and are discovered automatically when the database is initialized. +InvokeAI migrates its database with its own migrator. Migrations live in `invokeai/app/services/shared/sqlite_migrator/migrations/` and are discovered automatically when the database is initialized. Use a migration when a change modifies persisted database schema or persisted database data in a way that existing installs must receive on startup. +New migrations are portable: they run on every database backend, SQLite, MySQL and MariaDB. The migrations up to `PORTABLE_CUTOVER` (in `migration_loader.py`) run on SQLite only; they stay as written, and no new one of that kind is added. See [Database Layer](/development/guides/database-layer/) for the schema metadata a migration changes along with the database. + ## Naming New migration modules should use a date-stamped descriptive name: ```text -invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_06_30_add_example_table.py +invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_05_add_bananas.py ``` Date-stamped migration modules must expose a `build_migration()` builder: ```py -def build_migration() -> Migration: +def build_migration() -> PortableMigration: ... ``` The date prefix should be the date the migration is authored or merged, in `YYYY_MM_DD` format. Add a short snake_case description after the date. -Legacy numeric migration modules are still supported: - -```text -invokeai/app/services/shared/sqlite_migrator/migrations/migration_33.py -``` - -Numeric migration modules must expose the matching legacy builder: - -```py -def build_migration_33() -> Migration: - ... -``` - -The loader discovers both naming styles. Numeric modules are imported in numeric order for legacy compatibility. Date-stamped modules are imported in lexical order, but dependency metadata drives execution order. Non-matching migration module names, missing builders, unknown builder dependencies, and builders that do not return `Migration` fail at startup. +Legacy numeric migration modules (`migration_33.py`, with `build_migration_33()`) are still discovered, in numeric order. Date-stamped modules are imported in lexical order, but dependency metadata drives execution order. Non-matching migration module names, missing builders, unknown builder dependencies, and builders that do not return a migration fail at startup. ## IDs and Dependencies Every migration has a stable `id`. For new date-stamped migrations, use the module name without the `migration_` prefix: ```py -id="2026_06_30_add_example_table" -``` - -The loader enforces this match. If `migration_2026_06_30_add_example_table.py` returns a different ID, startup fails instead of persisting a typo that could later cause the migration to run again. - -Existing numeric migrations use IDs matching their file number: - -```py -id="migration_33" +id="2026_10_05_add_bananas" ``` -Legacy numeric migrations that define `from_version` and `to_version` may only depend on other legacy numeric migrations. New migrations should usually be date-stamped graph-only migrations, and should not define legacy versions unless a specific legacy compatibility requirement exists. - -For a migration that only needs an older schema point, depend on the migration that introduced the required schema or data: +The loader enforces this match. If `migration_2026_10_05_add_bananas.py` returns a different ID, startup fails instead of persisting a typo that could later cause the migration to run again. -```py -depends_on="migration_27" -``` +New production migrations always set `depends_on`, to the migration that introduced the schema or data they need, or else to the newest migration. Only a true root migration uses `depends_on=None`: a migration with no dependency is considered independently runnable, so omitting `depends_on` can let it run before the schema it expects exists. -New production migrations should always set `depends_on`. Only a true root migration should use `depends_on=None`. A migration with no dependency is considered independently runnable, so omitting `depends_on` for a normal schema or data change can allow it to run before the schema it expects exists. - -Dependencies drive execution order. If two migrations both depend on `migration_27`, either may run after `migration_27` unless one explicitly depends on the other. Do not rely on filename or lexical ordering to express a real dependency. - -Single-parent dependencies are currently supported. If a migration truly requires two independent branches, add a real dependency between those branches only when that ordering is semantically correct. Otherwise, update the migrator design and tests to support multiple dependencies. +Dependencies drive execution order. If two migrations both depend on the same migration, either may run first unless one explicitly depends on the other. Do not rely on filename or lexical ordering to express a real dependency. Single-parent dependencies are supported. ## Template ```py -import sqlite3 - -from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration - - -class AddExampleTableCallback: - def __call__(self, cursor: sqlite3.Cursor) -> None: - cursor.execute("ALTER TABLE images ADD COLUMN example TEXT;") - - -def build_migration() -> Migration: - return Migration( - id="2026_06_30_add_example_table", - depends_on="migration_32", - callback=AddExampleTableCallback(), +from sqlalchemy import Column, inspect + +from invokeai.app.services.shared.database.schema.metadata import default, inserted_at +from invokeai.app.services.shared.database.types import BigInt, Key +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import ( + PortableMigration, + PortableMigrationContext, +) + + +def _add_bananas(context: PortableMigrationContext) -> None: + if not inspect(context.conn).has_table("bananas"): + context.create_table( + "bananas", + Column("banana_id", Key(), primary_key=True), + Column("ripeness", BigInt(), nullable=False, server_default=default(0)), + inserted_at(), + ) + + +def build_migration() -> PortableMigration: + return PortableMigration( + id="2026_10_05_add_bananas", + depends_on="2026_10_04_add_db_locks", + callback=_add_bananas, ) ``` -Keep the callback self-contained. It receives an open SQLite cursor and must not commit or roll back the transaction. The migrator commits the migration and records it as applied only after the callback succeeds. If the callback raises, the migration transaction is rolled back and the migration is not recorded. +The callback changes the schema with Alembic's operations (`context.op.add_column`, `context.op.create_index`, `context.op.batch_alter_table`, ...) and reads and writes rows through `context.conn`. It does not commit or roll back. The migrator records the migration as applied in the same transaction, once the callback has succeeded. + +- **Create tables with `context.create_table()`**, not `op.create_table`. It creates a table exactly as the schema metadata does: InnoDB, the binary collation and DYNAMIC rows on a server whatever its defaults, and each backend's rules for defaults, generated columns and indexes. A foreign key names a table in `context.metadata`: one created earlier in the migration, or an existing one loaded with `Table(name, context.metadata, autoload_with=context.conn)`. +- **Make it idempotent.** On SQLite the migration runs in a transaction that a failure rolls back, DDL included. On MySQL and MariaDB every DDL statement commits as it runs, so a migration that failed halfway runs again from the start at the next startup and finds part of its work done. Check what is there before changing it (`inspect(context.conn)`). +- **Mind foreign keys on SQLite, and create no triggers.** There a portable migration runs with foreign keys off, so that rebuilding a table with `op.batch_alter_table` (SQLite's only way to alter most columns) does not delete, by cascade, the rows of the tables that reference it. They are checked before the migration commits: a migration that leaves rows referencing nothing fails, so one that deletes rows others reference deletes those itself. (Rows that referenced nothing before it ran do not count.) No database has triggers (`2026_10_07_drop_sqlite_triggers` dropped the last): the application sets what a trigger would, on every backend, and `test_schema_parity.py` fails when a migration creates one. +- **Spell out what Alembic does not know.** `op.add_column` applies none of the metadata's rules for a backend: a column that the metadata gives a SQLite-only default (`inserted_at()`, `updated_at()`) or a generated `NOT NULL` column needs its DDL written per backend (`context.conn.dialect.name`). +- **Change the schema metadata in the same commit.** New server databases are created from the metadata, so it must describe what the migration makes. `test_schema_parity.py` holds it to the migrated SQLite schema, and `test_portable_migrations.py` takes a server database from the cutover through every portable migration and compares it with one created from the metadata. ## Builder Dependencies @@ -105,35 +94,35 @@ Supported dependency names: Example: ```py -def build_migration(app_config: InvokeAIAppConfig, logger: Logger) -> Migration: - return Migration( - id="2026_06_30_normalize_model_paths", - depends_on="2026_06_30_add_example_table", +def build_migration(app_config: InvokeAIAppConfig, logger: Logger) -> PortableMigration: + return PortableMigration( + id="2026_10_05_normalize_model_paths", + depends_on="2026_10_05_add_bananas", callback=NormalizeModelPathsCallback(app_config=app_config, logger=logger), ) ``` Do not use `*args`, `**kwargs`, or positional-only parameters in migration builders. +A new MySQL or MariaDB database is not migrated: it is created at the newest schema, with the rows the migrations seed copied from a SQLite database that the migration chain builds in memory, in a temporary root. A migration's effects outside the database (files it moves or deletes) therefore happen on SQLite installs and on existing server databases, never when a server database is created. + ## Registration -Do not manually import or register migrations in `sqlite_util.py`. +Do not manually import or register migrations in `services/shared/database/startup.py`. -Database initialization builds a `MigrationBuildContext`, discovers migration modules, calls their builders, and registers the resulting `Migration` objects with `SqliteMigrator`. +Database initialization builds a `MigrationBuildContext`, discovers migration modules, calls their builders, and registers the resulting migrations with the `Migrator`. Manual registration is still available for tests: ```py -migrator.register_migration(Migration(...)) +migrator.register_migration(PortableMigration(...)) ``` ## Applied State and Compatibility -The migrator records stable migration IDs in the `applied_migrations` table. Legacy numeric versions are still written to the existing `migrations` table for migrations that define `to_version`. - -New date-stamped migrations should not define `from_version` or `to_version` unless there is a specific legacy compatibility reason. Use `id` and `depends_on` to define execution. +The migrator records stable migration IDs in the `applied_migrations` table. Legacy numeric versions are still written to the existing `migrations` table for the legacy migrations that define `to_version`. -Existing databases are bootstrapped from legacy numeric rows: +Existing SQLite databases are bootstrapped from legacy numeric rows: - legacy version `1` maps to `migration_1` - legacy version `2` maps to `migration_2` @@ -143,11 +132,11 @@ If a database contains an applied migration ID or legacy numeric version unknown For legacy migrations, the two metadata tables must agree. For example, `applied_migrations` may only record `migration_2` with `legacy_version = 2` if the legacy `migrations` table also contains version `2`. If these records are inconsistent, startup fails before user migration callbacks run. -When opening a file-backed legacy database for the first time after this migration system change, the migrator may need to create and populate `applied_migrations` even if no user migration callbacks need to run. This metadata bootstrap is backed up before the new metadata table is written. +When opening a file-backed legacy SQLite database for the first time after this migration system change, the migrator may need to create and populate `applied_migrations` even if no user migration callbacks need to run. This metadata bootstrap is backed up before the new metadata table is written. The migrator backs up a SQLite database file before it migrates it; it makes no backups of a server database. ## Tests -Add focused tests for each migration callback under: +Add focused tests for each migration under: ```text tests/app/services/shared/sqlite_migrator/migrations/ @@ -155,16 +144,15 @@ tests/app/services/shared/sqlite_migrator/migrations/ At minimum, cover: -- the schema or data change performed by the callback -- idempotent or already-migrated behavior when relevant +- the schema or data change performed by the migration +- running it again after it failed halfway, and already-migrated behavior where relevant - missing optional source tables or columns when the migration is expected to tolerate them - the builder's `id` and `depends_on` -- legacy `from_version` and `to_version`, if the migration defines them -Also run the migrator tests: +Tests that use the `empty_database` fixture (no tables) or the `database` fixture (the newest schema) run on SQLite and, with `INVOKEAI_TEST_DB_URL` set, on a MySQL or MariaDB server (see [Database Layer](/development/guides/database-layer/#testing)). Also run the migrator tests: ```bash pytest tests/app/services/shared/sqlite_migrator tests/test_sqlite_migrator.py ``` -These tests use in-memory databases or pytest temporary directories. They should not touch a real InvokeAI database. +These tests use in-memory databases, pytest temporary directories or a schema of their own on the test server. They should not touch a real InvokeAI database. diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index 2dcfa837812..60edab3e0bb 100644 --- a/docs/src/generated/settings.json +++ b/docs/src/generated/settings.json @@ -213,6 +213,17 @@ "type": "", "validation": {} }, + { + "category": "PATHS", + "default": null, + "description": "URL of a MySQL 8.4+ or MariaDB 10.11+ database to use instead of the SQLite database in `db_dir`, e.g. `mariadb+pymysql://invokeai:password@db.example/invokeai`. Needs the `mysql` extra. One InvokeAI process uses a database at a time. Read at startup only.", + "env_var": "INVOKEAI_DB_URL", + "literal_values": [], + "name": "db_url", + "required": false, + "type": "typing.Optional[str]", + "validation": {} + }, { "category": "PATHS", "default": "full", diff --git a/invokeai/app/AGENTS.md b/invokeai/app/AGENTS.md index 635b6896828..bd25fa32eb9 100644 --- a/invokeai/app/AGENTS.md +++ b/invokeai/app/AGENTS.md @@ -11,7 +11,11 @@ ## Persistence and lifecycle - Preserve atomicity, revision checks, idempotency, and explicit conflicts. Inspect migrations/old-record readers for stored-data changes; add needed migration/compatibility tests. -- Read `services/shared/sqlite/sqlite_database.py` and `services/project_records/project_records_sqlite.py`: nested transactions can prematurely commit outer operations. Use one transaction owner for multi-record writes, passing its cursor where supported. +- SQL lives only in `services/shared/database/`, reached through `Database.queries`; `tests/app/test_database_access_guard.py` enforces this. Query modules build their statements once, at module level with `bindparam()` (a statement built per call costs several times its execution), read rows by position, and build DTOs with `@mapped(...)` outside the transaction. A shape-dependent statement comes from a cached builder whose shapes a client cannot multiply (normalize its lists and flags, round its list lengths; see the guide). +- Use one transaction owner for multi-record writes: `with db.queries.transaction() as q:`, passing `q` to helpers. A second transaction on the same thread raises `NestedTransactionError`. See `services/shared/database/database.py`. +- SQLite serialises everything behind one lock; MySQL/MariaDB do not. Put an invariant like "only if still pending" into the statement (a conditional `UPDATE` with a checked row count), not into a read followed by a write. An invariant that spans rows (the last active administrator) gets a `DatabaseLock`, which every transaction changing what it reads acquires as its first call (`q.locks.acquire(...)`); its row is added by a migration. A decision about one row (may this board be claimed?) locks that row before reading what it decides on (a `@locking` query method, `SELECT ... FOR UPDATE`); every transaction changing what it reads writes or locks the same row, and rows are locked in one order (a project's before its board's). Run a unit whose work has no effects outside the database with `db.queries.run(work)`, which retries it after a deadlock. +- Tables are defined in `services/shared/database/schema/`, which must match the migrated SQLite schema exactly (`test_schema_parity.py`). Change the schema with a migration and the same change to the metadata, in one change. No database has triggers: queries set `updated_at` (also in upserts) and the queue's `started_at`, `completed_at`, `session_revision`. +- New migrations are `PortableMigration`s (Alembic operations; tables via `context.create_table()`) and idempotent: MySQL/MariaDB commit DDL as it runs, so a migration that failed halfway runs again. On SQLite they run with foreign keys off and are checked before commit. Migrations up to `PORTABLE_CUTOVER` stay SQLite-only and unchanged. See the Database Migrations guide. - Preserve project/account isolation, document bounds, queue admission/receipts, and lost-response retry safety. Never silently discard recovery data or resolve conflicts with last-write-wins. - Bound concurrency; support cancellation, timeouts/retries, and file/process/listener/task cleanup. Keep CPU/blocking I/O off the async server loop. - Preserve file/path validation, remote-fetch protections, and sanitized errors. Never modify real user models, databases, outputs, or credentials during tests/debugging. @@ -19,7 +23,7 @@ ## Efficiency and verification - Inspect queries for repeated lookups, full scans/materialization, missing bounds, and serialization. Measure representative data/query counts before adding caches or indexes. -- Use real temporary SQLite for transaction, migration, revision, and ownership tests; cover affected malformed input, cross-user access, cancellation, repeated requests, and partial failures. +- Use real temporary SQLite for transaction, migration, revision, and ownership tests; cover affected malformed input, cross-user access, cancellation, repeated requests, and partial failures. Tests on the database fixtures also run against MySQL and MariaDB when `INVOKEAI_TEST_DB_URL` is set (see `tests/fixtures/database.py`). - Read `tests/AGENTS.md`; use `tests/app/` or the existing owning suite. Run focused pytest and root Ruff; broaden to relevant service/router/invocation suites for milestones. - For graph execution, consult `services/shared/README.md`; test scheduling and invocation lifecycles through owning interfaces. - Report schema generation, migration validation, and unavailable integration/hardware checks. diff --git a/invokeai/app/api/dependencies.py b/invokeai/app/api/dependencies.py index faa3ec52432..f11b206427e 100644 --- a/invokeai/app/api/dependencies.py +++ b/invokeai/app/api/dependencies.py @@ -1,17 +1,18 @@ import asyncio from logging import Logger +from typing import Optional import torch from invokeai.app.services.app_settings import AppSettingsService from invokeai.app.services.auth.token_service import set_jwt_secret -from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage +from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage from invokeai.app.services.board_images.board_images_default import BoardImagesService -from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage -from invokeai.app.services.board_video_records.board_video_records_sqlite import SqliteBoardVideoRecordStorage +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage +from invokeai.app.services.board_video_records.board_video_records_default import BoardVideoRecordStorage from invokeai.app.services.boards.boards_default import BoardService from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService -from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ClientStatePersistenceSqlite +from invokeai.app.services.client_state_persistence.client_state_persistence_default import ClientStatePersistence from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.download.download_default import DownloadQueueService from invokeai.app.services.events.events_fastapievents import FastAPIEventService @@ -25,16 +26,16 @@ ) from invokeai.app.services.external_generation.startup import sync_configured_external_starter_models from invokeai.app.services.fonts.fonts_default import FontService -from invokeai.app.services.gallery.gallery_default import SqliteGalleryService +from invokeai.app.services.gallery.gallery_default import GalleryService from invokeai.app.services.gallery_maintenance.gallery_maintenance_default import GalleryMaintenanceService from invokeai.app.services.image_files.image_files_disk import DiskImageFileStorage from invokeai.app.services.image_index.image_index_default import ImageIndexService, warm_up_attention -from invokeai.app.services.image_index.image_index_records_sqlite import ImageIndexRecordsSqlite +from invokeai.app.services.image_index.image_index_records_default import ImageIndexRecords from invokeai.app.services.image_moves.image_moves_default import ImageMoveService -from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage from invokeai.app.services.images.images_default import ImageService from invokeai.app.services.intermediates.intermediates_default import IntermediatesService -from invokeai.app.services.intermediates.intermediates_records_sqlite import IntermediatesRecordsSqlite +from invokeai.app.services.intermediates.intermediates_records_default import IntermediatesRecords from invokeai.app.services.invocation_cache.invocation_cache_memory import MemoryInvocationCache from invokeai.app.services.invocation_services import InvocationServices from invokeai.app.services.invocation_stats.invocation_stats_default import InvocationStatsService @@ -42,31 +43,32 @@ from invokeai.app.services.model_images.model_images_default import ModelImageFileStorageDisk from invokeai.app.services.model_manager.model_manager_default import ModelManagerService from invokeai.app.services.model_records.model_records_sql import ModelRecordServiceSQL -from invokeai.app.services.model_relationship_records.model_relationship_records_sqlite import ( - SqliteModelRelationshipRecordStorage, +from invokeai.app.services.model_relationship_records.model_relationship_records_default import ( + ModelRelationshipRecordStorage, ) from invokeai.app.services.model_relationships.model_relationships_default import ModelRelationshipsService from invokeai.app.services.names.names_default import SimpleNameService from invokeai.app.services.object_serializer.object_serializer_disk import ObjectSerializerDisk from invokeai.app.services.object_serializer.object_serializer_forward_cache import ObjectSerializerForwardCache from invokeai.app.services.progress_previews.progress_previews_default import MemoryProgressPreviews -from invokeai.app.services.project_records.project_records_sqlite import ProjectRecordsSqlite +from invokeai.app.services.project_records.project_records_default import ProjectRecordsStorage from invokeai.app.services.session_processor.session_processor_default import ( DefaultSessionProcessor, DefaultSessionRunner, ) -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue -from invokeai.app.services.shared.sqlite.sqlite_util import init_db +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.startup import init_database from invokeai.app.services.style_preset_images.style_preset_images_disk import StylePresetImageFileStorageDisk -from invokeai.app.services.style_preset_records.style_preset_records_sqlite import SqliteStylePresetRecordsStorage -from invokeai.app.services.system_prompt_records.system_prompt_records_sqlite import SqliteSystemPromptRecordsStorage +from invokeai.app.services.style_preset_records.style_preset_records_default import StylePresetRecordsStorage +from invokeai.app.services.system_prompt_records.system_prompt_records_default import SystemPromptRecordsStorage from invokeai.app.services.urls.urls_default import LocalUrlService from invokeai.app.services.users.users_default import UserService from invokeai.app.services.video_files.video_files_disk import DiskVideoFileStorage -from invokeai.app.services.video_records.video_records_sqlite import SqliteVideoRecordStorage +from invokeai.app.services.video_records.video_records_default import VideoRecordStorage from invokeai.app.services.videos.videos_default import VideoService -from invokeai.app.services.wildcard_records.wildcard_records_sqlite import SqliteWildcardRecordsStorage -from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage +from invokeai.app.services.wildcard_records.wildcard_records_default import WildcardRecordsStorage +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage from invokeai.app.services.workflow_thumbnails.workflow_thumbnails_disk import WorkflowThumbnailFileStorageDisk from invokeai.backend.architectures import conditioning_safe_globals from invokeai.backend.architectures import validate as validate_architectures @@ -98,6 +100,7 @@ class ApiDependencies: """Contains and initializes all dependencies for the API""" invoker: Invoker + database: Optional[Database] = None @staticmethod def initialize( @@ -129,10 +132,10 @@ def initialize( style_presets_folder = config.style_presets_path workflow_thumbnails_folder = config.workflow_thumbnails_path - db = init_db(config=config, logger=logger, image_files=image_files) + database = init_database(config=config, logger=logger, image_files=image_files) # Initialize JWT secret from database - app_settings = AppSettingsService(db=db) + app_settings = AppSettingsService(database) jwt_secret = app_settings.get_jwt_secret() set_jwt_secret(jwt_secret) logger.info("JWT secret loaded from database") @@ -140,19 +143,19 @@ def initialize( configuration = config logger = logger - board_image_records = SqliteBoardImageRecordStorage(db=db) + board_image_records = BoardImageRecordStorage(database) board_images = BoardImagesService() - board_records = SqliteBoardRecordStorage(db=db) + board_records = BoardRecordStorage(database) boards = BoardService() events = FastAPIEventService(event_handler_id, loop=loop) bulk_download = BulkDownloadService() - image_records = SqliteImageRecordStorage(db=db) - image_moves = ImageMoveService(db=db, image_files=image_files, config=configuration, logger=logger) + image_records = ImageRecordStorage(database) + image_moves = ImageMoveService(database, image_files=image_files, config=configuration, logger=logger) images = ImageService() - video_records = SqliteVideoRecordStorage(db=db) + video_records = VideoRecordStorage(database) videos = VideoService() - board_video_records = SqliteBoardVideoRecordStorage(db=db) - gallery = SqliteGalleryService(db=db) + board_video_records = BoardVideoRecordStorage(database) + gallery = GalleryService(database) gallery_maintenance = GalleryMaintenanceService() invocation_cache = MemoryInvocationCache(max_cache_size=config.node_cache_size) tensors = ObjectSerializerForwardCache( @@ -175,7 +178,7 @@ def initialize( ), ) download_queue_service = DownloadQueueService(app_config=configuration, event_bus=events) - model_record_service = ModelRecordServiceSQL(db=db, logger=logger) + model_record_service = ModelRecordServiceSQL(database, logger=logger) model_manager = ModelManagerService.build_model_manager( app_config=configuration, model_record_service=model_record_service, @@ -195,26 +198,26 @@ def initialize( ) model_images_service = ModelImageFileStorageDisk(model_images_folder / "model_images") model_relationships = ModelRelationshipsService() - model_relationship_records = SqliteModelRelationshipRecordStorage(db=db) + model_relationship_records = ModelRelationshipRecordStorage(database) names = SimpleNameService() performance_statistics = InvocationStatsService() session_processor = DefaultSessionProcessor(session_runner=DefaultSessionRunner()) - session_queue = SqliteSessionQueue(db=db) + session_queue = SessionQueue(database) urls = LocalUrlService() - workflow_records = SqliteWorkflowRecordsStorage(db=db) - style_preset_records = SqliteStylePresetRecordsStorage(db=db) - wildcard_records = SqliteWildcardRecordsStorage(db=db) + workflow_records = WorkflowRecordsStorage(database) + style_preset_records = StylePresetRecordsStorage(database) + wildcard_records = WildcardRecordsStorage(database) style_preset_image_files = StylePresetImageFileStorageDisk(style_presets_folder / "images") - system_prompt_records = SqliteSystemPromptRecordsStorage(db=db) + system_prompt_records = SystemPromptRecordsStorage(database) workflow_thumbnails = WorkflowThumbnailFileStorageDisk(workflow_thumbnails_folder) - client_state_persistence = ClientStatePersistenceSqlite(db=db) - project_records = ProjectRecordsSqlite(db=db) - users = UserService(db=db) - image_index_records = ImageIndexRecordsSqlite(db=db) + client_state_persistence = ClientStatePersistence(database) + project_records = ProjectRecordsStorage(database) + users = UserService(database) + image_index_records = ImageIndexRecords(database) image_index = ImageIndexService() - intermediates = IntermediatesService(records=IntermediatesRecordsSqlite(db=db), logger=logger) + intermediates = IntermediatesService(records=IntermediatesRecords(database), logger=logger) fonts = FontService( - db=db, + database, fonts_dir=configuration.fonts_path, storage_dir=configuration.fonts_storage_path, logger=logger, @@ -229,7 +232,7 @@ def initialize( boards=boards, bulk_download=bulk_download, configuration=configuration, - database=db, + database=database, events=events, image_files=image_files, image_moves=image_moves, @@ -289,8 +292,13 @@ def initialize( model_manager=model_manager, logger=logger, ) + ApiDependencies.database = database @staticmethod def shutdown() -> None: if ApiDependencies.invoker: ApiDependencies.invoker.stop() + if ApiDependencies.database is not None: + # Releases a server database's instance lock, and checkpoints a SQLite database's write-ahead log. + ApiDependencies.database.dispose() + ApiDependencies.database = None diff --git a/invokeai/app/api/routers/app_info.py b/invokeai/app/api/routers/app_info.py index 6e8ae9b165c..c4cdfc51208 100644 --- a/invokeai/app/api/routers/app_info.py +++ b/invokeai/app/api/routers/app_info.py @@ -27,6 +27,7 @@ from invokeai.app.services.external_generation.startup import sync_configured_external_starter_models from invokeai.app.services.invocation_cache.invocation_cache_common import InvocationCacheStatus from invokeai.app.services.model_records.model_records_base import UnknownModelException +from invokeai.app.services.shared.database.startup import redacted_database_url from invokeai.backend.image_util.infill_methods.patchmatch import PatchMatch from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType from invokeai.backend.util.devices import TorchDevice @@ -238,6 +239,9 @@ def _redact_config_secrets(config: InvokeAIAppConfig) -> InvokeAIAppConfig: if config.download_proxy and "@" in config.download_proxy: updates["download_proxy"] = REDACTED_SECRET + if config.db_url: + updates["db_url"] = redacted_database_url(config.db_url) + return config.model_copy(update=updates) if updates else config diff --git a/invokeai/app/api/routers/board_images.py b/invokeai/app/api/routers/board_images.py index 4b0aca01ac9..1d935992c19 100644 --- a/invokeai/app/api/routers/board_images.py +++ b/invokeai/app/api/routers/board_images.py @@ -273,9 +273,9 @@ def add_images_to_board( # # Except that a name deleted between the ownership check and the insert lands # here too, and not as something recognizable: board_images.image_name is a - # foreign key onto images.image_name, so the INSERT fails with a bare - # sqlite3.IntegrityError. Nothing in the exception says "gone", so the record - # is probed instead — only on this path, so the happy path pays nothing. + # foreign key onto images.image_name, so the INSERT fails with a + # ForeignKeyViolation that does not say which key: the board or the image. So the + # record is probed instead — only on this path, so the happy path pays nothing. if not _image_record_exists(image_name): continue failed_images.add(image_name) diff --git a/invokeai/app/api/routers/boards.py b/invokeai/app/api/routers/boards.py index 6b0683b5a47..e072a0a874d 100644 --- a/invokeai/app/api/routers/boards.py +++ b/invokeai/app/api/routers/boards.py @@ -15,8 +15,7 @@ ) from invokeai.app.services.boards.boards_common import BoardDTO from invokeai.app.services.image_records.image_records_common import ImageCategory -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import MAX_PAGE_SIZE, OffsetPaginatedResults, SQLiteDirection boards_router = APIRouter(prefix="/v1/boards", tags=["boards"]) @@ -247,8 +246,8 @@ def list_boards( order_by: BoardRecordOrderBy = Query(default=BoardRecordOrderBy.CreatedAt, description="The attribute to order by"), direction: SQLiteDirection = Query(default=SQLiteDirection.Descending, description="The direction to order by"), all: Optional[bool] = Query(default=None, description="Whether to list all boards"), - offset: Optional[int] = Query(default=None, description="The page offset"), - limit: Optional[int] = Query(default=None, description="The number of boards per page"), + offset: Optional[int] = Query(default=None, ge=0, description="The page offset"), + limit: Optional[int] = Query(default=None, ge=0, le=MAX_PAGE_SIZE, description="The number of boards per page"), include_archived: bool = Query(default=False, description="Whether or not to include archived boards in list"), ) -> Union[OffsetPaginatedResults[BoardDTO], list[BoardDTO]]: """Gets a list of boards for the current user, including shared boards. Admin users see all boards.""" diff --git a/invokeai/app/api/routers/client_state.py b/invokeai/app/api/routers/client_state.py index 07790c7182a..dce8573f0f1 100644 --- a/invokeai/app/api/routers/client_state.py +++ b/invokeai/app/api/routers/client_state.py @@ -34,7 +34,7 @@ def get_client_state_by_key( def set_client_state( current_user: CurrentUserOrDefault, queue_id: str = Path(description="The queue id (ignored, kept for backwards compatibility)"), - key: str = Query(..., description="Key to set"), + key: str = Query(..., description="Key to set", max_length=255), value: str = Body(..., description="Stringified value to set"), ) -> str: """Sets the client state for the current user (or system user if not authenticated)""" diff --git a/invokeai/app/api/routers/gallery.py b/invokeai/app/api/routers/gallery.py index 06e0731f888..2f30d1eef2d 100644 --- a/invokeai/app/api/routers/gallery.py +++ b/invokeai/app/api/routers/gallery.py @@ -15,8 +15,7 @@ GalleryItemNamesResult, ) from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin -from invokeai.app.services.shared.pagination import MAX_PAGE_SIZE, OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import MAX_PAGE_SIZE, OffsetPaginatedResults, SQLiteDirection gallery_router = APIRouter(prefix="/v1/gallery", tags=["gallery"]) diff --git a/invokeai/app/api/routers/image_map.py b/invokeai/app/api/routers/image_map.py index c0af8f48684..6f205a285a8 100644 --- a/invokeai/app/api/routers/image_map.py +++ b/invokeai/app/api/routers/image_map.py @@ -434,7 +434,7 @@ async def get_image_map_points( user_id, is_admin = _scope(current_user) kinds = _served_kinds(include_videos) - # Both reads are SQLite work under the process-wide database lock, and on a large gallery + # Both reads are database work (on SQLite under the process-wide lock), and on a large gallery # the accessible listing is the expensive one — so they run off the event loop, like the # clustering below. current_items, record = await asyncio.to_thread( @@ -1093,7 +1093,7 @@ async def get_image_map_image_labels( if not vocabulary: raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail="The labeling vocabulary is empty") - # A BLOB read and a SQLite transaction; off the event loop like the + # A BLOB read and a database transaction; off the event loop like the # clustering endpoints, since this one fires per hovered point. # # A stored row whose blob length disagrees with its `dim` column raises out diff --git a/invokeai/app/api/routers/images.py b/invokeai/app/api/routers/images.py index 6968dde9771..874efd5f424 100644 --- a/invokeai/app/api/routers/images.py +++ b/invokeai/app/api/routers/images.py @@ -49,8 +49,7 @@ UnstarredImagesResult, ) from invokeai.app.services.intermediates.intermediates_base import IntermediatesCaller -from invokeai.app.services.shared.pagination import MAX_PAGE_SIZE, OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import MAX_PAGE_SIZE, OffsetPaginatedResults, SQLiteDirection from invokeai.app.util.controlnet_utils import heuristic_resize_fast from invokeai.backend.image_util.util import np_to_pil, pil_to_np diff --git a/invokeai/app/api/routers/model_manager.py b/invokeai/app/api/routers/model_manager.py index 39a26bd25b9..894274db4f7 100644 --- a/invokeai/app/api/routers/model_manager.py +++ b/invokeai/app/api/routers/model_manager.py @@ -42,7 +42,7 @@ UnknownModelException, ) from invokeai.app.services.orphaned_models import CONVERSION_SCRATCH_DIRNAME, OrphanedModelInfo -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.util.path_safety import is_plain_filename from invokeai.app.util.suppress_output import SuppressOutput from invokeai.backend.architectures import ArchitectureCapabilities, architecture_capabilities @@ -1823,12 +1823,9 @@ def get_orphaned_models(_: AdminUserOrDefault) -> list[OrphanedModelInfo]: """ from invokeai.app.services.orphaned_models import OrphanedModelsService - # Access the database through the model records service - model_records_service = ApiDependencies.invoker.services.model_manager.store - service = OrphanedModelsService( config=ApiDependencies.invoker.services.configuration, - db=model_records_service._db, # Access the database from model records service + store=ApiDependencies.invoker.services.model_manager.store, ) return service.find_orphaned_models() @@ -1849,12 +1846,9 @@ def delete_orphaned_models(request: DeleteOrphanedModelsRequest, _: AdminUserOrD """ from invokeai.app.services.orphaned_models import OrphanedModelsService - # Access the database through the model records service - model_records_service = ApiDependencies.invoker.services.model_manager.store - service = OrphanedModelsService( config=ApiDependencies.invoker.services.configuration, - db=model_records_service._db, # Access the database from model records service + store=ApiDependencies.invoker.services.model_manager.store, ) results = service.delete_orphaned_models(request.paths) diff --git a/invokeai/app/api/routers/model_relationships.py b/invokeai/app/api/routers/model_relationships.py index 3a882038c0f..f8945ad60b2 100644 --- a/invokeai/app/api/routers/model_relationships.py +++ b/invokeai/app/api/routers/model_relationships.py @@ -16,6 +16,7 @@ class ModelRelationshipCreateRequest(BaseModel): model_key_1: str = Field( ..., + max_length=255, description="The key of the first model in the relationship", examples=[ "aa3b247f-90c9-4416-bfcd-aeaa57a5339e", @@ -28,6 +29,7 @@ class ModelRelationshipCreateRequest(BaseModel): ) model_key_2: str = Field( ..., + max_length=255, description="The key of the second model in the relationship", examples=[ "3bb7c0eb-b6c8-469c-ad8c-4d69c06075e4", diff --git a/invokeai/app/api/routers/projects.py b/invokeai/app/api/routers/projects.py index a76bbcd8730..374a8c2d6aa 100644 --- a/invokeai/app/api/routers/projects.py +++ b/invokeai/app/api/routers/projects.py @@ -37,7 +37,9 @@ class ProjectCreateRequest(BaseModel): """Request body for creating a project.""" project_id: str | None = Field( - default=None, description="Client-generated project id (e.g. for imports); generated when omitted" + default=None, + max_length=255, + description="Client-generated project id (e.g. for imports); generated when omitted", ) board_id: str | None = Field( default=None, diff --git a/invokeai/app/api/routers/session_queue.py b/invokeai/app/api/routers/session_queue.py index c144d396f42..c64f32623b1 100644 --- a/invokeai/app/api/routers/session_queue.py +++ b/invokeai/app/api/routers/session_queue.py @@ -37,16 +37,16 @@ SessionQueueStatus, ) from invokeai.app.services.shared.graph import Graph, GraphExecutionState -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.services.video_records.video_records_common import VideoRecordNotFoundException session_queue_router = APIRouter(prefix="/v1/queue", tags=["queue"]) # Upper bound on the number of item ids a client may ask about in one request. Without it a -# caller can post tens of thousands of ids, which the SQLite layer would either expand past the -# per-statement bind limit or grind through in a long-running query. The list is meant to cover -# the rows a client actually has on screen, so this is far above any legitimate use. A limited id -# listing is capped at the same size, since its ids are fetched to be hydrated that way. +# caller can post tens of thousands of ids and make the database grind through a long-running +# query. The list is meant to cover the rows a client actually has on screen, so this is far +# above any legitimate use. A limited id listing is capped at the same size, since its ids are +# fetched to be hydrated that way. MAX_QUEUE_ITEM_IDS_PER_REQUEST = 1000 @@ -249,7 +249,7 @@ def get_queue_item_for_mutation(queue_id: str, item_id: int, current_user: Curre ) async def enqueue_batch( current_user: CurrentUserOrDefault, - queue_id: str = Path(description="The queue id to perform this operation on"), + queue_id: str = Path(description="The queue id to perform this operation on", max_length=255), batch: Batch = Body(description="Batch to process"), prepend: bool = Body(default=False, description="Whether or not to prepend this batch in the queue"), ) -> EnqueueBatchResult: diff --git a/invokeai/app/api/routers/videos.py b/invokeai/app/api/routers/videos.py index 7c22738f560..70bfa3456bc 100644 --- a/invokeai/app/api/routers/videos.py +++ b/invokeai/app/api/routers/videos.py @@ -41,8 +41,7 @@ from invokeai.app.api.routers.images import WorkflowAndGraphResponse from invokeai.app.invocations.fields import MetadataField, MetadataFieldValidator from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin, is_gallery_category -from invokeai.app.services.shared.pagination import MAX_PAGE_SIZE, OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import MAX_PAGE_SIZE, OffsetPaginatedResults, SQLiteDirection from invokeai.app.services.video_records.video_records_common import ( VideoNamesResult, VideoRecordChanges, diff --git a/invokeai/app/api/routers/virtual_boards.py b/invokeai/app/api/routers/virtual_boards.py index 2742b4ef96d..8274d5b1a2e 100644 --- a/invokeai/app/api/routers/virtual_boards.py +++ b/invokeai/app/api/routers/virtual_boards.py @@ -5,7 +5,7 @@ from invokeai.app.api.dependencies import ApiDependencies from invokeai.app.services.gallery.gallery_common import GalleryItemNamesResult from invokeai.app.services.image_records.image_records_common import ImageCategory, ImageNamesResult -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.services.virtual_boards.virtual_boards_common import VirtualSubBoardDTO virtual_boards_router = APIRouter(prefix="/v1/virtual_boards", tags=["virtual_boards"]) diff --git a/invokeai/app/api/routers/workflows.py b/invokeai/app/api/routers/workflows.py index d366f13e179..10cdf0d4cf3 100644 --- a/invokeai/app/api/routers/workflows.py +++ b/invokeai/app/api/routers/workflows.py @@ -10,8 +10,7 @@ from invokeai.app.api.auth_dependencies import CurrentUserOrDefault from invokeai.app.api.dependencies import ApiDependencies -from invokeai.app.services.shared.pagination import PaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import MAX_PAGE_SIZE, PaginatedResults, SQLiteDirection from invokeai.app.services.shared.workflow_call_compatibility import get_workflow_call_compatibility from invokeai.app.services.workflow_records.workflow_records_common import ( Workflow, @@ -57,7 +56,9 @@ def get_workflow( if not (is_default or is_owner or workflow.is_public or current_user.is_admin): raise HTTPException(status_code=403, detail="Not authorized to access this workflow") - thumbnail_url = ApiDependencies.invoker.services.workflow_thumbnails.get_url(workflow_id) + thumbnail_url = ApiDependencies.invoker.services.workflow_thumbnails.get_url( + workflow_id, category=workflow.workflow.meta.category + ) compatibility = get_workflow_call_compatibility( workflow=workflow.workflow.model_dump(), workflow_id=workflow.workflow_id, @@ -197,6 +198,7 @@ def create_workflow( workflow_id: Optional[str] = Body( default=None, embed=True, + max_length=255, description="A client-reserved UUID for the new record. Retrying the same creation with the same id returns " "the record already created for it; another owner's record or different content under that id is a 409.", ), @@ -242,8 +244,10 @@ def create_workflow( ) def list_workflows( current_user: CurrentUserOrDefault, - page: int = Query(default=0, description="The page to get"), - per_page: Optional[int] = Query(default=None, description="The number of workflows per page"), + page: int = Query(default=0, ge=0, description="The page to get"), + per_page: Optional[int] = Query( + default=None, ge=0, le=MAX_PAGE_SIZE, description="The number of workflows per page; all of them when omitted" + ), order_by: WorkflowRecordOrderBy = Query( default=WorkflowRecordOrderBy.Name, description="The attribute to order by" ), @@ -302,7 +306,9 @@ def list_workflows( continue workflows_with_thumbnails.append( WorkflowRecordListItemWithThumbnailDTO( - thumbnail_url=ApiDependencies.invoker.services.workflow_thumbnails.get_url(workflow.workflow_id), + thumbnail_url=ApiDependencies.invoker.services.workflow_thumbnails.get_url( + workflow.workflow_id, category=workflow.category + ), call_saved_workflow_compatibility=compatibility, **workflow.model_dump(), ) diff --git a/invokeai/app/api_app.py b/invokeai/app/api_app.py index 0b089bc0902..ee0b722fe0f 100644 --- a/invokeai/app/api_app.py +++ b/invokeai/app/api_app.py @@ -5,16 +5,19 @@ from collections.abc import Awaitable, Callable from contextlib import asynccontextmanager from pathlib import Path +from typing import Optional from fastapi import FastAPI, Request +from fastapi.exception_handlers import http_exception_handler from fastapi.middleware.cors import CORSMiddleware from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html -from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse +from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response from fastapi.security.utils import get_authorization_scheme_param from fastapi_events.handlers.local import local_handler from fastapi_events.middleware import EventHandlerASGIMiddleware from starlette.concurrency import run_in_threadpool from starlette.datastructures import Headers +from starlette.exceptions import HTTPException as StarletteHTTPException from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.middleware.gzip import GZipMiddleware, GZipResponder, IdentityResponder from starlette.types import ASGIApp, Message, Receive, Scope, Send @@ -52,6 +55,7 @@ ) from invokeai.app.api.sockets import SocketIO from invokeai.app.services.config.config_default import get_config +from invokeai.app.services.shared.database.errors import TransientDatabaseError from invokeai.app.util.custom_openapi import get_openapi_func from invokeai.backend.util.logging import InvokeAILogger from invokeai.frontend.cli.arg_parser import InvokeAIArgs @@ -110,6 +114,42 @@ async def lifespan(app: FastAPI): ) +@app.exception_handler(TransientDatabaseError) +async def database_busy(request: Request, error: TransientDatabaseError) -> JSONResponse: + return _database_busy(request, error) + + +@app.exception_handler(StarletteHTTPException) +async def http_error(request: Request, error: StarletteHTTPException) -> Response: + # Many routes answer every exception with an error of their own ("not found", "unexpected error"); one raised + # while handling a busy database still means the database is busy. + transient = _transient_cause(error) + if transient is not None: + return _database_busy(request, transient) + return await http_exception_handler(request, error) + + +def _database_busy(request: Request, error: TransientDatabaseError) -> JSONResponse: + # A lock held too long elsewhere, a race the call lost on every attempt, a server out of reach: the same + # request can succeed in a moment, so it is a busy server rather than a failed request. + logger.warning(f"{request.method} {request.url.path}: the database is busy: {error}") + return JSONResponse( + status_code=503, content={"detail": "The database is busy; try again"}, headers={"Retry-After": "1"} + ) + + +def _transient_cause(error: BaseException) -> Optional[TransientDatabaseError]: + """The busy database an exception was raised while handling, if any.""" + seen: set[int] = set() + cause: Optional[BaseException] = error.__cause__ or error.__context__ + while cause is not None and id(cause) not in seen: + if isinstance(cause, TransientDatabaseError): + return cause + seen.add(id(cause)) + cause = cause.__cause__ or cause.__context__ + return None + + class SlidingWindowTokenMiddleware(BaseHTTPMiddleware): """Refresh the JWT token on each authenticated response. diff --git a/invokeai/app/run_app.py b/invokeai/app/run_app.py index 4e513d210a8..a74a375c988 100644 --- a/invokeai/app/run_app.py +++ b/invokeai/app/run_app.py @@ -167,6 +167,10 @@ def run_app() -> None: try: loop.run_until_complete(server.serve()) + if not server.started: + # The app refused to start (its log says why, e.g. a database another process serves); uvicorn returns + # normally, so name the failure to whatever runs this process. 3 is uvicorn's own startup-failure code. + raise SystemExit(3) except KeyboardInterrupt: logger.info("InvokeAI shutting down...") # Gracefully shut down services (e.g. model download and install managers) so that any diff --git a/invokeai/app/services/app_settings/app_settings_service.py b/invokeai/app/services/app_settings/app_settings_service.py index 5580709ef65..ee89ec156d0 100644 --- a/invokeai/app/services/app_settings/app_settings_service.py +++ b/invokeai/app/services/app_settings/app_settings_service.py @@ -2,7 +2,7 @@ from typing import Optional -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database class AppSettingsService: @@ -12,13 +12,8 @@ class AppSettingsService: that needs to be persisted across restarts, such as JWT secrets. """ - def __init__(self, db: SqliteDatabase) -> None: - """Initialize the app settings service. - - Args: - db: The SQLite database instance - """ - self._db = db + def __init__(self, database: Database) -> None: + self._queries = database.queries def get(self, key: str) -> Optional[str]: """Get a setting value by key. @@ -29,32 +24,7 @@ def get(self, key: str) -> Optional[str]: Returns: The setting value if found, None otherwise """ - try: - with self._db.transaction() as cursor: - cursor.execute("SELECT value FROM app_settings WHERE key = ?;", (key,)) - row = cursor.fetchone() - return row[0] if row else None - except Exception: - return None - - def set(self, key: str, value: str) -> None: - """Set a setting value. - - Args: - key: The setting key - value: The setting value - """ - with self._db.transaction() as cursor: - cursor.execute( - """ - INSERT INTO app_settings (key, value) - VALUES (?, ?) - ON CONFLICT(key) DO UPDATE SET - value = excluded.value, - updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW'); - """, - (key, value), - ) + return self._queries.app_settings.get(key) def get_jwt_secret(self) -> str: """Get the JWT secret key from the database. diff --git a/invokeai/app/services/board_image_records/board_image_records_base.py b/invokeai/app/services/board_image_records/board_image_records_base.py index 34f74868cc1..e289d5ab7d8 100644 --- a/invokeai/app/services/board_image_records/board_image_records_base.py +++ b/invokeai/app/services/board_image_records/board_image_records_base.py @@ -55,17 +55,6 @@ def get_board_for_image( pass @abstractmethod - def get_image_count_for_board( - self, - board_id: str, - ) -> int: - """Gets the number of images for a board.""" - pass - - @abstractmethod - def get_asset_count_for_board( - self, - board_id: str, - ) -> int: - """Gets the number of assets for a board.""" + def get_counts_for_board(self, board_id: str) -> tuple[int, int]: + """Gets the numbers of images and of assets on a board, intermediates not counted.""" pass diff --git a/invokeai/app/services/board_image_records/board_image_records_default.py b/invokeai/app/services/board_image_records/board_image_records_default.py new file mode 100644 index 00000000000..bdbc240ba17 --- /dev/null +++ b/invokeai/app/services/board_image_records/board_image_records_default.py @@ -0,0 +1,54 @@ +from typing import Optional + +from invokeai.app.services.board_image_records.board_image_records_base import BoardImageRecordStorageBase +from invokeai.app.services.image_records.image_records_common import ImageCategory +from invokeai.app.services.shared.database.database import Database + + +class BoardImageRecordStorage(BoardImageRecordStorageBase): + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + + def add_image_to_board( + self, + board_id: str, + image_name: str, + ) -> None: + self._queries.board_images.add(board_id, image_name) + + def remove_image_from_board( + self, + image_name: str, + board_id: str, + ) -> int: + # Scoped to the board the caller was authorized against, not just the image. The + # routes read the image's board, authorize against *that* board, and only then + # remove; an unscoped DELETE would follow the image if it were moved in between, + # applying a decision taken about one board to a different one. Zero rows means the + # image left this board between the caller's read and this write. + return int(self._queries.board_images.remove(image_name, board_id)) + + def get_all_board_image_names_for_board( + self, + board_id: str, + categories: list[ImageCategory] | None, + is_intermediate: bool | None, + user_id: Optional[str] = None, + ) -> list[str]: + # `board_id` "none" means the images on no board; admins pass user_id=None to see every account's. + return self._queries.board_images.image_names( + None if board_id == "none" else board_id, + categories=categories, + is_intermediate=is_intermediate, + user_id=user_id, + ) + + def get_board_for_image( + self, + image_name: str, + ) -> Optional[str]: + return self._queries.board_images.board_of(image_name) + + def get_counts_for_board(self, board_id: str) -> tuple[int, int]: + return self._queries.board_images.counts(board_id) diff --git a/invokeai/app/services/board_image_records/board_image_records_sqlite.py b/invokeai/app/services/board_image_records/board_image_records_sqlite.py deleted file mode 100644 index 0bb95cfad90..00000000000 --- a/invokeai/app/services/board_image_records/board_image_records_sqlite.py +++ /dev/null @@ -1,211 +0,0 @@ -import sqlite3 -from typing import Optional, cast - -from invokeai.app.services.board_image_records.board_image_records_base import BoardImageRecordStorageBase -from invokeai.app.services.image_records.image_records_common import ( - ASSETS_CATEGORIES, - IMAGE_CATEGORIES, - ImageCategory, - ImageRecord, - deserialize_image_record, -) -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase - - -class SqliteBoardImageRecordStorage(BoardImageRecordStorageBase): - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - - def add_image_to_board( - self, - board_id: str, - image_name: str, - ) -> None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - INSERT INTO board_images (board_id, image_name) - VALUES (?, ?) - ON CONFLICT (image_name) DO UPDATE SET board_id = ?; - """, - (board_id, image_name, board_id), - ) - - def remove_image_from_board( - self, - image_name: str, - board_id: str, - ) -> int: - with self._db.transaction() as cursor: - # Scoped to the board the caller was authorized against, not just the image. The - # routes read the image's board, authorize against *that* board, and only then - # remove; an unscoped DELETE would follow the image if it were moved in between, - # applying a decision taken about one board to a different one. - cursor.execute( - """--sql - DELETE FROM board_images - WHERE image_name = ? AND board_id = ?; - """, - (image_name, board_id), - ) - # The row count is the only signal that the scope held. Zero rows means the image - # left this board between the caller's read and this write — swallowing that lets - # the route report a removal that did not happen. - return cursor.rowcount - - def get_images_for_board( - self, - board_id: str, - offset: int = 0, - limit: int = 10, - ) -> OffsetPaginatedResults[ImageRecord]: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT images.* - FROM board_images - INNER JOIN images ON board_images.image_name = images.image_name - WHERE board_images.board_id = ? - ORDER BY board_images.updated_at DESC; - """, - (board_id,), - ) - result = cast(list[sqlite3.Row], cursor.fetchall()) - images = [deserialize_image_record(dict(r)) for r in result] - - cursor.execute( - """--sql - SELECT COUNT(*) FROM images WHERE 1=1; - """ - ) - count = cast(int, cursor.fetchone()[0]) - - return OffsetPaginatedResults(items=images, offset=offset, limit=limit, total=count) - - def get_all_board_image_names_for_board( - self, - board_id: str, - categories: list[ImageCategory] | None, - is_intermediate: bool | None, - user_id: Optional[str] = None, - ) -> list[str]: - with self._db.transaction() as cursor: - params: list[str | bool] = [] - - # Base query is a join between images and board_images - stmt = """ - SELECT images.image_name - FROM images - LEFT JOIN board_images ON board_images.image_name = images.image_name - WHERE 1=1 - """ - - # Handle board_id filter - if board_id == "none": - stmt += """--sql - AND board_images.board_id IS NULL - """ - else: - stmt += """--sql - AND board_images.board_id = ? - """ - params.append(board_id) - - # Add the category filter - if categories is not None: - # Convert the enum values to unique list of strings - category_strings = [c.value for c in set(categories)] - # Create the correct length of placeholders - placeholders = ",".join("?" * len(category_strings)) - stmt += f"""--sql - AND images.image_category IN ( {placeholders} ) - """ - - # Unpack the included categories into the query params - for c in category_strings: - params.append(c) - - # Add the is_intermediate filter - if is_intermediate is not None: - stmt += """--sql - AND images.is_intermediate = ? - """ - params.append(is_intermediate) - - # Per-user filter — admins pass user_id=None to skip this clause. - if user_id is not None: - stmt += """--sql - AND images.user_id = ? - """ - params.append(user_id) - - # Put a ring on it - stmt += ";" - - cursor.execute(stmt, params) - - result = cast(list[sqlite3.Row], cursor.fetchall()) - image_names = [r[0] for r in result] - return image_names - - def get_board_for_image( - self, - image_name: str, - ) -> Optional[str]: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT board_id - FROM board_images - WHERE image_name = ?; - """, - (image_name,), - ) - result = cursor.fetchone() - if result is None: - return None - return cast(str, result[0]) - - def get_image_count_for_board(self, board_id: str) -> int: - with self._db.transaction() as cursor: - # Convert the enum values to unique list of strings - category_strings = [c.value for c in set(IMAGE_CATEGORIES)] - # Create the correct length of placeholders - placeholders = ",".join("?" * len(category_strings)) - cursor.execute( - f"""--sql - SELECT COUNT(*) - FROM board_images - -- Keep work proportional to this board's membership instead of allowing - -- the gallery index to make images the outer loop for this query. - CROSS JOIN images ON board_images.image_name = images.image_name - WHERE images.is_intermediate = FALSE AND images.image_category IN ( {placeholders} ) - AND board_images.board_id = ?; - """, - (*category_strings, board_id), - ) - count = cast(int, cursor.fetchone()[0]) - return count - - def get_asset_count_for_board(self, board_id: str) -> int: - with self._db.transaction() as cursor: - # Convert the enum values to unique list of strings - category_strings = [c.value for c in set(ASSETS_CATEGORIES)] - # Create the correct length of placeholders - placeholders = ",".join("?" * len(category_strings)) - cursor.execute( - f"""--sql - SELECT COUNT(*) - FROM board_images - -- Keep work proportional to this board's membership instead of allowing - -- the gallery index to make images the outer loop for this query. - CROSS JOIN images ON board_images.image_name = images.image_name - WHERE images.is_intermediate = FALSE AND images.image_category IN ( {placeholders} ) - AND board_images.board_id = ?; - """, - (*category_strings, board_id), - ) - count = cast(int, cursor.fetchone()[0]) - return count diff --git a/invokeai/app/services/board_records/board_records_base.py b/invokeai/app/services/board_records/board_records_base.py index 7eb98ccd09f..b16f871a042 100644 --- a/invokeai/app/services/board_records/board_records_base.py +++ b/invokeai/app/services/board_records/board_records_base.py @@ -2,8 +2,7 @@ from typing import Optional from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecord, BoardRecordOrderBy -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection class BoardRecordStorageBase(ABC): diff --git a/invokeai/app/services/board_records/board_records_common.py b/invokeai/app/services/board_records/board_records_common.py index 6a3ff508449..b0b8d9c4943 100644 --- a/invokeai/app/services/board_records/board_records_common.py +++ b/invokeai/app/services/board_records/board_records_common.py @@ -91,7 +91,9 @@ class BoardChanges(BaseModel, extra="forbid"): board_name: Optional[str] = Field( default=None, description="The board's new name.", max_length=BOARD_NAME_MAX_LENGTH ) - cover_image_name: Optional[str] = Field(default=None, description="The name of the board's new cover image.") + cover_image_name: Optional[str] = Field( + default=None, max_length=255, description="The name of the board's new cover image." + ) archived: Optional[bool] = Field(default=None, description="Whether or not the board is archived") board_visibility: Optional[BoardVisibility] = Field(default=None, description="The visibility of the board.") @@ -122,10 +124,3 @@ class BoardRecordProjectOwnedException(BoardRecordSaveException): def __init__(self, message="Board belongs to a project"): super().__init__(message) - - -class BoardRecordDeleteException(Exception): - """Raised when an board record cannot be deleted.""" - - def __init__(self, message="Board record not deleted"): - super().__init__(message) diff --git a/invokeai/app/services/board_records/board_records_default.py b/invokeai/app/services/board_records/board_records_default.py new file mode 100644 index 00000000000..1e7807a0e48 --- /dev/null +++ b/invokeai/app/services/board_records/board_records_default.py @@ -0,0 +1,133 @@ +from typing import Optional + +from invokeai.app.services.board_records.board_records_base import BoardRecordStorageBase +from invokeai.app.services.board_records.board_records_common import ( + BoardChanges, + BoardRecord, + BoardRecordNotFoundException, + BoardRecordOrderBy, + BoardRecordProjectOwnedException, + BoardRecordSaveException, +) +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection +from invokeai.app.util.misc import uuid_string + + +class BoardRecordStorage(BoardRecordStorageBase): + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + + def delete_if_unclaimed(self, board_id: str) -> bool: + return self._queries.boards.delete_if_unclaimed(board_id) + + def get_project_ids_for_boards(self, board_ids: list[str]) -> dict[str, str]: + if not board_ids: + return {} + # The caller is a board *listing*: an admin's `GET /boards/?all=true` passes every board on the install. + return self._queries.boards.project_ids(board_ids) + + def get_with_project_id(self, board_id: str) -> tuple[BoardRecord, Optional[str]]: + """The board and the project that claims it, in one query: `get_dto` needs both, for every board the + API returns.""" + found = self._queries.boards.get_with_project_id(board_id) + if found is None: + raise BoardRecordNotFoundException + return found + + def save( + self, + board_name: str, + user_id: str, + ) -> BoardRecord: + board_id = uuid_string() + + def insert(q: Queries) -> Optional[BoardRecord]: + q.boards.insert(board_id=board_id, board_name=board_name, user_id=user_id) + return q.boards.get(board_id) + + board = self._queries.run(insert) + if board is None: + raise BoardRecordSaveException + return board + + def get( + self, + board_id: str, + ) -> BoardRecord: + # A database error is deliberately NOT translated into BoardRecordNotFoundException. + # Translating it made the exception mean "no such board, OR the database is + # locked/corrupt/unreadable", and callers cannot tell those apart: the batch routes + # decide board write access off this read once per name, and a not-found answer there is + # a benign skip. A disk error would therefore drop names out of the response silently, + # reported neither as moved nor as failed. Mirrors ImageRecordStorage.get. + board = self._queries.boards.get(board_id) + if board is None: + raise BoardRecordNotFoundException + return board + + def is_board_shared_with_user(self, board_id: str, user_id: str) -> bool: + return self._queries.boards.is_shared_with(board_id, user_id) + + def get_shared_user_ids(self, board_id: str) -> list[str]: + return self._queries.boards.shared_user_ids(board_id) + + def update( + self, + board_id: str, + changes: BoardChanges, + ) -> BoardRecord: + def apply(q: Queries) -> Optional[BoardRecord]: + # One conditional write, not a read followed by a write. A project claim racing this statement + # either commits first and makes it change nothing, or waits until this update has committed; there + # is no stale DTO window in which project-owned state can be renamed, archived or published. + if not q.boards.update(board_id, changes): + if not q.boards.exists(board_id): + raise BoardRecordNotFoundException + raise BoardRecordProjectOwnedException + return q.boards.get(board_id) + + board = self._queries.run(apply) + if board is None: + raise BoardRecordNotFoundException + return board + + def get_many( + self, + user_id: str, + is_admin: bool, + order_by: BoardRecordOrderBy, + direction: SQLiteDirection, + offset: int = 0, + limit: int = 10, + include_archived: bool = False, + ) -> OffsetPaginatedResults[BoardRecord]: + # Admins see every board; other users their own, those shared with them, and shared or public ones. + boards, total = self._queries.boards.page( + user_id=user_id, + is_admin=is_admin, + order_by=order_by, + direction=direction, + offset=offset, + limit=limit, + include_archived=include_archived, + ) + return OffsetPaginatedResults[BoardRecord](items=boards, offset=offset, limit=limit, total=total) + + def get_all( + self, + user_id: str, + is_admin: bool, + order_by: BoardRecordOrderBy, + direction: SQLiteDirection, + include_archived: bool = False, + ) -> list[BoardRecord]: + return self._queries.boards.all( + user_id=user_id, + is_admin=is_admin, + order_by=order_by, + direction=direction, + include_archived=include_archived, + ) diff --git a/invokeai/app/services/board_records/board_records_sqlite.py b/invokeai/app/services/board_records/board_records_sqlite.py deleted file mode 100644 index 79c3fbb2e40..00000000000 --- a/invokeai/app/services/board_records/board_records_sqlite.py +++ /dev/null @@ -1,370 +0,0 @@ -import sqlite3 -from typing import Optional, Union, cast - -from invokeai.app.services.board_records.board_records_base import BoardRecordStorageBase -from invokeai.app.services.board_records.board_records_common import ( - BoardChanges, - BoardRecord, - BoardRecordDeleteException, - BoardRecordNotFoundException, - BoardRecordOrderBy, - BoardRecordProjectOwnedException, - BoardRecordSaveException, - deserialize_board_record, -) -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.util.misc import uuid_string - -# Board ids per `IN (...)`. Well under SQLITE_MAX_VARIABLE_NUMBER on every build, including the -# 999-variable default of older SQLite. -_BOARD_ID_QUERY_CHUNK = 500 - - -class SqliteBoardRecordStorage(BoardRecordStorageBase): - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - - def delete_if_unclaimed(self, board_id: str) -> bool: - with self._db.transaction() as cursor: - try: - cursor.execute( - """--sql - DELETE FROM boards - WHERE board_id = ? - AND NOT EXISTS ( - SELECT 1 FROM projects WHERE projects.board_id = boards.board_id - ); - """, - (board_id,), - ) - return cursor.rowcount > 0 - except Exception as e: - raise BoardRecordDeleteException from e - - def get_project_ids_for_boards(self, board_ids: list[str]) -> dict[str, str]: - if not board_ids: - return {} - - # Chunked because the caller is a board *listing*: an admin's `GET /boards/?all=true` - # passes every board on the install at once, and one bind parameter per board runs into - # SQLITE_MAX_VARIABLE_NUMBER — which fails the whole listing rather than degrading. - claimed: dict[str, str] = {} - with self._db.transaction() as cursor: - for start in range(0, len(board_ids), _BOARD_ID_QUERY_CHUNK): - chunk = board_ids[start : start + _BOARD_ID_QUERY_CHUNK] - placeholders = ", ".join("?" for _ in chunk) - cursor.execute( - f"""--sql - SELECT board_id, project_id FROM projects WHERE board_id IN ({placeholders}); - """, - tuple(chunk), - ) - claimed.update({row[0]: row[1] for row in cursor.fetchall()}) - return claimed - - def get_with_project_id(self, board_id: str) -> tuple[BoardRecord, Optional[str]]: - """The board and the project that claims it, in one query. - - `get_dto` needs both and is called on every authorization check, so asking for them - separately doubled the round trips — and the transactions, which take a re-entrant lock — - on the hottest read in the API. - """ - with self._db.transaction() as cursor: - try: - cursor.execute( - """--sql - SELECT boards.*, projects.project_id AS claimed_by_project_id - FROM boards - LEFT JOIN projects ON projects.board_id = boards.board_id - WHERE boards.board_id = ?; - """, - (board_id,), - ) - result = cast(Union[sqlite3.Row, None], cursor.fetchone()) - except sqlite3.Error as e: - raise BoardRecordNotFoundException from e - if result is None: - raise BoardRecordNotFoundException - row = dict(result) - project_id = row.pop("claimed_by_project_id", None) - return BoardRecord(**row), project_id - - def save( - self, - board_name: str, - user_id: str, - ) -> BoardRecord: - with self._db.transaction() as cursor: - try: - board_id = uuid_string() - cursor.execute( - """--sql - INSERT OR IGNORE INTO boards (board_id, board_name, user_id) - VALUES (?, ?, ?); - """, - (board_id, board_name, user_id), - ) - except sqlite3.Error as e: - raise BoardRecordSaveException from e - return self.get(board_id) - - def get( - self, - board_id: str, - ) -> BoardRecord: - # A sqlite3.Error is deliberately NOT translated into BoardRecordNotFoundException. - # Translating it made the exception mean "no such board, OR the database is - # locked/corrupt/unreadable", and callers cannot tell those apart: the batch routes - # decide board write access off this read once per name, and a not-found answer there is - # a benign skip. A disk error would therefore drop names out of the response silently, - # reported neither as moved nor as failed. Mirrors ImageRecordStorage.get. - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT * - FROM boards - WHERE board_id = ?; - """, - (board_id,), - ) - - result = cast(Union[sqlite3.Row, None], cursor.fetchone()) - if result is None: - raise BoardRecordNotFoundException - return BoardRecord(**dict(result)) - - def is_board_shared_with_user(self, board_id: str, user_id: str) -> bool: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT 1 - FROM shared_boards - WHERE board_id = ? AND user_id = ?; - """, - (board_id, user_id), - ) - - return cursor.fetchone() is not None - - def get_shared_user_ids(self, board_id: str) -> list[str]: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT user_id - FROM shared_boards - WHERE board_id = ?; - """, - (board_id,), - ) - - return [row["user_id"] for row in cursor.fetchall()] - - def update( - self, - board_id: str, - changes: BoardChanges, - ) -> BoardRecord: - with self._db.transaction() as cursor: - try: - changes_project_owned_state = ( - changes.board_name is not None - or changes.archived is not None - or changes.board_visibility is not None - ) - # One conditional write, not a read followed by a write. A project claim racing - # this statement either commits first and makes rowcount zero, or waits until this - # update has committed; there is no stale DTO window in which project-owned state - # can be renamed, archived or published. - cursor.execute( - """--sql - UPDATE boards - SET board_name = COALESCE(?, board_name), - cover_image_name = COALESCE(?, cover_image_name), - archived = COALESCE(?, archived), - board_visibility = COALESCE(?, board_visibility) - WHERE board_id = ? - AND ( - ? = FALSE - OR NOT EXISTS (SELECT 1 FROM projects WHERE projects.board_id = boards.board_id) - ); - """, - ( - changes.board_name, - changes.cover_image_name, - changes.archived, - changes.board_visibility.value if changes.board_visibility is not None else None, - board_id, - changes_project_owned_state, - ), - ) - - if cursor.rowcount == 0: - cursor.execute("SELECT 1 FROM boards WHERE board_id = ?;", (board_id,)) - if cursor.fetchone() is None: - raise BoardRecordNotFoundException - raise BoardRecordProjectOwnedException - - except sqlite3.Error as e: - raise BoardRecordSaveException from e - return self.get(board_id) - - def get_many( - self, - user_id: str, - is_admin: bool, - order_by: BoardRecordOrderBy, - direction: SQLiteDirection, - offset: int = 0, - limit: int = 10, - include_archived: bool = False, - ) -> OffsetPaginatedResults[BoardRecord]: - with self._db.transaction() as cursor: - # Build base query - admins see all boards, regular users see owned, shared, or public boards - if is_admin: - base_query = """ - SELECT DISTINCT boards.* - FROM boards - {archived_filter} - ORDER BY {order_by} {direction} - LIMIT ? OFFSET ?; - """ - - # Determine archived filter condition - archived_filter = "WHERE 1=1" if include_archived else "WHERE boards.archived = 0" - - final_query = base_query.format( - archived_filter=archived_filter, order_by=order_by.value, direction=direction.value - ) - - # Execute query to fetch boards - cursor.execute(final_query, (limit, offset)) - else: - base_query = """ - SELECT DISTINCT boards.* - FROM boards - LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id - WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public')) - {archived_filter} - ORDER BY {order_by} {direction} - LIMIT ? OFFSET ?; - """ - - # Determine archived filter condition - archived_filter = "" if include_archived else "AND boards.archived = 0" - - final_query = base_query.format( - archived_filter=archived_filter, order_by=order_by.value, direction=direction.value - ) - - # Execute query to fetch boards - cursor.execute(final_query, (user_id, user_id, limit, offset)) - - result = cast(list[sqlite3.Row], cursor.fetchall()) - boards = [deserialize_board_record(dict(r)) for r in result] - - # Determine count query - admins count all boards, regular users count accessible boards - if is_admin: - if include_archived: - count_query = """ - SELECT COUNT(DISTINCT boards.board_id) - FROM boards; - """ - else: - count_query = """ - SELECT COUNT(DISTINCT boards.board_id) - FROM boards - WHERE boards.archived = 0; - """ - cursor.execute(count_query) - else: - if include_archived: - count_query = """ - SELECT COUNT(DISTINCT boards.board_id) - FROM boards - LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id - WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public')); - """ - else: - count_query = """ - SELECT COUNT(DISTINCT boards.board_id) - FROM boards - LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id - WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public')) - AND boards.archived = 0; - """ - - # Execute count query - cursor.execute(count_query, (user_id, user_id)) - - count = cast(int, cursor.fetchone()[0]) - - return OffsetPaginatedResults[BoardRecord](items=boards, offset=offset, limit=limit, total=count) - - def get_all( - self, - user_id: str, - is_admin: bool, - order_by: BoardRecordOrderBy, - direction: SQLiteDirection, - include_archived: bool = False, - ) -> list[BoardRecord]: - with self._db.transaction() as cursor: - # Build query - admins see all boards, regular users see owned, shared, or public boards - if is_admin: - if order_by == BoardRecordOrderBy.Name: - base_query = """ - SELECT DISTINCT boards.* - FROM boards - {archived_filter} - ORDER BY LOWER(boards.board_name) {direction} - """ - else: - base_query = """ - SELECT DISTINCT boards.* - FROM boards - {archived_filter} - ORDER BY {order_by} {direction} - """ - - archived_filter = "WHERE 1=1" if include_archived else "WHERE boards.archived = 0" - - final_query = base_query.format( - archived_filter=archived_filter, order_by=order_by.value, direction=direction.value - ) - - cursor.execute(final_query) - else: - if order_by == BoardRecordOrderBy.Name: - base_query = """ - SELECT DISTINCT boards.* - FROM boards - LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id - WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public')) - {archived_filter} - ORDER BY LOWER(boards.board_name) {direction} - """ - else: - base_query = """ - SELECT DISTINCT boards.* - FROM boards - LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id - WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public')) - {archived_filter} - ORDER BY {order_by} {direction} - """ - - archived_filter = "" if include_archived else "AND boards.archived = 0" - - final_query = base_query.format( - archived_filter=archived_filter, order_by=order_by.value, direction=direction.value - ) - - cursor.execute(final_query, (user_id, user_id)) - - result = cast(list[sqlite3.Row], cursor.fetchall()) - boards = [deserialize_board_record(dict(r)) for r in result] - - return boards diff --git a/invokeai/app/services/board_video_records/board_video_records_base.py b/invokeai/app/services/board_video_records/board_video_records_base.py index f7e9fcdee7d..743e71ad6d2 100644 --- a/invokeai/app/services/board_video_records/board_video_records_base.py +++ b/invokeai/app/services/board_video_records/board_video_records_base.py @@ -38,11 +38,7 @@ def get_board_for_video(self, video_name: str) -> Optional[str]: pass @abstractmethod - def get_video_count_for_board(self, board_id: str) -> int: - """Gets the number of videos for a board.""" - pass - - @abstractmethod - def get_asset_video_count_for_board(self, board_id: str) -> int: - """Gets the number of asset-category (non-'general') videos for a board.""" + def get_counts_for_board(self, board_id: str) -> tuple[int, int]: + """Gets the numbers of videos and of asset-category (non-'general') videos on a board, intermediates not + counted.""" pass diff --git a/invokeai/app/services/board_video_records/board_video_records_default.py b/invokeai/app/services/board_video_records/board_video_records_default.py new file mode 100644 index 00000000000..6e256874012 --- /dev/null +++ b/invokeai/app/services/board_video_records/board_video_records_default.py @@ -0,0 +1,38 @@ +from typing import Optional + +from invokeai.app.services.board_video_records.board_video_records_base import BoardVideoRecordStorageBase +from invokeai.app.services.image_records.image_records_common import ImageCategory +from invokeai.app.services.shared.database.database import Database + + +class BoardVideoRecordStorage(BoardVideoRecordStorageBase): + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + + def add_video_to_board(self, board_id: str, video_name: str) -> None: + self._queries.board_videos.add(board_id, video_name) + + def remove_video_from_board(self, video_name: str) -> None: + self._queries.board_videos.remove(video_name) + + def get_all_board_video_names_for_board( + self, + board_id: str, + categories: list[ImageCategory] | None, + is_intermediate: bool | None, + user_id: Optional[str] = None, + ) -> list[str]: + # `board_id` "none" means the videos on no board; admins pass user_id=None to see every account's. + return self._queries.board_videos.video_names( + None if board_id == "none" else board_id, + categories=categories, + is_intermediate=is_intermediate, + user_id=user_id, + ) + + def get_board_for_video(self, video_name: str) -> Optional[str]: + return self._queries.board_videos.board_of(video_name) + + def get_counts_for_board(self, board_id: str) -> tuple[int, int]: + return self._queries.board_videos.counts(board_id) diff --git a/invokeai/app/services/board_video_records/board_video_records_sqlite.py b/invokeai/app/services/board_video_records/board_video_records_sqlite.py deleted file mode 100644 index 449796ad42c..00000000000 --- a/invokeai/app/services/board_video_records/board_video_records_sqlite.py +++ /dev/null @@ -1,113 +0,0 @@ -import sqlite3 -from typing import Optional, cast - -from invokeai.app.services.board_video_records.board_video_records_base import BoardVideoRecordStorageBase -from invokeai.app.services.image_records.image_records_common import ImageCategory -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase - - -class SqliteBoardVideoRecordStorage(BoardVideoRecordStorageBase): - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - - def add_video_to_board(self, board_id: str, video_name: str) -> None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - INSERT INTO board_videos (board_id, video_name) - VALUES (?, ?) - ON CONFLICT (video_name) DO UPDATE SET board_id = ?; - """, - (board_id, video_name, board_id), - ) - - def remove_video_from_board(self, video_name: str) -> None: - with self._db.transaction() as cursor: - cursor.execute( - "DELETE FROM board_videos WHERE video_name = ?;", - (video_name,), - ) - - def get_all_board_video_names_for_board( - self, - board_id: str, - categories: list[ImageCategory] | None, - is_intermediate: bool | None, - user_id: Optional[str] = None, - ) -> list[str]: - with self._db.transaction() as cursor: - params: list[str | bool] = [] - stmt = """ - SELECT videos.video_name - FROM videos - LEFT JOIN board_videos ON board_videos.video_name = videos.video_name - WHERE 1=1 - """ - if board_id == "none": - stmt += " AND board_videos.board_id IS NULL " - else: - stmt += " AND board_videos.board_id = ? " - params.append(board_id) - - if categories is not None: - category_strings = [c.value for c in set(categories)] - placeholders = ",".join("?" * len(category_strings)) - stmt += f" AND videos.video_category IN ( {placeholders} ) " - for c in category_strings: - params.append(c) - - if is_intermediate is not None: - stmt += " AND videos.is_intermediate = ? " - params.append(is_intermediate) - - if user_id is not None: - stmt += " AND videos.user_id = ? " - params.append(user_id) - - stmt += ";" - cursor.execute(stmt, params) - result = cast(list[sqlite3.Row], cursor.fetchall()) - return [r[0] for r in result] - - def get_board_for_video(self, video_name: str) -> Optional[str]: - with self._db.transaction() as cursor: - cursor.execute( - "SELECT board_id FROM board_videos WHERE video_name = ?;", - (video_name,), - ) - result = cursor.fetchone() - if result is None: - return None - return cast(str, result[0]) - - def get_video_count_for_board(self, board_id: str) -> int: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT COUNT(*) - FROM board_videos - INNER JOIN videos ON board_videos.video_name = videos.video_name - WHERE videos.is_intermediate = FALSE - AND board_videos.board_id = ?; - """, - (board_id,), - ) - count = cast(int, cursor.fetchone()[0]) - return count - - def get_asset_video_count_for_board(self, board_id: str) -> int: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT COUNT(*) - FROM board_videos - INNER JOIN videos ON board_videos.video_name = videos.video_name - WHERE videos.is_intermediate = FALSE - AND videos.video_category != 'general' - AND board_videos.board_id = ?; - """, - (board_id,), - ) - count = cast(int, cursor.fetchone()[0]) - return count diff --git a/invokeai/app/services/boards/boards_base.py b/invokeai/app/services/boards/boards_base.py index 4f1d36fb8ff..90f7fcefe87 100644 --- a/invokeai/app/services/boards/boards_base.py +++ b/invokeai/app/services/boards/boards_base.py @@ -2,8 +2,7 @@ from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecordOrderBy from invokeai.app.services.boards.boards_common import BoardDTO -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection class BoardServiceABC(ABC): diff --git a/invokeai/app/services/boards/boards_default.py b/invokeai/app/services/boards/boards_default.py index d8b6a18b696..1684083467f 100644 --- a/invokeai/app/services/boards/boards_default.py +++ b/invokeai/app/services/boards/boards_default.py @@ -4,8 +4,7 @@ from invokeai.app.services.boards.boards_base import BoardServiceABC from invokeai.app.services.boards.boards_common import BoardDTO, board_record_to_dto from invokeai.app.services.invoker import Invoker -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection class BoardService(BoardServiceABC): @@ -41,10 +40,8 @@ def _resolve_cover(self, board_id: str) -> tuple[Optional[str], Optional[str]]: def _get_counts(self, board_id: str) -> tuple[int, int, int, int]: """Return ``(image_count, video_count, asset_count, asset_video_count)`` for a board.""" - image_count = self.__invoker.services.board_image_records.get_image_count_for_board(board_id) - asset_count = self.__invoker.services.board_image_records.get_asset_count_for_board(board_id) - video_count = self.__invoker.services.board_video_records.get_video_count_for_board(board_id) - asset_video_count = self.__invoker.services.board_video_records.get_asset_video_count_for_board(board_id) + image_count, asset_count = self.__invoker.services.board_image_records.get_counts_for_board(board_id) + video_count, asset_video_count = self.__invoker.services.board_video_records.get_counts_for_board(board_id) return image_count, video_count, asset_count, asset_video_count def create( diff --git a/invokeai/app/services/client_state_persistence/client_state_persistence_default.py b/invokeai/app/services/client_state_persistence/client_state_persistence_default.py new file mode 100644 index 00000000000..01abf3758d7 --- /dev/null +++ b/invokeai/app/services/client_state_persistence/client_state_persistence_default.py @@ -0,0 +1,58 @@ +from invokeai.app.services.client_state_persistence.client_state_persistence_base import ClientStatePersistenceABC +from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.locks import DatabaseLock +from invokeai.app.services.shared.media_references import extract_media_references_from_json + + +class ClientStatePersistence(ClientStatePersistenceABC): + """ + Client state persistence on the application database. + This class stores client state data per user to prevent data leakage between users. + """ + + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + + def start(self, invoker: Invoker) -> None: + self._invoker = invoker + + def set_by_key(self, user_id: str, key: str, value: str) -> str: + # The legacy editor keeps its canvas layers and reference images only here, as + # intermediates; indexing them keeps cleanup from treating them as unused. Parsed before + # the transaction, which on SQLite holds the lock every other writer shares. + references = extract_media_references_from_json(value) + + def save(q: Queries) -> None: + # Shared with every write that makes media protected, exclusive for the intermediates cleanup's check and + # delete: the media this names cannot be deleted between that check and this commit. + q.locks.acquire(DatabaseLock.MEDIA_PROTECTION, shared=True) + q.client_state.set(user_id, key, value) + q.media_references.replace(owner_kind="client_state", user_id=user_id, owner_id=key, references=references) + + self._queries.run(save) + return value + + def get_by_key(self, user_id: str, key: str) -> str | None: + return self._queries.client_state.get(user_id, key) + + def get_keys_by_prefix(self, user_id: str, prefix: str) -> list[str]: + return self._queries.client_state.keys_with_prefix(user_id, prefix) + + def delete_by_key(self, user_id: str, key: str) -> None: + def delete_value(q: Queries) -> None: + # Only references of a value this transaction deleted: on a server, a value another transaction writes + # for the first time meanwhile keeps the references it writes. + if q.client_state.delete(user_id, key): + q.media_references.delete(owner_kind="client_state", user_id=user_id, owner_id=key) + + self._queries.run(delete_value) + + def delete(self, user_id: str) -> None: + def delete_all(q: Queries) -> None: + keys = q.client_state.delete_all(user_id) + q.media_references.delete_many(owner_kind="client_state", user_id=user_id, owner_ids=keys) + + self._queries.run(delete_all) diff --git a/invokeai/app/services/client_state_persistence/client_state_persistence_sqlite.py b/invokeai/app/services/client_state_persistence/client_state_persistence_sqlite.py deleted file mode 100644 index b27237c2e7e..00000000000 --- a/invokeai/app/services/client_state_persistence/client_state_persistence_sqlite.py +++ /dev/null @@ -1,96 +0,0 @@ -from invokeai.app.services.client_state_persistence.client_state_persistence_base import ClientStatePersistenceABC -from invokeai.app.services.invoker import Invoker -from invokeai.app.services.shared.media_references import ( - delete_media_references, - extract_media_references_from_json, - replace_media_references, -) -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase - - -class ClientStatePersistenceSqlite(ClientStatePersistenceABC): - """ - SQLite implementation for client state persistence. - This class stores client state data per user to prevent data leakage between users. - """ - - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - - def start(self, invoker: Invoker) -> None: - self._invoker = invoker - - def set_by_key(self, user_id: str, key: str, value: str) -> str: - # The legacy editor keeps its canvas layers and reference images only here, as - # intermediates; indexing them keeps cleanup from treating them as unused. Parsed before - # taking the database lock, which every other writer shares. - references = extract_media_references_from_json(value) - with self._db.transaction() as cursor: - cursor.execute( - """ - INSERT INTO client_state (user_id, key, value) - VALUES (?, ?, ?) - ON CONFLICT(user_id, key) DO UPDATE - SET value = excluded.value; - """, - (user_id, key, value), - ) - replace_media_references( - cursor, owner_kind="client_state", user_id=user_id, owner_id=key, references=references - ) - - return value - - def get_by_key(self, user_id: str, key: str) -> str | None: - with self._db.transaction() as cursor: - cursor.execute( - """ - SELECT value FROM client_state - WHERE user_id = ? AND key = ? - """, - (user_id, key), - ) - row = cursor.fetchone() - if row is None: - return None - return row[0] - - def get_keys_by_prefix(self, user_id: str, prefix: str) -> list[str]: - # Escape LIKE wildcards (%, _) and the escape char itself so callers can pass - # arbitrary strings as a literal prefix without accidental pattern matching. - escaped_prefix = prefix.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") - with self._db.transaction() as cursor: - cursor.execute( - """ - SELECT key FROM client_state - WHERE user_id = ? AND key LIKE ? ESCAPE '\\' - ORDER BY updated_at DESC - """, - (user_id, f"{escaped_prefix}%"), - ) - return [row[0] for row in cursor.fetchall()] - - def delete_by_key(self, user_id: str, key: str) -> None: - with self._db.transaction() as cursor: - cursor.execute( - """ - DELETE FROM client_state - WHERE user_id = ? AND key = ? - """, - (user_id, key), - ) - delete_media_references(cursor, owner_kind="client_state", user_id=user_id, owner_id=key) - - def delete(self, user_id: str) -> None: - with self._db.transaction() as cursor: - cursor.execute( - """ - DELETE FROM client_state - WHERE user_id = ? - """, - (user_id,), - ) - cursor.execute( - "DELETE FROM media_references WHERE owner_kind = 'client_state' AND user_id = ?;", (user_id,) - ) diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index 29ce47a5b78..9e6125f6ba2 100644 --- a/invokeai/app/services/config/config_default.py +++ b/invokeai/app/services/config/config_default.py @@ -92,6 +92,7 @@ class InvokeAIAppConfig(BaseSettings): download_cache_dir: Path to the directory that contains dynamically downloaded models. legacy_conf_dir: Path to directory of legacy checkpoint config files. db_dir: Path to InvokeAI databases directory. + db_url: URL of a MySQL 8.4+ or MariaDB 10.11+ database to use instead of the SQLite database in `db_dir`, e.g. `mariadb+pymysql://invokeai:password@db.example/invokeai`. Needs the `mysql` extra. One InvokeAI process uses a database at a time. Read at startup only. db_synchronous: SQLite durability setting. `full`, the default and what InvokeAI has always used, flushes every commit to disk. `normal` acknowledges commits without waiting for that flush - measured at roughly 12x shorter commits on an SSD - and cannot corrupt the database under WAL, which is why it is refused, with a warning, when WAL is unavailable for the database file. What `normal` gives up is the most recent transactions on a power loss or OS crash: a just-written image record or queue status, not the image file itself.
Valid values: `full`, `normal` outputs_dir: Path to directory for outputs. image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash` @@ -203,6 +204,7 @@ class InvokeAIAppConfig(BaseSettings): download_cache_dir: Path = Field(default=Path("models/.download_cache"), description="Path to the directory that contains dynamically downloaded models.") legacy_conf_dir: Path = Field(default=Path("configs"), description="Path to directory of legacy checkpoint config files.") db_dir: Path = Field(default=Path("databases"), description="Path to InvokeAI databases directory.") + db_url: Optional[str] = Field(default=None, description="URL of a MySQL 8.4+ or MariaDB 10.11+ database to use instead of the SQLite database in `db_dir`, e.g. `mariadb+pymysql://invokeai:password@db.example/invokeai`. Needs the `mysql` extra. One InvokeAI process uses a database at a time. Read at startup only.") db_synchronous: DB_SYNCHRONOUS = Field(default="full", description="SQLite durability setting. `full`, the default and what InvokeAI has always used, flushes every commit to disk. `normal` acknowledges commits without waiting for that flush - measured at roughly 12x shorter commits on an SSD - and cannot corrupt the database under WAL, which is why it is refused, with a warning, when WAL is unavailable for the database file. What `normal` gives up is the most recent transactions on a power loss or OS crash: a just-written image record or queue status, not the image file itself.") outputs_dir: Path = Field(default=Path("outputs"), description="Path to directory for outputs.") fonts_dir: Path = Field(default=Path("fonts"), description="Path to directory for custom fonts.") @@ -737,6 +739,17 @@ def ensure_fonts_dir(fonts_path: Path) -> None: logger.warning("Unable to initialize fonts directory at %s", fonts_path, exc_info=True) +def load_config_from_root(root: Path) -> InvokeAIAppConfig: + """The config of the install at `root`, from the environment and its `invokeai.yaml` as `get_config` reads them, + for the maintenance scripts. Unlike `get_config` it writes no example or default config file and copies no + legacy model configs; like it, it migrates an `invokeai.yaml` of an older schema.""" + config = InvokeAIAppConfig() + config._root = root + if config.config_file_path.exists(): + config.update_config(load_and_migrate_config(config.config_file_path), clobber=False) + return config + + @lru_cache(maxsize=1) def get_config() -> InvokeAIAppConfig: """Get the global singleton app config. diff --git a/invokeai/app/services/fonts/__init__.py b/invokeai/app/services/fonts/__init__.py index cd0dc1e5645..cf442974a6a 100644 --- a/invokeai/app/services/fonts/__init__.py +++ b/invokeai/app/services/fonts/__init__.py @@ -1,43 +1 @@ """Font library services.""" - -from invokeai.app.services.fonts.fonts_common import ( - FontAxis, - FontInstance, - FontRecord, - FontScope, - FontSource, - FontUploadResult, - FontValidationResult, -) -from invokeai.app.services.fonts.fonts_default import ( - FontChangedError, - FontDeleteForbiddenError, - FontForbiddenError, - FontInstanceError, - FontNotFoundError, - FontQuotaExceededError, - FontService, - FontServiceError, - FontStorageError, - FontValidationError, -) - -__all__ = [ - "FontAxis", - "FontChangedError", - "FontDeleteForbiddenError", - "FontForbiddenError", - "FontInstance", - "FontInstanceError", - "FontNotFoundError", - "FontQuotaExceededError", - "FontRecord", - "FontScope", - "FontService", - "FontServiceError", - "FontSource", - "FontStorageError", - "FontUploadResult", - "FontValidationError", - "FontValidationResult", -] diff --git a/invokeai/app/services/fonts/fonts_default.py b/invokeai/app/services/fonts/fonts_default.py index a6b53fe678d..385fe33acba 100644 --- a/invokeai/app/services/fonts/fonts_default.py +++ b/invokeai/app/services/fonts/fonts_default.py @@ -1,4 +1,4 @@ -"""SQLite-backed custom font library. +"""The custom font library. The service keeps the configured ``fonts_dir`` as a read-only, shared source and stores uploads in an application-owned directory. Font files are addressed by opaque IDs and @@ -36,7 +36,8 @@ FontUploadResult, FontValidationResult, ) -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries.fonts import MAX_SOURCE_PATH_LENGTH, DirectoryFont, UploadedFont SUPPORTED_FONT_EXTENSIONS = frozenset({".ttf", ".otf", ".woff", ".woff2"}) FONT_MEDIA_TYPES = { @@ -54,11 +55,6 @@ MAX_INSTANCE_CACHE_BYTES = 64 * 1024 * 1024 INSTANCE_JOB_TIMEOUT_SECONDS = 30.0 -_FONT_COLUMNS = """ - id, owner_id, family, label, style, weight, content_hash, scope, source, filename, - byte_size, axes_json, instances_json, storage_path, source_path -""" - class FontServiceError(RuntimeError): """Base class for errors raised by the font service.""" @@ -246,36 +242,6 @@ def _parse_font_worker(filename: str, data: bytes, max_bytes: int) -> FontValida return _parse_font_bytes(filename, data, max_bytes=max_bytes) -def _json_axes(axes: tuple[FontAxis, ...]) -> str: - return json.dumps([axis.__dict__ for axis in axes], separators=(",", ":"), sort_keys=True) - - -def _json_instances(instances: tuple[FontInstance, ...]) -> str: - return json.dumps([instance.__dict__ for instance in instances], separators=(",", ":"), sort_keys=True) - - -def _record_from_row(row: Any) -> FontRecord: - axes = tuple(FontAxis(**axis) for axis in json.loads(row[11])) - instances = tuple(FontInstance(**instance) for instance in json.loads(row[12])) - return FontRecord( - id=row[0], - family=row[2], - label=row[3], - style=row[4], - weight=int(row[5]), - content_hash=row[6], - scope=FontScope(row[7]), - source=FontSource(row[8]), - filename=row[9], - byte_size=int(row[10]), - axes=axes, - instances=instances, - owner_id=row[1], - storage_path=row[13], - source_path=row[14], - ) - - def _validation_from_record(record: FontRecord, filename: str) -> FontValidationResult: """Reuse indexed metadata when a directory file's content hash is unchanged.""" return FontValidationResult( @@ -291,6 +257,23 @@ def _validation_from_record(record: FontRecord, filename: str) -> FontValidation ) +def _indexed_as(record: FontRecord | None, font: DirectoryFont) -> bool: + """Whether the catalog already holds the directory font as found.""" + return ( + record is not None + and record.filename == font.filename + and record.source_path == font.source_path + and record.family == font.family + and record.label == font.label + and record.style == font.style + and record.weight == font.weight + and record.content_hash == font.content_hash + and record.byte_size == font.byte_size + and record.axes == font.axes + and record.instances == font.instances + ) + + def _read_bounded_file(path: Path, max_bytes: int) -> bytes: """Read at most ``max_bytes`` and reject a file that grows during the read.""" if path.stat().st_size > max_bytes: @@ -338,7 +321,7 @@ class FontService: def __init__( self, - db: SqliteDatabase, + database: Database, fonts_dir: Path, storage_dir: Path, logger: logging.Logger | None = None, @@ -356,7 +339,7 @@ def __init__( or storage_absolute.is_relative_to(fonts_absolute) ): raise ValueError("Font source and managed storage directories must not overlap") - self._db = db + self._queries = database.queries self._fonts_dir = fonts_dir self._storage_dir = storage_dir self._logger = logger or logging.getLogger(__name__) @@ -376,6 +359,8 @@ def __init__( # cleanup sweep must not observe an upload between writing its file and committing # its row, or it could remove a live upload. self._managed_storage_lock = threading.RLock() + # One scan at a time: two would race to insert the same new directory font. + self._rescan_lock = threading.Lock() self._instance_cache: OrderedDict[tuple[str, tuple[tuple[str, float], ...]], bytes] = OrderedDict() self._instance_cache_bytes = 0 self._cache_lock = threading.RLock() @@ -432,66 +417,46 @@ def upload( duplicate: FontRecord | None = None storage_path: str | None = None font_id = f"font_{uuid4().hex}" + # The duplicate and quota checks, the file and the row are published together: every upload, and the + # orphaned-file sweep, run under the managed-storage lock (one process serves the database). The file is + # written outside any transaction, so that no database transaction waits on its fsync. with self._managed_storage_lock: - try: - with self._db.transaction() as cursor: - cursor.execute("BEGIN IMMEDIATE;") - if scope == FontScope.PRIVATE: - cursor.execute( - f"SELECT {_FONT_COLUMNS} FROM fonts WHERE source = 'uploaded' AND scope = ? " - "AND owner_id = ? AND content_hash = ? LIMIT 1;", - (scope.value, owner_id, validation.content_hash), - ) - else: - cursor.execute( - f"SELECT {_FONT_COLUMNS} FROM fonts WHERE source = 'uploaded' AND scope = ? " - "AND owner_id IS NULL AND content_hash = ? LIMIT 1;", - (scope.value, validation.content_hash), - ) - row = cursor.fetchone() - if row is not None: - duplicate = _record_from_row(row) - else: - cursor.execute( - "SELECT COALESCE(SUM(byte_size), 0) FROM fonts WHERE source = 'uploaded' " - "AND scope = ? AND (owner_id IS ? OR (? = 'shared' AND owner_id IS NULL));", - (scope.value, owner_id, scope.value), - ) - current_bytes = int(cursor.fetchone()[0]) - if current_bytes + validation.byte_size > self._max_library_bytes: - raise FontQuotaExceededError("Font library storage quota exceeded") - - extension = Path(filename).suffix.lower() - storage_path = f"{font_id}{extension}" - target = _safe_relative_path(Path(storage_path), self._storage_dir) - self._write_atomic(target, data) - cursor.execute( - """--sql - INSERT INTO fonts ( - id, owner_id, scope, source, filename, storage_path, source_path, - family, label, style, weight, content_hash, byte_size, axes_json, instances_json - ) VALUES (?, ?, ?, 'uploaded', ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?); - """, - ( - font_id, - owner_id, - scope.value, - validation.filename, - storage_path, - validation.family, - validation.label, - validation.style, - validation.weight, - validation.content_hash, - validation.byte_size, - _json_axes(validation.axes), - _json_instances(validation.instances), - ), + duplicate, used_bytes = self._queries.fonts.upload_state( + owner_id=owner_id, content_hash=validation.content_hash + ) + if duplicate is None: + if used_bytes + validation.byte_size > self._max_library_bytes: + raise FontQuotaExceededError("Font library storage quota exceeded") + storage_path = f"{font_id}{Path(filename).suffix.lower()}" + try: + self._write_atomic(_safe_relative_path(Path(storage_path), self._storage_dir), data) + self._queries.fonts.insert_upload( + UploadedFont( + id=font_id, + owner_id=owner_id, + scope=scope, + filename=validation.filename, + storage_path=storage_path, + family=validation.family, + label=validation.label, + style=validation.style, + weight=validation.weight, + content_hash=validation.content_hash, + byte_size=validation.byte_size, + axes=validation.axes, + instances=validation.instances, ) - except Exception: - if storage_path is not None: - self._unlink_managed_file(storage_path) - raise + ) + except Exception: + # A lost connection can fail a commit that took effect: keep the file of a row that exists, or + # that cannot be checked; the startup sweep removes it if no row names it. + try: + indexed = self._queries.fonts.get(font_id) is not None + except Exception: + indexed = True + if not indexed: + self._unlink_managed_file(storage_path) + raise if duplicate is not None: return FontUploadResult(font=duplicate, created=False) @@ -504,35 +469,37 @@ def upload( def rescan_directory(self) -> int: """Index valid files in the configured directory and remove stale directory rows.""" + with self._rescan_lock: + return self._rescan_directory() + + def _rescan_directory(self) -> int: root = self._fonts_dir if not root.exists() or not root.is_dir() or root.is_symlink(): if root.is_symlink(): self._logger.warning("Skipping custom fonts directory %s: symlinks are not supported", root) - with self._db.transaction() as cursor: - cursor.execute("DELETE FROM fonts WHERE source = 'directory';") - removed = cursor.rowcount - if removed: + if self._queries.fonts.delete_directory_fonts(): self._bump_revision() return 0 - with self._db.transaction() as cursor: - cursor.execute(f"SELECT {_FONT_COLUMNS} FROM fonts WHERE source = 'directory';") - existing_records = { - record.source_path: record - for record in (_record_from_row(row) for row in cursor.fetchall()) - if record.source_path is not None - } + existing_records = { + record.source_path: record + for record in self._queries.fonts.directory_fonts() + if record.source_path is not None + } - discovered: list[tuple[str, FontValidationResult, str]] = [] + discovered: list[DirectoryFont] = [] for path in sorted(root.rglob("*")): if self._path_has_symlink_component(path, root): self._logger.warning("Skipping font path %s: symlinks are not supported", path) continue if not path.is_file() or path.suffix.lower() not in SUPPORTED_FONT_EXTENSIONS: continue + relative = path.relative_to(root).as_posix() + if len(relative) > MAX_SOURCE_PATH_LENGTH: + self._logger.warning("Skipping font file %s: its path is longer than the catalog holds", path) + continue try: data = _read_bounded_file(path, MAX_INDEXED_FONT_BYTES) - relative = path.relative_to(root).as_posix() content_hash = hashlib.sha256(data).hexdigest() existing = existing_records.get(relative) if existing is not None and existing.byte_size == len(data) and existing.content_hash == content_hash: @@ -548,74 +515,27 @@ def rescan_directory(self) -> int: continue # The path-derived ID lets a project distinguish two same-named files and remains # stable when a directory file is replaced. Its content hash detects that replacement. - font_id = f"directory_{hashlib.sha256(relative.encode('utf-8')).hexdigest()[:32]}" - discovered.append((font_id, metadata, relative)) - - discovered_ids = {font_id for font_id, _metadata, _relative in discovered} - existing_by_id = {record.id: record for record in existing_records.values()} - changed = set(existing_by_id) != discovered_ids - if not changed: - for font_id, metadata, relative in discovered: - existing = existing_by_id[font_id] - if ( - existing.filename != metadata.filename - or existing.source_path != relative - or existing.family != metadata.family - or existing.label != metadata.label - or existing.style != metadata.style - or existing.weight != metadata.weight - or existing.content_hash != metadata.content_hash - or existing.byte_size != metadata.byte_size - or existing.axes != metadata.axes - or existing.instances != metadata.instances - ): - changed = True - break - with self._db.transaction() as cursor: - for font_id, metadata, relative in discovered: - cursor.execute( - """--sql - INSERT INTO fonts ( - id, owner_id, scope, source, filename, storage_path, source_path, - family, label, style, weight, content_hash, byte_size, axes_json, instances_json - ) VALUES (?, NULL, 'shared', 'directory', ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?) - ON CONFLICT(id) DO UPDATE SET - filename = excluded.filename, - source_path = excluded.source_path, - family = excluded.family, - label = excluded.label, - style = excluded.style, - weight = excluded.weight, - content_hash = excluded.content_hash, - byte_size = excluded.byte_size, - axes_json = excluded.axes_json, - instances_json = excluded.instances_json; - """, - ( - font_id, - metadata.filename, - relative, - metadata.family, - metadata.label, - metadata.style, - metadata.weight, - metadata.content_hash, - metadata.byte_size, - _json_axes(metadata.axes), - _json_instances(metadata.instances), - ), - ) - - if discovered_ids: - placeholders = ",".join("?" for _ in discovered_ids) - cursor.execute( - f"DELETE FROM fonts WHERE source = 'directory' AND id NOT IN ({placeholders});", - tuple(discovered_ids), + discovered.append( + DirectoryFont( + id=f"directory_{hashlib.sha256(relative.encode('utf-8')).hexdigest()[:32]}", + filename=metadata.filename, + source_path=relative, + family=metadata.family, + label=metadata.label, + style=metadata.style, + weight=metadata.weight, + content_hash=metadata.content_hash, + byte_size=metadata.byte_size, + axes=metadata.axes, + instances=metadata.instances, ) - else: - cursor.execute("DELETE FROM fonts WHERE source = 'directory';") + ) - if changed: + existing_by_id = {record.id: record for record in existing_records.values()} + changed_fonts = [font for font in discovered if not _indexed_as(existing_by_id.get(font.id), font)] + discovered_ids = frozenset(font.id for font in discovered) + if changed_fonts or set(existing_by_id) != discovered_ids: + self._queries.fonts.replace_directory_fonts(changed_fonts, discovered_ids) self._invalidate_instance_cache() self._bump_revision() return len(discovered) @@ -636,43 +556,15 @@ def list( if scope not in (FontScope.ALL, FontScope.PRIVATE, FontScope.SHARED): raise ValueError("Invalid font scope") - conditions = [ - "(source = 'directory' OR scope = 'shared' OR (scope = 'private' AND owner_id = ?))", - ] - params: list[object] = [user_id] - if scope == FontScope.PRIVATE: - conditions.append("source = 'uploaded' AND scope = 'private' AND owner_id = ?") - params.append(user_id) - elif scope == FontScope.SHARED: - conditions.append("(source = 'directory' OR (source = 'uploaded' AND scope = 'shared'))") - if content_hash is not None: - if re.fullmatch(r"[a-f0-9]{64}", content_hash) is None: - raise ValueError("content_hash must be a lowercase SHA-256 hash") - conditions.append("content_hash = ?") - params.append(content_hash) - if search: - conditions.append( - "(family LIKE ? COLLATE NOCASE OR label LIKE ? COLLATE NOCASE OR filename LIKE ? COLLATE NOCASE)" - ) - term = f"%{search}%" - params.extend((term, term, term)) - where = " AND ".join(conditions) - with self._db.transaction() as cursor: - cursor.execute(f"SELECT COUNT(*) FROM fonts WHERE {where};", tuple(params)) - total = int(cursor.fetchone()[0]) - cursor.execute( - f"SELECT {_FONT_COLUMNS} FROM fonts WHERE {where} ORDER BY family COLLATE NOCASE, label COLLATE NOCASE, id LIMIT ? OFFSET ?;", - (*params, limit, offset), - ) - rows = cursor.fetchall() - return [_record_from_row(row) for row in rows], total + if content_hash is not None and re.fullmatch(r"[a-f0-9]{64}", content_hash) is None: + raise ValueError("content_hash must be a lowercase SHA-256 hash") + return self._queries.fonts.page( + user_id=user_id, scope=scope, content_hash=content_hash, search=search, offset=offset, limit=limit + ) def get(self, font_id: str) -> FontRecord | None: """Get an indexed font without applying account visibility.""" - with self._db.transaction() as cursor: - cursor.execute(f"SELECT {_FONT_COLUMNS} FROM fonts WHERE id = ?;", (font_id,)) - row = cursor.fetchone() - return _record_from_row(row) if row is not None else None + return self._queries.fonts.get(font_id) def get_accessible(self, *, user_id: str, font_id: str) -> FontRecord: """Get a font visible to the account, hiding private-resource existence.""" @@ -717,10 +609,8 @@ def delete(self, *, user_id: str, font_id: str, is_admin: bool = False) -> None: if record.scope == FontScope.SHARED and not is_admin: raise FontForbiddenError("Only administrators may delete a shared font") - with self._db.transaction() as cursor: - cursor.execute("DELETE FROM fonts WHERE id = ? AND source = 'uploaded';", (font_id,)) - if cursor.rowcount != 1: - raise FontNotFoundError("Font not found") + if not self._queries.fonts.delete_upload(font_id): + raise FontNotFoundError("Font not found") # A variable-font instance is derived from the source bytes. Drop every derived # value before unlinking the source so a later request cannot serve a deleted font # from memory while the filesystem cleanup is in progress. @@ -738,9 +628,7 @@ def cleanup_orphaned_files(self) -> None: self._invalidate_instance_cache() if not self._storage_dir.exists() or not self._storage_dir.is_dir() or self._storage_dir.is_symlink(): return - with self._db.transaction() as cursor: - cursor.execute("SELECT storage_path FROM fonts WHERE source = 'uploaded' AND storage_path IS NOT NULL;") - referenced = {row[0] for row in cursor.fetchall()} + referenced = self._queries.fonts.storage_paths() for path in self._storage_dir.iterdir(): if not path.is_file() or path.is_symlink(): continue @@ -761,13 +649,7 @@ def cleanup_orphaned_files(self) -> None: def prepare_user_cleanup(self, user_id: str) -> tuple[str, ...]: """Capture a user's private managed paths before its account row is deleted.""" with self._managed_storage_lock: - with self._db.transaction() as cursor: - cursor.execute( - "SELECT storage_path FROM fonts " - "WHERE source = 'uploaded' AND scope = 'private' AND owner_id = ? AND storage_path IS NOT NULL;", - (user_id,), - ) - storage_paths = tuple(row[0] for row in cursor.fetchall()) + storage_paths = self._queries.fonts.private_storage_paths(user_id) # A caller may delete the account through a separate service, which cascades # rows without notifying this process. Purge derived values before that race. self._invalidate_instance_cache() diff --git a/invokeai/app/services/gallery/gallery_base.py b/invokeai/app/services/gallery/gallery_base.py index 78047a5f8de..b82e71deeed 100644 --- a/invokeai/app/services/gallery/gallery_base.py +++ b/invokeai/app/services/gallery/gallery_base.py @@ -10,8 +10,7 @@ GalleryItemNamesResult, ) from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection from invokeai.app.services.virtual_boards.virtual_boards_common import VirtualSubBoardDTO diff --git a/invokeai/app/services/gallery/gallery_default.py b/invokeai/app/services/gallery/gallery_default.py index 5bd050380c2..9ed9d20fdc8 100644 --- a/invokeai/app/services/gallery/gallery_default.py +++ b/invokeai/app/services/gallery/gallery_default.py @@ -1,5 +1,4 @@ -import sqlite3 -from typing import Optional, Union, cast +from typing import Any, Optional, Sequence from invokeai.app.services.gallery.gallery_base import GalleryServiceABC from invokeai.app.services.gallery.gallery_common import ( @@ -13,26 +12,21 @@ ) from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin from invokeai.app.services.invoker import Invoker -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.video_records.video_records_common import MEDIA_ORIGIN_SQL_EXPR, coerce_media_origin +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries.gallery import Filters +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection +from invokeai.app.services.video_records.video_records_common import coerce_media_origin from invokeai.app.services.virtual_boards.virtual_boards_common import VirtualSubBoardDTO -class SqliteGalleryService(GalleryServiceABC): - """Implements a polymorphic gallery via UNION ALL across the `images` and `videos` tables. - - Filters are applied identically on each half. The two halves expose a common column set so - the result is shape-compatible (a literal `kind` discriminator + a `name` alias + duration/fps - that are NULL for images). - """ +class GalleryService(GalleryServiceABC): + """A gallery of images and videos as one list (see `queries/gallery.py`).""" __invoker: Invoker - def __init__(self, db: SqliteDatabase) -> None: + def __init__(self, database: Database) -> None: super().__init__() - self._db = db + self._queries = database.queries def start(self, invoker: Invoker) -> None: self.__invoker = invoker @@ -54,8 +48,7 @@ def list_items( created_to: Optional[str] = None, starred: Optional[bool] = None, ) -> OffsetPaginatedResults[GalleryItem]: - image_half, image_params, image_count_query = self._build_half( - kind="image", + filters = Filters( origin=origin, categories=categories, is_intermediate=is_intermediate, @@ -67,63 +60,15 @@ def list_items( created_to=created_to, starred=starred, ) - video_half, video_params, video_count_query = self._build_half( - kind="video", - origin=origin, - categories=categories, - is_intermediate=is_intermediate, - board_id=board_id, - search_term=search_term, - user_id=user_id, - is_admin=is_admin, - created_from=created_from, - created_to=created_to, - starred=starred, - ) - - order_clause = self._build_order_clause(starred_first, order_dir) - - # `media_origin` is joined back onto the CHOSEN page rather than selected inside the - # union halves. It is projected out of `videos.metadata`, a blob that lives on - # overflow pages, and the halves feed a sorter -- so selecting it there reads and - # JSON-parses every video row in the library before LIMIT, on a connection held - # behind a process-wide lock. Measured at 20k videos with 4 KB metadata that is - # ~+47ms on every gallery page; joining after LIMIT touches at most `limit` rows. - # (`ORDER BY` is repeated outside: a join over an ordered subquery does not preserve - # its order. Both clauses reference only columns the page already carries.) - union_query = f"""--sql - SELECT page.*, {MEDIA_ORIGIN_SQL_EXPR} - FROM ( - SELECT * FROM ( - {image_half} - UNION ALL - {video_half} - ) - {order_clause} - LIMIT ? OFFSET ? - ) AS page - LEFT JOIN videos ON page.kind = 'video' AND videos.video_name = page.name - {order_clause} - ; - """ - - with self._db.transaction() as cursor: - cursor.execute(union_query, image_params + video_params + [limit, offset]) - rows = cast(list[sqlite3.Row], cursor.fetchall()) - - cursor.execute(image_count_query, image_params) - image_count = cast(int, cursor.fetchone()[0]) - cursor.execute(video_count_query, video_params) - video_count = cast(int, cursor.fetchone()[0]) - - urls = self.__invoker.services.urls - items = [self._row_to_item(row, urls) for row in rows] - return OffsetPaginatedResults[GalleryItem]( - items=items, + rows, total = self._queries.gallery.page( + filters, offset=offset, limit=limit, - total=image_count + video_count, + starred_first=starred_first, + descending=order_dir == SQLiteDirection.Descending, ) + items = [self._to_item(row) for row in rows] + return OffsetPaginatedResults[GalleryItem](items=items, offset=offset, limit=limit, total=total) def get_item_location( self, @@ -142,8 +87,7 @@ def get_item_location( starred: Optional[bool] = None, ) -> Optional[GalleryItemLocation]: """Finds item's exact rank without building or returning the full name list.""" - image_half, image_params, _ = self._build_half( - kind="image", + filters = Filters( origin=origin, categories=categories, is_intermediate=is_intermediate, @@ -151,133 +95,27 @@ def get_item_location( search_term=search_term, user_id=user_id, is_admin=is_admin, - names_only=True, created_from=created_from, created_to=created_to, starred=starred, ) - video_half, video_params, _ = self._build_half( - kind="video", - origin=origin, - categories=categories, - is_intermediate=is_intermediate, - board_id=board_id, - search_term=search_term, - user_id=user_id, - is_admin=is_admin, - names_only=True, - created_from=created_from, - created_to=created_to, - starred=starred, - ) - - # Keep ranking expressions identical to paged listing order. Window rank avoids - # materializing every matching name in Python while still computing exact index and total. - order_clause = self._build_order_clause(starred_first=False, order_dir=order_dir) - query = f"""--sql - SELECT kind, name, item_index, total - FROM ( - SELECT - kind, - name, - ROW_NUMBER() OVER ({order_clause}) - 1 AS item_index, - COUNT(*) OVER () AS total - FROM ( - {image_half} - UNION ALL - {video_half} - ) + row = self._queries.gallery.location( + filters, kind=kind.value, name=name, descending=order_dir == SQLiteDirection.Descending ) - WHERE kind = ? AND name = ? - LIMIT 1; - """ - - with self._db.transaction() as cursor: - cursor.execute(query, image_params + video_params + [kind.value, name]) - row = cursor.fetchone() - if row is None: return None + found_kind, found_name, index, total = row return GalleryItemLocation( - kind=GalleryItemKind(row["kind"]), - name=row["name"], - index=row["item_index"], - total=row["total"], - ) - - def _query_name_rows( - self, - starred_first: bool, - order_dir: SQLiteDirection, - origin: Optional[ResourceOrigin], - categories: Optional[list[ImageCategory]], - is_intermediate: Optional[bool], - board_id: Optional[str], - search_term: Optional[str], - user_id: Optional[str], - is_admin: bool, - created_date: Optional[str], - created_from: Optional[str], - created_to: Optional[str], - starred: Optional[bool], - ) -> tuple[list[sqlite3.Row], int]: - """Runs the ordered name query and returns its rows plus the starred count. - - Shared by both name-list shapes so the deprecated `(kind, name)` variant and the flat - one can never drift apart in ordering or filtering. - """ - image_half, image_params, _ = self._build_half( - kind="image", - origin=origin, - categories=categories, - is_intermediate=is_intermediate, - board_id=board_id, - search_term=search_term, - user_id=user_id, - is_admin=is_admin, - names_only=True, - created_date=created_date, - created_from=created_from, - created_to=created_to, - starred=starred, - ) - video_half, video_params, _ = self._build_half( - kind="video", - origin=origin, - categories=categories, - is_intermediate=is_intermediate, - board_id=board_id, - search_term=search_term, - user_id=user_id, - is_admin=is_admin, - names_only=True, - created_date=created_date, - created_from=created_from, - created_to=created_to, - starred=starred, + kind=GalleryItemKind(found_kind), name=found_name, index=int(index), total=int(total) ) - order_clause = self._build_order_clause(starred_first, order_dir) - - union_query = f"""--sql - SELECT * FROM ( - {image_half} - UNION ALL - {video_half} + def _names(self, filters: Filters, starred_first: bool, order_dir: SQLiteDirection) -> tuple[Sequence[Any], int]: + """The ordered (kind, name, starred) rows and the starred count. Shared by both name-list shapes so the + deprecated `(kind, name)` variant and the flat one can never drift apart in ordering or filtering.""" + rows = self._queries.gallery.names( + filters, starred_first=starred_first, descending=order_dir == SQLiteDirection.Descending ) - {order_clause} - ; - """ - - with self._db.transaction() as cursor: - cursor.execute(union_query, image_params + video_params) - rows = cast(list[sqlite3.Row], cursor.fetchall()) - - starred_count = 0 - if starred_first: - starred_count = sum(1 for r in rows if r["starred"]) - - return rows, starred_count + return rows, (sum(1 for row in rows if row[2]) if starred_first else 0) def list_item_names( self, @@ -295,9 +133,7 @@ def list_item_names( created_to: Optional[str] = None, starred: Optional[bool] = None, ) -> GalleryItemNamesResult: - rows, starred_count = self._query_name_rows( - starred_first=starred_first, - order_dir=order_dir, + filters = Filters( origin=origin, categories=categories, is_intermediate=is_intermediate, @@ -310,7 +146,8 @@ def list_item_names( created_to=created_to, starred=starred, ) - refs = [GalleryItemRef(kind=GalleryItemKind(row["kind"]), name=row["name"]) for row in rows] + rows, starred_count = self._names(filters, starred_first, order_dir) + refs = [GalleryItemRef(kind=GalleryItemKind(kind), name=name) for kind, name, _ in rows] return GalleryItemNamesResult(items=refs, starred_count=starred_count, total_count=len(refs)) def get_item_names( @@ -329,9 +166,7 @@ def get_item_names( created_to: Optional[str] = None, starred: Optional[bool] = None, ) -> GalleryItemNames: - rows, starred_count = self._query_name_rows( - starred_first=starred_first, - order_dir=order_dir, + filters = Filters( origin=origin, categories=categories, is_intermediate=is_intermediate, @@ -344,100 +179,25 @@ def get_item_names( created_to=created_to, starred=starred, ) - # A list comprehension over the raw column, deliberately: building one model per row - # is what made the deprecated variant expensive. - names = [row["name"] for row in rows] + rows, starred_count = self._names(filters, starred_first, order_dir) + # The raw names, deliberately: building one model per row is what made the deprecated variant expensive. + names = [name for _, name, _ in rows] return GalleryItemNames(item_names=names, starred_count=starred_count, total_count=len(names)) - def get_dates( - self, - user_id: Optional[str] = None, - is_admin: bool = False, - ) -> list[VirtualSubBoardDTO]: - image_conditions = " AND images.is_intermediate = 0 " - video_conditions = " AND videos.is_intermediate = 0 " - image_params: list[Union[int, str, bool]] = [] - video_params: list[Union[int, str, bool]] = [] - - if user_id is not None and not is_admin: - image_conditions += " AND images.user_id = ? " - image_params.append(user_id) - video_conditions += " AND videos.user_id = ? " - video_params.append(user_id) - - union = f"""--sql - SELECT - images.created_at AS created_at, - 'image' AS kind, - images.image_name AS name, - images.image_category AS category - FROM images - WHERE 1=1 {image_conditions} - UNION ALL - SELECT - videos.created_at AS created_at, - 'video' AS kind, - videos.video_name AS name, - videos.video_category AS category - FROM videos - WHERE 1=1 {video_conditions} - """ - - counts_query = f"""--sql - SELECT - DATE(created_at) AS date, - SUM(CASE WHEN kind = 'image' AND category = 'general' THEN 1 ELSE 0 END) AS image_count, - SUM(CASE WHEN kind = 'image' AND category != 'general' THEN 1 ELSE 0 END) AS asset_count, - SUM(CASE WHEN kind = 'video' THEN 1 ELSE 0 END) AS video_count, - -- Same not-'general' predicate as the image asset_count above (rather than the - -- listing services' explicit asset-category allowlist), so image and video - -- counts stay consistent within this query. - SUM(CASE WHEN kind = 'video' AND category != 'general' THEN 1 ELSE 0 END) AS asset_video_count - FROM ({union}) - GROUP BY DATE(created_at) - ORDER BY date DESC; - """ - - # The cover is the newest item of each date. ROW_NUMBER with kind/name tie-breakers - # (rather than a bare-column MAX() aggregate) keeps the choice deterministic when - # several items share the maximum timestamp — otherwise the cover could flicker - # between equally-new items across refetches. - covers_query = f"""--sql - SELECT date, kind, name FROM ( - SELECT - DATE(created_at) AS date, - kind, - name, - ROW_NUMBER() OVER ( - PARTITION BY DATE(created_at) - ORDER BY created_at DESC, kind DESC, name DESC - ) AS rn - FROM ({union}) - ) - WHERE rn = 1; - """ - - with self._db.transaction() as cursor: - cursor.execute(counts_query, image_params + video_params) - count_rows = cast(list[sqlite3.Row], cursor.fetchall()) - cursor.execute(covers_query, image_params + video_params) - cover_rows = cast(list[sqlite3.Row], cursor.fetchall()) - - covers = {row["date"]: (row["kind"], row["name"]) for row in cover_rows} - + def get_dates(self, user_id: Optional[str] = None, is_admin: bool = False) -> list[VirtualSubBoardDTO]: + counts, covers = self._queries.gallery.date_counts(user_id if user_id is not None and not is_admin else None) boards: list[VirtualSubBoardDTO] = [] - for row in count_rows: - date = row["date"] - cover_kind, cover_name = covers.get(date, (None, None)) + for day in counts: + cover_kind, cover_name = covers.get(day.day, (None, None)) boards.append( VirtualSubBoardDTO( - virtual_board_id=f"by_date:{date}", - board_name=date, - date=date, - image_count=row["image_count"], - asset_count=row["asset_count"], - video_count=row["video_count"], - asset_video_count=row["asset_video_count"], + virtual_board_id=f"by_date:{day.day}", + board_name=day.day, + date=day.day, + image_count=day.images, + asset_count=day.assets, + video_count=day.videos, + asset_video_count=day.asset_videos, cover_image_name=cover_name if cover_kind == "image" else None, cover_video_name=cover_name if cover_kind == "video" else None, ) @@ -448,257 +208,39 @@ def get_board_media_summaries(self, board_ids: list[str]) -> dict[str, BoardMedi summaries = {board_id: BoardMediaSummary() for board_id in board_ids} if not board_ids: return summaries - - placeholders = ",".join("?" for _ in board_ids) - query = f"""--sql - SELECT - board_id, - SUM(CASE WHEN kind = 'image' AND category = 'general' THEN 1 ELSE 0 END) AS image_count, - SUM(CASE WHEN kind = 'image' AND category != 'general' THEN 1 ELSE 0 END) AS asset_count, - SUM(CASE WHEN kind = 'video' THEN 1 ELSE 0 END) AS video_count, - SUM(CASE WHEN kind = 'video' AND category != 'general' THEN 1 ELSE 0 END) AS asset_video_count, - MAX(CASE WHEN rank = 1 AND kind = 'image' THEN name END) AS cover_image_name, - MAX(CASE WHEN rank = 1 AND kind = 'video' THEN name END) AS cover_video_name - FROM ( - SELECT - board_id, - kind, - name, - category, - ROW_NUMBER() OVER ( - PARTITION BY board_id - ORDER BY starred DESC, created_at DESC, kind DESC, name DESC - ) AS rank - FROM ( - SELECT - board_images.board_id AS board_id, - 'image' AS kind, - images.image_name AS name, - images.image_category AS category, - images.starred AS starred, - images.created_at AS created_at - FROM board_images - INNER JOIN images ON board_images.image_name = images.image_name - WHERE images.is_intermediate = FALSE - AND board_images.board_id IN ({placeholders}) - UNION ALL - SELECT - board_videos.board_id AS board_id, - 'video' AS kind, - videos.video_name AS name, - videos.video_category AS category, - videos.starred AS starred, - videos.created_at AS created_at - FROM board_videos - INNER JOIN videos ON board_videos.video_name = videos.video_name - WHERE videos.is_intermediate = FALSE - AND board_videos.board_id IN ({placeholders}) - ) - ) - GROUP BY board_id; - """ - with self._db.transaction() as cursor: - cursor.execute(query, [*board_ids, *board_ids]) - rows = cast(list[sqlite3.Row], cursor.fetchall()) - for row in rows: - summaries[row["board_id"]] = BoardMediaSummary( - cover_image_name=row["cover_image_name"], - cover_video_name=row["cover_video_name"], - image_count=row["image_count"], - video_count=row["video_count"], - asset_count=row["asset_count"], - asset_video_count=row["asset_video_count"], + for summary in self._queries.gallery.board_summaries(board_ids): + summaries[summary.board_id] = BoardMediaSummary( + cover_image_name=summary.cover_image_name, + cover_video_name=summary.cover_video_name, + image_count=summary.images, + video_count=summary.videos, + asset_count=summary.assets, + asset_video_count=summary.asset_videos, ) return summaries - @staticmethod - def _build_order_clause(starred_first: bool, order_dir: SQLiteDirection) -> str: - """Ordering with `kind` + `name` tie-breakers. - - Images and videos created within the same timestamp granularity would otherwise - have no defined relative order — rows could reorder across refetches or shift - between offset pages. The tie-breakers follow `order_dir` so ascending and - descending remain exact mirror images. - """ - direction = order_dir.value - tie_breakers = f"kind {direction}, name {direction}" - if starred_first: - return f"ORDER BY starred DESC, created_at {direction}, {tie_breakers}" - return f"ORDER BY created_at {direction}, {tie_breakers}" - - def _build_half( - self, - kind: str, - origin: Optional[ResourceOrigin], - categories: Optional[list[ImageCategory]], - is_intermediate: Optional[bool], - board_id: Optional[str], - search_term: Optional[str], - user_id: Optional[str], - is_admin: bool, - names_only: bool = False, - created_date: Optional[str] = None, - created_from: Optional[str] = None, - created_to: Optional[str] = None, - starred: Optional[bool] = None, - ) -> tuple[str, list[Union[int, str, bool]], str]: - """Builds one half of the union (either `images` or `videos`). - - Returns `(query_with_select, params, count_query)`. Both halves emit the same columns so - UNION ALL is shape-compatible: `kind`, `name`, `width`, `height`, `category`, `starred`, - `is_intermediate`, `board_id`, `created_at`, `duration`, `fps`. (`media_origin` is - NOT one of them -- `list_items` joins it onto the chosen page instead; see there.) - - `names_only=True` selects only `kind`, `name`, `starred`, `created_at` (the minimum needed - for ordering + the counts result). - """ - if kind == "image": - base_table = "images" - join_table = "board_images" - name_col = "image_name" - category_col = "image_category" - origin_col = "image_origin" - duration_expr = "NULL" - fps_expr = "NULL" - elif kind == "video": - base_table = "videos" - join_table = "board_videos" - name_col = "video_name" - category_col = "video_category" - origin_col = "video_origin" - duration_expr = f"{base_table}.duration" - fps_expr = f"{base_table}.fps" - else: - raise ValueError(f"Unknown kind: {kind}") - - if board_id == "none": - from_clause = f"FROM {base_table}" - board_id_expr = "NULL" - elif board_id is not None: - # CROSS JOIN keeps explicit-board work proportional to board membership. - from_clause = ( - f"FROM {join_table} CROSS JOIN {base_table} ON {join_table}.{name_col} = {base_table}.{name_col}" - ) - board_id_expr = f"{join_table}.board_id" - elif names_only: - from_clause = f"FROM {base_table}" - board_id_expr = "NULL" - else: - from_clause = ( - f"FROM {base_table} LEFT JOIN {join_table} ON {join_table}.{name_col} = {base_table}.{name_col}" - ) - board_id_expr = f"{join_table}.board_id" - - if names_only: - select_cols = ( - f"'{kind}' AS kind, " - f"{base_table}.{name_col} AS name, " - f"{base_table}.starred AS starred, " - f"{base_table}.created_at AS created_at" - ) - else: - select_cols = ( - f"'{kind}' AS kind, " - f"{base_table}.{name_col} AS name, " - f"{base_table}.width AS width, " - f"{base_table}.height AS height, " - f"{base_table}.{category_col} AS category, " - f"{base_table}.starred AS starred, " - f"{base_table}.is_intermediate AS is_intermediate, " - f"{board_id_expr} AS board_id, " - f"{base_table}.created_at AS created_at, " - f"{duration_expr} AS duration, " - f"{fps_expr} AS fps" - ) - - conditions = "" - params: list[Union[int, str, bool]] = [] - - if origin is not None: - conditions += f" AND {base_table}.{origin_col} = ? " - params.append(origin.value) - - if categories is not None: - category_strings = [c.value for c in set(categories)] - placeholders = ",".join("?" * len(category_strings)) - conditions += f" AND {base_table}.{category_col} IN ( {placeholders} ) " - for c in category_strings: - params.append(c) - - if is_intermediate is not None: - conditions += f" AND {base_table}.is_intermediate = ? " - params.append(is_intermediate) - - if starred is not None: - conditions += f" AND {base_table}.starred = ? " - params.append(starred) - - if created_date is not None: - conditions += f" AND DATE({base_table}.created_at) = ? " - params.append(created_date) - - if created_from is not None: - conditions += f" AND {base_table}.created_at >= ? " - params.append(created_from) - - if created_to is not None: - conditions += f" AND {base_table}.created_at < DATE(?, '+1 day') " - params.append(created_to) - - if board_id == "none": - conditions += f""" AND NOT EXISTS ( - SELECT 1 - FROM {join_table} - WHERE {join_table}.{name_col} = {base_table}.{name_col} - ) """ - if user_id is not None and not is_admin: - conditions += f" AND {base_table}.user_id = ? " - params.append(user_id) - elif board_id is not None: - conditions += f" AND {join_table}.board_id = ? " - params.append(board_id) - elif user_id is not None and not is_admin: - # No board_id supplied — still enforce per-user isolation so - # non-admin callers cannot enumerate other users' items. - conditions += f" AND {base_table}.user_id = ? " - params.append(user_id) - - if search_term: - conditions += f" AND ({base_table}.metadata LIKE ? OR {base_table}.created_at LIKE ?) " - params.append(f"%{search_term.lower()}%") - params.append(f"%{search_term.lower()}%") - - half_query = f"SELECT {select_cols} {from_clause} WHERE 1=1 {conditions}" - count_query = f"SELECT COUNT(*) {from_clause} WHERE 1=1 {conditions}" - return half_query, params, count_query - - def _row_to_item(self, row: sqlite3.Row, urls) -> GalleryItem: - kind = GalleryItemKind(row["kind"]) - name = row["name"] - if kind == GalleryItemKind.IMAGE: - full_url = urls.get_image_url(name) - thumbnail_url = urls.get_image_url(name, thumbnail=True) - duration = None - fps = None - media_origin = None + def _to_item(self, row: Sequence[Any]) -> GalleryItem: + kind, name, width, height, category, starred, is_intermediate, board_id, created_at, duration, fps, origin = row + urls = self.__invoker.services.urls + item_kind = GalleryItemKind(kind) + if item_kind == GalleryItemKind.IMAGE: + full_url, thumbnail_url = urls.get_image_url(name), urls.get_image_url(name, thumbnail=True) + duration = fps = media_origin = None else: - full_url = urls.get_video_url(name) - thumbnail_url = urls.get_video_url(name, thumbnail=True) - duration = row["duration"] - fps = row["fps"] - media_origin = coerce_media_origin(row["media_origin"]) + full_url, thumbnail_url = urls.get_video_url(name), urls.get_video_url(name, thumbnail=True) + media_origin = coerce_media_origin(origin) return GalleryItem( - kind=kind, + kind=item_kind, name=name, full_url=full_url, thumbnail_url=thumbnail_url, - width=row["width"], - height=row["height"], - category=ImageCategory(row["category"]), - starred=bool(row["starred"]), - is_intermediate=bool(row["is_intermediate"]), - board_id=row["board_id"], - created_at=row["created_at"], + width=width, + height=height, + category=ImageCategory(category), + starred=bool(starred), + is_intermediate=bool(is_intermediate), + board_id=board_id, + created_at=created_at, duration=duration, fps=fps, media_origin=media_origin, diff --git a/invokeai/app/services/gallery_maintenance/gallery_maintenance_default.py b/invokeai/app/services/gallery_maintenance/gallery_maintenance_default.py index 0eaea62a4ca..aa3fcd6bbec 100644 --- a/invokeai/app/services/gallery_maintenance/gallery_maintenance_default.py +++ b/invokeai/app/services/gallery_maintenance/gallery_maintenance_default.py @@ -516,7 +516,11 @@ def _relative_parent(self, path: Path, root: Path) -> str: relative = self._relative(parent, root) return "" if relative == "." else relative - def _create_backup(self) -> str: + def _create_backup(self) -> str | None: + if self._services.database.dialect_name != "sqlite": + # A MySQL or MariaDB database is backed up with the server's own tools, by its operator. + self._logger.info("Gallery maintenance takes no backup of a server database; back it up with its tools") + return None backup_dir = self._services.configuration.db_path.parent / "backup" timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") destination = backup_dir / f"backup-{timestamp}-gallery-maintenance-{uuid.uuid4().hex}.db" @@ -526,7 +530,7 @@ def _create_backup(self) -> str: missing_directories.append(current) current = current.parent try: - self._services.database.backup_to(destination) + self._services.database.backup(destination) for directory in reversed(missing_directories): self.__fsync_directory(directory.parent) if not missing_directories: diff --git a/invokeai/app/services/image_index/cluster_labels.py b/invokeai/app/services/image_index/cluster_labels.py index 288095eb9fd..a43acc31a95 100644 --- a/invokeai/app/services/image_index/cluster_labels.py +++ b/invokeai/app/services/image_index/cluster_labels.py @@ -8,6 +8,7 @@ """ import hashlib +import unicodedata from pathlib import Path import numpy as np @@ -39,9 +40,9 @@ def normalize_custom_vocab_terms(raw: list[str]) -> list[str]: """Normalize user-supplied vocabulary terms for storage and embedding. - Whitespace is collapsed, terms are lowercased to match the bundled - vocabulary's convention, empties are dropped, and duplicates are removed - (first occurrence wins, order preserved). + Terms are NFKC-normalized, whitespace is collapsed, terms are lowercased to + match the bundled vocabulary's convention, empties are dropped, and + duplicates are removed (first occurrence wins, order preserved). Raises: ValueError: A term exceeds MAX_CUSTOM_VOCAB_TERM_LENGTH after @@ -58,7 +59,9 @@ def normalize_custom_vocab_terms(raw: list[str]) -> list[str]: # term lists could hash identically and serve each other's cached # embeddings. cleaned = "".join(ch for ch in entry if ch.isprintable() or ch.isspace()) - term = " ".join(cleaned.split()).lower() + # NFKC makes one term of the spellings MySQL's and MariaDB's case-insensitive collations equate (composed + # and decomposed accents, ligatures, full-width letters), so every backend stores the same terms. + term = " ".join(unicodedata.normalize("NFKC", cleaned).split()).lower() if not term: continue if len(term) > MAX_CUSTOM_VOCAB_TERM_LENGTH: diff --git a/invokeai/app/services/image_index/image_index_default.py b/invokeai/app/services/image_index/image_index_default.py index 16412d601e5..c8fabb680ea 100644 --- a/invokeai/app/services/image_index/image_index_default.py +++ b/invokeai/app/services/image_index/image_index_default.py @@ -25,6 +25,7 @@ from invokeai.app.services.image_records.image_records_common import ImageCategory from invokeai.app.services.images.images_common import ImageDTO from invokeai.app.services.session_queue.session_queue_common import DEFAULT_QUEUE_ID +from invokeai.app.services.shared.database.errors import TransientDatabaseError from invokeai.app.services.videos.videos_common import VideoDTO from invokeai.backend.model_manager.load.model_cache.model_cache import MODEL_LOAD_LOCK from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init @@ -1537,7 +1538,7 @@ def _failed_scope_retry_available(self, user_id: str, current_hash: str) -> bool terminal: /points answers `state: "empty", stale: false` (so the client stops asking) and every later job returns early (so the fit is never re-entered). The map stays blank until the accessible image set - changes, and the row is in SQLite, so restarting does not clear it. + changes, and the row is in the database, so restarting does not clear it. Retrying a failed scope once splits the difference the original code could not: a transient failure — a MemoryError while a generation holds @@ -1577,15 +1578,15 @@ def _process_projection(self, user_id: str, all_images: bool) -> None: items = records.list_accessible_embedded_items(scope_user, self._model_id) current_hash = scope_hash(self._model_id, items) cached = records.get_projection(user_id, self._model_id) - except Exception: + except Exception as e: # The DB reads are as failure-prone as the fit ("database is # locked" is a designed-for condition here), and the job was # already popped from the dedup map. Re-queue it rather than let # the generic worker handler swallow it: /refresh has already # answered `enqueued: true`, and a client waiting on # image_map_projection_ready would otherwise wait forever. - logger.exception(f"Image map: could not read the projection scope for user '{user_id}'; re-queueing") - self._requeue_projection(user_id, all_images) + logger.exception(f"Image map: could not read the projection scope for user '{user_id}'") + self._retry_projection_or_give_up(e, user_id, all_images, point_count=0) return retrying_failed_scope = ( @@ -1620,9 +1621,11 @@ def _process_projection(self, user_id: str, all_images: bool) -> None: # accessible embedding (hundreds of MB on a large gallery), which is # pure waste on the no-op path that /points and /refresh drive. found_items, matrix = records.get_embeddings(items, self._model_id) - except Exception: - logger.exception(f"Image map: could not read embeddings for user '{user_id}'; re-queueing") - self._requeue_projection(user_id, all_images) + except Exception as e: + logger.exception(f"Image map: could not read embeddings for user '{user_id}'") + self._retry_projection_or_give_up( + e, user_id, all_images, point_count=cached.point_count if cached is not None else 0 + ) return if self._stop_event.is_set(): @@ -1666,11 +1669,13 @@ def _process_projection(self, user_id: str, all_images: bool) -> None: found_items, coords, ) - except Exception: + except Exception as e: # Same reasoning as the read above: the fit is done but unsaved, so # re-queue rather than drop it and leave /refresh's promise unmet. - logger.exception(f"Image map: could not cache the projection for user '{user_id}'; re-queueing") - self._requeue_projection(user_id, all_images) + logger.exception(f"Image map: could not cache the projection for user '{user_id}'") + self._retry_projection_or_give_up( + e, user_id, all_images, point_count=cached.point_count if cached is not None else 0 + ) return # The budget moves only once the outcome is durable. Spending it before @@ -1699,6 +1704,18 @@ def _process_projection(self, user_id: str, all_images: bool) -> None: # what is purely a delivery failure. logger.exception(f"Image map: could not emit projection_ready for user '{user_id}'") + def _retry_projection_or_give_up(self, error: Exception, user_id: str, all_images: bool, point_count: int) -> None: + """Re-queues a projection whose database work failed for the moment (the database was busy, or the call lost + a race); gives up on one whose work failed for good, which would fail again on every retry, forever (a + projection larger than a server takes in one statement, for one). Giving up still tells a waiting client that + nothing newer is coming: `point_count` is what it is served.""" + assert self._invoker is not None + if isinstance(error, TransientDatabaseError): + self._invoker.services.logger.info(f"Image map: re-queueing the projection for user '{user_id}'") + self._requeue_projection(user_id, all_images) + return + self._invoker.services.events.emit_image_map_projection_ready(user_id=user_id, point_count=point_count) + def _requeue_projection(self, user_id: str, all_images: bool) -> None: """Put a popped job back, without clobbering a newer request for that user.""" with self._projection_lock: @@ -1818,11 +1835,9 @@ def _resolve_model_config(self, model_name: str) -> Optional["AnyModelConfig"]: configs = store.search_by_attr(model_name=model_name, model_type=model_type) if not configs: continue - # Names are not unique (the UNIQUE(name, base, type) constraint is gone), and - # `search_by_attr` orders by type/base/name/format — so for two otherwise identical - # rows the tiebreak is insertion order, and reinstalling one moves the winner. That - # changes the model hash, and `start()` then discards every embedding computed under - # the old one. Order by key instead: stable across reinstalls of the same set. + # Names are not unique (the UNIQUE(name, base, type) constraint is gone). Choose by key, + # which is stable across reinstalls of the same set: a different winner changes the + # model hash, and `start()` then discards every embedding computed under the old one. chosen = min(configs, key=lambda config: config.key) if len(configs) > 1 and (self._model_config is None or chosen.key != self._model_config.key): self._invoker.services.logger.warning( diff --git a/invokeai/app/services/image_index/image_index_records_base.py b/invokeai/app/services/image_index/image_index_records_base.py index c285e4e0356..24cceb5238d 100644 --- a/invokeai/app/services/image_index/image_index_records_base.py +++ b/invokeai/app/services/image_index/image_index_records_base.py @@ -17,11 +17,8 @@ class ImageIndexRecordsBase(ABC): embedding is deleted with its media by foreign key and eligibility and board access are answered from that kind's own tables. - Every method here is its own unit of work and must not be called from inside another - service's open transaction. `SqliteDatabase.transaction()` commits and rolls back the whole - shared connection rather than using savepoints, so a nested call either commits the outer - transaction's work early (on success) or discards it (on failure). No caller in the tree - nests these today; keep it that way. + Every method here runs in a transaction of its own, so none can be called inside another + transaction. """ @abstractmethod diff --git a/invokeai/app/services/image_index/image_index_records_default.py b/invokeai/app/services/image_index/image_index_records_default.py new file mode 100644 index 00000000000..41bd5c6f5dc --- /dev/null +++ b/invokeai/app/services/image_index/image_index_records_default.py @@ -0,0 +1,159 @@ +import json + +import numpy as np + +from invokeai.app.services.image_index.image_index_common import ( + EMBEDDING_DTYPE, + ImageIndexStatus, + IndexedItem, + MediaKind, + ProjectionRecord, + blob_to_coords, + blobs_to_embeddings, + coords_to_blob, + embedding_to_blob, +) +from invokeai.app.services.image_index.image_index_records_base import ImageIndexRecordsBase +from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.locks import DatabaseLock +from invokeai.backend.util.logging import InvokeAILogger + +# What `embedding_to_blob` writes. +_STORED_ENCODING = "float16" + + +class ImageIndexRecords(ImageIndexRecordsBase): + """Semantic image index storage.""" + + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + self._logger = InvokeAILogger.get_logger(self.__class__.__name__) + + def start(self, invoker: Invoker) -> None: + self._invoker = invoker + + def upsert_embedding(self, item: IndexedItem, model_id: str, embedding: np.ndarray) -> None: + blob = embedding_to_blob(embedding) + # The item may be deleted between being scheduled and embedded, which makes the embedding pointless. + if not self._queries.image_index.upsert_embedding(item, model_id, embedding.shape[0], blob, _STORED_ENCODING): + self._logger.debug(f"Skipped embedding for missing {item.kind} {item.name}") + + def get_embeddings(self, items: list[IndexedItem], model_id: str) -> tuple[list[IndexedItem], np.ndarray]: + # Dedupe while preserving order so repeated input items cannot + # double-count rows in downstream projection/similarity math. + items = list(dict.fromkeys(items)) + rows = self._queries.image_index.embeddings(items, model_id) + + # Read back in the caller's order: the returned matrix's rows align with it. + found_items = [item for item in items if item in rows] + + if not found_items: + return [], np.empty((0, 0), dtype=EMBEDDING_DTYPE) + + dim = rows[found_items[0]][0] + for item in found_items: + row_dim = rows[item][0] + if row_dim != dim: + raise ValueError(f"Inconsistent embedding dims for model {model_id}: found {row_dim} and {dim}") + matrix = blobs_to_embeddings( + [rows[item][1] for item in found_items], [rows[item][2] for item in found_items], dim + ) + return found_items, matrix + + def delete_embedding(self, item: IndexedItem) -> None: + self._queries.image_index.delete_embedding(item) + + def delete_embeddings_for_other_models(self, model_id: str) -> int: + return self._queries.image_index.delete_embeddings_for_other_models(model_id) + + def list_unembedded_items(self, model_id: str, limit: int) -> list[IndexedItem]: + if limit < 0: + # SQLite reads a negative LIMIT as unbounded, which would turn a backfill batch into a full-table load, + # and MySQL rejects it. Fail on a miscomputed limit instead. + raise ValueError(f"limit must be non-negative, got {limit}") + return self._queries.image_index.unembedded(model_id, limit) + + def count_index_status(self, model_id: str) -> ImageIndexStatus: + total, embedded = self._queries.image_index.status(model_id) + return ImageIndexStatus(total=total, embedded=embedded) + + def list_accessible_embedded_items(self, user_id: str | None, model_id: str) -> list[IndexedItem]: + # The projection scope_hash and semantic search both derive from this listing, so it must keep matching the + # gallery's access model. + return self._queries.image_index.accessible_embedded(user_id, model_id) + + def get_custom_vocab_terms(self) -> list[str]: + return self._queries.image_index.vocab_terms() + + def set_custom_vocab_terms(self, terms: list[str]) -> None: + def replace(q: Queries) -> None: + # On a server another replace would otherwise delete only what it sees committed, leaving both lists. + q.locks.acquire(DatabaseLock.IMAGE_INDEX_VOCABULARY) + q.image_index.replace_vocab_terms(terms) + + self._queries.run(replace) + + def get_projection(self, user_id: str, model_id: str) -> ProjectionRecord | None: + row = self._queries.image_index.projection(user_id, model_id) + if row is None: + return None + + names: list[str] = json.loads(row.image_names) + # A projection cached before videos were indexable has no kinds: every name in it is an + # image, which is what a NULL column means. Same reading for a row written by this code, + # where the column is always present. + kinds: list[MediaKind] = json.loads(row.item_kinds) if row.item_kinds is not None else ["image"] * len(names) + if len(kinds) != len(names): + # Unreachable from this code — both arrays are written by one statement from one + # list — but reachable by downgrading to a build that rewrites `image_names` and + # leaves `item_kinds` behind. Reported as "no cached projection" rather than + # raised: every reader is a request handler or the worker's own fit, so raising + # would 500 the map and spin the worker on a row that only `set_projection` + # repairs — which is never reached while the read keeps failing. + self._logger.warning( + f"Discarding the cached projection for user {user_id}: {len(names)} names but {len(kinds)} kinds" + ) + return None + + return ProjectionRecord( + user_id=user_id, + model_id=model_id, + scope_hash=row.scope_hash, + params=row.params, + point_count=row.point_count, + items=[IndexedItem(kind, name) for kind, name in zip(kinds, names, strict=True)], + coords=blob_to_coords(row.coords, row.point_count), + created_at=row.created_at, + updated_at=row.updated_at, + ) + + def set_projection( + self, + user_id: str, + model_id: str, + scope_hash: str, + params: str, + items: list[IndexedItem], + coords: np.ndarray, + ) -> None: + if len(items) != coords.shape[0]: + raise ValueError(f"Got {len(items)} items but {coords.shape[0]} coordinate rows") + stored = self._queries.image_index.set_projection( + user_id=user_id, + model_id=model_id, + scope_hash=scope_hash, + params=params, + point_count=len(items), + image_names=json.dumps([item.name for item in items]), + item_kinds=json.dumps([item.kind for item in items]), + coords=coords_to_blob(coords), + ) + # There is nobody left to serve the projection of an account deleted while it was computed. + if not stored: + self._logger.debug(f"Skipped projection for missing user {user_id}") + + def delete_projection(self, user_id: str, model_id: str) -> None: + self._queries.image_index.delete_projection(user_id, model_id) diff --git a/invokeai/app/services/image_index/image_index_records_sqlite.py b/invokeai/app/services/image_index/image_index_records_sqlite.py deleted file mode 100644 index e7d26778648..00000000000 --- a/invokeai/app/services/image_index/image_index_records_sqlite.py +++ /dev/null @@ -1,415 +0,0 @@ -import json -from dataclasses import dataclass - -import numpy as np - -from invokeai.app.services.board_records.board_records_common import BoardVisibility -from invokeai.app.services.image_index.image_index_common import ( - EMBEDDING_DTYPE, - ImageIndexStatus, - IndexedItem, - MediaKind, - ProjectionRecord, - blob_to_coords, - blobs_to_embeddings, - coords_to_blob, - embedding_to_blob, -) -from invokeai.app.services.image_index.image_index_records_base import ImageIndexRecordsBase -from invokeai.app.services.image_records.image_records_common import ImageCategory -from invokeai.app.services.invoker import Invoker -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.backend.util.logging import InvokeAILogger - -# SQLite's default variable limit is 999; stay well under it when chunking IN clauses. -_IN_CLAUSE_CHUNK = 500 - - -@dataclass(frozen=True) -class _Namespace: - """The tables and columns one media kind is indexed through. - - Images and videos are indexed by the same rules over different tables, so the rules are - written once and the SQL is built per namespace. Every field is a constant of this module; - nothing here is ever derived from a request. - """ - - kind: MediaKind - table: str - name_column: str - category_column: str - embeddings_table: str - board_table: str - - @property - def eligible_conditions(self) -> str: - """Conditions defining a "gallery" item, i.e. one worth indexing.""" - return f"{self.table}.is_intermediate = ? AND {self.table}.{self.category_column} = ?" - - -_IMAGES = _Namespace( - kind="image", - table="images", - name_column="image_name", - category_column="image_category", - embeddings_table="image_embeddings", - board_table="board_images", -) -_VIDEOS = _Namespace( - kind="video", - table="videos", - name_column="video_name", - category_column="video_category", - embeddings_table="video_embeddings", - board_table="board_videos", -) -# Images first, so every listing that concatenates namespaces has one stable order. -_NAMESPACES: tuple[_Namespace, ...] = (_IMAGES, _VIDEOS) -_NAMESPACE_BY_KIND: dict[MediaKind, _Namespace] = {namespace.kind: namespace for namespace in _NAMESPACES} - - -def _eligible_params() -> tuple[bool, str]: - return (False, ImageCategory.GENERAL.value) - - -class ImageIndexRecordsSqlite(ImageIndexRecordsBase): - """SQLite implementation of semantic image index storage.""" - - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - self._logger = InvokeAILogger.get_logger(self.__class__.__name__) - - def start(self, invoker: Invoker) -> None: - self._invoker = invoker - - def upsert_embedding(self, item: IndexedItem, model_id: str, embedding: np.ndarray) -> None: - namespace = _NAMESPACE_BY_KIND[item.kind] - blob = embedding_to_blob(embedding) - with self._db.transaction() as cursor: - # The item may be deleted between being scheduled and embedded, in which case the - # embedding is pointless and this is a no-op. `WHERE EXISTS` rather than catching the - # foreign-key IntegrityError: the guard and the insert are one atomic statement, and - # no exception is raised on the expected race, so nothing can leave the transaction - # in a state the caller did not ask for. An IntegrityError from here now means a - # genuine bug and is allowed to propagate. - cursor.execute( - f"""--sql - INSERT INTO {namespace.embeddings_table} ({namespace.name_column}, model_id, dim, embedding, encoding) - SELECT ?, ?, ?, ?, 'float16' - WHERE EXISTS ( - SELECT 1 FROM {namespace.table} WHERE {namespace.table}.{namespace.name_column} = ? - ) - ON CONFLICT ({namespace.name_column}, model_id) - DO UPDATE SET dim = excluded.dim, embedding = excluded.embedding, encoding = excluded.encoding; - """, - (item.name, model_id, embedding.shape[0], blob, item.name), - ) - if cursor.rowcount == 0: - self._logger.debug(f"Skipped embedding for missing {item.kind} {item.name}") - - def get_embeddings(self, items: list[IndexedItem], model_id: str) -> tuple[list[IndexedItem], np.ndarray]: - # Dedupe while preserving order so repeated input items cannot - # double-count rows in downstream projection/similarity math. - items = list(dict.fromkeys(items)) - rows: dict[IndexedItem, tuple[int, bytes, str]] = {} - - # Only the reads happen under the transaction. Deserialization and validation are done - # afterwards so a malformed row raises outside it: `transaction()` rolls the shared - # connection back on any exception, which would be an unpleasant side effect of what the - # caller asked to be a read. - with self._db.transaction() as cursor: - for namespace in _NAMESPACES: - names = [item.name for item in items if item.kind == namespace.kind] - for start in range(0, len(names), _IN_CLAUSE_CHUNK): - chunk = names[start : start + _IN_CLAUSE_CHUNK] - placeholders = ",".join("?" * len(chunk)) - cursor.execute( - f"""--sql - SELECT {namespace.name_column}, dim, embedding, encoding - FROM {namespace.embeddings_table} - WHERE model_id = ? AND {namespace.name_column} IN ({placeholders}); - """, - (model_id, *chunk), - ) - for name, dim, blob, encoding in cursor.fetchall(): - rows[IndexedItem(namespace.kind, name)] = (dim, blob, encoding) - - # Read back in the caller's order: the returned matrix's rows align with it. - found_items = [item for item in items if item in rows] - - if not found_items: - return [], np.empty((0, 0), dtype=EMBEDDING_DTYPE) - - dim = rows[found_items[0]][0] - for item in found_items: - row_dim = rows[item][0] - if row_dim != dim: - raise ValueError(f"Inconsistent embedding dims for model {model_id}: found {row_dim} and {dim}") - matrix = blobs_to_embeddings( - [rows[item][1] for item in found_items], [rows[item][2] for item in found_items], dim - ) - return found_items, matrix - - def delete_embedding(self, item: IndexedItem) -> None: - namespace = _NAMESPACE_BY_KIND[item.kind] - with self._db.transaction() as cursor: - cursor.execute( - f"""--sql - DELETE FROM {namespace.embeddings_table} WHERE {namespace.name_column} = ?; - """, - (item.name,), - ) - - def delete_embeddings_for_other_models(self, model_id: str) -> int: - deleted = 0 - with self._db.transaction() as cursor: - for namespace in _NAMESPACES: - cursor.execute( - f"""--sql - DELETE FROM {namespace.embeddings_table} WHERE model_id != ?; - """, - (model_id,), - ) - deleted += cursor.rowcount - return deleted - - def list_unembedded_items(self, model_id: str, limit: int) -> list[IndexedItem]: - if limit < 0: - # SQLite reads a negative LIMIT as unbounded, which would turn a backfill batch into - # a full-table load. Fail on a miscomputed limit instead of silently doing that. - raise ValueError(f"limit must be non-negative, got {limit}") - - # Each namespace contributes its own oldest `limit` rows, merged here so a batch is the - # globally oldest work rather than every image before any video. Taking `limit` from each - # is what makes the merge correct: whatever the merged prefix should contain is present - # in one of the two lists. - candidates: list[tuple[str, str, IndexedItem]] = [] - with self._db.transaction() as cursor: - for namespace in _NAMESPACES: - cursor.execute( - f"""--sql - SELECT {namespace.table}.{namespace.name_column}, {namespace.table}.created_at - FROM {namespace.table} - LEFT JOIN {namespace.embeddings_table} - ON {namespace.embeddings_table}.{namespace.name_column} - = {namespace.table}.{namespace.name_column} - AND {namespace.embeddings_table}.model_id = ? - WHERE {namespace.eligible_conditions} - AND {namespace.embeddings_table}.{namespace.name_column} IS NULL - ORDER BY {namespace.table}.created_at ASC, {namespace.table}.{namespace.name_column} ASC - LIMIT ?; - """, - (model_id, *_eligible_params(), limit), - ) - candidates.extend( - (created_at, name, IndexedItem(namespace.kind, name)) for name, created_at in cursor.fetchall() - ) - - candidates.sort(key=lambda candidate: (candidate[0], candidate[1])) - return [item for _, _, item in candidates[:limit]] - - def count_index_status(self, model_id: str) -> ImageIndexStatus: - total = 0 - embedded = 0 - with self._db.transaction() as cursor: - for namespace in _NAMESPACES: - cursor.execute( - f"""--sql - SELECT - COUNT(*), - COUNT({namespace.embeddings_table}.{namespace.name_column}) - FROM {namespace.table} - LEFT JOIN {namespace.embeddings_table} - ON {namespace.embeddings_table}.{namespace.name_column} - = {namespace.table}.{namespace.name_column} - AND {namespace.embeddings_table}.model_id = ? - WHERE {namespace.eligible_conditions}; - """, - (model_id, *_eligible_params()), - ) - namespace_total, namespace_embedded = cursor.fetchone() - total += namespace_total - embedded += namespace_embedded - return ImageIndexStatus(total=total, embedded=embedded) - - def list_accessible_embedded_items(self, user_id: str | None, model_id: str) -> list[IndexedItem]: - # Both clauses mirror the gallery "all" listing semantics - # (image_records_sqlite / video_records_sqlite): items on archived - # boards are hidden from every scope, and a scoped user sees their own - # unboarded items plus items on active boards they own, that are - # shared/public, or that were individually shared with them via - # shared_boards. The projection scope_hash and semantic search both - # derive from this listing, so any change here must keep matching the - # gallery's access model. - items: list[IndexedItem] = [] - with self._db.transaction() as cursor: - for namespace in _NAMESPACES: - params: list[object] = [model_id, *_eligible_params()] - if user_id is None: - # Administrative scope: everything except items on archived boards. - access_clause = f"""AND ( - {namespace.board_table}.board_id IS NULL - OR EXISTS ( - SELECT 1 FROM boards - WHERE boards.board_id = {namespace.board_table}.board_id - AND boards.archived = 0 - ) - )""" - else: - access_clause = f"""AND ( - ({namespace.board_table}.board_id IS NULL AND {namespace.table}.user_id = ?) - OR EXISTS ( - SELECT 1 FROM boards - WHERE boards.board_id = {namespace.board_table}.board_id - AND boards.archived = 0 - AND ( - boards.user_id = ? - OR boards.board_visibility IN (?, ?) - OR EXISTS ( - SELECT 1 FROM shared_boards - WHERE shared_boards.board_id = boards.board_id - AND shared_boards.user_id = ? - ) - ) - ) - )""" - params.extend( - [user_id, user_id, BoardVisibility.Shared.value, BoardVisibility.Public.value, user_id] - ) - - cursor.execute( - f"""--sql - SELECT DISTINCT {namespace.embeddings_table}.{namespace.name_column} - FROM {namespace.embeddings_table} - JOIN {namespace.table} - ON {namespace.table}.{namespace.name_column} - = {namespace.embeddings_table}.{namespace.name_column} - LEFT JOIN {namespace.board_table} - ON {namespace.board_table}.{namespace.name_column} = {namespace.table}.{namespace.name_column} - WHERE {namespace.embeddings_table}.model_id = ? - AND {namespace.eligible_conditions} - {access_clause} - ORDER BY {namespace.embeddings_table}.{namespace.name_column} ASC; - """, - params, - ) - items.extend(IndexedItem(namespace.kind, row[0]) for row in cursor.fetchall()) - return items - - def get_custom_vocab_terms(self) -> list[str]: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT term FROM image_index_vocab_terms ORDER BY term ASC; - """ - ) - return [row[0] for row in cursor.fetchall()] - - def set_custom_vocab_terms(self, terms: list[str]) -> None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - DELETE FROM image_index_vocab_terms; - """ - ) - # OR IGNORE: callers normalize away duplicates, but the NOCASE - # primary key would otherwise turn a case-variant pair that slipped - # through into an IntegrityError that rolls back the whole replace. - cursor.executemany( - """--sql - INSERT OR IGNORE INTO image_index_vocab_terms (term) VALUES (?); - """, - [(term,) for term in terms], - ) - - def get_projection(self, user_id: str, model_id: str) -> ProjectionRecord | None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT scope_hash, params, point_count, image_names, item_kinds, coords, created_at, updated_at - FROM image_projections - WHERE user_id = ? AND model_id = ?; - """, - (user_id, model_id), - ) - row = cursor.fetchone() - - if row is None: - return None - - names: list[str] = json.loads(row[3]) - # A projection cached before videos were indexable has no kinds: every name in it is an - # image, which is what a NULL column means. Same reading for a row written by this code, - # where the column is always present. - kinds: list[MediaKind] = json.loads(row[4]) if row[4] is not None else ["image"] * len(names) - if len(kinds) != len(names): - # Unreachable from this code — both arrays are written by one statement from one - # list — but reachable by downgrading to a build that rewrites `image_names` and - # leaves `item_kinds` behind. Reported as "no cached projection" rather than - # raised: every reader is a request handler or the worker's own fit, so raising - # would 500 the map and spin the worker on a row that only `set_projection` - # repairs — which is never reached while the read keeps failing. - self._logger.warning( - f"Discarding the cached projection for user {user_id}: {len(names)} names but {len(kinds)} kinds" - ) - return None - - return ProjectionRecord( - user_id=user_id, - model_id=model_id, - scope_hash=row[0], - params=row[1], - point_count=row[2], - items=[IndexedItem(kind, name) for kind, name in zip(kinds, names, strict=True)], - coords=blob_to_coords(row[5], row[2]), - created_at=row[6], - updated_at=row[7], - ) - - def set_projection( - self, - user_id: str, - model_id: str, - scope_hash: str, - params: str, - items: list[IndexedItem], - coords: np.ndarray, - ) -> None: - if len(items) != coords.shape[0]: - raise ValueError(f"Got {len(items)} items but {coords.shape[0]} coordinate rows") - blob = coords_to_blob(coords) - names = json.dumps([item.name for item in items]) - kinds = json.dumps([item.kind for item in items]) - with self._db.transaction() as cursor: - # No-op if the user was deleted while their projection was being computed; there is - # nobody left to serve it to. Guarded by `WHERE EXISTS` for the same reason as in - # `upsert_embedding`. - cursor.execute( - """--sql - INSERT INTO image_projections - (user_id, model_id, scope_hash, params, point_count, image_names, item_kinds, coords) - SELECT ?, ?, ?, ?, ?, ?, ?, ? - WHERE EXISTS (SELECT 1 FROM users WHERE users.user_id = ?) - ON CONFLICT (user_id, model_id) - DO UPDATE SET - scope_hash = excluded.scope_hash, - params = excluded.params, - point_count = excluded.point_count, - image_names = excluded.image_names, - item_kinds = excluded.item_kinds, - coords = excluded.coords; - """, - (user_id, model_id, scope_hash, params, len(items), names, kinds, blob, user_id), - ) - if cursor.rowcount == 0: - self._logger.debug(f"Skipped projection for missing user {user_id}") - - def delete_projection(self, user_id: str, model_id: str) -> None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - DELETE FROM image_projections WHERE user_id = ? AND model_id = ?; - """, - (user_id, model_id), - ) diff --git a/invokeai/app/services/image_moves/image_moves_default.py b/invokeai/app/services/image_moves/image_moves_default.py index a5e857f2eb8..29ca3569ba4 100644 --- a/invokeai/app/services/image_moves/image_moves_default.py +++ b/invokeai/app/services/image_moves/image_moves_default.py @@ -14,7 +14,9 @@ from invokeai.app.services.config import InvokeAIAppConfig from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase from invokeai.app.services.image_records.image_records_common import ImageCategory -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.image_moves import MoveItem from invokeai.app.util.thumbnails import make_thumbnail MoveJobState = Literal["planned", "moving", "moved", "committed", "error"] @@ -69,15 +71,28 @@ class UnreadableImageError(Exception): pass +def _journaled(move: PlannedImageMove) -> MoveItem: + return MoveItem( + image_name=move.image_name, + old_subfolder=move.old_subfolder, + new_subfolder=move.new_subfolder, + is_intermediate=move.is_intermediate, + old_path=str(move.old_path), + new_path=str(move.new_path), + old_thumbnail_path=str(move.old_thumbnail_path), + new_thumbnail_path=str(move.new_thumbnail_path), + ) + + class ImageMoveService: def __init__( self, - db: SqliteDatabase, + database: Database, image_files: ImageFileStorageBase, config: InvokeAIAppConfig, logger, ) -> None: - self._db = db + self._queries = database.queries self.image_files = image_files self._config = config self._logger = logger @@ -156,7 +171,7 @@ def reserve_gallery_maintenance(self) -> Iterator[None]: or (self._future is not None and not self._future.done()) ): raise ImageMoveJobAlreadyRunning("An image storage maintenance operation is already running") - if self._get_active_job_id() is not None: + if self._queries.image_moves.active_job_id() is not None: raise ImageMoveJobAlreadyRunning("An image move job is already active") self._gallery_maintenance_reserved = True @@ -187,7 +202,12 @@ def is_maintenance_active(self) -> bool: is_running = self._future is not None and not self._future.done() operation_reserved = self._future_operation is not None gallery_maintenance_reserved = self._gallery_maintenance_reserved - return gallery_maintenance_reserved or operation_reserved or is_running or self._get_active_job_id() is not None + return ( + gallery_maintenance_reserved + or operation_reserved + or is_running + or self._queries.image_moves.active_job_id() is not None + ) def _assert_no_active_queue_work(self) -> None: session_queue = self._session_queue @@ -210,7 +230,7 @@ def _start_background_operation( raise ImageMoveJobAlreadyRunning("An image storage maintenance operation is already running") if self._future_operation is not None or (self._future is not None and not self._future.done()): raise ImageMoveJobAlreadyRunning("An image move job is already running") - active_job_id = self._get_active_job_id() + active_job_id = self._queries.image_moves.active_job_id() if operation != "recovery" and active_job_id is not None: raise ImageMoveJobAlreadyRunning("An image move job is already active") self._last_background_error = None @@ -239,7 +259,7 @@ def _run_background_operation(self, operation: ImageMoveBackgroundOperation, tar def _record_background_error(self, message: str) -> None: with self._future_lock: self._last_background_error = message - active_job_id = self._get_active_job_id() + active_job_id = self._queries.image_moves.active_job_id() if active_job_id is not None: try: self.record_job_error_message(active_job_id, message) @@ -257,7 +277,7 @@ def _build_background_status_locked(self) -> ImageMoveBackgroundStatus: return ImageMoveBackgroundStatus( is_running=self._future is not None and not self._future.done(), operation=self._future_operation, - active_job_id=self._get_active_job_id(), + active_job_id=self._queries.image_moves.active_job_id(), latest_job=latest_job, last_error=self._last_background_error, needs_move_count=self.count_images_needing_move(), @@ -284,7 +304,7 @@ def move_all_images(self) -> ImageMoveResult: ) errors += plan_errors if not moves: - next_name = self._next_image_name(last_image_name) + next_name = self._queries.image_moves.next_image_name(last_image_name) if next_name is None: break last_image_name = next_name @@ -295,7 +315,7 @@ def move_all_images(self) -> ImageMoveResult: try: self.perform_filesystem_moves(job_id) committed += self.commit_database_updates(job_id) - errors += self._count_job_errors(job_id) + errors += self._queries.image_moves.error_count(job_id) except Exception as e: errors += 1 self.record_job_error_message(job_id, str(e)) @@ -305,15 +325,7 @@ def move_all_images(self) -> ImageMoveResult: return ImageMoveResult(planned=planned, committed=committed, errors=errors) def startup_recovery(self) -> ImageMoveResult: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT id FROM image_subfolder_move_jobs - WHERE state IN ('planned', 'moving', 'moved') - ORDER BY id; - """ - ) - job_ids = [cast(int, row[0]) for row in cursor.fetchall()] + job_ids = self._queries.image_moves.recoverable_job_ids() committed = 0 errors = 0 @@ -329,12 +341,12 @@ def startup_recovery(self) -> ImageMoveResult: except Exception as e: if self._is_unrecoverable_error(e): self.mark_job_unrecoverable(job_id, str(e)) - errors += max(1, self._count_job_errors(job_id)) + errors += max(1, self._queries.image_moves.error_count(job_id)) else: errors += 1 self.record_job_error_message(job_id, str(e)) else: - errors += self._count_job_errors(job_id) + errors += self._queries.image_moves.error_count(job_id) return ImageMoveResult(committed=committed, errors=errors) def plan_batch(self, last_image_name: str, limit: int) -> list[PlannedImageMove]: @@ -342,74 +354,34 @@ def plan_batch(self, last_image_name: str, limit: int) -> list[PlannedImageMove] return moves def count_images_needing_move(self) -> int: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT image_name, image_subfolder, image_category, is_intermediate, created_at - FROM images - WHERE deleted_at IS NULL; - """ - ) - rows = cursor.fetchall() - + # Runs over every image on each status poll, so it calls the subfolder rule directly. count = 0 - for row in rows: - old_subfolder = cast(str, row["image_subfolder"] or "") + for image_name, subfolder, category, is_intermediate, created_at in self._queries.image_moves.placements(): new_subfolder = self._get_new_subfolder( - image_name=cast(str, row["image_name"]), - image_category=ImageCategory(row["image_category"]), - is_intermediate=bool(row["is_intermediate"]), - created_at=row["created_at"], + image_name=image_name, + image_category=ImageCategory(category), + is_intermediate=is_intermediate, + created_at=created_at, ) - if new_subfolder != old_subfolder: + if new_subfolder != subfolder: count += 1 return count def _plan_batch( self, last_image_name: str, limit: int, record_missing_errors: bool ) -> tuple[list[PlannedImageMove], int]: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT image_name, image_subfolder, image_category, is_intermediate, created_at - FROM images - WHERE image_name > ? - AND deleted_at IS NULL - ORDER BY image_name - LIMIT ?; - """, - (last_image_name, limit), - ) - rows = cursor.fetchall() - moves: list[PlannedImageMove] = [] - for row in rows: - image_name = cast(str, row["image_name"]) - old_subfolder = cast(str, row["image_subfolder"] or "") + for image_name, subfolder, category, is_intermediate, created_at in self._queries.image_moves.placements_after( + last_image_name, limit + ): new_subfolder = self._get_new_subfolder( image_name=image_name, - image_category=ImageCategory(row["image_category"]), - is_intermediate=bool(row["is_intermediate"]), - created_at=row["created_at"], - ) - if new_subfolder == old_subfolder: - continue - moves.append( - PlannedImageMove( - image_name=image_name, - old_subfolder=old_subfolder, - new_subfolder=new_subfolder, - is_intermediate=bool(row["is_intermediate"]), - old_path=self.image_files.get_path(image_name, image_subfolder=old_subfolder), - new_path=self.image_files.get_path(image_name, image_subfolder=new_subfolder), - old_thumbnail_path=self.image_files.get_path( - image_name, thumbnail=True, image_subfolder=old_subfolder - ), - new_thumbnail_path=self.image_files.get_path( - image_name, thumbnail=True, image_subfolder=new_subfolder - ), - ) + image_category=ImageCategory(category), + is_intermediate=is_intermediate, + created_at=created_at, ) + if new_subfolder != subfolder: + moves.append(self._planned_move(image_name, subfolder, new_subfolder, is_intermediate)) errors = 0 if record_missing_errors: moves, errors = self._record_missing_source_errors(moves) @@ -419,88 +391,13 @@ def _plan_batch( def create_move_job(self, moves: Sequence[PlannedImageMove]) -> int: if not moves: raise ValueError("Cannot create an image move job with no items") - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT 1 - FROM image_subfolder_move_jobs - WHERE state NOT IN ('committed', 'error') - LIMIT 1; - """ - ) - if cursor.fetchone() is not None: - raise ValueError("Cannot create image move job while another active image move job exists") - cursor.execute("INSERT INTO image_subfolder_move_jobs (state) VALUES ('planned');") - job_id = cast(int, cursor.lastrowid) - cursor.executemany( - """--sql - INSERT INTO image_subfolder_move_items ( - job_id, - image_name, - old_subfolder, - new_subfolder, - is_intermediate, - old_path, - new_path, - old_thumbnail_path, - new_thumbnail_path, - state - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'planned'); - """, - [ - ( - job_id, - move.image_name, - move.old_subfolder, - move.new_subfolder, - int(move.is_intermediate), - str(move.old_path), - str(move.new_path), - str(move.old_thumbnail_path), - str(move.new_thumbnail_path), - ) - for move in moves - ], - ) - return job_id + job_id = self._queries.image_moves.create_job([_journaled(move) for move in moves]) + if job_id is None: + raise ValueError("Cannot create image move job while another active image move job exists") + return job_id def create_error_move_job(self, move: PlannedImageMove, message: str) -> int: - with self._db.transaction() as cursor: - cursor.execute( - "INSERT INTO image_subfolder_move_jobs (state, error_message) VALUES ('error', ?);", - (message,), - ) - job_id = cast(int, cursor.lastrowid) - cursor.execute( - """--sql - INSERT INTO image_subfolder_move_items ( - job_id, - image_name, - old_subfolder, - new_subfolder, - is_intermediate, - old_path, - new_path, - old_thumbnail_path, - new_thumbnail_path, - state, - error_message - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'error', ?); - """, - ( - job_id, - move.image_name, - move.old_subfolder, - move.new_subfolder, - int(move.is_intermediate), - str(move.old_path), - str(move.new_path), - str(move.old_thumbnail_path), - str(move.new_thumbnail_path), - message, - ), - ) - return job_id + return self._queries.image_moves.create_failed_job(_journaled(move), message) def preflight_moves(self, moves: Sequence[PlannedImageMove]) -> None: destinations: set[Path] = set() @@ -520,7 +417,7 @@ def preflight_moves(self, moves: Sequence[PlannedImageMove]) -> None: if move.new_thumbnail_path in thumbnail_destinations: raise ValueError(f"Duplicate destination thumbnail path: {move.new_thumbnail_path}") thumbnail_destinations.add(move.new_thumbnail_path) - if self._has_active_job_for_image(move.image_name): + if self._queries.image_moves.has_active_job_for_image(move.image_name): raise ValueError(f"Image {move.image_name} already has an active image move job") self._assert_same_filesystem(move.old_path, move.new_path) if move.old_thumbnail_path.exists(): @@ -542,10 +439,10 @@ def _record_missing_source_errors(self, moves: Sequence[PlannedImageMove]) -> tu return remaining_moves, errors def perform_filesystem_moves(self, job_id: int) -> None: - self._set_job_state(job_id, "moving") + self._queries.image_moves.set_job_state(job_id, "moving") self.complete_partial_filesystem_moves(job_id) self.cleanup_empty_source_dirs(job_id) - self._set_job_state(job_id, "moved") + self._queries.image_moves.set_job_state(job_id, "moved") def complete_partial_filesystem_moves(self, job_id: int) -> None: items = self._get_items(job_id, include_terminal=False) @@ -636,16 +533,7 @@ def _reconcile_destination_subfolder(self, item: PlannedImageMove) -> None: new_path = self.image_files.get_path(item.image_name, image_subfolder=item.new_subfolder) if old_path.exists() or not new_path.exists(): return - with self._db.transaction() as cursor: - cursor.execute( - """--sql - UPDATE images - SET image_subfolder = ? - WHERE image_name = ? - AND image_subfolder = ?; - """, - (item.new_subfolder, item.image_name, item.old_subfolder), - ) + self._queries.image_moves.repoint_image(item.image_name, item.old_subfolder, item.new_subfolder) def cleanup_empty_source_dirs(self, job_id: int) -> None: for item in self._get_items(job_id): @@ -659,154 +547,45 @@ def cleanup_empty_source_dirs(self, job_id: int) -> None: ) def commit_database_updates(self, job_id: int) -> int: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - UPDATE images - SET image_subfolder = ( - SELECT item.new_subfolder - FROM image_subfolder_move_items AS item - WHERE item.job_id = ? - AND item.image_name = images.image_name - ) - WHERE image_name IN ( - SELECT image_name - FROM image_subfolder_move_items - WHERE job_id = ? - AND state = 'moved' - ) - AND image_subfolder = ( - SELECT item.old_subfolder - FROM image_subfolder_move_items AS item - WHERE item.job_id = ? - AND item.image_name = images.image_name - ); - """, - (job_id, job_id, job_id), - ) - cursor.execute( - """--sql - SELECT COUNT(*) - FROM image_subfolder_move_items AS item - LEFT JOIN images ON images.image_name = item.image_name - WHERE item.job_id = ? - AND item.state = 'moved' - AND ( - images.image_name IS NULL - OR images.deleted_at IS NOT NULL - OR images.image_subfolder != item.new_subfolder - ); - """, - (job_id,), - ) - invalid_count = cast(int, cursor.fetchone()[0]) - if invalid_count: + def commit(q: Queries) -> int: + moved_count = q.image_moves.repoint_moved_images(job_id) + if q.image_moves.invalid_move_count(job_id): + # Raised inside the transaction, so that no record is repointed. raise RuntimeError(f"Image move job {job_id} failed commit validation") - cursor.execute( - "SELECT COUNT(*) FROM image_subfolder_move_items WHERE job_id = ? AND state = 'moved';", - (job_id,), - ) - moved_count = cast(int, cursor.fetchone()[0]) - cursor.execute( - """--sql - SELECT error_message - FROM image_subfolder_move_items - WHERE job_id = ? AND state = 'error' - ORDER BY image_name; - """, - (job_id,), - ) - error_rows = cursor.fetchall() - error_messages = [cast(str, row[0]) for row in error_rows if row[0]] - cursor.execute( - "UPDATE image_subfolder_move_items SET state = 'committed' WHERE job_id = ? AND state = 'moved';", - (job_id,), - ) - if not error_rows: - cursor.execute( - "UPDATE image_subfolder_move_jobs SET state = 'committed', error_message = NULL WHERE id = ?;", - (job_id,), - ) - else: - cursor.execute( - "UPDATE image_subfolder_move_jobs SET state = 'error', error_message = ? WHERE id = ?;", - ("\n".join(error_messages) or "One or more image move items could not be completed", job_id), + error_messages = q.image_moves.error_messages(job_id) + error_message = None + if error_messages: + error_message = "\n".join(message for message in error_messages if message) or ( + "One or more image move items could not be completed" ) - return moved_count + q.image_moves.finish_job(job_id, error_message=error_message) + return moved_count - def _count_job_errors(self, job_id: int) -> int: - with self._db.transaction() as cursor: - cursor.execute( - "SELECT COUNT(*) FROM image_subfolder_move_items WHERE job_id = ? AND state = 'error';", - (job_id,), - ) - return cast(int, cursor.fetchone()[0]) + return self._queries.run(commit) def mark_item_moved(self, job_id: int, image_name: str) -> None: - with self._db.transaction() as cursor: - cursor.execute( - "UPDATE image_subfolder_move_items SET state = 'moved' WHERE job_id = ? AND image_name = ?;", - (job_id, image_name), - ) + self._queries.image_moves.mark_item_moved(job_id, image_name) def record_job_error_message(self, job_id: int, message: str) -> None: - with self._db.transaction() as cursor: - cursor.execute( - "UPDATE image_subfolder_move_jobs SET error_message = ? WHERE id = ?;", - (message, job_id), - ) + self._queries.image_moves.set_job_error_message(job_id, message) def mark_item_unrecoverable(self, job_id: int, image_name: str, message: str) -> None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - UPDATE image_subfolder_move_items - SET state = 'error', error_message = ? - WHERE job_id = ? AND image_name = ?; - """, - (message, job_id, image_name), - ) + self._queries.image_moves.fail_item(job_id, image_name, message) def mark_job_unrecoverable(self, job_id: int, message: str) -> None: - with self._db.transaction() as cursor: - cursor.execute( - "UPDATE image_subfolder_move_jobs SET state = 'error', error_message = ? WHERE id = ?;", - (message, job_id), - ) - cursor.execute( - "UPDATE image_subfolder_move_items SET state = 'error', error_message = ? WHERE job_id = ?;", - (message, job_id), - ) + self._queries.image_moves.fail_job(job_id, message) def get_job(self, job_id: int) -> ImageMoveJob: - with self._db.transaction() as cursor: - cursor.execute( - "SELECT id, state, error_message FROM image_subfolder_move_jobs WHERE id = ?;", - (job_id,), - ) - row = cursor.fetchone() - if row is None: + job = self._queries.image_moves.job(job_id) + if job is None: raise ValueError(f"Image move job not found: {job_id}") - return ImageMoveJob( - id=cast(int, row["id"]), state=cast(MoveJobState, row["state"]), error_message=row["error_message"] - ) + return ImageMoveJob(id=job.id, state=cast(MoveJobState, job.state), error_message=job.error_message) def get_latest_job(self) -> ImageMoveJob | None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT id, state, error_message - FROM image_subfolder_move_jobs - ORDER BY id DESC - LIMIT 1; - """ - ) - row = cursor.fetchone() - if row is None: + job = self._queries.image_moves.latest_job() + if job is None: return None - return ImageMoveJob( - id=cast(int, row["id"]), state=cast(MoveJobState, row["state"]), error_message=row["error_message"] - ) + return ImageMoveJob(id=job.id, state=cast(MoveJobState, job.state), error_message=job.error_message) def _get_new_subfolder( self, image_name: str, image_category: ImageCategory, is_intermediate: bool, created_at: str | datetime @@ -824,83 +603,26 @@ def _get_new_subfolder( raise ValueError(f"Unknown image subfolder strategy: {strategy}") def _get_items(self, job_id: int, include_terminal: bool = True) -> list[PlannedImageMove]: - with self._db.transaction() as cursor: - query = """--sql - SELECT image_name, old_subfolder, new_subfolder, is_intermediate - FROM image_subfolder_move_items - WHERE job_id = ? - """ - params: tuple[object, ...] = (job_id,) - if not include_terminal: - query += " AND state NOT IN ('committed', 'error')" - query += " ORDER BY image_name;" - cursor.execute(query, params) - rows = cursor.fetchall() return [ - PlannedImageMove( - image_name=row["image_name"], - old_subfolder=row["old_subfolder"], - new_subfolder=row["new_subfolder"], - is_intermediate=bool(row["is_intermediate"]), - old_path=self.image_files.get_path(row["image_name"], image_subfolder=row["old_subfolder"]), - new_path=self.image_files.get_path(row["image_name"], image_subfolder=row["new_subfolder"]), - old_thumbnail_path=self.image_files.get_path( - row["image_name"], thumbnail=True, image_subfolder=row["old_subfolder"] - ), - new_thumbnail_path=self.image_files.get_path( - row["image_name"], thumbnail=True, image_subfolder=row["new_subfolder"] - ), + self._planned_move(image_name, old_subfolder, new_subfolder, is_intermediate) + for image_name, old_subfolder, new_subfolder, is_intermediate in self._queries.image_moves.items( + job_id, unfinished_only=not include_terminal ) - for row in rows ] - def _set_job_state(self, job_id: int, state: MoveJobState) -> None: - with self._db.transaction() as cursor: - cursor.execute("UPDATE image_subfolder_move_jobs SET state = ? WHERE id = ?;", (state, job_id)) - - def _has_active_job_for_image(self, image_name: str) -> bool: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT 1 - FROM image_subfolder_move_items AS item - JOIN image_subfolder_move_jobs AS job ON job.id = item.job_id - WHERE item.image_name = ? - AND job.state NOT IN ('committed', 'error') - LIMIT 1; - """, - (image_name,), - ) - return cursor.fetchone() is not None - - def _get_active_job_id(self) -> int | None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT id - FROM image_subfolder_move_jobs - WHERE state NOT IN ('committed', 'error') - ORDER BY id - LIMIT 1; - """ - ) - row = cursor.fetchone() - return None if row is None else cast(int, row["id"]) - - def _next_image_name(self, last_image_name: str) -> str | None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT image_name FROM images - WHERE image_name > ? - AND deleted_at IS NULL - ORDER BY image_name - LIMIT 1; - """, - (last_image_name,), - ) - row = cursor.fetchone() - return None if row is None else cast(str, row[0]) + def _planned_move( + self, image_name: str, old_subfolder: str, new_subfolder: str, is_intermediate: bool + ) -> PlannedImageMove: + return PlannedImageMove( + image_name=image_name, + old_subfolder=old_subfolder, + new_subfolder=new_subfolder, + is_intermediate=is_intermediate, + old_path=self.image_files.get_path(image_name, image_subfolder=old_subfolder), + new_path=self.image_files.get_path(image_name, image_subfolder=new_subfolder), + old_thumbnail_path=self.image_files.get_path(image_name, thumbnail=True, image_subfolder=old_subfolder), + new_thumbnail_path=self.image_files.get_path(image_name, thumbnail=True, image_subfolder=new_subfolder), + ) def _regenerate_thumbnail(self, image_path: Path, thumbnail_path: Path) -> None: thumbnail_path.parent.mkdir(parents=True, exist_ok=True) diff --git a/invokeai/app/services/image_records/image_records_base.py b/invokeai/app/services/image_records/image_records_base.py index 828f1c47203..541c66e8cfd 100644 --- a/invokeai/app/services/image_records/image_records_base.py +++ b/invokeai/app/services/image_records/image_records_base.py @@ -12,9 +12,7 @@ ResourceOrigin, ) from invokeai.app.services.shared.intermediate_delete import IntermediateDeleteGuard -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection -from invokeai.app.services.virtual_boards.virtual_boards_common import VirtualSubBoardDTO +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection class ImageRecordStorageBase(ABC): @@ -177,15 +175,6 @@ def get_image_names( """ pass - @abstractmethod - def get_image_dates( - self, - user_id: Optional[str] = None, - is_admin: bool = False, - ) -> list[VirtualSubBoardDTO]: - """Gets a list of dates with image counts, grouped by DATE(created_at).""" - pass - @abstractmethod def get_image_names_by_date( self, diff --git a/invokeai/app/services/image_records/image_records_common.py b/invokeai/app/services/image_records/image_records_common.py index 6436d1b003b..9d17cdf35fa 100644 --- a/invokeai/app/services/image_records/image_records_common.py +++ b/invokeai/app/services/image_records/image_records_common.py @@ -104,38 +104,6 @@ def __init__(self, message="Image record not saved"): super().__init__(message) -class ImageRecordDeleteException(Exception): - """Raised when an image record cannot be deleted.""" - - def __init__(self, message="Image record not deleted"): - super().__init__(message) - - -IMAGE_DTO_COLS = ", ".join( - [ - "images." + c - for c in [ - "image_name", - "image_origin", - "image_category", - "width", - "height", - "session_id", - "node_id", - "has_workflow", - "is_intermediate", - "created_at", - "updated_at", - "deleted_at", - "starred", - "image_subfolder", - "project_id", - "file_size_bytes", - ] - ] -) - - class ImageRecord(BaseModelExcludeNull): """Deserialized image record without metadata.""" diff --git a/invokeai/app/services/image_records/image_records_default.py b/invokeai/app/services/image_records/image_records_default.py new file mode 100644 index 00000000000..86a8c29e0ea --- /dev/null +++ b/invokeai/app/services/image_records/image_records_default.py @@ -0,0 +1,236 @@ +from collections.abc import Iterator +from datetime import datetime +from typing import Optional + +from invokeai.app.invocations.fields import MetadataField, MetadataFieldValidator +from invokeai.app.services.image_records.image_records_base import ImageRecordStorageBase +from invokeai.app.services.image_records.image_records_common import ( + ImageCategory, + ImageNamesResult, + ImageRecord, + ImageRecordChanges, + ImageRecordNotFoundException, + ImageRecordSaveException, + ResourceOrigin, +) +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.locks import DatabaseLock +from invokeai.app.services.shared.intermediate_delete import IntermediateDeleteGuard +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection + + +class ImageRecordStorage(ImageRecordStorageBase): + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + + def get(self, image_name: str) -> ImageRecord: + # A database error is deliberately NOT turned into ImageRecordNotFoundException: that made the exception mean + # "the row is absent, OR the database is locked/corrupt/unreadable". Callers that treat not-found as a benign + # outcome -- the concurrent-deletion skips in the images and board_images batch routes -- would then swallow a + # disk I/O error as a routine race and answer 200 with the name in no result list at all. + record = self._queries.images.get(image_name) + if record is None: + raise ImageRecordNotFoundException + return record + + def set_file_size_bytes(self, image_name: str, file_size_bytes: Optional[int]) -> None: + self._queries.images.set_file_size(image_name, file_size_bytes) + + def set_file_sizes_bytes(self, sizes: dict[str, int]) -> None: + self._queries.images.fill_file_sizes(sizes) + + def get_user_id(self, image_name: str) -> Optional[str]: + return self._queries.images.user_id(image_name) + + def get_metadata(self, image_name: str) -> Optional[MetadataField]: + # See get(): a database error must not masquerade as a missing row. + exists, metadata = self._queries.images.metadata(image_name) + if not exists: + raise ImageRecordNotFoundException + return MetadataFieldValidator.validate_json(metadata) if metadata is not None else None + + def exists(self, image_name: str) -> bool: + return self._queries.images.exists(image_name) + + def update( + self, + image_name: str, + changes: ImageRecordChanges, + ) -> None: + self._queries.images.update(image_name, changes) + + def get_many( + self, + offset: int = 0, + limit: int = 10, + starred_first: bool = True, + order_dir: SQLiteDirection = SQLiteDirection.Descending, + image_origin: Optional[ResourceOrigin] = None, + categories: Optional[list[ImageCategory]] = None, + is_intermediate: Optional[bool] = None, + board_id: Optional[str] = None, + search_term: Optional[str] = None, + created_from: Optional[str] = None, + created_to: Optional[str] = None, + user_id: Optional[str] = None, + is_admin: bool = False, + ) -> OffsetPaginatedResults[ImageRecord]: + images, total = self._queries.images.page( + offset=offset, + limit=limit, + starred_first=starred_first, + descending=order_dir == SQLiteDirection.Descending, + image_origin=image_origin, + categories=categories, + is_intermediate=is_intermediate, + board_id=board_id, + search_term=search_term, + created_from=created_from, + created_to=created_to, + user_id=user_id, + is_admin=is_admin, + ) + return OffsetPaginatedResults(items=images, offset=offset, limit=limit, total=total) + + def delete(self, image_name: str) -> None: + self._queries.images.delete(image_name) + + def delete_many(self, image_names: list[str]) -> None: + self._queries.images.delete_many(image_names) + + def iter_all_image_locations(self, batch_size: int = 500) -> Iterator[tuple[str, str]]: + """Yields every image location through bounded keyset pages, including hidden/intermediate rows.""" + if batch_size <= 0: + raise ValueError("batch_size must be positive") + last_image_name = "" + while True: + # A transaction per page: none is held while the caller works through it. + locations = self._queries.images.locations_after(last_image_name, batch_size) + yield from locations + if len(locations) < batch_size: + return + last_image_name = locations[-1][0] + + def get_subfolders(self, image_names: list[str]) -> dict[str, str]: + return self._queries.images.subfolders(image_names) + + def delete_intermediates_by_names( + self, image_names: list[str], guard: Optional[IntermediateDeleteGuard] = None + ) -> list[str]: + """Deletes the named image records, skipping any that are no longer intermediates. + + Returns the names whose records this call actually removed. Names that were already gone, and names whose + records survive because they are no longer intermediates, are both excluded -- the caller purges the files of + exactly the returned names and touches nothing else. + + ``guard`` narrows the names on this same transaction; see `IntermediateDeleteGuard`. + """ + + def delete(q: Queries) -> list[str]: + if guard is None: + return q.images.delete_intermediates(image_names) + # Held until the delete commits, so nothing becomes protected between the guard's answer and the delete. + q.locks.acquire(DatabaseLock.MEDIA_PROTECTION) + return q.images.delete_intermediates(guard(q, image_names)) + + return self._queries.run(delete) + + def save( + self, + image_name: str, + image_origin: ResourceOrigin, + image_category: ImageCategory, + width: int, + height: int, + has_workflow: bool, + is_intermediate: Optional[bool] = False, + starred: Optional[bool] = False, + session_id: Optional[str] = None, + node_id: Optional[str] = None, + metadata: Optional[str] = None, + user_id: Optional[str] = None, + image_subfolder: str = "", + project_id: Optional[str] = None, + ) -> datetime: + def insert(q: Queries) -> Optional[str]: + q.images.insert( + image_name=image_name, + image_origin=image_origin, + image_category=image_category, + width=width, + height=height, + has_workflow=has_workflow, + is_intermediate=is_intermediate, + starred=starred, + session_id=session_id, + node_id=node_id, + metadata=metadata, + user_id=user_id or "system", + image_subfolder=image_subfolder, + project_id=project_id, + ) + # Read back: an image of that name that exists already keeps its row, and its time is the one returned. + return q.images.created_at(image_name) + + created_at = self._queries.run(insert) + if created_at is None: + raise ImageRecordSaveException + return datetime.fromisoformat(created_at) + + def get_most_recent_image_for_board(self, board_id: str) -> Optional[ImageRecord]: + return self._queries.images.most_recent_on_board(board_id) + + def get_image_names( + self, + starred_first: bool = True, + order_dir: SQLiteDirection = SQLiteDirection.Descending, + image_origin: Optional[ResourceOrigin] = None, + categories: Optional[list[ImageCategory]] = None, + is_intermediate: Optional[bool] = None, + board_id: Optional[str] = None, + search_term: Optional[str] = None, + created_from: Optional[str] = None, + created_to: Optional[str] = None, + user_id: Optional[str] = None, + is_admin: bool = False, + ) -> ImageNamesResult: + image_names, starred_count = self._queries.images.names( + starred_first=starred_first, + descending=order_dir == SQLiteDirection.Descending, + image_origin=image_origin, + categories=categories, + is_intermediate=is_intermediate, + board_id=board_id, + search_term=search_term, + created_from=created_from, + created_to=created_to, + user_id=user_id, + is_admin=is_admin, + ) + return ImageNamesResult(image_names=image_names, starred_count=starred_count, total_count=len(image_names)) + + def get_image_names_by_date( + self, + date: str, + starred_first: bool = True, + order_dir: SQLiteDirection = SQLiteDirection.Descending, + categories: Optional[list[ImageCategory]] = None, + search_term: Optional[str] = None, + user_id: Optional[str] = None, + is_admin: bool = False, + ) -> ImageNamesResult: + # The day's images that are not intermediates: from its start up to the next day's. A date that is no ISO day + # matches nothing. + return self.get_image_names( + starred_first=starred_first, + order_dir=order_dir, + categories=categories, + is_intermediate=False, + search_term=search_term, + created_from=date, + created_to=date, + user_id=user_id, + is_admin=is_admin, + ) diff --git a/invokeai/app/services/image_records/image_records_sqlite.py b/invokeai/app/services/image_records/image_records_sqlite.py deleted file mode 100644 index a3e7ff2af80..00000000000 --- a/invokeai/app/services/image_records/image_records_sqlite.py +++ /dev/null @@ -1,853 +0,0 @@ -import sqlite3 -from collections.abc import Iterator -from dataclasses import dataclass -from datetime import datetime -from typing import Optional, Union, cast - -from invokeai.app.invocations.fields import MetadataField, MetadataFieldValidator -from invokeai.app.services.image_records.image_records_base import ImageRecordStorageBase -from invokeai.app.services.image_records.image_records_common import ( - ASSETS_CATEGORIES, - IMAGE_CATEGORIES, - IMAGE_DTO_COLS, - ImageCategory, - ImageNamesResult, - ImageRecord, - ImageRecordChanges, - ImageRecordDeleteException, - ImageRecordNotFoundException, - ImageRecordSaveException, - ResourceOrigin, - deserialize_image_record, -) -from invokeai.app.services.shared.intermediate_delete import IntermediateDeleteGuard -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.virtual_boards.virtual_boards_common import VirtualSubBoardDTO - - -@dataclass(frozen=True) -class _ImageQueryFilters: - image_origin: Optional[ResourceOrigin] = None - categories: Optional[list[ImageCategory]] = None - is_intermediate: Optional[bool] = None - board_id: Optional[str] = None - search_term: Optional[str] = None - created_from: Optional[str] = None - created_to: Optional[str] = None - user_id: Optional[str] = None - is_admin: bool = False - - -# Every uncategorized image plus every image on an active (non-archived) board: -# the board_id="all" visibility scope for admins and single-user installs. -_ALL_ACTIVE_BOARDS_CONDITION = """( - board_images.board_id IS NULL - OR EXISTS ( - SELECT 1 - FROM boards - WHERE boards.board_id = board_images.board_id - AND boards.archived = 0 - ) - )""" - - -def _build_image_query_conditions( - filters: _ImageQueryFilters, *, use_board_join: bool = True -) -> tuple[str, list[Union[int, str, bool]]]: - """Build shared image filters. - - Paginated DTO queries keep the outer board join because they return board - data. Names-only queries use correlated membership checks so SQLite can - scan images directly without multiplying rows or probing a redundant join. - """ - conditions: list[str] = [] - params: list[Union[int, str, bool]] = [] - - if filters.image_origin is not None: - conditions.append("images.image_origin = ?") - params.append(filters.image_origin.value) - - if filters.categories is not None: - category_strings = [category.value for category in set(filters.categories)] - placeholders = ",".join("?" * len(category_strings)) - conditions.append(f"images.image_category IN ( {placeholders} )") - params.extend(category_strings) - - if filters.is_intermediate is not None: - conditions.append("images.is_intermediate = ?") - params.append(filters.is_intermediate) - - if filters.board_id == "none": - conditions.append( - "board_images.board_id IS NULL" - if use_board_join - else """NOT EXISTS ( - SELECT 1 - FROM board_images - WHERE board_images.image_name = images.image_name - )""" - ) - if filters.user_id is not None and not filters.is_admin: - conditions.append("images.user_id = ?") - params.append(filters.user_id) - elif filters.board_id == "all": - if use_board_join: - if filters.is_admin: - conditions.append(_ALL_ACTIVE_BOARDS_CONDITION) - elif filters.user_id is not None: - conditions.append( - """( - (board_images.board_id IS NULL AND images.user_id = ?) - OR EXISTS ( - SELECT 1 - FROM boards - WHERE boards.board_id = board_images.board_id - AND boards.archived = 0 - AND ( - boards.user_id = ? - OR boards.board_visibility IN ('shared', 'public') - OR EXISTS ( - SELECT 1 - FROM shared_boards - WHERE shared_boards.board_id = boards.board_id - AND shared_boards.user_id = ? - ) - ) - ) - )""" - ) - params.extend([filters.user_id, filters.user_id, filters.user_id]) - else: - # Single-user mode has no current user; it reads the administrative scope. - conditions.append(_ALL_ACTIVE_BOARDS_CONDITION) - elif filters.is_admin or filters.user_id is None: - conditions.append( - """( - NOT EXISTS ( - SELECT 1 - FROM board_images - WHERE board_images.image_name = images.image_name - ) - OR EXISTS ( - SELECT 1 - FROM board_images - INNER JOIN boards ON boards.board_id = board_images.board_id - WHERE board_images.image_name = images.image_name - AND boards.archived = 0 - ) - )""" - ) - else: - conditions.append( - """( - ( - NOT EXISTS ( - SELECT 1 - FROM board_images - WHERE board_images.image_name = images.image_name - ) - AND images.user_id = ? - ) - OR EXISTS ( - SELECT 1 - FROM board_images - INNER JOIN boards ON boards.board_id = board_images.board_id - WHERE board_images.image_name = images.image_name - AND boards.archived = 0 - AND ( - boards.user_id = ? - OR boards.board_visibility IN ('shared', 'public') - OR EXISTS ( - SELECT 1 - FROM shared_boards - WHERE shared_boards.board_id = boards.board_id - AND shared_boards.user_id = ? - ) - ) - ) - )""" - ) - params.extend([filters.user_id, filters.user_id, filters.user_id]) - elif filters.board_id is not None: - conditions.append( - "board_images.board_id = ?" - if use_board_join - else """EXISTS ( - SELECT 1 - FROM board_images - WHERE board_images.image_name = images.image_name - AND board_images.board_id = ? - )""" - ) - params.append(filters.board_id) - elif filters.user_id is not None and not filters.is_admin: - conditions.append("images.user_id = ?") - params.append(filters.user_id) - - if filters.search_term: - conditions.append("(images.metadata LIKE ? OR images.created_at LIKE ?)") - search_pattern = f"%{filters.search_term.lower()}%" - params.extend([search_pattern, search_pattern]) - - if filters.created_from: - conditions.append("images.created_at >= ?") - params.append(filters.created_from) - - if filters.created_to: - conditions.append("images.created_at < DATE(?, '+1 day')") - params.append(filters.created_to) - - return "".join(f"\nAND {condition}" for condition in conditions), params - - -class SqliteImageRecordStorage(ImageRecordStorageBase): - # Conservative bound on bound parameters per statement. SQLITE_MAX_VARIABLE_NUMBER defaults to - # 999 on SQLite builds older than 3.32, and an image library can hold far more intermediates. - _MAX_SQL_VARIABLES = 500 - - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - - def get(self, image_name: str) -> ImageRecord: - # A sqlite3.Error is deliberately NOT translated into ImageRecordNotFoundException. - # This used to be caught and re-raised as not-found, which made the exception mean - # "the row is absent, OR the database is locked/corrupt/unreadable". Callers that - # treat not-found as a benign outcome — the concurrent-deletion skips in the images - # and board_images batch routes — would then swallow a disk I/O error as a routine - # race and answer 200 with the name in no result list at all. Let the storage error - # propagate: ImageService logs it and the route reports it as a real failure. - with self._db.transaction() as cursor: - cursor.execute( - f"""--sql - SELECT {IMAGE_DTO_COLS} FROM images - WHERE image_name = ?; - """, - (image_name,), - ) - - result = cast(Optional[sqlite3.Row], cursor.fetchone()) - - if not result: - raise ImageRecordNotFoundException - - return deserialize_image_record(dict(result)) - - def set_file_size_bytes(self, image_name: str, file_size_bytes: Optional[int]) -> None: - with self._db.transaction() as cursor: - try: - cursor.execute( - "UPDATE images SET file_size_bytes = ? WHERE image_name = ?;", - (file_size_bytes, image_name), - ) - except sqlite3.Error as e: - raise ImageRecordSaveException from e - - def set_file_sizes_bytes(self, sizes: dict[str, int]) -> None: - if not sizes: - return - with self._db.transaction() as cursor: - try: - # Backfill only fills gaps: the writer's own measurement, taken after the file exists, - # wins over one taken before it was written. - cursor.executemany( - "UPDATE images SET file_size_bytes = ? WHERE image_name = ? AND file_size_bytes IS NULL;", - [(size, name) for name, size in sizes.items()], - ) - except sqlite3.Error as e: - raise ImageRecordSaveException from e - - def get_user_id(self, image_name: str) -> Optional[str]: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT user_id FROM images - WHERE image_name = ?; - """, - (image_name,), - ) - result = cast(Optional[sqlite3.Row], cursor.fetchone()) - if not result: - return None - return cast(Optional[str], dict(result).get("user_id")) - - def get_metadata(self, image_name: str) -> Optional[MetadataField]: - # See get(): a storage error must not masquerade as a missing row. - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT metadata FROM images - WHERE image_name = ?; - """, - (image_name,), - ) - - result = cast(Optional[sqlite3.Row], cursor.fetchone()) - - if not result: - raise ImageRecordNotFoundException - - as_dict = dict(result) - metadata_raw = cast(Optional[str], as_dict.get("metadata", None)) - return MetadataFieldValidator.validate_json(metadata_raw) if metadata_raw is not None else None - - def exists(self, image_name: str) -> bool: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT 1 FROM images - WHERE image_name = ? - LIMIT 1; - """, - (image_name,), - ) - return cursor.fetchone() is not None - - def update( - self, - image_name: str, - changes: ImageRecordChanges, - ) -> None: - with self._db.transaction() as cursor: - try: - # Change the category of the image - if changes.image_category is not None: - cursor.execute( - """--sql - UPDATE images - SET image_category = ? - WHERE image_name = ?; - """, - (changes.image_category, image_name), - ) - - # Change the session associated with the image - if changes.session_id is not None: - cursor.execute( - """--sql - UPDATE images - SET session_id = ? - WHERE image_name = ?; - """, - (changes.session_id, image_name), - ) - - # Change the image's `is_intermediate`` flag - if changes.is_intermediate is not None: - cursor.execute( - """--sql - UPDATE images - SET is_intermediate = ? - WHERE image_name = ?; - """, - (changes.is_intermediate, image_name), - ) - - # Change the image's `starred`` state - if changes.starred is not None: - cursor.execute( - """--sql - UPDATE images - SET starred = ? - WHERE image_name = ?; - """, - (changes.starred, image_name), - ) - - except sqlite3.Error as e: - raise ImageRecordSaveException from e - - def get_many( - self, - offset: int = 0, - limit: int = 10, - starred_first: bool = True, - order_dir: SQLiteDirection = SQLiteDirection.Descending, - image_origin: Optional[ResourceOrigin] = None, - categories: Optional[list[ImageCategory]] = None, - is_intermediate: Optional[bool] = None, - board_id: Optional[str] = None, - search_term: Optional[str] = None, - created_from: Optional[str] = None, - created_to: Optional[str] = None, - user_id: Optional[str] = None, - is_admin: bool = False, - ) -> OffsetPaginatedResults[ImageRecord]: - with self._db.transaction() as cursor: - # Manually build two queries - one for the count, one for the records - count_query = """--sql - SELECT COUNT(*) - FROM images - LEFT JOIN board_images ON board_images.image_name = images.image_name - WHERE 1=1 - """ - - images_query = f"""--sql - SELECT {IMAGE_DTO_COLS} - FROM images - LEFT JOIN board_images ON board_images.image_name = images.image_name - WHERE 1=1 - """ - - query_conditions, query_params = _build_image_query_conditions( - _ImageQueryFilters( - image_origin=image_origin, - categories=categories, - is_intermediate=is_intermediate, - board_id=board_id, - search_term=search_term, - created_from=created_from, - created_to=created_to, - user_id=user_id, - is_admin=is_admin, - ) - ) - - if starred_first: - query_pagination = f"""--sql - ORDER BY images.starred DESC, images.created_at {order_dir.value} LIMIT ? OFFSET ? - """ - else: - query_pagination = f"""--sql - ORDER BY images.created_at {order_dir.value} LIMIT ? OFFSET ? - """ - - # Final images query with pagination - images_query += query_conditions + query_pagination + ";" - # Add all the parameters - images_params = query_params.copy() - # Add the pagination parameters - images_params.extend([limit, offset]) - - # Build the list of images, deserializing each row - cursor.execute(images_query, images_params) - result = cast(list[sqlite3.Row], cursor.fetchall()) - - images = [deserialize_image_record(dict(r)) for r in result] - - # Set up and execute the count query, without pagination - count_query += query_conditions + ";" - count_params = query_params.copy() - cursor.execute(count_query, count_params) - count = cast(int, cursor.fetchone()[0]) - - return OffsetPaginatedResults(items=images, offset=offset, limit=limit, total=count) - - def delete(self, image_name: str) -> None: - with self._db.transaction() as cursor: - try: - cursor.execute( - """--sql - DELETE FROM images - WHERE image_name = ?; - """, - (image_name,), - ) - except sqlite3.Error as e: - raise ImageRecordDeleteException from e - - def delete_many(self, image_names: list[str]) -> None: - with self._db.transaction() as cursor: - try: - placeholders = ",".join("?" for _ in image_names) - - # Construct the SQLite query with the placeholders - query = f"DELETE FROM images WHERE image_name IN ({placeholders})" - - # Execute the query with the list of IDs as parameters - cursor.execute(query, image_names) - - except sqlite3.Error as e: - raise ImageRecordDeleteException from e - - def get_subfolders(self, image_names: list[str]) -> dict[str, str]: - subfolders: dict[str, str] = {} - with self._db.transaction() as cursor: - for start in range(0, len(image_names), self._MAX_SQL_VARIABLES): - chunk = image_names[start : start + self._MAX_SQL_VARIABLES] - placeholders = ",".join("?" for _ in chunk) - cursor.execute( - f"SELECT image_name, image_subfolder FROM images WHERE image_name IN ({placeholders})", chunk - ) - for row in cursor.fetchall(): - subfolders[cast(str, row[0])] = cast(str, row[1]) - return subfolders - - def iter_all_image_locations(self, batch_size: int = 500) -> Iterator[tuple[str, str]]: - """Yields every image location through bounded keyset pages, including hidden/intermediate rows.""" - if batch_size <= 0: - raise ValueError("batch_size must be positive") - - last_image_name = "" - while True: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT image_name, image_subfolder - FROM images - WHERE image_name > ? - ORDER BY image_name - LIMIT ?; - """, - (last_image_name, batch_size), - ) - rows = cursor.fetchall() - - if not rows: - return - - locations = [(cast(str, row[0]), cast(str, row[1])) for row in rows] - for image_name, image_subfolder in locations: - last_image_name = image_name - yield image_name, image_subfolder - - if len(rows) < batch_size: - return - - def delete_intermediates_by_names( - self, image_names: list[str], guard: Optional[IntermediateDeleteGuard] = None - ) -> list[str]: - """Deletes the named image records, skipping any that are no longer intermediates. - - The ``is_intermediate`` predicate rides on the DELETE itself rather than on a preceding - SELECT, so an image promoted out of intermediate status keeps its record however the - promotion interleaves with this call. (Python's legacy sqlite3 transaction control opens a - transaction only before a write, so a SELECT here holds no read lock to rely on.) - - Returns the names whose records this call actually removed. Names that were already gone, and - names whose records survive because they are no longer intermediates, are both excluded — the - caller purges the files of exactly the returned names and touches nothing else. - - ``guard`` narrows each chunk on this same transaction; see `IntermediateDeleteGuard`. - """ - deleted: list[str] = [] - try: - with self._db.transaction() as cursor: - # Chunked to stay under SQLITE_MAX_VARIABLE_NUMBER; every chunk runs inside the one - # transaction above. - for start in range(0, len(image_names), self._MAX_SQL_VARIABLES): - chunk = image_names[start : start + self._MAX_SQL_VARIABLES] - if guard is not None: - chunk = guard(cursor, chunk) - if not chunk: - continue - placeholders = ",".join("?" for _ in chunk) - select_query = f"SELECT image_name FROM images WHERE image_name IN ({placeholders})" - - cursor.execute(select_query, chunk) - present_before = {cast(str, r[0]) for r in cursor.fetchall()} - cursor.execute( - f"DELETE FROM images WHERE image_name IN ({placeholders}) AND is_intermediate = TRUE", - chunk, - ) - cursor.execute(select_query, chunk) - present_after = {cast(str, r[0]) for r in cursor.fetchall()} - - deleted.extend(name for name in chunk if name in present_before and name not in present_after) - except sqlite3.Error as e: - # The try wraps the context manager so a failure in its commit is reported too. - raise ImageRecordDeleteException from e - return deleted - - def save( - self, - image_name: str, - image_origin: ResourceOrigin, - image_category: ImageCategory, - width: int, - height: int, - has_workflow: bool, - is_intermediate: Optional[bool] = False, - starred: Optional[bool] = False, - session_id: Optional[str] = None, - node_id: Optional[str] = None, - metadata: Optional[str] = None, - user_id: Optional[str] = None, - image_subfolder: str = "", - project_id: Optional[str] = None, - ) -> datetime: - with self._db.transaction() as cursor: - try: - cursor.execute( - """--sql - INSERT OR IGNORE INTO images ( - image_name, - image_origin, - image_category, - width, - height, - node_id, - session_id, - metadata, - is_intermediate, - starred, - has_workflow, - user_id, - image_subfolder, - project_id - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?); - """, - ( - image_name, - image_origin.value, - image_category.value, - width, - height, - node_id, - session_id, - metadata, - is_intermediate, - starred, - has_workflow, - user_id or "system", - image_subfolder, - project_id, - ), - ) - - cursor.execute( - """--sql - SELECT created_at - FROM images - WHERE image_name = ?; - """, - (image_name,), - ) - - created_at = datetime.fromisoformat(cursor.fetchone()[0]) - - except sqlite3.Error as e: - raise ImageRecordSaveException from e - return created_at - - def get_most_recent_image_for_board(self, board_id: str) -> Optional[ImageRecord]: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT images.* - FROM images - JOIN board_images ON images.image_name = board_images.image_name - WHERE board_images.board_id = ? - AND images.is_intermediate = FALSE - ORDER BY images.starred DESC, images.created_at DESC - LIMIT 1; - """, - (board_id,), - ) - - result = cast(Optional[sqlite3.Row], cursor.fetchone()) - - if result is None: - return None - - return deserialize_image_record(dict(result)) - - def get_image_names( - self, - starred_first: bool = True, - order_dir: SQLiteDirection = SQLiteDirection.Descending, - image_origin: Optional[ResourceOrigin] = None, - categories: Optional[list[ImageCategory]] = None, - is_intermediate: Optional[bool] = None, - board_id: Optional[str] = None, - search_term: Optional[str] = None, - created_from: Optional[str] = None, - created_to: Optional[str] = None, - user_id: Optional[str] = None, - is_admin: bool = False, - ) -> ImageNamesResult: - with self._db.transaction() as cursor: - query_conditions, query_params = _build_image_query_conditions( - _ImageQueryFilters( - image_origin=image_origin, - categories=categories, - is_intermediate=is_intermediate, - board_id=board_id, - search_term=search_term, - created_from=created_from, - created_to=created_to, - user_id=user_id, - is_admin=is_admin, - ), - use_board_join=False, - ) - - # Get starred count if starred_first is enabled - starred_count = 0 - if starred_first: - starred_count_query = f"""--sql - SELECT COUNT(*) - FROM images - WHERE images.starred = TRUE AND (1=1{query_conditions}) - """ - cursor.execute(starred_count_query, query_params) - starred_count = cast(int, cursor.fetchone()[0]) - - # Get all image names with proper ordering - if starred_first: - names_query = f"""--sql - SELECT images.image_name - FROM images - WHERE 1=1{query_conditions} - ORDER BY images.starred DESC, images.created_at {order_dir.value} - """ - else: - names_query = f"""--sql - SELECT images.image_name - FROM images - WHERE 1=1{query_conditions} - ORDER BY images.created_at {order_dir.value} - """ - - cursor.execute(names_query, query_params) - result = cast(list[sqlite3.Row], cursor.fetchall()) - image_names = [row[0] for row in result] - - return ImageNamesResult(image_names=image_names, starred_count=starred_count, total_count=len(image_names)) - - def get_image_dates( - self, - user_id: Optional[str] = None, - is_admin: bool = False, - ) -> list[VirtualSubBoardDTO]: - with self._db.transaction() as cursor: - query_conditions = "" - query_params: list[Union[int, str, bool]] = [] - - # Only non-intermediate images - query_conditions += """--sql - AND images.is_intermediate = 0 - """ - - # User isolation for non-admin users - if user_id is not None and not is_admin: - query_conditions += """--sql - AND images.user_id = ? - """ - query_params.append(user_id) - - # Derived from the category constants rather than `!= 'general'`, so a - # category belonging to neither view (OTHER, owned by canvas layers) is - # counted by neither. These are enum values, never user input. - image_categories_sql = ",".join(f"'{c.value}'" for c in IMAGE_CATEGORIES) - assets_categories_sql = ",".join(f"'{c.value}'" for c in ASSETS_CATEGORIES) - - query = f"""--sql - SELECT - DATE(images.created_at) as date, - SUM(CASE WHEN images.image_category IN ({image_categories_sql}) THEN 1 ELSE 0 END) as image_count, - SUM(CASE WHEN images.image_category IN ({assets_categories_sql}) THEN 1 ELSE 0 END) as asset_count, - ( - SELECT i2.image_name FROM images i2 - WHERE DATE(i2.created_at) = DATE(images.created_at) - AND i2.is_intermediate = 0 - ORDER BY i2.created_at DESC LIMIT 1 - ) as cover_image_name - FROM images - WHERE 1=1 - {query_conditions} - GROUP BY DATE(images.created_at) - ORDER BY date DESC; - """ - - cursor.execute(query, query_params) - result = cast(list[sqlite3.Row], cursor.fetchall()) - - return [ - VirtualSubBoardDTO( - virtual_board_id=f"by_date:{dict(row)['date']}", - board_name=dict(row)["date"], - date=dict(row)["date"], - image_count=dict(row)["image_count"], - asset_count=dict(row)["asset_count"], - cover_image_name=dict(row)["cover_image_name"], - ) - for row in result - ] - - def get_image_names_by_date( - self, - date: str, - starred_first: bool = True, - order_dir: SQLiteDirection = SQLiteDirection.Descending, - categories: Optional[list[ImageCategory]] = None, - search_term: Optional[str] = None, - user_id: Optional[str] = None, - is_admin: bool = False, - ) -> ImageNamesResult: - with self._db.transaction() as cursor: - query_conditions = "" - query_params: list[Union[int, str, bool]] = [] - - # Filter by date - query_conditions += """--sql - AND DATE(images.created_at) = ? - """ - query_params.append(date) - - # Only non-intermediate images - query_conditions += """--sql - AND images.is_intermediate = 0 - """ - - if categories is not None: - category_strings = [c.value for c in set(categories)] - placeholders = ",".join("?" * len(category_strings)) - query_conditions += f"""--sql - AND images.image_category IN ( {placeholders} ) - """ - for c in category_strings: - query_params.append(c) - - # User isolation for non-admin users - if user_id is not None and not is_admin: - query_conditions += """--sql - AND images.user_id = ? - """ - query_params.append(user_id) - - if search_term: - query_conditions += """--sql - AND ( - images.metadata LIKE ? - OR images.created_at LIKE ? - ) - """ - query_params.append(f"%{search_term.lower()}%") - query_params.append(f"%{search_term.lower()}%") - - # Get starred count if starred_first is enabled - starred_count = 0 - if starred_first: - starred_count_query = f"""--sql - SELECT COUNT(*) - FROM images - WHERE images.starred = TRUE AND (1=1{query_conditions}) - """ - cursor.execute(starred_count_query, query_params) - starred_count = cast(int, cursor.fetchone()[0]) - - # Get all image names with proper ordering - if starred_first: - names_query = f"""--sql - SELECT images.image_name - FROM images - WHERE 1=1{query_conditions} - ORDER BY images.starred DESC, images.created_at {order_dir.value} - """ - else: - names_query = f"""--sql - SELECT images.image_name - FROM images - WHERE 1=1{query_conditions} - ORDER BY images.created_at {order_dir.value} - """ - - cursor.execute(names_query, query_params) - result = cast(list[sqlite3.Row], cursor.fetchall()) - image_names = [row[0] for row in result] - - return ImageNamesResult(image_names=image_names, starred_count=starred_count, total_count=len(image_names)) diff --git a/invokeai/app/services/images/images_base.py b/invokeai/app/services/images/images_base.py index 14f98845c77..b83eac55b58 100644 --- a/invokeai/app/services/images/images_base.py +++ b/invokeai/app/services/images/images_base.py @@ -13,8 +13,7 @@ ) from invokeai.app.services.images.images_common import ImageDTO from invokeai.app.services.shared.intermediate_delete import IntermediateDeleteGuard, IntermediateDeleteResult -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection class ImageServiceABC(ABC): diff --git a/invokeai/app/services/images/images_default.py b/invokeai/app/services/images/images_default.py index f088f4bd56e..d4339f567bf 100644 --- a/invokeai/app/services/images/images_default.py +++ b/invokeai/app/services/images/images_default.py @@ -16,7 +16,6 @@ ImageNamesResult, ImageRecord, ImageRecordChanges, - ImageRecordDeleteException, ImageRecordNotFoundException, ImageRecordSaveException, InvalidImageCategoryException, @@ -33,8 +32,7 @@ JournaledDeleteAdapter, delete_journaled_intermediates, ) -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection class ImageService(ImageServiceABC): @@ -541,9 +539,6 @@ def delete(self, image_name: str): except ImageRecordNotFoundException: # Already deleted by another request; nothing here failed, so nothing to log. raise - except ImageRecordDeleteException: - self.__invoker.services.logger.error("Failed to delete image record") - raise except ImageFileDeleteException: self.__invoker.services.logger.error("Failed to delete image file") raise @@ -590,9 +585,6 @@ def delete_images_by_names(self, image_names: list[str]) -> tuple[list[str], lis log_error=self.__invoker.services.logger.error, ), ) - except ImageRecordDeleteException: - self.__invoker.services.logger.error("Failed to delete image records") - raise except ImageFileDeleteException: self.__invoker.services.logger.error("Failed to delete image files") raise @@ -623,9 +615,6 @@ def delete_intermediates_by_names( log_error=self.__invoker.services.logger.error, ), ) - except ImageRecordDeleteException: - self.__invoker.services.logger.error("Failed to delete image records") - raise except Exception as e: self.__invoker.services.logger.error("Problem deleting intermediate image records and files") raise e diff --git a/invokeai/app/services/intermediates/intermediates_default.py b/invokeai/app/services/intermediates/intermediates_default.py index 943da56f6da..2a72da54e40 100644 --- a/invokeai/app/services/intermediates/intermediates_default.py +++ b/invokeai/app/services/intermediates/intermediates_default.py @@ -49,8 +49,8 @@ IntermediatesUnavailableError, ) from invokeai.app.services.intermediates.intermediates_measurement import IntermediatesSizeMeasurer -from invokeai.app.services.intermediates.intermediates_records_sqlite import ( - IntermediatesRecordsSqlite, +from invokeai.app.services.intermediates.intermediates_records_default import ( + IntermediatesRecords, MediaKind, ReferenceOwner, ScopeTarget, @@ -111,7 +111,7 @@ class _Operation: class IntermediatesService(IntermediatesServiceBase): - def __init__(self, records: IntermediatesRecordsSqlite, logger: Optional[logging.Logger] = None) -> None: + def __init__(self, records: IntermediatesRecords, logger: Optional[logging.Logger] = None) -> None: self._records = records self._logger = logger or logging.getLogger(__name__) self._invoker: Optional[Invoker] = None @@ -129,6 +129,7 @@ def __init__(self, records: IntermediatesRecordsSqlite, logger: Optional[logging def start(self, invoker: Invoker) -> None: self._invoker = invoker self._stop.clear() + self._records.start() self._measurer.reset() def stop(self, invoker: Optional[Invoker] = None) -> None: diff --git a/invokeai/app/services/intermediates/intermediates_measurement.py b/invokeai/app/services/intermediates/intermediates_measurement.py index 5fd08c18096..819498aa07f 100644 --- a/invokeai/app/services/intermediates/intermediates_measurement.py +++ b/invokeai/app/services/intermediates/intermediates_measurement.py @@ -11,7 +11,7 @@ import time from typing import TYPE_CHECKING, Callable, Optional, cast -from invokeai.app.services.intermediates.intermediates_records_sqlite import IntermediatesRecordsSqlite, MediaKind +from invokeai.app.services.intermediates.intermediates_records_default import IntermediatesRecords, MediaKind if TYPE_CHECKING: from invokeai.app.services.invocation_services import InvocationServices @@ -24,7 +24,7 @@ class IntermediatesSizeMeasurer: def __init__( self, - records: IntermediatesRecordsSqlite, + records: IntermediatesRecords, services: Callable[[], "InvocationServices"], stop: threading.Event, logger: logging.Logger, diff --git a/invokeai/app/services/intermediates/intermediates_records_default.py b/invokeai/app/services/intermediates/intermediates_records_default.py new file mode 100644 index 00000000000..cc82c869797 --- /dev/null +++ b/invokeai/app/services/intermediates/intermediates_records_default.py @@ -0,0 +1,521 @@ +"""Storage of the intermediates manager. + +One classification expression decides what every intermediate is under the cleanup policy (see +`queries/intermediates.py`), whether it aggregates a summary, describes a preview, pages an operation or guards +a delete. + +A preview freezes the instant its recency is judged at; the operation it confirms and the guard on every +deleting transaction classify with that same recency cutoff and a live clock for everything else. Anything +created after the preview is newer than its cutoff, so it stays ``recent`` for the whole operation. Cached-output +holds are judged the same way by the operation itself, but the hold table is shared with live-clock callers +(summaries, cache hits) that may drop a hold once its grace has passed live; that only ever drops a hold whose +grace is over. + +Queue inputs are found by scanning each active item's stored session for media-name keys; the scan is cached +per item, so a queue of a thousand pending items is parsed once, not per query. +""" + +import re +import threading +from collections import defaultdict +from collections.abc import Callable +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from typing import Any, Collection, Iterator, NamedTuple, Optional, Sequence + +from invokeai.app.services.intermediates.intermediates_common import ( + BROWSER_HOLD_TTL_SECONDS, + MAX_BROWSER_HOLD_LEASES_PER_USER, + RECENT_GRACE_SECONDS, + IntermediatesCleanupMode, + IntermediatesKindCounts, +) +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.intermediates import ( + TARGET_SLOTS, + Classification, + MediaKind, + ScopeSpec, + WindowRow, +) +from invokeai.app.services.shared.database.queries.locks import DatabaseLock +from invokeai.app.services.shared.database.types import timestamp_text +from invokeai.app.services.shared.intermediate_delete import IntermediateDeleteGuard +from invokeai.app.services.shared.media_references import IMAGE_NAME_KEYS, VIDEO_NAME_KEYS, MediaReferences + +# One row of the manager: (owner, project), project None for the owner's unassigned intermediates. +ScopeTarget = tuple[str, Optional[str]] + +MEDIA_KINDS: tuple[MediaKind, ...] = ("image", "video") + + +def _name_pattern(keys: frozenset[str]) -> re.Pattern[str]: + alternatives = "|".join(sorted(re.escape(key) for key in keys)) + return re.compile(rf'"(?:{alternatives})"\s*:\s*"([^"\\]{{1,255}})"') + + +_IMAGE_NAME_RE = _name_pattern(IMAGE_NAME_KEYS) +_VIDEO_NAME_RE = _name_pattern(VIDEO_NAME_KEYS) + + +def _utc_now() -> datetime: + return datetime.now(timezone.utc) + + +class _Clock(NamedTuple): + """One instant for every statement of a classification.""" + + now: str + recent_cutoff: str + + +def _clock(recent_cutoff: Optional[str] = None) -> _Clock: + """The live clock, or the live clock judging recency as of an earlier preview.""" + now = _utc_now() + live_cutoff = timestamp_text(now - timedelta(seconds=RECENT_GRACE_SECONDS)) + return _Clock(timestamp_text(now), min(live_cutoff, recent_cutoff) if recent_cutoff is not None else live_cutoff) + + +@dataclass(frozen=True) +class _Classifying: + """What a classification binds: its clock and the media active queue items name, per kind.""" + + clock: _Clock + active_names: dict[MediaKind, frozenset[str]] + + def parameters(self, kind: MediaKind) -> dict[str, Any]: + return { + "now": self.clock.now, + "recent_cutoff": self.clock.recent_cutoff, + "active_names": self.active_names[kind], + } + + +class ReferenceOwner(NamedTuple): + """A saved document naming a media item, as `media_references` keys it.""" + + owner_kind: str + user_id: str + owner_id: str + + +def _add_count(counts: IntermediatesKindCounts, classification: Classification, n: int) -> None: + if classification == "safe": + counts.safe += n + elif classification == "referenced": + counts.referenced += n + elif classification == "active": + counts.active += n + else: + counts.recent += n + + +@dataclass +class ScopeCounts: + """Aggregated intermediates of one (owner, project) row for one media kind.""" + + counts: IntermediatesKindCounts = field(default_factory=IntermediatesKindCounts) + safe_bytes: int = 0 + referenced_bytes: int = 0 + unknown_size_count: int = 0 + + +@dataclass +class DeletableTotals: + count: int = 0 + referenced: int = 0 + measured_bytes: int = 0 + unknown_size_count: int = 0 + + +@dataclass +class ScopePreview: + """A preview's scope classified once, on one transaction and one clock, so its figures agree.""" + + clock: _Clock + # (owner, project) rows holding any intermediate in scope. + rows: set[tuple[str, Optional[str]]] = field(default_factory=set) + counts: dict[MediaKind, IntermediatesKindCounts] = field( + default_factory=lambda: {"image": IntermediatesKindCounts(), "video": IntermediatesKindCounts()} + ) + deletable: dict[MediaKind, DeletableTotals] = field( + default_factory=lambda: {"image": DeletableTotals(), "video": DeletableTotals()} + ) + # Force mode: the documents naming the deletable referenced items, with how many each names. + acknowledged: dict[ReferenceOwner, int] = field(default_factory=dict) + acknowledged_overflow: bool = False + + +def _scopes(*, user_id: Optional[str], targets: Optional[Sequence[ScopeTarget]], per_row: bool) -> list[ScopeSpec]: + """The statements' scopes that together cover a scope: selected rows, an owner's or everyone's intermediates. + + `per_row` gives each selected row a statement of its own, an index seek in creation order, where an OR over + many rows would sort the whole remaining scope on every window. + """ + if targets is not None: + if per_row: + return [ + ScopeSpec("project", targets=((owner, project),)) + if project is not None + else ScopeSpec("unassigned", user_id=owner) + for owner, project in targets + ] + return [ + ScopeSpec("targets", targets=tuple(targets[start : start + TARGET_SLOTS])) + for start in range(0, len(targets), TARGET_SLOTS) + ] + if user_id is not None: + return [ScopeSpec("owner", user_id=user_id)] + return [ScopeSpec("all")] + + +class IntermediatesRecords: + def __init__(self, database: Database) -> None: + self._queries = database.queries + # Media names an active queue item's session names, keyed by item id and stamped with the row's session + # revision, which changes with every session rewrite: a session rewritten while the item stays active (a + # workflow-call parent resuming with its child's outputs) is scanned again, a status change alone is not. + # Entries live as long as the item is active. Derived from committed rows only, so it is right whichever + # transaction rolls back; transactions on a server fill it side by side, hence the lock. + self._active_inputs: dict[int, tuple[int, frozenset[str], frozenset[str]]] = {} + self._active_inputs_lock = threading.Lock() + + def start(self) -> None: + """Drops the session holds and unmeasurable marks a previous process left behind: they were its own.""" + self._queries.intermediates.clear_process_state() + + def _expire_session_holds(self, clock: _Clock) -> None: + """Releases the holds of sessions that ended and drops those whose grace is over, in a short transaction of + its own: a classification that wrote them would hold their rows for as long as it reads, and a cleanup's + guard waiting on one of them would hold up every write that protects media.""" + self._queries.intermediates.expire_session_holds(now=clock.now, recent_cutoff=clock.recent_cutoff) + + def _classifying(self, q: Queries, clock: _Clock) -> _Classifying: + """Brings the scan of active queue items up to date, on the caller's transaction, and writes nothing. + + Runs on the caller's transaction, so that a classification and an enqueue it might race are ordered by the + database, never by a stale cache. Holds not expired yet only protect more: a hold of a session that ended + still counts as active until it is released. + """ + active = q.intermediates.active_items() + with self._active_inputs_lock: + known = {item: entry for item, entry in self._active_inputs.items() if active.get(item) == entry[0]} + changed = [item for item in active if item not in known] + for item, stamp, session in q.intermediates.sessions(changed): + text = session if isinstance(session, str) else "" + known[item] = ( + stamp, + frozenset(_IMAGE_NAME_RE.findall(text)), + frozenset(_VIDEO_NAME_RE.findall(text)), + ) + with self._active_inputs_lock: + self._active_inputs = known + images: set[str] = set() + videos: set[str] = set() + for _, image_names, video_names in known.values(): + images.update(image_names) + videos.update(video_names) + return _Classifying(clock, {"image": frozenset(images), "video": frozenset(videos)}) + + # region holds + + def hold_cached_media(self, session_id: str, references: MediaReferences) -> bool: + if references.is_empty(): + return True + + def hold(q: Queries) -> bool: + # Shared with every protecting write, exclusive for a cleanup's check and delete: the cleanup either + # deletes first (a cache miss), or sees this active session's hold. + q.locks.acquire(DatabaseLock.MEDIA_PROTECTION, shared=True) + if not q.intermediates.is_session_active(session_id): + return False + held: list[tuple[MediaKind, str]] = [] + wanted_by_kind: tuple[tuple[MediaKind, Collection[str]], ...] = ( + ("image", references.images), + ("video", references.videos), + ) + for kind, names in wanted_by_kind: + wanted = sorted(set(names)) + if wanted and q.intermediates.count_existing(kind, wanted) != len(wanted): + # Deletion won the race after the cache lookup. Do not return a stale output; invoking the + # node again creates fresh media instead. + return False + held.extend((kind, name) for name in wanted) + q.intermediates.hold_session_media(session_id, held) + return True + + return self._queries.run(hold) + + def replace_browser_hold(self, user_id: str, lease_id: str, images: Sequence[str], videos: Sequence[str]) -> None: + held_by_kind: tuple[tuple[MediaKind, Sequence[str]], ...] = (("image", images), ("video", videos)) + + def replace(q: Queries) -> None: + q.locks.acquire(DatabaseLock.MEDIA_PROTECTION, shared=True) + now = _utc_now() + q.intermediates.sweep_browser_holds(timestamp_text(now)) + q.intermediates.release_lease(user_id, lease_id) + expires_at = timestamp_text(now + timedelta(seconds=BROWSER_HOLD_TTL_SECONDS)) + for kind, names in held_by_kind: + if names: + q.intermediates.hold_for_lease( + kind, user_id=user_id, lease_id=lease_id, names=names, expires_at=expires_at + ) + # Past the cap, the leases refreshed longest ago stop protecting their media; an editor that is still + # open restores its lease on its next refresh. Newest first, ties by lease id. + others = sorted(lease for lease in q.intermediates.leases(user_id) if lease[0] != lease_id) + others.sort(key=lambda lease: lease[1], reverse=True) + lapsed = [lease for lease, _ in others[MAX_BROWSER_HOLD_LEASES_PER_USER - 1 :]] + if lapsed: + q.intermediates.release_leases(user_id, lapsed) + + self._queries.run(replace) + + def release_browser_hold(self, user_id: str, lease_id: str) -> None: + self._queries.intermediates.release_lease(user_id, lease_id) + + # endregion + + # region summary + + def summarize( + self, user_ids: Optional[Collection[str]], kinds: Sequence[MediaKind] = MEDIA_KINDS + ) -> dict[tuple[str, Optional[str]], dict[MediaKind, ScopeCounts]]: + """Aggregates the ``kinds`` intermediates of ``user_ids`` (None: everyone) by (owner, project) and classification.""" + scope = ScopeSpec("all") if user_ids is None else ScopeSpec("owners", users=frozenset(user_ids)) + clock = _clock() + self._expire_session_holds(clock) + + def summarize(q: Queries) -> dict[tuple[str, Optional[str]], dict[MediaKind, ScopeCounts]]: + rows: dict[tuple[str, Optional[str]], dict[MediaKind, ScopeCounts]] = defaultdict( + lambda: {"image": ScopeCounts(), "video": ScopeCounts()} + ) + classifying = self._classifying(q, clock) + for kind in kinds: + for row in q.intermediates.aggregate(kind, scope, **classifying.parameters(kind)): + counts = rows[(row.user_id, row.project_key)][kind] + _add_count(counts.counts, row.classification, row.count) + if row.classification == "safe": + counts.safe_bytes += row.measured_bytes + counts.unknown_size_count += row.unmeasured + elif row.classification == "referenced": + counts.referenced_bytes += row.measured_bytes + counts.unknown_size_count += row.unmeasured + return dict(rows) + + return self._queries.run(summarize, read_only=True) + + def get_projects(self, user_id: Optional[str]) -> dict[tuple[str, str], tuple[str, Optional[str]]]: + """Maps (owner, project) to (name, cover image): the newest durable image on the project's board.""" + return { + (owner, project_id): (name, cover) + for owner, project_id, name, cover in self._queries.intermediates.projects(user_id) + } + + def has_unmeasured_intermediates(self) -> bool: + return any(self._queries.intermediates.any_unmeasured(kind) for kind in MEDIA_KINDS) + + def next_unmeasured(self, kind: MediaKind, limit: int, *, min_age_seconds: int = 0) -> list[tuple[str, str]]: + """Unmeasured intermediates, oldest first; rows younger than ``min_age_seconds`` are left for later. + + A row is written before its file, so a brand-new one would measure as missing. Failures are marked + (`mark_unmeasurable`) rather than skipped by a bounded list, so later rows remain reachable. + """ + created_before = timestamp_text(_utc_now() - timedelta(seconds=min_age_seconds)) + return self._queries.intermediates.next_unmeasured(kind, limit=limit, created_before=created_before) + + def mark_unmeasurable(self, kind: MediaKind, names: Sequence[str]) -> None: + self._queries.intermediates.mark_unmeasurable(kind, names) + + # endregion + + # region scope + + def preview_scope( + self, + *, + user_id: Optional[str], + targets: Optional[Sequence[ScopeTarget]], + mode: IntermediatesCleanupMode, + is_admin: bool, + caller_user_id: str, + max_acknowledged: int, + ) -> ScopePreview: + """Counts every intermediate in scope and what a cleanup of ``mode`` may delete. + + ``targets`` narrows to selected rows; otherwise ``user_id`` narrows to an owner, and None means everyone. + One read-only transaction, which reads one snapshot, one instant and one view of active work serve every + figure, so the kept counts, the deletable totals and the acknowledged documents agree. + """ + scopes = _scopes(user_id=user_id, targets=targets, per_row=False) + clock = _clock() + self._expire_session_holds(clock) + + def preview(q: Queries) -> ScopePreview: + classifying = self._classifying(q, clock) + result = ScopePreview(clock=classifying.clock) + for kind in MEDIA_KINDS: + totals = result.deletable[kind] + for scope in scopes: + for row in q.intermediates.preview( + kind, + scope, + mode=mode, + is_admin=is_admin, + caller_user_id=caller_user_id, + **classifying.parameters(kind), + ): + result.rows.add((row.user_id, row.project_key)) + _add_count(result.counts[kind], row.classification, row.count) + if row.deletable: + totals.count += row.count + totals.measured_bytes += row.measured_bytes + totals.unknown_size_count += row.unmeasured + if row.classification == "referenced": + totals.referenced += row.count + if mode != "force": + continue + for owner_kind, owner_user, owner_id, n in q.intermediates.acknowledged( + kind, scope, is_admin=is_admin, caller_user_id=caller_user_id, **classifying.parameters(kind) + ): + owner = ReferenceOwner(owner_kind, owner_user, owner_id) + if owner not in result.acknowledged and len(result.acknowledged) >= max_acknowledged: + result.acknowledged_overflow = True + continue + result.acknowledged[owner] = result.acknowledged.get(owner, 0) + n + return result + + return self._queries.run(preview, read_only=True) + + def iter_deletable_batches( + self, + kind: MediaKind, + *, + user_id: Optional[str], + targets: Optional[Sequence[ScopeTarget]], + mode: IntermediatesCleanupMode, + is_admin: bool, + caller_user_id: str, + recent_cutoff: Optional[str], + limit: int, + ) -> Iterator[list[tuple[str, Optional[int]]]]: + """The deletable intermediates of a scope in creation order, a bounded window at a time. + + Each window classifies ``limit`` rows on one short transaction and yields the deletable ones (possibly + none, so the caller can re-check its authority between windows), so a long protected stretch never holds + the database for more than one window, and a row is classified once per operation however many are kept. + ``recent_cutoff`` freezes recency at the preview that was confirmed; None judges it live. + """ + for scope in _scopes(user_id=user_id, targets=targets, per_row=True): + after = ("", "") + while True: + clock = _clock(recent_cutoff) + self._expire_session_holds(clock) + rows = self._queries.run( + self._window(kind, scope, mode, is_admin, caller_user_id, clock, after, limit), read_only=True + ) + if not rows: + break + after = (rows[-1].created_at, rows[-1].name) + yield [(row.name, row.file_size_bytes) for row in rows if row.deletable] + + def _window( + self, + kind: MediaKind, + scope: ScopeSpec, + mode: IntermediatesCleanupMode, + is_admin: bool, + caller_user_id: str, + clock: _Clock, + after: tuple[str, str], + limit: int, + ) -> "Callable[[Queries], list[WindowRow]]": + def window(q: Queries) -> list[WindowRow]: + classifying = self._classifying(q, clock) + return q.intermediates.window( + kind, + scope, + mode=mode, + is_admin=is_admin, + caller_user_id=caller_user_id, + after=after, + limit=limit, + **classifying.parameters(kind), + ) + + return window + + def get_document_names(self, owners: Collection[ReferenceOwner]) -> dict[ReferenceOwner, str]: + """Names of the referencing documents, one query per kind; client state has none.""" + wanted: dict[str, set[ReferenceOwner]] = defaultdict(set) + for owner in owners: + if owner.owner_kind in ("project", "quarantined_project", "workflow"): + wanted[owner.owner_kind].add(owner) + + def names(q: Queries) -> dict[ReferenceOwner, str]: + found: dict[ReferenceOwner, str] = {} + for owner_kind, kind_owners in wanted.items(): + ids = sorted({owner.owner_id for owner in kind_owners}) + named = { + (owner, owner_id): name + for owner, owner_id, name in q.intermediates.document_names(owner_kind, ids) + if name is not None + } + for owner in kind_owners: + name = named.get((None if owner_kind == "workflow" else owner.user_id, owner.owner_id)) + if name is not None: + found[owner] = name + return found + + return self._queries.run(names, read_only=True) + + # endregion + + # region delete guard + + def make_delete_guard( + self, + kind: MediaKind, + *, + mode: IntermediatesCleanupMode, + allowed_user_ids: Optional[frozenset[str]], + caller_user_id: str, + is_admin: bool, + recent_cutoff: Optional[str] = None, + acknowledged: Optional[frozenset[ReferenceOwner]] = None, + ) -> IntermediateDeleteGuard: + """The final check of a cleanup batch, run on the deleting transaction. + + Re-applies the policy at the moment of deletion, judging recency as of ``recent_cutoff``: anything that + became active, referenced (safe mode), non-intermediate or another account's since the batch was read is + kept. A force clear deletes a referenced item only while every document naming it is in ``acknowledged``. + ``allowed_user_ids`` None means every account. + """ + acknowledged_documents = acknowledged or frozenset() + + def guard(q: Queries, names: Sequence[str]) -> list[str]: + classifying = self._classifying(q, _clock(recent_cutoff)) + named = set(names) + references: dict[str, set[ReferenceOwner]] = defaultdict(set) + if mode == "force": + for media_name, owner_kind, owner_user, owner_id in q.intermediates.reference_owners(kind, named): + references[media_name].add(ReferenceOwner(owner_kind, owner_user, owner_id)) + # Ownership is re-read here rather than trusted from the preview. + allowed = { + name + for name, owner in q.intermediates.still_deletable( + kind, + named, + mode=mode, + is_admin=is_admin, + caller_user_id=caller_user_id, + **classifying.parameters(kind), + ) + if (allowed_user_ids is None or owner in allowed_user_ids) + and references[name] <= acknowledged_documents + } + return [name for name in names if name in allowed] + + return guard + + # endregion diff --git a/invokeai/app/services/intermediates/intermediates_records_sqlite.py b/invokeai/app/services/intermediates/intermediates_records_sqlite.py deleted file mode 100644 index 82fee29ab2e..00000000000 --- a/invokeai/app/services/intermediates/intermediates_records_sqlite.py +++ /dev/null @@ -1,817 +0,0 @@ -"""Storage queries of the intermediates manager. - -One classification expression decides what every intermediate is under the cleanup policy, and it -is the same SQL whether it aggregates a summary, describes a preview, pages an operation or guards -a delete: - -- ``active``: produced or referenced by a pending, waiting or running queue item, or a - completed child of an active root workflow. -- ``recent``: created inside the grace window, so an in-flight browser upload is never collected - before it is promoted, referenced or enqueued; or a cached output whose consuming session ended - inside that window, since a cache hit reuses an old row. -- ``referenced``: named by a saved document (`media_references`). -- ``safe``: none of the above. - -A preview freezes the instant its recency is judged at; the operation it confirms and the guard -on every deleting transaction classify with that same recency cutoff and a live clock for -everything else. Anything created after the preview is newer than its cutoff, so it stays -``recent`` for the whole operation. Cached-output holds are judged the same way by the operation -itself, but the hold table is shared with live-clock callers (summaries, cache hits) that may -sweep a hold once its grace has passed live; that only ever drops a row whose grace is over. - -Queue inputs are found by scanning each active item's stored session for media-name keys; the -scan is cached per item so a queue of a thousand pending items is parsed once, not per query. -""" - -import re -import sqlite3 -from collections import defaultdict -from dataclasses import dataclass, field -from datetime import datetime, timedelta, timezone -from typing import Collection, Iterator, Literal, NamedTuple, Optional, Sequence, cast - -from invokeai.app.services.intermediates.intermediates_common import ( - BROWSER_HOLD_TTL_SECONDS, - MAX_BROWSER_HOLD_LEASES_PER_USER, - RECENT_GRACE_SECONDS, - IntermediatesCleanupMode, - IntermediatesKindCounts, -) -from invokeai.app.services.shared.intermediate_delete import IntermediateDeleteGuard -from invokeai.app.services.shared.media_references import IMAGE_NAME_KEYS, VIDEO_NAME_KEYS, MediaReferences -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase - -MediaKind = Literal["image", "video"] -Classification = Literal["safe", "referenced", "active", "recent"] -# One row of the manager: (owner, project), project None for the owner's unassigned intermediates. -ScopeTarget = tuple[str, Optional[str]] - -ACTIVE_QUEUE_STATUSES = ("pending", "in_progress", "waiting") - - -def _protected_queue_items_sql() -> str: - statuses = ", ".join(f"'{status}'" for status in ACTIVE_QUEUE_STATUSES) - # A child becomes completed before its outputs reach its waiting parent. Keep the - # whole completed subtree protected until the root finishes, including nested calls. - # CROSS JOIN pins the active roots first, then indexed child lookups, rather than - # scanning all completed history. Session JSON is fetched only for changed rows. - return f""" - SELECT item_id, session_id, status, session_revision FROM session_queue - WHERE status IN ({statuses}) - UNION ALL - SELECT child.item_id, child.session_id, child.status, child.session_revision - FROM session_queue root - CROSS JOIN session_queue child ON child.root_item_id = root.item_id - WHERE root.status IN ({statuses}) AND child.status = 'completed' - """ - - -# Keeps the OR-chain of a selection scope, and every IN list, under SQLITE_MAX_VARIABLE_NUMBER. -_MAX_SQL_VARIABLES = 500 -_TARGETS_PER_STATEMENT = 200 - - -def _name_pattern(keys: frozenset[str]) -> re.Pattern[str]: - alternatives = "|".join(sorted(re.escape(key) for key in keys)) - return re.compile(rf'"(?:{alternatives})"\s*:\s*"([^"\\]{{1,255}})"') - - -_IMAGE_NAME_RE = _name_pattern(IMAGE_NAME_KEYS) -_VIDEO_NAME_RE = _name_pattern(VIDEO_NAME_KEYS) - -_TABLES: dict[MediaKind, tuple[str, str, str]] = { - "image": ("images", "image_name", "image_subfolder"), - "video": ("videos", "video_name", "video_subfolder"), -} - - -def _utc_now() -> datetime: - return datetime.now(timezone.utc) - - -def _sql_timestamp(moment: datetime) -> str: - """The text `STRFTIME('%Y-%m-%d %H:%M:%f', ...)` stores, so bound instants compare with stored ones as strings.""" - return moment.strftime("%Y-%m-%d %H:%M:%S.") + f"{moment.microsecond // 1000:03d}" - - -class _Clock(NamedTuple): - """One instant for every statement of a classification: SQLite fixes `'NOW'` per statement, not per transaction.""" - - now: str - recent_cutoff: str - - -def _clock(recent_cutoff: Optional[str] = None) -> _Clock: - """The live clock, or the live clock judging recency as of an earlier preview.""" - now = _utc_now() - live_cutoff = _sql_timestamp(now - timedelta(seconds=RECENT_GRACE_SECONDS)) - return _Clock(_sql_timestamp(now), min(live_cutoff, recent_cutoff) if recent_cutoff is not None else live_cutoff) - - -class ReferenceOwner(NamedTuple): - """A saved document naming a media item, as `media_references` keys it.""" - - owner_kind: str - user_id: str - owner_id: str - - -# Media name → the documents naming it. -ReferenceOwners = dict[str, set[ReferenceOwner]] - - -def _add_count(counts: IntermediatesKindCounts, classification: Classification, n: int) -> None: - if classification == "safe": - counts.safe += n - elif classification == "referenced": - counts.referenced += n - elif classification == "active": - counts.active += n - else: - counts.recent += n - - -@dataclass -class ScopeCounts: - """Aggregated intermediates of one (owner, project) row for one media kind.""" - - counts: IntermediatesKindCounts = field(default_factory=IntermediatesKindCounts) - safe_bytes: int = 0 - referenced_bytes: int = 0 - unknown_size_count: int = 0 - - -@dataclass -class DeletableTotals: - count: int = 0 - referenced: int = 0 - measured_bytes: int = 0 - unknown_size_count: int = 0 - - -@dataclass -class ScopePreview: - """A preview's scope classified once, on one transaction and one clock, so its figures agree.""" - - clock: _Clock - # (owner, project) rows holding any intermediate in scope. - rows: set[tuple[str, Optional[str]]] = field(default_factory=set) - counts: dict[MediaKind, IntermediatesKindCounts] = field( - default_factory=lambda: {"image": IntermediatesKindCounts(), "video": IntermediatesKindCounts()} - ) - deletable: dict[MediaKind, DeletableTotals] = field( - default_factory=lambda: {"image": DeletableTotals(), "video": DeletableTotals()} - ) - # Force mode: the documents naming the deletable referenced items, with how many each names. - acknowledged: dict[ReferenceOwner, int] = field(default_factory=dict) - acknowledged_overflow: bool = False - - -class IntermediatesRecordsSqlite: - def __init__(self, db: SqliteDatabase) -> None: - self._db = db - # Media names an active queue item's session names, keyed by item id and stamped with the - # row's session revision, which a trigger bumps on every session rewrite: a session - # rewritten while the item stays active (a workflow-call parent resuming with its child's - # outputs) is re-scanned, a status change alone is not. Entries live as long as the item is - # active. This cache is derived from committed rows only, so it is safe across a - # rolled-back transaction; the temp table is not, which is why it is rebuilt on every call - # rather than skipped on an unchanged set. - self._active_inputs: dict[int, tuple[int, set[str], set[str]]] = {} - - @staticmethod - def _prepare_session_holds(cursor: sqlite3.Cursor, clock: _Clock) -> None: - # These holds share the cache's process lifetime. After restart the cache is empty and - # recovered queue sessions protect their serialized inputs through the normal scan. - # A cache hit reuses an old row, so the recency grace a fresh output gets is counted from - # when the consuming session is first seen finished instead. - cursor.execute( - "CREATE TEMP TABLE IF NOT EXISTS intermediates_session_media " - "(session_id TEXT, kind TEXT, name TEXT, released_at TEXT, PRIMARY KEY(session_id, kind, name)) " - "WITHOUT ROWID;" - ) - cursor.execute( - "UPDATE temp.intermediates_session_media SET released_at = ? " - "WHERE released_at IS NULL AND session_id NOT IN " - f"(SELECT session_id FROM ({_protected_queue_items_sql()}));", - (clock.now,), - ) - cursor.execute("DELETE FROM temp.intermediates_session_media WHERE released_at <= ?;", (clock.recent_cutoff,)) - - def hold_cached_media(self, session_id: str, references: MediaReferences) -> bool: - if references.is_empty(): - return True - with self._db.transaction() as cursor: - self._prepare_session_holds(cursor, _clock()) - statuses = ", ".join(f"'{status}'" for status in ACTIVE_QUEUE_STATUSES) - cursor.execute( - f"SELECT 1 FROM session_queue WHERE session_id = ? AND status IN ({statuses});", (session_id,) - ) - if cursor.fetchone() is None: - return False - held: list[tuple[str, str, str]] = [] - for kind, names in (("image", references.images), ("video", references.videos)): - table, name_column, _ = _TABLES[cast(MediaKind, kind)] - ordered = sorted(names) - for start in range(0, len(ordered), _MAX_SQL_VARIABLES): - chunk = ordered[start : start + _MAX_SQL_VARIABLES] - placeholders = ",".join("?" for _ in chunk) - cursor.execute(f"SELECT COUNT(*) FROM {table} WHERE {name_column} IN ({placeholders});", chunk) - if cursor.fetchone()[0] != len(chunk): - # Deletion won the race after the cache lookup. Do not return a stale - # output; invoking the node again will create fresh media instead. - return False - held.extend((session_id, kind, name) for name in ordered) - # The existence check and hold commit share the deleting transaction's lock. Cleanup - # either deletes first (a cache miss), or observes this active session's hold. - cursor.executemany("INSERT OR REPLACE INTO temp.intermediates_session_media VALUES (?, ?, ?, NULL);", held) - return True - - def replace_browser_hold(self, user_id: str, lease_id: str, images: Sequence[str], videos: Sequence[str]) -> None: - with self._db.transaction() as cursor: - # Indexed on expires_at, so the sweep touches only what expired. - cursor.execute( - "DELETE FROM intermediates_browser_holds WHERE expires_at <= STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW');" - ) - cursor.execute( - "DELETE FROM intermediates_browser_holds WHERE user_id = ? AND lease_id = ?;", (user_id, lease_id) - ) - cursor.execute("CREATE TEMP TABLE IF NOT EXISTS intermediates_hold_names (name TEXT PRIMARY KEY);") - for kind, names in (("image", images), ("video", videos)): - cursor.execute("DELETE FROM temp.intermediates_hold_names;") - cursor.executemany( - "INSERT OR IGNORE INTO temp.intermediates_hold_names(name) VALUES (?);", - [(name,) for name in names], - ) - table, name_column, _ = _TABLES[cast(MediaKind, kind)] - cursor.execute( - f"""--sql - INSERT INTO intermediates_browser_holds - (user_id, lease_id, media_kind, media_name, expires_at) - SELECT ?, ?, ?, m.{name_column}, - STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW', '+{BROWSER_HOLD_TTL_SECONDS} seconds') - FROM {table} m - JOIN temp.intermediates_hold_names h ON h.name = m.{name_column} - WHERE m.user_id = ? AND m.is_intermediate = TRUE; - """, - (user_id, lease_id, kind, user_id), - ) - # Past the cap, the leases refreshed longest ago stop protecting their media; an editor - # that is still open restores its lease on its next refresh. - cursor.execute( - """--sql - WITH leases AS ( - SELECT lease_id, MAX(expires_at) AS refreshed - FROM intermediates_browser_holds - WHERE user_id = ? - GROUP BY lease_id - ), lapsed AS ( - SELECT lease_id FROM leases - WHERE lease_id != ? - ORDER BY refreshed DESC, lease_id - LIMIT -1 OFFSET ? - ) - DELETE FROM intermediates_browser_holds - WHERE user_id = ? AND lease_id IN (SELECT lease_id FROM lapsed); - """, - (user_id, lease_id, MAX_BROWSER_HOLD_LEASES_PER_USER - 1, user_id), - ) - - def release_browser_hold(self, user_id: str, lease_id: str) -> None: - with self._db.transaction() as cursor: - cursor.execute( - "DELETE FROM intermediates_browser_holds WHERE user_id = ? AND lease_id = ?;", (user_id, lease_id) - ) - - # region policy expression - - @staticmethod - def _classification_sql(kind: MediaKind) -> str: - """Binds a `_Clock`'s two fields, in order, ahead of any later parameters of the statement.""" - table, name_column, _ = _TABLES[kind] - return f""" - CASE - WHEN ( - m.session_id IS NOT NULL - AND m.session_id IN (SELECT session_id FROM ({_protected_queue_items_sql()})) - ) OR m.{name_column} IN ( - SELECT name FROM temp.intermediates_active_media WHERE kind = '{kind}' - ) OR EXISTS ( - SELECT 1 FROM intermediates_browser_holds h - WHERE h.media_kind = '{kind}' AND h.media_name = m.{name_column} - AND h.expires_at > ? - ) THEN 'active' - WHEN m.created_at > ? - OR m.{name_column} IN ( - SELECT name FROM temp.intermediates_session_media - WHERE kind = '{kind}' AND released_at IS NOT NULL - ) THEN 'recent' - WHEN EXISTS ( - SELECT 1 FROM media_references r WHERE r.media_kind = '{kind}' AND r.media_name = m.{name_column} - ) THEN 'referenced' - ELSE 'safe' - END - """ - - @staticmethod - def _deletable_sql( - kind: MediaKind, mode: IntermediatesCleanupMode, *, is_admin: bool, caller_user_id: str - ) -> tuple[str, list[object]]: - """What a cleanup may delete, over the classified rows of `_scoped_sql`. - - Force mode adds referenced items, but another account's document is never the caller's to - break unless the caller administers the instance. Applied in SQL so a bounded page or - batch is a page of what the caller can actually delete. - """ - if mode == "safe": - return "cls = 'safe'", [] - if is_admin: - return "cls IN ('safe', 'referenced')", [] - return ( - "(cls = 'safe' OR (cls = 'referenced' AND NOT EXISTS (" - "SELECT 1 FROM media_references r WHERE r.media_kind = ? AND r.media_name = name AND r.user_id != ?)))", - [kind, caller_user_id], - ) - - def _prepare(self, cursor: sqlite3.Cursor, clock: Optional[_Clock] = None) -> _Clock: - """Makes the transaction's view of active work current, and returns the instant its classifications use. - - Runs on the caller's transaction so a classification and the enqueue it might race are - ordered by the database lock, never by a stale cache. - """ - clock = clock or _clock() - self._prepare_session_holds(cursor, clock) - cursor.execute( - """--sql - CREATE TEMP TABLE IF NOT EXISTS intermediates_active_media ( - kind TEXT NOT NULL, - name TEXT NOT NULL, - PRIMARY KEY (kind, name) - ) WITHOUT ROWID; - """ - ) - cursor.execute( - f"""--sql - SELECT item_id, COALESCE(session_revision, 0) - FROM ({_protected_queue_items_sql()}); - """ - ) - active = {cast(int, row[0]): cast(int, row[1]) for row in cursor.fetchall()} - - for stale in [item_id for item_id in self._active_inputs if item_id not in active]: - del self._active_inputs[stale] - changed = [ - item_id - for item_id, stamp in active.items() - if item_id not in self._active_inputs or self._active_inputs[item_id][0] != stamp - ] - for start in range(0, len(changed), _MAX_SQL_VARIABLES): - chunk = changed[start : start + _MAX_SQL_VARIABLES] - placeholders = ",".join("?" for _ in chunk) - cursor.execute( - f"""--sql - SELECT item_id, COALESCE(session_revision, 0), session - FROM session_queue WHERE item_id IN ({placeholders}); - """, - chunk, - ) - for item_id, stamp, session_json in cursor.fetchall(): - session_text = session_json if isinstance(session_json, str) else "" - self._active_inputs[cast(int, item_id)] = ( - cast(int, stamp), - set(_IMAGE_NAME_RE.findall(session_text)), - set(_VIDEO_NAME_RE.findall(session_text)), - ) - - cursor.execute("DELETE FROM temp.intermediates_active_media;") - rows: set[tuple[str, str]] = set() - for _, images, videos in self._active_inputs.values(): - rows.update(("image", name) for name in images) - rows.update(("video", name) for name in videos) - if rows: - cursor.executemany( - "INSERT OR IGNORE INTO temp.intermediates_active_media (kind, name) VALUES (?, ?);", sorted(rows) - ) - cursor.execute( - "INSERT OR IGNORE INTO temp.intermediates_active_media (kind, name) " - "SELECT kind, name FROM temp.intermediates_session_media WHERE released_at IS NULL;" - ) - return clock - - # endregion - - # region summary - - @classmethod - def _scoped_sql(cls, kind: MediaKind) -> str: - """Every intermediate of ``kind`` with its row and classification; callers append ``AND`` scope filters. - - Media whose project no longer exists, or that never had one, belong to the owner's - unassigned row (``project_key`` NULL). - """ - table, name_column, _ = _TABLES[kind] - return f""" - SELECT m.rowid AS rowid, m.{name_column} AS name, m.user_id AS user_id, - CASE WHEN p.project_id IS NULL THEN NULL ELSE m.project_id END AS project_key, - m.file_size_bytes AS file_size_bytes, - {cls._classification_sql(kind)} AS cls - FROM {table} m - LEFT JOIN projects p ON p.user_id = m.user_id AND p.project_id = m.project_id - WHERE m.is_intermediate = TRUE - """ - - @classmethod - def _aggregate( - cls, cursor: sqlite3.Cursor, clock: _Clock, kind: MediaKind, scope_clause: str, params: Sequence[object] - ) -> Iterator[tuple[str, Optional[str], Classification, int, int, int]]: - """(owner, project, classification, count, measured bytes, unmeasured count) groups within a scope.""" - cursor.execute( - f"""--sql - SELECT user_id, project_key, cls, COUNT(*), - SUM(COALESCE(file_size_bytes, 0)), - SUM(CASE WHEN file_size_bytes IS NULL THEN 1 ELSE 0 END) - FROM ({cls._scoped_sql(kind)} {scope_clause}) - GROUP BY user_id, project_key, cls; - """, - [*clock, *params], - ) - for owner, project_key, classification, n, total_bytes, unknown in cursor.fetchall(): - yield ( - cast(str, owner), - cast(Optional[str], project_key), - cast(Classification, classification), - cast(int, n), - cast(int, total_bytes), - cast(int, unknown), - ) - - def summarize( - self, user_ids: Optional[Collection[str]], kinds: Sequence[MediaKind] = ("image", "video") - ) -> dict[tuple[str, Optional[str]], dict[MediaKind, ScopeCounts]]: - """Aggregates the ``kinds`` intermediates of ``user_ids`` (None: everyone) by (owner, project) and classification.""" - rows: dict[tuple[str, Optional[str]], dict[MediaKind, ScopeCounts]] = defaultdict( - lambda: {"image": ScopeCounts(), "video": ScopeCounts()} - ) - owners = sorted(user_ids) if user_ids is not None else None - # Owner chunks partition the GROUP BY, so their rows never need merging. - clauses: list[tuple[str, list[object]]] = ( - [ - (f"AND m.user_id IN ({','.join('?' for _ in chunk)})", list(chunk)) - for chunk in ( - owners[start : start + _MAX_SQL_VARIABLES] for start in range(0, len(owners), _MAX_SQL_VARIABLES) - ) - ] - if owners is not None - else [("", [])] - ) - with self._db.transaction() as cursor: - clock = self._prepare(cursor) - for kind in kinds: - for clause, params in clauses: - for owner, project_key, cls, n, total_bytes, unknown in self._aggregate( - cursor, clock, kind, clause, params - ): - scope = rows[(owner, project_key)][kind] - _add_count(scope.counts, cls, n) - if cls == "safe": - scope.safe_bytes += total_bytes - scope.unknown_size_count += unknown - elif cls == "referenced": - scope.referenced_bytes += total_bytes - scope.unknown_size_count += unknown - return dict(rows) - - def get_projects(self, user_id: Optional[str]) -> dict[tuple[str, str], tuple[str, Optional[str]]]: - """Maps (owner, project) to (name, cover image): the newest durable image on the project's board.""" - owner_filter = "WHERE p.user_id = ?" if user_id is not None else "" - params: list[object] = [user_id] if user_id is not None else [] - with self._db.transaction() as cursor: - cursor.execute( - f"""--sql - SELECT p.user_id, p.project_id, p.name, - ( - SELECT bi.image_name FROM board_images bi - INNER JOIN images i ON i.image_name = bi.image_name - WHERE bi.board_id = p.board_id AND i.is_intermediate = FALSE - ORDER BY bi.created_at DESC LIMIT 1 - ) - FROM projects p {owner_filter}; - """, - params, - ) - return { - (cast(str, r[0]), cast(str, r[1])): (cast(str, r[2]), cast(Optional[str], r[3])) - for r in cursor.fetchall() - } - - def has_unmeasured_intermediates(self) -> bool: - with self._db.transaction() as cursor: - cursor.execute( - "CREATE TEMP TABLE IF NOT EXISTS intermediates_unmeasurable " - "(kind TEXT, name TEXT, PRIMARY KEY(kind, name));" - ) - for kind, (table, name_column, _) in _TABLES.items(): - cursor.execute( - f"SELECT 1 FROM {table} m WHERE m.is_intermediate = TRUE AND m.file_size_bytes IS NULL " - f"AND NOT EXISTS (SELECT 1 FROM temp.intermediates_unmeasurable u " - f"WHERE u.kind = ? AND u.name = m.{name_column}) LIMIT 1;", - (kind,), - ) - if cursor.fetchone() is not None: - return True - return False - - def next_unmeasured(self, kind: MediaKind, limit: int, *, min_age_seconds: int = 0) -> list[tuple[str, str]]: - """Unmeasured intermediates in insertion order; rows younger than ``min_age_seconds`` are left for later. - - A row is written before its file, so a brand-new one would measure as missing. Failures - are held in a temporary table instead of a bounded IN list, so later rows remain reachable. - """ - table, name_column, subfolder_column = _TABLES[kind] - with self._db.transaction() as cursor: - cursor.execute( - "CREATE TEMP TABLE IF NOT EXISTS intermediates_unmeasurable (kind TEXT, name TEXT, PRIMARY KEY(kind, name));" - ) - cursor.execute( - f"""--sql - SELECT m.{name_column}, m.{subfolder_column} FROM {table} m - WHERE m.is_intermediate = TRUE AND m.file_size_bytes IS NULL - AND m.created_at <= STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW', ?) - AND NOT EXISTS (SELECT 1 FROM temp.intermediates_unmeasurable u - WHERE u.kind = ? AND u.name = m.{name_column}) - ORDER BY m.rowid ASC - LIMIT ?; - """, - [f"-{min_age_seconds} seconds", kind, limit], - ) - return [(cast(str, r[0]), cast(str, r[1])) for r in cursor.fetchall()] - - def mark_unmeasurable(self, kind: MediaKind, names: Sequence[str]) -> None: - if not names: - return - with self._db.transaction() as cursor: - cursor.execute( - "CREATE TEMP TABLE IF NOT EXISTS intermediates_unmeasurable (kind TEXT, name TEXT, PRIMARY KEY(kind, name));" - ) - cursor.executemany( - "INSERT OR IGNORE INTO temp.intermediates_unmeasurable(kind, name) VALUES (?, ?);", - [(kind, name) for name in names], - ) - - # endregion - - # region scope - - @classmethod - def _scope_clauses( - cls, - *, - user_id: Optional[str], - targets: Optional[Sequence[ScopeTarget]], - per_statement: int = _TARGETS_PER_STATEMENT, - ) -> list[tuple[str, list[object]]]: - """``AND`` filters (with their parameters) that together cover a scope; each stays under the variable limit.""" - if targets is not None: - clauses: list[tuple[str, list[object]]] = [] - for start in range(0, len(targets), per_statement): - predicate, params = cls._scope_predicate(targets[start : start + per_statement]) - if predicate: - clauses.append((f"AND ({predicate})", params)) - return clauses - if user_id is not None: - return [("AND m.user_id = ?", [user_id])] - return [("", [])] - - @staticmethod - def _scope_predicate(targets: Sequence[ScopeTarget]) -> tuple[str, list[object]]: - clauses: list[str] = [] - params: list[object] = [] - for user_id, project_id in targets: - if project_id is None: - clauses.append("(m.user_id = ? AND p.project_id IS NULL)") - params.append(user_id) - else: - clauses.append("(m.user_id = ? AND m.project_id = ? AND p.project_id IS NOT NULL)") - params.extend((user_id, project_id)) - return " OR ".join(clauses), params - - def preview_scope( - self, - *, - user_id: Optional[str], - targets: Optional[Sequence[ScopeTarget]], - mode: IntermediatesCleanupMode, - is_admin: bool, - caller_user_id: str, - max_acknowledged: int, - ) -> ScopePreview: - """Counts every intermediate in scope and what a cleanup of ``mode`` may delete. - - ``targets`` narrows to selected rows; otherwise ``user_id`` narrows to an owner, and None - means everyone. One transaction, one instant and one view of active work serve every - figure, so the kept counts, the deletable totals and the acknowledged documents agree. - """ - clauses = self._scope_clauses(user_id=user_id, targets=targets) - with self._db.transaction() as cursor: - result = ScopePreview(clock=self._prepare(cursor)) - for kind in _TABLES: - deletable_sql, deletable_params = self._deletable_sql( - kind, mode, is_admin=is_admin, caller_user_id=caller_user_id - ) - totals = result.deletable[kind] - for clause, params in clauses: - # Parameters bind by textual position: the deletable predicate sits in the - # SELECT list, ahead of the classified subquery's clock and scope parameters. - cursor.execute( - f"""--sql - SELECT user_id, project_key, cls, {deletable_sql} AS deletable, COUNT(*), - SUM(COALESCE(file_size_bytes, 0)), - SUM(CASE WHEN file_size_bytes IS NULL THEN 1 ELSE 0 END) - FROM ({self._scoped_sql(kind)} {clause}) - GROUP BY user_id, project_key, cls, deletable; - """, - [*deletable_params, *result.clock, *params], - ) - for owner, project_key, cls, deletable, n, total_bytes, unknown in cursor.fetchall(): - result.rows.add((cast(str, owner), cast(Optional[str], project_key))) - _add_count(result.counts[kind], cast(Classification, cls), cast(int, n)) - if deletable: - totals.count += cast(int, n) - totals.measured_bytes += cast(int, total_bytes) - totals.unknown_size_count += cast(int, unknown) - if cls == "referenced": - totals.referenced += cast(int, n) - if mode == "force": - cursor.execute( - f"""--sql - SELECT r.owner_kind, r.user_id, r.owner_id, COUNT(*) - FROM ({self._scoped_sql(kind)} {clause}) s - JOIN media_references r ON r.media_kind = ? AND r.media_name = s.name - WHERE s.cls = 'referenced' AND {deletable_sql} - GROUP BY r.owner_kind, r.user_id, r.owner_id; - """, - [*result.clock, *params, kind, *deletable_params], - ) - for owner_kind, owner_user, owner_id, n in cursor.fetchall(): - owner = ReferenceOwner(str(owner_kind), str(owner_user), str(owner_id)) - if owner not in result.acknowledged and len(result.acknowledged) >= max_acknowledged: - result.acknowledged_overflow = True - continue - result.acknowledged[owner] = result.acknowledged.get(owner, 0) + cast(int, n) - return result - - @classmethod - def _window_sql(cls, kind: MediaKind, scope_clause: str, deletable_sql: str) -> str: - """One window of classified rows in insertion order. Binds ``[*deletable, *clock, *scope, after_rowid, limit]``.""" - return f"""--sql - SELECT rowid, name, file_size_bytes, {deletable_sql} AS deletable - FROM ({cls._scoped_sql(kind)} {scope_clause}) - WHERE rowid > ? - ORDER BY rowid - LIMIT ?; - """ - - def iter_deletable_batches( - self, - kind: MediaKind, - *, - user_id: Optional[str], - targets: Optional[Sequence[ScopeTarget]], - mode: IntermediatesCleanupMode, - is_admin: bool, - caller_user_id: str, - recent_cutoff: Optional[str], - limit: int, - ) -> Iterator[list[tuple[str, Optional[int]]]]: - """The deletable intermediates of a scope in insertion order, a bounded window at a time. - - Each window classifies ``limit`` rows on one short transaction and yields the deletable - ones (possibly none, so the caller can re-check its authority between windows), so a long - protected stretch never holds the database for more than one window, and a row is - classified once per operation however many are kept. Explicit rows are paged one at a - time: one (owner, project) per statement is an index seek in rowid order, where an OR over - many rows would sort the whole remaining scope on every window. ``recent_cutoff`` freezes - recency at the preview that was confirmed; None judges it live. - """ - deletable_sql, deletable_params = self._deletable_sql( - kind, mode, is_admin=is_admin, caller_user_id=caller_user_id - ) - window_sql = None - for clause, params in self._scope_clauses(user_id=user_id, targets=targets, per_statement=1): - window_sql = self._window_sql(kind, clause, deletable_sql) - after_rowid = 0 - while True: - with self._db.transaction() as cursor: - clock = self._prepare(cursor, _clock(recent_cutoff)) - cursor.execute(window_sql, [*deletable_params, *clock, *params, after_rowid, limit]) - rows = cursor.fetchall() - if not rows: - break - after_rowid = cast(int, rows[-1][0]) - yield [(cast(str, row[1]), cast(Optional[int], row[2])) for row in rows if row[3]] - - @staticmethod - def _reference_owners(cursor: sqlite3.Cursor, kind: MediaKind, names: Sequence[str]) -> ReferenceOwners: - owners: ReferenceOwners = {} - for start in range(0, len(names), _MAX_SQL_VARIABLES - 1): - chunk = list(names[start : start + _MAX_SQL_VARIABLES - 1]) - placeholders = ",".join("?" for _ in chunk) - cursor.execute( - f"SELECT media_name, owner_kind, user_id, owner_id FROM media_references " - f"WHERE media_kind = ? AND media_name IN ({placeholders});", - [kind, *chunk], - ) - for name, owner_kind, user_id, owner_id in cursor.fetchall(): - owners.setdefault(str(name), set()).add(ReferenceOwner(str(owner_kind), str(user_id), str(owner_id))) - return owners - - def get_document_names(self, owners: Collection[ReferenceOwner]) -> dict[ReferenceOwner, str]: - """Names of the referencing documents, one query per kind and chunk; client state has none.""" - queries = { - "project": "SELECT user_id, project_id, name FROM projects WHERE project_id IN ({});", - "quarantined_project": ( - "SELECT user_id, project_id, name FROM orphaned_projects_2026_08_06 WHERE project_id IN ({});" - ), - # Workflow ids are globally unique, so the owner is not part of the match. - "workflow": "SELECT NULL, workflow_id, name FROM workflow_library WHERE workflow_id IN ({});", - } - wanted: dict[str, set[ReferenceOwner]] = defaultdict(set) - for owner in owners: - if owner.owner_kind in queries: - wanted[owner.owner_kind].add(owner) - names: dict[ReferenceOwner, str] = {} - with self._db.transaction() as cursor: - for owner_kind, kind_owners in wanted.items(): - ids = sorted({owner.owner_id for owner in kind_owners}) - found: dict[tuple[Optional[str], str], str] = {} - for start in range(0, len(ids), _MAX_SQL_VARIABLES): - chunk = ids[start : start + _MAX_SQL_VARIABLES] - cursor.execute(queries[owner_kind].format(",".join("?" for _ in chunk)), chunk) - for user_id, owner_id, name in cursor.fetchall(): - if name is not None: - found[(cast(Optional[str], user_id), cast(str, owner_id))] = cast(str, name) - for owner in kind_owners: - name = found.get((None if owner_kind == "workflow" else owner.user_id, owner.owner_id)) - if name is not None: - names[owner] = name - return names - - # endregion - - # region delete guard - - def make_delete_guard( - self, - kind: MediaKind, - *, - mode: IntermediatesCleanupMode, - allowed_user_ids: Optional[frozenset[str]], - caller_user_id: str, - is_admin: bool, - recent_cutoff: Optional[str] = None, - acknowledged: Optional[frozenset[ReferenceOwner]] = None, - ) -> IntermediateDeleteGuard: - """The final check of a cleanup batch, run on the deleting transaction. - - Re-applies the policy at the moment of deletion, judging recency as of ``recent_cutoff``: - anything that became active, referenced (safe mode), non-intermediate or another - account's since the batch was read is kept. A force clear deletes a referenced item only - while every document naming it is in ``acknowledged``. ``allowed_user_ids`` None means - every account. - """ - table, name_column, _ = _TABLES[kind] - deletable_sql, deletable_params = self._deletable_sql( - kind, mode, is_admin=is_admin, caller_user_id=caller_user_id - ) - acknowledged_documents = acknowledged or frozenset() - - def guard(cursor: sqlite3.Cursor, names: Sequence[str]) -> list[str]: - clock = self._prepare(cursor, _clock(recent_cutoff)) - kept: list[str] = [] - for start in range(0, len(names), _MAX_SQL_VARIABLES): - chunk = list(names[start : start + _MAX_SQL_VARIABLES]) - placeholders = ",".join("?" for _ in chunk) - current_references = self._reference_owners(cursor, kind, chunk) if mode == "force" else {} - cursor.execute( - f"""--sql - SELECT name, user_id FROM ( - SELECT m.{name_column} AS name, m.user_id AS user_id, - {self._classification_sql(kind)} AS cls - FROM {table} m - WHERE m.{name_column} IN ({placeholders}) AND m.is_intermediate = TRUE - ) WHERE {deletable_sql}; - """, - [*clock, *chunk, *deletable_params], - ) - # Ownership is re-read here rather than trusted from the preview. - kept.extend( - cast(str, row[0]) - for row in cursor.fetchall() - if allowed_user_ids is None or cast(str, row[1]) in allowed_user_ids - if current_references.get(cast(str, row[0]), set()).issubset(acknowledged_documents) - ) - return kept - - return guard - - # endregion diff --git a/invokeai/app/services/invocation_services.py b/invokeai/app/services/invocation_services.py index fa7abe347ab..41de3b8d704 100644 --- a/invokeai/app/services/invocation_services.py +++ b/invokeai/app/services/invocation_services.py @@ -47,7 +47,7 @@ from invokeai.app.services.project_records.project_records_base import ProjectRecordsStorageBase from invokeai.app.services.session_processor.session_processor_base import SessionProcessorBase from invokeai.app.services.session_queue.session_queue_base import SessionQueueBase - from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase + from invokeai.app.services.shared.database.database import Database from invokeai.app.services.urls.urls_base import UrlServiceBase from invokeai.app.services.users.users_base import UserServiceBase from invokeai.app.services.video_files.video_files_base import VideoFileStorageBase @@ -75,7 +75,7 @@ def __init__( board_records: "BoardRecordStorageBase", bulk_download: "BulkDownloadBase", configuration: "InvokeAIAppConfig", - database: "SqliteDatabase", + database: "Database", events: "EventServiceBase", images: "ImageServiceABC", image_files: "ImageFileStorageBase", diff --git a/invokeai/app/services/model_records/__init__.py b/invokeai/app/services/model_records/__init__.py index 4fee477466d..c56e9e2724f 100644 --- a/invokeai/app/services/model_records/__init__.py +++ b/invokeai/app/services/model_records/__init__.py @@ -5,7 +5,6 @@ InvalidModelException, ModelRecordServiceBase, UnknownModelException, - ModelSummary, ModelRecordChanges, ModelRecordOrderBy, ) @@ -17,7 +16,6 @@ "DuplicateModelException", "InvalidModelException", "UnknownModelException", - "ModelSummary", "ModelRecordChanges", "ModelRecordOrderBy", ] diff --git a/invokeai/app/services/model_records/model_records_base.py b/invokeai/app/services/model_records/model_records_base.py index 4fa27d97d2a..655cadfe89b 100644 --- a/invokeai/app/services/model_records/model_records_base.py +++ b/invokeai/app/services/model_records/model_records_base.py @@ -5,12 +5,11 @@ from abc import ABC, abstractmethod from enum import Enum from pathlib import Path -from typing import Any, List, Optional, Set, Union +from typing import Any, List, Optional, Union -from pydantic import BaseModel, Field, field_validator +from pydantic import Field, field_validator -from invokeai.app.services.shared.pagination import PaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.util.model_exclude_null import BaseModelExcludeNull from invokeai.backend.model_manager.configs.controlnet import ControlAdapterDefaultSettings from invokeai.backend.model_manager.configs.external_api import ( @@ -76,18 +75,6 @@ class ModelRecordOrderBy(str, Enum): Path = "path" -class ModelSummary(BaseModel): - """A short summary of models for UI listing purposes.""" - - key: str = Field(description="model key") - type: ModelType = Field(description="model type") - base: BaseModelType = Field(description="base model") - format: ModelFormat = Field(description="model format") - name: str = Field(description="model name") - description: str = Field(description="short description of model") - tags: Set[str] = Field(description="tags associated with model") - - class ModelRecordChanges(BaseModelExcludeNull): """A set of changes to apply to a model.""" @@ -175,13 +162,11 @@ class ModelRecordServiceBase(ABC): @abstractmethod def add_model(self, config: AnyModelConfig) -> AnyModelConfig: """ - Add a model to the database. + Add a model to the database; the config's key becomes the record's key. - :param key: Unique key for the model - :param config: Model configuration record, either a dict with the - required fields or a ModelConfigBase instance. + :param config: Model configuration record. - Can raise DuplicateModelException and InvalidModelConfigException exceptions. + Raises DuplicateModelException when a model with the same path or key is installed. """ pass @@ -232,25 +217,8 @@ def get_model(self, key: str) -> AnyModelConfig: pass @abstractmethod - def get_model_by_hash(self, hash: str) -> AnyModelConfig: - """ - Retrieve the configuration for the indicated model. - - :param hash: Hash of model config to be fetched. - - Exceptions: UnknownModelException - """ - pass - - @abstractmethod - def list_models( - self, - page: int = 0, - per_page: int = 10, - order_by: ModelRecordOrderBy = ModelRecordOrderBy.Default, - direction: SQLiteDirection = SQLiteDirection.Ascending, - ) -> PaginatedResults[ModelSummary]: - """Return a paginated summary listing of each model in the database.""" + def get_model_paths(self) -> list[str]: + """Return the path of every model record, also of a record whose config no longer validates.""" pass @abstractmethod diff --git a/invokeai/app/services/model_records/model_records_sql.py b/invokeai/app/services/model_records/model_records_sql.py index e12f4e7de85..a351399159f 100644 --- a/invokeai/app/services/model_records/model_records_sql.py +++ b/invokeai/app/services/model_records/model_records_sql.py @@ -1,47 +1,35 @@ """ -SQL Implementation of the ModelRecordServiceBase API +SQL implementation of the ModelRecordServiceBase API, on the database layer. Typical usage: - from invokeai.backend.model_manager import ModelConfigStoreSQL - store = ModelConfigStoreSQL(sqlite_db) - config = dict( - path='/tmp/pokemon.bin', - name='old name', - base_model='sd-1', - type='embedding', - format='embedding_file', - ) - - # adding - the key becomes the model's "key" field - store.add_model('key1', config) - - # updating - config.name='new name' - store.update_model('key1', config) - - # checking for existence - if store.exists('key1'): + from invokeai.app.services.model_records import ModelRecordChanges, ModelRecordServiceSQL + store = ModelRecordServiceSQL(database, logger) + + # adding - the config's key becomes the record's key + store.add_model(config) + + # updating + store.update_model(config.key, ModelRecordChanges(name="new name")) + + # checking for existence + if store.exists(config.key): print("yes") - # fetching config - new_config = store.get_model('key1') - print(new_config.name, new_config.base) - assert new_config.key == 'key1' + # fetching a config + config = store.get_model(config.key) # deleting - store.del_model('key1') + store.del_model(config.key) # searching - configs = store.search_by_path(path='/tmp/pokemon.bin') - configs = store.search_by_hash('750a499f35e43b7e1b4d15c207aa2f01') - configs = store.search_by_attr(base_model='sd-2', model_type='main') + configs = store.search_by_path("/tmp/pokemon.bin") + configs = store.search_by_hash("750a499f35e43b7e1b4d15c207aa2f01") + configs = store.search_by_attr(base_model=BaseModelType.StableDiffusion2, model_type=ModelType.Main) """ import json import logging -import sqlite3 -from math import ceil from pathlib import Path from typing import List, Optional, Union @@ -53,12 +41,13 @@ ModelRecordChanges, ModelRecordOrderBy, ModelRecordServiceBase, - ModelSummary, UnknownModelException, ) -from invokeai.app.services.shared.pagination import PaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import UniqueViolation +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.models import MAX_KEY_LENGTH, MAX_PATH_LENGTH +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.backend.model_manager.configs.base import Config_Base from invokeai.backend.model_manager.configs.factory import AnyModelConfig, ModelConfigFactory from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType @@ -88,58 +77,43 @@ def _construct_config_for_type(fields: dict, target_type: ModelType) -> AnyModel ) +def _check_path_length(path: str) -> None: + if len(path) > MAX_PATH_LENGTH: + raise ValueError(f"A model's path can be at most {MAX_PATH_LENGTH} characters long") + + +def _parse(config: str) -> AnyModelConfig: + return ModelConfigFactory.from_dict(json.loads(config)) + + class ModelRecordServiceSQL(ModelRecordServiceBase): """Implementation of the ModelConfigStore ABC using a SQL database.""" - def __init__(self, db: SqliteDatabase, logger: logging.Logger): - """ - Initialize a new object from preexisting sqlite3 connection and threading lock objects. - - :param db: Sqlite connection object - """ + def __init__(self, database: Database, logger: logging.Logger): super().__init__() - self._db = db + self._queries = database.queries self._logger = logger def add_model(self, config: AnyModelConfig) -> AnyModelConfig: """ Add a model to the database. - :param key: Unique key for the model - :param config: Model configuration record, either a dict with the - required fields or a ModelConfigBase instance. + :param config: Model configuration record; its key becomes the record's key. - Can raise DuplicateModelException and InvalidModelConfigException exceptions. + Raises DuplicateModelException when a model with the same path or key is installed, and ValueError when + its key or path is longer than a record holds. """ - with self._db.transaction() as cursor: - try: - cursor.execute( - """--sql - INSERT INTO models ( - id, - config - ) - VALUES (?,?); - """, - ( - config.key, - config.model_dump_json(), - ), - ) - - except sqlite3.IntegrityError as e: - if "UNIQUE constraint failed" in str(e): - if "models.path" in str(e): - msg = f"A model with path '{config.path}' is already installed" - elif "models.name" in str(e): - msg = f"A model with name='{config.name}', type='{config.type}', base='{config.base}' is already installed" - else: - msg = f"A model with key '{config.key}' is already installed" - raise DuplicateModelException(msg) from e - else: - raise e - - return self.get_model(config.key) + if len(config.key) > MAX_KEY_LENGTH: + raise ValueError(f"A model's key can be at most {MAX_KEY_LENGTH} characters long") + _check_path_length(config.path) + stored = config.model_dump_json() + try: + self._queries.models.insert(config.key, stored) + except UniqueViolation as error: + if self._queries.models.at_path(str(config.path)): + raise DuplicateModelException(f"A model with path '{config.path}' is already installed") from error + raise DuplicateModelException(f"A model with key '{config.key}' is already installed") from error + return _parse(stored) def del_model(self, key: str) -> None: """ @@ -149,20 +123,16 @@ def del_model(self, key: str) -> None: Can raise an UnknownModelException """ - with self._db.transaction() as cursor: - cursor.execute( - """--sql - DELETE FROM models - WHERE id=?; - """, - (key,), - ) - if cursor.rowcount == 0: - raise UnknownModelException("model not found") + if not self._queries.models.delete(key): + raise UnknownModelException("model not found") def update_model(self, key: str, changes: ModelRecordChanges, allow_class_change: bool = False) -> AnyModelConfig: - with self._db.transaction() as cursor: - record = self.get_model(key) + def apply(q: Queries) -> str: + # Locked, so that a concurrent update applies its changes to this one's result, not beside it. + stored = q.models.lock(key) + if stored is None: + raise UnknownModelException("model not found") + record = _parse(stored) if allow_class_change: # The changes may cause the model config class to change. To handle this, we need to construct the new @@ -197,46 +167,28 @@ def update_model(self, key: str, changes: ModelRecordChanges, allow_class_change if stale_field not in changes.model_fields_set: fallback_dict.pop(stale_field, None) record = _construct_config_for_type(fallback_dict, changes.type) - - # If we get this far, the updated model config is valid, so we can save it to the database. - json_serialized = record.model_dump_json() else: # We are not allowing the model config class to change, so we can just update the existing instance in # place. If the changes are invalid for the existing class, an exception will be raised by pydantic. for field_name in changes.model_fields_set: setattr(record, field_name, getattr(changes, field_name)) - json_serialized = record.model_dump_json() - - cursor.execute( - """--sql - UPDATE models - SET - config=? - WHERE id=?; - """, - (json_serialized, key), - ) - if cursor.rowcount == 0: - raise UnknownModelException("model not found") - return self.get_model(key) + # If we get this far, the updated model config is valid, so we can save it to the database. + if "path" in changes.model_fields_set: + _check_path_length(record.path) + updated = record.model_dump_json() + q.models.save(key, updated) + return updated + + return _parse(self._queries.run(apply)) def replace_model(self, key: str, new_config: AnyModelConfig) -> AnyModelConfig: if key != new_config.key: raise ValueError("key does not match new_config.key") - with self._db.transaction() as cursor: - cursor.execute( - """--sql - UPDATE models - SET - config=? - WHERE id=?; - """, - (new_config.model_dump_json(), key), - ) - if cursor.rowcount == 0: - raise UnknownModelException("model not found") - return self.get_model(key) + stored = new_config.model_dump_json() + if not self._queries.models.save(key, stored): + raise UnknownModelException("model not found") + return _parse(stored) def get_model(self, key: str) -> AnyModelConfig: """ @@ -246,34 +198,10 @@ def get_model(self, key: str) -> AnyModelConfig: Exceptions: UnknownModelException """ - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT config FROM models - WHERE id=?; - """, - (key,), - ) - rows = cursor.fetchone() - if not rows: - raise UnknownModelException("model not found") - model = ModelConfigFactory.from_dict(json.loads(rows[0])) - return model - - def get_model_by_hash(self, hash: str) -> AnyModelConfig: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT config FROM models - WHERE hash=?; - """, - (hash,), - ) - rows = cursor.fetchone() - if not rows: + stored = self._queries.models.get(key) + if stored is None: raise UnknownModelException("model not found") - model = ModelConfigFactory.from_dict(json.loads(rows[0])) - return model + return _parse(stored) def exists(self, key: str) -> bool: """ @@ -281,16 +209,7 @@ def exists(self, key: str) -> bool: :param key: Unique key for the model to be deleted """ - with self._db.transaction() as cursor: - cursor.execute( - """--sql - select count(*) FROM models - WHERE id=?; - """, - (key,), - ) - count = cursor.fetchone()[0] - return count > 0 + return self._queries.models.exists(key) def search_by_attr( self, @@ -314,60 +233,28 @@ def search_by_attr( If none of the optional filters are passed, will return all models in the database. """ - with self._db.transaction() as cursor: - assert isinstance(order_by, ModelRecordOrderBy) - order_dir = "DESC" if direction == SQLiteDirection.Descending else "ASC" - ordering = { - ModelRecordOrderBy.Default: f"type {order_dir}, base COLLATE NOCASE {order_dir}, name COLLATE NOCASE {order_dir}, format", - ModelRecordOrderBy.Type: "type", - ModelRecordOrderBy.Base: "base COLLATE NOCASE", - ModelRecordOrderBy.Name: "name COLLATE NOCASE", - ModelRecordOrderBy.Format: "format", - ModelRecordOrderBy.Size: "IFNULL(json_extract(config, '$.file_size'), 0)", - ModelRecordOrderBy.DateAdded: "created_at", - ModelRecordOrderBy.DateModified: "updated_at", - ModelRecordOrderBy.Path: "path", - } - - where_clause: list[str] = [] - bindings: list[str] = [] - if model_name: - where_clause.append("name=?") - bindings.append(model_name) - if base_model: - where_clause.append("base=?") - bindings.append(base_model) - if model_type: - where_clause.append("type=?") - bindings.append(model_type) - if model_format: - where_clause.append("format=?") - bindings.append(model_format) - where = f"WHERE {' AND '.join(where_clause)}" if where_clause else "" - - cursor.execute( - f"""--sql - SELECT config - FROM models - {where} - ORDER BY {ordering[order_by]} {order_dir} -- using ? to bind doesn't work here for some reason; - """, - tuple(bindings), - ) - result = cursor.fetchall() + assert isinstance(order_by, ModelRecordOrderBy) + configs = self._queries.models.search( + name=model_name, + base=base_model, + model_type=model_type, + model_format=model_format, + order_by=order_by.value, + descending=direction == SQLiteDirection.Descending, + ) # Parse the model configs. results: list[AnyModelConfig] = [] - for row in result: + for config in configs: try: - model_config = ModelConfigFactory.from_dict(json.loads(row[0])) + model_config = _parse(config) except pydantic.ValidationError as e: # We catch this error so that the app can still run if there are invalid model configs in the database. # One reason that an invalid model config might be in the database is if someone had to rollback from a # newer version of the app that added a new model type. - row_data = f"{row[0][:64]}..." if len(row[0]) > 64 else row[0] + row_data = f"{config[:64]}..." if len(config) > 64 else config try: - name = json.loads(row[0]).get("name", "") + name = json.loads(config).get("name", "") except Exception: name = "" self._logger.warning( @@ -381,77 +268,11 @@ def search_by_attr( def search_by_path(self, path: Union[str, Path]) -> List[AnyModelConfig]: """Return models with the indicated path.""" - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT config FROM models - WHERE path=?; - """, - (str(path),), - ) - results = [ModelConfigFactory.from_dict(json.loads(x[0])) for x in cursor.fetchall()] - return results + return [_parse(config) for config in self._queries.models.at_path(str(path))] def search_by_hash(self, hash: str) -> List[AnyModelConfig]: """Return models with the indicated hash.""" - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT config FROM models - WHERE hash=?; - """, - (hash,), - ) - results = [ModelConfigFactory.from_dict(json.loads(x[0])) for x in cursor.fetchall()] - return results + return [_parse(config) for config in self._queries.models.with_hash(hash)] - def list_models( - self, - page: int = 0, - per_page: int = 10, - order_by: ModelRecordOrderBy = ModelRecordOrderBy.Default, - direction: SQLiteDirection = SQLiteDirection.Ascending, - ) -> PaginatedResults[ModelSummary]: - """Return a paginated summary listing of each model in the database.""" - with self._db.transaction() as cursor: - assert isinstance(order_by, ModelRecordOrderBy) - order_dir = "DESC" if direction == SQLiteDirection.Descending else "ASC" - ordering = { - ModelRecordOrderBy.Default: f"type {order_dir}, base COLLATE NOCASE {order_dir}, name COLLATE NOCASE {order_dir}, format", - ModelRecordOrderBy.Type: "type", - ModelRecordOrderBy.Base: "base COLLATE NOCASE", - ModelRecordOrderBy.Name: "name COLLATE NOCASE", - ModelRecordOrderBy.Format: "format", - ModelRecordOrderBy.Size: "IFNULL(json_extract(config, '$.file_size'), 0)", - ModelRecordOrderBy.DateAdded: "created_at", - ModelRecordOrderBy.DateModified: "updated_at", - ModelRecordOrderBy.Path: "path", - } - - # Lock so that the database isn't updated while we're doing the two queries. - # query1: get the total number of model configs - cursor.execute( - """--sql - select count(*) from models; - """, - (), - ) - total = int(cursor.fetchone()[0]) - - # query2: fetch key fields - cursor.execute( - f"""--sql - SELECT config - FROM models - ORDER BY {ordering[order_by]} {order_dir} -- using ? to bind doesn't work here for some reason - LIMIT ? - OFFSET ?; - """, - ( - per_page, - page * per_page, - ), - ) - rows = cursor.fetchall() - items = [ModelSummary.model_validate(dict(x)) for x in rows] - return PaginatedResults(page=page, pages=ceil(total / per_page), per_page=per_page, total=total, items=items) + def get_model_paths(self) -> list[str]: + return self._queries.models.paths() diff --git a/invokeai/app/services/model_relationship_records/model_relationship_records_base.py b/invokeai/app/services/model_relationship_records/model_relationship_records_base.py index 94fa179bc0a..5b029f29e79 100644 --- a/invokeai/app/services/model_relationship_records/model_relationship_records_base.py +++ b/invokeai/app/services/model_relationship_records/model_relationship_records_base.py @@ -18,8 +18,3 @@ def remove_model_relationship(self, model_key_1: str, model_key_2: str) -> None: def get_related_model_keys(self, model_key: str) -> list[str]: """Gets all models keys related to a given model key.""" pass - - @abstractmethod - def get_related_model_keys_batch(self, model_keys: list[str]) -> list[str]: - """Get related model keys for multiple models given a list of keys.""" - pass diff --git a/invokeai/app/services/model_relationship_records/model_relationship_records_default.py b/invokeai/app/services/model_relationship_records/model_relationship_records_default.py new file mode 100644 index 00000000000..85811cb95aa --- /dev/null +++ b/invokeai/app/services/model_relationship_records/model_relationship_records_default.py @@ -0,0 +1,21 @@ +from invokeai.app.services.model_relationship_records.model_relationship_records_base import ( + ModelRelationshipRecordStorageBase, +) +from invokeai.app.services.shared.database.database import Database + + +class ModelRelationshipRecordStorage(ModelRelationshipRecordStorageBase): + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + + def add_model_relationship(self, model_key_1: str, model_key_2: str) -> None: + if model_key_1 == model_key_2: + raise ValueError("Cannot relate a model to itself.") + self._queries.model_relationships.add(model_key_1, model_key_2) + + def remove_model_relationship(self, model_key_1: str, model_key_2: str) -> None: + self._queries.model_relationships.remove(model_key_1, model_key_2) + + def get_related_model_keys(self, model_key: str) -> list[str]: + return self._queries.model_relationships.related(model_key) diff --git a/invokeai/app/services/model_relationship_records/model_relationship_records_sqlite.py b/invokeai/app/services/model_relationship_records/model_relationship_records_sqlite.py deleted file mode 100644 index c12990b8c3a..00000000000 --- a/invokeai/app/services/model_relationship_records/model_relationship_records_sqlite.py +++ /dev/null @@ -1,55 +0,0 @@ -from invokeai.app.services.model_relationship_records.model_relationship_records_base import ( - ModelRelationshipRecordStorageBase, -) -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase - - -class SqliteModelRelationshipRecordStorage(ModelRelationshipRecordStorageBase): - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - - def add_model_relationship(self, model_key_1: str, model_key_2: str) -> None: - with self._db.transaction() as cursor: - if model_key_1 == model_key_2: - raise ValueError("Cannot relate a model to itself.") - a, b = sorted([model_key_1, model_key_2]) - cursor.execute( - "INSERT OR IGNORE INTO model_relationships (model_key_1, model_key_2) VALUES (?, ?)", - (a, b), - ) - - def remove_model_relationship(self, model_key_1: str, model_key_2: str) -> None: - with self._db.transaction() as cursor: - a, b = sorted([model_key_1, model_key_2]) - cursor.execute( - "DELETE FROM model_relationships WHERE model_key_1 = ? AND model_key_2 = ?", - (a, b), - ) - - def get_related_model_keys(self, model_key: str) -> list[str]: - with self._db.transaction() as cursor: - cursor.execute( - """ - SELECT model_key_2 FROM model_relationships WHERE model_key_1 = ? - UNION - SELECT model_key_1 FROM model_relationships WHERE model_key_2 = ? - """, - (model_key, model_key), - ) - result = [row[0] for row in cursor.fetchall()] - return result - - def get_related_model_keys_batch(self, model_keys: list[str]) -> list[str]: - with self._db.transaction() as cursor: - key_list = ",".join("?" for _ in model_keys) - cursor.execute( - f""" - SELECT model_key_2 FROM model_relationships WHERE model_key_1 IN ({key_list}) - UNION - SELECT model_key_1 FROM model_relationships WHERE model_key_2 IN ({key_list}) - """, - model_keys + model_keys, - ) - result = [row[0] for row in cursor.fetchall()] - return result diff --git a/invokeai/app/services/model_relationships/model_relationships_base.py b/invokeai/app/services/model_relationships/model_relationships_base.py index 1ea744a8dcb..7d79e648fb4 100644 --- a/invokeai/app/services/model_relationships/model_relationships_base.py +++ b/invokeai/app/services/model_relationships/model_relationships_base.py @@ -18,8 +18,3 @@ def remove_model_relationship(self, model_key_1: str, model_key_2: str) -> None: def get_related_model_keys(self, model_key: str) -> list[str]: """Gets all models keys related to a given model key.""" pass - - @abstractmethod - def get_related_model_keys_batch(self, model_keys: list[str]) -> list[str]: - """Get related model keys for multiple models.""" - pass diff --git a/invokeai/app/services/model_relationships/model_relationships_default.py b/invokeai/app/services/model_relationships/model_relationships_default.py index e4da482ff27..342b1ef47b8 100644 --- a/invokeai/app/services/model_relationships/model_relationships_default.py +++ b/invokeai/app/services/model_relationships/model_relationships_default.py @@ -26,6 +26,3 @@ def remove_relationship_from_models(self, model_1: AnyModelConfig, model_2: AnyM def get_related_keys_from_model(self, model: AnyModelConfig) -> list[str]: return self.get_related_model_keys(model.key) - - def get_related_model_keys_batch(self, model_keys: list[str]) -> list[str]: - return self.__invoker.services.model_relationship_records.get_related_model_keys_batch(model_keys) diff --git a/invokeai/app/services/orphaned_models/orphaned_models_service.py b/invokeai/app/services/orphaned_models/orphaned_models_service.py index d454c9b0018..6892b7482e1 100644 --- a/invokeai/app/services/orphaned_models/orphaned_models_service.py +++ b/invokeai/app/services/orphaned_models/orphaned_models_service.py @@ -4,10 +4,9 @@ in the database models table. """ -import json import shutil from pathlib import Path -from typing import Set +from typing import TYPE_CHECKING, Set from pydantic import BaseModel, Field @@ -19,7 +18,9 @@ delete_active_install_sentinel, is_recovery_protected_path, ) -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase + +if TYPE_CHECKING: + from invokeai.app.services.model_records import ModelRecordServiceBase # Scratch area for operations that must build a model on the models volume before registering it - # model conversion writes its diffusers copy here. It lives under the models root so the finished @@ -63,15 +64,15 @@ class OrphanedModelsService: ".git", } - def __init__(self, config: InvokeAIAppConfig, db: SqliteDatabase): + def __init__(self, config: InvokeAIAppConfig, store: "ModelRecordServiceBase"): """Initialize the service. Args: config: Application configuration containing models path - db: Database connection for querying registered models + store: The model records, whose files are not orphans """ self._config = config - self._db = db + self._store = store def find_orphaned_models(self) -> list[OrphanedModelInfo]: """Find all orphaned model directories. @@ -226,36 +227,27 @@ def _get_registered_model_directories(self, models_dir: Path) -> Set[Path]: """Get the set of all model directories from the database.""" model_directories = set() - with self._db.transaction() as cursor: - cursor.execute("SELECT config FROM models") - rows = cursor.fetchall() + # Every record's path, also of a record whose config no longer validates: its files are not orphans. + for path_str in self._store.get_model_paths(): + if not path_str: + continue + path = Path(path_str) - for row in rows: - try: - config = json.loads(row[0]) - if "path" in config and config["path"]: - path_str = config["path"] - path = Path(path_str) - - # If the path is relative, resolve it relative to models_dir - if not path.is_absolute(): - full_path = (models_dir / path).resolve() - else: - full_path = path.resolve() - - # Extract the top-level directory under models_dir - try: - rel_path = full_path.relative_to(models_dir) - if rel_path.parts: - top_level_dir = models_dir / rel_path.parts[0] - model_directories.add(top_level_dir.resolve()) - except ValueError: - # Path is not relative to models_dir - model_directories.add(full_path) - - except (json.JSONDecodeError, KeyError, TypeError): - # Skip invalid model configs - continue + # If the path is relative, resolve it relative to models_dir + if not path.is_absolute(): + full_path = (models_dir / path).resolve() + else: + full_path = path.resolve() + + # Extract the top-level directory under models_dir + try: + rel_path = full_path.relative_to(models_dir) + if rel_path.parts: + top_level_dir = models_dir / rel_path.parts[0] + model_directories.add(top_level_dir.resolve()) + except ValueError: + # Path is not relative to models_dir + model_directories.add(full_path) return model_directories diff --git a/invokeai/app/services/project_records/project_records_default.py b/invokeai/app/services/project_records/project_records_default.py new file mode 100644 index 00000000000..017b12536f4 --- /dev/null +++ b/invokeai/app/services/project_records/project_records_default.py @@ -0,0 +1,275 @@ +import json +import uuid +from typing import Any, Optional + +from invokeai.app.services.board_records.board_records_common import BOARD_NAME_MAX_LENGTH, BoardVisibility +from invokeai.app.services.project_records.project_records_base import ProjectRecordsStorageBase +from invokeai.app.services.project_records.project_records_common import ( + DEFAULT_PROJECT_CANVAS_SCHEMA_VERSION, + PROJECT_BOARD_SNAPSHOT_MAX_ITEMS, + PROJECT_DOCUMENT_MAX_BYTES, + ProjectBoardNotFoundError, + ProjectBoardSnapshotDTO, + ProjectBoardTooLargeError, + ProjectBoardUnavailableError, + ProjectCanvasSchemaDowngradeError, + ProjectCanvasSchemaUnsupportedError, + ProjectDocumentInvalidError, + ProjectDocumentTooLargeError, + ProjectRecordConflictError, + ProjectRecordDTO, + ProjectRecordExistsError, + ProjectRecordNotFoundError, + ProjectSummaryDTO, +) +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import UniqueViolation +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.locks import DatabaseLock +from invokeai.app.services.shared.media_references import extract_media_references +from invokeai.app.util.misc import uuid_string + + +def _serialize_project_document(data: dict[str, Any]) -> tuple[str, int]: + try: + document_json = json.dumps(data, ensure_ascii=False, allow_nan=False, separators=(",", ":")) + return document_json, len(document_json.encode("utf-8")) + except (TypeError, ValueError, UnicodeEncodeError) as error: + raise ProjectDocumentInvalidError from error + + +def _require_project_document_size(actual_bytes: int) -> None: + if actual_bytes > PROJECT_DOCUMENT_MAX_BYTES: + raise ProjectDocumentTooLargeError(actual_bytes=actual_bytes, max_bytes=PROJECT_DOCUMENT_MAX_BYTES) + + +def _require_supported_schema(*, project_id: str, minimum_version: int, client_maximum_version: int) -> None: + if client_maximum_version < minimum_version: + raise ProjectCanvasSchemaUnsupportedError( + project_id=project_id, + minimum_version=minimum_version, + client_maximum_version=client_maximum_version, + ) + + +def _record(summary: Optional[ProjectSummaryDTO], project_id: str, data: dict[str, Any]) -> ProjectRecordDTO: + """The record a write returns: its row as the write left it, read in the write's transaction, with the document + the write stored. Not read back: the document is the largest part of the row, and decoding it again would take + longer than the whole transaction.""" + if summary is None: + raise ProjectRecordNotFoundError(project_id) + return ProjectRecordDTO(**summary.model_dump(), data=data) + + +def _claim_board(q: Queries, *, user_id: str, board_id: str, board_name: str, project_id: str) -> None: + """Takes ownership of an existing private board for the project being created, renaming it after the project, + or explains why it cannot be taken. + + The board's row is locked before the claim looks at the board: a change to the board, or another project + claiming it, either committed before and is seen here, or waits until this transaction has ended and then finds + the board claimed. The project's row, if the account already has one, is locked before the board's, the order + every write of a project takes them in. + + The project being created is not counted as a claim on the board, so a client that re-sends a create it never + got an answer for is refused as `ProjectRecordExistsError` -- the truth. Otherwise it would be told its board + was unavailable, which reads as "somebody else took it" and is the opposite of what happened. A repeated create + is not an edge case: it is how a client establishes whether a request whose response was lost actually landed, + and the answer decides whether it deletes the media it uploaded. + """ + project_exists = q.projects.lock(user_id, project_id) is not None + board = q.boards.lock(board_id) + if board is None: + raise ProjectBoardNotFoundError(board_id) + owner, visibility = board + if owner != user_id: + raise ProjectBoardNotFoundError(board_id) + claimant = q.projects.claimant(board_id) + if ( + visibility != BoardVisibility.Private.value + or claimant not in (None, (user_id, project_id)) + or q.boards.is_shared(board_id) + ): + raise ProjectBoardUnavailableError(board_id) + # The board's claimant can be the project being created although its lock found nothing: the same create sent + # twice, both under way at once. Inserting it again would lock its row after the board's. + if project_exists or claimant is not None: + raise ProjectRecordExistsError(project_id) + # Un-archived as well as renamed. A project's board takes its archived state from the project, and + # `PATCH /boards/{id}` refuses to set it on a claimed board -- so a board that arrived archived would be + # invisible in every listing with no API left to fix it. + q.boards.rename(board_id, board_name, unarchive=True) + + +class ProjectRecordsStorage(ProjectRecordsStorageBase): + """Per-user project documents. A project and its board are one unit: every write touching both runs in one + transaction, so the board and the project commit or roll back together.""" + + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + + def create( + self, + user_id: str, + name: str, + data: dict[str, Any], + project_id: str | None = None, + board_id: str | None = None, + minimum_canvas_schema_version: int = DEFAULT_PROJECT_CANVAS_SCHEMA_VERSION, + max_canvas_schema_version: int = DEFAULT_PROJECT_CANVAS_SCHEMA_VERSION, + ) -> ProjectRecordDTO: + new_project_id = project_id or uuid.uuid4().hex + _require_supported_schema( + project_id=new_project_id, + minimum_version=minimum_canvas_schema_version, + client_maximum_version=max_canvas_schema_version, + ) + document_json, document_bytes = _serialize_project_document(data) + _require_project_document_size(document_bytes) + references = extract_media_references(data) + board_name = name[:BOARD_NAME_MAX_LENGTH] + + def insert(q: Queries) -> Optional[ProjectSummaryDTO]: + # Shared with every write that makes media protected, exclusive for the intermediates cleanup's check and + # delete: the media this names cannot be deleted between that check and this commit. + q.locks.acquire(DatabaseLock.MEDIA_PROTECTION, shared=True) + if board_id is None: + project_board_id = uuid_string() + q.boards.insert(board_id=project_board_id, board_name=board_name, user_id=user_id) + else: + _claim_board(q, user_id=user_id, board_id=board_id, board_name=board_name, project_id=new_project_id) + project_board_id = board_id + try: + q.projects.insert( + project_id=new_project_id, + user_id=user_id, + name=name, + data=document_json, + board_id=project_board_id, + minimum_canvas_schema_version=minimum_canvas_schema_version, + ) + except UniqueViolation as error: + # The account has a project with this id, created meanwhile or, when the create brings no board, + # before. It cannot be the board: a new board is the project's alone, and a claimed one is locked and + # claimed by no other project. The board insert or rename rolls back with the project insert, so a + # refused create changes no board. + raise ProjectRecordExistsError(new_project_id) from error + # Indexed in the document's own transaction: a reader never sees a saved project whose media the + # intermediates cleanup does not know about. + q.media_references.replace( + owner_kind="project", user_id=user_id, owner_id=new_project_id, references=references + ) + return q.projects.summary(user_id, new_project_id) + + return _record(self._queries.run(insert), new_project_id, data) + + def get( + self, + user_id: str, + project_id: str, + max_canvas_schema_version: int = DEFAULT_PROJECT_CANVAS_SCHEMA_VERSION, + ) -> ProjectRecordDTO: + record = self._queries.projects.get(user_id, project_id) + if record is None: + raise ProjectRecordNotFoundError(project_id) + _require_supported_schema( + project_id=project_id, + minimum_version=record.minimum_canvas_schema_version, + client_maximum_version=max_canvas_schema_version, + ) + return record + + def list(self, user_id: str) -> list[ProjectSummaryDTO]: + return self._queries.projects.summaries(user_id) + + def update( + self, + user_id: str, + project_id: str, + expected_revision: int, + name: str, + data: dict[str, Any], + minimum_canvas_schema_version: int | None = None, + max_canvas_schema_version: int = DEFAULT_PROJECT_CANVAS_SCHEMA_VERSION, + ) -> ProjectRecordDTO: + document_json, document_bytes = _serialize_project_document(data) + references = extract_media_references(data) + + def save(q: Queries) -> Optional[ProjectSummaryDTO]: + q.locks.acquire(DatabaseLock.MEDIA_PROTECTION, shared=True) + # Locked before anything is checked, so the checks and the write apply to the same row, whatever other + # transactions save, delete or create under this id meanwhile. + project = q.projects.lock(user_id, project_id) + if project is None: + raise ProjectRecordNotFoundError(project_id) + current_minimum_version = project.minimum_canvas_schema_version + _require_supported_schema( + project_id=project_id, + minimum_version=current_minimum_version, + client_maximum_version=max_canvas_schema_version, + ) + if project.revision != expected_revision: + raise ProjectRecordConflictError(project_id, expected_revision, project.revision) + next_minimum_version = ( + current_minimum_version if minimum_canvas_schema_version is None else minimum_canvas_schema_version + ) + if next_minimum_version < current_minimum_version: + raise ProjectCanvasSchemaDowngradeError( + project_id=project_id, + current_version=current_minimum_version, + requested_version=next_minimum_version, + ) + _require_supported_schema( + project_id=project_id, + minimum_version=next_minimum_version, + client_maximum_version=max_canvas_schema_version, + ) + _require_project_document_size(document_bytes) + + q.projects.save( + user_id=user_id, + project_id=project_id, + name=name, + data=document_json, + minimum_canvas_schema_version=next_minimum_version, + ) + q.boards.rename(project.board_id, name[:BOARD_NAME_MAX_LENGTH]) + q.media_references.replace( + owner_kind="project", user_id=user_id, owner_id=project_id, references=references + ) + return q.projects.summary(user_id, project_id) + + return _record(self._queries.run(save), project_id, data) + + def delete(self, user_id: str, project_id: str) -> None: + def delete(q: Queries) -> None: + # The project's row, then its board's, as every write of a project locks them: the board deleted below + # is the one of the project deleted here. + project = q.projects.lock(user_id, project_id) + if project is None: + return + q.boards.lock(project.board_id) + q.projects.delete(user_id, project_id) + q.media_references.delete(owner_kind="project", user_id=user_id, owner_id=project_id) + # Only after the project is gone: a claimed board cannot be deleted. Deleting the board deletes its + # memberships, returning the media to Uncategorized; the image and video records and their files are + # untouched. + q.boards.delete_if_unclaimed(project.board_id) + + self._queries.run(delete) + + def get_board_snapshot(self, user_id: str, project_id: str) -> ProjectBoardSnapshotDTO: + # Bounded, unlike the board's counts: the caller holds the whole answer in memory and so does this, and + # the route is reachable by anyone with a project id. + items = self._queries.projects.board_items(user_id, project_id, PROJECT_BOARD_SNAPSHOT_MAX_ITEMS + 1) + if items is None: + raise ProjectRecordNotFoundError(project_id) + if len(items) > PROJECT_BOARD_SNAPSHOT_MAX_ITEMS: + raise ProjectBoardTooLargeError(project_id, PROJECT_BOARD_SNAPSHOT_MAX_ITEMS) + return ProjectBoardSnapshotDTO(items=items) + + def get_board_id(self, user_id: str, project_id: str) -> str: + board_id = self._queries.projects.board_id(user_id, project_id) + if board_id is None: + raise ProjectRecordNotFoundError(project_id) + return board_id diff --git a/invokeai/app/services/project_records/project_records_sqlite.py b/invokeai/app/services/project_records/project_records_sqlite.py deleted file mode 100644 index 4f8ea95f8d7..00000000000 --- a/invokeai/app/services/project_records/project_records_sqlite.py +++ /dev/null @@ -1,526 +0,0 @@ -import json -import sqlite3 -import uuid -from typing import Any, Optional - -from invokeai.app.services.board_records.board_records_common import BOARD_NAME_MAX_LENGTH, BoardVisibility -from invokeai.app.services.image_records.image_records_common import ASSETS_CATEGORIES, IMAGE_CATEGORIES -from invokeai.app.services.invoker import Invoker -from invokeai.app.services.project_records.project_records_base import ProjectRecordsStorageBase -from invokeai.app.services.project_records.project_records_common import ( - DEFAULT_PROJECT_CANVAS_SCHEMA_VERSION, - PROJECT_BOARD_SNAPSHOT_MAX_ITEMS, - PROJECT_DOCUMENT_MAX_BYTES, - ProjectBoardItemDTO, - ProjectBoardNotFoundError, - ProjectBoardSnapshotDTO, - ProjectBoardTooLargeError, - ProjectBoardUnavailableError, - ProjectCanvasSchemaDowngradeError, - ProjectCanvasSchemaUnsupportedError, - ProjectDocumentInvalidError, - ProjectDocumentTooLargeError, - ProjectRecordConflictError, - ProjectRecordDTO, - ProjectRecordExistsError, - ProjectRecordNotFoundError, - ProjectSummaryDTO, -) -from invokeai.app.services.shared.media_references import ( - delete_media_references, - extract_media_references, - replace_media_references, -) -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.util.misc import uuid_string - -# The categories a board view lists, straight from the gallery's own definition of them. Spelling -# them out in the SQL below would put a fourth copy of this list in the least discoverable place -# there is: `other` — the canvas's private category — is in neither, which is exactly why the -# snapshot must not carry it, and that reasoning belongs with the constants that encode it. -_VISIBLE_BOARD_CATEGORIES = tuple(category.value for category in (*IMAGE_CATEGORIES, *ASSETS_CATEGORIES)) -_VISIBLE_CATEGORY_PLACEHOLDERS = ", ".join("?" for _ in _VISIBLE_BOARD_CATEGORIES) - - -def _serialize_project_document(data: dict[str, Any]) -> tuple[str, int]: - try: - document_json = json.dumps(data, ensure_ascii=False, allow_nan=False, separators=(",", ":")) - return document_json, len(document_json.encode("utf-8")) - except (TypeError, ValueError, UnicodeEncodeError) as error: - raise ProjectDocumentInvalidError from error - - -def _require_project_document_size(actual_bytes: int) -> None: - if actual_bytes > PROJECT_DOCUMENT_MAX_BYTES: - raise ProjectDocumentTooLargeError(actual_bytes=actual_bytes, max_bytes=PROJECT_DOCUMENT_MAX_BYTES) - - -class ProjectRecordsSqlite(ProjectRecordsStorageBase): - """SQLite implementation of per-user project document storage. - - A project and its board are one unit: every write touching both happens inside a single - `self._db.transaction()`. - - That transaction is why this class talks to `boards` in raw SQL rather than through - `BoardService`. `SqliteDatabase.transaction()` guards a re-entrant lock and commits on the way - out of *every* level, so calling another service from inside an open transaction would commit - this one's half-finished work — including `self.get()`. - """ - - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - - def start(self, invoker: Invoker) -> None: - self._invoker = invoker - - def create( - self, - user_id: str, - name: str, - data: dict[str, Any], - project_id: str | None = None, - board_id: str | None = None, - minimum_canvas_schema_version: int = DEFAULT_PROJECT_CANVAS_SCHEMA_VERSION, - max_canvas_schema_version: int = DEFAULT_PROJECT_CANVAS_SCHEMA_VERSION, - ) -> ProjectRecordDTO: - project_id = project_id or uuid.uuid4().hex - self._require_supported_schema( - project_id=project_id, - minimum_version=minimum_canvas_schema_version, - client_maximum_version=max_canvas_schema_version, - ) - document_json, document_bytes = _serialize_project_document(data) - _require_project_document_size(document_bytes) - references = extract_media_references(data) - - try: - with self._db.transaction() as cursor: - if board_id is None: - resolved_board_id = self._insert_board(cursor, user_id=user_id, name=name) - else: - self._claim_board(cursor, user_id=user_id, board_id=board_id, name=name, project_id=project_id) - resolved_board_id = board_id - - cursor.execute( - """--sql - INSERT INTO projects ( - project_id, user_id, name, data, board_id, minimum_canvas_schema_version - ) - VALUES (?, ?, ?, ?, ?, ?); - """, - ( - project_id, - user_id, - name, - document_json, - resolved_board_id, - minimum_canvas_schema_version, - ), - ) - # Indexed on the same transaction as the document: a reader can never see a saved - # project whose referenced media the intermediates cleanup does not know about. - replace_media_references( - cursor, - owner_kind="project", - user_id=user_id, - owner_id=project_id, - references=references, - ) - record = self._fetch_record(cursor, user_id=user_id, project_id=project_id) - except sqlite3.IntegrityError as e: - # The board insert/rename rolls back with the project insert, so a rejected create never - # leaves an orphan board behind. - classified = self._classify_integrity_error(e, project_id=project_id, board_id=board_id) - - if classified is None: - raise - raise classified from e - - if record is None: - raise ProjectRecordNotFoundError(project_id) - - return record - - def get( - self, - user_id: str, - project_id: str, - max_canvas_schema_version: int = DEFAULT_PROJECT_CANVAS_SCHEMA_VERSION, - ) -> ProjectRecordDTO: - with self._db.transaction() as cursor: - record = self._fetch_record(cursor, user_id=user_id, project_id=project_id) - - if record is None: - raise ProjectRecordNotFoundError(project_id) - - self._require_supported_schema( - project_id=project_id, - minimum_version=record.minimum_canvas_schema_version, - client_maximum_version=max_canvas_schema_version, - ) - - return record - - def list(self, user_id: str) -> list[ProjectSummaryDTO]: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT project_id, board_id, name, revision, minimum_canvas_schema_version, - created_at, updated_at - FROM projects - WHERE user_id = ? - -- rowid breaks ties between rows created in the same millisecond, - -- keeping the listing in true insertion order. - ORDER BY created_at ASC, rowid ASC; - """, - (user_id,), - ) - rows = cursor.fetchall() - - return [ - ProjectSummaryDTO( - project_id=row[0], - board_id=row[1], - name=row[2], - revision=row[3], - minimum_canvas_schema_version=row[4], - created_at=row[5], - updated_at=row[6], - ) - for row in rows - ] - - def update( - self, - user_id: str, - project_id: str, - expected_revision: int, - name: str, - data: dict[str, Any], - minimum_canvas_schema_version: int | None = None, - max_canvas_schema_version: int = DEFAULT_PROJECT_CANVAS_SCHEMA_VERSION, - ) -> ProjectRecordDTO: - document_json, document_bytes = _serialize_project_document(data) - references = extract_media_references(data) - - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT revision, minimum_canvas_schema_version - FROM projects - WHERE user_id = ? AND project_id = ?; - """, - (user_id, project_id), - ) - row = cursor.fetchone() - - if row is None: - raise ProjectRecordNotFoundError(project_id) - - current_revision, current_minimum_version = row - self._require_supported_schema( - project_id=project_id, - minimum_version=current_minimum_version, - client_maximum_version=max_canvas_schema_version, - ) - - if current_revision != expected_revision: - raise ProjectRecordConflictError(project_id, expected_revision, current_revision) - - next_minimum_version = ( - current_minimum_version if minimum_canvas_schema_version is None else minimum_canvas_schema_version - ) - - if next_minimum_version < current_minimum_version: - raise ProjectCanvasSchemaDowngradeError( - project_id=project_id, - current_version=current_minimum_version, - requested_version=next_minimum_version, - ) - - self._require_supported_schema( - project_id=project_id, - minimum_version=next_minimum_version, - client_maximum_version=max_canvas_schema_version, - ) - _require_project_document_size(document_bytes) - - cursor.execute( - """--sql - UPDATE projects - SET name = ?, data = ?, minimum_canvas_schema_version = ?, revision = revision + 1 - WHERE user_id = ? AND project_id = ? AND revision = ?; - """, - ( - name, - document_json, - next_minimum_version, - user_id, - project_id, - expected_revision, - ), - ) - - if cursor.rowcount == 0: - # Another SQLite connection can win after the read above but before this - # compare-and-swap. Re-read inside the transaction so the refusal reports the - # actual winning revision and never renames the board. - cursor.execute( - """--sql - SELECT revision FROM projects - WHERE user_id = ? AND project_id = ?; - """, - (user_id, project_id), - ) - raced_row = cursor.fetchone() - - if raced_row is None: - raise ProjectRecordNotFoundError(project_id) - - raise ProjectRecordConflictError(project_id, expected_revision, raced_row[0]) - - cursor.execute( - """--sql - UPDATE boards - SET board_name = ? - WHERE board_id = (SELECT board_id FROM projects WHERE user_id = ? AND project_id = ?); - """, - (name[:BOARD_NAME_MAX_LENGTH], user_id, project_id), - ) - replace_media_references( - cursor, - owner_kind="project", - user_id=user_id, - owner_id=project_id, - references=references, - ) - record = self._fetch_record(cursor, user_id=user_id, project_id=project_id) - - if record is None: - raise ProjectRecordNotFoundError(project_id) - - return record - - @staticmethod - def _fetch_record(cursor: sqlite3.Cursor, *, user_id: str, project_id: str) -> ProjectRecordDTO | None: - """Read a project through the caller's transaction so a write returns exactly its own row.""" - cursor.execute( - """--sql - SELECT project_id, board_id, name, data, revision, minimum_canvas_schema_version, - created_at, updated_at - FROM projects - WHERE user_id = ? AND project_id = ?; - """, - (user_id, project_id), - ) - row = cursor.fetchone() - - if row is None: - return None - - return ProjectRecordDTO( - project_id=row[0], - board_id=row[1], - name=row[2], - data=json.loads(row[3]), - revision=row[4], - minimum_canvas_schema_version=row[5], - created_at=row[6], - updated_at=row[7], - ) - - @staticmethod - def _require_supported_schema(*, project_id: str, minimum_version: int, client_maximum_version: int) -> None: - if client_maximum_version < minimum_version: - raise ProjectCanvasSchemaUnsupportedError( - project_id=project_id, - minimum_version=minimum_version, - client_maximum_version=client_maximum_version, - ) - - def delete(self, user_id: str, project_id: str) -> None: - with self._db.transaction() as cursor: - row = self._board_row(cursor, user_id=user_id, project_id=project_id) - - if row is None: - return - - cursor.execute( - """--sql - DELETE FROM projects - WHERE user_id = ? AND project_id = ?; - """, - (user_id, project_id), - ) - delete_media_references(cursor, owner_kind="project", user_id=user_id, owner_id=project_id) - # Only after the project is gone: the board FK is RESTRICT, so a claimed board cannot be - # deleted. Deleting the board cascades its memberships, returning the media to - # Uncategorized; the image and video records and their files are untouched. - cursor.execute( - """--sql - DELETE FROM boards WHERE board_id = ?; - """, - (row[0],), - ) - - def get_board_snapshot(self, user_id: str, project_id: str) -> ProjectBoardSnapshotDTO: - with self._db.transaction() as cursor: - row = self._board_row(cursor, user_id=user_id, project_id=project_id) - - if row is None: - raise ProjectRecordNotFoundError(project_id) - - # One union over both namespaces, shaped like `get_board_media_summaries`. The filters - # are what make the result match what the gallery shows for this board: intermediates - # are hidden, and the categories come from IMAGE_CATEGORIES/ASSETS_CATEGORIES, which - # `other` — the canvas's private category — is deliberately in neither of. - # - # `deleted_at` is deliberately *not* filtered. Soft delete is unused and no other - # board query consults it, so filtering here would make the snapshot disagree with the - # counts shown beside the board. - # - # Bounded, unlike the counts: the caller holds the whole answer in memory and so does - # this, and the route is reachable by anyone with a project id. - cursor.execute( - f"""--sql - SELECT 'image' AS kind, images.image_name AS name, - images.image_category AS category, images.starred AS starred - FROM board_images - INNER JOIN images ON board_images.image_name = images.image_name - WHERE board_images.board_id = ? - AND images.is_intermediate = FALSE - AND images.image_category IN ({_VISIBLE_CATEGORY_PLACEHOLDERS}) - UNION ALL - SELECT 'video' AS kind, videos.video_name AS name, - videos.video_category AS category, videos.starred AS starred - FROM board_videos - INNER JOIN videos ON board_videos.video_name = videos.video_name - WHERE board_videos.board_id = ? - AND videos.is_intermediate = FALSE - AND videos.video_category IN ({_VISIBLE_CATEGORY_PLACEHOLDERS}) - ORDER BY kind ASC, name ASC - LIMIT ?; - """, - ( - row[0], - *_VISIBLE_BOARD_CATEGORIES, - row[0], - *_VISIBLE_BOARD_CATEGORIES, - PROJECT_BOARD_SNAPSHOT_MAX_ITEMS + 1, - ), - ) - rows = cursor.fetchall() - - if len(rows) > PROJECT_BOARD_SNAPSHOT_MAX_ITEMS: - raise ProjectBoardTooLargeError(project_id, PROJECT_BOARD_SNAPSHOT_MAX_ITEMS) - - return ProjectBoardSnapshotDTO( - items=[ProjectBoardItemDTO(kind=r[0], name=r[1], category=r[2], starred=bool(r[3])) for r in rows] - ) - - def get_board_id(self, user_id: str, project_id: str) -> str: - with self._db.transaction() as cursor: - row = self._board_row(cursor, user_id=user_id, project_id=project_id) - - if row is None: - raise ProjectRecordNotFoundError(project_id) - - return row[0] - - def _board_row(self, cursor: sqlite3.Cursor, *, user_id: str, project_id: str) -> Optional[sqlite3.Row]: - """The project's board row, or `None` when this user has no such project. - - Cursor-taking rather than transaction-opening, so callers stay in charge of the unit of work - this class's docstring describes. - """ - cursor.execute( - """--sql - SELECT board_id FROM projects WHERE user_id = ? AND project_id = ?; - """, - (user_id, project_id), - ) - - return cursor.fetchone() - - def _insert_board(self, cursor: sqlite3.Cursor, *, user_id: str, name: str) -> str: - board_id = uuid_string() - cursor.execute( - """--sql - INSERT INTO boards (board_id, board_name, user_id, board_visibility, archived) - VALUES (?, ?, ?, ?, FALSE); - """, - (board_id, name[:BOARD_NAME_MAX_LENGTH], user_id, BoardVisibility.Private.value), - ) - return board_id - - def _claim_board(self, cursor: sqlite3.Cursor, *, user_id: str, board_id: str, name: str, project_id: str) -> None: - """Take ownership of an existing private board, or explain why it cannot be taken. - - One query answers every way a claim can fail. The `UNIQUE` constraint on `projects.board_id` - settles the race this check cannot see: two imports claiming the same board both pass here, - and exactly one of them survives the insert. - - The project being created is excluded from the "already claimed" test, so a client that - re-sends a create it never got an answer for reaches the project insert and is refused there, - as `ProjectRecordExistsError` — the truth. Without the exclusion it would be told its board - was unavailable, which reads as "somebody else took it" and is the opposite of what happened. - A repeated create is not an edge case: it is how a client establishes whether a request whose - response was lost actually landed, and the answer decides whether it deletes the media it - uploaded. - """ - cursor.execute( - """--sql - SELECT - b.user_id, - b.board_visibility, - EXISTS(SELECT 1 FROM shared_boards s WHERE s.board_id = b.board_id), - EXISTS( - SELECT 1 FROM projects p - WHERE p.board_id = b.board_id AND NOT (p.user_id = ? AND p.project_id = ?) - ) - FROM boards b - WHERE b.board_id = ?; - """, - (user_id, project_id, board_id), - ) - row = cursor.fetchone() - - if row is None or row[0] != user_id: - raise ProjectBoardNotFoundError(board_id) - - if row[1] != BoardVisibility.Private.value or row[2] or row[3]: - raise ProjectBoardUnavailableError(board_id) - - # Un-archived as well as renamed. A project's board takes its archived state from the - # project, and `PATCH /boards/{id}` refuses to set it on a claimed board — so a board that - # arrived archived would be invisible in every listing with no API left to fix it. - cursor.execute( - """--sql - UPDATE boards SET board_name = ?, archived = FALSE WHERE board_id = ?; - """, - (name[:BOARD_NAME_MAX_LENGTH], board_id), - ) - - @staticmethod - def _classify_integrity_error( - error: sqlite3.IntegrityError, *, project_id: str, board_id: Optional[str] - ) -> Optional[Exception]: - """Name which constraint a rejected insert violated, or `None` to let it through unchanged. - - Two are expected: the composite primary key, meaning the user already has this project id, - and `projects.board_id`, meaning another project claimed the board between `_claim_board` - looking and this insert committing. - - Decided on `sqlite_errorname` rather than on the message text. A substring match would call - a foreign-key failure — a user row that vanished, say — "project already exists", sending - the client to pick a new id forever over a problem no new id can fix. Anything unrecognized - is re-raised as itself rather than guessed at. - """ - errorname = getattr(error, "sqlite_errorname", "") - - if errorname == "SQLITE_CONSTRAINT_UNIQUE" and board_id is not None: - return ProjectBoardUnavailableError(board_id) - if errorname == "SQLITE_CONSTRAINT_PRIMARYKEY": - return ProjectRecordExistsError(project_id) - return None diff --git a/invokeai/app/services/session_processor/workflow_call_batch.py b/invokeai/app/services/session_processor/workflow_call_batch.py index fe492fbb320..d2416c058b3 100644 --- a/invokeai/app/services/session_processor/workflow_call_batch.py +++ b/invokeai/app/services/session_processor/workflow_call_batch.py @@ -22,7 +22,7 @@ ) from invokeai.app.services.shared.execution_state_migration import load_execution_state from invokeai.app.services.shared.graph import GraphExecutionState, WorkflowCallFrame -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.services.shared.workflow_graph_builder import ( UnsupportedWorkflowNodeError, apply_workflow_inputs_to_workflow, diff --git a/invokeai/app/services/session_queue/session_queue_base.py b/invokeai/app/services/session_queue/session_queue_base.py index 65eac12d193..20aa4b278aa 100644 --- a/invokeai/app/services/session_queue/session_queue_base.py +++ b/invokeai/app/services/session_queue/session_queue_base.py @@ -27,8 +27,7 @@ SessionQueueStatus, ) from invokeai.app.services.shared.graph import GraphExecutionState -from invokeai.app.services.shared.pagination import CursorPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import CursorPaginatedResults, SQLiteDirection @dataclass(frozen=True) diff --git a/invokeai/app/services/session_queue/session_queue_common.py b/invokeai/app/services/session_queue/session_queue_common.py index 6ca268a7379..ef749015f53 100644 --- a/invokeai/app/services/session_queue/session_queue_common.py +++ b/invokeai/app/services/session_queue/session_queue_common.py @@ -101,7 +101,7 @@ class BatchDatum(BaseModel): class Batch(BaseModel): - batch_id: str = Field(default_factory=uuid_string, description="The ID of the batch") + batch_id: str = Field(default_factory=uuid_string, max_length=255, description="The ID of the batch") idempotency_key: str | None = Field( default=None, min_length=1, diff --git a/invokeai/app/services/session_queue/session_queue_default.py b/invokeai/app/services/session_queue/session_queue_default.py new file mode 100644 index 00000000000..65c3b6fa006 --- /dev/null +++ b/invokeai/app/services/session_queue/session_queue_default.py @@ -0,0 +1,1178 @@ +import asyncio +import hashlib +import json +from datetime import timedelta +from typing import Any, Literal, Optional + +from pydantic_core import to_jsonable_python + +from invokeai.app.services.invoker import Invoker +from invokeai.app.services.session_queue.session_queue_base import SessionQueueBase, WorkflowCallChildCompletion +from invokeai.app.services.session_queue.session_queue_common import ( + DEFAULT_QUEUE_ID, + QUEUE_ITEM_STATUS, + Batch, + BatchStatus, + CancelAllExceptCurrentResult, + CancelByBatchIDsResult, + CancelByDestinationResult, + CancelByQueueIDResult, + ClearResult, + DeleteAllExceptCurrentResult, + DeleteByDestinationResult, + EnqueueBatchReceipt, + EnqueueBatchResult, + EnqueueIdempotencyConflictError, + EnqueueProjectNotFoundError, + EnqueueReceiptLimitError, + IsEmptyResult, + IsFullResult, + ItemIdsResult, + NodeFieldValue, + PruneResult, + RetryItemsResult, + SessionQueueCountsByDestination, + SessionQueueItem, + SessionQueueItemChangedError, + SessionQueueItemNotFoundError, + SessionQueueItemSummary, + SessionQueueStatus, + TooManySessionsError, + ValueToInsertTuple, + calc_session_count, + prepare_values_to_insert, + uuid_string, +) +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.locks import DatabaseLock +from invokeai.app.services.shared.database.queries.session_queue import Scope, SettledReceipt +from invokeai.app.services.shared.execution_state_migration import dump_execution_state +from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from invokeai.app.services.shared.pagination import CursorPaginatedResults, SQLiteDirection + +MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPTS_PER_OWNER = 10_000 +MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPT_BYTES_PER_OWNER = 64 * 1024 * 1024 +MAX_ENQUEUE_RECEIPTS_PER_OWNER = 100_000 +MAX_ENQUEUE_RECEIPT_BYTES_PER_OWNER = 128 * 1024 * 1024 +ACKNOWLEDGED_ENQUEUE_RECEIPT_RETENTION = timedelta(days=7) + +# How far past the fairness-chosen candidate (in item_id distance) the device-affinity swap may +# look for a warm-model item. This bounds two things at once: the scoring query's cost (at most +# this many session blobs are scanned per dequeue) and how long a cold item can be deferred — the +# starved candidate's item_id never changes, so newly enqueued warm items eventually fall outside +# the window and the cold item runs after at most ~this many swaps. +AFFINITY_MAX_LOOKAHEAD = 32 + +_FINISHED = ("completed", "failed", "canceled") + +# Every write that adds items takes the queue's admission lock, so that two of them neither both fit into its last +# free places nor both write one enqueue receipt, and shares the media protection: the items' sessions make the media +# they name active, which the intermediates cleanup must not delete between its check and its delete. +_ADMISSION = (DatabaseLock.SESSION_QUEUE_ADMISSION,) +_PROTECTING_MEDIA = (DatabaseLock.MEDIA_PROTECTION,) + + +def _session_json(session: GraphExecutionState) -> str: + # Persisted sessions resume execution across queue boundaries. + return json.dumps(dump_execution_state(session), default=to_jsonable_python) + + +def _decode_enqueue_receipt(batch_id: str, requested: int, enqueued: int, item_ids_json: str) -> EnqueueBatchReceipt: + item_ids = json.loads(item_ids_json) + if not isinstance(item_ids, list) or not all(isinstance(item_id, int) for item_id in item_ids): + raise RuntimeError("Stored queue enqueue receipt contains invalid item ids") + return EnqueueBatchReceipt(batch_id=batch_id, requested=requested, enqueued=enqueued, item_ids=item_ids) + + +class SessionQueue(SessionQueueBase): + __invoker: Invoker + + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + + def start(self, invoker: Invoker) -> None: + self.__invoker = invoker + self._set_in_progress_to_canceled() + config = self.__invoker.services.configuration + if config.clear_queue_on_startup: + clear_result = self.clear(DEFAULT_QUEUE_ID) + if clear_result.deleted > 0: + self.__invoker.services.logger.info(f"Cleared all {clear_result.deleted} queue items") + return + + if config.max_queue_history is not None: + deleted = self._prune_terminal_to_limit(DEFAULT_QUEUE_ID, config.max_queue_history) + if deleted > 0: + self.__invoker.services.logger.info( + f"Pruned {deleted} completed/failed/canceled queue items (kept up to {config.max_queue_history})" + ) + + def _set_in_progress_to_canceled(self) -> None: + """ + Sets all in_progress or waiting queue items to canceled. Run on app startup, not associated with any queue. + This is necessary because the invoker may have been killed while processing a queue item or while a parent + queue item was suspended waiting on a child workflow execution. + """ + self._queries.session_queue.cancel_interrupted() + + def _prune_terminal_to_limit(self, queue_id: str, keep: int) -> int: + """Prune terminal items (completed/failed/canceled) to keep at most N most-recent items.""" + return self._queries.session_queue.prune(queue_id, user_id=None, keep=keep) + + async def enqueue_batch( + self, queue_id: str, batch: Batch, prepend: bool, user_id: str = "system" + ) -> EnqueueBatchResult: + # The route awaits this method, but every operation below is synchronous database/CPU work. Keep the + # complete enqueue in one worker operation, off the event loop. + return await asyncio.to_thread(self._enqueue_batch, queue_id, batch, prepend, user_id) + + def acknowledge_enqueue(self, queue_id: str, idempotency_key: str, user_id: str = "system") -> None: + self._queries.session_queue.acknowledge_receipt(queue_id, user_id, idempotency_key) + + def get_enqueue_receipt( + self, queue_id: str, idempotency_key: str, user_id: str = "system" + ) -> EnqueueBatchReceipt | None: + receipt = self._queries.session_queue.enqueue_receipt(queue_id, user_id, idempotency_key) + return _decode_enqueue_receipt(*receipt) if receipt is not None else None + + def _enqueue_batch(self, queue_id: str, batch: Batch, prepend: bool, user_id: str) -> EnqueueBatchResult: + requested_count = calc_session_count(batch=batch) + payload_hash = ( + hashlib.sha256( + json.dumps( + { + "batch": batch.model_dump(mode="json", exclude={"batch_id", "idempotency_key"}), + "prepend": prepend, + }, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + ).hexdigest() + if batch.idempotency_key is not None + else None + ) + key = batch.idempotency_key + + def settle(receipt: SettledReceipt) -> EnqueueBatchResult: + if receipt.payload_hash != payload_hash: + raise EnqueueIdempotencyConflictError(f"Idempotency key {key} is already used by another submission") + result = _decode_enqueue_receipt(receipt.batch_id, receipt.requested, receipt.enqueued, receipt.item_ids) + return EnqueueBatchResult( + queue_id=queue_id, + requested=result.requested, + enqueued=result.enqueued, + batch=batch.model_copy(update={"batch_id": result.batch_id}), + priority=receipt.priority, + item_ids=result.item_ids, + ) + + def require_project(q: Queries) -> None: + if batch.project_id is not None and q.projects.summary(user_id, batch.project_id) is None: + raise EnqueueProjectNotFoundError(batch.project_id) + + def survey(q: Queries) -> SettledReceipt | int: + """A settled receipt, or the pending items' count to size the values to prepare.""" + if key is not None and (receipt := q.session_queue.settled_receipt(queue_id, user_id, key)) is not None: + return receipt + require_project(q) + return q.session_queue.pending_count(queue_id) + + surveyed = self._queries.run(survey, read_only=True) + if isinstance(surveyed, SettledReceipt): + return settle(surveyed) + + # Preparing the values is CPU work, done outside a transaction; admission below takes what still fits. + max_queue_size = self.__invoker.services.configuration.max_queue_size + prepared_values = prepare_values_to_insert( + queue_id=queue_id, + batch=batch, + priority=0, + max_new_queue_items=max(0, max_queue_size - surveyed), + user_id=user_id, + ) + + def admit(q: Queries) -> SettledReceipt | EnqueueBatchResult: + q.locks.acquire(*_ADMISSION, also_shared=_PROTECTING_MEDIA) + if key is not None: + q.session_queue.delete_acknowledged_receipts(user_id, ACKNOWLEDGED_ENQUEUE_RECEIPT_RETENTION) + if (receipt := q.session_queue.settled_receipt(queue_id, user_id, key)) is not None: + return receipt + require_project(q) + + max_new_queue_items = max(0, max_queue_size - q.session_queue.pending_count(queue_id)) + priority = q.session_queue.top_pending_priority(queue_id) + 1 if prepend else 0 + values_to_insert = prepared_values[:max_new_queue_items] + if priority != 0: + values_to_insert = [(*value[:5], priority, *value[6:]) for value in values_to_insert] + enqueued_count = len(values_to_insert) + accepted_batch = batch + if enqueued_count > 0: + accepted_batch_id = batch.batch_id + while q.session_queue.batch_taken(queue_id, user_id, accepted_batch_id): + accepted_batch_id = uuid_string() + if accepted_batch_id != batch.batch_id: + accepted_batch = batch.model_copy(update={"batch_id": accepted_batch_id}) + values_to_insert = [(*value[:3], accepted_batch_id, *value[4:]) for value in values_to_insert] + usage = None + if key is not None and enqueued_count > 0: + usage = q.session_queue.receipt_usage(user_id) + if usage.count >= MAX_ENQUEUE_RECEIPTS_PER_OWNER: + raise EnqueueReceiptLimitError("Enqueue receipts exceed the count limit") + if usage.unacknowledged_count >= MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPTS_PER_OWNER: + raise EnqueueReceiptLimitError("Too many unacknowledged enqueue requests") + if usage.bytes >= MAX_ENQUEUE_RECEIPT_BYTES_PER_OWNER: + raise EnqueueReceiptLimitError("Enqueue receipts exceed the storage limit") + if usage.unacknowledged_bytes >= MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPT_BYTES_PER_OWNER: + raise EnqueueReceiptLimitError("Unacknowledged enqueue receipts exceed the storage limit") + q.session_queue.insert_items(values_to_insert) + item_ids = q.session_queue.batch_item_ids(queue_id, user_id, accepted_batch.batch_id) + if usage is not None: + receipt_byte_size = len( + json.dumps( + [ + queue_id, + user_id, + key, + payload_hash, + accepted_batch.batch_id, + requested_count, + enqueued_count, + priority, + item_ids, + ], + ensure_ascii=False, + separators=(",", ":"), + ).encode("utf-8") + ) + if usage.unacknowledged_bytes + receipt_byte_size > MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPT_BYTES_PER_OWNER: + raise EnqueueReceiptLimitError("Unacknowledged enqueue receipts exceed the storage limit") + if usage.bytes + receipt_byte_size > MAX_ENQUEUE_RECEIPT_BYTES_PER_OWNER: + raise EnqueueReceiptLimitError("Enqueue receipts exceed the storage limit") + q.session_queue.insert_receipt( + { + "queue_id": queue_id, + "user_id": user_id, + "idempotency_key": key, + "payload_hash": payload_hash, + "batch_id": accepted_batch.batch_id, + "requested": requested_count, + "enqueued": enqueued_count, + "priority": priority, + "item_ids": json.dumps(item_ids, separators=(",", ":")), + "byte_size": receipt_byte_size, + } + ) + return EnqueueBatchResult( + queue_id=queue_id, + requested=requested_count, + enqueued=enqueued_count, + batch=accepted_batch, + priority=priority, + item_ids=item_ids, + ) + + admitted = self._queries.run(admit) + if isinstance(admitted, SettledReceipt): + return settle(admitted) + self.__invoker.services.events.emit_batch_enqueued(admitted, user_id=user_id) + return admitted + + def dequeue(self, device: Optional[str] = None) -> Optional[SessionQueueItem]: + config = self.__invoker.services.configuration + use_round_robin = config.multiuser and config.session_queue_mode == "round_robin" + + # Snapshot the claiming device's warm models first: the lookup touches the ModelCache lock, which other + # threads may hold across long operations (VRAM transfers, cache clears). A slightly stale snapshot is fine + # for a heuristic. An explicitly configured session_queue_mode=FIFO is a request for strict insertion order, + # so it opts out of affinity reordering entirely (the setting defaults to round_robin, which keeps affinity + # active for single-user installs even though they use the FIFO order). + if config.session_queue_mode == "round_robin": + resident_model_keys = self._get_device_resident_model_keys(device) + else: + resident_model_keys = set() + + # Several workers (multi-GPU) may pick the same candidate: the claim moves it to in_progress only if it is + # still pending, so exactly one of them gets it and the others pick again. + while True: + raw_queue_item = self._queries.session_queue.next_to_run(round_robin=use_round_robin) + if raw_queue_item is None: + return None + queue_item, readable = self._hydrate_queue_item(raw_queue_item, quarantine=True) + if not readable: + continue + queue_item = self._apply_device_affinity(queue_item, resident_model_keys) + claimed = self._queries.session_queue.claim(queue_item.item_id, device) + if claimed is None: + continue + # Patch the item already read instead of re-reading (and re-parsing the session graph of) its row. The + # claim's columns include the device recorded, so the UI can label the item by GPU. + for field_name, value in claimed.items(): + setattr(queue_item, field_name, value) + self._emit_status_changed(queue_item) + return queue_item + + @staticmethod + def _make_unreadable_queue_item(raw_queue_item: dict[str, Any], error: Exception) -> SessionQueueItem: + """Build a metadata-preserving placeholder for an unreadable runtime snapshot.""" + placeholder = SessionQueueItem.model_construct(**raw_queue_item) + # A placeholder has no trustworthy execution result. Never expose a corrupt or newer snapshot as complete. + placeholder.status = "failed" + placeholder.session = GraphExecutionState(graph=Graph()) + placeholder.workflow = None + placeholder.field_values = None + placeholder._snapshot_readable = False + message = f"Unable to load execution state: {error}" + placeholder.error_type = type(error).__name__ + placeholder.error_message = message + placeholder.error_traceback = message + return placeholder + + def _hydrate_queue_item(self, raw_queue_item: dict[str, Any], *, quarantine: bool) -> tuple[SessionQueueItem, bool]: + """Hydrate one queue row without letting an unreadable snapshot break queue access.""" + try: + return SessionQueueItem.queue_item_from_dict(raw_queue_item), True + except (TypeError, ValueError) as exc: + if quarantine: + return self._quarantine_unreadable_queue_item(raw_queue_item, exc), False + return self._make_unreadable_queue_item(raw_queue_item, exc), False + + def _project_queue_item_for_read(self, raw_queue_item: dict[str, Any]) -> SessionQueueItem: + """Read queue metadata and response results without rebuilding runtime execution state.""" + try: + return SessionQueueItem.queue_item_from_dict(raw_queue_item, hydrate_runtime=False) + except (TypeError, ValueError) as exc: + return self._make_unreadable_queue_item(raw_queue_item, exc) + + def _get_queue_item_for_api(self, item_id: int) -> SessionQueueItem: + """Read queue metadata and response results without rebuilding execution runtime state.""" + raw_queue_item = self._queries.session_queue.item(item_id) + if raw_queue_item is None: + raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}") + return self._project_queue_item_for_read(raw_queue_item) + + def _quarantine_unreadable_queue_item(self, raw_queue_item: dict[str, Any], error: Exception) -> SessionQueueItem: + """Fail a pending row whose runtime snapshot is newer than this worker can read. + + The real session cannot be hydrated, so use a minimal in-memory placeholder only for the + status transition/event. The persisted session remains untouched for postmortem recovery. + """ + + placeholder = self._make_unreadable_queue_item(raw_queue_item, error) + placeholder.status = "pending" + return self._set_queue_item_status( + item_id=placeholder.item_id, + status="failed", + error_type=placeholder.error_type, + error_message=placeholder.error_message, + error_traceback=placeholder.error_traceback, + queue_item=placeholder, + ) + + def _apply_device_affinity(self, candidate: SessionQueueItem, resident_keys: set[str]) -> SessionQueueItem: + """Swap the fairness-chosen candidate for a nearby same-user, same-priority pending item + whose models are already cached on the claiming device, if one exists. + + Cross-device model reloads are expensive (tens of seconds for large models), so when a user + has queued a mix of models, preferring an item whose models are warm on the freeing GPU cuts + thrash. Fairness is preserved by construction: round-robin decides *which user* is served and + priority ordering decides *which tier* of their items is eligible; this heuristic only + reorders within that user's equal-priority pending items, and only within + AFFINITY_MAX_LOOKAHEAD of the candidate's item_id, so a cold item's deferral is bounded. + The caller passes an empty key set to disable affinity (legacy single-device mode, explicit + FIFO mode, or cache introspection unavailable). + """ + if not resident_keys: + return candidate + # Model keys are UUID strings that appear verbatim in the session JSON, so residency can be + # scored with substring matches — no need to parse each candidate's session. Sort for + # deterministic parameter binding. + warm = self._queries.session_queue.warmest( + user_id=candidate.user_id, + priority=candidate.priority, + after_item_id=candidate.item_id, + lookahead=AFFINITY_MAX_LOOKAHEAD, + resident_keys=sorted(resident_keys), + ) + if warm is None or warm["item_id"] == candidate.item_id: + # No warm-model item for this user (or the candidate already is one) — keep the + # fairness-chosen candidate. + return candidate + queue_item, readable = self._hydrate_queue_item(warm, quarantine=True) + return queue_item if readable else candidate + + def _get_device_resident_model_keys(self, device: Optional[str]) -> set[str]: + """Best-effort lookup of the model keys currently cached for the given generation device.""" + if device is None: + return set() + try: + cache = self.__invoker.services.model_manager.load.ram_caches.get(device) + if cache is None: + return set() + return set(cache.cached_model_keys()) + except Exception: + # Affinity is purely an optimization — dequeue must never fail because cache + # introspection did (e.g. model manager not fully started, or mocked in tests). + return set() + + def get_next(self, queue_id: str, origin_prefix: Optional[str] = None) -> Optional[SessionQueueItem]: + raw_queue_item = self._queries.session_queue.item_by_status(queue_id, "pending", origin_prefix) + return self._hydrate_queue_item(raw_queue_item, quarantine=False)[0] if raw_queue_item is not None else None + + def get_current(self, queue_id: str, origin_prefix: Optional[str] = None) -> Optional[SessionQueueItem]: + raw_queue_item = self._queries.session_queue.item_by_status(queue_id, "in_progress", origin_prefix) + return self._hydrate_queue_item(raw_queue_item, quarantine=False)[0] if raw_queue_item is not None else None + + def _get_queue_item_by_status_for_api( + self, queue_id: str, status: Literal["pending", "in_progress"], origin_prefix: Optional[str] + ) -> Optional[SessionQueueItem]: + raw_queue_item = self._queries.session_queue.item_by_status(queue_id, status, origin_prefix) + return self._project_queue_item_for_read(raw_queue_item) if raw_queue_item is not None else None + + def get_current_for_api(self, queue_id: str, origin_prefix: Optional[str] = None) -> Optional[SessionQueueItem]: + return self._get_queue_item_by_status_for_api(queue_id, "in_progress", origin_prefix) + + def get_next_for_api(self, queue_id: str, origin_prefix: Optional[str] = None) -> Optional[SessionQueueItem]: + return self._get_queue_item_by_status_for_api(queue_id, "pending", origin_prefix) + + def _set_queue_item_status( + self, + item_id: int, + status: QUEUE_ITEM_STATUS, + error_type: Optional[str] = None, + error_message: Optional[str] = None, + error_traceback: Optional[str] = None, + device: Optional[str] = None, + queue_item: Optional[SessionQueueItem] = None, + ) -> SessionQueueItem: + return self._transition_queue_item_status( + item_id=item_id, + status=status, + error_type=error_type, + error_message=error_message, + error_traceback=error_traceback, + device=device, + queue_item=queue_item, + )[0] + + def _transition_queue_item_status( + self, + item_id: int, + status: QUEUE_ITEM_STATUS, + error_type: Optional[str] = None, + error_message: Optional[str] = None, + error_traceback: Optional[str] = None, + device: Optional[str] = None, + queue_item: Optional[SessionQueueItem] = None, + ) -> tuple[SessionQueueItem, bool]: + """Move a queue item to `status` unless it is already finished (completed, failed or + canceled), returning the item and whether THIS call performed the transition. + + The guard is the UPDATE's own condition, so two callers racing to cancel the same row cannot + both observe themselves as the one that canceled it — exactly one sees `transitioned=True` + (the bulk-cancel counters rely on this), on every backend. When no transition happens, the + item is returned unchanged and no status-changed event is emitted; a vanished row raises + SessionQueueItemNotFoundError. + """ + if queue_item is not None and queue_item.item_id != item_id: + raise ValueError(f"Queue item {queue_item.item_id} does not match requested item {item_id}") + + exists, changed = self._queries.session_queue.transition( + item_id, + status, + error_type=error_type, + error_message=error_message, + error_traceback=error_traceback, + device=device, + ) + if not exists: + raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}") + if changed is None: + # Already finished: return it unchanged (get_queue_item raises if it was deleted meanwhile). + return self.get_queue_item(item_id), False + + if queue_item is None: + queue_item = self.get_queue_item(item_id) + else: + # `device` is among the changed columns, so a caller that supplied its own queue_item still + # sees the device this transition recorded rather than a stale value. + for field_name, value in changed.items(): + setattr(queue_item, field_name, value) + self._emit_status_changed(queue_item) + return queue_item, True + + def _emit_status_changed(self, queue_item: SessionQueueItem) -> None: + batch_status = self.get_batch_status(queue_id=queue_item.queue_id, batch_id=queue_item.batch_id) + # The QueueItemStatusChangedEvent ships to user:{queue_item.user_id} and admin rooms. + # acting_user_id ensures the embedded current-item identifiers are redacted when the + # in-progress item belongs to someone else, while leaving aggregate counts global. + # Doing this inside get_queue_status guarantees the redaction decision and the + # embedded identifiers come from the same lightweight metadata snapshot, eliminating the + # race where a second read could find None and skip scrubbing stale identifiers. + # user_id additionally embeds the owner's per-user counts so the owner's client can + # apply the event's queue_status optimistically without waiting for a refetch; the + # sanitized companion nulls them before reaching anyone else. + queue_status = self.get_queue_status( + queue_id=queue_item.queue_id, user_id=queue_item.user_id, acting_user_id=queue_item.user_id + ) + self.__invoker.services.events.emit_queue_item_status_changed(queue_item, batch_status, queue_status) + + def _get_workflow_call_descendant_ids(self, item_id: int) -> list[int]: + return self._queries.session_queue.descendant_ids(item_id) + + def _get_workflow_call_chain_item_ids(self, item_id: int) -> list[int]: + chain_item_ids = self._queries.session_queue.chain_item_ids(item_id) + if chain_item_ids is None: + raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}, or one of its ancestors") + return chain_item_ids + + def is_empty(self, queue_id: str) -> IsEmptyResult: + return IsEmptyResult(is_empty=self._queries.session_queue.count(queue_id) == 0) + + def is_full(self, queue_id: str) -> IsFullResult: + max_queue_size = self.__invoker.services.configuration.max_queue_size + return IsFullResult(is_full=self._queries.session_queue.count(queue_id) >= max_queue_size) + + def _cancel_in_progress(self, scope: Scope) -> list[int]: + """Cancel every in-progress item of the scope, emitting a cancel event for each. + + The bulk cancels leave in-progress items out of their single UPDATE, because a running item + must be canceled via `_transition_queue_item_status()` so that its `QueueItemStatusChangedEvent` + is emitted — the session processor responds to that event by setting the cancel event of the + worker running that exact item_id. With multiple workers (multi-GPU) more than one item can be + in_progress at once, so each is canceled individually. + + Returns the item ids of the in-progress items actually canceled. + """ + canceled: list[int] = [] + for item_id in self._queries.session_queue.in_progress(scope): + # Count only the items THIS call actually moved to 'canceled'. An item that finished + # meanwhile — including one canceled by a concurrent bulk request that selected the same + # row — is a no-op transition and must not be counted again. The transition raises if the + # row vanished entirely (a concurrent clear/delete); such an item needs no cancellation, so + # skip it rather than failing the whole bulk operation. + try: + _, transitioned = self._transition_queue_item_status(item_id, "canceled") + if transitioned: + canceled.append(item_id) + except SessionQueueItemNotFoundError: + continue + return canceled + + def _emit_queue_items_canceled(self, queue_id: str, item_ids_by_user: dict[str, list[int]]) -> None: + """Emits queue_items_canceled for a bulk cancel/delete, unless nothing was affected — + an empty result must not broadcast a pointless refetch signal to every client. + + Bulk cancel/delete operations change many rows in one statement and therefore emit no + per-item queue_item_status_changed events; this lets every affected owner (and everyone's + badge counts) refresh.""" + if item_ids_by_user: + self.__invoker.services.events.emit_queue_items_canceled(queue_id, item_ids_by_user) + + def clear(self, queue_id: str, user_id: Optional[str] = None) -> ClearResult: + # Cancel every in-progress item in scope BEFORE deleting rows, so each running worker is + # signaled to stop via its item's own status-changed event. With multiple workers (multi-GPU) + # more than one item can be in_progress at once, and a user-scoped clear must cancel all of + # that user's running items — and ONLY that user's: other users' rows are out of scope and + # their workers must keep running. See delete_by_destination for the same pattern. + scope = Scope(queue_id, user_id=user_id) + self._cancel_in_progress(scope) + deleted = self._queries.session_queue.delete_scope(scope) + self.__invoker.services.events.emit_queue_cleared(queue_id, user_id) + return ClearResult(deleted=sum(len(item_ids) for item_ids in deleted.values())) + + def delete_queue_items_by_id(self, item_ids: list[int]) -> None: + self._queries.session_queue.delete_items(item_ids) + + def prune(self, queue_id: str, user_id: Optional[str] = None) -> PruneResult: + return PruneResult(deleted=self._queries.session_queue.prune(queue_id, user_id=user_id)) + + def cancel_queue_item(self, item_id: int) -> SessionQueueItem: + chain_item_ids = self._get_workflow_call_chain_item_ids(item_id) + canceled_item: SessionQueueItem | None = None + for chain_item_id in chain_item_ids: + queue_item = self._set_queue_item_status(item_id=chain_item_id, status="canceled") + if chain_item_id == item_id: + canceled_item = queue_item + assert canceled_item is not None + return canceled_item + + def delete_queue_item(self, item_id: int) -> None: + """Deletes a session queue item""" + chain_item_ids = self._get_workflow_call_chain_item_ids(item_id) + if any(self.get_queue_item(chain_item_id).status not in _FINISHED for chain_item_id in chain_item_ids): + self.cancel_queue_item(item_id) + self.delete_queue_items_by_id(chain_item_ids) + + def complete_queue_item(self, item_id: int, queue_item: Optional[SessionQueueItem] = None) -> SessionQueueItem: + return self._set_queue_item_status(item_id=item_id, status="completed", queue_item=queue_item) + + def suspend_queue_item(self, item_id: int, queue_item: Optional[SessionQueueItem] = None) -> SessionQueueItem: + return self._set_queue_item_status(item_id=item_id, status="waiting", queue_item=queue_item) + + def resume_queue_item(self, item_id: int, queue_item: Optional[SessionQueueItem] = None) -> SessionQueueItem: + return self._set_queue_item_status(item_id=item_id, status="pending", queue_item=queue_item) + + def fail_queue_item( + self, + item_id: int, + error_type: str, + error_message: str, + error_traceback: str, + ) -> SessionQueueItem: + return self._set_queue_item_status( + item_id=item_id, + status="failed", + error_type=error_type, + error_message=error_message, + error_traceback=error_traceback, + ) + + def cancel_by_batch_ids( + self, queue_id: str, batch_ids: list[str], user_id: Optional[str] = None + ) -> CancelByBatchIDsResult: + scope = Scope(queue_id, user_id=user_id, batch_ids=batch_ids) + canceled_item_ids_by_user = self._queries.session_queue.cancel_waiting_work(scope) + count = sum(len(item_ids) for item_ids in canceled_item_ids_by_user.values()) + # Cancel every in-progress item of the same scope (multi-GPU: possibly several at once). Each + # cancel emits its own per-item queue_item_status_changed, so the bulk event below need not + # include them. + count += len(self._cancel_in_progress(scope)) + self._emit_queue_items_canceled(queue_id, canceled_item_ids_by_user) + return CancelByBatchIDsResult(canceled=count) + + def cancel_by_destination( + self, queue_id: str, destination: str, user_id: Optional[str] = None + ) -> CancelByDestinationResult: + scope = Scope(queue_id, user_id=user_id, destination=destination) + canceled_item_ids_by_user = self._queries.session_queue.cancel_waiting_work(scope) + count = sum(len(item_ids) for item_ids in canceled_item_ids_by_user.values()) + count += len(self._cancel_in_progress(scope)) + self._emit_queue_items_canceled(queue_id, canceled_item_ids_by_user) + return CancelByDestinationResult(canceled=count) + + def delete_by_destination( + self, queue_id: str, destination: str, user_id: Optional[str] = None + ) -> DeleteByDestinationResult: + # Cancel every in-progress item first so each running worker is signaled to stop before we + # delete its row. With multiple workers (multi-GPU) more than one item can be in_progress; + # canceling only one would leave the others running (and then failing to update a deleted row). + scope = Scope(queue_id, user_id=user_id, destination=destination) + canceled_in_progress_ids = set(self._cancel_in_progress(scope)) + deleted_item_ids_by_user = self._queries.session_queue.delete_scope(scope) + count = sum(len(item_ids) for item_ids in deleted_item_ids_by_user.values()) + # The in-progress items canceled above each emitted their own per-item + # queue_item_status_changed, so the bulk event below must not signal them a second time. + # They are left out of the event, but not of the deletion or the returned count. + for owner_user_id, item_ids in list(deleted_item_ids_by_user.items()): + remaining = [item_id for item_id in item_ids if item_id not in canceled_in_progress_ids] + if remaining: + deleted_item_ids_by_user[owner_user_id] = remaining + else: + del deleted_item_ids_by_user[owner_user_id] + self._emit_queue_items_canceled(queue_id, deleted_item_ids_by_user) + return DeleteByDestinationResult(deleted=count) + + def delete_all_except_current(self, queue_id: str, user_id: Optional[str] = None) -> DeleteAllExceptCurrentResult: + # The chains of the items running now are read in the transaction that deletes, so none can + # start running in between and lose its children. + scope = Scope(queue_id, user_id=user_id, except_current=True) + deleted_item_ids_by_user = self._queries.session_queue.delete_scope(scope, waiting_work_only=True) + self._emit_queue_items_canceled(queue_id, deleted_item_ids_by_user) + return DeleteAllExceptCurrentResult( + deleted=sum(len(item_ids) for item_ids in deleted_item_ids_by_user.values()) + ) + + def cancel_by_queue_id( + self, queue_id: str, user_id: Optional[str] = None, origin_prefix: Optional[str] = None + ) -> CancelByQueueIDResult: + scope = Scope(queue_id, user_id=user_id, origin_prefix=origin_prefix) + canceled_item_ids_by_user = self._queries.session_queue.cancel_waiting_work(scope) + count = sum(len(item_ids) for item_ids in canceled_item_ids_by_user.values()) + count += len(self._cancel_in_progress(scope)) + self._emit_queue_items_canceled(queue_id, canceled_item_ids_by_user) + return CancelByQueueIDResult(canceled=count) + + def cancel_all_except_current( + self, queue_id: str, user_id: Optional[str] = None, origin_prefix: Optional[str] = None + ) -> CancelAllExceptCurrentResult: + scope = Scope(queue_id, user_id=user_id, origin_prefix=origin_prefix, except_current=True) + canceled_item_ids_by_user = self._queries.session_queue.cancel_waiting_work(scope) + self._emit_queue_items_canceled(queue_id, canceled_item_ids_by_user) + return CancelAllExceptCurrentResult( + canceled=sum(len(item_ids) for item_ids in canceled_item_ids_by_user.values()) + ) + + def _get_queue_item_with_load_status(self, item_id: int) -> tuple[SessionQueueItem, bool]: + raw_queue_item = self._queries.session_queue.item(item_id) + if raw_queue_item is None: + raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}") + return self._hydrate_queue_item(raw_queue_item, quarantine=False) + + def get_queue_item(self, item_id: int) -> SessionQueueItem: + return self._get_queue_item_with_load_status(item_id)[0] + + def get_queue_item_for_api(self, item_id: int) -> SessionQueueItem: + return self._get_queue_item_for_api(item_id) + + def get_queue_item_workflow_json(self, item_id: int) -> str | None: + return self._queries.session_queue.workflow(item_id) + + def _save_session(self, item_id: int, session: GraphExecutionState, *, only_unfinished: bool) -> bool: + session_json = _session_json(session) + + def save(q: Queries) -> Optional[bool]: + # A session names the media its item uses, which this makes active. + q.locks.acquire(*_PROTECTING_MEDIA, shared=True) + return q.session_queue.set_session(item_id, session_json, only_unfinished=only_unfinished) + + saved = self._queries.run(save) + if saved is None: + raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}") + return saved + + def save_queue_item_session(self, item_id: int, session: GraphExecutionState) -> None: + self._save_session(item_id, session, only_unfinished=False) + + def _save_queue_item_session_if_active(self, item_id: int, session: GraphExecutionState) -> bool: + """Persist a session only while its queue item is non-terminal; whether it was persisted.""" + return self._save_session(item_id, session, only_unfinished=True) + + def set_queue_item_session(self, item_id: int, session: GraphExecutionState) -> SessionQueueItem: + self.save_queue_item_session(item_id, session) + return self.get_queue_item(item_id) + + def record_workflow_call_child_completion( + self, parent_item_id: int, child_item_id: int, output_values: dict[str, Any] + ) -> WorkflowCallChildCompletion | None: + def record(q: Queries) -> WorkflowCallChildCompletion | None: + # The parent's row stays locked from the read to the write, so two children completing at + # once each record into the session the other wrote. + q.locks.acquire(*_PROTECTING_MEDIA, shared=True) + row = q.session_queue.lock_item(parent_item_id) + if row is None: + raise SessionQueueItemNotFoundError(f"No queue item with id {parent_item_id}") + parent_queue_item, readable = self._hydrate_queue_item(row, quarantine=False) + if not readable: + raise ValueError("Unable to record workflow call child completion for an unreadable parent session.") + if parent_queue_item.status in _FINISHED: + return None + + execution = parent_queue_item.session.waiting_workflow_call_execution + if execution is not None and child_item_id in execution.completed_child_item_ids: + return None + generic_update = parent_queue_item.session.record_generic_child_completion(child_item_id, output_values) + if generic_update is not None and not generic_update.changed: + return None + legacy_should_resume, legacy_values = ( + parent_queue_item.session.record_waiting_workflow_call_child_completion(child_item_id, output_values) + ) + if generic_update is None: + should_resume_parent, aggregated_values = legacy_should_resume, legacy_values + else: + should_resume_parent = generic_update.status == "completed" + aggregated_values = { + key: values[0] if len(values) == 1 else values + for key, values in generic_update.aggregated_outputs.items() + } + if generic_update.status == "completed" and aggregated_values != legacy_values: + raise ValueError("Generic child aggregation disagrees with workflow-call aggregation.") + + q.session_queue.set_session(parent_item_id, _session_json(parent_queue_item.session), only_unfinished=False) + return WorkflowCallChildCompletion( + parent_queue_item=parent_queue_item, + should_resume=should_resume_parent, + aggregated_values=aggregated_values, + ) + + return self._queries.run(record) + + def _child_values( + self, + parent_queue_item: SessionQueueItem, + child_session: GraphExecutionState, + field_values: list[NodeFieldValue] | None, + ) -> dict[str, Any]: + workflow_call_execution = parent_queue_item.session.waiting_workflow_call_execution + assert workflow_call_execution is not None + return { + "queue_id": parent_queue_item.queue_id, + "session": _session_json(child_session), + "session_id": child_session.id, + "batch_id": parent_queue_item.batch_id, + "field_values": json.dumps(field_values, default=to_jsonable_python) if field_values is not None else None, + "priority": parent_queue_item.priority, + "workflow": None, + "origin": parent_queue_item.origin, + "destination": parent_queue_item.destination, + "retried_from_item_id": None, + "user_id": parent_queue_item.user_id, + "project_id": parent_queue_item.project_id, + "workflow_call_id": workflow_call_execution.id, + "parent_item_id": parent_queue_item.item_id, + "parent_session_id": parent_queue_item.session_id, + "root_item_id": parent_queue_item.root_item_id or parent_queue_item.item_id, + "workflow_call_depth": workflow_call_execution.depth, + "status": "pending", + } + + def enqueue_workflow_call_children( + self, + parent_queue_item: SessionQueueItem, + child_sessions: list[tuple[GraphExecutionState, list[NodeFieldValue] | None]], + ) -> list[SessionQueueItem]: + if parent_queue_item.session.waiting_workflow_call_execution is None: + raise ValueError("Parent queue item is missing active workflow call execution metadata.") + if not child_sessions: + raise ValueError("Workflow call must enqueue at least one child execution.") + child_values = [ + self._child_values(parent_queue_item, child_session, field_values) + for child_session, field_values in child_sessions + ] + max_queue_size = self.__invoker.services.configuration.max_queue_size + + def enqueue(q: Queries) -> list[int]: + q.locks.acquire(*_ADMISSION, also_shared=_PROTECTING_MEDIA) + parent_status = q.session_queue.lock_status(parent_queue_item.item_id) + if parent_status is None: + raise SessionQueueItemNotFoundError(f"No queue item with id {parent_queue_item.item_id}") + if parent_status in _FINISHED: + raise ValueError("Cannot enqueue workflow call children for a terminal parent queue item.") + if q.session_queue.pending_count(parent_queue_item.queue_id) + len(child_values) > max_queue_size: + raise TooManySessionsError( + "call_saved_workflow exceeds remaining queue capacity for child workflow executions" + ) + child_item_ids = [q.session_queue.insert_item(values) for values in child_values] + + # A copy: a unit run again after a lost race inserts children with other ids, which the parent's session + # must not already name. + waiting_session = parent_queue_item.session.model_copy(deep=True) + waiting_session.set_waiting_workflow_call_child_item_ids(child_item_ids) + waiting_session_json = _session_json(waiting_session) + if not q.session_queue.wait_for_children( + parent_queue_item.item_id, + session=waiting_session_json, + expected_session=getattr(parent_queue_item, "_session_json", None) or waiting_session_json, + ): + raise SessionQueueItemChangedError("Parent queue item changed while enqueuing workflow call children") + return child_item_ids + + child_item_ids = self._queries.run(enqueue) + parent_queue_item.session.set_waiting_workflow_call_child_item_ids(child_item_ids) + parent_queue_item.status = "waiting" + child_queue_items = [self.get_queue_item(item_id) for item_id in child_item_ids] + for queue_item in [self.get_queue_item(parent_queue_item.item_id), *child_queue_items]: + self._emit_status_changed(queue_item) + return child_queue_items + + def enqueue_workflow_call_child( + self, + parent_queue_item: SessionQueueItem, + child_session: GraphExecutionState, + field_values: list[NodeFieldValue] | None = None, + ) -> SessionQueueItem: + if parent_queue_item.session.waiting_workflow_call_execution is None: + raise ValueError("Parent queue item is missing active workflow call execution metadata.") + values = self._child_values(parent_queue_item, child_session, field_values) + max_queue_size = self.__invoker.services.configuration.max_queue_size + + def enqueue(q: Queries) -> int: + q.locks.acquire(*_ADMISSION, also_shared=_PROTECTING_MEDIA) + parent_status = q.session_queue.lock_status(parent_queue_item.item_id) + if parent_status is None: + raise SessionQueueItemNotFoundError(f"No queue item with id {parent_queue_item.item_id}") + if parent_status in _FINISHED: + raise ValueError("Cannot enqueue workflow call child for a terminal parent queue item.") + if q.session_queue.pending_count(parent_queue_item.queue_id) >= max_queue_size: + raise TooManySessionsError( + "call_saved_workflow exceeds remaining queue capacity for child workflow executions" + ) + return q.session_queue.insert_item(values) + + queue_item = self.get_queue_item(self._queries.run(enqueue)) + self._emit_status_changed(queue_item) + return queue_item + + def cancel_workflow_call_children( + self, workflow_call_id: str, exclude_item_ids: set[int] | None = None + ) -> list[int]: + exclude_item_ids = exclude_item_ids or set() + item_ids_with_descendants: list[int] = [] + for item_id in self._queries.session_queue.workflow_call_item_ids(workflow_call_id): + item_ids_with_descendants.append(item_id) + item_ids_with_descendants.extend(self._get_workflow_call_descendant_ids(item_id)) + canceled_item_ids: list[int] = [] + for item_id in dict.fromkeys(item_ids_with_descendants): + if item_id in exclude_item_ids: + continue + if self.get_queue_item(item_id).status in _FINISHED: + continue + self._set_queue_item_status(item_id=item_id, status="canceled") + canceled_item_ids.append(item_id) + return canceled_item_ids + + def list_queue_items( + self, + queue_id: str, + limit: int, + priority: int, + cursor: Optional[int] = None, + status: Optional[QUEUE_ITEM_STATUS] = None, + destination: Optional[str] = None, + ) -> CursorPaginatedResults[SessionQueueItem]: + raw_queue_items = self._queries.session_queue.page( + queue_id, + limit=limit + 1, + after=(priority, cursor) if cursor is not None else None, + status=status, + destination=destination, + ) + items = [self._hydrate_queue_item(raw_queue_item, quarantine=False)[0] for raw_queue_item in raw_queue_items] + has_more = len(items) > limit + return CursorPaginatedResults(items=items[:limit], limit=limit, has_more=has_more) + + def list_all_queue_items( + self, + queue_id: str, + destination: Optional[str] = None, + ) -> list[SessionQueueItem]: + """Gets all queue items with fully rehydrated runtime sessions.""" + raw_queue_items = self._queries.session_queue.all_items(queue_id, destination) + return [self._hydrate_queue_item(raw_queue_item, quarantine=False)[0] for raw_queue_item in raw_queue_items] + + def list_all_queue_items_for_api( + self, + queue_id: str, + destination: Optional[str] = None, + ) -> list[SessionQueueItem]: + """Gets response-shaped queue items without rebuilding runtime execution state.""" + raw_queue_items = self._queries.session_queue.all_items(queue_id, destination) + return [self._project_queue_item_for_read(raw_queue_item) for raw_queue_item in raw_queue_items] + + def get_queue_item_ids( + self, + queue_id: str, + order_dir: SQLiteDirection = SQLiteDirection.Descending, + user_id: Optional[str] = None, + origin_prefix: Optional[str] = None, + limit: Optional[int] = None, + ) -> ItemIdsResult: + item_ids = self._queries.session_queue.item_ids( + queue_id, + descending=order_dir == SQLiteDirection.Descending, + user_id=user_id, + origin_prefix=origin_prefix, + limit=limit, + ) + return ItemIdsResult(item_ids=item_ids, total_count=len(item_ids)) + + def has_active_queue_work(self) -> bool: + return self._queries.session_queue.has_active_work() + + def get_queue_item_summaries_by_ids(self, queue_id: str, item_ids: list[int]) -> list[SessionQueueItemSummary]: + if not item_ids: + return [] + summaries_by_id = { + raw_summary["item_id"]: SessionQueueItemSummary.queue_item_summary_from_dict(raw_summary) + for raw_summary in self._queries.session_queue.summaries(queue_id, item_ids) + } + return [summaries_by_id[item_id] for item_id in item_ids if item_id in summaries_by_id] + + def get_queue_status( + self, + queue_id: str, + user_id: Optional[str] = None, + acting_user_id: Optional[str] = None, + origin_prefix: Optional[str] = None, + is_admin: bool = False, + ) -> SessionQueueStatus: + # Aggregate counts are global across all users within the requested scope. With user_id, the + # same snapshot also yields that user's own counts, for the per-user portion of the badge; + # they are returned in separate fields and never replace the global counts. Only the current + # item's identifiers are read, not a full SessionQueueItem: this runs on every status poll and + # every status change, and hydrating the item would deserialize its whole session graph. + counts, own_counts, current_item = self._queries.session_queue.status_counts(queue_id, user_id, origin_prefix) + + user_pending: Optional[int] = None + user_in_progress: Optional[int] = None + if user_id is not None: + user_pending = own_counts.get("pending", 0) + user_in_progress = own_counts.get("in_progress", 0) + + # Redaction is decided from the same current_item snapshot used to embed identifiers, + # so a concurrent transition (e.g. B finishing while A's status changes) cannot leave + # stale identifiers in the result. The aggregate counts stay global; only the current + # item's identifiers are gated. acting_user_id (event path) takes precedence over + # user_id (API path) when deciding the redaction owner; either being None means a + # global caller who may see the current item. is_admin disables redaction outright so + # admin callers can pass their user_id (for the per-user counts) without losing + # visibility of other users' current item. + owner_user_id = user_id if acting_user_id is None else acting_user_id + current_item_id = None + current_session_id = None + current_batch_id = None + if current_item is not None and (is_admin or owner_user_id is None or current_item.user_id == owner_user_id): + current_item_id = current_item.item_id + current_session_id = current_item.session_id + current_batch_id = current_item.batch_id + + return SessionQueueStatus( + queue_id=queue_id, + item_id=current_item_id, + session_id=current_session_id, + batch_id=current_batch_id, + pending=counts.get("pending", 0), + in_progress=counts.get("in_progress", 0), + waiting=counts.get("waiting", 0), + completed=counts.get("completed", 0), + failed=counts.get("failed", 0), + canceled=counts.get("canceled", 0), + total=sum(counts.values()), + user_pending=user_pending, + user_in_progress=user_in_progress, + ) + + def get_batch_status(self, queue_id: str, batch_id: str, user_id: Optional[str] = None) -> BatchStatus: + counts, origin, destination = self._queries.session_queue.batch_counts(queue_id, batch_id, user_id) + return BatchStatus( + batch_id=batch_id, + origin=origin, + destination=destination, + queue_id=queue_id, + pending=counts.get("pending", 0), + in_progress=counts.get("in_progress", 0), + waiting=counts.get("waiting", 0), + completed=counts.get("completed", 0), + failed=counts.get("failed", 0), + canceled=counts.get("canceled", 0), + total=sum(counts.values()), + ) + + def get_counts_by_destination( + self, queue_id: str, destination: str, user_id: Optional[str] = None + ) -> SessionQueueCountsByDestination: + counts = self._queries.session_queue.destination_counts(queue_id, destination, user_id) + return SessionQueueCountsByDestination( + queue_id=queue_id, + destination=destination, + pending=counts.get("pending", 0), + in_progress=counts.get("in_progress", 0), + waiting=counts.get("waiting", 0), + completed=counts.get("completed", 0), + failed=counts.get("failed", 0), + canceled=counts.get("canceled", 0), + total=sum(counts.values()), + ) + + def retry_items_by_id(self, queue_id: str, item_ids: list[int]) -> RetryItemsResult: + """Retries the given queue items""" + max_queue_size = self.__invoker.services.configuration.max_queue_size + + def survey(q: Queries) -> list[tuple[int, str, ValueToInsertTuple]]: + """The failed or canceled roots to clone, each once, with its owner and the clone's values. Reading and + validating the graphs is the costly part, done without holding the queue's admission.""" + clones: list[tuple[int, str, ValueToInsertTuple]] = [] + seen_root_item_ids: set[int] = set() + hydrated_queue_items: dict[int, SessionQueueItem] = {} + + def read(item_id: int) -> Optional[SessionQueueItem]: + """The item with its runtime state, which the retry validates its graph with.""" + if item_id not in hydrated_queue_items: + raw_queue_item = q.session_queue.item(item_id) + if raw_queue_item is None: + return None + hydrated_queue_items[item_id] = self._hydrate_queue_item(raw_queue_item, quarantine=False)[0] + return hydrated_queue_items[item_id] + + for item_id in item_ids: + queue_item = read(item_id) + if queue_item is None or queue_item.queue_id != queue_id: + continue + if queue_item.status not in ("failed", "canceled"): + continue + + root_item_id = queue_item.root_item_id or queue_item.item_id + if root_item_id in seen_root_item_ids: + continue + seen_root_item_ids.add(root_item_id) + + root_queue_item = read(root_item_id) + if root_queue_item is None: + raise SessionQueueItemNotFoundError(f"No queue item with id {root_item_id}") + if not root_queue_item._snapshot_readable: + continue + if root_queue_item.status not in ("failed", "canceled"): + continue + + field_values_json = ( + json.dumps(root_queue_item.field_values, default=to_jsonable_python) + if root_queue_item.field_values + else None + ) + workflow_json = ( + json.dumps(root_queue_item.workflow, default=to_jsonable_python) + if root_queue_item.workflow + else None + ) + # Validate the graph before dumping a fresh empty execution state. The full read above already + # rehydrated runtime state for contract-compatible validation and recovery semantics. + root_graph = Graph.model_validate( + root_queue_item.session.graph.model_dump(mode="python", warnings=False), strict=False + ) + cloned_session = GraphExecutionState(graph=root_graph) + retried_from_item_id = ( + root_queue_item.retried_from_item_id + if root_queue_item.retried_from_item_id is not None + else root_queue_item.item_id + ) + clones.append( + ( + root_item_id, + root_queue_item.user_id, + ( + root_queue_item.queue_id, + _session_json(cloned_session), + cloned_session.id, + root_queue_item.batch_id, + field_values_json, + root_queue_item.priority, + workflow_json, + root_queue_item.origin, + root_queue_item.destination, + retried_from_item_id, + root_queue_item.user_id, + root_queue_item.project_id, + ), + ) + ) + return clones + + clones = self._queries.run(survey, read_only=True) + + def admit(q: Queries) -> Optional[list[tuple[int, str, ValueToInsertTuple]]]: + """The clones the queue had room for, inserted; None when it is full.""" + q.locks.acquire(*_ADMISSION, also_shared=_PROTECTING_MEDIA) + max_new_queue_items = max_queue_size - q.session_queue.pending_count(queue_id) + if max_new_queue_items <= 0: + return None + admitted = clones[:max_new_queue_items] + q.session_queue.insert_items([values for _, _, values in admitted]) + return admitted + + admitted = self._queries.run(admit) + if admitted is None: + return RetryItemsResult(queue_id=queue_id, retried_item_ids=[]) + retried_item_ids_by_user: dict[str, list[int]] = {} + for root_item_id, user_id, _ in admitted: + retried_item_ids_by_user.setdefault(user_id, []).append(root_item_id) + retry_result = RetryItemsResult( + queue_id=queue_id, retried_item_ids=[root_item_id for root_item_id, _, _ in admitted] + ) + self.__invoker.services.events.emit_queue_items_retried( + retry_result, + user_ids=list(dict.fromkeys(user_id for _, user_id, _ in admitted)), + retried_item_ids_by_user=retried_item_ids_by_user, + ) + return retry_result diff --git a/invokeai/app/services/session_queue/session_queue_sqlite.py b/invokeai/app/services/session_queue/session_queue_sqlite.py deleted file mode 100644 index 68c858798fd..00000000000 --- a/invokeai/app/services/session_queue/session_queue_sqlite.py +++ /dev/null @@ -1,2389 +0,0 @@ -import asyncio -import hashlib -import json -import sqlite3 -import threading -from collections.abc import Sequence -from typing import Any, Literal, Optional, Union, cast - -from pydantic_core import to_jsonable_python - -from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue.session_queue_base import SessionQueueBase, WorkflowCallChildCompletion -from invokeai.app.services.session_queue.session_queue_common import ( - DEFAULT_QUEUE_ID, - QUEUE_ITEM_STATUS, - Batch, - BatchStatus, - CancelAllExceptCurrentResult, - CancelByBatchIDsResult, - CancelByDestinationResult, - CancelByQueueIDResult, - ClearResult, - DeleteAllExceptCurrentResult, - DeleteByDestinationResult, - EnqueueBatchReceipt, - EnqueueBatchResult, - EnqueueIdempotencyConflictError, - EnqueueProjectNotFoundError, - EnqueueReceiptLimitError, - IsEmptyResult, - IsFullResult, - ItemIdsResult, - NodeFieldValue, - PruneResult, - RetryItemsResult, - SessionQueueCountsByDestination, - SessionQueueItem, - SessionQueueItemChangedError, - SessionQueueItemNotFoundError, - SessionQueueItemSummary, - SessionQueueStatus, - TooManySessionsError, - ValueToInsertTuple, - calc_session_count, - prepare_values_to_insert, - uuid_string, -) -from invokeai.app.services.shared.execution_state_migration import ( - dump_execution_state, -) -from invokeai.app.services.shared.graph import Graph, GraphExecutionState -from invokeai.app.services.shared.pagination import CursorPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase - -# Maximum number of ids bound into a single `IN (...)` clause. SQLite's compile-time bind limit is -# 999 on builds older than 3.32 and 32766 on newer ones; staying under the lower figure (leaving -# room for the other bind params in the statement) keeps the queries portable across both. -SQLITE_MAX_BIND_PARAMS_PER_CHUNK = 900 -MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPTS_PER_OWNER = 10_000 -MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPT_BYTES_PER_OWNER = 64 * 1024 * 1024 -MAX_ENQUEUE_RECEIPTS_PER_OWNER = 100_000 -MAX_ENQUEUE_RECEIPT_BYTES_PER_OWNER = 128 * 1024 * 1024 -ACKNOWLEDGED_ENQUEUE_RECEIPT_RETENTION_DAYS = 7 - -# A completed child is still recovery state for its active root, including the interval -# before its returned media is persisted in its parent. Both history-pruning paths must -# retain it; intermediates cleanup also uses these rows to protect the root's media. -PRUNABLE_QUEUE_ITEMS_SQL = """ - status IN ('completed', 'failed', 'canceled') - AND (root_item_id IS NULL OR root_item_id NOT IN ( - SELECT item_id FROM session_queue WHERE status IN ('pending', 'in_progress', 'waiting') - )) -""" - -# Round-robin dequeue (multiuser fairness): pick the next pending item from the user who was -# least-recently served. -# -# The "whose turn is it" ordering key is each candidate user's most recent started_at. We compute -# it with a *correlated* MAX subquery rather than a GROUP BY over all started rows: the candidate -# set is one row per user with pending work, and each MAX(started_at) WHERE user_id = ? is -# satisfied by an indexed seek (idx_session_queue_user_started_at) instead of scanning the full -# retained queue history. This keeps dequeue cost proportional to the number of active users, not -# to total history (which is unbounded by default). MAX() ignores NULL started_at values, so users -# with only pending items fall back to the epoch via COALESCE and are served first. -# -# Kept as a module constant so the scaling test can EXPLAIN QUERY PLAN the exact production SQL. -ROUND_ROBIN_DEQUEUE_QUERY = """--sql - WITH user_next_item AS ( - -- For each user, select their single best pending item (highest priority, then oldest). - SELECT - user_id, - item_id, - ROW_NUMBER() OVER ( - PARTITION BY user_id - ORDER BY priority DESC, item_id ASC - ) AS rn - FROM session_queue - WHERE status = 'pending' - ) - SELECT - sq.*, - u.display_name AS user_display_name, - u.email AS user_email - FROM session_queue sq - LEFT JOIN users u ON sq.user_id = u.user_id - JOIN user_next_item uni ON sq.item_id = uni.item_id AND uni.rn = 1 - ORDER BY - COALESCE( - ( - SELECT MAX(served.started_at) - FROM session_queue served - WHERE served.user_id = sq.user_id - ), - '1970-01-01' - ) ASC, - sq.item_id ASC - LIMIT 1 - """ - -# Upper bound on how many resident model keys the device-affinity scoring query binds. Real -# per-device caches hold at most a handful of models; this only guards the SQL parameter count. -MAX_AFFINITY_MODEL_KEYS = 50 - -# How far past the fairness-chosen candidate (in item_id distance) the device-affinity swap may -# look for a warm-model item. This bounds two things at once: the scoring query's cost (at most -# this many session blobs are scanned per dequeue) and how long a cold item can be deferred — the -# starved candidate's item_id never changes, so newly enqueued warm items eventually fall outside -# the window and the cold item runs after at most ~this many swaps. -AFFINITY_MAX_LOOKAHEAD = 32 - -# FIFO dequeue (single-user mode, or round_robin explicitly disabled): strict priority then -# insertion order. -FIFO_DEQUEUE_QUERY = """--sql - SELECT - sq.*, - u.display_name as user_display_name, - u.email as user_email - FROM session_queue sq - LEFT JOIN users u ON sq.user_id = u.user_id - WHERE sq.status = 'pending' - ORDER BY - sq.priority DESC, - sq.item_id ASC - LIMIT 1 - """ - -# The database keeps no table statistics, so the planner takes `queue_id = ?` as selective. A -# statement that filters `queue_id = ?` and reads only idx_session_queue_listing's columns, or one -# that looks rows up by `item_id IN (...)` beside `queue_id = ?`, is moved onto that index and scans -# the queue's whole retained history. Likewise `user_id = ?` can win over a selective column such as -# batch_id and read the user's whole history. Guard such a term with a unary + (which keeps it off -# every index) or give the statement a selective indexed predicate the planner prefers. -# -# This capacity check runs on every enqueue and must cost the pending backlog, not the history: with -# the guard, the status index finds the pending items. -PENDING_QUEUE_ITEM_COUNT_QUERY = "SELECT count(*) FROM session_queue WHERE +queue_id = ? AND status = 'pending';" - - -class SqliteSessionQueue(SessionQueueBase): - __invoker: Invoker - - # Serializes the select-candidate-then-claim sequence in `dequeue()`. The DB connection's - # RLock serializes individual statements, but the gap between selecting the next pending item - # and marking it 'in_progress' is a race: with multiple session-processor workers (multi-GPU), - # two workers could select the same item. Holding this lock across the whole claim prevents it. - _dequeue_lock = threading.Lock() - - def start(self, invoker: Invoker) -> None: - self.__invoker = invoker - self._set_in_progress_to_canceled() - config = self.__invoker.services.configuration - if config.clear_queue_on_startup: - clear_result = self.clear(DEFAULT_QUEUE_ID) - if clear_result.deleted > 0: - self.__invoker.services.logger.info(f"Cleared all {clear_result.deleted} queue items") - return - - if config.max_queue_history is not None: - deleted = self._prune_terminal_to_limit(DEFAULT_QUEUE_ID, config.max_queue_history) - if deleted > 0: - self.__invoker.services.logger.info( - f"Pruned {deleted} completed/failed/canceled queue items (kept up to {config.max_queue_history})" - ) - - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - - def _set_in_progress_to_canceled(self) -> None: - """ - Sets all in_progress or waiting queue items to canceled. Run on app startup, not associated with any queue. - This is necessary because the invoker may have been killed while processing a queue item or while a parent - queue item was suspended waiting on a child workflow execution. - """ - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT item_id - FROM session_queue - WHERE status = 'in_progress' - OR status = 'waiting'; - """ - ) - interrupted_item_ids = [row[0] for row in cast(list[sqlite3.Row], cursor.fetchall())] - item_ids_to_cancel: set[int] = set() - for item_id in interrupted_item_ids: - item_ids_to_cancel.update(self._get_workflow_call_chain_item_ids(item_id)) - if not item_ids_to_cancel: - return - with self._db.transaction() as cursor: - placeholders = ",".join("?" for _ in item_ids_to_cancel) - cursor.execute( - f"""--sql - UPDATE session_queue - SET status = 'canceled', - status_sequence = COALESCE(status_sequence, 0) + 1 - WHERE item_id IN ({placeholders}) - AND ( - status = 'pending' - OR status = 'in_progress' - OR status = 'waiting' - ); - """, - tuple(item_ids_to_cancel), - ) - - def _prune_terminal_to_limit(self, queue_id: str, keep: int) -> int: - """Prune terminal items (completed/failed/canceled) to keep at most N most-recent items.""" - with self._db.transaction() as cursor: - where = f"""--sql - WHERE - queue_id = ? - AND {PRUNABLE_QUEUE_ITEMS_SQL} - """ - cursor.execute( - f"""--sql - DELETE - FROM session_queue - {where} - AND item_id NOT IN ( - SELECT item_id - FROM session_queue - {where} - ORDER BY COALESCE(completed_at, updated_at, created_at) DESC, item_id DESC - LIMIT ? - ); - """, - (queue_id, queue_id, keep), - ) - count = cursor.rowcount - return count - - def _get_current_queue_size(self, queue_id: str) -> int: - """Gets the current number of pending queue items""" - with self._db.transaction() as cursor: - cursor.execute(PENDING_QUEUE_ITEM_COUNT_QUERY, (queue_id,)) - count = cast(int, cursor.fetchone()[0]) - return count - - def _get_highest_priority(self, queue_id: str) -> int: - """Gets the highest priority value in the queue""" - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT MAX(priority) - FROM session_queue - WHERE - queue_id = ? - AND status = 'pending' - """, - (queue_id,), - ) - priority = cast(Union[int, None], cursor.fetchone()[0]) or 0 - return priority - - async def enqueue_batch( - self, queue_id: str, batch: Batch, prepend: bool, user_id: str = "system" - ) -> EnqueueBatchResult: - # The route awaits this method, but every operation below is synchronous SQLite/CPU - # work. Keep the complete transaction in one worker operation; otherwise the queue-size - # check before the first await and the insert/event work after later awaits execute on the - # event loop. - return await asyncio.to_thread(self._enqueue_batch, queue_id, batch, prepend, user_id) - - def acknowledge_enqueue(self, queue_id: str, idempotency_key: str, user_id: str = "system") -> None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - UPDATE session_queue_enqueue_receipts - SET acknowledged_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW') - WHERE queue_id = ? AND user_id = ? AND idempotency_key = ? - """, - (queue_id, user_id, idempotency_key), - ) - - def get_enqueue_receipt( - self, queue_id: str, idempotency_key: str, user_id: str = "system" - ) -> EnqueueBatchReceipt | None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT batch_id, requested, enqueued, item_ids - FROM session_queue_enqueue_receipts - WHERE queue_id = ? AND user_id = ? AND idempotency_key = ? - """, - (queue_id, user_id, idempotency_key), - ) - receipt = cursor.fetchone() - return self._decode_enqueue_receipt(receipt) if receipt is not None else None - - @staticmethod - def _decode_enqueue_receipt(receipt: sqlite3.Row) -> EnqueueBatchReceipt: - item_ids = json.loads(receipt["item_ids"]) - if not isinstance(item_ids, list) or not all(isinstance(item_id, int) for item_id in item_ids): - raise RuntimeError("Stored queue enqueue receipt contains invalid item ids") - return EnqueueBatchReceipt( - batch_id=receipt["batch_id"], - requested=receipt["requested"], - enqueued=receipt["enqueued"], - item_ids=item_ids, - ) - - def _enqueue_batch(self, queue_id: str, batch: Batch, prepend: bool, user_id: str) -> EnqueueBatchResult: - requested_count = calc_session_count(batch=batch) - payload_hash = ( - hashlib.sha256( - json.dumps( - { - "batch": batch.model_dump(mode="json", exclude={"batch_id", "idempotency_key"}), - "prepend": prepend, - }, - ensure_ascii=False, - separators=(",", ":"), - sort_keys=True, - ).encode("utf-8") - ).hexdigest() - if batch.idempotency_key is not None - else None - ) - - def get_receipt(cursor: sqlite3.Cursor) -> sqlite3.Row | None: - if batch.idempotency_key is None: - return None - cursor.execute( - """--sql - SELECT payload_hash, batch_id, requested, enqueued, priority, item_ids - FROM session_queue_enqueue_receipts - WHERE queue_id = ? AND user_id = ? AND idempotency_key = ? - """, - (queue_id, user_id, batch.idempotency_key), - ) - return cast(sqlite3.Row | None, cursor.fetchone()) - - def settle_receipt(receipt: sqlite3.Row) -> EnqueueBatchResult: - if receipt["payload_hash"] != payload_hash: - raise EnqueueIdempotencyConflictError( - f"Idempotency key {batch.idempotency_key} is already used by another submission" - ) - result = self._decode_enqueue_receipt(receipt) - return EnqueueBatchResult( - queue_id=queue_id, - requested=result.requested, - enqueued=result.enqueued, - batch=batch.model_copy(update={"batch_id": result.batch_id}), - priority=receipt["priority"], - item_ids=result.item_ids, - ) - - def require_project(cursor: sqlite3.Cursor) -> None: - if batch.project_id is None: - return - cursor.execute( - "SELECT 1 FROM projects WHERE user_id = ? AND project_id = ?;", - (user_id, batch.project_id), - ) - if cursor.fetchone() is None: - raise EnqueueProjectNotFoundError(batch.project_id) - - with self._db.transaction() as cursor: - if batch.idempotency_key is not None: - receipt = get_receipt(cursor) - if receipt is not None: - return settle_receipt(receipt) - require_project(cursor) - cursor.execute(PENDING_QUEUE_ITEM_COUNT_QUERY, (queue_id,)) - preliminary_queue_size = cast(int, cursor.fetchone()[0]) - - max_queue_size = self.__invoker.services.configuration.max_queue_size - preliminary_capacity = max(0, max_queue_size - preliminary_queue_size) - prepared_values = prepare_values_to_insert( - queue_id=queue_id, - batch=batch, - priority=0, - max_new_queue_items=preliminary_capacity, - user_id=user_id, - ) - - with self._db.transaction() as cursor: - if batch.idempotency_key is not None: - cursor.execute( - """--sql - DELETE FROM session_queue_enqueue_receipts - WHERE user_id = ? - AND acknowledged_at IS NOT NULL - AND acknowledged_at <= STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW', ?) - """, - (user_id, f"-{ACKNOWLEDGED_ENQUEUE_RECEIPT_RETENTION_DAYS} days"), - ) - receipt = get_receipt(cursor) - if receipt is not None: - return settle_receipt(receipt) - - require_project(cursor) - - cursor.execute(PENDING_QUEUE_ITEM_COUNT_QUERY, (queue_id,)) - current_queue_size = cast(int, cursor.fetchone()[0]) - max_new_queue_items = max(0, max_queue_size - current_queue_size) - priority = 0 - if prepend: - cursor.execute( - "SELECT MAX(priority) FROM session_queue WHERE queue_id = ? AND status = 'pending';", - (queue_id,), - ) - priority = (cast(int | None, cursor.fetchone()[0]) or 0) + 1 - values_to_insert = prepared_values[:max_new_queue_items] - if priority != 0: - values_to_insert = [(*value[:5], priority, *value[6:]) for value in values_to_insert] - enqueued_count = len(values_to_insert) - accepted_batch = batch - if enqueued_count > 0: - accepted_batch_id = batch.batch_id - while True: - # The unary + on queue_id and user_id keeps this, and the readback below, on - # idx_session_queue_batch_id; unguarded, the planner picks a user_id index and - # reads the user's whole history (see PENDING_QUEUE_ITEM_COUNT_QUERY). - cursor.execute( - """--sql - SELECT 1 FROM session_queue - WHERE +queue_id = ? AND +user_id = ? AND batch_id = ? - LIMIT 1; - """, - (queue_id, user_id, accepted_batch_id), - ) - if cursor.fetchone() is None: - break - accepted_batch_id = uuid_string() - if accepted_batch_id != batch.batch_id: - accepted_batch = batch.model_copy(update={"batch_id": accepted_batch_id}) - values_to_insert = [(*value[:3], accepted_batch_id, *value[4:]) for value in values_to_insert] - receipt_bytes = 0 - unacknowledged_bytes = 0 - if batch.idempotency_key is not None and enqueued_count > 0: - cursor.execute( - """--sql - SELECT - COUNT(*), - COALESCE(SUM(byte_size), 0), - COALESCE(SUM(CASE WHEN acknowledged_at IS NULL THEN 1 ELSE 0 END), 0), - COALESCE(SUM(CASE WHEN acknowledged_at IS NULL THEN byte_size ELSE 0 END), 0) - FROM session_queue_enqueue_receipts - WHERE user_id = ? - """, - (user_id,), - ) - receipt_count, receipt_bytes, unacknowledged_count, unacknowledged_bytes = cursor.fetchone() - if cast(int, receipt_count) >= MAX_ENQUEUE_RECEIPTS_PER_OWNER: - raise EnqueueReceiptLimitError("Enqueue receipts exceed the count limit") - if cast(int, unacknowledged_count) >= MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPTS_PER_OWNER: - raise EnqueueReceiptLimitError("Too many unacknowledged enqueue requests") - if cast(int, receipt_bytes) >= MAX_ENQUEUE_RECEIPT_BYTES_PER_OWNER: - raise EnqueueReceiptLimitError("Enqueue receipts exceed the storage limit") - if cast(int, unacknowledged_bytes) >= MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPT_BYTES_PER_OWNER: - raise EnqueueReceiptLimitError("Unacknowledged enqueue receipts exceed the storage limit") - cursor.executemany( - """--sql - INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id, user_id, project_id) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - values_to_insert, - ) - cursor.execute( - """--sql - SELECT item_id - FROM session_queue - WHERE +queue_id = ? AND +user_id = ? AND batch_id = ? - ORDER BY item_id ASC; - """, - (queue_id, user_id, accepted_batch.batch_id), - ) - item_ids = [row[0] for row in cursor.fetchall()] - if batch.idempotency_key is not None and enqueued_count > 0: - item_ids_json = json.dumps(item_ids, separators=(",", ":")) - receipt_byte_size = len( - json.dumps( - [ - queue_id, - user_id, - batch.idempotency_key, - payload_hash, - accepted_batch.batch_id, - requested_count, - enqueued_count, - priority, - item_ids, - ], - ensure_ascii=False, - separators=(",", ":"), - ).encode("utf-8") - ) - if ( - cast(int, unacknowledged_bytes) + receipt_byte_size - > MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPT_BYTES_PER_OWNER - ): - raise EnqueueReceiptLimitError("Unacknowledged enqueue receipts exceed the storage limit") - if cast(int, receipt_bytes) + receipt_byte_size > MAX_ENQUEUE_RECEIPT_BYTES_PER_OWNER: - raise EnqueueReceiptLimitError("Enqueue receipts exceed the storage limit") - cursor.execute( - """--sql - INSERT INTO session_queue_enqueue_receipts ( - queue_id, user_id, idempotency_key, payload_hash, batch_id, - requested, enqueued, priority, item_ids, byte_size - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?); - """, - ( - queue_id, - user_id, - batch.idempotency_key, - payload_hash, - accepted_batch.batch_id, - requested_count, - enqueued_count, - priority, - item_ids_json, - receipt_byte_size, - ), - ) - enqueue_result = EnqueueBatchResult( - queue_id=queue_id, - requested=requested_count, - enqueued=enqueued_count, - batch=accepted_batch, - priority=priority, - item_ids=item_ids, - ) - self.__invoker.services.events.emit_batch_enqueued(enqueue_result, user_id=user_id) - return enqueue_result - - def dequeue(self, device: Optional[str] = None) -> Optional[SessionQueueItem]: - config = self.__invoker.services.configuration - use_round_robin = config.multiuser and config.session_queue_mode == "round_robin" - - query = ROUND_ROBIN_DEQUEUE_QUERY if use_round_robin else FIFO_DEQUEUE_QUERY - - # Snapshot the claiming device's warm models BEFORE taking the dequeue lock: the lookup - # touches the ModelCache lock, which other threads may hold across long operations (VRAM - # transfers, cache clears). A slightly stale snapshot is fine for a heuristic; stalling - # every worker's dequeue is not. An explicitly configured session_queue_mode=FIFO is a - # request for strict insertion order, so it opts out of affinity reordering entirely - # (the setting defaults to round_robin, which keeps affinity active for single-user - # installs even though they use the FIFO query). - if config.session_queue_mode == "round_robin": - resident_model_keys = self._get_device_resident_model_keys(device) - else: - resident_model_keys = set() - - # Hold the dequeue lock across the select-then-claim so concurrent workers (multi-GPU) - # cannot select and claim the same pending item. `_set_queue_item_status` already no-ops - # if the item was concurrently moved to a terminal state (e.g. canceled), so we only need - # to guard against two dequeues racing for the same pending row. - with self._dequeue_lock: - while True: - with self._db.transaction() as cursor: - cursor.execute(query) - result = cast(Union[sqlite3.Row, None], cursor.fetchone()) - if result is None: - return None - raw_result = dict(result) - queue_item, readable = self._hydrate_queue_item(raw_result, quarantine=True) - if not readable: - continue - queue_item = self._apply_device_affinity(queue_item, resident_model_keys) - # Record the claiming worker's device so the UI can label the item by GPU. Passing the - # item we already materialized lets _set_queue_item_status patch it in place instead of - # re-reading (and re-parsing the session graph of) the row we just read. - queue_item = self._set_queue_item_status( - item_id=queue_item.item_id, status="in_progress", device=device, queue_item=queue_item - ) - return queue_item - - @staticmethod - def _make_unreadable_queue_item(raw_queue_item: dict[str, Any], error: Exception) -> SessionQueueItem: - """Build a metadata-preserving placeholder for an unreadable runtime snapshot.""" - placeholder = SessionQueueItem.model_construct(**raw_queue_item) - # A placeholder has no trustworthy execution result. Never expose a corrupt or newer snapshot as complete. - placeholder.status = "failed" - placeholder.session = GraphExecutionState(graph=Graph()) - placeholder.workflow = None - placeholder.field_values = None - placeholder._snapshot_readable = False - message = f"Unable to load execution state: {error}" - placeholder.error_type = type(error).__name__ - placeholder.error_message = message - placeholder.error_traceback = message - return placeholder - - def _hydrate_queue_item(self, raw_queue_item: dict[str, Any], *, quarantine: bool) -> tuple[SessionQueueItem, bool]: - """Hydrate one queue row without letting an unreadable snapshot break queue access.""" - try: - return SessionQueueItem.queue_item_from_dict(raw_queue_item), True - except (TypeError, ValueError) as exc: - if quarantine: - return self._quarantine_unreadable_queue_item(raw_queue_item, exc), False - return self._make_unreadable_queue_item(raw_queue_item, exc), False - - def _project_queue_item_for_read(self, raw_queue_item: dict[str, Any]) -> SessionQueueItem: - """Read queue metadata and response results without rebuilding runtime execution state.""" - try: - return SessionQueueItem.queue_item_from_dict(raw_queue_item, hydrate_runtime=False) - except (TypeError, ValueError) as exc: - return self._make_unreadable_queue_item(raw_queue_item, exc) - - def _get_queue_item_for_read( - self, - item_id: int, - *, - cursor: sqlite3.Cursor | None = None, - hydrate_runtime: bool = False, - ) -> SessionQueueItem: - """Read queue metadata and response results without rebuilding execution runtime state.""" - if cursor is None: - with self._db.transaction() as transaction_cursor: - return self._get_queue_item_for_read( - item_id, cursor=transaction_cursor, hydrate_runtime=hydrate_runtime - ) - - cursor.execute( - """--sql - SELECT - sq.*, - u.display_name AS user_display_name, - u.email AS user_email - FROM session_queue sq - LEFT JOIN users u ON sq.user_id = u.user_id - WHERE sq.item_id = ? - """, - (item_id,), - ) - result = cast(Union[sqlite3.Row, None], cursor.fetchone()) - if result is None: - raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}") - raw_queue_item = dict(result) - if hydrate_runtime: - return self._hydrate_queue_item(raw_queue_item, quarantine=False)[0] - return self._project_queue_item_for_read(raw_queue_item) - - def _get_queue_item_for_api(self, item_id: int, *, cursor: sqlite3.Cursor | None = None) -> SessionQueueItem: - """Read one queue row with either full runtime hydration or response projection.""" - return self._get_queue_item_for_read(item_id, cursor=cursor, hydrate_runtime=False) - - def _get_queue_item_for_retry(self, item_id: int, *, cursor: sqlite3.Cursor | None = None) -> SessionQueueItem: - """Read the graph and retry metadata with full runtime hydration.""" - return self._get_queue_item_for_read(item_id, cursor=cursor, hydrate_runtime=True) - - def _quarantine_unreadable_queue_item(self, raw_queue_item: dict[str, Any], error: Exception) -> SessionQueueItem: - """Fail a pending row whose runtime snapshot is newer than this worker can read. - - The real session cannot be hydrated, so use a minimal in-memory placeholder only for the - status transition/event. The persisted session remains untouched for postmortem recovery. - """ - - placeholder = self._make_unreadable_queue_item(raw_queue_item, error) - placeholder.status = "pending" - return self._set_queue_item_status( - item_id=placeholder.item_id, - status="failed", - error_type=placeholder.error_type, - error_message=placeholder.error_message, - error_traceback=placeholder.error_traceback, - queue_item=placeholder, - ) - - def _apply_device_affinity(self, candidate: SessionQueueItem, resident_keys: set[str]) -> SessionQueueItem: - """Swap the fairness-chosen candidate for a nearby same-user, same-priority pending item - whose models are already cached on the claiming device, if one exists. - - Cross-device model reloads are expensive (tens of seconds for large models), so when a user - has queued a mix of models, preferring an item whose models are warm on the freeing GPU cuts - thrash. Fairness is preserved by construction: round-robin decides *which user* is served and - priority ordering decides *which tier* of their items is eligible; this heuristic only - reorders within that user's equal-priority pending items, and only within - AFFINITY_MAX_LOOKAHEAD of the candidate's item_id, so a cold item's deferral is bounded. - The caller passes an empty key set to disable affinity (legacy single-device mode, explicit - FIFO mode, or cache introspection unavailable). - """ - if not resident_keys: - return candidate - # Model keys are UUID strings that appear verbatim in the session JSON, so residency can be - # scored with substring matches — no need to parse each candidate's session. Sort for - # deterministic parameter binding; cap to bound the query if a cache is unexpectedly large. - keys = sorted(resident_keys)[:MAX_AFFINITY_MODEL_KEYS] - score = " + ".join(["(instr(sq.session, ?) > 0)"] * len(keys)) - with self._db.transaction() as cursor: - cursor.execute( - f"""--sql - SELECT - sq.*, - u.display_name AS user_display_name, - u.email AS user_email, - ({score}) AS affinity - FROM session_queue sq - LEFT JOIN users u ON sq.user_id = u.user_id - WHERE sq.status = 'pending' - AND sq.user_id IS ? - AND sq.priority = ? - AND sq.item_id >= ? - AND sq.item_id <= ? - ORDER BY affinity DESC, sq.item_id ASC - LIMIT 1 - """, - ( - *keys, - candidate.user_id, - candidate.priority, - candidate.item_id, - candidate.item_id + AFFINITY_MAX_LOOKAHEAD, - ), - ) - row = cast(Union[sqlite3.Row, None], cursor.fetchone()) - if row is None: - return candidate - row_dict = dict(row) - if not row_dict.pop("affinity", 0) or row_dict["item_id"] == candidate.item_id: - # No warm-model item for this user (or the candidate already is one) — keep the - # fairness-chosen candidate. - return candidate - queue_item, readable = self._hydrate_queue_item(row_dict, quarantine=True) - return queue_item if readable else candidate - - def _get_device_resident_model_keys(self, device: Optional[str]) -> set[str]: - """Best-effort lookup of the model keys currently cached for the given generation device.""" - if device is None: - return set() - try: - cache = self.__invoker.services.model_manager.load.ram_caches.get(device) - if cache is None: - return set() - return set(cache.cached_model_keys()) - except Exception: - # Affinity is purely an optimization — dequeue must never fail because cache - # introspection did (e.g. model manager not fully started, or mocked in tests). - return set() - - def get_next(self, queue_id: str, origin_prefix: Optional[str] = None) -> Optional[SessionQueueItem]: - with self._db.transaction() as cursor: - query = """--sql - SELECT - sq.*, - u.display_name as user_display_name, - u.email as user_email - FROM session_queue sq - LEFT JOIN users u ON sq.user_id = u.user_id - WHERE - sq.queue_id = ? - AND sq.status = 'pending' - """ - params = [queue_id] - if origin_prefix is not None: - query += """--sql - AND sq.origin LIKE ? - """ - params.append(f"{origin_prefix}%") - query += """--sql - ORDER BY - sq.priority DESC, - sq.created_at ASC - LIMIT 1 - """ - cursor.execute(query, params) - result = cast(Union[sqlite3.Row, None], cursor.fetchone()) - if result is None: - return None - return self._hydrate_queue_item(dict(result), quarantine=False)[0] - - def get_current(self, queue_id: str, origin_prefix: Optional[str] = None) -> Optional[SessionQueueItem]: - with self._db.transaction() as cursor: - query = """--sql - SELECT - sq.*, - u.display_name as user_display_name, - u.email as user_email - FROM session_queue sq - LEFT JOIN users u ON sq.user_id = u.user_id - WHERE - sq.queue_id = ? - AND sq.status = 'in_progress' - """ - params = [queue_id] - if origin_prefix is not None: - query += """--sql - AND sq.origin LIKE ? - """ - params.append(f"{origin_prefix}%") - query += """--sql - LIMIT 1 - """ - cursor.execute(query, params) - result = cast(Union[sqlite3.Row, None], cursor.fetchone()) - if result is None: - return None - return self._hydrate_queue_item(dict(result), quarantine=False)[0] - - def _get_queue_item_by_status_for_api( - self, queue_id: str, status: Literal["pending", "in_progress"], origin_prefix: Optional[str] - ) -> Optional[SessionQueueItem]: - query = """--sql - SELECT - sq.*, - u.display_name as user_display_name, - u.email as user_email - FROM session_queue sq - LEFT JOIN users u ON sq.user_id = u.user_id - WHERE - sq.queue_id = ? - AND sq.status = ? - """ - params: list[str] = [queue_id, status] - if origin_prefix is not None: - query += " AND sq.origin LIKE ?" - params.append(f"{origin_prefix}%") - if status == "pending": - query += " ORDER BY sq.priority DESC, sq.created_at ASC" - query += " LIMIT 1" - with self._db.transaction() as cursor: - cursor.execute(query, params) - result = cast(Union[sqlite3.Row, None], cursor.fetchone()) - if result is None: - return None - return self._project_queue_item_for_read(dict(result)) - - def get_current_for_api(self, queue_id: str, origin_prefix: Optional[str] = None) -> Optional[SessionQueueItem]: - return self._get_queue_item_by_status_for_api(queue_id, "in_progress", origin_prefix) - - def get_next_for_api(self, queue_id: str, origin_prefix: Optional[str] = None) -> Optional[SessionQueueItem]: - return self._get_queue_item_by_status_for_api(queue_id, "pending", origin_prefix) - - def _set_queue_item_status( - self, - item_id: int, - status: QUEUE_ITEM_STATUS, - error_type: Optional[str] = None, - error_message: Optional[str] = None, - error_traceback: Optional[str] = None, - device: Optional[str] = None, - queue_item: Optional[SessionQueueItem] = None, - ) -> SessionQueueItem: - return self._transition_queue_item_status( - item_id=item_id, - status=status, - error_type=error_type, - error_message=error_message, - error_traceback=error_traceback, - device=device, - queue_item=queue_item, - )[0] - - def _transition_queue_item_status( - self, - item_id: int, - status: QUEUE_ITEM_STATUS, - error_type: Optional[str] = None, - error_message: Optional[str] = None, - error_traceback: Optional[str] = None, - device: Optional[str] = None, - queue_item: Optional[SessionQueueItem] = None, - ) -> tuple[SessionQueueItem, bool]: - """Move a queue item to `status` unless it is already finished (completed, failed or - canceled), returning the item and whether THIS call performed the transition. - - The terminal guard and the UPDATE run in one transaction, so two callers racing to cancel - the same row cannot both observe themselves as the one that canceled it — exactly one sees - `transitioned=True` (the bulk-cancel counters rely on this). Note the atomicity comes from - SqliteDatabase.transaction() holding its process-wide lock for the whole block, not from - the SQL: the guard is a Python check between the SELECT and the UPDATE, valid only while - every writer goes through that lock. When no transition happens, the item is returned - unchanged and no status-changed event is emitted; a vanished row raises - SessionQueueItemNotFoundError. - """ - if queue_item is not None and queue_item.item_id != item_id: - raise ValueError(f"Queue item {queue_item.item_id} does not match requested item {item_id}") - - updated_status_row: sqlite3.Row | None = None - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT status FROM session_queue WHERE item_id = ? - """, - (item_id,), - ) - row = cursor.fetchone() - if row is None: - raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}") - current_status = row[0] - - # Only update if not already finished (completed, failed or canceled). - if current_status not in ("completed", "failed", "canceled"): - cursor.execute( - """--sql - UPDATE session_queue - SET status = ?, status_sequence = COALESCE(status_sequence, 0) + 1, error_type = ?, error_message = ?, error_traceback = ?, device = COALESCE(?, device) - WHERE item_id = ? - """, - (status, error_type, error_message, error_traceback, device, item_id), - ) - cursor.execute( - """--sql - SELECT - status, - status_sequence, - error_type, - error_message, - error_traceback, - created_at, - updated_at, - started_at, - completed_at, - device - FROM session_queue - WHERE item_id = ? - """, - (item_id,), - ) - updated_status_row = cast(sqlite3.Row, cursor.fetchone()) - - if updated_status_row is None: - # Already finished (return it unchanged) or deleted (get_queue_item raises). - return self.get_queue_item(item_id), False - - if queue_item is None: - queue_item = self.get_queue_item(item_id) - else: - # `device` is in the SELECT above, so a caller that supplied its own queue_item still - # sees the device this claim just recorded rather than a stale value. - for field_name in updated_status_row.keys(): - setattr(queue_item, field_name, updated_status_row[field_name]) - batch_status = self.get_batch_status(queue_id=queue_item.queue_id, batch_id=queue_item.batch_id) - # The QueueItemStatusChangedEvent ships to user:{queue_item.user_id} and admin rooms. - # acting_user_id ensures the embedded current-item identifiers are redacted when the - # in-progress item belongs to someone else, while leaving aggregate counts global. - # Doing this inside get_queue_status guarantees the redaction decision and the - # embedded identifiers come from the same lightweight metadata snapshot, eliminating the - # race where a second read could find None and skip scrubbing stale identifiers. - # user_id additionally embeds the owner's per-user counts so the owner's client can - # apply the event's queue_status optimistically without waiting for a refetch; the - # sanitized companion nulls them before reaching anyone else. - queue_status = self.get_queue_status( - queue_id=queue_item.queue_id, user_id=queue_item.user_id, acting_user_id=queue_item.user_id - ) - - self.__invoker.services.events.emit_queue_item_status_changed(queue_item, batch_status, queue_status) - return queue_item, True - - def _get_workflow_call_child_ids(self, item_id: int) -> list[int]: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT item_id - FROM session_queue - WHERE parent_item_id = ? - ORDER BY item_id ASC - """, - (item_id,), - ) - rows = cast(list[sqlite3.Row], cursor.fetchall()) - return [row[0] for row in rows] - - def _get_workflow_call_descendant_ids(self, item_id: int) -> list[int]: - descendant_ids: list[int] = [] - queue: list[int] = [item_id] - while queue: - current_item_id = queue.pop(0) - child_ids = self._get_workflow_call_child_ids(current_item_id) - descendant_ids.extend(child_ids) - queue.extend(child_ids) - return descendant_ids - - def _get_workflow_call_ancestor_ids(self, item_id: int) -> list[int]: - ancestor_ids: list[int] = [] - current_item_id = item_id - while True: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT parent_item_id - FROM session_queue - WHERE item_id = ? - """, - (current_item_id,), - ) - row = cursor.fetchone() - if row is None: - raise SessionQueueItemNotFoundError(f"No queue item with id {current_item_id}") - parent_item_id = row[0] - if parent_item_id is None: - break - ancestor_ids.append(parent_item_id) - current_item_id = parent_item_id - return ancestor_ids - - def _get_workflow_call_chain_item_ids(self, item_id: int) -> list[int]: - ancestor_ids = self._get_workflow_call_ancestor_ids(item_id) - root_item_id = ancestor_ids[-1] if ancestor_ids else item_id - descendant_ids = self._get_workflow_call_descendant_ids(root_item_id) - chain_item_ids = ancestor_ids + [item_id] + descendant_ids - deduped_chain_item_ids = list(dict.fromkeys(chain_item_ids)) - return deduped_chain_item_ids - - def _get_current_workflow_call_chain_item_ids(self, queue_id: str) -> set[int]: - """Item ids in the workflow-call chains of every currently executing item. - - With multiple GPU workers, several items can be in_progress at once. Every one of them is - "current": the "except current" operations must leave each active worker's complete - ancestor/descendant chain intact, not just the chain of one arbitrarily selected item. - Waiting chains with no in-progress item remain cancellable backlog, exactly as in - single-worker mode. - """ - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT item_id - FROM session_queue - WHERE queue_id = ? AND status = 'in_progress' - ORDER BY item_id ASC - """, - (queue_id,), - ) - in_progress_ids = [cast(int, row[0]) for row in cursor.fetchall()] - - if in_progress_ids: - chain_item_ids: set[int] = set() - for item_id in in_progress_ids: - chain_item_ids.update(self._get_workflow_call_chain_item_ids(item_id)) - return chain_item_ids - - # Nothing is in progress. A chain can still be mid-execution in the gap between one child - # finishing and the next being dequeued (parent 'waiting', next child 'pending'); protect - # the chain of the next item that would be dequeued, as before. - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT child.item_id - FROM session_queue child - JOIN session_queue parent ON parent.item_id = child.parent_item_id - WHERE - child.queue_id = ? - AND child.status = 'pending' - AND parent.status = 'waiting' - ORDER BY child.priority DESC, child.created_at ASC, child.item_id ASC - LIMIT 1 - """, - (queue_id,), - ) - row = cursor.fetchone() - if row is None: - return set() - return set(self._get_workflow_call_chain_item_ids(cast(int, row[0]))) - - def is_empty(self, queue_id: str) -> IsEmptyResult: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT count(*) - FROM session_queue - WHERE queue_id = ? - """, - (queue_id,), - ) - is_empty = cast(int, cursor.fetchone()[0]) == 0 - return IsEmptyResult(is_empty=is_empty) - - def is_full(self, queue_id: str) -> IsFullResult: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT count(*) - FROM session_queue - WHERE queue_id = ? - """, - (queue_id,), - ) - max_queue_size = self.__invoker.services.configuration.max_queue_size - is_full = cast(int, cursor.fetchone()[0]) >= max_queue_size - return IsFullResult(is_full=is_full) - - def clear(self, queue_id: str, user_id: Optional[str] = None) -> ClearResult: - user_filter = "AND user_id = ?" if user_id is not None else "" - # Cancel every in-progress item in scope BEFORE deleting rows, so each running - # worker is signaled to stop via its item's own status-changed event. With - # multiple workers (multi-GPU) more than one item can be in_progress at once, and - # a user-scoped clear must cancel all of that user's running items — and ONLY - # that user's: other users' rows are out of scope and their workers must keep - # running. See delete_by_destination for the same pattern. - # + keeps the in-progress lookup on a status index rather than the listing index's walk of - # the whole queue (see PENDING_QUEUE_ITEM_COUNT_QUERY). The delete visits every row in scope - # whichever index finds them, so its plan is left to the planner. - params: list[Any] = [queue_id] - if user_id is not None: - params.append(user_id) - self._cancel_in_progress_matching(f"+queue_id == ? {user_filter}", params) - - with self._db.transaction() as cursor: - cursor.execute( - f"""--sql - DELETE - FROM session_queue - WHERE queue_id = ? - {user_filter} - """, - tuple(params), - ) - count = cursor.rowcount - self.__invoker.services.events.emit_queue_cleared(queue_id, user_id) - return ClearResult(deleted=count) - - def delete_queue_items_by_id(self, item_ids: list[int]) -> None: - if not item_ids: - return - placeholders = ", ".join(["?" for _ in item_ids]) - with self._db.transaction() as cursor: - cursor.execute( - f"""--sql - DELETE - FROM session_queue - WHERE item_id IN ({placeholders}) - """, - tuple(item_ids), - ) - - def prune(self, queue_id: str, user_id: Optional[str] = None) -> PruneResult: - with self._db.transaction() as cursor: - # Build WHERE clause with optional user_id filter - user_filter = "AND user_id = ?" if user_id is not None else "" - where = f"""--sql - WHERE - queue_id = ? - AND {PRUNABLE_QUEUE_ITEMS_SQL} - {user_filter} - """ - params: list[Any] = [queue_id] - if user_id is not None: - params.append(user_id) - - cursor.execute( - f"""--sql - DELETE - FROM session_queue - {where}; - """, - tuple(params), - ) - count = cursor.rowcount - return PruneResult(deleted=count) - - def cancel_queue_item(self, item_id: int) -> SessionQueueItem: - chain_item_ids = self._get_workflow_call_chain_item_ids(item_id) - canceled_item: SessionQueueItem | None = None - for chain_item_id in chain_item_ids: - queue_item = self._set_queue_item_status(item_id=chain_item_id, status="canceled") - if chain_item_id == item_id: - canceled_item = queue_item - assert canceled_item is not None - return canceled_item - - def delete_queue_item(self, item_id: int) -> None: - """Deletes a session queue item""" - chain_item_ids = self._get_workflow_call_chain_item_ids(item_id) - if any( - self.get_queue_item(chain_item_id).status not in {"completed", "failed", "canceled"} - for chain_item_id in chain_item_ids - ): - self.cancel_queue_item(item_id) - self.delete_queue_items_by_id(chain_item_ids) - - def complete_queue_item(self, item_id: int, queue_item: Optional[SessionQueueItem] = None) -> SessionQueueItem: - queue_item = self._set_queue_item_status(item_id=item_id, status="completed", queue_item=queue_item) - return queue_item - - def suspend_queue_item(self, item_id: int, queue_item: Optional[SessionQueueItem] = None) -> SessionQueueItem: - queue_item = self._set_queue_item_status(item_id=item_id, status="waiting", queue_item=queue_item) - return queue_item - - def resume_queue_item(self, item_id: int, queue_item: Optional[SessionQueueItem] = None) -> SessionQueueItem: - queue_item = self._set_queue_item_status(item_id=item_id, status="pending", queue_item=queue_item) - return queue_item - - def fail_queue_item( - self, - item_id: int, - error_type: str, - error_message: str, - error_traceback: str, - ) -> SessionQueueItem: - queue_item = self._set_queue_item_status( - item_id=item_id, - status="failed", - error_type=error_type, - error_message=error_message, - error_traceback=error_traceback, - ) - return queue_item - - def _cancel_in_progress_matching(self, match_filter: str, params: list[Any]) -> list[int]: - """Cancel every in-progress item matching `match_filter`, emitting a cancel event for each. - - The bulk-cancel methods exclude in-progress items from their single UPDATE statement, because - a running item must be canceled via `_set_queue_item_status()` so that its - `QueueItemStatusChangedEvent` is emitted — the session processor responds to that event by - setting the cancel event of the worker running that exact item_id. With multiple workers - (multi-GPU) more than one item can be in_progress at once, so each matching item is canceled - individually here rather than relying on a single `get_current()` (which returns only one). - - `match_filter` is a WHERE fragment without the leading WHERE (e.g. - "queue_id == ? AND batch_id IN (?, ?)"); `params` are its bound values. - - Returns the item ids of the in-progress items actually canceled. - """ - with self._db.transaction() as cursor: - cursor.execute( - f"""--sql - SELECT item_id - FROM session_queue - WHERE status == 'in_progress' AND {match_filter}; - """, - tuple(params), - ) - # Cancel in id order; sorted here for the reason given in _collect_item_ids_by_user. - item_ids = sorted(row[0] for row in cursor.fetchall()) - - canceled: list[int] = [] - for item_id in item_ids: - # Count only the items THIS call actually moved to 'canceled'. An item that finished - # between the SELECT and now — including one canceled by a concurrent bulk request - # that selected the same row — is a no-op transition and must not be counted again. - # The transition raises if the row vanished entirely (a concurrent clear/delete); - # such an item needs no cancellation, so skip it rather than failing the whole bulk - # operation. - try: - _, transitioned = self._transition_queue_item_status(item_id, "canceled") - if transitioned: - canceled.append(item_id) - except SessionQueueItemNotFoundError: - continue - return canceled - - def _collect_item_ids_by_user( - self, cursor: sqlite3.Cursor, where: str, params: Sequence[Any] - ) -> dict[str, list[int]]: - """Groups the item ids matched by `where` by their owner's user id. - - Bulk cancel/delete operations mutate many rows in a single SQL statement and therefore - emit no per-item queue_item_status_changed events. Call this with the operation's WHERE - clause before mutating, then pass the result to _emit_queue_items_canceled after the - transaction commits so every affected owner (and everyone's badge counts) can refresh. - """ - cursor.execute( - f"""--sql - SELECT item_id, user_id - FROM session_queue - {where}; - """, - tuple(params), - ) - item_ids_by_user: dict[str, list[int]] = {} - # Ascending ids, sorted here rather than in SQL: the bulk WHEREs are answered by a status, - # batch or listing index as each caller chooses, and an ORDER BY would move the plan onto - # an index that reads the scope's whole history. - for item_id, owner_user_id in sorted(tuple(row) for row in cursor.fetchall()): - item_ids_by_user.setdefault(owner_user_id, []).append(item_id) - return item_ids_by_user - - def _emit_queue_items_canceled(self, queue_id: str, item_ids_by_user: dict[str, list[int]]) -> None: - """Emits queue_items_canceled for a bulk cancel/delete, unless nothing was affected — - an empty result must not broadcast a pointless refetch signal to every client.""" - if item_ids_by_user: - self.__invoker.services.events.emit_queue_items_canceled(queue_id, item_ids_by_user) - - def cancel_by_batch_ids( - self, queue_id: str, batch_ids: list[str], user_id: Optional[str] = None - ) -> CancelByBatchIDsResult: - placeholders = ", ".join(["?" for _ in batch_ids]) - # Build the match filter (with optional user_id filter) shared by the bulk update and the - # in-progress cancellation below. - # + keeps the batch_id index in charge (see PENDING_QUEUE_ITEM_COUNT_QUERY). - user_filter = "AND +user_id = ?" if user_id is not None else "" - match_filter = f"+queue_id == ? AND batch_id IN ({placeholders}) {user_filter}" - params: list[Any] = [queue_id] + batch_ids - if user_id is not None: - params.append(user_id) - - with self._db.transaction() as cursor: - where = f"""--sql - WHERE {match_filter} - AND status != 'canceled' - AND status != 'completed' - AND status != 'failed' - -- In-progress items are canceled individually below so each worker is signaled. - AND status != 'in_progress' - """ - canceled_item_ids_by_user = self._collect_item_ids_by_user(cursor, where, params) - count = sum(len(item_ids) for item_ids in canceled_item_ids_by_user.values()) - cursor.execute( - f"""--sql - UPDATE session_queue - SET status = 'canceled', - status_sequence = COALESCE(status_sequence, 0) + 1 - {where}; - """, - tuple(params), - ) - - # Cancel every in-progress item matching the same filter (multi-GPU: possibly several at - # once). Each cancel emits its own per-item queue_item_status_changed, so the bulk event - # below need not include them (the WHERE above already excludes in-progress rows). - count += len(self._cancel_in_progress_matching(match_filter, params)) - - self._emit_queue_items_canceled(queue_id, canceled_item_ids_by_user) - return CancelByBatchIDsResult(canceled=count) - - def cancel_by_destination( - self, queue_id: str, destination: str, user_id: Optional[str] = None - ) -> CancelByDestinationResult: - user_filter = "AND user_id = ?" if user_id is not None else "" - match_filter = f"queue_id == ? AND destination == ? {user_filter}" - params: list[Any] = [queue_id, destination] - if user_id is not None: - params.append(user_id) - - with self._db.transaction() as cursor: - where = f"""--sql - WHERE {match_filter} - AND status != 'canceled' - AND status != 'completed' - AND status != 'failed' - -- In-progress items are canceled individually below so each worker is signaled. - AND status != 'in_progress' - """ - canceled_item_ids_by_user = self._collect_item_ids_by_user(cursor, where, params) - count = sum(len(item_ids) for item_ids in canceled_item_ids_by_user.values()) - cursor.execute( - f"""--sql - UPDATE session_queue - SET status = 'canceled', - status_sequence = COALESCE(status_sequence, 0) + 1 - {where}; - """, - tuple(params), - ) - - # Cancel every in-progress item matching the same filter (multi-GPU: possibly several at - # once). Each cancel emits its own per-item queue_item_status_changed, so the bulk event - # below need not include them (the WHERE above already excludes in-progress rows). - count += len(self._cancel_in_progress_matching(match_filter, params)) - - self._emit_queue_items_canceled(queue_id, canceled_item_ids_by_user) - return CancelByDestinationResult(canceled=count) - - def delete_by_destination( - self, queue_id: str, destination: str, user_id: Optional[str] = None - ) -> DeleteByDestinationResult: - user_filter = "AND user_id = ?" if user_id is not None else "" - match_filter = f"queue_id == ? AND destination == ? {user_filter}" - params: list[Any] = [queue_id, destination] - if user_id is not None: - params.append(user_id) - - # Cancel every in-progress item first so each running worker is signaled to stop before we - # delete its row. With multiple workers (multi-GPU) more than one item can be in_progress; - # canceling only get_current() would leave the others running (and then failing to update a - # deleted row). See _cancel_in_progress_matching. - canceled_in_progress_ids = set(self._cancel_in_progress_matching(match_filter, params)) - - with self._db.transaction() as cursor: - where = f"""--sql - WHERE - queue_id == ? - AND destination == ? - {user_filter} - """ - deleted_item_ids_by_user = self._collect_item_ids_by_user(cursor, where, params) - count = sum(len(item_ids) for item_ids in deleted_item_ids_by_user.values()) - # The in-progress items canceled above each emitted their own per-item - # queue_item_status_changed, so the bulk event below must not signal them a second - # time. Their rows still match the destination WHERE (the cancel only flipped their - # status), so they are excluded from the collected ids — but not from the DELETE or - # the returned count. - if canceled_in_progress_ids: - for owner_user_id in list(deleted_item_ids_by_user): - remaining = [ - i for i in deleted_item_ids_by_user[owner_user_id] if i not in canceled_in_progress_ids - ] - if remaining: - deleted_item_ids_by_user[owner_user_id] = remaining - else: - del deleted_item_ids_by_user[owner_user_id] - cursor.execute( - f"""--sql - DELETE FROM session_queue - {where}; - """, - tuple(params), - ) - self._emit_queue_items_canceled(queue_id, deleted_item_ids_by_user) - return DeleteByDestinationResult(deleted=count) - - def delete_all_except_current(self, queue_id: str, user_id: Optional[str] = None) -> DeleteAllExceptCurrentResult: - current_chain_item_ids = self._get_current_workflow_call_chain_item_ids(queue_id) - with self._db.transaction() as cursor: - # Build WHERE clause with optional user_id filter - user_filter = "AND user_id = ?" if user_id is not None else "" - current_chain_filter = "" - if current_chain_item_ids: - placeholders = ", ".join(["?" for _ in current_chain_item_ids]) - current_chain_filter = f"AND item_id NOT IN ({placeholders})" - # + keeps the statements on a status index (see PENDING_QUEUE_ITEM_COUNT_QUERY). - where = f"""--sql - WHERE - +queue_id == ? - AND status IN ('pending', 'waiting') - {user_filter} - {current_chain_filter} - """ - params: list[Any] = [queue_id] - if user_id is not None: - params.append(user_id) - params.extend(current_chain_item_ids) - - deleted_item_ids_by_user = self._collect_item_ids_by_user(cursor, where, params) - count = sum(len(item_ids) for item_ids in deleted_item_ids_by_user.values()) - cursor.execute( - f"""--sql - DELETE - FROM session_queue - {where}; - """, - tuple(params), - ) - self._emit_queue_items_canceled(queue_id, deleted_item_ids_by_user) - return DeleteAllExceptCurrentResult(deleted=count) - - def cancel_by_queue_id( - self, queue_id: str, user_id: Optional[str] = None, origin_prefix: Optional[str] = None - ) -> CancelByQueueIDResult: - user_filter = "AND user_id = ?" if user_id is not None else "" - origin_filter = "AND origin LIKE ?" if origin_prefix is not None else "" - # + keeps the statements on a status index, which finds the live rows, rather than the - # listing index, which walks the queue's whole history (see PENDING_QUEUE_ITEM_COUNT_QUERY). - match_filter = f"+queue_id == ? {user_filter} {origin_filter}" - params: list[Any] = [queue_id] - if user_id is not None: - params.append(user_id) - if origin_prefix is not None: - params.append(f"{origin_prefix}%") - - with self._db.transaction() as cursor: - # In-progress items are canceled individually below so each worker is signaled. - where = f"""--sql - WHERE {match_filter} - AND status IN ('pending', 'waiting') - """ - canceled_item_ids_by_user = self._collect_item_ids_by_user(cursor, where, params) - count = sum(len(item_ids) for item_ids in canceled_item_ids_by_user.values()) - cursor.execute( - f"""--sql - UPDATE session_queue - SET status = 'canceled', - status_sequence = COALESCE(status_sequence, 0) + 1 - {where}; - """, - tuple(params), - ) - - # Cancel every in-progress item in the queue (multi-GPU: possibly several at once). Each - # cancel emits its own per-item queue_item_status_changed; the bulk event below covers the - # silently-updated rows. - count += len(self._cancel_in_progress_matching(match_filter, params)) - self._emit_queue_items_canceled(queue_id, canceled_item_ids_by_user) - return CancelByQueueIDResult(canceled=count) - - def cancel_all_except_current( - self, queue_id: str, user_id: Optional[str] = None, origin_prefix: Optional[str] = None - ) -> CancelAllExceptCurrentResult: - current_chain_item_ids = self._get_current_workflow_call_chain_item_ids(queue_id) - with self._db.transaction() as cursor: - # Build WHERE clause with optional user_id and origin_prefix filters - user_filter = "AND user_id = ?" if user_id is not None else "" - origin_filter = "AND origin LIKE ?" if origin_prefix is not None else "" - current_chain_filter = "" - if current_chain_item_ids: - placeholders = ", ".join(["?" for _ in current_chain_item_ids]) - current_chain_filter = f"AND item_id NOT IN ({placeholders})" - # + keeps the statements on a status index (see PENDING_QUEUE_ITEM_COUNT_QUERY). - where = f"""--sql - WHERE - +queue_id == ? - AND status IN ('pending', 'waiting') - {user_filter} - {origin_filter} - {current_chain_filter} - """ - params: list[Any] = [queue_id] - if user_id is not None: - params.append(user_id) - if origin_prefix is not None: - params.append(f"{origin_prefix}%") - params.extend(current_chain_item_ids) - - canceled_item_ids_by_user = self._collect_item_ids_by_user(cursor, where, params) - count = sum(len(item_ids) for item_ids in canceled_item_ids_by_user.values()) - cursor.execute( - f"""--sql - UPDATE session_queue - SET status = 'canceled', - status_sequence = COALESCE(status_sequence, 0) + 1 - {where}; - """, - tuple(params), - ) - self._emit_queue_items_canceled(queue_id, canceled_item_ids_by_user) - return CancelAllExceptCurrentResult(canceled=count) - - def _get_queue_item_with_load_status(self, item_id: int) -> tuple[SessionQueueItem, bool]: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT - sq.*, - u.display_name as user_display_name, - u.email as user_email - FROM session_queue sq - LEFT JOIN users u ON sq.user_id = u.user_id - WHERE sq.item_id = ? - """, - (item_id,), - ) - result = cast(Union[sqlite3.Row, None], cursor.fetchone()) - if result is None: - raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}") - return self._hydrate_queue_item(dict(result), quarantine=False) - - def get_queue_item(self, item_id: int) -> SessionQueueItem: - return self._get_queue_item_with_load_status(item_id)[0] - - def get_queue_item_for_api(self, item_id: int) -> SessionQueueItem: - return self._get_queue_item_for_api(item_id) - - def get_queue_item_workflow_json(self, item_id: int) -> str | None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT workflow - FROM session_queue - WHERE item_id = ? - """, - (item_id,), - ) - result = cast(Union[sqlite3.Row, None], cursor.fetchone()) - return cast(str | None, result["workflow"]) if result is not None else None - - def save_queue_item_session(self, item_id: int, session: GraphExecutionState) -> None: - with self._db.transaction() as cursor: - # Use exclude_none so we don't end up with a bunch of nulls in the graph - this can cause validation errors - # when the graph is loaded. Persisted sessions are used to resume execution across queue boundaries. - session_json = json.dumps(dump_execution_state(session), default=to_jsonable_python) - cursor.execute( - """--sql - UPDATE session_queue - SET session = ? - WHERE item_id = ? - """, - (session_json, item_id), - ) - if cursor.rowcount == 0: - raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}") - - def _save_queue_item_session_if_active(self, item_id: int, session: GraphExecutionState) -> bool: - """Persist a session only while its queue item is non-terminal. - - This is an internal race guard. The existing transaction lock makes the status check and - write atomic for all queue mutations in this process without adding a persistence column. - """ - session_json = json.dumps(dump_execution_state(session), default=to_jsonable_python) - with self._db.transaction() as cursor: - cursor.execute( - """--sql - UPDATE session_queue - SET session = ? - WHERE item_id = ? AND status NOT IN ('completed', 'failed', 'canceled') - """, - (session_json, item_id), - ) - if cursor.rowcount != 0: - return True - cursor.execute("SELECT 1 FROM session_queue WHERE item_id = ?", (item_id,)) - if cursor.fetchone() is None: - raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}") - return False - - def set_queue_item_session(self, item_id: int, session: GraphExecutionState) -> SessionQueueItem: - self.save_queue_item_session(item_id, session) - return self.get_queue_item(item_id) - - def record_workflow_call_child_completion( - self, parent_item_id: int, child_item_id: int, output_values: dict[str, Any] - ) -> WorkflowCallChildCompletion | None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT * - FROM session_queue - WHERE item_id = ? - """, - (parent_item_id,), - ) - row = cast(sqlite3.Row | None, cursor.fetchone()) - if row is None: - raise SessionQueueItemNotFoundError(f"No queue item with id {parent_item_id}") - parent_queue_item, readable = self._hydrate_queue_item(dict(row), quarantine=False) - if not readable: - raise ValueError("Unable to record workflow call child completion for an unreadable parent session.") - if parent_queue_item.status in ("completed", "failed", "canceled"): - return None - - execution = parent_queue_item.session.waiting_workflow_call_execution - if execution is not None and child_item_id in execution.completed_child_item_ids: - return None - generic_update = parent_queue_item.session.record_generic_child_completion(child_item_id, output_values) - if generic_update is not None and not generic_update.changed: - return None - legacy_should_resume, legacy_values = ( - parent_queue_item.session.record_waiting_workflow_call_child_completion(child_item_id, output_values) - ) - if generic_update is None: - should_resume_parent, aggregated_values = legacy_should_resume, legacy_values - else: - should_resume_parent = generic_update.status == "completed" - aggregated_values = { - key: values[0] if len(values) == 1 else values - for key, values in generic_update.aggregated_outputs.items() - } - if generic_update.status == "completed" and aggregated_values != legacy_values: - raise ValueError("Generic child aggregation disagrees with workflow-call aggregation.") - - session_json = json.dumps(dump_execution_state(parent_queue_item.session), default=to_jsonable_python) - cursor.execute( - """--sql - UPDATE session_queue - SET session = ? - WHERE item_id = ? - """, - (session_json, parent_item_id), - ) - if cursor.rowcount == 0: - raise SessionQueueItemNotFoundError(f"No queue item with id {parent_item_id}") - - return WorkflowCallChildCompletion( - parent_queue_item=parent_queue_item, - should_resume=should_resume_parent, - aggregated_values=aggregated_values, - ) - - def enqueue_workflow_call_children( - self, - parent_queue_item: SessionQueueItem, - child_sessions: list[tuple[GraphExecutionState, list[NodeFieldValue] | None]], - ) -> list[SessionQueueItem]: - workflow_call_execution = parent_queue_item.session.waiting_workflow_call_execution - if workflow_call_execution is None: - raise ValueError("Parent queue item is missing active workflow call execution metadata.") - if not child_sessions: - raise ValueError("Workflow call must enqueue at least one child execution.") - - serialized_children = [ - ( - json.dumps(dump_execution_state(child_session), default=to_jsonable_python), - json.dumps(field_values, default=to_jsonable_python) if field_values is not None else None, - ) - for child_session, field_values in child_sessions - ] - root_item_id = parent_queue_item.root_item_id or parent_queue_item.item_id - child_item_ids: list[int] = [] - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT status - FROM session_queue - WHERE item_id = ? - """, - (parent_queue_item.item_id,), - ) - parent_status_row = cursor.fetchone() - if parent_status_row is None: - raise SessionQueueItemNotFoundError(f"No queue item with id {parent_queue_item.item_id}") - if parent_status_row[0] in ("completed", "failed", "canceled"): - raise ValueError("Cannot enqueue workflow call children for a terminal parent queue item.") - - cursor.execute(PENDING_QUEUE_ITEM_COUNT_QUERY, (parent_queue_item.queue_id,)) - pending_count = cast(int, cursor.fetchone()[0]) - max_queue_size = self.__invoker.services.configuration.max_queue_size - if pending_count + len(child_sessions) > max_queue_size: - raise TooManySessionsError( - "call_saved_workflow exceeds remaining queue capacity for child workflow executions" - ) - - for (session_json, field_values_json), (child_session, _field_values) in zip( - serialized_children, child_sessions, strict=True - ): - cursor.execute( - """--sql - INSERT INTO session_queue ( - queue_id, - session, - session_id, - batch_id, - field_values, - priority, - workflow, - origin, - destination, - retried_from_item_id, - user_id, - workflow_call_id, - parent_item_id, - parent_session_id, - root_item_id, - workflow_call_depth, - project_id, - status - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'pending') - """, - ( - parent_queue_item.queue_id, - session_json, - child_session.id, - parent_queue_item.batch_id, - field_values_json, - parent_queue_item.priority, - None, - parent_queue_item.origin, - parent_queue_item.destination, - None, - parent_queue_item.user_id, - workflow_call_execution.id, - parent_queue_item.item_id, - parent_queue_item.session_id, - root_item_id, - workflow_call_execution.depth, - parent_queue_item.project_id, - ), - ) - child_item_ids.append(cast(int, cursor.lastrowid)) - - parent_queue_item.session.set_waiting_workflow_call_child_item_ids(child_item_ids) - session_json = json.dumps(dump_execution_state(parent_queue_item.session), default=to_jsonable_python) - cursor.execute( - """--sql - UPDATE session_queue - SET session = ?, - status = 'waiting', status_sequence = COALESCE(status_sequence, 0) + 1 - WHERE item_id = ? AND session = ? - AND status NOT IN ('completed', 'failed', 'canceled') - """, - ( - session_json, - parent_queue_item.item_id, - getattr(parent_queue_item, "_session_json", None) - or json.dumps(dump_execution_state(parent_queue_item.session), default=to_jsonable_python), - ), - ) - if cursor.rowcount == 0: - raise SessionQueueItemChangedError("Parent queue item changed while enqueuing workflow call children") - - parent_queue_item.status = "waiting" - child_queue_items = [self.get_queue_item(item_id) for item_id in child_item_ids] - for queue_item in [self.get_queue_item(parent_queue_item.item_id), *child_queue_items]: - batch_status = self.get_batch_status(queue_id=queue_item.queue_id, batch_id=queue_item.batch_id) - queue_status = self.get_queue_status( - queue_id=queue_item.queue_id, user_id=queue_item.user_id, acting_user_id=queue_item.user_id - ) - self.__invoker.services.events.emit_queue_item_status_changed(queue_item, batch_status, queue_status) - return child_queue_items - - def enqueue_workflow_call_child( - self, - parent_queue_item: SessionQueueItem, - child_session: GraphExecutionState, - field_values: list[NodeFieldValue] | None = None, - ) -> SessionQueueItem: - workflow_call_execution = parent_queue_item.session.waiting_workflow_call_execution - if workflow_call_execution is None: - raise ValueError("Parent queue item is missing active workflow call execution metadata.") - - session_json = json.dumps(dump_execution_state(child_session), default=to_jsonable_python) - field_values_json = json.dumps(field_values, default=to_jsonable_python) if field_values is not None else None - root_item_id = parent_queue_item.root_item_id or parent_queue_item.item_id - - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT status - FROM session_queue - WHERE item_id = ? - """, - (parent_queue_item.item_id,), - ) - parent_row = cursor.fetchone() - if parent_row is None: - raise SessionQueueItemNotFoundError(f"No queue item with id {parent_queue_item.item_id}") - if parent_row[0] in ("completed", "failed", "canceled"): - raise ValueError("Cannot enqueue workflow call child for a terminal parent queue item.") - - cursor.execute(PENDING_QUEUE_ITEM_COUNT_QUERY, (parent_queue_item.queue_id,)) - pending_count = cast(int, cursor.fetchone()[0]) - if pending_count >= self.__invoker.services.configuration.max_queue_size: - raise TooManySessionsError( - "call_saved_workflow exceeds remaining queue capacity for child workflow executions" - ) - - cursor.execute( - """--sql - INSERT INTO session_queue ( - queue_id, - session, - session_id, - batch_id, - field_values, - priority, - workflow, - origin, - destination, - retried_from_item_id, - user_id, - project_id, - workflow_call_id, - parent_item_id, - parent_session_id, - root_item_id, - workflow_call_depth, - status - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'pending') - """, - ( - parent_queue_item.queue_id, - session_json, - child_session.id, - parent_queue_item.batch_id, - field_values_json, - parent_queue_item.priority, - None, - parent_queue_item.origin, - parent_queue_item.destination, - None, - parent_queue_item.user_id, - parent_queue_item.project_id, - workflow_call_execution.id, - parent_queue_item.item_id, - parent_queue_item.session_id, - root_item_id, - workflow_call_execution.depth, - ), - ) - item_id = cursor.lastrowid - - queue_item = self.get_queue_item(item_id) - batch_status = self.get_batch_status(queue_id=queue_item.queue_id, batch_id=queue_item.batch_id) - queue_status = self.get_queue_status( - queue_id=queue_item.queue_id, user_id=queue_item.user_id, acting_user_id=queue_item.user_id - ) - self.__invoker.services.events.emit_queue_item_status_changed(queue_item, batch_status, queue_status) - return queue_item - - def cancel_workflow_call_children( - self, workflow_call_id: str, exclude_item_ids: set[int] | None = None - ) -> list[int]: - exclude_item_ids = exclude_item_ids or set() - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT item_id - FROM session_queue - WHERE workflow_call_id = ? - ORDER BY item_id ASC - """, - (workflow_call_id,), - ) - item_ids = [row[0] for row in cast(list[sqlite3.Row], cursor.fetchall())] - item_ids_with_descendants: list[int] = [] - for item_id in item_ids: - item_ids_with_descendants.append(item_id) - item_ids_with_descendants.extend(self._get_workflow_call_descendant_ids(item_id)) - item_ids = list(dict.fromkeys(item_ids_with_descendants)) - canceled_item_ids: list[int] = [] - for item_id in item_ids: - if item_id in exclude_item_ids: - continue - queue_item = self.get_queue_item(item_id) - if queue_item.status in {"completed", "failed", "canceled"}: - continue - self._set_queue_item_status(item_id=item_id, status="canceled") - canceled_item_ids.append(item_id) - return canceled_item_ids - - def list_queue_items( - self, - queue_id: str, - limit: int, - priority: int, - cursor: Optional[int] = None, - status: Optional[QUEUE_ITEM_STATUS] = None, - destination: Optional[str] = None, - ) -> CursorPaginatedResults[SessionQueueItem]: - with self._db.transaction() as cursor_: - item_id = cursor - query = """--sql - SELECT * - FROM session_queue - WHERE queue_id = ? - """ - params: list[Union[str, int]] = [queue_id] - - if status is not None: - query += """--sql - AND status = ? - """ - params.append(status) - - if destination is not None: - query += """---sql - AND destination = ? - """ - params.append(destination) - - if item_id is not None: - query += """--sql - AND ((priority < ?) OR (priority = ? AND item_id > ?)) - """ - params.extend([priority, priority, item_id]) - - query += """--sql - ORDER BY - priority DESC, - item_id ASC - LIMIT ? - """ - params.append(limit + 1) - cursor_.execute(query, params) - results = cast(list[sqlite3.Row], cursor_.fetchall()) - items = [self._hydrate_queue_item(dict(result), quarantine=False)[0] for result in results] - has_more = False - if len(items) > limit: - # remove the extra item - items.pop() - has_more = True - return CursorPaginatedResults(items=items, limit=limit, has_more=has_more) - - def _list_all_queue_item_rows( - self, - queue_id: str, - destination: Optional[str] = None, - ) -> list[sqlite3.Row]: - with self._db.transaction() as cursor: - query = """--sql - SELECT - sq.*, - u.display_name as user_display_name, - u.email as user_email - FROM session_queue sq - LEFT JOIN users u ON sq.user_id = u.user_id - WHERE sq.queue_id = ? - """ - params: list[Union[str, int]] = [queue_id] - - if destination is not None: - query += """---sql - AND sq.destination = ? - """ - params.append(destination) - - query += """--sql - ORDER BY - sq.priority DESC, - sq.item_id ASC - ; - """ - cursor.execute(query, params) - return cast(list[sqlite3.Row], cursor.fetchall()) - - def list_all_queue_items( - self, - queue_id: str, - destination: Optional[str] = None, - ) -> list[SessionQueueItem]: - """Gets all queue items with fully rehydrated runtime sessions.""" - results = self._list_all_queue_item_rows(queue_id=queue_id, destination=destination) - return [self._hydrate_queue_item(dict(result), quarantine=False)[0] for result in results] - - def list_all_queue_items_for_api( - self, - queue_id: str, - destination: Optional[str] = None, - ) -> list[SessionQueueItem]: - """Gets response-shaped queue items without rebuilding runtime execution state.""" - results = self._list_all_queue_item_rows(queue_id=queue_id, destination=destination) - return [self._project_queue_item_for_read(dict(result)) for result in results] - - def get_queue_item_ids( - self, - queue_id: str, - order_dir: SQLiteDirection = SQLiteDirection.Descending, - user_id: Optional[str] = None, - origin_prefix: Optional[str] = None, - limit: Optional[int] = None, - ) -> ItemIdsResult: - with self._db.transaction() as cursor_: - query = """--sql - SELECT item_id - FROM session_queue - WHERE queue_id = ? - """ - query_params: list[str | int] = [queue_id] - - if user_id is not None: - query += " AND user_id = ?" - query_params.append(user_id) - - if origin_prefix is not None: - query += " AND origin LIKE ?" - query_params.append(f"{origin_prefix}%") - - # Items enqueued together share a created_at; they stay in enqueue order in either - # direction. idx_session_queue_listing, scanned backwards, yields the newest-first order. - query += f" ORDER BY created_at {order_dir.value}, item_id ASC" - - # Either direction walks the index in order (oldest first sorts only items that share a - # created_at), so a limited read stops once `limit` entries pass the filters and costs the - # window rather than the history. A scope with fewer matches than `limit` (an origin prefix - # with few or no items) still walks the queue's whole index, and one whose matches lie far - # back walks it until it reaches them: origin is filtered on each entry, not used to seek. - if limit is not None: - query += " LIMIT ?" - query_params.append(limit) - - cursor_.execute(query, query_params) - result = cast(list[sqlite3.Row], cursor_.fetchall()) - item_ids = [row[0] for row in result] - - return ItemIdsResult(item_ids=item_ids, total_count=len(item_ids)) - - def get_queue_item_summaries_by_ids(self, queue_id: str, item_ids: list[int]) -> list[SessionQueueItemSummary]: - if not item_ids: - return [] - - rows: list[sqlite3.Row] = [] - with self._db.transaction() as cursor: - # Each id becomes one bind parameter, so a single IN (...) would blow past SQLite's - # per-statement variable limit for large id lists. Query in chunks instead - callers - # are bounded at the API layer, but this keeps any caller from hitting that ceiling. - # The unary + keeps queue_id off idx_session_queue_listing, so the requested primary - # keys are looked up rather than the queue's history scanned (see - # PENDING_QUEUE_ITEM_COUNT_QUERY for the rule). - for chunk_start in range(0, len(item_ids), SQLITE_MAX_BIND_PARAMS_PER_CHUNK): - chunk = item_ids[chunk_start : chunk_start + SQLITE_MAX_BIND_PARAMS_PER_CHUNK] - placeholders = ", ".join("?" for _ in chunk) - cursor.execute( - f"""--sql - SELECT - sq.item_id, - sq.created_at, - sq.status, - sq.device, - sq.started_at, - sq.completed_at, - sq.origin, - sq.destination, - sq.batch_id, - sq.user_id, - u.display_name AS user_display_name, - u.email AS user_email, - sq.field_values, - sq.parent_item_id - FROM session_queue sq - LEFT JOIN users u ON sq.user_id = u.user_id - WHERE +sq.queue_id = ? AND sq.item_id IN ({placeholders}) - """, - (queue_id, *chunk), - ) - rows.extend(cast(list[sqlite3.Row], cursor.fetchall())) - - summaries_by_id = { - row["item_id"]: SessionQueueItemSummary.queue_item_summary_from_dict(dict(row)) for row in rows - } - return [summaries_by_id[item_id] for item_id in item_ids if item_id in summaries_by_id] - - def has_active_queue_work(self) -> bool: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT 1 - FROM session_queue - WHERE status IN ('pending', 'in_progress') - LIMIT 1 - """ - ) - return cursor.fetchone() is not None - - def get_queue_status( - self, - queue_id: str, - user_id: Optional[str] = None, - acting_user_id: Optional[str] = None, - origin_prefix: Optional[str] = None, - is_admin: bool = False, - ) -> SessionQueueStatus: - with self._db.transaction() as cursor: - # Aggregate counts are global across all users within the requested scope. When - # user_id is provided, the same pass also counts that user's own items, so the caller - # can render the per-user portion of the badge. These are returned in the separate - # user_pending/user_in_progress fields and never replace the global counts. Counting - # both in one statement reads the scope's history once instead of twice. - params: list[str] = [] - user_count = "" - if user_id is not None: - user_count = ", SUM(CASE WHEN user_id = ? THEN 1 ELSE 0 END)" - params.append(user_id) - query = f"""--sql - SELECT status, count(*){user_count} - FROM session_queue - WHERE queue_id = ? - """ - params.append(queue_id) - - if origin_prefix is not None: - query += " AND origin LIKE ?" - params.append(f"{origin_prefix}%") - - query += " GROUP BY status" - cursor.execute(query, params) - counts_result = cast(list[sqlite3.Row], cursor.fetchall()) - - # Only the four identifier columns, not a full SessionQueueItem: this runs on - # every status poll, and hydrating the item would deserialize its whole session - # graph. The origin filter mirrors the aggregate-count query above so a - # scoped caller never sees another scope's current item. - current_item_query = """--sql - SELECT item_id, session_id, batch_id, user_id - FROM session_queue - WHERE queue_id = ? AND status = 'in_progress' - """ - current_item_params: list[str] = [queue_id] - - if origin_prefix is not None: - current_item_query += " AND origin LIKE ?" - current_item_params.append(f"{origin_prefix}%") - - current_item_query += " LIMIT 1" - cursor.execute(current_item_query, current_item_params) - current_item = cast(Union[sqlite3.Row, None], cursor.fetchone()) - - total = sum(row[1] or 0 for row in counts_result) - counts: dict[str, int] = {row[0]: row[1] for row in counts_result} - - user_pending: Optional[int] = None - user_in_progress: Optional[int] = None - if user_id is not None: - user_counts: dict[str, int] = {row[0]: row[2] for row in counts_result} - user_pending = user_counts.get("pending", 0) - user_in_progress = user_counts.get("in_progress", 0) - - # Redaction is decided from the same current_item snapshot used to embed identifiers, - # so a concurrent transition (e.g. B finishing while A's status changes) cannot leave - # stale identifiers in the result. The aggregate counts stay global; only the current - # item's identifiers are gated. acting_user_id (event path) takes precedence over - # user_id (API path) when deciding the redaction owner; either being None means a - # global caller who may see the current item. is_admin disables redaction outright so - # admin callers can pass their user_id (for the per-user counts) without losing - # visibility of other users' current item. - owner_user_id = user_id if acting_user_id is None else acting_user_id - current_item_id = None - current_session_id = None - current_batch_id = None - if current_item is not None and (is_admin or owner_user_id is None or current_item["user_id"] == owner_user_id): - current_item_id = current_item["item_id"] - current_session_id = current_item["session_id"] - current_batch_id = current_item["batch_id"] - - return SessionQueueStatus( - queue_id=queue_id, - item_id=current_item_id, - session_id=current_session_id, - batch_id=current_batch_id, - pending=counts.get("pending", 0), - in_progress=counts.get("in_progress", 0), - waiting=counts.get("waiting", 0), - completed=counts.get("completed", 0), - failed=counts.get("failed", 0), - canceled=counts.get("canceled", 0), - total=total, - user_pending=user_pending, - user_in_progress=user_in_progress, - ) - - def get_batch_status(self, queue_id: str, batch_id: str, user_id: Optional[str] = None) -> BatchStatus: - with self._db.transaction() as cursor: - query = """--sql - SELECT status, count(*), origin, destination - FROM session_queue - WHERE queue_id = ? AND batch_id = ? - """ - params: list[str] = [queue_id, batch_id] - if user_id is not None: - # + keeps the batch_id index in charge (see PENDING_QUEUE_ITEM_COUNT_QUERY). - query += " AND +user_id = ?" - params.append(user_id) - query += " GROUP BY status" - cursor.execute(query, params) - result = cast(list[sqlite3.Row], cursor.fetchall()) - total = sum(row[1] or 0 for row in result) - counts: dict[str, int] = {row[0]: row[1] for row in result} - origin = result[0]["origin"] if result else None - destination = result[0]["destination"] if result else None - - return BatchStatus( - batch_id=batch_id, - origin=origin, - destination=destination, - queue_id=queue_id, - pending=counts.get("pending", 0), - in_progress=counts.get("in_progress", 0), - waiting=counts.get("waiting", 0), - completed=counts.get("completed", 0), - failed=counts.get("failed", 0), - canceled=counts.get("canceled", 0), - total=total, - ) - - def get_counts_by_destination( - self, queue_id: str, destination: str, user_id: Optional[str] = None - ) -> SessionQueueCountsByDestination: - with self._db.transaction() as cursor: - query = """--sql - SELECT status, count(*) - FROM session_queue - WHERE queue_id = ? AND destination = ? - """ - params: list[str] = [queue_id, destination] - if user_id is not None: - query += " AND user_id = ?" - params.append(user_id) - query += " GROUP BY status" - cursor.execute(query, params) - counts_result = cast(list[sqlite3.Row], cursor.fetchall()) - - total = sum(row[1] or 0 for row in counts_result) - counts: dict[str, int] = {row[0]: row[1] for row in counts_result} - - return SessionQueueCountsByDestination( - queue_id=queue_id, - destination=destination, - pending=counts.get("pending", 0), - in_progress=counts.get("in_progress", 0), - waiting=counts.get("waiting", 0), - completed=counts.get("completed", 0), - failed=counts.get("failed", 0), - canceled=counts.get("canceled", 0), - total=total, - ) - - def retry_items_by_id(self, queue_id: str, item_ids: list[int]) -> RetryItemsResult: - """Retries the given queue items""" - with self._db.transaction() as cursor: - values_to_insert: list[ValueToInsertTuple] = [] - retried_root_item_ids: list[int] = [] - retried_user_ids: list[str] = [] - retried_item_ids_by_user: dict[str, list[int]] = {} - seen_root_item_ids: set[int] = set() - hydrated_queue_items: dict[int, SessionQueueItem] = {} - max_new_queue_items = self.__invoker.services.configuration.max_queue_size - self._get_current_queue_size( - queue_id - ) - - if max_new_queue_items <= 0: - return RetryItemsResult(queue_id=queue_id, retried_item_ids=[]) - - for item_id in item_ids: - if item_id in hydrated_queue_items: - queue_item = hydrated_queue_items[item_id] - else: - try: - queue_item = self._get_queue_item_for_retry(item_id, cursor=cursor) - except SessionQueueItemNotFoundError: - continue - hydrated_queue_items[item_id] = queue_item - if queue_item.queue_id != queue_id: - continue - - if queue_item.status not in ("failed", "canceled"): - continue - - root_item_id = queue_item.root_item_id or queue_item.item_id - if root_item_id in seen_root_item_ids: - continue - seen_root_item_ids.add(root_item_id) - - if root_item_id in hydrated_queue_items: - root_queue_item = hydrated_queue_items[root_item_id] - else: - root_queue_item = self._get_queue_item_for_retry(root_item_id, cursor=cursor) - hydrated_queue_items[root_item_id] = root_queue_item - if not root_queue_item._snapshot_readable: - continue - if root_queue_item.status not in ("failed", "canceled"): - continue - - retried_root_item_ids.append(root_item_id) - retried_user_ids.append(root_queue_item.user_id) - retried_item_ids_by_user.setdefault(root_queue_item.user_id, []).append(root_item_id) - - field_values_json = ( - json.dumps(root_queue_item.field_values, default=to_jsonable_python) - if root_queue_item.field_values - else None - ) - workflow_json = ( - json.dumps(root_queue_item.workflow, default=to_jsonable_python) - if root_queue_item.workflow - else None - ) - # Validate the graph before dumping a fresh empty execution state. The full retry read above already - # rehydrated runtime state for contract-compatible validation and recovery semantics. - root_graph = Graph.model_validate( - root_queue_item.session.graph.model_dump(mode="python", warnings=False), strict=False - ) - cloned_session = GraphExecutionState(graph=root_graph) - cloned_session_json = json.dumps(dump_execution_state(cloned_session), default=to_jsonable_python) - - retried_from_item_id = ( - root_queue_item.retried_from_item_id - if root_queue_item.retried_from_item_id is not None - else root_queue_item.item_id - ) - - value_to_insert: ValueToInsertTuple = ( - root_queue_item.queue_id, - cloned_session_json, - cloned_session.id, - root_queue_item.batch_id, - field_values_json, - root_queue_item.priority, - workflow_json, - root_queue_item.origin, - root_queue_item.destination, - retried_from_item_id, - root_queue_item.user_id, - root_queue_item.project_id, - ) - values_to_insert.append(value_to_insert) - - if len(values_to_insert) >= max_new_queue_items: - break - - cursor.executemany( - """--sql - INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id, user_id, project_id) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - values_to_insert, - ) - - retry_result = RetryItemsResult( - queue_id=queue_id, - retried_item_ids=retried_root_item_ids, - ) - self.__invoker.services.events.emit_queue_items_retried( - retry_result, - user_ids=list(dict.fromkeys(retried_user_ids)), - retried_item_ids_by_user=retried_item_ids_by_user, - ) - return retry_result diff --git a/invokeai/app/services/shared/sqlite/__init__.py b/invokeai/app/services/shared/database/__init__.py similarity index 100% rename from invokeai/app/services/shared/sqlite/__init__.py rename to invokeai/app/services/shared/database/__init__.py diff --git a/invokeai/app/services/shared/database/copy.py b/invokeai/app/services/shared/database/copy.py new file mode 100644 index 00000000000..af3f5355e82 --- /dev/null +++ b/invokeai/app/services/shared/database/copy.py @@ -0,0 +1,400 @@ +"""Copying the rows of the application's tables from one database to another, whatever their backends: the +bootstrap of a server database, and `invoke-db-copy`, which moves an install's SQLite database to a server.""" + +import hashlib +import json +import unicodedata +from collections.abc import Callable, Iterable, Iterator, Sequence +from typing import Any, NamedTuple, Optional + +from sqlalchemy import Connection, Insert, LargeBinary, Table, cast, func, insert, inspect, select + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.dialect import insert_ignore +from invokeai.app.services.shared.database.schema import metadata +from invokeai.app.services.shared.database.schema.image_index import image_index_vocab_terms +from invokeai.app.services.shared.database.schema.migrator import applied_migrations, migrations +from invokeai.app.services.shared.database.schema.models import models +from invokeai.app.services.shared.database.types import TIMESTAMP_LENGTH, Key, NoCaseKey, Timestamp + +# Rows are read a few at a time, and written in parts of one INSERT each; a part's bytes keep that statement well +# within a server's max_allowed_packet, which startup requires to be 64 MiB at least. +_READ_ROWS = 50 +_BATCH_ROWS = 1000 +_BATCH_BYTES = 16 * 1024 * 1024 + +# A row as the copy reads it: the values of the table's stored (not generated) columns, by name. +Row = dict[str, Any] +Transform = Callable[[Table, Row], Row] + +# The migrator's records, written last of all: a database without them is refused as incomplete. +_RECORDS = [migrations, applied_migrations] +_DATA = [table for table in metadata.sorted_tables if table not in _RECORDS] +# Terms a server's case-insensitive collation equates, which SQLite's ASCII-only NOCASE keeps apart: the server +# keeps one of them, which is all the vocabulary needs. +_MERGED_ON_A_SERVER = (image_index_vocab_terms,) + + +def copy_rows( + source: Database, + target: Database, + *, + tables: Optional[Sequence[Table]] = None, + transform: Optional[Transform] = None, + keep_first: Sequence[Table] = (), + progress: Optional[Callable[[str], None]] = None, +) -> dict[str, int]: + """Copies every row of the application tables (by default all of them) from `source` into `target`. + + Tables are copied in foreign key order, in batches of a transaction each, and a table the source lacks is + skipped. The target computes the generated columns itself. The ids an AUTOINCREMENT table generates next continue + after the highest one the source ever issued, deleted rows included: such a table issues no id twice, across the + copy too, in either direction. (Other integer keys continue as the target's backend does: SQLite's rowid tables + reuse the highest id after it is deleted.) + + :param transform: Applied to each row before it is written. + :param keep_first: Tables where a row the target holds equal to another one is skipped, not refused. + :param progress: Called with each table's name before it is copied. + :return: The number of rows read, per table. + """ + wanted = set(tables if tables is not None else metadata.sorted_tables) + copied: dict[str, int] = {} + with source.begin(write=False) as src: + present = set(inspect(src).get_table_names()) + issued = _sqlite_sequences(src) if source.dialect_name == "sqlite" else _server_sequences(src) + for table in metadata.sorted_tables: + if table not in wanted or table.name not in present: + continue + if progress is not None: + progress(table.name) + statement: Insert = insert_ignore(target.dialect_name, table) if table in keep_first else insert(table) + rows = _rows(src, table) + if transform is not None: + rows = (transform(table, row) for row in rows) + copied[table.name] = 0 + for part in _parts(rows): + with target.begin(write=True) as dst: + dst.execute(statement, part) + copied[table.name] += len(part) + if table.autoincrement_column is not None: + _continue_ids(target, table, issued.get(table.name, 0)) + return copied + + +class Problem(NamedTuple): + """Rows the target would refuse, or that the copy changes.""" + + table: str + count: int + what: str + + +def find_orphans(database: Database) -> list[Problem]: + """The rows of a SQLite database whose foreign keys name no row: a server refuses to store them.""" + with database.begin(write=False) as conn: + found = conn.exec_driver_sql("PRAGMA foreign_key_check").all() + counts: dict[tuple[str, str], int] = {} + for table, _rowid, parent, _key in found: + counts[(table, parent)] = counts.get((table, parent), 0) + 1 + return [ + Problem(table, count, f"with a reference to a missing row of {parent}") + for (table, parent), count in sorted(counts.items()) + ] + + +def delete_orphans(database: Database) -> list[Problem]: + """Removes what a server would refuse from a SQLite database whose foreign keys name missing rows, until none is + left. A reference the schema clears when its row is deleted (ON DELETE SET NULL) is cleared; any other orphan row + is deleted, with what the database's own rules delete with it (ON DELETE CASCADE). For a snapshot only: never + the user's database. + + :return: What changed, per table. + """ + before = _row_counts(database) + cleared: dict[str, int] = {} + while True: + with database.begin(write=True) as conn: + found = {(row[0], row[1], row[3]) for row in conn.exec_driver_sql("PRAGMA foreign_key_check").all()} + if not found: + break + changed = 0 + for table, rowid, key in sorted(found): + name = conn.dialect.identifier_preparer.quote(table) + columns = _clearable_reference(conn, table, key) + if columns is not None: + quote = conn.dialect.identifier_preparer.quote + assignments = ", ".join(f"{quote(column)} = NULL" for column in columns) + count = conn.exec_driver_sql(f"UPDATE {name} SET {assignments} WHERE rowid = ?", (rowid,)).rowcount + cleared[table] = cleared.get(table, 0) + count + else: + count = conn.exec_driver_sql(f"DELETE FROM {name} WHERE rowid = ?", (rowid,)).rowcount + changed += count + if not changed: + raise RuntimeError(f"Could not remove the rows whose foreign keys name missing rows: {sorted(found)}") + after = _row_counts(database) + changes = [ + Problem(table, count, "with a reference to a missing row, which was cleared") + for table, count in cleared.items() + ] + changes += [ + Problem(table, count - after.get(table, 0), "left out") + for table, count in before.items() + if count > after.get(table, 0) + ] + return sorted(changes) + + +def find_oversized(database: Database, max_allowed_packet: int) -> list[Problem]: + """The rows of a SQLite database a server cannot store: text longer than its column holds there (computed + columns included, which the server computes from the copy), and rows larger than half a statement the server + takes (`max_allowed_packet`), which leaves room for the driver's escaping.""" + problems: list[Problem] = [] + with database.begin(write=False) as conn: + present = set(inspect(conn).get_table_names()) + for table in metadata.sorted_tables: + if table.name not in present: + continue + for column in table.columns: + length = _length_of(column.type) + if length is None: + continue + # A server's VARCHAR(n) counts characters, as SQLite's length() of text does. + too_long = select(func.count()).select_from(table).where(func.length(column) > length) + count = conn.execute(too_long).scalar_one() + if count: + article = "an" if column.name[0] in "aeiou" else "a" + what = f"with {article} {column.name} longer than {length} characters" + problems.append(Problem(table.name, int(count), what)) + sizes = [func.coalesce(func.length(cast(column, LargeBinary)), 0) for column in _stored_columns(table)] + too_large = ( + select(func.count()).select_from(table).where(sum(sizes[1:], sizes[0]) > max_allowed_packet // 2) + ) + count = conn.execute(too_large).scalar_one() + if count: + problems.append(Problem(table.name, int(count), "larger than half the server's max_allowed_packet")) + return problems + + +def normalize_for_a_server(table: Table, row: Row) -> Row: + """A row as a server can store it: a model config's `file_size` a whole number of bytes. (SQLite stores what a + JSON number holds; a server computes the integer `file_size` column from it and refuses a fraction.)""" + if table is models: + config = json.loads(row["config"]) + size = config.get("file_size") + if isinstance(size, float): + config["file_size"] = round(size) + return {**row, "config": json.dumps(config)} + return row + + +def count_normalized(database: Database) -> list[Problem]: + """The rows `normalize_for_a_server` changes.""" + with database.begin(write=False) as conn: + fractional = conn.execute( + select(func.count()).where(func.json_type(models.c.config, "$.file_size") == "real") + ).scalar_one() + return ( + [Problem(models.name, int(fractional), "with a file_size of a fraction of a byte, rounded")] + if (fractional) + else [] + ) + + +def copy_database( + source: Database, target: Database, progress: Optional[Callable[[str], None]] = None +) -> dict[str, int]: + """Creates the schema in the empty `target` and copies the rows of `source` into it, all but the migrator's + records: without those, the app refuses the target as incomplete. `copy_records` adds them once `verify_copy` + found the copy complete. + + :param progress: Called with each table's name before it is copied. + """ + with target.begin(write=True) as conn: + metadata.create_all(conn) + return copy_rows( + source, + target, + tables=_DATA, + transform=normalize_for_a_server, + keep_first=_MERGED_ON_A_SERVER, + progress=progress, + ) + + +def copy_records(source: Database, target: Database) -> dict[str, int]: + """Copies the migrator's records, which make the target a database the app takes.""" + return copy_rows(source, target, tables=_RECORDS) + + +def verify_copy(source: Database, target: Database, *, records: bool = False) -> list[str]: + """Compares each table `copy_database` (or with `records`, `copy_records`) wrote with the source's, row by row; + what differs, described. A table whose rows a server merges (see `_MERGED_ON_A_SERVER`) matches when the target + holds only rows of the source, and one of each set of rows the server equates.""" + tables = _RECORDS if records else _DATA + expected = _row_digests(source, tables, normalize_for_a_server) + actual = _row_digests(target, tables) + mismatches: list[str] = [] + for table in tables: + want, got = expected.get(table.name, []), actual.get(table.name, []) + if table in _MERGED_ON_A_SERVER: + if not set(got) <= set(want) or _folded_terms(target) != _folded_terms(source): + mismatches.append(f"{table.name}: rows or contents differ") + elif sorted(want) != sorted(got): + mismatches.append(f"{table.name}: rows or contents differ ({len(want)} copied, {len(got)} held)") + return mismatches + + +def merged_rows(source: Database, target: Database) -> dict[str, int]: + """How many rows fewer than the source the target holds, per table whose rows a server merges.""" + result: dict[str, int] = {} + for table in _MERGED_ON_A_SERVER: + fewer = _count(source, table) - _count(target, table) + if fewer: + result[table.name] = fewer + return result + + +def _count(database: Database, table: Table) -> int: + with database.begin(write=False) as conn: + return int(conn.execute(select(func.count()).select_from(table)).scalar_one()) + + +def _folded_terms(database: Database) -> set[str]: + # The case and compatibility forms that the servers' case-insensitive collations equate. + with database.begin(write=False) as conn: + terms = conn.execute(select(image_index_vocab_terms.c.term)).scalars().all() + # Ignorable code points (a soft hyphen, a zero-width space) are format characters, which they ignore. + return { + "".join(c for c in unicodedata.normalize("NFKC", term).casefold() if unicodedata.category(c) != "Cf") + for term in terms + } + + +def _row_digests( + database: Database, tables: Sequence[Table], transform: Optional[Transform] = None +) -> dict[str, list[bytes]]: + result: dict[str, list[bytes]] = {} + with database.begin(write=False) as conn: + present = set(inspect(conn).get_table_names()) + for table in tables: + if table.name not in present: + continue + digests: list[bytes] = [] + for row in _rows(conn, table): + stored = transform(table, row) if transform is not None else row + digests.append(hashlib.sha256(_canonical(stored)).digest()) + result[table.name] = digests + return result + + +def _row_counts(database: Database) -> dict[str, int]: + with database.begin(write=False) as conn: + present = set(inspect(conn).get_table_names()) + return { + table.name: int(conn.execute(select(func.count()).select_from(table)).scalar_one()) + for table in metadata.sorted_tables + if table.name in present + } + + +def _clearable_reference(conn: Connection, table: str, key: int) -> Optional[list[str]]: + """The columns of a SQLite foreign key that the schema sets to NULL when its row is deleted, if it does.""" + name = conn.dialect.identifier_preparer.quote(table) + parts = [row for row in conn.exec_driver_sql(f"PRAGMA foreign_key_list({name})").all() if row[0] == key] + if not parts or parts[0][6] != "SET NULL": + return None + return [row[3] for row in parts] + + +def _rows(conn: Connection, table: Table) -> Iterator[Row]: + """The table's rows, a few at a time in memory: a project document alone may be 32 MiB.""" + columns = _stored_columns(table) + statement = select(*columns).order_by(*table.primary_key.columns).execution_options(yield_per=_READ_ROWS) + for row in conn.execute(statement): + yield row._asdict() + + +def _parts(rows: Iterable[Row]) -> Iterator[list[Row]]: + """The rows in parts of at most `_BATCH_ROWS` rows and `_BATCH_BYTES`, unless one row alone is larger.""" + part: list[Row] = [] + size = 0 + for row in rows: + row_size = _size_of(row) + if part and (size + row_size > _BATCH_BYTES or len(part) == _BATCH_ROWS): + yield part + part, size = [], 0 + part.append(row) + size += row_size + if part: + yield part + + +def _size_of(row: Row) -> int: + """About the bytes a row takes in a statement: its text encoded, and its binary values escaped at worst.""" + size = 0 + for value in row.values(): + if isinstance(value, str): + size += len(value.encode()) + elif isinstance(value, (bytes, bytearray, memoryview)): + size += 2 * len(value) + return size + + +def _stored_columns(table: Table) -> list[Any]: + return [column for column in table.columns if column.computed is None] + + +def _length_of(column_type: Any) -> Optional[int]: + """The characters a server's column of this type holds, when it holds fewer than a long text.""" + if isinstance(column_type, (Key, NoCaseKey)): + return column_type.length + if isinstance(column_type, Timestamp): + return TIMESTAMP_LENGTH + return None + + +def _canonical(row: Row) -> bytes: + def value(item: Any) -> Any: + if isinstance(item, (bytes, bytearray, memoryview)): + return {"bytes": bytes(item).hex()} + if isinstance(item, bool): + return int(item) + if isinstance(item, float) and item.is_integer(): + return int(item) + return item + + return json.dumps({name: value(item) for name, item in sorted(row.items())}, default=str).encode() + + +def _server_sequences(conn: Connection) -> dict[str, int]: + """The highest id each AUTO_INCREMENT table of a MySQL or MariaDB database has issued, deleted rows included.""" + if conn.dialect.name == "mysql": + # MySQL serves these statistics from a cache that it refreshes once a day by default. + conn.exec_driver_sql("SET SESSION information_schema_stats_expiry = 0") + rows = conn.exec_driver_sql( + "SELECT TABLE_NAME, AUTO_INCREMENT FROM information_schema.TABLES " + "WHERE TABLE_SCHEMA = DATABASE() AND AUTO_INCREMENT IS NOT NULL" + ).all() + return {str(name): int(next_id) - 1 for name, next_id in rows} + + +def _sqlite_sequences(conn: Connection) -> dict[str, int]: + """The highest id each AUTOINCREMENT table has issued, deleted rows included.""" + if conn.exec_driver_sql("SELECT 1 FROM sqlite_master WHERE name = 'sqlite_sequence'").first() is None: + return {} + return {row[0]: row[1] for row in conn.exec_driver_sql("SELECT name, seq FROM sqlite_sequence").all()} + + +def _continue_ids(target: Database, table: Table, issued: int) -> None: + column = table.autoincrement_column + assert column is not None + with target.begin(write=True) as conn: + highest = max(conn.execute(select(func.max(column))).scalar() or 0, issued) + if target.dialect_name != "sqlite": + name = conn.dialect.identifier_preparer.format_table(table) + conn.exec_driver_sql(f"ALTER TABLE {name} AUTO_INCREMENT = {highest + 1}") + elif table.dialect_options["sqlite"]["autoincrement"]: + # Only an AUTOINCREMENT table never reuses an id; SQLite keeps its high-water mark here. + parameters = (highest, table.name) + if conn.exec_driver_sql("UPDATE sqlite_sequence SET seq = ? WHERE name = ?", parameters).rowcount == 0: + conn.exec_driver_sql("INSERT INTO sqlite_sequence (seq, name) VALUES (?, ?)", parameters) diff --git a/invokeai/app/services/shared/database/database.py b/invokeai/app/services/shared/database/database.py new file mode 100644 index 00000000000..ac75b61965b --- /dev/null +++ b/invokeai/app/services/shared/database/database.py @@ -0,0 +1,415 @@ +import os +import random +import sqlite3 +import threading +import time +from collections.abc import Callable, Iterator +from contextlib import AbstractContextManager, closing, contextmanager, nullcontext +from dataclasses import dataclass +from logging import Logger +from pathlib import Path +from typing import Optional, TypeVar, get_args + +from sqlalchemy import URL, Connection, Engine, make_url +from sqlalchemy.exc import DBAPIError +from sqlalchemy.exc import TimeoutError as PoolTimeoutError + +from invokeai.app.services.config.config_default import DB_SYNCHRONOUS +from invokeai.app.services.shared.database.engines import ( + SERVER_DIALECTS, + WRITE_INTENT, + create_mysql_engine, + create_sqlite_engine, +) +from invokeai.app.services.shared.database.errors import ( + ConflictError, + DatabaseUnavailableError, + NestedTransactionError, + translate_error, +) +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.session_lock import InstanceLock + +R = TypeVar("R") + +# A transaction that lost a race (a deadlock) is run again from the start this many times in all before the +# error reaches the caller. +_CONFLICT_ATTEMPTS = 3 +_CONFLICT_BASE_DELAY_SECONDS = 0.05 + +# The one server driver, whose error numbers `translate_error` reads. +_SUPPORTED_SERVER_DRIVER = "pymysql" + + +def parse_database_url(url: str) -> URL: + """The parts of a database URL; raises `ValueError` without repeating the URL, which holds a password.""" + try: + parsed = make_url(url) + except Exception: + parsed = None + # An `@` in the password that is not percent-encoded ends up in the host. + if parsed is None or "@" in (parsed.host or ""): + raise ValueError( + "The database URL cannot be read. Percent-encode the characters @ : / % # in its user name and password, " + "e.g. @ as %40." + ) + return parsed + + +@dataclass(frozen=True) +class SqliteConnection: + """The single connection of a SQLite database, the lock that serialises all work on it, and its file.""" + + conn: sqlite3.Connection + lock: threading.RLock + path: Optional[Path] + + +def _refuse_stray_transaction(sqlite: SqliteConnection) -> None: + """Names the cause when a statement left a transaction open on the raw SQLite connection. + + Python's sqlite3 begins a transaction before a data-modifying statement and keeps it open until + commit(), so a write made on the raw connection without committing would otherwise make the next + BEGIN fail with "cannot start a transaction within a transaction" -- and its rollback discard that write. + """ + if sqlite.conn.in_transaction: + raise RuntimeError( + "The SQLite connection has a transaction open that no database transaction owns: a statement ran " + "on the raw connection without being committed" + ) + + +class Database: + """The application database: owns its connections and every transaction on them. + + Queries run through `queries`, either one call per transaction (`db.queries..(...)`) + or several calls in one unit of work (`with db.queries.transaction() as q: ...`). + + SQLite is a single connection behind a process-wide re-entrant lock that is held for each whole + transaction, which serialises all database work exactly as before this layer existed. MySQL and MariaDB + use a connection pool and no process lock, so their invariants rest on the SQL itself. + + A thread may have one transaction open per database at a time; opening a second raises + `NestedTransactionError` (see there for why nesting is not joined). + """ + + def __init__(self, engine: Engine, logger: Logger, *, sqlite: Optional[SqliteConnection] = None) -> None: + self._engine = engine + self._logger = logger + self._sqlite = sqlite + # The SQLite lock serialises every transaction, so one connection object serves them all. Checking + # one out of the pool per transaction would cost about as much as the statement of a point read. + self._sqlite_connection: Optional[Connection] = engine.connect() if sqlite is not None else None + self._thread_state = threading.local() + self._disposed = False + # The instance lock of a server database (see `hold_instance_lock`). + self._instance_lock: Optional[InstanceLock] = None + self.queries = Queries(self) + + @classmethod + def open_sqlite( + cls, + db_path: Optional[Path], + logger: Logger, + *, + verbose: bool = False, + synchronous: DB_SYNCHRONOUS = "full", + ) -> "Database": + """Opens a SQLite database file, or an in-memory database when `db_path` is None. + + :param verbose: Log every SQL statement at debug level. + :param synchronous: SQLite `synchronous` setting. Defaults to `full`, what InvokeAI has always used. + `normal` is refused, with a warning, for a database file whose journal mode did not become WAL. + """ + # Called directly by the user-management commands and tests, not only through the validated + # config, and `synchronous` is interpolated into a PRAGMA below. `DB_SYNCHRONOUS` is a static + # annotation, so the value is checked here, before anything is opened. + if synchronous not in get_args(DB_SYNCHRONOUS): + raise ValueError(f"Invalid synchronous setting {synchronous!r}, expected one of {get_args(DB_SYNCHRONOUS)}") + + if db_path is None: + logger.info("Initializing in-memory database") + else: + db_path.parent.mkdir(parents=True, exist_ok=True) + logger.info(f"Initializing database at {db_path}") + + conn = sqlite3.connect(database=db_path or ":memory:", check_same_thread=False) + # Rows of the raw cursors (the migrator's); queries get plain tuples (see `_begin_sqlite`). + conn.row_factory = sqlite3.Row + + if verbose: + # On the connection rather than the engine, so it also covers the migrator's raw cursors. + conn.set_trace_callback(logger.debug) + + conn.execute("PRAGMA foreign_keys = ON;") + + # Enable Write-Ahead Logging (WAL) mode for better concurrency. The statement reports the mode + # that was actually established: WAL needs shared memory, so it can fail to engage -- quietly -- + # on network filesystems, some container volume drivers and `SQLITE_OMIT_WAL` builds. + journal_mode = str(conn.execute("PRAGMA journal_mode = WAL;").fetchone()[0]).lower() + + # Set a busy timeout to prevent database lockups during writes + conn.execute("PRAGMA busy_timeout = 5000;") # 5 seconds + + # Durability. SQLite's own default is `full`, which fsyncs on every commit; `normal` trades the + # last transactions on a power loss or OS crash for roughly 12x shorter commits. Shorter commits + # matter twice here, because every write holds the lock that serialises all database work. + # + # `normal` cannot corrupt the database only *because* of WAL; on a rollback journal it can, so a + # database file that did not get WAL keeps `full` instead. In-memory databases report `memory` + # and have no durability to trade, so the requirement does not apply to them. + # + # Refused rather than fatal: the setting is a performance preference, and starting with stronger + # durability than asked for costs commit time, where starting with weaker durability than the + # documented guarantee risks the user's database. A warning names the mode so the cause is + # visible instead of merely slow. + if synchronous == "normal" and db_path is not None and journal_mode != "wal": + logger.warning( + f"Journal mode is '{journal_mode}', not WAL, so db_synchronous='normal' was refused and " + "'full' used instead: without WAL that setting can corrupt the database on power loss." + ) + synchronous = "full" + + # PRAGMA does not take bind parameters, so the value is interpolated. It is checked against + # DB_SYNCHRONOUS at the top of this method. + conn.execute(f"PRAGMA synchronous = {synchronous.upper()};") + + return cls( + create_sqlite_engine(conn), logger, sqlite=SqliteConnection(conn=conn, lock=threading.RLock(), path=db_path) + ) + + @classmethod + def open_url(cls, url: str, logger: Logger) -> "Database": + """Opens a MySQL or MariaDB database through PyMySQL, e.g. `mariadb+pymysql://user:password@host/invokeai`.""" + parsed = parse_database_url(url) + if parsed.get_backend_name() not in SERVER_DIALECTS or parsed.get_driver_name() != _SUPPORTED_SERVER_DRIVER: + raise ValueError( + f"Unsupported database URL scheme '{parsed.drivername}', expected " + f"{' or '.join(f'{backend}+{_SUPPORTED_SERVER_DRIVER}' for backend in SERVER_DIALECTS)}" + ) + logger.info(f"Connecting to database {parsed.render_as_string(hide_password=True)}") + try: + engine = create_mysql_engine(parsed) + except ModuleNotFoundError as e: + if e.name != _SUPPORTED_SERVER_DRIVER: + raise + raise ValueError( + "The MySQL and MariaDB driver is not installed: install InvokeAI with its `mysql` extra, e.g. " + 'pip install "invokeai[mysql]"' + ) from None + return cls(engine, logger) + + @property + def dialect_name(self) -> str: + """`sqlite`, `mysql` or `mariadb`.""" + return self._engine.dialect.name + + @property + def engine(self) -> Engine: + """The SQLAlchemy engine, for instrumentation (event listeners) in tests and benchmarks, and for a + connection of its own outside any transaction, as the migrator's lock needs on a server. + + On SQLite, connecting through it raises: the one connection is held by this database. Use `begin()`. + """ + return self._engine + + @property + def logger(self) -> Logger: + return self._logger + + @property + def sqlite(self) -> SqliteConnection: + """The raw SQLite connection, its lock and its file: for the migrator, whose older migrations take a cursor. + + Raises on any other backend. + """ + if self._sqlite is None: + raise RuntimeError(f"The {self.dialect_name} database has no SQLite connection") + return self._sqlite + + @contextmanager + def begin(self, *, write: bool) -> Iterator[Connection]: + """One transaction on one connection, committed when the block exits normally and rolled back otherwise. + + For the database layer itself (query modules, the migrator, the copy tool); services run their + queries through `queries`. Driver errors with a backend-neutral meaning are raised as the matching + `DatabaseError` subclass. + + :param write: Whether the transaction may write. On SQLite a writer begins with BEGIN IMMEDIATE. + """ + self._claim_thread() + try: + with self._exclusive(): + # Translated here, around the checkout too: a server that cannot be reached fails it. + try: + with self._connection() as conn: + began_sqlite = False + try: + if self._sqlite is not None: + self._begin_sqlite(write) + began_sqlite = True + else: + conn.execution_options(**{WRITE_INTENT: write}) + with conn.begin(): + yield conn + finally: + if began_sqlite: + self._end_sqlite() + except (DBAPIError, sqlite3.Error) as error: + translated = translate_error(error, self.dialect_name) + if translated is None: + raise + raise translated from error + except PoolTimeoutError as error: + raise DatabaseUnavailableError(f"No database connection came free in time: {error}") from error + finally: + self._release_thread() + + def _begin_sqlite(self, write: bool) -> None: + """Begins the transaction that the SQLAlchemy transaction around it then commits or rolls back. + + The driver begins a transaction only before data-modifying statements: a read would see a new + snapshot per statement, and DDL would commit as it runs. So it is begun here, explicitly, and on the + driver connection: SQLAlchemy's own begin is a no-op for this driver, and issuing the BEGIN through + SQLAlchemy would cost as much again as the point read it begins. + """ + assert self._sqlite is not None + _refuse_stray_transaction(self._sqlite) + self._sqlite.conn.execute("BEGIN IMMEDIATE" if write else "BEGIN") + # SQLAlchemy builds its own rows; the driver need not build `sqlite3.Row` objects first. (A third + # of the time of a large fetch, measured.) The lock is held, so no raw cursor runs meanwhile. + self._sqlite.conn.row_factory = None + + def _end_sqlite(self) -> None: + assert self._sqlite is not None + self._sqlite.conn.row_factory = sqlite3.Row + if self._sqlite.conn.in_transaction: + # When the COMMIT itself fails, SQLAlchemy skips the driver rollback, and this connection never + # returns to a pool whose reset would roll it back instead. + self._sqlite.conn.rollback() + + def run(self, work: Callable[[Connection], R], *, write: bool) -> R: + """Runs `work` in a transaction of its own, again from the start when it loses a race (`ConflictError`). + + `work` must have no effects outside the transaction, since a retry repeats it. + """ + + def transaction() -> R: + with self.begin(write=write) as conn: + return work(conn) + + return self.retry_conflicts(transaction) + + def retry_conflicts(self, transaction: Callable[[], R]) -> R: + """Calls `transaction`, which runs one transaction, again when it loses a race (`ConflictError`): three + attempts in all, with a short random pause between them.""" + attempt = 1 + while True: + try: + return transaction() + except ConflictError as error: + if attempt >= _CONFLICT_ATTEMPTS: + raise + self._logger.warning( + f"Retrying a database transaction that lost a race " + f"(attempt {attempt + 1} of {_CONFLICT_ATTEMPTS}): {error}" + ) + time.sleep(_CONFLICT_BASE_DELAY_SECONDS * attempt * random.uniform(0.5, 1.5)) + attempt += 1 + + def clean(self) -> None: + """Reclaims the free pages of a SQLite database file, reporting the freed space. + + A no-op for in-memory and server databases. + """ + if self._sqlite is None or self._sqlite.path is None: + return + sqlite = self._sqlite + assert sqlite.path is not None + try: + # Serialize with transactions: services with their own threads (e.g. the image index worker) + # may have a statement in progress on this shared connection, and VACUUM refuses to run + # concurrently with one ("cannot VACUUM - SQL statements in progress"). + with sqlite.lock: + initial_db_size = sqlite.path.stat().st_size + sqlite.conn.execute("VACUUM;") + final_db_size = sqlite.path.stat().st_size + freed_space_in_mb = round((initial_db_size - final_db_size) / 1024 / 1024, 2) + if freed_space_in_mb > 0: + self._logger.info(f"Cleaned database (freed {freed_space_in_mb}MB)") + except Exception as e: + self._logger.error(f"Error cleaning database: {e}") + raise + + def backup(self, destination: Path) -> None: + """Copies a SQLite database into the file `destination` with SQLite's online backup, which is consistent + while the database is in use and includes what its write-ahead log holds. (A copy of the database file is + not.) A server database is backed up by its operator; this raises there. + + It holds the database's lock while it copies, so it is for startup, tools and maintenance, not for a running + app's requests. It creates the destination's directory, and refuses a destination that exists (the file is + reserved before copying, so one made meanwhile is not overwritten either) and a caller inside a transaction of + its own, whose uncommitted changes the copy would hold. A failed backup leaves no file.""" + sqlite = self.sqlite + destination = Path(destination) + destination.parent.mkdir(parents=True, exist_ok=True) + os.close(os.open(destination, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)) + try: + with sqlite.lock: + if sqlite.conn.in_transaction: + raise NestedTransactionError("A backup cannot be taken inside a transaction") + with closing(sqlite3.connect(destination)) as target: + sqlite.conn.backup(target) + except BaseException: + for path in (destination, Path(f"{destination}-wal"), Path(f"{destination}-shm")): + path.unlink(missing_ok=True) + raise + + def hold_instance_lock(self) -> None: + """Takes the lock of a server database that one process at a time holds, for as long as this database is + open; raises `DatabaseInUseError` when another process holds it. A no-op on SQLite. + + The lock belongs to the session of a connection of its own, which the server ends, and the lock with it, + when the process does. A thread checks it and takes it again should the server end that session sooner. + """ + if self._sqlite is None and self._instance_lock is None: + self._instance_lock = InstanceLock(self._engine, self._logger) + + def dispose(self) -> None: + """Closes every connection, the instance lock's included; transactions begun afterwards are refused. Safe to + call more than once. + + Closing a SQLite connection also checkpoints its write-ahead log into the database file. + """ + with self._exclusive(): + self._disposed = True + if self._instance_lock is not None: + self._instance_lock.close() + self._instance_lock = None + if self._sqlite_connection is not None: + self._sqlite_connection.close() + self._engine.dispose() + if self._sqlite is not None: + self._sqlite.conn.close() + + def _exclusive(self) -> AbstractContextManager[object]: + return self._sqlite.lock if self._sqlite is not None else nullcontext() + + def _connection(self) -> AbstractContextManager[Connection]: + """The SQLite connection object, held under the lock; on a server, a connection from the pool.""" + if self._sqlite_connection is not None: + return nullcontext(self._sqlite_connection) + return self._engine.connect() + + def _claim_thread(self) -> None: + if self._disposed: + raise RuntimeError("The database has been closed") + if getattr(self._thread_state, "open", False): + raise NestedTransactionError( + "This thread already has a transaction open on this database; " + "run the work on that transaction instead of opening another" + ) + self._thread_state.open = True + + def _release_thread(self) -> None: + self._thread_state.open = False diff --git a/invokeai/app/services/shared/database/dialect.py b/invokeai/app/services/shared/database/dialect.py new file mode 100644 index 00000000000..53caceed6ab --- /dev/null +++ b/invokeai/app/services/shared/database/dialect.py @@ -0,0 +1,431 @@ +"""SQL constructs that each backend spells differently, compiled for the backend in use.""" + +import json +import re +from collections.abc import Iterable, Sequence +from typing import Any + +from sqlalchemy import Boolean, FromClause, Insert, Integer, Join, String, Table, UniqueConstraint, literal_column +from sqlalchemy.dialects import mysql, postgresql, sqlite +from sqlalchemy.ext.compiler import compiles +from sqlalchemy.sql.compiler import SQLCompiler +from sqlalchemy.sql.elements import ColumnElement +from sqlalchemy.sql.functions import FunctionElement +from sqlalchemy.types import TypeDecorator + +from invokeai.app.services.shared.database.engines import ( + MARIADB_BINARY_COLLATION, + MYSQL_BINARY_COLLATION, + SERVER_DIALECTS, +) + +_SIMPLE_JSON_PATH = re.compile(r"\$(\.[A-Za-z_][A-Za-z0-9_]*)+") +_LIKE_ESCAPE = "\\" + + +def _require_only_primary_key(table: Table) -> None: + # MySQL and MariaDB resolve a conflict with any unique key, the other backends only one on the key named. A + # unique index of the primary key's own columns (some migrated SQLite tables have one) is that same key. + if not table.primary_key.columns: + raise ValueError(f"Table {table.name} has no primary key") + key = set(table.primary_key.columns) + unique_keys = [ + set(constraint.columns) for constraint in table.constraints if isinstance(constraint, UniqueConstraint) + ] + unique_keys += [set(index.columns) for index in table.indexes if index.unique] + if any(columns != key for columns in unique_keys): + raise ValueError(f"Table {table.name} has a unique key besides its primary key") + + +def upsert(dialect_name: str, table: Table, *, update: Sequence[str]) -> Insert: + """An INSERT into `table` that, where a row with the same primary key exists, sets that row's `update` + columns to the values it would have inserted instead. + + Execute it with the row's values, keyed by column name. A column's `onupdate` does not apply to the update: + list `updated_at` in `update` and pass its value. Every unique key of the table must be on the primary key's + columns, because MySQL and MariaDB update the row a conflict with any unique key finds, the others only on the + key named. + """ + _require_only_primary_key(table) + if dialect_name in SERVER_DIALECTS: + on_server = mysql.insert(table) + return on_server.on_duplicate_key_update({name: on_server.inserted[name] for name in update}) + if dialect_name == "sqlite": + on_sqlite = sqlite.insert(table) + return on_sqlite.on_conflict_do_update( + index_elements=list(table.primary_key.columns), + set_={name: on_sqlite.excluded[name] for name in update}, + ) + if dialect_name == "postgresql": + on_postgresql = postgresql.insert(table) + return on_postgresql.on_conflict_do_update( + index_elements=list(table.primary_key.columns), + set_={name: on_postgresql.excluded[name] for name in update}, + ) + raise ValueError(f"No upsert for the {dialect_name} dialect") + + +def insert_ignore(dialect_name: str, table: Table) -> Insert: + """An INSERT into `table` that leaves a row with the same primary key as it is instead of failing. + + It skips that conflict only: a NULL, a failed CHECK and a missing foreign key still raise. (SQLite's `INSERT + OR IGNORE` would skip the first two, MySQL's `INSERT IGNORE` all three.) Its row count does not tell whether + the row was inserted: a skipped row counts 0 on SQLite and 1 on MySQL and MariaDB, which count the rows found. + Every unique key must be on the primary key's columns, as for `upsert`. + """ + _require_only_primary_key(table) + if dialect_name in SERVER_DIALECTS: + # MySQL has no DO NOTHING; setting a key column to its own value changes nothing. + key = next(iter(table.primary_key.columns)) + return mysql.insert(table).on_duplicate_key_update({key.name: key}) + if dialect_name == "sqlite": + return sqlite.insert(table).on_conflict_do_nothing(index_elements=list(table.primary_key.columns)) + if dialect_name == "postgresql": + return postgresql.insert(table).on_conflict_do_nothing(index_elements=list(table.primary_key.columns)) + raise ValueError(f"No insert_ignore for the {dialect_name} dialect") + + +def sql_true() -> ColumnElement[bool]: + """TRUE as the predicates of SQLite's partial indexes spell it (`WHERE is_intermediate = TRUE`). + + SQLAlchemy renders `true()` as 1 on SQLite, and SQLite's planner uses a partial index only for a query term + that matches its predicate as written: `is_intermediate = 1` leaves `WHERE is_intermediate = TRUE` unused. + """ + return literal_column("TRUE", Boolean()) + + +def fixed_limit(count: int) -> ColumnElement[int]: + """A LIMIT written into the statement, for a count that is the same on every execution, such as a lookup's 1. + + SQLite runs a statement whose LIMIT is a bound parameter 10-20 µs slower, on every execution and with the same + plan, than one whose LIMIT is a literal. A count the caller chooses stays a bound parameter: each literal would + be a statement of its own. + """ + return literal_column(str(int(count)), Integer) + + +class JsonValue(ColumnElement[Any]): + """The value at a member path of a JSON document held in a text column, as SQLite's `json_extract()` gives it. + + SQL NULL for a JSON null and for a missing member, and the text of a string, alike on every backend; with + `integer`, the whole number of an integer member (a server refuses a fraction). Booleans, arrays, objects + and numbers written with an exponent or trailing zeros come back differently per backend (SQLite gives 1 + for true, a server 'true'), so only strings and integers are compared. + + For the expressions of generated columns. MySQL and MariaDB need the CASE: `JSON_UNQUOTE()` turns a JSON + null into the text 'null'. + """ + + inherit_cache = False + + def __init__(self, document: str, path: str, *, integer: bool = False) -> None: + if _SIMPLE_JSON_PATH.fullmatch(path) is None: + raise ValueError(f"Unsupported JSON path {path!r}: only member paths such as '$.meta.category' are") + self.document = document + self.path = path + self.integer = integer + + +@compiles(JsonValue, "sqlite") +def _compile_json_value_sqlite(element: JsonValue, compiler: SQLCompiler, **kw: Any) -> str: + return f"json_extract({compiler.preparer.quote(element.document)}, {_literal(compiler, element.path)})" + + +@compiles(JsonValue, "mysql") +@compiles(JsonValue, "mariadb") +def _compile_json_value_mysql(element: JsonValue, compiler: SQLCompiler, **kw: Any) -> str: + extracted = f"JSON_EXTRACT({compiler.preparer.quote(element.document)}, {_literal(compiler, element.path)})" + value = f"CASE WHEN JSON_TYPE({extracted}) = 'NULL' THEN NULL ELSE JSON_UNQUOTE({extracted}) END" + return f"CAST({value} AS SIGNED)" if element.integer else value + + +@compiles(JsonValue, "postgresql") +def _compile_json_value_postgresql(element: JsonValue, compiler: SQLCompiler, **kw: Any) -> str: + members = "{" + ",".join(element.path.split(".")[1:]) + "}" + value = f"(CAST({compiler.preparer.quote(element.document)} AS JSONB) #>> {_literal(compiler, members)})" + return f"CAST({value} AS BIGINT)" if element.integer else value + + +class JsonString(FunctionElement[str]): + """The string at a member path of a JSON document in a text column, for a query; SQL NULL when the document is + no valid JSON, the member is missing, or it holds anything but a string. + + The same on every backend: SQLite's `json_extract()` would give 1 for true where a server gives 'true'. + """ + + inherit_cache = True + name = "json_string" + + def __init__(self, document: ColumnElement[Any], path: str) -> None: + if _SIMPLE_JSON_PATH.fullmatch(path) is None: + raise ValueError(f"Unsupported JSON path {path!r}: only member paths such as '$.meta.category' are") + super().__init__(document, literal_column(f"'{path}'")) + + +@compiles(JsonString, "sqlite") +def _compile_json_string_sqlite(element: JsonString, compiler: SQLCompiler, **kw: Any) -> str: + document, path = (compiler.process(clause, **kw) for clause in element.clauses) + # Nested, so that json_type() never reads a document json_valid() rejected: it would raise. + member = f"CASE WHEN json_type({document}, {path}) = 'text' THEN json_extract({document}, {path}) END" + return f"CASE WHEN json_valid({document}) THEN {member} END" + + +@compiles(JsonString, "mysql") +@compiles(JsonString, "mariadb") +def _compile_json_string_mysql(element: JsonString, compiler: SQLCompiler, **kw: Any) -> str: + document, path = (compiler.process(clause, **kw) for clause in element.clauses) + extracted = f"JSON_EXTRACT({document}, {path})" + # JSON_TYPE() answers in utf8mb4_bin, which MySQL will not compare with a literal in the connection's collation. + member = f"CASE WHEN JSON_TYPE({extracted}) = 'STRING' COLLATE utf8mb4_bin THEN JSON_UNQUOTE({extracted}) END" + return f"CASE WHEN JSON_VALID({document}) THEN {member} END" + + +class CaseInsensitiveLike(FunctionElement[bool]): + """`expression LIKE pattern`, ignoring case, with backslash as the escape character in `pattern` (see + `like_prefix`). + + SQLite's own LIKE, which ignores the case of ASCII letters, there; LOWER() on both sides elsewhere, which + folds the case of every letter. (Lowering on SQLite would cost as much as the scan it filters.) + """ + + inherit_cache = True + type = Boolean() + name = "case_insensitive_like" + # A condition as it stands: otherwise a WHERE clause compares it with 1, which keeps SQLite from turning a + # LIKE on an index with its collation into a range of that index. + _is_implicitly_boolean = True + + def __init__(self, expression: ColumnElement[str], pattern: ColumnElement[str]) -> None: + super().__init__(expression, pattern) + + +class ContainsText(FunctionElement[bool]): + """Whether `haystack` contains `needle` as written, case and all; NULL for a NULL needle. `instr()` on SQLite, + MySQL and MariaDB, `strpos()` on Postgres: a substring test that needs no pattern escaping.""" + + inherit_cache = True + type = Boolean() + name = "contains_text" + _is_implicitly_boolean = True + + def __init__(self, haystack: ColumnElement[str], needle: ColumnElement[str]) -> None: + super().__init__(haystack, needle) + + +@compiles(ContainsText) +def _compile_contains_text(element: ContainsText, compiler: SQLCompiler, **kw: Any) -> str: + haystack, needle = (compiler.process(clause, **kw) for clause in element.clauses) + return f"(instr({haystack}, {needle}) > 0)" + + +@compiles(ContainsText, "postgresql") +def _compile_contains_text_postgresql(element: ContainsText, compiler: SQLCompiler, **kw: Any) -> str: + haystack, needle = (compiler.process(clause, **kw) for clause in element.clauses) + return f"(strpos({haystack}, {needle}) > 0)" + + +def _escape_like(text: str) -> str: + escaped = text.replace(_LIKE_ESCAPE, _LIKE_ESCAPE * 2) + return escaped.replace("%", _LIKE_ESCAPE + "%").replace("_", _LIKE_ESCAPE + "_") + + +def like_prefix(prefix: str) -> str: + """The `CaseInsensitiveLike` pattern of the values that start with `prefix`; `%`, `_` and `\\` in it match + only themselves.""" + return _escape_like(prefix) + "%" + + +def like_contains(text: str) -> str: + """The `CaseInsensitiveLike` pattern of the values that contain `text`; `%`, `_` and `\\` in it match only + themselves.""" + return "%" + _escape_like(text) + "%" + + +@compiles(CaseInsensitiveLike, "sqlite") +def _compile_case_insensitive_like_sqlite(element: CaseInsensitiveLike, compiler: SQLCompiler, **kw: Any) -> str: + expression, pattern = (compiler.process(clause, **kw) for clause in element.clauses) + return f"{expression} LIKE {pattern} ESCAPE {_literal(compiler, _LIKE_ESCAPE)}" + + +@compiles(CaseInsensitiveLike) +def _compile_case_insensitive_like(element: CaseInsensitiveLike, compiler: SQLCompiler, **kw: Any) -> str: + expression, pattern = (compiler.process(clause, **kw) for clause in element.clauses) + return f"lower({expression}) LIKE lower({pattern}) ESCAPE {_literal(compiler, _LIKE_ESCAPE)}" + + +# The longest string a bound set holds on a server: media names are at most 255 characters. +BOUND_SET_MAX_LENGTH = 255 + + +def bound_set(values: Iterable[str] | Iterable[int]) -> str: + """The parameter of an `InBoundSet`: the values as one JSON array.""" + return json.dumps(sorted(set(values))) + + +class InBoundSet(FunctionElement[bool]): + """`expression IN` the strings or integers of a set bound as one parameter (`bound_set`), strings compared as + stored, case and all. + + A set of any size is one statement and one parameter, where an IN list would bind every value and be a + statement per length: for sets held in process memory, such as the media names of active queue items, and long + lists a caller passes. Strings are at most `BOUND_SET_MAX_LENGTH` characters long. SQLite reads the array with + `json_each()`, MySQL and MariaDB with `JSON_TABLE()`, strings in the tables' binary collation; an integer + expression reads integers. + """ + + inherit_cache = True + name = "in_bound_set" + _is_implicitly_boolean = True + + def __init__(self, expression: ColumnElement[Any], values: ColumnElement[str]) -> None: + super().__init__(expression, values) + + +def _of_integers(element: InBoundSet) -> bool: + column_type = next(iter(element.clauses)).type + if isinstance(column_type, TypeDecorator): + column_type = column_type.impl_instance + return isinstance(column_type, Integer) + + +@compiles(InBoundSet, "sqlite") +def _compile_in_bound_set_sqlite(element: InBoundSet, compiler: SQLCompiler, **kw: Any) -> str: + expression, values = (compiler.process(clause, **kw) for clause in element.clauses) + return f"{expression} IN (SELECT value FROM json_each({values}))" + + +@compiles(InBoundSet, "mysql") +@compiles(InBoundSet, "mariadb") +def _compile_in_bound_set_server(element: InBoundSet, compiler: SQLCompiler, **kw: Any) -> str: + expression, values = (compiler.process(clause, **kw) for clause in element.clauses) + if _of_integers(element): + column = "value BIGINT PATH '$'" + else: + collation = MARIADB_BINARY_COLLATION if compiler.dialect.name == "mariadb" else MYSQL_BINARY_COLLATION + column = f"value VARCHAR({BOUND_SET_MAX_LENGTH}) CHARACTER SET utf8mb4 COLLATE {collation} PATH '$'" + return f"{expression} IN (SELECT bound.value FROM JSON_TABLE({values}, '$[*]' COLUMNS ({column})) AS bound)" + + +@compiles(InBoundSet, "postgresql") +def _compile_in_bound_set_postgresql(element: InBoundSet, compiler: SQLCompiler, **kw: Any) -> str: + expression, values = (compiler.process(clause, **kw) for clause in element.clauses) + value = "CAST(value AS BIGINT)" if _of_integers(element) else "value" + return f"{expression} IN (SELECT {value} FROM jsonb_array_elements_text(CAST({values} AS JSONB)) AS value)" + + +class KeysetAfter(FunctionElement[bool]): + """`(first, second) > (first_value, second_value)`: the rows after a keyset position, in that order. + + A row value on SQLite, whose planner starts an index range at it; spelled out elsewhere, since MariaDB does + not start a range at a row value (`first > a OR (first = a AND second > b)`). + """ + + inherit_cache = True + name = "keyset_after" + _is_implicitly_boolean = True + + def __init__( + self, + first: ColumnElement[Any], + second: ColumnElement[Any], + first_value: ColumnElement[Any], + second_value: ColumnElement[Any], + ) -> None: + super().__init__(first, second, first_value, second_value) + + +@compiles(KeysetAfter, "sqlite") +def _compile_keyset_after_sqlite(element: KeysetAfter, compiler: SQLCompiler, **kw: Any) -> str: + first, second, first_value, second_value = (compiler.process(clause, **kw) for clause in element.clauses) + return f"({first}, {second}) > ({first_value}, {second_value})" + + +@compiles(KeysetAfter) +def _compile_keyset_after(element: KeysetAfter, compiler: SQLCompiler, **kw: Any) -> str: + first, second, first_value, second_value = (compiler.process(clause, **kw) for clause in element.clauses) + return f"({first} > {first_value} OR ({first} = {first_value} AND {second} > {second_value}))" + + +class CaseInsensitiveOrder(FunctionElement[str]): + """`expression` as an ORDER BY key that ignores case. + + SQLite's NOCASE collation there, which folds the case of ASCII letters as its LOWER() does; LOWER() elsewhere, + which folds the case of every letter. (LOWER() on SQLite is a function call per row: a fifth of the time it + takes to order 500 models, measured.) + """ + + inherit_cache = True + name = "case_insensitive_order" + + def __init__(self, expression: ColumnElement[str]) -> None: + super().__init__(expression) + + +@compiles(CaseInsensitiveOrder, "sqlite") +def _compile_case_insensitive_order_sqlite(element: CaseInsensitiveOrder, compiler: SQLCompiler, **kw: Any) -> str: + (expression,) = (compiler.process(clause, **kw) for clause in element.clauses) + return f"{expression} COLLATE NOCASE" + + +@compiles(CaseInsensitiveOrder) +def _compile_case_insensitive_order(element: CaseInsensitiveOrder, compiler: SQLCompiler, **kw: Any) -> str: + (expression,) = (compiler.process(clause, **kw) for clause in element.clauses) + return f"lower({expression})" + + +class Unindexed(FunctionElement[Any]): + """`column` in a WHERE term that SQLite must not look up through an index: `+column` there, `column` elsewhere. + + SQLite keeps no table statistics, so it takes `queue_id = ?` (or `user_id = ?`) for a selective term and can + walk an index on it through a queue's whole history, when another term (a status, a batch, an id) finds the few + rows wanted. The unary plus keeps the term off every index. + """ + + inherit_cache = True + name = "unindexed" + + def __init__(self, column: ColumnElement[Any]) -> None: + super().__init__(column) + self.type = column.type + + +@compiles(Unindexed, "sqlite") +def _compile_unindexed_sqlite(element: Unindexed, compiler: SQLCompiler, **kw: Any) -> str: + (column,) = (compiler.process(clause, **kw) for clause in element.clauses) + return f"+{column}" + + +@compiles(Unindexed) +def _compile_unindexed(element: Unindexed, compiler: SQLCompiler, **kw: Any) -> str: + (column,) = (compiler.process(clause, **kw) for clause in element.clauses) + return column + + +class OrderedJoin(Join): + """An inner join that SQLite plans with its left side as the outer loop. + + SQLite never reorders the tables of a CROSS JOIN, so it renders as `left CROSS JOIN right ON ...` there: a + planner hint that keeps the work proportional to the left side (a board's membership, say, rather than every + image an index would offer first). Other backends get a plain JOIN and choose the order themselves. Inner only: + a CROSS JOIN has no outer form. + """ + + inherit_cache = True + + def __init__(self, left: FromClause, right: FromClause, onclause: ColumnElement[bool]) -> None: + super().__init__(left, right, onclause) + + +@compiles(OrderedJoin, "sqlite") +def _compile_ordered_join_sqlite(element: OrderedJoin, compiler: SQLCompiler, **kw: Any) -> str: + # `SQLCompiler.visit_join` with CROSS JOIN for JOIN. The ON clause tells SQLAlchemy's linter that the two sides + # are joined, as it does for a plain join. + kw.pop("asfrom", None) + assert element.onclause is not None + left = compiler.process(element.left, asfrom=True, **kw) + right = compiler.process(element.right, asfrom=True, **kw) + return f"{left} CROSS JOIN {right} ON {compiler.process(element.onclause, **kw)}" + + +def _literal(compiler: SQLCompiler, value: str) -> str: + return compiler.render_literal_value(value, String()) diff --git a/invokeai/app/services/shared/database/engines.py b/invokeai/app/services/shared/database/engines.py new file mode 100644 index 00000000000..30effb665e1 --- /dev/null +++ b/invokeai/app/services/shared/database/engines.py @@ -0,0 +1,138 @@ +"""Engine construction for each supported backend.""" + +import sqlite3 +from typing import Any + +from sqlalchemy import URL, Connection, Engine, create_engine, event +from sqlalchemy.engine import ExceptionContext +from sqlalchemy.pool import AssertionPool + +# The dialect names of the server backends. MariaDB must be reached through a `mariadb+` URL: SQLAlchemy names +# the dialect after the URL, and the schema's per-flavor DDL (collations, generated columns) follows that name. +SERVER_DIALECTS = ("mysql", "mariadb") + +# Execution option carrying a transaction's write intent from `Database.begin()` to the MySQL begin hook. +WRITE_INTENT = "invokeai_write_intent" + +# Pinned per connection so behaviour does not depend on how the server was configured. Strict mode turns +# truncation and invalid values into errors; ONLY_FULL_GROUP_BY rejects the ambiguous grouping SQLite allows. +MYSQL_SQL_MODE = ( + "STRICT_TRANS_TABLES,ONLY_FULL_GROUP_BY,NO_ZERO_IN_DATE,NO_ZERO_DATE,ERROR_FOR_DIVISION_BY_ZERO," + "NO_ENGINE_SUBSTITUTION" +) + +# Binary, NO PAD collations: they compare byte for byte and treat trailing spaces as significant, as +# SQLite's default BINARY collation does. (`utf8mb4_bin` is PAD SPACE, so 'a ' would equal 'a'.) +MYSQL_BINARY_COLLATION = "utf8mb4_0900_bin" +MARIADB_BINARY_COLLATION = "utf8mb4_nopad_bin" + +# PyMySQL's SERVER_STATUS_IN_TRANS flag: the server reports an open transaction on this connection. +_SERVER_STATUS_IN_TRANS = 0x0001 + + +def create_sqlite_engine(conn: sqlite3.Connection) -> Engine: + """An engine over exactly one existing SQLite connection. + + The SQLite database is one connection behind a process-wide lock, and `Database` holds the one + SQLAlchemy connection over it; work goes through `Database.begin()`, which takes the lock and begins + every transaction explicitly. + + The pool allows a single checkout, which that connection holds, so connecting through the engine + directly raises. A second connection object over the same driver connection would run outside the lock, + and closing it would roll back whatever transaction happened to be open on the shared connection. + """ + engine = create_engine( + "sqlite+pysqlite://", creator=lambda: conn, poolclass=AssertionPool, pool_reset_on_return=None + ) + event.listen(engine, "handle_error", _never_disconnected) + return engine + + +def _never_disconnected(context: ExceptionContext) -> None: + # SQLAlchemy takes an exception that is not an `Exception` (KeyboardInterrupt, SystemExit) raised during + # a statement for a lost connection and closes it. A local file cannot be disconnected, and closing the + # only connection would end the database -- an in-memory one with all its data. The transaction around + # the statement rolls back instead. + context.is_disconnect = False + + +def create_mysql_engine(url: URL) -> Engine: + """A pooled engine for MySQL or MariaDB. + + Reads run at the server's REPEATABLE READ, so a read that issues several statements sees one snapshot, + as SQLite's readers do. Writers switch to READ COMMITTED, which takes no gap locks and so deadlocks far + less under concurrent writes; their invariants are enforced by the statements themselves. + + Each checkout pings the server, so a connection the server dropped while idle is replaced instead of + failing the first statement. + """ + engine = create_engine( + url, + pool_size=10, + max_overflow=10, + pool_timeout=30, + pool_pre_ping=True, + # Below the servers' default `wait_timeout` of eight hours, with margin for proxies that cut sooner. + pool_recycle=1800, + # The reset hook below rolls back only when needed. + pool_reset_on_return=None, + isolation_level="REPEATABLE READ", + connect_args={"charset": "utf8mb4"}, + ) + # First, so the session settings are in place before SQLAlchemy inspects the connection (it reads + # sql_mode, e.g. ANSI_QUOTES, to decide how to quote identifiers). + event.listen(engine, "connect", _configure_mysql_session, insert=True) + expects_mariadb = engine.dialect.name == "mariadb" + + def set_collation(dbapi_connection: Any, connection_record: object) -> None: + _set_mysql_collation(dbapi_connection, expects_mariadb) + + # Last, because SQLAlchemy's own connection setup sends a plain `SET NAMES`, resetting the collation. + event.listen(engine, "connect", set_collation) + event.listen(engine, "begin", _begin_mysql_transaction) + event.listen(engine, "reset", _reset_mysql_connection) + return engine + + +def _configure_mysql_session(dbapi_connection: Any, connection_record: object) -> None: + cursor = dbapi_connection.cursor() + try: + cursor.execute(f"SET SESSION sql_mode = '{MYSQL_SQL_MODE}'") + cursor.execute("SET SESSION time_zone = '+00:00'") + cursor.execute("SET SESSION innodb_lock_wait_timeout = 10") + # Metadata locks (DDL waiting for open transactions) default to waiting a year. + cursor.execute("SET SESSION lock_wait_timeout = 60") + finally: + cursor.close() + + +def _set_mysql_collation(dbapi_connection: Any, expects_mariadb: bool) -> None: + is_mariadb = "mariadb" in dbapi_connection.get_server_info().lower() + if is_mariadb and not expects_mariadb: + # (The reverse, a mariadb+ URL for MySQL, SQLAlchemy refuses itself.) + raise ValueError( + "The database server is MariaDB, but the URL names MySQL: use a mariadb+pymysql:// URL, so that " + "tables are created with MariaDB's collations and column definitions" + ) + # How literals and parameters compare; columns compare by their table's collation. + collation = MARIADB_BINARY_COLLATION if is_mariadb else MYSQL_BINARY_COLLATION + cursor = dbapi_connection.cursor() + try: + cursor.execute(f"SET NAMES utf8mb4 COLLATE {collation}") + finally: + cursor.close() + + +def _begin_mysql_transaction(connection: Connection) -> None: + if connection.get_execution_options().get(WRITE_INTENT, True): + # Applies to the next transaction only, which is the one beginning now. + connection.exec_driver_sql("SET TRANSACTION ISOLATION LEVEL READ COMMITTED") + + +def _reset_mysql_connection(dbapi_connection: Any, connection_record: object, reset_state: Any) -> None: + # The pool's own reset sends ROLLBACK on every return: one more round trip per call. A connection comes + # back with a transaction open only after something went wrong, so roll back only then. + if reset_state.terminate_only: + return + if getattr(dbapi_connection, "server_status", _SERVER_STATUS_IN_TRANS) & _SERVER_STATUS_IN_TRANS: + dbapi_connection.rollback() diff --git a/invokeai/app/services/shared/database/errors.py b/invokeai/app/services/shared/database/errors.py new file mode 100644 index 00000000000..687354856c6 --- /dev/null +++ b/invokeai/app/services/shared/database/errors.py @@ -0,0 +1,154 @@ +"""Database errors that do not depend on the backend. + +The database layer raises these instead of driver exceptions, so a service can tell a duplicate key from +a missing reference without knowing which database it runs on. Errors are classified by the driver's +error code, not by message text: messages differ between backends and change between versions. The one +exception is a code SQLite shares between two causes, told apart by SQLite's fixed message. +""" + +import sqlite3 +from typing import Optional + +from sqlalchemy.exc import DBAPIError + + +class DatabaseError(Exception): + """Base class of the errors raised by the database layer.""" + + +class IntegrityViolation(DatabaseError): + """A constraint rejected the statement. + + Raised inside a transaction, it ends that transaction: catch it there only to raise an error of your + own (see `TransactionFailedError`). + """ + + +class UniqueViolation(IntegrityViolation): + """A primary key or unique constraint already holds the value.""" + + +class ForeignKeyViolation(IntegrityViolation): + """A foreign key names a missing row, or a restricting reference blocks the delete.""" + + +class CheckViolation(IntegrityViolation): + """A CHECK constraint rejected the row.""" + + +class NotNullViolation(IntegrityViolation): + """A NOT NULL column received no value.""" + + +class TransientDatabaseError(DatabaseError): + """Concurrent work kept the transaction from completing; it is rolled back, and may be run again.""" + + +class ConflictError(TransientDatabaseError): + """The transaction lost a race: a deadlock, or a snapshot that another writer invalidated. + + Running it again from the start usually succeeds; a single query call is retried automatically. + """ + + +class LockTimeoutError(TransientDatabaseError): + """The transaction gave up waiting for a lock that other work held for too long. + + Not retried automatically: it has already waited, and waiting again would also hold up everything + queued behind it. + """ + + +class DatabaseUnavailableError(TransientDatabaseError): + """No connection to the database server could be made, or none came free in time. Not retried automatically. + + A connection lost during a statement is not this: the server also closes the connection of a statement larger + than it takes, which no retry would change. + """ + + +class NestedTransactionError(DatabaseError): + """A transaction was opened on a thread that already has one open on the same database. + + Nesting is refused rather than joined: an inner commit would commit the outer transaction's work + early, and on a pooled backend the inner transaction would run on a second connection that neither + sees the outer one's writes nor can wait for its locks. Use the outer transaction's queries instead. + """ + + +class DatabaseInUseError(DatabaseError): + """Another InvokeAI process uses this server database. + + One process at a time serves a database: at startup a process cancels the queue items left running, and + syncs its bundled workflows and style presets, which would undo or collide with another process's work. + """ + + +class ReadOnlyTransactionError(DatabaseError): + """A write was attempted inside a read-only transaction.""" + + +class TransactionFailedError(DatabaseError): + """The transaction was used after one of its statements failed, or left without raising. + + Backends disagree on what a failed statement leaves behind -- SQLite and MySQL undo just that statement + (MySQL a whole deadlocked transaction), Postgres refuses everything after it -- so no transaction + continues past one: it takes no further calls and rolls back instead of committing. + """ + + +# Extended result codes, which Python's sqlite3 reports in `sqlite_errorcode`. +_SQLITE_CODES: dict[int, type[DatabaseError]] = { + 1555: UniqueViolation, # SQLITE_CONSTRAINT_PRIMARYKEY + 2067: UniqueViolation, # SQLITE_CONSTRAINT_UNIQUE + 787: ForeignKeyViolation, # SQLITE_CONSTRAINT_FOREIGNKEY + 275: CheckViolation, # SQLITE_CONSTRAINT_CHECK + 1299: NotNullViolation, # SQLITE_CONSTRAINT_NOTNULL + 517: ConflictError, # SQLITE_BUSY_SNAPSHOT: a read snapshot cannot become a write after another commit +} +# Primary result codes, matched after the extended ones; their other extended variants share the low byte. +_SQLITE_PRIMARY_CODES: dict[int, type[DatabaseError]] = { + 5: LockTimeoutError, # SQLITE_BUSY, after the busy timeout ran out + 6: LockTimeoutError, # SQLITE_LOCKED +} +# SQLite enforces ON DELETE RESTRICT through a trigger program, so a violation carries the code of a trigger's +# RAISE(); only SQLite's fixed foreign key message tells the two apart. +_SQLITE_CONSTRAINT_TRIGGER = 1811 +_SQLITE_FOREIGN_KEY_MESSAGE = "FOREIGN KEY constraint failed" + +# MySQL and MariaDB server error numbers, the first argument of the PyMySQL exception. +_MYSQL_ERRORS: dict[int, type[DatabaseError]] = { + 1062: UniqueViolation, # ER_DUP_ENTRY + 1216: ForeignKeyViolation, # ER_NO_REFERENCED_ROW + 1217: ForeignKeyViolation, # ER_ROW_IS_REFERENCED + 1451: ForeignKeyViolation, # ER_ROW_IS_REFERENCED_2 + 1452: ForeignKeyViolation, # ER_NO_REFERENCED_ROW_2 + 3819: CheckViolation, # ER_CHECK_CONSTRAINT_VIOLATED (MySQL) + 4025: CheckViolation, # ER_CONSTRAINT_FAILED (MariaDB) + 1048: NotNullViolation, # ER_BAD_NULL_ERROR + 1364: NotNullViolation, # ER_NO_DEFAULT_FOR_FIELD + 1213: ConflictError, # ER_LOCK_DEADLOCK + 1205: LockTimeoutError, # ER_LOCK_WAIT_TIMEOUT + 2003: DatabaseUnavailableError, # CR_CONN_HOST_ERROR +} + + +def translate_error(error: BaseException, dialect_name: str) -> Optional[DatabaseError]: + """The backend-neutral error for a driver error, or None when it has no backend-neutral meaning. + + :param error: A SQLAlchemy `DBAPIError`, or the driver's own exception. + """ + original = error.orig if isinstance(error, DBAPIError) else error + if dialect_name == "sqlite" and isinstance(original, sqlite3.Error): + code = getattr(original, "sqlite_errorcode", None) + if code is None: + return None + if code == _SQLITE_CONSTRAINT_TRIGGER and str(original) == _SQLITE_FOREIGN_KEY_MESSAGE: + return ForeignKeyViolation(str(original)) + neutral = _SQLITE_CODES.get(code) or _SQLITE_PRIMARY_CODES.get(code & 0xFF) + return neutral(str(original)) if neutral is not None else None + if dialect_name in ("mysql", "mariadb"): + args = getattr(original, "args", ()) + if args and isinstance(args[0], int) and args[0] in _MYSQL_ERRORS: + return _MYSQL_ERRORS[args[0]](str(original)) + return None diff --git a/invokeai/app/services/shared/database/queries/__init__.py b/invokeai/app/services/shared/database/queries/__init__.py new file mode 100644 index 00000000000..e6aca788ece --- /dev/null +++ b/invokeai/app/services/shared/database/queries/__init__.py @@ -0,0 +1,195 @@ +"""Every query of the application, one module per domain. + +A query module subclasses `QueryModule` (see `base`). Each of its methods takes the connection right after +`self` and is decorated with `read` or `write`, which supply it, so callers never see a connection: + + user = db.queries.users.get(user_id) # a transaction of its own + with db.queries.transaction() as q: # one transaction for several calls + q.client_state.set(user_id, key, value) + q.media_references.replace(...) + +A query method returns plain values, or rows read completely that `mapped` turns into DTOs once its own +transaction has ended. It does not call other query modules -- composition belongs to the caller's +transaction -- and has no effects outside the database, because a call that loses a race is run again from +the start. It passes execution options per statement (`conn.execute(statement, execution_options=...)`), never +to the connection: on SQLite every transaction shares one connection object, so an option set on it would +stay for all later transactions. +""" + +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from functools import cached_property +from typing import TYPE_CHECKING, Optional, Self, TypeVar + +from invokeai.app.services.shared.database.errors import NestedTransactionError, TransactionFailedError +from invokeai.app.services.shared.database.queries.app_settings import AppSettingQueries +from invokeai.app.services.shared.database.queries.base import OwnTransaction, QueryScope, SharedTransaction +from invokeai.app.services.shared.database.queries.board_images import BoardImageQueries +from invokeai.app.services.shared.database.queries.board_videos import BoardVideoQueries +from invokeai.app.services.shared.database.queries.boards import BoardQueries +from invokeai.app.services.shared.database.queries.client_state import ClientStateQueries +from invokeai.app.services.shared.database.queries.fonts import FontQueries +from invokeai.app.services.shared.database.queries.gallery import GalleryQueries +from invokeai.app.services.shared.database.queries.image_index import ImageIndexQueries +from invokeai.app.services.shared.database.queries.image_moves import ImageMoveQueries +from invokeai.app.services.shared.database.queries.images import ImageQueries +from invokeai.app.services.shared.database.queries.intermediates import IntermediateQueries +from invokeai.app.services.shared.database.queries.locks import LockQueries +from invokeai.app.services.shared.database.queries.media_references import MediaReferenceQueries +from invokeai.app.services.shared.database.queries.model_relationships import ModelRelationshipQueries +from invokeai.app.services.shared.database.queries.models import ModelQueries +from invokeai.app.services.shared.database.queries.projects import ProjectQueries +from invokeai.app.services.shared.database.queries.session_queue import SessionQueueQueries +from invokeai.app.services.shared.database.queries.style_presets import StylePresetQueries +from invokeai.app.services.shared.database.queries.system_prompts import SystemPromptQueries +from invokeai.app.services.shared.database.queries.users import UserQueries +from invokeai.app.services.shared.database.queries.videos import VideoQueries +from invokeai.app.services.shared.database.queries.wildcards import WildcardQueries +from invokeai.app.services.shared.database.queries.workflows import WorkflowQueries + +if TYPE_CHECKING: + from invokeai.app.services.shared.database.database import Database + +R = TypeVar("R") + + +class Queries: + """Every query of the application, grouped by domain as `queries..(...)`. + + `Database.queries` runs each call in a transaction of its own. The queries yielded by `transaction()` + share one transaction, which commits when the block exits. A domain's query module is created on first + use, once per `Queries`. + """ + + def __init__(self, database: "Database", scope: Optional[QueryScope] = None) -> None: + self._database = database + self._scope: QueryScope = scope if scope is not None else OwnTransaction(database) + + @cached_property + def app_settings(self) -> AppSettingQueries: + return AppSettingQueries(self._scope) + + @cached_property + def board_images(self) -> BoardImageQueries: + return BoardImageQueries(self._scope) + + @cached_property + def board_videos(self) -> BoardVideoQueries: + return BoardVideoQueries(self._scope) + + @cached_property + def boards(self) -> BoardQueries: + return BoardQueries(self._scope) + + @cached_property + def client_state(self) -> ClientStateQueries: + return ClientStateQueries(self._scope) + + @cached_property + def fonts(self) -> FontQueries: + return FontQueries(self._scope) + + @cached_property + def gallery(self) -> GalleryQueries: + return GalleryQueries(self._scope) + + @cached_property + def image_index(self) -> ImageIndexQueries: + return ImageIndexQueries(self._scope) + + @cached_property + def image_moves(self) -> ImageMoveQueries: + return ImageMoveQueries(self._scope) + + @cached_property + def images(self) -> ImageQueries: + return ImageQueries(self._scope) + + @cached_property + def intermediates(self) -> IntermediateQueries: + return IntermediateQueries(self._scope) + + @cached_property + def locks(self) -> LockQueries: + return LockQueries(self._scope) + + @cached_property + def media_references(self) -> MediaReferenceQueries: + return MediaReferenceQueries(self._scope) + + @cached_property + def model_relationships(self) -> ModelRelationshipQueries: + return ModelRelationshipQueries(self._scope) + + @cached_property + def models(self) -> ModelQueries: + return ModelQueries(self._scope) + + @cached_property + def projects(self) -> ProjectQueries: + return ProjectQueries(self._scope) + + @cached_property + def session_queue(self) -> SessionQueueQueries: + return SessionQueueQueries(self._scope) + + @cached_property + def style_presets(self) -> StylePresetQueries: + return StylePresetQueries(self._scope) + + @cached_property + def system_prompts(self) -> SystemPromptQueries: + return SystemPromptQueries(self._scope) + + @cached_property + def users(self) -> UserQueries: + return UserQueries(self._scope) + + @cached_property + def videos(self) -> VideoQueries: + return VideoQueries(self._scope) + + @cached_property + def wildcards(self) -> WildcardQueries: + return WildcardQueries(self._scope) + + @cached_property + def workflows(self) -> WorkflowQueries: + return WorkflowQueries(self._scope) + + @contextmanager + def transaction(self, *, read_only: bool = False) -> Iterator[Self]: + """Runs every call on the yielded queries in one transaction, committed when the block exits normally. + + A read-only transaction sees one consistent snapshot and refuses writes. The block is not retried + (`run()` retries work that has no effects outside the database). After a statement failed, the + transaction takes no further calls and does not commit: leave the block by raising (a caught + `DatabaseError` can be turned into a domain error there). + """ + if isinstance(self._scope, SharedTransaction): + raise NestedTransactionError("These queries already belong to a transaction") + with self._database.begin(write=not read_only) as conn: + scope = SharedTransaction(conn, writable=not read_only, dialect_name=self._database.dialect_name) + try: + yield type(self)(self._database, scope) + if scope.failed: + raise TransactionFailedError( + "A statement of this transaction failed and the error was not raised further; " + "the transaction was rolled back" + ) + finally: + scope.close() + + def run(self, work: Callable[[Self], R], *, read_only: bool = False) -> R: + """Runs `work` with the queries of one transaction, as `transaction()` does, and again from the start + when the transaction loses a race (`ConflictError`). + + On MySQL and MariaDB, transactions that write the same rows can deadlock, and the server aborts one of + them. `work` must have no effects outside the database, since a retry repeats it. + """ + + def transaction() -> R: + with self.transaction(read_only=read_only) as q: + return work(q) + + return self._database.retry_conflicts(transaction) diff --git a/invokeai/app/services/shared/database/queries/app_settings.py b/invokeai/app/services/shared/database/queries/app_settings.py new file mode 100644 index 00000000000..5edff149111 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/app_settings.py @@ -0,0 +1,16 @@ +"""Server-wide settings kept in the database.""" + +from typing import Optional + +from sqlalchemy import Connection, bindparam, select + +from invokeai.app.services.shared.database.queries.base import QueryModule, read +from invokeai.app.services.shared.database.schema.app_settings import app_settings + +_GET = select(app_settings.c.value).where(app_settings.c.key == bindparam("key")) + + +class AppSettingQueries(QueryModule): + @read + def get(self, conn: Connection, key: str) -> Optional[str]: + return conn.execute(_GET, {"key": key}).scalar_one_or_none() diff --git a/invokeai/app/services/shared/database/queries/base.py b/invokeai/app/services/shared/database/queries/base.py new file mode 100644 index 00000000000..7b99ef957ae --- /dev/null +++ b/invokeai/app/services/shared/database/queries/base.py @@ -0,0 +1,155 @@ +"""What every query module is built from: `QueryModule`, the `read`, `write` and `mapped` decorators, and their scopes.""" + +import functools +from collections.abc import Callable +from datetime import date, timedelta +from typing import TYPE_CHECKING, Concatenate, ParamSpec, Protocol, TypeVar + +from sqlalchemy import Connection +from sqlalchemy.exc import DBAPIError + +from invokeai.app.services.shared.database.errors import ( + ReadOnlyTransactionError, + TransactionFailedError, + translate_error, +) + +if TYPE_CHECKING: + from invokeai.app.services.shared.database.database import Database + +P = ParamSpec("P") +R = TypeVar("R") +T = TypeVar("T") +M = TypeVar("M", bound="QueryModule") + +# The most values one `IN (...)` list binds; longer lists are queried in chunks of this size. SQLite builds +# before 3.32 accept at most 999 parameters per statement. +IN_CHUNK = 500 + + +class QueryScope(Protocol): + """Where the calls of a query module run: in transactions of their own, or in a shared one.""" + + def run(self, work: Callable[[Connection], R], *, write: bool) -> R: ... + + +class OwnTransaction: + """Each call opens and commits a transaction of its own, retried when it loses a race.""" + + def __init__(self, database: "Database") -> None: + self._database = database + + def run(self, work: Callable[[Connection], R], *, write: bool) -> R: + return self._database.run(work, write=write) + + +class SharedTransaction: + """Calls run on the connection of an open transaction.""" + + def __init__(self, conn: Connection, *, writable: bool, dialect_name: str) -> None: + self._conn = conn + self._writable = writable + self._dialect_name = dialect_name + self._open = True + self.failed = False + # The calls run in this transaction so far, the running one included. + self.calls = 0 + + def run(self, work: Callable[[Connection], R], *, write: bool) -> R: + self.calls += 1 + if not self._open: + raise RuntimeError("These queries belong to a transaction that has already ended") + if self.failed: + raise TransactionFailedError("A statement of this transaction failed; leave the transaction") + if write and not self._writable: + raise ReadOnlyTransactionError("A read-only transaction cannot write") + try: + return work(self._conn) + except DBAPIError as error: + # Backends disagree on what a failed statement leaves behind: SQLite and MySQL undo just the + # statement (MySQL undoes a whole deadlocked transaction), Postgres refuses the rest of the + # transaction. So none continues; and the error is translated here, not only at the end, so a + # caller inside the transaction can catch it to raise its own. + self.failed = True + translated = translate_error(error, self._dialect_name) + if translated is None: + raise + raise translated from error + + def close(self) -> None: + self._open = False + + +class QueryModule: + """Base of the per-domain query modules. + + Build each statement once, at module level, with `bindparam()` for its values, and execute it with + them. Building a statement costs several times what executing it does: SQLAlchemy compiles it once, + but derives its cache key on every execution of a new statement object. + """ + + def __init__(self, scope: QueryScope) -> None: + self._scope = scope + + +def read(method: Callable[Concatenate[M, Connection, P], R]) -> Callable[Concatenate[M, P], R]: + """Marks a query method that only reads. Run on its own, it gets a read transaction (on SQLite, BEGIN).""" + + @functools.wraps(method) + def call(self: M, /, *args: P.args, **kwargs: P.kwargs) -> R: + return self._scope.run(lambda conn: method(self, conn, *args, **kwargs), write=False) + + return call + + +def write(method: Callable[Concatenate[M, Connection, P], R]) -> Callable[Concatenate[M, P], R]: + """Marks a query method that writes. Run on its own, it gets a write transaction (on SQLite, BEGIN IMMEDIATE).""" + + @functools.wraps(method) + def call(self: M, /, *args: P.args, **kwargs: P.kwargs) -> R: + return self._scope.run(lambda conn: method(self, conn, *args, **kwargs), write=True) + + return call + + +def locking(method: Callable[Concatenate[M, Connection, P], R]) -> Callable[Concatenate[M, P], R]: + """Marks a query method that locks rows until its transaction ends (`SELECT ... FOR UPDATE`; a SQLite transaction + excludes every other one already). It runs as a write, and only in a `transaction()`: in a transaction of its + own, the lock would be released as soon as it was taken.""" + + @functools.wraps(method) + def call(self: M, /, *args: P.args, **kwargs: P.kwargs) -> R: + if not isinstance(self._scope, SharedTransaction): + raise RuntimeError(f"{method.__qualname__} locks rows, which only a transaction() holds") + return self._scope.run(lambda conn: method(self, conn, *args, **kwargs), write=True) + + return call + + +def mapped(mapper: Callable[[T], R]) -> Callable[[Callable[Concatenate[M, P], T]], Callable[Concatenate[M, P], R]]: + """Applies `mapper` to what a `read` or `write` method returns, outside its transaction when it has one of its own. + + For building DTOs from rows: on SQLite a transaction holds the database's process-wide lock, and validating a + DTO costs more than reading its row. Put it above `@read` or `@write`, on a method that returns its rows read + completely (`first()`, `all()`): they outlive the transaction. + """ + + def decorate(method: Callable[Concatenate[M, P], T]) -> Callable[Concatenate[M, P], R]: + @functools.wraps(method) + def call(self: M, /, *args: P.args, **kwargs: P.kwargs) -> R: + return mapper(method(self, *args, **kwargs)) + + return call + + return decorate + + +def day_after(day: str) -> str: + """The day after an ISO day (`YYYY-MM-DD`), as text that timestamps of that day sort before. An invalid day, or + the last one there is, gives the empty text, which no timestamp sorts before: SQLite's DATE() gave NULL for both, + which matched nothing.""" + try: + parsed = date.fromisoformat(day) + return (parsed + timedelta(days=1)).isoformat() if parsed.isoformat() == day else "" + except (ValueError, OverflowError): + return "" diff --git a/invokeai/app/services/shared/database/queries/board_access.py b/invokeai/app/services/shared/database/queries/board_access.py new file mode 100644 index 00000000000..42271800b69 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/board_access.py @@ -0,0 +1,17 @@ +"""Which boards an account may read: one rule for every query that lists boards or media through them.""" + +from sqlalchemy import ColumnElement, exists, literal, or_ + +from invokeai.app.services.board_records.board_records_common import BoardVisibility +from invokeai.app.services.shared.database.schema.boards import boards, shared_boards + +_SHARED_VISIBILITIES = (literal(BoardVisibility.Shared.value), literal(BoardVisibility.Public.value)) + + +def readable_board(user_id: ColumnElement[str]) -> ColumnElement[bool]: + """Boards an account that is not an administrator may read: its own, those shared or public, and those shared + with it. Archived boards are not excluded here: listings decide that themselves.""" + shared_with_account = exists().where( + shared_boards.c.board_id == boards.c.board_id, shared_boards.c.user_id == user_id + ) + return or_(boards.c.user_id == user_id, boards.c.board_visibility.in_(_SHARED_VISIBILITIES), shared_with_account) diff --git a/invokeai/app/services/shared/database/queries/board_images.py b/invokeai/app/services/shared/database/queries/board_images.py new file mode 100644 index 00000000000..cb206dc2efe --- /dev/null +++ b/invokeai/app/services/shared/database/queries/board_images.py @@ -0,0 +1,109 @@ +"""Which board each image is on: at most one, since the image name is the key.""" + +import functools +from collections.abc import Sequence +from typing import Any, Optional + +from sqlalchemy import Connection, Insert, Select, bindparam, case, delete, false, func, literal, select + +from invokeai.app.services.image_records.image_records_common import ( + ASSETS_CATEGORIES, + IMAGE_CATEGORIES, + ImageCategory, +) +from invokeai.app.services.shared.database.dialect import OrderedJoin, upsert +from invokeai.app.services.shared.database.queries.base import QueryModule, read, write +from invokeai.app.services.shared.database.schema.boards import board_images +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.types import now_text + +_BOARD_OF = select(board_images.c.board_id).where(board_images.c.image_name == bindparam("image_name")) +_REMOVE = delete(board_images).where( + board_images.c.image_name == bindparam("image_name"), board_images.c.board_id == bindparam("board_id") +) + + +def _in(categories: Sequence[ImageCategory | str]) -> list[Any]: + # Bound one by one: a list would be an expanding parameter, rendered anew on every execution. + return [literal(value) for value in sorted({ImageCategory(category).value for category in categories})] + + +def _counted(categories: list[ImageCategory]) -> Any: + return func.count(case((images.c.image_category.in_(_in(categories)), 1))) + + +# Both counts in one pass. From the board's membership to its images, not the other way round: an index on +# `images` would otherwise make every image of the gallery the outer loop of a board's count. +_COUNTS = ( + select(_counted(IMAGE_CATEGORIES), _counted(ASSETS_CATEGORIES)) + .select_from(OrderedJoin(board_images, images, board_images.c.image_name == images.c.image_name)) + .where(images.c.is_intermediate == false(), board_images.c.board_id == bindparam("board_id")) +) + + +@functools.cache +def _add(dialect_name: str) -> Insert: + return upsert(dialect_name, board_images, update=("board_id", "updated_at")) + + +@functools.cache +def _names(off_board: bool, categories: Optional[tuple[str, ...]], by_intermediate: bool, by_user: bool) -> Select[Any]: + statement = select(images.c.image_name).select_from( + images.outerjoin(board_images, board_images.c.image_name == images.c.image_name) + ) + if off_board: + statement = statement.where(board_images.c.board_id.is_(None)) + else: + statement = statement.where(board_images.c.board_id == bindparam("board_id")) + if categories is not None: + statement = statement.where(images.c.image_category.in_(_in(categories))) + if by_intermediate: + statement = statement.where(images.c.is_intermediate == bindparam("is_intermediate")) + if by_user: + statement = statement.where(images.c.user_id == bindparam("user_id")) + return statement + + +class BoardImageQueries(QueryModule): + @write + def add(self, conn: Connection, board_id: str, image_name: str) -> None: + """Puts the image on the board, taking it off the board it was on.""" + conn.execute( + _add(conn.dialect.name), {"board_id": board_id, "image_name": image_name, "updated_at": now_text()} + ) + + @write + def remove(self, conn: Connection, image_name: str, board_id: str) -> bool: + """Takes the image off this board; whether it was on it.""" + return conn.execute(_REMOVE, {"image_name": image_name, "board_id": board_id}).rowcount > 0 + + @read + def board_of(self, conn: Connection, image_name: str) -> Optional[str]: + return conn.execute(_BOARD_OF, {"image_name": image_name}).scalar_one_or_none() + + @read + def image_names( + self, + conn: Connection, + board_id: Optional[str], + *, + categories: Optional[list[ImageCategory]], + is_intermediate: Optional[bool], + user_id: Optional[str], + ) -> list[str]: + """The names of the images on a board, or on none (`board_id` None), that match the filters given.""" + statement = _names( + board_id is None, + None if categories is None else tuple(sorted({category.value for category in categories})), + is_intermediate is not None, + user_id is not None, + ) + parameters = {"board_id": board_id, "is_intermediate": is_intermediate, "user_id": user_id} + return list(conn.execute(statement, parameters).scalars().all()) + + @read + def counts(self, conn: Connection, board_id: str) -> tuple[int, int]: + """The board's images that are not intermediates: those of the image categories, and those of the asset + categories.""" + counts = conn.execute(_COUNTS, {"board_id": board_id}).one() + return counts[0], counts[1] diff --git a/invokeai/app/services/shared/database/queries/board_videos.py b/invokeai/app/services/shared/database/queries/board_videos.py new file mode 100644 index 00000000000..06516d32ccf --- /dev/null +++ b/invokeai/app/services/shared/database/queries/board_videos.py @@ -0,0 +1,90 @@ +"""Which board each video is on: at most one, since the video name is the key.""" + +import functools +from typing import Any, Optional + +from sqlalchemy import Connection, Insert, Select, bindparam, case, delete, false, func, literal, select + +from invokeai.app.services.image_records.image_records_common import ImageCategory +from invokeai.app.services.shared.database.dialect import upsert +from invokeai.app.services.shared.database.queries.base import QueryModule, read, write +from invokeai.app.services.shared.database.schema.boards import board_videos +from invokeai.app.services.shared.database.schema.videos import videos +from invokeai.app.services.shared.database.types import now_text + +_BOARD_OF = select(board_videos.c.board_id).where(board_videos.c.video_name == bindparam("video_name")) +_REMOVE = delete(board_videos).where(board_videos.c.video_name == bindparam("video_name")) +# Both counts in one pass: every video, and those of a category other than general (assets). +_COUNTS = ( + select(func.count(), func.count(case((videos.c.video_category != literal(ImageCategory.GENERAL.value), 1)))) + .select_from(board_videos.join(videos, board_videos.c.video_name == videos.c.video_name)) + .where(videos.c.is_intermediate == false(), board_videos.c.board_id == bindparam("board_id")) +) + + +@functools.cache +def _add(dialect_name: str) -> Insert: + return upsert(dialect_name, board_videos, update=("board_id", "updated_at")) + + +@functools.cache +def _names(off_board: bool, categories: Optional[tuple[str, ...]], by_intermediate: bool, by_user: bool) -> Select[Any]: + statement = select(videos.c.video_name).select_from( + videos.outerjoin(board_videos, board_videos.c.video_name == videos.c.video_name) + ) + if off_board: + statement = statement.where(board_videos.c.board_id.is_(None)) + else: + statement = statement.where(board_videos.c.board_id == bindparam("board_id")) + if categories is not None: + # Bound one by one: a list would be an expanding parameter, rendered anew on every execution. + statement = statement.where(videos.c.video_category.in_([literal(value) for value in categories])) + if by_intermediate: + statement = statement.where(videos.c.is_intermediate == bindparam("is_intermediate")) + if by_user: + statement = statement.where(videos.c.user_id == bindparam("user_id")) + return statement + + +class BoardVideoQueries(QueryModule): + @write + def add(self, conn: Connection, board_id: str, video_name: str) -> None: + """Puts the video on the board, taking it off the board it was on.""" + conn.execute( + _add(conn.dialect.name), {"board_id": board_id, "video_name": video_name, "updated_at": now_text()} + ) + + @write + def remove(self, conn: Connection, video_name: str) -> None: + """Takes the video off whichever board it is on.""" + conn.execute(_REMOVE, {"video_name": video_name}) + + @read + def board_of(self, conn: Connection, video_name: str) -> Optional[str]: + return conn.execute(_BOARD_OF, {"video_name": video_name}).scalar_one_or_none() + + @read + def video_names( + self, + conn: Connection, + board_id: Optional[str], + *, + categories: Optional[list[ImageCategory]], + is_intermediate: Optional[bool], + user_id: Optional[str], + ) -> list[str]: + """The names of the videos on a board, or on none (`board_id` None), that match the filters given.""" + statement = _names( + board_id is None, + None if categories is None else tuple(sorted({category.value for category in categories})), + is_intermediate is not None, + user_id is not None, + ) + parameters = {"board_id": board_id, "is_intermediate": is_intermediate, "user_id": user_id} + return list(conn.execute(statement, parameters).scalars().all()) + + @read + def counts(self, conn: Connection, board_id: str) -> tuple[int, int]: + """The board's videos that are not intermediates: all of them, and those of a category other than general.""" + counts = conn.execute(_COUNTS, {"board_id": board_id}).one() + return counts[0], counts[1] diff --git a/invokeai/app/services/shared/database/queries/boards.py b/invokeai/app/services/shared/database/queries/boards.py new file mode 100644 index 00000000000..08d8d55a489 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/boards.py @@ -0,0 +1,299 @@ +"""Boards and who they are shared with.""" + +import functools +import itertools +from collections.abc import Sequence +from typing import Any, Optional + +from sqlalchemy import ( + ColumnElement, + Connection, + Row, + Select, + Update, + bindparam, + delete, + exists, + false, + func, + insert, + literal, + select, + update, +) + +from invokeai.app.services.board_records.board_records_common import ( + BoardChanges, + BoardRecord, + BoardRecordOrderBy, + BoardVisibility, +) +from invokeai.app.services.shared.database.dialect import CaseInsensitiveOrder +from invokeai.app.services.shared.database.queries.base import IN_CHUNK, QueryModule, locking, mapped, read, write +from invokeai.app.services.shared.database.queries.board_access import readable_board +from invokeai.app.services.shared.database.schema.boards import boards, shared_boards +from invokeai.app.services.shared.database.schema.projects import projects +from invokeai.app.services.shared.pagination import SQLiteDirection + +_BOARD_COLUMNS = ( + boards.c.board_id, + boards.c.board_name, + boards.c.user_id, + boards.c.created_at, + boards.c.updated_at, + boards.c.deleted_at, + boards.c.cover_image_name, + boards.c.archived, + boards.c.board_visibility, +) +_CLAIMED = exists().where(projects.c.board_id == boards.c.board_id) + +_GET = select(*_BOARD_COLUMNS).where(boards.c.board_id == bindparam("board_id")) +_GET_WITH_PROJECT_ID = ( + select(*_BOARD_COLUMNS, projects.c.project_id) + .select_from(boards.outerjoin(projects, projects.c.board_id == boards.c.board_id)) + .where(boards.c.board_id == bindparam("board_id")) +) +_EXISTS = select(literal(1)).where(boards.c.board_id == bindparam("board_id")) +_LOCK = ( + select(boards.c.user_id, boards.c.board_visibility) + .where(boards.c.board_id == bindparam("board_id")) + .with_for_update() +) +_PROJECT_IDS = select(projects.c.board_id, projects.c.project_id).where( + projects.c.board_id.in_(bindparam("board_ids", expanding=True)) +) +_INSERT = insert(boards) +_DELETE_IF_UNCLAIMED = delete(boards).where(boards.c.board_id == bindparam("board_id"), ~_CLAIMED) +_IS_SHARED_WITH = select(literal(1)).where( + shared_boards.c.board_id == bindparam("board_id"), shared_boards.c.user_id == bindparam("user_id") +) +_SHARED_USER_IDS = select(shared_boards.c.user_id).where(shared_boards.c.board_id == bindparam("board_id")) +_IS_SHARED = select(exists().where(shared_boards.c.board_id == bindparam("board_id"))) +_RENAME = ( + update(boards) + .where(boards.c.board_id == bindparam("target_board_id")) + .values(board_name=bindparam("new_board_name")) +) +_RENAME_AND_UNARCHIVE = _RENAME.values(archived=false()) + + +def _board(row: Sequence[Any]) -> BoardRecord: + board_id, board_name, user_id, created_at, updated_at, deleted_at, cover_image_name, archived, visibility = row + try: + board_visibility = BoardVisibility(visibility) + except ValueError: + board_visibility = BoardVisibility.Private + return BoardRecord( + board_id=board_id, + board_name=board_name, + user_id=user_id, + created_at=created_at, + updated_at=updated_at, + deleted_at=deleted_at, + cover_image_name=cover_image_name, + archived=archived, + board_visibility=board_visibility, + ) + + +def _board_or_none(row: Optional[Sequence[Any]]) -> Optional[BoardRecord]: + return _board(row) if row is not None else None + + +def _board_and_project_id(row: Optional[Sequence[Any]]) -> Optional[tuple[BoardRecord, Optional[str]]]: + return (_board(row[:-1]), row[-1]) if row is not None else None + + +def _boards(rows: Sequence[Sequence[Any]]) -> list[BoardRecord]: + return [_board(row) for row in rows] + + +def _boards_and_total(result: tuple[Sequence[Sequence[Any]], int]) -> tuple[list[BoardRecord], int]: + rows, total = result + return _boards(rows), total + + +@functools.cache +def _update(changes_owned_state: bool) -> Update: + """Sets the given fields (a NULL leaves one as it is). A change to what a project owns -- the name, the + archived flag, the visibility -- applies only to a board no project claims.""" + statement = ( + update(boards) + .where(boards.c.board_id == bindparam("target_board_id")) + .values( + board_name=func.coalesce(bindparam("new_board_name", type_=boards.c.board_name.type), boards.c.board_name), + cover_image_name=func.coalesce( + bindparam("new_cover_image_name", type_=boards.c.cover_image_name.type), boards.c.cover_image_name + ), + archived=func.coalesce(bindparam("new_archived", type_=boards.c.archived.type), boards.c.archived), + board_visibility=func.coalesce( + bindparam("new_board_visibility", type_=boards.c.board_visibility.type), boards.c.board_visibility + ), + ) + ) + return statement.where(~_CLAIMED) if changes_owned_state else statement + + +def _listing(statement: Select[Any], is_admin: bool, include_archived: bool) -> Select[Any]: + if not is_admin: + statement = statement.where(readable_board(bindparam("user_id"))) + if not include_archived: + statement = statement.where(boards.c.archived == false()) + return statement + + +@functools.cache +def _list( + is_admin: bool, + include_archived: bool, + order_by: BoardRecordOrderBy, + direction: SQLiteDirection, + paged: bool, +) -> Select[Any]: + # The enums are also strings, which share their cache entries: compare by value, never by identity. + # `get_all` has always sorted names case-insensitively, a page of `get_many` by the names as stored. + key: ColumnElement[Any] = boards.c.board_name if order_by == BoardRecordOrderBy.Name else boards.c.created_at + if order_by == BoardRecordOrderBy.Name and not paged: + key = CaseInsensitiveOrder(boards.c.board_name) + # The board id breaks ties, so that equal keys keep one order from page to page. + ordering = [key, boards.c.board_id] + statement = _listing(select(*_BOARD_COLUMNS), is_admin, include_archived).order_by( + *(column.desc() if direction == SQLiteDirection.Descending else column.asc() for column in ordering) + ) + if paged: + statement = statement.limit(bindparam("limit")).offset(bindparam("offset")) + return statement + + +@functools.cache +def _count(is_admin: bool, include_archived: bool) -> Select[Any]: + return _listing(select(func.count()).select_from(boards), is_admin, include_archived) + + +class BoardQueries(QueryModule): + @mapped(_board_or_none) + @read + def get(self, conn: Connection, board_id: str) -> Optional[Row[Any]]: + return conn.execute(_GET, {"board_id": board_id}).first() + + @mapped(_board_and_project_id) + @read + def get_with_project_id(self, conn: Connection, board_id: str) -> Optional[Row[Any]]: + """The board and the id of the project that claims it, if one does.""" + return conn.execute(_GET_WITH_PROJECT_ID, {"board_id": board_id}).first() + + @read + def exists(self, conn: Connection, board_id: str) -> bool: + return conn.execute(_EXISTS, {"board_id": board_id}).first() is not None + + @locking + def lock(self, conn: Connection, board_id: str) -> Optional[tuple[str, str]]: + """Locks the board's row until the transaction ends: its owner and visibility, or None when there is no + such board. + + Locked before a transaction reads what it decides on, the row makes every transaction that changes the + board -- its owner, visibility or sharing, a project claiming it, its deletion -- wait until this one has + ended, and this one's later reads see what such a transaction committed before. Lock the board's project + first, if there is one: rows are locked in that order. + """ + row: Optional[Sequence[Any]] = conn.execute(_LOCK, {"board_id": board_id}).first() + if row is None: + return None + user_id, board_visibility = row + return user_id, board_visibility + + @read + def project_ids(self, conn: Connection, board_ids: Sequence[str]) -> dict[str, str]: + """The id of the project that claims each of these boards, for the boards a project claims.""" + claimed: dict[str, str] = {} + for chunk in itertools.batched(board_ids, IN_CHUNK): + claimed.update((row[0], row[1]) for row in conn.execute(_PROJECT_IDS, {"board_ids": list(chunk)}).all()) + return claimed + + @mapped(_boards_and_total) + @read + def page( + self, + conn: Connection, + *, + user_id: str, + is_admin: bool, + order_by: BoardRecordOrderBy, + direction: SQLiteDirection, + offset: int, + limit: int, + include_archived: bool, + ) -> tuple[Sequence[Row[Any]], int]: + """A page of the boards the user sees, and how many there are in all.""" + parameters = {"user_id": user_id, "limit": limit, "offset": offset} + rows = conn.execute(_list(is_admin, include_archived, order_by, direction, True), parameters).all() + total: int = conn.execute(_count(is_admin, include_archived), parameters).scalar_one() + return rows, total + + @mapped(_boards) + @read + def all( + self, + conn: Connection, + *, + user_id: str, + is_admin: bool, + order_by: BoardRecordOrderBy, + direction: SQLiteDirection, + include_archived: bool, + ) -> Sequence[Row[Any]]: + """Every board the user sees.""" + statement = _list(is_admin, include_archived, order_by, direction, False) + return conn.execute(statement, {"user_id": user_id}).all() + + @read + def is_shared_with(self, conn: Connection, board_id: str, user_id: str) -> bool: + return conn.execute(_IS_SHARED_WITH, {"board_id": board_id, "user_id": user_id}).first() is not None + + @read + def shared_user_ids(self, conn: Connection, board_id: str) -> list[str]: + return list(conn.execute(_SHARED_USER_IDS, {"board_id": board_id}).scalars().all()) + + @read + def is_shared(self, conn: Connection, board_id: str) -> bool: + """Whether the board is shared with any account.""" + return bool(conn.execute(_IS_SHARED, {"board_id": board_id}).scalar_one()) + + @write + def insert(self, conn: Connection, *, board_id: str, board_name: str, user_id: str) -> None: + conn.execute(_INSERT, {"board_id": board_id, "board_name": board_name, "user_id": user_id}) + + @write + def update(self, conn: Connection, board_id: str, changes: BoardChanges) -> bool: + """Applies the changes; whether a board took them. A board a project claims keeps its name, archived flag + and visibility, so a change to one of those leaves it as it is.""" + changes_owned_state = ( + changes.board_name is not None or changes.archived is not None or changes.board_visibility is not None + ) + result = conn.execute( + _update(changes_owned_state), + { + "target_board_id": board_id, + "new_board_name": changes.board_name, + "new_cover_image_name": changes.cover_image_name, + "new_archived": changes.archived, + "new_board_visibility": ( + changes.board_visibility.value if changes.board_visibility is not None else None + ), + }, + ) + return result.rowcount > 0 + + @write + def rename(self, conn: Connection, board_id: str, board_name: str, *, unarchive: bool = False) -> None: + """Renames the board, and with `unarchive` un-archives it, whether or not a project claims it: for the + projects storage, whose boards take their names from their projects. (`update` leaves the name of a + claimed board as it is.)""" + statement = _RENAME_AND_UNARCHIVE if unarchive else _RENAME + conn.execute(statement, {"target_board_id": board_id, "new_board_name": board_name}) + + @write + def delete_if_unclaimed(self, conn: Connection, board_id: str) -> bool: + """Deletes the board unless a project claims it; whether it did.""" + return conn.execute(_DELETE_IF_UNCLAIMED, {"board_id": board_id}).rowcount > 0 diff --git a/invokeai/app/services/shared/database/queries/client_state.py b/invokeai/app/services/shared/database/queries/client_state.py new file mode 100644 index 00000000000..dcb74d5b6f5 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/client_state.py @@ -0,0 +1,73 @@ +"""Per-user state of the web client: a value per key and account.""" + +import functools +import itertools +from typing import Optional + +from sqlalchemy import Connection, Insert, bindparam, delete, select + +from invokeai.app.services.shared.database.dialect import CaseInsensitiveLike, like_prefix, upsert +from invokeai.app.services.shared.database.queries.base import IN_CHUNK, QueryModule, read, write +from invokeai.app.services.shared.database.schema.client_state import client_state +from invokeai.app.services.shared.database.types import now_text + +_GET = select(client_state.c.value).where( + client_state.c.user_id == bindparam("user_id"), client_state.c.key == bindparam("key") +) +# Case-insensitive, as SQLite's LIKE always matched these prefixes. Most recently written first. +_KEYS_WITH_PREFIX = ( + select(client_state.c.key) + .where( + client_state.c.user_id == bindparam("user_id"), + CaseInsensitiveLike(client_state.c.key, bindparam("pattern")), + ) + .order_by(client_state.c.updated_at.desc()) +) +_DELETE = delete(client_state).where( + client_state.c.user_id == bindparam("user_id"), client_state.c.key == bindparam("key") +) +_LOCK_KEYS = select(client_state.c.key).where(client_state.c.user_id == bindparam("user_id")).with_for_update() +_DELETE_KEYS = delete(client_state).where( + client_state.c.user_id == bindparam("user_id"), client_state.c.key.in_(bindparam("keys", expanding=True)) +) + + +@functools.cache +def _set(dialect_name: str) -> Insert: + return upsert(dialect_name, client_state, update=("value", "updated_at")) + + +class ClientStateQueries(QueryModule): + @read + def get(self, conn: Connection, user_id: str, key: str) -> Optional[str]: + return conn.execute(_GET, {"user_id": user_id, "key": key}).scalar_one_or_none() + + @read + def keys_with_prefix(self, conn: Connection, user_id: str, prefix: str) -> list[str]: + """The account's keys that start with `prefix`, ignoring case; `%` and `_` in it match only themselves.""" + return list( + conn.execute(_KEYS_WITH_PREFIX, {"user_id": user_id, "pattern": like_prefix(prefix)}).scalars().all() + ) + + @write + def set(self, conn: Connection, user_id: str, key: str, value: str) -> None: + conn.execute( + _set(conn.dialect.name), {"user_id": user_id, "key": key, "value": value, "updated_at": now_text()} + ) + + @write + def delete(self, conn: Connection, user_id: str, key: str) -> bool: + """Deletes the value; whether there was one.""" + return conn.execute(_DELETE, {"user_id": user_id, "key": key}).rowcount > 0 + + @write + def delete_all(self, conn: Connection, user_id: str) -> list[str]: + """Deletes the account's values and returns their keys. + + Only values this transaction locked: on a server, a key that another transaction writes meanwhile, for + the first time, is neither deleted nor returned. + """ + keys: list[str] = list(conn.execute(_LOCK_KEYS, {"user_id": user_id}).scalars().all()) + for chunk in itertools.batched(keys, IN_CHUNK): + conn.execute(_DELETE_KEYS, {"user_id": user_id, "keys": list(chunk)}) + return keys diff --git a/invokeai/app/services/shared/database/queries/fonts.py b/invokeai/app/services/shared/database/queries/fonts.py new file mode 100644 index 00000000000..0a71331a889 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/fonts.py @@ -0,0 +1,388 @@ +"""The font catalog: fonts users uploaded (private to their owner, or shared), and the fonts of the configured font +directory, which are shared and read-only.""" + +import functools +import itertools +import json +from collections.abc import Sequence +from typing import Any, NamedTuple, Optional + +from sqlalchemy import ( + ColumnElement, + Connection, + Row, + Select, + and_, + bindparam, + delete, + func, + insert, + literal, + or_, + select, + update, +) + +from invokeai.app.services.fonts.fonts_common import FontAxis, FontInstance, FontRecord, FontScope, FontSource +from invokeai.app.services.shared.database.dialect import ( + CaseInsensitiveLike, + CaseInsensitiveOrder, + fixed_limit, + like_contains, +) +from invokeai.app.services.shared.database.queries.base import IN_CHUNK, QueryModule, mapped, read, write +from invokeai.app.services.shared.database.schema.fonts import fonts +from invokeai.app.services.shared.database.schema.metadata import PATH_LENGTH +from invokeai.app.services.shared.database.types import now_text + +_F = fonts.c +_COLUMNS = ( + _F.id, + _F.owner_id, + _F.family, + _F.label, + _F.style, + _F.weight, + _F.content_hash, + _F.scope, + _F.source, + _F.filename, + _F.byte_size, + _F.axes_json, + _F.instances_json, + _F.storage_path, + _F.source_path, +) + +_UPLOADED = _F.source == literal("uploaded") +_DIRECTORY = _F.source == literal("directory") +_PRIVATE = _F.scope == literal("private") +_SHARED = _F.scope == literal("shared") +_OWN_PRIVATE = and_(_UPLOADED, _PRIVATE, _F.owner_id == bindparam("owner")) +_SHARED_UPLOAD = and_(_UPLOADED, _SHARED, _F.owner_id.is_(None)) + +_FONT = select(*_COLUMNS).where(_F.id == bindparam("font")) +_PRIVATE_DUPLICATE = select(*_COLUMNS).where(_OWN_PRIVATE, _F.content_hash == bindparam("hash")).limit(fixed_limit(1)) +_SHARED_DUPLICATE = select(*_COLUMNS).where(_SHARED_UPLOAD, _F.content_hash == bindparam("hash")).limit(fixed_limit(1)) +_PRIVATE_BYTES = select(func.coalesce(func.sum(_F.byte_size), 0)).where(_OWN_PRIVATE) +_SHARED_BYTES = select(func.coalesce(func.sum(_F.byte_size), 0)).where(_SHARED_UPLOAD) +_INSERT = insert(fonts) +_DELETE_UPLOAD = delete(fonts).where(_F.id == bindparam("font"), _UPLOADED) + +_DIRECTORY_FONTS = select(*_COLUMNS).where(_DIRECTORY) +_DIRECTORY_IDS = select(_F.id).where(_DIRECTORY) +_UPDATE_DIRECTORY_FONT = ( + update(fonts) + .where(_F.id == bindparam("font"), _DIRECTORY) + .values( + filename=bindparam("filename"), + source_path=bindparam("source_path"), + family=bindparam("family"), + label=bindparam("label"), + style=bindparam("style"), + weight=bindparam("weight"), + content_hash=bindparam("content_hash"), + byte_size=bindparam("byte_size"), + axes_json=bindparam("axes_json"), + instances_json=bindparam("instances_json"), + updated_at=bindparam("now"), + ) +) +_DELETE_DIRECTORY_FONTS = delete(fonts).where(_DIRECTORY, _F.id.in_(bindparam("ids", expanding=True))) +_DELETE_ALL_DIRECTORY_FONTS = delete(fonts).where(_DIRECTORY) + +_STORAGE_PATHS = select(_F.storage_path).where(_UPLOADED, _F.storage_path.is_not(None)) +_PRIVATE_STORAGE_PATHS = select(_F.storage_path).where(_OWN_PRIVATE, _F.storage_path.is_not(None)) + +# The longest path, relative to the font directory, that the catalog holds. (A server holds a key column only up +# to its declared length.) +MAX_SOURCE_PATH_LENGTH = PATH_LENGTH + + +class DirectoryFont(NamedTuple): + """A font file the directory scan found, as the catalog stores it.""" + + id: str + filename: str + source_path: str + family: str + label: str + style: str + weight: int + content_hash: str + byte_size: int + axes: tuple[FontAxis, ...] + instances: tuple[FontInstance, ...] + + +class UploadedFont(NamedTuple): + id: str + owner_id: Optional[str] + scope: FontScope + filename: str + storage_path: str + family: str + label: str + style: str + weight: int + content_hash: str + byte_size: int + axes: tuple[FontAxis, ...] + instances: tuple[FontInstance, ...] + + +class _Shape(NamedTuple): + scope: FontScope + by_hash: bool + searched: bool + + +def _json_axes(axes: tuple[FontAxis, ...]) -> str: + return json.dumps([axis.__dict__ for axis in axes], separators=(",", ":"), sort_keys=True) + + +def _json_instances(instances: tuple[FontInstance, ...]) -> str: + return json.dumps([instance.__dict__ for instance in instances], separators=(",", ":"), sort_keys=True) + + +_SCOPES = {scope.value: scope for scope in FontScope} +_SOURCES = {source.value: source for source in FontSource} +# What `_json_axes(())` and `_json_instances(())` write, and the columns' default: most fonts are not variable. +_NONE = "[]" + + +def _record(row: Row[Any]) -> FontRecord: + # A page builds a hundred of these, which costs more than reading their rows. + ( + font_id, + owner_id, + family, + label, + style, + weight, + content_hash, + scope, + source, + filename, + byte_size, + axes_json, + instances_json, + storage_path, + source_path, + ) = row + return FontRecord( + id=font_id, + family=family, + label=label, + style=style, + weight=int(weight), + content_hash=content_hash, + scope=_SCOPES[scope], + source=_SOURCES[source], + filename=filename, + byte_size=int(byte_size), + axes=() if axes_json == _NONE else tuple(FontAxis(**axis) for axis in json.loads(axes_json)), + instances=( + () if instances_json == _NONE else tuple(FontInstance(**item) for item in json.loads(instances_json)) + ), + owner_id=owner_id, + storage_path=storage_path, + source_path=source_path, + ) + + +def _record_or_none(row: Optional[Row[Any]]) -> Optional[FontRecord]: + return _record(row) if row is not None else None + + +def _records(rows: Sequence[Row[Any]]) -> list[FontRecord]: + return [_record(row) for row in rows] + + +def _upload_state(state: tuple[Optional[Row[Any]], int]) -> tuple[Optional[FontRecord], int]: + duplicate, used_bytes = state + return _record_or_none(duplicate), used_bytes + + +def _page(page: tuple[Sequence[Row[Any]], int]) -> tuple[list[FontRecord], int]: + rows, total = page + return _records(rows), total + + +def _metadata_values(font: DirectoryFont | UploadedFont) -> dict[str, Any]: + return { + "filename": font.filename, + "family": font.family, + "label": font.label, + "style": font.style, + "weight": font.weight, + "content_hash": font.content_hash, + "byte_size": font.byte_size, + "axes_json": _json_axes(font.axes), + "instances_json": _json_instances(font.instances), + } + + +def _conditions(shape: _Shape) -> list[ColumnElement[bool]]: + # Directory fonts and shared uploads are everyone's; a private upload only its owner's. + conditions: list[ColumnElement[bool]] = [or_(_DIRECTORY, _SHARED, and_(_PRIVATE, _F.owner_id == bindparam("user")))] + if shape.scope == FontScope.PRIVATE: + conditions.append(and_(_UPLOADED, _PRIVATE, _F.owner_id == bindparam("user"))) + elif shape.scope == FontScope.SHARED: + conditions.append(or_(_DIRECTORY, and_(_UPLOADED, _SHARED))) + if shape.by_hash: + conditions.append(_F.content_hash == bindparam("hash")) + if shape.searched: + pattern = bindparam("pattern") + conditions.append( + or_( + CaseInsensitiveLike(_F.family, pattern), + CaseInsensitiveLike(_F.label, pattern), + CaseInsensitiveLike(_F.filename, pattern), + ) + ) + return conditions + + +@functools.cache +def _listing(shape: _Shape) -> Select[Any]: + return ( + select(*_COLUMNS) + .where(*_conditions(shape)) + .order_by(CaseInsensitiveOrder(_F.family), CaseInsensitiveOrder(_F.label), _F.id) + .limit(bindparam("limit")) + .offset(bindparam("offset")) + ) + + +@functools.cache +def _count(shape: _Shape) -> Select[Any]: + return select(func.count()).select_from(fonts).where(*_conditions(shape)) + + +class FontQueries(QueryModule): + @mapped(_record_or_none) + @read + def get(self, conn: Connection, font_id: str) -> Optional[Row[Any]]: + """The font, whoever may see it.""" + return conn.execute(_FONT, {"font": font_id}).first() + + @mapped(_page) + @read + def page( + self, + conn: Connection, + *, + user_id: str, + scope: FontScope, + content_hash: Optional[str], + search: Optional[str], + offset: int, + limit: int, + ) -> tuple[Sequence[Row[Any]], int]: + """`limit` of the fonts the account may see, from `offset`, by family and label ignoring case, then id; and + how many there are. `scope` keeps only the account's private uploads, or only the shared fonts; `search` + keeps the fonts whose family, label or file name contains it, ignoring case.""" + shape = _Shape(scope, content_hash is not None, bool(search)) + parameters: dict[str, Any] = {"user": user_id} + if content_hash is not None: + parameters["hash"] = content_hash + if search: + parameters["pattern"] = like_contains(search) + rows = conn.execute(_listing(shape), {**parameters, "limit": limit, "offset": offset}).all() + if len(rows) < limit and (rows or offset == 0): + # The page ends the listing, so it tells how many there are. + return rows, offset + len(rows) + return rows, int(conn.execute(_count(shape), parameters).scalar_one()) + + @mapped(_upload_state) + @read + def upload_state( + self, conn: Connection, *, owner_id: Optional[str], content_hash: str + ) -> tuple[Optional[Row[Any]], int]: + """The upload of the same content in the same place (the owner's private uploads, or the shared ones with + no owner), if there is one; and how many bytes the uploads in that place take.""" + if owner_id is None: + duplicate = conn.execute(_SHARED_DUPLICATE, {"hash": content_hash}).first() + used = conn.execute(_SHARED_BYTES).scalar_one() + else: + duplicate = conn.execute(_PRIVATE_DUPLICATE, {"owner": owner_id, "hash": content_hash}).first() + used = conn.execute(_PRIVATE_BYTES, {"owner": owner_id}).scalar_one() + return duplicate, int(used) + + @write + def insert_upload(self, conn: Connection, font: UploadedFont) -> None: + conn.execute( + _INSERT, + { + "id": font.id, + "owner_id": font.owner_id, + "scope": font.scope.value, + "source": FontSource.UPLOADED.value, + "storage_path": font.storage_path, + "source_path": None, + **_metadata_values(font), + }, + ) + + @write + def delete_upload(self, conn: Connection, font_id: str) -> bool: + """Deletes the uploaded font; whether there was one.""" + return conn.execute(_DELETE_UPLOAD, {"font": font_id}).rowcount == 1 + + @mapped(_records) + @read + def directory_fonts(self, conn: Connection) -> Sequence[Row[Any]]: + return conn.execute(_DIRECTORY_FONTS).all() + + @write + def replace_directory_fonts( + self, conn: Connection, changed: Sequence[DirectoryFont], keep_ids: frozenset[str] + ) -> None: + """Makes the directory fonts those of `keep_ids`: inserts or updates the `changed` ones, and deletes the + others.""" + existing = set(conn.execute(_DIRECTORY_IDS).scalars().all()) + now = now_text() + inserted = [font for font in changed if font.id not in existing] + updated = [font for font in changed if font.id in existing] + if inserted: + conn.execute( + _INSERT, + [ + { + "id": font.id, + "owner_id": None, + "scope": FontScope.SHARED.value, + "source": FontSource.DIRECTORY.value, + "storage_path": None, + "source_path": font.source_path, + **_metadata_values(font), + } + for font in inserted + ], + ) + if updated: + conn.execute( + _UPDATE_DIRECTORY_FONT, + [ + {"font": font.id, "source_path": font.source_path, "now": now, **_metadata_values(font)} + for font in updated + ], + ) + for chunk in itertools.batched(sorted(existing - keep_ids), IN_CHUNK): + conn.execute(_DELETE_DIRECTORY_FONTS, {"ids": list(chunk)}) + + @write + def delete_directory_fonts(self, conn: Connection) -> int: + """Deletes every directory font; how many there were.""" + return conn.execute(_DELETE_ALL_DIRECTORY_FONTS).rowcount + + @read + def storage_paths(self, conn: Connection) -> set[str]: + """The files of every uploaded font, relative to the managed storage.""" + return set(conn.execute(_STORAGE_PATHS).scalars().all()) + + @read + def private_storage_paths(self, conn: Connection, owner_id: str) -> tuple[str, ...]: + """The files of the account's private uploads.""" + return tuple(conn.execute(_PRIVATE_STORAGE_PATHS, {"owner": owner_id}).scalars().all()) diff --git a/invokeai/app/services/shared/database/queries/gallery.py b/invokeai/app/services/shared/database/queries/gallery.py new file mode 100644 index 00000000000..c29e039ab5d --- /dev/null +++ b/invokeai/app/services/shared/database/queries/gallery.py @@ -0,0 +1,429 @@ +"""The gallery: images and videos listed as one, through a UNION ALL of two halves with the same columns. + +The filters apply alike to each half. A literal `kind` tells the halves apart, and duration and fps are NULL for +images. Lists break ties by kind and name in the listing's direction, so that items of one moment keep one order +from page to page. +""" + +import functools +import itertools +from collections.abc import Sequence +from typing import Any, NamedTuple, Optional + +from sqlalchemy import ( + ColumnElement, + Connection, + FromClause, + Row, + Select, + Subquery, + Table, + and_, + bindparam, + case, + exists, + false, + func, + literal, + null, + or_, + select, + union_all, +) + +from invokeai.app.services.image_records.image_records_common import ( + ASSETS_CATEGORIES, + IMAGE_CATEGORIES, + ImageCategory, + ResourceOrigin, +) +from invokeai.app.services.shared.database.dialect import ( + CaseInsensitiveLike, + JsonString, + OrderedJoin, + fixed_limit, + like_contains, +) +from invokeai.app.services.shared.database.queries.base import IN_CHUNK, QueryModule, day_after, read +from invokeai.app.services.shared.database.schema.boards import board_images, board_videos +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.videos import videos + + +class _Half(NamedTuple): + media: Table + membership: Table + name: str + category: str + origin: str + + +_HALVES = { + "image": _Half(images, board_images, "image_name", "image_category", "image_origin"), + "video": _Half(videos, board_videos, "video_name", "video_category", "video_origin"), +} + + +class _Shape(NamedTuple): + """Which filters a listing has: never their values, which are bound, so that the shapes are few.""" + + origin: bool + # The categories' values, sorted; None for no category filter. + categories: Optional[tuple[str, ...]] + intermediate: bool + starred: bool + # "any" (no board filter), "none" (on no board) or "one". + board: str + # Whether the listing is limited to one account's items; else everything, as for an administrator. + scoped: bool + search: bool + created_date: bool + created_from: bool + created_to: bool + + +class Filters(NamedTuple): + """A gallery listing's filters, as the service passes them.""" + + origin: Optional[ResourceOrigin] = None + categories: Optional[Sequence[ImageCategory]] = None + is_intermediate: Optional[bool] = None + board_id: Optional[str] = None + search_term: Optional[str] = None + user_id: Optional[str] = None + is_admin: bool = False + created_date: Optional[str] = None + created_from: Optional[str] = None + created_to: Optional[str] = None + starred: Optional[bool] = None + + def shape(self) -> _Shape: + return _Shape( + origin=self.origin is not None, + categories=None + if self.categories is None + else tuple(sorted({ImageCategory(c).value for c in self.categories})), + intermediate=self.is_intermediate is not None, + starred=self.starred is not None, + board="any" if self.board_id is None else "none" if self.board_id == "none" else "one", + scoped=self.user_id is not None and not self.is_admin, + search=bool(self.search_term), + created_date=self.created_date is not None, + created_from=self.created_from is not None, + created_to=self.created_to is not None, + ) + + def parameters(self) -> dict[str, Any]: + return { + "origin": self.origin.value if self.origin is not None else None, + "is_intermediate": self.is_intermediate, + "starred": self.starred, + "board_id": self.board_id, + "user_id": self.user_id, + # Lowered here, as the legacy services did: SQLite's LIKE folds the case of ASCII letters only. + "pattern": like_contains(self.search_term.lower()) if self.search_term else None, + "created_date": self.created_date, + "created_date_before": day_after(self.created_date) if self.created_date is not None else None, + "created_from": self.created_from, + "created_before": day_after(self.created_to) if self.created_to is not None else None, + } + + +def _from(kind: str, shape: _Shape, with_board: bool) -> tuple[FromClause, ColumnElement[Any]]: + """The half's FROM and its board id: a board's items come through its memberships, so that the work stays + proportional to the board; a listing of every item joins the memberships only for the board id it shows.""" + half = _HALVES[kind] + media, membership = half.media, half.membership + on = membership.c[half.name] == media.c[half.name] + if shape.board == "one": + return OrderedJoin(membership, media, on), membership.c.board_id + if shape.board == "any" and with_board: + return media.outerjoin(membership, on), membership.c.board_id + return media, null() + + +def _conditions(kind: str, shape: _Shape) -> list[ColumnElement[bool]]: + half = _HALVES[kind] + media, membership = half.media, half.membership + conditions: list[ColumnElement[bool]] = [] + if shape.origin: + conditions.append(media.c[half.origin] == bindparam("origin")) + if shape.categories is not None: + conditions.append(media.c[half.category].in_([literal(category) for category in shape.categories])) + if shape.intermediate: + conditions.append(media.c.is_intermediate == bindparam("is_intermediate")) + if shape.starred: + conditions.append(media.c.starred == bindparam("starred")) + if shape.created_date: + conditions.append(media.c.created_at >= bindparam("created_date")) + conditions.append(media.c.created_at < bindparam("created_date_before")) + if shape.created_from: + conditions.append(media.c.created_at >= bindparam("created_from")) + if shape.created_to: + conditions.append(media.c.created_at < bindparam("created_before")) + own = media.c.user_id == bindparam("user_id") + if shape.board == "none": + conditions.append(~exists(select(literal(1)).where(membership.c[half.name] == media.c[half.name]))) + if shape.scoped: + conditions.append(own) + elif shape.board == "one": + conditions.append(membership.c.board_id == bindparam("board_id")) + elif shape.scoped: + # Without a board, still only the account's own items, so that it cannot list every account's. + conditions.append(own) + if shape.search: + pattern: ColumnElement[str] = bindparam("pattern") + conditions.append( + or_(CaseInsensitiveLike(media.c.metadata, pattern), CaseInsensitiveLike(media.c.created_at, pattern)) + ) + return conditions + + +def _half(kind: str, shape: _Shape, names_only: bool) -> Select[Any]: + half = _HALVES[kind] + media = half.media + source, board_id = _from(kind, shape, with_board=not names_only) + if names_only: + columns: list[ColumnElement[Any]] = [ + literal(kind).label("kind"), + media.c[half.name].label("name"), + media.c.starred.label("starred"), + media.c.created_at.label("created_at"), + ] + else: + columns = [ + literal(kind).label("kind"), + media.c[half.name].label("name"), + media.c.width.label("width"), + media.c.height.label("height"), + media.c[half.category].label("category"), + media.c.starred.label("starred"), + media.c.is_intermediate.label("is_intermediate"), + board_id.label("board_id"), + media.c.created_at.label("created_at"), + (media.c.duration if kind == "video" else null()).label("duration"), + (media.c.fps if kind == "video" else null()).label("fps"), + ] + return select(*columns).select_from(source).where(*_conditions(kind, shape)) + + +def _ordering(items: Subquery, starred_first: bool, descending: bool) -> list[ColumnElement[Any]]: + keys = [items.c.created_at, items.c.kind, items.c.name] + ordering = [key.desc() if descending else key.asc() for key in keys] + return [items.c.starred.desc(), *ordering] if starred_first else ordering + + +@functools.lru_cache(maxsize=256) +def _page(shape: _Shape, starred_first: bool, descending: bool) -> Select[Any]: + """A page of items. `media_origin` is joined onto the chosen page rather than selected inside the halves: it is + read out of `videos.metadata`, and selecting it there would read and parse every video before LIMIT.""" + items = union_all(_half("image", shape, False), _half("video", shape, False)).subquery("items") + page = ( + select(items) + .order_by(*_ordering(items, starred_first, descending)) + .limit(bindparam("limit")) + .offset(bindparam("offset")) + .subquery("page") + ) + video_of_page = and_(page.c.kind == literal("video"), videos.c.video_name == page.c.name) + return ( + select(*page.c, JsonString(videos.c.metadata, "$.media_origin").label("media_origin")) + .select_from(page.outerjoin(videos, video_of_page)) + .order_by(*_ordering(page, starred_first, descending)) + ) + + +@functools.lru_cache(maxsize=256) +def _count(shape: _Shape) -> Select[Any]: + """The items of both halves that match, in one statement.""" + counts = [] + for kind in _HALVES: + # Without the board id a listing shows: a membership join cannot change the count, and MySQL does not + # eliminate it. + source, _ = _from(kind, shape, with_board=False) + counts.append(select(func.count()).select_from(source).where(*_conditions(kind, shape)).scalar_subquery()) + return select(counts[0] + counts[1]) + + +@functools.lru_cache(maxsize=256) +def _names(shape: _Shape, starred_first: bool, descending: bool) -> Select[Any]: + items = union_all(_half("image", shape, True), _half("video", shape, True)).subquery("items") + return select(items.c.kind, items.c.name, items.c.starred).order_by(*_ordering(items, starred_first, descending)) + + +def _in(column: ColumnElement[Any], categories: Sequence[ImageCategory]) -> ColumnElement[bool]: + return column.in_([literal(category.value) for category in categories]) + + +def _counted(items: Subquery) -> list[ColumnElement[int]]: + """Image and asset counts as the gallery's views list them (`IMAGE_CATEGORIES`, `ASSETS_CATEGORIES`), so that a + count never includes canvas-owned images (OTHER), which neither view shows; and every video, and those of a + category other than general, as a board counts its videos.""" + is_image, is_video = items.c.kind == literal("image"), items.c.kind == literal("video") + return [ + func.sum(case((and_(is_image, _in(items.c.category, IMAGE_CATEGORIES)), 1), else_=0)), + func.sum(case((and_(is_image, _in(items.c.category, ASSETS_CATEGORIES)), 1), else_=0)), + func.sum(case((is_video, 1), else_=0)), + func.sum(case((and_(is_video, items.c.category != literal(ImageCategory.GENERAL.value)), 1), else_=0)), + ] + + +def _dated(kind: str, scoped: bool) -> Select[Any]: + half = _HALVES[kind] + media = half.media + conditions = [media.c.is_intermediate == false()] + if kind == "image": + # Only what the gallery's views show: canvas-owned images (OTHER) neither make a day nor cover one. + conditions.append(_in(media.c[half.category], [*IMAGE_CATEGORIES, *ASSETS_CATEGORIES])) + if scoped: + conditions.append(media.c.user_id == bindparam("user_id")) + return select( + # The day of a canonical timestamp: its first ten characters, on every backend. + func.substr(media.c.created_at, 1, 10).label("day"), + literal(kind).label("kind"), + media.c[half.name].label("name"), + media.c[half.category].label("category"), + media.c.created_at.label("created_at"), + ).where(*conditions) + + +@functools.cache +def _date_counts(scoped: bool) -> Select[Any]: + items = union_all(_dated("image", scoped), _dated("video", scoped)).subquery("items") + return select(items.c.day, *_counted(items)).group_by(items.c.day).order_by(items.c.day.desc()) + + +@functools.cache +def _date_covers(scoped: bool) -> Select[Any]: + """The newest item of each day; kind and name break ties, so that the cover does not flicker between refetches.""" + items = union_all(_dated("image", scoped), _dated("video", scoped)).subquery("items") + newest_first = func.row_number().over( + partition_by=items.c.day, order_by=[items.c.created_at.desc(), items.c.kind.desc(), items.c.name.desc()] + ) + ranked = select(items.c.day, items.c.kind, items.c.name, newest_first.label("rank")).subquery("ranked") + return select(ranked.c.day, ranked.c.kind, ranked.c.name).where(ranked.c.rank == 1) + + +def _on_board(kind: str) -> Select[Any]: + half = _HALVES[kind] + media, membership = half.media, half.membership + return ( + select( + membership.c.board_id.label("board_id"), + literal(kind).label("kind"), + media.c[half.name].label("name"), + media.c[half.category].label("category"), + media.c.starred.label("starred"), + media.c.created_at.label("created_at"), + ) + .select_from(membership.join(media, membership.c[half.name] == media.c[half.name])) + .where(media.c.is_intermediate == false(), membership.c.board_id.in_(bindparam("board_ids", expanding=True))) + ) + + +def _board_summaries() -> Select[Any]: + items = union_all(_on_board("image"), _on_board("video")).subquery("items") + cover_first = func.row_number().over( + partition_by=items.c.board_id, + order_by=[items.c.starred.desc(), items.c.created_at.desc(), items.c.kind.desc(), items.c.name.desc()], + ) + ranked = select(*items.c, cover_first.label("rank")).subquery("ranked") + cover = ranked.c.rank == 1 + return select( + ranked.c.board_id, + *_counted(ranked), + func.max(case((and_(cover, ranked.c.kind == literal("image")), ranked.c.name))), + func.max(case((and_(cover, ranked.c.kind == literal("video")), ranked.c.name))), + ).group_by(ranked.c.board_id) + + +_BOARD_SUMMARIES = _board_summaries() + + +class DateCounts(NamedTuple): + day: str + images: int + assets: int + videos: int + asset_videos: int + + +class BoardSummary(NamedTuple): + board_id: str + images: int + assets: int + videos: int + asset_videos: int + cover_image_name: Optional[str] + cover_video_name: Optional[str] + + +@functools.lru_cache(maxsize=256) +def _location(shape: _Shape, descending: bool) -> Select[Any]: + """An item's zero-based index in the listing `_names` orders (starred items not first), and how many items match: + ranked in the database, so no list of names is built. (Bind names other than the filters' own.)""" + items = union_all(_half("image", shape, True), _half("video", shape, True)).subquery("items") + ranked = select( + items.c.kind, + items.c.name, + (func.row_number().over(order_by=_ordering(items, False, descending)) - 1).label("item_index"), + func.count().over().label("total"), + ).subquery("ranked") + return ( + select(ranked.c.kind, ranked.c.name, ranked.c.item_index, ranked.c.total) + .where(ranked.c.kind == bindparam("location_kind"), ranked.c.name == bindparam("location_name")) + .limit(fixed_limit(1)) + ) + + +class GalleryQueries(QueryModule): + @read + def page( + self, conn: Connection, filters: Filters, *, offset: int, limit: int, starred_first: bool, descending: bool + ) -> tuple[Sequence[Row[Any]], int]: + """A page of items, with the columns of `_half` and `media_origin`, and how many match in all.""" + shape = filters.shape() + parameters = filters.parameters() + rows = conn.execute( + _page(shape, starred_first, descending), {**parameters, "limit": limit, "offset": offset} + ).all() + return rows, int(conn.execute(_count(shape), parameters).scalar_one()) + + @read + def names(self, conn: Connection, filters: Filters, *, starred_first: bool, descending: bool) -> Sequence[Row[Any]]: + """(kind, name, starred) of every item that matches, in order.""" + return conn.execute(_names(filters.shape(), starred_first, descending), filters.parameters()).all() + + @read + def location( + self, conn: Connection, filters: Filters, *, kind: str, name: str, descending: bool + ) -> Optional[Row[Any]]: + """(kind, name, index, total) of the named item in the listing, or None when the filters leave it out.""" + parameters = {**filters.parameters(), "location_kind": kind, "location_name": name} + return conn.execute(_location(filters.shape(), descending), parameters).first() + + @read + def date_counts( + self, conn: Connection, user_id: Optional[str] + ) -> tuple[list[DateCounts], dict[str, tuple[str, str]]]: + """The counts of each day that has items that are not intermediates (only `user_id`'s, if given), newest + first, and each day's cover as (kind, name).""" + scoped = user_id is not None + parameters = {"user_id": user_id} + counts = [ + DateCounts(str(r[0]), int(r[1] or 0), int(r[2] or 0), int(r[3] or 0), int(r[4] or 0)) + for r in conn.execute(_date_counts(scoped), parameters).all() + ] + covers = {str(r[0]): (str(r[1]), str(r[2])) for r in conn.execute(_date_covers(scoped), parameters).all()} + return counts, covers + + @read + def board_summaries(self, conn: Connection, board_ids: Sequence[str]) -> list[BoardSummary]: + """Counts and covers of the boards' items that are not intermediates; a board without any has no row.""" + summaries: list[BoardSummary] = [] + for chunk in itertools.batched(board_ids, IN_CHUNK // 2): + for r in conn.execute(_BOARD_SUMMARIES, {"board_ids": list(chunk)}).all(): + summaries.append( + BoardSummary(str(r[0]), int(r[1] or 0), int(r[2] or 0), int(r[3] or 0), int(r[4] or 0), r[5], r[6]) + ) + return summaries diff --git a/invokeai/app/services/shared/database/queries/image_index.py b/invokeai/app/services/shared/database/queries/image_index.py new file mode 100644 index 00000000000..266dc813d16 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/image_index.py @@ -0,0 +1,342 @@ +"""The image index: embeddings of gallery images and videos under each embedding model, the accounts' cached map +projections, and the custom labeling vocabulary. + +Only gallery items are indexed: images and videos in the general category that are not intermediates. Each kind has +its own embedding table, so that an embedding is deleted with its media by foreign key. Listings of both kinds give +the images first. +""" + +import functools +from collections.abc import Sequence +from typing import Any, NamedTuple, Optional + +from sqlalchemy import ( + ColumnElement, + Connection, + Delete, + Row, + Select, + Table, + and_, + bindparam, + delete, + exists, + false, + func, + literal, + or_, + select, + union_all, +) + +from invokeai.app.services.image_index.image_index_common import IndexedItem, MediaKind +from invokeai.app.services.image_records.image_records_common import ImageCategory +from invokeai.app.services.shared.database.dialect import InBoundSet, bound_set, insert_ignore, upsert +from invokeai.app.services.shared.database.queries.base import QueryModule, mapped, read, write +from invokeai.app.services.shared.database.queries.board_access import readable_board +from invokeai.app.services.shared.database.schema.boards import board_images, board_videos, boards +from invokeai.app.services.shared.database.schema.image_index import ( + image_embeddings, + image_index_vocab_terms, + image_projections, + video_embeddings, +) +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.users import users +from invokeai.app.services.shared.database.schema.videos import videos +from invokeai.app.services.shared.database.types import now_text + + +class _Kind(NamedTuple): + media: Table + embeddings: Table + membership: Table + name: str + category: str + + +# Images first: every listing of both kinds has this one order. +_KINDS: dict[MediaKind, _Kind] = { + "image": _Kind(images, image_embeddings, board_images, "image_name", "image_category"), + "video": _Kind(videos, video_embeddings, board_videos, "video_name", "video_category"), +} + +_P = image_projections.c +_THE_PROJECTION = and_(_P.user_id == bindparam("user_id"), _P.model_id == bindparam("model_id")) +_PROJECTION = select( + _P.scope_hash, _P.params, _P.point_count, _P.image_names, _P.item_kinds, _P.coords, _P.created_at, _P.updated_at +).where(_THE_PROJECTION) +_DELETE_PROJECTION = delete(image_projections).where(_THE_PROJECTION) +# Shared, so that the account is not deleted before the projection written for it is. +_LOCK_USER = select(literal(1)).where(users.c.user_id == bindparam("user_id")).with_for_update(read=True) + +_V = image_index_vocab_terms.c +_VOCAB_TERMS = select(_V.term).order_by(_V.term) +_CLEAR_VOCAB_TERMS = delete(image_index_vocab_terms) + + +def _eligible(kind: _Kind) -> list[ColumnElement[bool]]: + """The conditions of a gallery item, the only kind worth indexing.""" + return [ + kind.media.c.is_intermediate == false(), + kind.media.c[kind.category] == literal(ImageCategory.GENERAL.value), + ] + + +def _with_embedding(kind: _Kind) -> Any: + """The media with their embedding under the bound model, if they have one.""" + embeddings = kind.embeddings + on = and_(embeddings.c[kind.name] == kind.media.c[kind.name], embeddings.c.model_id == bindparam("model_id")) + return kind.media.outerjoin(embeddings, on) + + +@functools.cache +def _lock_media(kind: MediaKind) -> Select[Any]: + # Shared, so that the media is not deleted before the embedding written for it is. + media, name = _KINDS[kind].media, _KINDS[kind].name + return select(literal(1)).where(media.c[name] == bindparam("name")).with_for_update(read=True) + + +@functools.cache +def _upsert_embedding(kind: MediaKind, dialect_name: str) -> Any: + return upsert(dialect_name, _KINDS[kind].embeddings, update=["dim", "embedding", "encoding"]) + + +def _named_embeddings(kind: MediaKind) -> Select[Any]: + embeddings, name = _KINDS[kind].embeddings, _KINDS[kind].name + return select( + literal(kind), embeddings.c[name], embeddings.c.dim, embeddings.c.embedding, embeddings.c.encoding + ).where(embeddings.c.model_id == bindparam("model_id"), InBoundSet(embeddings.c[name], bindparam(f"{kind}_names"))) + + +# Both kinds in one statement, each kind's names in one parameter: an IN list would bind every name, in statements +# of IN_CHUNK. +_EMBEDDINGS = union_all(*(_named_embeddings(kind) for kind in _KINDS)) + + +@functools.cache +def _delete_embedding(kind: MediaKind) -> Delete: + embeddings, name = _KINDS[kind].embeddings, _KINDS[kind].name + return delete(embeddings).where(embeddings.c[name] == bindparam("name")) + + +@functools.cache +def _delete_other_models(kind: MediaKind) -> Delete: + embeddings = _KINDS[kind].embeddings + return delete(embeddings).where(embeddings.c.model_id != bindparam("model_id")) + + +@functools.cache +def _unembedded(kind: MediaKind) -> Select[Any]: + """The kind's oldest eligible items without an embedding under the model.""" + the_kind = _KINDS[kind] + media = the_kind.media + return ( + select(media.c[the_kind.name], media.c.created_at) + .select_from(_with_embedding(the_kind)) + .where(*_eligible(the_kind), the_kind.embeddings.c[the_kind.name].is_(None)) + .order_by(media.c.created_at, media.c[the_kind.name]) + .limit(bindparam("limit")) + ) + + +@functools.cache +def _status(kind: MediaKind) -> Select[Any]: + """The kind's eligible items, and how many of them have an embedding under the model.""" + the_kind = _KINDS[kind] + return ( + select(func.count(), func.count(the_kind.embeddings.c[the_kind.name])) + .select_from(_with_embedding(the_kind)) + .where(*_eligible(the_kind)) + ) + + +@functools.cache +def _accessible(kind: MediaKind, scoped: bool) -> Select[Any]: + """The kind's embedded items an account can list, by name: as the gallery's listing of every image, items on + archived boards are hidden from every scope, and `scoped` lists only the account's own items on no board and those + on boards it may read.""" + the_kind = _KINDS[kind] + media, embeddings, membership, name = the_kind.media, the_kind.embeddings, the_kind.membership, the_kind.name + listed_board = [boards.c.archived == false()] + if scoped: + listed_board.append(readable_board(bindparam("user_id"))) + on_listed_board = exists(select(literal(1)).where(boards.c.board_id == membership.c.board_id, *listed_board)) + on_no_board = membership.c.board_id.is_(None) + if scoped: + on_no_board = and_(on_no_board, media.c.user_id == bindparam("user_id")) + source = embeddings.join(media, media.c[name] == embeddings.c[name]).outerjoin( + membership, membership.c[name] == media.c[name] + ) + # Each item once without DISTINCT: an item is on one board at most, and has one embedding under the model. + return ( + select(embeddings.c[name]) + .select_from(source) + .where(embeddings.c.model_id == bindparam("model_id"), *_eligible(the_kind), or_(on_no_board, on_listed_board)) + .order_by(embeddings.c[name]) + ) + + +@functools.cache +def _set_projection(dialect_name: str) -> Any: + replaced = ["scope_hash", "params", "point_count", "image_names", "item_kinds", "coords", "updated_at"] + return upsert(dialect_name, image_projections, update=replaced) + + +@functools.cache +def _insert_vocab_term(dialect_name: str) -> Any: + # Callers pass distinct terms, but the case-insensitive key would make a pair differing in case fail the replace. + return insert_ignore(dialect_name, image_index_vocab_terms) + + +class ProjectionRow(NamedTuple): + """A cached projection as stored: its item names and kinds as JSON arrays, its coordinates as float32 bytes.""" + + scope_hash: str + params: str + point_count: int + image_names: str + # None for a projection cached before videos were indexed, all of whose items are images. + item_kinds: Optional[str] + coords: bytes + created_at: str + updated_at: str + + +def _found_embeddings(rows: Sequence[Sequence[Any]]) -> dict[IndexedItem, tuple[int, bytes, str]]: + return {IndexedItem(kind, name): (dim, embedding, encoding) for kind, name, dim, embedding, encoding in rows} + + +def _items(names: list[tuple[MediaKind, Sequence[str]]]) -> list[IndexedItem]: + return [IndexedItem(kind, name) for kind, kind_names in names for name in kind_names] + + +class ImageIndexQueries(QueryModule): + @write + def upsert_embedding( + self, conn: Connection, item: IndexedItem, model_id: str, dim: int, embedding: bytes, encoding: str + ) -> bool: + """Stores the item's embedding, in the float type `encoding` names, under the model, replacing one stored + before. Whether the item exists: an item deleted meanwhile gets none.""" + if conn.execute(_lock_media(item.kind), {"name": item.name}).first() is None: + return False + values = { + _KINDS[item.kind].name: item.name, + "model_id": model_id, + "dim": dim, + "embedding": embedding, + "encoding": encoding, + } + conn.execute(_upsert_embedding(item.kind, conn.dialect.name), values) + return True + + @mapped(_found_embeddings) + @read + def embeddings(self, conn: Connection, items: Sequence[IndexedItem], model_id: str) -> Sequence[Row[Any]]: + """The dimension, bytes and encoding of each item's embedding under the model, for the items that have one.""" + if not items: + return [] + parameters = {f"{kind}_names": bound_set(item.name for item in items if item.kind == kind) for kind in _KINDS} + return conn.execute(_EMBEDDINGS, {"model_id": model_id, **parameters}).all() + + @write + def delete_embedding(self, conn: Connection, item: IndexedItem) -> None: + """Deletes the item's embeddings under every model.""" + conn.execute(_delete_embedding(item.kind), {"name": item.name}) + + @write + def delete_embeddings_for_other_models(self, conn: Connection, model_id: str) -> int: + """Deletes the embeddings of every model but this one; how many.""" + return sum(conn.execute(_delete_other_models(kind), {"model_id": model_id}).rowcount for kind in _KINDS) + + @read + def unembedded(self, conn: Connection, model_id: str, limit: int) -> list[IndexedItem]: + """The `limit` oldest eligible items of both kinds without an embedding under the model, oldest first.""" + candidates: list[tuple[str, str, IndexedItem]] = [] + for kind in _KINDS: + rows = conn.execute(_unembedded(kind), {"model_id": model_id, "limit": limit}).all() + candidates.extend((created_at, name, IndexedItem(kind, name)) for name, created_at in rows) + # Each kind's own oldest `limit` hold the oldest `limit` of both. The sort is stable: images first on a tie. + candidates.sort(key=lambda candidate: (candidate[0], candidate[1])) + return [item for _, _, item in candidates[:limit]] + + @read + def status(self, conn: Connection, model_id: str) -> tuple[int, int]: + """The eligible items of both kinds, and how many of them have an embedding under the model.""" + total = embedded = 0 + for kind in _KINDS: + kind_total, kind_embedded = conn.execute(_status(kind), {"model_id": model_id}).one() + total += kind_total + embedded += kind_embedded + return total, embedded + + @mapped(_items) + @read + def accessible_embedded( + self, conn: Connection, user_id: Optional[str], model_id: str + ) -> list[tuple[MediaKind, Sequence[str]]]: + """The embedded items the account can list, or with None every item not on an archived board: images by name, + then videos by name.""" + parameters = {"model_id": model_id, "user_id": user_id} + return [ + (kind, conn.execute(_accessible(kind, user_id is not None), parameters).scalars().all()) for kind in _KINDS + ] + + @read + def vocab_terms(self, conn: Connection) -> list[str]: + """The custom vocabulary, in the case-insensitive order of its terms.""" + return list(conn.execute(_VOCAB_TERMS).scalars().all()) + + @write + def replace_vocab_terms(self, conn: Connection, terms: Sequence[str]) -> None: + """Replaces the custom vocabulary with these terms; of terms differing only in case, the first is kept.""" + conn.execute(_CLEAR_VOCAB_TERMS) + if terms: + # Stamped once: the column's default would take the time for every term. + now = now_text() + rows = [{"term": term, "created_at": now} for term in terms] + conn.execute(_insert_vocab_term(conn.dialect.name), rows) + + @read + def projection(self, conn: Connection, user_id: str, model_id: str) -> Optional[ProjectionRow]: + row = conn.execute(_PROJECTION, {"user_id": user_id, "model_id": model_id}).first() + return ProjectionRow(*row) if row is not None else None + + @write + def set_projection( + self, + conn: Connection, + *, + user_id: str, + model_id: str, + scope_hash: str, + params: str, + point_count: int, + image_names: str, + item_kinds: str, + coords: bytes, + ) -> bool: + """Stores the account's projection, replacing one stored before. Whether the account exists: one deleted + meanwhile gets none.""" + if conn.execute(_LOCK_USER, {"user_id": user_id}).first() is None: + return False + now = now_text() + values = { + "user_id": user_id, + "model_id": model_id, + "scope_hash": scope_hash, + "params": params, + "point_count": point_count, + "image_names": image_names, + "item_kinds": item_kinds, + "coords": coords, + # One stamp for both: on a first insert the projection was created when it was last updated. + "created_at": now, + "updated_at": now, + } + conn.execute(_set_projection(conn.dialect.name), values) + return True + + @write + def delete_projection(self, conn: Connection, user_id: str, model_id: str) -> None: + conn.execute(_DELETE_PROJECTION, {"user_id": user_id, "model_id": model_id}) diff --git a/invokeai/app/services/shared/database/queries/image_moves.py b/invokeai/app/services/shared/database/queries/image_moves.py new file mode 100644 index 00000000000..2ce75c8ab75 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/image_moves.py @@ -0,0 +1,298 @@ +"""Image moves: jobs that move image files between subfolders, journaled per image so that a move interrupted +anywhere can be finished, and the images' records repointed, after a restart. + +A job is active until it is committed or failed (`error`); there is at most one active job. +""" + +from collections.abc import Sequence +from typing import Any, NamedTuple, Optional + +from sqlalchemy import ( + Connection, + Row, + and_, + bindparam, + func, + insert, + literal, + or_, + select, + update, +) + +from invokeai.app.services.shared.database.dialect import fixed_limit +from invokeai.app.services.shared.database.queries.base import QueryModule, mapped, read, write +from invokeai.app.services.shared.database.schema.image_moves import ( + image_subfolder_move_items, + image_subfolder_move_jobs, +) +from invokeai.app.services.shared.database.schema.images import images + +_J = image_subfolder_move_jobs.c +_M = image_subfolder_move_items.c +_I = images.c + +_FINISHED_JOB = _J.state.in_([literal("committed"), literal("error")]) +_FINISHED_ITEM = _M.state.in_([literal("committed"), literal("error")]) + +_RECOVERABLE_JOB_IDS = ( + select(_J.id).where(_J.state.in_([literal("planned"), literal("moving"), literal("moved")])).order_by(_J.id) +) +_ACTIVE_JOB_ID = select(_J.id).where(~_FINISHED_JOB).order_by(_J.id).limit(fixed_limit(1)) +_JOB = select(_J.id, _J.state, _J.error_message).where(_J.id == bindparam("job")) +_LATEST_JOB = select(_J.id, _J.state, _J.error_message).order_by(_J.id.desc()).limit(fixed_limit(1)) +_INSERT_JOB = insert(image_subfolder_move_jobs) +_INSERT_ITEMS = insert(image_subfolder_move_items) +_SET_JOB_STATE = update(image_subfolder_move_jobs).where(_J.id == bindparam("job")).values(state=bindparam("new_state")) +_SET_JOB_ERROR = ( + update(image_subfolder_move_jobs).where(_J.id == bindparam("job")).values(error_message=bindparam("message")) +) +_FAIL_JOB = ( + update(image_subfolder_move_jobs) + .where(_J.id == bindparam("job")) + .values(state=literal("error"), error_message=bindparam("message")) +) +_FINISH_JOB = ( + update(image_subfolder_move_jobs) + .where(_J.id == bindparam("job")) + .values(state=bindparam("new_state"), error_message=bindparam("message")) +) + +_THE_ITEM = and_(_M.job_id == bindparam("job"), _M.image_name == bindparam("image")) +_ITEM_COLUMNS = (_M.image_name, _M.old_subfolder, _M.new_subfolder, _M.is_intermediate) +_ITEMS = select(*_ITEM_COLUMNS).where(_M.job_id == bindparam("job")).order_by(_M.image_name) +_UNFINISHED_ITEMS = _ITEMS.where(~_FINISHED_ITEM) +_MARK_ITEM_MOVED = update(image_subfolder_move_items).where(_THE_ITEM).values(state=literal("moved")) +_FAIL_ITEM = ( + update(image_subfolder_move_items) + .where(_THE_ITEM) + .values(state=literal("error"), error_message=bindparam("message")) +) +_FAIL_ITEMS = ( + update(image_subfolder_move_items) + .where(_M.job_id == bindparam("job")) + .values(state=literal("error"), error_message=bindparam("message")) +) +_ITEM_ERRORS = ( + select(func.count()) + .select_from(image_subfolder_move_items) + .where(_M.job_id == bindparam("job"), _M.state == literal("error")) +) +_ERROR_MESSAGES = ( + select(_M.error_message).where(_M.job_id == bindparam("job"), _M.state == literal("error")).order_by(_M.image_name) +) +_MOVED_COUNT = ( + select(func.count()) + .select_from(image_subfolder_move_items) + .where(_M.job_id == bindparam("job"), _M.state == literal("moved")) +) +_COMMIT_ITEMS = ( + update(image_subfolder_move_items) + .where(_M.job_id == bindparam("job"), _M.state == literal("moved")) + .values(state=literal("committed")) +) +# The moved items whose image is gone, deleted or not where the move put it. +_INVALID_MOVES = ( + select(func.count()) + .select_from(image_subfolder_move_items.outerjoin(images, _I.image_name == _M.image_name)) + .where( + _M.job_id == bindparam("job"), + _M.state == literal("moved"), + or_(_I.image_name.is_(None), _I.deleted_at.is_not(None), _I.image_subfolder != _M.new_subfolder), + ) +) +_ACTIVE_JOB_FOR_IMAGE = ( + select(literal(1)) + .select_from(image_subfolder_move_items.join(image_subfolder_move_jobs, _J.id == _M.job_id)) + .where(_M.image_name == bindparam("image"), ~_FINISHED_JOB) + .limit(fixed_limit(1)) +) + +_LIVE_IMAGE = _I.deleted_at.is_(None) +_PLACEMENT = (_I.image_name, _I.image_subfolder, _I.image_category, _I.is_intermediate, _I.created_at) +_PLACEMENTS = select(*_PLACEMENT).where(_LIVE_IMAGE) +_PLACEMENTS_AFTER = ( + select(*_PLACEMENT) + .where(_I.image_name > bindparam("after"), _LIVE_IMAGE) + .order_by(_I.image_name) + .limit(bindparam("limit")) +) +_NEXT_IMAGE_NAME = ( + select(_I.image_name) + .where(_I.image_name > bindparam("after"), _LIVE_IMAGE) + .order_by(_I.image_name) + .limit(fixed_limit(1)) +) +# Repoints an image the move found at its new subfolder, unless something else repointed it meanwhile. +_REPOINT = ( + update(images) + .where(_I.image_name == bindparam("target"), _I.image_subfolder == bindparam("old_subfolder")) + .values(image_subfolder=bindparam("new_subfolder")) +) +# The same for every image a job moved, in one statement that goes from the job's items to their images by key +# (`UPDATE ... FROM` on SQLite, a multi-table UPDATE on a server). With subqueries on the items instead, MariaDB +# scans every image; with a statement per image, a server takes a round trip per image. +_REPOINT_MOVED = ( + update(images) + .where( + _M.job_id == bindparam("job"), + _M.state == literal("moved"), + _I.image_name == _M.image_name, + _I.image_subfolder == _M.old_subfolder, + ) + .values(image_subfolder=_M.new_subfolder) +) + + +class MoveJob(NamedTuple): + id: int + state: str + error_message: Optional[str] + + +class MoveItem(NamedTuple): + image_name: str + old_subfolder: str + new_subfolder: str + is_intermediate: bool + old_path: str + new_path: str + old_thumbnail_path: str + new_thumbnail_path: str + + +# Where an image's files are, and what decides where they belong: (image name, subfolder, category, intermediate, +# created at). Plain tuples, as status polls read every image. +Placement = tuple[str, str, str, bool, str] + + +def _placements(rows: Sequence[Row[Any]]) -> list[Placement]: + return [(row[0], row[1], row[2], bool(row[3]), row[4]) for row in rows] + + +def _items(rows: Sequence[Row[Any]]) -> list[tuple[str, str, str, bool]]: + return [(row[0], row[1], row[2], bool(row[3])) for row in rows] + + +def _item_values(job_id: int, item: MoveItem, state: str, error_message: Optional[str] = None) -> dict[str, Any]: + return {**item._asdict(), "job_id": job_id, "state": state, "error_message": error_message} + + +class ImageMoveQueries(QueryModule): + @read + def recoverable_job_ids(self, conn: Connection) -> list[int]: + """The jobs a restart finishes: planned, or interrupted while moving or before their records were repointed.""" + return [int(job_id) for job_id in conn.execute(_RECOVERABLE_JOB_IDS).scalars().all()] + + @read + def active_job_id(self, conn: Connection) -> Optional[int]: + return conn.execute(_ACTIVE_JOB_ID).scalar() + + @read + def job(self, conn: Connection, job_id: int) -> Optional[MoveJob]: + row = conn.execute(_JOB, {"job": job_id}).first() + return MoveJob(row[0], row[1], row[2]) if row is not None else None + + @read + def latest_job(self, conn: Connection) -> Optional[MoveJob]: + row = conn.execute(_LATEST_JOB).first() + return MoveJob(row[0], row[1], row[2]) if row is not None else None + + @mapped(_placements) + @read + def placements(self, conn: Connection) -> Sequence[Row[Any]]: + """Every live image's placement.""" + return conn.execute(_PLACEMENTS).all() + + @mapped(_placements) + @read + def placements_after(self, conn: Connection, after: str, limit: int) -> Sequence[Row[Any]]: + """Up to `limit` live images' placements, by name, after the name `after`.""" + return conn.execute(_PLACEMENTS_AFTER, {"after": after, "limit": limit}).all() + + @read + def next_image_name(self, conn: Connection, after: str) -> Optional[str]: + return conn.execute(_NEXT_IMAGE_NAME, {"after": after}).scalar() + + @write + def create_job(self, conn: Connection, items: Sequence[MoveItem]) -> Optional[int]: + """Journals a planned job of these items; its id, or None while another job is active. The move service + serialises its callers, so the check and the insert need no lock.""" + if conn.execute(_ACTIVE_JOB_ID).first() is not None: + return None + job_id = int(conn.execute(_INSERT_JOB, {"state": "planned"}).inserted_primary_key[0]) + conn.execute(_INSERT_ITEMS, [_item_values(job_id, item, "planned") for item in items]) + return job_id + + @write + def create_failed_job(self, conn: Connection, item: MoveItem, message: str) -> int: + """Journals a job that failed before it started, for its one item.""" + job_id = int(conn.execute(_INSERT_JOB, {"state": "error", "error_message": message}).inserted_primary_key[0]) + conn.execute(_INSERT_ITEMS, [_item_values(job_id, item, "error", message)]) + return job_id + + @mapped(_items) + @read + def items(self, conn: Connection, job_id: int, *, unfinished_only: bool) -> Sequence[Row[Any]]: + """The job's items as (image name, old subfolder, new subfolder, intermediate), by name; with + `unfinished_only`, those neither committed nor failed.""" + return conn.execute(_UNFINISHED_ITEMS if unfinished_only else _ITEMS, {"job": job_id}).all() + + @read + def has_active_job_for_image(self, conn: Connection, image_name: str) -> bool: + return conn.execute(_ACTIVE_JOB_FOR_IMAGE, {"image": image_name}).first() is not None + + @read + def error_count(self, conn: Connection, job_id: int) -> int: + return int(conn.execute(_ITEM_ERRORS, {"job": job_id}).scalar_one()) + + @write + def set_job_state(self, conn: Connection, job_id: int, state: str) -> None: + conn.execute(_SET_JOB_STATE, {"job": job_id, "new_state": state}) + + @write + def set_job_error_message(self, conn: Connection, job_id: int, message: str) -> None: + conn.execute(_SET_JOB_ERROR, {"job": job_id, "message": message}) + + @write + def fail_job(self, conn: Connection, job_id: int, message: str) -> None: + """Fails the job and every one of its items with the message.""" + conn.execute(_FAIL_JOB, {"job": job_id, "message": message}) + conn.execute(_FAIL_ITEMS, {"job": job_id, "message": message}) + + @write + def mark_item_moved(self, conn: Connection, job_id: int, image_name: str) -> None: + conn.execute(_MARK_ITEM_MOVED, {"job": job_id, "image": image_name}) + + @write + def fail_item(self, conn: Connection, job_id: int, image_name: str, message: str) -> None: + conn.execute(_FAIL_ITEM, {"job": job_id, "image": image_name, "message": message}) + + @write + def repoint_image(self, conn: Connection, image_name: str, old_subfolder: str, new_subfolder: str) -> None: + """Points the image's record at its new subfolder, if it still points at the old one.""" + conn.execute(_REPOINT, {"target": image_name, "old_subfolder": old_subfolder, "new_subfolder": new_subfolder}) + + @write + def repoint_moved_images(self, conn: Connection, job_id: int) -> int: + """Points the records of the job's moved images at their new subfolders, unless something else repointed + them meanwhile; how many items the job moved.""" + conn.execute(_REPOINT_MOVED, {"job": job_id}) + return int(conn.execute(_MOVED_COUNT, {"job": job_id}).scalar_one()) + + @read + def invalid_move_count(self, conn: Connection, job_id: int) -> int: + """How many of the job's moved items do not have their image where the move put it.""" + return int(conn.execute(_INVALID_MOVES, {"job": job_id}).scalar_one()) + + @read + def error_messages(self, conn: Connection, job_id: int) -> list[Optional[str]]: + """The messages of the job's failed items, by image name; None for one without a message.""" + return list(conn.execute(_ERROR_MESSAGES, {"job": job_id}).scalars().all()) + + @write + def finish_job(self, conn: Connection, job_id: int, *, error_message: Optional[str]) -> None: + """Commits the job's moved items, and the job, or fails the job with the message.""" + conn.execute(_COMMIT_ITEMS, {"job": job_id}) + state = "committed" if error_message is None else "error" + conn.execute(_FINISH_JOB, {"job": job_id, "new_state": state, "message": error_message}) diff --git a/invokeai/app/services/shared/database/queries/images.py b/invokeai/app/services/shared/database/queries/images.py new file mode 100644 index 00000000000..4ac4580dec6 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/images.py @@ -0,0 +1,529 @@ +"""Image records: one row per image file, with its generation metadata, listed by the gallery.""" + +import functools +import itertools +from collections.abc import Sequence +from typing import Any, NamedTuple, Optional + +from sqlalchemy import ( + ColumnElement, + Connection, + Row, + Select, + Update, + and_, + bindparam, + delete, + exists, + false, + func, + literal, + or_, + select, + true, + update, +) + +from invokeai.app.services.image_records.image_records_common import ( + ImageCategory, + ImageRecord, + ImageRecordChanges, + ResourceOrigin, + deserialize_image_record, +) +from invokeai.app.services.shared.database.dialect import ( + CaseInsensitiveLike, + fixed_limit, + insert_ignore, + like_contains, +) +from invokeai.app.services.shared.database.queries.base import IN_CHUNK, QueryModule, day_after, mapped, read, write +from invokeai.app.services.shared.database.queries.board_access import readable_board +from invokeai.app.services.shared.database.schema.boards import board_images, boards +from invokeai.app.services.shared.database.schema.images import images + +_I = images.c +_BI = board_images.c +_B = boards.c + +# The columns of an `ImageRecord`, which has no metadata and no owner. +_RECORD = ( + _I.image_name, + _I.image_origin, + _I.image_category, + _I.width, + _I.height, + _I.session_id, + _I.node_id, + _I.has_workflow, + _I.is_intermediate, + _I.created_at, + _I.updated_at, + _I.deleted_at, + _I.starred, + _I.image_subfolder, + _I.project_id, + _I.file_size_bytes, +) +_RECORD_NAMES = tuple(column.name for column in _RECORD) + +_GET = select(*_RECORD).where(_I.image_name == bindparam("image_name")) +_GET_USER_ID = select(_I.user_id).where(_I.image_name == bindparam("image_name")) +_GET_METADATA = select(_I.metadata).where(_I.image_name == bindparam("image_name")) +_GET_CREATED_AT = select(_I.created_at).where(_I.image_name == bindparam("image_name")) +_EXISTS = select(literal(1)).where(_I.image_name == bindparam("image_name")) +_SUBFOLDERS = select(_I.image_name, _I.image_subfolder).where( + _I.image_name.in_(bindparam("image_names", expanding=True)) +) +_MOST_RECENT_ON_BOARD = ( + select(*_RECORD) + .join(board_images, _BI.image_name == _I.image_name) + .where(_BI.board_id == bindparam("board_id"), _I.is_intermediate == false()) + .order_by(_I.starred.desc(), _I.created_at.desc(), _I.image_name.desc()) + .limit(fixed_limit(1)) +) +_SET_FILE_SIZE = ( + update(images) + .where(_I.image_name == bindparam("target_image_name")) + .values(file_size_bytes=bindparam("new_file_size_bytes")) +) +# A backfill fills gaps only: the writer's own measurement, taken after the file exists, wins over one taken before. +_FILL_FILE_SIZE = ( + update(images) + .where(_I.image_name == bindparam("target_image_name"), _I.file_size_bytes.is_(None)) + .values(file_size_bytes=bindparam("new_file_size_bytes")) +) +_DELETE = delete(images).where(_I.image_name == bindparam("image_name")) +# Every image's location, a page at a time in name order, intermediates included. +_LOCATIONS_AFTER = ( + select(_I.image_name, _I.image_subfolder) + .where(_I.image_name > bindparam("after")) + .order_by(_I.image_name) + .limit(bindparam("limit")) +) +_DELETE_MANY = delete(images).where(_I.image_name.in_(bindparam("image_names", expanding=True))) +# Locked, so that an image promoted out of the intermediates meanwhile is either seen promoted or deleted before it is. +_LOCK_INTERMEDIATES = ( + select(_I.image_name) + .where(_I.image_name.in_(bindparam("image_names", expanding=True)), _I.is_intermediate == true()) + .with_for_update() +) +_DELETE_INTERMEDIATES = delete(images).where( + _I.image_name.in_(bindparam("image_names", expanding=True)), _I.is_intermediate == true() +) + + +@functools.cache +def _insert(dialect_name: str) -> Any: + # A name that is taken keeps its row, as SQLite's INSERT OR IGNORE did. + return insert_ignore(dialect_name, images) + + +@functools.cache +def _update(fields: tuple[str, ...]) -> Update: + """Sets these fields of the image; one statement per set of fields, of which there are few.""" + return ( + update(images) + .where(_I.image_name == bindparam("target_image_name")) + .values({field: bindparam(f"new_{field}") for field in fields}) + ) + + +class _Shape(NamedTuple): + """Which filters a listing has: never their values, which are bound, so that the shapes are few.""" + + origin: bool + # The categories' values, sorted; None for no category filter. Each value is bound on its own: SQLAlchemy compiles + # lists of one length once, whatever they hold. + categories: Optional[tuple[str, ...]] + intermediate: bool + # "any" (no board filter), "none" (on no board), "all" (on no board or a readable one) or "one". + board: str + # Whether the listing is limited to what one account may see; else everything, as for an administrator. + scoped: bool + search: bool + created_from: bool + created_to: bool + + +def _shape( + *, + image_origin: Optional[ResourceOrigin], + categories: Optional[Sequence[ImageCategory]], + is_intermediate: Optional[bool], + board_id: Optional[str], + search_term: Optional[str], + created_from: Optional[str], + created_to: Optional[str], + user_id: Optional[str], + is_admin: bool, +) -> _Shape: + board = "any" if board_id is None else board_id if board_id in ("none", "all") else "one" + return _Shape( + origin=image_origin is not None, + categories=None if categories is None else tuple(sorted({ImageCategory(c).value for c in categories})), + intermediate=is_intermediate is not None, + board=board, + scoped=user_id is not None and not is_admin, + search=bool(search_term), + created_from=created_from is not None, + created_to=created_to is not None, + ) + + +def _parameters( + *, + image_origin: Optional[ResourceOrigin], + is_intermediate: Optional[bool], + board_id: Optional[str], + search_term: Optional[str], + created_from: Optional[str], + created_to: Optional[str], + user_id: Optional[str], +) -> dict[str, Any]: + return { + "image_origin": image_origin.value if image_origin is not None else None, + "is_intermediate": is_intermediate, + "board_id": board_id, + # Lowered here, as the legacy storage did: SQLite's LIKE folds the case of ASCII letters only. + "pattern": like_contains(search_term.lower()) if search_term else None, + "created_from": created_from, + "created_before": day_after(created_to) if created_to is not None else None, + "user_id": user_id, + } + + +def _listed_board(scoped: bool) -> list[ColumnElement[bool]]: + """The conditions on a board whose images the scope lists: active, and with `scoped` one the account may read + (its own, shared with everyone, or shared with the account).""" + if not scoped: + return [_B.archived == false()] + return [_B.archived == false(), readable_board(bindparam("user_id"))] + + +def _conditions(shape: _Shape, joined: bool) -> list[ColumnElement[bool]]: + """The listing's conditions. `joined` lists through a LEFT JOIN of the board memberships, for the board data a + page of records returns; else the memberships are checked per image, so that SQLite scans the images alone.""" + conditions: list[ColumnElement[bool]] = [] + if shape.origin: + conditions.append(_I.image_origin == bindparam("image_origin")) + if shape.categories is not None: + conditions.append(_I.image_category.in_([literal(category) for category in shape.categories])) + if shape.intermediate: + conditions.append(_I.is_intermediate == bindparam("is_intermediate")) + + own = _I.user_id == bindparam("user_id") + membership = select(literal(1)).where(_BI.image_name == _I.image_name) + on_no_board = _BI.board_id.is_(None) if joined else ~exists(membership) + if shape.board == "none": + conditions.append(on_no_board) + if shape.scoped: + conditions.append(own) + elif shape.board == "all": + if joined: + board = exists(select(literal(1)).where(_B.board_id == _BI.board_id, *_listed_board(shape.scoped))) + else: + memberships = board_images.join(boards, _B.board_id == _BI.board_id) + board = exists( + select(literal(1)) + .select_from(memberships) + .where(_BI.image_name == _I.image_name, *_listed_board(shape.scoped)) + ) + conditions.append(or_(and_(on_no_board, own) if shape.scoped else on_no_board, board)) + elif shape.board == "one": + on_the_board = _BI.board_id == bindparam("board_id") + conditions.append(on_the_board if joined else exists(membership.where(on_the_board))) + elif shape.scoped: + conditions.append(own) + + if shape.search: + pattern: ColumnElement[str] = bindparam("pattern") + conditions.append(or_(CaseInsensitiveLike(_I.metadata, pattern), CaseInsensitiveLike(_I.created_at, pattern))) + if shape.created_from: + conditions.append(_I.created_at >= bindparam("created_from")) + if shape.created_to: + conditions.append(_I.created_at < bindparam("created_before")) + return conditions + + +def _ordering(starred_first: bool, descending: bool) -> list[ColumnElement[Any]]: + # The name breaks ties, so that images of the same moment keep one order from page to page. + ordering: list[ColumnElement[Any]] = ( + [_I.created_at.desc(), _I.image_name.desc()] if descending else [_I.created_at, _I.image_name] + ) + return [_I.starred.desc(), *ordering] if starred_first else ordering + + +_JOINED_FROM = images.outerjoin(board_images, _BI.image_name == _I.image_name) + + +@functools.lru_cache(maxsize=256) +def _page(shape: _Shape, starred_first: bool, descending: bool) -> Select[Any]: + """A page of records; the number of shapes is bounded, the cache keeps the recent ones.""" + return ( + select(*_RECORD) + .select_from(_JOINED_FROM) + .where(*_conditions(shape, joined=True)) + .order_by(*_ordering(starred_first, descending)) + .limit(bindparam("limit")) + .offset(bindparam("offset")) + ) + + +@functools.lru_cache(maxsize=256) +def _count(shape: _Shape) -> Select[Any]: + return select(func.count()).select_from(_JOINED_FROM).where(*_conditions(shape, joined=True)) + + +@functools.lru_cache(maxsize=256) +def _names(shape: _Shape, starred_first: bool, descending: bool) -> Select[Any]: + return ( + select(_I.image_name).where(*_conditions(shape, joined=False)).order_by(*_ordering(starred_first, descending)) + ) + + +@functools.lru_cache(maxsize=256) +def _starred_count(shape: _Shape) -> Select[Any]: + return select(func.count()).select_from(images).where(_I.starred == true(), *_conditions(shape, joined=False)) + + +def _record(row: Sequence[Any]) -> ImageRecord: + return deserialize_image_record(dict(zip(_RECORD_NAMES, row, strict=True))) + + +def _record_or_none(row: Optional[Sequence[Any]]) -> Optional[ImageRecord]: + return _record(row) if row is not None else None + + +def _records(page: tuple[Sequence[Sequence[Any]], int]) -> tuple[list[ImageRecord], int]: + rows, total = page + return [_record(row) for row in rows], total + + +class ImageQueries(QueryModule): + @mapped(_record_or_none) + @read + def get(self, conn: Connection, image_name: str) -> Optional[Row[Any]]: + return conn.execute(_GET, {"image_name": image_name}).first() + + @read + def user_id(self, conn: Connection, image_name: str) -> Optional[str]: + """The image's owner; None also for an image that does not exist.""" + return conn.execute(_GET_USER_ID, {"image_name": image_name}).scalar() + + @read + def metadata(self, conn: Connection, image_name: str) -> tuple[bool, Optional[str]]: + """Whether the image exists, and its metadata as stored: JSON text, or None.""" + row = conn.execute(_GET_METADATA, {"image_name": image_name}).first() + return (False, None) if row is None else (True, row[0]) + + @read + def exists(self, conn: Connection, image_name: str) -> bool: + return conn.execute(_EXISTS, {"image_name": image_name}).first() is not None + + @read + def subfolders(self, conn: Connection, image_names: Sequence[str]) -> dict[str, str]: + """The subfolder of each named image that exists.""" + subfolders: dict[str, str] = {} + for chunk in itertools.batched(image_names, IN_CHUNK): + subfolders.update((row[0], row[1]) for row in conn.execute(_SUBFOLDERS, {"image_names": list(chunk)}).all()) + return subfolders + + @read + def locations_after(self, conn: Connection, after: str, limit: int) -> list[tuple[str, str]]: + """The name and subfolder of up to `limit` images, in name order, after the image named `after`.""" + rows = conn.execute(_LOCATIONS_AFTER, {"after": after, "limit": limit}).all() + return [(str(name), str(subfolder)) for name, subfolder in rows] + + @mapped(_record_or_none) + @read + def most_recent_on_board(self, conn: Connection, board_id: str) -> Optional[Row[Any]]: + """The board's cover: its newest image that is not an intermediate, a starred one first.""" + return conn.execute(_MOST_RECENT_ON_BOARD, {"board_id": board_id}).first() + + @mapped(_records) + @read + def page( + self, + conn: Connection, + *, + offset: int, + limit: int, + starred_first: bool, + descending: bool, + image_origin: Optional[ResourceOrigin], + categories: Optional[Sequence[ImageCategory]], + is_intermediate: Optional[bool], + board_id: Optional[str], + search_term: Optional[str], + created_from: Optional[str], + created_to: Optional[str], + user_id: Optional[str], + is_admin: bool, + ) -> tuple[Sequence[Row[Any]], int]: + """A page of the records that match, and how many match in all.""" + shape = _shape( + image_origin=image_origin, + categories=categories, + is_intermediate=is_intermediate, + board_id=board_id, + search_term=search_term, + created_from=created_from, + created_to=created_to, + user_id=user_id, + is_admin=is_admin, + ) + parameters = _parameters( + image_origin=image_origin, + is_intermediate=is_intermediate, + board_id=board_id, + search_term=search_term, + created_from=created_from, + created_to=created_to, + user_id=user_id, + ) + rows = conn.execute( + _page(shape, starred_first, descending), {**parameters, "limit": limit, "offset": offset} + ).all() + return rows, conn.execute(_count(shape), parameters).scalar_one() + + @read + def names( + self, + conn: Connection, + *, + starred_first: bool, + descending: bool, + image_origin: Optional[ResourceOrigin], + categories: Optional[Sequence[ImageCategory]], + is_intermediate: Optional[bool], + board_id: Optional[str], + search_term: Optional[str], + created_from: Optional[str], + created_to: Optional[str], + user_id: Optional[str], + is_admin: bool, + ) -> tuple[list[str], int]: + """The names of the images that match, in order, and how many of them are starred (with `starred_first`; + else 0).""" + shape = _shape( + image_origin=image_origin, + categories=categories, + is_intermediate=is_intermediate, + board_id=board_id, + search_term=search_term, + created_from=created_from, + created_to=created_to, + user_id=user_id, + is_admin=is_admin, + ) + parameters = _parameters( + image_origin=image_origin, + is_intermediate=is_intermediate, + board_id=board_id, + search_term=search_term, + created_from=created_from, + created_to=created_to, + user_id=user_id, + ) + starred = conn.execute(_starred_count(shape), parameters).scalar_one() if starred_first else 0 + names = list(conn.execute(_names(shape, starred_first, descending), parameters).scalars().all()) + return names, starred + + @write + def insert( + self, + conn: Connection, + *, + image_name: str, + image_origin: ResourceOrigin, + image_category: ImageCategory, + width: int, + height: int, + has_workflow: bool, + is_intermediate: Optional[bool], + starred: Optional[bool], + session_id: Optional[str], + node_id: Optional[str], + metadata: Optional[str], + user_id: str, + image_subfolder: str, + project_id: Optional[str], + created_at: Optional[str] = None, + ) -> None: + """Adds the image, unless an image of that name exists. It was created now, unless `created_at` (a timestamp + text, see `types.timestamp_text`) says when.""" + conn.execute( + _insert(conn.dialect.name), + { + "image_name": image_name, + "image_origin": image_origin.value, + "image_category": image_category.value, + "width": width, + "height": height, + "node_id": node_id, + "session_id": session_id, + "metadata": metadata, + "is_intermediate": is_intermediate, + "starred": starred, + "has_workflow": has_workflow, + "user_id": user_id, + "image_subfolder": image_subfolder, + "project_id": project_id, + **({"created_at": created_at} if created_at is not None else {}), + }, + ) + + @read + def created_at(self, conn: Connection, image_name: str) -> Optional[str]: + return conn.execute(_GET_CREATED_AT, {"image_name": image_name}).scalar() + + @write + def update(self, conn: Connection, image_name: str, changes: ImageRecordChanges) -> None: + """Applies the changes to the image's category, session, intermediate flag and star.""" + fields: dict[str, Any] = {} + if changes.image_category is not None: + fields["image_category"] = ImageCategory(changes.image_category).value + if changes.session_id is not None: + fields["session_id"] = changes.session_id + if changes.is_intermediate is not None: + fields["is_intermediate"] = changes.is_intermediate + if changes.starred is not None: + fields["starred"] = changes.starred + if fields: + values = {f"new_{field}": value for field, value in fields.items()} + conn.execute(_update(tuple(sorted(fields))), {"target_image_name": image_name, **values}) + + @write + def set_file_size(self, conn: Connection, image_name: str, file_size_bytes: Optional[int]) -> None: + conn.execute(_SET_FILE_SIZE, {"target_image_name": image_name, "new_file_size_bytes": file_size_bytes}) + + @write + def fill_file_sizes(self, conn: Connection, sizes: dict[str, int]) -> None: + """Records the sizes of the images whose size is not known yet.""" + if sizes: + conn.execute( + _FILL_FILE_SIZE, + [{"target_image_name": name, "new_file_size_bytes": size} for name, size in sizes.items()], + ) + + @write + def delete(self, conn: Connection, image_name: str) -> None: + conn.execute(_DELETE, {"image_name": image_name}) + + @write + def delete_many(self, conn: Connection, image_names: Sequence[str]) -> None: + for chunk in itertools.batched(image_names, IN_CHUNK): + conn.execute(_DELETE_MANY, {"image_names": list(chunk)}) + + @write + def delete_intermediates(self, conn: Connection, image_names: Sequence[str]) -> list[str]: + """Deletes those of the named images that are intermediates; the names deleted, in the given order.""" + deleted: list[str] = [] + for chunk in itertools.batched(image_names, IN_CHUNK): + names = list(chunk) + locked: set[str] = set(conn.execute(_LOCK_INTERMEDIATES, {"image_names": names}).scalars().all()) + if locked: + conn.execute(_DELETE_INTERMEDIATES, {"image_names": sorted(locked)}) + deleted.extend(name for name in names if name in locked) + return deleted diff --git a/invokeai/app/services/shared/database/queries/intermediates.py b/invokeai/app/services/shared/database/queries/intermediates.py new file mode 100644 index 00000000000..304a77639ef --- /dev/null +++ b/invokeai/app/services/shared/database/queries/intermediates.py @@ -0,0 +1,783 @@ +"""Intermediates cleanup: what every intermediate is under the cleanup policy, and the state that decides it. + +One classification expression decides it, the same whether it aggregates a summary, describes a preview, pages +an operation or guards a delete: + +- ``active``: produced or referenced by a pending, waiting or running queue item, or a completed child of an + active root workflow; held by a running session that reuses it as a cached output; or leased by a browser tab. +- ``recent``: created inside the grace window, or held by a session that ended inside it. +- ``referenced``: named by a saved document (`media_references`). +- ``safe``: none of the above. + +The caller binds the clock (`now`, `recent_cutoff`) and the names its scan of active queue sessions found +(`active_names`, a `bound_set`), so that every statement of one classification sees the same instant and the +same active work. +""" + +import functools +from collections.abc import Collection, Sequence +from typing import Any, Literal, NamedTuple, Optional + +from sqlalchemy import ( + ColumnElement, + Connection, + Select, + String, + Subquery, + Table, + and_, + bindparam, + case, + delete, + exists, + false, + func, + literal, + null, + or_, + select, + union_all, + update, +) + +from invokeai.app.services.shared.database.dialect import ( + InBoundSet, + KeysetAfter, + OrderedJoin, + bound_set, + fixed_limit, + insert_ignore, + sql_true, + upsert, +) +from invokeai.app.services.shared.database.queries.base import IN_CHUNK, QueryModule, read, write +from invokeai.app.services.shared.database.queries.session_queue import ACTIVE_QUEUE_STATUSES +from invokeai.app.services.shared.database.schema.boards import board_images +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.intermediates import ( + intermediates_browser_holds, + intermediates_session_holds, + intermediates_unmeasurable, +) +from invokeai.app.services.shared.database.schema.media_references import media_references +from invokeai.app.services.shared.database.schema.projects import orphaned_projects_2026_08_06, projects +from invokeai.app.services.shared.database.schema.session_queue import session_queue +from invokeai.app.services.shared.database.schema.videos import videos +from invokeai.app.services.shared.database.schema.workflows import workflow_library + +MediaKind = Literal["image", "video"] +Classification = Literal["safe", "referenced", "active", "recent"] +CleanupMode = Literal["safe", "force"] + +# The selected rows of a preview, per statement. A statement has a power of two of slots, the spare ones bound to +# NULL, so that a few statements serve every selection while each row is checked against few slots. +TARGET_SLOTS = 32 + +_MEDIA: dict[MediaKind, tuple[Table, str, str]] = { + "image": (images, "image_name", "image_subfolder"), + "video": (videos, "video_name", "video_subfolder"), +} +_Q = session_queue.c +_H = intermediates_session_holds.c +_B = intermediates_browser_holds.c +_U = intermediates_unmeasurable.c +_R = media_references.c +_P = projects.c + + +def _protected_items() -> Subquery: + """The active queue items, and the completed items of an active root, which keep their outputs until it ends. + + A child is completed before its outputs reach its waiting parent; the whole completed subtree stays protected + until the root finishes. The join starts from the active roots, then looks the children up by index, rather + than scanning all completed history. + """ + statuses = [literal(status) for status in ACTIVE_QUEUE_STATUSES] + root = session_queue.alias("root") + child = session_queue.alias("child") + direct = select(_Q.item_id, _Q.session_id, _Q.session_revision).where(_Q.status.in_(statuses)) + nested = ( + select(child.c.item_id, child.c.session_id, child.c.session_revision) + .select_from(OrderedJoin(root, child, child.c.root_item_id == root.c.item_id)) + .where(root.c.status.in_(statuses), child.c.status == literal("completed")) + ) + return union_all(direct, nested).subquery("protected") + + +_PROTECTED = _protected_items() +_ACTIVE_ITEMS = select(_PROTECTED.c.item_id, func.coalesce(_PROTECTED.c.session_revision, 0)) +_SESSIONS = select(_Q.item_id, func.coalesce(_Q.session_revision, 0), _Q.session).where( + _Q.item_id.in_(bindparam("item_ids", expanding=True)) +) +_SESSION_IS_ACTIVE = select(literal(1)).where( + _Q.session_id == bindparam("session_id"), _Q.status.in_([literal(status) for status in ACTIVE_QUEUE_STATUSES]) +) + +# A session hold is released the first time its session is seen finished, and dropped once its grace is over. +_RELEASE_SESSION_HOLDS = ( + update(intermediates_session_holds) + .where(_H.released_at.is_(None), _H.session_id.not_in(select(_protected_items().c.session_id))) + .values(released_at=bindparam("now")) +) +_DROP_RELEASED_SESSION_HOLDS = delete(intermediates_session_holds).where(_H.released_at <= bindparam("recent_cutoff")) +_CLEAR_SESSION_HOLDS = delete(intermediates_session_holds) + +_SWEEP_BROWSER_HOLDS = delete(intermediates_browser_holds).where(_B.expires_at <= bindparam("now")) +_RELEASE_LEASE = delete(intermediates_browser_holds).where( + _B.user_id == bindparam("user_id"), _B.lease_id == bindparam("lease_id") +) +_LEASES = select(_B.lease_id, func.max(_B.expires_at)).where(_B.user_id == bindparam("user_id")).group_by(_B.lease_id) +_RELEASE_LEASES = delete(intermediates_browser_holds).where( + _B.user_id == bindparam("user_id"), _B.lease_id.in_(bindparam("lease_ids", expanding=True)) +) + +_CLEAR_UNMEASURABLE = delete(intermediates_unmeasurable) + +_DOCUMENT_NAMES = { + "project": select(_P.user_id, _P.project_id, _P.name).where(_P.project_id.in_(bindparam("ids", expanding=True))), + "quarantined_project": select( + orphaned_projects_2026_08_06.c.user_id, + orphaned_projects_2026_08_06.c.project_id, + orphaned_projects_2026_08_06.c.name, + ).where(orphaned_projects_2026_08_06.c.project_id.in_(bindparam("ids", expanding=True))), + # Workflow ids are unique across accounts, so the owner is not part of the match. + "workflow": select(null(), workflow_library.c.workflow_id, workflow_library.c.name).where( + workflow_library.c.workflow_id.in_(bindparam("ids", expanding=True)) + ), +} + + +def _is_intermediate(media: Table) -> ColumnElement[bool]: + # As the partial indexes on intermediates spell it, so that SQLite uses them. + return media.c.is_intermediate == sql_true() + + +def _classification(kind: MediaKind) -> ColumnElement[str]: + media, name_column, _ = _MEDIA[kind] + name = media.c[name_column] + held = select(_H.media_name).where(_H.media_kind == literal(kind)) + leased = exists( + select(literal(1)).where( + _B.media_kind == literal(kind), _B.media_name == name, _B.expires_at > bindparam("now") + ) + ) + active = or_( + and_(media.c.session_id.is_not(None), media.c.session_id.in_(select(_protected_items().c.session_id))), + InBoundSet(name, bindparam("active_names", type_=String())), + name.in_(held.where(_H.released_at.is_(None))), + leased, + ) + recent = or_(media.c.created_at > bindparam("recent_cutoff"), name.in_(held.where(_H.released_at.is_not(None)))) + referenced = exists(select(literal(1)).where(_R.media_kind == literal(kind), _R.media_name == name)) + return case( + (active, literal("active")), + (recent, literal("recent")), + (referenced, literal("referenced")), + else_=literal("safe"), + ) + + +ScopeKind = Literal["all", "owner", "owners", "project", "unassigned", "targets"] + + +class ScopeShape(NamedTuple): + """What a statement's scope looks like: its kind, and for selected rows how many slots it has.""" + + kind: ScopeKind + slots: int = 0 + + +class ScopeSpec(NamedTuple): + """Which intermediates a statement covers. Statements are built per kind of scope, its values are bound. + + - "all"; + - "owner": `user_id`'s; + - "owners": those of the accounts in `users`; + - "project": the media of `targets[0]`, an (owner, project) of a project that exists; + - "unassigned": `user_id`'s media of no existing project; + - "targets": up to `TARGET_SLOTS` (owner, project) rows, a None project for the owner's unassigned row. + """ + + kind: ScopeKind + user_id: Optional[str] = None + users: frozenset[str] = frozenset() + targets: tuple[tuple[str, Optional[str]], ...] = () + + @property + def shape(self) -> ScopeShape: + if self.kind != "targets": + return ScopeShape(self.kind) + if not 0 < len(self.targets) <= TARGET_SLOTS: + raise ValueError(f"1 to {TARGET_SLOTS} rows a statement") + return ScopeShape("targets", 1 << (len(self.targets) - 1).bit_length()) + + def parameters(self) -> dict[str, Any]: + if self.kind in ("owner", "unassigned"): + return {"scope_user": self.user_id} + if self.kind == "project": + owner, project = self.targets[0] + return {"scope_user": owner, "scope_project": project} + if self.kind == "owners": + return {"scope_users": bound_set(self.users)} + if self.kind == "targets": + slots = [*self.targets, *[(None, None)] * (self.shape.slots - len(self.targets))] + parameters: dict[str, Any] = {} + for i, (owner, project) in enumerate(slots): + parameters[f"target_user_{i}"] = owner + parameters[f"target_project_{i}"] = project + return parameters + return {} + + +def _scope_conditions(shape: ScopeShape, media: Table) -> list[ColumnElement[bool]]: + scope = shape.kind + project_exists = _P.project_id.is_not(None) + if scope == "owner": + return [media.c.user_id == bindparam("scope_user")] + if scope == "owners": + return [InBoundSet(media.c.user_id, bindparam("scope_users", type_=String()))] + if scope == "project": + return [ + media.c.user_id == bindparam("scope_user"), + media.c.project_id == bindparam("scope_project"), + project_exists, + ] + if scope == "unassigned": + return [media.c.user_id == bindparam("scope_user"), _P.project_id.is_(None)] + if scope == "targets": + slots = [] + for i in range(shape.slots): + target_project = bindparam(f"target_project_{i}", type_=String()) + slots.append( + and_( + media.c.user_id == bindparam(f"target_user_{i}", type_=String()), + or_( + and_(target_project.is_(None), _P.project_id.is_(None)), + and_(media.c.project_id == target_project, project_exists), + ), + ) + ) + return [or_(*slots)] + return [] + + +def _scoped(kind: MediaKind, scope: ScopeShape) -> Subquery: + """Every intermediate of `kind` in `scope`, with its (owner, project) row and its classification. + + Media whose project no longer exists, or that never had one, belong to the owner's unassigned row. + """ + media, name_column, _ = _MEDIA[kind] + owned_project = and_(_P.user_id == media.c.user_id, _P.project_id == media.c.project_id) + return ( + select( + media.c[name_column].label("name"), + media.c.user_id.label("user_id"), + case((_P.project_id.is_(None), null()), else_=media.c.project_id).label("project_key"), + media.c.file_size_bytes.label("file_size_bytes"), + media.c.created_at.label("created_at"), + _classification(kind).label("cls"), + ) + .select_from(media.outerjoin(projects, owned_project)) + .where(_is_intermediate(media), *_scope_conditions(scope, media)) + .subquery("scoped") + ) + + +def _deletable(kind: MediaKind, classified: Subquery, mode: CleanupMode, is_admin: bool) -> ColumnElement[int]: + """1 for what a cleanup of `mode` may delete, else 0. + + Force mode adds referenced items, but another account's document is never the caller's to break unless the + caller administers the instance (`caller_user_id`). + """ + cls = classified.c.cls + if mode == "safe": + allowed: ColumnElement[bool] = cls == literal("safe") + elif is_admin: + allowed = cls.in_([literal("safe"), literal("referenced")]) + else: + others_document = exists( + select(literal(1)).where( + _R.media_kind == literal(kind), + _R.media_name == classified.c.name, + _R.user_id != bindparam("caller_user_id"), + ) + ) + allowed = or_(cls == literal("safe"), and_(cls == literal("referenced"), ~others_document)) + return case((allowed, 1), else_=0) + + +def _with_deletable(kind: MediaKind, scope: ScopeShape, mode: CleanupMode, is_admin: bool) -> Subquery: + scoped = _scoped(kind, scope) + return select(*scoped.c, _deletable(kind, scoped, mode, is_admin).label("deletable")).subquery("classified") + + +@functools.cache +def _aggregate(kind: MediaKind, scope: ScopeShape) -> Select[Any]: + s = _scoped(kind, scope) + return select( + s.c.user_id, + s.c.project_key, + s.c.cls, + func.count(), + func.sum(func.coalesce(s.c.file_size_bytes, 0)), + func.sum(case((s.c.file_size_bytes.is_(None), 1), else_=0)), + ).group_by(s.c.user_id, s.c.project_key, s.c.cls) + + +@functools.cache +def _preview(kind: MediaKind, scope: ScopeShape, mode: CleanupMode, is_admin: bool) -> Select[Any]: + c = _with_deletable(kind, scope, mode, is_admin) + return select( + c.c.user_id, + c.c.project_key, + c.c.cls, + c.c.deletable, + func.count(), + func.sum(func.coalesce(c.c.file_size_bytes, 0)), + func.sum(case((c.c.file_size_bytes.is_(None), 1), else_=0)), + ).group_by(c.c.user_id, c.c.project_key, c.c.cls, c.c.deletable) + + +@functools.cache +def _acknowledged(kind: MediaKind, scope: ScopeShape, is_admin: bool) -> Select[Any]: + """The documents naming the referenced items a force cleanup may delete, with how many each names.""" + c = _with_deletable(kind, scope, "force", is_admin) + return ( + select(_R.owner_kind, _R.user_id, _R.owner_id, func.count()) + .select_from(c.join(media_references, and_(_R.media_kind == literal(kind), _R.media_name == c.c.name))) + .where(c.c.cls == literal("referenced"), c.c.deletable == 1) + .group_by(_R.owner_kind, _R.user_id, _R.owner_id) + ) + + +@functools.cache +def _window(kind: MediaKind, scope: ScopeShape, mode: CleanupMode, is_admin: bool) -> Select[Any]: + """The next rows of a scope in creation order, after (`after_created_at`, `after_name`), classified. + + The page is chosen from the plain rows, by an index seek, before anything is classified: a server cannot merge a + derived table whose columns have subqueries, and classifying the whole remaining scope on every window would + make a cleanup quadratic in its scope. + """ + media, name_column, _ = _MEDIA[kind] + name = media.c[name_column] + after = KeysetAfter( + media.c.created_at, + name, + bindparam("after_created_at", type_=String()), + bindparam("after_name", type_=String()), + ) + owned_project = and_(_P.user_id == media.c.user_id, _P.project_id == media.c.project_id) + page = ( + select(name.label("name")) + .select_from(media.outerjoin(projects, owned_project)) + .where(_is_intermediate(media), *_scope_conditions(scope, media), after) + .order_by(media.c.created_at, name) + .limit(bindparam("limit")) + .subquery("page") + ) + classified = ( + select( + name.label("name"), + media.c.file_size_bytes.label("file_size_bytes"), + media.c.created_at.label("created_at"), + _classification(kind).label("cls"), + ) + .select_from(media.join(page, page.c.name == name)) + .subquery("classified") + ) + return select( + classified.c.name, + classified.c.file_size_bytes, + _deletable(kind, classified, mode, is_admin), + classified.c.created_at, + ).order_by(classified.c.created_at, classified.c.name) + + +@functools.cache +def _still_deletable(kind: MediaKind, mode: CleanupMode, is_admin: bool) -> Select[Any]: + """Which of the named intermediates a cleanup may delete now, with their owners.""" + media, name_column, _ = _MEDIA[kind] + name = media.c[name_column] + classified = ( + select(name.label("name"), media.c.user_id.label("user_id"), _classification(kind).label("cls")) + .where(InBoundSet(name, bindparam("names", type_=String())), _is_intermediate(media)) + .subquery("named") + ) + return select(classified.c.name, classified.c.user_id).where(_deletable(kind, classified, mode, is_admin) == 1) + + +@functools.cache +def _count_existing(kind: MediaKind) -> Select[Any]: + media, name_column, _ = _MEDIA[kind] + return select(func.count()).where(InBoundSet(media.c[name_column], bindparam("names", type_=String()))) + + +@functools.cache +def _hold_for_lease(kind: MediaKind, dialect_name: str) -> Any: + """Leases the named intermediates the account owns. A refresh of the same lease that ran meanwhile has inserted + the same rows on a server, where writers do not wait for each other here: they get the later expiry.""" + media, name_column, _ = _MEDIA[kind] + leased = select( + bindparam("user_id", type_=String()), + bindparam("lease_id", type_=String()), + literal(kind), + media.c[name_column], + bindparam("expires_at", type_=String()), + ).where( + InBoundSet(media.c[name_column], bindparam("names", type_=String())), + media.c.user_id == bindparam("user_id", type_=String()), + _is_intermediate(media), + ) + return upsert(dialect_name, intermediates_browser_holds, update=["expires_at"]).from_select( + ["user_id", "lease_id", "media_kind", "media_name", "expires_at"], leased + ) + + +@functools.cache +def _reference_owners(kind: MediaKind) -> Select[Any]: + return select(_R.media_name, _R.owner_kind, _R.user_id, _R.owner_id).where( + _R.media_kind == literal(kind), InBoundSet(_R.media_name, bindparam("names", type_=String())) + ) + + +def _unmeasured(kind: MediaKind) -> tuple[Table, list[ColumnElement[bool]]]: + media, name_column, _ = _MEDIA[kind] + marked = exists(select(literal(1)).where(_U.media_kind == literal(kind), _U.media_name == media.c[name_column])) + return media, [_is_intermediate(media), media.c.file_size_bytes.is_(None), ~marked] + + +@functools.cache +def _any_unmeasured(kind: MediaKind) -> Select[Any]: + media, conditions = _unmeasured(kind) + return select(literal(1)).select_from(media).where(*conditions).limit(fixed_limit(1)) + + +@functools.cache +def _next_unmeasured(kind: MediaKind) -> Select[Any]: + media, conditions = _unmeasured(kind) + _, name_column, subfolder_column = _MEDIA[kind] + return ( + select(media.c[name_column], media.c[subfolder_column]) + .where(*conditions, media.c.created_at <= bindparam("created_before")) + .order_by(media.c.created_at, media.c[name_column]) + .limit(bindparam("limit")) + ) + + +@functools.cache +def _hold(dialect_name: str) -> Any: + # A hold taken again starts over: held while its session runs. + return upsert(dialect_name, intermediates_session_holds, update=["released_at"]) + + +@functools.cache +def _mark_unmeasurable(dialect_name: str) -> Any: + return insert_ignore(dialect_name, intermediates_unmeasurable) + + +_COVER = ( + select(board_images.c.image_name) + .join(images, images.c.image_name == board_images.c.image_name) + .where(board_images.c.board_id == _P.board_id, images.c.is_intermediate == false()) + .order_by(board_images.c.created_at.desc(), board_images.c.image_name.desc()) + .limit(fixed_limit(1)) + .scalar_subquery() +) +_PROJECTS = select(_P.user_id, _P.project_id, _P.name, _COVER) +_OWNERS_PROJECTS = _PROJECTS.where(_P.user_id == bindparam("user_id")) + + +class AggregateRow(NamedTuple): + user_id: str + project_key: Optional[str] + classification: Classification + count: int + measured_bytes: int + unmeasured: int + + +class PreviewRow(NamedTuple): + user_id: str + project_key: Optional[str] + classification: Classification + deletable: bool + count: int + measured_bytes: int + unmeasured: int + + +class WindowRow(NamedTuple): + name: str + file_size_bytes: Optional[int] + deletable: bool + created_at: str + + +def _clock_parameters(now: str, recent_cutoff: str, active_names: Collection[str]) -> dict[str, Any]: + return {"now": now, "recent_cutoff": recent_cutoff, "active_names": bound_set(active_names)} + + +class IntermediateQueries(QueryModule): + # region active work + + @read + def active_items(self, conn: Connection) -> dict[int, int]: + """The protected queue items, each with the revision of its stored session.""" + return {int(row[0]): int(row[1]) for row in conn.execute(_ACTIVE_ITEMS).all()} + + @read + def sessions(self, conn: Connection, item_ids: Collection[int]) -> list[tuple[int, int, Optional[str]]]: + """(item, revision, session JSON) of each named queue item.""" + rows: list[tuple[int, int, Optional[str]]] = [] + ordered = sorted(item_ids) + for start in range(0, len(ordered), IN_CHUNK): + chunk = ordered[start : start + IN_CHUNK] + rows.extend((int(r[0]), int(r[1]), r[2]) for r in conn.execute(_SESSIONS, {"item_ids": chunk}).all()) + return rows + + @read + def is_session_active(self, conn: Connection, session_id: str) -> bool: + return conn.execute(_SESSION_IS_ACTIVE, {"session_id": session_id}).first() is not None + + @read + def count_existing(self, conn: Connection, kind: MediaKind, names: Collection[str]) -> int: + """How many of the named media exist.""" + return int(conn.execute(_count_existing(kind), {"names": bound_set(names)}).scalar_one()) + + # endregion + + # region holds + + @write + def expire_session_holds(self, conn: Connection, *, now: str, recent_cutoff: str) -> None: + """Releases the holds of sessions no longer active, and drops those whose grace is over.""" + conn.execute(_RELEASE_SESSION_HOLDS, {"now": now}) + conn.execute(_DROP_RELEASED_SESSION_HOLDS, {"recent_cutoff": recent_cutoff}) + + @write + def hold_session_media(self, conn: Connection, session_id: str, held: Sequence[tuple[MediaKind, str]]) -> None: + if held: + conn.execute( + _hold(conn.dialect.name), + [ + {"session_id": session_id, "media_kind": kind, "media_name": name, "released_at": None} + for kind, name in held + ], + ) + + @write + def clear_process_state(self, conn: Connection) -> None: + """Empties the session holds and the unmeasurable marks, which belong to the process that wrote them.""" + conn.execute(_CLEAR_SESSION_HOLDS) + conn.execute(_CLEAR_UNMEASURABLE) + + @write + def sweep_browser_holds(self, conn: Connection, now: str) -> None: + conn.execute(_SWEEP_BROWSER_HOLDS, {"now": now}) + + @write + def release_lease(self, conn: Connection, user_id: str, lease_id: str) -> None: + conn.execute(_RELEASE_LEASE, {"user_id": user_id, "lease_id": lease_id}) + + @write + def hold_for_lease( + self, conn: Connection, kind: MediaKind, *, user_id: str, lease_id: str, names: Collection[str], expires_at: str + ) -> None: + """Leases those of the named media that are the account's intermediates.""" + conn.execute( + _hold_for_lease(kind, conn.dialect.name), + {"user_id": user_id, "lease_id": lease_id, "names": bound_set(names), "expires_at": expires_at}, + ) + + @read + def leases(self, conn: Connection, user_id: str) -> list[tuple[str, str]]: + """The account's leases, each with when it was last refreshed (its latest expiry).""" + return [(str(r[0]), str(r[1])) for r in conn.execute(_LEASES, {"user_id": user_id}).all()] + + @write + def release_leases(self, conn: Connection, user_id: str, lease_ids: Sequence[str]) -> None: + for start in range(0, len(lease_ids), IN_CHUNK): + conn.execute(_RELEASE_LEASES, {"user_id": user_id, "lease_ids": list(lease_ids[start : start + IN_CHUNK])}) + + # endregion + + # region classification + + @read + def aggregate( + self, + conn: Connection, + kind: MediaKind, + scope: ScopeSpec, + *, + now: str, + recent_cutoff: str, + active_names: Collection[str], + ) -> list[AggregateRow]: + """(owner, project, classification) groups of a scope's intermediates, with counts and sizes.""" + parameters = {**_clock_parameters(now, recent_cutoff, active_names), **scope.parameters()} + return [ + AggregateRow(str(r[0]), r[1], r[2], int(r[3]), int(r[4] or 0), int(r[5] or 0)) + for r in conn.execute(_aggregate(kind, scope.shape), parameters).all() + ] + + @read + def preview( + self, + conn: Connection, + kind: MediaKind, + scope: ScopeSpec, + *, + mode: CleanupMode, + is_admin: bool, + caller_user_id: str, + now: str, + recent_cutoff: str, + active_names: Collection[str], + ) -> list[PreviewRow]: + parameters = { + **_clock_parameters(now, recent_cutoff, active_names), + **scope.parameters(), + "caller_user_id": caller_user_id, + } + return [ + PreviewRow(str(r[0]), r[1], r[2], bool(r[3]), int(r[4]), int(r[5] or 0), int(r[6] or 0)) + for r in conn.execute(_preview(kind, scope.shape, mode, is_admin), parameters).all() + ] + + @read + def acknowledged( + self, + conn: Connection, + kind: MediaKind, + scope: ScopeSpec, + *, + is_admin: bool, + caller_user_id: str, + now: str, + recent_cutoff: str, + active_names: Collection[str], + ) -> list[tuple[str, str, str, int]]: + """(owner kind, owner account, owner id, items) of the documents a force cleanup of the scope breaks.""" + parameters = { + **_clock_parameters(now, recent_cutoff, active_names), + **scope.parameters(), + "caller_user_id": caller_user_id, + } + return [ + (str(r[0]), str(r[1]), str(r[2]), int(r[3])) + for r in conn.execute(_acknowledged(kind, scope.shape, is_admin), parameters).all() + ] + + @read + def window( + self, + conn: Connection, + kind: MediaKind, + scope: ScopeSpec, + *, + mode: CleanupMode, + is_admin: bool, + caller_user_id: str, + after: tuple[str, str], + limit: int, + now: str, + recent_cutoff: str, + active_names: Collection[str], + ) -> list[WindowRow]: + parameters = { + **_clock_parameters(now, recent_cutoff, active_names), + **scope.parameters(), + "caller_user_id": caller_user_id, + "after_created_at": after[0], + "after_name": after[1], + "limit": limit, + } + return [ + WindowRow(str(r[0]), r[1], bool(r[2]), str(r[3])) + for r in conn.execute(_window(kind, scope.shape, mode, is_admin), parameters).all() + ] + + @read + def still_deletable( + self, + conn: Connection, + kind: MediaKind, + names: Collection[str], + *, + mode: CleanupMode, + is_admin: bool, + caller_user_id: str, + now: str, + recent_cutoff: str, + active_names: Collection[str], + ) -> list[tuple[str, str]]: + """(name, owner) of those of the named intermediates a cleanup may delete now.""" + parameters = { + **_clock_parameters(now, recent_cutoff, active_names), + "names": bound_set(names), + "caller_user_id": caller_user_id, + } + return [(str(r[0]), str(r[1])) for r in conn.execute(_still_deletable(kind, mode, is_admin), parameters).all()] + + @read + def reference_owners( + self, conn: Connection, kind: MediaKind, names: Collection[str] + ) -> list[tuple[str, str, str, str]]: + """(media, owner kind, owner account, owner id) of each document naming one of the named media.""" + return [ + (str(r[0]), str(r[1]), str(r[2]), str(r[3])) + for r in conn.execute(_reference_owners(kind), {"names": bound_set(names)}).all() + ] + + # endregion + + # region documents and projects + + @read + def projects(self, conn: Connection, user_id: Optional[str]) -> list[tuple[str, str, str, Optional[str]]]: + """(owner, project, name, cover) of the account's projects (None: everyone's); the cover is the newest image + on the project's board that is not an intermediate.""" + statement = _OWNERS_PROJECTS if user_id is not None else _PROJECTS + return [(str(r[0]), str(r[1]), str(r[2]), r[3]) for r in conn.execute(statement, {"user_id": user_id}).all()] + + @read + def document_names( + self, conn: Connection, owner_kind: str, ids: Sequence[str] + ) -> list[tuple[Optional[str], str, Optional[str]]]: + """(owner, id, name) of the documents of `owner_kind` with the given ids; workflows have no owner here.""" + rows: list[tuple[Optional[str], str, Optional[str]]] = [] + for start in range(0, len(ids), IN_CHUNK): + chunk = list(ids[start : start + IN_CHUNK]) + rows.extend( + (r[0], str(r[1]), r[2]) for r in conn.execute(_DOCUMENT_NAMES[owner_kind], {"ids": chunk}).all() + ) + return rows + + # endregion + + # region measurement + + @read + def any_unmeasured(self, conn: Connection, kind: MediaKind) -> bool: + return conn.execute(_any_unmeasured(kind)).first() is not None + + @read + def next_unmeasured( + self, conn: Connection, kind: MediaKind, *, limit: int, created_before: str + ) -> list[tuple[str, str]]: + """(name, subfolder) of unmeasured intermediates created by `created_before`, oldest first.""" + statement = _next_unmeasured(kind) + return [ + (str(r[0]), str(r[1])) + for r in conn.execute(statement, {"limit": limit, "created_before": created_before}).all() + ] + + @write + def mark_unmeasurable(self, conn: Connection, kind: MediaKind, names: Sequence[str]) -> None: + if names: + conn.execute( + _mark_unmeasurable(conn.dialect.name), [{"media_kind": kind, "media_name": name} for name in names] + ) + + # endregion diff --git a/invokeai/app/services/shared/database/queries/locks.py b/invokeai/app/services/shared/database/queries/locks.py new file mode 100644 index 00000000000..4d2f4c40490 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/locks.py @@ -0,0 +1,76 @@ +"""Named locks, for invariants that span rows. + +On MySQL and MariaDB transactions run side by side, so a check followed by a write ("is another active +administrator left?") can interleave with another transaction's. Every transaction that changes what such a +check reads takes the check's lock first, and the second waits for the first to end. Its reads then see what +the first committed, because writers run at READ COMMITTED (see `engines.py`): each statement reads the newest +committed rows. On SQLite all transactions are serialised already; the lock rows are still read there, so a +missing one is noticed on every backend. +""" + +import functools +from collections.abc import Iterable +from enum import StrEnum +from typing import Any + +from sqlalchemy import Connection, Select, literal, select + +from invokeai.app.services.shared.database.queries.base import QueryModule, SharedTransaction, write +from invokeai.app.services.shared.database.schema.locks import db_locks + + +class DatabaseLock(StrEnum): + """A lock, by the work it serialises. Each has a row in `db_locks`, which a migration adds.""" + + # Changes to who is an active administrator: counting them and changing one is one step. + ADMIN_ACCOUNTS = "admin_accounts" + # A replace of the image index's custom vocabulary: deleting every term and inserting the new ones is one step. + IMAGE_INDEX_VOCABULARY = "image_index_vocabulary" + # What protects media from the intermediates cleanup: a cleanup takes it exclusively for its check and delete, + # a write that makes media protected (a reference, a hold) shares it. + MEDIA_PROTECTION = "media_protection" + # What the session queue admits: counting its pending items and adding some is one step, as is finding an + # enqueue request's receipt and writing it. + SESSION_QUEUE_ADMISSION = "session_queue_admission" + + +@functools.cache +def _lock(names: tuple[str, ...], shared: bool) -> Select[Any]: + """Locks the named rows, in name order, so that two transactions that take several of the same locks take them in + the same order. One statement per set of names, of which there are few: an IN list bound value by value costs half + what an expanding one does, on every transaction that takes a lock.""" + statement = select(db_locks.c.name).where(db_locks.c.name.in_([literal(name) for name in names])) + return statement.order_by(db_locks.c.name).with_for_update(read=shared) + + +class LockQueries(QueryModule): + @write + def acquire( + self, conn: Connection, *locks: DatabaseLock, shared: bool = False, also_shared: Iterable[DatabaseLock] = () + ) -> None: + """Holds `locks` until the transaction ends. + + It must be the first call of a `transaction()`: a lock taken outside one would be released at once, and + one taken after other rows were read or locked would guard reads made without it and could deadlock + against a transaction that holds it. Raises `RuntimeError` otherwise, and when a lock has no row: + locking nothing would guard nothing, silently. + + `shared` locks share with each other and wait only for, and hold off, an exclusive lock of the same name. + `also_shared` adds shared locks to exclusive ones; every lock is still taken in name order. + """ + if not isinstance(self._scope, SharedTransaction) or self._scope.calls != 1: + raise RuntimeError("Database locks are taken by the first call of a transaction, all at once") + modes = {lock.value: shared for lock in locks} + modes.update({lock.value: True for lock in also_shared if lock.value not in modes}) + names = tuple(sorted(modes)) + if not names: + raise ValueError("No lock to acquire") + # SQLite has no row locks (its transaction excludes every other): one read finds the rows there. + if len(set(modes.values())) == 1 or conn.dialect.name == "sqlite": + locked: set[str] = set(conn.execute(_lock(names, modes[names[0]])).scalars().all()) + else: + locked = set() + for name in names: + locked.update(conn.execute(_lock((name,), modes[name])).scalars().all()) + if missing := [name for name in names if name not in locked]: + raise RuntimeError(f"No row in db_locks for {', '.join(missing)}: the migration that adds it has not run") diff --git a/invokeai/app/services/shared/database/queries/media_references.py b/invokeai/app/services/shared/database/queries/media_references.py new file mode 100644 index 00000000000..d7cecf90307 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/media_references.py @@ -0,0 +1,69 @@ +"""The index of which saved documents reference which media (see `invokeai.app.services.shared.media_references`). + +A document's references are written in the transaction that saves or deletes the document, after the document's +own row. On a server, that row's lock keeps two saves of one document from interleaving their references. A +delete removes the references only of documents whose rows it deleted: a document written for the first time +while the delete runs has no row the delete could have locked, and must keep the references it writes. +""" + +import itertools +from collections.abc import Collection + +from sqlalchemy import Connection, bindparam, delete, insert + +from invokeai.app.services.shared.database.queries.base import IN_CHUNK, QueryModule, write +from invokeai.app.services.shared.database.schema.media_references import media_references +from invokeai.app.services.shared.media_references import MediaReferenceOwnerKind, MediaReferences + +_DELETE_OWNER = delete(media_references).where( + media_references.c.owner_kind == bindparam("owner_kind"), + media_references.c.user_id == bindparam("user_id"), + media_references.c.owner_id == bindparam("owner_id"), +) +_DELETE_OWNERS = delete(media_references).where( + media_references.c.owner_kind == bindparam("owner_kind"), + media_references.c.user_id == bindparam("user_id"), + media_references.c.owner_id.in_(bindparam("owner_ids", expanding=True)), +) +_DELETE_OWNED_BY = delete(media_references).where( + media_references.c.user_id == bindparam("user_id"), + media_references.c.owner_kind.in_(bindparam("owner_kinds", expanding=True)), +) +_INSERT = insert(media_references) + + +class MediaReferenceQueries(QueryModule): + @write + def replace( + self, + conn: Connection, + *, + owner_kind: MediaReferenceOwnerKind, + user_id: str, + owner_id: str, + references: MediaReferences, + ) -> None: + """Makes the index for one document equal to `references`.""" + owner = {"owner_kind": owner_kind, "user_id": user_id, "owner_id": owner_id} + conn.execute(_DELETE_OWNER, owner) + rows = [{**owner, "media_kind": "image", "media_name": name} for name in sorted(references.images)] + rows.extend({**owner, "media_kind": "video", "media_name": name} for name in sorted(references.videos)) + if rows: + conn.execute(_INSERT, rows) + + @write + def delete(self, conn: Connection, *, owner_kind: MediaReferenceOwnerKind, user_id: str, owner_id: str) -> None: + conn.execute(_DELETE_OWNER, {"owner_kind": owner_kind, "user_id": user_id, "owner_id": owner_id}) + + @write + def delete_many( + self, conn: Connection, *, owner_kind: MediaReferenceOwnerKind, user_id: str, owner_ids: Collection[str] + ) -> None: + for chunk in itertools.batched(owner_ids, IN_CHUNK): + conn.execute(_DELETE_OWNERS, {"owner_kind": owner_kind, "user_id": user_id, "owner_ids": list(chunk)}) + + @write + def delete_owned_by(self, conn: Connection, user_id: str, *owner_kinds: MediaReferenceOwnerKind) -> None: + """Deletes the references of every document of these kinds that the account owns, for documents deleted + with the account: deleting its row locks theirs, and keeps new ones from being written meanwhile.""" + conn.execute(_DELETE_OWNED_BY, {"user_id": user_id, "owner_kinds": list(owner_kinds)}) diff --git a/invokeai/app/services/shared/database/queries/model_relationships.py b/invokeai/app/services/shared/database/queries/model_relationships.py new file mode 100644 index 00000000000..f746cbf27bc --- /dev/null +++ b/invokeai/app/services/shared/database/queries/model_relationships.py @@ -0,0 +1,47 @@ +"""Which installed models belong together: pairs of related models, each stored once.""" + +import functools + +from sqlalchemy import Connection, Insert, bindparam, delete, select, union + +from invokeai.app.services.shared.database.dialect import insert_ignore +from invokeai.app.services.shared.database.queries.base import QueryModule, read, write +from invokeai.app.services.shared.database.schema.models import model_relationships + +_R = model_relationships.c + +_DELETE = delete(model_relationships).where( + _R.model_key_1 == bindparam("model_key_1"), _R.model_key_2 == bindparam("model_key_2") +) +# A pair is stored once, so a model's relatives are on either side of it. UNION drops duplicates. +_RELATED = union( + select(_R.model_key_2.label("related_key")).where(_R.model_key_1 == bindparam("key")), + select(_R.model_key_1.label("related_key")).where(_R.model_key_2 == bindparam("key")), +).order_by("related_key") + + +@functools.cache +def _add(dialect_name: str) -> Insert: + return insert_ignore(dialect_name, model_relationships) + + +def _pair(model_key_1: str, model_key_2: str) -> dict[str, str]: + # The smaller key first, so that a pair has one row whichever way round it is named. + first, second = sorted((model_key_1, model_key_2)) + return {"model_key_1": first, "model_key_2": second} + + +class ModelRelationshipQueries(QueryModule): + @write + def add(self, conn: Connection, model_key_1: str, model_key_2: str) -> None: + """Relates the two models, unless they are already.""" + conn.execute(_add(conn.dialect.name), _pair(model_key_1, model_key_2)) + + @write + def remove(self, conn: Connection, model_key_1: str, model_key_2: str) -> None: + conn.execute(_DELETE, _pair(model_key_1, model_key_2)) + + @read + def related(self, conn: Connection, model_key: str) -> list[str]: + """The keys of the models related to this one, in order.""" + return list(conn.execute(_RELATED, {"key": model_key}).scalars().all()) diff --git a/invokeai/app/services/shared/database/queries/models.py b/invokeai/app/services/shared/database/queries/models.py new file mode 100644 index 00000000000..1a14f18c2b1 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/models.py @@ -0,0 +1,120 @@ +"""Installed models: each one's config, a JSON document whose members the other columns are generated from.""" + +import functools +from typing import Any, Optional + +from sqlalchemy import ColumnElement, Connection, Select, bindparam, delete, insert, literal, select, update + +from invokeai.app.services.shared.database.dialect import CaseInsensitiveOrder +from invokeai.app.services.shared.database.queries.base import QueryModule, locking, read, write +from invokeai.app.services.shared.database.schema.models import models + +_M = models.c + +# The longest key and path a model can have: a server stores no longer one in their columns. +MAX_KEY_LENGTH: int = _M.id.type.length +MAX_PATH_LENGTH: int = _M.path.type.length + +_GET = select(_M.config).where(_M.id == bindparam("key")) +_LOCK = _GET.with_for_update() +_EXISTS = select(literal(1)).where(_M.id == bindparam("key")) +_AT_PATH = select(_M.config).where(_M.path == bindparam("path")) +# A hash is not unique. The oldest model of a file comes first, as SQLite's table order had it, and the key breaks +# ties: the route that recalls a model by its hash takes the first. +_WITH_HASH = select(_M.config).where(_M.hash == bindparam("hash")).order_by(_M.created_at, _M.id) +_PATHS = select(_M.path) +_INSERT = insert(models) +_SAVE = update(models).where(_M.id == bindparam("key")).values(config=bindparam("config")) +_DELETE = delete(models).where(_M.id == bindparam("key")) + +# The search filters, in the order `search` takes them. +_FILTERS = (_M.name, _M.base, _M.type, _M.format) +# Keyed by the values of `ModelRecordOrderBy`, which is not imported: its module loads the model manager. Names and +# bases sort ignoring case. +_ORDERINGS: dict[str, tuple[ColumnElement[Any], ...]] = { + "default": (_M.type, CaseInsensitiveOrder(_M.base), CaseInsensitiveOrder(_M.name), _M.format), + "type": (_M.type,), + "base": (CaseInsensitiveOrder(_M.base),), + "name": (CaseInsensitiveOrder(_M.name),), + "format": (_M.format,), + "size": (_M.file_size,), + "created_at": (_M.created_at,), + "updated_at": (_M.updated_at,), + "path": (_M.path,), +} + + +@functools.cache +def _search(filtered: tuple[bool, ...], order_by: str, descending: bool) -> Select[Any]: + """The configs of the models matching the filters that are set, in order; few combinations of filters and + orders exist.""" + statement = select(_M.config).where( + *(column == bindparam(f"filter_{column.name}") for column, on in zip(_FILTERS, filtered, strict=True) if on) + ) + # The key breaks ties, so that equal values keep one order. + ordering = (*_ORDERINGS[order_by], _M.id) + return statement.order_by(*(column.desc() if descending else column.asc() for column in ordering)) + + +class ModelQueries(QueryModule): + @read + def get(self, conn: Connection, key: str) -> Optional[str]: + """The model's config.""" + return conn.execute(_GET, {"key": key}).scalar() + + @locking + def lock(self, conn: Connection, key: str) -> Optional[str]: + """The model's config, its row locked until the transaction ends.""" + return conn.execute(_LOCK, {"key": key}).scalar() + + @read + def exists(self, conn: Connection, key: str) -> bool: + return conn.execute(_EXISTS, {"key": key}).first() is not None + + @read + def search( + self, + conn: Connection, + *, + name: Optional[str], + base: Optional[str], + model_type: Optional[str], + model_format: Optional[str], + order_by: str, + descending: bool, + ) -> list[str]: + """The configs of the models matching every filter given (an empty one matches all), ordered by + `order_by`, a value of `ModelRecordOrderBy`.""" + values = (name, base, model_type, model_format) + statement = _search(tuple(bool(value) for value in values), order_by, descending) + parameters = {f"filter_{column.name}": value for column, value in zip(_FILTERS, values, strict=True) if value} + return list(conn.execute(statement, parameters).scalars().all()) + + @read + def at_path(self, conn: Connection, path: str) -> list[str]: + """The configs of the models at the path: none or one, since a path is unique.""" + return list(conn.execute(_AT_PATH, {"path": path}).scalars().all()) + + @read + def with_hash(self, conn: Connection, hash: str) -> list[str]: + """The configs of the models whose file has the hash.""" + return list(conn.execute(_WITH_HASH, {"hash": hash}).scalars().all()) + + @read + def paths(self, conn: Connection) -> list[str]: + """The path of every model, also of one whose config no longer validates.""" + return list(conn.execute(_PATHS).scalars().all()) + + @write + def insert(self, conn: Connection, key: str, config: str) -> None: + conn.execute(_INSERT, {"id": key, "config": config}) + + @write + def save(self, conn: Connection, key: str, config: str) -> bool: + """Replaces the model's config; whether the model exists.""" + return conn.execute(_SAVE, {"key": key, "config": config}).rowcount > 0 + + @write + def delete(self, conn: Connection, key: str) -> bool: + """Deletes the model, and its relationships with it; whether there was one.""" + return conn.execute(_DELETE, {"key": key}).rowcount > 0 diff --git a/invokeai/app/services/shared/database/queries/projects.py b/invokeai/app/services/shared/database/queries/projects.py new file mode 100644 index 00000000000..9fb3151fabb --- /dev/null +++ b/invokeai/app/services/shared/database/queries/projects.py @@ -0,0 +1,261 @@ +"""Canvas projects: their documents, revisions, and the board each one claims.""" + +import json +from collections.abc import Sequence +from typing import Any, NamedTuple, Optional + +from sqlalchemy import ( + ColumnElement, + Connection, + Row, + bindparam, + delete, + false, + insert, + literal, + select, + union_all, + update, +) + +from invokeai.app.services.image_records.image_records_common import ASSETS_CATEGORIES, IMAGE_CATEGORIES +from invokeai.app.services.project_records.project_records_common import ( + ProjectBoardItemDTO, + ProjectRecordDTO, + ProjectSummaryDTO, +) +from invokeai.app.services.shared.database.queries.base import QueryModule, locking, mapped, read, write +from invokeai.app.services.shared.database.schema.boards import board_images, board_videos +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.projects import projects +from invokeai.app.services.shared.database.schema.videos import videos + + +class LockedProject(NamedTuple): + """A project's row as a transaction locked it.""" + + revision: int + minimum_canvas_schema_version: int + board_id: str + + +_SUMMARY_COLUMNS = ( + projects.c.project_id, + projects.c.board_id, + projects.c.name, + projects.c.revision, + projects.c.minimum_canvas_schema_version, + projects.c.created_at, + projects.c.updated_at, +) +_THE_PROJECT = (projects.c.user_id == bindparam("user_id"), projects.c.project_id == bindparam("project_id")) + + +_GET = select(*_SUMMARY_COLUMNS, projects.c.data).where(*_THE_PROJECT) +_SUMMARY = select(*_SUMMARY_COLUMNS).where(*_THE_PROJECT) +# Oldest first. Projects created in the same millisecond come in the order of their ids. +_SUMMARIES = ( + select(*_SUMMARY_COLUMNS) + .where(projects.c.user_id == bindparam("user_id")) + .order_by(projects.c.created_at, projects.c.project_id) +) +_BOARD_ID = select(projects.c.board_id).where(*_THE_PROJECT) +_LOCK = ( + select(projects.c.revision, projects.c.minimum_canvas_schema_version, projects.c.board_id) + .where(*_THE_PROJECT) + .with_for_update() +) +_CLAIMANT = select(projects.c.user_id, projects.c.project_id).where(projects.c.board_id == bindparam("board_id")) +_INSERT = insert(projects) +_SAVE = ( + update(projects) + .where(projects.c.user_id == bindparam("target_user_id"), projects.c.project_id == bindparam("target_project_id")) + .values( + name=bindparam("new_name"), + data=bindparam("new_data"), + minimum_canvas_schema_version=bindparam("new_minimum_canvas_schema_version"), + revision=projects.c.revision + 1, + ) +) +_DELETE = delete(projects).where(*_THE_PROJECT) + + +def _shown_by_the_gallery(category: ColumnElement[str]) -> ColumnElement[bool]: + # `other`, the canvas's private category, is in neither list: the gallery never shows it on a board. + return category.in_([literal(each.value) for each in (*IMAGE_CATEGORIES, *ASSETS_CATEGORIES)]) + + +# What the gallery shows on a board, of both kinds. `deleted_at` is not filtered: soft delete is unused, and no +# other board query consults it, so filtering here would make the snapshot disagree with the board's counts. +_ITEMS = union_all( + select( + literal("image").label("kind"), + images.c.image_name.label("name"), + images.c.image_category.label("category"), + images.c.starred.label("starred"), + ) + .select_from(board_images.join(images, board_images.c.image_name == images.c.image_name)) + .where( + board_images.c.board_id == bindparam("board_id"), + images.c.is_intermediate == false(), + _shown_by_the_gallery(images.c.image_category), + ), + select(literal("video"), videos.c.video_name, videos.c.video_category, videos.c.starred) + .select_from(board_videos.join(videos, board_videos.c.video_name == videos.c.video_name)) + .where( + board_videos.c.board_id == bindparam("board_id"), + videos.c.is_intermediate == false(), + _shown_by_the_gallery(videos.c.video_category), + ), +) +_BOARD_ITEMS = _ITEMS.order_by(_ITEMS.selected_columns.kind, _ITEMS.selected_columns.name).limit(bindparam("limit")) + + +def _summary(row: Sequence[Any]) -> ProjectSummaryDTO: + project_id, board_id, name, revision, minimum_canvas_schema_version, created_at, updated_at = row + return ProjectSummaryDTO( + project_id=project_id, + board_id=board_id, + name=name, + revision=revision, + minimum_canvas_schema_version=minimum_canvas_schema_version, + created_at=created_at, + updated_at=updated_at, + ) + + +def _summary_or_none(row: Optional[Sequence[Any]]) -> Optional[ProjectSummaryDTO]: + return _summary(row) if row is not None else None + + +def _summaries(rows: Sequence[Sequence[Any]]) -> list[ProjectSummaryDTO]: + return [_summary(row) for row in rows] + + +def _record_or_none(row: Optional[Sequence[Any]]) -> Optional[ProjectRecordDTO]: + if row is None: + return None + project_id, board_id, name, revision, minimum_canvas_schema_version, created_at, updated_at, data = row + return ProjectRecordDTO( + project_id=project_id, + board_id=board_id, + name=name, + revision=revision, + minimum_canvas_schema_version=minimum_canvas_schema_version, + created_at=created_at, + updated_at=updated_at, + data=json.loads(data), + ) + + +def _board_items_or_none(rows: Optional[Sequence[Sequence[Any]]]) -> Optional[list[ProjectBoardItemDTO]]: + if rows is None: + return None + return [ + ProjectBoardItemDTO(kind=kind, name=name, category=category, starred=starred) + for kind, name, category, starred in rows + ] + + +class ProjectQueries(QueryModule): + @mapped(_record_or_none) + @read + def get(self, conn: Connection, user_id: str, project_id: str) -> Optional[Row[Any]]: + """The project with its document.""" + return conn.execute(_GET, {"user_id": user_id, "project_id": project_id}).first() + + @mapped(_summary_or_none) + @read + def summary(self, conn: Connection, user_id: str, project_id: str) -> Optional[Row[Any]]: + """The project without its document.""" + return conn.execute(_SUMMARY, {"user_id": user_id, "project_id": project_id}).first() + + @mapped(_summaries) + @read + def summaries(self, conn: Connection, user_id: str) -> Sequence[Row[Any]]: + """The account's projects without their documents, oldest first.""" + return conn.execute(_SUMMARIES, {"user_id": user_id}).all() + + @read + def board_id(self, conn: Connection, user_id: str, project_id: str) -> Optional[str]: + return conn.execute(_BOARD_ID, {"user_id": user_id, "project_id": project_id}).scalar_one_or_none() + + @locking + def lock(self, conn: Connection, user_id: str, project_id: str) -> Optional[LockedProject]: + """Locks the project's row until the transaction ends: what a write of the project decides on, or None when + there is no such project. Lock it before its board: rows are locked in that order.""" + row: Optional[Sequence[Any]] = conn.execute(_LOCK, {"user_id": user_id, "project_id": project_id}).first() + if row is None: + return None + revision, minimum_canvas_schema_version, board_id = row + return LockedProject(revision, minimum_canvas_schema_version, board_id) + + @read + def claimant(self, conn: Connection, board_id: str) -> Optional[tuple[str, str]]: + """The account and id of the project that claims the board, if one does.""" + row: Optional[Sequence[Any]] = conn.execute(_CLAIMANT, {"board_id": board_id}).first() + if row is None: + return None + user_id, project_id = row + return user_id, project_id + + @mapped(_board_items_or_none) + @read + def board_items(self, conn: Connection, user_id: str, project_id: str, limit: int) -> Optional[Sequence[Row[Any]]]: + """Up to `limit` of what the gallery shows on the project's board, by kind and then name; None when there + is no such project.""" + board_id = conn.execute(_BOARD_ID, {"user_id": user_id, "project_id": project_id}).scalar_one_or_none() + if board_id is None: + return None + return conn.execute(_BOARD_ITEMS, {"board_id": board_id, "limit": limit}).all() + + @write + def insert( + self, + conn: Connection, + *, + project_id: str, + user_id: str, + name: str, + data: str, + board_id: str, + minimum_canvas_schema_version: int, + ) -> None: + conn.execute( + _INSERT, + { + "project_id": project_id, + "user_id": user_id, + "name": name, + "data": data, + "board_id": board_id, + "minimum_canvas_schema_version": minimum_canvas_schema_version, + }, + ) + + @write + def save( + self, + conn: Connection, + *, + user_id: str, + project_id: str, + name: str, + data: str, + minimum_canvas_schema_version: int, + ) -> None: + """Writes the project as its next revision. The caller has locked it and checked the revision it saves over.""" + conn.execute( + _SAVE, + { + "target_user_id": user_id, + "target_project_id": project_id, + "new_name": name, + "new_data": data, + "new_minimum_canvas_schema_version": minimum_canvas_schema_version, + }, + ) + + @write + def delete(self, conn: Connection, user_id: str, project_id: str) -> None: + conn.execute(_DELETE, {"user_id": user_id, "project_id": project_id}) diff --git a/invokeai/app/services/shared/database/queries/session_queue.py b/invokeai/app/services/shared/database/queries/session_queue.py new file mode 100644 index 00000000000..f543bf63eb5 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/session_queue.py @@ -0,0 +1,1075 @@ +"""The session queue: the items users enqueued, their workflow-call chains, and the counts the queue's views show. + +An item is read as a dict of its row (every column of `session_queue`) with its owner's display name and email, the +form `SessionQueueItem.queue_item_from_dict()` and `SessionQueueItemSummary.queue_item_summary_from_dict()` take. +An origin prefix matches as SQLite's `LIKE` always has, ignoring case; its `%` and `_` are taken as written. +""" + +import functools +import itertools +from collections.abc import Iterable, Sequence +from datetime import datetime, timedelta, timezone +from typing import Any, NamedTuple, Optional + +from sqlalchemy import ( + ColumnElement, + Connection, + Row, + Select, + String, + Update, + and_, + bindparam, + case, + delete, + func, + insert, + literal, + or_, + select, + update, +) + +from invokeai.app.services.shared.database.dialect import ( + CaseInsensitiveLike, + ContainsText, + InBoundSet, + Unindexed, + bound_set, + fixed_limit, + like_prefix, +) +from invokeai.app.services.shared.database.queries.base import IN_CHUNK, QueryModule, locking, mapped, read, write +from invokeai.app.services.shared.database.schema.session_queue import session_queue, session_queue_enqueue_receipts +from invokeai.app.services.shared.database.schema.users import users +from invokeai.app.services.shared.database.types import now_text, timestamp_text + +_Q = session_queue.c +_R = session_queue_enqueue_receipts.c + +_ITEM = (*session_queue.c, users.c.display_name.label("user_display_name"), users.c.email.label("user_email")) +_ITEM_NAMES = tuple(column.name for column in _ITEM) +_ITEMS = session_queue.outerjoin(users, users.c.user_id == _Q.user_id) + +_SUMMARY = ( + _Q.item_id, + _Q.created_at, + _Q.status, + _Q.device, + _Q.started_at, + _Q.completed_at, + _Q.origin, + _Q.destination, + _Q.batch_id, + _Q.user_id, + users.c.display_name.label("user_display_name"), + users.c.email.label("user_email"), + _Q.field_values, + _Q.parent_item_id, +) +_SUMMARY_NAMES = tuple(column.name for column in _SUMMARY) + +_IN_QUEUE = _Q.queue_id == bindparam("queue_id") +# The same term kept off every index on SQLite: for statements whose rows a status, batch or item id finds. The queue's +# listing index (queue, creation time, ...) would otherwise walk the queue's whole history for them, as SQLite, keeping +# no statistics, takes `queue_id = ?` for selective. +_IN_QUEUE_UNINDEXED = Unindexed(_Q.queue_id) == bindparam("queue_id") +_PENDING = _Q.status == literal("pending") +_IN_PROGRESS = _Q.status == literal("in_progress") + +_RECEIPT = select(_R.batch_id, _R.requested, _R.enqueued, _R.item_ids).where( + _R.queue_id == bindparam("queue_id"), + _R.user_id == bindparam("user_id"), + _R.idempotency_key == bindparam("idempotency_key"), +) +_ITEM_BY_ID = select(*_ITEM).select_from(_ITEMS).where(_Q.item_id == bindparam("item_id")) +_WORKFLOW = select(_Q.workflow).where(_Q.item_id == bindparam("item_id")) +_PARENT = select(_Q.parent_item_id).where(_Q.item_id == bindparam("item_id")) +_CHILDREN = select(_Q.item_id).where(_Q.parent_item_id == bindparam("item_id")).order_by(_Q.item_id) +_IN_PROGRESS_IDS = select(_Q.item_id).where(_IN_QUEUE_UNINDEXED, _IN_PROGRESS).order_by(_Q.item_id) +_PARENT_ROW = session_queue.alias("parent") +# The child that would be dequeued next while its parent waits: its chain is still running between two children. +_NEXT_WAITING_CHILD = ( + select(_Q.item_id) + .select_from(session_queue.join(_PARENT_ROW, _PARENT_ROW.c.item_id == _Q.parent_item_id)) + .where(_IN_QUEUE_UNINDEXED, _PENDING, _PARENT_ROW.c.status == literal("waiting")) + .order_by(_Q.priority.desc(), _Q.created_at, _Q.item_id) + .limit(fixed_limit(1)) +) +_COUNT = select(func.count()).where(_IN_QUEUE) +_SUMMARIES = ( + select(*_SUMMARY).select_from(_ITEMS).where(_IN_QUEUE_UNINDEXED, InBoundSet(_Q.item_id, bindparam("item_ids"))) +) + + +def _origin_matches() -> ColumnElement[bool]: + return CaseInsensitiveLike(_Q.origin, bindparam("origin_pattern")) + + +def _origin_pattern(origin_prefix: Optional[str]) -> Optional[str]: + return like_prefix(origin_prefix) if origin_prefix is not None else None + + +@functools.cache +def _item_by_status(status: str, by_origin: bool) -> Select[Any]: + """The pending item to run next, or an item in progress (the earliest enqueued of several).""" + conditions = [_IN_QUEUE_UNINDEXED, _Q.status == literal(status)] + if by_origin: + conditions.append(_origin_matches()) + ordering = [_Q.priority.desc(), _Q.created_at, _Q.item_id] if status == "pending" else [_Q.item_id] + return select(_Q.item_id).where(*conditions).order_by(*ordering).limit(fixed_limit(1)) + + +@functools.cache +def _all_items(by_destination: bool) -> Select[Any]: + conditions = [_IN_QUEUE_UNINDEXED] + if by_destination: + conditions.append(_Q.destination == bindparam("destination")) + return select(*_ITEM).select_from(_ITEMS).where(*conditions).order_by(_Q.priority.desc(), _Q.item_id) + + +@functools.cache +def _page(by_status: bool, by_destination: bool, after: bool) -> Select[Any]: + """A page of the queue, highest priority first, then in enqueue order; with `after`, the items after an item. + + The priority index pages it on SQLite: the listing index would read and sort the whole queue for each page.""" + conditions = [_IN_QUEUE_UNINDEXED] + if by_status: + conditions.append(_Q.status == bindparam("status")) + if by_destination: + conditions.append(_Q.destination == bindparam("destination")) + if after: + after_priority: ColumnElement[int] = bindparam("after_priority") + conditions.append( + or_( + _Q.priority < after_priority, + and_(_Q.priority == after_priority, _Q.item_id > bindparam("after_item_id")), + ) + ) + return ( + select(*_ITEM) + .select_from(_ITEMS) + .where(*conditions) + .order_by(_Q.priority.desc(), _Q.item_id) + .limit(bindparam("limit")) + ) + + +@functools.cache +def _item_ids(descending: bool, by_user: bool, by_origin: bool, limited: bool) -> Select[Any]: + """The ids, in an order the listing index yields on SQLite without a sort (newest first), or sorting only the items + that share a creation time (oldest first); so a limited read stops after its window.""" + conditions = [_IN_QUEUE] + if by_user: + conditions.append(_Q.user_id == bindparam("user_id")) + if by_origin: + conditions.append(_origin_matches()) + # The item id breaks ties, in enqueue order either way, as SQLite returned them before: a batch's items share + # their creation time to the millisecond, and the queue's views show a batch in that order. + ordering = [_Q.created_at.desc() if descending else _Q.created_at, _Q.item_id] + statement = select(_Q.item_id).where(*conditions).order_by(*ordering) + return statement.limit(bindparam("limit")) if limited else statement + + +@functools.cache +def _status_counts(with_own: bool, by_origin: bool) -> Select[Any]: + """Every status's count in the queue and, `with_own`, in the account's part of it, in one pass: on SQLite the + listing index holds every column read, so no row is.""" + conditions = [_IN_QUEUE] + if by_origin: + conditions.append(_origin_matches()) + own = func.sum(case((_Q.user_id == bindparam("user_id"), 1), else_=0)) if with_own else literal(0) + return select(_Q.status, func.count(), own).where(*conditions).group_by(_Q.status) + + +@functools.cache +def _current(by_origin: bool) -> Select[Any]: + """The identifiers of the item in progress, the earliest enqueued of several.""" + conditions = [_IN_QUEUE_UNINDEXED, _IN_PROGRESS] + if by_origin: + conditions.append(_origin_matches()) + return ( + select(_Q.item_id, _Q.session_id, _Q.batch_id, _Q.user_id) + .where(*conditions) + .order_by(_Q.item_id) + .limit(fixed_limit(1)) + ) + + +@functools.cache +def _batch_counts(by_user: bool) -> Select[Any]: + # The batch index finds a batch's few rows; an index on the queue or the account would read their whole history. + conditions = [_IN_QUEUE_UNINDEXED, _Q.batch_id == bindparam("batch_id")] + if by_user: + conditions.append(Unindexed(_Q.user_id) == bindparam("user_id")) + # A batch's items share their origin and destination: its children and retries copy them. + return ( + select(_Q.status, func.count(), func.min(_Q.origin), func.min(_Q.destination)) + .where(*conditions) + .group_by(_Q.status) + ) + + +@functools.cache +def _destination_counts(by_user: bool) -> Select[Any]: + conditions = [_IN_QUEUE, _Q.destination == bindparam("destination")] + if by_user: + conditions.append(_Q.user_id == bindparam("user_id")) + return select(_Q.status, func.count()).where(*conditions).group_by(_Q.status) + + +def _row_of(conn: Connection, item_id: Optional[int]) -> Optional[Row[Any]]: + return conn.execute(_ITEM_BY_ID, {"item_id": item_id}).first() if item_id is not None else None + + +def _item(row: Sequence[Any]) -> dict[str, Any]: + return dict(zip(_ITEM_NAMES, row, strict=True)) + + +def _item_or_none(row: Optional[Sequence[Any]]) -> Optional[dict[str, Any]]: + return _item(row) if row is not None else None + + +def _items(rows: Sequence[Sequence[Any]]) -> list[dict[str, Any]]: + return [_item(row) for row in rows] + + +def _summaries(rows: Sequence[Sequence[Any]]) -> list[dict[str, Any]]: + return [dict(zip(_SUMMARY_NAMES, row, strict=True)) for row in rows] + + +def _counts(rows: Sequence[Sequence[Any]]) -> dict[str, int]: + return {str(row[0]): int(row[1]) for row in rows} + + +# The walks below stop at an item they have seen: the app never links a cycle (a parent is always enqueued before its +# children), but a corrupt one must not hold the transaction, and on SQLite every other thread, forever. + + +def _descendants(conn: Connection, item_id: int) -> list[int]: + """Every descendant of the item, breadth first.""" + descendants: list[int] = [] + seen = {item_id} + frontier = [item_id] + while frontier: + children: Sequence[int] = conn.execute(_CHILDREN, {"item_id": frontier.pop(0)}).scalars().all() + new = [child for child in children if child not in seen] + seen.update(new) + descendants.extend(new) + frontier.extend(new) + return descendants + + +def _chain(conn: Connection, item_id: int, *, reachable: bool = False) -> Optional[list[int]]: + """The item's ancestors from its parent up, the item, and every descendant of the root, breadth first. None if the + item does not exist, or an ancestor does not and not `reachable`; with `reachable`, the highest ancestor that + exists stands for the root.""" + ancestors: list[int] = [] + current = item_id + while True: + row = conn.execute(_PARENT, {"item_id": current}).first() + if row is None: + if current == item_id or not reachable: + return None + ancestors.pop() + break + parent = row[0] + if parent is None or parent == item_id or parent in ancestors: + break + ancestors.append(int(parent)) + current = int(parent) + root = ancestors[-1] if ancestors else item_id + return list(dict.fromkeys([*ancestors, item_id, *_descendants(conn, root)])) + + +class EnqueueReceipt(NamedTuple): + batch_id: str + requested: int + enqueued: int + # The enqueued items' ids, as a JSON array. + item_ids: str + + +class CurrentItem(NamedTuple): + item_id: int + session_id: str + batch_id: str + user_id: str + + +class StatusCounts(NamedTuple): + """A queue's counts by status, the account's own, and the item in progress, read in one snapshot.""" + + counts: dict[str, int] + own_counts: dict[str, int] + current: Optional[CurrentItem] + + +class BatchCounts(NamedTuple): + counts: dict[str, int] + origin: Optional[str] + destination: Optional[str] + + +# --- Writes --------------------------------------------------------------------------------------------------------- +# +# A status change is one conditional UPDATE: it applies only to an item that is not finished, and its row count says +# whether this call made it. These statements also set `updated_at`, `started_at`, `completed_at` and +# `session_revision`: no database trigger does. + +_TERMINAL = ("completed", "failed", "canceled") +# Statuses of an item that has not finished; the intermediates cleanup protects their media. +ACTIVE_QUEUE_STATUSES = ("pending", "in_progress", "waiting") +_WAITING_WORK = ("pending", "waiting") +_NOT_FINISHED = _Q.status.not_in([literal(status) for status in _TERMINAL]) +_NEXT_SEQUENCE = func.coalesce(_Q.status_sequence, 0) + 1 +_NEXT_REVISION = _Q.session_revision + 1 +_THE_ITEM = _Q.item_id == bindparam("target_item_id") + +_STATUS_COLUMNS = ( + _Q.status, + _Q.status_sequence, + _Q.error_type, + _Q.error_message, + _Q.error_traceback, + _Q.created_at, + _Q.updated_at, + _Q.started_at, + _Q.completed_at, + _Q.device, +) +_STATUS_NAMES = tuple(column.name for column in _STATUS_COLUMNS) +_STATUS_ROW = select(*_STATUS_COLUMNS).where(_Q.item_id == bindparam("item_id")) +_ITEM_EXISTS = select(literal(1)).where(_Q.item_id == bindparam("item_id")) +_LOCK_STATUS = select(_Q.status).where(_Q.item_id == bindparam("item_id")).with_for_update() +# Locks the item's own row only: on a server a FOR UPDATE over a join would lock the owner's account row as well. +_LOCK_ITEM = select(*session_queue.c).where(_Q.item_id == bindparam("item_id")).with_for_update() +_LOCK_ITEM_NAMES = tuple(column.name for column in session_queue.c) + +_CLAIM = ( + update(session_queue) + .where(_THE_ITEM, _PENDING) + .values( + status=literal("in_progress"), + status_sequence=_NEXT_SEQUENCE, + device=func.coalesce(bindparam("new_device", type_=String()), _Q.device), + started_at=bindparam("now"), + ) +) + +_RANKED_PENDING = ( + select( + _Q.item_id, + func.row_number().over(partition_by=_Q.user_id, order_by=(_Q.priority.desc(), _Q.item_id)).label("rank"), + ) + .where(_PENDING) + .cte("user_next_item") +) +_SERVED = session_queue.alias("served") +# The account served longest ago first: its latest start, a MAX over the account's started items (on SQLite an index +# seek; on a server a scan of the account's history in its index). Accounts never served sort first. +_LAST_SERVED = select(func.max(_SERVED.c.started_at)).where(_SERVED.c.user_id == _Q.user_id).scalar_subquery() +# The next item's id; its row is read by id. Sorting whole rows, sessions included, made MySQL copy every pending row +# into a temporary table on each dequeue. +_ROUND_ROBIN_NEXT = ( + select(_Q.item_id) + .select_from( + session_queue.join(_RANKED_PENDING, and_(_RANKED_PENDING.c.item_id == _Q.item_id, _RANKED_PENDING.c.rank == 1)) + ) + .order_by(func.coalesce(_LAST_SERVED, literal("1970-01-01")), _Q.item_id) + .limit(fixed_limit(1)) +) +_FIFO_NEXT = select(_Q.item_id).where(_PENDING).order_by(_Q.priority.desc(), _Q.item_id).limit(fixed_limit(1)) + +# How many resident model keys the affinity score takes: real caches hold a handful. +MAX_AFFINITY_MODEL_KEYS = 50 + + +@functools.cache +def _affinity(slots: int) -> Select[Any]: + """The account's pending item of the same priority within the lookahead whose session names the most resident + models. Keys are bound in `slots` parameters, a power of two, the spare ones NULL, which score nothing.""" + score = sum( + (case((ContainsText(_Q.session, bindparam(f"key_{i}", type_=String())), 1), else_=0) for i in range(slots)), + start=literal(0), + ) + return ( + select(*_ITEM, score.label("affinity")) + .select_from(_ITEMS) + .where( + _PENDING, + _Q.user_id.is_not_distinct_from(bindparam("user_id", type_=String())), + _Q.priority == bindparam("priority"), + _Q.item_id.between(bindparam("first_item_id"), bindparam("last_item_id")), + ) + .order_by(score.desc(), _Q.item_id) + .limit(fixed_limit(1)) + ) + + +def _slots(count: int) -> int: + return 1 << max(count - 1, 0).bit_length() + + +@functools.cache +def _transition(stamp: Optional[str]) -> Update: + """Moves an unfinished item to a status, stamping `started_at` or `completed_at` as the status calls for.""" + values: dict[str, Any] = { + "status": bindparam("new_status"), + "status_sequence": _NEXT_SEQUENCE, + "error_type": bindparam("new_error_type", type_=String()), + "error_message": bindparam("new_error_message", type_=String()), + "error_traceback": bindparam("new_error_traceback", type_=String()), + "device": func.coalesce(bindparam("new_device", type_=String()), _Q.device), + } + if stamp is not None: + values[stamp] = bindparam("now") + return update(session_queue).where(_THE_ITEM, _NOT_FINISHED).values(values) + + +def _stamp(status: str) -> Optional[str]: + if status == "in_progress": + return "started_at" + return "completed_at" if status in _TERMINAL else None + + +# By id lists rather than bound sets: MariaDB before 11.1 runs an UPDATE or DELETE with an IN subquery as a scan of +# the table that evaluates the subquery for every row. +_CANCEL_IDS = ( + update(session_queue) + .where(_Q.item_id.in_(bindparam("item_ids", expanding=True)), _Q.status.in_(bindparam("statuses", expanding=True))) + .values(status=literal("canceled"), status_sequence=_NEXT_SEQUENCE, completed_at=bindparam("now")) +) +_DELETE_IDS = delete(session_queue).where(_Q.item_id.in_(bindparam("item_ids", expanding=True))) +_INTERRUPTED = select(_Q.item_id).where(_Q.status.in_([literal("in_progress"), literal("waiting")])) + +_SET_SESSION = ( + update(session_queue).where(_THE_ITEM).values(session=bindparam("new_session"), session_revision=_NEXT_REVISION) +) +_SET_ACTIVE_SESSION = _SET_SESSION.where(_NOT_FINISHED) +# Suspends the parent for the children it just enqueued, unless its session changed meanwhile or it finished. +_WAIT_FOR_CHILDREN = ( + update(session_queue) + .where(_THE_ITEM, _NOT_FINISHED, _Q.session == bindparam("expected_session")) + .values( + session=bindparam("new_session"), + session_revision=_NEXT_REVISION, + status=literal("waiting"), + status_sequence=_NEXT_SEQUENCE, + ) +) +_WORKFLOW_CALL_ITEMS = ( + select(_Q.item_id).where(_Q.workflow_call_id == bindparam("workflow_call_id")).order_by(_Q.item_id) +) + +_INSERT_ITEM = insert(session_queue) +# Every enqueue counts the pending backlog, which the status index finds, not the queue's history. +_PENDING_COUNT = select(func.count()).where(_IN_QUEUE_UNINDEXED, _PENDING) +_TOP_PENDING_PRIORITY = select(func.max(_Q.priority)).where(_IN_QUEUE_UNINDEXED, _PENDING) +# A batch is a few rows of the batch index; an index on the account would read the account's whole history. +_IN_BATCH = and_( + _IN_QUEUE_UNINDEXED, + Unindexed(_Q.user_id) == bindparam("user_id"), + _Q.batch_id == bindparam("batch_id"), +) +_BATCH_TAKEN = select(literal(1)).where(_IN_BATCH).limit(fixed_limit(1)) +_BATCH_ITEM_IDS = select(_Q.item_id).where(_IN_BATCH).order_by(_Q.item_id) +# Whether any queue has work to finish, for maintenance that must not run beside it. +_ACTIVE_WORK = ( + select(literal(1)).where(_Q.status.in_([literal("pending"), literal("in_progress")])).limit(fixed_limit(1)) +) + +_THE_RECEIPT = and_( + _R.queue_id == bindparam("queue_id"), + _R.user_id == bindparam("user_id"), + _R.idempotency_key == bindparam("idempotency_key"), +) +_SETTLED_RECEIPT = select(_R.payload_hash, _R.batch_id, _R.requested, _R.enqueued, _R.priority, _R.item_ids).where( + _THE_RECEIPT +) +_ACKNOWLEDGE = ( + update(session_queue_enqueue_receipts) + .where( + _R.queue_id == bindparam("target_queue_id"), + _R.user_id == bindparam("target_user_id"), + _R.idempotency_key == bindparam("target_idempotency_key"), + ) + .values(acknowledged_at=bindparam("now")) +) +_DELETE_ACKNOWLEDGED = delete(session_queue_enqueue_receipts).where( + _R.user_id == bindparam("user_id"), _R.acknowledged_at.is_not(None), _R.acknowledged_at <= bindparam("before") +) +_UNACKNOWLEDGED = _R.acknowledged_at.is_(None) +_RECEIPT_USAGE = select( + func.count(), + func.coalesce(func.sum(_R.byte_size), 0), + func.coalesce(func.sum(case((_UNACKNOWLEDGED, 1), else_=0)), 0), + func.coalesce(func.sum(case((_UNACKNOWLEDGED, _R.byte_size), else_=0)), 0), +).where(_R.user_id == bindparam("user_id")) +_INSERT_RECEIPT = insert(session_queue_enqueue_receipts) + +_ACTIVE_ITEM_IDS = select(_Q.item_id).where(_Q.status.in_([literal(status) for status in ACTIVE_QUEUE_STATUSES])) +# History: finished items that no active workflow call still needs. A completed child is recovery state for its active +# root until the root ends, and the intermediates cleanup protects the root's media through it. +_PRUNABLE = and_( + _Q.status.in_([literal(status) for status in _TERMINAL]), + or_(_Q.root_item_id.is_(None), _Q.root_item_id.not_in(_ACTIVE_ITEM_IDS)), +) +_LATEST_FIRST = (func.coalesce(_Q.completed_at, _Q.updated_at, _Q.created_at).desc(), _Q.item_id.desc()) + + +class Scope(NamedTuple): + """The items a bulk cancel or delete applies to: the queue's, narrowed by every filter given.""" + + queue_id: str + user_id: Optional[str] = None + origin_prefix: Optional[str] = None + destination: Optional[str] = None + batch_ids: Optional[Sequence[str]] = None + # Leave out the workflow-call chains of the items running when the scope is changed. + except_current: bool = False + + def parameters(self) -> dict[str, Any]: + return { + "queue_id": self.queue_id, + "user_id": self.user_id, + "origin_pattern": _origin_pattern(self.origin_prefix), + "destination": self.destination, + "batch_ids": sorted(set(self.batch_ids)) if self.batch_ids is not None else None, + } + + +@functools.cache +def _scoped( + by_user: bool, + by_origin: bool, + by_destination: bool, + by_batch: bool, + statuses: tuple[str, ...], + locked: bool, + ordered: bool, +) -> Select[Any]: + """The (item id, owner) of the scope's items of these statuses (all, for none), locked for a change if `locked`; + unlocked, only their ids. (`user_id` and `status` follow the session in the row, which SQLite reaches through + the session's overflow pages: a read of ids alone lets its status index answer the status.) + + The rows come through a status index (or, for batches, the batch index), not the queue's listing index, which + would walk the queue's history. With `ordered` they come in id order: a server locks them in that order, which + keeps two such changes from deadlocking. SQLite locks no rows, and an ORDER BY there would move the plan onto an + index that reads the scope's whole history, so its callers sort. + """ + conditions: list[ColumnElement[bool]] = [_IN_QUEUE_UNINDEXED] + if by_user: + # Beside a batch, the batch index finds the rows; beside a status, the status index leads with the account. + user: ColumnElement[Any] = Unindexed(_Q.user_id) if by_batch else _Q.user_id + conditions.append(user == bindparam("user_id")) + if by_origin: + conditions.append(_origin_matches()) + if by_destination: + conditions.append(_Q.destination == bindparam("destination")) + if by_batch: + # An IN list rather than a bound set: callers name a batch or a few, which the batch index finds. + conditions.append(_Q.batch_id.in_(bindparam("batch_ids", expanding=True))) + if statuses: + conditions.append(_Q.status.in_([literal(status) for status in statuses])) + statement = select(_Q.item_id) if not locked else select(_Q.item_id, _Q.user_id) + statement = statement.where(*conditions) + if ordered: + statement = statement.order_by(_Q.item_id) + return statement.with_for_update() if locked else statement + + +def _scope_shape(scope: Scope) -> tuple[bool, bool, bool, bool]: + return ( + scope.user_id is not None, + scope.origin_prefix is not None, + scope.destination is not None, + scope.batch_ids is not None, + ) + + +def _locked_scope(conn: Connection, scope: Scope, statuses: tuple[str, ...]) -> Sequence[Sequence[Any]]: + """The (item id, owner) of the scope's items of these statuses, locked until the transaction ends. + + The running chains are read after the lock, and left out then: a claim of a locked item waits for this + transaction, and a child enqueue in flight has committed when the lock is granted, so the chains read are those + the change applies against. (Read before the lock, an item could start running, or a running item enqueue its + children, between the read and the change.)""" + statement = _scoped(*_scope_shape(scope), statuses, True, conn.dialect.name != "sqlite") + rows = sorted(conn.execute(statement, scope.parameters()).all()) + if scope.except_current and rows: + running = _current_chain(conn, scope.queue_id) + rows = [row for row in rows if int(row[0]) not in running] + return rows + + +def _cancel(conn: Connection, item_ids: Sequence[int], statuses: Sequence[str]) -> int: + """Cancels the items of these statuses, by id; how many.""" + now = now_text() + canceled = 0 + for chunk in itertools.batched(item_ids, IN_CHUNK): + parameters = {"item_ids": list(chunk), "statuses": list(statuses), "now": now} + canceled += conn.execute(_CANCEL_IDS, parameters).rowcount + return canceled + + +def _delete(conn: Connection, item_ids: Iterable[int]) -> int: + """Deletes the items, by id; how many.""" + return sum( + conn.execute(_DELETE_IDS, {"item_ids": list(chunk)}).rowcount for chunk in itertools.batched(item_ids, IN_CHUNK) + ) + + +@functools.cache +def _prunable(by_user: bool, latest: bool) -> Select[Any]: + """The ids of the queue's history (the account's, if given); with `latest`, its `keep` latest items. (Found + through the status index: the listing index holds no `root_item_id`, so it would read every row as well.)""" + conditions = [_IN_QUEUE_UNINDEXED, _PRUNABLE] + if by_user: + conditions.append(_Q.user_id == bindparam("user_id")) + statement = select(_Q.item_id).where(*conditions) + return statement.order_by(*_LATEST_FIRST).limit(bindparam("keep")) if latest else statement + + +def _by_owner(rows: Sequence[Sequence[Any]]) -> dict[str, list[int]]: + by_owner: dict[str, list[int]] = {} + for item_id, owner in rows: + by_owner.setdefault(owner, []).append(int(item_id)) + return by_owner + + +def _current_chain(conn: Connection, queue_id: str) -> set[int]: + in_progress: Sequence[int] = conn.execute(_IN_PROGRESS_IDS, {"queue_id": queue_id}).scalars().all() + if not in_progress: + waiting_child = conn.execute(_NEXT_WAITING_CHILD, {"queue_id": queue_id}).scalar() + in_progress = [waiting_child] if waiting_child is not None else [] + chain: set[int] = set() + for item_id in in_progress: + # A running item whose parent row is gone still protects what of its chain remains. + chain.update(_chain(conn, item_id, reachable=True) or []) + return chain + + +def _status_row(row: Optional[Sequence[Any]]) -> Optional[dict[str, Any]]: + return dict(zip(_STATUS_NAMES, row, strict=True)) if row is not None else None + + +class Transition(NamedTuple): + """What a status change found: whether the item exists, and its status columns if this call changed them.""" + + exists: bool + changed: Optional[dict[str, Any]] + + +class SettledReceipt(NamedTuple): + payload_hash: str + batch_id: str + requested: int + enqueued: int + priority: int + # The enqueued items' ids, as a JSON array. + item_ids: str + + +class ReceiptUsage(NamedTuple): + count: int + bytes: int + unacknowledged_count: int + unacknowledged_bytes: int + + +ITEM_COLUMNS = ( + "queue_id", + "session", + "session_id", + "batch_id", + "field_values", + "priority", + "workflow", + "origin", + "destination", + "retried_from_item_id", + "user_id", + "project_id", +) +"""The columns of the tuples `prepare_values_to_insert()` builds, in their order.""" + + +class SessionQueueQueries(QueryModule): + @read + def enqueue_receipt( + self, conn: Connection, queue_id: str, user_id: str, idempotency_key: str + ) -> Optional[EnqueueReceipt]: + parameters = {"queue_id": queue_id, "user_id": user_id, "idempotency_key": idempotency_key} + row = conn.execute(_RECEIPT, parameters).first() + return EnqueueReceipt(row[0], row[1], row[2], row[3]) if row is not None else None + + @mapped(_item_or_none) + @read + def item(self, conn: Connection, item_id: int) -> Optional[Row[Any]]: + return conn.execute(_ITEM_BY_ID, {"item_id": item_id}).first() + + @mapped(_item_or_none) + @read + def item_by_status( + self, conn: Connection, queue_id: str, status: str, origin_prefix: Optional[str] + ) -> Optional[Row[Any]]: + """The `pending` item that runs next (highest priority, then oldest), or the `in_progress` one first + claimed.""" + statement = _item_by_status(status, origin_prefix is not None) + parameters = {"queue_id": queue_id, "origin_pattern": _origin_pattern(origin_prefix)} + return _row_of(conn, conn.execute(statement, parameters).scalar()) + + @read + def workflow(self, conn: Connection, item_id: int) -> Optional[str]: + """The item's workflow as stored: JSON text, or None also for an item that does not exist.""" + return conn.execute(_WORKFLOW, {"item_id": item_id}).scalar() + + @mapped(_items) + @read + def all_items(self, conn: Connection, queue_id: str, destination: Optional[str]) -> Sequence[Row[Any]]: + """Every item of the queue, or of its destination, highest priority first, then in enqueue order.""" + parameters = {"queue_id": queue_id, "destination": destination} + return conn.execute(_all_items(destination is not None), parameters).all() + + @read + def count(self, conn: Connection, queue_id: str) -> int: + """How many items the queue holds, of every status.""" + return int(conn.execute(_COUNT, {"queue_id": queue_id}).scalar_one()) + + @mapped(_items) + @read + def page( + self, + conn: Connection, + queue_id: str, + *, + limit: int, + after: Optional[tuple[int, int]], + status: Optional[str], + destination: Optional[str], + ) -> Sequence[Row[Any]]: + """Up to `limit` items, highest priority first, then in enqueue order; with `after` (a priority and an item + id), those that come after that item. Every filter applies to every page.""" + statement = _page(status is not None, destination is not None, after is not None) + parameters = { + "queue_id": queue_id, + "status": status, + "destination": destination, + "after_priority": after[0] if after is not None else None, + "after_item_id": after[1] if after is not None else None, + "limit": limit, + } + return conn.execute(statement, parameters).all() + + @read + def item_ids( + self, + conn: Connection, + queue_id: str, + *, + descending: bool, + user_id: Optional[str], + origin_prefix: Optional[str], + limit: Optional[int] = None, + ) -> list[int]: + """The ids of the queue's items (the account's, with the origin, if given) by creation time; with `limit`, the + first `limit` of that order.""" + statement = _item_ids(descending, user_id is not None, origin_prefix is not None, limit is not None) + parameters = { + "queue_id": queue_id, + "user_id": user_id, + "origin_pattern": _origin_pattern(origin_prefix), + "limit": limit, + } + return list(conn.execute(statement, parameters).scalars().all()) + + @mapped(_summaries) + @read + def summaries(self, conn: Connection, queue_id: str, item_ids: Sequence[int]) -> Sequence[Row[Any]]: + """The summary columns of the named items of the queue that exist, in no order.""" + return conn.execute(_SUMMARIES, {"queue_id": queue_id, "item_ids": bound_set(item_ids)}).all() + + @read + def status_counts( + self, conn: Connection, queue_id: str, user_id: Optional[str], origin_prefix: Optional[str] + ) -> StatusCounts: + by_origin = origin_prefix is not None + parameters = {"queue_id": queue_id, "user_id": user_id, "origin_pattern": _origin_pattern(origin_prefix)} + rows = conn.execute(_status_counts(user_id is not None, by_origin), parameters).all() + current = conn.execute(_current(by_origin), parameters).first() + return StatusCounts( + counts={str(row[0]): int(row[1]) for row in rows}, + own_counts={str(row[0]): int(row[2]) for row in rows if row[2]}, + current=CurrentItem(current[0], current[1], current[2], current[3]) if current is not None else None, + ) + + @read + def batch_counts(self, conn: Connection, queue_id: str, batch_id: str, user_id: Optional[str]) -> BatchCounts: + parameters = {"queue_id": queue_id, "batch_id": batch_id, "user_id": user_id} + rows = conn.execute(_batch_counts(user_id is not None), parameters).all() + return BatchCounts( + counts=_counts(rows), + origin=rows[0][2] if rows else None, + destination=rows[0][3] if rows else None, + ) + + @mapped(_counts) + @read + def destination_counts( + self, conn: Connection, queue_id: str, destination: str, user_id: Optional[str] + ) -> Sequence[Row[Any]]: + parameters = {"queue_id": queue_id, "destination": destination, "user_id": user_id} + return conn.execute(_destination_counts(user_id is not None), parameters).all() + + @read + def chain_item_ids(self, conn: Connection, item_id: int) -> Optional[list[int]]: + """The item's workflow-call chain: its ancestors from its parent up, the item, and every descendant of the + root, breadth first. None if the item, or an ancestor, does not exist.""" + return _chain(conn, item_id) + + @read + def descendant_ids(self, conn: Connection, item_id: int) -> list[int]: + """Every descendant of the item, breadth first.""" + return _descendants(conn, item_id) + + @read + def current_chain_item_ids(self, conn: Connection, queue_id: str) -> set[int]: + """The ids in the workflow-call chains of every item in progress, or, with none in progress, in the chain of + the child that runs next while its parent waits.""" + return _current_chain(conn, queue_id) + + # --- Dequeue ---------------------------------------------------------------------------------------------------- + + @mapped(_item_or_none) + @read + def next_to_run(self, conn: Connection, *, round_robin: bool) -> Optional[Row[Any]]: + """The pending item to run next: in round robin, the best of the account served longest ago; else the highest + priority, then the oldest.""" + return _row_of(conn, conn.execute(_ROUND_ROBIN_NEXT if round_robin else _FIFO_NEXT).scalar()) + + @read + def warmest( + self, + conn: Connection, + *, + user_id: Optional[str], + priority: int, + after_item_id: int, + lookahead: int, + resident_keys: Sequence[str], + ) -> Optional[dict[str, Any]]: + """The account's pending item of this priority, from `after_item_id` to `lookahead` items after it, whose + session names the most of the resident model keys; None when none names any.""" + keys = list(resident_keys)[:MAX_AFFINITY_MODEL_KEYS] + if not keys: + return None + slots = _slots(len(keys)) + parameters: dict[str, Any] = {f"key_{i}": key for i, key in enumerate(keys)} + parameters.update({f"key_{i}": None for i in range(len(keys), slots)}) + parameters.update( + { + "user_id": user_id, + "priority": priority, + "first_item_id": after_item_id, + "last_item_id": after_item_id + lookahead, + } + ) + row = conn.execute(_affinity(slots), parameters).first() + if row is None or not row[-1]: + return None + return _item(row[:-1]) + + @mapped(_status_row) + @write + def claim(self, conn: Connection, item_id: int, device: Optional[str]) -> Optional[Row[Any]]: + """Moves the item from pending to in progress for `device`; its status columns, or None if it is no longer + pending (another worker claimed it, or it was canceled).""" + parameters = {"target_item_id": item_id, "new_device": device, "now": now_text()} + if conn.execute(_CLAIM, parameters).rowcount == 0: + return None + return conn.execute(_STATUS_ROW, {"item_id": item_id}).first() + + # --- Status changes --------------------------------------------------------------------------------------------- + + @write + def transition( + self, + conn: Connection, + item_id: int, + status: str, + *, + error_type: Optional[str] = None, + error_message: Optional[str] = None, + error_traceback: Optional[str] = None, + device: Optional[str] = None, + ) -> Transition: + """Moves the item to `status` unless it is finished (completed, failed or canceled).""" + parameters = { + "target_item_id": item_id, + "new_status": status, + "new_error_type": error_type, + "new_error_message": error_message, + "new_error_traceback": error_traceback, + "new_device": device, + "now": now_text(), + } + if conn.execute(_transition(_stamp(status)), parameters).rowcount == 0: + return Transition(exists=conn.execute(_ITEM_EXISTS, {"item_id": item_id}).first() is not None, changed=None) + return Transition(exists=True, changed=_status_row(conn.execute(_STATUS_ROW, {"item_id": item_id}).first())) + + @write + def cancel_interrupted(self, conn: Connection) -> None: + """Cancels every item in progress or waiting, and the unfinished items of their workflow-call chains: none can + resume after a restart.""" + chain: set[int] = set() + for item_id in conn.execute(_INTERRUPTED).scalars().all(): + chain.update(_chain(conn, item_id, reachable=True) or [item_id]) + _cancel(conn, sorted(chain), ACTIVE_QUEUE_STATUSES) + + @read + def in_progress(self, conn: Connection, scope: Scope) -> list[int]: + """The ids of the scope's items in progress.""" + statement = _scoped(*_scope_shape(scope), ("in_progress",), False, conn.dialect.name != "sqlite") + return sorted(int(row[0]) for row in conn.execute(statement, scope.parameters()).all()) + + @read + def has_active_work(self, conn: Connection) -> bool: + """Whether any queue has an item pending or in progress.""" + return conn.execute(_ACTIVE_WORK).first() is not None + + @write + def cancel_waiting_work(self, conn: Connection, scope: Scope) -> dict[str, list[int]]: + """Cancels the scope's pending and waiting items; the ids it canceled, by owner.""" + rows = _locked_scope(conn, scope, _WAITING_WORK) + _cancel(conn, [int(row[0]) for row in rows], _WAITING_WORK) + return _by_owner(rows) + + @write + def delete_scope(self, conn: Connection, scope: Scope, *, waiting_work_only: bool = False) -> dict[str, list[int]]: + """Deletes the scope's items, or only its pending and waiting ones; the ids it deleted, by owner.""" + rows = _locked_scope(conn, scope, _WAITING_WORK if waiting_work_only else ()) + _delete(conn, (int(row[0]) for row in rows)) + return _by_owner(rows) + + @write + def delete_items(self, conn: Connection, item_ids: Sequence[int]) -> None: + _delete(conn, item_ids) + + @write + def prune(self, conn: Connection, queue_id: str, *, user_id: Optional[str], keep: Optional[int] = None) -> int: + """Deletes the queue's history (the account's, if given), all of it or all but its `keep` latest items; how + many items it deleted.""" + parameters = {"queue_id": queue_id, "user_id": user_id, "keep": keep} + prunable: set[int] = set(conn.execute(_prunable(user_id is not None, False), parameters).scalars().all()) + if keep is not None: + prunable.difference_update(conn.execute(_prunable(user_id is not None, True), parameters).scalars().all()) + # Counted as deleted: a clear or prune running meanwhile may have deleted some of them first. + return _delete(conn, sorted(prunable)) + + # --- Sessions and workflow calls -------------------------------------------------------------------------------- + + @write + def set_session(self, conn: Connection, item_id: int, session: str, *, only_unfinished: bool) -> Optional[bool]: + """Stores the item's session, with `only_unfinished` only while it is not finished. Whether it was stored; None + if the item does not exist.""" + statement = _SET_ACTIVE_SESSION if only_unfinished else _SET_SESSION + if conn.execute(statement, {"target_item_id": item_id, "new_session": session}).rowcount != 0: + return True + return False if conn.execute(_ITEM_EXISTS, {"item_id": item_id}).first() is not None else None + + @locking + def lock_status(self, conn: Connection, item_id: int) -> Optional[str]: + """The item's status, its row locked until the transaction ends; None if it does not exist.""" + return conn.execute(_LOCK_STATUS, {"item_id": item_id}).scalar() + + @locking + def lock_item(self, conn: Connection, item_id: int) -> Optional[dict[str, Any]]: + """The item's row (without its owner's name), locked until the transaction ends; None if it does not exist.""" + row = conn.execute(_LOCK_ITEM, {"item_id": item_id}).first() + return dict(zip(_LOCK_ITEM_NAMES, row, strict=True)) if row is not None else None + + @write + def insert_item(self, conn: Connection, values: dict[str, Any]) -> int: + """Adds an item; its id.""" + now = now_text() + inserted = conn.execute(_INSERT_ITEM, {"created_at": now, "updated_at": now, **values}) + return int(inserted.inserted_primary_key[0]) + + @write + def wait_for_children(self, conn: Connection, item_id: int, *, session: str, expected_session: str) -> bool: + """Suspends the item with its new session for the children it enqueued; False if its session is no longer + `expected_session` or it finished meanwhile.""" + parameters = {"target_item_id": item_id, "new_session": session, "expected_session": expected_session} + return conn.execute(_WAIT_FOR_CHILDREN, parameters).rowcount != 0 + + @read + def workflow_call_item_ids(self, conn: Connection, workflow_call_id: str) -> list[int]: + return [ + int(item_id) + for item_id in conn.execute(_WORKFLOW_CALL_ITEMS, {"workflow_call_id": workflow_call_id}).scalars().all() + ] + + # --- Enqueueing ------------------------------------------------------------------------------------------------- + + @read + def pending_count(self, conn: Connection, queue_id: str) -> int: + return int(conn.execute(_PENDING_COUNT, {"queue_id": queue_id}).scalar_one()) + + @read + def top_pending_priority(self, conn: Connection, queue_id: str) -> int: + """The highest priority of the queue's pending items; 0 for none.""" + return int(conn.execute(_TOP_PENDING_PRIORITY, {"queue_id": queue_id}).scalar() or 0) + + @read + def batch_taken(self, conn: Connection, queue_id: str, user_id: str, batch_id: str) -> bool: + parameters = {"queue_id": queue_id, "user_id": user_id, "batch_id": batch_id} + return conn.execute(_BATCH_TAKEN, parameters).first() is not None + + @write + def insert_items(self, conn: Connection, rows: Sequence[Sequence[Any]]) -> None: + """Adds items from tuples of `ITEM_COLUMNS`, stamped with one time.""" + if rows: + now = now_text() + values = [ + {**dict(zip(ITEM_COLUMNS, row, strict=True)), "created_at": now, "updated_at": now} for row in rows + ] + conn.execute(_INSERT_ITEM, values) + + @read + def batch_item_ids(self, conn: Connection, queue_id: str, user_id: str, batch_id: str) -> list[int]: + parameters = {"queue_id": queue_id, "user_id": user_id, "batch_id": batch_id} + return [int(item_id) for item_id in conn.execute(_BATCH_ITEM_IDS, parameters).scalars().all()] + + @read + def settled_receipt( + self, conn: Connection, queue_id: str, user_id: str, idempotency_key: str + ) -> Optional[SettledReceipt]: + parameters = {"queue_id": queue_id, "user_id": user_id, "idempotency_key": idempotency_key} + row = conn.execute(_SETTLED_RECEIPT, parameters).first() + return SettledReceipt(row[0], row[1], row[2], row[3], row[4], row[5]) if row is not None else None + + @write + def delete_acknowledged_receipts(self, conn: Connection, user_id: str, older_than: timedelta) -> None: + """Deletes the account's receipts acknowledged longer ago than `older_than`.""" + before = timestamp_text(datetime.now(timezone.utc) - older_than) + conn.execute(_DELETE_ACKNOWLEDGED, {"user_id": user_id, "before": before}) + + @read + def receipt_usage(self, conn: Connection, user_id: str) -> ReceiptUsage: + count, size, unacknowledged, unacknowledged_size = conn.execute(_RECEIPT_USAGE, {"user_id": user_id}).one() + return ReceiptUsage(int(count), int(size), int(unacknowledged), int(unacknowledged_size)) + + @write + def insert_receipt(self, conn: Connection, values: dict[str, Any]) -> None: + conn.execute(_INSERT_RECEIPT, {"created_at": now_text(), **values}) + + @write + def acknowledge_receipt(self, conn: Connection, queue_id: str, user_id: str, idempotency_key: str) -> None: + parameters = { + "target_queue_id": queue_id, + "target_user_id": user_id, + "target_idempotency_key": idempotency_key, + "now": now_text(), + } + conn.execute(_ACKNOWLEDGE, parameters) diff --git a/invokeai/app/services/shared/database/queries/style_presets.py b/invokeai/app/services/shared/database/queries/style_presets.py new file mode 100644 index 00000000000..8cab8318299 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/style_presets.py @@ -0,0 +1,134 @@ +"""Style presets: the accounts' saved prompt templates, and the bundled ones (type `default`).""" + +import functools +from collections.abc import Mapping, Sequence +from typing import Any, Optional + +from sqlalchemy import ( + Connection, + Row, + Select, + Update, + bindparam, + delete, + insert, + literal, + or_, + select, + true, + update, +) + +from invokeai.app.services.shared.database.dialect import CaseInsensitiveOrder +from invokeai.app.services.shared.database.queries.base import QueryModule, mapped, read, write +from invokeai.app.services.shared.database.schema.style_presets import style_presets +from invokeai.app.services.style_preset_records.style_preset_records_common import ( + PresetType, + StylePresetChanges, + StylePresetRecordDTO, + StylePresetWithoutId, +) + +_P = style_presets.c +_COLUMNS = (_P.id, _P.name, _P.preset_data, _P.type, _P.user_id, _P.is_public) +_NAMES = tuple(column.name for column in _COLUMNS) + +_GET = select(*_COLUMNS).where(_P.id == bindparam("style_preset_id")) +_INSERT = insert(style_presets) +_DELETE = delete(style_presets).where(_P.id == bindparam("style_preset_id")) +_DELETE_BUNDLED = delete(style_presets).where(_P.type == literal(PresetType.Default.value)) + + +@functools.cache +def _update(fields: tuple[str, ...]) -> Update: + """Sets these fields of the preset; one statement per set of fields, of which there are few.""" + return ( + update(style_presets) + .where(_P.id == bindparam("target_style_preset_id")) + .values({field: bindparam(f"new_{field}") for field in fields}) + ) + + +@functools.cache +def _list(visible_to_all: bool, scoped: bool, typed: bool) -> Select[Any]: + """The presets an account sees, in name order: an administrator every one, anyone else the bundled and shared + ones, and with `scoped` also its own.""" + statement = select(*_COLUMNS) + if not visible_to_all: + visible = [_P.type == literal(PresetType.Default.value), _P.is_public == true()] + if scoped: + visible.append(_P.user_id == bindparam("user_id")) + statement = statement.where(or_(*visible)) + if typed: + statement = statement.where(_P.type == bindparam("type")) + # The id breaks ties, so that equal names keep one order. + return statement.order_by(CaseInsensitiveOrder(_P.name), _P.id) + + +def _preset(row: Sequence[Any]) -> StylePresetRecordDTO: + return StylePresetRecordDTO.from_dict(dict(zip(_NAMES, row, strict=True))) + + +def _preset_or_none(row: Optional[Sequence[Any]]) -> Optional[StylePresetRecordDTO]: + return _preset(row) if row is not None else None + + +def _presets(rows: Sequence[Sequence[Any]]) -> list[StylePresetRecordDTO]: + return [_preset(row) for row in rows] + + +class StylePresetQueries(QueryModule): + @mapped(_preset_or_none) + @read + def get(self, conn: Connection, style_preset_id: str) -> Optional[Row[Any]]: + return conn.execute(_GET, {"style_preset_id": style_preset_id}).first() + + @mapped(_presets) + @read + def all( + self, conn: Connection, *, type: Optional[PresetType], user_id: Optional[str], is_admin: bool + ) -> Sequence[Row[Any]]: + """The presets the account sees, of `type` if given, in name order.""" + statement = _list(is_admin, user_id is not None, type is not None) + return conn.execute(statement, {"user_id": user_id, "type": type}).all() + + @write + def insert(self, conn: Connection, presets: Mapping[str, StylePresetWithoutId], user_id: str) -> None: + """Adds the presets, keyed by their new ids, as the account's.""" + if presets: + conn.execute( + _INSERT, + [ + { + "id": style_preset_id, + "name": preset.name, + "preset_data": preset.preset_data.model_dump_json(), + "type": preset.type.value, + "user_id": user_id, + "is_public": preset.is_public, + } + for style_preset_id, preset in presets.items() + ], + ) + + @write + def update(self, conn: Connection, style_preset_id: str, changes: StylePresetChanges) -> None: + """Applies the changes to the preset's name, data and visibility; its type stays as it is.""" + fields: dict[str, Any] = {} + if changes.name is not None: + fields["name"] = changes.name + if changes.preset_data is not None: + fields["preset_data"] = changes.preset_data.model_dump_json() + if changes.is_public is not None: + fields["is_public"] = changes.is_public + if fields: + values = {f"new_{field}": value for field, value in fields.items()} + conn.execute(_update(tuple(sorted(fields))), {"target_style_preset_id": style_preset_id, **values}) + + @write + def delete(self, conn: Connection, style_preset_id: str) -> None: + conn.execute(_DELETE, {"style_preset_id": style_preset_id}) + + @write + def delete_bundled(self, conn: Connection) -> None: + conn.execute(_DELETE_BUNDLED) diff --git a/invokeai/app/services/shared/database/queries/system_prompts.py b/invokeai/app/services/shared/database/queries/system_prompts.py new file mode 100644 index 00000000000..7a40a683be8 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/system_prompts.py @@ -0,0 +1,149 @@ +"""System prompts for the prompt-expansion models: the accounts' own, and shared ones.""" + +import functools +from collections.abc import Sequence +from typing import Any, Optional + +from sqlalchemy import ( + Connection, + Delete, + Row, + Select, + Update, + bindparam, + delete, + insert, + literal, + or_, + select, + true, + update, +) + +from invokeai.app.services.shared.database.dialect import CaseInsensitiveOrder +from invokeai.app.services.shared.database.queries.base import QueryModule, mapped, read, write +from invokeai.app.services.shared.database.schema.system_prompts import system_prompts +from invokeai.app.services.shared.database.types import now_text +from invokeai.app.services.system_prompt_records.system_prompt_records_common import ( + SystemPromptChanges, + SystemPromptRecordDTO, + SystemPromptWithoutId, +) + +_S = system_prompts.c +_COLUMNS = (_S.id, _S.name, _S.content, _S.max_tokens, _S.user_id, _S.is_public, _S.created_at, _S.updated_at) +_NAMES = tuple(column.name for column in _COLUMNS) + +_GET = select(*_COLUMNS).where(_S.id == bindparam("system_prompt_id")) +_INSERT = insert(system_prompts) +_ALL = select(*_COLUMNS).order_by(CaseInsensitiveOrder(_S.name), _S.id) +# An account's own prompts and the shared ones. The id breaks ties of equal names. +_VISIBLE = ( + select(*_COLUMNS) + .where(or_(_S.user_id == bindparam("user_id"), _S.is_public == true())) + .order_by(CaseInsensitiveOrder(_S.name), _S.id) +) + + +@functools.cache +def _owned(scoped: bool) -> Select[Any]: + """The prompt, if it exists and, with `scoped`, the account owns it.""" + statement = select(literal(1)).where(_S.id == bindparam("system_prompt_id")) + return statement.where(_S.user_id == bindparam("user_id")) if scoped else statement + + +@functools.cache +def _update(fields: tuple[str, ...], scoped: bool) -> Update: + """Sets these fields of the prompt, with `scoped` only if the account owns it; few sets of fields are possible.""" + statement = ( + update(system_prompts) + .where(_S.id == bindparam("target_system_prompt_id")) + .values({field: bindparam(f"new_{field}") for field in fields}) + ) + return statement.where(_S.user_id == bindparam("owner_id")) if scoped else statement + + +@functools.cache +def _delete(scoped: bool) -> Delete: + statement = delete(system_prompts).where(_S.id == bindparam("system_prompt_id")) + return statement.where(_S.user_id == bindparam("user_id")) if scoped else statement + + +def _prompt(row: Sequence[Any]) -> SystemPromptRecordDTO: + return SystemPromptRecordDTO.from_dict(dict(zip(_NAMES, row, strict=True))) + + +def _prompt_or_none(row: Optional[Sequence[Any]]) -> Optional[SystemPromptRecordDTO]: + return _prompt(row) if row is not None else None + + +def _prompts(rows: Sequence[Sequence[Any]]) -> list[SystemPromptRecordDTO]: + return [_prompt(row) for row in rows] + + +class SystemPromptQueries(QueryModule): + @mapped(_prompt_or_none) + @read + def get(self, conn: Connection, system_prompt_id: str) -> Optional[Row[Any]]: + return conn.execute(_GET, {"system_prompt_id": system_prompt_id}).first() + + @mapped(_prompts) + @read + def all(self, conn: Connection, user_id: Optional[str]) -> Sequence[Row[Any]]: + """Every prompt, in name order; with `user_id`, the account's own and the shared ones.""" + if user_id is None: + return conn.execute(_ALL).all() + return conn.execute(_VISIBLE, {"user_id": user_id}).all() + + @write + def insert( + self, conn: Connection, system_prompt_id: str, prompt: SystemPromptWithoutId, user_id: str, is_public: bool + ) -> str: + """Adds the prompt; the time it was created at, which is also the time it was last updated.""" + now = now_text() + conn.execute( + _INSERT, + { + "id": system_prompt_id, + "name": prompt.name, + "content": prompt.content, + "max_tokens": prompt.max_tokens, + "user_id": user_id, + "is_public": is_public, + "created_at": now, + "updated_at": now, + }, + ) + return now + + @write + def update( + self, conn: Connection, system_prompt_id: str, changes: SystemPromptChanges, user_id: Optional[str] + ) -> bool: + """Applies the changes, with `user_id` only if the account owns the prompt; whether the prompt exists (and + the account owns it). Without changes, nothing is written.""" + fields: dict[str, Any] = {} + if changes.name is not None: + fields["name"] = changes.name + if changes.content is not None: + fields["content"] = changes.content + if changes.is_public is not None: + fields["is_public"] = changes.is_public + # `max_tokens` is the one field whose null is a value rather than "no change": it means "drop back to the + # endpoint default". `is not None` would make a cap unclearable once set, so the presence of the key in the + # request body decides instead. + if "max_tokens" in changes.model_fields_set: + fields["max_tokens"] = changes.max_tokens + scoped = user_id is not None + if not fields: + parameters = {"system_prompt_id": system_prompt_id, "user_id": user_id} + return conn.execute(_owned(scoped), parameters).first() is not None + values = {f"new_{field}": value for field, value in fields.items()} + parameters = {"target_system_prompt_id": system_prompt_id, "owner_id": user_id, **values} + return conn.execute(_update(tuple(sorted(fields)), scoped), parameters).rowcount > 0 + + @write + def delete(self, conn: Connection, system_prompt_id: str, user_id: Optional[str]) -> bool: + """Deletes the prompt, with `user_id` only if the account owns it; whether one was deleted.""" + parameters = {"system_prompt_id": system_prompt_id, "user_id": user_id} + return conn.execute(_delete(user_id is not None), parameters).rowcount > 0 diff --git a/invokeai/app/services/shared/database/queries/users.py b/invokeai/app/services/shared/database/queries/users.py new file mode 100644 index 00000000000..0c9a02060f0 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/users.py @@ -0,0 +1,180 @@ +"""Accounts.""" + +import itertools +from collections.abc import Collection, Sequence +from datetime import datetime +from typing import Any, Optional + +from sqlalchemy import Connection, Row, bindparam, delete, func, insert, select, true, update + +from invokeai.app.services.shared.database.dialect import fixed_limit +from invokeai.app.services.shared.database.queries.base import IN_CHUNK, QueryModule, mapped, read, write +from invokeai.app.services.shared.database.schema.users import users +from invokeai.app.services.users.users_common import UserDTO + +# Every query that yields a `UserDTO` selects these columns and maps them with `_user`. A projection copied +# per query drifts: `get_many` once missed `token_epoch` and reported every account it returned as never +# having rotated its password. +_USER_COLUMNS = ( + users.c.user_id, + users.c.email, + users.c.display_name, + users.c.is_admin, + users.c.is_active, + users.c.created_at, + users.c.updated_at, + users.c.last_login_at, + users.c.token_epoch, +) +_IS_ACTIVE_ADMIN = (users.c.is_admin == true(), users.c.is_active == true()) + +_GET = select(*_USER_COLUMNS).where(users.c.user_id == bindparam("user_id")) +_GET_MANY = select(*_USER_COLUMNS).where(users.c.user_id.in_(bindparam("user_ids", expanding=True))) +_GET_BY_EMAIL = select(*_USER_COLUMNS).where(users.c.email == bindparam("email")) +_GET_CREDENTIALS = select(*_USER_COLUMNS, users.c.password_hash).where(users.c.email == bindparam("email")) +_PAGE = ( + select(*_USER_COLUMNS) + .order_by(users.c.created_at.desc(), users.c.user_id) + .limit(bindparam("limit")) + .offset(bindparam("offset")) +) +_COUNT_ACTIVE_ADMINS = select(func.count()).select_from(users).where(*_IS_ACTIVE_ADMIN) +_FIRST_ACTIVE_ADMIN_EMAIL = ( + select(users.c.email).where(*_IS_ACTIVE_ADMIN).order_by(users.c.created_at, users.c.user_id).limit(fixed_limit(1)) +) +_INSERT = insert(users) +# An UPDATE reserves bind names that equal column names for its SET clause. +_RECORD_LOGIN = update(users).where(users.c.user_id == bindparam("target_user_id")) +_DELETE = delete(users).where(users.c.user_id == bindparam("user_id")) + + +def _user(row: Sequence[Any]) -> UserDTO: + # Unpacked rather than read by name: a lookup by name on a row costs ten times one by position. + user_id, email, display_name, is_admin, is_active, created_at, updated_at, last_login_at, token_epoch = row + return UserDTO( + user_id=user_id, + email=email, + display_name=display_name, + is_admin=is_admin, + is_active=is_active, + created_at=datetime.fromisoformat(created_at), + updated_at=datetime.fromisoformat(updated_at), + last_login_at=datetime.fromisoformat(last_login_at) if last_login_at else None, + token_epoch=token_epoch, + ) + + +def _user_or_none(row: Optional[Sequence[Any]]) -> Optional[UserDTO]: + return _user(row) if row is not None else None + + +def _users(rows: Sequence[Sequence[Any]]) -> list[UserDTO]: + return [_user(row) for row in rows] + + +def _users_by_id(rows: Sequence[Sequence[Any]]) -> dict[str, UserDTO]: + return {user.user_id: user for user in map(_user, rows)} + + +def _credentials(row: Optional[Sequence[Any]]) -> Optional[tuple[UserDTO, str]]: + return (_user(row[:-1]), row[-1]) if row is not None else None + + +class UserQueries(QueryModule): + @mapped(_user_or_none) + @read + def get(self, conn: Connection, user_id: str) -> Optional[Row[Any]]: + return conn.execute(_GET, {"user_id": user_id}).first() + + @mapped(_users_by_id) + @read + def get_many(self, conn: Connection, user_ids: Collection[str]) -> list[Row[Any]]: + """The accounts with these ids, keyed by id; unknown ids are absent.""" + rows: list[Row[Any]] = [] + for chunk in itertools.batched(dict.fromkeys(user_ids), IN_CHUNK): + rows.extend(conn.execute(_GET_MANY, {"user_ids": list(chunk)}).all()) + return rows + + @mapped(_user_or_none) + @read + def get_by_email(self, conn: Connection, email: str) -> Optional[Row[Any]]: + return conn.execute(_GET_BY_EMAIL, {"email": email}).first() + + @mapped(_credentials) + @read + def get_credentials(self, conn: Connection, email: str) -> Optional[Row[Any]]: + """The account with this email and its password hash.""" + return conn.execute(_GET_CREDENTIALS, {"email": email}).first() + + @mapped(_users) + @read + def page(self, conn: Connection, limit: int, offset: int) -> Sequence[Row[Any]]: + """Accounts, newest first.""" + return conn.execute(_PAGE, {"limit": limit, "offset": offset}).all() + + @read + def count_active_admins(self, conn: Connection) -> int: + return conn.execute(_COUNT_ACTIVE_ADMINS).scalar_one() + + @read + def first_active_admin_email(self, conn: Connection) -> Optional[str]: + return conn.execute(_FIRST_ACTIVE_ADMIN_EMAIL).scalar_one_or_none() + + @write + def insert( + self, + conn: Connection, + *, + user_id: str, + email: str, + display_name: Optional[str], + password_hash: str, + is_admin: bool, + ) -> None: + conn.execute( + _INSERT, + { + "user_id": user_id, + "email": email, + "display_name": display_name, + "password_hash": password_hash, + "is_admin": is_admin, + }, + ) + + @write + def update( + self, + conn: Connection, + user_id: str, + *, + display_name: Optional[str] = None, + password_hash: Optional[str] = None, + is_admin: Optional[bool] = None, + is_active: Optional[bool] = None, + ) -> None: + """Changes the given fields; None leaves a field as it is. A new password hash also raises the account's + token epoch, which revokes every token issued before it.""" + changes: dict[str, Any] = {} + if display_name is not None: + changes["display_name"] = display_name + if password_hash is not None: + changes["password_hash"] = password_hash + # In SQL, so concurrent changes cannot read, increment and write over each other. + changes["token_epoch"] = users.c.token_epoch + 1 + if is_admin is not None: + changes["is_admin"] = is_admin + if is_active is not None: + changes["is_active"] = is_active + if changes: + conn.execute(update(users).where(users.c.user_id == user_id).values(changes)) + + @write + def record_login(self, conn: Connection, user_id: str, at: datetime) -> None: + # ISO 8601 with its UTC offset, as last_login_at has always been stored: the API reports it as an aware + # time. + conn.execute(_RECORD_LOGIN, {"target_user_id": user_id, "last_login_at": at.isoformat()}) + + @write + def delete(self, conn: Connection, user_id: str) -> None: + conn.execute(_DELETE, {"user_id": user_id}) diff --git a/invokeai/app/services/shared/database/queries/videos.py b/invokeai/app/services/shared/database/queries/videos.py new file mode 100644 index 00000000000..6e049974741 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/videos.py @@ -0,0 +1,468 @@ +"""Video records: one row per video file, with its generation metadata.""" + +import functools +import itertools +from collections.abc import Sequence +from typing import Any, NamedTuple, Optional + +from sqlalchemy import ( + ColumnElement, + Connection, + Row, + Select, + Update, + bindparam, + delete, + false, + func, + literal, + or_, + select, + true, + update, +) + +from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin +from invokeai.app.services.shared.database.dialect import ( + CaseInsensitiveLike, + JsonString, + fixed_limit, + insert_ignore, + like_contains, +) +from invokeai.app.services.shared.database.queries.base import IN_CHUNK, QueryModule, mapped, read, write +from invokeai.app.services.shared.database.schema.boards import board_videos +from invokeai.app.services.shared.database.schema.videos import videos +from invokeai.app.services.video_records.video_records_common import ( + VideoRecord, + VideoRecordChanges, + deserialize_video_record, +) + +_V = videos.c +_BV = board_videos.c + +# The columns of a `VideoRecord`, and the one member of the metadata the clients need on every row: it marks an +# upload converted from an audio file. Only it is read, so that listings do not carry whole metadata documents. +_RECORD = ( + _V.video_name, + _V.video_origin, + _V.video_category, + _V.width, + _V.height, + _V.duration, + _V.fps, + _V.session_id, + _V.node_id, + _V.has_workflow, + _V.is_intermediate, + _V.created_at, + _V.updated_at, + _V.deleted_at, + _V.starred, + _V.video_subfolder, + _V.project_id, + _V.file_size_bytes, + JsonString(_V.metadata, "$.media_origin").label("media_origin"), +) +_RECORD_NAMES = tuple(column.name for column in _RECORD) + +_GET = select(*_RECORD).where(_V.video_name == bindparam("video_name")) +_GET_USER_ID = select(_V.user_id).where(_V.video_name == bindparam("video_name")) +_GET_METADATA = select(_V.metadata).where(_V.video_name == bindparam("video_name")) +_GET_CREATED_AT = select(_V.created_at).where(_V.video_name == bindparam("video_name")) +_EXISTS = select(literal(1)).where(_V.video_name == bindparam("video_name")) +_SUBFOLDERS = select(_V.video_name, _V.video_subfolder).where( + _V.video_name.in_(bindparam("video_names", expanding=True)) +) +_MOST_RECENT_ON_BOARD = ( + select(*_RECORD) + .join(board_videos, _BV.video_name == _V.video_name) + .where(_BV.board_id == bindparam("board_id"), _V.is_intermediate == false()) + .order_by(_V.starred.desc(), _V.created_at.desc(), _V.video_name.desc()) + .limit(fixed_limit(1)) +) +_SET_FILE_SIZE = ( + update(videos) + .where(_V.video_name == bindparam("target_video_name")) + .values(file_size_bytes=bindparam("new_file_size_bytes")) +) +# A backfill fills gaps only: the writer's own measurement, taken after the file exists, wins over one taken before. +_FILL_FILE_SIZE = ( + update(videos) + .where(_V.video_name == bindparam("target_video_name"), _V.file_size_bytes.is_(None)) + .values(file_size_bytes=bindparam("new_file_size_bytes")) +) +_DELETE = delete(videos).where(_V.video_name == bindparam("video_name")) +_DELETE_MANY = delete(videos).where(_V.video_name.in_(bindparam("video_names", expanding=True))) +# Locked, so that a video promoted out of the intermediates meanwhile is either seen promoted or deleted before it is. +_LOCK_INTERMEDIATES = ( + select(_V.video_name) + .where(_V.video_name.in_(bindparam("video_names", expanding=True)), _V.is_intermediate == true()) + .with_for_update() +) +_DELETE_INTERMEDIATES = delete(videos).where( + _V.video_name.in_(bindparam("video_names", expanding=True)), _V.is_intermediate == true() +) + + +@functools.cache +def _insert(dialect_name: str) -> Any: + # A name that is taken keeps its row, as SQLite's INSERT OR IGNORE did. + return insert_ignore(dialect_name, videos) + + +@functools.cache +def _update(fields: tuple[str, ...]) -> Update: + """Sets these fields of the video; one statement per set of fields, of which there are few.""" + return ( + update(videos) + .where(_V.video_name == bindparam("target_video_name")) + .values({field: bindparam(f"new_{field}") for field in fields}) + ) + + +class _Shape(NamedTuple): + """Which filters a listing has: never their values, which are bound, so that the shapes are few.""" + + origin: bool + # The categories' values, sorted; None for no category filter. Each value is bound on its own: SQLAlchemy compiles + # lists of one length once, whatever they hold. + categories: Optional[tuple[str, ...]] + intermediate: bool + # "any" (no board filter), "none" (on no board) or "one". + board: str + # Whether the listing is limited to what one account may see; else everything, as for an administrator. + scoped: bool + search: bool + + +def _shape( + *, + video_origin: Optional[ResourceOrigin], + categories: Optional[Sequence[ImageCategory]], + is_intermediate: Optional[bool], + board_id: Optional[str], + search_term: Optional[str], + user_id: Optional[str], + is_admin: bool, +) -> _Shape: + return _Shape( + origin=video_origin is not None, + categories=None if categories is None else tuple(sorted({ImageCategory(c).value for c in categories})), + intermediate=is_intermediate is not None, + board="any" if board_id is None else "none" if board_id == "none" else "one", + scoped=user_id is not None and not is_admin, + search=bool(search_term), + ) + + +def _parameters( + *, + video_origin: Optional[ResourceOrigin], + is_intermediate: Optional[bool], + board_id: Optional[str], + search_term: Optional[str], + user_id: Optional[str], +) -> dict[str, Any]: + return { + "video_origin": video_origin.value if video_origin is not None else None, + "is_intermediate": is_intermediate, + "board_id": board_id, + # Lowered here, as the legacy storage did: SQLite's LIKE folds the case of ASCII letters only. + "pattern": like_contains(search_term.lower()) if search_term else None, + "user_id": user_id, + } + + +def _conditions(shape: _Shape) -> list[ColumnElement[bool]]: + """The listing's conditions, on the videos LEFT JOINed with their board memberships.""" + conditions: list[ColumnElement[bool]] = [] + if shape.origin: + conditions.append(_V.video_origin == bindparam("video_origin")) + if shape.categories is not None: + conditions.append(_V.video_category.in_([literal(category) for category in shape.categories])) + if shape.intermediate: + conditions.append(_V.is_intermediate == bindparam("is_intermediate")) + own = _V.user_id == bindparam("user_id") + if shape.board == "none": + conditions.append(_BV.board_id.is_(None)) + if shape.scoped: + conditions.append(own) + elif shape.board == "one": + conditions.append(_BV.board_id == bindparam("board_id")) + elif shape.scoped: + # Without a board, still only the account's own videos, so that it cannot list every account's. + conditions.append(own) + if shape.search: + pattern: ColumnElement[str] = bindparam("pattern") + conditions.append(or_(CaseInsensitiveLike(_V.metadata, pattern), CaseInsensitiveLike(_V.created_at, pattern))) + return conditions + + +def _ordering(starred_first: bool, descending: bool) -> list[ColumnElement[Any]]: + ordering: list[ColumnElement[Any]] = ( + [_V.created_at.desc(), _V.video_name.desc()] if descending else [_V.created_at, _V.video_name] + ) + return [_V.starred.desc(), *ordering] if starred_first else ordering + + +_JOINED_FROM = videos.outerjoin(board_videos, _BV.video_name == _V.video_name) + + +@functools.lru_cache(maxsize=128) +def _page(shape: _Shape, starred_first: bool, descending: bool) -> Select[Any]: + """A page of records; the number of shapes is bounded, the cache keeps the recent ones.""" + return ( + select(*_RECORD) + .select_from(_JOINED_FROM) + .where(*_conditions(shape)) + .order_by(*_ordering(starred_first, descending)) + .limit(bindparam("limit")) + .offset(bindparam("offset")) + ) + + +@functools.lru_cache(maxsize=128) +def _count(shape: _Shape) -> Select[Any]: + return select(func.count()).select_from(_JOINED_FROM).where(*_conditions(shape)) + + +@functools.lru_cache(maxsize=128) +def _names(shape: _Shape, starred_first: bool, descending: bool) -> Select[Any]: + return ( + select(_V.video_name) + .select_from(_JOINED_FROM) + .where(*_conditions(shape)) + .order_by(*_ordering(starred_first, descending)) + ) + + +@functools.lru_cache(maxsize=128) +def _starred_count(shape: _Shape) -> Select[Any]: + return select(func.count()).select_from(_JOINED_FROM).where(_V.starred == true(), *_conditions(shape)) + + +def _record(row: Sequence[Any]) -> VideoRecord: + return deserialize_video_record(dict(zip(_RECORD_NAMES, row, strict=True))) + + +def _record_or_none(row: Optional[Sequence[Any]]) -> Optional[VideoRecord]: + return _record(row) if row is not None else None + + +def _records(page: tuple[Sequence[Sequence[Any]], int]) -> tuple[list[VideoRecord], int]: + rows, total = page + return [_record(row) for row in rows], total + + +class VideoQueries(QueryModule): + @mapped(_record_or_none) + @read + def get(self, conn: Connection, video_name: str) -> Optional[Row[Any]]: + return conn.execute(_GET, {"video_name": video_name}).first() + + @read + def user_id(self, conn: Connection, video_name: str) -> Optional[str]: + """The video's owner; None also for a video that does not exist.""" + return conn.execute(_GET_USER_ID, {"video_name": video_name}).scalar() + + @read + def metadata(self, conn: Connection, video_name: str) -> tuple[bool, Optional[str]]: + """Whether the video exists, and its metadata as stored: JSON text, or None.""" + row = conn.execute(_GET_METADATA, {"video_name": video_name}).first() + return (False, None) if row is None else (True, row[0]) + + @read + def exists(self, conn: Connection, video_name: str) -> bool: + return conn.execute(_EXISTS, {"video_name": video_name}).first() is not None + + @read + def subfolders(self, conn: Connection, video_names: Sequence[str]) -> dict[str, str]: + """The subfolder of each named video that exists.""" + subfolders: dict[str, str] = {} + for chunk in itertools.batched(video_names, IN_CHUNK): + subfolders.update((row[0], row[1]) for row in conn.execute(_SUBFOLDERS, {"video_names": list(chunk)}).all()) + return subfolders + + @mapped(_record_or_none) + @read + def most_recent_on_board(self, conn: Connection, board_id: str) -> Optional[Row[Any]]: + """The board's cover: its newest video that is not an intermediate, a starred one first.""" + return conn.execute(_MOST_RECENT_ON_BOARD, {"board_id": board_id}).first() + + @mapped(_records) + @read + def page( + self, + conn: Connection, + *, + offset: int, + limit: int, + starred_first: bool, + descending: bool, + video_origin: Optional[ResourceOrigin], + categories: Optional[Sequence[ImageCategory]], + is_intermediate: Optional[bool], + board_id: Optional[str], + search_term: Optional[str], + user_id: Optional[str], + is_admin: bool, + ) -> tuple[Sequence[Row[Any]], int]: + """A page of the records that match, and how many match in all.""" + shape = _shape( + video_origin=video_origin, + categories=categories, + is_intermediate=is_intermediate, + board_id=board_id, + search_term=search_term, + user_id=user_id, + is_admin=is_admin, + ) + parameters = _parameters( + video_origin=video_origin, + is_intermediate=is_intermediate, + board_id=board_id, + search_term=search_term, + user_id=user_id, + ) + rows = conn.execute( + _page(shape, starred_first, descending), {**parameters, "limit": limit, "offset": offset} + ).all() + return rows, conn.execute(_count(shape), parameters).scalar_one() + + @read + def names( + self, + conn: Connection, + *, + starred_first: bool, + descending: bool, + video_origin: Optional[ResourceOrigin], + categories: Optional[Sequence[ImageCategory]], + is_intermediate: Optional[bool], + board_id: Optional[str], + search_term: Optional[str], + user_id: Optional[str], + is_admin: bool, + ) -> tuple[list[str], int]: + """The names of the videos that match, in order, and how many of them are starred (with `starred_first`; + else 0).""" + shape = _shape( + video_origin=video_origin, + categories=categories, + is_intermediate=is_intermediate, + board_id=board_id, + search_term=search_term, + user_id=user_id, + is_admin=is_admin, + ) + parameters = _parameters( + video_origin=video_origin, + is_intermediate=is_intermediate, + board_id=board_id, + search_term=search_term, + user_id=user_id, + ) + starred = conn.execute(_starred_count(shape), parameters).scalar_one() if starred_first else 0 + names = list(conn.execute(_names(shape, starred_first, descending), parameters).scalars().all()) + return names, starred + + @write + def insert( + self, + conn: Connection, + *, + video_name: str, + video_origin: ResourceOrigin, + video_category: ImageCategory, + width: int, + height: int, + duration: float, + fps: Optional[float], + has_workflow: bool, + is_intermediate: Optional[bool], + starred: Optional[bool], + session_id: Optional[str], + node_id: Optional[str], + metadata: Optional[str], + user_id: str, + video_subfolder: str, + project_id: Optional[str], + ) -> None: + """Adds the video, unless a video of that name exists.""" + conn.execute( + _insert(conn.dialect.name), + { + "video_name": video_name, + "video_origin": video_origin.value, + "video_category": video_category.value, + "width": width, + "height": height, + "duration": float(duration), + "fps": float(fps) if fps is not None else None, + "node_id": node_id, + "session_id": session_id, + "metadata": metadata, + "is_intermediate": is_intermediate, + "starred": starred, + "has_workflow": has_workflow, + "user_id": user_id, + "video_subfolder": video_subfolder, + "project_id": project_id, + }, + ) + + @read + def created_at(self, conn: Connection, video_name: str) -> Optional[str]: + return conn.execute(_GET_CREATED_AT, {"video_name": video_name}).scalar() + + @write + def update(self, conn: Connection, video_name: str, changes: VideoRecordChanges) -> None: + """Applies the changes to the video's category, session, intermediate flag and star.""" + fields: dict[str, Any] = {} + if changes.video_category is not None: + fields["video_category"] = ImageCategory(changes.video_category).value + if changes.session_id is not None: + fields["session_id"] = changes.session_id + if changes.is_intermediate is not None: + fields["is_intermediate"] = changes.is_intermediate + if changes.starred is not None: + fields["starred"] = changes.starred + if fields: + values = {f"new_{field}": value for field, value in fields.items()} + conn.execute(_update(tuple(sorted(fields))), {"target_video_name": video_name, **values}) + + @write + def set_file_size(self, conn: Connection, video_name: str, file_size_bytes: Optional[int]) -> None: + conn.execute(_SET_FILE_SIZE, {"target_video_name": video_name, "new_file_size_bytes": file_size_bytes}) + + @write + def fill_file_sizes(self, conn: Connection, sizes: dict[str, int]) -> None: + """Records the sizes of the videos whose size is not known yet.""" + if sizes: + conn.execute( + _FILL_FILE_SIZE, + [{"target_video_name": name, "new_file_size_bytes": size} for name, size in sizes.items()], + ) + + @write + def delete(self, conn: Connection, video_name: str) -> None: + conn.execute(_DELETE, {"video_name": video_name}) + + @write + def delete_many(self, conn: Connection, video_names: Sequence[str]) -> None: + for chunk in itertools.batched(video_names, IN_CHUNK): + conn.execute(_DELETE_MANY, {"video_names": list(chunk)}) + + @write + def delete_intermediates(self, conn: Connection, video_names: Sequence[str]) -> list[str]: + """Deletes those of the named videos that are intermediates; the names deleted, in the given order.""" + deleted: list[str] = [] + for chunk in itertools.batched(video_names, IN_CHUNK): + names = list(chunk) + locked: set[str] = set(conn.execute(_LOCK_INTERMEDIATES, {"video_names": names}).scalars().all()) + if locked: + conn.execute(_DELETE_INTERMEDIATES, {"video_names": sorted(locked)}) + deleted.extend(name for name in names if name in locked) + return deleted diff --git a/invokeai/app/services/shared/database/queries/wildcards.py b/invokeai/app/services/shared/database/queries/wildcards.py new file mode 100644 index 00000000000..e411734e74e --- /dev/null +++ b/invokeai/app/services/shared/database/queries/wildcards.py @@ -0,0 +1,84 @@ +"""Wildcards: an account's named lists of values that a prompt can draw from.""" + +import functools +import json +from collections.abc import Sequence +from typing import Any, Optional + +from sqlalchemy import Connection, Row, Update, bindparam, delete, insert, select, update + +from invokeai.app.services.shared.database.queries.base import QueryModule, mapped, read, write +from invokeai.app.services.shared.database.schema.wildcards import wildcards +from invokeai.app.services.wildcard_records.wildcard_records_common import ( + WildcardChanges, + WildcardRecordDTO, + WildcardWithoutId, +) + +_W = wildcards.c +_COLUMNS = (_W.id, _W.name, _W.values_json, _W.user_id) + +_GET = select(*_COLUMNS).where(_W.id == bindparam("wildcard_id")) +# Names are unique per account, so name order has no ties. +_OWNED = select(*_COLUMNS).where(_W.user_id == bindparam("user_id")).order_by(_W.name) +_INSERT = insert(wildcards) +_DELETE = delete(wildcards).where(_W.id == bindparam("wildcard_id")) + + +@functools.cache +def _update(fields: tuple[str, ...]) -> Update: + """Sets these fields of the wildcard; one statement per set of fields, of which there are few.""" + return ( + update(wildcards) + .where(_W.id == bindparam("target_wildcard_id")) + .values({field: bindparam(f"new_{field}") for field in fields}) + ) + + +def _wildcard(row: Sequence[Any]) -> WildcardRecordDTO: + wildcard_id, name, values_json, user_id = row + return WildcardRecordDTO.from_dict({"id": wildcard_id, "name": name, "values": values_json, "user_id": user_id}) + + +def _wildcard_or_none(row: Optional[Sequence[Any]]) -> Optional[WildcardRecordDTO]: + return _wildcard(row) if row is not None else None + + +def _wildcards(rows: Sequence[Sequence[Any]]) -> list[WildcardRecordDTO]: + return [_wildcard(row) for row in rows] + + +class WildcardQueries(QueryModule): + @mapped(_wildcard_or_none) + @read + def get(self, conn: Connection, wildcard_id: str) -> Optional[Row[Any]]: + return conn.execute(_GET, {"wildcard_id": wildcard_id}).first() + + @mapped(_wildcards) + @read + def owned_by(self, conn: Connection, user_id: str) -> Sequence[Row[Any]]: + """The account's wildcards, in name order.""" + return conn.execute(_OWNED, {"user_id": user_id}).all() + + @write + def insert(self, conn: Connection, wildcard_id: str, wildcard: WildcardWithoutId, user_id: str) -> None: + conn.execute( + _INSERT, + {"id": wildcard_id, "name": wildcard.name, "values_json": json.dumps(wildcard.values), "user_id": user_id}, + ) + + @write + def update(self, conn: Connection, wildcard_id: str, changes: WildcardChanges) -> None: + """Applies the changes to the wildcard's name and values.""" + fields: dict[str, Any] = {} + if changes.name is not None: + fields["name"] = changes.name + if changes.values is not None: + fields["values_json"] = json.dumps(changes.values) + if fields: + values = {f"new_{field}": value for field, value in fields.items()} + conn.execute(_update(tuple(sorted(fields))), {"target_wildcard_id": wildcard_id, **values}) + + @write + def delete(self, conn: Connection, wildcard_id: str) -> None: + conn.execute(_DELETE, {"wildcard_id": wildcard_id}) diff --git a/invokeai/app/services/shared/database/queries/workflows.py b/invokeai/app/services/shared/database/queries/workflows.py new file mode 100644 index 00000000000..737870a3b25 --- /dev/null +++ b/invokeai/app/services/shared/database/queries/workflows.py @@ -0,0 +1,415 @@ +"""The workflow library: the accounts' own workflows (category `user`) and the bundled ones (`default`).""" + +import functools +import itertools +from collections.abc import Collection, Sequence +from typing import Any, Literal, NamedTuple, Optional, Unpack + +from sqlalchemy import ( + BindParameter, + ColumnElement, + Connection, + Row, + Select, + Update, + bindparam, + case, + delete, + false, + func, + insert, + literal, + or_, + select, + true, + update, +) + +from invokeai.app.services.shared.database.dialect import CaseInsensitiveLike, like_contains +from invokeai.app.services.shared.database.queries.base import IN_CHUNK, QueryModule, locking, mapped, read, write +from invokeai.app.services.shared.database.schema.workflows import workflow_library +from invokeai.app.services.shared.database.types import now_text +from invokeai.app.services.shared.pagination import SQLiteDirection +from invokeai.app.services.workflow_records.workflow_records_common import ( + WorkflowCategory, + WorkflowRecordDTO, + WorkflowRecordDTOBase, + WorkflowRecordListItemDTO, + WorkflowRecordListItemDTOValidator, + WorkflowRecordOrderBy, +) + + +class LockedWorkflow(NamedTuple): + """A workflow's row as a transaction locked it: what a write of the workflow decides on.""" + + category: str + user_id: Optional[str] + revision: int + # The document, when the lock was taken to rewrite it. + document: Optional[str] = None + + +class _Shape(NamedTuple): + """Which conditions a listing or a count has; their values are bound.""" + + categories: tuple[WorkflowCategory, ...] + # How many tags a workflow is looked for under, any of them. + tags: int + has_been_opened: Optional[bool] + # Whether a text is looked for in the name, description and tags. + searching: bool + # Whether it covers only the account's own workflows and the bundled ones. + scoped: bool + is_public: Optional[bool] + + +# A client chooses how many tags it sends, and every number would be a statement of its own that SQLAlchemy keeps +# compiled. Up to this many, a statement has the next power of two of tag slots, the ones left over bound to NULL, +# which matches nothing; a listing for more tags builds and compiles its statement for the call alone, and counts take +# the tags this many at a time. +_TAG_SLOTS = 32 + + +def _slots(count: int) -> int: + """The tag slots of a statement for `count` tags: the next power of two up to `_TAG_SLOTS`, beyond it `count`.""" + return count if count == 0 or count > _TAG_SLOTS else 1 << (count - 1).bit_length() + + +def _flag(value: Optional[bool]) -> Optional[bool]: + # Shapes are cache keys, and 1 == True: normalized, whatever a caller passes for a flag picks the right statement. + return None if value is None else bool(value) + + +def _where( + *, + categories: Optional[Sequence[WorkflowCategory]] = None, + tags: Optional[Sequence[str]] = None, + has_been_opened: Optional[bool] = None, + query: Optional[str] = None, + user_id: Optional[str] = None, + is_public: Optional[bool] = None, +) -> tuple[_Shape, dict[str, Any]]: + """The conditions' shape, and their values: tags and the text looked for match as they are, ignoring case.""" + searched = query.strip() if query else "" + tags = list(tags or ()) + slots = _slots(len(tags)) + shape = _Shape( + categories=tuple(sorted({WorkflowCategory(category) for category in categories or ()})), + tags=slots, + has_been_opened=_flag(has_been_opened), + searching=bool(searched), + scoped=user_id is not None, + is_public=_flag(is_public), + ) + parameters: dict[str, Any] = { + f"tag_{i}": like_contains(tags[i].strip()) if i < len(tags) else None for i in range(slots) + } + if searched: + parameters["query"] = like_contains(searched) + if user_id is not None: + parameters["user_id"] = user_id + return shape, parameters + + +_W = workflow_library.c +_SUMMARY_COLUMNS = ( + _W.workflow_id, + _W.name, + _W.created_at, + _W.updated_at, + _W.opened_at, + _W.last_run_at, + _W.user_id, + _W.is_public, + _W.revision, +) +_LIST_COLUMNS = (*_SUMMARY_COLUMNS, _W.description, _W.category, _W.tags) +_SUMMARY_NAMES = tuple(column.name for column in _SUMMARY_COLUMNS) +_LIST_NAMES = tuple(column.name for column in _LIST_COLUMNS) +_ORDER_COLUMNS = { + WorkflowRecordOrderBy.CreatedAt: _W.created_at, + WorkflowRecordOrderBy.UpdatedAt: _W.updated_at, + WorkflowRecordOrderBy.OpenedAt: _W.opened_at, + WorkflowRecordOrderBy.Name: _W.name, + WorkflowRecordOrderBy.IsPublic: _W.is_public, +} + +_THE_WORKFLOW = _W.workflow_id == bindparam("workflow_id") +_GET = select(*_SUMMARY_COLUMNS, _W.workflow).where(_THE_WORKFLOW) +_SUMMARY = select(*_SUMMARY_COLUMNS).where(_THE_WORKFLOW) +_DOCUMENTS = select(_W.workflow_id, _W.workflow).where(_W.workflow_id.in_(bindparam("workflow_ids", expanding=True))) +_LOCK = select(_W.category, _W.user_id, _W.revision).where(_THE_WORKFLOW).with_for_update() +_LOCK_WITH_DOCUMENT = select(_W.category, _W.user_id, _W.revision, _W.workflow).where(_THE_WORKFLOW).with_for_update() +_BUNDLED = select(_W.workflow_id, _W.name).where(_W.category == literal(WorkflowCategory.Default.value)) +_INSERT = insert(workflow_library) +_SAVE = ( + update(workflow_library) + .where(_W.workflow_id == bindparam("target_workflow_id")) + .values(workflow=bindparam("new_workflow"), revision=_W.revision + 1) +) +_SET_PUBLIC = ( + update(workflow_library) + .where(_W.workflow_id == bindparam("target_workflow_id")) + .values(workflow=bindparam("new_workflow"), is_public=bindparam("new_is_public")) +) +_DELETE = delete(workflow_library).where(_THE_WORKFLOW) + + +def _conditions(shape: _Shape) -> list[ColumnElement[bool]]: + """The conditions of a shape; the statements built from them are cached, not they.""" + conditions: list[ColumnElement[bool]] = [] + if shape.categories: + conditions.append(_W.category.in_([literal(category.value) for category in shape.categories])) + if shape.tags: + conditions.append(or_(*(CaseInsensitiveLike(_W.tags, bindparam(f"tag_{i}")) for i in range(shape.tags)))) + if shape.has_been_opened is not None: + conditions.append(_W.opened_at.is_not(None) if shape.has_been_opened else _W.opened_at.is_(None)) + if shape.searching: + query: BindParameter[str] = bindparam("query") + conditions.append( + or_( + CaseInsensitiveLike(_W.name, query), + CaseInsensitiveLike(_W.description, query), + CaseInsensitiveLike(_W.tags, query), + ) + ) + if shape.scoped: + # The account's own workflows, and the bundled ones every account sees. + conditions.append( + or_(_W.user_id == bindparam("user_id"), _W.category == literal(WorkflowCategory.Default.value)) + ) + if shape.is_public is not None: + conditions.append(_W.is_public == (true() if shape.is_public else false())) + return conditions + + +@functools.lru_cache(maxsize=256) +def _list(shape: _Shape, order_by: WorkflowRecordOrderBy, direction: SQLiteDirection, paged: bool) -> Select[Any]: + # The enums are also strings, which share their cache entries: compare by value, never by identity. The workflow + # id breaks ties, so that equal keys keep one order from page to page. + ordering = (_ORDER_COLUMNS[order_by], _W.workflow_id) + statement = ( + select(*_LIST_COLUMNS) + .where(*_conditions(shape)) + .order_by(*(column.desc() if direction == SQLiteDirection.Descending else column.asc() for column in ordering)) + ) + if paged: + statement = statement.limit(bindparam("limit")).offset(bindparam("offset")) + return statement + + +@functools.lru_cache(maxsize=256) +def _count(shape: _Shape) -> Select[Any]: + return select(func.count()).select_from(workflow_library).where(*_conditions(shape)) + + +@functools.lru_cache(maxsize=256) +def _tag_counts(slots: int, shape: _Shape) -> Select[Unpack[tuple[Any, ...]]]: + """How many of the workflows `shape` covers have each tag, in one pass.""" + counts: list[ColumnElement[int]] = [ + func.count(case((CaseInsensitiveLike(_W.tags, bindparam(f"count_{i}")), 1))) for i in range(slots) + ] + return select(*counts).where(*_conditions(shape)) + + +@functools.lru_cache(maxsize=64) +def _category_count(shape: _Shape) -> Select[Any]: + # One category per count, not one count grouped by category: the generated columns are virtual on SQLite, and + # grouping reads each workflow's category from its document where a search by category reads it from the index. + return ( + select(func.count()) + .select_from(workflow_library) + .where(*_conditions(shape), _W.category == bindparam("category")) + ) + + +@functools.lru_cache(maxsize=64) +def _tag_lists(shape: _Shape) -> Select[Any]: + return select(_W.tags).distinct().where(_W.tags.is_not(None), _W.tags != literal(""), *_conditions(shape)) + + +@functools.cache +def _mark(column: str, scoped: bool) -> Update: + """Sets a timestamp column of the workflow; with `scoped`, only when the account owns it.""" + statement = update(workflow_library).where(_W.workflow_id == bindparam("target_workflow_id")) + if scoped: + statement = statement.where(_W.user_id == bindparam("owner_id")) + return statement.values({column: bindparam("now")}) + + +def _record_or_none(row: Optional[Sequence[Any]]) -> Optional[WorkflowRecordDTO]: + if row is None: + return None + *summary, workflow = row + return WorkflowRecordDTO.from_dict({**dict(zip(_SUMMARY_NAMES, summary, strict=False)), "workflow": workflow}) + + +def _summary_or_none(row: Optional[Sequence[Any]]) -> Optional[WorkflowRecordDTOBase]: + return WorkflowRecordDTOBase(**dict(zip(_SUMMARY_NAMES, row, strict=False))) if row is not None else None + + +def _items_and_total(result: tuple[Sequence[Sequence[Any]], int]) -> tuple[list[WorkflowRecordListItemDTO], int]: + rows, total = result + items = [ + WorkflowRecordListItemDTOValidator.validate_python(dict(zip(_LIST_NAMES, row, strict=False))) for row in rows + ] + return items, total + + +class WorkflowQueries(QueryModule): + @mapped(_record_or_none) + @read + def get(self, conn: Connection, workflow_id: str) -> Optional[Row[Any]]: + """The workflow with its document.""" + return conn.execute(_GET, {"workflow_id": workflow_id}).first() + + @mapped(_summary_or_none) + @read + def summary(self, conn: Connection, workflow_id: str) -> Optional[Row[Any]]: + """The workflow without its document.""" + return conn.execute(_SUMMARY, {"workflow_id": workflow_id}).first() + + @read + def documents(self, conn: Connection, workflow_ids: Collection[str]) -> dict[str, str]: + """The documents of those of these workflows that exist.""" + found: dict[str, str] = {} + for chunk in itertools.batched(workflow_ids, IN_CHUNK): + found.update((row[0], row[1]) for row in conn.execute(_DOCUMENTS, {"workflow_ids": list(chunk)}).all()) + return found + + @read + def bundled(self, conn: Connection) -> list[tuple[str, str]]: + """The id and name of every bundled workflow stored.""" + rows: Sequence[Sequence[Any]] = conn.execute(_BUNDLED).all() + return [(workflow_id, name) for workflow_id, name in rows] + + @locking + def lock(self, conn: Connection, workflow_id: str, *, with_document: bool = False) -> Optional[LockedWorkflow]: + """Locks the workflow's row until the transaction ends, or returns None when there is no such workflow.""" + statement = _LOCK_WITH_DOCUMENT if with_document else _LOCK + row: Optional[Sequence[Any]] = conn.execute(statement, {"workflow_id": workflow_id}).first() + return LockedWorkflow(*row) if row is not None else None + + @mapped(_items_and_total) + @read + def page( + self, + conn: Connection, + *, + order_by: WorkflowRecordOrderBy, + direction: SQLiteDirection, + offset: int, + limit: Optional[int], + categories: Optional[Sequence[WorkflowCategory]], + tags: Optional[Sequence[str]], + has_been_opened: Optional[bool], + query: Optional[str], + user_id: Optional[str], + is_public: Optional[bool], + ) -> tuple[Sequence[Row[Any]], int]: + """The workflows the filters keep, in order, `limit` of them from `offset` (all of them with None), and how + many the filters keep in all. `user_id` keeps the account's own workflows and the bundled ones.""" + shape, parameters = _where( + categories=categories, + tags=tags, + has_been_opened=has_been_opened, + query=query, + user_id=user_id, + is_public=is_public, + ) + if shape.tags <= _TAG_SLOTS: + listing, count, options = _list(shape, order_by, direction, limit is not None), _count(shape), None + else: + listing = _list.__wrapped__(shape, order_by, direction, limit is not None) + count, options = _count.__wrapped__(shape), {"compiled_cache": None} + rows = conn.execute(listing, {**parameters, "limit": limit, "offset": offset}, execution_options=options).all() + total: int = conn.execute(count, parameters, execution_options=options).scalar_one() + return rows, total + + @read + def tag_counts( + self, + conn: Connection, + tags: Sequence[str], + *, + categories: Optional[Sequence[WorkflowCategory]], + has_been_opened: Optional[bool], + user_id: Optional[str], + is_public: Optional[bool], + ) -> list[int]: + """How many of the workflows the filters keep have each of `tags`.""" + shape, parameters = _where( + categories=categories, has_been_opened=has_been_opened, user_id=user_id, is_public=is_public + ) + counts: list[int] = [] + for chunk in itertools.batched(tags, _TAG_SLOTS): + slots = _slots(len(chunk)) + patterns = {f"count_{i}": like_contains(chunk[i].strip()) if i < len(chunk) else None for i in range(slots)} + row: Sequence[Any] = conn.execute(_tag_counts(slots, shape), {**parameters, **patterns}).one() + counts.extend(row[: len(chunk)]) + return counts + + @read + def category_counts( + self, + conn: Connection, + categories: Sequence[WorkflowCategory], + *, + has_been_opened: Optional[bool], + user_id: Optional[str], + is_public: Optional[bool], + ) -> dict[str, int]: + """How many of the workflows the filters keep are in each of `categories`.""" + shape, parameters = _where( + categories=categories, has_been_opened=has_been_opened, user_id=user_id, is_public=is_public + ) + statement = _category_count(shape) + return { + category.value: conn.execute(statement, {**parameters, "category": category.value}).scalar_one() + for category in dict.fromkeys(WorkflowCategory(category) for category in categories) + } + + @read + def tag_lists( + self, + conn: Connection, + *, + categories: Optional[Sequence[WorkflowCategory]], + user_id: Optional[str], + is_public: Optional[bool], + ) -> list[str]: + """The distinct comma-separated tags of the workflows the filters keep, of those that have tags.""" + shape, parameters = _where(categories=categories, user_id=user_id, is_public=is_public) + return list(conn.execute(_tag_lists(shape), parameters).scalars().all()) + + @write + def insert(self, conn: Connection, *, workflow_id: str, workflow: str, user_id: str, is_public: bool) -> None: + """Adds a workflow at revision 1.""" + conn.execute( + _INSERT, {"workflow_id": workflow_id, "workflow": workflow, "user_id": user_id, "is_public": is_public} + ) + + @write + def save(self, conn: Connection, workflow_id: str, workflow: str) -> None: + """Writes the workflow's document as its next revision.""" + conn.execute(_SAVE, {"target_workflow_id": workflow_id, "new_workflow": workflow}) + + @write + def set_public(self, conn: Connection, workflow_id: str, workflow: str, is_public: bool) -> None: + """Writes whether the workflow is shared, with its document carrying the matching tag; the revision stays.""" + conn.execute( + _SET_PUBLIC, {"target_workflow_id": workflow_id, "new_workflow": workflow, "new_is_public": is_public} + ) + + @write + def mark( + self, conn: Connection, column: Literal["opened_at", "last_run_at"], workflow_id: str, user_id: Optional[str] + ) -> None: + """Sets the workflow's `opened_at` or `last_run_at` to now; with `user_id`, only if the account owns it.""" + parameters = {"target_workflow_id": workflow_id, "now": now_text(), "owner_id": user_id} + conn.execute(_mark(column, user_id is not None), parameters) + + @write + def delete(self, conn: Connection, workflow_id: str) -> None: + conn.execute(_DELETE, {"workflow_id": workflow_id}) diff --git a/invokeai/app/services/shared/database/schema/__init__.py b/invokeai/app/services/shared/database/schema/__init__.py new file mode 100644 index 00000000000..66343487d8f --- /dev/null +++ b/invokeai/app/services/shared/database/schema/__init__.py @@ -0,0 +1,33 @@ +"""The application schema as SQLAlchemy metadata, one module per domain. + +It is the schema the SQLite migrations build (`tests/app/services/shared/database/test_schema_parity.py` holds +the two to each other), and the schema a server database is created with. A schema change is a migration and +the matching change here, together. +""" + +# Every domain module is imported, so that `metadata` holds every table. +from invokeai.app.services.shared.database.schema import ( # noqa: F401 + app_settings, + boards, + client_state, + fonts, + image_index, + image_moves, + images, + intermediates, + locks, + media_references, + migrator, + models, + projects, + session_queue, + style_presets, + system_prompts, + users, + videos, + wildcards, + workflows, +) +from invokeai.app.services.shared.database.schema.metadata import metadata + +__all__ = ["metadata"] diff --git a/invokeai/app/services/shared/database/schema/app_settings.py b/invokeai/app/services/shared/database/schema/app_settings.py new file mode 100644 index 00000000000..987eb05887d --- /dev/null +++ b/invokeai/app/services/shared/database/schema/app_settings.py @@ -0,0 +1,14 @@ +"""Server-wide settings the application keeps in the database.""" + +from sqlalchemy import Column + +from invokeai.app.services.shared.database.schema.metadata import inserted_at, table, updated_at +from invokeai.app.services.shared.database.types import Key, LongText + +app_settings = table( + "app_settings", + Column("key", Key(), primary_key=True), + Column("value", LongText(), nullable=False), + inserted_at(), + updated_at(), +) diff --git a/invokeai/app/services/shared/database/schema/boards.py b/invokeai/app/services/shared/database/schema/boards.py new file mode 100644 index 00000000000..21dd3f7ce1b --- /dev/null +++ b/invokeai/app/services/shared/database/schema/boards.py @@ -0,0 +1,71 @@ +"""Boards, who they are shared with, and the media on them.""" + +from sqlalchemy import Boolean, Column, ForeignKey, Index + +from invokeai.app.services.shared.database.schema.metadata import ( + ENUM_LENGTH, + USER_ID_LENGTH, + default, + inserted_at, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import Key, LongText, Timestamp + +boards = table( + "boards", + Column("board_id", Key(), primary_key=True), + Column("board_name", LongText(), nullable=False), + Column("cover_image_name", Key(), ForeignKey("images.image_name", ondelete="SET NULL")), + inserted_at(), + updated_at(), + Column("deleted_at", Timestamp()), # unused + Column("archived", Boolean(), server_default=default(False)), + Column("user_id", Key(USER_ID_LENGTH), server_default=default("system")), + Column("is_public", Boolean(), nullable=False, server_default=default(False)), + Column("board_visibility", Key(ENUM_LENGTH), nullable=False, server_default=default("private")), +) + +Index("idx_boards_board_visibility", boards.c.board_visibility) +Index("idx_boards_created_at", boards.c.created_at) +Index("idx_boards_is_public", boards.c.is_public) +Index("idx_boards_user_id", boards.c.user_id) + +shared_boards = table( + "shared_boards", + Column("board_id", Key(), ForeignKey("boards.board_id", ondelete="CASCADE"), primary_key=True), + Column("user_id", Key(USER_ID_LENGTH), ForeignKey("users.user_id", ondelete="CASCADE"), primary_key=True), + Column("can_edit", Boolean(), nullable=False, server_default=default(False)), + inserted_at("shared_at"), +) + +# The primary key covers it. Only SQLite has it, where a migration created it. +Index("idx_shared_boards_board_id", shared_boards.c.board_id).ddl_if(dialect="sqlite") +Index("idx_shared_boards_user_id", shared_boards.c.user_id) + +# An image is on at most one board: its name is the primary key. +board_images = table( + "board_images", + Column("board_id", Key(), ForeignKey("boards.board_id", ondelete="CASCADE"), nullable=False), + Column("image_name", Key(), ForeignKey("images.image_name", ondelete="CASCADE"), primary_key=True), + inserted_at(), + updated_at(), + Column("deleted_at", Timestamp()), # unused +) + +# `idx_board_images_board_id_created_at` covers it. Only SQLite has it, where a migration created it. +Index("idx_board_images_board_id", board_images.c.board_id).ddl_if(dialect="sqlite") +Index("idx_board_images_board_id_created_at", board_images.c.board_id, board_images.c.created_at) + +board_videos = table( + "board_videos", + Column("board_id", Key(), ForeignKey("boards.board_id", ondelete="CASCADE"), nullable=False), + Column("video_name", Key(), ForeignKey("videos.video_name", ondelete="CASCADE"), primary_key=True), + inserted_at(), + updated_at(), + Column("deleted_at", Timestamp()), # unused +) + +# `idx_board_videos_board_id_created_at` covers it. Only SQLite has it, where a migration created it. +Index("idx_board_videos_board_id", board_videos.c.board_id).ddl_if(dialect="sqlite") +Index("idx_board_videos_board_id_created_at", board_videos.c.board_id, board_videos.c.created_at) diff --git a/invokeai/app/services/shared/database/schema/client_state.py b/invokeai/app/services/shared/database/schema/client_state.py new file mode 100644 index 00000000000..55ee4d0770e --- /dev/null +++ b/invokeai/app/services/shared/database/schema/client_state.py @@ -0,0 +1,22 @@ +"""Per-user state of the web client.""" + +from sqlalchemy import Column, ForeignKey, Index + +from invokeai.app.services.shared.database.schema.metadata import ( + CURRENT_TIMESTAMP, + USER_ID_LENGTH, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import Key, LongText + +client_state = table( + "client_state", + Column("user_id", Key(USER_ID_LENGTH), ForeignKey("users.user_id", ondelete="CASCADE"), primary_key=True), + Column("key", Key(), primary_key=True), + Column("value", LongText(), nullable=False), + updated_at(sqlite_default=CURRENT_TIMESTAMP), +) + +# The primary key covers it. Only SQLite has it, where a migration created it. +Index("idx_client_state_user_id", client_state.c.user_id).ddl_if(dialect="sqlite") diff --git a/invokeai/app/services/shared/database/schema/fonts.py b/invokeai/app/services/shared/database/schema/fonts.py new file mode 100644 index 00000000000..cdaa4a06100 --- /dev/null +++ b/invokeai/app/services/shared/database/schema/fonts.py @@ -0,0 +1,71 @@ +"""The font catalog: fonts users uploaded, and fonts found in the configured font directory.""" + +from sqlalchemy import CheckConstraint, Column, ForeignKey, Index, text + +from invokeai.app.services.shared.database.schema.metadata import ( + ENUM_LENGTH, + PATH_LENGTH, + USER_ID_LENGTH, + default, + inserted_at, + table, +) +from invokeai.app.services.shared.database.types import BigInt, Key, LongText + +fonts = table( + "fonts", + Column("id", Key(), primary_key=True), + Column("owner_id", Key(USER_ID_LENGTH), ForeignKey("users.user_id", ondelete="CASCADE")), + Column("scope", Key(ENUM_LENGTH), nullable=False), + Column("source", Key(ENUM_LENGTH), nullable=False), + Column("filename", LongText(), nullable=False), + Column("storage_path", LongText()), + # Relative to the font directory. + Column("source_path", Key(PATH_LENGTH)), + # From the font file's name table. + Column("family", LongText(), nullable=False), + Column("label", LongText(), nullable=False), + Column("style", LongText(), nullable=False), + Column("weight", BigInt(), nullable=False), + # SHA-256, hex. + Column("content_hash", Key(64), nullable=False), + Column("byte_size", BigInt(), nullable=False), + Column("axes_json", LongText(), nullable=False, server_default=default("[]")), + Column("instances_json", LongText(), nullable=False, server_default=default("[]")), + inserted_at(), + # Set by the font service; no trigger ever maintained it. + inserted_at("updated_at"), + CheckConstraint("scope IN ('private', 'shared')", name="scope"), + CheckConstraint("source IN ('uploaded', 'directory')", name="source"), + CheckConstraint("weight >= 1 AND weight <= 1000", name="weight"), + CheckConstraint("length(content_hash) = 64", name="content_hash"), + CheckConstraint("byte_size > 0", name="byte_size"), + CheckConstraint( + """ + (source = 'directory' AND owner_id IS NULL AND scope = 'shared' AND storage_path IS NULL AND source_path IS NOT NULL) + OR + (source = 'uploaded' AND storage_path IS NOT NULL AND source_path IS NULL AND + ((scope = 'private' AND owner_id IS NOT NULL) OR (scope = 'shared' AND owner_id IS NULL))) + """, + name="source_fields", + ), +) + +# Partial on SQLite only. On a server it covers every row, which is the same constraint: by the CHECK above, only +# directory fonts have a source path, and NULLs never collide. +Index( + "idx_fonts_directory_source_path", + fonts.c.source_path, + unique=True, + sqlite_where=text("source = 'directory'"), +) +Index("idx_fonts_uploaded_hash", fonts.c.source, fonts.c.scope, fonts.c.owner_id, fonts.c.content_hash) +# For listing by family, case-insensitively. Only SQLite has it: a server can neither index a column under +# another collation nor sort by a prefix index, and `idx_fonts_uploaded_hash` serves the filter there. +Index( + "idx_fonts_visible_private", + fonts.c.source, + fonts.c.scope, + fonts.c.owner_id, + fonts.c.family.collate("NOCASE"), +).ddl_if(dialect="sqlite") diff --git a/invokeai/app/services/shared/database/schema/image_index.py b/invokeai/app/services/shared/database/schema/image_index.py new file mode 100644 index 00000000000..4127afbc2bf --- /dev/null +++ b/invokeai/app/services/shared/database/schema/image_index.py @@ -0,0 +1,66 @@ +"""The image index: embeddings of images and videos, cached map projections, and custom vocabulary.""" + +from sqlalchemy import CheckConstraint, Column, ForeignKey, Index + +from invokeai.app.services.shared.database.schema.metadata import ( + USER_ID_LENGTH, + default, + inserted_at, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import BigInt, Blob, Key, LongText, NoCaseKey + +image_embeddings = table( + "image_embeddings", + Column("image_name", Key(), ForeignKey("images.image_name", ondelete="CASCADE"), primary_key=True), + # The embedding model's content hash, not its install key. + Column("model_id", Key(), primary_key=True), + Column("dim", BigInt(), nullable=False), + # L2-normalized, in the little-endian float type `encoding` names: dim * 2 or dim * 4 bytes. + Column("embedding", Blob(), nullable=False), + inserted_at(), + Column("encoding", LongText(), nullable=False, server_default=default("float32")), + CheckConstraint("encoding IN ('float32', 'float16')", name="encoding"), +) + +Index("idx_image_embeddings_model_id", image_embeddings.c.model_id) + +video_embeddings = table( + "video_embeddings", + Column("video_name", Key(), ForeignKey("videos.video_name", ondelete="CASCADE"), primary_key=True), + # The embedding model's content hash, not its install key. + Column("model_id", Key(), primary_key=True), + Column("dim", BigInt(), nullable=False), + # L2-normalized, in the little-endian float type `encoding` names: dim * 2 or dim * 4 bytes. + Column("embedding", Blob(), nullable=False), + inserted_at(), + Column("encoding", LongText(), nullable=False, server_default=default("float32")), + CheckConstraint("encoding IN ('float32', 'float16')", name="encoding"), +) + +Index("idx_video_embeddings_model_id", video_embeddings.c.model_id) + +image_projections = table( + "image_projections", + Column("user_id", Key(USER_ID_LENGTH), ForeignKey("users.user_id", ondelete="CASCADE"), primary_key=True), + Column("model_id", Key(), primary_key=True), + # Fingerprint of the media the projection was computed over: a different one makes it stale. + Column("scope_hash", LongText(), nullable=False), + # The projection parameters (JSON). + Column("params", LongText(), nullable=False), + Column("point_count", BigInt(), nullable=False), + # Media names (JSON array), row-aligned with `coords`. + Column("image_names", LongText(), nullable=False), + # float32, shape (point_count, 2). + Column("coords", Blob(), nullable=False), + inserted_at(), + updated_at(), + Column("item_kinds", LongText()), +) + +image_index_vocab_terms = table( + "image_index_vocab_terms", + Column("term", NoCaseKey(), primary_key=True), + inserted_at(), +) diff --git a/invokeai/app/services/shared/database/schema/image_moves.py b/invokeai/app/services/shared/database/schema/image_moves.py new file mode 100644 index 00000000000..83547c387ad --- /dev/null +++ b/invokeai/app/services/shared/database/schema/image_moves.py @@ -0,0 +1,45 @@ +"""Jobs that move images between subfolders, and what each did to each image.""" + +from sqlalchemy import Boolean, CheckConstraint, Column, ForeignKey, Index + +from invokeai.app.services.shared.database.schema.metadata import ( + ENUM_LENGTH, + default, + inserted_at, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import BigInt, Key, LongText + +image_subfolder_move_jobs = table( + "image_subfolder_move_jobs", + Column("id", BigInt(), primary_key=True, nullable=True), + Column("state", LongText(), nullable=False), + inserted_at(), + updated_at(), + Column("error_message", LongText()), + CheckConstraint("state IN ('planned', 'moving', 'moved', 'committed', 'error')", name="state"), +) + +image_subfolder_move_items = table( + "image_subfolder_move_items", + # No ON DELETE: the items are the audit trail of a job's moves, never deleted on their own. + Column("job_id", BigInt(), ForeignKey("image_subfolder_move_jobs.id"), primary_key=True), + # Deleting an image takes its move history with it. + Column("image_name", Key(), ForeignKey("images.image_name", ondelete="CASCADE"), primary_key=True), + Column("old_subfolder", LongText(), nullable=False), + Column("new_subfolder", LongText(), nullable=False), + Column("is_intermediate", Boolean(), nullable=False, server_default=default(False)), + Column("old_path", LongText()), + Column("new_path", LongText()), + Column("old_thumbnail_path", LongText()), + Column("new_thumbnail_path", LongText()), + Column("state", Key(ENUM_LENGTH), nullable=False), + Column("error_message", LongText()), + CheckConstraint("state IN ('planned', 'moved', 'committed', 'error')", name="state"), +) + +Index("idx_image_subfolder_move_items_image_name", image_subfolder_move_items.c.image_name) +Index( + "idx_image_subfolder_move_items_job_state", image_subfolder_move_items.c.job_id, image_subfolder_move_items.c.state +) diff --git a/invokeai/app/services/shared/database/schema/images.py b/invokeai/app/services/shared/database/schema/images.py new file mode 100644 index 00000000000..fc86ce40b9b --- /dev/null +++ b/invokeai/app/services/shared/database/schema/images.py @@ -0,0 +1,99 @@ +"""Image records; the files themselves are on disk.""" + +from sqlalchemy import Boolean, Column, Index, text + +from invokeai.app.services.shared.database.schema.metadata import ( + ENUM_LENGTH, + ON_SERVERS, + USER_ID_LENGTH, + default, + inserted_at, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import BigInt, Key, LongText, Timestamp + +images = table( + "images", + Column("image_name", Key(), primary_key=True), + Column("image_origin", Key(ENUM_LENGTH), nullable=False), + Column("image_category", Key(ENUM_LENGTH), nullable=False), + Column("width", BigInt(), nullable=False), + Column("height", BigInt(), nullable=False), + Column("session_id", LongText()), + Column("node_id", LongText()), + Column("metadata", LongText()), + Column("is_intermediate", Boolean(), server_default=default(False)), + inserted_at(), + updated_at(), + Column("deleted_at", Timestamp()), # unused + Column("starred", Boolean(), server_default=default(False)), + Column("has_workflow", Boolean(), server_default=default(False)), + Column("user_id", Key(USER_ID_LENGTH), server_default=default("system")), + Column("image_subfolder", LongText(), nullable=False, server_default=default("")), + Column("project_id", Key()), + Column("file_size_bytes", BigInt()), +) + +Index("idx_images_created_at", images.c.created_at) +Index("idx_images_image_category", images.c.image_category) +# The primary key covers it. Only SQLite has it, where a migration created it. +Index("idx_images_image_name", images.c.image_name, unique=True).ddl_if(dialect="sqlite") +Index("idx_images_image_origin", images.c.image_origin) +# The intermediates of a scope, all of them, and those whose size is still unknown, each in creation order (the +# keyset of the cleanup's windows). Partial indexes on SQLite; a server has none, so its variants lead with the +# condition's columns instead, which narrows a query with the same condition to the same rows. +Index( + "idx_images_intermediate_scope", + images.c.user_id, + images.c.project_id, + images.c.created_at, + images.c.image_name, + sqlite_where=text("is_intermediate = TRUE"), +).ddl_if(dialect="sqlite") +Index( + "idx_images_intermediate_scope", + images.c.is_intermediate, + images.c.user_id, + images.c.project_id, + images.c.created_at, + images.c.image_name, +).ddl_if(dialect=ON_SERVERS) +Index( + "idx_images_intermediates_owner", + images.c.user_id, + images.c.created_at, + images.c.image_name, + sqlite_where=text("is_intermediate = TRUE"), +).ddl_if(dialect="sqlite") +Index( + "idx_images_intermediates_owner", + images.c.is_intermediate, + images.c.user_id, + images.c.created_at, + images.c.image_name, +).ddl_if(dialect=ON_SERVERS) +Index( + "idx_images_intermediates_created", + images.c.created_at, + images.c.image_name, + sqlite_where=text("is_intermediate = TRUE"), +).ddl_if(dialect="sqlite") +Index("idx_images_intermediates_created", images.c.is_intermediate, images.c.created_at, images.c.image_name).ddl_if( + dialect=ON_SERVERS +) +Index("idx_images_starred", images.c.starred) +Index( + "idx_images_unmeasured_intermediates", + images.c.created_at, + images.c.image_name, + sqlite_where=text("is_intermediate = TRUE AND file_size_bytes IS NULL"), +).ddl_if(dialect="sqlite") +Index( + "idx_images_unmeasured_intermediates", + images.c.is_intermediate, + images.c.file_size_bytes, + images.c.created_at, + images.c.image_name, +).ddl_if(dialect=ON_SERVERS) +Index("idx_images_user_id", images.c.user_id) diff --git a/invokeai/app/services/shared/database/schema/intermediates.py b/invokeai/app/services/shared/database/schema/intermediates.py new file mode 100644 index 00000000000..0bbf6c4e452 --- /dev/null +++ b/invokeai/app/services/shared/database/schema/intermediates.py @@ -0,0 +1,52 @@ +"""State of the intermediates cleanup: what keeps intermediate media from being cleaned up, and what cannot be measured. + +The browser holds outlive the process. The session holds and the unmeasurable marks are the process's own +state, kept in tables so that every connection sees them; the cleanup service empties them when it starts. +""" + +from sqlalchemy import Column, Index + +from invokeai.app.services.shared.database.schema.metadata import ENUM_LENGTH, USER_ID_LENGTH, table +from invokeai.app.services.shared.database.types import Key, Timestamp + +# Leases of browser tabs on intermediate media they still show. +intermediates_browser_holds = table( + "intermediates_browser_holds", + Column("user_id", Key(USER_ID_LENGTH), primary_key=True), + Column("lease_id", Key(), primary_key=True), + Column("media_kind", Key(ENUM_LENGTH), primary_key=True), + Column("media_name", Key(), primary_key=True), + Column("expires_at", Timestamp("TEXT"), nullable=False), +) + +Index("idx_intermediates_browser_holds_expires_at", intermediates_browser_holds.c.expires_at) +Index( + "idx_intermediates_browser_holds_media", + intermediates_browser_holds.c.media_kind, + intermediates_browser_holds.c.media_name, + intermediates_browser_holds.c.expires_at, +) + +# Cached outputs a running session reuses: held while the session is active, then for the recency grace from +# `released_at`, when the session was first seen finished. +intermediates_session_holds = table( + "intermediates_session_holds", + Column("session_id", Key(), primary_key=True), + Column("media_kind", Key(ENUM_LENGTH), primary_key=True), + Column("media_name", Key(), primary_key=True), + Column("released_at", Timestamp("TEXT")), +) + +Index( + "idx_intermediates_session_holds_media", + intermediates_session_holds.c.media_kind, + intermediates_session_holds.c.media_name, + intermediates_session_holds.c.released_at, +) + +# Intermediates whose file could not be measured, so that the measurement moves on to later ones. +intermediates_unmeasurable = table( + "intermediates_unmeasurable", + Column("media_kind", Key(ENUM_LENGTH), primary_key=True), + Column("media_name", Key(), primary_key=True), +) diff --git a/invokeai/app/services/shared/database/schema/locks.py b/invokeai/app/services/shared/database/schema/locks.py new file mode 100644 index 00000000000..b1b7d92f753 --- /dev/null +++ b/invokeai/app/services/shared/database/schema/locks.py @@ -0,0 +1,12 @@ +"""Named locks: one row per lock, which a transaction locks to serialise work with every other connection.""" + +from sqlalchemy import Column + +from invokeai.app.services.shared.database.schema.metadata import ENUM_LENGTH, table +from invokeai.app.services.shared.database.types import Key + +db_locks = table( + "db_locks", + # A `DatabaseLock`; migrations add the rows. + Column("name", Key(ENUM_LENGTH), primary_key=True), +) diff --git a/invokeai/app/services/shared/database/schema/media_references.py b/invokeai/app/services/shared/database/schema/media_references.py new file mode 100644 index 00000000000..47b8a76d840 --- /dev/null +++ b/invokeai/app/services/shared/database/schema/media_references.py @@ -0,0 +1,21 @@ +"""Which documents (projects, workflows, client state) reference which media, so that referenced media is kept.""" + +from sqlalchemy import Column, Index + +from invokeai.app.services.shared.database.schema.metadata import ENUM_LENGTH, USER_ID_LENGTH, table +from invokeai.app.services.shared.database.types import Key + +media_references = table( + "media_references", + # A MediaReferenceOwnerKind. + Column("owner_kind", Key(ENUM_LENGTH), primary_key=True), + # The owning account: project ids are unique per user, not globally. + Column("user_id", Key(USER_ID_LENGTH), primary_key=True), + # A project id, a workflow id or a client state key. + Column("owner_id", Key(), primary_key=True), + # 'image' or 'video'. + Column("media_kind", Key(ENUM_LENGTH), primary_key=True), + Column("media_name", Key(), primary_key=True), +) + +Index("idx_media_references_media", media_references.c.media_kind, media_references.c.media_name) diff --git a/invokeai/app/services/shared/database/schema/metadata.py b/invokeai/app/services/shared/database/schema/metadata.py new file mode 100644 index 00000000000..9d250292b1a --- /dev/null +++ b/invokeai/app/services/shared/database/schema/metadata.py @@ -0,0 +1,161 @@ +"""The metadata every table belongs to, and the rules its DDL follows on each backend. + +Which type a text column gets: + +- `Key` where a server must index the whole value: a primary key, a unique constraint, a foreign key, and + indexes on identifiers and enumerated values. A server bounds it; an index there holds 768 characters in all. +- `LongText` everywhere else, unbounded. An index on free text (a name, a description) covers a prefix of it on a + server (`LONG_TEXT_INDEX_PREFIX`), so such text has no length limit there either. + +Values longer than their `Key` are refused on a server, where SQLite stores them; the lengths leave room above +what the application itself generates. + +A primary key column that a migration declared without NOT NULL is `nullable=True`, so that SQLite declares it +alike; a server makes every primary key column NOT NULL regardless. + +An index that differs between backends is a pair of `Index` objects with the same name, each limited to its +backends with `ddl_if`. Only `create_all()` and `Index.create()` honor that (`Table.to_metadata()` drops it), so +tables are created from this metadata, never from copies of its tables. +""" + +from typing import Any + +from sqlalchemy import CheckConstraint, Column, MetaData, Table, TextClause, text +from sqlalchemy.ext.compiler import compiles +from sqlalchemy.schema import CreateColumn, SchemaItem +from sqlalchemy.sql.compiler import DDLCompiler + +from invokeai.app.services.shared.database.engines import ( + MARIADB_BINARY_COLLATION, + MYSQL_BINARY_COLLATION, + SERVER_DIALECTS, +) +from invokeai.app.services.shared.database.types import Timestamp, now_text + +metadata = MetaData( + naming_convention={ + "pk": "pk_%(table_name)s", + "uq": "uq_%(table_name)s_%(column_0_N_name)s", + "ck": "ck_%(table_name)s_%(constraint_name)s", + "fk": "fk_%(table_name)s_%(column_0_N_name)s_%(referred_table_name)s", + } +) + +# Every table on a server is InnoDB and compares text byte for byte, whatever the server's defaults. Keys of up to +# 3072 bytes need the DYNAMIC row format, the default a server may have been configured away from. +_SERVER_TABLE_OPTIONS: dict[str, Any] = { + "mysql_engine": "InnoDB", + "mysql_charset": "utf8mb4", + "mysql_collate": MYSQL_BINARY_COLLATION, + "mysql_row_format": "DYNAMIC", + "mariadb_engine": "InnoDB", + "mariadb_charset": "utf8mb4", + "mariadb_collate": MARIADB_BINARY_COLLATION, + "mariadb_row_format": "DYNAMIC", +} + +# `ddl_if(dialect=ON_SERVERS)`: only on the server backends. (`ddl_if` takes a tuple of dialect names, as +# documented, though its annotation names one.) +ON_SERVERS: Any = SERVER_DIALECTS + +# Lengths of `Key` text other than the default 255. User ids are generated (a uuid, or 'system'); kept short so +# that keys with several text columns fit a server's index. +USER_ID_LENGTH = 64 +ENUM_LENGTH = 32 # one of a fixed set of values: a status, a kind, a scope +PATH_LENGTH = 768 # a file path under a unique index, which a server can hold only whole + +# An index on a LONGTEXT column covers this many leading characters on a server. +LONG_TEXT_PREFIX_LENGTH = 255 +LONG_TEXT_INDEX_PREFIX: dict[str, Any] = { + "mysql_length": LONG_TEXT_PREFIX_LENGTH, + "mariadb_length": LONG_TEXT_PREFIX_LENGTH, +} + +# The column defaults of the SQLite schema's timestamps. Servers have none: the application sets timestamps. +STRFTIME_NOW = "STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')" +CURRENT_TIMESTAMP = "CURRENT_TIMESTAMP" + +# `Column.info` keys: a default rendered on SQLite only, and a generated column whose NOT NULL MariaDB gets as a +# CHECK constraint instead. +_SQLITE_DEFAULT = "invokeai_sqlite_default" +_MARIADB_NOT_NULL_CHECK = "invokeai_mariadb_not_null_check" + + +def default(value: str | int | float | bool) -> TextClause: + """A constant column default. + + A string is parenthesized, as an expression default, because a server accepts a literal default for + LONGTEXT only in that form. SQLite reports the default as the bare literal either way. + """ + if isinstance(value, bool): + return text("TRUE" if value else "FALSE") + if isinstance(value, str): + return text("('" + value.replace("'", "''") + "')") + return text(repr(value)) + + +def table(name: str, *items: SchemaItem, **options: Any) -> Table: + """A table of the application schema.""" + return define_table(metadata, name, *items, **options) + + +def define_table(target: MetaData, name: str, *items: SchemaItem, **options: Any) -> Table: + """A table in `target` that follows the schema's rules: the server table options, and on MariaDB a CHECK + for each generated column that may not be NULL. Migrations create tables with it too (see + `PortableMigrationContext.create_table`), so that a migrated server database matches a created one. + """ + created = Table(name, target, *items, **{**_SERVER_TABLE_OPTIONS, **options}) + for column in created.columns: + if column.computed is not None and not column.nullable: + # MariaDB rejects NOT NULL on a generated column (see `_create_mariadb_column`). + column.info[_MARIADB_NOT_NULL_CHECK] = True + created.append_constraint( + CheckConstraint(column.is_not(None), name=f"{column.name}_not_null").ddl_if(dialect="mariadb") + ) + return created + + +def inserted_at( + name: str = "created_at", *, declared: str = "DATETIME", sqlite_default: str = STRFTIME_NOW +) -> Column[str]: + """A timestamp the application sets when it inserts the row.""" + return Column(name, Timestamp(declared), nullable=False, default=now_text, info={_SQLITE_DEFAULT: sqlite_default}) + + +def updated_at(name: str = "updated_at", *, sqlite_default: str = STRFTIME_NOW) -> Column[str]: + """A timestamp the application sets when it inserts the row, and when an `update()` statement changes it. + + An upsert does not set it: SQLAlchemy leaves `onupdate` out of `ON CONFLICT DO UPDATE` and `ON DUPLICATE KEY + UPDATE`, so an upsert names it in its update values. + """ + return Column( + name, + Timestamp(), + nullable=False, + default=now_text, + onupdate=now_text, + info={_SQLITE_DEFAULT: sqlite_default}, + ) + + +@compiles(CreateColumn, "sqlite") +def _create_sqlite_column(create: CreateColumn, compiler: DDLCompiler, **kw: Any) -> str: + spec: str = compiler.visit_create_column(create, **kw) # type: ignore[no-untyped-call] + sqlite_default = create.element.info.get(_SQLITE_DEFAULT) + return spec if sqlite_default is None else f"{spec} DEFAULT ({sqlite_default})" + + +@compiles(CreateColumn, "mariadb") +def _create_mariadb_column(create: CreateColumn, compiler: DDLCompiler, **kw: Any) -> str: + column = create.element + if column.computed is None or not column.info.get(_MARIADB_NOT_NULL_CHECK): + spec: str = compiler.visit_create_column(create, **kw) # type: ignore[no-untyped-call] + return spec + # MariaDB rejects NOT NULL on a generated column, so it is left out here and `table()` adds a CHECK instead. + return " ".join( + ( + compiler.preparer.format_column(column), + compiler.dialect.type_compiler_instance.process(column.type, type_expression=column), + compiler.process(column.computed), + ) + ) diff --git a/invokeai/app/services/shared/database/schema/migrator.py b/invokeai/app/services/shared/database/schema/migrator.py new file mode 100644 index 00000000000..944a14302e9 --- /dev/null +++ b/invokeai/app/services/shared/database/schema/migrator.py @@ -0,0 +1,21 @@ +"""The migrator's records of which migrations a database has had.""" + +from sqlalchemy import Column + +from invokeai.app.services.shared.database.schema.metadata import inserted_at, table +from invokeai.app.services.shared.database.types import BigInt, Key + +# The schema version of the numbered migrations, the last of which, 34, ended them. +migrations = table( + "migrations", + Column("version", BigInt(), primary_key=True, nullable=True, autoincrement=False), + inserted_at("migrated_at"), +) + +applied_migrations = table( + "applied_migrations", + Column("migration_id", Key(), primary_key=True, nullable=True), + # The numbered version of a migration from before ids, for databases those versions migrated. + Column("legacy_version", BigInt(), unique=True), + inserted_at("migrated_at"), +) diff --git a/invokeai/app/services/shared/database/schema/models.py b/invokeai/app/services/shared/database/schema/models.py new file mode 100644 index 00000000000..891cd2f29ab --- /dev/null +++ b/invokeai/app/services/shared/database/schema/models.py @@ -0,0 +1,61 @@ +"""Installed models, and which of them belong together.""" + +from sqlalchemy import Column, Computed, ForeignKey, Index, UniqueConstraint + +from invokeai.app.services.shared.database.dialect import JsonValue +from invokeai.app.services.shared.database.schema.metadata import ( + LONG_TEXT_INDEX_PREFIX, + PATH_LENGTH, + inserted_at, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import BigInt, Key, LongText + + +def _config_member(name: str, column_type: Key | LongText | BigInt, *, nullable: bool = False) -> Column[object]: + """A generated column holding the member of `config` with the same name.""" + return Column( + name, + column_type, + Computed(JsonValue("config", f"$.{name}", integer=isinstance(column_type, BigInt))), + nullable=nullable, + ) + + +models = table( + "models", + Column("id", Key(), primary_key=True), + _config_member("hash", LongText()), + _config_member("base", LongText()), + _config_member("type", LongText()), + # Unique, so a server indexes it whole, which bounds its length there. + _config_member("path", Key(PATH_LENGTH)), + _config_member("format", LongText()), + _config_member("name", LongText()), + _config_member("description", LongText(), nullable=True), + _config_member("source", LongText()), + _config_member("source_type", LongText()), + _config_member("source_api_response", LongText(), nullable=True), + _config_member("trigger_phrases", LongText(), nullable=True), + _config_member("file_size", BigInt()), + # The whole config (JSON), with the members of its model type's subclass. + Column("config", LongText(), nullable=False), + inserted_at(), + updated_at(), + UniqueConstraint("path"), +) + +Index("base_index", models.c.base, **LONG_TEXT_INDEX_PREFIX) +Index("name_index", models.c.name, **LONG_TEXT_INDEX_PREFIX) +Index("type_index", models.c.type, **LONG_TEXT_INDEX_PREFIX) + +model_relationships = table( + "model_relationships", + # Keys of related models: model_key_1 < model_key_2, so that each pair is stored once. + Column("model_key_1", Key(), ForeignKey("models.id", ondelete="CASCADE"), primary_key=True), + Column("model_key_2", Key(), ForeignKey("models.id", ondelete="CASCADE"), primary_key=True), + inserted_at(declared="TEXT DATETIME"), +) + +Index("keyx_model_relationships_model_key_2", model_relationships.c.model_key_2) diff --git a/invokeai/app/services/shared/database/schema/projects.py b/invokeai/app/services/shared/database/schema/projects.py new file mode 100644 index 00000000000..06014242517 --- /dev/null +++ b/invokeai/app/services/shared/database/schema/projects.py @@ -0,0 +1,46 @@ +"""Canvas projects, and the documents of projects that lost their owner.""" + +from sqlalchemy import CheckConstraint, Column, ForeignKey, PrimaryKeyConstraint + +from invokeai.app.services.shared.database.schema.metadata import ( + USER_ID_LENGTH, + default, + inserted_at, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import BigInt, Key, LongText, Timestamp + +projects = table( + "projects", + # Generated by the client; unique per user, not globally. + Column("project_id", Key(), nullable=False), + Column("user_id", Key(USER_ID_LENGTH), ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False), + Column("name", LongText(), nullable=False), + # The client's project document (JSON). + Column("data", LongText(), nullable=False), + # The project's own board. Unique: at most one project per board. RESTRICT: a claimed board cannot be deleted. + Column("board_id", Key(), ForeignKey("boards.board_id", ondelete="RESTRICT"), nullable=False, unique=True), + # Incremented on every update, for optimistic concurrency. + Column("revision", BigInt(), nullable=False, server_default=default(1)), + inserted_at(), + updated_at(), + Column("minimum_canvas_schema_version", BigInt(), nullable=False, server_default=default(2)), + PrimaryKeyConstraint("user_id", "project_id"), + CheckConstraint("minimum_canvas_schema_version >= 1", name="minimum_canvas_schema_version"), +) + +# Where the project boards migration put projects whose owner no longer existed, whole, for repair. Free of +# foreign keys, so it can hold what `projects` cannot. Created by that migration only when it had such projects. +orphaned_projects_2026_08_06 = table( + "orphaned_projects_2026_08_06", + Column("project_id", Key(), nullable=False), + Column("user_id", Key(USER_ID_LENGTH), nullable=False), + Column("name", LongText(), nullable=False), + Column("data", LongText(), nullable=False), + Column("revision", BigInt(), nullable=False), + Column("created_at", Timestamp(), nullable=False), + Column("updated_at", Timestamp(), nullable=False), + inserted_at("quarantined_at"), + PrimaryKeyConstraint("user_id", "project_id"), +) diff --git a/invokeai/app/services/shared/database/schema/session_queue.py b/invokeai/app/services/shared/database/schema/session_queue.py new file mode 100644 index 00000000000..0f1a48a3944 --- /dev/null +++ b/invokeai/app/services/shared/database/schema/session_queue.py @@ -0,0 +1,109 @@ +"""The session queue, and the receipts that make enqueueing idempotent.""" + +from sqlalchemy import Column, ForeignKey, Index + +from invokeai.app.services.shared.database.schema.metadata import ( + ENUM_LENGTH, + USER_ID_LENGTH, + default, + inserted_at, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import BigInt, Key, LongText, Timestamp + +session_queue = table( + "session_queue", + # Never reused (AUTOINCREMENT on SQLite): it orders the queue and pages through it. + Column("item_id", BigInt(), primary_key=True, nullable=True), + Column("batch_id", Key(), nullable=False), + Column("queue_id", LongText(), nullable=False), + Column("session_id", Key(), nullable=False, unique=True), + Column("field_values", LongText()), + Column("session", LongText(), nullable=False), + Column("status", Key(ENUM_LENGTH), nullable=False, server_default=default("pending")), + # Higher runs first. + Column("priority", BigInt(), nullable=False, server_default=default(0)), + Column("error_traceback", LongText()), + inserted_at(), + updated_at(), + Column("started_at", Timestamp()), + Column("completed_at", Timestamp()), + Column("workflow", LongText()), + Column("error_type", LongText()), + Column("error_message", LongText()), + Column("origin", LongText()), + Column("destination", LongText()), + Column("retried_from_item_id", BigInt()), + Column("user_id", Key(USER_ID_LENGTH), server_default=default("system")), + Column("status_sequence", BigInt(), server_default=default(0)), + Column("device", LongText()), + Column("workflow_call_id", Key()), + Column("parent_item_id", BigInt()), + Column("parent_session_id", Key()), + Column("root_item_id", BigInt()), + Column("workflow_call_depth", BigInt()), + Column("project_id", LongText()), + # Counts the writes of `session`, so a writer can tell whether it still holds the latest one. + Column("session_revision", BigInt(), nullable=False, server_default=default(0)), + sqlite_autoincrement=True, +) + +Index("idx_session_queue_batch_id", session_queue.c.batch_id) +Index("idx_session_queue_created_priority", session_queue.c.priority) +# Covered by the longer indexes starting with the same column, as is `idx_session_queue_user_id`. Only SQLite +# has these, where migrations created them. +Index("idx_session_queue_created_status", session_queue.c.status).ddl_if(dialect="sqlite") +# The primary key and the unique constraint cover these. Only SQLite has them, where migrations created them. +Index("idx_session_queue_item_id", session_queue.c.item_id, unique=True).ddl_if(dialect="sqlite") +# A queue's listing and its status counts, answered from the index alone (2026_10_04_add_session_queue_listing_index). +# Only SQLite has it: on a server `queue_id` and `origin` are long texts, which no index holds whole. +Index( + "idx_session_queue_listing", + session_queue.c.queue_id, + session_queue.c.created_at, + session_queue.c.item_id.desc(), + session_queue.c.user_id, + session_queue.c.origin, + session_queue.c.status, +).ddl_if(dialect="sqlite") +Index("idx_session_queue_session_id", session_queue.c.session_id, unique=True).ddl_if(dialect="sqlite") +Index("idx_session_queue_parent_item_id", session_queue.c.parent_item_id) +Index("idx_session_queue_parent_session_id", session_queue.c.parent_session_id) +Index("idx_session_queue_root_item_id", session_queue.c.root_item_id) +# The round-robin dequeue: pending items per user, best first. +Index( + "idx_session_queue_round_robin_pending", + session_queue.c.status, + session_queue.c.user_id, + session_queue.c.priority.desc(), + session_queue.c.item_id.asc(), +) +Index("idx_session_queue_user_id", session_queue.c.user_id).ddl_if(dialect="sqlite") +Index("idx_session_queue_user_started_at", session_queue.c.user_id, session_queue.c.started_at) +Index("idx_session_queue_workflow_call_depth", session_queue.c.workflow_call_depth) +Index("idx_session_queue_workflow_call_id", session_queue.c.workflow_call_id) + +# What an enqueue request with an idempotency key did, so that sending it again does not enqueue it twice. +session_queue_enqueue_receipts = table( + "session_queue_enqueue_receipts", + Column("queue_id", Key(), primary_key=True), + Column("user_id", Key(USER_ID_LENGTH), ForeignKey("users.user_id", ondelete="CASCADE"), primary_key=True), + Column("idempotency_key", Key(), primary_key=True), + Column("payload_hash", LongText(), nullable=False), + Column("batch_id", LongText(), nullable=False), + Column("requested", BigInt(), nullable=False), + Column("enqueued", BigInt(), nullable=False), + Column("priority", BigInt(), nullable=False), + Column("item_ids", LongText(), nullable=False), + Column("byte_size", BigInt(), nullable=False), + inserted_at(), + Column("acknowledged_at", Timestamp()), +) + +Index( + "idx_session_queue_enqueue_receipts_owner_ack", + session_queue_enqueue_receipts.c.user_id, + session_queue_enqueue_receipts.c.acknowledged_at, + session_queue_enqueue_receipts.c.byte_size, +) diff --git a/invokeai/app/services/shared/database/schema/style_presets.py b/invokeai/app/services/shared/database/schema/style_presets.py new file mode 100644 index 00000000000..a451924a74d --- /dev/null +++ b/invokeai/app/services/shared/database/schema/style_presets.py @@ -0,0 +1,29 @@ +"""Style presets: saved prompt templates.""" + +from sqlalchemy import Boolean, Column, Index + +from invokeai.app.services.shared.database.schema.metadata import ( + LONG_TEXT_INDEX_PREFIX, + USER_ID_LENGTH, + default, + inserted_at, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import Key, LongText + +style_presets = table( + "style_presets", + Column("id", Key(), primary_key=True), + Column("name", LongText(), nullable=False), + Column("preset_data", LongText(), nullable=False), + Column("type", LongText(), nullable=False, server_default=default("user")), + inserted_at(), + updated_at(), + Column("user_id", Key(USER_ID_LENGTH), server_default=default("system")), + Column("is_public", Boolean(), nullable=False, server_default=default(False)), +) + +Index("idx_style_presets_is_public", style_presets.c.is_public) +Index("idx_style_presets_name", style_presets.c.name, **LONG_TEXT_INDEX_PREFIX) +Index("idx_style_presets_user_id", style_presets.c.user_id) diff --git a/invokeai/app/services/shared/database/schema/system_prompts.py b/invokeai/app/services/shared/database/schema/system_prompts.py new file mode 100644 index 00000000000..9974ecb1df5 --- /dev/null +++ b/invokeai/app/services/shared/database/schema/system_prompts.py @@ -0,0 +1,28 @@ +"""System prompts for the prompt-expansion models.""" + +from sqlalchemy import Boolean, Column, Index + +from invokeai.app.services.shared.database.schema.metadata import ( + LONG_TEXT_INDEX_PREFIX, + USER_ID_LENGTH, + default, + inserted_at, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import BigInt, Key, LongText + +system_prompts = table( + "system_prompts", + Column("id", Key(), primary_key=True), + Column("name", LongText(), nullable=False), + Column("content", LongText(), nullable=False), + Column("user_id", Key(USER_ID_LENGTH), nullable=False, server_default=default("system")), + Column("is_public", Boolean(), nullable=False, server_default=default(False)), + inserted_at(), + updated_at(), + Column("max_tokens", BigInt()), +) + +Index("idx_system_prompts_name", system_prompts.c.name, **LONG_TEXT_INDEX_PREFIX) +Index("idx_system_prompts_user_id", system_prompts.c.user_id) diff --git a/invokeai/app/services/shared/database/schema/users.py b/invokeai/app/services/shared/database/schema/users.py new file mode 100644 index 00000000000..82ea897a580 --- /dev/null +++ b/invokeai/app/services/shared/database/schema/users.py @@ -0,0 +1,63 @@ +"""Accounts, and the session and invitation tables of an earlier design that nothing uses.""" + +from sqlalchemy import Boolean, Column, ForeignKey, Index + +from invokeai.app.services.shared.database.schema.metadata import ( + USER_ID_LENGTH, + default, + inserted_at, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import BigInt, Key, LongText, Timestamp + +users = table( + "users", + Column("user_id", Key(USER_ID_LENGTH), primary_key=True), + Column("email", Key(), nullable=False, unique=True), + Column("display_name", LongText()), + Column("password_hash", LongText(), nullable=False), + Column("is_admin", Boolean(), nullable=False, server_default=default(False)), + Column("is_active", Boolean(), nullable=False, server_default=default(True)), + inserted_at(), + updated_at(), + Column("last_login_at", Timestamp()), + # Raised to revoke every token issued to the account so far. + Column("token_epoch", BigInt(), nullable=False, server_default=default(0)), +) + +# The unique constraint covers it. Only SQLite has it, where a migration created it. +Index("idx_users_email", users.c.email).ddl_if(dialect="sqlite") +Index("idx_users_is_active", users.c.is_active) +Index("idx_users_is_admin", users.c.is_admin) + +user_sessions = table( + "user_sessions", + Column("session_id", Key(), primary_key=True), + Column("user_id", Key(USER_ID_LENGTH), ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False), + Column("token_hash", Key(), nullable=False), + Column("expires_at", Timestamp(), nullable=False), + inserted_at(), + inserted_at("last_activity_at"), +) + +Index("idx_user_sessions_expires_at", user_sessions.c.expires_at) +Index("idx_user_sessions_token_hash", user_sessions.c.token_hash) +Index("idx_user_sessions_user_id", user_sessions.c.user_id) + +user_invitations = table( + "user_invitations", + Column("invitation_id", Key(), primary_key=True), + Column("email", Key(), nullable=False), + Column("invited_by", Key(USER_ID_LENGTH), ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False), + Column("invitation_code", Key(), nullable=False, unique=True), + Column("is_admin", Boolean(), nullable=False, server_default=default(False)), + Column("expires_at", Timestamp(), nullable=False), + Column("used_at", Timestamp()), + inserted_at(), +) + +Index("idx_user_invitations_email", user_invitations.c.email) +Index("idx_user_invitations_expires_at", user_invitations.c.expires_at) +# The unique constraint covers it. Only SQLite has it, where a migration created it. +Index("idx_user_invitations_invitation_code", user_invitations.c.invitation_code).ddl_if(dialect="sqlite") diff --git a/invokeai/app/services/shared/database/schema/videos.py b/invokeai/app/services/shared/database/schema/videos.py new file mode 100644 index 00000000000..557e257e886 --- /dev/null +++ b/invokeai/app/services/shared/database/schema/videos.py @@ -0,0 +1,100 @@ +"""Video records; the files themselves are on disk.""" + +from sqlalchemy import Boolean, Column, Index, text + +from invokeai.app.services.shared.database.schema.metadata import ( + ENUM_LENGTH, + ON_SERVERS, + USER_ID_LENGTH, + default, + inserted_at, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import BigInt, Key, LongText, Real, Timestamp + +videos = table( + "videos", + Column("video_name", Key(), primary_key=True), + Column("video_origin", Key(ENUM_LENGTH), nullable=False), + Column("video_category", Key(ENUM_LENGTH), nullable=False), + Column("width", BigInt(), nullable=False), + Column("height", BigInt(), nullable=False), + Column("duration", Real(), nullable=False, server_default=default(0.0)), + Column("fps", Real()), + Column("session_id", LongText()), + Column("node_id", LongText()), + Column("metadata", LongText()), + Column("is_intermediate", Boolean(), server_default=default(False)), + Column("starred", Boolean(), server_default=default(False)), + Column("has_workflow", Boolean(), server_default=default(False)), + # No foreign key to `users`, as for images: deleting a user leaves their media for an admin to review. + Column("user_id", Key(USER_ID_LENGTH), nullable=False, server_default=default("system")), + Column("video_subfolder", LongText(), nullable=False, server_default=default("")), + inserted_at(), + updated_at(), + Column("deleted_at", Timestamp()), # unused + Column("project_id", Key()), + Column("file_size_bytes", BigInt()), +) + +Index("idx_videos_created_at", videos.c.created_at) +# As for images: partial indexes on SQLite, and server variants that lead with the condition's columns. +Index( + "idx_videos_intermediate_scope", + videos.c.user_id, + videos.c.project_id, + videos.c.created_at, + videos.c.video_name, + sqlite_where=text("is_intermediate = TRUE"), +).ddl_if(dialect="sqlite") +Index( + "idx_videos_intermediate_scope", + videos.c.is_intermediate, + videos.c.user_id, + videos.c.project_id, + videos.c.created_at, + videos.c.video_name, +).ddl_if(dialect=ON_SERVERS) +Index( + "idx_videos_intermediates_owner", + videos.c.user_id, + videos.c.created_at, + videos.c.video_name, + sqlite_where=text("is_intermediate = TRUE"), +).ddl_if(dialect="sqlite") +Index( + "idx_videos_intermediates_owner", + videos.c.is_intermediate, + videos.c.user_id, + videos.c.created_at, + videos.c.video_name, +).ddl_if(dialect=ON_SERVERS) +Index( + "idx_videos_intermediates_created", + videos.c.created_at, + videos.c.video_name, + sqlite_where=text("is_intermediate = TRUE"), +).ddl_if(dialect="sqlite") +Index("idx_videos_intermediates_created", videos.c.is_intermediate, videos.c.created_at, videos.c.video_name).ddl_if( + dialect=ON_SERVERS +) +Index("idx_videos_starred", videos.c.starred) +Index( + "idx_videos_unmeasured_intermediates", + videos.c.created_at, + videos.c.video_name, + sqlite_where=text("is_intermediate = TRUE AND file_size_bytes IS NULL"), +).ddl_if(dialect="sqlite") +Index( + "idx_videos_unmeasured_intermediates", + videos.c.is_intermediate, + videos.c.file_size_bytes, + videos.c.created_at, + videos.c.video_name, +).ddl_if(dialect=ON_SERVERS) +Index("idx_videos_user_id", videos.c.user_id) +Index("idx_videos_video_category", videos.c.video_category) +# The primary key covers it. Only SQLite has it, where a migration created it. +Index("idx_videos_video_name", videos.c.video_name, unique=True).ddl_if(dialect="sqlite") +Index("idx_videos_video_origin", videos.c.video_origin) diff --git a/invokeai/app/services/shared/database/schema/wildcards.py b/invokeai/app/services/shared/database/schema/wildcards.py new file mode 100644 index 00000000000..205b4d00f7c --- /dev/null +++ b/invokeai/app/services/shared/database/schema/wildcards.py @@ -0,0 +1,25 @@ +"""Wildcards: named lists of values that a prompt can draw from.""" + +from sqlalchemy import Column, ForeignKey, Index + +from invokeai.app.services.shared.database.schema.metadata import ( + USER_ID_LENGTH, + default, + inserted_at, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import Key, LongText + +wildcards = table( + "wildcards", + Column("id", Key(), primary_key=True), + # At most 128 characters (validated by the service). + Column("name", Key(), nullable=False), + Column("values_json", LongText(), nullable=False, server_default=default("[]")), + Column("user_id", Key(USER_ID_LENGTH), ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False), + inserted_at(), + updated_at(), +) + +Index("idx_wildcards_user_id_name", wildcards.c.user_id, wildcards.c.name, unique=True) diff --git a/invokeai/app/services/shared/database/schema/workflows.py b/invokeai/app/services/shared/database/schema/workflows.py new file mode 100644 index 00000000000..664cb202cc4 --- /dev/null +++ b/invokeai/app/services/shared/database/schema/workflows.py @@ -0,0 +1,41 @@ +"""The workflow library.""" + +from sqlalchemy import Boolean, Column, Computed, Index + +from invokeai.app.services.shared.database.dialect import JsonValue +from invokeai.app.services.shared.database.schema.metadata import ( + LONG_TEXT_INDEX_PREFIX, + USER_ID_LENGTH, + default, + inserted_at, + table, + updated_at, +) +from invokeai.app.services.shared.database.types import BigInt, Key, LongText, Timestamp + +workflow_library = table( + "workflow_library", + Column("workflow_id", Key(), primary_key=True), + Column("workflow", LongText(), nullable=False), + inserted_at(), + updated_at(), + # Members of `workflow`, for filtering and sorting. + Column("category", LongText(), Computed(JsonValue("workflow", "$.meta.category")), nullable=False), + Column("name", LongText(), Computed(JsonValue("workflow", "$.name")), nullable=False), + Column("description", LongText(), Computed(JsonValue("workflow", "$.description")), nullable=False), + Column("tags", LongText(), Computed(JsonValue("workflow", "$.tags"))), + Column("opened_at", Timestamp()), + Column("user_id", Key(USER_ID_LENGTH), server_default=default("system")), + Column("is_public", Boolean(), nullable=False, server_default=default(False)), + Column("last_run_at", Timestamp()), + Column("revision", BigInt(), nullable=False, server_default=default(1)), +) + +Index("idx_workflow_library_category", workflow_library.c.category, **LONG_TEXT_INDEX_PREFIX) +Index("idx_workflow_library_created_at", workflow_library.c.created_at) +Index("idx_workflow_library_description", workflow_library.c.description, **LONG_TEXT_INDEX_PREFIX) +Index("idx_workflow_library_is_public", workflow_library.c.is_public) +Index("idx_workflow_library_name", workflow_library.c.name, **LONG_TEXT_INDEX_PREFIX) +Index("idx_workflow_library_opened_at", workflow_library.c.opened_at) +Index("idx_workflow_library_updated_at", workflow_library.c.updated_at) +Index("idx_workflow_library_user_id", workflow_library.c.user_id) diff --git a/invokeai/app/services/shared/database/session_lock.py b/invokeai/app/services/shared/database/session_lock.py new file mode 100644 index 00000000000..a2096fc2fb1 --- /dev/null +++ b/invokeai/app/services/shared/database/session_lock.py @@ -0,0 +1,141 @@ +"""A lock of a MySQL or MariaDB server (`GET_LOCK`), held by the session of a connection of its own.""" + +import hashlib +import threading +from logging import Logger +from typing import Optional + +from sqlalchemy import Connection, Engine +from sqlalchemy.exc import DBAPIError + +from invokeai.app.services.shared.database.errors import DatabaseInUseError + + +class SessionLock: + """A server-wide lock on one database, for one purpose, held by the session of a connection that does nothing + else. The server releases it when that session ends: when the lock is closed, when the process ends, and also + when the connection is cut off (an idle timeout, a restart, a network failure). So a holder checks with `held()` + rather than assuming it. + """ + + def __init__(self, engine: Engine, purpose: str) -> None: + self._conn: Connection = engine.connect() + try: + self.database_name = str(self._conn.exec_driver_sql("SELECT DATABASE()").scalar_one()) + self._conn.commit() + except BaseException: + self.close() + raise + # Lock names are server-wide and at most 64 characters long. + self.name = f"invokeai.{purpose}.{hashlib.sha1(self.database_name.encode()).hexdigest()}" + + def take(self, timeout_seconds: int) -> bool: + """Takes the lock, waiting at most `timeout_seconds`; whether it did. Raises when the server could not try.""" + taken = self._conn.exec_driver_sql("SELECT GET_LOCK(%s, %s)", (self.name, timeout_seconds)).scalar() + self._conn.commit() + if taken is None: + raise RuntimeError(f"The database server could not take the lock {self.name}") + return taken == 1 + + def held(self) -> bool: + """Whether this lock's session still holds it. A check also keeps the session from going idle.""" + try: + held = self._conn.exec_driver_sql("SELECT IS_USED_LOCK(%s) = CONNECTION_ID()", (self.name,)).scalar() + self._conn.commit() + except DBAPIError: + return False + return held == 1 + + def set_idle_timeout(self, seconds: int) -> None: + """Lets the server end this lock's session, and release the lock, after `seconds` without a statement: how + long a lock outlives a host that stopped without closing its connections.""" + self._conn.exec_driver_sql(f"SET SESSION wait_timeout = {int(seconds)}") + self._conn.commit() + + def close(self) -> None: + """Releases the lock by ending its session. (A connection returned to the pool would keep both.)""" + self._conn.invalidate() + self._conn.close() + + +# How often a running process checks that it still holds its instance lock, which also keeps the lock's session +# from going idle. +INSTANCE_LOCK_CHECK_SECONDS = 60 +# How long the server keeps the instance lock of a process whose host stopped without closing its connection. +INSTANCE_LOCK_IDLE_TIMEOUT_SECONDS = 600 + + +class InstanceLock: + """The lock one InvokeAI process at a time holds on a server database, for as long as it serves it. + + A thread checks it every `INSTANCE_LOCK_CHECK_SECONDS`, and takes it again when the server has ended its + session; when another process has taken it meanwhile, it logs an error, as two processes then serve the database. + """ + + def __init__(self, engine: Engine, logger: Logger) -> None: + self._engine = engine + self._logger = logger + self._guard = threading.Lock() + self._closed = threading.Event() + lock, refused = self._take() + if lock is None: + raise DatabaseInUseError( + f"Another InvokeAI process uses the database {refused.database_name!r}. Stop it first: one process at " + "a time serves a database. If none runs, the server still holds the lock of one that stopped without " + f"closing its connection, and releases it within {INSTANCE_LOCK_IDLE_TIMEOUT_SECONDS // 60} minutes; " + f"or end that connection now: SELECT IS_USED_LOCK('{refused.name}') names the one to KILL." + ) + self._lock: Optional[SessionLock] = lock + threading.Thread(target=self._keep, name="database-instance-lock", daemon=True).start() + + def check(self) -> None: + """Verifies the lock, and takes it again when the server has released it.""" + with self._guard: + if self._closed.is_set() or (self._lock is not None and self._lock.held()): + return + if self._lock is not None: + self._lock.close() + self._lock = None + self._logger.warning( + "The database server ended the session holding this process's lock on the database" + ) + try: + lock, refused = self._take() + except Exception as e: # noqa: BLE001 - checked again on the next round + self._logger.error(f"Could not take the lock on the database again: {e}") + return + if lock is None: + self._logger.error( + f"Another InvokeAI process has taken the lock on the database {refused.database_name!r}: two " + "processes now serve it, and undo each other's work. Stop one of them." + ) + return + self._lock = lock + self._logger.info("Took the lock on the database again") + + def close(self) -> None: + self._closed.set() + with self._guard: + if self._lock is not None: + self._lock.close() + self._lock = None + + def _take(self) -> tuple[Optional[SessionLock], SessionLock]: + """The lock if this process took it, and the lock it tried (closed when another process holds it).""" + lock = SessionLock(self._engine, "instance") + try: + if lock.take(0): + lock.set_idle_timeout(INSTANCE_LOCK_IDLE_TIMEOUT_SECONDS) + return lock, lock + except BaseException: + lock.close() + raise + lock.close() + return None, lock + + def _keep(self) -> None: + while not self._closed.wait(INSTANCE_LOCK_CHECK_SECONDS): + try: + self.check() + except Exception as e: # noqa: BLE001 - the thread outlives one failed round + self._logger.error(f"Could not check the lock on the database: {e}") diff --git a/invokeai/app/services/shared/database/startup.py b/invokeai/app/services/shared/database/startup.py new file mode 100644 index 00000000000..c724c8e21ff --- /dev/null +++ b/invokeai/app/services/shared/database/startup.py @@ -0,0 +1,152 @@ +"""Opening the database a config names: the SQLite file in `db_dir`, or the MySQL or MariaDB database of `db_url`.""" + +import re +from logging import Logger +from typing import Optional + +from sqlalchemy import inspect + +from invokeai.app.services.config.config_default import InvokeAIAppConfig +from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase +from invokeai.app.services.shared.database.database import Database, parse_database_url +from invokeai.app.services.shared.sqlite_migrator.migration_loader import MigrationBuildContext, build_migrations +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import Migrator, NoImageFiles + +# The oldest servers the layer is tested on and relies on (window functions, `JSON_TABLE`, `FOR SHARE`, NO PAD +# binary collations). +MINIMUM_SERVER_VERSIONS = {"mysql": (8, 4), "mariadb": (10, 11)} +# A project document is up to 32 MiB, and a statement carries it with room to spare; MariaDB's default is 16 MiB. +MINIMUM_MAX_ALLOWED_PACKET = 64 * 1024 * 1024 + + +class DatabaseSetupError(RuntimeError): + """The configured database cannot be used as it is set up; the message says what to change.""" + + +def open_database(config: InvokeAIAppConfig, logger: Logger) -> Database: + """Opens the database the config names, as it is: no checks, no migration.""" + if config.db_url: + if config.use_memory_db: + raise DatabaseSetupError("`use_memory_db` and `db_url` exclude each other: set one of them") + return Database.open_url(config.db_url, logger) + db_path = None if config.use_memory_db else config.db_path + return Database.open_sqlite(db_path, logger, verbose=config.log_sql, synchronous=config.db_synchronous) + + +def init_database(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileStorageBase) -> Database: + """Opens the app's database and brings it to the newest schema. + + A server database is checked first, and this process takes its instance lock before it migrates or anything + else starts: one process at a time serves a database. + + :param image_files: The image files service, which some migrations need. + """ + database = open_database(config, logger) + try: + if database.dialect_name != "sqlite": + check_server(database) + database.hold_instance_lock() + _warn_when_new_beside_sqlite(database, config, logger) + _migrator(database, config, logger, image_files).run_migrations() + except BaseException: + database.dispose() + raise + return database + + +def open_migrated_database(config: InvokeAIAppConfig, logger: Logger) -> Database: + """Opens the database for a tool that runs beside or instead of the app: it must be at the newest schema, which + the app brings it to when it starts.""" + if not config.db_url and not config.use_memory_db and not config.db_path.exists(): + # Opening it would create an empty database that the app, configured elsewhere, would never use. + raise DatabaseSetupError(f"There is no database at {config.db_path}: is this the InvokeAI root you meant?") + database = open_database(config, logger) + try: + if _migrator(database, config, logger, NoImageFiles()).has_pending_migrations(): + raise DatabaseSetupError( + "The database has not been updated to this version of InvokeAI yet: start InvokeAI once, then run " + "this again" + ) + except BaseException: + database.dispose() + raise + return database + + +def open_copy_target(url: str, logger: Logger) -> tuple[Database, int]: + """Opens the server database `invoke-db-copy` fills: checked as the app checks it, held as the app holds it, + and empty. Also its `max_allowed_packet`.""" + database = Database.open_url(url, logger) + try: + max_allowed_packet = check_server(database) + database.hold_instance_lock() + with database.begin(write=False) as conn: + if inspect(conn).get_table_names(): + raise DatabaseSetupError( + "The target database is not empty: copy into a new, empty database, so nothing in it is mixed " + "with or overwritten by the copy. If an earlier copy failed, or InvokeAI has started on it, drop " + "it and create it again empty" + ) + except BaseException: + database.dispose() + raise + return database, max_allowed_packet + + +def redacted_database_url(url: Optional[str]) -> Optional[str]: + """The URL with its password masked, for logs and the runtime config.""" + if not url: + return url + try: + return parse_database_url(url).render_as_string(hide_password=True) + except ValueError: + # Not a URL that could be opened either; show none of it. + return "***" + + +def _migrator( + database: Database, config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileStorageBase +) -> Migrator: + migrator = Migrator(database) + for migration in build_migrations(MigrationBuildContext(app_config=config, logger=logger, image_files=image_files)): + migrator.register_migration(migration) + return migrator + + +def check_server(database: Database) -> int: + """Refuses a server too old, or one that takes too small a statement; its `max_allowed_packet`.""" + with database.begin(write=False) as conn: + version_text = str(conn.exec_driver_sql("SELECT VERSION()").scalar_one()) + max_allowed_packet = int(conn.exec_driver_sql("SELECT @@max_allowed_packet").scalar_one()) + minimum = MINIMUM_SERVER_VERSIONS[database.dialect_name] + version = _version_of(version_text) + if version < minimum: + product = "MariaDB" if database.dialect_name == "mariadb" else "MySQL" + raise DatabaseSetupError( + f"{product} {version_text} is too old: InvokeAI needs {product} {minimum[0]}.{minimum[1]} or newer" + ) + if max_allowed_packet < MINIMUM_MAX_ALLOWED_PACKET: + raise DatabaseSetupError( + f"The server's max_allowed_packet is {max_allowed_packet // (1024 * 1024)} MiB; InvokeAI needs at least " + f"{MINIMUM_MAX_ALLOWED_PACKET // (1024 * 1024)} MiB to store large projects. Set it in the server's " + "configuration, e.g. `max_allowed_packet=64M` under [mysqld], and restart the server." + ) + return max_allowed_packet + + +def _version_of(version_text: str) -> tuple[int, int]: + match = re.match(r"(\d+)\.(\d+)", version_text) + if match is None: + raise DatabaseSetupError(f"The database server reports a version InvokeAI cannot read: {version_text!r}") + return int(match.group(1)), int(match.group(2)) + + +def _warn_when_new_beside_sqlite(database: Database, config: InvokeAIAppConfig, logger: Logger) -> None: + with database.begin(write=False) as conn: + empty = not inspect(conn).get_table_names() + if empty and config.db_path.exists(): + logger.warning( + f"The database {redacted_database_url(config.db_url)} is new, and InvokeAI starts it empty, but there is " + f"a SQLite database at {config.db_path}. To keep its boards, images, workflows and settings, stop " + "InvokeAI and copy it over with `invoke-db-copy` before using the new database." + ) diff --git a/invokeai/app/services/shared/database/types.py b/invokeai/app/services/shared/database/types.py new file mode 100644 index 00000000000..3c7ffc76e3b --- /dev/null +++ b/invokeai/app/services/shared/database/types.py @@ -0,0 +1,199 @@ +"""Column types: each renders the declared type of the migrated SQLite schema there, and its equivalent on a server. + +SQLite keeps a column's declared type as written, and the schema metadata reproduces the migrated schema +exactly, so on SQLite these types render the migrations' declarations verbatim. A server needs bounded types +where SQLite has none: an index key holds at most 3072 bytes on MySQL and MariaDB, 768 utf8mb4 characters. + +| Type | SQLite | MySQL / MariaDB | +|---|---|---| +| `Key(n=255)`: text a key, a foreign key or an index covers whole | TEXT | VARCHAR(n) | +| `NoCaseKey(n=255)`: the same, compared case-insensitively | TEXT COLLATE NOCASE | VARCHAR(n), case-insensitive | +| `LongText`: any other text, JSON documents among it | TEXT | LONGTEXT | +| `Timestamp(declared)`: canonical UTC text (see `now_text`) | `declared` (DATETIME) | VARCHAR(32) | +| `BigInt` | INTEGER, so a primary key is the rowid | BIGINT | +| `Real` | REAL | DOUBLE | +| `Blob` | BLOB | LONGBLOB | + +Booleans are SQLAlchemy's `Boolean`: BOOLEAN on SQLite, TINYINT(1) on a server. +""" + +import functools +from collections.abc import Callable +from datetime import datetime, timezone +from typing import Any, TypeVar, cast + +from sqlalchemy import BigInteger, Double, Integer, LargeBinary, String, Text +from sqlalchemy.dialects.mysql import LONGBLOB, LONGTEXT +from sqlalchemy.engine import Dialect +from sqlalchemy.types import REAL, TypeDecorator, TypeEngine, UserDefinedType + +from invokeai.app.services.shared.database.engines import SERVER_DIALECTS + +# Case-insensitive, accent- and space-sensitive (NO PAD) Unicode collations. SQLite's NOCASE folds the case of +# ASCII letters only; these compare by the Unicode collation algorithm, so they also equate the case of other +# letters, full-width and compatibility forms ('Ｃａｔ' = 'cat', 'fi' = 'fi'), and precomposed and decomposed +# accents, and they ignore ignorable code points (a soft hyphen, a zero-width space). +MYSQL_NOCASE_COLLATION = "utf8mb4_0900_as_ci" +MARIADB_NOCASE_COLLATION = "utf8mb4_uca1400_nopad_as_ci" + +# 'YYYY-MM-DD HH:MM:SS.fff' needs 23 characters. Some rows hold ISO 8601 text instead, at most 32 characters: +# `datetime.isoformat()` with microseconds and a UTC offset. +TIMESTAMP_LENGTH = 32 + +_Load = TypeVar("_Load", bound=Callable[[Any, Dialect], TypeEngine[Any]]) + + +def _per_backend(load: _Load) -> _Load: + """Keeps the type that `load` gives each backend, per column type. + + SQLAlchemy asks for it again on every execution of an expanding IN, and adapts every new type it gets anew, + which cost about 19 µs a statement. + """ + + @functools.wraps(load) + def cached(self: Any, dialect: Dialect) -> TypeEngine[Any]: + types: dict[tuple[type, str], TypeEngine[Any]] = self.__dict__.setdefault("_backend_types", {}) + key = (type(dialect), dialect.name) + if key not in types: + types[key] = load(self, dialect) + return types[key] + + return cast(_Load, cached) + + +def timestamp_text(moment: datetime) -> str: + """A UTC moment as the canonical timestamp text, 'YYYY-MM-DD HH:MM:SS.fff'. + + The format of SQLite's `STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')`, the schema's column default: timestamps + sort and compare as text, on every backend. The application is the only clock, so a server's clock and + time zone never enter a timestamp. + """ + return moment.strftime("%Y-%m-%d %H:%M:%S.%f")[:23] + + +def now_text() -> str: + """The current UTC time as the canonical timestamp text (see `timestamp_text`).""" + return timestamp_text(datetime.now(timezone.utc)) + + +class _Declared(UserDefinedType[str]): + """A type rendered as the given declaration, verbatim; values pass through as text.""" + + cache_ok = True + + def __init__(self, declaration: str) -> None: + self.declaration = declaration + + def get_col_spec(self, **kw: Any) -> str: + return self.declaration + + +class Key(TypeDecorator[str]): + """Text that a key, a foreign key or an index covers whole: TEXT on SQLite, VARCHAR(length) on a server.""" + + impl = Text + cache_ok = True + + def __init__(self, length: int = 255) -> None: + super().__init__() + self.length = length + + @_per_backend + def load_dialect_impl(self, dialect: Dialect) -> TypeEngine[Any]: + if dialect.name in SERVER_DIALECTS: + return dialect.type_descriptor(String(self.length)) + return dialect.type_descriptor(Text()) + + +class NoCaseKey(TypeDecorator[str]): + """Key text compared case-insensitively: SQLite's NOCASE, a case-insensitive collation on a server. + + PostgreSQL has no such collation by default; there it renders SQLite's, which a PostgreSQL server refuses. + """ + + impl = Text + cache_ok = True + + def __init__(self, length: int = 255) -> None: + super().__init__() + self.length = length + + @_per_backend + def load_dialect_impl(self, dialect: Dialect) -> TypeEngine[Any]: + if dialect.name == "mysql": + return dialect.type_descriptor(String(self.length, collation=MYSQL_NOCASE_COLLATION)) + if dialect.name == "mariadb": + return dialect.type_descriptor(String(self.length, collation=MARIADB_NOCASE_COLLATION)) + return dialect.type_descriptor(Text(collation="NOCASE")) + + +class LongText(TypeDecorator[str]): + """Text of any length that no index covers: TEXT on SQLite, LONGTEXT on a server.""" + + impl = Text + cache_ok = True + + @_per_backend + def load_dialect_impl(self, dialect: Dialect) -> TypeEngine[Any]: + if dialect.name in SERVER_DIALECTS: + return dialect.type_descriptor(LONGTEXT()) + return dialect.type_descriptor(Text()) + + +class Timestamp(TypeDecorator[str]): + """A point in time as canonical UTC text (see `now_text`). + + :param declared: The declared type on SQLite, as the migrations wrote it: DATETIME, or TEXT for a few. + """ + + impl = String + cache_ok = True + + def __init__(self, declared: str = "DATETIME") -> None: + super().__init__(TIMESTAMP_LENGTH) + self.declared = declared + + @_per_backend + def load_dialect_impl(self, dialect: Dialect) -> TypeEngine[Any]: + if dialect.name == "sqlite": + return dialect.type_descriptor(_Declared(self.declared)) + return dialect.type_descriptor(String(TIMESTAMP_LENGTH)) + + +class BigInt(TypeDecorator[int]): + """An integer: INTEGER on SQLite, where an INTEGER primary key is the rowid; BIGINT elsewhere.""" + + impl = BigInteger + cache_ok = True + + @_per_backend + def load_dialect_impl(self, dialect: Dialect) -> TypeEngine[Any]: + if dialect.name == "sqlite": + return dialect.type_descriptor(Integer()) + return dialect.type_descriptor(BigInteger()) + + +class Real(TypeDecorator[float]): + """A double-precision float: REAL on SQLite, DOUBLE elsewhere.""" + + impl = Double + cache_ok = True + + @_per_backend + def load_dialect_impl(self, dialect: Dialect) -> TypeEngine[Any]: + if dialect.name == "sqlite": + return dialect.type_descriptor(REAL()) + return dialect.type_descriptor(Double()) + + +class Blob(TypeDecorator[bytes]): + """Bytes of any length: BLOB on SQLite, LONGBLOB on a server.""" + + impl = LargeBinary + cache_ok = True + + @_per_backend + def load_dialect_impl(self, dialect: Dialect) -> TypeEngine[Any]: + if dialect.name in SERVER_DIALECTS: + return dialect.type_descriptor(LONGBLOB()) + return dialect.type_descriptor(LargeBinary()) diff --git a/invokeai/app/services/shared/intermediate_delete.py b/invokeai/app/services/shared/intermediate_delete.py index 7c1f75bf3c4..7c65cc2334d 100644 --- a/invokeai/app/services/shared/intermediate_delete.py +++ b/invokeai/app/services/shared/intermediate_delete.py @@ -1,9 +1,11 @@ """Conditional intermediate deletion shared by the image and video services. -A guard runs on the deleting transaction's cursor, immediately before the `DELETE`, and returns -the subset of the candidate names that may still go. Because the database serializes writers -through one connection and lock, nothing — a project save, an enqueue, a promotion — can make a -name protected between the guard's answer and the record's removal. +A guard runs on the deleting transaction, immediately before the `DELETE`, and returns the subset +of the candidate names that may still go. That transaction holds the `media_protection` lock, which +every write that makes media protected (a reference, a hold) shares: nothing — a project save, a +browser hold, a reused cached output — can make a name protected between the guard's answer and the +record's removal. (On SQLite the database serializes all writers anyway; the session queue, which makes +media active as well, runs only there until it is ported, and then shares the lock too.) Records go first and files second, with a durable journal spanning the two. Staging files before a conditional record delete would need a restore for every survivor, and a concurrent delete of a @@ -12,14 +14,16 @@ recovery finishes the purge for every journalled name whose record is gone. """ -import sqlite3 from collections.abc import Callable from dataclasses import dataclass, field -from typing import Optional, Protocol, Sequence +from typing import TYPE_CHECKING, Optional, Protocol, Sequence + +if TYPE_CHECKING: + from invokeai.app.services.shared.database.queries import Queries class IntermediateDeleteGuard(Protocol): - def __call__(self, cursor: sqlite3.Cursor, names: Sequence[str]) -> list[str]: ... + def __call__(self, q: "Queries", names: Sequence[str]) -> list[str]: ... @dataclass diff --git a/invokeai/app/services/shared/invocation_context.py b/invokeai/app/services/shared/invocation_context.py index 2be9fc01032..ecfec580948 100644 --- a/invokeai/app/services/shared/invocation_context.py +++ b/invokeai/app/services/shared/invocation_context.py @@ -19,7 +19,7 @@ from invokeai.app.services.model_records.model_records_base import UnknownModelException from invokeai.app.services.session_processor.session_processor_common import ProgressImage from invokeai.app.services.shared.execution_effects import ExecutionEffectsRecorder, ExecutionInterface -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.services.videos.videos_common import VideoDTO from invokeai.app.services.wildcard_records.wildcard_records_common import build_wildcard_manager from invokeai.app.util.step_callback import diffusion_step_callback diff --git a/invokeai/app/services/shared/media_references.py b/invokeai/app/services/shared/media_references.py index 1ae4253cf33..5367a240615 100644 --- a/invokeai/app/services/shared/media_references.py +++ b/invokeai/app/services/shared/media_references.py @@ -11,12 +11,12 @@ Extraction is deliberately over-inclusive: every string under an image- or video-name key counts, history included. A false reference retains a file; a missed one deletes something in use. -Applied migrations call the table DDL, the extractors and `replace_media_references`: keep their DDL -and signatures frozen, since changing them changes what a fresh install's migrations do. +Applied migrations call the extractors: keep their behaviour frozen, since changing it changes what a fresh +install's migrations do. The migrations' table DDL and writes are frozen in `sqlite_migrator/migrations/ +_media_references_v1.py`; the application writes through `database.queries.media_references`. """ import json -import sqlite3 from dataclasses import dataclass, field from typing import Literal @@ -78,63 +78,3 @@ def extract_media_references_from_json(document_json: str) -> MediaReferences: return extract_media_references(json.loads(document_json)) except ValueError: return MediaReferences() - - -def replace_media_references( - cursor: sqlite3.Cursor, - *, - owner_kind: MediaReferenceOwnerKind, - user_id: str, - owner_id: str, - references: MediaReferences, -) -> None: - """Makes the index for one owner equal to `references`, on the caller's transaction.""" - delete_media_references(cursor, owner_kind=owner_kind, user_id=user_id, owner_id=owner_id) - rows: list[tuple[str, str, str, str, str]] = [] - rows.extend((owner_kind, user_id, owner_id, "image", name) for name in sorted(references.images)) - rows.extend((owner_kind, user_id, owner_id, "video", name) for name in sorted(references.videos)) - if rows: - cursor.executemany( - """--sql - INSERT OR IGNORE INTO media_references (owner_kind, user_id, owner_id, media_kind, media_name) - VALUES (?, ?, ?, ?, ?); - """, - rows, - ) - - -def delete_media_references( - cursor: sqlite3.Cursor, *, owner_kind: MediaReferenceOwnerKind, user_id: str, owner_id: str -) -> None: - cursor.execute( - """--sql - DELETE FROM media_references - WHERE owner_kind = ? AND user_id = ? AND owner_id = ?; - """, - (owner_kind, user_id, owner_id), - ) - - -def create_media_references_table(cursor: sqlite3.Cursor) -> None: - """DDL shared by the migration and tests that need the table without the whole chain.""" - cursor.execute( - """--sql - CREATE TABLE IF NOT EXISTS media_references ( - -- A MediaReferenceOwnerKind - owner_kind TEXT NOT NULL, - -- The owning account; project ids are unique per user, not globally. - user_id TEXT NOT NULL, - owner_id TEXT NOT NULL, - -- 'image' or 'video' - media_kind TEXT NOT NULL, - media_name TEXT NOT NULL, - PRIMARY KEY (owner_kind, user_id, owner_id, media_kind, media_name) - ); - """ - ) - cursor.execute( - """--sql - CREATE INDEX IF NOT EXISTS idx_media_references_media - ON media_references(media_kind, media_name); - """ - ) diff --git a/invokeai/app/services/shared/pagination.py b/invokeai/app/services/shared/pagination.py index a2dd4fc8145..6193c69e0de 100644 --- a/invokeai/app/services/shared/pagination.py +++ b/invokeai/app/services/shared/pagination.py @@ -1,7 +1,10 @@ +from enum import Enum from typing import Generic, TypeVar from pydantic import BaseModel, Field +from invokeai.app.util.metaenum import MetaEnum + GenericBaseModel = TypeVar("GenericBaseModel", bound=BaseModel) # Upper bound for offset-paginated `limit` query params. Values flow verbatim into @@ -10,6 +13,14 @@ MAX_PAGE_SIZE = 1000 +class SQLiteDirection(str, Enum, metaclass=MetaEnum): + # The direction of an ordering. (Its name is part of the API schema, from when only SQLite backed it; a docstring + # would add a description to that schema.) + + Ascending = "ASC" + Descending = "DESC" + + class CursorPaginatedResults(BaseModel, Generic[GenericBaseModel]): """ Cursor-paginated results diff --git a/invokeai/app/services/shared/sqlite/sqlite_common.py b/invokeai/app/services/shared/sqlite/sqlite_common.py deleted file mode 100644 index 25206952011..00000000000 --- a/invokeai/app/services/shared/sqlite/sqlite_common.py +++ /dev/null @@ -1,10 +0,0 @@ -from enum import Enum - -from invokeai.app.util.metaenum import MetaEnum - -sqlite_memory = ":memory:" - - -class SQLiteDirection(str, Enum, metaclass=MetaEnum): - Ascending = "ASC" - Descending = "DESC" diff --git a/invokeai/app/services/shared/sqlite/sqlite_database.py b/invokeai/app/services/shared/sqlite/sqlite_database.py deleted file mode 100644 index 29bd0a32ac7..00000000000 --- a/invokeai/app/services/shared/sqlite/sqlite_database.py +++ /dev/null @@ -1,165 +0,0 @@ -import os -import sqlite3 -import threading -from collections.abc import Generator -from contextlib import contextmanager -from logging import Logger -from pathlib import Path -from typing import get_args - -from invokeai.app.services.config.config_default import DB_SYNCHRONOUS -from invokeai.app.services.shared.sqlite.sqlite_common import sqlite_memory - - -class SqliteDatabase: - """ - Manages a connection to an SQLite database. - - :param db_path: Path to the database file. If None, an in-memory database is used. - :param logger: Logger to use for logging. - :param verbose: Whether to log SQL statements. Provides `logger.debug` as the SQLite trace callback. - :param synchronous: SQLite `synchronous` setting. Defaults to `full`, what InvokeAI has always used. - `normal` is refused, with a warning, for a database file whose journal mode did not become WAL. - - This is a light wrapper around the `sqlite3` module, providing a few conveniences: - - The database file is written to disk if it does not exist. - - Foreign key constraints are enabled by default. - - The connection is configured to use the `sqlite3.Row` row factory. - - In addition to the constructor args, the instance provides the following attributes and methods: - - `conn`: A `sqlite3.Connection` object. Note that the connection must never be closed if the database is in-memory. - - `lock`: A shared re-entrant lock, used to approximate thread safety. - - `clean()`: Runs the SQL `VACUUM;` command and reports on the freed space. - """ - - def __init__( - self, - db_path: Path | None, - logger: Logger, - verbose: bool = False, - synchronous: DB_SYNCHRONOUS = "full", - ) -> None: - """Initializes the database. This is used internally by the class constructor.""" - # This constructor is called directly (user management commands, tests), not only through the - # validated config, and `synchronous` is interpolated into a PRAGMA below. `DB_SYNCHRONOUS` is - # a static annotation, so the value is checked here, before anything is opened. - if synchronous not in get_args(DB_SYNCHRONOUS): - raise ValueError(f"Invalid synchronous setting {synchronous!r}, expected one of {get_args(DB_SYNCHRONOUS)}") - - self._logger = logger - self._db_path = db_path - self._verbose = verbose - self._lock = threading.RLock() - - if not self._db_path: - logger.info("Initializing in-memory database") - else: - self._db_path.parent.mkdir(parents=True, exist_ok=True) - self._logger.info(f"Initializing database at {self._db_path}") - - self._conn = sqlite3.connect(database=self._db_path or sqlite_memory, check_same_thread=False) - self._conn.row_factory = sqlite3.Row - - if self._verbose: - self._conn.set_trace_callback(self._logger.debug) - - # Enable foreign key constraints - self._conn.execute("PRAGMA foreign_keys = ON;") - - # Enable Write-Ahead Logging (WAL) mode for better concurrency. The statement reports the mode - # that was actually established: WAL needs shared memory, so it can fail to engage -- quietly -- - # on network filesystems, some container volume drivers and `SQLITE_OMIT_WAL` builds. - journal_mode = str(self._conn.execute("PRAGMA journal_mode = WAL;").fetchone()[0]).lower() - - # Set a busy timeout to prevent database lockups during writes - self._conn.execute("PRAGMA busy_timeout = 5000;") # 5 seconds - - # Durability. SQLite's own default is `full`, which fsyncs on every commit; `normal` trades the - # last transactions on a power loss or OS crash for roughly 12x shorter commits. Shorter commits - # matter twice here, because every write holds the lock that serialises all database work. - # - # `normal` cannot corrupt the database only *because* of WAL; on a rollback journal it can, so a - # database file that did not get WAL keeps `full` instead. In-memory databases report `memory` - # and have no durability to trade, so the requirement does not apply to them. - # - # Refused rather than fatal: the setting is a performance preference, and starting with stronger - # durability than asked for costs commit time, where starting with weaker durability than the - # documented guarantee risks the user's database. A warning names the mode so the cause is - # visible instead of merely slow. - if synchronous == "normal" and self._db_path is not None and journal_mode != "wal": - self._logger.warning( - f"Journal mode is '{journal_mode}', not WAL, so db_synchronous='normal' was refused and " - "'full' used instead: without WAL that setting can corrupt the database on power loss." - ) - synchronous = "full" - - # PRAGMA does not take bind parameters, so the value is interpolated. It is checked against - # DB_SYNCHRONOUS at the top of this constructor. - self._conn.execute(f"PRAGMA synchronous = {synchronous.upper()};") - - def clean(self) -> None: - """ - Cleans the database by running the VACUUM command, reporting on the freed space. - """ - # No need to clean in-memory database - if not self._db_path: - return - try: - # Serialize with transaction() users: services with their own - # threads (e.g. the image index worker) may have a statement in - # progress on this shared connection, and VACUUM refuses to run - # concurrently with one ("cannot VACUUM - SQL statements in - # progress"). - with self._lock, self._conn as conn: - initial_db_size = Path(self._db_path).stat().st_size - conn.execute("VACUUM;") - conn.commit() - final_db_size = Path(self._db_path).stat().st_size - freed_space_in_mb = round((initial_db_size - final_db_size) / 1024 / 1024, 2) - if freed_space_in_mb > 0: - self._logger.info(f"Cleaned database (freed {freed_space_in_mb}MB)") - except Exception as e: - self._logger.error(f"Error cleaning database: {e}") - raise - - def backup_to(self, destination: Path) -> None: - """Create a consistent SQLite backup at a new destination. - - The connection backup API includes committed WAL pages. The shared lock keeps this - connection's transaction users from interleaving with the copy, and the destination is - reserved exclusively so an existing backup can never be overwritten. - """ - destination = Path(destination) - destination.parent.mkdir(parents=True, exist_ok=True) - descriptor = os.open(destination, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600) - os.close(descriptor) - try: - with self._lock: - backup_connection = sqlite3.connect(destination) - try: - self._conn.backup(backup_connection) - backup_connection.commit() - finally: - backup_connection.close() - except Exception: - destination.unlink(missing_ok=True) - Path(f"{destination}-wal").unlink(missing_ok=True) - Path(f"{destination}-shm").unlink(missing_ok=True) - raise - - @contextmanager - def transaction(self) -> Generator[sqlite3.Cursor, None, None]: - """ - Thread-safe context manager for DB work. - Acquires the RLock, yields a Cursor, then commits or rolls back. - """ - with self._lock: - cursor = self._conn.cursor() - try: - yield cursor - self._conn.commit() - except Exception: - self._conn.rollback() - raise - finally: - cursor.close() diff --git a/invokeai/app/services/shared/sqlite/sqlite_util.py b/invokeai/app/services/shared/sqlite/sqlite_util.py deleted file mode 100644 index 5f6125b5ad1..00000000000 --- a/invokeai/app/services/shared/sqlite/sqlite_util.py +++ /dev/null @@ -1,32 +0,0 @@ -from logging import Logger - -from invokeai.app.services.config.config_default import InvokeAIAppConfig -from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.shared.sqlite_migrator.migration_loader import MigrationBuildContext, build_migrations -from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator - - -def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileStorageBase) -> SqliteDatabase: - """ - Initializes the SQLite database. - - :param config: The app config - :param logger: The logger - :param image_files: The image files service (used by migration 2) - - This function: - - Instantiates a :class:`SqliteDatabase` - - Instantiates a :class:`SqliteMigrator` and registers all migrations - - Runs all migrations - """ - db_path = None if config.use_memory_db else config.db_path - db = SqliteDatabase(db_path=db_path, logger=logger, verbose=config.log_sql, synchronous=config.db_synchronous) - - migrator = SqliteMigrator(db=db) - migration_context = MigrationBuildContext(app_config=config, logger=logger, image_files=image_files) - for migration in build_migrations(migration_context): - migrator.register_migration(migration) - migrator.run_migrations() - - return db diff --git a/invokeai/app/services/shared/sqlite_migrator/migration_loader.py b/invokeai/app/services/shared/sqlite_migrator/migration_loader.py index c32e829892e..1c8361d80d0 100644 --- a/invokeai/app/services/shared/sqlite_migrator/migration_loader.py +++ b/invokeai/app/services/shared/sqlite_migrator/migration_loader.py @@ -10,9 +10,14 @@ from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase -from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import MigrationBase DEFAULT_MIGRATIONS_PACKAGE = "invokeai.app.services.shared.sqlite_migrator.migrations" + +# The last migration that runs on SQLite only. Every migration after it is portable (a `portable_callback`), so +# that it runs on every database backend. +PORTABLE_CUTOVER = "2026_10_01_add_anima_variant" + _MIGRATION_MODULE_RE = re.compile(r"^migration_(\d+)$") _DATED_MIGRATION_MODULE_RE = re.compile(r"^migration_(\d{4}_\d{2}_\d{2}_[a-z0-9][a-z0-9_]*)$") @@ -39,7 +44,7 @@ def get_dependency(self, name: str) -> Any: raise MigrationLoaderError(f"Migration builder requested unknown dependency '{name}'") -MigrationBuilder = Callable[..., Migration] +MigrationBuilder = Callable[..., MigrationBase] @dataclass(frozen=True) @@ -91,12 +96,14 @@ def discover_migration_builders(package_name: str = DEFAULT_MIGRATIONS_PACKAGE) return sorted(discovered, key=lambda discovered_builder: discovered_builder.sort_key) -def build_migrations(context: MigrationBuildContext, package_name: str = DEFAULT_MIGRATIONS_PACKAGE) -> list[Migration]: - migrations: list[Migration] = [] +def build_migrations( + context: MigrationBuildContext, package_name: str = DEFAULT_MIGRATIONS_PACKAGE +) -> list[MigrationBase]: + migrations: list[MigrationBase] = [] for discovered_builder in discover_migration_builders(package_name): kwargs = _get_builder_kwargs(discovered_builder.builder, context, discovered_builder.module_name) migration = discovered_builder.builder(**kwargs) - if not isinstance(migration, Migration): + if not isinstance(migration, MigrationBase): raise MigrationLoaderError( f"Migration builder '{discovered_builder.builder.__name__}' in " f"'{discovered_builder.module_name}' must return Migration" diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/_media_references_v1.py b/invokeai/app/services/shared/sqlite_migrator/migrations/_media_references_v1.py new file mode 100644 index 00000000000..06d6c189009 --- /dev/null +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/_media_references_v1.py @@ -0,0 +1,69 @@ +"""The `media_references` table as the SQLite-only migrations create and fill it, frozen with them. + +Changing this changes what a fresh install's migrations do. The application writes through +`database.queries.media_references`. +""" + +import sqlite3 + +from invokeai.app.services.shared.media_references import MediaReferenceOwnerKind, MediaReferences + + +def replace_media_references( + cursor: sqlite3.Cursor, + *, + owner_kind: MediaReferenceOwnerKind, + user_id: str, + owner_id: str, + references: MediaReferences, +) -> None: + """Makes the index for one owner equal to `references`, on the caller's transaction.""" + delete_media_references(cursor, owner_kind=owner_kind, user_id=user_id, owner_id=owner_id) + rows: list[tuple[str, str, str, str, str]] = [] + rows.extend((owner_kind, user_id, owner_id, "image", name) for name in sorted(references.images)) + rows.extend((owner_kind, user_id, owner_id, "video", name) for name in sorted(references.videos)) + if rows: + cursor.executemany( + """--sql + INSERT OR IGNORE INTO media_references (owner_kind, user_id, owner_id, media_kind, media_name) + VALUES (?, ?, ?, ?, ?); + """, + rows, + ) + + +def delete_media_references( + cursor: sqlite3.Cursor, *, owner_kind: MediaReferenceOwnerKind, user_id: str, owner_id: str +) -> None: + cursor.execute( + """--sql + DELETE FROM media_references + WHERE owner_kind = ? AND user_id = ? AND owner_id = ?; + """, + (owner_kind, user_id, owner_id), + ) + + +def create_media_references_table(cursor: sqlite3.Cursor) -> None: + """DDL shared by the migration and tests that need the table without the whole chain.""" + cursor.execute( + """--sql + CREATE TABLE IF NOT EXISTS media_references ( + -- A MediaReferenceOwnerKind + owner_kind TEXT NOT NULL, + -- The owning account; project ids are unique per user, not globally. + user_id TEXT NOT NULL, + owner_id TEXT NOT NULL, + -- 'image' or 'video' + media_kind TEXT NOT NULL, + media_name TEXT NOT NULL, + PRIMARY KEY (owner_kind, user_id, owner_id, media_kind, media_name) + ); + """ + ) + cursor.execute( + """--sql + CREATE INDEX IF NOT EXISTS idx_media_references_media + ON media_references(media_kind, media_name); + """ + ) diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_09_22_add_intermediates_management.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_09_22_add_intermediates_management.py index a1ce015c169..ffd459b248f 100644 --- a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_09_22_add_intermediates_management.py +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_09_22_add_intermediates_management.py @@ -18,10 +18,8 @@ import sqlite3 from logging import Logger -from invokeai.app.services.shared.media_references import ( - create_media_references_table, - extract_media_references, -) +from invokeai.app.services.shared.media_references import extract_media_references +from invokeai.app.services.shared.sqlite_migrator.migrations._media_references_v1 import create_media_references_table from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_09_23_index_client_state_references.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_09_23_index_client_state_references.py index 80e1a59ed46..17ff3237597 100644 --- a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_09_23_index_client_state_references.py +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_09_23_index_client_state_references.py @@ -10,11 +10,8 @@ import sqlite3 from logging import Logger -from invokeai.app.services.shared.media_references import ( - MediaReferenceOwnerKind, - extract_media_references_from_json, - replace_media_references, -) +from invokeai.app.services.shared.media_references import MediaReferenceOwnerKind, extract_media_references_from_json +from invokeai.app.services.shared.sqlite_migrator.migrations._media_references_v1 import replace_media_references from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_04_add_db_locks.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_04_add_db_locks.py new file mode 100644 index 00000000000..74b3d44450d --- /dev/null +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_04_add_db_locks.py @@ -0,0 +1,35 @@ +"""Add `db_locks`, the rows a transaction locks to serialise work with other connections. + +SQLite serialises all writes through its single connection; on MySQL and MariaDB, invariants that span rows +(the last active administrator, for one) need a lock that every transaction changing them takes first. A +lock is a row: `SELECT ... FOR UPDATE` on it waits for the transaction that holds it to end. +""" + +from sqlalchemy import Column, column, insert, inspect, select, table + +from invokeai.app.services.shared.database.types import Key +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import ( + PortableMigration, + PortableMigrationContext, +) + +_db_locks = table("db_locks", column("name")) + +# Literal rather than imported: a migration keeps meaning what it meant when it was written. +_ADMIN_ACCOUNTS = "admin_accounts" + + +def _add_db_locks(context: PortableMigrationContext) -> None: + if not inspect(context.conn).has_table("db_locks"): + context.create_table("db_locks", Column("name", Key(32), primary_key=True)) + # On a server the table may exist without the row: its creation committed, and the run then failed. + if context.conn.execute(select(_db_locks.c.name).where(_db_locks.c.name == _ADMIN_ACCOUNTS)).first() is None: + context.conn.execute(insert(_db_locks).values(name=_ADMIN_ACCOUNTS)) + + +def build_migration() -> PortableMigration: + return PortableMigration( + id="2026_10_04_add_db_locks", + depends_on="2026_10_01_add_anima_variant", + callback=_add_db_locks, + ) diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_04_add_session_queue_listing_index.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_04_add_session_queue_listing_index.py index 21e767ddab7..38f433cb0ae 100644 --- a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_04_add_session_queue_listing_index.py +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_04_add_session_queue_listing_index.py @@ -17,27 +17,37 @@ lands at the left edge and splits pages half-full: about 1.6x the index size once history grows. The database keeps no table statistics, so this index attracts statements that filter only on -`queue_id`; see `PENDING_QUEUE_ITEM_COUNT_QUERY` in the session queue service for the guard rule. +`queue_id`; the queue's queries keep such terms off it with `dialect.Unindexed`. + +SQLite only. On MySQL and MariaDB `queue_id` and `origin` are long texts, which no index holds whole, and a server's +planner weighs its indexes by their statistics. """ -import sqlite3 +from sqlalchemy import inspect, text + +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import ( + PortableMigration, + PortableMigrationContext, +) -from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration +_NAME = "idx_session_queue_listing" -class AddSessionQueueListingIndexCallback: - def __call__(self, cursor: sqlite3.Cursor) -> None: - cursor.execute( - """--sql - CREATE INDEX IF NOT EXISTS idx_session_queue_listing - ON session_queue (queue_id, created_at ASC, item_id DESC, user_id, origin, status); - """ - ) +def _add_session_queue_listing_index(context: PortableMigrationContext) -> None: + if context.conn.dialect.name != "sqlite": + return + if any(index["name"] == _NAME for index in inspect(context.conn).get_indexes("session_queue")): + return + context.op.create_index( + _NAME, + "session_queue", + ["queue_id", "created_at", text("item_id DESC"), "user_id", "origin", "status"], + ) -def build_migration() -> Migration: - return Migration( +def build_migration() -> PortableMigration: + return PortableMigration( id="2026_10_04_add_session_queue_listing_index", depends_on="2026_10_01_add_anima_variant", - callback=AddSessionQueueListingIndexCallback(), + callback=_add_session_queue_listing_index, ) diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_06_add_intermediates_state.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_06_add_intermediates_state.py new file mode 100644 index 00000000000..fc3c07fc587 --- /dev/null +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_06_add_intermediates_state.py @@ -0,0 +1,82 @@ +"""Keep the intermediates cleanup's state in tables, and let its windows page by creation time. + +SQLite's TEMP tables belong to the one connection that created them; MySQL and MariaDB run each transaction on a +connection of a pool. The holds of cached outputs and the marks of unmeasurable files move into ordinary tables, +which the cleanup service empties when it starts. A `media_protection` lock serialises a cleanup's deleting +transaction with the saves that reference media. The cleanup pages intermediates by (creation time, name), since a +server has no rowid, so its indexes end with those columns. +""" + +from typing import Optional + +from sqlalchemy import Column, Index, column, insert, inspect, select, table, text + +from invokeai.app.services.shared.database.types import Key, Timestamp +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import ( + PortableMigration, + PortableMigrationContext, +) + +_db_locks = table("db_locks", column("name")) + +# Literals rather than imports: a migration keeps meaning what it meant when it was written. +_MEDIA_PROTECTION = "media_protection" +_ENUM_LENGTH = 32 +_INTERMEDIATE = "is_intermediate = TRUE" +_UNMEASURED = "is_intermediate = TRUE AND file_size_bytes IS NULL" + + +def _index(context: PortableMigrationContext, name: str, table_name: str, columns: list[str], where: str) -> None: + """Gives the table the index on SQLite (partial, on `where`) or on a server (leading with `where`'s columns).""" + on_sqlite = context.conn.dialect.name == "sqlite" + if not on_sqlite: + leading = ["is_intermediate", "file_size_bytes"] if "file_size_bytes" in where else ["is_intermediate"] + columns = [*leading, *columns] + existing: Optional[list[str]] = next( + (index["column_names"] for index in inspect(context.conn).get_indexes(table_name) if index["name"] == name), + None, + ) + if existing == columns: + return + if existing is not None: + context.op.drop_index(name, table_name=table_name) + if on_sqlite: + context.op.create_index(name, table_name, columns, sqlite_where=text(where)) + else: + context.op.create_index(name, table_name, columns) + + +def _add_intermediates_state(context: PortableMigrationContext) -> None: + tables = inspect(context.conn).get_table_names() + if "intermediates_session_holds" not in tables: + context.create_table( + "intermediates_session_holds", + Column("session_id", Key(), primary_key=True), + Column("media_kind", Key(_ENUM_LENGTH), primary_key=True), + Column("media_name", Key(), primary_key=True), + Column("released_at", Timestamp("TEXT")), + Index("idx_intermediates_session_holds_media", "media_kind", "media_name", "released_at"), + ) + if "intermediates_unmeasurable" not in tables: + context.create_table( + "intermediates_unmeasurable", + Column("media_kind", Key(_ENUM_LENGTH), primary_key=True), + Column("media_name", Key(), primary_key=True), + ) + if context.conn.execute(select(_db_locks.c.name).where(_db_locks.c.name == _MEDIA_PROTECTION)).first() is None: + context.conn.execute(insert(_db_locks).values(name=_MEDIA_PROTECTION)) + + for kind, name_column in (("images", "image_name"), ("videos", "video_name")): + keyset = ["created_at", name_column] + _index(context, f"idx_{kind}_intermediate_scope", kind, ["user_id", "project_id", *keyset], _INTERMEDIATE) + _index(context, f"idx_{kind}_intermediates_owner", kind, ["user_id", *keyset], _INTERMEDIATE) + _index(context, f"idx_{kind}_intermediates_created", kind, keyset, _INTERMEDIATE) + _index(context, f"idx_{kind}_unmeasured_intermediates", kind, keyset, _UNMEASURED) + + +def build_migration() -> PortableMigration: + return PortableMigration( + id="2026_10_06_add_intermediates_state", + depends_on="2026_10_04_add_db_locks", + callback=_add_intermediates_state, + ) diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_07_add_image_index_vocabulary_lock.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_07_add_image_index_vocabulary_lock.py new file mode 100644 index 00000000000..9f21d764ae1 --- /dev/null +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_07_add_image_index_vocabulary_lock.py @@ -0,0 +1,31 @@ +"""Add the `image_index_vocabulary` lock, which a replace of the image index's custom vocabulary takes first. + +The replace deletes every term and inserts the new ones. On MySQL and MariaDB two replaces running side by side +would each delete what the other has not committed yet, and leave the terms of both. +""" + +from sqlalchemy import column, insert, select, table + +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import ( + PortableMigration, + PortableMigrationContext, +) + +_db_locks = table("db_locks", column("name")) + +# Literal rather than imported: a migration keeps meaning what it meant when it was written. +_IMAGE_INDEX_VOCABULARY = "image_index_vocabulary" + + +def _add_image_index_vocabulary_lock(context: PortableMigrationContext) -> None: + lock = select(_db_locks.c.name).where(_db_locks.c.name == _IMAGE_INDEX_VOCABULARY) + if context.conn.execute(lock).first() is None: + context.conn.execute(insert(_db_locks).values(name=_IMAGE_INDEX_VOCABULARY)) + + +def build_migration() -> PortableMigration: + return PortableMigration( + id="2026_10_07_add_image_index_vocabulary_lock", + depends_on="2026_10_06_add_intermediates_state", + callback=_add_image_index_vocabulary_lock, + ) diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_07_add_session_queue_admission_lock.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_07_add_session_queue_admission_lock.py new file mode 100644 index 00000000000..f6abe0d02bc --- /dev/null +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_07_add_session_queue_admission_lock.py @@ -0,0 +1,32 @@ +"""Add the `session_queue_admission` lock, which every write that adds items to the session queue takes first. + +The queue holds at most `max_queue_size` pending items, and an enqueue request with an idempotency key is settled by +its receipt. On MySQL and MariaDB two enqueues running side by side would each count the pending items and look for +the receipt before the other inserts, so both could fit into the last free place, or both write the same receipt. +""" + +from sqlalchemy import column, insert, select, table + +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import ( + PortableMigration, + PortableMigrationContext, +) + +_db_locks = table("db_locks", column("name")) + +# Literal rather than imported: a migration keeps meaning what it meant when it was written. +_SESSION_QUEUE_ADMISSION = "session_queue_admission" + + +def _add_session_queue_admission_lock(context: PortableMigrationContext) -> None: + lock = select(_db_locks.c.name).where(_db_locks.c.name == _SESSION_QUEUE_ADMISSION) + if context.conn.execute(lock).first() is None: + context.conn.execute(insert(_db_locks).values(name=_SESSION_QUEUE_ADMISSION)) + + +def build_migration() -> PortableMigration: + return PortableMigration( + id="2026_10_07_add_session_queue_admission_lock", + depends_on="2026_10_07_add_image_index_vocabulary_lock", + callback=_add_session_queue_admission_lock, + ) diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_07_drop_sqlite_triggers.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_07_drop_sqlite_triggers.py new file mode 100644 index 00000000000..e0c5d983861 --- /dev/null +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2026_10_07_drop_sqlite_triggers.py @@ -0,0 +1,54 @@ +"""Drop the SQLite triggers: the application sets what they set, on every backend. + +The triggers stamped `updated_at` on update, `started_at` and `completed_at` on queue status changes, and counted +`session_revision` when a queue item's session changed. MySQL and MariaDB databases never had them, so the +application sets those columns itself; with both writing on SQLite, the triggers only overwrote the application's +values, and in a format of their own for boards and client state (`CURRENT_TIMESTAMP`, without milliseconds). + +One behaviour changes with them: `session_revision` now counts every write of a session, as it does on servers, +rather than only those that change its text. Its one reader compares it for equality. +""" + +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import ( + PortableMigration, + PortableMigrationContext, +) + +# Literal rather than read from the database: a migration keeps meaning what it meant when it was written. +_TRIGGERS = ( + "models_updated_at", + "style_presets", + "tg_app_settings_updated_at", + "tg_board_images_updated_at", + "tg_board_videos_updated_at", + "tg_boards_updated_at", + "tg_client_state_updated_at", + "tg_image_projections_updated_at", + "tg_image_subfolder_move_jobs_updated_at", + "tg_images_updated_at", + "tg_projects_updated_at", + "tg_session_queue_completed_at", + "tg_session_queue_session_revision", + "tg_session_queue_started_at", + "tg_session_queue_updated_at", + "tg_system_prompts_updated_at", + "tg_users_updated_at", + "tg_videos_updated_at", + "tg_wildcards_updated_at", + "tg_workflow_library_updated_at", +) + + +def _drop_sqlite_triggers(context: PortableMigrationContext) -> None: + if context.conn.dialect.name != "sqlite": + return + for name in _TRIGGERS: + context.conn.exec_driver_sql(f'DROP TRIGGER IF EXISTS "{name}"') + + +def build_migration() -> PortableMigration: + return PortableMigration( + id="2026_10_07_drop_sqlite_triggers", + depends_on="2026_10_07_add_session_queue_admission_lock", + callback=_drop_sqlite_triggers, + ) diff --git a/invokeai/app/services/shared/sqlite_migrator/sqlite_migrator_common.py b/invokeai/app/services/shared/sqlite_migrator/sqlite_migrator_common.py index 1bbb8d6f1e9..31b477f18be 100644 --- a/invokeai/app/services/shared/sqlite_migrator/sqlite_migrator_common.py +++ b/invokeai/app/services/shared/sqlite_migrator/sqlite_migrator_common.py @@ -1,7 +1,16 @@ import sqlite3 -from typing import Optional, Protocol, runtime_checkable +from dataclasses import dataclass, field +from logging import Logger +from typing import Any, Optional, Protocol, runtime_checkable +from alembic.operations import Operations from pydantic import BaseModel, ConfigDict, Field, model_validator +from sqlalchemy import Connection, MetaData, Table +from sqlalchemy.schema import SchemaItem +from typing_extensions import Self + +from invokeai.app.services.shared.database.schema.metadata import define_table +from invokeai.app.services.shared.database.schema.metadata import metadata as schema_metadata @runtime_checkable @@ -20,6 +29,50 @@ class MigrateCallback(Protocol): def __call__(self, cursor: sqlite3.Cursor) -> None: ... +@dataclass(frozen=True) +class PortableMigrationContext: + """What a portable migration works with. + + :param conn: The connection of the migration's transaction. + :param op: Alembic's operations (`op.add_column`, `op.create_index`, ...), bound to `conn`. + :param logger: The logger. + :param metadata: The tables this migration creates. A foreign key names a table here: one created earlier + in the migration, or an existing one loaded with `Table(name, context.metadata, autoload_with=conn)`. + """ + + conn: Connection + op: Operations + logger: Logger + metadata: MetaData = field(default_factory=lambda: MetaData(naming_convention=schema_metadata.naming_convention)) + + def create_table(self, name: str, *items: SchemaItem, **options: Any) -> Table: + """Creates a table the way the schema metadata creates its tables. Use it instead of `op.create_table`. + + It gets the server table options (InnoDB, the binary collation, DYNAMIC rows, whatever the server's + defaults) and each backend's rules for defaults, generated columns and per-backend indexes, which + Alembic's operations know nothing of. Items are given as to `schema.metadata.table()`, with indexes + naming their columns. + """ + created = define_table(self.metadata, name, *items, **options) + created.create(self.conn) + return created + + +@runtime_checkable +class PortableMigrateCallback(Protocol): + """A callback that performs a migration on any database backend, with Alembic's operations. + + It runs in a transaction the migrator commits. On MySQL and MariaDB every DDL statement commits on its own, + though, so a portable migration must be idempotent: run again after it failed halfway, it finds part of its + work done and completes the rest. On SQLite it runs with foreign keys off, so that rebuilding a table + (Alembic's batch mode) does not cascade into the tables that reference it; they are checked before the + migration commits, so rows it orphans fail it. The schema metadata in `shared/database/schema/` gets the + same change in the same commit, because new server databases are created from it. + """ + + def __call__(self, context: PortableMigrationContext) -> None: ... + + class MigrationError(RuntimeError): """Raised when a migration fails.""" @@ -28,70 +81,18 @@ class MigrationVersionError(ValueError): """Raised when a migration version is invalid.""" -class Migration(BaseModel): +class MigrationBase(BaseModel): """ - Represents a migration for a SQLite database. + What every migration has: a stable ID, and the migration that must run first. :param from_version: The legacy database version on which this migration may be run :param to_version: The legacy database version that results from this migration :param id: The stable migration ID. Legacy migrations default to ``migration_{to_version}``. :param depends_on: The stable ID of the migration that must run first. - :param migrate_callback: The callback to run to perform the migration Migrations are executed according to their stable ID dependencies. Existing legacy migrations also keep ``from_version`` and ``to_version`` so older numeric migration state can be mapped to applied migration IDs. - New graph-only migrations may omit legacy versions, but must provide an explicit ``id``. - - Migration callbacks will be provided an open cursor to the database. They should not commit their - transaction; this is handled by the migrator. - - It is suggested to use a class to define the migration callback and a builder function to create - the :class:`Migration`. This allows the callback to be provided with additional dependencies and - keeps things tidy, as all migration logic is self-contained. - - Example: - ```py - # Define the migration callback class - class Migration1Callback: - # This migration needs a logger, so we define a class that accepts a logger in its constructor. - def __init__(self, image_files: ImageFileStorageBase) -> None: - self._image_files = ImageFileStorageBase - - # This dunder method allows the instance of the class to be called like a function. - def __call__(self, cursor: sqlite3.Cursor) -> None: - self._add_with_banana_column(cursor) - self._do_something_with_images(cursor) - - def _add_with_banana_column(self, cursor: sqlite3.Cursor) -> None: - \"""Adds the with_banana column to the sushi table.\""" - # Execute SQL using the cursor, taking care to *not commit* a transaction - cursor.execute('ALTER TABLE sushi ADD COLUMN with_banana BOOLEAN DEFAULT TRUE;') - - def _do_something_with_images(self, cursor: sqlite3.Cursor) -> None: - \"""Does something with the image files service.\""" - self._image_files.get(...) - - # Define the migration builder function. This function creates an instance of the migration callback - # class and returns a Migration. - def build_migration_1(image_files: ImageFileStorageBase) -> Migration: - \"""Builds the migration from database version 0 to 1. - Requires the image files service to... - \""" - - migration_1 = Migration( - from_version=0, - to_version=1, - migrate_callback=Migration1Callback(image_files=image_files), - ) - - return migration_1 - - # Register the migration after all dependencies have been initialized - db = SqliteDatabase(db_path, logger) - migrator = SqliteMigrator(db) - migrator.register_migration(build_migration_1(image_files)) - migrator.run_migrations() - ``` + New graph-only migrations omit legacy versions, and must provide an explicit ``id``. """ from_version: Optional[int] = Field( @@ -102,10 +103,9 @@ def build_migration_1(image_files: ImageFileStorageBase) -> Migration: ) id: Optional[str] = Field(default=None, description="Stable migration ID") depends_on: Optional[str] = Field(default=None, description="Stable ID of the migration dependency") - callback: MigrateCallback = Field(description="The callback to run to perform the migration") @model_validator(mode="after") - def validate_versions_and_ids(self) -> "Migration": + def validate_versions_and_ids(self) -> Self: """Validates legacy versions and derives stable IDs for legacy migrations.""" has_from_version = self.from_version is not None has_to_version = self.to_version is not None @@ -139,6 +139,55 @@ def sort_key(self) -> tuple[int, int, str]: model_config = ConfigDict(arbitrary_types_allowed=True) +class Migration(MigrationBase): + """ + A migration that runs on SQLite only, given an open cursor. Every migration up to ``PORTABLE_CUTOVER`` is one; + new migrations are portable (:class:`PortableMigration`). + + :param callback: The callback to run to perform the migration. It is provided an open cursor, and does not + commit; the migrator does. + """ + + callback: MigrateCallback = Field(description="The callback to run to perform the migration") + + +class PortableMigration(MigrationBase): + """ + A migration that runs on every database backend: a callback using Alembic's operations (see + :class:`PortableMigrateCallback`). Every migration after ``PORTABLE_CUTOVER`` is one. + + Example: + ```py + def _add_bananas(context: PortableMigrationContext) -> None: + # Idempotent: on MySQL and MariaDB, a failed run may have created the table already. + if not inspect(context.conn).has_table("bananas"): + context.create_table( + "bananas", + Column("banana_id", Key(), primary_key=True), + Column("ripeness", BigInt(), nullable=False, server_default=default(0)), + inserted_at(), + ) + + + def build_migration() -> PortableMigration: + return PortableMigration( + id="2026_10_05_add_bananas", + depends_on="2026_10_01_add_anima_variant", + callback=_add_bananas, + ) + ``` + """ + + callback: PortableMigrateCallback = Field(description="The callback to run to perform the migration") + + @model_validator(mode="after") + def validate_no_legacy_version(self) -> Self: + """A portable migration is identified by its ID alone.""" + if self.from_version is not None or self.to_version is not None: + raise MigrationVersionError("a portable migration has no legacy version") + return self + + class MigrationSet: """ A set of Migrations. Performs validation during migration registration and provides utility methods. @@ -149,9 +198,9 @@ class MigrationSet: """ def __init__(self) -> None: - self._migrations: set[Migration] = set() + self._migrations: set[MigrationBase] = set() - def register(self, migration: Migration) -> None: + def register(self, migration: MigrationBase) -> None: """Registers a migration.""" migration_from_already_registered = migration.from_version is not None and any( m.from_version == migration.from_version for m in self._migrations if m.from_version is not None @@ -166,7 +215,7 @@ def register(self, migration: Migration) -> None: raise MigrationVersionError("Migration with id already registered") self._migrations.add(migration) - def get(self, from_version: int) -> Optional[Migration]: + def get(self, from_version: int) -> Optional[MigrationBase]: """Gets the migration that may be run on the given database version.""" # register() ensures that there is only one migration with a given from_version, so this is safe. return next((m for m in self._migrations if m.from_version == from_version), None) @@ -210,7 +259,7 @@ def validate_dependency_graph(self) -> None: visiting: set[str] = set() visited: set[str] = set() - def visit(migration: Migration) -> None: + def visit(migration: MigrationBase) -> None: migration_id = migration.id if migration_id is None: raise MigrationError("Migration is missing id") @@ -227,7 +276,7 @@ def visit(migration: Migration) -> None: for migration in self._migrations: visit(migration) - def get_migration_plan(self, applied_migration_ids: set[str]) -> list[Migration]: + def get_migration_plan(self, applied_migration_ids: set[str]) -> list[MigrationBase]: """Gets a deterministic migration plan from the set of applied migration IDs.""" self.validate_dependency_graph() known_migration_ids = set(self.migrations_by_id) @@ -236,7 +285,7 @@ def get_migration_plan(self, applied_migration_ids: set[str]) -> list[Migration] unknown_ids = ", ".join(sorted(unknown_applied_ids)) raise MigrationError(f"Database contains unknown applied migration IDs: {unknown_ids}") - plan: list[Migration] = [] + plan: list[MigrationBase] = [] planned_or_applied_ids = set(applied_migration_ids) remaining = { migration.id: migration for migration in self._migrations if migration.id not in applied_migration_ids @@ -276,9 +325,9 @@ def latest_version(self) -> int: return latest_version or 0 @property - def migrations(self) -> tuple[Migration, ...]: + def migrations(self) -> tuple[MigrationBase, ...]: return tuple(sorted(self._migrations, key=lambda migration: migration.sort_key)) @property - def migrations_by_id(self) -> dict[str, Migration]: + def migrations_by_id(self) -> dict[str, MigrationBase]: return {migration.id or "": migration for migration in self._migrations} diff --git a/invokeai/app/services/shared/sqlite_migrator/sqlite_migrator_impl.py b/invokeai/app/services/shared/sqlite_migrator/sqlite_migrator_impl.py index 27c0921c738..162cb29f02d 100644 --- a/invokeai/app/services/shared/sqlite_migrator/sqlite_migrator_impl.py +++ b/invokeai/app/services/shared/sqlite_migrator/sqlite_migrator_impl.py @@ -1,53 +1,86 @@ +import logging import sqlite3 -from contextlib import closing +import tempfile +from collections import Counter +from collections.abc import Iterator +from contextlib import contextmanager from datetime import datetime from pathlib import Path -from typing import Optional - -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration, MigrationError, MigrationSet - - -class SqliteMigrator: +from typing import NoReturn, Optional + +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy import Connection, insert, inspect, select +from sqlalchemy.exc import DBAPIError + +from invokeai.app.services.config.config_default import DefaultInvokeAIAppConfig +from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase +from invokeai.app.services.shared.database.copy import copy_rows +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema import metadata +from invokeai.app.services.shared.database.schema.migrator import applied_migrations, migrations +from invokeai.app.services.shared.database.session_lock import SessionLock +from invokeai.app.services.shared.sqlite_migrator.migration_loader import MigrationBuildContext, build_migrations +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import ( + Migration, + MigrationBase, + MigrationError, + MigrationSet, + PortableMigration, + PortableMigrationContext, +) + +# How long the migration of a server database waits for another process's migration of it to finish. +MIGRATION_LOCK_TIMEOUT_SECONDS = 600 + + +class Migrator: """ - Manages migrations for a SQLite database. + Brings the database up to date by running the migrations it has not had, in dependency order. - :param db: The instance of :class:`SqliteDatabase` to migrate. + :param database: The database to migrate. - Migrations should be registered with :meth:`register_migration`, either directly or via the migration loader. + Migrations are registered with :meth:`register_migration`, either directly or via the migration loader. They are planned by stable migration ID dependencies and recorded in the ``applied_migrations`` table. Legacy numeric versions are still written for migrations that define ``to_version``. - Each migration is run in a transaction. If a migration fails, the transaction is rolled back. + A SQLite database runs every migration, the legacy cursor migrations and the portable ones, each in a + transaction of its own that a failure rolls back, after a backup of the database file. + + A MySQL or MariaDB database is created at the newest schema instead (see :meth:`_bootstrap`), so only the + portable migrations added after its creation run on it. A server-wide lock keeps two processes from migrating + the same database at once. DDL commits as it runs there, so a failed portable migration is not rolled back: + its id is recorded only once it succeeds, and it runs again, from the start, the next time. Example Usage: ```py - db = SqliteDatabase(db_path="my_db.db", logger=logger) - migrator = SqliteMigrator(db=db) - migrator.register_migration(build_migration_1()) - migrator.register_migration(build_migration_2()) + migrator = Migrator(database) + for migration in build_migrations(migration_context): + migrator.register_migration(migration) migrator.run_migrations() ``` """ backup_path: Optional[Path] = None - def __init__(self, db: SqliteDatabase) -> None: - self._db = db - self._logger = db._logger + def __init__(self, database: Database) -> None: + self._database = database + self._logger = database.logger self._migration_set = MigrationSet() self._backup_path: Optional[Path] = None - def register_migration(self, migration: Migration) -> None: + def register_migration(self, migration: MigrationBase) -> None: """Registers a migration.""" self._migration_set.register(migration) self._logger.debug(f"Registered migration {migration.from_version} -> {migration.to_version}") def run_migrations(self) -> bool: - """Migrates the database to the latest version.""" + """Migrates the database to the latest version. Returns whether a migration ran or the database was created.""" # This throws if there is a problem. self._migration_set.validate_dependency_graph() - cursor = self._db._conn.cursor() + if self._database.dialect_name != "sqlite": + return self._run_server_migrations() + cursor = self._database.sqlite.conn.cursor() self._validate_existing_applied_migrations(cursor=cursor) self._create_migrations_table(cursor=cursor) self._validate_existing_legacy_migrations(cursor=cursor) @@ -80,21 +113,31 @@ def _backup_db(self) -> None: """Makes a backup of the db if it is a file db and a backup has not already been made.""" if self._backup_path is not None: return - if self._db._db_path is not None: + db_path = self._database.sqlite.path + if db_path is not None: timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") - self._backup_path = self._db._db_path.parent / f"{self._db._db_path.stem}_backup_{timestamp}.db" + self._backup_path = db_path.parent / f"{db_path.stem}_backup_{timestamp}.db" + # A backup of the same second exists when the database was migrated again right away, e.g. on a restart + # after a failed migration: keep it, and number this one. + attempt = 1 + while self._backup_path.exists(): + self._backup_path = db_path.parent / f"{db_path.stem}_backup_{timestamp}-{attempt}.db" + attempt += 1 self._logger.info(f"Backing up database to {str(self._backup_path)}") - with closing(sqlite3.connect(self._backup_path)) as backup_conn: - self._db._conn.backup(backup_conn) + self._database.backup(self._backup_path) else: self._logger.info("Using in-memory database, no backup needed") - def _run_migration(self, migration: Migration) -> None: - """Runs a single migration.""" + def _run_migration(self, migration: MigrationBase) -> None: + """Runs a single migration on the SQLite database.""" + if isinstance(migration, PortableMigration): + self._run_portable_migration(migration) + return + assert isinstance(migration, Migration) try: # Using sqlite3.Connection as a context manager commits a the transaction on exit, or rolls it back if an # exception is raised. - with self._db._conn as conn: + with self._database.sqlite.conn as conn: cursor = conn.cursor() # Begun explicitly, because the context manager above only commits or rolls back a # transaction that is already open — it does not start one. The connection runs in @@ -311,3 +354,221 @@ def _get_current_version(cls, cursor: sqlite3.Cursor) -> int: if "no such table" in str(e): return 0 raise + + def _run_portable_migration(self, migration: PortableMigration) -> None: + """Runs a portable migration and records it in one transaction (on a server, its DDL commits as it runs). + + On SQLite, foreign keys are off while it runs: rebuilding a table (Alembic's batch mode) drops the old + one, which would otherwise delete the rows of every table that references it, by cascade. The pragma + cannot change inside a transaction, so it is set before, under the database's lock, and the foreign + keys are checked before the transaction commits (see `_check_foreign_keys`). + """ + self._logger.debug(f"Running migration '{migration.id}'") + try: + if self._database.dialect_name != "sqlite": + self._apply_portable_migration(migration) + else: + sqlite = self._database.sqlite + with sqlite.lock: + sqlite.conn.execute("PRAGMA foreign_keys = OFF") + try: + self._apply_portable_migration(migration) + finally: + sqlite.conn.execute("PRAGMA foreign_keys = ON") + except Exception as e: + msg = f"Error running migration '{migration.id}': {e}" + self._logger.error(msg) + raise MigrationError(msg) from e + self._logger.debug(f"Successfully ran migration '{migration.id}'") + + def _apply_portable_migration(self, migration: PortableMigration) -> None: + with self._database.begin(write=True) as conn: + violations_before = self._foreign_key_violations(conn) + if violations_before: + self._logger.warning( + "The database holds rows whose foreign keys match nothing (table, referenced table): " + f"{dict(violations_before)}" + ) + migration.callback( + PortableMigrationContext( + conn=conn, op=Operations(MigrationContext.configure(conn)), logger=self._logger + ) + ) + self._check_foreign_keys(conn, violations_before) + conn.execute(insert(applied_migrations).values(migration_id=migration.id)) + + def _check_foreign_keys(self, conn: Connection, before: Optional[Counter[tuple[str, str]]]) -> None: + """Fails a SQLite migration that leaves rows whose foreign keys match nothing. + + Only rows the migration orphaned count: a database can hold orphans from before (tools that wrote with + foreign keys off did leave some), and those must not stop every migration from then on. So the migration + fails when there are more such rows after it than before. They are counted, not identified: a rebuilt + table renumbers its rows, and a renamed one reports them under its new name. + """ + if before is None: + return + after = self._foreign_key_violations(conn) + if after is None: + raise MigrationError("The migration leaves a foreign key that names no table or key") + if after.total() > before.total(): + raise MigrationError( + f"The migration leaves rows whose foreign keys match nothing (table, referenced table): {dict(after)}" + ) + + def _foreign_key_violations(self, conn: Connection) -> Optional[Counter[tuple[str, str]]]: + """On SQLite, the rows whose foreign keys match nothing, per table and referenced table. + + None when SQLite cannot check (a foreign key in the database names no table or key), or on a server, + which enforces foreign keys as rows are written. + """ + if self._database.dialect_name != "sqlite": + return None + try: + rows = conn.exec_driver_sql("PRAGMA foreign_key_check").all() + except DBAPIError as e: + self._logger.warning(f"The database's foreign keys could not be checked: {e}") + return None + return Counter((str(row[0]), str(row[2])) for row in rows) + + def has_pending_migrations(self) -> bool: + """Whether the database lacks a registered migration, without changing it. A database without the migrator's + records (a new one, or one from before them) lacks them all. Raises for a server database with tables but + no records, which the app refuses as well.""" + if self._database.dialect_name != "sqlite" and self._server_has_tables(): + applied_ids = self._server_applied_migration_ids() + return bool(self._migration_set.get_migration_plan(applied_migration_ids=applied_ids)) + with self._database.begin(write=False) as conn: + if not inspect(conn).has_table(applied_migrations.name): + return self._migration_set.count > 0 + applied = set(conn.execute(select(applied_migrations.c.migration_id)).scalars()) + return bool(self._migration_set.get_migration_plan(applied_migration_ids=applied)) + + def _run_server_migrations(self) -> bool: + with self._server_migration_lock() as lock: + bootstrapped = False + if not self._server_has_tables(): + self._bootstrap() + bootstrapped = True + plan = self._migration_set.get_migration_plan(applied_migration_ids=self._server_applied_migration_ids()) + portable: list[PortableMigration] = [] + for migration in plan: + if not isinstance(migration, PortableMigration): + raise MigrationError( + f"Migration '{migration.id}' runs on SQLite only, and this {self._database.dialect_name} " + "database has not had it" + ) + portable.append(migration) + if portable: + self._logger.info("Database update needed") + for migration in portable: + lock.verify() + self._run_portable_migration(migration) + if portable: + self._logger.info("Database updated successfully") + return bootstrapped or bool(portable) + + @contextmanager + def _server_migration_lock(self) -> Iterator["_ServerMigrationLock"]: + """Holds a lock that one process at a time can hold for this database, on a connection of its own.""" + session_lock = SessionLock(self._database.engine, "migrate") + try: + lock = _ServerMigrationLock(session_lock) + lock.acquire(self._logger) + yield lock + finally: + session_lock.close() + + def _server_has_tables(self) -> bool: + with self._database.begin(write=False) as conn: + return bool(inspect(conn).get_table_names()) + + def _server_applied_migration_ids(self) -> set[str]: + with self._database.begin(write=False) as conn: + applied_ids: set[str] = set() + if inspect(conn).has_table(applied_migrations.name): + applied_ids = set(conn.execute(select(applied_migrations.c.migration_id)).scalars()) + if not applied_ids: + # The records are written last, so this is an interrupted creation, if not another application's tables. + raise MigrationError( + "The database has tables but no record of the migrations it has had: it is not an InvokeAI " + "database, or its creation was interrupted. Use an empty database." + ) + return applied_ids + + def _bootstrap(self) -> None: + """Creates the newest schema in the empty server database, with the rows the migrations seed. + + The rows come from a reference: an in-memory SQLite database that the migration chain builds, in a + temporary root, so the migrations' clean-ups of old files touch nothing real. The migrator's records go in + last, so a creation that is interrupted leaves a database `_server_applied_migration_ids` refuses. + """ + self._logger.info("Creating the database schema") + # The reference's migrations report clean-ups of a root that is not the user's. + quiet = self._logger.getChild("reference") + quiet.setLevel(logging.WARNING) + with tempfile.TemporaryDirectory() as root: + # Settings from the environment or a config file would point the clean-ups at real directories. + config = DefaultInvokeAIAppConfig() + config._root = Path(root) + reference = Database.open_sqlite(None, quiet) + try: + reference_migrator = Migrator(reference) + context = MigrationBuildContext(app_config=config, logger=quiet, image_files=NoImageFiles()) + for migration in build_migrations(context): + reference_migrator.register_migration(migration) + reference_migrator.run_migrations() + + with self._database.begin(write=True) as conn: + metadata.create_all(conn) + records = [migrations, applied_migrations] + copy_rows(reference, self._database, tables=[t for t in metadata.sorted_tables if t not in records]) + # `applied_migrations` last of all: a database without its rows is refused as incomplete. + copy_rows(reference, self._database, tables=[migrations]) + copy_rows(reference, self._database, tables=[applied_migrations]) + finally: + reference.dispose() + + +def _no_image_files(*args: object, **kwargs: object) -> NoReturn: + raise RuntimeError("No image files were given: a migration that reads them cannot run here") + + +# For the migrations that read image files, where none runs or there are no images: the reference database of a +# bootstrap, and the check for pending migrations. +NoImageFiles = type( + "NoImageFiles", + (ImageFileStorageBase,), + dict.fromkeys(ImageFileStorageBase.__abstractmethods__, _no_image_files), +) + + +_LOCK_LOST = ( + "The connection holding the migration lock was lost, so another process may be migrating the database; start again" +) + + +class _ServerMigrationLock: + """The lock on migrating one server database, which one process at a time holds. It is verified before each + migration rather than assumed: the server releases it with its connection, even one cut off while idle.""" + + def __init__(self, lock: SessionLock) -> None: + self._lock = lock + + def acquire(self, logger: logging.Logger) -> None: + try: + if self._lock.take(0): + return + logger.info("Waiting for another process to finish migrating the database") + if self._lock.take(MIGRATION_LOCK_TIMEOUT_SECONDS): + return + except RuntimeError as e: + raise MigrationError(str(e)) from e + raise MigrationError( + f"Another process has been migrating this database for {MIGRATION_LOCK_TIMEOUT_SECONDS} s; " + "let it finish before starting again" + ) + + def verify(self) -> None: + """Raises if the lock is no longer held.""" + if not self._lock.held(): + raise MigrationError(_LOCK_LOST) diff --git a/invokeai/app/services/style_preset_records/style_preset_records_common.py b/invokeai/app/services/style_preset_records/style_preset_records_common.py index 7e465128f30..be183bcf587 100644 --- a/invokeai/app/services/style_preset_records/style_preset_records_common.py +++ b/invokeai/app/services/style_preset_records/style_preset_records_common.py @@ -3,14 +3,17 @@ import csv import json from enum import Enum -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional import pydantic -from fastapi import UploadFile from pydantic import AliasChoices, BaseModel, ConfigDict, Field, TypeAdapter from invokeai.app.util.metaenum import MetaEnum +if TYPE_CHECKING: + # Only the router passes an upload; the database layer imports this module for its records. + from fastapi import UploadFile + class StylePresetNotFoundError(Exception): """Raised when a style preset is not found""" @@ -88,7 +91,7 @@ class InvalidPresetImportDataError(ValueError): pass -async def parse_presets_from_file(file: UploadFile) -> list[StylePresetWithoutId]: +async def parse_presets_from_file(file: "UploadFile") -> list[StylePresetWithoutId]: """Parses style presets from a file. The file must be a CSV or JSON file. If CSV, the file must have the following columns: diff --git a/invokeai/app/services/style_preset_records/style_preset_records_default.py b/invokeai/app/services/style_preset_records/style_preset_records_default.py new file mode 100644 index 00000000000..b3493d87e14 --- /dev/null +++ b/invokeai/app/services/style_preset_records/style_preset_records_default.py @@ -0,0 +1,79 @@ +import json +from pathlib import Path +from typing import Optional + +from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.style_preset_records.style_preset_records_base import StylePresetRecordsStorageBase +from invokeai.app.services.style_preset_records.style_preset_records_common import ( + PresetType, + StylePresetChanges, + StylePresetNotFoundError, + StylePresetRecordDTO, + StylePresetWithoutId, +) +from invokeai.app.util.misc import uuid_string + +# System user id used for default / shipped presets and for legacy rows pre-dating +# the per-user ownership columns added in migration 27. +SYSTEM_USER_ID = "system" + + +def _found(preset: Optional[StylePresetRecordDTO], style_preset_id: str) -> StylePresetRecordDTO: + if preset is None: + raise StylePresetNotFoundError(f"Style preset with id {style_preset_id} not found") + return preset + + +class StylePresetRecordsStorage(StylePresetRecordsStorageBase): + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + + def start(self, invoker: Invoker) -> None: + self._invoker = invoker + self._sync_default_style_presets() + + def get(self, style_preset_id: str) -> StylePresetRecordDTO: + """Gets a style preset by ID.""" + return _found(self._queries.style_presets.get(style_preset_id), style_preset_id) + + def create(self, style_preset: StylePresetWithoutId, user_id: str) -> StylePresetRecordDTO: + style_preset_id = uuid_string() + self._queries.style_presets.insert({style_preset_id: style_preset}, user_id) + # The record is what was stored, so it is not read back. + return StylePresetRecordDTO(**style_preset.model_dump(), id=style_preset_id, user_id=user_id) + + def create_many(self, style_presets: list[StylePresetWithoutId], user_id: str) -> None: + self._queries.style_presets.insert({uuid_string(): style_preset for style_preset in style_presets}, user_id) + + def update(self, style_preset_id: str, changes: StylePresetChanges) -> StylePresetRecordDTO: + def apply(q: Queries) -> Optional[StylePresetRecordDTO]: + q.style_presets.update(style_preset_id, changes) + return q.style_presets.get(style_preset_id) + + return _found(self._queries.run(apply), style_preset_id) + + def delete(self, style_preset_id: str) -> None: + self._queries.style_presets.delete(style_preset_id) + + def get_many( + self, + type: PresetType | None = None, + user_id: str | None = None, + is_admin: bool = False, + ) -> list[StylePresetRecordDTO]: + # Visible to non-admin: own + default + public. + return self._queries.style_presets.all(type=type, user_id=user_id, is_admin=is_admin) + + def _sync_default_style_presets(self) -> None: + """Syncs default style presets to the database. Internal use only.""" + with open(Path(__file__).parent / Path("default_style_presets.json"), "r") as file: + bundled = {uuid_string(): StylePresetWithoutId.model_validate(preset) for preset in json.load(file)} + + def sync(q: Queries) -> None: + q.style_presets.delete_bundled() + q.style_presets.insert(bundled, SYSTEM_USER_ID) + + self._queries.run(sync) diff --git a/invokeai/app/services/style_preset_records/style_preset_records_sqlite.py b/invokeai/app/services/style_preset_records/style_preset_records_sqlite.py deleted file mode 100644 index 03397133ae9..00000000000 --- a/invokeai/app/services/style_preset_records/style_preset_records_sqlite.py +++ /dev/null @@ -1,188 +0,0 @@ -import json -from pathlib import Path - -from invokeai.app.services.invoker import Invoker -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.style_preset_records.style_preset_records_base import StylePresetRecordsStorageBase -from invokeai.app.services.style_preset_records.style_preset_records_common import ( - PresetType, - StylePresetChanges, - StylePresetNotFoundError, - StylePresetRecordDTO, - StylePresetWithoutId, -) -from invokeai.app.util.misc import uuid_string - -# System user id used for default / shipped presets and for legacy rows pre-dating -# the per-user ownership columns added in migration 27. -SYSTEM_USER_ID = "system" - - -class SqliteStylePresetRecordsStorage(StylePresetRecordsStorageBase): - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - - def start(self, invoker: Invoker) -> None: - self._invoker = invoker - self._sync_default_style_presets() - - def get(self, style_preset_id: str) -> StylePresetRecordDTO: - """Gets a style preset by ID.""" - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT * - FROM style_presets - WHERE id = ?; - """, - (style_preset_id,), - ) - row = cursor.fetchone() - if row is None: - raise StylePresetNotFoundError(f"Style preset with id {style_preset_id} not found") - return StylePresetRecordDTO.from_dict(dict(row)) - - def create(self, style_preset: StylePresetWithoutId, user_id: str) -> StylePresetRecordDTO: - style_preset_id = uuid_string() - with self._db.transaction() as cursor: - cursor.execute( - """--sql - INSERT OR IGNORE INTO style_presets ( - id, - name, - preset_data, - type, - user_id, - is_public - ) - VALUES (?, ?, ?, ?, ?, ?); - """, - ( - style_preset_id, - style_preset.name, - style_preset.preset_data.model_dump_json(), - style_preset.type, - user_id, - 1 if style_preset.is_public else 0, - ), - ) - return self.get(style_preset_id) - - def create_many(self, style_presets: list[StylePresetWithoutId], user_id: str) -> None: - with self._db.transaction() as cursor: - for style_preset in style_presets: - style_preset_id = uuid_string() - cursor.execute( - """--sql - INSERT OR IGNORE INTO style_presets ( - id, - name, - preset_data, - type, - user_id, - is_public - ) - VALUES (?, ?, ?, ?, ?, ?); - """, - ( - style_preset_id, - style_preset.name, - style_preset.preset_data.model_dump_json(), - style_preset.type, - user_id, - 1 if style_preset.is_public else 0, - ), - ) - - return None - - def update(self, style_preset_id: str, changes: StylePresetChanges) -> StylePresetRecordDTO: - with self._db.transaction() as cursor: - if changes.name is not None: - cursor.execute( - """--sql - UPDATE style_presets - SET name = ? - WHERE id = ?; - """, - (changes.name, style_preset_id), - ) - - if changes.preset_data is not None: - cursor.execute( - """--sql - UPDATE style_presets - SET preset_data = ? - WHERE id = ?; - """, - (changes.preset_data.model_dump_json(), style_preset_id), - ) - - if changes.is_public is not None: - cursor.execute( - """--sql - UPDATE style_presets - SET is_public = ? - WHERE id = ?; - """, - (1 if changes.is_public else 0, style_preset_id), - ) - - return self.get(style_preset_id) - - def delete(self, style_preset_id: str) -> None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - DELETE from style_presets - WHERE id = ?; - """, - (style_preset_id,), - ) - return None - - def get_many( - self, - type: PresetType | None = None, - user_id: str | None = None, - is_admin: bool = False, - ) -> list[StylePresetRecordDTO]: - clauses: list[str] = [] - params: list[object] = [] - - if not is_admin: - # Visible to non-admin: own + default + public. - visibility = "(type = 'default' OR is_public = 1" - if user_id is not None: - visibility += " OR user_id = ?" - params.append(user_id) - visibility += ")" - clauses.append(visibility) - - if type is not None: - clauses.append("type = ?") - params.append(type) - - where = f"WHERE {' AND '.join(clauses)} " if clauses else "" - query = f"SELECT * FROM style_presets {where}ORDER BY LOWER(name) ASC" - - with self._db.transaction() as cursor: - cursor.execute(query, params) - rows = cursor.fetchall() - return [StylePresetRecordDTO.from_dict(dict(row)) for row in rows] - - def _sync_default_style_presets(self) -> None: - """Syncs default style presets to the database. Internal use only.""" - with self._db.transaction() as cursor: - cursor.execute( - """--sql - DELETE FROM style_presets - WHERE type = "default"; - """ - ) - with open(Path(__file__).parent / Path("default_style_presets.json"), "r") as file: - presets = json.load(file) - for preset in presets: - style_preset = StylePresetWithoutId.model_validate(preset) - self.create(style_preset, user_id=SYSTEM_USER_ID) diff --git a/invokeai/app/services/system_prompt_records/system_prompt_records_default.py b/invokeai/app/services/system_prompt_records/system_prompt_records_default.py new file mode 100644 index 00000000000..ad883ebf3a7 --- /dev/null +++ b/invokeai/app/services/system_prompt_records/system_prompt_records_default.py @@ -0,0 +1,79 @@ +from typing import Optional + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.system_prompt_records.system_prompt_records_base import ( + SystemPromptRecordsStorageBase, +) +from invokeai.app.services.system_prompt_records.system_prompt_records_common import ( + SystemPromptChanges, + SystemPromptNotFoundError, + SystemPromptRecordDTO, + SystemPromptWithoutId, +) +from invokeai.app.util.misc import uuid_string + + +def _not_found(system_prompt_id: str) -> SystemPromptNotFoundError: + return SystemPromptNotFoundError(f"System prompt with id {system_prompt_id} not found") + + +def _found(prompt: Optional[SystemPromptRecordDTO], system_prompt_id: str) -> SystemPromptRecordDTO: + if prompt is None: + raise _not_found(system_prompt_id) + return prompt + + +class SystemPromptRecordsStorage(SystemPromptRecordsStorageBase): + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + + def get(self, system_prompt_id: str) -> SystemPromptRecordDTO: + return _found(self._queries.system_prompts.get(system_prompt_id), system_prompt_id) + + def create( + self, + system_prompt: SystemPromptWithoutId, + user_id: str, + is_public: bool = False, + ) -> SystemPromptRecordDTO: + system_prompt_id = uuid_string() + created_at = self._queries.system_prompts.insert(system_prompt_id, system_prompt, user_id, is_public) + # The record is what was stored, so it is not read back. + return SystemPromptRecordDTO( + **system_prompt.model_dump(), + id=system_prompt_id, + user_id=user_id, + is_public=is_public, + created_at=created_at, + updated_at=created_at, + ) + + def update( + self, + system_prompt_id: str, + changes: SystemPromptChanges, + user_id: Optional[str] = None, + ) -> SystemPromptRecordDTO: + def apply(q: Queries) -> Optional[SystemPromptRecordDTO]: + # A prompt that does not exist and, when scoped, one the account does not own are alike not found. + if not q.system_prompts.update(system_prompt_id, changes, user_id): + raise _not_found(system_prompt_id) + return q.system_prompts.get(system_prompt_id) + + return _found(self._queries.run(apply), system_prompt_id) + + def delete(self, system_prompt_id: str, user_id: Optional[str] = None) -> None: + """Delete a prompt, optionally scoped to an owner. + + Raises `SystemPromptNotFoundError` when nothing was deleted -- either the id does not + exist or (when scoped) it is not owned by `user_id`. Deleting silently would let the + single-user router report success for an id that `GET` 404s on, and would make the + symmetry with `update()` (which does raise) a trap for the next caller. + """ + if not self._queries.system_prompts.delete(system_prompt_id, user_id): + raise _not_found(system_prompt_id) + + def get_many(self, user_id: Optional[str] = None) -> list[SystemPromptRecordDTO]: + return self._queries.system_prompts.all(user_id) diff --git a/invokeai/app/services/system_prompt_records/system_prompt_records_sqlite.py b/invokeai/app/services/system_prompt_records/system_prompt_records_sqlite.py deleted file mode 100644 index e7f05887ddb..00000000000 --- a/invokeai/app/services/system_prompt_records/system_prompt_records_sqlite.py +++ /dev/null @@ -1,141 +0,0 @@ -from typing import Optional - -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.system_prompt_records.system_prompt_records_base import ( - SystemPromptRecordsStorageBase, -) -from invokeai.app.services.system_prompt_records.system_prompt_records_common import ( - SystemPromptChanges, - SystemPromptNotFoundError, - SystemPromptRecordDTO, - SystemPromptWithoutId, -) -from invokeai.app.util.misc import uuid_string - - -class SqliteSystemPromptRecordsStorage(SystemPromptRecordsStorageBase): - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - - def get(self, system_prompt_id: str) -> SystemPromptRecordDTO: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT * FROM system_prompts WHERE id = ?; - """, - (system_prompt_id,), - ) - row = cursor.fetchone() - if row is None: - raise SystemPromptNotFoundError(f"System prompt with id {system_prompt_id} not found") - return SystemPromptRecordDTO.from_dict(dict(row)) - - def create( - self, - system_prompt: SystemPromptWithoutId, - user_id: str, - is_public: bool = False, - ) -> SystemPromptRecordDTO: - system_prompt_id = uuid_string() - with self._db.transaction() as cursor: - cursor.execute( - """--sql - INSERT INTO system_prompts (id, name, content, user_id, is_public, max_tokens) - VALUES (?, ?, ?, ?, ?, ?); - """, - ( - system_prompt_id, - system_prompt.name, - system_prompt.content, - user_id, - is_public, - system_prompt.max_tokens, - ), - ) - return self.get(system_prompt_id) - - def update( - self, - system_prompt_id: str, - changes: SystemPromptChanges, - user_id: Optional[str] = None, - ) -> SystemPromptRecordDTO: - with self._db.transaction() as cursor: - # Confirm the row exists and (if scoped) is owned by the caller — distinguishes 404 from 403. - if user_id is not None: - cursor.execute( - "SELECT 1 FROM system_prompts WHERE id = ? AND user_id = ?;", - (system_prompt_id, user_id), - ) - else: - cursor.execute("SELECT 1 FROM system_prompts WHERE id = ?;", (system_prompt_id,)) - if cursor.fetchone() is None: - raise SystemPromptNotFoundError(f"System prompt with id {system_prompt_id} not found") - - scope_clause = " AND user_id = ?" if user_id is not None else "" - scope_args: tuple = (user_id,) if user_id is not None else () - - if changes.name is not None: - cursor.execute( - f"UPDATE system_prompts SET name = ? WHERE id = ?{scope_clause};", - (changes.name, system_prompt_id, *scope_args), - ) - if changes.content is not None: - cursor.execute( - f"UPDATE system_prompts SET content = ? WHERE id = ?{scope_clause};", - (changes.content, system_prompt_id, *scope_args), - ) - if changes.is_public is not None: - cursor.execute( - f"UPDATE system_prompts SET is_public = ? WHERE id = ?{scope_clause};", - (changes.is_public, system_prompt_id, *scope_args), - ) - # `max_tokens` is the one field whose null is a value rather than "no change": it means - # "drop back to the endpoint default". `is not None` would make a cap unclearable once - # set, so the presence of the key in the request body decides instead. - if "max_tokens" in changes.model_fields_set: - cursor.execute( - f"UPDATE system_prompts SET max_tokens = ? WHERE id = ?{scope_clause};", - (changes.max_tokens, system_prompt_id, *scope_args), - ) - return self.get(system_prompt_id) - - def delete(self, system_prompt_id: str, user_id: Optional[str] = None) -> None: - """Delete a prompt, optionally scoped to an owner. - - Raises `SystemPromptNotFoundError` when nothing was deleted -- either the id does not - exist or (when scoped) it is not owned by `user_id`. Deleting silently would let the - single-user router report success for an id that `GET` 404s on, and would make the - symmetry with `update()` (which does raise) a trap for the next caller. - """ - with self._db.transaction() as cursor: - if user_id is not None: - cursor.execute( - "DELETE FROM system_prompts WHERE id = ? AND user_id = ?;", - (system_prompt_id, user_id), - ) - else: - cursor.execute("DELETE FROM system_prompts WHERE id = ?;", (system_prompt_id,)) - if cursor.rowcount == 0: - raise SystemPromptNotFoundError(f"System prompt with id {system_prompt_id} not found") - - def get_many(self, user_id: Optional[str] = None) -> list[SystemPromptRecordDTO]: - with self._db.transaction() as cursor: - if user_id is not None: - cursor.execute( - """--sql - SELECT * FROM system_prompts - WHERE user_id = ? OR is_public = TRUE - ORDER BY LOWER(name) ASC; - """, - (user_id,), - ) - else: - cursor.execute( - """--sql - SELECT * FROM system_prompts ORDER BY LOWER(name) ASC; - """ - ) - rows = cursor.fetchall() - return [SystemPromptRecordDTO.from_dict(dict(row)) for row in rows] diff --git a/invokeai/app/services/users/users_common.py b/invokeai/app/services/users/users_common.py index 706597e52cf..eb0cb6a7f86 100644 --- a/invokeai/app/services/users/users_common.py +++ b/invokeai/app/services/users/users_common.py @@ -5,6 +5,9 @@ from pydantic import BaseModel, Field, field_validator from pydantic_core import PydanticCustomError +# The longest address SMTP carries (RFC 5321), and the longest email-validator accepts. +MAX_EMAIL_LENGTH = 254 + def validate_email_with_special_domains(email: str) -> str: """Validate email address, allowing special-use domains like .local for testing. @@ -55,6 +58,13 @@ def validate_email_with_special_domains(email: str) -> str: "Email address must have both local and domain parts", ) + # email-validator stops before its length check when it refuses the domain. + if len(email) > MAX_EMAIL_LENGTH: + raise PydanticCustomError( + "value_error", + f"Email address must be at most {MAX_EMAIL_LENGTH} characters", + ) + # Allow localhost and domains with dots if domain == "localhost" or "." in domain: return email diff --git a/invokeai/app/services/users/users_default.py b/invokeai/app/services/users/users_default.py index bc3ac08f33e..29094af9762 100644 --- a/invokeai/app/services/users/users_default.py +++ b/invokeai/app/services/users/users_default.py @@ -1,13 +1,14 @@ -"""Default SQLite implementation of user service.""" +"""Default implementation of the user service.""" -import sqlite3 from collections.abc import Sequence from datetime import datetime, timezone -from typing import Any from uuid import uuid4 from invokeai.app.services.auth.password_utils import hash_password, validate_password_strength, verify_password -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import UniqueViolation +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.locks import DatabaseLock from invokeai.app.services.users.users_base import UserServiceBase from invokeai.app.services.users.users_common import ( LAST_ADMIN_DETAIL, @@ -20,46 +21,12 @@ UserUpdateRequest, ) -# Bound-parameter chunk size for `get_many`. SQLite's SQLITE_MAX_VARIABLE_NUMBER is 32766 -# on modern builds but 999 on older ones; stay under the smaller limit. -USER_LOOKUP_CHUNK_SIZE = 900 - -# Every query that produces a `UserDTO` selects exactly these columns, in this order, and -# maps them with `_user_dto_from_row`. The projection and its mapper are kept together -# deliberately: they were previously copied out per query, and a field added to the DTO -# afterwards had to be added to each copy by hand. `get_many` was missed, so it reported -# `token_epoch` as the model default (0) for every user it returned, indistinguishable -# from an account that had never rotated its password. -_USER_DTO_COLUMNS = ( - "user_id, email, display_name, is_admin, is_active, created_at, updated_at, last_login_at, token_epoch" -) - - -def _user_dto_from_row(row: Sequence[Any]) -> UserDTO: - """Build a `UserDTO` from a row selected with `_USER_DTO_COLUMNS`.""" - return UserDTO( - user_id=row[0], - email=row[1], - display_name=row[2], - is_admin=bool(row[3]), - is_active=bool(row[4]), - created_at=datetime.fromisoformat(row[5]), - updated_at=datetime.fromisoformat(row[6]), - last_login_at=datetime.fromisoformat(row[7]) if row[7] else None, - token_epoch=int(row[8]), - ) - class UserService(UserServiceBase): - """SQLite-based user service.""" + """User service on the application database.""" - def __init__(self, db: SqliteDatabase): - """Initialize user service. - - Args: - db: SQLite database instance - """ - self._db = db + def __init__(self, database: Database): + self._queries = database.queries def create(self, user_data: UserCreateRequest, strict_password_checking: bool = True) -> UserDTO: """Create a new user.""" @@ -73,8 +40,8 @@ def _create( ) -> UserDTO: """Insert a user, optionally conditional on no administrator existing yet. - `require_no_admin` is evaluated on the cursor that performs the INSERT so that the check - and the write are one atomic step - see `create_admin` for why that matters. + `require_no_admin` is checked in the transaction that inserts, so that the check and the write are one + atomic step - see `create_admin` for why that matters. """ # Validate password strength if strict_password_checking: @@ -85,103 +52,50 @@ def _create( raise ValueError("Password cannot be empty") # Check if email already exists - if self.get_by_email(user_data.email) is not None: + if self._queries.users.get_by_email(user_data.email) is not None: raise ValueError(f"User with email {user_data.email} already exists") user_id = str(uuid4()) - # Hash before opening the transaction: hashing is deliberately slow, and the database is a - # single connection behind a process-wide lock. + # Hash before opening the transaction: hashing is deliberately slow, and on SQLite a transaction holds + # the database's process-wide lock. password_hash = hash_password(user_data.password) - with self._db.transaction() as cursor: - if require_no_admin: - # BEGIN IMMEDIATE takes the write lock up front, so this count cannot change - # before the INSERT below commits. In-process callers are additionally - # serialized by the database's shared RLock; the explicit lock also covers a - # second process (invoke-useradd --admin) writing during the setup window. - cursor.execute("BEGIN IMMEDIATE") - # Same predicate as has_admin(), read on the cursor that performs the write. - cursor.execute("SELECT COUNT(*) FROM users WHERE is_admin = TRUE AND is_active = TRUE") - row = cursor.fetchone() - if row and row[0] > 0: + def insert_user(q: Queries) -> UserDTO | None: + if user_data.is_admin: + # A setup that runs meanwhile waits here until this transaction ends, then counts its admin. + q.locks.acquire(DatabaseLock.ADMIN_ACCOUNTS) + if require_no_admin and q.users.count_active_admins() > 0: raise ValueError("Admin user already exists") try: - cursor.execute( - """ - INSERT INTO users (user_id, email, display_name, password_hash, is_admin) - VALUES (?, ?, ?, ?, ?) - """, - (user_id, user_data.email, user_data.display_name, password_hash, user_data.is_admin), + q.users.insert( + user_id=user_id, + email=user_data.email, + display_name=user_data.display_name, + password_hash=password_hash, + is_admin=user_data.is_admin, ) - except sqlite3.IntegrityError as e: + except UniqueViolation as e: raise ValueError(f"Failed to create user: {e}") from e + return q.users.get(user_id) - user = self.get(user_id) + user = self._queries.run(insert_user) if user is None: raise RuntimeError("Failed to retrieve created user") return user def get(self, user_id: str) -> UserDTO | None: """Get user by ID.""" - with self._db.transaction() as cursor: - cursor.execute( - f""" - SELECT {_USER_DTO_COLUMNS} - FROM users - WHERE user_id = ? - """, - (user_id,), - ) - row = cursor.fetchone() - - if row is None: - return None - - return _user_dto_from_row(row) + return self._queries.users.get(user_id) def get_many(self, user_ids: Sequence[str]) -> dict[str, UserDTO]: """Get users by ID, keyed by user_id. Unknown ids are absent from the result.""" - unique_ids = list(dict.fromkeys(user_ids)) - if not unique_ids: + if not user_ids: return {} - - users: dict[str, UserDTO] = {} - # Chunked so a caller with many distinct owners can't exceed SQLite's bound-parameter - # limit (999 on builds predating 3.32). - for start in range(0, len(unique_ids), USER_LOOKUP_CHUNK_SIZE): - chunk = unique_ids[start : start + USER_LOOKUP_CHUNK_SIZE] - placeholders = ",".join("?" for _ in chunk) - with self._db.transaction() as cursor: - cursor.execute( - f""" - SELECT {_USER_DTO_COLUMNS} - FROM users - WHERE user_id IN ({placeholders}) - """, - chunk, - ) - rows = cursor.fetchall() - for row in rows: - users[row[0]] = _user_dto_from_row(row) - return users + return self._queries.users.get_many(user_ids) def get_by_email(self, email: str) -> UserDTO | None: """Get user by email.""" - with self._db.transaction() as cursor: - cursor.execute( - f""" - SELECT {_USER_DTO_COLUMNS} - FROM users - WHERE email = ? - """, - (email,), - ) - row = cursor.fetchone() - - if row is None: - return None - - return _user_dto_from_row(row) + return self._queries.users.get_by_email(email) def update(self, user_id: str, changes: UserUpdateRequest, strict_password_checking: bool = True) -> UserDTO: """Update user.""" @@ -203,48 +117,28 @@ def update(self, user_id: str, changes: UserUpdateRequest, strict_password_check elif not changes.password: raise ValueError("Password cannot be empty") - # Build update query dynamically based on provided fields - updates: list[str] = [] - params: list[str | bool | int] = [] - - if changes.display_name is not None: - updates.append("display_name = ?") - params.append(changes.display_name) - - if changes.password is not None: - updates.append("password_hash = ?") - params.append(hash_password(changes.password)) - # Rotating the password revokes every token issued under the old one. A JWT is - # self-contained, so without this a stolen token survives the password change - # meant to evict the thief — and sliding-window refresh renews it indefinitely. - # Computed in SQL so concurrent bumps can't read-modify-write over each other. - updates.append("token_epoch = token_epoch + 1") - - if changes.is_admin is not None: - updates.append("is_admin = ?") - params.append(changes.is_admin) - - if changes.is_active is not None: - updates.append("is_active = ?") - params.append(changes.is_active) - - if not updates: + if (changes.display_name, changes.password, changes.is_admin, changes.is_active) == (None, None, None, None): return user - params.append(user_id) - query = f"UPDATE users SET {', '.join(updates)} WHERE user_id = ?" - - with self._db.transaction() as cursor: - # BEGIN IMMEDIATE takes the write lock up front, so the guard below reads a count - # no other connection can change before this transaction commits. Without it the - # SELECT would run in autocommit and a second process (invoke-usermod) could slip - # its own write in between. In-process callers are additionally serialized by the - # database's shared RLock. - cursor.execute("BEGIN IMMEDIATE") - self._assert_not_last_admin(cursor, user_id, is_admin=changes.is_admin, is_active=changes.is_active) - cursor.execute(query, params) + # Rotating the password revokes every token issued under the old one (the query raises the token + # epoch). A JWT is self-contained, so without that a stolen token survives the password change meant + # to evict the thief - and sliding-window refresh renews it indefinitely. + password_hash = hash_password(changes.password) if changes.password is not None else None + + def apply_changes(q: Queries) -> UserDTO | None: + if changes.is_admin is not None or changes.is_active is not None: + q.locks.acquire(DatabaseLock.ADMIN_ACCOUNTS) + self._assert_not_last_admin(q, user_id, is_admin=changes.is_admin, is_active=changes.is_active) + q.users.update( + user_id, + display_name=changes.display_name, + password_hash=password_hash, + is_admin=changes.is_admin, + is_active=changes.is_active, + ) + return q.users.get(user_id) - updated_user = self.get(user_id) + updated_user = self._queries.run(apply_changes) if updated_user is None: raise RuntimeError("Failed to retrieve updated user") return updated_user @@ -257,37 +151,23 @@ def delete(self, user_id: str) -> None: self._assert_system_user_protected(user_id, is_deleting=True) - with self._db.transaction() as cursor: - # See the note in `update`: the guard and the write share one write-locked - # transaction so the admin count cannot change between them. - cursor.execute("BEGIN IMMEDIATE") - self._assert_not_last_admin(cursor, user_id, is_deleting=True) - cursor.execute("DELETE FROM users WHERE user_id = ?", (user_id,)) + def delete_user(q: Queries) -> None: + q.locks.acquire(DatabaseLock.ADMIN_ACCOUNTS) + self._assert_not_last_admin(q, user_id, is_deleting=True) + q.users.delete(user_id) # Projects and client state cascade with the account; the media references they held # would otherwise protect a departed account's intermediates from every cleanup forever. # Workflows and quarantined projects do not cascade, so their references stay with them. - cursor.execute( - "DELETE FROM media_references WHERE user_id = ? AND owner_kind IN ('project', 'client_state')", - (user_id,), - ) + q.media_references.delete_owned_by(user_id, "project", "client_state") + + self._queries.run(delete_user) def authenticate(self, email: str, password: str) -> UserDTO | None: """Authenticate user credentials.""" - with self._db.transaction() as cursor: - # `password_hash` is appended rather than inserted, so the DTO columns keep the - # positions `_user_dto_from_row` expects. - cursor.execute( - f""" - SELECT {_USER_DTO_COLUMNS}, password_hash - FROM users - WHERE email = ? - """, - (email,), - ) - row = cursor.fetchone() - - if row is None: + credentials = self._queries.users.get_credentials(email) + if credentials is None: return None + user, password_hash = credentials # The system account is not a login account. It owns everything migrated from # before multiuser support, so a token bearing `user_id="system"` reads and writes @@ -296,31 +176,21 @@ def authenticate(self, email: str, password: str) -> UserDTO | None: # usable hash, and the migration that clears it cannot reach a row damaged by # direct SQL afterwards. Refusing here makes "system cannot authenticate" hold # regardless of what the row contains. - if row[0] == SYSTEM_USER_ID: + if user.user_id == SYSTEM_USER_ID: return None - password_hash = row[9] if not verify_password(password, password_hash): return None - # Update last login time last_login_at = datetime.now(timezone.utc) - with self._db.transaction() as cursor: - cursor.execute( - "UPDATE users SET last_login_at = ? WHERE user_id = ?", - (last_login_at.isoformat(), row[0]), - ) + self._queries.users.record_login(user.user_id, last_login_at) # Report the login just recorded rather than the one the row was read with. - return _user_dto_from_row(row).model_copy(update={"last_login_at": last_login_at}) + return user.model_copy(update={"last_login_at": last_login_at}) def has_admin(self) -> bool: """Check if any admin user exists.""" - with self._db.transaction() as cursor: - cursor.execute("SELECT COUNT(*) FROM users WHERE is_admin = TRUE AND is_active = TRUE") - row = cursor.fetchone() - count = row[0] if row else 0 - return bool(count > 0) + return self._queries.users.count_active_admins() > 0 def create_admin(self, user_data: UserCreateRequest, strict_password_checking: bool = True) -> UserDTO: """Create the first admin user (for initial setup). @@ -342,44 +212,19 @@ def create_admin(self, user_data: UserCreateRequest, strict_password_checking: b def list_users(self, limit: int = 100, offset: int = 0) -> list[UserDTO]: """List all users.""" - with self._db.transaction() as cursor: - cursor.execute( - f""" - SELECT {_USER_DTO_COLUMNS} - FROM users - ORDER BY created_at DESC - LIMIT ? OFFSET ? - """, - (limit, offset), - ) - rows = cursor.fetchall() - - return [_user_dto_from_row(row) for row in rows] + return self._queries.users.page(limit, offset) def get_admin_email(self) -> str | None: """Get the email address of the first active admin user.""" - with self._db.transaction() as cursor: - cursor.execute( - """ - SELECT email FROM users - WHERE is_admin = TRUE AND is_active = TRUE - ORDER BY created_at ASC - LIMIT 1 - """, - ) - row = cursor.fetchone() - return row[0] if row else None + return self._queries.users.first_active_admin_email() def count_admins(self) -> int: """Count active admin users.""" - with self._db.transaction() as cursor: - cursor.execute("SELECT COUNT(*) FROM users WHERE is_admin = TRUE AND is_active = TRUE") - row = cursor.fetchone() - return int(row[0]) if row else 0 + return self._queries.users.count_active_admins() def _assert_not_last_admin( self, - cursor: sqlite3.Cursor, + q: Queries, user_id: str, *, is_deleting: bool = False, @@ -388,31 +233,28 @@ def _assert_not_last_admin( ) -> None: """Reject a change that would drop the number of active administrators to zero. - Must be called on the cursor of the transaction that performs the write, after that - transaction has taken its write lock — see the ``BEGIN IMMEDIATE`` in the callers. - Reading the count in a separate transaction is what made this a TOCTOU: two callers - could each observe two admins and each proceed to remove one. + Called in the transaction that performs the write, which holds the `ADMIN_ACCOUNTS` lock: every change + to who is an active administrator takes it first. Reading the count in a separate transaction is what + made this a TOCTOU: two callers could each observe two admins and each proceed to remove one. On SQLite + the transaction alone serializes them; on a server the lock does. ``is_admin``/``is_active`` are the *requested* values, where ``None`` means "not being changed" — deletion is signalled separately by ``is_deleting`` rather than by both being ``None``, which would also describe a rename. Only a change that actually revokes administrator status is checked, so renaming the last admin stays allowed. """ - cursor.execute("SELECT is_admin, is_active FROM users WHERE user_id = ?", (user_id,)) - row = cursor.fetchone() - if row is None: + user = q.users.get(user_id) + if user is None: return # An admin who is already inactive is not counted, so removing them changes nothing. - if not (bool(row[0]) and bool(row[1])): + if not (user.is_admin and user.is_active): return if not (is_deleting or is_admin is False or is_active is False): return - cursor.execute("SELECT COUNT(*) FROM users WHERE is_admin = TRUE AND is_active = TRUE") - count_row = cursor.fetchone() - if (int(count_row[0]) if count_row else 0) <= 1: + if q.users.count_active_admins() <= 1: raise LastAdministratorError(LAST_ADMIN_DETAIL) def _assert_system_user_protected( diff --git a/invokeai/app/services/video_records/video_records_base.py b/invokeai/app/services/video_records/video_records_base.py index fd3dd0d067f..5be03be4e01 100644 --- a/invokeai/app/services/video_records/video_records_base.py +++ b/invokeai/app/services/video_records/video_records_base.py @@ -5,8 +5,7 @@ from invokeai.app.invocations.fields import MetadataField from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin from invokeai.app.services.shared.intermediate_delete import IntermediateDeleteGuard -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection from invokeai.app.services.video_records.video_records_common import ( VideoNamesResult, VideoRecord, diff --git a/invokeai/app/services/video_records/video_records_common.py b/invokeai/app/services/video_records/video_records_common.py index 5c8afcabf25..0e17480aa39 100644 --- a/invokeai/app/services/video_records/video_records_common.py +++ b/invokeai/app/services/video_records/video_records_common.py @@ -26,26 +26,6 @@ def __init__(self, message="Video record not saved"): super().__init__(message) -class VideoRecordDeleteException(Exception): - """Raised when a video record cannot be deleted.""" - - def __init__(self, message="Video record not deleted"): - super().__init__(message) - - -# The `media_origin` marker, projected out of the `metadata` JSON blob. Kept as a bare -# expression so the polymorphic gallery query can alias it into its own UNION half. -# -# The `json_valid` guard is not decoration: `json_extract` RAISES on unparseable text, and -# this expression now runs on every row of every video listing. An unguarded call would let -# a single malformed blob fail the whole page rather than one video -- and the column is -# plain TEXT with no CHECK constraint, so nothing but convention keeps one out. Every -# in-tree writer goes through `MetadataField`, so this is insurance, not a live bug. -MEDIA_ORIGIN_JSON_EXPR = ( - "CASE WHEN json_valid(videos.metadata) THEN json_extract(videos.metadata, '$.media_origin') END" -) -MEDIA_ORIGIN_SQL_EXPR = f"{MEDIA_ORIGIN_JSON_EXPR} AS media_origin" - # The longest marker worth carrying. Only `audio_upload` has meaning today, but the field is # an open vocabulary, so this is a sanity bound rather than an allowlist. It matters because # upload metadata is client-supplied and unbounded, and this one key now rides EVERY row of @@ -80,39 +60,6 @@ def coerce_media_origin(value: Any) -> Optional[str]: return value if MEDIA_ORIGIN_PATTERN.match(value) else None -VIDEO_DTO_COLS = ", ".join( - [ - "videos." + c - for c in [ - "video_name", - "video_origin", - "video_category", - "width", - "height", - "duration", - "fps", - "session_id", - "node_id", - "has_workflow", - "is_intermediate", - "created_at", - "updated_at", - "deleted_at", - "starred", - "video_subfolder", - "project_id", - "file_size_bytes", - ] - ] - # `media_origin` is not a column: it is the one key of the `metadata` JSON blob the - # frontend needs on every row (it marks an upload the ingest converter wrapped from an - # audio file). Extracting just that key keeps listings from carrying whole metadata - # blobs. `json_extract` yields NULL for a NULL or non-object blob; what it yields for a - # non-string value is `coerce_media_origin`'s problem, not the query's. - + [MEDIA_ORIGIN_SQL_EXPR] -) - - class VideoRecord(BaseModelExcludeNull): """Deserialized video record: the columns, plus the one `media_origin` key projected out of metadata.""" diff --git a/invokeai/app/services/video_records/video_records_default.py b/invokeai/app/services/video_records/video_records_default.py new file mode 100644 index 00000000000..4ec0a7a3b8a --- /dev/null +++ b/invokeai/app/services/video_records/video_records_default.py @@ -0,0 +1,180 @@ +from datetime import datetime +from typing import Optional + +from invokeai.app.invocations.fields import MetadataField, MetadataFieldValidator +from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.locks import DatabaseLock +from invokeai.app.services.shared.intermediate_delete import IntermediateDeleteGuard +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection +from invokeai.app.services.video_records.video_records_base import VideoRecordStorageBase +from invokeai.app.services.video_records.video_records_common import ( + VideoNamesResult, + VideoRecord, + VideoRecordChanges, + VideoRecordNotFoundException, + VideoRecordSaveException, +) + + +class VideoRecordStorage(VideoRecordStorageBase): + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + + def get(self, video_name: str) -> VideoRecord: + # A database error is deliberately NOT turned into VideoRecordNotFoundException: it would make the exception + # mean "the row is absent, OR the database is unreadable". `_assert_video_read_access` answers 404 on a + # not-found, and the clients drop their references to a video on that 404 -- so a locked database would + # silently clear the user's workflow fields. + record = self._queries.videos.get(video_name) + if record is None: + raise VideoRecordNotFoundException + return record + + def get_subfolders(self, video_names: list[str]) -> dict[str, str]: + return self._queries.videos.subfolders(video_names) + + def delete_intermediates_by_names( + self, video_names: list[str], guard: Optional[IntermediateDeleteGuard] = None + ) -> list[str]: + def delete(q: Queries) -> list[str]: + if guard is None: + return q.videos.delete_intermediates(video_names) + # Held until the delete commits, so nothing becomes protected between the guard's answer and the delete. + q.locks.acquire(DatabaseLock.MEDIA_PROTECTION) + return q.videos.delete_intermediates(guard(q, video_names)) + + return self._queries.run(delete) + + def set_file_size_bytes(self, video_name: str, file_size_bytes: Optional[int]) -> None: + self._queries.videos.set_file_size(video_name, file_size_bytes) + + def set_file_sizes_bytes(self, sizes: dict[str, int]) -> None: + self._queries.videos.fill_file_sizes(sizes) + + def get_user_id(self, video_name: str) -> Optional[str]: + return self._queries.videos.user_id(video_name) + + def get_most_recent_video_for_board(self, board_id: str) -> Optional[VideoRecord]: + return self._queries.videos.most_recent_on_board(board_id) + + def exists(self, video_name: str) -> bool: + return self._queries.videos.exists(video_name) + + def get_metadata(self, video_name: str) -> Optional[MetadataField]: + # See get(): a database error must not masquerade as a missing row. + exists, metadata = self._queries.videos.metadata(video_name) + if not exists: + raise VideoRecordNotFoundException + return MetadataFieldValidator.validate_json(metadata) if metadata is not None else None + + def update(self, video_name: str, changes: VideoRecordChanges) -> None: + self._queries.videos.update(video_name, changes) + + def get_many( + self, + offset: int = 0, + limit: int = 10, + starred_first: bool = True, + order_dir: SQLiteDirection = SQLiteDirection.Descending, + video_origin: Optional[ResourceOrigin] = None, + categories: Optional[list[ImageCategory]] = None, + is_intermediate: Optional[bool] = None, + board_id: Optional[str] = None, + search_term: Optional[str] = None, + user_id: Optional[str] = None, + is_admin: bool = False, + ) -> OffsetPaginatedResults[VideoRecord]: + videos, total = self._queries.videos.page( + offset=offset, + limit=limit, + starred_first=starred_first, + descending=order_dir == SQLiteDirection.Descending, + video_origin=video_origin, + categories=categories, + is_intermediate=is_intermediate, + board_id=board_id, + search_term=search_term, + user_id=user_id, + is_admin=is_admin, + ) + return OffsetPaginatedResults(items=videos, offset=offset, limit=limit, total=total) + + def delete(self, video_name: str) -> None: + self._queries.videos.delete(video_name) + + def delete_many(self, video_names: list[str]) -> None: + self._queries.videos.delete_many(video_names) + + def save( + self, + video_name: str, + video_origin: ResourceOrigin, + video_category: ImageCategory, + width: int, + height: int, + duration: float, + fps: Optional[float], + has_workflow: bool, + is_intermediate: Optional[bool] = False, + starred: Optional[bool] = False, + session_id: Optional[str] = None, + node_id: Optional[str] = None, + metadata: Optional[str] = None, + user_id: Optional[str] = None, + video_subfolder: str = "", + project_id: Optional[str] = None, + ) -> datetime: + def insert(q: Queries) -> Optional[str]: + q.videos.insert( + video_name=video_name, + video_origin=video_origin, + video_category=video_category, + width=width, + height=height, + duration=duration, + fps=fps, + has_workflow=has_workflow, + is_intermediate=is_intermediate, + starred=starred, + session_id=session_id, + node_id=node_id, + metadata=metadata, + user_id=user_id or "system", + video_subfolder=video_subfolder, + project_id=project_id, + ) + # Read back: a video of that name that exists already keeps its row, and its time is the one returned. + return q.videos.created_at(video_name) + + created_at = self._queries.run(insert) + if created_at is None: + raise VideoRecordSaveException + return datetime.fromisoformat(created_at) + + def get_video_names( + self, + starred_first: bool = True, + order_dir: SQLiteDirection = SQLiteDirection.Descending, + video_origin: Optional[ResourceOrigin] = None, + categories: Optional[list[ImageCategory]] = None, + is_intermediate: Optional[bool] = None, + board_id: Optional[str] = None, + search_term: Optional[str] = None, + user_id: Optional[str] = None, + is_admin: bool = False, + ) -> VideoNamesResult: + video_names, starred_count = self._queries.videos.names( + starred_first=starred_first, + descending=order_dir == SQLiteDirection.Descending, + video_origin=video_origin, + categories=categories, + is_intermediate=is_intermediate, + board_id=board_id, + search_term=search_term, + user_id=user_id, + is_admin=is_admin, + ) + return VideoNamesResult(video_names=video_names, starred_count=starred_count, total_count=len(video_names)) diff --git a/invokeai/app/services/video_records/video_records_sqlite.py b/invokeai/app/services/video_records/video_records_sqlite.py deleted file mode 100644 index d5c6df6801b..00000000000 --- a/invokeai/app/services/video_records/video_records_sqlite.py +++ /dev/null @@ -1,463 +0,0 @@ -import sqlite3 -from datetime import datetime -from typing import Optional, Union, cast - -from invokeai.app.invocations.fields import MetadataField, MetadataFieldValidator -from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin -from invokeai.app.services.shared.intermediate_delete import IntermediateDeleteGuard -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.video_records.video_records_base import VideoRecordStorageBase -from invokeai.app.services.video_records.video_records_common import ( - VIDEO_DTO_COLS, - VideoNamesResult, - VideoRecord, - VideoRecordChanges, - VideoRecordDeleteException, - VideoRecordNotFoundException, - VideoRecordSaveException, - deserialize_video_record, -) - - -class SqliteVideoRecordStorage(VideoRecordStorageBase): - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - - def get(self, video_name: str) -> VideoRecord: - # A sqlite3.Error is deliberately NOT translated into VideoRecordNotFoundException, for - # the same reason as SqliteImageRecordStorage.get: it would make the exception mean "the - # row is absent, OR the database is unreadable". `_assert_video_read_access` now answers - # 404 on a positive not-found, and the clients drop their references to a video on that - # 404 — so a locked database would silently clear the user's workflow fields. - with self._db.transaction() as cursor: - cursor.execute( - f"""--sql - SELECT {VIDEO_DTO_COLS} FROM videos - WHERE video_name = ?; - """, - (video_name,), - ) - result = cast(Optional[sqlite3.Row], cursor.fetchone()) - - if not result: - raise VideoRecordNotFoundException - return deserialize_video_record(dict(result)) - - _MAX_SQL_VARIABLES = 500 - - def get_subfolders(self, video_names: list[str]) -> dict[str, str]: - subfolders: dict[str, str] = {} - with self._db.transaction() as cursor: - for start in range(0, len(video_names), self._MAX_SQL_VARIABLES): - chunk = video_names[start : start + self._MAX_SQL_VARIABLES] - placeholders = ",".join("?" for _ in chunk) - cursor.execute( - f"SELECT video_name, video_subfolder FROM videos WHERE video_name IN ({placeholders})", chunk - ) - for row in cursor.fetchall(): - subfolders[cast(str, row[0])] = cast(str, row[1]) - return subfolders - - def delete_intermediates_by_names( - self, video_names: list[str], guard: Optional[IntermediateDeleteGuard] = None - ) -> list[str]: - deleted: list[str] = [] - try: - with self._db.transaction() as cursor: - for start in range(0, len(video_names), self._MAX_SQL_VARIABLES): - chunk = video_names[start : start + self._MAX_SQL_VARIABLES] - if guard is not None: - chunk = guard(cursor, chunk) - if not chunk: - continue - placeholders = ",".join("?" for _ in chunk) - select_query = f"SELECT video_name FROM videos WHERE video_name IN ({placeholders})" - cursor.execute(select_query, chunk) - present_before = {cast(str, r[0]) for r in cursor.fetchall()} - cursor.execute( - f"DELETE FROM videos WHERE video_name IN ({placeholders}) AND is_intermediate = TRUE", - chunk, - ) - cursor.execute(select_query, chunk) - present_after = {cast(str, r[0]) for r in cursor.fetchall()} - deleted.extend(name for name in chunk if name in present_before and name not in present_after) - except sqlite3.Error as e: - raise VideoRecordDeleteException from e - return deleted - - def set_file_size_bytes(self, video_name: str, file_size_bytes: Optional[int]) -> None: - with self._db.transaction() as cursor: - try: - cursor.execute( - "UPDATE videos SET file_size_bytes = ? WHERE video_name = ?;", - (file_size_bytes, video_name), - ) - except sqlite3.Error as e: - raise VideoRecordSaveException from e - - def set_file_sizes_bytes(self, sizes: dict[str, int]) -> None: - if not sizes: - return - with self._db.transaction() as cursor: - try: - # Backfill only fills gaps: the writer's own measurement, taken after the file exists, - # wins over one taken before it was written. - cursor.executemany( - "UPDATE videos SET file_size_bytes = ? WHERE video_name = ? AND file_size_bytes IS NULL;", - [(size, name) for name, size in sizes.items()], - ) - except sqlite3.Error as e: - raise VideoRecordSaveException from e - - def get_user_id(self, video_name: str) -> Optional[str]: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT user_id FROM videos - WHERE video_name = ?; - """, - (video_name,), - ) - result = cast(Optional[sqlite3.Row], cursor.fetchone()) - if not result: - return None - return cast(Optional[str], dict(result).get("user_id")) - - def get_most_recent_video_for_board(self, board_id: str) -> Optional[VideoRecord]: - with self._db.transaction() as cursor: - cursor.execute( - f"""--sql - SELECT {VIDEO_DTO_COLS} - FROM videos - JOIN board_videos ON videos.video_name = board_videos.video_name - WHERE board_videos.board_id = ? - AND videos.is_intermediate = FALSE - ORDER BY videos.starred DESC, videos.created_at DESC, videos.video_name DESC - LIMIT 1; - """, - (board_id,), - ) - result = cast(Optional[sqlite3.Row], cursor.fetchone()) - if result is None: - return None - return deserialize_video_record(dict(result)) - - def exists(self, video_name: str) -> bool: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT 1 FROM videos - WHERE video_name = ?; - """, - (video_name,), - ) - return cursor.fetchone() is not None - - def get_metadata(self, video_name: str) -> Optional[MetadataField]: - with self._db.transaction() as cursor: - try: - cursor.execute( - """--sql - SELECT metadata FROM videos - WHERE video_name = ?; - """, - (video_name,), - ) - result = cast(Optional[sqlite3.Row], cursor.fetchone()) - except sqlite3.Error as e: - raise VideoRecordNotFoundException from e - - if not result: - raise VideoRecordNotFoundException - - as_dict = dict(result) - metadata_raw = cast(Optional[str], as_dict.get("metadata", None)) - return MetadataFieldValidator.validate_json(metadata_raw) if metadata_raw is not None else None - - def update(self, video_name: str, changes: VideoRecordChanges) -> None: - with self._db.transaction() as cursor: - try: - if changes.video_category is not None: - cursor.execute( - "UPDATE videos SET video_category = ? WHERE video_name = ?;", - (changes.video_category.value, video_name), - ) - if changes.session_id is not None: - cursor.execute( - "UPDATE videos SET session_id = ? WHERE video_name = ?;", - (changes.session_id, video_name), - ) - if changes.is_intermediate is not None: - cursor.execute( - "UPDATE videos SET is_intermediate = ? WHERE video_name = ?;", - (changes.is_intermediate, video_name), - ) - if changes.starred is not None: - cursor.execute( - "UPDATE videos SET starred = ? WHERE video_name = ?;", - (changes.starred, video_name), - ) - except sqlite3.Error as e: - raise VideoRecordSaveException from e - - def get_many( - self, - offset: int = 0, - limit: int = 10, - starred_first: bool = True, - order_dir: SQLiteDirection = SQLiteDirection.Descending, - video_origin: Optional[ResourceOrigin] = None, - categories: Optional[list[ImageCategory]] = None, - is_intermediate: Optional[bool] = None, - board_id: Optional[str] = None, - search_term: Optional[str] = None, - user_id: Optional[str] = None, - is_admin: bool = False, - ) -> OffsetPaginatedResults[VideoRecord]: - with self._db.transaction() as cursor: - count_query = """--sql - SELECT COUNT(*) - FROM videos - LEFT JOIN board_videos ON board_videos.video_name = videos.video_name - WHERE 1=1 - """ - videos_query = f"""--sql - SELECT {VIDEO_DTO_COLS} - FROM videos - LEFT JOIN board_videos ON board_videos.video_name = videos.video_name - WHERE 1=1 - """ - - query_conditions = "" - query_params: list[Union[int, str, bool]] = [] - - if video_origin is not None: - query_conditions += " AND videos.video_origin = ? " - query_params.append(video_origin.value) - - if categories is not None: - category_strings = [c.value for c in set(categories)] - placeholders = ",".join("?" * len(category_strings)) - query_conditions += f" AND videos.video_category IN ( {placeholders} ) " - for c in category_strings: - query_params.append(c) - - if is_intermediate is not None: - query_conditions += " AND videos.is_intermediate = ? " - query_params.append(is_intermediate) - - if board_id == "none": - query_conditions += " AND board_videos.board_id IS NULL " - if user_id is not None and not is_admin: - query_conditions += " AND videos.user_id = ? " - query_params.append(user_id) - elif board_id is not None: - query_conditions += " AND board_videos.board_id = ? " - query_params.append(board_id) - elif user_id is not None and not is_admin: - # No board_id supplied — still enforce per-user isolation so - # non-admin callers cannot enumerate other users' videos. - query_conditions += " AND videos.user_id = ? " - query_params.append(user_id) - - if search_term: - query_conditions += " AND (videos.metadata LIKE ? OR videos.created_at LIKE ?) " - query_params.append(f"%{search_term.lower()}%") - query_params.append(f"%{search_term.lower()}%") - - if starred_first: - query_pagination = ( - f" ORDER BY videos.starred DESC, videos.created_at {order_dir.value}, " - f"videos.video_name {order_dir.value} LIMIT ? OFFSET ? " - ) - else: - query_pagination = ( - f" ORDER BY videos.created_at {order_dir.value}, videos.video_name {order_dir.value} " - "LIMIT ? OFFSET ? " - ) - - videos_query += query_conditions + query_pagination + ";" - videos_params = query_params.copy() - videos_params.extend([limit, offset]) - cursor.execute(videos_query, videos_params) - result = cast(list[sqlite3.Row], cursor.fetchall()) - videos = [deserialize_video_record(dict(r)) for r in result] - - count_query += query_conditions + ";" - cursor.execute(count_query, query_params.copy()) - count = cast(int, cursor.fetchone()[0]) - - return OffsetPaginatedResults(items=videos, offset=offset, limit=limit, total=count) - - def delete(self, video_name: str) -> None: - with self._db.transaction() as cursor: - try: - cursor.execute("DELETE FROM videos WHERE video_name = ?;", (video_name,)) - except sqlite3.Error as e: - raise VideoRecordDeleteException from e - - def delete_many(self, video_names: list[str]) -> None: - with self._db.transaction() as cursor: - try: - placeholders = ",".join("?" for _ in video_names) - cursor.execute(f"DELETE FROM videos WHERE video_name IN ({placeholders})", video_names) - except sqlite3.Error as e: - raise VideoRecordDeleteException from e - - def save( - self, - video_name: str, - video_origin: ResourceOrigin, - video_category: ImageCategory, - width: int, - height: int, - duration: float, - fps: Optional[float], - has_workflow: bool, - is_intermediate: Optional[bool] = False, - starred: Optional[bool] = False, - session_id: Optional[str] = None, - node_id: Optional[str] = None, - metadata: Optional[str] = None, - user_id: Optional[str] = None, - video_subfolder: str = "", - project_id: Optional[str] = None, - ) -> datetime: - with self._db.transaction() as cursor: - try: - cursor.execute( - """--sql - INSERT OR IGNORE INTO videos ( - video_name, - video_origin, - video_category, - width, - height, - duration, - fps, - node_id, - session_id, - metadata, - is_intermediate, - starred, - has_workflow, - user_id, - video_subfolder, - project_id - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?); - """, - ( - video_name, - video_origin.value, - video_category.value, - width, - height, - float(duration), - float(fps) if fps is not None else None, - node_id, - session_id, - metadata, - is_intermediate, - starred, - has_workflow, - user_id or "system", - video_subfolder, - project_id, - ), - ) - - cursor.execute( - "SELECT created_at FROM videos WHERE video_name = ?;", - (video_name,), - ) - created_at = datetime.fromisoformat(cursor.fetchone()[0]) - except sqlite3.Error as e: - raise VideoRecordSaveException from e - return created_at - - def get_video_names( - self, - starred_first: bool = True, - order_dir: SQLiteDirection = SQLiteDirection.Descending, - video_origin: Optional[ResourceOrigin] = None, - categories: Optional[list[ImageCategory]] = None, - is_intermediate: Optional[bool] = None, - board_id: Optional[str] = None, - search_term: Optional[str] = None, - user_id: Optional[str] = None, - is_admin: bool = False, - ) -> VideoNamesResult: - with self._db.transaction() as cursor: - query_conditions = "" - query_params: list[Union[int, str, bool]] = [] - - if video_origin is not None: - query_conditions += " AND videos.video_origin = ? " - query_params.append(video_origin.value) - - if categories is not None: - category_strings = [c.value for c in set(categories)] - placeholders = ",".join("?" * len(category_strings)) - query_conditions += f" AND videos.video_category IN ( {placeholders} ) " - for c in category_strings: - query_params.append(c) - - if is_intermediate is not None: - query_conditions += " AND videos.is_intermediate = ? " - query_params.append(is_intermediate) - - if board_id == "none": - query_conditions += " AND board_videos.board_id IS NULL " - if user_id is not None and not is_admin: - query_conditions += " AND videos.user_id = ? " - query_params.append(user_id) - elif board_id is not None: - query_conditions += " AND board_videos.board_id = ? " - query_params.append(board_id) - elif user_id is not None and not is_admin: - # No board_id supplied — still enforce per-user isolation so - # non-admin callers cannot enumerate other users' videos. - query_conditions += " AND videos.user_id = ? " - query_params.append(user_id) - - if search_term: - query_conditions += " AND (videos.metadata LIKE ? OR videos.created_at LIKE ?) " - query_params.append(f"%{search_term.lower()}%") - query_params.append(f"%{search_term.lower()}%") - - starred_count = 0 - if starred_first: - cursor.execute( - f"""--sql - SELECT COUNT(*) - FROM videos - LEFT JOIN board_videos ON board_videos.video_name = videos.video_name - WHERE videos.starred = TRUE AND (1=1{query_conditions}) - """, - query_params, - ) - starred_count = cast(int, cursor.fetchone()[0]) - - order_clause = ( - f" ORDER BY videos.starred DESC, videos.created_at {order_dir.value}, " - f"videos.video_name {order_dir.value} " - if starred_first - else f" ORDER BY videos.created_at {order_dir.value}, videos.video_name {order_dir.value} " - ) - cursor.execute( - f"""--sql - SELECT videos.video_name - FROM videos - LEFT JOIN board_videos ON board_videos.video_name = videos.video_name - WHERE 1=1{query_conditions} - {order_clause} - """, - query_params, - ) - result = cast(list[sqlite3.Row], cursor.fetchall()) - video_names = [row[0] for row in result] - return VideoNamesResult(video_names=video_names, starred_count=starred_count, total_count=len(video_names)) diff --git a/invokeai/app/services/videos/videos_base.py b/invokeai/app/services/videos/videos_base.py index 4bf0b5d2a5b..62f069936c5 100644 --- a/invokeai/app/services/videos/videos_base.py +++ b/invokeai/app/services/videos/videos_base.py @@ -7,8 +7,7 @@ from invokeai.app.invocations.fields import MetadataField from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin from invokeai.app.services.shared.intermediate_delete import IntermediateDeleteGuard, IntermediateDeleteResult -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection from invokeai.app.services.video_records.video_records_common import ( VideoNamesResult, VideoRecord, diff --git a/invokeai/app/services/videos/videos_default.py b/invokeai/app/services/videos/videos_default.py index b24cd772e94..635d878e3ad 100644 --- a/invokeai/app/services/videos/videos_default.py +++ b/invokeai/app/services/videos/videos_default.py @@ -19,8 +19,7 @@ JournaledDeleteAdapter, delete_journaled_intermediates, ) -from invokeai.app.services.shared.pagination import OffsetPaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import OffsetPaginatedResults, SQLiteDirection from invokeai.app.services.video_files.video_files_common import ( VideoFileDeleteException, VideoFileNotFoundException, @@ -30,7 +29,6 @@ VideoNamesResult, VideoRecord, VideoRecordChanges, - VideoRecordDeleteException, VideoRecordNotFoundException, VideoRecordSaveException, ) @@ -386,11 +384,6 @@ def delete(self, video_name: str) -> None: except Exception as cleanup_error: self.__invoker.services.logger.error(f"Failed to purge staged video files: {cleanup_error}") self._on_deleted(video_name) - except VideoRecordDeleteException: - if token is not None: - self.__invoker.services.video_files.rollback_delete(token) - self.__invoker.services.logger.error("Failed to delete video record") - raise except VideoFileDeleteException: self.__invoker.services.logger.error("Failed to delete video file") raise @@ -456,9 +449,6 @@ def delete_videos_by_names(self, video_names: list[str]) -> tuple[list[str], lis log_error=self.__invoker.services.logger.error, ), ) - except VideoRecordDeleteException: - self.__invoker.services.logger.error("Failed to delete video records") - raise except VideoFileDeleteException: self.__invoker.services.logger.error("Failed to delete video files") raise diff --git a/invokeai/app/services/wildcard_records/wildcard_records_common.py b/invokeai/app/services/wildcard_records/wildcard_records_common.py index c182f0dd1b5..41b74302254 100644 --- a/invokeai/app/services/wildcard_records/wildcard_records_common.py +++ b/invokeai/app/services/wildcard_records/wildcard_records_common.py @@ -1,10 +1,12 @@ import json import re -from typing import Any +from typing import TYPE_CHECKING, Any -from dynamicprompts.wildcards import WildcardManager from pydantic import BaseModel, Field, TypeAdapter, field_validator +if TYPE_CHECKING: + from dynamicprompts.wildcards import WildcardManager + class WildcardNotFoundError(Exception): """Raised when a wildcard is not found""" @@ -105,7 +107,7 @@ def from_dict(cls, data: dict[str, Any]) -> "WildcardRecordDTO": WildcardRecordDTOValidator = TypeAdapter(WildcardRecordDTO) -def build_wildcard_manager(wildcards: list[WildcardRecordDTO]) -> WildcardManager: +def build_wildcard_manager(wildcards: list[WildcardRecordDTO]) -> "WildcardManager": """Builds the manager that resolves `__name__` against a user's wildcards. `root_map` accepts in-memory value lists, so wildcards resolve without ever touching the @@ -115,5 +117,8 @@ def build_wildcard_manager(wildcards: list[WildcardRecordDTO]) -> WildcardManage Constructed per request. The tree is built lazily from a handful of small lists, so there is nothing here worth caching and invalidating against edits. """ + # Imported here: the database layer imports this module for its records, and only expanding a prompt needs it. + from dynamicprompts.wildcards import WildcardManager + root_map = {"": [{wildcard.name: wildcard.values for wildcard in wildcards if wildcard.values}]} return WildcardManager(root_map=root_map) diff --git a/invokeai/app/services/wildcard_records/wildcard_records_default.py b/invokeai/app/services/wildcard_records/wildcard_records_default.py new file mode 100644 index 00000000000..f444eb266d1 --- /dev/null +++ b/invokeai/app/services/wildcard_records/wildcard_records_default.py @@ -0,0 +1,60 @@ +from typing import Optional + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import UniqueViolation +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.wildcard_records.wildcard_records_base import WildcardRecordsStorageBase +from invokeai.app.services.wildcard_records.wildcard_records_common import ( + WildcardChanges, + WildcardNameConflictError, + WildcardNotFoundError, + WildcardRecordDTO, + WildcardWithoutId, +) +from invokeai.app.util.misc import uuid_string + + +def _found(wildcard: Optional[WildcardRecordDTO], wildcard_id: str) -> WildcardRecordDTO: + if wildcard is None: + raise WildcardNotFoundError(f"Wildcard with id {wildcard_id} not found") + return wildcard + + +def _name_taken(name: Optional[str]) -> WildcardNameConflictError: + return WildcardNameConflictError(f"A wildcard named '{name}' already exists") + + +class WildcardRecordsStorage(WildcardRecordsStorageBase): + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + + def get(self, wildcard_id: str) -> WildcardRecordDTO: + return _found(self._queries.wildcards.get(wildcard_id), wildcard_id) + + def create(self, wildcard: WildcardWithoutId, user_id: str) -> WildcardRecordDTO: + wildcard_id = uuid_string() + try: + self._queries.wildcards.insert(wildcard_id, wildcard, user_id) + except UniqueViolation as error: + # The (user_id, name) unique index is the authority on uniqueness, so a concurrent create loses here + # rather than in a check-then-insert race. (An owner that does not exist is a ForeignKeyViolation.) + raise _name_taken(wildcard.name) from error + # The record is what was stored, so it is not read back. + return WildcardRecordDTO(**wildcard.model_dump(), id=wildcard_id, user_id=user_id) + + def update(self, wildcard_id: str, changes: WildcardChanges) -> WildcardRecordDTO: + def apply(q: Queries) -> Optional[WildcardRecordDTO]: + try: + q.wildcards.update(wildcard_id, changes) + except UniqueViolation as error: + raise _name_taken(changes.name) from error + return q.wildcards.get(wildcard_id) + + return _found(self._queries.run(apply), wildcard_id) + + def delete(self, wildcard_id: str) -> None: + self._queries.wildcards.delete(wildcard_id) + + def get_many(self, user_id: str) -> list[WildcardRecordDTO]: + return self._queries.wildcards.owned_by(user_id) diff --git a/invokeai/app/services/wildcard_records/wildcard_records_sqlite.py b/invokeai/app/services/wildcard_records/wildcard_records_sqlite.py deleted file mode 100644 index c598b09cb60..00000000000 --- a/invokeai/app/services/wildcard_records/wildcard_records_sqlite.py +++ /dev/null @@ -1,108 +0,0 @@ -import json -import sqlite3 - -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.wildcard_records.wildcard_records_base import WildcardRecordsStorageBase -from invokeai.app.services.wildcard_records.wildcard_records_common import ( - WildcardChanges, - WildcardNameConflictError, - WildcardNotFoundError, - WildcardRecordDTO, - WildcardWithoutId, -) -from invokeai.app.util.misc import uuid_string - - -def _is_name_conflict(error: sqlite3.IntegrityError) -> bool: - """Whether an integrity failure is the duplicate name and not, say, an unknown owner. - - The table also carries a users foreign key, and reporting a missing owner as - "that name is taken" would send the caller chasing the wrong thing. - """ - return "UNIQUE constraint failed" in str(error) - - -class SqliteWildcardRecordsStorage(WildcardRecordsStorageBase): - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - - def get(self, wildcard_id: str) -> WildcardRecordDTO: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT id, name, values_json AS "values", user_id - FROM wildcards - WHERE id = ?; - """, - (wildcard_id,), - ) - row = cursor.fetchone() - if row is None: - raise WildcardNotFoundError(f"Wildcard with id {wildcard_id} not found") - return WildcardRecordDTO.from_dict(dict(row)) - - def create(self, wildcard: WildcardWithoutId, user_id: str) -> WildcardRecordDTO: - wildcard_id = uuid_string() - try: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - INSERT INTO wildcards (id, name, values_json, user_id) - VALUES (?, ?, ?, ?); - """, - (wildcard_id, wildcard.name, json.dumps(wildcard.values), user_id), - ) - except sqlite3.IntegrityError as e: - # The (user_id, name) unique index is the authority on uniqueness, so a concurrent - # create loses here rather than in a check-then-insert race. - if not _is_name_conflict(e): - raise - raise WildcardNameConflictError(f"A wildcard named '{wildcard.name}' already exists") from e - return self.get(wildcard_id) - - def update(self, wildcard_id: str, changes: WildcardChanges) -> WildcardRecordDTO: - try: - with self._db.transaction() as cursor: - if changes.name is not None: - cursor.execute( - """--sql - UPDATE wildcards SET name = ? WHERE id = ?; - """, - (changes.name, wildcard_id), - ) - if changes.values is not None: - cursor.execute( - """--sql - UPDATE wildcards SET values_json = ? WHERE id = ?; - """, - (json.dumps(changes.values), wildcard_id), - ) - except sqlite3.IntegrityError as e: - if not _is_name_conflict(e): - raise - raise WildcardNameConflictError(f"A wildcard named '{changes.name}' already exists") from e - return self.get(wildcard_id) - - def delete(self, wildcard_id: str) -> None: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - DELETE FROM wildcards WHERE id = ?; - """, - (wildcard_id,), - ) - - def get_many(self, user_id: str) -> list[WildcardRecordDTO]: - with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT id, name, values_json AS "values", user_id - FROM wildcards - WHERE user_id = ? - ORDER BY name ASC; - """, - (user_id,), - ) - rows = cursor.fetchall() - return [WildcardRecordDTO.from_dict(dict(row)) for row in rows] diff --git a/invokeai/app/services/workflow_records/workflow_records_base.py b/invokeai/app/services/workflow_records/workflow_records_base.py index d7c3196f1f2..ac146c63be6 100644 --- a/invokeai/app/services/workflow_records/workflow_records_base.py +++ b/invokeai/app/services/workflow_records/workflow_records_base.py @@ -1,8 +1,7 @@ from abc import ABC, abstractmethod from typing import Optional -from invokeai.app.services.shared.pagination import PaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import PaginatedResults, SQLiteDirection from invokeai.app.services.workflow_records.workflow_records_common import ( WORKFLOW_LIBRARY_DEFAULT_USER_ID, Workflow, diff --git a/invokeai/app/services/workflow_records/workflow_records_default.py b/invokeai/app/services/workflow_records/workflow_records_default.py new file mode 100644 index 00000000000..a3a024161bc --- /dev/null +++ b/invokeai/app/services/workflow_records/workflow_records_default.py @@ -0,0 +1,339 @@ +import uuid +from pathlib import Path +from typing import Optional, Union + +from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import UniqueViolation +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.locks import DatabaseLock +from invokeai.app.services.shared.media_references import extract_media_references_from_json +from invokeai.app.services.shared.pagination import PaginatedResults, SQLiteDirection +from invokeai.app.services.workflow_records.workflow_records_base import WorkflowRecordsStorageBase +from invokeai.app.services.workflow_records.workflow_records_common import ( + WORKFLOW_LIBRARY_DEFAULT_USER_ID, + Workflow, + WorkflowAccessDeniedError, + WorkflowCategory, + WorkflowIdConflictError, + WorkflowImmutableError, + WorkflowNotFoundError, + WorkflowRecordDTO, + WorkflowRecordDTOBase, + WorkflowRecordListItemDTO, + WorkflowRecordOrderBy, + WorkflowRevisionConflictError, + WorkflowValidator, + WorkflowWithoutID, +) +from invokeai.app.util.misc import uuid_string + +_BUNDLED_WORKFLOWS = Path(__file__).parent / "default_workflows" + + +class _IdTaken(Exception): + """A creation's id belongs to a stored workflow.""" + + +def _not_found(workflow_id: str) -> WorkflowNotFoundError: + return WorkflowNotFoundError(f"Workflow with id {workflow_id} not found") + + +def _owner(user_id: Optional[str]) -> str: + return user_id if user_id is not None else WORKFLOW_LIBRARY_DEFAULT_USER_ID + + +def _record(summary: Optional[WorkflowRecordDTOBase], workflow: Workflow) -> WorkflowRecordDTO: + """The record a write returns: its row as the write left it, read in the write's transaction, with the workflow + the write stored.""" + if summary is None: + raise _not_found(workflow.id) + return WorkflowRecordDTO(**summary.model_dump(), workflow=workflow) + + +def _same_submission(existing: WorkflowRecordDTO, workflow: Workflow, user_id: str) -> WorkflowRecordDTO: + """A retried creation is accepted only when the record it finds is this owner's identical submission. + + Any other record under the id is a conflict; its content and owner are not revealed to the caller. + """ + if existing.user_id != user_id or existing.workflow.model_dump() != workflow.model_dump(): + raise WorkflowIdConflictError(workflow.id) + return existing + + +class WorkflowRecordsStorage(WorkflowRecordsStorageBase): + def __init__(self, database: Database) -> None: + super().__init__() + self._queries = database.queries + + def start(self, invoker: Invoker) -> None: + self._invoker = invoker + self._sync_default_workflows() + + def get(self, workflow_id: str) -> WorkflowRecordDTO: + """Gets a workflow by ID.""" + record = self._queries.workflows.get(workflow_id) + if record is None: + raise _not_found(workflow_id) + return record + + def create( + self, + workflow: WorkflowWithoutID, + user_id: str = WORKFLOW_LIBRARY_DEFAULT_USER_ID, + is_public: bool = False, + workflow_id: Optional[str] = None, + ) -> WorkflowRecordDTO: + if workflow.meta.category is WorkflowCategory.Default: + raise ValueError("Default workflows cannot be created via this method") + if workflow_id is not None: + try: + uuid.UUID(workflow_id) + except ValueError as e: + raise ValueError("A reserved workflow id must be a UUID") from e + + workflow_with_id = Workflow(**workflow.model_dump(), id=workflow_id or uuid_string()) + document_json = workflow_with_id.model_dump_json() + references = extract_media_references_from_json(document_json) + + def insert(q: Queries) -> Optional[WorkflowRecordDTOBase]: + # Shared with every write that makes media protected, exclusive for the intermediates cleanup's check and + # delete: the media this names cannot be deleted between that check and this commit. + q.locks.acquire(DatabaseLock.MEDIA_PROTECTION, shared=True) + try: + q.workflows.insert( + workflow_id=workflow_with_id.id, workflow=document_json, user_id=user_id, is_public=is_public + ) + except UniqueViolation as error: + # Only a reserved id can collide; a generated one is unique for practical purposes. + raise _IdTaken from error + q.media_references.replace( + owner_kind="workflow", user_id=user_id, owner_id=workflow_with_id.id, references=references + ) + return q.workflows.summary(workflow_with_id.id) + + try: + return _record(self._queries.run(insert), workflow_with_id) + except _IdTaken as error: + if workflow_id is None: + raise WorkflowIdConflictError(workflow_with_id.id) from error + # A creation sent again, to learn whether it landed -- possibly while the first one was still under way -- + # finds the record it made; anything else stored under the id is a conflict. + existing = self._queries.workflows.get(workflow_with_id.id) + if existing is None: + raise WorkflowIdConflictError(workflow_with_id.id) from error + return _same_submission(existing, workflow_with_id, user_id) + + def update( + self, workflow: Workflow, user_id: Optional[str] = None, expected_revision: Optional[int] = None + ) -> WorkflowRecordDTO: + if workflow.meta.category is WorkflowCategory.Default: + raise ValueError("A workflow cannot be updated into the default category") + + document_json = workflow.model_dump_json() + references = extract_media_references_from_json(document_json) + + def save(q: Queries) -> Optional[WorkflowRecordDTOBase]: + q.locks.acquire(DatabaseLock.MEDIA_PROTECTION, shared=True) + # Locked before anything is checked: the checks and the write apply to the same row. `category` is + # generated from the stored JSON, so it still describes the record as it is, not as the request would + # rewrite it. + stored = q.workflows.lock(workflow.id) + if stored is None: + raise _not_found(workflow.id) + if stored.category == WorkflowCategory.Default.value: + raise WorkflowImmutableError(workflow.id) + owner = _owner(stored.user_id) + if user_id is not None and owner != user_id: + raise WorkflowAccessDeniedError(workflow.id) + if expected_revision is not None and stored.revision != expected_revision: + raise WorkflowRevisionConflictError(workflow.id, expected_revision, stored.revision) + q.workflows.save(workflow.id, document_json) + q.media_references.replace( + owner_kind="workflow", user_id=owner, owner_id=workflow.id, references=references + ) + return q.workflows.summary(workflow.id) + + return _record(self._queries.run(save), workflow) + + def delete(self, workflow_id: str, user_id: Optional[str] = None) -> None: + def delete(q: Queries) -> None: + stored = q.workflows.lock(workflow_id) + if stored is None: + raise _not_found(workflow_id) + if stored.category == WorkflowCategory.Default.value: + raise WorkflowImmutableError(workflow_id) + owner = _owner(stored.user_id) + if user_id is not None and owner != user_id: + raise WorkflowAccessDeniedError(workflow_id) + q.workflows.delete(workflow_id) + # Every write of a workflow indexes its references under its owner. + q.media_references.delete(owner_kind="workflow", user_id=owner, owner_id=workflow_id) + + self._queries.run(delete) + + def update_is_public(self, workflow_id: str, is_public: bool, user_id: Optional[str] = None) -> WorkflowRecordDTO: + """Updates the is_public field of a workflow and manages the 'shared' tag automatically.""" + + def change(q: Queries) -> tuple[Optional[WorkflowRecordDTOBase], Union[Workflow, str]]: + # The document is read and rewritten under the row's lock. A concurrent workflow edit must not land + # between the read and this full-document write: its reference index would then describe the edit while + # the stored workflow names the previous assets. + stored = q.workflows.lock(workflow_id, with_document=True) + if stored is None or stored.document is None: + raise _not_found(workflow_id) + if stored.category != WorkflowCategory.User.value or (user_id is not None and stored.user_id != user_id): + return q.workflows.summary(workflow_id), stored.document + workflow = Workflow.model_validate_json(stored.document) + tags_list = [t.strip() for t in workflow.tags.split(",") if t.strip()] if workflow.tags else [] + if is_public and "shared" not in tags_list: + tags_list.append("shared") + elif not is_public and "shared" in tags_list: + tags_list.remove("shared") + workflow = workflow.model_copy(update={"tags": ", ".join(tags_list)}) + # Visibility is bookkeeping: the `shared` tag rewrite does not advance the content revision, so an editor + # holding the previous revision is not asked to resolve a conflict it cannot see. + q.workflows.set_public(workflow_id, workflow.model_dump_json(), is_public) + return q.workflows.summary(workflow_id), workflow + + summary, workflow = self._queries.run(change) + return _record( + summary, workflow if isinstance(workflow, Workflow) else WorkflowValidator.validate_json(workflow) + ) + + def get_many( + self, + order_by: WorkflowRecordOrderBy, + direction: SQLiteDirection, + categories: Optional[list[WorkflowCategory]], + page: int = 0, + per_page: Optional[int] = None, + query: Optional[str] = None, + tags: Optional[list[str]] = None, + has_been_opened: Optional[bool] = None, + user_id: Optional[str] = None, + is_public: Optional[bool] = None, + ) -> PaginatedResults[WorkflowRecordListItemDTO]: + workflows, total = self._queries.workflows.page( + order_by=order_by, + direction=direction, + offset=page * per_page if per_page else 0, + limit=per_page if per_page else None, + categories=categories, + tags=tags, + has_been_opened=has_been_opened, + query=query, + user_id=user_id, + is_public=is_public, + ) + pages = total // per_page + (total % per_page > 0) if per_page else 1 # Without pagination, one page. + return PaginatedResults( + items=workflows, + page=page, + per_page=per_page if per_page else total, + pages=pages, + total=total, + ) + + def counts_by_tag( + self, + tags: list[str], + categories: Optional[list[WorkflowCategory]] = None, + has_been_opened: Optional[bool] = None, + user_id: Optional[str] = None, + is_public: Optional[bool] = None, + ) -> dict[str, int]: + if not tags: + return {} + counts = self._queries.workflows.tag_counts( + tags, categories=categories, has_been_opened=has_been_opened, user_id=user_id, is_public=is_public + ) + return dict(zip(tags, counts, strict=True)) + + def counts_by_category( + self, + categories: list[WorkflowCategory], + has_been_opened: Optional[bool] = None, + user_id: Optional[str] = None, + is_public: Optional[bool] = None, + ) -> dict[str, int]: + if not categories: + return {} + return self._queries.workflows.category_counts( + categories, has_been_opened=has_been_opened, user_id=user_id, is_public=is_public + ) + + def update_opened_at(self, workflow_id: str, user_id: Optional[str] = None) -> None: + self._queries.workflows.mark("opened_at", workflow_id, user_id) + + def update_last_run_at(self, workflow_id: str, user_id: Optional[str] = None) -> None: + self._queries.workflows.mark("last_run_at", workflow_id, user_id) + + def get_all_tags( + self, + categories: Optional[list[WorkflowCategory]] = None, + user_id: Optional[str] = None, + is_public: Optional[bool] = None, + ) -> list[str]: + tag_lists = self._queries.workflows.tag_lists(categories=categories, user_id=user_id, is_public=is_public) + # Each workflow's tags are one comma-separated string. + return sorted({tag.strip() for tag_list in tag_lists for tag in tag_list.split(",") if tag.strip()}) + + def _sync_default_workflows(self) -> None: + """Syncs default workflows to the database. Internal use only. + + An enhancement might be to only update workflows that have changed. This would require stable default + workflow IDs, and properly incrementing the workflow version. + + It's much simpler to just replace them all with whichever workflows are in the directory. + + The downside is that the `updated_at` and `opened_at` timestamps for default workflows are meaningless, as + they are overwritten every time the server starts. + """ + bundled: list[Workflow] = [] + for path in _BUNDLED_WORKFLOWS.glob("*.json"): + workflow = WorkflowValidator.validate_json(path.read_bytes()) + assert workflow.id.startswith("default_"), ( + f'Invalid default workflow ID (must start with "default_"): {workflow.id}' + ) + assert workflow.meta.category is WorkflowCategory.Default, ( + f"Invalid default workflow category: {workflow.meta.category}" + ) + bundled.append(workflow) + bundled_ids = {workflow.id for workflow in bundled} + + def sync(q: Queries) -> tuple[list[tuple[str, str]], list[Workflow], list[Workflow]]: + stored = q.workflows.documents(bundled_ids) + obsolete = [ + (workflow_id, name) for workflow_id, name in q.workflows.bundled() if workflow_id not in bundled_ids + ] + missing = [workflow for workflow in bundled if workflow.id not in stored] + changed = [ + workflow + for workflow in bundled + if workflow.id in stored and workflow != WorkflowValidator.validate_json(stored[workflow.id]) + ] + for workflow_id, _name in obsolete: + # Not `delete`, which refuses bundled workflows. + q.workflows.delete(workflow_id) + for workflow in missing: + q.workflows.insert( + workflow_id=workflow.id, + workflow=workflow.model_dump_json(), + user_id=WORKFLOW_LIBRARY_DEFAULT_USER_ID, + is_public=False, + ) + for workflow in changed: + # Not `update`, which refuses bundled workflows. A changed bundle is new content, so its revision + # advances like any other write. + q.workflows.save(workflow.id, workflow.model_dump_json()) + return obsolete, missing, changed + + obsolete, missing, changed = self._queries.run(sync) + logger = self._invoker.services.logger + for workflow_id, name in obsolete: + logger.debug(f"Deleting obsolete default workflow {name} ({workflow_id})") + for workflow in missing: + logger.debug(f"Adding missing default workflow {workflow.name} ({workflow.id})") + for workflow in changed: + logger.debug(f"Updating library workflow {workflow.name} ({workflow.id})") diff --git a/invokeai/app/services/workflow_records/workflow_records_sqlite.py b/invokeai/app/services/workflow_records/workflow_records_sqlite.py deleted file mode 100644 index 1e028a3e489..00000000000 --- a/invokeai/app/services/workflow_records/workflow_records_sqlite.py +++ /dev/null @@ -1,693 +0,0 @@ -import sqlite3 -import uuid -from pathlib import Path -from typing import Optional - -from invokeai.app.services.invoker import Invoker -from invokeai.app.services.shared.media_references import ( - MediaReferences, - extract_media_references_from_json, - replace_media_references, -) -from invokeai.app.services.shared.pagination import PaginatedResults -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.workflow_records.workflow_records_base import WorkflowRecordsStorageBase -from invokeai.app.services.workflow_records.workflow_records_common import ( - WORKFLOW_LIBRARY_DEFAULT_USER_ID, - Workflow, - WorkflowAccessDeniedError, - WorkflowCategory, - WorkflowIdConflictError, - WorkflowImmutableError, - WorkflowNotFoundError, - WorkflowRecordDTO, - WorkflowRecordListItemDTO, - WorkflowRecordListItemDTOValidator, - WorkflowRecordOrderBy, - WorkflowRevisionConflictError, - WorkflowValidator, - WorkflowWithoutID, -) -from invokeai.app.util.misc import uuid_string - -SQL_TIME_FORMAT = "%Y-%m-%d %H:%M:%f" - -_RECORD_COLUMNS = ( - "workflow_id, workflow, name, created_at, updated_at, opened_at, last_run_at, user_id, is_public, revision" -) - - -class SqliteWorkflowRecordsStorage(WorkflowRecordsStorageBase): - def __init__(self, db: SqliteDatabase) -> None: - super().__init__() - self._db = db - - def start(self, invoker: Invoker) -> None: - self._invoker = invoker - self._sync_default_workflows() - - def get(self, workflow_id: str) -> WorkflowRecordDTO: - """Gets a workflow by ID.""" - with self._db.transaction() as cursor: - return self._get_on_cursor(cursor, workflow_id) - - @staticmethod - def _get_on_cursor(cursor: sqlite3.Cursor, workflow_id: str) -> WorkflowRecordDTO: - """Reads one record on the caller's transaction; `get()` on an open transaction would commit it early.""" - cursor.execute( - f"SELECT {_RECORD_COLUMNS} FROM workflow_library WHERE workflow_id = ?;", - (workflow_id,), - ) - row = cursor.fetchone() - if row is None: - raise WorkflowNotFoundError(f"Workflow with id {workflow_id} not found") - return WorkflowRecordDTO.from_dict(dict(row)) - - def create( - self, - workflow: WorkflowWithoutID, - user_id: str = WORKFLOW_LIBRARY_DEFAULT_USER_ID, - is_public: bool = False, - workflow_id: Optional[str] = None, - ) -> WorkflowRecordDTO: - if workflow.meta.category is WorkflowCategory.Default: - raise ValueError("Default workflows cannot be created via this method") - if workflow_id is not None: - try: - uuid.UUID(workflow_id) - except ValueError as e: - raise ValueError("A reserved workflow id must be a UUID") from e - - workflow_with_id = Workflow(**workflow.model_dump(), id=workflow_id or uuid_string()) - document_json = workflow_with_id.model_dump_json() - references = extract_media_references_from_json(document_json) - with self._db.transaction() as cursor: - if workflow_id is not None: - existing = self._match_reserved_record(cursor, workflow_with_id, user_id) - if existing is not None: - return existing - try: - cursor.execute( - """--sql - INSERT INTO workflow_library ( - workflow_id, - workflow, - user_id, - is_public, - revision - ) - VALUES (?, ?, ?, ?, 1); - """, - (workflow_with_id.id, document_json, user_id, is_public), - ) - except sqlite3.IntegrityError as e: - # Only a reserved id can collide; a generated one is unique for practical purposes. - raise WorkflowIdConflictError(workflow_with_id.id) from e - self._index_references(cursor, workflow_with_id.id, references, user_id=user_id) - return self._get_on_cursor(cursor, workflow_with_id.id) - - @staticmethod - def _match_reserved_record(cursor: sqlite3.Cursor, workflow: Workflow, user_id: str) -> Optional[WorkflowRecordDTO]: - """A retried creation is accepted only when the record it finds is this owner's identical submission. - - Any other record under the id is a conflict; its content and owner are not revealed to the caller. - """ - cursor.execute( - f"SELECT {_RECORD_COLUMNS} FROM workflow_library WHERE workflow_id = ?;", - (workflow.id,), - ) - row = cursor.fetchone() - if row is None: - return None - existing = WorkflowRecordDTO.from_dict(dict(row)) - if existing.user_id != user_id or existing.workflow.model_dump() != workflow.model_dump(): - raise WorkflowIdConflictError(workflow.id) - return existing - - def update( - self, workflow: Workflow, user_id: Optional[str] = None, expected_revision: Optional[int] = None - ) -> WorkflowRecordDTO: - if workflow.meta.category is WorkflowCategory.Default: - raise ValueError("A workflow cannot be updated into the default category") - - document_json = workflow.model_dump_json() - references = extract_media_references_from_json(document_json) - with self._db.transaction() as cursor: - # `category` is generated from the stored JSON, so it still describes the record as it is, not as - # the request would rewrite it. - cursor.execute( - "SELECT category, user_id, revision FROM workflow_library WHERE workflow_id = ?;", - (workflow.id,), - ) - row = cursor.fetchone() - if row is None: - raise WorkflowNotFoundError(f"Workflow with id {workflow.id} not found") - stored_category, owner_id, current_revision = row - if stored_category == WorkflowCategory.Default.value: - raise WorkflowImmutableError(workflow.id) - owner = str(owner_id) if owner_id is not None else WORKFLOW_LIBRARY_DEFAULT_USER_ID - if user_id is not None and owner != user_id: - raise WorkflowAccessDeniedError(workflow.id) - if expected_revision is not None and current_revision != expected_revision: - raise WorkflowRevisionConflictError(workflow.id, expected_revision, current_revision) - - cursor.execute( - """--sql - UPDATE workflow_library - SET workflow = ?, revision = revision + 1 - WHERE workflow_id = ? AND revision = ?; - """, - (document_json, workflow.id, current_revision), - ) - if cursor.rowcount == 0: - # Another connection won between the read and this compare-and-swap. - cursor.execute("SELECT revision FROM workflow_library WHERE workflow_id = ?;", (workflow.id,)) - raced = cursor.fetchone() - if raced is None: - raise WorkflowNotFoundError(f"Workflow with id {workflow.id} not found") - raise WorkflowRevisionConflictError(workflow.id, expected_revision or current_revision, raced[0]) - self._index_references(cursor, workflow.id, references, user_id=owner) - return self._get_on_cursor(cursor, workflow.id) - - def delete(self, workflow_id: str, user_id: Optional[str] = None) -> None: - with self._db.transaction() as cursor: - cursor.execute( - "SELECT category, user_id FROM workflow_library WHERE workflow_id = ?;", - (workflow_id,), - ) - row = cursor.fetchone() - if row is None: - raise WorkflowNotFoundError(f"Workflow with id {workflow_id} not found") - stored_category, owner_id = row - if stored_category == WorkflowCategory.Default.value: - raise WorkflowImmutableError(workflow_id) - owner = str(owner_id) if owner_id is not None else WORKFLOW_LIBRARY_DEFAULT_USER_ID - if user_id is not None and owner != user_id: - raise WorkflowAccessDeniedError(workflow_id) - cursor.execute("DELETE from workflow_library WHERE workflow_id = ?;", (workflow_id,)) - # The row is gone, so its owner is the only fact left to delete by; every owner's - # rows for this id are dropped since workflow ids are globally unique. - cursor.execute( - "DELETE FROM media_references WHERE owner_kind = 'workflow' AND owner_id = ?;", - (workflow_id,), - ) - return None - - def update_is_public(self, workflow_id: str, is_public: bool, user_id: Optional[str] = None) -> WorkflowRecordDTO: - """Updates the is_public field of a workflow and manages the 'shared' tag automatically.""" - # Read and rewrite the current document under one transaction. A concurrent workflow edit - # must not land between the read and this full-document write: its reference index would - # then describe the edit while the stored workflow names the previous assets. - with self._db.transaction() as cursor: - if user_id is not None: - cursor.execute( - "SELECT workflow FROM workflow_library WHERE workflow_id = ? AND category = 'user' AND user_id = ?;", - (workflow_id, user_id), - ) - else: - cursor.execute( - "SELECT workflow FROM workflow_library WHERE workflow_id = ? AND category = 'user';", (workflow_id,) - ) - row = cursor.fetchone() - if row is not None: - workflow = Workflow.model_validate_json(row[0]) - tags_list = [t.strip() for t in workflow.tags.split(",") if t.strip()] if workflow.tags else [] - if is_public and "shared" not in tags_list: - tags_list.append("shared") - elif not is_public and "shared" in tags_list: - tags_list.remove("shared") - updated_workflow = workflow.model_copy(update={"tags": ", ".join(tags_list)}) - # Visibility is bookkeeping: the `shared` tag rewrite does not advance the content revision, so - # an editor holding the previous revision is not asked to resolve a conflict it cannot see. - cursor.execute( - """--sql - UPDATE workflow_library - SET workflow = ?, is_public = ? WHERE workflow_id = ?; - """, - (updated_workflow.model_dump_json(), is_public, workflow_id), - ) - return self._get_on_cursor(cursor, workflow_id) - - @staticmethod - def _index_references( - cursor: sqlite3.Cursor, workflow_id: str, references: MediaReferences, *, user_id: str - ) -> None: - """Records the media a library workflow names, on the caller's transaction.""" - replace_media_references( - cursor, owner_kind="workflow", user_id=user_id, owner_id=workflow_id, references=references - ) - - def get_many( - self, - order_by: WorkflowRecordOrderBy, - direction: SQLiteDirection, - categories: Optional[list[WorkflowCategory]], - page: int = 0, - per_page: Optional[int] = None, - query: Optional[str] = None, - tags: Optional[list[str]] = None, - has_been_opened: Optional[bool] = None, - user_id: Optional[str] = None, - is_public: Optional[bool] = None, - ) -> PaginatedResults[WorkflowRecordListItemDTO]: - with self._db.transaction() as cursor: - # sanitize! - assert order_by in WorkflowRecordOrderBy - assert direction in SQLiteDirection - - # We will construct the query dynamically based on the query params - - # The main query to get the workflows / counts - main_query = """ - SELECT - workflow_id, - category, - name, - description, - created_at, - updated_at, - opened_at, - last_run_at, - tags, - user_id, - is_public, - revision - FROM workflow_library - """ - count_query = "SELECT COUNT(*) FROM workflow_library" - - # Start with an empty list of conditions and params - conditions: list[str] = [] - params: list[str | int] = [] - - if categories: - # Categories is a list of WorkflowCategory enum values, and a single string in the DB - - # Ensure all categories are valid (is this necessary?) - assert all(c in WorkflowCategory for c in categories) - - # Construct a placeholder string for the number of categories - placeholders = ", ".join("?" for _ in categories) - - # Construct the condition string & params - category_condition = f"category IN ({placeholders})" - category_params = [category.value for category in categories] - - conditions.append(category_condition) - params.extend(category_params) - - if tags: - # Tags is a list of strings, and a single string in the DB - # The string in the DB has no guaranteed format - - # Construct a list of conditions for each tag - tags_conditions = ["tags LIKE ?" for _ in tags] - tags_conditions_joined = " OR ".join(tags_conditions) - tags_condition = f"({tags_conditions_joined})" - - # And the params for the tags, case-insensitive - tags_params = [f"%{t.strip()}%" for t in tags] - - conditions.append(tags_condition) - params.extend(tags_params) - - if has_been_opened: - conditions.append("opened_at IS NOT NULL") - elif has_been_opened is False: - conditions.append("opened_at IS NULL") - - # Ignore whitespace in the query - stripped_query = query.strip() if query else None - if stripped_query: - # Construct a wildcard query for the name, description, and tags - wildcard_query = "%" + stripped_query + "%" - query_condition = "(name LIKE ? OR description LIKE ? OR tags LIKE ?)" - - conditions.append(query_condition) - params.extend([wildcard_query, wildcard_query, wildcard_query]) - - if user_id is not None: - # Scope to the given user but always include default workflows - conditions.append("(user_id = ? OR category = 'default')") - params.append(user_id) - - if is_public is True: - conditions.append("is_public = TRUE") - elif is_public is False: - conditions.append("is_public = FALSE") - - if conditions: - # If there are conditions, add a WHERE clause and then join the conditions - main_query += " WHERE " - count_query += " WHERE " - - all_conditions = " AND ".join(conditions) - main_query += all_conditions - count_query += all_conditions - - # After this point, the query and params differ for the main query and the count query - main_params = params.copy() - count_params = params.copy() - - # Main query also gets ORDER BY and LIMIT/OFFSET - main_query += f" ORDER BY {order_by.value} {direction.value}" - - if per_page: - main_query += " LIMIT ? OFFSET ?" - main_params.extend([per_page, page * per_page]) - - # Put a ring on it - main_query += ";" - count_query += ";" - - cursor.execute(main_query, main_params) - rows = cursor.fetchall() - workflows = [WorkflowRecordListItemDTOValidator.validate_python(dict(row)) for row in rows] - - cursor.execute(count_query, count_params) - total = cursor.fetchone()[0] - - if per_page: - pages = total // per_page + (total % per_page > 0) - else: - pages = 1 # If no pagination, there is only one page - - return PaginatedResults( - items=workflows, - page=page, - per_page=per_page if per_page else total, - pages=pages, - total=total, - ) - - def counts_by_tag( - self, - tags: list[str], - categories: Optional[list[WorkflowCategory]] = None, - has_been_opened: Optional[bool] = None, - user_id: Optional[str] = None, - is_public: Optional[bool] = None, - ) -> dict[str, int]: - if not tags: - return {} - - with self._db.transaction() as cursor: - result: dict[str, int] = {} - # Base conditions for categories and selected tags - base_conditions: list[str] = [] - base_params: list[str | int] = [] - - # Add category conditions - if categories: - assert all(c in WorkflowCategory for c in categories) - placeholders = ", ".join("?" for _ in categories) - base_conditions.append(f"category IN ({placeholders})") - base_params.extend([category.value for category in categories]) - - if has_been_opened: - base_conditions.append("opened_at IS NOT NULL") - elif has_been_opened is False: - base_conditions.append("opened_at IS NULL") - - if user_id is not None: - # Scope to the given user but always include default workflows - base_conditions.append("(user_id = ? OR category = 'default')") - base_params.append(user_id) - - if is_public is True: - base_conditions.append("is_public = TRUE") - elif is_public is False: - base_conditions.append("is_public = FALSE") - - # For each tag to count, run a separate query - for tag in tags: - # Start with the base conditions - conditions = base_conditions.copy() - params = base_params.copy() - - # Add this specific tag condition - conditions.append("tags LIKE ?") - params.append(f"%{tag.strip()}%") - - # Construct the full query - stmt = """--sql - SELECT COUNT(*) - FROM workflow_library - """ - - if conditions: - stmt += " WHERE " + " AND ".join(conditions) - - cursor.execute(stmt, params) - count = cursor.fetchone()[0] - result[tag] = count - - return result - - def counts_by_category( - self, - categories: list[WorkflowCategory], - has_been_opened: Optional[bool] = None, - user_id: Optional[str] = None, - is_public: Optional[bool] = None, - ) -> dict[str, int]: - with self._db.transaction() as cursor: - result: dict[str, int] = {} - # Base conditions for categories - base_conditions: list[str] = [] - base_params: list[str | int] = [] - - # Add category conditions - if categories: - assert all(c in WorkflowCategory for c in categories) - placeholders = ", ".join("?" for _ in categories) - base_conditions.append(f"category IN ({placeholders})") - base_params.extend([category.value for category in categories]) - - if has_been_opened: - base_conditions.append("opened_at IS NOT NULL") - elif has_been_opened is False: - base_conditions.append("opened_at IS NULL") - - if user_id is not None: - # Scope to the given user but always include default workflows - base_conditions.append("(user_id = ? OR category = 'default')") - base_params.append(user_id) - - if is_public is True: - base_conditions.append("is_public = TRUE") - elif is_public is False: - base_conditions.append("is_public = FALSE") - - # For each category to count, run a separate query - for category in categories: - # Start with the base conditions - conditions = base_conditions.copy() - params = base_params.copy() - - # Add this specific category condition - conditions.append("category = ?") - params.append(category.value) - - # Construct the full query - stmt = """--sql - SELECT COUNT(*) - FROM workflow_library - """ - - if conditions: - stmt += " WHERE " + " AND ".join(conditions) - - cursor.execute(stmt, params) - count = cursor.fetchone()[0] - result[category.value] = count - - return result - - def update_opened_at(self, workflow_id: str, user_id: Optional[str] = None) -> None: - with self._db.transaction() as cursor: - if user_id is not None: - cursor.execute( - f"""--sql - UPDATE workflow_library - SET opened_at = STRFTIME('{SQL_TIME_FORMAT}', 'NOW') - WHERE workflow_id = ? AND user_id = ?; - """, - (workflow_id, user_id), - ) - else: - cursor.execute( - f"""--sql - UPDATE workflow_library - SET opened_at = STRFTIME('{SQL_TIME_FORMAT}', 'NOW') - WHERE workflow_id = ?; - """, - (workflow_id,), - ) - - def update_last_run_at(self, workflow_id: str, user_id: Optional[str] = None) -> None: - with self._db.transaction() as cursor: - if user_id is not None: - cursor.execute( - f"""--sql - UPDATE workflow_library - SET last_run_at = STRFTIME('{SQL_TIME_FORMAT}', 'NOW') - WHERE workflow_id = ? AND user_id = ?; - """, - (workflow_id, user_id), - ) - else: - cursor.execute( - f"""--sql - UPDATE workflow_library - SET last_run_at = STRFTIME('{SQL_TIME_FORMAT}', 'NOW') - WHERE workflow_id = ?; - """, - (workflow_id,), - ) - - def get_all_tags( - self, - categories: Optional[list[WorkflowCategory]] = None, - user_id: Optional[str] = None, - is_public: Optional[bool] = None, - ) -> list[str]: - with self._db.transaction() as cursor: - conditions: list[str] = [] - params: list[str] = [] - - # Only get workflows that have tags - conditions.append("tags IS NOT NULL AND tags != ''") - - if categories: - assert all(c in WorkflowCategory for c in categories) - placeholders = ", ".join("?" for _ in categories) - conditions.append(f"category IN ({placeholders})") - params.extend([category.value for category in categories]) - - if user_id is not None: - # Scope to the given user but always include default workflows - conditions.append("(user_id = ? OR category = 'default')") - params.append(user_id) - - if is_public is True: - conditions.append("is_public = TRUE") - elif is_public is False: - conditions.append("is_public = FALSE") - - stmt = """--sql - SELECT DISTINCT tags - FROM workflow_library - """ - - if conditions: - stmt += " WHERE " + " AND ".join(conditions) - - cursor.execute(stmt, params) - rows = cursor.fetchall() - - # Parse comma-separated tags and collect unique tags - all_tags: set[str] = set() - - for row in rows: - tags_value = row[0] - if tags_value and isinstance(tags_value, str): - # Tags are stored as comma-separated string - for tag in tags_value.split(","): - tag_stripped = tag.strip() - if tag_stripped: - all_tags.add(tag_stripped) - - return sorted(all_tags) - - def _sync_default_workflows(self) -> None: - """Syncs default workflows to the database. Internal use only.""" - - """ - An enhancement might be to only update workflows that have changed. This would require stable - default workflow IDs, and properly incrementing the workflow version. - - It's much simpler to just replace them all with whichever workflows are in the directory. - - The downside is that the `updated_at` and `opened_at` timestamps for default workflows are - meaningless, as they are overwritten every time the server starts. - """ - - with self._db.transaction() as cursor: - workflows_from_file: list[Workflow] = [] - workflows_to_update: list[Workflow] = [] - workflows_to_add: list[Workflow] = [] - workflows_dir = Path(__file__).parent / Path("default_workflows") - workflow_paths = workflows_dir.glob("*.json") - for path in workflow_paths: - bytes_ = path.read_bytes() - workflow_from_file = WorkflowValidator.validate_json(bytes_) - - assert workflow_from_file.id.startswith("default_"), ( - f'Invalid default workflow ID (must start with "default_"): {workflow_from_file.id}' - ) - - assert workflow_from_file.meta.category is WorkflowCategory.Default, ( - f"Invalid default workflow category: {workflow_from_file.meta.category}" - ) - - workflows_from_file.append(workflow_from_file) - - # Read on this cursor: `get()` would open a nested transaction and commit this one early. - cursor.execute("SELECT workflow FROM workflow_library WHERE workflow_id = ?;", (workflow_from_file.id,)) - row = cursor.fetchone() - if row is None: - self._invoker.services.logger.debug( - f"Adding missing default workflow {workflow_from_file.name} ({workflow_from_file.id})" - ) - workflows_to_add.append(workflow_from_file) - continue - if workflow_from_file != WorkflowValidator.validate_json(row[0]): - self._invoker.services.logger.debug( - f"Updating library workflow {workflow_from_file.name} ({workflow_from_file.id})" - ) - workflows_to_update.append(workflow_from_file) - - cursor.execute("SELECT workflow_id, name FROM workflow_library WHERE category = 'default';") - library_workflows_from_db = cursor.fetchall() - - workflows_from_file_ids = [w.id for w in workflows_from_file] - - for workflow_id, name in library_workflows_from_db: - if workflow_id not in workflows_from_file_ids: - self._invoker.services.logger.debug(f"Deleting obsolete default workflow {name} ({workflow_id})") - # We cannot use the `delete` method here, as it only deletes non-default workflows - cursor.execute( - """--sql - DELETE from workflow_library - WHERE workflow_id = ?; - """, - (workflow_id,), - ) - - for w in workflows_to_add: - # We cannot use the `create` method here, as it only creates non-default workflows - cursor.execute( - """--sql - INSERT INTO workflow_library ( - workflow_id, - workflow, - revision - ) - VALUES (?, ?, 1); - """, - (w.id, w.model_dump_json()), - ) - - for w in workflows_to_update: - # We cannot use the `update` method here, as it refuses default workflows. A changed bundle is - # new content, so its revision advances like any other write. - cursor.execute( - """--sql - UPDATE workflow_library - SET workflow = ?, revision = revision + 1 - WHERE workflow_id = ?; - """, - (w.model_dump_json(), w.id), - ) diff --git a/invokeai/app/services/workflow_thumbnails/workflow_thumbnails_base.py b/invokeai/app/services/workflow_thumbnails/workflow_thumbnails_base.py index f51d200dea1..f6b18e62363 100644 --- a/invokeai/app/services/workflow_thumbnails/workflow_thumbnails_base.py +++ b/invokeai/app/services/workflow_thumbnails/workflow_thumbnails_base.py @@ -1,20 +1,27 @@ from abc import ABC, abstractmethod from pathlib import Path +from typing import Optional from PIL import Image +from invokeai.app.services.workflow_records.workflow_records_common import WorkflowCategory + class WorkflowThumbnailServiceBase(ABC): """Base class for workflow thumbnail services""" @abstractmethod - def get_path(self, workflow_id: str, with_hash: bool = True) -> Path: - """Gets the path to a workflow thumbnail""" + def get_path(self, workflow_id: str, with_hash: bool = True, category: Optional[WorkflowCategory] = None) -> Path: + """Gets the path to a workflow thumbnail. A caller that knows the workflow's category passes it, sparing a + read of the workflow.""" pass @abstractmethod - def get_url(self, workflow_id: str, with_hash: bool = True) -> str | None: - """Gets the URL of a workflow thumbnail""" + def get_url( + self, workflow_id: str, with_hash: bool = True, category: Optional[WorkflowCategory] = None + ) -> str | None: + """Gets the URL of a workflow thumbnail. A caller that knows the workflow's category passes it, sparing a + read of the workflow.""" pass @abstractmethod diff --git a/invokeai/app/services/workflow_thumbnails/workflow_thumbnails_disk.py b/invokeai/app/services/workflow_thumbnails/workflow_thumbnails_disk.py index 3fbfa7607fe..19eec2ab08f 100644 --- a/invokeai/app/services/workflow_thumbnails/workflow_thumbnails_disk.py +++ b/invokeai/app/services/workflow_thumbnails/workflow_thumbnails_disk.py @@ -1,4 +1,5 @@ from pathlib import Path +from typing import Optional from PIL import Image from PIL.Image import Image as PILImageType @@ -41,9 +42,10 @@ def save(self, workflow_id: str, image: PILImageType) -> None: except Exception as e: raise WorkflowThumbnailFileSaveException from e - def get_path(self, workflow_id: str, with_hash: bool = True) -> Path: - workflow = self._invoker.services.workflow_records.get(workflow_id).workflow - if workflow.meta.category is WorkflowCategory.Default: + def get_path(self, workflow_id: str, with_hash: bool = True, category: Optional[WorkflowCategory] = None) -> Path: + if category is None: + category = self._invoker.services.workflow_records.get(workflow_id).workflow.meta.category + if category == WorkflowCategory.Default: default_thumbnails_dir = Path(__file__).parent / Path("default_workflow_thumbnails") path = default_thumbnails_dir / (workflow_id + ".png") else: @@ -51,8 +53,10 @@ def get_path(self, workflow_id: str, with_hash: bool = True) -> Path: return path - def get_url(self, workflow_id: str, with_hash: bool = True) -> str | None: - path = self.get_path(workflow_id) + def get_url( + self, workflow_id: str, with_hash: bool = True, category: Optional[WorkflowCategory] = None + ) -> str | None: + path = self.get_path(workflow_id, category=category) if not self._validate_path(path): return diff --git a/invokeai/app/util/database_copy.py b/invokeai/app/util/database_copy.py new file mode 100644 index 00000000000..97540c97b57 --- /dev/null +++ b/invokeai/app/util/database_copy.py @@ -0,0 +1,324 @@ +"""`invoke-db-copy`: copies an install's SQLite database into a new MySQL or MariaDB database, or such a database back +into a new SQLite file.""" + +import argparse +import os +import shutil +import sys +import tempfile +from logging import Logger +from pathlib import Path +from typing import Optional + +from invokeai.app.services.config.config_default import InvokeAIAppConfig, load_config_from_root +from invokeai.app.services.shared.database.copy import ( + Problem, + copy_database, + copy_records, + count_normalized, + delete_orphans, + find_orphans, + find_oversized, + merged_rows, + verify_copy, +) +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.startup import ( + open_copy_target, + open_migrated_database, + redacted_database_url, +) +from invokeai.backend.util.logging import InvokeAILogger + +_START_AGAIN = ( + "The target holds part of the copy, which InvokeAI refuses. Drop the target database, create it again empty, and " + "run invoke-db-copy again." +) + + +def main() -> None: + parser = argparse.ArgumentParser( + prog="invoke-db-copy", + description=( + "Copy an InvokeAI install's SQLite database into a new, empty MySQL or MariaDB database; or, with " + "--to-sqlite, such a database back into a new SQLite file." + ), + epilog=( + "Stop InvokeAI first: what it writes during the copy is not copied. The source is not changed. The server " + "database is the one `db_url` names in invokeai.yaml or INVOKEAI_DB_URL, unless --target or --source " + "names another." + ), + ) + parser.add_argument( + "--target", + help="URL of the target, e.g. mariadb+pymysql://user:pw@host/db (default: db_url). A URL here is visible to " + "other users of this computer and kept in the shell's history; prefer db_url", + ) + parser.add_argument( + "--to-sqlite", + type=Path, + metavar="FILE", + help="Copy the server database into this new SQLite file instead, e.g. to go back to SQLite", + ) + parser.add_argument( + "--source", + help="With --to-sqlite: URL of the server database (default: db_url). A URL here is visible to other users of " + "this computer and kept in the shell's history; prefer db_url", + ) + parser.add_argument("--root", type=Path, help="InvokeAI root directory (default: as InvokeAI finds it)") + parser.add_argument( + "--check", action="store_true", help="Check the source and the target and report, without copying" + ) + parser.add_argument( + "--orphans", + choices=("fail", "skip"), + default="fail", + help="Rows whose foreign keys name no row, which a server refuses: stop (default), or leave them out", + ) + args = parser.parse_args() + if args.to_sqlite is not None: + if args.target is not None or args.orphans == "skip": + parser.error("--to-sqlite takes --source, not --target or --orphans: a server database has no orphans") + sys.exit(copy_to_sqlite(args.to_sqlite, source_url=args.source, root=args.root, check_only=args.check)) + if args.source is not None: + parser.error("--source belongs to --to-sqlite; the source of a copy to a server is the install's SQLite file") + sys.exit(copy(args.target, root=args.root, check_only=args.check, skip_orphans=args.orphans == "skip")) + + +def copy(target_url: Optional[str], *, root: Optional[Path], check_only: bool, skip_orphans: bool) -> int: + """Runs the copy and reports on it; the process's exit code. + + :param target_url: The target; by default, the database the install's config names. + """ + config = _install_config(root) + target_url = target_url or config.db_url + if not target_url: + print("Name the target: set db_url in invokeai.yaml or INVOKEAI_DB_URL, or pass --target.") + return 1 + # The source is the install's SQLite database, whatever db_url says. + source_config = config.model_copy(update={"db_url": None, "use_memory_db": False}) + logger = InvokeAILogger.get_logger("invoke-db-copy") + print(f"Source: {source_config.db_path}") + print(f"Target: {redacted_database_url(target_url)}") + try: + source = open_migrated_database(source_config, logger) + try: + target, max_allowed_packet = open_copy_target(target_url, logger) + except BaseException: + source.dispose() + raise + except Exception as e: + # A database not set up for it, a URL of another kind, a server out of reach, the `mysql` extra missing. + print(f"\nCannot copy: {e}") + return 1 + + try: + try: + # A consistent snapshot to work from, beside the database, which it is as large as: orphans can be left + # out of it, and the source stays as it is. + snapshot_dir = Path(tempfile.mkdtemp(prefix="invoke-db-copy-", dir=source_config.db_path.parent)) + except BaseException: + source.dispose() + raise + try: + snapshot_path = snapshot_dir / "snapshot.db" + try: + source.backup(snapshot_path) + finally: + source.dispose() + snapshot = Database.open_sqlite(snapshot_path, logger) + try: + return _copy(snapshot, target, max_allowed_packet, check_only=check_only, skip_orphans=skip_orphans) + finally: + snapshot.dispose() + finally: + shutil.rmtree(snapshot_dir, ignore_errors=True) + if snapshot_dir.exists(): + print( + f"\nCould not delete the snapshot at {snapshot_dir}, which holds a copy of the database: delete it." + ) + except Exception as e: + # Before anything is copied (the copy reports its own failures): the snapshot, or a check of it. + print(f"\nCannot copy: {e}\nNothing was copied.") + return 1 + finally: + target.dispose() + + +def _copy( + snapshot: Database, target: Database, max_allowed_packet: int, *, check_only: bool, skip_orphans: bool +) -> int: + blocking: list[Problem] = [] + orphans = find_orphans(snapshot) + _report("Rows whose foreign keys name a missing row", orphans) + if orphans and skip_orphans: + _report("Left out of the copy", delete_orphans(snapshot)) + elif orphans: + blocking.extend(orphans) + print(" Copy with --orphans skip to leave them out.") + oversized = find_oversized(snapshot, max_allowed_packet) + _report("Rows the target cannot store", oversized) + blocking.extend(oversized) + _report("Values the copy changes", count_normalized(snapshot)) + + if blocking: + print("\nNothing was copied.") + return 1 + if check_only: + print("\nThe check found nothing that stops a copy. Nothing was copied.") + return 0 + + print("\nCopying...") + try: + copied = copy_database(snapshot, target, progress=lambda table: print(f" {table}", flush=True)) + print(f"Copied {_rows(sum(copied.values()))} of {len(copied)} tables. Verifying...") + mismatches = verify_copy(snapshot, target) + if not mismatches: + copy_records(snapshot, target) + mismatches = verify_copy(snapshot, target, records=True) + merged = merged_rows(snapshot, target) + except Exception as e: + print(f"\nThe copy failed: {e}\n{_START_AGAIN}") + return 1 + if mismatches: + print("\nThe copy does not match its source:") + for mismatch in mismatches: + print(f" {mismatch}") + print(_START_AGAIN) + return 1 + for table, fewer in merged.items(): + print(f" {table}: {_rows(fewer)} the target treats as equal to others were merged.") + print("\nDone: every table matches its source. InvokeAI uses the target once db_url names it.") + return 0 + + +def copy_to_sqlite(target_path: Path, *, source_url: Optional[str], root: Optional[Path], check_only: bool) -> int: + """Copies a server database into a new SQLite file and reports on it; the process's exit code. + + The file appears only once the copy is complete and verified, and an existing file is never overwritten. + + :param source_url: The server database; by default, the one the install's config names. + """ + config = _install_config(root) + source_url = source_url or config.db_url + if not source_url: + print("Name the source: set db_url in invokeai.yaml or INVOKEAI_DB_URL, or pass --source.") + return 1 + target_path = target_path.absolute() + logger = InvokeAILogger.get_logger("invoke-db-copy") + print(f"Source: {redacted_database_url(source_url)}") + print(f"Target: {target_path}") + if target_path.exists(): + print( + f"\n{target_path} exists. Name a new file: the copy overwrites no database. To replace the install's " + "database, rename it first." + ) + return 1 + if not target_path.parent.is_dir(): + print(f"\nThere is no directory {target_path.parent} to write the copy into.") + return 1 + try: + source_config = config.model_copy(update={"db_url": source_url, "use_memory_db": False}) + source = open_migrated_database(source_config, logger) + try: + # Held while copying, so that no InvokeAI process changes the database meanwhile. + source.hold_instance_lock() + except BaseException: + source.dispose() + raise + except Exception as e: + # A database InvokeAI uses or has not updated, a URL of another kind, a server out of reach. + print(f"\nCannot copy: {e}") + return 1 + + try: + if check_only: + print("\nThe check found nothing that stops a copy. Nothing was copied.") + return 0 + return _copy_to_sqlite(source, target_path, logger) + finally: + source.dispose() + + +def _copy_to_sqlite(source: Database, target_path: Path, logger: Logger) -> int: + # Written beside the target and moved into place once verified: a failed copy leaves no file that a later run, + # or InvokeAI, could take for a database. + try: + work_dir = Path(tempfile.mkdtemp(prefix="invoke-db-copy-", dir=target_path.parent)) + except OSError as e: + print(f"\nCannot copy: {e}") + return 1 + try: + work_path = work_dir / target_path.name + print("\nCopying...") + target = Database.open_sqlite(work_path, logger) + try: + copied = copy_database(source, target, progress=lambda table: print(f" {table}", flush=True)) + print(f"Copied {_rows(sum(copied.values()))} of {len(copied)} tables. Verifying...") + mismatches = verify_copy(source, target) + if not mismatches: + copy_records(source, target) + mismatches = verify_copy(source, target, records=True) + finally: + target.dispose() + if mismatches: + print("\nThe copy does not match its source, so it was not kept:") + for mismatch in mismatches: + print(f" {mismatch}") + return 1 + if not _move_into_place(work_path, target_path): + print(f"\n{target_path} was made while copying, so the copy was not kept: it overwrites no file.") + return 1 + except Exception as e: + print(f"\nThe copy failed, so it was not kept: {e}") + return 1 + finally: + shutil.rmtree(work_dir, ignore_errors=True) + if work_dir.exists(): + print(f"\nCould not delete {work_dir}, which holds a copy of the database: delete it.") + print( + f"\nDone: every table matches its source. To use {target_path.name}, with InvokeAI stopped: put it in place of " + "the install's database (renaming that one first), and remove db_url from invokeai.yaml and the environment." + ) + return 0 + + +def _move_into_place(work_path: Path, target_path: Path) -> bool: + """Moves the copy to `target_path` unless a file is there; whether it did.""" + try: + # A link, unlike a rename, refuses a path that exists, without a moment between checking and moving. + os.link(work_path, target_path) + except FileExistsError: + return False + except OSError: + # A file system without hard links (FAT, some network shares): check, then move. + if target_path.exists() or target_path.is_symlink(): + return False + os.replace(work_path, target_path) + return True + # The work file's other name goes with its directory, which the caller deletes. + return True + + +def _install_config(root: Optional[Path]) -> InvokeAIAppConfig: + if root is None: + from invokeai.app.services.config import get_config + + root = get_config().root_path + return load_config_from_root(root) + + +def _report(heading: str, problems: list[Problem]) -> None: + if not problems: + return + print(f"\n{heading}:") + for problem in problems: + print(f" {problem.table}: {_rows(problem.count)} {problem.what}") + + +def _rows(count: int) -> str: + return f"{count} row" if count == 1 else f"{count} rows" + + +if __name__ == "__main__": + main() diff --git a/invokeai/app/util/user_management.py b/invokeai/app/util/user_management.py index e9c78082797..4ce27cdfe73 100644 --- a/invokeai/app/util/user_management.py +++ b/invokeai/app/util/user_management.py @@ -35,13 +35,22 @@ ) +def _cli_config() -> Any: + """The install's config, from the environment and its `invokeai.yaml` as the app reads them: where its database + is (`db_dir`, `db_url`), and its other paths.""" + from invokeai.app.services.config import get_config + from invokeai.app.services.config.config_default import load_config_from_root + + return load_config_from_root(get_config().root_path) + + def _prepare_deleted_user_fonts(config: Any, db: Any, user_id: str, logger: Any) -> tuple[Any | None, tuple[str, ...]]: """Snapshot private managed paths before the user-row foreign-key cascade.""" try: from invokeai.app.services.fonts.fonts_default import FontService fonts = FontService( - db=db, + db, fonts_dir=config.fonts_path, storage_dir=config.fonts_storage_path, logger=logger, @@ -72,8 +81,7 @@ def _cleanup_deleted_user_fonts(fonts: Any | None, user_id: str, storage_paths: def _add_user_interactive() -> bool: """Add a user interactively by prompting for details.""" from invokeai.app.services.auth.password_utils import validate_password_strength - from invokeai.app.services.config import get_config - from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase + from invokeai.app.services.shared.database.startup import open_migrated_database from invokeai.app.services.users.users_common import UserCreateRequest from invokeai.app.services.users.users_default import UserService from invokeai.backend.util.logging import InvokeAILogger @@ -106,8 +114,8 @@ def _add_user_interactive() -> bool: is_admin = is_admin_input in ("y", "yes") try: - config = get_config() - db = SqliteDatabase(config.db_path, InvokeAILogger.get_logger(), synchronous=config.db_synchronous) + config = _cli_config() + db = open_migrated_database(config, InvokeAILogger.get_logger()) user_service = UserService(db) user_data = UserCreateRequest(email=email, display_name=display_name, password=password, is_admin=is_admin) @@ -135,8 +143,7 @@ def _add_user_interactive() -> bool: def _add_user_cli(email: str, password: str, display_name: str | None = None, is_admin: bool = False) -> bool: """Add a user via CLI arguments.""" from invokeai.app.services.auth.password_utils import validate_password_strength - from invokeai.app.services.config import get_config - from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase + from invokeai.app.services.shared.database.startup import open_migrated_database from invokeai.app.services.users.users_common import UserCreateRequest from invokeai.app.services.users.users_default import UserService from invokeai.backend.util.logging import InvokeAILogger @@ -147,8 +154,8 @@ def _add_user_cli(email: str, password: str, display_name: str | None = None, is return False try: - config = get_config() - db = SqliteDatabase(config.db_path, InvokeAILogger.get_logger(), synchronous=config.db_synchronous) + config = _cli_config() + db = open_migrated_database(config, InvokeAILogger.get_logger()) user_service = UserService(db) user_data = UserCreateRequest(email=email, display_name=display_name, password=password, is_admin=is_admin) @@ -209,8 +216,7 @@ def useradd() -> None: def _delete_user_interactive() -> bool: """Delete a user interactively by prompting for email.""" - from invokeai.app.services.config import get_config - from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase + from invokeai.app.services.shared.database.startup import open_migrated_database from invokeai.app.services.users.users_default import UserService from invokeai.backend.util.logging import InvokeAILogger @@ -222,9 +228,9 @@ def _delete_user_interactive() -> bool: return False try: - config = get_config() + config = _cli_config() logger = InvokeAILogger.get_logger(config=config) - db = SqliteDatabase(config.db_path, logger, synchronous=config.db_synchronous) + db = open_migrated_database(config, logger) user_service = UserService(db) user = user_service.get_by_email(email) @@ -264,15 +270,14 @@ def _delete_user_interactive() -> bool: def _delete_user_cli(email: str, force: bool = False) -> bool: """Delete a user via CLI arguments.""" - from invokeai.app.services.config import get_config - from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase + from invokeai.app.services.shared.database.startup import open_migrated_database from invokeai.app.services.users.users_default import UserService from invokeai.backend.util.logging import InvokeAILogger try: - config = get_config() + config = _cli_config() logger = InvokeAILogger.get_logger(config=config) - db = SqliteDatabase(config.db_path, logger, synchronous=config.db_synchronous) + db = open_migrated_database(config, logger) user_service = UserService(db) user = user_service.get_by_email(email) @@ -341,17 +346,15 @@ def userdel() -> None: def _list_users_table() -> bool: """List all users in a formatted table.""" - from invokeai.app.services.config import get_config - from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase + from invokeai.app.services.shared.database.startup import open_migrated_database from invokeai.app.services.users.users_default import UserService from invokeai.backend.util.logging import InvokeAILogger - config = get_config() + config = _cli_config() logger = InvokeAILogger.get_logger(config=config) - db = SqliteDatabase(config.db_path, logger, synchronous=config.db_synchronous) - user_service = UserService(db) - try: + db = open_migrated_database(config, logger) + user_service = UserService(db) users = user_service.list_users() if not users: @@ -381,17 +384,15 @@ def _list_users_table() -> bool: def _list_users_json() -> bool: """List all users in JSON format.""" - from invokeai.app.services.config import get_config - from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase + from invokeai.app.services.shared.database.startup import open_migrated_database from invokeai.app.services.users.users_default import UserService from invokeai.backend.util.logging import InvokeAILogger - config = get_config() + config = _cli_config() logger = InvokeAILogger.get_logger(config=config) - db = SqliteDatabase(config.db_path, logger, synchronous=config.db_synchronous) - user_service = UserService(db) - try: + db = open_migrated_database(config, logger) + user_service = UserService(db) users = user_service.list_users() users_data = [ @@ -448,8 +449,7 @@ def userlist() -> None: def _modify_user_interactive() -> bool: """Modify a user interactively by prompting for details.""" from invokeai.app.services.auth.password_utils import validate_password_strength - from invokeai.app.services.config import get_config - from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase + from invokeai.app.services.shared.database.startup import open_migrated_database from invokeai.app.services.users.users_common import UserUpdateRequest from invokeai.app.services.users.users_default import UserService from invokeai.backend.util.logging import InvokeAILogger @@ -462,8 +462,8 @@ def _modify_user_interactive() -> bool: return False try: - config = get_config() - db = SqliteDatabase(config.db_path, InvokeAILogger.get_logger(), synchronous=config.db_synchronous) + config = _cli_config() + db = open_migrated_database(config, InvokeAILogger.get_logger()) user_service = UserService(db) user = user_service.get_by_email(email) @@ -550,8 +550,7 @@ def _modify_user_cli( ) -> bool: """Modify a user via CLI arguments.""" from invokeai.app.services.auth.password_utils import validate_password_strength - from invokeai.app.services.config import get_config - from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase + from invokeai.app.services.shared.database.startup import open_migrated_database from invokeai.app.services.users.users_common import UserUpdateRequest from invokeai.app.services.users.users_default import UserService from invokeai.backend.util.logging import InvokeAILogger @@ -563,8 +562,8 @@ def _modify_user_cli( return False try: - config = get_config() - db = SqliteDatabase(config.db_path, InvokeAILogger.get_logger(), synchronous=config.db_synchronous) + config = _cli_config() + db = open_migrated_database(config, InvokeAILogger.get_logger()) user_service = UserService(db) user = user_service.get_by_email(email) diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json index 2178e9612b0..75719b63cec 100644 --- a/invokeai/frontend/api/openapi.json +++ b/invokeai/frontend/api/openapi.json @@ -10633,7 +10633,8 @@ "schema": { "anyOf": [ { - "type": "integer" + "type": "integer", + "minimum": 0 }, { "type": "null" @@ -10651,7 +10652,9 @@ "schema": { "anyOf": [ { - "type": "integer" + "type": "integer", + "maximum": 1000, + "minimum": 0 }, { "type": "null" @@ -12164,6 +12167,7 @@ "required": true, "schema": { "type": "string", + "maxLength": 255, "description": "The queue id to perform this operation on", "title": "Queue Id" }, @@ -14022,6 +14026,7 @@ "required": false, "schema": { "type": "integer", + "minimum": 0, "description": "The page to get", "default": 0, "title": "Page" @@ -14035,16 +14040,18 @@ "schema": { "anyOf": [ { - "type": "integer" + "type": "integer", + "maximum": 1000, + "minimum": 0 }, { "type": "null" } ], - "description": "The number of workflows per page", + "description": "The number of workflows per page; all of them when omitted", "title": "Per Page" }, - "description": "The number of workflows per page" + "description": "The number of workflows per page; all of them when omitted" }, { "name": "order_by", @@ -15610,6 +15617,7 @@ "required": true, "schema": { "type": "string", + "maxLength": 255, "description": "Key to set", "title": "Key" }, @@ -20519,6 +20527,7 @@ "properties": { "batch_id": { "type": "string", + "maxLength": 255, "title": "Batch Id", "description": "The ID of the batch" }, @@ -21080,7 +21089,8 @@ "cover_image_name": { "anyOf": [ { - "type": "string" + "type": "string", + "maxLength": 255 }, { "type": "null" @@ -21485,7 +21495,8 @@ "workflow_id": { "anyOf": [ { - "type": "string" + "type": "string", + "maxLength": 255 }, { "type": "null" @@ -59090,6 +59101,18 @@ "description": "Path to InvokeAI databases directory.", "default": "databases" }, + "db_url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Db Url", + "description": "URL of a MySQL 8.4+ or MariaDB 10.11+ database to use instead of the SQLite database in `db_dir`, e.g. `mariadb+pymysql://invokeai:password@db.example/invokeai`. Needs the `mysql` extra. One InvokeAI process uses a database at a time. Read at startup only." + }, "db_synchronous": { "type": "string", "enum": ["full", "normal"], @@ -59765,7 +59788,7 @@ "additionalProperties": false, "type": "object", "title": "InvokeAIAppConfig", - "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n show_donation_link: Show the Donate to InvokeAI link in the frontend menus.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n db_synchronous: SQLite durability setting. `full`, the default and what InvokeAI has always used, flushes every commit to disk. `normal` acknowledges commits without waiting for that flush - measured at roughly 12x shorter commits on an SSD - and cannot corrupt the database under WAL, which is why it is refused, with a warning, when WAL is unavailable for the database file. What `normal` gives up is the most recent transactions on a power loss or OS crash: a just-written image record or queue status, not the image file itself.
Valid values: `full`, `normal`\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n offload_on_architecture_switch: When a session uses main models of another architecture than the previous one on the same GPU, move every model it does not use from VRAM to RAM before it runs. Models both sessions use stay. Switching back reloads the moved models from RAM, which costs a few seconds per switch when both architectures would have fit in VRAM together; turn this off if you alternate between such models.\n clear_vram_after_session: After every queue item (each image of a batch), move the models held in VRAM to RAM and return the freed memory to the GPU driver. The models stay cached in RAM; the next item reloads them from there, which costs a few seconds (more with keep_ram_copy_of_weights off). Changes nothing where VRAM is system memory (CPU, MPS, integrated GPUs). Turn this on if a generation that runs fine the first time runs out of memory, or slows down sharply because the GPU spills into shared memory, when it is repeated.\n fp8_compute: Keep ComfyUI 'scaled fp8' checkpoints quantized instead of dequantizing them at load, and run their matmuls on the fp8 tensor cores (requires an Ada/SM 8.9 or newer NVIDIA GPU, or a ROCm GPU that passes an fp8 matmul probe; falls back automatically otherwise). Speeds up denoising and keeps the transformer at fp8 size - FP8 Storage alone also keeps it at that size, without the speedup - but quantizes activations as well, so images will differ from previous versions at the same seed. Reproducibility also requires the model to be FULLY resident in VRAM: a layer whose weights are still in RAM falls back to the dequantized path, and since which layers are resident shifts from run to run, the same seed then yields visibly different images. For repeatable output, ensure the model loads at 100% (e.g. enable_partial_loading=false with enough free VRAM).\n fp8_compute_full_precision_hints: Honor the per-layer 'full_precision_matrix_mult' flags that some scaled-fp8 checkpoints ship. Those layers then dequantize on every forward instead of using the fp8 tensor cores, which can cost a large part of the fp8_compute speedup - on checkpoints that mark many layers, most of it. Set to false to run every quantized layer on the fp8 tensor cores, ignoring the producer's instruction; faster, but the marked layers were flagged as numerically sensitive, so quality may suffer. Only has an effect when fp8_compute is enabled.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. Unset on a ROCm build of PyTorch on Windows, \"expandable_segments:True\" is used.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, `xpu:N` (where N is a device number)\n precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n noise_dtype: The dtype seeded noise is drawn in for SD1.5/SDXL, SD3, FLUX.1, FLUX.2, CogView4 and the Z-Image seed variance enhancer. `float32` draws the same noise on every platform. `float16` is the half-precision draw of earlier versions: on macOS it keeps the images your seeds gave before this update; on Windows and Linux both values give (nearly) the same images, and images made before the update cannot be reproduced there.
Valid values: `float32`, `float16`\n wan_memory_optimization: Enable experimental Wan memory optimizations at the cost of slower generation.\n pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it.\n auto_tiled_decode: Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there.\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n allow_private_download_urls: Allow the download queue to fetch from loopback, link-local and private-network addresses. Disabled by default so that a download URL cannot be used to reach services that are only reachable from the server. Enable this only if you install models from a mirror on your own network.\n download_proxy: Optional HTTP proxy for model downloads. The proxy must enforce the public-address policy because proxy-side DNS cannot be checked by InvokeAI.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n image_index_enabled: Maintain a semantic embedding index of gallery images and videos, used by the image map and semantic search features.\n image_index_model: Name of the installed CLIP Vision or SigLIP model used to embed gallery images and video thumbnails. Changing the model discards embeddings computed by the previous model.\n image_index_device: Set to `cpu` to compute image embeddings on the CPU with a service-local copy of the model - avoids VRAM use and lets indexing run during generations. Any other value is ignored: embeddings otherwise run on the model cache's device, pausing while generations are in progress.\n image_index_batch_size: Number of gallery items embedded per batch by the image index worker.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_atlascloud_api_key: API key for Atlas Cloud image generation.\n external_atlascloud_base_url: Base URL override for Atlas Cloud image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set.\n http_compression_level: Compression level for gzipped HTTP API responses. 0 disables response compression entirely, 1 is fastest, 9 (the default) is smallest. Compression runs on the event loop and blocks the whole server while it works, and level 9 costs about 5.5x the time of level 1 for 0.4 percentage points of extra compression, so lowering this makes the app noticeably more responsive on large libraries. Set to 0 when a reverse proxy already compresses responses." + "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n show_donation_link: Show the Donate to InvokeAI link in the frontend menus.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n db_url: URL of a MySQL 8.4+ or MariaDB 10.11+ database to use instead of the SQLite database in `db_dir`, e.g. `mariadb+pymysql://invokeai:password@db.example/invokeai`. Needs the `mysql` extra. One InvokeAI process uses a database at a time. Read at startup only.\n db_synchronous: SQLite durability setting. `full`, the default and what InvokeAI has always used, flushes every commit to disk. `normal` acknowledges commits without waiting for that flush - measured at roughly 12x shorter commits on an SSD - and cannot corrupt the database under WAL, which is why it is refused, with a warning, when WAL is unavailable for the database file. What `normal` gives up is the most recent transactions on a power loss or OS crash: a just-written image record or queue status, not the image file itself.
Valid values: `full`, `normal`\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n offload_on_architecture_switch: When a session uses main models of another architecture than the previous one on the same GPU, move every model it does not use from VRAM to RAM before it runs. Models both sessions use stay. Switching back reloads the moved models from RAM, which costs a few seconds per switch when both architectures would have fit in VRAM together; turn this off if you alternate between such models.\n clear_vram_after_session: After every queue item (each image of a batch), move the models held in VRAM to RAM and return the freed memory to the GPU driver. The models stay cached in RAM; the next item reloads them from there, which costs a few seconds (more with keep_ram_copy_of_weights off). Changes nothing where VRAM is system memory (CPU, MPS, integrated GPUs). Turn this on if a generation that runs fine the first time runs out of memory, or slows down sharply because the GPU spills into shared memory, when it is repeated.\n fp8_compute: Keep ComfyUI 'scaled fp8' checkpoints quantized instead of dequantizing them at load, and run their matmuls on the fp8 tensor cores (requires an Ada/SM 8.9 or newer NVIDIA GPU, or a ROCm GPU that passes an fp8 matmul probe; falls back automatically otherwise). Speeds up denoising and keeps the transformer at fp8 size - FP8 Storage alone also keeps it at that size, without the speedup - but quantizes activations as well, so images will differ from previous versions at the same seed. Reproducibility also requires the model to be FULLY resident in VRAM: a layer whose weights are still in RAM falls back to the dequantized path, and since which layers are resident shifts from run to run, the same seed then yields visibly different images. For repeatable output, ensure the model loads at 100% (e.g. enable_partial_loading=false with enough free VRAM).\n fp8_compute_full_precision_hints: Honor the per-layer 'full_precision_matrix_mult' flags that some scaled-fp8 checkpoints ship. Those layers then dequantize on every forward instead of using the fp8 tensor cores, which can cost a large part of the fp8_compute speedup - on checkpoints that mark many layers, most of it. Set to false to run every quantized layer on the fp8 tensor cores, ignoring the producer's instruction; faster, but the marked layers were flagged as numerically sensitive, so quality may suffer. Only has an effect when fp8_compute is enabled.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. Unset on a ROCm build of PyTorch on Windows, \"expandable_segments:True\" is used.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, `xpu:N` (where N is a device number)\n precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n noise_dtype: The dtype seeded noise is drawn in for SD1.5/SDXL, SD3, FLUX.1, FLUX.2, CogView4 and the Z-Image seed variance enhancer. `float32` draws the same noise on every platform. `float16` is the half-precision draw of earlier versions: on macOS it keeps the images your seeds gave before this update; on Windows and Linux both values give (nearly) the same images, and images made before the update cannot be reproduced there.
Valid values: `float32`, `float16`\n wan_memory_optimization: Enable experimental Wan memory optimizations at the cost of slower generation.\n pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it.\n auto_tiled_decode: Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there.\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n allow_private_download_urls: Allow the download queue to fetch from loopback, link-local and private-network addresses. Disabled by default so that a download URL cannot be used to reach services that are only reachable from the server. Enable this only if you install models from a mirror on your own network.\n download_proxy: Optional HTTP proxy for model downloads. The proxy must enforce the public-address policy because proxy-side DNS cannot be checked by InvokeAI.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n image_index_enabled: Maintain a semantic embedding index of gallery images and videos, used by the image map and semantic search features.\n image_index_model: Name of the installed CLIP Vision or SigLIP model used to embed gallery images and video thumbnails. Changing the model discards embeddings computed by the previous model.\n image_index_device: Set to `cpu` to compute image embeddings on the CPU with a service-local copy of the model - avoids VRAM use and lets indexing run during generations. Any other value is ignored: embeddings otherwise run on the model cache's device, pausing while generations are in progress.\n image_index_batch_size: Number of gallery items embedded per batch by the image index worker.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_atlascloud_api_key: API key for Atlas Cloud image generation.\n external_atlascloud_base_url: Base URL override for Atlas Cloud image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set.\n http_compression_level: Compression level for gzipped HTTP API responses. 0 disables response compression entirely, 1 is fastest, 9 (the default) is smallest. Compression runs on the event loop and blocks the whole server while it works, and level 9 costs about 5.5x the time of level 1 for 0.4 percentage points of extra compression, so lowering this makes the app noticeably more responsive on large libraries. Set to 0 when a reverse proxy already compresses responses." }, "InvokeAIAppConfigWithSetFields": { "properties": { @@ -87241,6 +87264,7 @@ "properties": { "model_key_1": { "type": "string", + "maxLength": 255, "title": "Model Key 1", "description": "The key of the first model in the relationship", "examples": [ @@ -87254,6 +87278,7 @@ }, "model_key_2": { "type": "string", + "maxLength": 255, "title": "Model Key 2", "description": "The key of the second model in the relationship", "examples": [ @@ -90631,7 +90656,8 @@ "project_id": { "anyOf": [ { - "type": "string" + "type": "string", + "maxLength": 255 }, { "type": "null" diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts index 07069ed8464..3184eb09546 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -22564,6 +22564,7 @@ export type components = { * download_cache_dir: Path to the directory that contains dynamically downloaded models. * legacy_conf_dir: Path to directory of legacy checkpoint config files. * db_dir: Path to InvokeAI databases directory. + * db_url: URL of a MySQL 8.4+ or MariaDB 10.11+ database to use instead of the SQLite database in `db_dir`, e.g. `mariadb+pymysql://invokeai:password@db.example/invokeai`. Needs the `mysql` extra. One InvokeAI process uses a database at a time. Read at startup only. * db_synchronous: SQLite durability setting. `full`, the default and what InvokeAI has always used, flushes every commit to disk. `normal` acknowledges commits without waiting for that flush - measured at roughly 12x shorter commits on an SSD - and cannot corrupt the database under WAL, which is why it is refused, with a warning, when WAL is unavailable for the database file. What `normal` gives up is the most recent transactions on a power loss or OS crash: a just-written image record or queue status, not the image file itself.
Valid values: `full`, `normal` * outputs_dir: Path to directory for outputs. * image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash` @@ -22772,6 +22773,11 @@ export type components = { * @default databases */ db_dir?: string; + /** + * Db Url + * @description URL of a MySQL 8.4+ or MariaDB 10.11+ database to use instead of the SQLite database in `db_dir`, e.g. `mariadb+pymysql://invokeai:password@db.example/invokeai`. Needs the `mysql` extra. One InvokeAI process uses a database at a time. Read at startup only. + */ + db_url?: string | null; /** * Db Synchronous * @description SQLite durability setting. `full`, the default and what InvokeAI has always used, flushes every commit to disk. `normal` acknowledges commits without waiting for that flush - measured at roughly 12x shorter commits on an SSD - and cannot corrupt the database under WAL, which is why it is refused, with a warning, when WAL is unavailable for the database file. What `normal` gives up is the most recent transactions on a power loss or OS crash: a just-written image record or queue status, not the image file itself. @@ -58977,7 +58983,7 @@ export interface operations { query?: { /** @description The page to get */ page?: number; - /** @description The number of workflows per page */ + /** @description The number of workflows per page; all of them when omitted */ per_page?: number | null; /** @description The attribute to order by */ order_by?: components["schemas"]["WorkflowRecordOrderBy"]; diff --git a/invokeai/frontend/install/import_images.py b/invokeai/frontend/install/import_images.py index 3a005f7f38f..5de0251859e 100644 --- a/invokeai/frontend/install/import_images.py +++ b/invokeai/frontend/install/import_images.py @@ -5,24 +5,32 @@ import datetime import glob import json -import locale import os import re import shutil -import sqlite3 from pathlib import Path +from typing import Optional import PIL import PIL.ImageOps import PIL.PngImagePlugin -import yaml from prompt_toolkit import prompt from prompt_toolkit.completion import PathCompleter from prompt_toolkit.key_binding import KeyBindings from prompt_toolkit.shortcuts import message_dialog -from invokeai.app.services.config.config_default import get_config -from invokeai.app.util.misc import uuid_string +from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage +from invokeai.app.services.board_records.board_records_common import BoardRecordOrderBy +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage +from invokeai.app.services.config.config_default import InvokeAIAppConfig, get_config, load_config_from_root +from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import DatabaseError +from invokeai.app.services.shared.database.startup import open_migrated_database, redacted_database_url +from invokeai.app.services.shared.database.types import timestamp_text +from invokeai.app.services.shared.pagination import SQLiteDirection +from invokeai.backend.util.logging import InvokeAILogger app_config = get_config() @@ -46,14 +54,15 @@ def __init__(self): TIMESTAMP_STRING = datetime.datetime.utcnow().strftime("%Y%m%dT%H%M%SZ") - INVOKE_DIRNAME = "invokeai" YAML_FILENAME = "invokeai.yaml" - DATABASE_FILENAME = "invokeai.db" database_path = None database_backup_dir = None outputs_path = None thumbnail_path = None + synchronous = "full" + # The install's config, when its paths are used rather than ones the user gave. + app_config: Optional[InvokeAIAppConfig] = None def find_and_load(self): """Find the yaml config file and load""" @@ -69,22 +78,17 @@ def confirm_and_load(self, invoke_root): """Validate a yaml path exists, confirms the user wants to use it and loads config.""" yaml_path = os.path.join(invoke_root, self.YAML_FILENAME) if os.path.exists(yaml_path): - db_dir, outdir = self.load_paths_from_yaml(yaml_path) - if os.path.isabs(db_dir): - database_path = os.path.join(db_dir, self.DATABASE_FILENAME) - else: - database_path = os.path.join(invoke_root, db_dir, self.DATABASE_FILENAME) - - if os.path.isabs(outdir): - outputs_path = os.path.join(outdir, "images") - else: - outputs_path = os.path.join(invoke_root, outdir, "images") + config = load_config_from_root(Path(invoke_root)) + self.synchronous = config.db_synchronous + database_path = str(config.db_path) + outputs_path = str(config.outputs_path / "images") - db_exists = os.path.exists(database_path) + # A server database is checked when it is opened. + db_exists = bool(config.db_url) or os.path.exists(database_path) outdir_exists = os.path.exists(outputs_path) text = f"Found {self.YAML_FILENAME} file at {yaml_path}:" - text += f"\n Database : {database_path}" + text += f"\n Database : {redacted_database_url(config.db_url) or database_path}" text += f"\n Outputs : {outputs_path}" text += "\n\nUse these paths for import (yes) or choose different ones (no) [Yn]: " @@ -92,6 +96,7 @@ def confirm_and_load(self, invoke_root): if (prompt(text).strip() or "Y").upper().startswith("Y"): self.database_path = database_path self.outputs_path = outputs_path + self.app_config = config return True else: return False @@ -140,18 +145,6 @@ def confirm_and_load_from_user(self): return - def load_paths_from_yaml(self, yaml_path): - """Load an Invoke AI yaml file and get the database and outputs paths.""" - try: - with open(yaml_path, "rt", encoding=locale.getpreferredencoding()) as file: - yamlinfo = yaml.safe_load(file) - db_dir = yamlinfo.get("InvokeAI", {}).get("Paths", {}).get("db_dir", None) - outdir = yamlinfo.get("InvokeAI", {}).get("Paths", {}).get("outdir", None) - return db_dir, outdir - except Exception: - print(f"Failed to load paths from yaml file! {yaml_path}!") - return None, None - class ImportStats: """DTO for tracking work progress.""" @@ -374,79 +367,100 @@ def split_prompt(self, raw_prompt: str): class DatabaseMapper: - """Class to abstract database functionality.""" + """The script's work on the database, through the app's records.""" + + # The account that owns what the import adds, as for a single-user install. + OWNER = "system" - def __init__(self, database_path, database_backup_dir): + def __init__( + self, database_path, database_backup_dir, synchronous="full", app_config: Optional[InvokeAIAppConfig] = None + ): self.database_path = database_path self.database_backup_dir = database_backup_dir - self.connection = None - self.cursor = None + self.synchronous = synchronous + self.app_config = app_config + self.database: Optional[Database] = None def connect(self): - """Open connection to the database.""" - self.connection = sqlite3.connect(self.database_path) - self.cursor = self.connection.cursor() + """Open the database: the install's, as its config names it, or the SQLite file the user gave.""" + logger = InvokeAILogger.get_logger("import_images") + if self.app_config is not None: + self.database = open_migrated_database(self.app_config, logger) + else: + self.database = Database.open_sqlite(Path(self.database_path), logger, synchronous=self.synchronous) + + def _boards(self): + assert self.database is not None + return BoardRecordStorage(self.database).get_all( + user_id=self.OWNER, + is_admin=True, + order_by=BoardRecordOrderBy.CreatedAt, + direction=SQLiteDirection.Ascending, + include_archived=True, + ) def get_board_names(self): """Get a list of the current board names from the database.""" - sql_get_board_name = "SELECT board_name FROM boards" - self.cursor.execute(sql_get_board_name) - rows = self.cursor.fetchall() - return [row[0] for row in rows] + return [board.board_name for board in self._boards()] def does_image_exist(self, image_name): """Check database if a image name already exists and return a boolean.""" - sql_get_image_by_name = f"SELECT image_name FROM images WHERE image_name='{image_name}'" - self.cursor.execute(sql_get_image_by_name) - rows = self.cursor.fetchall() - return True if len(rows) > 0 else False - - def add_new_image_to_database(self, filename, width, height, metadata, modified_date_string): - """Add an image to the database.""" - sql_add_image = f"""INSERT INTO images (image_name, image_origin, image_category, width, height, session_id, node_id, metadata, is_intermediate, created_at, updated_at) -VALUES ('{filename}', 'internal', 'general', {width}, {height}, null, null, '{metadata}', 0, '{modified_date_string}', '{modified_date_string}')""" - self.cursor.execute(sql_add_image) - self.connection.commit() + assert self.database is not None + return ImageRecordStorage(self.database).exists(image_name) + + def add_new_image_to_database(self, filename, width, height, metadata, modified_date: datetime.datetime): + """Add an image to the database, created when its file was last modified.""" + assert self.database is not None + self.database.queries.images.insert( + image_name=filename, + image_origin=ResourceOrigin.INTERNAL, + image_category=ImageCategory.GENERAL, + width=width, + height=height, + has_workflow=False, + is_intermediate=False, + starred=False, + session_id=None, + node_id=None, + metadata=metadata, + user_id=self.OWNER, + image_subfolder="", + project_id=None, + created_at=timestamp_text(modified_date), + ) def get_board_id_with_create(self, board_name): - """Get the board id for supplied name, and create the board if one does not exist.""" - sql_find_board = f"SELECT board_id FROM boards WHERE board_name='{board_name}' COLLATE NOCASE" - self.cursor.execute(sql_find_board) - rows = self.cursor.fetchall() - if len(rows) > 0: - return rows[0][0] - else: - board_date_string = datetime.datetime.utcnow().date().isoformat() - new_board_id = uuid_string() - sql_insert_board = f"INSERT INTO boards (board_id, board_name, created_at, updated_at) VALUES ('{new_board_id}', '{board_name}', '{board_date_string}', '{board_date_string}')" - self.cursor.execute(sql_insert_board) - self.connection.commit() - return new_board_id + """The id of the oldest board of that name, whoever owns it, ignoring case; a new board if there is none.""" + assert self.database is not None + for board in self._boards(): + if board.board_name.lower() == board_name.lower(): + return board.board_id + return BoardRecordStorage(self.database).save(board_name, self.OWNER).board_id def add_image_to_board(self, filename, board_id): """Add an image mapping to a board.""" - add_datetime_str = datetime.datetime.utcnow().isoformat() - sql_add_image_to_board = f"""INSERT INTO board_images (board_id, image_name, created_at, updated_at) - VALUES ('{board_id}', '{filename}', '{add_datetime_str}', '{add_datetime_str}')""" - self.cursor.execute(sql_add_image_to_board) - self.connection.commit() + assert self.database is not None + BoardImageRecordStorage(self.database).add_image_to_board(board_id=board_id, image_name=filename) def disconnect(self): - """Disconnect from the db, cleaning up connections and cursors.""" - if self.cursor is not None: - self.cursor.close() - if self.connection is not None: - self.connection.close() + """Close the database.""" + if self.database is not None: + self.database.dispose() + self.database = None def backup(self, timestamp_string): """Take a backup of the database.""" + assert self.database is not None + if self.database.dialect_name != "sqlite": + print("The database is a server database: take a backup of it with the server's tools before going on.") + return if not os.path.exists(self.database_backup_dir): print(f"Database backup directory {self.database_backup_dir} does not exist -> creating...", end="") os.makedirs(self.database_backup_dir) print("Done!") database_backup_path = os.path.join(self.database_backup_dir, f"backup-{timestamp_string}-invokeai.db") print(f"Making DB Backup at {database_backup_path}...", end="") - shutil.copy2(self.database_path, database_backup_path) + self.database.backup(Path(database_backup_path)) print("Done!") @@ -692,9 +706,19 @@ def process(self): config = Config() config.find_and_load() - db_mapper = DatabaseMapper(config.database_path, config.database_backup_dir) + db_mapper = DatabaseMapper( + config.database_path, + config.database_backup_dir, + synchronous=config.synchronous, + app_config=config.app_config, + ) db_mapper.connect() + try: + self._import(config, db_mapper) + finally: + db_mapper.disconnect() + def _import(self, config: "Config", db_mapper: DatabaseMapper): import_dir, is_recurse, import_file_list = self.get_import_file_list() ImportStats.count_source_files = len(import_file_list) @@ -745,7 +769,7 @@ def process(self): for filepath in import_file_list: try: self.import_image(filepath, board_name_option, db_mapper, config) - except sqlite3.Error as sql_ex: + except DatabaseError as sql_ex: print(f"A database related exception was found processing {filepath}, will continue to next file. ") print("Exception detail:") print(sql_ex) diff --git a/pyproject.toml b/pyproject.toml index 5dd3e631510..1731b68f978 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,6 +85,12 @@ dependencies = [ "pydantic-settings", "pydantic", "python-socketio", + # Schema migrations that run on every database backend are written with Alembic's operations (no Alembic + # environment or revision scripts). + "alembic>=1.20,<2", + # The database layer compiles every query through SQLAlchemy Core. Keep the minor-version bound: dialect + # compilation and the pysqlite transaction handling it relies on are exactly what changes between minors. + "sqlalchemy>=2.1,<2.2", "uvicorn[standard]", # Auxiliary dependencies, pinned only if necessary. @@ -184,6 +190,8 @@ dependencies = [ "torchvision==0.22.1; sys_platform == 'darwin' or (sys_platform == 'linux' and platform_machine == 'aarch64')", ] +# Driver for running the application database on MySQL or MariaDB instead of SQLite. +"mysql" = ["pymysql>=1.2.3,<2"] "onnx" = ["onnxruntime"] "onnx-cuda" = ["onnxruntime-gpu"] "onnx-directml" = ["onnxruntime-directml"] @@ -299,6 +307,7 @@ explicit = true "invoke-userdel" = "invokeai.app.util.user_management:userdel" "invoke-userlist" = "invokeai.app.util.user_management:userlist" "invoke-usermod" = "invokeai.app.util.user_management:usermod" +"invoke-db-copy" = "invokeai.app.util.database_copy:main" [project.urls] "Homepage" = "https://invoke.ai/" @@ -363,7 +372,11 @@ addopts = "--strict-markers --dist loadfile -m \"not slow\"" markers = [ "slow: Marks tests as slow. Disabled by default. To run all tests, use -m \"\". To run only slow tests, use -m \"slow\".", "timeout: Marks the timeout override.", + "uses_database: Runs on the `empty_database` fixture, which targets INVOKEAI_TEST_DB_URL when it is set and SQLite otherwise. Applied by the collection hook in tests/conftest.py.", + "sqlite_only: Depends on SQLite internals (query plans, PRAGMAs, the single connection); skipped when INVOKEAI_TEST_DB_URL targets another backend.", ] +# A cache-key or other SQLAlchemy warning is a bug in the query layer, not noise. +filterwarnings = ["error::sqlalchemy.exc.SAWarning"] [tool.coverage.run] branch = true source = ["invokeai"] diff --git a/scripts/benchmark_database.py b/scripts/benchmark_database.py new file mode 100644 index 00000000000..129df875f63 --- /dev/null +++ b/scripts/benchmark_database.py @@ -0,0 +1,651 @@ +#!/usr/bin/env python +"""Times representative database operations through the services, to compare two versions of the code. + +Run it once per version and compare the results, e.g. a branch against its base: + + PYTHONPATH= python scripts/benchmark_database.py --json base.json + PYTHONPATH= python scripts/benchmark_database.py --json head.json + python scripts/benchmark_database.py --compare base.json head.json + +The database is a SQLite file in a temporary directory, built by the real migrations and seeded through +the services, so what is timed is the production path including its transaction handling. Each operation +runs once to warm up and then repeatedly; the median and p95 time per call are reported, and the number +of SQL statements per call, counted in a separate pass with statement logging on so that the logging +does not distort the timings. Timings vary between runs on a busy machine: compare runs made back to +back, and repeat a comparison before trusting a small difference. + +A change that ports a service to the database layer renames the storage this script constructs, so the base of +such a change runs its own copy of the script: `python /scripts/benchmark_database.py`. +Operations only the branch times show without a base value. +""" + +import argparse +import json +import logging +import random +import statistics +import sys +import tempfile +import time +import uuid +from collections.abc import Callable +from pathlib import Path +from typing import Any +from unittest import mock + +# A median may grow by this fraction or by this many milliseconds, whichever is larger. The floor is the fixed +# cost of a call through the query layer: SQLAlchemy Core and an explicit transaction make a point read on +# SQLite about 11 microseconds slower than a raw cursor did (measured). Calls that do real work stay bound by +# the fraction. +BUDGET_FRACTION = 0.10 +BUDGET_FLOOR_MS = 0.05 + +ACCOUNTS = 200 +CLIENT_STATE_KEYS = 200 +PROJECTS = 50 +WORKFLOWS = 300 +# One account's library: its own style presets beside the bundled and the shared ones, its system prompts, and the +# wildcards its prompts draw from. +STYLE_PRESETS = 50 +SYSTEM_PROMPTS = 20 +WILDCARDS = 25 +MODELS = 500 +VIDEOS = 2_000 +WORKFLOW_TAGS = ["sdxl", "flux", "upscale", "inpaint", "video", "portrait", "landscape", "controlnet"] + + +class _StatementCounter(logging.Handler): + """Counts the statements SQLite reports through the trace callback a verbose database installs.""" + + def __init__(self) -> None: + super().__init__(level=logging.DEBUG) + self.count = 0 + + def handle(self, record: logging.LogRecord) -> bool: + self.count += 1 + return True + + +class Services: + def __init__(self, db: Any) -> None: + from invokeai.app.services.board_image_records.board_image_records_default import ( + BoardImageRecordStorage, + ) + from invokeai.app.services.board_records.board_records_default import BoardRecordStorage + from invokeai.app.services.board_video_records.board_video_records_default import BoardVideoRecordStorage + from invokeai.app.services.client_state_persistence.client_state_persistence_default import ( + ClientStatePersistence, + ) + from invokeai.app.services.gallery.gallery_default import GalleryService + from invokeai.app.services.image_records.image_records_default import ImageRecordStorage + from invokeai.app.services.model_records import ModelRecordServiceSQL + from invokeai.app.services.project_records.project_records_default import ProjectRecordsStorage + from invokeai.app.services.style_preset_records.style_preset_records_default import StylePresetRecordsStorage + from invokeai.app.services.system_prompt_records.system_prompt_records_default import ( + SystemPromptRecordsStorage, + ) + from invokeai.app.services.users.users_default import UserService + from invokeai.app.services.video_records.video_records_default import VideoRecordStorage + from invokeai.app.services.wildcard_records.wildcard_records_default import WildcardRecordsStorage + from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage + + # A `Database`, or the cursor facade of a base from before it was removed. + self.database = getattr(db, "database", db) + self.image_records = ImageRecordStorage(self.database) + self.video_records = VideoRecordStorage(self.database) + self.board_records = BoardRecordStorage(self.database) + self.board_image_records = BoardImageRecordStorage(self.database) + self.board_video_records = BoardVideoRecordStorage(self.database) + self.gallery = GalleryService(self.database) + self.users = UserService(self.database) + self.client_state = ClientStatePersistence(self.database) + self.project_records = ProjectRecordsStorage(self.database) + self.workflow_records = WorkflowRecordsStorage(self.database) + self.style_preset_records = StylePresetRecordsStorage(self.database) + self.system_prompt_records = SystemPromptRecordsStorage(self.database) + self.wildcard_records = WildcardRecordsStorage(self.database) + self.model_records = ModelRecordServiceSQL(self.database, logging.getLogger("benchmark_database.quiet")) + # Listing gallery items builds their URLs through the invoker's URL service. + from invokeai.app.services.urls.urls_default import LocalUrlService + + invoker = mock.Mock() + invoker.services.urls = LocalUrlService() + invoker.services.logger = logging.getLogger("benchmark_database.quiet") + self.gallery.start(invoker) + # Starting the library stores the bundled workflows and style presets. + self.workflow_records.start(invoker) + self.style_preset_records.start(invoker) + + +def _metadata(rng: random.Random) -> str: + words = ["portrait", "landscape", "cinematic", "volumetric", "light", "detailed", "film", "grain", "neon"] + prompt = " ".join(rng.choice(words) for _ in range(120)) + return json.dumps( + { + "generation_mode": "txt2img", + "positive_prompt": prompt, + "negative_prompt": prompt[:400], + "model": {"key": str(uuid.UUID(int=rng.getrandbits(128))), "name": "Some Model", "base": "flux"}, + "seed": rng.getrandbits(32), + "steps": 30, + "cfg_scale": 3.5, + "width": 1024, + "height": 1024, + "scheduler": "euler", + "loras": [{"key": str(uuid.UUID(int=rng.getrandbits(128))), "weight": 0.75} for _ in range(3)], + } + ) + + +def _project_document(rng: random.Random, names: list[str]) -> dict[str, Any]: + """A canvas of 20 layers, each an image and five brush strokes: about 25 KB of JSON referencing 20 images.""" + layers = [ + { + "id": f"layer_{i}", + "type": "raster_layer", + "isEnabled": True, + "opacity": 1.0, + "position": {"x": rng.randrange(-512, 512), "y": rng.randrange(-512, 512)}, + "objects": [ + {"id": f"image_{i}", "type": "image", "image": {"image_name": rng.choice(names), "width": 1024}}, + *( + { + "id": f"line_{i}_{j}", + "type": "brush_line", + "strokeWidth": 50, + "color": {"r": rng.randrange(256), "g": rng.randrange(256), "b": rng.randrange(256), "a": 1}, + "points": [rng.randrange(1024) for _ in range(40)], + } + for j in range(5) + ), + ], + } + for i in range(20) + ] + return {"canvas": {"layers": layers, "bbox": {"x": 0, "y": 0, "width": 1024, "height": 1024}}} + + +def _workflow_template() -> Any: + """The bundled workflow of median size, as an account's workflow.""" + from invokeai.app.services.workflow_records import workflow_records_common + + bundled = Path(workflow_records_common.__file__).parent / "default_workflows" + paths = sorted(bundled.glob("*.json"), key=lambda path: path.stat().st_size) + workflow = json.loads(paths[len(paths) // 2].read_text(encoding="utf-8")) + workflow.pop("id", None) + workflow["meta"]["category"] = "user" + return workflow_records_common.WorkflowWithoutIDValidator.validate_python(workflow) + + +def _model_config(i: int) -> Any: + """A main model, a VAE or an embedding, by turns.""" + from invokeai.backend.model_manager.configs.main import Main_Diffusers_SDXL_Config + from invokeai.backend.model_manager.configs.textual_inversion import TI_File_SD1_Config + from invokeai.backend.model_manager.configs.vae import VAE_Diffusers_SD1_Config + from invokeai.backend.model_manager.taxonomy import ( + BaseModelType, + ModelRepoVariant, + ModelSourceType, + ModelType, + ModelVariantType, + SchedulerPredictionType, + ) + + common = { + "key": str(uuid.UUID(int=i)), + "path": f"/models/{i}/model", + "name": f"Model {i}", + "hash": f"blake3:{i:064x}", + "file_size": 2_000_000 + i, + "source": f"https://example.com/models/{i}", + "source_type": ModelSourceType.Url, + } + if i % 3 == 0: + return Main_Diffusers_SDXL_Config( + **common, + base=BaseModelType.StableDiffusionXL, + type=ModelType.Main, + variant=ModelVariantType.Normal, + prediction_type=SchedulerPredictionType.Epsilon, + repo_variant=ModelRepoVariant.Default, + ) + if i % 3 == 1: + return VAE_Diffusers_SD1_Config( + **common, base=BaseModelType.StableDiffusion1, type=ModelType.VAE, repo_variant=ModelRepoVariant.Default + ) + return TI_File_SD1_Config(**common, base=BaseModelType.StableDiffusion1, type=ModelType.TextualInversion) + + +def _seed(services: Services, images: int, boards: int, rng: random.Random) -> list[str]: + from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin + from invokeai.app.services.style_preset_records.style_preset_records_common import ( + PresetData, + PresetType, + StylePresetWithoutId, + ) + from invokeai.app.services.system_prompt_records.system_prompt_records_common import SystemPromptWithoutId + from invokeai.app.services.wildcard_records.wildcard_records_common import WildcardWithoutId + + board_ids = [services.board_records.save(board_name=f"Board {i}", user_id="system").board_id for i in range(boards)] + names: list[str] = [] + for i in range(images): + name = f"{uuid.UUID(int=rng.getrandbits(128))}.png" + services.image_records.save( + image_name=name, + image_origin=ResourceOrigin.INTERNAL, + image_category=ImageCategory.GENERAL, + width=1024, + height=1024, + has_workflow=False, + is_intermediate=i % 20 == 0, + starred=i % 50 == 0, + metadata=_metadata(rng), + user_id="system", + ) + names.append(name) + if i % 2 == 0: + services.board_image_records.add_image_to_board(board_id=board_ids[i % boards], image_name=name) + for i in range(VIDEOS): + name = f"{uuid.UUID(int=rng.getrandbits(128))}.mp4" + services.video_records.save( + video_name=name, + video_origin=ResourceOrigin.INTERNAL, + video_category=ImageCategory.GENERAL, + width=1280, + height=720, + duration=5.0, + fps=24.0, + has_workflow=False, + is_intermediate=i % 20 == 0, + starred=i % 50 == 0, + # Every row's listing reads the media origin out of its metadata. + metadata=json.dumps({"media_origin": "audio_upload"}) if i % 10 == 0 else _metadata(rng), + user_id="system", + ) + if i % 2 == 0: + services.board_video_records.add_video_to_board(board_id=board_ids[i % boards], video_name=name) + # Accounts are inserted directly: the service would hash a password for each, which takes most of a second. + with services.database.queries.transaction() as q: + for i in range(ACCOUNTS): + q.users.insert( + user_id=f"user-{i}", + email=f"user{i}@example.com", + display_name=f"User {i}", + password_hash="-", + is_admin=False, + ) + for i in range(CLIENT_STATE_KEYS): + services.client_state.set_by_key( + "system", f"canvas_snapshot:{i}", json.dumps({"imageName": names[i % len(names)]}) + ) + # Each project claims one of the oldest boards, so the boards stay as many as asked for and the newest one, + # whose board operations are timed, belongs to no project. + for i in range(min(PROJECTS, boards - 1)): + services.project_records.create("system", f"Project {i}", _project_document(rng, names), board_id=board_ids[i]) + template = _workflow_template() + for i in range(WORKFLOWS): + tags = ", ".join(rng.sample(WORKFLOW_TAGS, rng.randint(1, 3))) + workflow = template.model_copy(update={"name": f"Workflow {i}", "tags": tags}) + services.workflow_records.create(workflow, user_id="system", is_public=i % 10 == 0) + for i in range(STYLE_PRESETS): + data = PresetData(positive_prompt=f"style {i}, {{prompt}}, highly detailed", negative_prompt="blurry, lowres") + preset = StylePresetWithoutId(name=f"Preset {i}", preset_data=data, type=PresetType.User, is_public=i % 5 == 0) + services.style_preset_records.create(preset, user_id="user-0" if i % 2 else "user-1") + for i in range(SYSTEM_PROMPTS): + prompt = SystemPromptWithoutId(name=f"Prompt {i}", content="Expand the prompt. " * 20) + services.system_prompt_records.create(prompt, user_id="user-0", is_public=i % 4 == 0) + values = [f"value {i}" for i in range(20)] + for i in range(WILDCARDS): + services.wildcard_records.create(WildcardWithoutId(name=f"set{i}", values=values), user_id="user-0") + for i in range(MODELS): + services.model_records.add_model(_model_config(i)) + return names + + +def _operations( + services: Services, names: list[str], rng: random.Random +) -> dict[str, tuple[Callable[[], object], int]]: + """Operation name -> (one call, how many calls to time).""" + from invokeai.app.services.board_records.board_records_common import BoardRecordOrderBy + from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin + from invokeai.app.services.model_records import ModelRecordChanges + + try: + from invokeai.app.services.shared.pagination import SQLiteDirection + except ImportError: # A base from before it moved there. + from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection # type: ignore[no-redef] + from invokeai.app.services.style_preset_records.style_preset_records_common import StylePresetChanges + from invokeai.app.services.system_prompt_records.system_prompt_records_common import SystemPromptChanges + from invokeai.app.services.wildcard_records.wildcard_records_common import WildcardWithoutId + from invokeai.app.services.workflow_records.workflow_records_common import ( + Workflow, + WorkflowCategory, + WorkflowRecordOrderBy, + ) + + sample = iter(rng.choices(names, k=100_000)) + a_board = services.board_records.get_all( + user_id="system", is_admin=True, order_by=BoardRecordOrderBy.CreatedAt, direction=SQLiteDirection.Descending + )[0].board_id + + def board_dto_queries() -> object: + # The queries of `BoardService.get_dto`: the board with the project claiming it, and its counts. + return ( + services.board_records.get_with_project_id(a_board), + services.board_image_records.get_counts_for_board(a_board), + services.board_video_records.get_counts_for_board(a_board), + ) + + project_ids = [summary.project_id for summary in services.project_records.list("system")] + autosave_documents = [_project_document(rng, names) for _ in range(10)] + revision: list[int] = [] + + def autosave() -> object: + # One project saved again and again, each time from the revision the last save returned. The timed and the + # counted pass save the same project: each reads its revision on its first call, the warm-up. + if not revision: + revision.append(services.project_records.get("system", project_ids[0]).revision) + record = services.project_records.update( + "system", project_ids[0], revision[0], "Autosaved", autosave_documents[revision[0] % 10] + ) + revision[0] = record.revision + return record + + library = [WorkflowCategory.User] + workflows = services.workflow_records.get_many( + WorkflowRecordOrderBy.CreatedAt, SQLiteDirection.Ascending, library, user_id="system" + ).items + workflow_ids = [workflow.workflow_id for workflow in workflows] + edited = services.workflow_records.get(workflow_ids[0]).workflow + edits = [Workflow(**edited.model_copy(update={"notes": f"Edit {i}"}).model_dump()) for i in range(10)] + saves = iter(range(10**9)) + + def library_page() -> object: + # What `GET /v1/workflows` asks of the storage for a page of 50: the page, then each workflow, which the + # route checks for being callable. + page = services.workflow_records.get_many( + WorkflowRecordOrderBy.UpdatedAt, SQLiteDirection.Descending, library, page=0, per_page=50 + ) + return [services.workflow_records.get(item.workflow_id) for item in page.items] + + own_presets = [ + preset.id for preset in services.style_preset_records.get_many(user_id="user-0") if preset.user_id == "user-0" + ] + own_prompts = [ + prompt.id for prompt in services.system_prompt_records.get_many(user_id="user-0") if prompt.user_id == "user-0" + ] + wildcard_values = [f"value {i}" for i in range(20)] + model_keys = [str(uuid.UUID(int=i)) for i in range(MODELS)] + + general = [ImageCategory.GENERAL] + video_names = services.video_records.get_video_names(user_id="system", is_admin=True).video_names + + def save_image() -> object: + return services.image_records.save( + image_name=f"{uuid.uuid4()}.png", + image_origin=ResourceOrigin.INTERNAL, + image_category=ImageCategory.GENERAL, + width=1024, + height=1024, + has_workflow=False, + metadata=_metadata(rng), + user_id="system", + ) + + return { + "image_records.get": (lambda: services.image_records.get(next(sample)), 1000), + "image_records.get_many(page of 100)": ( + lambda: services.image_records.get_many( + limit=100, categories=general, is_intermediate=False, user_id="system", is_admin=True + ), + 50, + ), + "image_records.get_image_names(all)": ( + lambda: services.image_records.get_image_names( + categories=general, is_intermediate=False, user_id="system", is_admin=True + ), + 10, + ), + # A non-admin's "all" view: images on no board that it owns, and those on boards it may read. + "image_records.get_image_names(all, account)": ( + lambda: services.image_records.get_image_names( + categories=general, is_intermediate=False, board_id="all", user_id="user-1", is_admin=False + ), + 10, + ), + "video_records.get": (lambda: services.video_records.get(rng.choice(video_names)), 1000), + "video_records.get_many(page of 100)": ( + lambda: services.video_records.get_many( + limit=100, categories=general, is_intermediate=False, user_id="system", is_admin=True + ), + 50, + ), + "video_records.get_video_names(all)": ( + lambda: services.video_records.get_video_names( + categories=general, is_intermediate=False, user_id="system", is_admin=True + ), + 20, + ), + "gallery.get_item_names(all)": ( + lambda: services.gallery.get_item_names( + categories=general, is_intermediate=False, user_id="system", is_admin=True + ), + 10, + ), + "gallery.list_items(page of 100)": ( + lambda: services.gallery.list_items( + limit=100, categories=general, is_intermediate=False, user_id="system", is_admin=True + ), + 50, + ), + "board_records.get_all": ( + lambda: services.board_records.get_all( + user_id="system", + is_admin=True, + order_by=BoardRecordOrderBy.CreatedAt, + direction=SQLiteDirection.Descending, + ), + 50, + ), + "boards.get_dto (record and counts)": (board_dto_queries, 500), + "image_records.save": (save_image, 200), + # Every authenticated request reads its account. + "users.get": (lambda: services.users.get(f"user-{rng.randrange(ACCOUNTS)}"), 1000), + "users.list_users(page of 100)": (lambda: services.users.list_users(limit=100), 50), + "client_state.get_keys_by_prefix(all)": ( + lambda: services.client_state.get_keys_by_prefix("system", "canvas_snapshot:"), + 200, + ), + "client_state.set_by_key": ( + lambda: services.client_state.set_by_key("system", "canvas", json.dumps({"imageName": next(sample)})), + 200, + ), + "projects.get (25 KB document)": (lambda: services.project_records.get("system", rng.choice(project_ids)), 500), + "projects.list": (lambda: services.project_records.list("system"), 200), + "projects.update (autosave)": (autosave, 200), + "workflows.get": (lambda: services.workflow_records.get(rng.choice(workflow_ids)), 500), + "workflows.library page (list + get each)": (library_page, 20), + "workflows.get_many(page of 50)": ( + lambda: services.workflow_records.get_many( + WorkflowRecordOrderBy.Name, SQLiteDirection.Ascending, library, page=0, per_page=50, user_id="system" + ), + 50, + ), + "workflows.update": (lambda: services.workflow_records.update(edits[next(saves) % 10]), 200), + "workflows.counts_by_tag(5)": ( + lambda: services.workflow_records.counts_by_tag(WORKFLOW_TAGS[:5], library, user_id="system"), + 50, + ), + # A prompt-template node reads its preset, a text-LLM node its system prompt, and every prompt expansion + # reads the account's wildcards. + "style_presets.get (prompt template)": ( + lambda: services.style_preset_records.get(own_presets[0]), + 1000, + ), + "style_presets.get_many (account)": ( + lambda: services.style_preset_records.get_many(user_id="user-0"), + 200, + ), + "style_presets.update": ( + lambda: services.style_preset_records.update( + own_presets[1], StylePresetChanges(name=f"Preset {next(saves)}", type=None) + ), + 200, + ), + "system_prompts.get (text LLM)": (lambda: services.system_prompt_records.get(own_prompts[0]), 1000), + "system_prompts.update": ( + lambda: services.system_prompt_records.update( + own_prompts[1], SystemPromptChanges(content=f"Expand {next(saves)}"), user_id="user-0" + ), + 200, + ), + "wildcards.get_many (dynamic prompts)": (lambda: services.wildcard_records.get_many("user-0"), 500), + "wildcards.create": ( + lambda: services.wildcard_records.create( + WildcardWithoutId(name=f"new{uuid.uuid4().hex}", values=wildcard_values), user_id="user-0" + ), + 200, + ), + # Loading a model reads its config, and the model manager lists every model. + "models.get_model": (lambda: services.model_records.get_model(rng.choice(model_keys)), 1000), + "models.exists": (lambda: services.model_records.exists(rng.choice(model_keys)), 1000), + "models.search_by_attr(all)": (lambda: services.model_records.search_by_attr(), 20), + "models.update_model": ( + lambda: services.model_records.update_model( + model_keys[0], ModelRecordChanges(description=f"Edit {next(saves)}") + ), + 200, + ), + } + + +def _time(call: Callable[[], object], repeats: int) -> tuple[float, float]: + call() + durations: list[float] = [] + for _ in range(repeats): + started = time.perf_counter_ns() + call() + durations.append((time.perf_counter_ns() - started) / 1e6) + durations.sort() + return statistics.median(durations), durations[min(len(durations) - 1, int(len(durations) * 0.95))] + + +def run(images: int, boards: int, seed: int) -> dict[str, Any]: + import invokeai.app + from invokeai.app.services.config.config_default import InvokeAIAppConfig + from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase + + try: + from invokeai.app.services.shared.database.database import Database + from invokeai.app.services.shared.database.startup import init_database + + def open_database(path: Path, logger: logging.Logger, verbose: bool = False) -> Any: + return Database.open_sqlite(path, logger, verbose=verbose) + + except ImportError: # A base from before the cursor facade was removed. + from invokeai.app.services.shared.sqlite.sqlite_database import ( # type: ignore[no-redef] + SqliteDatabase as open_database, + ) + from invokeai.app.services.shared.sqlite.sqlite_util import init_db as init_database + + quiet = logging.getLogger("benchmark_database.quiet") + quiet.setLevel(logging.WARNING) + counting = logging.getLogger("benchmark_database.statements") + counting.setLevel(logging.DEBUG) + counting.propagate = False + counter = _StatementCounter() + counting.addHandler(counter) + + # Code from before the database layer cannot close its connection, and Windows refuses to delete an + # open database file: such a run leaves its temporary directory behind. + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + # Seeding commits tens of thousands of times; `normal` keeps that short. The timed connection + # below uses the default `full`, so commits are timed as an install makes them. + config = InvokeAIAppConfig(db_dir=Path(tmp), db_synchronous="normal") + # The migrations clean up files under the root (legacy caches and models): never a real install's. + config._root = Path(tmp) + seeding_db = init_database(config=config, logger=quiet, image_files=mock.Mock(spec=ImageFileStorageBase)) + rng = random.Random(seed) + started = time.perf_counter() + names = _seed(Services(seeding_db), images, boards, rng) + seeded_in = time.perf_counter() - started + + timed_db = open_database(config.db_path, quiet) + counted_db = open_database(config.db_path, counting, verbose=True) + timed = _operations(Services(timed_db), names, random.Random(seed + 1)) + counted = _operations(Services(counted_db), names, random.Random(seed + 1)) + + results: dict[str, dict[str, float]] = {} + for name, (call, repeats) in timed.items(): + median_ms, p95_ms = _time(call, repeats) + count_call, _ = counted[name] + count_call() # warm up the same way the timed pass did + counter.count = 0 + count_call() + results[name] = {"median_ms": median_ms, "p95_ms": p95_ms, "statements": counter.count, "calls": repeats} + + for db in (seeding_db, timed_db, counted_db): + database = getattr(db, "database", db) + if hasattr(database, "dispose"): + database.dispose() + + return { + # `invokeai.app`, because an editable install leaves the top-level package without a `__file__`. + "code": str(Path(invokeai.app.__file__).parent.parent), + "python": sys.version.split()[0], + "images": images, + "boards": boards, + "seeded_in_s": round(seeded_in, 1), + "operations": results, + } + + +def compare(base: dict[str, Any], head: dict[str, Any]) -> int: + print(f"base: {base['code']}\nhead: {head['code']}\n") + print(f"{'operation':40} {'base ms':>9} {'head ms':>9} {'change':>8} {'stmts':>9} budget") + over_budget = 0 + for name, after in head["operations"].items(): + before = base["operations"].get(name) + if before is None: + print(f"{name:40} {'-':>9} {after['median_ms']:9.3f}") + continue + change = (after["median_ms"] - before["median_ms"]) / before["median_ms"] + allowed = max(before["median_ms"] * BUDGET_FRACTION, BUDGET_FLOOR_MS) + within = after["median_ms"] - before["median_ms"] <= allowed + over_budget += not within + statements = f"{before['statements']}->{after['statements']}" + verdict = "ok" if within else "OVER" + print( + f"{name:40} {before['median_ms']:9.3f} {after['median_ms']:9.3f} {change:+8.1%} {statements:>9} {verdict}" + ) + return 1 if over_budget else 0 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--images", type=int, default=20_000) + parser.add_argument("--boards", type=int, default=200) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--json", type=Path, help="write the results to this file") + parser.add_argument("--compare", nargs=2, type=Path, metavar=("BASE", "HEAD"), help="compare two result files") + args = parser.parse_args() + if args.boards < 2: + parser.error("--boards must be at least 2: projects claim the oldest boards, and the newest stays unclaimed") + + if args.compare: + base, head = (json.loads(path.read_text()) for path in args.compare) + return compare(base, head) + + results = run(args.images, args.boards, args.seed) + for name, result in results["operations"].items(): + print( + f"{name:40} median {result['median_ms']:8.3f} ms p95 {result['p95_ms']:8.3f} ms " + f"{result['statements']} statements" + ) + if args.json: + args.json.write_text(json.dumps(results, indent=2)) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/remove_orphaned_models.py b/scripts/remove_orphaned_models.py deleted file mode 100755 index 3b28f59c9ed..00000000000 --- a/scripts/remove_orphaned_models.py +++ /dev/null @@ -1,464 +0,0 @@ -#!/usr/bin/env python -"""Script to remove orphaned model files from INVOKEAI_ROOT directory. - -Orphaned models are ones that appear in the INVOKEAI_ROOT/models directory, -but which are not referenced in the database `models` table. -""" - -import argparse -import datetime -import json -import locale -import os -import shutil -import sqlite3 -from pathlib import Path -from typing import Set - -import yaml - - -class ConfigMapper: - """Configuration loader for InvokeAI paths.""" - - YAML_FILENAME = "invokeai.yaml" - DATABASE_FILENAME = "invokeai.db" - DEFAULT_DB_DIR = "databases" - DEFAULT_MODELS_DIR = "models" - - def __init__(self): - self.database_path = None - self.database_backup_dir = None - self.models_path = None - - def load(self, root_path: Path) -> bool: - """Load configuration from root directory.""" - yaml_path = root_path / self.YAML_FILENAME - if not yaml_path.exists(): - print(f"Unable to find {self.YAML_FILENAME} at {yaml_path}!") - return False - - db_dir, models_dir = self._load_paths_from_yaml_file(yaml_path) - - if db_dir is None: - db_dir = self.DEFAULT_DB_DIR - print(f"The {self.YAML_FILENAME} file was found but is missing the db_dir setting! Defaulting to {db_dir}") - - if models_dir is None: - models_dir = self.DEFAULT_MODELS_DIR - print( - f"The {self.YAML_FILENAME} file was found but is missing the models_dir setting! Defaulting to {models_dir}" - ) - - # Set database path - if os.path.isabs(db_dir): - self.database_path = Path(db_dir) / self.DATABASE_FILENAME - else: - self.database_path = root_path / db_dir / self.DATABASE_FILENAME - - self.database_backup_dir = self.database_path.parent / "backup" - - # Set models path - if os.path.isabs(models_dir): - self.models_path = Path(models_dir) - else: - self.models_path = root_path / models_dir - - db_exists = self.database_path.exists() - models_exists = self.models_path.exists() - - print(f"Found {self.YAML_FILENAME} file at {yaml_path}:") - print(f" Database : {self.database_path} - {'Exists!' if db_exists else 'Not Found!'}") - print(f" Models : {self.models_path} - {'Exists!' if models_exists else 'Not Found!'}") - - if db_exists and models_exists: - return True - else: - print( - "\nOne or more paths specified in invokeai.yaml do not exist. Please inspect/correct the configuration." - ) - return False - - def _load_paths_from_yaml_file(self, yaml_path: Path): - """Load paths from YAML configuration file.""" - try: - with open(yaml_path, "rt", encoding=locale.getpreferredencoding()) as file: - yamlinfo = yaml.safe_load(file) - db_dir = yamlinfo.get("InvokeAI", {}).get("Paths", {}).get("db_dir", None) - models_dir = yamlinfo.get("InvokeAI", {}).get("Paths", {}).get("models_dir", None) - return db_dir, models_dir - except Exception as e: - print(f"Failed to load paths from yaml file! {yaml_path}! Error: {e}") - return None, None - - -class DatabaseMapper: - """Class to abstract database functionality.""" - - def __init__(self, database_path: Path, database_backup_dir: Path): - self.database_path = database_path - self.database_backup_dir = database_backup_dir - self.connection = None - self.cursor = None - - def backup(self, timestamp_string: str): - """Take a backup of the database.""" - if not self.database_backup_dir.exists(): - print(f"Database backup directory {self.database_backup_dir} does not exist -> creating...", end="") - self.database_backup_dir.mkdir(parents=True, exist_ok=True) - print("Done!") - - database_backup_path = self.database_backup_dir / f"backup-{timestamp_string}-invokeai.db" - print(f"Making DB Backup at {database_backup_path}...", end="") - shutil.copy2(self.database_path, database_backup_path) - print("Done!") - - def connect(self): - """Open connection to the database.""" - self.connection = sqlite3.connect(str(self.database_path)) - self.cursor = self.connection.cursor() - - def get_all_model_directories(self, models_dir: Path) -> Set[Path]: - """Get the set of all model directories from the database. - - A model directory is the top-level directory under models/ that contains - the model files. If the path in the database is just a directory, that's - the model directory. If it's a file path, we extract the first directory - component. - - Args: - models_dir: The root models directory path. Relative paths from the database - will be resolved relative to this directory. - - Returns: - Set of absolute Path objects for model directories. - """ - sql_get_models = "SELECT config FROM models" - self.cursor.execute(sql_get_models) - rows = self.cursor.fetchall() - model_directories = set() - for row in rows: - try: - config = json.loads(row[0]) - if "path" in config and config["path"]: - path_str = config["path"] - # Convert to Path object - path = Path(path_str) - - # If the path is relative, resolve it relative to models_dir - # If it's absolute, use it as-is - if not path.is_absolute(): - full_path = (models_dir / path).resolve() - else: - full_path = path.resolve() - - # Extract the top-level directory under models_dir - # This handles both cases: - # 1. path is "model-id" -> model-id is the directory - # 2. path is "model-id/file.safetensors" -> model-id is the directory - try: - # Get the relative path from models_dir - rel_path = full_path.relative_to(models_dir) - # Get the first component (top-level directory) - if rel_path.parts: - top_level_dir = models_dir / rel_path.parts[0] - model_directories.add(top_level_dir.resolve()) - except ValueError: - # Path is not relative to models_dir, use the path itself - # This handles absolute paths outside models_dir - model_directories.add(full_path) - - except (json.JSONDecodeError, KeyError, TypeError) as e: - print(f"Warning: Failed to parse model config: {e}") - continue - return model_directories - - def disconnect(self): - """Disconnect from the database.""" - if self.cursor is not None: - self.cursor.close() - if self.connection is not None: - self.connection.close() - - -class ModelFileMapper: - """Class to handle model file system operations.""" - - # Common model file extensions - MODEL_EXTENSIONS = { - ".safetensors", - ".ckpt", - ".pt", - ".pth", - ".bin", - ".onnx", - } - - # Directories to skip during scan - SKIP_DIRS = { - ".download_cache", - ".convert_cache", - "__pycache__", - ".git", - } - - def __init__(self, models_path: Path): - self.models_path = models_path - - def get_all_model_directories(self) -> Set[Path]: - """ - Get all directories in the models path that contain model files. - Returns a set of directory paths that contain at least one model file. - """ - model_dirs = set() - - for item in self.models_path.rglob("*"): - # Skip directories we don't want to scan - if any(skip_dir in item.parts for skip_dir in self.SKIP_DIRS): - continue - - if item.is_file() and item.suffix.lower() in self.MODEL_EXTENSIONS: - # Add the parent directory of the model file - model_dirs.add(item.parent) - - return model_dirs - - def get_all_model_files(self) -> Set[Path]: - """Get all model files in the models directory.""" - model_files = set() - - for item in self.models_path.rglob("*"): - # Skip directories we don't want to scan - if any(skip_dir in item.parts for skip_dir in self.SKIP_DIRS): - continue - - if item.is_file() and item.suffix.lower() in self.MODEL_EXTENSIONS: - model_files.add(item.resolve()) - - return model_files - - def remove_file(self, file_path: Path): - """Remove a single model file.""" - try: - file_path.unlink() - print(f" Deleted file: {file_path}") - except Exception as e: - print(f" Error deleting {file_path}: {e}") - - def remove_directory_if_empty(self, directory: Path): - """Remove a directory if it's empty (after removing files).""" - try: - if directory.exists() and not any(directory.iterdir()): - directory.rmdir() - print(f" Deleted empty directory: {directory}") - except Exception as e: - print(f" Error removing directory {directory}: {e}") - - -class OrphanedModelsApp: - """Main application class for removing orphaned model files.""" - - def __init__(self, delete_without_confirm: bool = False): - self.delete_without_confirm = delete_without_confirm - self.orphaned_count = 0 - - def find_orphaned_files_by_directory( - self, file_mapper: ModelFileMapper, db_mapper: DatabaseMapper, models_path: Path - ) -> dict[Path, list[Path]]: - """Find orphaned files grouped by their parent directory. - - A file is orphaned if it's NOT under any model directory registered in the database. - Model directories are extracted from the database paths - if a path is - 'model-id/file.safetensors', then 'model-id' is the model directory and ALL files - under it belong to that model. - """ - print("\nScanning models directory for orphaned models...") - - # Get all model files on disk - disk_model_files = file_mapper.get_all_model_files() - print(f"Found {len(disk_model_files)} model directories on disk") - - # Get all model directories from database - db_model_directories = db_mapper.get_all_model_directories(models_path) - print(f"Found {len(db_model_directories)} model directories in database") - - # Find orphaned files (files on disk but not under any registered model directory) - orphaned_files = set() - for disk_file in disk_model_files: - # Check if this file is under any registered model directory - is_under_model_dir = False - for model_dir in db_model_directories: - try: - # Check if disk_file is under model_dir - disk_file.relative_to(model_dir) - is_under_model_dir = True - break - except ValueError: - # Not under this model directory, continue checking - continue - - if not is_under_model_dir: - orphaned_files.add(disk_file) - - # Group orphaned files by their parent directory - orphaned_dirs = {} - for orphaned_file in orphaned_files: - parent = orphaned_file.parent - if parent not in orphaned_dirs: - orphaned_dirs[parent] = [] - orphaned_dirs[parent].append(orphaned_file) - - return orphaned_dirs - - def ask_to_continue(self) -> bool: - """Ask user whether they want to continue with the operation.""" - while True: - try: - input_choice = input("\nDo you wish to delete these models? (Y or N) [N]: ") - # Default to 'N' if user presses Enter without input - if input_choice.strip() == "": - return False - if str.lower(input_choice) == "y": - return True - if str.lower(input_choice) == "n": - return False - print("Please enter Y or N") - except (KeyboardInterrupt, EOFError): - return False - - def remove_orphaned_models(self, config: ConfigMapper, file_mapper: ModelFileMapper, db_mapper: DatabaseMapper): - """Remove orphaned model directories.""" - print("\n" + "=" * 80) - print("= Remove Orphaned Model Files") - print("=" * 80) - print("\nThis operation will find model files in the models directory that are not") - print("referenced in the database and remove them.") - print() - print(f"Database File Path : {config.database_path}") - print(f"Models Directory : {config.models_path}") - print() - print("Notes:") - print("- A database backup will be created before any changes") - print("- Model files not referenced in the database will be permanently deleted") - print("- This operation cannot be undone (except by restoring the deleted files)") - print() - - # Connect to database and find orphaned files - db_mapper.connect() - try: - orphaned_dirs = self.find_orphaned_files_by_directory(file_mapper, db_mapper, config.models_path) - - if not orphaned_dirs: - print("\nNo orphaned model files found!") - return - - print(f"\nFound {len(orphaned_dirs)} directories with orphaned model files:") - print() - - for directory, files in sorted(orphaned_dirs.items()): - print(f"Directory: {directory}") - for file in sorted(files): - print(f" - {file.name}") - print() - - self.orphaned_count = sum(len(files) for files in orphaned_dirs.values()) - print(f"Total orphans: {self.orphaned_count}") - - # Ask for confirmation unless --delete flag is used - if not self.delete_without_confirm: - if not self.ask_to_continue(): - print("\nOperation cancelled by user.") - self.orphaned_count = 0 # Reset count since no files were removed - return - - # Create database backup with timestamp - timestamp_string = datetime.datetime.now(datetime.UTC).strftime("%Y%m%dT%H%M%SZ") - db_mapper.backup(timestamp_string) - - # Delete the orphaned files - print("\nDeleting orphaned model files...") - for directory, files in sorted(orphaned_dirs.items()): - for file in sorted(files): - file_mapper.remove_file(file) - # After removing files, clean up the directory if it's now empty - file_mapper.remove_directory_if_empty(directory) - - finally: - db_mapper.disconnect() - - def main(self, root_path: Path): - """Main entry point.""" - print("\n" + "=" * 80) - print("Orphaned Model Files Cleanup for InvokeAI") - print("=" * 80 + "\n") - - config_mapper = ConfigMapper() - if not config_mapper.load(root_path): - print("\nInvalid configuration...exiting.\n") - return 1 - - file_mapper = ModelFileMapper(config_mapper.models_path) - db_mapper = DatabaseMapper(config_mapper.database_path, config_mapper.database_backup_dir) - - try: - self.remove_orphaned_models(config_mapper, file_mapper, db_mapper) - except KeyboardInterrupt: - print("\n\nOperation cancelled by user.") - return 1 - except Exception as e: - print(f"\n\nError during operation: {e}") - import traceback - - traceback.print_exc() - return 1 - - print("\n" + "=" * 80) - print("= Operation Complete") - print("=" * 80) - print(f"\nOrphaned model files removed: {self.orphaned_count}") - print() - - return 0 - - -def main(): - """Command-line entry point.""" - parser = argparse.ArgumentParser( - description="Remove orphaned model files from InvokeAI installation", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=""" -This script finds and removes model files that exist in the models directory -but are not referenced in the InvokeAI database. This can happen if: -- Models were manually deleted from the database -- The database was reset but model files were kept -- Files were manually copied into the models directory - -By default, the script will list orphaned files and ask for confirmation -before deleting them. -""", - ) - parser.add_argument( - "--root", - type=Path, - default=os.environ.get("INVOKEAI_ROOT", "."), - help="InvokeAI root directory (default: $INVOKEAI_ROOT or current directory)", - ) - parser.add_argument( - "--delete", - action="store_true", - help="Delete orphan model files without asking for confirmation", - ) - args = parser.parse_args() - - # Resolve the root path - root_path = Path(args.root).resolve() - if not root_path.exists(): - print(f"Error: Root directory does not exist: {root_path}") - return 1 - - app = OrphanedModelsApp(delete_without_confirm=args.delete) - return app.main(root_path) - - -if __name__ == "__main__": - exit(main()) diff --git a/tests/app/routers/conftest.py b/tests/app/routers/conftest.py index 08de579f71f..1221872b69e 100644 --- a/tests/app/routers/conftest.py +++ b/tests/app/routers/conftest.py @@ -31,6 +31,7 @@ from invokeai.app.api.dependencies import ApiDependencies from invokeai.app.api_app import app from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.users.users_common import UserCreateRequest @@ -106,13 +107,13 @@ def _create_board(client: TestClient, token: str, name: str = "Loose+Board") -> @pytest.fixture -def enable_multiuser(monkeypatch: Any, mock_invoker: Invoker): +def enable_multiuser(monkeypatch: Any, mock_invoker: Invoker, mock_sqlite_database: Database): """Enable multiuser mode and patch ApiDependencies across the routers covered by router-level tests. Replaces None-valued services with MagicMocks so that routes can run end-to-end. """ - from invokeai.app.services.style_preset_records.style_preset_records_sqlite import ( - SqliteStylePresetRecordsStorage, + from invokeai.app.services.style_preset_records.style_preset_records_default import ( + StylePresetRecordsStorage, ) mock_invoker.services.configuration.multiuser = True @@ -126,9 +127,7 @@ def enable_multiuser(monkeypatch: Any, mock_invoker: Invoker): # Style preset records uses a real SQLite-backed storage on the same in-memory # database that image_records was wired up against. This lets cross-user tests # exercise the actual filter SQL instead of asserting on MagicMock calls. - mock_invoker.services.style_preset_records = SqliteStylePresetRecordsStorage( - db=mock_invoker.services.image_records._db - ) + mock_invoker.services.style_preset_records = StylePresetRecordsStorage(mock_sqlite_database) # Required by board_image_records-touching helpers in some tests. if mock_invoker.services.board_images is None: diff --git a/tests/app/routers/test_app_info.py b/tests/app/routers/test_app_info.py index 3dc41d06518..7b8b24b86af 100644 --- a/tests/app/routers/test_app_info.py +++ b/tests/app/routers/test_app_info.py @@ -1,9 +1,11 @@ import os +from collections.abc import Callable from pathlib import Path from typing import Any from unittest.mock import Mock import pytest +from fastapi import HTTPException from fastapi.testclient import TestClient from invokeai.app.api.dependencies import ApiDependencies @@ -14,6 +16,7 @@ from invokeai.app.services.external_generation.external_generation_common import ExternalProviderStatus from invokeai.app.services.image_files.image_subfolder_strategy import DateStrategy, create_subfolder_strategy from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.errors import LockTimeoutError from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig, ExternalModelCapabilities from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType @@ -604,3 +607,42 @@ def test_reset_external_provider_config_rejects_non_admin_users( def _get_provider_config(payload: list[dict[str, Any]], provider_id: str) -> dict[str, Any]: return next(item for item in payload if item["provider_id"] == provider_id) + + +def _busy() -> None: + raise LockTimeoutError("lock wait timeout exceeded") + + +def _busy_behind_a_routes_own_error() -> None: + # As many routes do: every exception becomes an answer of their own. + try: + _busy() + except Exception: + raise HTTPException(status_code=404, detail="Board not found") + + +@pytest.mark.parametrize("route", [_busy, _busy_behind_a_routes_own_error]) +def test_a_busy_database_answers_503_with_a_retry_hint(client: TestClient, route: Callable[[], None]) -> None: + app.add_api_route("/api/v1/test/busy", route) + try: + response = client.get("/api/v1/test/busy") + finally: + app.router.routes.pop() + + assert response.status_code == 503 + assert response.headers["Retry-After"] == "1" + assert response.json() == {"detail": "The database is busy; try again"} + + +def test_a_routes_own_error_is_answered_as_it_is(client: TestClient) -> None: + def missing() -> None: + raise HTTPException(status_code=404, detail="Board not found") + + app.add_api_route("/api/v1/test/missing", missing) + try: + response = client.get("/api/v1/test/missing") + finally: + app.router.routes.pop() + + assert response.status_code == 404 + assert response.json() == {"detail": "Board not found"} diff --git a/tests/app/routers/test_board_images_batch_races.py b/tests/app/routers/test_board_images_batch_races.py index 90c05dc0063..6188d220ba2 100644 --- a/tests/app/routers/test_board_images_batch_races.py +++ b/tests/app/routers/test_board_images_batch_races.py @@ -20,6 +20,7 @@ from invokeai.app.api.dependencies import ApiDependencies from invokeai.app.api_app import app from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database class MockApiDependencies(ApiDependencies): @@ -193,36 +194,30 @@ def write_access(board_id: str, current_user: object) -> None: assert set(body["failed_images"]) == {"second.png", "third.png"} -def test_sqlite_remove_returns_the_row_count() -> None: +def test_remove_returns_the_row_count(database: Database) -> None: """The storage layer itself: the row count is the only signal the scoped DELETE's scope held, and a `None` return silently restores report-a-removal-that-did-not-happen at the - route (`None == 0` is False). Stubbed at the cursor in the manner of the image-records - storage tests.""" - from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage - - storage = SqliteBoardImageRecordStorage.__new__(SqliteBoardImageRecordStorage) - - class _Cursor: - rowcount = 0 - - def execute(self, *args: object, **kwargs: object) -> None: - pass - - class _Db: - def transaction(self): - from contextlib import contextmanager - - @contextmanager - def _cm(): - yield _Cursor() - - return _cm() - - storage._db = _Db() # pyright: ignore[reportAttributeAccessIssue] + route (`None == 0` is False).""" + from sqlalchemy import insert + + from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage + from invokeai.app.services.board_records.board_records_default import BoardRecordStorage + from invokeai.app.services.shared.database.schema.images import images + + boards = BoardRecordStorage(database) + on_board, other_board = boards.save("P", "system").board_id, boards.save("Q", "system").board_id + with database.begin(write=True) as conn: + conn.execute( + insert(images).values( + image_name="raced.png", image_origin="internal", image_category="general", width=1, height=1 + ) + ) + storage = BoardImageRecordStorage(database) + storage.add_image_to_board(on_board, "raced.png") - assert storage.remove_image_from_board("raced.png", "board-p") == 0 - _Cursor.rowcount = 1 - assert storage.remove_image_from_board("ok.png", "board-p") == 1 + assert storage.remove_image_from_board("raced.png", other_board) == 0 + assert storage.remove_image_from_board("raced.png", on_board) == 1 + assert storage.remove_image_from_board("raced.png", on_board) == 0 @pytest.mark.parametrize( diff --git a/tests/app/routers/test_boards_multiuser.py b/tests/app/routers/test_boards_multiuser.py index b184c61deb8..a7061acf7b6 100644 --- a/tests/app/routers/test_boards_multiuser.py +++ b/tests/app/routers/test_boards_multiuser.py @@ -13,7 +13,9 @@ from invokeai.app.api.routers.boards import delete_board from invokeai.app.api_app import app from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.users.users_common import UserCreateRequest +from tests.fixtures.sqlite_database import sqlite_cursor class MockApiDependencies(ApiDependencies): @@ -82,7 +84,7 @@ def enable_multiuser_for_tests(monkeypatch: Any, mock_invoker: Invoker): # so the route doesn't hit AttributeError on the None placeholders in mock_services. mock_board_video_records = MagicMock() mock_board_video_records.get_all_board_video_names_for_board.return_value = [] - mock_board_video_records.get_video_count_for_board.return_value = 0 + mock_board_video_records.get_counts_for_board.return_value = (0, 0) mock_invoker.services.board_video_records = mock_board_video_records # The board service also consults video_records for cover-image selection (most recent video). mock_video_records = MagicMock() @@ -997,12 +999,12 @@ def test_generic_update_of_a_project_board_is_refused( def test_setting_only_the_cover_of_a_project_board_is_still_allowed( - client: TestClient, mock_invoker: Invoker, user1_token: str + client: TestClient, mock_invoker: Invoker, mock_sqlite_database: Database, user1_token: str ): """The cover is a display detail with no bearing on the project relationship.""" board_id = _claim_board_for_a_project(client, mock_invoker, user1_token, "Cover+Board") # `boards.cover_image_name` is a foreign key, so the image has to actually exist. - with mock_invoker.services.board_records._db.transaction() as cursor: + with sqlite_cursor(mock_sqlite_database) as cursor: cursor.execute( "INSERT INTO images (image_name, image_origin, image_category, width, height)" " VALUES ('cover.png', 'internal', 'general', 64, 64);" diff --git a/tests/app/routers/test_fonts.py b/tests/app/routers/test_fonts.py index e3fb2d179e6..f8d0fc404d7 100644 --- a/tests/app/routers/test_fonts.py +++ b/tests/app/routers/test_fonts.py @@ -9,12 +9,13 @@ from invokeai.app.services.fonts.fonts_default import FontService from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database @pytest.fixture -def font_service(mock_invoker: Invoker, tmp_path: Path) -> FontService: +def font_service(mock_invoker: Invoker, mock_sqlite_database: Database, tmp_path: Path) -> FontService: service = FontService( - db=mock_invoker.services.image_records._db, + mock_sqlite_database, fonts_dir=tmp_path / "fonts", storage_dir=tmp_path / "uploaded-fonts", ) diff --git a/tests/app/routers/test_gallery_item_location.py b/tests/app/routers/test_gallery_item_location.py index 1fc8687017d..f8cd3383ce3 100644 --- a/tests/app/routers/test_gallery_item_location.py +++ b/tests/app/routers/test_gallery_item_location.py @@ -9,6 +9,7 @@ import httpx import pytest from fastapi import FastAPI, HTTPException, status +from sqlalchemy import update from invokeai.app.api.auth_dependencies import get_current_user_or_default from invokeai.app.api.routers import gallery as gallery_router_module @@ -17,7 +18,9 @@ from invokeai.app.services.gallery.gallery_common import GalleryItemKind, GalleryItemLocation from invokeai.app.services.image_records.image_records_common import ImageCategory, ImageRecordChanges, ResourceOrigin from invokeai.app.services.invoker import Invoker -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.videos import videos +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.services.video_records.video_records_common import VideoRecordChanges @@ -50,8 +53,9 @@ def _save_video(invoker: Invoker, name: str, user_id: str) -> None: def _set_created_at(invoker: Invoker, table: str, name_column: str, name: str, created_at: str) -> None: - with invoker.services.image_records._db.transaction() as cursor: - cursor.execute(f"UPDATE {table} SET created_at = ? WHERE {name_column} = ?", (created_at, name)) + media = {"images": images, "videos": videos}[table] + with invoker.services.database.begin(write=True) as conn: + conn.execute(update(media).where(media.c[name_column] == name).values(created_at=created_at)) def _locate( diff --git a/tests/app/routers/test_gallery_maintenance.py b/tests/app/routers/test_gallery_maintenance.py index 38cef480bc9..0dc4f2a61e2 100644 --- a/tests/app/routers/test_gallery_maintenance.py +++ b/tests/app/routers/test_gallery_maintenance.py @@ -230,9 +230,10 @@ def test_disconnected_execute_keeps_reservation_until_completion_without_retry( ) -> None: """An abandoned HTTP response must not release the operation's image-storage reservation.""" service = _patch_dependencies(monkeypatch, mock_invoker) - image_moves = ImageMoveService(db=Mock(), image_files=Mock(), config=Mock(), logger=Mock()) + database = Mock() + database.queries.image_moves.active_job_id.return_value = None + image_moves = ImageMoveService(database=database, image_files=Mock(), config=Mock(), logger=Mock()) image_moves._invoker = mock_invoker - monkeypatch.setattr(image_moves, "_get_active_job_id", lambda: None) monkeypatch.setattr(image_moves, "_assert_no_active_queue_work", lambda: None) mock_invoker.services.image_moves = image_moves started = Event() diff --git a/tests/app/routers/test_image_map.py b/tests/app/routers/test_image_map.py index 1594e1db7da..7fe0531e9c9 100644 --- a/tests/app/routers/test_image_map.py +++ b/tests/app/routers/test_image_map.py @@ -11,32 +11,33 @@ import pytest from fastapi.testclient import TestClient from PIL import Image +from sqlalchemy import insert from invokeai.app.api.dependencies import ApiDependencies from invokeai.app.api_app import app -from invokeai.app.services.board_video_records.board_video_records_sqlite import SqliteBoardVideoRecordStorage +from invokeai.app.services.board_video_records.board_video_records_default import BoardVideoRecordStorage from invokeai.app.services.config.config_default import InvokeAIAppConfig -from invokeai.app.services.gallery.gallery_default import SqliteGalleryService +from invokeai.app.services.gallery.gallery_default import GalleryService from invokeai.app.services.image_index.image_index_base import ImageIndexServiceBase from invokeai.app.services.image_index.image_index_common import ( ImageIndexStatus, IndexedItem, ) -from invokeai.app.services.image_index.image_index_records_sqlite import ImageIndexRecordsSqlite +from invokeai.app.services.image_index.image_index_records_default import ImageIndexRecords from invokeai.app.services.image_index.projection import ( DEFAULT_CLUSTER_MIN_SAMPLES, cluster_with_diagnostics, scope_hash, ) from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin -from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage from invokeai.app.services.invocation_services import InvocationServices from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.image_index import image_embeddings from invokeai.app.services.users.users_common import UserCreateRequest -from invokeai.app.services.video_records.video_records_sqlite import SqliteVideoRecordStorage +from invokeai.app.services.video_records.video_records_default import VideoRecordStorage from invokeai.app.services.videos.videos_default import VideoService -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database MODEL_ID = "test-model-hash" DIM = 4 @@ -65,7 +66,7 @@ class FakeImageIndexService(ImageIndexServiceBase): def __init__(self, model_id: str | None = MODEL_ID) -> None: self._model_id = model_id - self.index_records: ImageIndexRecordsSqlite | None = None + self.index_records: ImageIndexRecords | None = None self.projection_requests: list[tuple[str, bool]] = [] self.spent_failed_scopes: dict[str, str] = {} self.search_calls: list[tuple[str | None, int]] = [] @@ -184,7 +185,7 @@ def image_index_service() -> FakeImageIndexService: return FakeImageIndexService() -def _video_service(thumbnails: Path, video_records: SqliteVideoRecordStorage) -> VideoService: +def _video_service(thumbnails: Path, video_records: VideoRecordStorage) -> VideoService: """A video service that resolves thumbnails to real files, which is all these endpoints read. It goes through the record store first, like the real one: a name with no video raises @@ -205,36 +206,37 @@ def get_path(video_name: str, thumbnail: bool = False) -> str: @pytest.fixture -def mock_services(image_index_service: FakeImageIndexService, tmp_path: Path) -> InvocationServices: - from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage - from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage +def mock_services( + image_index_service: FakeImageIndexService, tmp_path: Path, mock_sqlite_database: Database +) -> InvocationServices: + from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage + from invokeai.app.services.board_records.board_records_default import BoardRecordStorage from invokeai.app.services.boards.boards_default import BoardService from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService - from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ( - ClientStatePersistenceSqlite, + from invokeai.app.services.client_state_persistence.client_state_persistence_default import ( + ClientStatePersistence, ) from invokeai.app.services.images.images_default import ImageService from invokeai.app.services.invocation_cache.invocation_cache_memory import MemoryInvocationCache from invokeai.app.services.invocation_stats.invocation_stats_default import InvocationStatsService - from invokeai.app.services.project_records.project_records_sqlite import ProjectRecordsSqlite + from invokeai.app.services.project_records.project_records_default import ProjectRecordsStorage from invokeai.app.services.users.users_default import UserService from tests.test_nodes import TestEventService configuration = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0) - logger = InvokeAILogger.get_logger() - db = create_mock_sqlite_database(configuration, logger) + db = mock_sqlite_database services = InvocationServices( - board_image_records=SqliteBoardImageRecordStorage(db=db), + board_image_records=BoardImageRecordStorage(db), board_images=None, # type: ignore - board_records=SqliteBoardRecordStorage(db=db), + board_records=BoardRecordStorage(db), boards=BoardService(), bulk_download=BulkDownloadService(), configuration=configuration, database=db, events=TestEventService(), image_files=None, # type: ignore - image_records=SqliteImageRecordStorage(db=db), + image_records=ImageRecordStorage(db), images=ImageService(), invocation_cache=MemoryInvocationCache(max_cache_size=0), logger=logging, # type: ignore @@ -254,17 +256,17 @@ def mock_services(image_index_service: FakeImageIndexService, tmp_path: Path) -> workflow_thumbnails=None, # type: ignore model_relationship_records=None, # type: ignore model_relationships=None, # type: ignore - client_state_persistence=ClientStatePersistenceSqlite(db=db), - project_records=ProjectRecordsSqlite(db=db), + client_state_persistence=ClientStatePersistence(db), + project_records=ProjectRecordsStorage(db), users=UserService(db), wildcard_records=None, # type: ignore system_prompt_records=None, # type: ignore - videos=_video_service(tmp_path, video_records := SqliteVideoRecordStorage(db=db)), + videos=_video_service(tmp_path, video_records := VideoRecordStorage(db)), video_files=None, # type: ignore video_records=video_records, - board_video_records=SqliteBoardVideoRecordStorage(db=db), - gallery=SqliteGalleryService(db=db), - image_index_records=(index_records := ImageIndexRecordsSqlite(db=db)), + board_video_records=BoardVideoRecordStorage(db), + gallery=GalleryService(db), + image_index_records=(index_records := ImageIndexRecords(db)), image_index=image_index_service, intermediates=None, # type: ignore external_generation=None, # type: ignore @@ -289,7 +291,7 @@ def client(monkeypatch, mock_invoker: Invoker) -> TestClient: return TestClient(app) -def _records(mock_invoker: Invoker) -> ImageIndexRecordsSqlite: +def _records(mock_invoker: Invoker) -> ImageIndexRecords: return mock_invoker.services.image_index_records @@ -1723,7 +1725,13 @@ def test_multiuser_image_labels_enforce_read_access(multiuser, mock_invoker: Inv assert response.status_code == 403 -def _seed_degenerate_embedding(mock_invoker: Invoker, image_name: str, vector: np.ndarray) -> None: +def _insert_embedding_row(db: Database, image_name: str, dim: int, blob: bytes) -> None: + """Write an embedding row straight to the table, bypassing the writer's guards.""" + with db.begin(write=True) as conn: + conn.execute(insert(image_embeddings).values(image_name=image_name, model_id=MODEL_ID, dim=dim, embedding=blob)) + + +def _seed_degenerate_embedding(mock_invoker: Invoker, db: Database, image_name: str, vector: np.ndarray) -> None: """Write an embedding blob straight to the table, bypassing the writer's guards. `embedding_to_blob` refuses non-finite and all-zero vectors, but rows @@ -1731,13 +1739,8 @@ def _seed_degenerate_embedding(mock_invoker: Invoker, image_name: str, vector: n validates only the blob's length. """ _save_unembedded_image(mock_invoker, image_name) - records = _records(mock_invoker) blob = np.ascontiguousarray(vector, dtype=np.float32).tobytes() - with records._db.transaction() as cursor: - cursor.execute( - "INSERT INTO image_embeddings (image_name, model_id, dim, embedding) VALUES (?, ?, ?, ?);", - (image_name, MODEL_ID, vector.shape[0], blob), - ) + _insert_embedding_row(db, image_name, vector.shape[0], blob) @pytest.mark.parametrize( @@ -1748,7 +1751,7 @@ def _seed_degenerate_embedding(mock_invoker: Invoker, image_name: str, vector: n ], ) def test_image_labels_refuse_a_degenerate_stored_embedding( - mock_invoker: Invoker, client: TestClient, vector: np.ndarray + mock_invoker: Invoker, client: TestClient, mock_sqlite_database: Database, vector: np.ndarray ) -> None: """A degenerate row must not yield confident nonsense. @@ -1757,7 +1760,7 @@ def test_image_labels_refuse_a_degenerate_stored_embedding( phrases presented as this image's tags, with a `score` that serializes as JSON null against a schema declaring it a float. """ - _seed_degenerate_embedding(mock_invoker, "degenerate.png", vector) + _seed_degenerate_embedding(mock_invoker, mock_sqlite_database, "degenerate.png", vector) response = client.get("/api/v1/image_map/image_labels", params={"image_name": "degenerate.png"}) @@ -1768,7 +1771,9 @@ def test_image_labels_refuse_a_degenerate_stored_embedding( assert response.json()["detail"] == "This item's stored embedding cannot be labeled" -def test_image_labels_corrupt_blob_is_404_not_500(mock_invoker: Invoker, client: TestClient) -> None: +def test_image_labels_corrupt_blob_is_404_not_500( + mock_invoker: Invoker, client: TestClient, mock_sqlite_database: Database +) -> None: """A row whose blob length disagrees with its `dim` column. `blob_to_embedding` raises on it. Unhandled that is a 500, and because @@ -1776,12 +1781,7 @@ def test_image_labels_corrupt_blob_is_404_not_500(mock_invoker: Invoker, client: a second for as long as the point is hovered. """ _save_unembedded_image(mock_invoker, "corrupt.png") - records = _records(mock_invoker) - with records._db.transaction() as cursor: - cursor.execute( - "INSERT INTO image_embeddings (image_name, model_id, dim, embedding) VALUES (?, ?, ?, ?);", - ("corrupt.png", MODEL_ID, DIM, np.zeros(DIM - 1, dtype=np.float32).tobytes()), - ) + _insert_embedding_row(mock_sqlite_database, "corrupt.png", DIM, np.zeros(DIM - 1, dtype=np.float32).tobytes()) response = client.get("/api/v1/image_map/image_labels", params={"image_name": "corrupt.png"}) @@ -1789,17 +1789,12 @@ def test_image_labels_corrupt_blob_is_404_not_500(mock_invoker: Invoker, client: def test_image_labels_dim_mismatch_with_the_vocabulary_is_404_not_500( - mock_invoker: Invoker, client: TestClient + mock_invoker: Invoker, client: TestClient, mock_sqlite_database: Database ) -> None: """A stored embedding from a different-width model than the vocab matrix.""" _save_unembedded_image(mock_invoker, "wide.png") - records = _records(mock_invoker) wide = np.ones(DIM + 3, dtype=np.float32) - with records._db.transaction() as cursor: - cursor.execute( - "INSERT INTO image_embeddings (image_name, model_id, dim, embedding) VALUES (?, ?, ?, ?);", - ("wide.png", MODEL_ID, wide.shape[0], wide.tobytes()), - ) + _insert_embedding_row(mock_sqlite_database, "wide.png", wide.shape[0], wide.tobytes()) response = client.get("/api/v1/image_map/image_labels", params={"image_name": "wide.png"}) diff --git a/tests/app/routers/test_image_moves.py b/tests/app/routers/test_image_moves.py index 3fc814d7f32..0be0d44bb88 100644 --- a/tests/app/routers/test_image_moves.py +++ b/tests/app/routers/test_image_moves.py @@ -8,28 +8,27 @@ from invokeai.app.api.dependencies import ApiDependencies from invokeai.app.api_app import app from invokeai.app.services.auth.token_service import set_jwt_secret -from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage -from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage -from invokeai.app.services.board_video_records.board_video_records_sqlite import SqliteBoardVideoRecordStorage +from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage +from invokeai.app.services.board_video_records.board_video_records_default import BoardVideoRecordStorage from invokeai.app.services.boards.boards_default import BoardService from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService -from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ClientStatePersistenceSqlite +from invokeai.app.services.client_state_persistence.client_state_persistence_default import ClientStatePersistence from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.image_moves.image_moves_default import ImageMoveJobAlreadyRunning, ImageMoveQueueActive -from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage from invokeai.app.services.images.images_default import ImageService from invokeai.app.services.invocation_cache.invocation_cache_memory import MemoryInvocationCache from invokeai.app.services.invocation_services import InvocationServices from invokeai.app.services.invocation_stats.invocation_stats_default import InvocationStatsService from invokeai.app.services.invoker import Invoker -from invokeai.app.services.project_records.project_records_sqlite import ProjectRecordsSqlite +from invokeai.app.services.project_records.project_records_default import ProjectRecordsStorage +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.users.users_common import UserCreateRequest from invokeai.app.services.users.users_default import UserService -from invokeai.app.services.video_records.video_records_sqlite import SqliteVideoRecordStorage -from invokeai.app.services.wildcard_records.wildcard_records_sqlite import SqliteWildcardRecordsStorage -from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database +from invokeai.app.services.video_records.video_records_default import VideoRecordStorage +from invokeai.app.services.wildcard_records.wildcard_records_default import WildcardRecordsStorage +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage from tests.test_nodes import TestEventService @@ -46,22 +45,21 @@ def client() -> TestClient: @pytest.fixture -def mock_services() -> InvocationServices: +def mock_services(mock_sqlite_database: Database) -> InvocationServices: configuration = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0) - logger = InvokeAILogger.get_logger() - db = create_mock_sqlite_database(configuration, logger) + db = mock_sqlite_database image_moves = MagicMock() return InvocationServices( - board_image_records=SqliteBoardImageRecordStorage(db=db), + board_image_records=BoardImageRecordStorage(db), board_images=None, # type: ignore - board_records=SqliteBoardRecordStorage(db=db), + board_records=BoardRecordStorage(db), boards=BoardService(), bulk_download=BulkDownloadService(), configuration=configuration, database=db, events=TestEventService(), image_files=None, # type: ignore - image_records=SqliteImageRecordStorage(db=db), + image_records=ImageRecordStorage(db), images=ImageService(), invocation_cache=MemoryInvocationCache(max_cache_size=0), logger=logging, # type: ignore @@ -74,7 +72,7 @@ def mock_services() -> InvocationServices: session_processor=None, # type: ignore session_queue=None, # type: ignore urls=None, # type: ignore - workflow_records=SqliteWorkflowRecordsStorage(db=db), + workflow_records=WorkflowRecordsStorage(db), tensors=None, # type: ignore conditioning=None, # type: ignore style_preset_records=None, # type: ignore @@ -82,14 +80,14 @@ def mock_services() -> InvocationServices: workflow_thumbnails=None, # type: ignore model_relationship_records=None, # type: ignore model_relationships=None, # type: ignore - client_state_persistence=ClientStatePersistenceSqlite(db=db), - project_records=ProjectRecordsSqlite(db=db), + client_state_persistence=ClientStatePersistence(db), + project_records=ProjectRecordsStorage(db), users=UserService(db), - wildcard_records=SqliteWildcardRecordsStorage(db=db), + wildcard_records=WildcardRecordsStorage(db), videos=None, # type: ignore video_files=None, # type: ignore - video_records=SqliteVideoRecordStorage(db=db), - board_video_records=SqliteBoardVideoRecordStorage(db=db), + video_records=VideoRecordStorage(db), + board_video_records=BoardVideoRecordStorage(db), gallery=None, # type: ignore image_index_records=None, # type: ignore image_index=None, # type: ignore diff --git a/tests/app/routers/test_images.py b/tests/app/routers/test_images.py index cb73ef5d6b6..b07d76ed357 100644 --- a/tests/app/routers/test_images.py +++ b/tests/app/routers/test_images.py @@ -1,6 +1,4 @@ import os -import sqlite3 -from contextlib import contextmanager from pathlib import Path from typing import Any from unittest.mock import MagicMock @@ -18,6 +16,8 @@ from invokeai.app.services.image_records.image_records_common import ImageNamesResult, ImageRecordNotFoundException from invokeai.app.services.images.images_common import ImageDTO from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import LockTimeoutError from invokeai.app.services.shared.pagination import MAX_PAGE_SIZE, OffsetPaginatedResults @@ -471,7 +471,7 @@ def failing_get_dto(image_name: str): def test_delete_image_db_fault_during_lookup_returns_500_not_404( - monkeypatch: Any, mock_invoker: Invoker, tmp_path: Path, client: TestClient + monkeypatch: Any, mock_invoker: Invoker, mock_sqlite_database: Database, tmp_path: Path, client: TestClient ) -> None: """A database fault while reading the record is a 500, driven through the real record store. @@ -483,13 +483,13 @@ def test_delete_image_db_fault_during_lookup_returns_500_not_404( storage = prepare_delete_image_test(monkeypatch, mock_invoker, tmp_path) _save_deletable_image(mock_invoker, storage, "del.png") - # Break the table out from under the query. Any sqlite3.Error would do; this one is deterministic. + # Break the table out from under the query. Any database error would do; this one is deterministic. records = mock_invoker.services.image_records - records._db._conn.execute("ALTER TABLE images RENAME TO images_moved;") + mock_sqlite_database.sqlite.conn.execute("ALTER TABLE images RENAME TO images_moved;") try: response = client.delete("/api/v1/images/i/del.png") finally: - records._db._conn.execute("ALTER TABLE images_moved RENAME TO images;") + mock_sqlite_database.sqlite.conn.execute("ALTER TABLE images_moved RENAME TO images;") assert response.status_code == 500 assert response.json()["detail"] == "Failed to delete image" @@ -498,24 +498,35 @@ def test_delete_image_db_fault_during_lookup_returns_500_not_404( assert storage.get_path("del.png").exists() -def test_delete_image_db_failure_returns_500_and_restores_files( - monkeypatch: Any, mock_invoker: Invoker, tmp_path: Path, client: TestClient +@pytest.mark.parametrize( + ("failure", "status", "detail"), + [ + (LockTimeoutError("database is locked"), 503, "The database is busy; try again"), + (RuntimeError("disk I/O error"), 500, "Failed to delete image"), + ], +) +def test_delete_image_db_failure_is_reported_and_restores_files( + monkeypatch: Any, + mock_invoker: Invoker, + tmp_path: Path, + client: TestClient, + failure: Exception, + status: int, + detail: str, ) -> None: - from invokeai.app.services.image_records.image_records_common import ImageRecordDeleteException - storage = prepare_delete_image_test(monkeypatch, mock_invoker, tmp_path) _save_deletable_image(mock_invoker, storage, "del.png") def failing_delete(image_name: str) -> None: - raise ImageRecordDeleteException() + raise failure monkeypatch.setattr(mock_invoker.services.image_records, "delete", failing_delete) response = client.delete("/api/v1/images/i/del.png") - # The route must report the failure, not a success-shaped empty payload. - assert response.status_code == 500 - assert response.json()["detail"] == "Failed to delete image" + # The route must report the failure, not a success-shaped empty payload; a busy database as such. + assert response.status_code == status + assert response.json()["detail"] == detail # The staged files must be rolled back: image and thumbnail restored, record intact. assert storage.get_path("del.png").exists() assert storage.get_path("del.png", thumbnail=True).exists() @@ -710,27 +721,6 @@ def get_dto(image_name: str) -> MagicMock: assert body["failed_images"] == [] -def test_image_records_get_does_not_disguise_a_storage_error_as_not_found(monkeypatch: Any) -> None: - """The narrowing itself: a sqlite3.Error out of the SELECT must stay a sqlite3.Error.""" - from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage - - storage = SqliteImageRecordStorage.__new__(SqliteImageRecordStorage) - - class _Cursor: - def execute(self, *args: Any, **kwargs: Any) -> None: - raise sqlite3.OperationalError("database disk image is malformed") - - class _Db: - @contextmanager - def transaction(self): - yield _Cursor() - - storage._db = _Db() # pyright: ignore[reportAttributeAccessIssue] - - with pytest.raises(sqlite3.OperationalError): - storage.get("a.png") - - def test_delete_still_reports_a_genuine_storage_failure( monkeypatch: Any, mock_invoker: Invoker, client: TestClient ) -> None: diff --git a/tests/app/routers/test_intermediates.py b/tests/app/routers/test_intermediates.py index 3552533e8de..701b5631465 100644 --- a/tests/app/routers/test_intermediates.py +++ b/tests/app/routers/test_intermediates.py @@ -12,6 +12,8 @@ from invokeai.app.services.image_files.image_files_disk import DiskImageFileStorage from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database +from tests.fixtures.sqlite_database import sqlite_cursor @pytest.fixture @@ -36,7 +38,7 @@ def _user_id(mock_invoker: Invoker, email: str) -> str: return user.user_id -def _seed_intermediate(mock_invoker: Invoker, name: str, user_id: str) -> None: +def _seed_intermediate(mock_invoker: Invoker, database: Database, name: str, user_id: str) -> None: mock_invoker.services.image_records.save( image_name=name, image_origin=ResourceOrigin.INTERNAL, @@ -47,7 +49,7 @@ def _seed_intermediate(mock_invoker: Invoker, name: str, user_id: str) -> None: is_intermediate=True, user_id=user_id, ) - with mock_invoker.services.image_records._db.transaction() as cursor: + with sqlite_cursor(database) as cursor: cursor.execute("UPDATE images SET created_at = '2020-01-01 00:00:00.000' WHERE image_name = ?;", (name,)) mock_invoker.services.image_files.save(image=Image.new("RGB", (8, 8)), image_name=name) @@ -83,10 +85,15 @@ def test_every_route_requires_authentication(enable_multiuser: Any, client: Test def test_non_admins_see_only_their_rows_and_cannot_widen_scope( - storage_ready: None, mock_invoker: Invoker, client: TestClient, user1_token: str, user2_token: str + storage_ready: None, + mock_invoker: Invoker, + mock_sqlite_database: Database, + client: TestClient, + user1_token: str, + user2_token: str, ) -> None: - _seed_intermediate(mock_invoker, "u1.png", _user_id(mock_invoker, "user1@test.com")) - _seed_intermediate(mock_invoker, "u2.png", _user_id(mock_invoker, "user2@test.com")) + _seed_intermediate(mock_invoker, mock_sqlite_database, "u1.png", _user_id(mock_invoker, "user1@test.com")) + _seed_intermediate(mock_invoker, mock_sqlite_database, "u2.png", _user_id(mock_invoker, "user2@test.com")) summary = client.get("/api/v1/intermediates/summary", headers=_auth(user1_token)) assert summary.status_code == 200 @@ -117,10 +124,16 @@ def test_non_admins_see_only_their_rows_and_cannot_widen_scope( def test_preview_operation_and_list_flow( - storage_ready: None, mock_invoker: Invoker, client: TestClient, user1_token: str, user2_token: str, admin_token: str + storage_ready: None, + mock_invoker: Invoker, + mock_sqlite_database: Database, + client: TestClient, + user1_token: str, + user2_token: str, + admin_token: str, ) -> None: user1 = _user_id(mock_invoker, "user1@test.com") - _seed_intermediate(mock_invoker, "u1.png", user1) + _seed_intermediate(mock_invoker, mock_sqlite_database, "u1.png", user1) preview = client.post( "/api/v1/intermediates/previews", @@ -184,10 +197,15 @@ def test_mutations_are_refused_during_image_storage_maintenance( def test_hold_is_account_scoped_at_the_http_boundary( - storage_ready: None, mock_invoker: Invoker, client: TestClient, user1_token: str, user2_token: str + storage_ready: None, + mock_invoker: Invoker, + mock_sqlite_database: Database, + client: TestClient, + user1_token: str, + user2_token: str, ) -> None: user1 = _user_id(mock_invoker, "user1@test.com") - _seed_intermediate(mock_invoker, "held.png", user1) + _seed_intermediate(mock_invoker, mock_sqlite_database, "held.png", user1) assert ( client.put( "/api/v1/intermediates/holds/tab", json={"images": ["held.png"]}, headers=_auth(user1_token) @@ -214,11 +232,16 @@ def test_hold_lease_ids_are_single_tokens(storage_ready: None, client: TestClien def test_legacy_clear_keeps_its_shape_and_the_safety_policy( - storage_ready: None, mock_invoker: Invoker, client: TestClient, user1_token: str, admin_token: str + storage_ready: None, + mock_invoker: Invoker, + mock_sqlite_database: Database, + client: TestClient, + user1_token: str, + admin_token: str, ) -> None: user1 = _user_id(mock_invoker, "user1@test.com") - _seed_intermediate(mock_invoker, "safe.png", user1) - _seed_intermediate(mock_invoker, "referenced.png", user1) + _seed_intermediate(mock_invoker, mock_sqlite_database, "safe.png", user1) + _seed_intermediate(mock_invoker, mock_sqlite_database, "referenced.png", user1) mock_invoker.services.project_records.create(user1, "P", {"imageName": "referenced.png"}) # The count is what the clear would delete, so the legacy button settles at zero. diff --git a/tests/app/routers/test_media_copy.py b/tests/app/routers/test_media_copy.py index 75c9af2f32d..0d86164d688 100644 --- a/tests/app/routers/test_media_copy.py +++ b/tests/app/routers/test_media_copy.py @@ -22,8 +22,10 @@ from invokeai.app.services.images.images_common import ImageDTO from invokeai.app.services.invoker import Invoker from invokeai.app.services.names.names_default import SimpleNameService +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.urls.urls_default import LocalUrlService from tests.app.routers.conftest import _auth, _create_board +from tests.fixtures.sqlite_database import sqlite_cursor SOURCE_METADATA = {"positive_prompt": "a cat", "seed": 12345} SOURCE_WORKFLOW = '{"name": "a workflow"}' @@ -34,9 +36,9 @@ def _owner_id(client: TestClient, token: str) -> str: return board.json()["user_id"] -def _insert_image_record(mock_invoker: Invoker, name: str, user_id: str, category: str = "control") -> None: +def _insert_image_record(db: Database, name: str, user_id: str, category: str = "control") -> None: """A real row, so the route's read-access check runs against real data.""" - with mock_invoker.services.board_records._db.transaction() as cursor: + with sqlite_cursor(db) as cursor: cursor.execute( "INSERT INTO images (image_name, image_origin, image_category, width, height, user_id)" " VALUES (?, 'internal', ?, 64, 64, ?);", @@ -219,12 +221,17 @@ def test_copying_into_a_board_you_cannot_write_is_refused( def test_copying_someone_elses_image_is_refused_per_name( - client: TestClient, mock_invoker: Invoker, user1_token: str, user2_token: str, real_images: DiskImageFileStorage + client: TestClient, + mock_invoker: Invoker, + mock_sqlite_database: Database, + user1_token: str, + user2_token: str, + real_images: DiskImageFileStorage, ): mine = _owner_id(client, user1_token) theirs = _owner_id(client, user2_token) source = _create_source_image(mock_invoker, mine) - _insert_image_record(mock_invoker, "theirs.png", theirs) + _insert_image_record(mock_sqlite_database, "theirs.png", theirs) response = _copy_images(client, user1_token, image_names=[source.image_name, "theirs.png"]) @@ -245,9 +252,9 @@ def test_a_batch_larger_than_the_cap_is_refused( assert response.status_code == status.HTTP_422_UNPROCESSABLE_CONTENT -def _insert_video_record(mock_invoker: Invoker, name: str, user_id: str) -> None: +def _insert_video_record(db: Database, name: str, user_id: str) -> None: """A real row, so the board service's cover resolution keeps working around it.""" - with mock_invoker.services.board_records._db.transaction() as cursor: + with sqlite_cursor(db) as cursor: cursor.execute( "INSERT INTO videos (video_name, video_origin, video_category, width, height, duration, fps, user_id)" " VALUES (?, 'internal', 'general', 640, 480, 2.5, 24.0, ?);", @@ -256,11 +263,11 @@ def _insert_video_record(mock_invoker: Invoker, name: str, user_id: str) -> None def test_copying_a_video_delegates_copy_invariants_to_the_video_service( - client: TestClient, mock_invoker: Invoker, user1_token: str + client: TestClient, mock_invoker: Invoker, mock_sqlite_database: Database, user1_token: str ): user_id = _owner_id(client, user1_token) board_id = _create_board(client, user1_token, "Video+Target") - _insert_video_record(mock_invoker, "src.mp4", user_id) + _insert_video_record(mock_sqlite_database, "src.mp4", user_id) videos = MagicMock() created = MagicMock() @@ -281,13 +288,13 @@ def test_copying_a_video_delegates_copy_invariants_to_the_video_service( def test_a_video_copy_that_missed_its_board_is_reported_as_failed( - client: TestClient, mock_invoker: Invoker, user1_token: str + client: TestClient, mock_invoker: Invoker, mock_sqlite_database: Database, user1_token: str ): """`create` treats board attachment as best-effort, which is right for a generation and wrong here: the caller is about to remap a document onto the name we return.""" user_id = _owner_id(client, user1_token) board_id = _create_board(client, user1_token, "Video+Target") - _insert_video_record(mock_invoker, "src.mp4", user_id) + _insert_video_record(mock_sqlite_database, "src.mp4", user_id) videos = MagicMock() videos.copy.side_effect = RuntimeError("copy missed board") @@ -302,9 +309,11 @@ def test_a_video_copy_that_missed_its_board_is_reported_as_failed( assert response.json() == {"copied": [], "failed": ["src.mp4"]} -def test_copying_someone_elses_video_is_refused_per_name(client: TestClient, mock_invoker: Invoker, user1_token: str): +def test_copying_someone_elses_video_is_refused_per_name( + client: TestClient, mock_invoker: Invoker, mock_sqlite_database: Database, user1_token: str +): _owner_id(client, user1_token) - _insert_video_record(mock_invoker, "theirs.mp4", "somebody-else") + _insert_video_record(mock_sqlite_database, "theirs.mp4", "somebody-else") videos = MagicMock() mock_invoker.services.videos = videos diff --git a/tests/app/routers/test_model_manager.py b/tests/app/routers/test_model_manager.py index 0de53e29d55..ae1e5f7b957 100644 --- a/tests/app/routers/test_model_manager.py +++ b/tests/app/routers/test_model_manager.py @@ -14,7 +14,8 @@ ExternalModelCapabilities, ExternalModelPanelSchema, ) -from invokeai.backend.model_manager.taxonomy import ModelType +from invokeai.backend.model_manager.configs.textual_inversion import TI_File_SD1_Config +from invokeai.backend.model_manager.taxonomy import ModelSourceType, ModelType @pytest.fixture(autouse=True, scope="module") @@ -1413,3 +1414,40 @@ def test_an_encoded_key_reaches_the_image_route_intact(client: TestClient, tmp_p assert response.status_code == 200 deps.invoker.services.model_images.get_path.assert_called_once_with("X?y") + + +def test_the_orphan_routes_scan_with_the_model_records( + monkeypatch: Any, client: TestClient, mm2_model_manager: Any, mm2_app_config: Any +) -> None: + models_path: Path = mm2_app_config.models_path + for directory in ("registered", "unregistered"): + (models_path / directory).mkdir(parents=True, exist_ok=True) + (models_path / directory / "model.safetensors").write_bytes(b"not really a model") + mm2_model_manager.store.add_model( + TI_File_SD1_Config( + path="registered/model.safetensors", + name="registered", + hash="ABC123", + file_size=18, + source="test/source/", + source_type=ModelSourceType.Path, + ) + ) + services = type("Services", (), {})() + services.model_manager = mm2_model_manager + services.configuration = mm2_app_config + invoker = DummyInvoker(services) + monkeypatch.setattr("invokeai.app.api.routers.model_manager.ApiDependencies", MockApiDependencies(invoker)) + monkeypatch.setattr("invokeai.app.api.auth_dependencies.ApiDependencies", MockApiDependencies(invoker)) + + found = client.get("/api/v2/models/sync/orphaned") + + assert found.status_code == 200, found.text + reported = {orphan["path"] for orphan in found.json()} + assert "unregistered" in reported and "registered" not in reported + + deleted = client.request("DELETE", "/api/v2/models/sync/orphaned", json={"paths": ["unregistered"]}) + + assert deleted.status_code == 200, deleted.text + assert deleted.json() == {"deleted": ["unregistered"], "errors": {}} + assert (models_path / "registered").exists() and not (models_path / "unregistered").exists() diff --git a/tests/app/routers/test_model_manager_authorization.py b/tests/app/routers/test_model_manager_authorization.py index cf1cc911829..3a90e70f730 100644 --- a/tests/app/routers/test_model_manager_authorization.py +++ b/tests/app/routers/test_model_manager_authorization.py @@ -400,9 +400,11 @@ def test_runtime_config_redacts_secrets() -> None: external_gemini_base_url="https://example.invalid", remote_api_tokens=[URLRegexTokenPair(url_regex="example.com", token="bearer-secret")], download_proxy="http://proxy-user:proxy-secret@proxy.example:3128", + db_url="mariadb+pymysql://invokeai:db-secret@db.example/invokeai", ) redacted = _redact_config_secrets(config) + assert redacted.db_url == "mariadb+pymysql://invokeai:***@db.example/invokeai" assert redacted.external_openai_api_key == REDACTED_SECRET assert redacted.remote_api_tokens is not None assert redacted.remote_api_tokens[0].token == REDACTED_SECRET diff --git a/tests/app/routers/test_multiuser_authorization.py b/tests/app/routers/test_multiuser_authorization.py index 14a3d55c9aa..cf12b881833 100644 --- a/tests/app/routers/test_multiuser_authorization.py +++ b/tests/app/routers/test_multiuser_authorization.py @@ -22,12 +22,12 @@ from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.invocation_services import InvocationServices from invokeai.app.services.invoker import Invoker -from invokeai.app.services.project_records.project_records_sqlite import ProjectRecordsSqlite +from invokeai.app.services.project_records.project_records_default import ProjectRecordsStorage from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem, SessionQueueItemSummary +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.users.users_common import UserCreateRequest -from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage +from tests.fixtures.sqlite_database import sqlite_cursor, sqlite_cursor_of class MockApiDependencies(ApiDependencies): @@ -74,42 +74,41 @@ def _mock_urls() -> MagicMock: @pytest.fixture -def mock_services() -> InvocationServices: - from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage - from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage - from invokeai.app.services.board_video_records.board_video_records_sqlite import SqliteBoardVideoRecordStorage +def mock_services(mock_sqlite_database: Database) -> InvocationServices: + from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage + from invokeai.app.services.board_records.board_records_default import BoardRecordStorage + from invokeai.app.services.board_video_records.board_video_records_default import BoardVideoRecordStorage from invokeai.app.services.boards.boards_default import BoardService from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService - from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ( - ClientStatePersistenceSqlite, + from invokeai.app.services.client_state_persistence.client_state_persistence_default import ( + ClientStatePersistence, ) - from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage + from invokeai.app.services.image_records.image_records_default import ImageRecordStorage from invokeai.app.services.images.images_default import ImageService from invokeai.app.services.invocation_cache.invocation_cache_memory import MemoryInvocationCache from invokeai.app.services.invocation_stats.invocation_stats_default import InvocationStatsService - from invokeai.app.services.system_prompt_records.system_prompt_records_sqlite import ( - SqliteSystemPromptRecordsStorage, + from invokeai.app.services.system_prompt_records.system_prompt_records_default import ( + SystemPromptRecordsStorage, ) from invokeai.app.services.users.users_default import UserService - from invokeai.app.services.video_records.video_records_sqlite import SqliteVideoRecordStorage - from invokeai.app.services.wildcard_records.wildcard_records_sqlite import SqliteWildcardRecordsStorage + from invokeai.app.services.video_records.video_records_default import VideoRecordStorage + from invokeai.app.services.wildcard_records.wildcard_records_default import WildcardRecordsStorage from tests.test_nodes import TestEventService configuration = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0) - logger = InvokeAILogger.get_logger() - db = create_mock_sqlite_database(configuration, logger) + db = mock_sqlite_database return InvocationServices( - board_image_records=SqliteBoardImageRecordStorage(db=db), + board_image_records=BoardImageRecordStorage(db), board_images=None, # type: ignore - board_records=SqliteBoardRecordStorage(db=db), + board_records=BoardRecordStorage(db), boards=BoardService(), bulk_download=BulkDownloadService(), configuration=configuration, database=db, events=TestEventService(), image_files=None, # type: ignore - image_records=SqliteImageRecordStorage(db=db), + image_records=ImageRecordStorage(db), images=ImageService(), invocation_cache=MemoryInvocationCache(max_cache_size=0), logger=logging, # type: ignore @@ -126,24 +125,24 @@ def mock_services() -> InvocationServices: # every name was skipped before the ownership check ran, and the test passed no # matter what the route did. The returns must be strings; ImageDTO validates them. urls=_mock_urls(), - workflow_records=SqliteWorkflowRecordsStorage(db=db), + workflow_records=WorkflowRecordsStorage(db), tensors=None, # type: ignore conditioning=None, # type: ignore style_preset_records=None, # type: ignore style_preset_image_files=None, # type: ignore - system_prompt_records=SqliteSystemPromptRecordsStorage(db=db), + system_prompt_records=SystemPromptRecordsStorage(db), workflow_thumbnails=None, # type: ignore model_relationship_records=None, # type: ignore model_relationships=None, # type: ignore - client_state_persistence=ClientStatePersistenceSqlite(db=db), - project_records=ProjectRecordsSqlite(db=db), + client_state_persistence=ClientStatePersistence(db), + project_records=ProjectRecordsStorage(db), users=UserService(db), - wildcard_records=SqliteWildcardRecordsStorage(db=db), + wildcard_records=WildcardRecordsStorage(db), external_generation=None, # type: ignore videos=None, # type: ignore video_files=None, # type: ignore - video_records=SqliteVideoRecordStorage(db=db), - board_video_records=SqliteBoardVideoRecordStorage(db=db), + video_records=VideoRecordStorage(db), + board_video_records=BoardVideoRecordStorage(db), gallery=None, # type: ignore image_index_records=None, # type: ignore image_index=None, # type: ignore @@ -249,10 +248,9 @@ def _share_board(client: TestClient, token: str, board_id: str) -> None: assert r.status_code == status.HTTP_201_CREATED -def _share_board_with_user(mock_invoker: Invoker, board_id: str, user_id: str) -> None: +def _share_board_with_user(db: Database, board_id: str, user_id: str) -> None: """Insert an explicit per-user share row directly into the shared_boards table.""" - board_records: Any = mock_invoker.services.board_records - with board_records._db.transaction() as cursor: + with sqlite_cursor(db) as cursor: cursor.execute( "INSERT OR IGNORE INTO shared_boards (board_id, user_id) VALUES (?, ?)", (board_id, user_id), @@ -281,7 +279,7 @@ def _insert_pending_queue_item(session_queue: Any, user_id: str, queue_id: str = graph.add_node(PromptTestInvocation(id="prompt", prompt="test")) session = GraphExecutionState(graph=graph) session_json = session.model_dump_json(warnings=False, exclude_none=True) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -1132,7 +1130,12 @@ def test_shared_board_image_readable_by_other_user( assert r.status_code != status.HTTP_403_FORBIDDEN def test_explicit_share_grants_image_read_to_member( - self, client: TestClient, mock_invoker: Invoker, user1_token: str, user2_token: str + self, + client: TestClient, + mock_invoker: Invoker, + mock_sqlite_database: Database, + user1_token: str, + user2_token: str, ): """An image on a private board explicitly shared with user2 should be readable by user2 — matching the visibility the board_id="all" listing already grants.""" @@ -1146,7 +1149,7 @@ def test_explicit_share_grants_image_read_to_member( r = client.get("/api/v1/images/i/explicit-share-img", headers=_auth(user2_token)) assert r.status_code == status.HTTP_403_FORBIDDEN - _share_board_with_user(mock_invoker, board_id, user2.user_id) + _share_board_with_user(mock_sqlite_database, board_id, user2.user_id) r = client.get("/api/v1/images/i/explicit-share-img", headers=_auth(user2_token)) assert r.status_code != status.HTTP_403_FORBIDDEN @@ -1209,14 +1212,19 @@ def test_list_images_shared_board_allowed_for_non_owner( assert r.status_code == status.HTTP_200_OK def test_list_images_explicitly_shared_private_board_allowed_for_member( - self, client: TestClient, mock_invoker: Invoker, user1_token: str, user2_token: str + self, + client: TestClient, + mock_invoker: Invoker, + mock_sqlite_database: Database, + user1_token: str, + user2_token: str, ): """User2 should be able to list a private board explicitly shared with them, on both the DTO and names endpoints.""" user2 = mock_invoker.services.users.get_by_email("user2@test.com") assert user2 is not None board_id = _create_board(client, user1_token, "Explicit Share Enum Board") - _share_board_with_user(mock_invoker, board_id, user2.user_id) + _share_board_with_user(mock_sqlite_database, board_id, user2.user_id) r = client.get(f"/api/v1/images/?board_id={board_id}", headers=_auth(user2_token)) assert r.status_code == status.HTTP_200_OK @@ -1225,7 +1233,12 @@ def test_list_images_explicitly_shared_private_board_allowed_for_member( assert r.status_code == status.HTTP_200_OK def test_explicitly_shared_private_board_dto_and_image_names_allowed_for_member( - self, client: TestClient, mock_invoker: Invoker, user1_token: str, user2_token: str + self, + client: TestClient, + mock_invoker: Invoker, + mock_sqlite_database: Database, + user1_token: str, + user2_token: str, ): """User2 should be able to fetch the board DTO and board image_names of a private board explicitly shared with them; a non-member (user1's other boards aside) stays 403.""" @@ -1238,7 +1251,7 @@ def test_explicitly_shared_private_board_dto_and_image_names_allowed_for_member( r = client.get(f"/api/v1/boards/{board_id}/image_names", headers=_auth(user2_token)) assert r.status_code == status.HTTP_403_FORBIDDEN - _share_board_with_user(mock_invoker, board_id, user2.user_id) + _share_board_with_user(mock_sqlite_database, board_id, user2.user_id) r = client.get(f"/api/v1/boards/{board_id}", headers=_auth(user2_token)) assert r.status_code == status.HTTP_200_OK @@ -1704,7 +1717,7 @@ class TestWorkflowListScoping: """Tests that listing workflows in multiuser mode does not filter out default workflows.""" def test_default_workflows_visible_when_listing_user_and_default( - self, client: TestClient, mock_invoker: Invoker, user1_token: str + self, client: TestClient, mock_invoker: Invoker, mock_sqlite_database: Database, user1_token: str ): """When categories=['user','default'], default workflows must still appear even though user_id_filter is set to the current user (default workflows belong to 'system').""" @@ -1732,7 +1745,7 @@ def test_default_workflows_visible_when_listing_user_and_default( ) wf_with_id = Workflow(**default_wf.model_dump(), id=uuid_string()) # Insert directly via DB since the create API rejects default workflows - with mock_invoker.services.workflow_records._db.transaction() as cursor: + with sqlite_cursor(mock_sqlite_database) as cursor: cursor.execute( "INSERT INTO workflow_library (workflow_id, workflow, user_id) VALUES (?, ?, ?)", (wf_with_id.id, wf_with_id.model_dump_json(), "system"), @@ -1755,7 +1768,7 @@ def test_default_workflows_visible_when_listing_user_and_default( assert "user" in categories_found def test_default_workflows_visible_when_no_category_filter( - self, client: TestClient, mock_invoker: Invoker, user1_token: str + self, client: TestClient, mock_invoker: Invoker, mock_sqlite_database: Database, user1_token: str ): """When no categories filter is given, default workflows should still appear.""" from invokeai.app.services.workflow_records.workflow_records_common import ( @@ -1781,7 +1794,7 @@ def test_default_workflows_visible_when_no_category_filter( form_fields=[], ) wf_with_id = Workflow(**default_wf.model_dump(), id=uuid_string()) - with mock_invoker.services.workflow_records._db.transaction() as cursor: + with sqlite_cursor(mock_sqlite_database) as cursor: cursor.execute( "INSERT INTO workflow_library (workflow_id, workflow, user_id) VALUES (?, ?, ?)", (wf_with_id.id, wf_with_id.model_dump_json(), "system"), @@ -2158,7 +2171,12 @@ def test_recall_shared_board_image_allowed( assert r.status_code != status.HTTP_403_FORBIDDEN def test_recall_explicitly_shared_board_image_allowed_for_member( - self, client: TestClient, mock_invoker: Invoker, user1_token: str, user2_token: str + self, + client: TestClient, + mock_invoker: Invoker, + mock_sqlite_database: Database, + user1_token: str, + user2_token: str, ): """User2 should be able to reference an image on a private board explicitly shared with them.""" user1 = mock_invoker.services.users.get_by_email("user1@test.com") @@ -2177,7 +2195,7 @@ def test_recall_explicitly_shared_board_image_allowed_for_member( ) assert r.status_code == status.HTTP_403_FORBIDDEN - _share_board_with_user(mock_invoker, board_id, user2.user_id) + _share_board_with_user(mock_sqlite_database, board_id, user2.user_id) r = client.post( "/api/v1/recall/default", @@ -2370,16 +2388,15 @@ def test_session_queue_status_has_user_fields(self): assert scoped.pending == 5 # global, unchanged assert scoped.user_pending == 2 # this user's share - def _setup_queue_router(self, mock_invoker: Invoker): + def _setup_queue_router(self, mock_invoker: Invoker, db: Database): """Wire a real session queue and a stub processor into the invoker the router uses, so GET /queue/{queue_id}/status exercises the real service contract.""" from unittest.mock import MagicMock from invokeai.app.services.session_processor.session_processor_common import SessionProcessorStatus - from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue + from invokeai.app.services.session_queue.session_queue_default import SessionQueue - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) + queue = SessionQueue(db) queue.start(mock_invoker) mock_invoker.services.session_queue = queue @@ -2389,13 +2406,18 @@ def _setup_queue_router(self, mock_invoker: Invoker): return queue def test_get_queue_status_route_returns_global_and_user_counts( - self, setup_jwt_secret: None, enable_multiuser: Any, mock_invoker: Invoker, client: TestClient + self, + setup_jwt_secret: None, + enable_multiuser: Any, + mock_invoker: Invoker, + mock_sqlite_database: Database, + client: TestClient, ): """Regression test: GET /api/v1/queue/{queue_id}/status must return 200 (not 500) and the expected global and per-user counts for both non-admin and admin callers. Previously the router called get_queue_status() with a keyword the service did not accept, raising a TypeError that the broad except turned into a 500 for every status request.""" - queue = self._setup_queue_router(mock_invoker) + queue = self._setup_queue_router(mock_invoker, mock_sqlite_database) user1_id = _create_user(mock_invoker, "user1@test.com", "User One") user2_id = _create_user(mock_invoker, "user2@test.com", "User Two") diff --git a/tests/app/routers/test_privilege_revocation.py b/tests/app/routers/test_privilege_revocation.py index d510ee592ce..83cca4490a2 100644 --- a/tests/app/routers/test_privilege_revocation.py +++ b/tests/app/routers/test_privilege_revocation.py @@ -23,10 +23,9 @@ from invokeai.app.services.events.events_common import UserAccessChangedEvent from invokeai.app.services.invocation_services import InvocationServices from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.users.users_common import UserCreateRequest -from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage class MockApiDependencies(ApiDependencies): @@ -49,41 +48,40 @@ def client(): @pytest.fixture -def mock_services() -> InvocationServices: - from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage - from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage - from invokeai.app.services.board_video_records.board_video_records_sqlite import SqliteBoardVideoRecordStorage +def mock_services(mock_sqlite_database: Database) -> InvocationServices: + from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage + from invokeai.app.services.board_records.board_records_default import BoardRecordStorage + from invokeai.app.services.board_video_records.board_video_records_default import BoardVideoRecordStorage from invokeai.app.services.boards.boards_default import BoardService from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService - from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ( - ClientStatePersistenceSqlite, + from invokeai.app.services.client_state_persistence.client_state_persistence_default import ( + ClientStatePersistence, ) - from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage + from invokeai.app.services.image_records.image_records_default import ImageRecordStorage from invokeai.app.services.images.images_default import ImageService from invokeai.app.services.invocation_cache.invocation_cache_memory import MemoryInvocationCache from invokeai.app.services.invocation_stats.invocation_stats_default import InvocationStatsService - from invokeai.app.services.system_prompt_records.system_prompt_records_sqlite import ( - SqliteSystemPromptRecordsStorage, + from invokeai.app.services.system_prompt_records.system_prompt_records_default import ( + SystemPromptRecordsStorage, ) from invokeai.app.services.users.users_default import UserService - from invokeai.app.services.video_records.video_records_sqlite import SqliteVideoRecordStorage + from invokeai.app.services.video_records.video_records_default import VideoRecordStorage from tests.test_nodes import TestEventService configuration = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0) - logger = InvokeAILogger.get_logger() - db = create_mock_sqlite_database(configuration, logger) + db = mock_sqlite_database return InvocationServices( - board_image_records=SqliteBoardImageRecordStorage(db=db), + board_image_records=BoardImageRecordStorage(db), board_images=None, # type: ignore - board_records=SqliteBoardRecordStorage(db=db), + board_records=BoardRecordStorage(db), boards=BoardService(), bulk_download=BulkDownloadService(), configuration=configuration, database=db, events=TestEventService(), image_files=None, # type: ignore - image_records=SqliteImageRecordStorage(db=db), + image_records=ImageRecordStorage(db), images=ImageService(), invocation_cache=MemoryInvocationCache(max_cache_size=0), logger=logging, # type: ignore @@ -95,22 +93,22 @@ def mock_services() -> InvocationServices: session_processor=None, # type: ignore session_queue=None, # type: ignore urls=None, # type: ignore - workflow_records=SqliteWorkflowRecordsStorage(db=db), + workflow_records=WorkflowRecordsStorage(db), tensors=None, # type: ignore conditioning=None, # type: ignore style_preset_records=None, # type: ignore style_preset_image_files=None, # type: ignore - system_prompt_records=SqliteSystemPromptRecordsStorage(db=db), + system_prompt_records=SystemPromptRecordsStorage(db), workflow_thumbnails=None, # type: ignore model_relationship_records=None, # type: ignore model_relationships=None, # type: ignore - client_state_persistence=ClientStatePersistenceSqlite(db=db), + client_state_persistence=ClientStatePersistence(db), users=UserService(db), external_generation=None, # type: ignore videos=None, # type: ignore video_files=None, # type: ignore - video_records=SqliteVideoRecordStorage(db=db), - board_video_records=SqliteBoardVideoRecordStorage(db=db), + video_records=VideoRecordStorage(db), + board_video_records=BoardVideoRecordStorage(db), gallery=None, # type: ignore wildcard_records=None, # type: ignore project_records=None, # type: ignore diff --git a/tests/app/routers/test_projects.py b/tests/app/routers/test_projects.py index 34e25fda5da..70e1615e085 100644 --- a/tests/app/routers/test_projects.py +++ b/tests/app/routers/test_projects.py @@ -12,8 +12,10 @@ from fastapi.testclient import TestClient from invokeai.app.services.invoker import Invoker -from invokeai.app.services.project_records import project_records_sqlite +from invokeai.app.services.project_records import project_records_default +from invokeai.app.services.shared.database.database import Database from tests.app.routers.conftest import _auth, _create_board +from tests.fixtures.sqlite_database import sqlite_cursor def _create_project(client: TestClient, token: str, **body: Any): @@ -36,7 +38,7 @@ def test_creating_a_project_creates_and_returns_its_board(client: TestClient, us def test_oversized_project_documents_return_a_structured_413( client: TestClient, user1_token: str, monkeypatch: pytest.MonkeyPatch ): - monkeypatch.setattr(project_records_sqlite, "PROJECT_DOCUMENT_MAX_BYTES", 16) + monkeypatch.setattr(project_records_default, "PROJECT_DOCUMENT_MAX_BYTES", 16) response = _create_project(client, user1_token, data={"value": "x" * 32}) @@ -53,7 +55,7 @@ def test_oversized_project_updates_return_a_structured_413( client: TestClient, user1_token: str, monkeypatch: pytest.MonkeyPatch ): created = _create_project(client, user1_token).json() - monkeypatch.setattr(project_records_sqlite, "PROJECT_DOCUMENT_MAX_BYTES", 16) + monkeypatch.setattr(project_records_default, "PROJECT_DOCUMENT_MAX_BYTES", 16) response = client.put( f"/api/v1/projects/{created['project_id']}", @@ -266,9 +268,11 @@ def test_deleting_a_project_deletes_its_board(client: TestClient, user1_token: s assert board.status_code == status.HTTP_404_NOT_FOUND -def test_deleting_a_project_leaves_its_media_uncategorized(client: TestClient, mock_invoker: Invoker, user1_token: str): +def test_deleting_a_project_leaves_its_media_uncategorized( + client: TestClient, mock_invoker: Invoker, mock_sqlite_database: Database, user1_token: str +): created = _create_project(client, user1_token).json() - with mock_invoker.services.board_records._db.transaction() as cursor: + with sqlite_cursor(mock_sqlite_database) as cursor: cursor.execute( "INSERT INTO images (image_name, image_origin, image_category, width, height)" " VALUES ('kept.png', 'internal', 'general', 64, 64);" @@ -279,7 +283,7 @@ def test_deleting_a_project_leaves_its_media_uncategorized(client: TestClient, m client.delete(f"/api/v1/projects/{created['project_id']}", headers=_auth(user1_token)) - with mock_invoker.services.board_records._db.transaction() as cursor: + with sqlite_cursor(mock_sqlite_database) as cursor: cursor.execute("SELECT COUNT(*) FROM images WHERE image_name = 'kept.png';") assert cursor.fetchone()[0] == 1 cursor.execute("SELECT COUNT(*) FROM board_images WHERE image_name = 'kept.png';") @@ -287,10 +291,10 @@ def test_deleting_a_project_leaves_its_media_uncategorized(client: TestClient, m def test_the_board_snapshot_lists_the_projects_visible_media( - client: TestClient, mock_invoker: Invoker, user1_token: str + client: TestClient, mock_invoker: Invoker, mock_sqlite_database: Database, user1_token: str ): created = _create_project(client, user1_token).json() - with mock_invoker.services.board_records._db.transaction() as cursor: + with sqlite_cursor(mock_sqlite_database) as cursor: for name, category, intermediate in ( ("shown.png", "general", False), ("asset.png", "control", False), diff --git a/tests/app/routers/test_session_queue_item_id_limits.py b/tests/app/routers/test_session_queue_item_id_limits.py index 48cc95221dc..c4083259610 100644 --- a/tests/app/routers/test_session_queue_item_id_limits.py +++ b/tests/app/routers/test_session_queue_item_id_limits.py @@ -1,31 +1,27 @@ """Guards the batch bound on the queue summary route. -`item_summaries_by_ids` takes a client-supplied list of ids and the SQLite layer binds one -parameter per id. Unbounded, a client could post tens of thousands of ids: past SQLite's -per-statement variable limit the query raises `OperationalError`, which the route reports as a -generic HTTP 500, and even below that limit it is an invitation to make the server do arbitrary -work per request. The route caps the list so oversized requests are rejected by validation -instead. +`item_summaries_by_ids` takes a client-supplied list of ids. Unbounded, a client could post tens +of thousands of ids and make the server do arbitrary work per request. The route caps the list so +oversized requests are rejected by validation instead. The id listing the client hydrates from takes an optional limit with the same bound. A limited listing returns the head of the full order and counts only the ids it returns: counting every match would read the queue's whole history, which is what the limit exists to avoid. """ -import json import uuid -from typing import cast from unittest.mock import MagicMock import pytest from fastapi.testclient import TestClient -from pydantic_core import to_jsonable_python +from sqlalchemy import insert from invokeai.app.api.dependencies import ApiDependencies from invokeai.app.api.routers.session_queue import MAX_QUEUE_ITEM_IDS_PER_REQUEST from invokeai.app.api_app import app from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.schema.session_queue import session_queue as session_queue_table from invokeai.app.services.shared.graph import Graph, GraphExecutionState SUMMARIES_ROUTE = "/api/v1/queue/default/item_summaries_by_ids" @@ -86,21 +82,22 @@ def test_full_items_by_ids_rejects_oversized_id_lists(mock_queue_invoker: MagicM @pytest.fixture def listed_queue(monkeypatch: pytest.MonkeyPatch, mock_invoker: Invoker) -> list[int]: """A real queue of five items in two projects, the ids in insertion (and creation) order.""" - session_queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) - mock_invoker.services.session_queue = session_queue + database = mock_invoker.services.database + mock_invoker.services.session_queue = SessionQueue(database) monkeypatch.setattr(ApiDependencies, "invoker", mock_invoker, raising=False) - session = json.dumps(to_jsonable_python(GraphExecutionState(graph=Graph()).model_dump())) + session = GraphExecutionState(graph=Graph()).model_dump_json() item_ids = [] for index, origin in enumerate(["webv2:p:a:q:1", "webv2:p:b:q:2", "webv2:p:a:q:3", "webv2:p:a:q:4", "x"]): - with session_queue._db.transaction() as cursor: - cursor.execute( - """--sql - INSERT INTO session_queue (queue_id, session, session_id, batch_id, created_at, origin) - VALUES ('default', ?, ?, 'batch', ?, ?); - """, - (session, str(uuid.uuid4()), f"2026-10-01 10:00:0{index}.000", origin), - ) - item_ids.append(cast(int, cursor.lastrowid)) + values = { + "queue_id": "default", + "session": session, + "session_id": str(uuid.uuid4()), + "batch_id": "batch", + "created_at": f"2026-10-01 10:00:0{index}.000", + "origin": origin, + } + with database.begin(write=True) as conn: + item_ids.append(int(conn.execute(insert(session_queue_table).values(values)).inserted_primary_key[0])) return item_ids diff --git a/tests/app/routers/test_session_queue_workflow_call.py b/tests/app/routers/test_session_queue_workflow_call.py index b8432fb4a73..3137277c056 100644 --- a/tests/app/routers/test_session_queue_workflow_call.py +++ b/tests/app/routers/test_session_queue_workflow_call.py @@ -14,12 +14,12 @@ from invokeai.app.services.invocation_services import InvocationServices from invokeai.app.services.invoker import Invoker from invokeai.app.services.session_processor.session_processor_common import SessionProcessorStatus -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph, GraphExecutionState from invokeai.app.services.users.users_common import UserCreateRequest -from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage +from tests.fixtures.sqlite_database import sqlite_cursor_of class MockApiDependencies(ApiDependencies): @@ -42,40 +42,39 @@ def client(): @pytest.fixture -def mock_services() -> InvocationServices: - from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage - from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage - from invokeai.app.services.board_video_records.board_video_records_sqlite import SqliteBoardVideoRecordStorage +def mock_services(mock_sqlite_database: Database) -> InvocationServices: + from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage + from invokeai.app.services.board_records.board_records_default import BoardRecordStorage + from invokeai.app.services.board_video_records.board_video_records_default import BoardVideoRecordStorage from invokeai.app.services.boards.boards_default import BoardService from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService - from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ( - ClientStatePersistenceSqlite, + from invokeai.app.services.client_state_persistence.client_state_persistence_default import ( + ClientStatePersistence, ) - from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage + from invokeai.app.services.image_records.image_records_default import ImageRecordStorage from invokeai.app.services.images.images_default import ImageService from invokeai.app.services.invocation_cache.invocation_cache_memory import MemoryInvocationCache from invokeai.app.services.invocation_stats.invocation_stats_default import InvocationStatsService - from invokeai.app.services.project_records.project_records_sqlite import ProjectRecordsSqlite + from invokeai.app.services.project_records.project_records_default import ProjectRecordsStorage from invokeai.app.services.users.users_default import UserService - from invokeai.app.services.video_records.video_records_sqlite import SqliteVideoRecordStorage - from invokeai.app.services.wildcard_records.wildcard_records_sqlite import SqliteWildcardRecordsStorage + from invokeai.app.services.video_records.video_records_default import VideoRecordStorage + from invokeai.app.services.wildcard_records.wildcard_records_default import WildcardRecordsStorage from tests.test_nodes import TestEventService configuration = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0) - logger = InvokeAILogger.get_logger() - db = create_mock_sqlite_database(configuration, logger) + db = mock_sqlite_database return InvocationServices( - board_image_records=SqliteBoardImageRecordStorage(db=db), + board_image_records=BoardImageRecordStorage(db), board_images=None, # type: ignore - board_records=SqliteBoardRecordStorage(db=db), + board_records=BoardRecordStorage(db), boards=BoardService(), bulk_download=BulkDownloadService(), configuration=configuration, database=db, events=TestEventService(), image_files=None, # type: ignore - image_records=SqliteImageRecordStorage(db=db), + image_records=ImageRecordStorage(db), images=ImageService(), invocation_cache=MemoryInvocationCache(max_cache_size=0), logger=logging, # type: ignore @@ -87,7 +86,7 @@ def mock_services() -> InvocationServices: session_processor=None, # type: ignore session_queue=None, # type: ignore urls=None, # type: ignore - workflow_records=SqliteWorkflowRecordsStorage(db=db), + workflow_records=WorkflowRecordsStorage(db), tensors=None, # type: ignore conditioning=None, # type: ignore style_preset_records=None, # type: ignore @@ -95,16 +94,16 @@ def mock_services() -> InvocationServices: workflow_thumbnails=None, # type: ignore model_relationship_records=None, # type: ignore model_relationships=None, # type: ignore - client_state_persistence=ClientStatePersistenceSqlite(db=db), - project_records=ProjectRecordsSqlite(db=db), + client_state_persistence=ClientStatePersistence(db), + project_records=ProjectRecordsStorage(db), users=UserService(db), - wildcard_records=SqliteWildcardRecordsStorage(db=db), + wildcard_records=WildcardRecordsStorage(db), external_generation=None, # type: ignore system_prompt_records=None, # type: ignore videos=None, # type: ignore video_files=None, # type: ignore - video_records=SqliteVideoRecordStorage(db=db), - board_video_records=SqliteBoardVideoRecordStorage(db=db), + video_records=VideoRecordStorage(db), + board_video_records=BoardVideoRecordStorage(db), gallery=None, # type: ignore image_index_records=None, # type: ignore image_index=None, # type: ignore @@ -113,9 +112,9 @@ def mock_services() -> InvocationServices: @pytest.fixture -def mock_invoker(mock_services: InvocationServices) -> Invoker: +def mock_invoker(mock_services: InvocationServices, mock_sqlite_database: Database) -> Invoker: invoker = Invoker(services=mock_services) - queue = SqliteSessionQueue(db=mock_services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_services.session_queue = queue mock_services.session_processor = MagicMock() mock_services.session_processor.get_status.return_value = SessionProcessorStatus( @@ -171,7 +170,7 @@ def user2_token(enable_multiuser: Any, mock_invoker: Invoker, client: TestClient def _insert_queue_item( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, *, queue_id: str = "default", user_id: str, @@ -184,7 +183,7 @@ def _insert_queue_item( workflow_call_depth: int | None = None, ) -> int: session = session or GraphExecutionState(graph=Graph()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( diff --git a/tests/app/routers/test_style_presets.py b/tests/app/routers/test_style_presets.py index 9289b80f137..1f25976e601 100644 --- a/tests/app/routers/test_style_presets.py +++ b/tests/app/routers/test_style_presets.py @@ -1,6 +1,6 @@ """Router-level tests for /api/v1/style_presets. -Backed by a real SqliteStylePresetRecordsStorage from the shared conftest, so SQL +Backed by a real StylePresetRecordsStorage from the shared conftest, so SQL filtering and ownership rules are exercised end-to-end. style_preset_image_files remains a MagicMock — file IO is not under test here. diff --git a/tests/app/routers/test_system_prompts_multiuser.py b/tests/app/routers/test_system_prompts_multiuser.py index 6c2f5959678..7360473f75d 100644 --- a/tests/app/routers/test_system_prompts_multiuser.py +++ b/tests/app/routers/test_system_prompts_multiuser.py @@ -19,9 +19,8 @@ from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.invocation_services import InvocationServices from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.users.users_common import UserCreateRequest -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database class MockApiDependencies(ApiDependencies): @@ -44,39 +43,38 @@ def client(): @pytest.fixture -def mock_services() -> InvocationServices: - from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage - from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage +def mock_services(mock_sqlite_database: Database) -> InvocationServices: + from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage + from invokeai.app.services.board_records.board_records_default import BoardRecordStorage from invokeai.app.services.boards.boards_default import BoardService from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService - from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ( - ClientStatePersistenceSqlite, + from invokeai.app.services.client_state_persistence.client_state_persistence_default import ( + ClientStatePersistence, ) - from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage + from invokeai.app.services.image_records.image_records_default import ImageRecordStorage from invokeai.app.services.images.images_default import ImageService from invokeai.app.services.invocation_cache.invocation_cache_memory import MemoryInvocationCache from invokeai.app.services.invocation_stats.invocation_stats_default import InvocationStatsService - from invokeai.app.services.system_prompt_records.system_prompt_records_sqlite import ( - SqliteSystemPromptRecordsStorage, + from invokeai.app.services.system_prompt_records.system_prompt_records_default import ( + SystemPromptRecordsStorage, ) from invokeai.app.services.users.users_default import UserService from tests.test_nodes import TestEventService configuration = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0) - logger = InvokeAILogger.get_logger() - db = create_mock_sqlite_database(configuration, logger) + db = mock_sqlite_database return InvocationServices( - board_image_records=SqliteBoardImageRecordStorage(db=db), + board_image_records=BoardImageRecordStorage(db), board_images=None, # type: ignore - board_records=SqliteBoardRecordStorage(db=db), + board_records=BoardRecordStorage(db), boards=BoardService(), bulk_download=BulkDownloadService(), configuration=configuration, database=db, events=TestEventService(), image_files=None, # type: ignore - image_records=SqliteImageRecordStorage(db=db), + image_records=ImageRecordStorage(db), images=ImageService(), invocation_cache=MemoryInvocationCache(max_cache_size=0), logger=logging, # type: ignore @@ -93,11 +91,11 @@ def mock_services() -> InvocationServices: conditioning=None, # type: ignore style_preset_records=None, # type: ignore style_preset_image_files=None, # type: ignore - system_prompt_records=SqliteSystemPromptRecordsStorage(db=db), + system_prompt_records=SystemPromptRecordsStorage(db), workflow_thumbnails=None, # type: ignore model_relationship_records=None, # type: ignore model_relationships=None, # type: ignore - client_state_persistence=ClientStatePersistenceSqlite(db=db), + client_state_persistence=ClientStatePersistence(db), project_records=None, # type: ignore users=UserService(db), wildcard_records=None, # type: ignore diff --git a/tests/app/routers/test_system_prompts_single_user.py b/tests/app/routers/test_system_prompts_single_user.py index e6c5a8a90b5..89c923e7e0c 100644 --- a/tests/app/routers/test_system_prompts_single_user.py +++ b/tests/app/routers/test_system_prompts_single_user.py @@ -16,8 +16,7 @@ from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.invocation_services import InvocationServices from invokeai.app.services.invoker import Invoker -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database +from invokeai.app.services.shared.database.database import Database class MockApiDependencies(ApiDependencies): @@ -33,39 +32,38 @@ def client(): @pytest.fixture -def mock_services() -> InvocationServices: - from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage - from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage +def mock_services(mock_sqlite_database: Database) -> InvocationServices: + from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage + from invokeai.app.services.board_records.board_records_default import BoardRecordStorage from invokeai.app.services.boards.boards_default import BoardService from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService - from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ( - ClientStatePersistenceSqlite, + from invokeai.app.services.client_state_persistence.client_state_persistence_default import ( + ClientStatePersistence, ) - from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage + from invokeai.app.services.image_records.image_records_default import ImageRecordStorage from invokeai.app.services.images.images_default import ImageService from invokeai.app.services.invocation_cache.invocation_cache_memory import MemoryInvocationCache from invokeai.app.services.invocation_stats.invocation_stats_default import InvocationStatsService - from invokeai.app.services.system_prompt_records.system_prompt_records_sqlite import ( - SqliteSystemPromptRecordsStorage, + from invokeai.app.services.system_prompt_records.system_prompt_records_default import ( + SystemPromptRecordsStorage, ) from invokeai.app.services.users.users_default import UserService from tests.test_nodes import TestEventService configuration = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0) - logger = InvokeAILogger.get_logger() - db = create_mock_sqlite_database(configuration, logger) + db = mock_sqlite_database return InvocationServices( - board_image_records=SqliteBoardImageRecordStorage(db=db), + board_image_records=BoardImageRecordStorage(db), board_images=None, # type: ignore - board_records=SqliteBoardRecordStorage(db=db), + board_records=BoardRecordStorage(db), boards=BoardService(), bulk_download=BulkDownloadService(), configuration=configuration, database=db, events=TestEventService(), image_files=None, # type: ignore - image_records=SqliteImageRecordStorage(db=db), + image_records=ImageRecordStorage(db), images=ImageService(), invocation_cache=MemoryInvocationCache(max_cache_size=0), logger=logging, # type: ignore @@ -82,11 +80,11 @@ def mock_services() -> InvocationServices: conditioning=None, # type: ignore style_preset_records=None, # type: ignore style_preset_image_files=None, # type: ignore - system_prompt_records=SqliteSystemPromptRecordsStorage(db=db), + system_prompt_records=SystemPromptRecordsStorage(db), workflow_thumbnails=None, # type: ignore model_relationship_records=None, # type: ignore model_relationships=None, # type: ignore - client_state_persistence=ClientStatePersistenceSqlite(db=db), + client_state_persistence=ClientStatePersistence(db), project_records=None, # type: ignore users=UserService(db), wildcard_records=None, # type: ignore diff --git a/tests/app/routers/test_videos_multiuser.py b/tests/app/routers/test_videos_multiuser.py index 77a1029bd68..a72d10cc51b 100644 --- a/tests/app/routers/test_videos_multiuser.py +++ b/tests/app/routers/test_videos_multiuser.py @@ -104,7 +104,7 @@ def enable_multiuser_for_videos(monkeypatch: Any, mock_invoker: Invoker): # The board service computes video_count + cover_video_name on every get_dto/update; # an unconfigured MagicMock returns nested MagicMocks that fail Pydantic validation and # the boards route swallows the exception as a 404. Pin sane defaults. - mock_invoker.services.board_video_records.get_video_count_for_board.return_value = 0 + mock_invoker.services.board_video_records.get_counts_for_board.return_value = (0, 0) mock_invoker.services.video_records = MagicMock() mock_invoker.services.video_records.get_most_recent_video_for_board.return_value = None mock_invoker.services.board_images = MagicMock() diff --git a/tests/app/routers/test_workflows_multiuser.py b/tests/app/routers/test_workflows_multiuser.py index 5fe87492f86..fc59eaf193a 100644 --- a/tests/app/routers/test_workflows_multiuser.py +++ b/tests/app/routers/test_workflows_multiuser.py @@ -13,11 +13,10 @@ from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.invocation_services import InvocationServices from invokeai.app.services.invoker import Invoker -from invokeai.app.services.project_records.project_records_sqlite import ProjectRecordsSqlite +from invokeai.app.services.project_records.project_records_default import ProjectRecordsStorage +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.users.users_common import UserCreateRequest -from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage class MockApiDependencies(ApiDependencies): @@ -57,42 +56,41 @@ def client(): @pytest.fixture -def mock_services() -> InvocationServices: - from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage - from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage - from invokeai.app.services.board_video_records.board_video_records_sqlite import SqliteBoardVideoRecordStorage +def mock_services(mock_sqlite_database: Database) -> InvocationServices: + from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage + from invokeai.app.services.board_records.board_records_default import BoardRecordStorage + from invokeai.app.services.board_video_records.board_video_records_default import BoardVideoRecordStorage from invokeai.app.services.boards.boards_default import BoardService from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService - from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ( - ClientStatePersistenceSqlite, + from invokeai.app.services.client_state_persistence.client_state_persistence_default import ( + ClientStatePersistence, ) - from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage + from invokeai.app.services.image_records.image_records_default import ImageRecordStorage from invokeai.app.services.images.images_default import ImageService from invokeai.app.services.invocation_cache.invocation_cache_memory import MemoryInvocationCache from invokeai.app.services.invocation_stats.invocation_stats_default import InvocationStatsService - from invokeai.app.services.system_prompt_records.system_prompt_records_sqlite import ( - SqliteSystemPromptRecordsStorage, + from invokeai.app.services.system_prompt_records.system_prompt_records_default import ( + SystemPromptRecordsStorage, ) from invokeai.app.services.users.users_default import UserService - from invokeai.app.services.video_records.video_records_sqlite import SqliteVideoRecordStorage - from invokeai.app.services.wildcard_records.wildcard_records_sqlite import SqliteWildcardRecordsStorage + from invokeai.app.services.video_records.video_records_default import VideoRecordStorage + from invokeai.app.services.wildcard_records.wildcard_records_default import WildcardRecordsStorage from tests.test_nodes import TestEventService configuration = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0) - logger = InvokeAILogger.get_logger() - db = create_mock_sqlite_database(configuration, logger) + db = mock_sqlite_database return InvocationServices( - board_image_records=SqliteBoardImageRecordStorage(db=db), + board_image_records=BoardImageRecordStorage(db), board_images=None, # type: ignore - board_records=SqliteBoardRecordStorage(db=db), + board_records=BoardRecordStorage(db), boards=BoardService(), bulk_download=BulkDownloadService(), configuration=configuration, database=db, events=TestEventService(), image_files=None, # type: ignore - image_records=SqliteImageRecordStorage(db=db), + image_records=ImageRecordStorage(db), images=ImageService(), invocation_cache=MemoryInvocationCache(max_cache_size=0), logger=logging, # type: ignore @@ -104,24 +102,24 @@ def mock_services() -> InvocationServices: session_processor=None, # type: ignore session_queue=None, # type: ignore urls=None, # type: ignore - workflow_records=SqliteWorkflowRecordsStorage(db=db), + workflow_records=WorkflowRecordsStorage(db), tensors=None, # type: ignore conditioning=None, # type: ignore style_preset_records=None, # type: ignore style_preset_image_files=None, # type: ignore - system_prompt_records=SqliteSystemPromptRecordsStorage(db=db), + system_prompt_records=SystemPromptRecordsStorage(db), workflow_thumbnails=None, # type: ignore model_relationship_records=None, # type: ignore model_relationships=None, # type: ignore - client_state_persistence=ClientStatePersistenceSqlite(db=db), - project_records=ProjectRecordsSqlite(db=db), + client_state_persistence=ClientStatePersistence(db), + project_records=ProjectRecordsStorage(db), users=UserService(db), - wildcard_records=SqliteWildcardRecordsStorage(db=db), + wildcard_records=WildcardRecordsStorage(db), external_generation=None, # type: ignore videos=None, # type: ignore video_files=None, # type: ignore - video_records=SqliteVideoRecordStorage(db=db), - board_video_records=SqliteBoardVideoRecordStorage(db=db), + video_records=VideoRecordStorage(db), + board_video_records=BoardVideoRecordStorage(db), gallery=None, # type: ignore image_index_records=None, # type: ignore image_index=None, # type: ignore diff --git a/tests/app/routers/test_workflows_revisions.py b/tests/app/routers/test_workflows_revisions.py index 0f1b892cb1d..3ed329b1845 100644 --- a/tests/app/routers/test_workflows_revisions.py +++ b/tests/app/routers/test_workflows_revisions.py @@ -7,7 +7,7 @@ from fastapi.testclient import TestClient from invokeai.app.services.invoker import Invoker -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.services.workflow_records.workflow_records_common import WorkflowCategory, WorkflowRecordOrderBy from tests.app.routers.test_workflows_multiuser import WORKFLOW_BODY, create_workflow diff --git a/tests/app/services/app_settings/test_app_settings_service.py b/tests/app/services/app_settings/test_app_settings_service.py new file mode 100644 index 00000000000..d92ef0aa218 --- /dev/null +++ b/tests/app/services/app_settings/test_app_settings_service.py @@ -0,0 +1,18 @@ +import pytest +from sqlalchemy import delete + +from invokeai.app.services.app_settings import AppSettingsService +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.app_settings import app_settings + + +def test_the_jwt_secret_a_new_database_starts_with_is_read(database: Database) -> None: + assert len(AppSettingsService(database).get_jwt_secret()) == 64 + + +def test_a_database_without_a_jwt_secret_is_reported(database: Database) -> None: + with database.begin(write=True) as conn: + conn.execute(delete(app_settings).where(app_settings.c.key == "jwt_secret")) + + with pytest.raises(RuntimeError, match="JWT secret not found"): + AppSettingsService(database).get_jwt_secret() diff --git a/tests/app/services/auth/test_data_isolation.py b/tests/app/services/auth/test_data_isolation.py index 45a538a6beb..8b273400534 100644 --- a/tests/app/services/auth/test_data_isolation.py +++ b/tests/app/services/auth/test_data_isolation.py @@ -15,7 +15,7 @@ from invokeai.app.api_app import app from invokeai.app.services.board_records.board_records_common import BoardRecordOrderBy from invokeai.app.services.invoker import Invoker -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.services.users.users_common import UserCreateRequest diff --git a/tests/app/services/auth/test_performance.py b/tests/app/services/auth/test_performance.py index 77c1d1cbfbc..63c4bc6d5e5 100644 --- a/tests/app/services/auth/test_performance.py +++ b/tests/app/services/auth/test_performance.py @@ -15,46 +15,20 @@ import time from concurrent.futures import ThreadPoolExecutor, as_completed -from logging import Logger import pytest from invokeai.app.services.auth.password_utils import hash_password, verify_password from invokeai.app.services.auth.token_service import TokenData, create_access_token, verify_token -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.users.users_common import UserCreateRequest from invokeai.app.services.users.users_default import UserService @pytest.fixture -def logger() -> Logger: - """Create a logger for testing.""" - return Logger("test_performance") - - -@pytest.fixture -def user_service(logger: Logger) -> UserService: - """Create a user service with in-memory database for testing.""" - db = SqliteDatabase(db_path=None, logger=logger, verbose=False) - - # Create users table - db._conn.execute(""" - CREATE TABLE users ( - user_id TEXT NOT NULL PRIMARY KEY, - email TEXT NOT NULL UNIQUE, - display_name TEXT, - password_hash TEXT NOT NULL, - is_admin BOOLEAN NOT NULL DEFAULT FALSE, - is_active BOOLEAN NOT NULL DEFAULT TRUE, - 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')), - last_login_at DATETIME, - token_epoch INTEGER NOT NULL DEFAULT 0 - ); - """) - db._conn.commit() - - return UserService(db) +def user_service(database: Database) -> UserService: + """Create a user service on a database at the newest schema.""" + return UserService(database) @pytest.mark.slow diff --git a/tests/app/services/boards/test_board_storage.py b/tests/app/services/boards/test_board_storage.py new file mode 100644 index 00000000000..71f7f607c1d --- /dev/null +++ b/tests/app/services/boards/test_board_storage.py @@ -0,0 +1,431 @@ +"""Board records and board memberships on every database backend.""" + +import pytest +from sqlalchemy import insert, select + +from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage +from invokeai.app.services.board_records.board_records_common import ( + BoardChanges, + BoardRecordNotFoundException, + BoardRecordOrderBy, + BoardRecordProjectOwnedException, + BoardVisibility, +) +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage +from invokeai.app.services.board_video_records.board_video_records_default import BoardVideoRecordStorage +from invokeai.app.services.image_records.image_records_common import ImageCategory +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import board_images as board_image_queries +from invokeai.app.services.shared.database.queries import board_videos as board_video_queries +from invokeai.app.services.shared.database.queries import boards as board_queries +from invokeai.app.services.shared.database.schema.boards import board_images as board_images_table +from invokeai.app.services.shared.database.schema.boards import board_videos as board_videos_table +from invokeai.app.services.shared.database.schema.boards import boards as boards_table +from invokeai.app.services.shared.database.schema.boards import shared_boards +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.projects import projects +from invokeai.app.services.shared.database.schema.users import users +from invokeai.app.services.shared.database.schema.videos import videos +from invokeai.app.services.shared.pagination import SQLiteDirection + +ALICE = "alice" +BOB = "bob" +CAROL = "carol" + + +@pytest.fixture +def board_records(database: Database) -> BoardRecordStorage: + with database.begin(write=True) as conn: + conn.execute( + insert(users), + [{"user_id": user, "email": f"{user}@test.com", "password_hash": "-"} for user in (ALICE, BOB, CAROL)], + ) + return BoardRecordStorage(database) + + +@pytest.fixture +def board_images(database: Database) -> BoardImageRecordStorage: + return BoardImageRecordStorage(database) + + +@pytest.fixture +def board_videos(database: Database) -> BoardVideoRecordStorage: + return BoardVideoRecordStorage(database) + + +def _claim(database: Database, board_id: str, project_id: str = "project") -> None: + with database.begin(write=True) as conn: + conn.execute( + insert(projects).values(project_id=project_id, user_id=ALICE, name="P", data="{}", board_id=board_id) + ) + + +def _image( + database: Database, + name: str, + *, + category: ImageCategory = ImageCategory.GENERAL, + intermediate: bool = False, + user_id: str = ALICE, +) -> str: + with database.begin(write=True) as conn: + conn.execute( + insert(images).values( + image_name=name, + image_origin="internal", + image_category=category.value, + width=1, + height=1, + is_intermediate=intermediate, + user_id=user_id, + ) + ) + return name + + +def _video( + database: Database, + name: str, + *, + category: ImageCategory = ImageCategory.GENERAL, + intermediate: bool = False, + user_id: str = ALICE, +) -> str: + with database.begin(write=True) as conn: + conn.execute( + insert(videos).values( + video_name=name, + video_origin="internal", + video_category=category.value, + width=1, + height=1, + is_intermediate=intermediate, + user_id=user_id, + ) + ) + return name + + +# region board records + + +def test_a_saved_board_is_read_back_with_the_project_that_claims_it( + database: Database, board_records: BoardRecordStorage +) -> None: + free = board_records.save("Free", ALICE) + claimed = board_records.save("Claimed", ALICE) + _claim(database, claimed.board_id) + + assert board_records.get(free.board_id) == free + assert board_records.get_with_project_id(free.board_id) == (free, None) + assert board_records.get_with_project_id(claimed.board_id) == (claimed, "project") + assert board_records.get_project_ids_for_boards([free.board_id, claimed.board_id]) == {claimed.board_id: "project"} + with pytest.raises(BoardRecordNotFoundException): + board_records.get("missing") + with pytest.raises(BoardRecordNotFoundException): + board_records.get_with_project_id("missing") + + +def test_project_ids_are_looked_up_in_chunks( + database: Database, board_records: BoardRecordStorage, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(board_queries, "IN_CHUNK", 2) + board_ids = [board_records.save(f"Board {index}", ALICE).board_id for index in range(5)] + for index, board_id in enumerate(board_ids): + _claim(database, board_id, f"project {index}") + + assert board_records.get_project_ids_for_boards(board_ids) == { + board_id: f"project {index}" for index, board_id in enumerate(board_ids) + } + + +def test_users_list_their_own_shared_and_public_boards(database: Database, board_records: BoardRecordStorage) -> None: + board_records.save("Own", ALICE) + shared_with_alice = board_records.save("Shared with Alice", BOB).board_id + shared_with_carol = board_records.save("Shared with Carol", BOB).board_id + shared = board_records.save("Shared", BOB).board_id + public = board_records.save("Public", BOB).board_id + board_records.save("Private", BOB) + archived = board_records.save("Archived", ALICE).board_id + board_records.update(shared, BoardChanges(board_visibility=BoardVisibility.Shared)) + board_records.update(public, BoardChanges(board_visibility=BoardVisibility.Public)) + board_records.update(archived, BoardChanges(archived=True)) + with database.begin(write=True) as conn: + conn.execute( + insert(shared_boards), + [{"board_id": shared_with_alice, "user_id": ALICE}, {"board_id": shared_with_carol, "user_id": CAROL}], + ) + + def listed(*, is_admin: bool, include_archived: bool = False) -> set[str]: + boards = board_records.get_all( + ALICE, is_admin, BoardRecordOrderBy.CreatedAt, SQLiteDirection.Descending, include_archived + ) + page = board_records.get_many( + ALICE, is_admin, BoardRecordOrderBy.CreatedAt, SQLiteDirection.Descending, 0, 100, include_archived + ) + assert [board.board_id for board in page.items] == [board.board_id for board in boards] + assert page.total == len(boards) + return {board.board_name for board in boards} + + assert listed(is_admin=False) == {"Own", "Shared with Alice", "Shared", "Public"} + assert listed(is_admin=False, include_archived=True) == {"Own", "Shared with Alice", "Shared", "Public", "Archived"} + assert listed(is_admin=True) == {"Own", "Shared with Alice", "Shared with Carol", "Shared", "Public", "Private"} + assert board_records.is_board_shared_with_user(shared_with_carol, CAROL) is True + assert board_records.is_board_shared_with_user(shared_with_carol, ALICE) is False + assert board_records.get_shared_user_ids(shared_with_carol) == [CAROL] + + +def _insert_boards(database: Database, boards: list[tuple[str, str, str]]) -> None: + """Boards of the first account, as (board_id, board_name, created_at).""" + with database.begin(write=True) as conn: + conn.execute( + insert(boards_table), + [ + {"board_id": board_id, "board_name": name, "user_id": ALICE, "created_at": created_at} + for board_id, name, created_at in boards + ], + ) + + +def test_boards_list_by_name_ignoring_case(database: Database, board_records: BoardRecordStorage) -> None: + for name in ("cherry", "Banana", "apple"): + board_records.save(name, ALICE) + + by_name = board_records.get_all(ALICE, False, BoardRecordOrderBy.Name, SQLiteDirection.Ascending) + + assert [board.board_name for board in by_name] == ["apple", "Banana", "cherry"] + + +def test_boards_list_by_age(database: Database, board_records: BoardRecordStorage) -> None: + # Ages that agree with neither the order of insertion nor that of the ids. + _insert_boards( + database, + [ + ("board-b", "second", "2026-01-02 00:00:00.000"), + ("board-c", "third", "2026-01-01 00:00:00.000"), + ("board-a", "first", "2026-01-03 00:00:00.000"), + ], + ) + + by_age = board_records.get_all(ALICE, False, BoardRecordOrderBy.CreatedAt, SQLiteDirection.Descending) + + assert [board.board_id for board in by_age] == ["board-a", "board-b", "board-c"] + + +def test_boards_of_one_age_list_in_the_order_of_their_ids_on_every_page( + database: Database, board_records: BoardRecordStorage +) -> None: + # Inserted out of id order, with names in yet another order. + _insert_boards( + database, + [ + ("board-b", "Banana", "2026-01-01 00:00:00.000"), + ("board-c", "apple", "2026-01-01 00:00:00.000"), + ("board-a", "cherry", "2026-01-01 00:00:00.000"), + ], + ) + + by_age = board_records.get_all(ALICE, False, BoardRecordOrderBy.CreatedAt, SQLiteDirection.Descending) + pages = [ + board.board_id + for offset in range(3) + for board in board_records.get_many( + ALICE, False, BoardRecordOrderBy.CreatedAt, SQLiteDirection.Descending, offset, 1 + ).items + ] + + assert [board.board_id for board in by_age] == ["board-c", "board-b", "board-a"] + assert pages == ["board-c", "board-b", "board-a"] + + +def test_a_change_keeps_what_it_does_not_name(database: Database, board_records: BoardRecordStorage) -> None: + board = board_records.save("Board", ALICE).board_id + _image(database, "cover.png") + board_records.update(board, BoardChanges(cover_image_name="cover.png")) + + changed = board_records.update(board, BoardChanges(board_name="Renamed", archived=True)) + + assert (changed.board_name, changed.archived, changed.cover_image_name) == ("Renamed", True, "cover.png") + + +def test_a_project_board_keeps_what_its_project_owns(database: Database, board_records: BoardRecordStorage) -> None: + board = board_records.save("Project", ALICE) + _claim(database, board.board_id) + _image(database, "cover.png") + + for changes in ( + BoardChanges(board_name="Renamed"), + BoardChanges(archived=True), + BoardChanges(board_visibility=BoardVisibility.Public), + ): + with pytest.raises(BoardRecordProjectOwnedException): + board_records.update(board.board_id, changes) + updated = board_records.update(board.board_id, BoardChanges(cover_image_name="cover.png")) + + assert (updated.board_name, updated.archived, updated.board_visibility) == ( + "Project", + False, + BoardVisibility.Private, + ) + assert updated.cover_image_name == "cover.png" + with pytest.raises(BoardRecordNotFoundException): + board_records.update("missing", BoardChanges(board_name="Renamed")) + + +def test_only_a_board_no_project_claims_is_deleted(database: Database, board_records: BoardRecordStorage) -> None: + free = board_records.save("Free", ALICE).board_id + claimed = board_records.save("Claimed", ALICE).board_id + _claim(database, claimed) + + assert board_records.delete_if_unclaimed(free) is True + assert board_records.delete_if_unclaimed(claimed) is False + with pytest.raises(BoardRecordNotFoundException): + board_records.get(free) + assert board_records.get(claimed).board_id == claimed + + +# endregion + +# region board memberships + + +def test_an_image_is_on_one_board_at_a_time( + database: Database, board_records: BoardRecordStorage, board_images: BoardImageRecordStorage +) -> None: + first = board_records.save("First", ALICE).board_id + second = board_records.save("Second", ALICE).board_id + image = _image(database, "image.png") + + board_images.add_image_to_board(first, image) + board_images.add_image_to_board(second, image) + + assert board_images.get_board_for_image(image) == second + assert board_images.remove_image_from_board(image, first) == 0 + assert board_images.remove_image_from_board(image, second) == 1 + assert board_images.get_board_for_image(image) is None + + +def test_moving_to_another_board_renews_when_the_membership_was_updated( + database: Database, + board_records: BoardRecordStorage, + board_images: BoardImageRecordStorage, + board_videos: BoardVideoRecordStorage, + monkeypatch: pytest.MonkeyPatch, +) -> None: + first = board_records.save("First", ALICE).board_id + second = board_records.save("Second", ALICE).board_id + image = _image(database, "moved.png") + video = _video(database, "moved.mp4") + for module in (board_image_queries, board_video_queries): + monkeypatch.setattr(module, "now_text", lambda: "2030-01-01 00:00:00.000") + board_images.add_image_to_board(first, image) + board_videos.add_video_to_board(first, video) + + for module in (board_image_queries, board_video_queries): + monkeypatch.setattr(module, "now_text", lambda: "2030-01-02 00:00:00.000") + board_images.add_image_to_board(second, image) + board_videos.add_video_to_board(second, video) + + with database.begin(write=False) as conn: + image_row = conn.execute(select(board_images_table.c.board_id, board_images_table.c.updated_at)).one() + video_row = conn.execute(select(board_videos_table.c.board_id, board_videos_table.c.updated_at)).one() + assert tuple(image_row) == (second, "2030-01-02 00:00:00.000") + assert tuple(video_row) == (second, "2030-01-02 00:00:00.000") + + +def test_image_names_and_counts_follow_the_filters( + database: Database, board_records: BoardRecordStorage, board_images: BoardImageRecordStorage +) -> None: + board = board_records.save("Board", ALICE).board_id + on_board = [ + _image(database, "general.png"), + _image(database, "control.png", category=ImageCategory.CONTROL), + _image(database, "mask.png", category=ImageCategory.MASK), + _image(database, "intermediate.png", intermediate=True), + _image(database, "bobs.png", user_id=BOB), + ] + for name in on_board: + board_images.add_image_to_board(board, name) + _image(database, "loose.png") + + def names(board_id: str, **filters: object) -> set[str]: + filters = {"categories": None, "is_intermediate": None, **filters} + return set(board_images.get_all_board_image_names_for_board(board_id, **filters)) # type: ignore[arg-type] + + assert names(board) == set(on_board) + assert names("none") == {"loose.png"} + assert names(board, categories=[ImageCategory.CONTROL, ImageCategory.MASK]) == {"control.png", "mask.png"} + assert names(board, is_intermediate=True) == {"intermediate.png"} + assert names(board, user_id=BOB) == {"bobs.png"} + + +def test_image_counts_are_per_board_and_per_kind( + database: Database, board_records: BoardRecordStorage, board_images: BoardImageRecordStorage +) -> None: + board = board_records.save("Board", ALICE).board_id + other = board_records.save("Other", ALICE).board_id + for name, category, intermediate in ( + ("general.png", ImageCategory.GENERAL, False), + ("control.png", ImageCategory.CONTROL, False), + ("mask.png", ImageCategory.MASK, False), + ("user.png", ImageCategory.USER, False), + ("intermediate.png", ImageCategory.CONTROL, True), + ): + board_images.add_image_to_board(board, _image(database, name, category=category, intermediate=intermediate)) + board_images.add_image_to_board(other, _image(database, "elsewhere.png")) + + assert board_images.get_counts_for_board(board) == (1, 3) + assert board_images.get_counts_for_board(other) == (1, 0) + + +def test_a_video_is_on_one_board_at_a_time_and_counted_by_category( + database: Database, board_records: BoardRecordStorage, board_videos: BoardVideoRecordStorage +) -> None: + first = board_records.save("First", ALICE).board_id + second = board_records.save("Second", ALICE).board_id + general = _video(database, "general.mp4") + asset = _video(database, "asset.mp4", category=ImageCategory.USER) + intermediate = _video(database, "intermediate.mp4", intermediate=True) + for name in (general, asset, intermediate): + board_videos.add_video_to_board(first, name) + board_videos.add_video_to_board(second, general) + + assert board_videos.get_board_for_video(general) == second + assert board_videos.get_counts_for_board(first) == (1, 1) # asset.mp4 + assert board_videos.get_counts_for_board(second) == (1, 0) + + board_videos.remove_video_from_board(general) + + assert board_videos.get_board_for_video(general) is None + + +def test_video_names_follow_the_filters( + database: Database, board_records: BoardRecordStorage, board_videos: BoardVideoRecordStorage +) -> None: + """Deleting a board with its media takes the names from here: the account filter keeps the videos of another + account out of it.""" + board = board_records.save("Board", ALICE).board_id + for name, category, intermediate, user in ( + ("general.mp4", ImageCategory.GENERAL, False, ALICE), + ("asset.mp4", ImageCategory.USER, False, ALICE), + ("intermediate.mp4", ImageCategory.GENERAL, True, ALICE), + ("bobs.mp4", ImageCategory.GENERAL, False, BOB), + ): + video = _video(database, name, category=category, intermediate=intermediate, user_id=user) + board_videos.add_video_to_board(board, video) + _video(database, "alices_loose.mp4") + _video(database, "bobs_loose.mp4", user_id=BOB) + + def names(board_id: str, **filters: object) -> set[str]: + filters = {"categories": None, "is_intermediate": None, **filters} + return set(board_videos.get_all_board_video_names_for_board(board_id, **filters)) # type: ignore[arg-type] + + assert names(board) == {"general.mp4", "asset.mp4", "intermediate.mp4", "bobs.mp4"} + assert names("none") == {"alices_loose.mp4", "bobs_loose.mp4"} + assert names("none", user_id=BOB) == {"bobs_loose.mp4"} + assert names(board, user_id=BOB) == {"bobs.mp4"} + assert names(board, categories=[ImageCategory.USER]) == {"asset.mp4"} + assert names(board, is_intermediate=True) == {"intermediate.mp4"} + + +# endregion diff --git a/tests/app/services/boards/test_boards_default.py b/tests/app/services/boards/test_boards_default.py index d2ee6b7c214..a2791f2ee71 100644 --- a/tests/app/services/boards/test_boards_default.py +++ b/tests/app/services/boards/test_boards_default.py @@ -2,9 +2,16 @@ from types import SimpleNamespace from unittest.mock import MagicMock +import pytest + from invokeai.app.services.board_records.board_records_common import BoardRecordOrderBy +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage from invokeai.app.services.invoker import Invoker -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import LockTimeoutError +from invokeai.app.services.shared.database.queries.boards import BoardQueries +from invokeai.app.services.shared.pagination import SQLiteDirection +from tests.fixtures.sqlite_database import sqlite_cursor def test_board_cover_uses_gallery_tie_breakers(mock_invoker: Invoker) -> None: @@ -35,9 +42,8 @@ def test_board_listing_fetches_media_summaries_once(mock_invoker: Invoker) -> No mock_invoker.services.gallery.get_board_media_summaries = MagicMock(return_value=summaries) # type: ignore[attr-defined] mock_invoker.services.image_records.get_most_recent_image_for_board = MagicMock() mock_invoker.services.video_records.get_most_recent_video_for_board = MagicMock() - mock_invoker.services.board_image_records.get_image_count_for_board = MagicMock() - mock_invoker.services.board_image_records.get_asset_count_for_board = MagicMock() - mock_invoker.services.board_video_records.get_video_count_for_board = MagicMock() + mock_invoker.services.board_image_records.get_counts_for_board = MagicMock() + mock_invoker.services.board_video_records.get_counts_for_board = MagicMock() result = mock_invoker.services.boards.get_many( user_id="user", @@ -53,9 +59,8 @@ def test_board_listing_fetches_media_summaries_once(mock_invoker: Invoker) -> No assert set(mock_invoker.services.gallery.get_board_media_summaries.call_args.args[0]) == set(board_ids) # type: ignore[attr-defined] mock_invoker.services.image_records.get_most_recent_image_for_board.assert_not_called() mock_invoker.services.video_records.get_most_recent_video_for_board.assert_not_called() - mock_invoker.services.board_image_records.get_image_count_for_board.assert_not_called() - mock_invoker.services.board_image_records.get_asset_count_for_board.assert_not_called() - mock_invoker.services.board_video_records.get_video_count_for_board.assert_not_called() + mock_invoker.services.board_image_records.get_counts_for_board.assert_not_called() + mock_invoker.services.board_video_records.get_counts_for_board.assert_not_called() def test_admin_board_listing_batches_owner_lookup(mock_invoker: Invoker) -> None: @@ -131,8 +136,32 @@ def test_non_admin_board_listing_skips_owner_lookup(mock_invoker: Invoker) -> No mock_invoker.services.users.get_many.assert_not_called() # type: ignore[attr-defined] -def test_board_records_get_does_not_disguise_a_storage_error_as_not_found() -> None: - """A sqlite3.Error out of the SELECT must stay a sqlite3.Error. +def test_a_board_dto_counts_each_kind_of_media(mock_invoker: Invoker, mock_sqlite_database: Database) -> None: + services = mock_invoker.services + board_id = services.board_records.save("Board", "system").board_id + with sqlite_cursor(mock_sqlite_database) as cursor: + cursor.executemany( + "INSERT INTO images (image_name, image_origin, image_category, width, height) VALUES (?, 'internal', ?, 1, 1)", + [("general.png", "general"), ("control.png", "control"), ("mask.png", "mask")], + ) + cursor.executemany( + "INSERT INTO videos (video_name, video_origin, video_category, width, height) VALUES (?, 'internal', ?, 1, 1)", + [("first.mp4", "general"), ("second.mp4", "general"), ("asset.mp4", "user")], + ) + for name in ("general.png", "control.png", "mask.png"): + services.board_image_records.add_image_to_board(board_id, name) + for name in ("first.mp4", "second.mp4", "asset.mp4"): + services.board_video_records.add_video_to_board(board_id, name) + + dto = services.boards.get_dto(board_id) + + assert (dto.image_count, dto.asset_count, dto.video_count, dto.asset_video_count) == (1, 2, 3, 1) + + +def test_board_records_get_does_not_disguise_a_storage_error_as_not_found( + database: Database, monkeypatch: pytest.MonkeyPatch +) -> None: + """A database error out of the SELECT must stay a database error, here a lock held too long elsewhere. Translating it made BoardRecordNotFoundException mean "no such board, OR the database is unreadable", and the board-image batch routes cannot tell those apart: they decide write @@ -141,26 +170,11 @@ def test_board_records_get_does_not_disguise_a_storage_error_as_not_found() -> N and the client would keep showing them as moved until the next refresh. Mirrors test_image_records_get_does_not_disguise_a_storage_error_as_not_found. """ - import sqlite3 - from contextlib import contextmanager - from typing import Any - - import pytest - - from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage - - storage = SqliteBoardRecordStorage.__new__(SqliteBoardRecordStorage) - - class _Cursor: - def execute(self, *args: Any, **kwargs: Any) -> None: - raise sqlite3.OperationalError("database disk image is malformed") - class _Db: - @contextmanager - def transaction(self): - yield _Cursor() + def locked(self: BoardQueries, board_id: str) -> None: + raise LockTimeoutError("database is locked") - storage._db = _Db() # pyright: ignore[reportAttributeAccessIssue] + monkeypatch.setattr(BoardQueries, "get", locked) - with pytest.raises(sqlite3.OperationalError): - storage.get("board-1") + with pytest.raises(LockTimeoutError): + BoardRecordStorage(database).get("board-1") diff --git a/tests/app/services/client_state_persistence/test_client_state_persistence.py b/tests/app/services/client_state_persistence/test_client_state_persistence.py new file mode 100644 index 00000000000..633ba067420 --- /dev/null +++ b/tests/app/services/client_state_persistence/test_client_state_persistence.py @@ -0,0 +1,204 @@ +"""Client state persistence on every database backend.""" + +import threading +from collections.abc import Callable +from typing import Any + +import pytest +from sqlalchemy import insert, select + +from invokeai.app.services.client_state_persistence.client_state_persistence_default import ClientStatePersistence +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries.client_state import ClientStateQueries +from invokeai.app.services.shared.database.schema.client_state import client_state +from invokeai.app.services.shared.database.schema.media_references import media_references +from invokeai.app.services.users.users_common import SYSTEM_USER_ID, UserCreateRequest +from invokeai.app.services.users.users_default import UserService +from tests.fixtures.database import external_test_db_url + +server_only = pytest.mark.skipif( + external_test_db_url() is None, reason="needs a MySQL or MariaDB server (INVOKEAI_TEST_DB_URL)" +) + +# Names an image and a video; client state values hold whatever JSON the web client writes. +DOCUMENT = '{"layers": [{"image_name": "a.png"}], "video": {"video_name": "v.mp4"}}' + + +@pytest.fixture +def store(database: Database) -> ClientStatePersistence: + return ClientStatePersistence(database) + + +@pytest.fixture +def other_user(database: Database) -> str: + request = UserCreateRequest(email="other@test.com", display_name="Other", password="Password123") + return UserService(database).create(request, strict_password_checking=False).user_id + + +def _insert(database: Database, user_id: str, key: str, updated_at: str) -> None: + with database.begin(write=True) as conn: + conn.execute(insert(client_state).values(user_id=user_id, key=key, value="{}", updated_at=updated_at)) + + +def _references(database: Database) -> set[tuple[Any, ...]]: + """Every row of the reference index: (owner_kind, user_id, owner_id, media_kind, media_name).""" + with database.begin(write=False) as conn: + return {tuple(row) for row in conn.execute(select(media_references))} + + +def _client_state_reference(user_id: str, key: str, media_kind: str, media_name: str) -> tuple[str, ...]: + return ("client_state", user_id, key, media_kind, media_name) + + +def test_a_value_is_kept_per_account_and_replaced_by_the_next(store: ClientStatePersistence, other_user: str) -> None: + store.set_by_key(SYSTEM_USER_ID, "key", "first") + store.set_by_key(SYSTEM_USER_ID, "key", "second") + store.set_by_key(other_user, "key", "theirs") + + assert store.get_by_key(SYSTEM_USER_ID, "key") == "second" + assert store.get_by_key(other_user, "key") == "theirs" + assert store.get_by_key(SYSTEM_USER_ID, "missing") is None + + +def test_replacing_a_value_renews_when_it_was_written(database: Database, store: ClientStatePersistence) -> None: + _insert(database, SYSTEM_USER_ID, "key", "2000-01-01 00:00:00.000") + + store.set_by_key(SYSTEM_USER_ID, "key", "new") + + with database.begin(write=False) as conn: + assert conn.execute(select(client_state.c.updated_at)).scalar_one() > "2000-01-01 00:00:00.000" + + +def test_keys_with_a_prefix_list_most_recently_written_first_ignoring_case( + database: Database, store: ClientStatePersistence, other_user: str +) -> None: + _insert(database, SYSTEM_USER_ID, "snap:old", "2026-01-01 00:00:00.000") + _insert(database, SYSTEM_USER_ID, "SNAP:upper", "2026-01-03 00:00:00.000") + _insert(database, SYSTEM_USER_ID, "snap:new", "2026-01-02 00:00:00.000") + _insert(database, SYSTEM_USER_ID, "other:key", "2026-01-04 00:00:00.000") + _insert(database, other_user, "snap:theirs", "2026-01-05 00:00:00.000") + + assert store.get_keys_by_prefix(SYSTEM_USER_ID, "snap:") == ["SNAP:upper", "snap:new", "snap:old"] + + +def test_like_wildcards_in_a_prefix_match_only_themselves(store: ClientStatePersistence) -> None: + for key in ("a%b", "axb", "a_c", "ayc", "a\\d", "aed"): + store.set_by_key(SYSTEM_USER_ID, key, "{}") + + assert store.get_keys_by_prefix(SYSTEM_USER_ID, "a%") == ["a%b"] + assert store.get_keys_by_prefix(SYSTEM_USER_ID, "a_") == ["a_c"] + assert store.get_keys_by_prefix(SYSTEM_USER_ID, "a\\") == ["a\\d"] + + +def test_media_references_follow_the_values_that_name_them(database: Database, store: ClientStatePersistence) -> None: + def references(*names: tuple[str, str, str]) -> set[tuple[str, ...]]: + return {_client_state_reference(SYSTEM_USER_ID, key, kind, name) for key, kind, name in names} + + store.set_by_key(SYSTEM_USER_ID, "canvas", DOCUMENT) + store.set_by_key(SYSTEM_USER_ID, "staging", '{"imageName": "c.png"}') + assert _references(database) == references( + ("canvas", "image", "a.png"), ("canvas", "video", "v.mp4"), ("staging", "image", "c.png") + ) + + store.set_by_key(SYSTEM_USER_ID, "canvas", '{"imageName": "b.png"}') + assert _references(database) == references(("canvas", "image", "b.png"), ("staging", "image", "c.png")) + + store.delete_by_key(SYSTEM_USER_ID, "canvas") + assert _references(database) == references(("staging", "image", "c.png")) + + store.delete(SYSTEM_USER_ID) + assert _references(database) == set() + assert store.get_keys_by_prefix(SYSTEM_USER_ID, "") == [] + + +def test_one_accounts_client_state_leaves_other_accounts_alone( + database: Database, store: ClientStatePersistence, other_user: str +) -> None: + # The web client uses the same keys for every account. + store.set_by_key(other_user, "canvas", DOCUMENT) + store.set_by_key(other_user, "other", DOCUMENT) + store.set_by_key(SYSTEM_USER_ID, "canvas", DOCUMENT) + project_reference = ("project", SYSTEM_USER_ID, "project", "image", "p.png") + with database.begin(write=True) as conn: + conn.execute( + insert(media_references).values(dict(zip(media_references.c.keys(), project_reference, strict=False))) + ) + + store.set_by_key(SYSTEM_USER_ID, "canvas", "{}") + store.delete_by_key(SYSTEM_USER_ID, "canvas") + store.set_by_key(SYSTEM_USER_ID, "other", DOCUMENT) + store.delete(SYSTEM_USER_ID) + + assert store.get_by_key(other_user, "canvas") == DOCUMENT + assert store.get_by_key(other_user, "other") == DOCUMENT + assert _references(database) == { + _client_state_reference(other_user, key, kind, name) + for key in ("canvas", "other") + for kind, name in (("image", "a.png"), ("video", "v.mp4")) + } | {project_reference} + + +def _pause_after(monkeypatch: pytest.MonkeyPatch, method: str, while_paused: Callable[[], None]) -> Callable[[], None]: + """Makes `ClientStateQueries.` run `while_paused` on another thread right after it, inside its + transaction, and wait for it. Returns a function to call afterwards, which raises what that thread raised.""" + real = getattr(ClientStateQueries, method) + finished = threading.Event() + errors: list[BaseException] = [] + + def run() -> None: + try: + while_paused() + except BaseException as e: # noqa: BLE001 - raised again by the returned function + errors.append(e) + finally: + finished.set() + + def then_pause(self: ClientStateQueries, *args: Any, **kwargs: Any) -> Any: + result = real(self, *args, **kwargs) + threading.Thread(target=run).start() + assert finished.wait(timeout=30) + return result + + monkeypatch.setattr(ClientStateQueries, method, then_pause) + + def check() -> None: + assert finished.is_set() + if errors: + raise errors[0] + + return check + + +@server_only +def test_deleting_a_key_while_it_is_first_written_leaves_the_new_value_its_references( + database: Database, store: ClientStatePersistence, monkeypatch: pytest.MonkeyPatch +) -> None: + # On a server the delete finds no row to lock, so the first write goes ahead while the delete's transaction is + # still open. (SQLite runs one transaction at a time; there the write would wait for the delete.) + check = _pause_after(monkeypatch, "delete", lambda: store.set_by_key(SYSTEM_USER_ID, "canvas", DOCUMENT)) + + store.delete_by_key(SYSTEM_USER_ID, "canvas") + check() + + assert store.get_by_key(SYSTEM_USER_ID, "canvas") == DOCUMENT + assert _references(database) == { + _client_state_reference(SYSTEM_USER_ID, "canvas", "image", "a.png"), + _client_state_reference(SYSTEM_USER_ID, "canvas", "video", "v.mp4"), + } + + +@server_only +def test_deleting_all_values_while_a_new_key_is_written_leaves_that_key_its_references( + database: Database, store: ClientStatePersistence, monkeypatch: pytest.MonkeyPatch +) -> None: + store.set_by_key(SYSTEM_USER_ID, "old", '{"imageName": "old.png"}') + check = _pause_after(monkeypatch, "delete_all", lambda: store.set_by_key(SYSTEM_USER_ID, "new", DOCUMENT)) + + store.delete(SYSTEM_USER_ID) + check() + + assert store.get_keys_by_prefix(SYSTEM_USER_ID, "") == ["new"] + assert _references(database) == { + _client_state_reference(SYSTEM_USER_ID, "new", "image", "a.png"), + _client_state_reference(SYSTEM_USER_ID, "new", "video", "v.mp4"), + } diff --git a/tests/app/services/fonts/test_fonts_service.py b/tests/app/services/fonts/test_fonts_service.py index b1ab4bcd9a3..be2513bf6f5 100644 --- a/tests/app/services/fonts/test_fonts_service.py +++ b/tests/app/services/fonts/test_fonts_service.py @@ -5,14 +5,14 @@ import subprocess import threading from pathlib import Path -from unittest.mock import Mock import pytest from fontTools.fontBuilder import FontBuilder from fontTools.pens.ttGlyphPen import TTGlyphPen from fontTools.ttLib.tables.TupleVariation import TupleVariation +from sqlalchemy import select -from invokeai.app.services.config.config_default import InvokeAIAppConfig +from invokeai.app.services.fonts import fonts_default from invokeai.app.services.fonts.fonts_common import FontScope, FontSource, FontUploadResult from invokeai.app.services.fonts.fonts_default import ( MAX_INDEXED_FONT_BYTES, @@ -23,18 +23,19 @@ FontService, FontValidationError, ) -from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase -from invokeai.app.services.shared.sqlite.sqlite_util import init_db +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import fonts as font_queries +from invokeai.app.services.shared.database.queries.fonts import FontQueries, UploadedFont +from invokeai.app.services.shared.database.schema.fonts import fonts as fonts_table from invokeai.app.services.users.users_common import UserCreateRequest from invokeai.app.services.users.users_default import UserService +from tests.fixtures.races import when_called @pytest.fixture -def font_service(tmp_path: Path) -> FontService: - config = InvokeAIAppConfig(use_memory_db=True) - db = init_db(config, logging.getLogger("font-tests"), Mock(spec=ImageFileStorageBase)) +def font_service(tmp_path: Path, database: Database) -> FontService: service = FontService( - db=db, + database, fonts_dir=tmp_path / "fonts", storage_dir=tmp_path / "uploaded-fonts", logger=logging.getLogger("font-tests"), @@ -48,8 +49,8 @@ def font_bytes() -> bytes: return (Path(__file__).parents[4] / "invokeai" / "assets" / "fonts" / "inter" / "Inter-Regular.ttf").read_bytes() -def _create_user(service: FontService, email: str) -> str: - users = UserService(service._db) +def _create_user(database: Database, email: str) -> str: + users = UserService(database) return users.create(UserCreateRequest(email=email, password="TestPass123", display_name=email)).user_id @@ -104,9 +105,11 @@ def _variable_font_bytes() -> bytes: return output.getvalue() -def test_private_uploads_are_deduplicated_and_isolated(font_service: FontService, font_bytes: bytes) -> None: - alice = _create_user(font_service, "alice-fonts@test.com") - bob = _create_user(font_service, "bob-fonts@test.com") +def test_private_uploads_are_deduplicated_and_isolated( + font_service: FontService, database: Database, font_bytes: bytes +) -> None: + alice = _create_user(database, "alice-fonts@test.com") + bob = _create_user(database, "bob-fonts@test.com") first = font_service.upload(user_id=alice, filename="Inter-Regular.ttf", data=font_bytes) duplicate = font_service.upload(user_id=alice, filename="renamed.ttf", data=font_bytes) @@ -121,17 +124,22 @@ def test_private_uploads_are_deduplicated_and_isolated(font_service: FontService def test_user_cleanup_removes_private_files_but_preserves_shared_files( - font_service: FontService, font_bytes: bytes + font_service: FontService, database: Database, font_bytes: bytes ) -> None: - users = UserService(font_service._db) - user_id = _create_user(font_service, "deleted-font-owner@test.com") + users = UserService(database) + user_id = _create_user(database, "deleted-font-owner@test.com") + + other_user_id = _create_user(database, "kept-font-owner@test.com") private = font_service.upload(user_id=user_id, filename="Private.ttf", data=font_bytes).font shared = font_service.upload(user_id="system", filename="Shared.ttf", data=font_bytes, scope=FontScope.SHARED).font + other = font_service.upload(user_id=other_user_id, filename="Other.ttf", data=font_bytes).font assert private.storage_path is not None assert shared.storage_path is not None + assert other.storage_path is not None private_path = font_service._storage_dir / private.storage_path shared_path = font_service._storage_dir / shared.storage_path + other_path = font_service._storage_dir / other.storage_path assert private_path.is_file() assert shared_path.is_file() @@ -139,14 +147,36 @@ def test_user_cleanup_removes_private_files_but_preserves_shared_files( users.delete(user_id) font_service.cleanup_user(user_id, font_storage_paths) + # A private upload is the owner's own copy, even of content shared already. + assert other.scope == FontScope.PRIVATE and other.id != shared.id + assert font_storage_paths == (private.storage_path,) assert not private_path.exists() assert shared_path.is_file() + assert other_path.is_file() assert font_service.get(private.id) is None assert font_service.get(shared.id) is not None + assert font_service.get(other.id) is not None + + +def test_orphan_sweep_removes_only_unreferenced_managed_files(font_service: FontService, font_bytes: bytes) -> None: + uploaded = font_service.upload(user_id="system", filename="Inter.ttf", data=font_bytes).font + assert uploaded.storage_path is not None + orphan = font_service._storage_dir / f"font_{'0' * 32}.ttf" + interrupted = font_service._storage_dir / f".font_{'1' * 32}.ttf.{'2' * 32}.tmp" + operator_file = font_service._storage_dir / "notes.txt" + for path in (orphan, interrupted, operator_file): + path.write_bytes(b"x") + + font_service.cleanup_orphaned_files() + + assert (font_service._storage_dir / uploaded.storage_path).is_file() + assert not orphan.exists() + assert not interrupted.exists() + assert operator_file.is_file() def test_shared_uploads_are_visible_and_directory_fonts_are_read_only( - font_service: FontService, font_bytes: bytes, tmp_path: Path + font_service: FontService, database: Database, font_bytes: bytes, tmp_path: Path ) -> None: shared = font_service.upload(user_id="system", filename="Inter.ttf", data=font_bytes, scope=FontScope.SHARED) assert shared.font.scope == FontScope.SHARED @@ -161,15 +191,20 @@ def test_shared_uploads_are_visible_and_directory_fonts_are_read_only( with pytest.raises(FontDeleteForbiddenError): font_service.delete(user_id="system", font_id=directory[0].id, is_admin=True) + # Deleting an upload never deletes a directory font of the same id. + assert database.queries.fonts.delete_upload(directory[0].id) is False + assert font_service.get(directory[0].id) is not None -def test_validation_and_quota_do_not_leave_files(font_service: FontService, font_bytes: bytes) -> None: +def test_validation_and_quota_do_not_leave_files( + font_service: FontService, database: Database, font_bytes: bytes +) -> None: validated = font_service.validate("Inter-Regular.ttf", font_bytes) assert validated.content_hash assert font_service.list(user_id="system")[1] == 0 tiny_quota = FontService( - db=font_service._db, + database, fonts_dir=font_service._fonts_dir, storage_dir=font_service._storage_dir / "quota", max_library_bytes=len(font_bytes) - 1, @@ -195,9 +230,11 @@ def test_rescan_removes_missing_directory_rows_and_detects_replacement( with pytest.raises(FontChangedError): font_service.read_file(user_id="system", font_id=record.id, expected_hash=record.content_hash) + revision = font_service.revision directory_file.unlink() assert font_service.rescan_directory() == 0 assert font_service.list(user_id="system")[1] == 0 + assert font_service.revision > revision def test_rescan_clears_index_when_directory_is_removed( @@ -332,10 +369,11 @@ def test_variable_instances_are_pinned_and_cached(font_service: FontService, tmp def test_deleted_font_cannot_repopulate_instance_cache_after_worker_finishes( - font_service: FontService, font_bytes: bytes, monkeypatch: pytest.MonkeyPatch + font_service: FontService, database: Database, font_bytes: bytes, monkeypatch: pytest.MonkeyPatch ) -> None: - user_id = _create_user(font_service, "inflight-font-owner@test.com") + user_id = _create_user(database, "inflight-font-owner@test.com") uploaded = font_service.upload(user_id=user_id, filename="Variable.ttf", data=_variable_font_bytes()).font + assert [instance.name for instance in uploaded.instances] == ["Bold"] font_service.instance( user_id=user_id, font_id=uploaded.id, @@ -420,3 +458,217 @@ def kill(self) -> None: # This call proves the worker admission and active-slot semaphores were both released. uploaded = font_service.upload(user_id="system", filename="Inter.ttf", data=font_bytes) assert uploaded.created is True + + +def _catalog_upload( + database: Database, font_id: str, family: str, *, owner_id: str | None, label: str = "Regular", digit: str = "a" +) -> None: + """Catalogs an upload without a file: listing reads the catalog only.""" + database.queries.fonts.insert_upload( + UploadedFont( + id=font_id, + owner_id=owner_id, + scope=FontScope.PRIVATE if owner_id is not None else FontScope.SHARED, + filename=f"{font_id}.ttf", + storage_path=f"{font_id}.ttf", + family=family, + label=label, + style="normal", + weight=400, + content_hash=digit * 64, + byte_size=100, + axes=(), + instances=(), + ) + ) + + +def test_listing_filters_by_account_scope_and_hash_and_orders_ignoring_case( + font_service: FontService, database: Database, font_bytes: bytes, tmp_path: Path +) -> None: + alice = _create_user(database, "alice-listing@test.com") + bob = _create_user(database, "bob-listing@test.com") + _catalog_upload(database, "font_b", "beta", owner_id=alice) + _catalog_upload(database, "font_a2", "alpha", owner_id=alice, label="Bold", digit="b") + _catalog_upload(database, "font_a1", "Alpha", owner_id=None, label="bold", digit="c") + _catalog_upload(database, "font_a0", "ALPHA", owner_id=None, label="Bold", digit="d") + _catalog_upload(database, "font_bob", "Aardvark", owner_id=bob, digit="e") + directory_file = tmp_path / "fonts" / "Inter.ttf" + directory_file.parent.mkdir(parents=True) + directory_file.write_bytes(font_bytes) + assert font_service.rescan_directory() == 1 + inter = font_service.list(user_id=bob, scope=FontScope.SHARED, search="inter")[0][0] + + def ids(**kwargs: object) -> tuple[list[str], int]: + records, total = font_service.list(user_id=alice, **kwargs) # type: ignore[arg-type] + return [record.id for record in records], total + + # Family, then label, ignoring case; then id. Bob's private font is not Alice's to see. + assert ids() == (["font_a0", "font_a1", "font_a2", "font_b", inter.id], 5) + assert ids(offset=1, limit=2) == (["font_a1", "font_a2"], 5) + assert ids(offset=4, limit=10) == ([inter.id], 5) + assert ids(offset=9, limit=10) == ([], 5) + assert ids(limit=0) == ([], 5) + assert ids(scope=FontScope.PRIVATE) == (["font_a2", "font_b"], 2) + assert ids(scope=FontScope.SHARED) == (["font_a0", "font_a1", inter.id], 3) + assert ids(content_hash="c" * 64) == (["font_a1"], 1) + assert ids(content_hash="e" * 64) == ([], 0) + # Full pages, which count the fonts with a statement of their own. + assert ids(scope=FontScope.PRIVATE, limit=1) == (["font_a2"], 2) + assert ids(scope=FontScope.SHARED, offset=1, limit=1) == (["font_a1"], 3) + assert ids(content_hash="c" * 64, limit=0) == ([], 1) + assert font_service.list(user_id=bob, scope=FontScope.PRIVATE)[1] == 1 + + +def test_listing_search_matches_text_ignoring_case(font_service: FontService, database: Database) -> None: + _catalog_upload(database, "font_percent", "Grotesk 100%", owner_id=None, digit="a") + _catalog_upload(database, "font_thousand", "Grotesk 1000", owner_id=None, digit="b") + _catalog_upload(database, "font_label", "Serif", owner_id=None, label="Condensed_Bold", digit="c") + _catalog_upload(database, "font_space", "Serif", owner_id=None, label="Condensed Bold", digit="d") + + def found(search: str) -> list[str]: + return [record.id for record in font_service.list(user_id="system", search=search)[0]] + + assert found("grotesk") == ["font_percent", "font_thousand"] + # `%` and `_` match only themselves. + assert found("100%") == ["font_percent"] + assert found("condensed_") == ["font_label"] + assert found("FONT_SPACE.TTF") == ["font_space"] + first, total = font_service.list(user_id="system", search="grotesk", limit=1) + assert ([record.id for record in first], total) == (["font_percent"], 2) + + +def test_quota_counts_each_place_of_uploads_on_its_own( + font_service: FontService, database: Database, font_bytes: bytes +) -> None: + alice = _create_user(database, "alice-quota@test.com") + bob = _create_user(database, "bob-quota@test.com") + one_font = FontService( + database, + fonts_dir=font_service._fonts_dir, + storage_dir=font_service._storage_dir, + max_library_bytes=len(font_bytes), + ) + try: + assert one_font.upload(user_id=alice, filename="Inter.ttf", data=font_bytes).created + # Neither Alice's upload counts toward the shared fonts, nor do they count toward Bob's. + assert one_font.upload(user_id=alice, filename="Inter.ttf", data=font_bytes, scope=FontScope.SHARED).created + assert one_font.upload(user_id=bob, filename="Inter.ttf", data=font_bytes).created + with pytest.raises(FontQuotaExceededError): + one_font.upload(user_id=alice, filename="Variable.ttf", data=_variable_font_bytes()) + with pytest.raises(FontQuotaExceededError): + one_font.upload(user_id=bob, filename="Variable.ttf", data=_variable_font_bytes(), scope=FontScope.SHARED) + finally: + one_font.stop() + + +def test_rescan_changes_the_catalog_only_when_the_directory_changed( + font_service: FontService, font_bytes: bytes, tmp_path: Path +) -> None: + directory_file = tmp_path / "fonts" / "Face.ttf" + directory_file.parent.mkdir(parents=True) + directory_file.write_bytes(font_bytes) + assert font_service.rescan_directory() == 1 + indexed = font_service.list(user_id="system")[0] + revision = font_service.revision + + assert font_service.rescan_directory() == 1 + assert font_service.revision == revision + assert font_service.list(user_id="system")[0] == indexed + + directory_file.write_bytes(_variable_font_bytes()) + assert font_service.rescan_directory() == 1 + assert font_service.revision > revision + [replaced] = font_service.list(user_id="system")[0] + assert replaced.id == indexed[0].id + assert (replaced.family, replaced.source_path) == ("Fixture Variable", "Face.ttf") + assert replaced.content_hash != indexed[0].content_hash + assert [axis.tag for axis in replaced.axes] == ["wght"] + assert [instance.name for instance in replaced.instances] == ["Bold"] + # The catalog holds the variable font as found, so another scan changes nothing. + revision = font_service.revision + assert font_service.rescan_directory() == 1 + assert font_service.revision == revision + + +def test_rescan_updates_changed_fonts_and_deletes_only_missing_ones( + font_service: FontService, database: Database, font_bytes: bytes, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + root = tmp_path / "fonts" + root.mkdir() + for name in ("Missing.ttf", "Changed.ttf", "Kept.ttf"): + (root / name).write_bytes(font_bytes) + assert font_service.rescan_directory() == 3 + before = _directory_rows(database) + + monkeypatch.setattr(font_queries, "now_text", lambda: "2031-02-03 04:05:06.789") + (root / "Missing.ttf").unlink() + (root / "Changed.ttf").write_bytes(_variable_font_bytes()) + assert font_service.rescan_directory() == 2 + + after = _directory_rows(database) + assert set(after) == {"Changed.ttf", "Kept.ttf"} + assert after["Kept.ttf"] == before["Kept.ttf"] + assert after["Changed.ttf"][0] == "Fixture Variable" + assert after["Changed.ttf"][1] == "2031-02-03 04:05:06.789" + + +def _directory_rows(database: Database) -> dict[str, tuple[str, str]]: + columns = fonts_table.c + statement = select(columns.source_path, columns.family, columns.updated_at).where(columns.source == "directory") + with database.begin(write=False) as conn: + return {path: (family, updated_at) for path, family, updated_at in conn.execute(statement).all()} + + +def test_rescan_skips_a_path_longer_than_the_catalog_holds( + font_service: FontService, font_bytes: bytes, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + root = tmp_path / "fonts" + (root / "nested").mkdir(parents=True) + (root / "Short.ttf").write_bytes(font_bytes) + (root / "nested" / "Long.ttf").write_bytes(font_bytes) + monkeypatch.setattr(fonts_default, "MAX_SOURCE_PATH_LENGTH", len("nested/Long.tt")) + + assert font_service.rescan_directory() == 1 + assert [record.source_path for record in font_service.list(user_id="system")[0]] == ["Short.ttf"] + + +def test_upload_removes_its_file_only_when_its_row_is_missing( + font_service: FontService, font_bytes: bytes, monkeypatch: pytest.MonkeyPatch +) -> None: + real_insert = FontQueries.insert_upload + + def refused(self: FontQueries, font: UploadedFont) -> None: + raise RuntimeError("refused") + + monkeypatch.setattr(FontQueries, "insert_upload", refused) + with pytest.raises(RuntimeError, match="refused"): + font_service.upload(user_id="system", filename="Inter.ttf", data=font_bytes) + assert not list(font_service._storage_dir.glob("font_*")) + + # A commit that took effect but reported a failure, as when the connection drops during it. + def lost_reply(self: FontQueries, font: UploadedFont) -> None: + real_insert(self, font) + raise ConnectionError("lost") + + monkeypatch.setattr(FontQueries, "insert_upload", lost_reply) + with pytest.raises(ConnectionError): + font_service.upload(user_id="system", filename="Inter.ttf", data=font_bytes) + [record] = font_service.list(user_id="system")[0] + assert record.storage_path is not None + assert (font_service._storage_dir / record.storage_path).is_file() + + +def test_rescans_run_one_at_a_time( + font_service: FontService, font_bytes: bytes, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + directory_file = tmp_path / "fonts" / "Inter.ttf" + directory_file.parent.mkdir(parents=True) + directory_file.write_bytes(font_bytes) + assert font_service.rescan_directory() == 1 + + # The second scan starts once the first has read the catalog, and waits for it to finish. + ended = when_called(monkeypatch, FontQueries, "directory_fonts", font_service.rescan_directory) + assert font_service.rescan_directory() == 1 + assert ended() == [] + assert font_service.list(user_id="system")[1] == 1 diff --git a/tests/app/services/gallery/test_gallery_default.py b/tests/app/services/gallery/test_gallery_default.py index 5f39115522d..1e21f412ad5 100644 --- a/tests/app/services/gallery/test_gallery_default.py +++ b/tests/app/services/gallery/test_gallery_default.py @@ -1,4 +1,4 @@ -"""Regression tests for SqliteGalleryService multiuser isolation and date-based +"""Regression tests for the gallery service's multiuser isolation and date-based virtual boards. Covers JPPhoto's code-review findings (PR #9163): @@ -17,44 +17,46 @@ from collections.abc import Callable from types import SimpleNamespace -from typing import TypeVar +from typing import Any, TypeVar import pytest +from sqlalchemy import update -from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage -from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage -from invokeai.app.services.board_video_records.board_video_records_sqlite import SqliteBoardVideoRecordStorage -from invokeai.app.services.config.config_default import InvokeAIAppConfig +from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage +from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardVisibility +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage +from invokeai.app.services.board_video_records.board_video_records_default import BoardVideoRecordStorage from invokeai.app.services.gallery.gallery_common import GalleryItemKind -from invokeai.app.services.gallery.gallery_default import SqliteGalleryService +from invokeai.app.services.gallery.gallery_default import GalleryService from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin -from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import gallery as gallery_queries +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.videos import videos +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.services.urls.urls_default import LocalUrlService -from invokeai.app.services.video_records.video_records_sqlite import SqliteVideoRecordStorage -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database +from invokeai.app.services.video_records.video_records_default import VideoRecordStorage +from tests.fixtures.database import capture_statements, explain_query_plan @pytest.fixture -def services(): - config = InvokeAIAppConfig(use_memory_db=True) - logger = InvokeAILogger.get_logger(config=config) - db = create_mock_sqlite_database(config, logger) - gallery = SqliteGalleryService(db=db) +def services(database: Database): + gallery = GalleryService(database) gallery.start(SimpleNamespace(services=SimpleNamespace(urls=LocalUrlService()))) # type: ignore[arg-type] return { + "db": database, "gallery": gallery, - "images": SqliteImageRecordStorage(db=db), - "videos": SqliteVideoRecordStorage(db=db), - "boards": SqliteBoardRecordStorage(db=db), - "board_images": SqliteBoardImageRecordStorage(db=db), - "board_videos": SqliteBoardVideoRecordStorage(db=db), + "images": ImageRecordStorage(database), + "videos": VideoRecordStorage(database), + "boards": BoardRecordStorage(database), + "board_images": BoardImageRecordStorage(database), + "board_videos": BoardVideoRecordStorage(database), } def _save_image( - store: SqliteImageRecordStorage, + store: ImageRecordStorage, name: str, user_id: str, category: ImageCategory = ImageCategory.GENERAL, @@ -72,7 +74,7 @@ def _save_image( def _save_video( - store: SqliteVideoRecordStorage, + store: VideoRecordStorage, name: str, user_id: str, category: ImageCategory = ImageCategory.GENERAL, @@ -95,16 +97,11 @@ def _save_video( def _capture_plan(services, call: Callable[[], T], statement_marker: str) -> tuple[T, str, list[str]]: - db = services["images"]._db - statements: list[str] = [] - db._conn.set_trace_callback(statements.append) - try: + database = services["db"] + with capture_statements(database) as statements: result = call() - finally: - db._conn.set_trace_callback(None) - statement = next(statement for statement in statements if statement_marker in statement) - details = [row[3] for row in db._conn.execute(f"EXPLAIN QUERY PLAN {statement}").fetchall()] - return result, statement, details + statement, parameters = next((sql, parameters) for sql, parameters in statements if statement_marker in sql) + return result, statement, explain_query_plan(database, statement, parameters) @pytest.fixture @@ -152,15 +149,17 @@ def test_explicit_none_board_still_isolates(self, seeded) -> None: def _backdate(services, table: str, name_col: str, name: str, created_at: str) -> None: """Rewrites created_at so tests can build multi-date galleries (save() always stamps now).""" - db = services["images"]._db - with db.transaction() as cursor: - cursor.execute(f"UPDATE {table} SET created_at = ? WHERE {name_col} = ?", (created_at, name)) + _update(services, table, name_col, name, created_at=created_at) def _star(services, table: str, name_col: str, name: str) -> None: - db = services["images"]._db - with db.transaction() as cursor: - cursor.execute(f"UPDATE {table} SET starred = 1 WHERE {name_col} = ?", (name,)) + _update(services, table, name_col, name, starred=True) + + +def _update(services, table: str, name_col: str, name: str, **values: object) -> None: + media = {"images": images, "videos": videos}[table] + with services["db"].begin(write=True) as conn: + conn.execute(update(media).where(media.c[name_col] == name).values(**values)) def _start_gallery_for_item_results(services) -> None: @@ -491,25 +490,11 @@ def test_returns_counts_and_deterministic_covers_in_one_result(self, services) - services["board_videos"].add_video_to_board(populated.board_id, "cover.mp4") services["board_videos"].add_video_to_board(populated.board_id, "intermediate.mp4") services["board_videos"].add_video_to_board(populated.board_id, "uploaded.mp4") - with services["images"]._db.transaction() as cursor: - cursor.execute( - "UPDATE images SET starred = 1, created_at = ? WHERE image_name = ?", - ("2026-01-05 12:00:00", "cover.png"), - ) - cursor.execute( - "UPDATE videos SET starred = 1, created_at = ? WHERE video_name = ?", - ("2026-01-05 12:00:00", "cover.mp4"), - ) - cursor.execute( - "UPDATE videos SET is_intermediate = 1 WHERE video_name = ?", - ("intermediate.mp4",), - ) + _update(services, "images", "image_name", "cover.png", starred=True, created_at="2026-01-05 12:00:00") + _update(services, "videos", "video_name", "cover.mp4", starred=True, created_at="2026-01-05 12:00:00") + _update(services, "videos", "video_name", "intermediate.mp4", is_intermediate=True) - summaries, _, details = _capture_plan( - services, - lambda: services["gallery"].get_board_media_summaries([populated.board_id, empty.board_id]), - "ROW_NUMBER() OVER", - ) + summaries = services["gallery"].get_board_media_summaries([populated.board_id, empty.board_id]) assert summaries[populated.board_id].image_count == 1 assert summaries[populated.board_id].video_count == 2 @@ -522,6 +507,19 @@ def test_returns_counts_and_deterministic_covers_in_one_result(self, services) - assert summaries[empty.board_id].asset_video_count == 0 assert summaries[empty.board_id].cover_image_name is None assert summaries[empty.board_id].cover_video_name is None + + @pytest.mark.sqlite_only + def test_summaries_start_from_the_boards_memberships(self, services) -> None: + board = services["boards"].save("Small", "alice") + _save_image(services["images"], "image.png", user_id="alice") + _save_video(services["videos"], "video.mp4", user_id="alice") + services["board_images"].add_image_to_board(board.board_id, "image.png") + services["board_videos"].add_video_to_board(board.board_id, "video.mp4") + + _, _, details = _capture_plan( + services, lambda: services["gallery"].get_board_media_summaries([board.board_id]), "row_number() OVER" + ) + assert next(i for i, detail in enumerate(details) if "board_images" in detail) < next( i for i, detail in enumerate(details) if "images" in detail and "board_images" not in detail ) @@ -530,6 +528,7 @@ def test_returns_counts_and_deterministic_covers_in_one_result(self, services) - ) +@pytest.mark.sqlite_only class TestGalleryQueryPlans: def test_default_name_list_omits_membership_join_and_query_hints(self, services) -> None: _save_image(services["images"], "image.png", user_id="alice") @@ -551,7 +550,7 @@ def test_default_name_list_omits_membership_join_and_query_hints(self, services) } assert result.total_count == 2 assert result.starred_count == 0 - assert "LEFT JOIN board_images" not in statement + assert "JOIN board_images" not in statement assert "INDEXED BY" not in statement assert "NOT INDEXED" not in statement @@ -692,7 +691,7 @@ def test_paginated_item_path_keeps_result_shape_and_avoids_gallery_index(self, s assert len(result.items) == 1 assert result.items[0].kind in {GalleryItemKind.IMAGE, GalleryItemKind.VIDEO} assert result.items[0].full_url - assert "LEFT JOIN board_images" in statement + assert "LEFT OUTER JOIN board_images" in statement assert "INDEXED BY" not in statement assert "NOT INDEXED" not in statement @@ -701,13 +700,13 @@ def test_board_image_count_starts_from_membership(self, services) -> None: _save_image(services["images"], "image.png", user_id="alice") services["board_images"].add_image_to_board(board.board_id, "image.png") - count, _, details = _capture_plan( + counts, _, details = _capture_plan( services, - lambda: services["board_images"].get_image_count_for_board(board.board_id), - "SELECT COUNT(*)", + lambda: services["board_images"].get_counts_for_board(board.board_id), + "FROM board_images CROSS JOIN images", ) - assert count == 1 + assert counts == (1, 0) assert next(i for i, detail in enumerate(details) if "board_images" in detail) < next( i for i, detail in enumerate(details) if "images" in detail and "board_images" not in detail ) @@ -751,7 +750,7 @@ def test_same_timestamp_order_is_deterministic(self, services) -> None: ] def test_ascending_is_mirror_of_descending(self, services) -> None: - from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection + from invokeai.app.services.shared.pagination import SQLiteDirection _seed_same_timestamp(services) gallery = services["gallery"] @@ -832,8 +831,7 @@ def test_location_uses_board_category_view_search_date_star_and_user_filters(sel ]: _backdate(services, table, name_col, name, created_at) _star(services, "videos", "video_name", "general.mp4") - with services["images"]._db.transaction() as cursor: - cursor.execute("UPDATE images SET is_intermediate = 1 WHERE image_name = ?", ("control.png",)) + _update(services, "images", "image_name", "control.png", is_intermediate=True) common = {"user_id": "alice", "is_admin": False} filter_sets = [ @@ -890,7 +888,7 @@ def test_missing_target_returns_none(self, services) -> None: ) -def _save_marked_video(store: SqliteVideoRecordStorage, name: str, user_id: str, metadata: str | None) -> None: +def _save_marked_video(store: VideoRecordStorage, name: str, user_id: str, metadata: str | None) -> None: store.save( video_name=name, video_origin=ResourceOrigin.EXTERNAL, @@ -958,3 +956,314 @@ def test_a_malformed_metadata_blob_does_not_fail_the_page(self, services) -> Non "bad.mp4": None, "wrapped.mp4": "audio_upload", } + + +class TestCanvasOwnedImagesAreNotCounted: + """OTHER is the category of images a canvas layer owns: neither the images view nor the assets view lists them, + so neither a day's nor a board's counts may include them, or a badge disagrees with what it opens.""" + + def test_a_day_counts_only_what_its_views_list(self, services) -> None: + for name, category in ( + ("generated.png", ImageCategory.GENERAL), + ("uploaded.png", ImageCategory.USER), + ("painted.png", ImageCategory.OTHER), + ): + _save_image(services["images"], name, user_id="alice", category=category) + _backdate(services, "images", "image_name", name, "2026-01-05 12:00:00.000") + + (day,) = services["gallery"].get_dates(user_id="alice", is_admin=False) + + assert (day.image_count, day.asset_count) == (1, 1) + + def test_a_board_counts_only_what_its_views_list(self, services) -> None: + board = services["boards"].save("Project", "alice") + for name, category in ( + ("generated.png", ImageCategory.GENERAL), + ("uploaded.png", ImageCategory.USER), + ("painted.png", ImageCategory.OTHER), + ): + _save_image(services["images"], name, user_id="alice", category=category) + services["board_images"].add_image_to_board(board.board_id, name) + + summary = services["gallery"].get_board_media_summaries([board.board_id])[board.board_id] + + assert (summary.image_count, summary.asset_count) == (1, 1) + # The same counts as the board's own badge. + assert services["board_images"].get_counts_for_board(board.board_id) == (1, 1) + + +class TestFilters: + def test_search_ignores_case_and_matches_percent_and_underscore_literally(self, services) -> None: + for name, prompt in (("hit.png", "a 100% Cat_Photo"), ("near.png", "a 100x catXphoto")): + services["images"].save( + image_name=name, + image_origin=ResourceOrigin.INTERNAL, + image_category=ImageCategory.GENERAL, + width=8, + height=8, + has_workflow=False, + metadata=f'{{"prompt": "{prompt}"}}', + user_id="alice", + ) + + result = services["gallery"].list_item_names(search_term="100% cat_", is_admin=True) + + assert [item.name for item in result.items] == ["hit.png"] + + def test_a_search_ignores_the_case_of_letters_beyond_ascii(self, services) -> None: + services["images"].save( + image_name="apples.png", + image_origin=ResourceOrigin.INTERNAL, + image_category=ImageCategory.GENERAL, + width=8, + height=8, + has_workflow=False, + metadata='{"prompt": "frische äpfel"}', + user_id="alice", + ) + + result = services["gallery"].list_item_names(search_term="ÄPFEL", is_admin=True) + + assert [item.name for item in result.items] == ["apples.png"] + + def test_the_categories_filter_applies_to_both_kinds_in_pages_and_names(self, services) -> None: + for name, category in (("general.png", ImageCategory.GENERAL), ("control.png", ImageCategory.CONTROL)): + _save_image(services["images"], name, user_id="alice", category=category) + for name, category in (("general.mp4", ImageCategory.GENERAL), ("user.mp4", ImageCategory.USER)): + _save_video(services["videos"], name, user_id="alice", category=category) + assets = [ImageCategory.CONTROL, ImageCategory.USER] + + names = services["gallery"].list_item_names(categories=assets, is_admin=True) + page = services["gallery"].list_items(limit=10, categories=assets, is_admin=True) + + assert {item.name for item in names.items} == {"control.png", "user.mp4"} + assert {item.name for item in page.items} == {"control.png", "user.mp4"} and page.total == 2 + + def test_the_origin_filter_applies_to_both_kinds(self, services) -> None: + _save_image(services["images"], "internal.png", user_id="alice") + _save_video(services["videos"], "internal.mp4", user_id="alice") + services["videos"].save( + video_name="external.mp4", + video_origin=ResourceOrigin.EXTERNAL, + video_category=ImageCategory.GENERAL, + width=8, + height=8, + duration=1.0, + fps=8.0, + has_workflow=False, + user_id="alice", + ) + + result = services["gallery"].list_items(limit=10, origin=ResourceOrigin.EXTERNAL, is_admin=True) + + assert [item.name for item in result.items] == ["external.mp4"] and result.total == 1 + + @pytest.mark.parametrize("day", ["2026-1-05", "not a day", "2026-01-05 00:00", "9999-12-31", ""]) + def test_a_day_that_is_no_iso_day_or_the_last_one_matches_nothing(self, services, day: str) -> None: + _save_image(services["images"], "image.png", user_id="alice") + _backdate(services, "images", "image_name", "image.png", "2026-01-05 12:00:00.000") + + assert services["gallery"].list_item_names(created_date=day, is_admin=True).items == [] + assert services["gallery"].list_items(limit=10, created_to=day, is_admin=True).total == 0 + + def test_pages_walk_the_listing_and_a_page_of_none_still_counts(self, services) -> None: + for index, name in enumerate(("a.png", "b.png")): + _save_image(services["images"], name, user_id="alice") + _backdate(services, "images", "image_name", name, f"2026-01-0{index + 1} 00:00:00.000") + _save_video(services["videos"], "c.mp4", user_id="alice") + _backdate(services, "videos", "video_name", "c.mp4", "2026-01-03 00:00:00.000") + names = [item.name for item in services["gallery"].list_item_names(is_admin=True).items] + + pages = [services["gallery"].list_items(offset=offset, limit=1, is_admin=True) for offset in range(3)] + count_only = services["gallery"].list_items(limit=0, is_admin=True) + + assert [[item.name for item in page.items] for page in pages] == [[name] for name in names] + assert {page.total for page in pages} == {3} + assert count_only.items == [] and count_only.total == 3 + + def test_a_board_lists_only_its_own_items(self, services) -> None: + mine, other = services["boards"].save("Mine", "alice"), services["boards"].save("Other", "alice") + for name, board in (("mine.png", mine), ("other.png", other)): + _save_image(services["images"], name, user_id="alice") + services["board_images"].add_image_to_board(board.board_id, name) + for name, board in (("mine.mp4", mine), ("other.mp4", other)): + _save_video(services["videos"], name, user_id="alice") + services["board_videos"].add_video_to_board(board.board_id, name) + + names = services["gallery"].list_item_names(board_id=mine.board_id, is_admin=True) + page = services["gallery"].list_items(limit=10, board_id=mine.board_id, is_admin=True) + + assert {item.name for item in names.items} == {"mine.png", "mine.mp4"} + assert {item.name for item in page.items} == {"mine.png", "mine.mp4"} and page.total == 2 + assert {item.board_id for item in page.items} == {mine.board_id} + + +class TestReviewedBehaviour: + """Pages, names, dates and summaries in the cases the review of the port found unpinned.""" + + def test_shared_board_lists_owners_items_for_viewer(self, services) -> None: + board = services["boards"].save("Bob's", "bob") + services["boards"].update(board.board_id, BoardChanges(board_visibility=BoardVisibility.Shared)) + _save_image(services["images"], "bob.png", user_id="bob") + _save_video(services["videos"], "bob.mp4", user_id="bob") + services["board_images"].add_image_to_board(board.board_id, "bob.png") + services["board_videos"].add_video_to_board(board.board_id, "bob.mp4") + _start_gallery_for_item_results(services) + g = services["gallery"] + + names = g.get_item_names(board_id=board.board_id, user_id="alice", is_admin=False) + page = g.list_items(limit=10, board_id=board.board_id, user_id="alice", is_admin=False) + + assert set(names.item_names) == {"bob.png", "bob.mp4"} + assert {i.name for i in page.items} == {"bob.png", "bob.mp4"} and page.total == 2 + + def test_intermediates_are_left_out_of_pages_names_and_totals(self, services) -> None: + for name in ("kept.png", "temp.png"): + _save_image(services["images"], name, user_id="alice") + for name in ("kept.mp4", "temp.mp4"): + _save_video(services["videos"], name, user_id="alice") + _update(services, "images", "image_name", "temp.png", is_intermediate=True) + _update(services, "videos", "video_name", "temp.mp4", is_intermediate=True) + _start_gallery_for_item_results(services) + g = services["gallery"] + + names = g.get_item_names(is_intermediate=False, user_id="alice", is_admin=False) + page = g.list_items(limit=10, is_intermediate=False, user_id="alice", is_admin=False) + only = g.get_item_names(is_intermediate=True, user_id="alice", is_admin=False) + + assert set(names.item_names) == {"kept.png", "kept.mp4"} + assert {i.name for i in page.items} == {"kept.png", "kept.mp4"} and page.total == 2 + assert set(only.item_names) == {"temp.png", "temp.mp4"} + + def test_pages_follow_direction_and_starred_first_like_names(self, services) -> None: + for name in ("a.png", "b.png"): + _save_image(services["images"], name, user_id="alice") + for name in ("a.mp4", "b.mp4"): + _save_video(services["videos"], name, user_id="alice") + _save_image(services["images"], "old-starred.png", user_id="alice") + for table, col, name in [ + ("images", "image_name", "a.png"), + ("images", "image_name", "b.png"), + ("videos", "video_name", "a.mp4"), + ("videos", "video_name", "b.mp4"), + ]: + _backdate(services, table, col, name, "2026-01-05 12:00:00.000") + _backdate(services, "images", "image_name", "old-starred.png", "2026-01-01 12:00:00.000") + _star(services, "images", "image_name", "old-starred.png") + _start_gallery_for_item_results(services) + g = services["gallery"] + + for order_dir in (SQLiteDirection.Ascending, SQLiteDirection.Descending): + for starred_first in (True, False): + kwargs: dict[str, Any] = { + "order_dir": order_dir, + "starred_first": starred_first, + "user_id": "alice", + "is_admin": False, + } + expected = g.get_item_names(**kwargs).item_names + page = [i.name for i in g.list_items(limit=10, **kwargs).items] + walked = [g.list_items(offset=o, limit=2, **kwargs).items for o in (0, 2, 4)] + assert page == expected, (order_dir, starred_first) + assert [i.name for p in walked for i in p] == expected, (order_dir, starred_first) + assert ( + g.get_item_names(order_dir=SQLiteDirection.Ascending, starred_first=False, user_id="alice").item_names[0] + == "old-starred.png" + ) + + def test_a_days_cover_is_the_viewers_own(self, services) -> None: + _save_image(services["images"], "alice.png", user_id="alice") + _save_image(services["images"], "bob.png", user_id="bob") + _backdate(services, "images", "image_name", "alice.png", "2026-01-05 09:00:00") + _backdate(services, "images", "image_name", "bob.png", "2026-01-05 10:00:00") + + (day,) = services["gallery"].get_dates(user_id="alice", is_admin=False) + + assert day.cover_image_name == "alice.png" + + def test_intermediates_neither_make_nor_count_nor_cover_a_day(self, services) -> None: + _save_image(services["images"], "kept.png", user_id="alice") + _save_video(services["videos"], "temp.mp4", user_id="alice") + _save_image(services["images"], "lone-temp.png", user_id="alice") + _backdate(services, "images", "image_name", "kept.png", "2026-01-05 09:00:00") + _backdate(services, "videos", "video_name", "temp.mp4", "2026-01-05 10:00:00") + _backdate(services, "images", "image_name", "lone-temp.png", "2026-01-06 10:00:00") + _update(services, "videos", "video_name", "temp.mp4", is_intermediate=True) + _update(services, "images", "image_name", "lone-temp.png", is_intermediate=True) + + days = services["gallery"].get_dates(user_id="alice", is_admin=False) + + assert [(d.date, d.image_count, d.video_count, d.cover_image_name) for d in days] == [ + ("2026-01-05", 1, 0, "kept.png") + ] + + def test_unboarded_pages_and_totals_are_the_viewers_own(self, services) -> None: + _save_image(services["images"], "alice.png", user_id="alice") + _save_video(services["videos"], "bob.mp4", user_id="bob") + _start_gallery_for_item_results(services) + + page = services["gallery"].list_items(limit=10, board_id="none", user_id="alice", is_admin=False) + + assert [i.name for i in page.items] == ["alice.png"] and page.total == 1 + + def test_page_items_carry_their_columns(self, services) -> None: + board = services["boards"].save("B", "alice") + services["videos"].save( + video_name="clip.mp4", + video_origin=ResourceOrigin.INTERNAL, + video_category=ImageCategory.GENERAL, + width=640, + height=360, + duration=2.5, + fps=24.0, + has_workflow=False, + is_intermediate=False, + user_id="alice", + ) + services["board_videos"].add_video_to_board(board.board_id, "clip.mp4") + _start_gallery_for_item_results(services) + + (item,) = services["gallery"].list_items(limit=10, user_id="alice", is_admin=False).items + + assert (item.width, item.height, item.duration, item.fps, item.board_id) == ( + 640, + 360, + 2.5, + 24.0, + board.board_id, + ) + + def test_board_summaries_span_chunks(self, services, monkeypatch) -> None: + monkeypatch.setattr(gallery_queries, "IN_CHUNK", 2) + boards = [services["boards"].save(f"B{i}", "alice") for i in range(3)] + for i, board in enumerate(boards): + _save_image(services["images"], f"{i}.png", user_id="alice") + services["board_images"].add_image_to_board(board.board_id, f"{i}.png") + + summaries = services["gallery"].get_board_media_summaries([b.board_id for b in boards]) + + assert [summaries[b.board_id].cover_image_name for b in boards] == ["0.png", "1.png", "2.png"] + + def test_board_cover_is_the_newest_without_stars(self, services) -> None: + board = services["boards"].save("B", "alice") + for name in ("z-old.png", "a-new.png"): + _save_image(services["images"], name, user_id="alice") + services["board_images"].add_image_to_board(board.board_id, name) + _backdate(services, "images", "image_name", "z-old.png", "2026-01-01 00:00:00") + _backdate(services, "images", "image_name", "a-new.png", "2026-01-02 00:00:00") + + summary = services["gallery"].get_board_media_summaries([board.board_id])[board.board_id] + + assert summary.cover_image_name == "a-new.png" + + def test_canvas_owned_image_neither_makes_nor_covers_a_day(self, services) -> None: + _save_image(services["images"], "generated.png", user_id="alice") + _save_image(services["images"], "painted.png", user_id="alice", category=ImageCategory.OTHER) + _save_image(services["images"], "lone-painted.png", user_id="alice", category=ImageCategory.OTHER) + _backdate(services, "images", "image_name", "generated.png", "2026-01-05 09:00:00") + _backdate(services, "images", "image_name", "painted.png", "2026-01-05 10:00:00") + _backdate(services, "images", "image_name", "lone-painted.png", "2026-01-06 10:00:00") + + days = services["gallery"].get_dates(user_id="alice", is_admin=False) + + assert [(d.date, d.cover_image_name) for d in days] == [("2026-01-05", "generated.png")] diff --git a/tests/app/services/gallery_maintenance/test_gallery_maintenance_service.py b/tests/app/services/gallery_maintenance/test_gallery_maintenance_service.py index d33be313a3f..eccd9c44322 100644 --- a/tests/app/services/gallery_maintenance/test_gallery_maintenance_service.py +++ b/tests/app/services/gallery_maintenance/test_gallery_maintenance_service.py @@ -8,10 +8,11 @@ import numpy as np import pytest from PIL import Image +from sqlalchemy import ColumnElement, Table, func, select -from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage +from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage from invokeai.app.services.board_records.board_records_common import BoardChanges -from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.gallery_maintenance.gallery_maintenance_common import ( GalleryMaintenanceError, @@ -21,16 +22,18 @@ from invokeai.app.services.gallery_maintenance.gallery_maintenance_default import GalleryMaintenanceService from invokeai.app.services.image_files.image_files_disk import DiskImageFileStorage from invokeai.app.services.image_index.image_index_common import IndexedItem -from invokeai.app.services.image_index.image_index_records_sqlite import ImageIndexRecordsSqlite +from invokeai.app.services.image_index.image_index_records_default import ImageIndexRecords from invokeai.app.services.image_records.image_records_common import ( ImageCategory, ImageRecordNotFoundException, ResourceOrigin, ) -from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage -from invokeai.app.services.video_records.video_records_sqlite import SqliteVideoRecordStorage -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.boards import board_images +from invokeai.app.services.shared.database.schema.image_index import image_embeddings, video_embeddings +from invokeai.app.services.shared.database.schema.image_moves import image_subfolder_move_items +from invokeai.app.services.video_records.video_records_default import VideoRecordStorage class _ImageMutationGate: @@ -41,13 +44,12 @@ def reserve_gallery_maintenance(self): @pytest.fixture def maintenance( - tmp_path: Path, -) -> tuple[GalleryMaintenanceService, object, SqliteImageRecordStorage, DiskImageFileStorage]: + tmp_path: Path, database: Database +) -> tuple[GalleryMaintenanceService, Database, ImageRecordStorage, DiskImageFileStorage]: config = InvokeAIAppConfig(use_memory_db=True) config._root = tmp_path - logger = InvokeAILogger.get_logger(config=config) - db = create_mock_sqlite_database(config, logger) - records = SqliteImageRecordStorage(db) + db = database + records = ImageRecordStorage(db) files = DiskImageFileStorage(config.outputs_path / "images") invoker = MagicMock() invoker.services.configuration = config @@ -55,10 +57,10 @@ def maintenance( invoker.services.database = db invoker.services.image_records = records invoker.services.image_files = files - invoker.services.image_index_records = ImageIndexRecordsSqlite(db) - invoker.services.video_records = SqliteVideoRecordStorage(db) - invoker.services.board_records = SqliteBoardRecordStorage(db) - invoker.services.board_image_records = SqliteBoardImageRecordStorage(db) + invoker.services.image_index_records = ImageIndexRecords(db) + invoker.services.video_records = VideoRecordStorage(db) + invoker.services.board_records = BoardRecordStorage(db) + invoker.services.board_image_records = BoardImageRecordStorage(db) invoker.services.images = MagicMock() invoker.services.image_moves = _ImageMutationGate() files.start(invoker) @@ -67,8 +69,16 @@ def maintenance( return service, db, records, files +def _count(database: Database, table: Table, where: ColumnElement[bool] | None = None) -> int: + statement = select(func.count()).select_from(table) + if where is not None: + statement = statement.where(where) + with database.begin(write=False) as conn: + return conn.execute(statement).scalar_one() + + def _save_image_record( - records: SqliteImageRecordStorage, + records: ImageRecordStorage, image_name: str, image_subfolder: str = "", *, @@ -159,19 +169,11 @@ def test_remove_missing_archives_thumbnail_and_cascades_all_image_relations(main assert not thumbnail_path.exists() archived_thumbnail = Path(result.archive_path) / "thumbnails" / "archived" / "board" / "missing.webp" assert archived_thumbnail.read_bytes() == b"recoverable thumbnail" - with db.transaction() as cursor: - assert cursor.execute("SELECT COUNT(*) FROM board_images WHERE image_name='missing.png';").fetchone()[0] == 0 - assert ( - cursor.execute("SELECT COUNT(*) FROM image_embeddings WHERE image_name='missing.png';").fetchone()[0] == 0 - ) - assert ( - cursor.execute( - "SELECT COUNT(*) FROM image_subfolder_move_items WHERE image_name='missing.png';" - ).fetchone()[0] - == 0 - ) - assert cursor.execute("SELECT COUNT(*) FROM image_embeddings WHERE image_name='kept.png';").fetchone()[0] == 1 - assert cursor.execute("SELECT COUNT(*) FROM video_embeddings WHERE video_name='kept.mp4';").fetchone()[0] == 1 + assert _count(db, board_images, board_images.c.image_name == "missing.png") == 0 + assert _count(db, image_embeddings, image_embeddings.c.image_name == "missing.png") == 0 + assert _count(db, image_subfolder_move_items, image_subfolder_move_items.c.image_name == "missing.png") == 0 + assert _count(db, image_embeddings, image_embeddings.c.image_name == "kept.png") == 1 + assert _count(db, video_embeddings, video_embeddings.c.video_name == "kept.mp4") == 1 assert IndexedItem( kind="image", name="missing.png" ) not in invoker.services.image_index_records.list_accessible_embedded_items(None, "model-a") @@ -179,7 +181,11 @@ def test_remove_missing_archives_thumbnail_and_cascades_all_image_relations(main assert _embedding_bytes(invoker.services, "video", "kept.mp4") == kept_video_embedding assert _embedding_bytes(invoker.services, "image", "archived-board-kept.png") == archived_board_embedding invoker.services.images.notify_deleted.assert_called_once_with("missing.png") - assert Path(result.backup_path).is_file() + if db.dialect_name == "sqlite": + assert result.backup_path is not None and Path(result.backup_path).is_file() + else: + # A server database is backed up by its operator, not by gallery maintenance. + assert result.backup_path is None def test_archive_untracked_preserves_nested_paths_and_never_archives_tracked_or_video_files(maintenance) -> None: @@ -219,7 +225,7 @@ def test_archive_untracked_preserves_nested_paths_and_never_archives_tracked_or_ assert _embedding_bytes(service._invoker.services, "image", "tracked.png") == embedding_before -def test_regenerate_missing_thumbnail_is_record_scoped_and_refreshes_cached_size(maintenance) -> None: +def test_regenerate_missing_thumbnail_is_record_scoped_and_refreshes_cached_size(maintenance, monkeypatch) -> None: service, db, records, files = maintenance _save_image_record(records, "missing-thumb.png", "nested") original_bytes = _write_png(files, "missing-thumb.png", "nested") @@ -238,15 +244,16 @@ def test_regenerate_missing_thumbnail_is_record_scoped_and_refreshes_cached_size assert preview.affected_count == 2 assert preview.skipped_count == 1 - original_backup = service._invoker.services.database.backup_to + original_backup = service._create_backup - def create_thumbnail_during_backup(destination: Path) -> None: - original_backup(destination) + def create_thumbnail_during_backup() -> str | None: + backup_path = original_backup() raced_thumbnail = files.get_path("raced-thumb.png", thumbnail=True, image_subfolder="nested") raced_thumbnail.parent.mkdir(parents=True, exist_ok=True) raced_thumbnail.write_bytes(b"created after preview scan") + return backup_path - service._invoker.services.database.backup_to = create_thumbnail_during_backup + monkeypatch.setattr(service, "_create_backup", create_thumbnail_during_backup) result = service.execute(GalleryMaintenanceOperation.REGENERATE_THUMBNAILS, preview.fingerprint) @@ -264,8 +271,7 @@ def create_thumbnail_during_backup(destination: Path) -> None: "missing-thumb.png", image_subfolder="nested" ) assert _embedding_bytes(service._invoker.services, "image", "missing-thumb.png") == embedding_before - with db.transaction() as cursor: - assert cursor.execute("SELECT COUNT(*) FROM image_embeddings;").fetchone()[0] == 1 + assert _count(db, image_embeddings) == 1 def test_execute_rejects_changed_preview_without_mutating_newly_untracked_file(maintenance) -> None: @@ -279,14 +285,15 @@ def test_execute_rejects_changed_preview_without_mutating_newly_untracked_file(m assert files.get_path("arrived-after-preview.png").read_bytes() == added_bytes -def test_backup_failure_prevents_record_removal_and_thumbnail_archival(maintenance) -> None: +@pytest.mark.sqlite_only # Only a SQLite database is backed up by gallery maintenance. +def test_backup_failure_prevents_record_removal_and_thumbnail_archival(maintenance, monkeypatch) -> None: service, _db, records, files = maintenance _save_image_record(records, "missing.png", "recoverable") thumbnail = files.get_path("missing.png", thumbnail=True, image_subfolder="recoverable") thumbnail.parent.mkdir(parents=True) thumbnail.write_bytes(b"recoverable thumbnail") preview = service.preview(GalleryMaintenanceOperation.REMOVE_MISSING) - service._invoker.services.database.backup_to = MagicMock(side_effect=OSError("disk full")) + monkeypatch.setattr(service._invoker.services.database, "backup", MagicMock(side_effect=OSError("disk full"))) with pytest.raises(GalleryMaintenanceError, match="database backup failed"): service.execute(GalleryMaintenanceOperation.REMOVE_MISSING, preview.fingerprint) @@ -388,21 +395,22 @@ def test_remove_missing_preserves_record_when_same_basename_image_exists_elsewhe assert files.get_path("duplicate.png", image_subfolder="actual").read_bytes() == elsewhere -def test_archive_untracked_executes_only_confirmed_files_with_unchanged_stats(maintenance) -> None: +def test_archive_untracked_executes_only_confirmed_files_with_unchanged_stats(maintenance, monkeypatch) -> None: service, _db, _records, files = maintenance confirmed = _write_png(files, "confirmed.png") changed_path = files.get_path("changed.png") changed_path.parent.mkdir(parents=True, exist_ok=True) changed_path.write_bytes(b"before preview") preview = service.preview(GalleryMaintenanceOperation.ARCHIVE_UNTRACKED) - original_backup = service._invoker.services.database.backup_to + original_backup = service._create_backup - def change_storage(destination: Path) -> None: - original_backup(destination) + def change_storage() -> str | None: + backup_path = original_backup() changed_path.write_bytes(b"changed after preview") _write_png(files, "arrived.png") + return backup_path - service._invoker.services.database.backup_to = change_storage + monkeypatch.setattr(service, "_create_backup", change_storage) result = service.execute(GalleryMaintenanceOperation.ARCHIVE_UNTRACKED, preview.fingerprint) @@ -428,17 +436,18 @@ def test_remove_missing_counts_committed_deletion_when_followup_check_fails(main assert result.status == "partial" -def test_remove_missing_keeps_record_when_original_appears_during_backup(maintenance) -> None: +def test_remove_missing_keeps_record_when_original_appears_during_backup(maintenance, monkeypatch) -> None: service, _db, records, files = maintenance _save_image_record(records, "appeared.png", "during-backup") preview = service.preview(GalleryMaintenanceOperation.REMOVE_MISSING) - original_backup = service._invoker.services.database.backup_to + original_backup = service._create_backup - def create_original(destination: Path) -> None: - original_backup(destination) + def create_original() -> str | None: + backup_path = original_backup() _write_png(files, "appeared.png", "during-backup") + return backup_path - service._invoker.services.database.backup_to = create_original + monkeypatch.setattr(service, "_create_backup", create_original) result = service.execute(GalleryMaintenanceOperation.REMOVE_MISSING, preview.fingerprint) @@ -774,6 +783,7 @@ def track_source_quarantine(path: Path, target: Path) -> Path: assert parent_syncs and parent_syncs[0] > index +@pytest.mark.sqlite_only # Only a SQLite database is backed up by gallery maintenance. def test_backup_directory_sync_failure_prevents_gallery_deletion(maintenance, monkeypatch) -> None: service, _db, records, files = maintenance _save_image_record(records, "missing.png", "recoverable") @@ -786,7 +796,7 @@ def test_backup_directory_sync_failure_prevents_gallery_deletion(maintenance, mo preview = service.preview(GalleryMaintenanceOperation.REMOVE_MISSING) backup_directory = service._invoker.services.configuration.db_path.parent / "backup" database_directory = backup_directory.parent - backup_to = service._invoker.services.database.backup_to + backup = service._invoker.services.database.backup directory_fsync = GalleryMaintenanceService._GalleryMaintenanceService__fsync_directory real_mkdir = os.mkdir events: list[tuple[str, Path]] = [] @@ -805,7 +815,7 @@ def track_mkdir(path, *args, **kwargs): def create_backup(destination: Path) -> None: nonlocal completed_backup - backup_to(destination) + backup(destination) completed_backup = destination events.append(("backup", destination)) @@ -818,7 +828,7 @@ def fail_backup_directory_sync(directory: Path) -> None: raise OSError("injected backup directory-entry sync failure") directory_fsync(directory) - service._invoker.services.database.backup_to = create_backup + monkeypatch.setattr(service._invoker.services.database, "backup", create_backup) monkeypatch.setattr(os, "mkdir", track_mkdir) monkeypatch.setattr( GalleryMaintenanceService, diff --git a/tests/app/services/image_index/test_cluster_labels.py b/tests/app/services/image_index/test_cluster_labels.py index 4e78e8bed81..256b2af92b5 100644 --- a/tests/app/services/image_index/test_cluster_labels.py +++ b/tests/app/services/image_index/test_cluster_labels.py @@ -27,6 +27,16 @@ def test_normalize_custom_vocab_terms_lowercases_collapses_and_dedupes() -> None assert normalized == ["golden retriever", "zebra"] +def test_normalize_custom_vocab_terms_folds_the_spellings_server_collations_equate() -> None: + # MySQL and MariaDB store the vocabulary under a case-insensitive collation that equates these spellings, so + # without the fold a server would silently keep only the first of each pair where SQLite keeps both. + decomposed = "café" + assert normalize_custom_vocab_terms(["café", decomposed, "fish", "ｆｉｓｈ"]) == [ + "café", + "fish", + ] + + def test_normalize_custom_vocab_terms_strips_control_characters() -> None: # A NUL inside a term would make the fingerprint's NUL-joined phrase # stream ambiguous: vocab_fingerprint(["cat", "dog"]) equals diff --git a/tests/app/services/image_index/test_image_index_records.py b/tests/app/services/image_index/test_image_index_records.py index 11cd4a33b18..9a0302e3d23 100644 --- a/tests/app/services/image_index/test_image_index_records.py +++ b/tests/app/services/image_index/test_image_index_records.py @@ -1,15 +1,13 @@ """Tests for the image index records service: embedding storage, eligibility, access scoping, projections.""" -import sqlite3 - import numpy as np import pytest +from sqlalchemy import func, insert, select, update -from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage +from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardVisibility -from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage -from invokeai.app.services.board_video_records.board_video_records_sqlite import SqliteBoardVideoRecordStorage -from invokeai.app.services.config.config_default import InvokeAIAppConfig +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage +from invokeai.app.services.board_video_records.board_video_records_default import BoardVideoRecordStorage from invokeai.app.services.image_index import image_index_common from invokeai.app.services.image_index.image_index_common import ( IndexedItem, @@ -19,18 +17,23 @@ coords_to_blob, embedding_to_blob, ) -from invokeai.app.services.image_index.image_index_records_sqlite import ( - _IN_CLAUSE_CHUNK, - ImageIndexRecordsSqlite, -) +from invokeai.app.services.image_index.image_index_records_default import ImageIndexRecords from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin -from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries.locks import DatabaseLock +from invokeai.app.services.shared.database.schema.boards import shared_boards +from invokeai.app.services.shared.database.schema.image_index import ( + image_embeddings, + image_projections, + video_embeddings, +) +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.videos import videos from invokeai.app.services.users.users_common import UserCreateRequest from invokeai.app.services.users.users_default import UserService -from invokeai.app.services.video_records.video_records_sqlite import SqliteVideoRecordStorage -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database +from invokeai.app.services.video_records.video_records_default import VideoRecordStorage +from tests.fixtures.races import while_in_flight SYSTEM_USER_ID = "system" MODEL_ID = "model-hash-1" @@ -39,44 +42,38 @@ @pytest.fixture -def db() -> SqliteDatabase: - config = InvokeAIAppConfig(use_memory_db=True) - return create_mock_sqlite_database(config=config, logger=InvokeAILogger.get_logger()) +def image_records(database: Database) -> ImageRecordStorage: + return ImageRecordStorage(database) @pytest.fixture -def image_records(db: SqliteDatabase) -> SqliteImageRecordStorage: - return SqliteImageRecordStorage(db=db) +def board_records(database: Database) -> BoardRecordStorage: + return BoardRecordStorage(database) @pytest.fixture -def board_records(db: SqliteDatabase) -> SqliteBoardRecordStorage: - return SqliteBoardRecordStorage(db=db) +def board_image_records(database: Database) -> BoardImageRecordStorage: + return BoardImageRecordStorage(database) @pytest.fixture -def board_image_records(db: SqliteDatabase) -> SqliteBoardImageRecordStorage: - return SqliteBoardImageRecordStorage(db=db) +def video_records(database: Database) -> VideoRecordStorage: + return VideoRecordStorage(database) @pytest.fixture -def video_records(db: SqliteDatabase) -> SqliteVideoRecordStorage: - return SqliteVideoRecordStorage(db=db) +def board_video_records(database: Database) -> BoardVideoRecordStorage: + return BoardVideoRecordStorage(database) @pytest.fixture -def board_video_records(db: SqliteDatabase) -> SqliteBoardVideoRecordStorage: - return SqliteBoardVideoRecordStorage(db=db) +def index_records(database: Database) -> ImageIndexRecords: + return ImageIndexRecords(database) @pytest.fixture -def index_records(db: SqliteDatabase) -> ImageIndexRecordsSqlite: - return ImageIndexRecordsSqlite(db=db) - - -@pytest.fixture -def other_user_id(db: SqliteDatabase) -> str: - users = UserService(db=db) +def other_user_id(database: Database) -> str: + users = UserService(database) user = users.create( UserCreateRequest(email="other@example.com", display_name="Other", password="TestPass123", is_admin=False) ) @@ -84,7 +81,7 @@ def other_user_id(db: SqliteDatabase) -> str: def _save_image( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, image_name: str, user_id: str = SYSTEM_USER_ID, is_intermediate: bool = False, @@ -103,7 +100,7 @@ def _save_image( def _save_video( - video_records: SqliteVideoRecordStorage, + video_records: VideoRecordStorage, video_name: str, user_id: str = SYSTEM_USER_ID, is_intermediate: bool = False, @@ -123,6 +120,18 @@ def _save_video( ) +def _set_created_at(database: Database, kind: str, name: str, created_at: str) -> None: + media, column = (images, "image_name") if kind == "image" else (videos, "video_name") + with database.begin(write=True) as conn: + conn.execute(update(media).where(media.c[column] == name).values(created_at=created_at)) + + +def _share_board(database: Database, board_id: str, user_id: str) -> None: + """No service writes shared_boards yet.""" + with database.begin(write=True) as conn: + conn.execute(insert(shared_boards).values(board_id=board_id, user_id=user_id, can_edit=False)) + + def imgs(*names: str) -> list[IndexedItem]: """The image-namespace items for these names.""" return [IndexedItem("image", name) for name in names] @@ -325,8 +334,18 @@ def test_coords_from_blob_are_writable() -> None: # --- Embedding CRUD --- +def _stored(database: Database, item: IndexedItem) -> tuple[str, int]: + """The encoding and byte length of the item's stored embedding.""" + table, name = (image_embeddings, "image_name") if item.kind == "image" else (video_embeddings, "video_name") + with database.begin(write=False) as conn: + row = conn.execute( + select(table.c.encoding, func.length(table.c.embedding)).where(table.c[name] == item.name) + ).one() + return str(row[0]), int(row[1]) + + def test_upsert_and_get_roundtrip( - db: SqliteDatabase, image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite + image_records: ImageRecordStorage, index_records: ImageIndexRecords, database: Database ) -> None: _save_image(image_records, "a.png") _save_image(image_records, "b.png") @@ -334,12 +353,7 @@ def test_upsert_and_get_roundtrip( index_records.upsert_embedding(IndexedItem("image", "a.png"), MODEL_ID, va) index_records.upsert_embedding(IndexedItem("image", "b.png"), MODEL_ID, vb) - with db.transaction() as cursor: - cursor.execute( - "SELECT encoding, LENGTH(embedding) FROM image_embeddings WHERE image_name = ? AND model_id = ?;", - ("a.png", MODEL_ID), - ) - assert tuple(cursor.fetchone()) == ("float16", DIM * 2) + assert _stored(database, IndexedItem("image", "a.png")) == ("float16", DIM * 2) names, matrix = index_records.get_embeddings(imgs("b.png", "a.png", "missing.png"), MODEL_ID) @@ -350,7 +364,7 @@ def test_upsert_and_get_roundtrip( def test_upsert_replaces_existing_embedding( - image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite + image_records: ImageRecordStorage, index_records: ImageIndexRecords ) -> None: _save_image(image_records, "a.png") index_records.upsert_embedding(IndexedItem("image", "a.png"), MODEL_ID, _vec(1)) @@ -363,7 +377,7 @@ def test_upsert_replaces_existing_embedding( def test_get_embeddings_empty_input_and_no_matches( - image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite + image_records: ImageRecordStorage, index_records: ImageIndexRecords ) -> None: names, matrix = index_records.get_embeddings([], MODEL_ID) assert names == [] @@ -376,21 +390,25 @@ def test_get_embeddings_empty_input_and_no_matches( assert matrix.shape == (0, 0) -def test_get_embeddings_chunks_large_requests( - image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite +def test_get_embeddings_returns_a_large_request_in_the_callers_order( + image_records: ImageRecordStorage, video_records: VideoRecordStorage, index_records: ImageIndexRecords ) -> None: - # Straddle the IN-clause chunk boundary so the multi-chunk path really runs. Derived from - # the constant rather than hardcoded, so raising the chunk size cannot silently reduce this - # to a single-chunk test. - count = _IN_CLAUSE_CHUNK + 1 + # More names than an IN list would bind, of both kinds, requested interleaved and against the order the + # database reads them in: the rows must come back aligned with the request. + count = 600 + requested = [] for i in range(count): - _save_image(image_records, f"img-{i:04d}.png") - index_records.upsert_embedding(IndexedItem("image", f"img-{i:04d}.png"), MODEL_ID, _vec(i)) + if i % 2: + _save_video(video_records, f"item-{i:04d}.mp4") + requested.append(IndexedItem("video", f"item-{i:04d}.mp4")) + else: + _save_image(image_records, f"item-{i:04d}.png") + requested.append(IndexedItem("image", f"item-{i:04d}.png")) + index_records.upsert_embedding(requested[-1], MODEL_ID, _vec(i)) + requested.reverse() - # Reverse order: SQLite would return each chunk in its own order, so this fails unless the - # implementation reorders results back to the caller's sequence within every chunk. - requested = [IndexedItem("image", f"img-{i:04d}.png") for i in reversed(range(count))] names, matrix = index_records.get_embeddings(requested, MODEL_ID) + assert names == requested assert matrix.shape == (count, DIM) # Rows must still align with the names after chunking and reordering. @@ -399,7 +417,7 @@ def test_get_embeddings_chunks_large_requests( def test_get_embeddings_rejects_inconsistent_dims( - image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite, db: SqliteDatabase + image_records: ImageRecordStorage, index_records: ImageIndexRecords, database: Database ) -> None: # The ABC promises a failure on mixed dims under one model_id. Write a short vector behind # the service's back, since upsert_embedding alone cannot produce the inconsistency. @@ -407,10 +425,11 @@ def test_get_embeddings_rejects_inconsistent_dims( _save_image(image_records, "b.png") index_records.upsert_embedding(IndexedItem("image", "a.png"), MODEL_ID, _vec(1)) short = np.ones(DIM // 2, dtype=np.float32) - with db.transaction() as cursor: - cursor.execute( - "INSERT INTO image_embeddings (image_name, model_id, dim, embedding) VALUES (?, ?, ?, ?);", - ("b.png", MODEL_ID, DIM // 2, embedding_to_blob(short)), + with database.begin(write=True) as conn: + conn.execute( + insert(image_embeddings).values( + image_name="b.png", model_id=MODEL_ID, dim=DIM // 2, embedding=embedding_to_blob(short) + ) ) with pytest.raises(ValueError, match="Inconsistent embedding dims"): @@ -418,7 +437,7 @@ def test_get_embeddings_rejects_inconsistent_dims( def test_get_embeddings_deduplicates_input_names( - image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite + image_records: ImageRecordStorage, index_records: ImageIndexRecords ) -> None: _save_image(image_records, "a.png") index_records.upsert_embedding(IndexedItem("image", "a.png"), MODEL_ID, _vec(1)) @@ -429,50 +448,68 @@ def test_get_embeddings_deduplicates_input_names( assert matrix.shape == (1, DIM) -def test_upsert_embedding_for_deleted_image_is_noop(index_records: ImageIndexRecordsSqlite) -> None: +def test_upsert_embedding_for_deleted_image_is_noop(index_records: ImageIndexRecords) -> None: # The image was deleted (or never existed) by the time the write lands; - # the FK violation must not escape as a raw sqlite3 error. + # the foreign key must not fail the write. index_records.upsert_embedding(IndexedItem("image", "gone.png"), MODEL_ID, _vec(1)) assert index_records.get_embeddings(imgs("gone.png"), MODEL_ID)[0] == [] -def test_set_projection_for_deleted_user_is_noop(index_records: ImageIndexRecordsSqlite) -> None: +def test_set_projection_for_deleted_user_is_noop(index_records: ImageIndexRecords) -> None: index_records.set_projection( "no-such-user", MODEL_ID, "hash-1", "{}", imgs("a.png"), np.zeros((1, 2), dtype=np.float32) ) assert index_records.get_projection("no-such-user", MODEL_ID) is None -def test_skipped_write_raises_nothing_and_leaves_the_connection_usable( - image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite, db: SqliteDatabase +def test_a_skipped_write_leaves_the_next_one_working( + image_records: ImageRecordStorage, index_records: ImageIndexRecords ) -> None: - # The missing-parent case is handled by `WHERE EXISTS`, not by provoking and swallowing an - # IntegrityError. That matters because `db.transaction()` commits/rolls back the whole - # shared connection: a swallowed error would leave the connection mid-statement-failure, and - # an escaping one would roll back. Neither may happen, and the next write must still work. + # The missing parent is checked before the write, not provoked as a foreign-key error and swallowed. index_records.upsert_embedding(IndexedItem("image", "gone.png"), MODEL_ID, _vec(1)) index_records.set_projection( "no-such-user", MODEL_ID, "hash-1", "{}", imgs("a.png"), np.zeros((1, 2), dtype=np.float32) ) - assert db._conn.in_transaction is False - _save_image(image_records, "real.png") index_records.upsert_embedding(IndexedItem("image", "real.png"), MODEL_ID, _vec(2)) assert index_records.get_embeddings(imgs("real.png"), MODEL_ID)[0] == imgs("real.png") -def test_skipped_write_does_not_raise_integrity_error(index_records: ImageIndexRecordsSqlite) -> None: - # Pin the mechanism, not just the outcome: if the guard regressed to relying on the foreign - # key, this would raise instead of no-opping. - try: - index_records.upsert_embedding(IndexedItem("image", "gone.png"), MODEL_ID, _vec(1)) - except sqlite3.IntegrityError as exc: # pragma: no cover - the assertion is the point - pytest.fail(f"missing-parent write must be a no-op, not an IntegrityError: {exc}") +def test_an_image_deleted_while_its_embedding_is_written_gets_none( + database: Database, image_records: ImageRecordStorage, index_records: ImageIndexRecords +) -> None: + # On a server the image's deletion is in flight when the embedding is written: the write waits for it and then + # finds no image, rather than failing on the foreign key. + _save_image(image_records, "a.png") + + errors = while_in_flight( + database, + lambda q: q.images.delete("a.png"), + lambda: index_records.upsert_embedding(IndexedItem("image", "a.png"), MODEL_ID, _vec(1)), + ) + + assert errors == [] + assert index_records.get_embeddings(imgs("a.png"), MODEL_ID)[0] == [] + + +def test_an_account_deleted_while_its_projection_is_written_gets_none( + database: Database, index_records: ImageIndexRecords, other_user_id: str +) -> None: + errors = while_in_flight( + database, + lambda q: q.users.delete(other_user_id), + lambda: index_records.set_projection( + other_user_id, MODEL_ID, "hash-1", "{}", imgs("a.png"), np.zeros((1, 2), dtype=np.float32) + ), + ) + + assert errors == [] + assert index_records.get_projection(other_user_id, MODEL_ID) is None def test_delete_embedding_removes_all_models( - image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite + image_records: ImageRecordStorage, index_records: ImageIndexRecords ) -> None: _save_image(image_records, "a.png") index_records.upsert_embedding(IndexedItem("image", "a.png"), MODEL_ID, _vec(1)) @@ -485,7 +522,7 @@ def test_delete_embedding_removes_all_models( def test_image_delete_cascades_to_embeddings( - db: SqliteDatabase, image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite + database: Database, image_records: ImageRecordStorage, index_records: ImageIndexRecords ) -> None: _save_image(image_records, "a.png") index_records.upsert_embedding(IndexedItem("image", "a.png"), MODEL_ID, _vec(1)) @@ -496,7 +533,7 @@ def test_image_delete_cascades_to_embeddings( def test_delete_embeddings_for_other_models( - image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite + image_records: ImageRecordStorage, index_records: ImageIndexRecords ) -> None: _save_image(image_records, "a.png") index_records.upsert_embedding(IndexedItem("image", "a.png"), MODEL_ID, _vec(1)) @@ -513,7 +550,7 @@ def test_delete_embeddings_for_other_models( def test_list_unembedded_skips_ineligible_and_embedded( - image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite + image_records: ImageRecordStorage, index_records: ImageIndexRecords ) -> None: _save_image(image_records, "eligible.png") _save_image(image_records, "embedded.png") @@ -527,7 +564,7 @@ def test_list_unembedded_skips_ineligible_and_embedded( def test_list_unembedded_respects_limit_and_returns_oldest_first( - image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite, db: SqliteDatabase + image_records: ImageRecordStorage, index_records: ImageIndexRecords, database: Database ) -> None: # `created_at` has millisecond resolution, so five back-to-back saves almost always tie and # the `image_name ASC` tie-break alone decides the result — which would let a reversed @@ -535,12 +572,7 @@ def test_list_unembedded_respects_limit_and_returns_oldest_first( # alphabetical order so the two orderings disagree and only `created_at` can satisfy this. for i in range(5): _save_image(image_records, f"img-{i}.png") - with db.transaction() as cursor: - for i in range(5): - cursor.execute( - "UPDATE images SET created_at = ? WHERE image_name = ?;", - (f"2026-01-0{5 - i} 00:00:00.000", f"img-{i}.png"), - ) + _set_created_at(database, "image", f"img-{i}.png", f"2026-01-0{5 - i} 00:00:00.000") batch = index_records.list_unembedded_items(MODEL_ID, limit=3) @@ -549,7 +581,24 @@ def test_list_unembedded_respects_limit_and_returns_oldest_first( assert batch == imgs("img-4.png", "img-3.png", "img-2.png") -def test_list_unembedded_rejects_negative_limit(index_records: ImageIndexRecordsSqlite) -> None: +def test_an_embedding_under_another_model_leaves_an_item_to_embed( + image_records: ImageRecordStorage, video_records: VideoRecordStorage, index_records: ImageIndexRecords +) -> None: + # After a model switch every item is work again, until the old model's rows are deleted. + _save_image(image_records, "a.png") + _save_video(video_records, "clip.mp4") + index_records.upsert_embedding(IndexedItem("image", "a.png"), OTHER_MODEL_ID, _vec(1)) + index_records.upsert_embedding(IndexedItem("video", "clip.mp4"), OTHER_MODEL_ID, _vec(2)) + + assert index_records.list_unembedded_items(MODEL_ID, limit=10) == [ + IndexedItem("image", "a.png"), + IndexedItem("video", "clip.mp4"), + ] + status = index_records.count_index_status(MODEL_ID) + assert (status.total, status.embedded) == (2, 0) + + +def test_list_unembedded_rejects_negative_limit(index_records: ImageIndexRecords) -> None: # SQLite reads a negative LIMIT as unbounded, which would turn a bounded backfill batch into # a full-table load. with pytest.raises(ValueError, match="non-negative"): @@ -557,13 +606,13 @@ def test_list_unembedded_rejects_negative_limit(index_records: ImageIndexRecords def test_list_unembedded_zero_limit_returns_nothing( - image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite + image_records: ImageRecordStorage, index_records: ImageIndexRecords ) -> None: _save_image(image_records, "a.png") assert index_records.list_unembedded_items(MODEL_ID, limit=0) == [] -def test_count_index_status(image_records: SqliteImageRecordStorage, index_records: ImageIndexRecordsSqlite) -> None: +def test_count_index_status(image_records: ImageRecordStorage, index_records: ImageIndexRecords) -> None: _save_image(image_records, "a.png") _save_image(image_records, "b.png") _save_image(image_records, "intermediate.png", is_intermediate=True) @@ -580,10 +629,10 @@ def test_count_index_status(image_records: SqliteImageRecordStorage, index_recor def test_accessible_images_scoping( - image_records: SqliteImageRecordStorage, - board_records: SqliteBoardRecordStorage, - board_image_records: SqliteBoardImageRecordStorage, - index_records: ImageIndexRecordsSqlite, + image_records: ImageRecordStorage, + board_records: BoardRecordStorage, + board_image_records: BoardImageRecordStorage, + index_records: ImageIndexRecords, other_user_id: str, ) -> None: # System user's images: one unboarded, one on a private board, one shared, one public. @@ -629,11 +678,11 @@ def test_accessible_images_scoping( def test_accessible_images_includes_individually_shared_boards( - db: SqliteDatabase, - image_records: SqliteImageRecordStorage, - board_records: SqliteBoardRecordStorage, - board_image_records: SqliteBoardImageRecordStorage, - index_records: ImageIndexRecordsSqlite, + database: Database, + image_records: ImageRecordStorage, + board_records: BoardRecordStorage, + board_image_records: BoardImageRecordStorage, + index_records: ImageIndexRecords, other_user_id: str, ) -> None: # A private board individually shared with other_user via shared_boards @@ -643,15 +692,11 @@ def test_accessible_images_includes_individually_shared_boards( index_records.upsert_embedding(IndexedItem("image", "own-individually-shared.png"), MODEL_ID, _vec(1)) board_id = board_records.save("Individually shared", SYSTEM_USER_ID).board_id board_image_records.add_image_to_board(board_id, "own-individually-shared.png") - with db.transaction() as cursor: - cursor.execute( - "INSERT INTO shared_boards (board_id, user_id, can_edit) VALUES (?, ?, ?);", - (board_id, other_user_id, False), - ) + _share_board(database, board_id, other_user_id) assert index_records.list_accessible_embedded_items(other_user_id, MODEL_ID) == imgs("own-individually-shared.png") # A third party without the share still cannot see it. - users = UserService(db=db) + users = UserService(database) third_user = users.create( UserCreateRequest(email="third@example.com", display_name="Third", password="TestPass123", is_admin=False) ) @@ -659,15 +704,15 @@ def test_accessible_images_includes_individually_shared_boards( def test_accessible_images_includes_boards_owned_by_user( - image_records: SqliteImageRecordStorage, - board_records: SqliteBoardRecordStorage, - board_image_records: SqliteBoardImageRecordStorage, - index_records: ImageIndexRecordsSqlite, + image_records: ImageRecordStorage, + board_records: BoardRecordStorage, + board_image_records: BoardImageRecordStorage, + index_records: ImageIndexRecords, other_user_id: str, ) -> None: # An image uploaded by another user onto a board the system user OWNS is # accessible to the board owner — matching the gallery "all" listing - # (image_records_sqlite), which grants access via boards.user_id even + # (image_records_default), which grants access via boards.user_id even # without shared/public visibility or a shared_boards row. _save_image(image_records, "theirs-on-my-board.png", user_id=other_user_id) index_records.upsert_embedding(IndexedItem("image", "theirs-on-my-board.png"), MODEL_ID, _vec(1)) @@ -681,10 +726,10 @@ def test_accessible_images_includes_boards_owned_by_user( def test_accessible_images_excludes_archived_boards( - image_records: SqliteImageRecordStorage, - board_records: SqliteBoardRecordStorage, - board_image_records: SqliteBoardImageRecordStorage, - index_records: ImageIndexRecordsSqlite, + image_records: ImageRecordStorage, + board_records: BoardRecordStorage, + board_image_records: BoardImageRecordStorage, + index_records: ImageIndexRecords, other_user_id: str, ) -> None: # Archived boards hide their images from every scope, mirroring the @@ -712,10 +757,10 @@ def test_accessible_images_excludes_archived_boards( def test_accessible_images_returns_boarded_images_once( - image_records: SqliteImageRecordStorage, - board_records: SqliteBoardRecordStorage, - board_image_records: SqliteBoardImageRecordStorage, - index_records: ImageIndexRecordsSqlite, + image_records: ImageRecordStorage, + board_records: BoardRecordStorage, + board_image_records: BoardImageRecordStorage, + index_records: ImageIndexRecords, ) -> None: # `board_images` has PRIMARY KEY (image_name), so an image is on at most one board and the # board join cannot multiply rows. This pins that the join stays single-valued: if @@ -729,11 +774,33 @@ def test_accessible_images_returns_boarded_images_once( assert index_records.list_accessible_embedded_items(SYSTEM_USER_ID, MODEL_ID) == imgs("a.png") +def test_an_embedded_item_that_left_the_gallery_is_not_accessible( + database: Database, + image_records: ImageRecordStorage, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, +) -> None: + # Nothing deletes an embedding when its item becomes an intermediate or leaves the general category, so the + # listing itself must keep such items off the map and out of search, in every scope. + for name in ("kept.png", "now-intermediate.png", "now-mask.png"): + _save_image(image_records, name) + index_records.upsert_embedding(IndexedItem("image", name), MODEL_ID, _vec(len(name))) + _save_video(video_records, "now-intermediate.mp4") + index_records.upsert_embedding(IndexedItem("video", "now-intermediate.mp4"), MODEL_ID, _vec(1)) + with database.begin(write=True) as conn: + conn.execute(update(images).where(images.c.image_name == "now-intermediate.png").values(is_intermediate=True)) + conn.execute(update(images).where(images.c.image_name == "now-mask.png").values(image_category="mask")) + conn.execute(update(videos).where(videos.c.video_name == "now-intermediate.mp4").values(is_intermediate=True)) + + for user_id in (SYSTEM_USER_ID, None): + assert index_records.list_accessible_embedded_items(user_id, MODEL_ID) == imgs("kept.png") + + def test_accessible_images_are_filtered_by_model_id( - image_records: SqliteImageRecordStorage, - board_records: SqliteBoardRecordStorage, - board_image_records: SqliteBoardImageRecordStorage, - index_records: ImageIndexRecordsSqlite, + image_records: ImageRecordStorage, + board_records: BoardRecordStorage, + board_image_records: BoardImageRecordStorage, + index_records: ImageIndexRecords, other_user_id: str, ) -> None: # The scope hash is derived from this listing while the embedding matrix is fetched by @@ -757,11 +824,11 @@ def test_accessible_images_are_filtered_by_model_id( def test_accessible_images_exclude_individually_shared_archived_board( - image_records: SqliteImageRecordStorage, - board_records: SqliteBoardRecordStorage, - board_image_records: SqliteBoardImageRecordStorage, - index_records: ImageIndexRecordsSqlite, - db: SqliteDatabase, + image_records: ImageRecordStorage, + board_records: BoardRecordStorage, + board_image_records: BoardImageRecordStorage, + index_records: ImageIndexRecords, + database: Database, other_user_id: str, ) -> None: # Archiving is tested for visibility-shared boards; the shared_boards grant is a separate @@ -770,11 +837,7 @@ def test_accessible_images_exclude_individually_shared_archived_board( index_records.upsert_embedding(IndexedItem("image", "granted.png"), MODEL_ID, _vec(1)) board_id = board_records.save("Private", SYSTEM_USER_ID).board_id board_image_records.add_image_to_board(board_id, "granted.png") - with db.transaction() as cursor: - cursor.execute( - "INSERT INTO shared_boards (board_id, user_id) VALUES (?, ?);", - (board_id, other_user_id), - ) + _share_board(database, board_id, other_user_id) assert index_records.list_accessible_embedded_items(other_user_id, MODEL_ID) == imgs("granted.png") @@ -786,7 +849,7 @@ def test_accessible_images_exclude_individually_shared_archived_board( # --- Projections --- -def test_projection_roundtrip(index_records: ImageIndexRecordsSqlite) -> None: +def test_projection_roundtrip(index_records: ImageIndexRecords) -> None: coords = np.array([[0.5, -1.5], [2.0, 3.0]], dtype=np.float32) index_records.set_projection( SYSTEM_USER_ID, MODEL_ID, "hash-1", '{"n_neighbors": 15}', imgs("a.png", "b.png"), coords @@ -802,22 +865,45 @@ def test_projection_roundtrip(index_records: ImageIndexRecordsSqlite) -> None: assert np.array_equal(record.coords, coords) -def test_projection_upsert_replaces(index_records: ImageIndexRecordsSqlite) -> None: +def test_projection_upsert_replaces(index_records: ImageIndexRecords) -> None: index_records.set_projection( SYSTEM_USER_ID, MODEL_ID, "hash-1", "{}", imgs("a.png"), np.zeros((1, 2), dtype=np.float32) ) index_records.set_projection( - SYSTEM_USER_ID, MODEL_ID, "hash-2", "{}", imgs("a.png", "b.png"), np.ones((2, 2), dtype=np.float32) + SYSTEM_USER_ID, MODEL_ID, "hash-2", '{"n": 2}', imgs("a.png", "b.png"), np.ones((2, 2), dtype=np.float32) ) record = index_records.get_projection(SYSTEM_USER_ID, MODEL_ID) assert record is not None assert record.scope_hash == "hash-2" + assert record.params == '{"n": 2}' assert record.point_count == 2 + assert record.items == imgs("a.png", "b.png") + assert np.array_equal(record.coords, np.ones((2, 2), dtype=np.float32)) + + +def test_replacing_a_projection_stamps_when_and_keeps_since_when( + database: Database, index_records: ImageIndexRecords +) -> None: + index_records.set_projection( + SYSTEM_USER_ID, MODEL_ID, "hash-1", "{}", imgs("a.png"), np.zeros((1, 2), dtype=np.float32) + ) + old = "2020-01-01 00:00:00.000" + with database.begin(write=True) as conn: + conn.execute(update(image_projections).values(created_at=old, updated_at=old)) + index_records.set_projection( + SYSTEM_USER_ID, MODEL_ID, "hash-2", "{}", imgs("a.png"), np.zeros((1, 2), dtype=np.float32) + ) -def test_projection_missing_and_delete_idempotent(index_records: ImageIndexRecordsSqlite) -> None: + record = index_records.get_projection(SYSTEM_USER_ID, MODEL_ID) + assert record is not None + assert record.created_at == old + assert record.updated_at > old + + +def test_projection_missing_and_delete_idempotent(index_records: ImageIndexRecords) -> None: assert index_records.get_projection(SYSTEM_USER_ID, MODEL_ID) is None index_records.delete_projection(SYSTEM_USER_ID, MODEL_ID) # no-op @@ -828,14 +914,14 @@ def test_projection_missing_and_delete_idempotent(index_records: ImageIndexRecor assert index_records.get_projection(SYSTEM_USER_ID, MODEL_ID) is None -def test_projection_rejects_mismatched_lengths(index_records: ImageIndexRecordsSqlite) -> None: +def test_projection_rejects_mismatched_lengths(index_records: ImageIndexRecords) -> None: with pytest.raises(ValueError): index_records.set_projection( SYSTEM_USER_ID, MODEL_ID, "hash-1", "{}", imgs("a.png"), np.zeros((2, 2), dtype=np.float32) ) -def test_custom_vocab_terms_roundtrip_sorted_with_replace_semantics(index_records: ImageIndexRecordsSqlite) -> None: +def test_custom_vocab_terms_roundtrip_sorted_with_replace_semantics(index_records: ImageIndexRecords) -> None: assert index_records.get_custom_vocab_terms() == [] index_records.set_custom_vocab_terms(["zebra", "aardvark"]) @@ -849,20 +935,37 @@ def test_custom_vocab_terms_roundtrip_sorted_with_replace_semantics(index_record assert index_records.get_custom_vocab_terms() == [] -def test_custom_vocab_case_variant_duplicates_cannot_fail_the_replace(index_records: ImageIndexRecordsSqlite) -> None: +def test_custom_vocab_case_variant_duplicates_cannot_fail_the_replace(index_records: ImageIndexRecords) -> None: # Callers normalize, but the NOCASE primary key is a backstop: a # case-variant pair must collapse to one row, not roll back the replace. index_records.set_custom_vocab_terms(["Zebra", "zebra"]) - assert len(index_records.get_custom_vocab_terms()) == 1 + assert index_records.get_custom_vocab_terms() == ["Zebra"] + + +def test_a_replace_of_the_vocabulary_waits_for_one_in_flight( + database: Database, index_records: ImageIndexRecords +) -> None: + # On a server two replaces side by side would each delete none of the other's uncommitted terms and keep both + # lists; the replace in flight here holds the vocabulary's lock, which the other must wait for. + index_records.set_custom_vocab_terms(["alpha"]) + + errors = while_in_flight( + database, + lambda q: q.locks.acquire(DatabaseLock.IMAGE_INDEX_VOCABULARY), + lambda: index_records.set_custom_vocab_terms(["beta"]), + ) + + assert errors == [] + assert index_records.get_custom_vocab_terms() == ["beta"] # --- Videos --- def test_videos_are_stored_counted_and_listed_alongside_images( - image_records: SqliteImageRecordStorage, - video_records: SqliteVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + image_records: ImageRecordStorage, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, ) -> None: _save_image(image_records, "a.png") _save_video(video_records, "clip.mp4") @@ -882,24 +985,27 @@ def test_videos_are_stored_counted_and_listed_alongside_images( def test_records_read_mixed_legacy_and_fp16_encodings_across_media_kinds( - db: SqliteDatabase, - image_records: SqliteImageRecordStorage, - video_records: SqliteVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + database: Database, + image_records: ImageRecordStorage, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, ) -> None: _save_image(image_records, "legacy.png") _save_video(video_records, "half.mp4") legacy = _vec(1) half = _vec(2) - with db.transaction() as cursor: - cursor.execute( - "INSERT INTO image_embeddings (image_name, model_id, dim, embedding) VALUES (?, ?, ?, ?);", - ("legacy.png", MODEL_ID, DIM, legacy.tobytes()), + with database.begin(write=True) as conn: + # A row from before fp16 storage: the column's default names its encoding. + conn.execute( + insert(image_embeddings).values( + image_name="legacy.png", model_id=MODEL_ID, dim=DIM, embedding=legacy.tobytes() + ) ) - cursor.execute( - "INSERT INTO video_embeddings (video_name, model_id, dim, embedding, encoding) VALUES (?, ?, ?, ?, ?);", - ("half.mp4", MODEL_ID, DIM, embedding_to_blob(half), "float16"), + conn.execute( + insert(video_embeddings).values( + video_name="half.mp4", model_id=MODEL_ID, dim=DIM, embedding=embedding_to_blob(half), encoding="float16" + ) ) items, matrix = index_records.get_embeddings( @@ -943,27 +1049,24 @@ def test_embedding_matrix_decode_keeps_rows_aligned_across_encodings_and_chunks( def test_image_and_video_upserts_set_fp16_encoding( - db: SqliteDatabase, - image_records: SqliteImageRecordStorage, - video_records: SqliteVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + database: Database, + image_records: ImageRecordStorage, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, ) -> None: _save_image(image_records, "a.png") _save_video(video_records, "clip.mp4") index_records.upsert_embedding(IndexedItem("image", "a.png"), MODEL_ID, _vec(1)) index_records.upsert_embedding(IndexedItem("video", "clip.mp4"), MODEL_ID, _vec(2)) - with db.transaction() as cursor: - cursor.execute("SELECT encoding, LENGTH(embedding) FROM image_embeddings WHERE image_name = 'a.png';") - assert tuple(cursor.fetchone()) == ("float16", DIM * 2) - cursor.execute("SELECT encoding, LENGTH(embedding) FROM video_embeddings WHERE video_name = 'clip.mp4';") - assert tuple(cursor.fetchone()) == ("float16", DIM * 2) + assert _stored(database, IndexedItem("image", "a.png")) == ("float16", DIM * 2) + assert _stored(database, IndexedItem("video", "clip.mp4")) == ("float16", DIM * 2) def test_an_image_and_a_video_are_different_items( - image_records: SqliteImageRecordStorage, - video_records: SqliteVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + image_records: ImageRecordStorage, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, ) -> None: # Names are generated with kind-specific extensions, so this cannot happen by accident — # but the namespaces must not be able to read or overwrite each other's rows. @@ -991,7 +1094,7 @@ def test_an_image_and_a_video_are_different_items( def test_ineligible_videos_are_not_indexable( - video_records: SqliteVideoRecordStorage, index_records: ImageIndexRecordsSqlite + video_records: VideoRecordStorage, index_records: ImageIndexRecords ) -> None: _save_video(video_records, "eligible.mp4") _save_video(video_records, "intermediate.mp4", is_intermediate=True) @@ -1002,10 +1105,10 @@ def test_ineligible_videos_are_not_indexable( def test_unembedded_listing_merges_both_kinds_oldest_first( - db: SqliteDatabase, - image_records: SqliteImageRecordStorage, - video_records: SqliteVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + database: Database, + image_records: ImageRecordStorage, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, ) -> None: # Interleaved creation times: listing one namespace and then the other would put every # image before any video, so a video generated hours ago would wait behind newer images. @@ -1015,14 +1118,13 @@ def test_unembedded_listing_merges_both_kinds_oldest_first( _save_video(video_records, "second.mp4") _save_image(image_records, "third.png") _save_video(video_records, "fourth.mp4") - with db.transaction() as cursor: - for table, column, name, created_at in ( - ("images", "image_name", "first.png", "2026-01-01 00:00:01.000"), - ("videos", "video_name", "second.mp4", "2026-01-01 00:00:02.000"), - ("images", "image_name", "third.png", "2026-01-01 00:00:03.000"), - ("videos", "video_name", "fourth.mp4", "2026-01-01 00:00:04.000"), - ): - cursor.execute(f"UPDATE {table} SET created_at = ? WHERE {column} = ?;", (created_at, name)) + for kind, name, created_at in ( + ("image", "first.png", "2026-01-01 00:00:01.000"), + ("video", "second.mp4", "2026-01-01 00:00:02.000"), + ("image", "third.png", "2026-01-01 00:00:03.000"), + ("video", "fourth.mp4", "2026-01-01 00:00:04.000"), + ): + _set_created_at(database, kind, name, created_at) batch = index_records.list_unembedded_items(MODEL_ID, limit=3) @@ -1034,7 +1136,7 @@ def test_unembedded_listing_merges_both_kinds_oldest_first( def test_deleting_a_video_deletes_its_embedding( - video_records: SqliteVideoRecordStorage, index_records: ImageIndexRecordsSqlite + video_records: VideoRecordStorage, index_records: ImageIndexRecords ) -> None: _save_video(video_records, "clip.mp4") index_records.upsert_embedding(IndexedItem("video", "clip.mp4"), MODEL_ID, _vec(1)) @@ -1044,7 +1146,7 @@ def test_deleting_a_video_deletes_its_embedding( assert index_records.get_embeddings(vids("clip.mp4"), MODEL_ID)[0] == [] -def test_embedding_a_missing_video_is_a_noop(index_records: ImageIndexRecordsSqlite) -> None: +def test_embedding_a_missing_video_is_a_noop(index_records: ImageIndexRecords) -> None: # The video may be deleted between being scheduled and embedded; the write is dropped # rather than raising a foreign-key error at the worker. index_records.upsert_embedding(IndexedItem("video", "gone.mp4"), MODEL_ID, _vec(1)) @@ -1052,10 +1154,10 @@ def test_embedding_a_missing_video_is_a_noop(index_records: ImageIndexRecordsSql def test_video_access_scoping_follows_board_membership( - video_records: SqliteVideoRecordStorage, - board_records: SqliteBoardRecordStorage, - board_video_records: SqliteBoardVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + video_records: VideoRecordStorage, + board_records: BoardRecordStorage, + board_video_records: BoardVideoRecordStorage, + index_records: ImageIndexRecords, other_user_id: str, ) -> None: _save_video(video_records, "own-unboarded.mp4") @@ -1078,10 +1180,10 @@ def test_video_access_scoping_follows_board_membership( def test_videos_on_archived_boards_are_hidden_from_every_scope( - video_records: SqliteVideoRecordStorage, - board_records: SqliteBoardRecordStorage, - board_video_records: SqliteBoardVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + video_records: VideoRecordStorage, + board_records: BoardRecordStorage, + board_video_records: BoardVideoRecordStorage, + index_records: ImageIndexRecords, ) -> None: _save_video(video_records, "archived.mp4") index_records.upsert_embedding(IndexedItem("video", "archived.mp4"), MODEL_ID, _vec(1)) @@ -1094,9 +1196,9 @@ def test_videos_on_archived_boards_are_hidden_from_every_scope( def test_accessible_listing_orders_images_before_videos( - image_records: SqliteImageRecordStorage, - video_records: SqliteVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + image_records: ImageRecordStorage, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, ) -> None: # The listing feeds the projection scope hash and the search matrix, both of which need a # stable order across calls. @@ -1111,7 +1213,7 @@ def test_accessible_listing_orders_images_before_videos( ] -def test_projection_roundtrips_mixed_items(index_records: ImageIndexRecordsSqlite) -> None: +def test_projection_roundtrips_mixed_items(index_records: ImageIndexRecords) -> None: coords = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) items = [IndexedItem("image", "a.png"), IndexedItem("video", "clip.mp4")] @@ -1122,19 +1224,21 @@ def test_projection_roundtrips_mixed_items(index_records: ImageIndexRecordsSqlit assert record.items == items -def test_projection_without_stored_kinds_reads_as_images( - db: SqliteDatabase, index_records: ImageIndexRecordsSqlite -) -> None: +def test_projection_without_stored_kinds_reads_as_images(database: Database, index_records: ImageIndexRecords) -> None: # Rows written before videos were indexable have no kinds column. They are all-image # projections, and must keep serving rather than being discarded on upgrade. coords = coords_to_blob(np.array([[1.0, 2.0]], dtype=np.float32)) - with db.transaction() as cursor: - cursor.execute( - """--sql - INSERT INTO image_projections (user_id, model_id, scope_hash, params, point_count, image_names, coords) - VALUES (?, ?, ?, ?, ?, ?, ?); - """, - (SYSTEM_USER_ID, MODEL_ID, "hash-1", "{}", 1, '["legacy.png"]', coords), + with database.begin(write=True) as conn: + conn.execute( + insert(image_projections).values( + user_id=SYSTEM_USER_ID, + model_id=MODEL_ID, + scope_hash="hash-1", + params="{}", + point_count=1, + image_names='["legacy.png"]', + coords=coords, + ) ) record = index_records.get_projection(SYSTEM_USER_ID, MODEL_ID) @@ -1144,9 +1248,9 @@ def test_projection_without_stored_kinds_reads_as_images( def test_delete_embeddings_for_other_models_covers_both_kinds( - image_records: SqliteImageRecordStorage, - video_records: SqliteVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + image_records: ImageRecordStorage, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, ) -> None: _save_image(image_records, "a.png") _save_video(video_records, "clip.mp4") diff --git a/tests/app/services/image_index/test_image_index_service.py b/tests/app/services/image_index/test_image_index_service.py index a47efd83c90..36a3d201baa 100644 --- a/tests/app/services/image_index/test_image_index_service.py +++ b/tests/app/services/image_index/test_image_index_service.py @@ -25,15 +25,16 @@ _POLL_SECONDS, ImageIndexService, ) -from invokeai.app.services.image_index.image_index_records_sqlite import ImageIndexRecordsSqlite +from invokeai.app.services.image_index.image_index_records_default import ImageIndexRecords from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin -from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage from invokeai.app.services.images.images_common import image_record_to_dto from invokeai.app.services.images.images_default import ImageService from invokeai.app.services.model_load.model_load_default import ModelLoadService from invokeai.app.services.model_records.model_records_sql import ModelRecordServiceSQL -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.video_records.video_records_sqlite import SqliteVideoRecordStorage +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import LockTimeoutError +from invokeai.app.services.video_records.video_records_default import VideoRecordStorage from invokeai.app.services.videos.videos_common import VideoDTO, video_record_to_dto from invokeai.app.services.videos.videos_default import VideoService from invokeai.backend.model_manager.configs.clip_vision import CLIPVision_Diffusers_Config @@ -94,19 +95,19 @@ def _fake_encode(images: list[Image.Image]) -> np.ndarray: @pytest.fixture -def db() -> SqliteDatabase: +def db() -> Database: config = InvokeAIAppConfig(use_memory_db=True) return create_mock_sqlite_database(config=config, logger=InvokeAILogger.get_logger()) @pytest.fixture -def image_records(db: SqliteDatabase) -> SqliteImageRecordStorage: - return SqliteImageRecordStorage(db=db) +def image_records(db: Database) -> ImageRecordStorage: + return ImageRecordStorage(db) @pytest.fixture -def index_records(db: SqliteDatabase) -> ImageIndexRecordsSqlite: - return ImageIndexRecordsSqlite(db=db) +def index_records(db: Database) -> ImageIndexRecords: + return ImageIndexRecords(db) @pytest.fixture(params=[CLIPVision_Diffusers_Config, SigLIP_Diffusers_Config], ids=["clip", "siglip"]) @@ -133,8 +134,8 @@ def images_service() -> ImageService: @pytest.fixture -def video_records(db: SqliteDatabase) -> SqliteVideoRecordStorage: - return SqliteVideoRecordStorage(db=db) +def video_records(db: Database) -> VideoRecordStorage: + return VideoRecordStorage(db) @pytest.fixture @@ -157,14 +158,14 @@ def get_path(video_name: str, thumbnail: bool = False) -> str: def _make_invoker( images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, enabled: bool = True, - image_records: SqliteImageRecordStorage | None = None, + image_records: ImageRecordStorage | None = None, device: str | None = "cpu", session_queue: object | None = None, model_manager: object | None = None, videos_service: VideoService | None = None, - video_records: SqliteVideoRecordStorage | None = None, + video_records: VideoRecordStorage | None = None, ) -> SimpleNamespace: config = InvokeAIAppConfig( use_memory_db=True, @@ -207,7 +208,7 @@ def accelerator_host(monkeypatch: pytest.MonkeyPatch) -> None: def _save_image( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, image_name: str, is_intermediate: bool = False, image_category: ImageCategory = ImageCategory.GENERAL, @@ -225,7 +226,7 @@ def _save_image( def _save_video( - video_records: SqliteVideoRecordStorage, + video_records: VideoRecordStorage, video_name: str, is_intermediate: bool = False, video_category: ImageCategory = ImageCategory.GENERAL, @@ -244,12 +245,12 @@ def _save_video( ) -def _video_dto_for(video_records: SqliteVideoRecordStorage, video_name: str) -> VideoDTO: +def _video_dto_for(video_records: VideoRecordStorage, video_name: str) -> VideoDTO: record = video_records.get(video_name) return video_record_to_dto(record, video_url="http://x/v.mp4", thumbnail_url="http://x/v.webp", board_id=None) -def _dto_for(image_records: SqliteImageRecordStorage, image_name: str): +def _dto_for(image_records: ImageRecordStorage, image_name: str): record = image_records.get(image_name) return image_record_to_dto(record, image_url="http://x/i.png", thumbnail_url="http://x/t.png", board_id=None) @@ -262,7 +263,7 @@ def test_constructor_requires_matched_test_seams() -> None: def test_disabled_service_is_inert( - images_service: ImageService, index_records: ImageIndexRecordsSqlite, service: ImageIndexService + images_service: ImageService, index_records: ImageIndexRecords, service: ImageIndexService ) -> None: invoker = _make_invoker(images_service, index_records, enabled=False) service.start(invoker) @@ -274,9 +275,9 @@ def test_disabled_service_is_inert( def test_backfill_indexes_preexisting_eligible_images( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: for i in range(6): @@ -298,9 +299,9 @@ def test_backfill_indexes_preexisting_eligible_images( def test_on_changed_indexes_new_eligible_image_and_skips_ineligible( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: invoker = _make_invoker(images_service, index_records) @@ -318,9 +319,9 @@ def test_on_changed_indexes_new_eligible_image_and_skips_ineligible( def test_unloadable_image_is_skipped_and_backfill_completes( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: _save_image(image_records, "good.png") @@ -341,9 +342,9 @@ def get_pil_image(image_name: str) -> Image.Image: def test_backfill_logs_what_it_indexed_and_that_it_finished( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, caplog: pytest.LogCaptureFixture, ) -> None: @@ -367,9 +368,9 @@ def test_backfill_logs_what_it_indexed_and_that_it_finished( def test_restart_over_an_indexed_gallery_logs_nothing( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, caplog: pytest.LogCaptureFixture, ) -> None: @@ -388,9 +389,9 @@ def test_restart_over_an_indexed_gallery_logs_nothing( def test_backfill_reports_images_it_could_not_embed( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, caplog: pytest.LogCaptureFixture, ) -> None: @@ -421,9 +422,9 @@ def get_pil_image(image_name: str) -> Image.Image: def test_transient_encode_failure_is_retried_to_success( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: calls = {"n": 0} @@ -489,7 +490,7 @@ def search_by_attr(model_name=None, model_type=None): def test_try_activate_picks_up_a_model_installed_after_startup( encoder_config: CLIPVision_Diffusers_Config | SigLIP_Diffusers_Config, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, monkeypatch: pytest.MonkeyPatch, ) -> None: # The encoder is usually installed from the image map itself, long after @@ -546,7 +547,7 @@ def resolve(self, model_name: str): def test_failed_late_activation_rolls_back_rather_than_wedging_the_service( encoder_config: CLIPVision_Diffusers_Config | SigLIP_Diffusers_Config, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, monkeypatch: pytest.MonkeyPatch, ) -> None: # Publishing the model before the worker exists would leave the service @@ -596,11 +597,11 @@ def launch(self, invoker): @pytest.mark.parametrize("normalize_loaded_path", [False, True], ids=["catalog-path", "normalized-loaded-path"]) def test_encoder_metadata_edits_keep_worker_indexing_without_map_requests( normalize_loaded_path: bool, - db: SqliteDatabase, + db: Database, encoder_config: CLIPVision_Diffusers_Config | SigLIP_Diffusers_Config, - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, monkeypatch: pytest.MonkeyPatch, ) -> None: models_dir = Path(encoder_config.path).parent @@ -659,9 +660,9 @@ def encode(self, images): def test_deleted_encoder_becomes_unavailable_and_reinstall_resumes_without_duplicate_callbacks( encoder_config: CLIPVision_Diffusers_Config | SigLIP_Diffusers_Config, - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, monkeypatch: pytest.MonkeyPatch, ) -> None: installed = [encoder_config] @@ -719,9 +720,9 @@ def test_encoder_replacement_waits_for_inflight_embedding_and_discards_retired_b replacement_kind: str, background_batch: bool, encoder_config: CLIPVision_Diffusers_Config | SigLIP_Diffusers_Config, - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, monkeypatch: pytest.MonkeyPatch, ) -> None: installed = [encoder_config] @@ -821,9 +822,9 @@ def _write_clip_vision_weights(path: Path, seed: int) -> None: def test_reinstalled_encoder_keeping_its_key_embeds_with_the_new_weights( tmp_path: Path, - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, monkeypatch: pytest.MonkeyPatch, ) -> None: # The shared model cache is keyed by model key alone. Reinstalling the encoder in place changes @@ -891,9 +892,9 @@ def test_reinstalled_encoder_keeping_its_key_embeds_with_the_new_weights( def test_worker_detected_replacement_resumes_indexing_without_requests( encoder_config: CLIPVision_Diffusers_Config | SigLIP_Diffusers_Config, - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, monkeypatch: pytest.MonkeyPatch, ) -> None: # Deleting and re-adding the encoder gives it a new key. With the image map @@ -930,7 +931,7 @@ def test_worker_detected_replacement_resumes_indexing_without_requests( def test_stop_releases_a_worker_waiting_for_request_users_to_drain( encoder_config: CLIPVision_Diffusers_Config | SigLIP_Diffusers_Config, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, monkeypatch: pytest.MonkeyPatch, ) -> None: installed = [encoder_config] @@ -956,9 +957,9 @@ def test_stop_releases_a_worker_waiting_for_request_users_to_drain( def test_projection_finishing_after_encoder_removal_does_not_publish( encoder_config: CLIPVision_Diffusers_Config | SigLIP_Diffusers_Config, - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, monkeypatch: pytest.MonkeyPatch, ) -> None: installed = [encoder_config] @@ -995,7 +996,7 @@ def project(matrix): def test_late_activation_survives_a_model_store_failure( images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, monkeypatch: pytest.MonkeyPatch, ) -> None: # These endpoints reported `model_missing` before they resolved anything; @@ -1018,9 +1019,9 @@ def explode(self, model_name): def test_broken_encoder_leaves_images_pending_rather_than_quarantined( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: service = ImageIndexService(encode_fn=lambda images: np.zeros((1,), dtype=np.float32), model_id=MODEL_ID) try: @@ -1040,9 +1041,9 @@ def test_broken_encoder_leaves_images_pending_rather_than_quarantined( def test_status_event_reports_failures_so_a_settled_index_is_not_mistaken_for_complete( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """`pending` excludes failures, so on its own it cannot express "gave up on some". @@ -1074,9 +1075,9 @@ def get_pil_image(image_name: str) -> Image.Image: def test_index_recovers_from_an_encoder_outage_without_a_restart( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """The point of separating systemic from per-image failure. @@ -1115,9 +1116,9 @@ def encode(images: list[Image.Image]) -> np.ndarray: def test_sustained_storage_failure_does_not_quarantine_images( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """A database that is down is no more the images' fault than a missing model is. @@ -1153,9 +1154,9 @@ def flaky_upsert(item: IndexedItem, model_id: str, embedding: np.ndarray) -> Non def test_batch_failure_is_charged_to_the_images_when_the_encoder_is_healthy( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """The other half: a poisonous image must still be quarantined so the backlog can advance. @@ -1198,9 +1199,9 @@ def test_systemic_backoff_grows_and_is_capped(service: ImageIndexService) -> Non def test_zero_norm_embedding_fails_only_its_own_image( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """A degenerate encoder row must not cost the rest of its batch their embeddings. @@ -1230,9 +1231,9 @@ def encode(images: list[Image.Image]) -> np.ndarray: def test_start_discards_only_other_models_embeddings( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """The one destructive operation in the service: prove what it does and does not delete. @@ -1257,9 +1258,9 @@ def test_start_discards_only_other_models_embeddings( def test_disabled_service_does_not_discard_embeddings( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """Turning the feature off must not destroy an index built while it was on.""" _save_image(image_records, "a.png") @@ -1274,9 +1275,9 @@ def test_disabled_service_does_not_discard_embeddings( def test_worker_waits_for_generation_to_finish_when_not_on_cpu( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, accelerator_host: None, ) -> None: """The VRAM contract: off the CPU path, embedding must pause while a generation runs. @@ -1303,9 +1304,9 @@ def test_worker_waits_for_generation_to_finish_when_not_on_cpu( def test_generation_wait_does_not_block_shutdown( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, accelerator_host: None, ) -> None: """A generation that never ends must not stop the worker from honouring stop().""" @@ -1323,9 +1324,9 @@ def test_generation_wait_does_not_block_shutdown( def test_projection_does_not_wait_for_an_in_progress_generation( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, accelerator_host: None, ) -> None: """A projection reads stored embeddings only — no encoder, no GPU — so it has no @@ -1362,9 +1363,9 @@ def test_projection_does_not_wait_for_an_in_progress_generation( def test_a_partially_stored_batch_does_not_escalate_the_backoff( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """A batch that stores ANYTHING resets the systemic-failure counter, even though the batch also failed. @@ -1408,9 +1409,9 @@ def flaky_upsert(item, model_id, embedding): def test_unparseable_device_is_ignored_rather_than_wedging_the_worker( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """`image_index_device` is free-form config with no validator. @@ -1430,7 +1431,7 @@ def test_unparseable_device_is_ignored_rather_than_wedging_the_worker( def test_empty_model_name_does_not_resolve_to_an_arbitrary_model( images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """`search_by_attr` drops its name predicate for a falsy name. @@ -1452,7 +1453,7 @@ def test_empty_model_name_does_not_resolve_to_an_arbitrary_model( def test_duplicate_model_names_resolve_deterministically( images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """Two models can share a name, and the pick must not move when one is reinstalled. @@ -1476,9 +1477,9 @@ def resolve(order: list[SimpleNamespace]) -> SimpleNamespace: def test_status_event_emitted( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: _save_image(image_records, "a.png") @@ -1499,9 +1500,9 @@ def _status_events(invoker) -> list[ImageIndexStatusEvent]: def test_on_changed_emits_pending_status_before_embedding( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: invoker = _make_invoker(images_service, index_records) @@ -1519,9 +1520,9 @@ def test_on_changed_emits_pending_status_before_embedding( def test_on_deleted_emits_status( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: _save_image(image_records, "a.png") @@ -1538,9 +1539,9 @@ def test_on_deleted_emits_status( def test_permanently_failed_image_still_reaches_quiescence( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: """A permanently-failing image must not wedge pending above zero forever.""" @@ -1567,9 +1568,9 @@ def get_pil_image(image_name: str) -> Image.Image: def test_upsert_failure_routes_through_retry_to_success( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: """A raise while storing embeddings must feed the retry path, not strand the image.""" @@ -1596,9 +1597,9 @@ def flaky_upsert(item: IndexedItem, model_id: str, embedding: np.ndarray) -> Non def test_upsert_value_error_also_routes_through_retry( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: """ValueError must reach the retry path like any other raise. @@ -1631,9 +1632,9 @@ def flaky_upsert(item: IndexedItem, model_id: str, embedding: np.ndarray) -> Non def test_normalizable_extreme_magnitudes_are_not_dropped( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """Tiny-but-normalizable vectors must not be misread as degenerate. @@ -1659,9 +1660,9 @@ def encode(images: list[Image.Image]) -> np.ndarray: def test_ineligible_transition_clears_failure_bookkeeping( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: """An image that leaves eligibility must stop counting against `failed`.""" @@ -1682,9 +1683,9 @@ def test_ineligible_transition_clears_failure_bookkeeping( def test_owner_poke_emitted_at_quiescence( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: """Owners get a counts-free user-routed poke once their embeds settle.""" @@ -1703,7 +1704,7 @@ def test_owner_poke_emitted_at_quiescence( def test_stop_joins_worker( - images_service: ImageService, index_records: ImageIndexRecordsSqlite, service: ImageIndexService + images_service: ImageService, index_records: ImageIndexRecords, service: ImageIndexService ) -> None: service.start(_make_invoker(images_service, index_records)) assert service._worker is not None and service._worker.is_alive() @@ -1717,9 +1718,9 @@ def test_stop_joins_worker( def test_projection_job_computes_and_caches( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: from invokeai.app.services.events.events_common import ImageMapProjectionReadyEvent @@ -1749,9 +1750,9 @@ def test_projection_job_computes_and_caches( def test_projection_failure_caches_empty_result_instead_of_looping( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, monkeypatch, ) -> None: @@ -1788,9 +1789,9 @@ def broken_umap(embeddings, seed=42): def test_a_permanently_failing_projection_is_retried_once_not_every_request( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, monkeypatch, ) -> None: @@ -1824,9 +1825,9 @@ def broken_umap(embeddings, seed=42): def test_a_cached_row_with_no_finite_points_is_a_failed_fit_not_a_result( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, monkeypatch, ) -> None: @@ -1881,9 +1882,9 @@ def broken_umap(embeddings, seed=42): def test_a_lost_projection_write_does_not_burn_the_retry( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, monkeypatch, ) -> None: @@ -1921,7 +1922,7 @@ def test_a_lost_projection_write_does_not_burn_the_retry( def failing_set_projection(*args, **kwargs): writes["n"] += 1 if writes["n"] == 1: - raise RuntimeError("database is locked") + raise LockTimeoutError("database is locked") return real_set_projection(*args, **kwargs) monkeypatch.setattr(index_records, "set_projection", failing_set_projection) @@ -1937,9 +1938,9 @@ def failing_set_projection(*args, **kwargs): def test_an_explicit_refresh_restores_a_spent_retry( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, monkeypatch, ) -> None: @@ -1994,9 +1995,9 @@ def test_projection_request_dedup_is_last_writer_wins(service: ImageIndexService def test_systemic_embedding_outage_does_not_starve_projections( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """The emergent interaction between "never retire on a systemic failure" and "projections only at quiescence". @@ -2040,9 +2041,9 @@ def broken_encode(images: list[Image.Image]) -> np.ndarray: def test_search_similar_ranks_by_cosine_and_respects_scope( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: service.start(_make_invoker(images_service, index_records)) @@ -2075,7 +2076,7 @@ def unit(index: int, mix: float = 0.0) -> np.ndarray: def test_embed_image_normalizes_and_requires_running_service( - images_service: ImageService, index_records: ImageIndexRecordsSqlite, service: ImageIndexService + images_service: ImageService, index_records: ImageIndexRecords, service: ImageIndexService ) -> None: probe = Image.new("RGB", (4, 4)) @@ -2090,7 +2091,7 @@ def test_embed_image_normalizes_and_requires_running_service( def test_embed_image_retries_once_after_a_failed_encode( - images_service: ImageService, index_records: ImageIndexRecordsSqlite + images_service: ImageService, index_records: ImageIndexRecords ) -> None: # A failed load evicts the model from the RAM cache so the next attempt # rebuilds it from disk; embed_image must make that second attempt itself. @@ -2122,10 +2123,49 @@ def always_failing(images): service.stop() +def test_a_projection_whose_database_work_fails_for_good_is_given_up_once( + image_records: ImageRecordStorage, + images_service: ImageService, + index_records: ImageIndexRecords, +) -> None: + """A failure that is not the database being busy (a statement the server refuses, say) fails the same way on + every retry: re-queueing it would spin the worker forever. The waiting client still hears that nothing newer is + coming.""" + from invokeai.app.services.events.events_common import ImageMapProjectionReadyEvent + + calls = {"n": 0} + + def refused(*args, **kwargs): + calls["n"] += 1 + raise RuntimeError("the server refused the statement") + + index_records.list_accessible_embedded_items = refused # type: ignore[method-assign] + + service = ImageIndexService(encode_fn=_fake_encode, model_id=MODEL_ID) + invoker = _make_invoker(images_service, index_records) + try: + _save_image(image_records, "img-0.png") + service.start(invoker) + _wait_until(lambda: not service._backfill_pending.is_set()) + + service.request_projection("system") + + _wait_until( + lambda: any(isinstance(e, ImageMapProjectionReadyEvent) for e in invoker.services.events.events), + timeout=30.0, + ) + time.sleep(0.5) + assert calls["n"] == 1 + ready = [e for e in invoker.services.events.events if isinstance(e, ImageMapProjectionReadyEvent)] + assert [(e.user_id, e.point_count) for e in ready] == [("system", 0)] + finally: + service.stop() + + def test_projection_request_is_requeued_when_the_database_read_fails( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """The job is popped from the dedup map before the work runs. @@ -2139,7 +2179,7 @@ def test_projection_request_is_requeued_when_the_database_read_fails( def flaky_list(user_id, model_id): calls["n"] += 1 if calls["n"] == 1: - raise RuntimeError("database is locked") + raise LockTimeoutError("database is locked") return real_list(user_id, model_id) index_records.list_accessible_embedded_items = flaky_list # type: ignore[method-assign] @@ -2161,9 +2201,9 @@ def flaky_list(user_id, model_id): def test_unchanged_scope_does_not_recompute_the_projection( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, ) -> None: """Repeat requests over an unchanged gallery must not re-run the fit. @@ -2249,7 +2289,7 @@ def test_projection_job_is_popped_before_running(service: ImageIndexService) -> def test_embed_text_unavailable_without_model_config( - images_service: ImageService, index_records: ImageIndexRecordsSqlite, service: ImageIndexService + images_service: ImageService, index_records: ImageIndexRecords, service: ImageIndexService ) -> None: from invokeai.app.services.image_index.image_index_base import TextSearchUnavailableError @@ -2517,16 +2557,16 @@ def test_search_similar_returns_empty_when_not_running(service: ImageIndexServic def test_search_similar_scopes_to_the_requesting_user( - db: SqliteDatabase, - image_records: SqliteImageRecordStorage, + db: Database, + image_records: ImageRecordStorage, images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: from invokeai.app.services.users.users_common import UserCreateRequest from invokeai.app.services.users.users_default import UserService - other_user = UserService(db=db).create( + other_user = UserService(db).create( UserCreateRequest(email="scoped@example.com", display_name="Scoped", password="TestPass123", is_admin=False) ) service.start(_make_invoker(images_service, index_records)) @@ -2986,7 +3026,7 @@ def _invalidate() -> None: def test_invalidate_vocab_rebuilds_with_the_new_terms_on_the_worker( images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, tmp_path, monkeypatch: pytest.MonkeyPatch, @@ -3020,7 +3060,7 @@ def test_invalidate_vocab_rebuilds_with_the_new_terms_on_the_worker( def test_invalidate_vocab_retries_a_failed_build( images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, tmp_path, monkeypatch: pytest.MonkeyPatch, @@ -3083,7 +3123,7 @@ def test_a_cache_with_a_mismatched_row_count_is_discarded_and_re_embedded( def test_a_transient_custom_terms_read_failure_keeps_the_rebuild_queued( images_service: ImageService, - index_records: ImageIndexRecordsSqlite, + index_records: ImageIndexRecords, service: ImageIndexService, tmp_path, monkeypatch: pytest.MonkeyPatch, @@ -3121,11 +3161,11 @@ def _flaky() -> list[str]: def test_backfill_indexes_preexisting_videos( - image_records: SqliteImageRecordStorage, + image_records: ImageRecordStorage, images_service: ImageService, videos_service: VideoService, - video_records: SqliteVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: _save_image(image_records, "img.png") @@ -3145,8 +3185,8 @@ def test_backfill_indexes_preexisting_videos( def test_new_video_is_indexed_from_its_thumbnail( images_service: ImageService, videos_service: VideoService, - video_records: SqliteVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: invoker = _make_invoker(images_service, index_records, videos_service=videos_service, video_records=video_records) @@ -3166,8 +3206,8 @@ def test_new_video_is_indexed_from_its_thumbnail( def test_video_thumbnail_is_what_gets_embedded( images_service: ImageService, videos_service: VideoService, - video_records: SqliteVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, ) -> None: # The embedding must come from the video's own thumbnail, not from the image service (which # here returns a different colour) — a wrong source would still produce a plausible vector. @@ -3204,8 +3244,8 @@ def distance(a: tuple[int, ...], b: tuple[int, ...]) -> int: def test_unreadable_video_thumbnail_is_charged_to_the_video( images_service: ImageService, videos_service: VideoService, - video_records: SqliteVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: # A video whose thumbnail never got written must not stall the backfill; it retires after @@ -3224,8 +3264,8 @@ def test_unreadable_video_thumbnail_is_charged_to_the_video( def test_deleted_video_is_forgotten( images_service: ImageService, videos_service: VideoService, - video_records: SqliteVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: invoker = _make_invoker(images_service, index_records, videos_service=videos_service, video_records=video_records) @@ -3252,9 +3292,9 @@ def test_deleted_video_is_forgotten( def test_video_owner_is_poked_when_its_embedding_lands( images_service: ImageService, videos_service: VideoService, - video_records: SqliteVideoRecordStorage, - image_records: SqliteImageRecordStorage, - index_records: ImageIndexRecordsSqlite, + video_records: VideoRecordStorage, + image_records: ImageRecordStorage, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: # The status event is admin-only, so this per-user poke is the only signal a non-admin @@ -3284,8 +3324,8 @@ def test_video_owner_is_poked_when_its_embedding_lands( def test_video_leaving_eligibility_clears_its_failure_bookkeeping( images_service: ImageService, videos_service: VideoService, - video_records: SqliteVideoRecordStorage, - index_records: ImageIndexRecordsSqlite, + video_records: VideoRecordStorage, + index_records: ImageIndexRecords, service: ImageIndexService, ) -> None: # A video that stops being a gallery item must stop counting against `failed`, or it skews diff --git a/tests/app/services/image_moves/test_image_moves_default.py b/tests/app/services/image_moves/test_image_moves_default.py index 0bd78605612..e5524e88249 100644 --- a/tests/app/services/image_moves/test_image_moves_default.py +++ b/tests/app/services/image_moves/test_image_moves_default.py @@ -6,33 +6,33 @@ import pytest from PIL import Image +from sqlalchemy import insert, select, update from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.image_files.image_files_disk import DiskImageFileStorage from invokeai.app.services.image_moves.image_moves_default import ( + ImageMoveJob, ImageMoveJobAlreadyRunning, ImageMoveQueueActive, ImageMoveService, UnreadableImageError, ) from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin -from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.shared.sqlite.sqlite_util import init_db +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.image_moves import ( + image_subfolder_move_items, + image_subfolder_move_jobs, +) +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.session_queue import session_queue from invokeai.backend.util.logging import InvokeAILogger -def _build_db(tmp_path: Path) -> SqliteDatabase: - logger = InvokeAILogger.get_logger() - config = InvokeAIAppConfig(use_memory_db=True) - config._root = tmp_path - image_files = DiskImageFileStorage(tmp_path / "images") - return init_db(config=config, logger=logger, image_files=image_files) - - def _save_record( - records: SqliteImageRecordStorage, + database: Database, + records: ImageRecordStorage, image_name: str, subfolder: str, created_at: str, @@ -48,13 +48,18 @@ def _save_record( is_intermediate=is_intermediate, image_subfolder=subfolder, ) - with records._db.transaction() as cursor: - cursor.execute("UPDATE images SET created_at = ? WHERE image_name = ?;", (created_at, image_name)) + _set_record(database, image_name, created_at=created_at) + + +def _set_record(database: Database, image_name: str, **values: str) -> None: + with database.begin(write=True) as conn: + conn.execute(update(images).where(images.c.image_name == image_name).values(**values)) def _save_image( + database: Database, service: ImageMoveService, - records: SqliteImageRecordStorage, + records: ImageRecordStorage, image_name: str, subfolder: str, created_at: str, @@ -62,6 +67,7 @@ def _save_image( is_intermediate: bool = False, ) -> None: _save_record( + database, records, image_name=image_name, subfolder=subfolder, @@ -78,38 +84,35 @@ def _corrupt_png_idat(path: Path) -> None: path.write_bytes(data) -def _service(tmp_path: Path, strategy: str = "date") -> tuple[ImageMoveService, SqliteImageRecordStorage]: - db = _build_db(tmp_path) - records = SqliteImageRecordStorage(db=db) +def _service(tmp_path: Path, database: Database, strategy: str = "date") -> tuple[ImageMoveService, ImageRecordStorage]: + records = ImageRecordStorage(database) storage = DiskImageFileStorage(tmp_path / "images") invoker = MagicMock() invoker.services.configuration.pil_compress_level = 6 storage.start(invoker) config = InvokeAIAppConfig(use_memory_db=True, image_subfolder_strategy=strategy) config._root = tmp_path - service = ImageMoveService(db=db, image_files=storage, config=config, logger=InvokeAILogger.get_logger()) + service = ImageMoveService(database, image_files=storage, config=config, logger=InvokeAILogger.get_logger()) return service, records -def _job_item_states(service: ImageMoveService, job_id: int) -> dict[str, str]: - with service._db.transaction() as cursor: - cursor.execute( - "SELECT image_name, state FROM image_subfolder_move_items WHERE job_id = ? ORDER BY image_name;", - (job_id,), - ) - return {row["image_name"]: row["state"] for row in cursor.fetchall()} +def _job_item_states(database: Database, job_id: int) -> dict[str, str]: + items = image_subfolder_move_items.c + statement = select(items.image_name, items.state).where(items.job_id == job_id).order_by(items.image_name) + with database.begin(write=False) as conn: + return dict(conn.execute(statement).all()) -def _job_states(service: ImageMoveService) -> dict[int, str]: - with service._db.transaction() as cursor: - cursor.execute("SELECT id, state FROM image_subfolder_move_jobs ORDER BY id;") - return {row["id"]: row["state"] for row in cursor.fetchall()} +def _job_states(database: Database) -> dict[int, str]: + jobs = image_subfolder_move_jobs.c + with database.begin(write=False) as conn: + return dict(conn.execute(select(jobs.id, jobs.state).order_by(jobs.id)).all()) -def test_move_all_images_uses_created_at_for_date_strategy(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_move_all_images_uses_created_at_for_date_strategy(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-a.png" - _save_record(records, image_name=image_name, subfolder="", created_at="2024-02-03 04:05:06.000") + _save_record(database, records, image_name=image_name, subfolder="", created_at="2024-02-03 04:05:06.000") service.image_files.save(Image.new("RGB", (16, 16), "red"), image_name=image_name) result = service.move_all_images() @@ -122,10 +125,11 @@ def test_move_all_images_uses_created_at_for_date_strategy(tmp_path: Path) -> No assert not service.image_files.get_path(image_name, image_subfolder="").exists() -def test_missing_intermediate_source_file_is_treated_as_success(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_missing_intermediate_source_file_is_treated_as_success(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "missing-intermediate.png" _save_record( + database, records, image_name=image_name, subfolder="", @@ -143,11 +147,12 @@ def test_missing_intermediate_source_file_is_treated_as_success(tmp_path: Path) assert service.get_latest_job().state == "committed" -def test_missing_intermediate_source_file_removes_orphaned_thumbnail(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_missing_intermediate_source_file_removes_orphaned_thumbnail(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "missing-intermediate-with-thumbnail.png" old_subfolder = "old/intermediate" _save_image( + database, service, records, image_name=image_name, @@ -168,10 +173,11 @@ def test_missing_intermediate_source_file_removes_orphaned_thumbnail(tmp_path: P assert records.get(image_name).image_subfolder == "2024/02/04" -def test_missing_non_intermediate_source_file_still_fails(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_missing_non_intermediate_source_file_still_fails(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "missing-general.png" _save_record( + database, records, image_name=image_name, subfolder="", @@ -183,18 +189,21 @@ def test_missing_non_intermediate_source_file_still_fails(tmp_path: Path) -> Non service.plan_batch(last_image_name="", limit=100) -def test_move_all_images_continues_after_missing_non_intermediate_source_file(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_move_all_images_continues_after_missing_non_intermediate_source_file( + tmp_path: Path, database: Database +) -> None: + service, records = _service(tmp_path, database, strategy="date") missing_image_name = "missing-general.png" valid_image_name = "valid-general.png" _save_record( + database, records, image_name=missing_image_name, subfolder="", created_at="2024-02-04 04:05:06.000", is_intermediate=False, ) - _save_image(service, records, valid_image_name, "", "2024-02-05 04:05:06.000", "blue") + _save_image(database, service, records, valid_image_name, "", "2024-02-05 04:05:06.000", "blue") result = service.move_all_images() @@ -202,14 +211,14 @@ def test_move_all_images_continues_after_missing_non_intermediate_source_file(tm assert result.committed == 1 assert records.get(missing_image_name).image_subfolder == "" assert records.get(valid_image_name).image_subfolder == "2024/02/05" - assert "error" in _job_states(service).values() - assert "committed" in _job_states(service).values() + assert "error" in _job_states(database).values() + assert "committed" in _job_states(database).values() -def test_move_recovers_existing_16_bit_destination_without_thumbnail(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_move_recovers_existing_16_bit_destination_without_thumbnail(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "large-16-bit.png" - _save_record(records, image_name=image_name, subfolder="", created_at="2024-02-04 04:05:06.000") + _save_record(database, records, image_name=image_name, subfolder="", created_at="2024-02-04 04:05:06.000") old_path = service.image_files.get_path(image_name) old_path.parent.mkdir(parents=True, exist_ok=True) image = Image.new("I;16", (1024, 1024), 32768) @@ -233,8 +242,8 @@ def test_move_recovers_existing_16_bit_destination_without_thumbnail(tmp_path: P image.close() -def test_regenerate_thumbnail_closes_temp_file_before_writing(tmp_path: Path) -> None: - service, _records = _service(tmp_path, strategy="date") +def test_regenerate_thumbnail_closes_temp_file_before_writing(tmp_path: Path, database: Database) -> None: + service, _records = _service(tmp_path, database, strategy="date") class TrackingTempFile: name = str(tmp_path / "thumbnail.tmp") @@ -273,10 +282,10 @@ def save(_path: Path, format: str) -> None: assert temp_file.closed -def test_move_does_not_relocate_source_when_thumbnail_generation_fails(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_move_does_not_relocate_source_when_thumbnail_generation_fails(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "thumbnail-retry.png" - _save_record(records, image_name=image_name, subfolder="", created_at="2024-02-05 04:05:06.000") + _save_record(database, records, image_name=image_name, subfolder="", created_at="2024-02-05 04:05:06.000") old_path = service.image_files.get_path(image_name) old_path.parent.mkdir(parents=True, exist_ok=True) Image.new("RGB", (32, 32), "red").save(old_path, format="PNG") @@ -302,10 +311,11 @@ def test_move_does_not_relocate_source_when_thumbnail_generation_fails(tmp_path: assert service.get_job(job_id).state == "committed" -def test_recovery_treats_missing_intermediate_source_file_as_success(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_recovery_treats_missing_intermediate_source_file_as_success(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "missing-intermediate-recovery.png" _save_record( + database, records, image_name=image_name, subfolder="", @@ -321,13 +331,13 @@ def test_recovery_treats_missing_intermediate_source_file_as_success(tmp_path: P assert recovered.errors == 0 assert records.get(image_name).image_subfolder == "2024/02/05" assert service.get_job(job_id).state == "committed" - assert _job_item_states(service, job_id) == {image_name: "committed"} + assert _job_item_states(database, job_id) == {image_name: "committed"} -def test_startup_recovery_commits_after_files_moved_but_db_not_updated(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_startup_recovery_commits_after_files_moved_but_db_not_updated(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-b.png" - _save_record(records, image_name=image_name, subfolder="", created_at="2025-06-07 08:09:10.000") + _save_record(database, records, image_name=image_name, subfolder="", created_at="2025-06-07 08:09:10.000") service.image_files.save(Image.new("RGB", (16, 16), "blue"), image_name=image_name) moves = service.plan_batch(last_image_name="", limit=100) @@ -343,10 +353,10 @@ def test_startup_recovery_commits_after_files_moved_but_db_not_updated(tmp_path: assert service.get_job(job_id).state == "committed" -def test_status_reports_unplanned_images_after_recovery(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") - _save_image(service, records, "image-recovered.png", "", "2024-02-03 04:05:06.000", "red") - _save_image(service, records, "image-unplanned.png", "", "2024-02-04 04:05:06.000", "blue") +def test_status_reports_unplanned_images_after_recovery(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") + _save_image(database, service, records, "image-recovered.png", "", "2024-02-03 04:05:06.000", "red") + _save_image(database, service, records, "image-unplanned.png", "", "2024-02-04 04:05:06.000", "blue") moves = service.plan_batch(last_image_name="", limit=1) job_id = service.create_move_job(moves) service.perform_filesystem_moves(job_id) @@ -363,11 +373,13 @@ def test_status_reports_unplanned_images_after_recovery(tmp_path: Path) -> None: assert status.needs_move_count == 1 -def test_cleanup_empty_source_directories_after_move(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_cleanup_empty_source_directories_after_move(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-c.png" old_subfolder = "old/nested" - _save_record(records, image_name=image_name, subfolder=old_subfolder, created_at="2024-11-12 01:02:03.000") + _save_record( + database, records, image_name=image_name, subfolder=old_subfolder, created_at="2024-11-12 01:02:03.000" + ) service.image_files.save(Image.new("RGB", (16, 16), "green"), image_name=image_name, image_subfolder=old_subfolder) old_parent = service.image_files.get_path(image_name, image_subfolder=old_subfolder).parent old_thumb_parent = service.image_files.get_path(image_name, thumbnail=True, image_subfolder=old_subfolder).parent @@ -381,8 +393,8 @@ def test_cleanup_empty_source_directories_after_move(tmp_path: Path) -> None: @pytest.mark.skipif(not hasattr(os, "symlink"), reason="symlinks are not supported on this platform") -def test_cleanup_empty_source_directories_stays_within_symlinked_root(tmp_path: Path) -> None: - service, _records = _service(tmp_path, strategy="date") +def test_cleanup_empty_source_directories_stays_within_symlinked_root(tmp_path: Path, database: Database) -> None: + service, _records = _service(tmp_path, database, strategy="date") real_root = tmp_path / "real-root" linked_root = tmp_path / "linked-root" sibling = tmp_path / "sibling" @@ -403,11 +415,11 @@ def test_cleanup_empty_source_directories_stays_within_symlinked_root(tmp_path: assert not (real_root / "old").exists() -def test_startup_recovery_cleans_empty_source_directories(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_startup_recovery_cleans_empty_source_directories(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-recovery-cleanup.png" old_subfolder = "old/recovery" - _save_image(service, records, image_name, old_subfolder, "2024-11-13 01:02:03.000", "green") + _save_image(database, service, records, image_name, old_subfolder, "2024-11-13 01:02:03.000", "green") moves = service.plan_batch(last_image_name="", limit=100) job_id = service.create_move_job(moves) move = moves[0] @@ -427,10 +439,10 @@ def test_startup_recovery_cleans_empty_source_directories(tmp_path: Path) -> Non assert service.get_job(job_id).state == "committed" -def test_preflight_rejects_active_uncommitted_job_for_same_image(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_preflight_rejects_active_uncommitted_job_for_same_image(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-d.png" - _save_record(records, image_name=image_name, subfolder="", created_at="2024-01-02 03:04:05.000") + _save_record(database, records, image_name=image_name, subfolder="", created_at="2024-01-02 03:04:05.000") service.image_files.save(Image.new("RGB", (16, 16), "yellow"), image_name=image_name) moves = service.plan_batch(last_image_name="", limit=100) @@ -440,10 +452,10 @@ def test_preflight_rejects_active_uncommitted_job_for_same_image(tmp_path: Path) service.plan_batch(last_image_name="", limit=100) -def test_create_move_job_rejects_second_active_job_from_stale_plan(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_create_move_job_rejects_second_active_job_from_stale_plan(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-active-race.png" - _save_image(service, records, image_name, "", "2024-01-03 03:04:05.000", "yellow") + _save_image(database, service, records, image_name, "", "2024-01-03 03:04:05.000", "yellow") stale_plan_a = service.plan_batch(last_image_name="", limit=100) stale_plan_b = service.plan_batch(last_image_name="", limit=100) @@ -453,10 +465,10 @@ def test_create_move_job_rejects_second_active_job_from_stale_plan(tmp_path: Pat service.create_move_job(stale_plan_b) -def test_startup_recovery_completes_planned_job_before_any_file_move(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_startup_recovery_completes_planned_job_before_any_file_move(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-e.png" - _save_image(service, records, image_name, "", "2024-03-04 05:06:07.000", "purple") + _save_image(database, service, records, image_name, "", "2024-03-04 05:06:07.000", "purple") moves = service.plan_batch(last_image_name="", limit=100) job_id = service.create_move_job(moves) @@ -470,13 +482,13 @@ def test_startup_recovery_completes_planned_job_before_any_file_move(tmp_path: P assert recovered_twice.errors == 0 assert records.get(image_name).image_subfolder == "2024/03/04" assert service.get_job(job_id).state == "committed" - assert _job_item_states(service, job_id) == {image_name: "committed"} + assert _job_item_states(database, job_id) == {image_name: "committed"} -def test_background_recovery_can_start_when_journal_job_is_active(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_background_recovery_can_start_when_journal_job_is_active(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-background-recovery.png" - _save_image(service, records, image_name, "", "2024-03-05 05:06:07.000", "purple") + _save_image(database, service, records, image_name, "", "2024-03-05 05:06:07.000", "purple") job_id = service.create_move_job(service.plan_batch(last_image_name="", limit=100)) status = service.start_background_recovery() @@ -492,10 +504,10 @@ def test_background_recovery_can_start_when_journal_job_is_active(tmp_path: Path assert service.get_job(job_id).state == "committed" -def test_start_runs_recovery_before_normal_operation(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_start_runs_recovery_before_normal_operation(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-startup-recovery.png" - _save_image(service, records, image_name, "", "2024-03-05 05:06:07.000", "purple") + _save_image(database, service, records, image_name, "", "2024-03-05 05:06:07.000", "purple") job_id = service.create_move_job(service.plan_batch(last_image_name="", limit=100)) service.perform_filesystem_moves(job_id) @@ -506,10 +518,10 @@ def test_start_runs_recovery_before_normal_operation(tmp_path: Path) -> None: assert service.is_maintenance_active() is False -def test_start_leaves_maintenance_active_when_recovery_remains_incomplete(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_start_leaves_maintenance_active_when_recovery_remains_incomplete(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-startup-recovery-retry.png" - _save_image(service, records, image_name, "", "2024-03-05 05:06:07.000", "purple") + _save_image(database, service, records, image_name, "", "2024-03-05 05:06:07.000", "purple") service.create_move_job(service.plan_batch(last_image_name="", limit=100)) with patch.object(service, "complete_partial_filesystem_moves", side_effect=OSError("temporary failure")): @@ -520,8 +532,10 @@ def test_start_leaves_maintenance_active_when_recovery_remains_incomplete(tmp_pa @pytest.mark.parametrize(("pending", "in_progress"), [(1, 0), (0, 1)]) -def test_background_move_rejects_active_queue_work(tmp_path: Path, pending: int, in_progress: int) -> None: - service, _records = _service(tmp_path, strategy="date") +def test_background_move_rejects_active_queue_work( + tmp_path: Path, database: Database, pending: int, in_progress: int +) -> None: + service, _records = _service(tmp_path, database, strategy="date") invoker = MagicMock() invoker.services.session_queue.has_active_queue_work.return_value = pending > 0 or in_progress > 0 service.start(invoker) @@ -530,8 +544,8 @@ def test_background_move_rejects_active_queue_work(tmp_path: Path, pending: int, service.start_background_move_all() -def test_background_move_is_reserved_before_queue_check(tmp_path: Path) -> None: - service, _records = _service(tmp_path, strategy="date") +def test_background_move_is_reserved_before_queue_check(tmp_path: Path, database: Database) -> None: + service, _records = _service(tmp_path, database, strategy="date") invoker = MagicMock() def has_active_queue_work() -> bool: @@ -547,8 +561,10 @@ def has_active_queue_work() -> bool: assert service.is_maintenance_active() is False -def test_gallery_maintenance_reservation_blocks_moves_and_releases_after_failure(tmp_path: Path) -> None: - service, _records = _service(tmp_path, strategy="date") +def test_gallery_maintenance_reservation_blocks_moves_and_releases_after_failure( + tmp_path: Path, database: Database +) -> None: + service, _records = _service(tmp_path, database, strategy="date") try: with service.reserve_gallery_maintenance(): @@ -567,8 +583,10 @@ def test_gallery_maintenance_reservation_blocks_moves_and_releases_after_failure service.stop() -def test_gallery_maintenance_reservation_refuses_active_queue_and_releases_flag(tmp_path: Path) -> None: - service, _records = _service(tmp_path, strategy="date") +def test_gallery_maintenance_reservation_refuses_active_queue_and_releases_flag( + tmp_path: Path, database: Database +) -> None: + service, _records = _service(tmp_path, database, strategy="date") invoker = MagicMock() def has_active_queue_work() -> bool: @@ -590,18 +608,21 @@ def has_active_queue_work() -> bool: @pytest.mark.parametrize("status", ["pending", "in_progress"]) -def test_gallery_maintenance_reservation_refuses_active_work_on_custom_queue(tmp_path: Path, status: str) -> None: - service, _records = _service(tmp_path, strategy="date") - queue = SqliteSessionQueue(db=service._db) - with service._db.transaction() as cursor: - cursor.execute( - """--sql - INSERT INTO session_queue (batch_id, queue_id, session_id, session, status) - VALUES ('custom-batch', 'custom-queue', 'custom-session', '{}', ?) - """, - (status,), +def test_gallery_maintenance_reservation_refuses_active_work_on_custom_queue( + tmp_path: Path, database: Database, status: str +) -> None: + service, _records = _service(tmp_path, database, strategy="date") + with database.begin(write=True) as conn: + conn.execute( + insert(session_queue).values( + batch_id="custom-batch", + queue_id="custom-queue", + session_id="custom-session", + session="{}", + status=status, + ) ) - service.set_session_queue(queue) + service.set_session_queue(SessionQueue(database)) try: with pytest.raises(ImageMoveQueueActive, match="queue work is active"): @@ -613,10 +634,12 @@ def test_gallery_maintenance_reservation_refuses_active_work_on_custom_queue(tmp service.stop() -def test_maintenance_is_active_while_background_job_or_uncommitted_journal_exists(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_maintenance_is_active_while_background_job_or_uncommitted_journal_exists( + tmp_path: Path, database: Database +) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-maintenance-active.png" - _save_image(service, records, image_name, "", "2024-03-05 05:06:07.000", "purple") + _save_image(database, service, records, image_name, "", "2024-03-05 05:06:07.000", "purple") service.create_move_job(service.plan_batch(last_image_name="", limit=100)) assert service.is_maintenance_active() is True @@ -635,8 +658,8 @@ def wait_for_release() -> None: service._future.result(timeout=5) -def test_background_worker_error_is_exposed_in_status(tmp_path: Path) -> None: - service, _records = _service(tmp_path, strategy="date") +def test_background_worker_error_is_exposed_in_status(tmp_path: Path, database: Database) -> None: + service, _records = _service(tmp_path, database, strategy="date") started_worker = threading.Event() release_worker = threading.Event() @@ -659,10 +682,10 @@ def raise_error() -> None: assert status.last_error == "background failed" -def test_stop_waits_for_active_background_job_without_recording_error(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_stop_waits_for_active_background_job_without_recording_error(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-background-stop.png" - _save_image(service, records, image_name, "", "2024-03-05 05:06:07.000", "purple") + _save_image(database, service, records, image_name, "", "2024-03-05 05:06:07.000", "purple") job_id = service.create_move_job(service.plan_batch(last_image_name="", limit=100)) release_worker = threading.Event() @@ -683,10 +706,10 @@ def wait_for_shutdown() -> None: assert service.get_background_status().last_error is None -def test_startup_recovery_completes_partial_multi_image_move(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") - _save_image(service, records, "image-f.png", "", "2024-04-05 06:07:08.000", "orange") - _save_image(service, records, "image-g.png", "", "2024-04-06 06:07:08.000", "cyan") +def test_startup_recovery_completes_partial_multi_image_move(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") + _save_image(database, service, records, "image-f.png", "", "2024-04-05 06:07:08.000", "orange") + _save_image(database, service, records, "image-g.png", "", "2024-04-06 06:07:08.000", "cyan") moves = service.plan_batch(last_image_name="", limit=100) job_id = service.create_move_job(moves) @@ -706,22 +729,20 @@ def test_startup_recovery_completes_partial_multi_image_move(tmp_path: Path) -> assert records.get("image-f.png").image_subfolder == "2024/04/05" assert records.get("image-g.png").image_subfolder == "2024/04/06" assert service.get_job(job_id).state == "committed" - assert _job_item_states(service, job_id) == {"image-f.png": "committed", "image-g.png": "committed"} + assert _job_item_states(database, job_id) == {"image-f.png": "committed", "image-g.png": "committed"} -def test_startup_recovery_marks_committed_after_db_update_but_before_journal_commit(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_startup_recovery_marks_committed_after_db_update_but_before_journal_commit( + tmp_path: Path, database: Database +) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-h.png" - _save_image(service, records, image_name, "", "2024-05-06 07:08:09.000", "pink") + _save_image(database, service, records, image_name, "", "2024-05-06 07:08:09.000", "pink") moves = service.plan_batch(last_image_name="", limit=100) job_id = service.create_move_job(moves) service.perform_filesystem_moves(job_id) - with service._db.transaction() as cursor: - cursor.execute( - "UPDATE images SET image_subfolder = ? WHERE image_name = ?;", - ("2024/05/06", image_name), - ) + _set_record(database, image_name, image_subfolder="2024/05/06") recovered_once = service.startup_recovery() recovered_twice = service.startup_recovery() @@ -732,13 +753,15 @@ def test_startup_recovery_marks_committed_after_db_update_but_before_journal_com assert recovered_twice.errors == 0 assert records.get(image_name).image_subfolder == "2024/05/06" assert service.get_job(job_id).state == "committed" - assert _job_item_states(service, job_id) == {image_name: "committed"} + assert _job_item_states(database, job_id) == {image_name: "committed"} -def test_startup_recovery_marks_error_when_both_old_and_new_full_size_files_exist(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_startup_recovery_marks_error_when_both_old_and_new_full_size_files_exist( + tmp_path: Path, database: Database +) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-i.png" - _save_image(service, records, image_name, "", "2024-07-08 09:10:11.000", "red") + _save_image(database, service, records, image_name, "", "2024-07-08 09:10:11.000", "red") moves = service.plan_batch(last_image_name="", limit=100) job_id = service.create_move_job(moves) move = moves[0] @@ -751,13 +774,15 @@ def test_startup_recovery_marks_error_when_both_old_and_new_full_size_files_exis assert recovered.errors == 1 assert records.get(image_name).image_subfolder == "" assert service.get_job(job_id).state == "error" - assert _job_item_states(service, job_id) == {image_name: "error"} + assert _job_item_states(database, job_id) == {image_name: "error"} -def test_startup_recovery_marks_error_when_neither_old_nor_new_full_size_file_exists(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_startup_recovery_marks_error_when_neither_old_nor_new_full_size_file_exists( + tmp_path: Path, database: Database +) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-j.png" - _save_image(service, records, image_name, "", "2024-08-09 10:11:12.000", "blue") + _save_image(database, service, records, image_name, "", "2024-08-09 10:11:12.000", "blue") moves = service.plan_batch(last_image_name="", limit=100) job_id = service.create_move_job(moves) moves[0].old_path.unlink() @@ -768,15 +793,15 @@ def test_startup_recovery_marks_error_when_neither_old_nor_new_full_size_file_ex assert recovered.errors == 1 assert records.get(image_name).image_subfolder == "" assert service.get_job(job_id).state == "error" - assert _job_item_states(service, job_id) == {image_name: "error"} + assert _job_item_states(database, job_id) == {image_name: "error"} -def test_startup_recovery_isolates_corrupt_image_and_commits_other_items(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_startup_recovery_isolates_corrupt_image_and_commits_other_items(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") good_image_name = "image-a-good.png" corrupt_image_name = "image-b-corrupt.png" - _save_image(service, records, good_image_name, "", "2024-09-10 11:12:13.000", "white") - _save_image(service, records, corrupt_image_name, "", "2024-09-11 11:12:13.000", "black") + _save_image(database, service, records, good_image_name, "", "2024-09-10 11:12:13.000", "white") + _save_image(database, service, records, corrupt_image_name, "", "2024-09-11 11:12:13.000", "black") moves = service.plan_batch(last_image_name="", limit=100) corrupt_move = next(move for move in moves if move.image_name == corrupt_image_name) @@ -791,7 +816,7 @@ def test_startup_recovery_isolates_corrupt_image_and_commits_other_items(tmp_pat assert service.is_maintenance_active() is False assert service.get_job(job_id).state == "error" assert corrupt_image_name in (service.get_job(job_id).error_message or "") - assert _job_item_states(service, job_id) == { + assert _job_item_states(database, job_id) == { good_image_name: "committed", corrupt_image_name: "error", } @@ -803,10 +828,12 @@ def test_startup_recovery_isolates_corrupt_image_and_commits_other_items(tmp_pat assert not corrupt_move.new_path.exists() -def test_startup_recovery_marks_truncated_image_error_without_retrying_forever(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_startup_recovery_marks_truncated_image_error_without_retrying_forever( + tmp_path: Path, database: Database +) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-truncated.png" - _save_image(service, records, image_name, "", "2024-09-12 11:12:13.000", "white") + _save_image(database, service, records, image_name, "", "2024-09-12 11:12:13.000", "white") move = service.plan_batch(last_image_name="", limit=100)[0] move.old_thumbnail_path.unlink() job_id = service.create_move_job([move]) @@ -822,16 +849,16 @@ def test_startup_recovery_marks_truncated_image_error_without_retrying_forever(t assert recovered.errors == 1 assert service.is_maintenance_active() is False assert service.get_job(job_id).state == "error" - assert _job_item_states(service, job_id) == {image_name: "error"} + assert _job_item_states(database, job_id) == {image_name: "error"} assert records.get(image_name).image_subfolder == "" assert move.old_path.exists() assert not move.new_path.exists() -def test_move_all_images_marks_corrupt_idat_as_unrecoverable(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_move_all_images_marks_corrupt_idat_as_unrecoverable(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-corrupt-idat.png" - _save_image(service, records, image_name, "", "2024-09-16 11:12:13.000", "white") + _save_image(database, service, records, image_name, "", "2024-09-16 11:12:13.000", "white") move = service.plan_batch(last_image_name="", limit=100)[0] move.old_thumbnail_path.unlink() _corrupt_png_idat(move.old_path) @@ -847,11 +874,12 @@ def test_move_all_images_marks_corrupt_idat_as_unrecoverable(tmp_path: Path) -> assert not move.new_path.exists() -def test_move_all_images_counts_and_reports_each_unrecoverable_item(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_move_all_images_counts_and_reports_each_unrecoverable_item(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_names = ["image-corrupt-a.png", "image-corrupt-b.png"] for index, image_name in enumerate(image_names): _save_image( + database, service, records, image_name, @@ -874,10 +902,12 @@ def test_move_all_images_counts_and_reports_each_unrecoverable_item(tmp_path: Pa assert all(records.get(image_name).image_subfolder == "" for image_name in image_names) -def test_startup_recovery_keeps_db_consistent_when_thumbnail_regeneration_fails(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_startup_recovery_keeps_db_consistent_when_thumbnail_regeneration_fails( + tmp_path: Path, database: Database +) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-thumbnail-recovery.png" - _save_image(service, records, image_name, "", "2024-09-13 11:12:13.000", "white") + _save_image(database, service, records, image_name, "", "2024-09-13 11:12:13.000", "white") move = service.plan_batch(last_image_name="", limit=100)[0] move.new_thumbnail_path.parent.mkdir(parents=True, exist_ok=True) copy2(move.old_thumbnail_path, move.new_thumbnail_path) @@ -898,15 +928,21 @@ def test_startup_recovery_keeps_db_consistent_when_thumbnail_regeneration_fails( assert not move.new_path.exists() -def test_startup_recovery_reconciles_db_after_destination_thumbnail_failure(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +@pytest.mark.parametrize(("repointed_meanwhile", "expected_subfolder"), [(False, "2024/09/15"), (True, "elsewhere")]) +def test_startup_recovery_reconciles_db_after_destination_thumbnail_failure( + tmp_path: Path, database: Database, repointed_meanwhile: bool, expected_subfolder: str +) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-destination-recovery.png" - _save_image(service, records, image_name, "", "2024-09-15 11:12:13.000", "white") + _save_image(database, service, records, image_name, "", "2024-09-15 11:12:13.000", "white") move = service.plan_batch(last_image_name="", limit=100)[0] move.old_thumbnail_path.unlink() move.new_path.parent.mkdir(parents=True, exist_ok=True) move.old_path.replace(move.new_path) job_id = service.create_move_job([move]) + if repointed_meanwhile: + # Reconciling repoints a record only while it still names the old subfolder. + _set_record(database, image_name, image_subfolder="elsewhere") with patch.object( service, @@ -919,16 +955,16 @@ def test_startup_recovery_reconciles_db_after_destination_thumbnail_failure(tmp_ assert recovered.errors == 1 assert service.is_maintenance_active() is False assert service.get_job(job_id).state == "error" - assert _job_item_states(service, job_id) == {image_name: "error"} - assert records.get(image_name).image_subfolder == "2024/09/15" + assert _job_item_states(database, job_id) == {image_name: "error"} + assert records.get(image_name).image_subfolder == expected_subfolder assert not move.old_path.exists() assert move.new_path.exists() -def test_startup_recovery_finalizes_job_after_item_error_was_recorded(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_startup_recovery_finalizes_job_after_item_error_was_recorded(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-recorded-error.png" - _save_image(service, records, image_name, "", "2024-09-14 11:12:13.000", "white") + _save_image(database, service, records, image_name, "", "2024-09-14 11:12:13.000", "white") move = service.plan_batch(last_image_name="", limit=100)[0] job_id = service.create_move_job([move]) service.mark_item_unrecoverable(job_id, image_name, "image is corrupt") @@ -939,13 +975,13 @@ def test_startup_recovery_finalizes_job_after_item_error_was_recorded(tmp_path: assert recovered.errors == 1 assert service.is_maintenance_active() is False assert service.get_job(job_id).state == "error" - assert _job_item_states(service, job_id) == {image_name: "error"} + assert _job_item_states(database, job_id) == {image_name: "error"} -def test_startup_recovery_keeps_job_recoverable_after_ordinary_exception(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_startup_recovery_keeps_job_recoverable_after_ordinary_exception(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-k.png" - _save_image(service, records, image_name, "", "2024-09-10 11:12:13.000", "white") + _save_image(database, service, records, image_name, "", "2024-09-10 11:12:13.000", "white") job_id = service.create_move_job(service.plan_batch(last_image_name="", limit=100)) with patch.object(service, "complete_partial_filesystem_moves", side_effect=OSError("temporary failure")): @@ -965,10 +1001,12 @@ def test_startup_recovery_keeps_job_recoverable_after_ordinary_exception(tmp_pat assert service.get_job(job_id).state == "committed" -def test_startup_recovery_regenerates_thumbnail_when_old_and_new_thumbnails_exist(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_startup_recovery_regenerates_thumbnail_when_old_and_new_thumbnails_exist( + tmp_path: Path, database: Database +) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-l.png" - _save_image(service, records, image_name, "", "2024-10-11 12:13:14.000", "black") + _save_image(database, service, records, image_name, "", "2024-10-11 12:13:14.000", "black") moves = service.plan_batch(last_image_name="", limit=100) job_id = service.create_move_job(moves) move = moves[0] @@ -987,19 +1025,19 @@ def test_startup_recovery_regenerates_thumbnail_when_old_and_new_thumbnails_exis assert service.get_job(job_id).state == "committed" -def test_preflight_rejects_duplicate_thumbnail_destination_paths(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") - _save_image(service, records, "same-name.jpg", "", "2024-12-13 14:15:16.000", "red") - _save_image(service, records, "same-name.png", "", "2024-12-13 14:15:16.000", "green") +def test_preflight_rejects_duplicate_thumbnail_destination_paths(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") + _save_image(database, service, records, "same-name.jpg", "", "2024-12-13 14:15:16.000", "red") + _save_image(database, service, records, "same-name.png", "", "2024-12-13 14:15:16.000", "green") with pytest.raises(ValueError, match="Duplicate destination thumbnail path"): service.plan_batch(last_image_name="", limit=100) -def test_successful_filesystem_move_fsyncs_files_and_directories(tmp_path: Path) -> None: - service, records = _service(tmp_path, strategy="date") +def test_successful_filesystem_move_fsyncs_files_and_directories(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") image_name = "image-m.png" - _save_image(service, records, image_name, "", "2025-01-02 03:04:05.000", "blue") + _save_image(database, service, records, image_name, "", "2025-01-02 03:04:05.000", "blue") job_id = service.create_move_job(service.plan_batch(last_image_name="", limit=100)) with ( @@ -1017,8 +1055,8 @@ def test_successful_filesystem_move_fsyncs_files_and_directories(tmp_path: Path) fsync_dir.assert_any_call(moved.old_thumbnail_path.parent) -def test_fsync_dir_ignores_platform_close_failures(tmp_path: Path) -> None: - service, _records = _service(tmp_path, strategy="date") +def test_fsync_dir_ignores_platform_close_failures(tmp_path: Path, database: Database) -> None: + service, _records = _service(tmp_path, database, strategy="date") with ( patch("invokeai.app.services.image_moves.image_moves_default.os.open", return_value=123), @@ -1034,8 +1072,8 @@ def test_fsync_dir_ignores_platform_close_failures(tmp_path: Path) -> None: service._fsync_dir(tmp_path) -def test_fsync_file_ignores_platform_fsync_failures(tmp_path: Path) -> None: - service, _records = _service(tmp_path, strategy="date") +def test_fsync_file_ignores_platform_fsync_failures(tmp_path: Path, database: Database) -> None: + service, _records = _service(tmp_path, database, strategy="date") path = tmp_path / "image.png" path.write_bytes(b"test") @@ -1044,3 +1082,126 @@ def test_fsync_file_ignores_platform_fsync_failures(tmp_path: Path) -> None: side_effect=OSError(9, "Bad file descriptor"), ): service._fsync_file(path) + + +def test_commit_rolls_back_every_repoint_when_a_record_was_repointed_meanwhile( + tmp_path: Path, database: Database +) -> None: + service, records = _service(tmp_path, database, strategy="date") + # The same day, so that one statement repoints both. + _save_image(database, service, records, "image-kept.png", "", "2024-03-01 01:02:03.000", "red") + _save_image(database, service, records, "image-raced.png", "", "2024-03-01 04:05:06.000", "blue") + moves = service.plan_batch(last_image_name="", limit=100) + job_id = service.create_move_job(moves) + service.perform_filesystem_moves(job_id) + _set_record(database, "image-raced.png", image_subfolder="elsewhere") + + with pytest.raises(RuntimeError, match="failed commit validation"): + service.commit_database_updates(job_id) + + # The commit repoints a record only while it still names the old subfolder, and validates the job as a whole. + assert records.get("image-raced.png").image_subfolder == "elsewhere" + assert records.get("image-kept.png").image_subfolder == "" + assert _job_item_states(database, job_id) == {"image-kept.png": "moved", "image-raced.png": "moved"} + # Records and files disagree until the job is finished, so it stays active. + assert service.get_job(job_id).state == "moved" + assert service.is_maintenance_active() + with pytest.raises(ValueError, match="active image move job"): + service.create_move_job(moves) + + +def test_commit_repoints_each_image_to_its_own_subfolder(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") + _save_image(database, service, records, "image-1.png", "", "2024-03-01 01:02:03.000", "red") + _save_image(database, service, records, "image-2.png", "", "2024-03-01 04:05:06.000", "blue") + _save_image(database, service, records, "image-3.png", "old", "2024-03-01 07:08:09.000", "green") + _save_image(database, service, records, "image-4.png", "", "2024-03-02 01:02:03.000", "pink") + _save_image(database, service, records, "image-5.png", "2024/03/03", "2024-03-03 01:02:03.000", "gray") + job_id = service.create_move_job(service.plan_batch(last_image_name="", limit=100)) + service.perform_filesystem_moves(job_id) + + assert service.commit_database_updates(job_id) == 4 + + names = [f"image-{i}.png" for i in range(1, 6)] + assert {name: records.get(name).image_subfolder for name in names} == { + "image-1.png": "2024/03/01", + "image-2.png": "2024/03/01", + "image-3.png": "2024/03/01", + "image-4.png": "2024/03/02", + "image-5.png": "2024/03/03", + } + + +def test_commit_rejects_a_moved_image_whose_record_was_deleted(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") + _save_image(database, service, records, "image-deleted.png", "", "2024-03-03 01:02:03.000", "red") + job_id = service.create_move_job(service.plan_batch(last_image_name="", limit=100)) + service.perform_filesystem_moves(job_id) + _set_record(database, "image-deleted.png", deleted_at="2024-03-04 00:00:00.000") + + with pytest.raises(RuntimeError, match="failed commit validation"): + service.commit_database_updates(job_id) + + assert service.get_job(job_id).state == "moved" + + +def test_deleted_records_are_neither_counted_nor_planned(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") + _save_image(database, service, records, "a-deleted.png", "", "2024-03-05 01:02:03.000", "red") + _save_image(database, service, records, "b-live.png", "", "2024-03-06 01:02:03.000", "blue") + _set_record(database, "a-deleted.png", deleted_at="2024-03-07 00:00:00.000") + + assert service.count_images_needing_move() == 1 + assert [move.image_name for move in service.plan_batch(last_image_name="", limit=100)] == ["b-live.png"] + assert database.queries.image_moves.next_image_name("") == "b-live.png" + + +def test_finished_jobs_are_not_recovered(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") + _save_image(database, service, records, "image-f.png", "", "2024-03-08 01:02:03.000", "red") + move = service.plan_batch(last_image_name="", limit=100)[0] + failed_job_id = service.create_error_move_job(move, "source vanished") + + assert database.queries.image_moves.recoverable_job_ids() == [] + job_id = service.create_move_job([move]) + assert database.queries.image_moves.recoverable_job_ids() == [job_id] + assert service.get_latest_job() == ImageMoveJob(id=job_id, state="planned", error_message=None) + + service.record_job_error_message(job_id, "temporary failure") + service.move_all_images() + + assert database.queries.image_moves.recoverable_job_ids() == [] + assert service.get_job(failed_job_id).error_message == "source vanished" + # A committed job no longer reports what went wrong before. + assert service.get_job(job_id) == ImageMoveJob(id=job_id, state="committed", error_message=None) + + +def test_unrecoverable_job_fails_every_item_and_releases_its_images(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") + _save_image(database, service, records, "image-g.png", "", "2024-03-09 01:02:03.000", "red") + _save_image(database, service, records, "image-h.png", "", "2024-03-10 01:02:03.000", "blue") + job_id = service.create_move_job(service.plan_batch(last_image_name="", limit=100)) + + service.mark_job_unrecoverable(job_id, "disk gone") + + assert service.get_job(job_id) == ImageMoveJob(id=job_id, state="error", error_message="disk gone") + assert _job_item_states(database, job_id) == {"image-g.png": "error", "image-h.png": "error"} + assert not service.is_maintenance_active() + assert len(service.plan_batch(last_image_name="", limit=100)) == 2 + + +def test_commit_reports_item_errors_in_image_name_order(tmp_path: Path, database: Database) -> None: + service, records = _service(tmp_path, database, strategy="date") + _save_image(database, service, records, "image-i.png", "", "2024-03-11 01:02:03.000", "red") + _save_image(database, service, records, "image-j.png", "", "2024-03-12 01:02:03.000", "blue") + # Journaled out of name order, so that only the commit's ordering puts the messages in it. + job_id = service.create_move_job(list(reversed(service.plan_batch(last_image_name="", limit=100)))) + service.mark_item_unrecoverable(job_id, "image-j.png", "second") + service.mark_item_unrecoverable(job_id, "image-i.png", "first") + + assert service.commit_database_updates(job_id) == 0 + + assert service.get_job(job_id) == ImageMoveJob(id=job_id, state="error", error_message="first\nsecond") + # Failed items moved nothing, so their records stay where they were. + assert records.get("image-i.png").image_subfolder == "" + assert records.get("image-j.png").image_subfolder == "" diff --git a/tests/app/services/image_records/test_image_records.py b/tests/app/services/image_records/test_image_records.py new file mode 100644 index 00000000000..625b50e7874 --- /dev/null +++ b/tests/app/services/image_records/test_image_records.py @@ -0,0 +1,1041 @@ +"""Image records on every database backend. + +Verifies that image_subfolder round-trips through save(), get(), get_many() and get_subfolders(), that +get_many()/get_image_names() enforce per-user ownership isolation, and how intermediates are deleted. +""" + +from collections.abc import Sequence +from typing import Any, Optional + +import pytest +from sqlalchemy import insert, select, true, update +from sqlalchemy.exc import DBAPIError + +from invokeai.app.invocations.fields import MetadataFieldValidator +from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage +from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardVisibility +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage +from invokeai.app.services.image_records.image_records_common import ( + ImageCategory, + ImageNamesResult, + ImageRecordChanges, + ImageRecordNotFoundException, + ResourceOrigin, +) +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries import images as image_queries +from invokeai.app.services.shared.database.schema.boards import shared_boards +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.users import users +from invokeai.app.services.shared.pagination import SQLiteDirection +from tests.fixtures.database import capture_statements, explain_query_plan +from tests.fixtures.races import while_in_flight +from tests.fixtures.sqlite_database import sqlite_cursor + +Stores = tuple[ImageRecordStorage, BoardRecordStorage, BoardImageRecordStorage] + + +@pytest.fixture +def store(database: Database) -> ImageRecordStorage: + return ImageRecordStorage(database) + + +@pytest.fixture +def stores(database: Database) -> Stores: + """Image, board, and board-image storages sharing one database.""" + return ImageRecordStorage(database), BoardRecordStorage(database), BoardImageRecordStorage(database) + + +def _save( + store: ImageRecordStorage, + name: str, + subfolder: str = "", + is_intermediate: bool = False, + user_id: Optional[str] = None, + category: ImageCategory = ImageCategory.GENERAL, + metadata: Optional[str] = None, +) -> None: + store.save( + image_name=name, + image_origin=ResourceOrigin.INTERNAL, + image_category=category, + width=64, + height=64, + has_workflow=False, + is_intermediate=is_intermediate, + image_subfolder=subfolder, + user_id=user_id, + metadata=metadata, + ) + + +def _set_created_at(database: Database, name: str, created_at: str) -> None: + """created_at is written when the record is saved; tests set it directly.""" + with database.begin(write=True) as conn: + conn.execute(update(images).where(images.c.image_name == name).values(created_at=created_at)) + + +def _intermediates(database: Database) -> list[tuple[str, str]]: + with database.begin(write=False) as conn: + statement = select(images.c.image_name, images.c.image_subfolder).where(images.c.is_intermediate == true()) + return [(row[0], row[1]) for row in conn.execute(statement.order_by(images.c.image_name))] + + +def _capture_names_plan( + database: Database, store: ImageRecordStorage, **kwargs: Any +) -> tuple[ImageNamesResult, str, list[str]]: + with capture_statements(database) as statements: + result = store.get_image_names(**kwargs) + statement, parameters = next( + (statement, parameters) + for statement, parameters in statements + if statement.lstrip().startswith("SELECT images.image_name") + ) + return result, statement, explain_query_plan(database, statement, parameters) + + +class TestImageSubfolderRoundTrip: + """save() -> get() preserves image_subfolder.""" + + def test_default_empty_subfolder(self, store: ImageRecordStorage) -> None: + _save(store, "img_default.png") + record = store.get("img_default.png") + assert record.image_subfolder == "" + + def test_custom_subfolder(self, store: ImageRecordStorage) -> None: + _save(store, "img_sub.png", subfolder="2026/04/11") + record = store.get("img_sub.png") + assert record.image_subfolder == "2026/04/11" + + def test_nested_subfolder(self, store: ImageRecordStorage) -> None: + _save(store, "img_nested.png", subfolder="a/b/c/d") + record = store.get("img_nested.png") + assert record.image_subfolder == "a/b/c/d" + + +class TestGetManySubfolder: + """get_many() deserializes image_subfolder for every row.""" + + def test_get_many_returns_subfolders(self, store: ImageRecordStorage) -> None: + _save(store, "flat.png", subfolder="") + _save(store, "dated.png", subfolder="2026/01") + _save(store, "hashed.png", subfolder="ab") + + result = store.get_many(limit=10, order_dir=SQLiteDirection.Ascending) + by_name = {r.image_name: r.image_subfolder for r in result.items} + + assert by_name["flat.png"] == "" + assert by_name["dated.png"] == "2026/01" + assert by_name["hashed.png"] == "ab" + + +class TestImageRecordExists: + def test_exists_returns_true_for_saved_image(self, store: ImageRecordStorage) -> None: + _save(store, "exists.png") + + assert store.exists("exists.png") is True + + def test_exists_returns_false_for_missing_image(self, store: ImageRecordStorage) -> None: + assert store.exists("missing.png") is False + + +class TestSave: + def test_saving_a_taken_name_keeps_the_record_and_returns_its_time( + self, database: Database, store: ImageRecordStorage + ) -> None: + _save(store, "taken.png", subfolder="first", user_id="alice") + _set_created_at(database, "taken.png", "2026-01-02 03:04:05.678") + + created_at = store.save( + image_name="taken.png", + image_origin=ResourceOrigin.EXTERNAL, + image_category=ImageCategory.USER, + width=8, + height=8, + has_workflow=True, + image_subfolder="second", + user_id="bob", + ) + + assert created_at.isoformat(sep=" ", timespec="milliseconds") == "2026-01-02 03:04:05.678" + record = store.get("taken.png") + assert (record.image_subfolder, record.image_origin, store.get_user_id("taken.png")) == ( + "first", + ResourceOrigin.INTERNAL, + "alice", + ) + + def test_an_image_without_an_owner_belongs_to_the_system(self, store: ImageRecordStorage) -> None: + _save(store, "unowned.png") + + assert store.get_user_id("unowned.png") == "system" + assert store.get_user_id("missing.png") is None + + +class TestFields: + def test_save_stores_every_field_and_update_changes_every_field(self, store: ImageRecordStorage) -> None: + store.save( + image_name="full.png", + image_origin=ResourceOrigin.EXTERNAL, + image_category=ImageCategory.CONTROL, + width=640, + height=480, + has_workflow=True, + is_intermediate=True, + starred=True, + session_id="session-1", + node_id="node-1", + metadata='{"seed": 1}', + user_id="user-1", + image_subfolder="sub/dir", + project_id="project-1", + ) + + saved = store.get("full.png") + store.update( + "full.png", + ImageRecordChanges( + image_category=ImageCategory.USER, session_id="session-2", is_intermediate=False, starred=False + ), + ) + updated = store.get("full.png") + + assert ( + saved.image_origin, + saved.image_category, + saved.width, + saved.height, + saved.has_workflow, + saved.is_intermediate, + saved.starred, + saved.session_id, + saved.node_id, + saved.image_subfolder, + saved.project_id, + ) == ( + ResourceOrigin.EXTERNAL, + ImageCategory.CONTROL, + 640, + 480, + True, + True, + True, + "session-1", + "node-1", + "sub/dir", + "project-1", + ) + assert store.get_user_id("full.png") == "user-1" + assert store.get_metadata("full.png") == MetadataFieldValidator.validate_json('{"seed": 1}') + assert (updated.image_category, updated.session_id, updated.is_intermediate, updated.starred) == ( + ImageCategory.USER, + "session-2", + False, + False, + ) + assert updated.node_id == "node-1" and updated.has_workflow is True + + +class TestFilters: + def test_the_origin_filter_narrows_pages_and_names_alike(self, store: ImageRecordStorage) -> None: + _save(store, "internal.png") + store.save( + image_name="external.png", + image_origin=ResourceOrigin.EXTERNAL, + image_category=ImageCategory.GENERAL, + width=8, + height=8, + has_workflow=False, + ) + + page = store.get_many(limit=10, image_origin=ResourceOrigin.EXTERNAL) + names = store.get_image_names(image_origin=ResourceOrigin.EXTERNAL) + + assert [r.image_name for r in page.items] == names.image_names == ["external.png"] + assert page.total == 1 + + def test_a_size_backfill_fills_only_sizes_not_known_yet(self, store: ImageRecordStorage) -> None: + _save(store, "unknown.png") + _save(store, "measured.png") + store.set_file_size_bytes("measured.png", 100) + + store.set_file_sizes_bytes({"unknown.png": 7, "measured.png": 9}) + + assert [store.get(name).file_size_bytes for name in ("unknown.png", "measured.png")] == [7, 100] + + +class TestGetSubfolders: + """get_subfolders() maps the named rows to their on-disk subfolders without touching them.""" + + def test_returns_subfolders_of_existing_rows_only(self, store: ImageRecordStorage) -> None: + _save(store, "keep.png", subfolder="general", is_intermediate=False) + _save(store, "tmp1.png", subfolder="intermediate", is_intermediate=True) + + assert store.get_subfolders(["keep.png", "tmp1.png", "missing.png"]) == { + "keep.png": "general", + "tmp1.png": "intermediate", + } + assert store.get("tmp1.png").image_subfolder == "intermediate" + + def test_intermediates_are_deleted_via_delete_intermediates_by_names( + self, database: Database, store: ImageRecordStorage + ) -> None: + _save(store, "tmp.png", subfolder="x", is_intermediate=True) + deleted = store.delete_intermediates_by_names([name for name, _ in _intermediates(database)]) + + assert deleted == ["tmp.png"] + with pytest.raises(ImageRecordNotFoundException): + store.get("tmp.png") + + +@pytest.mark.sqlite_only # A broken table would stay broken for every later test on a shared server schema. +class TestQueryFaultsAreNotNotFound: + """A failing query means the database is unavailable, not that the image is missing. + + Reporting a query fault as "not found" propagates all the way to the API, where it becomes a 404 + and tells the frontend to drop a live image from its cache. + """ + + def _break_the_images_table(self, database: Database) -> None: + with database.begin(write=True) as conn: + conn.exec_driver_sql("ALTER TABLE images RENAME TO images_moved") + + def test_get_raises_the_db_error_not_not_found(self, database: Database, store: ImageRecordStorage) -> None: + _save(store, "live.png") + self._break_the_images_table(database) + + with pytest.raises(DBAPIError): + store.get("live.png") + + def test_get_metadata_raises_the_db_error_not_not_found( + self, database: Database, store: ImageRecordStorage + ) -> None: + _save(store, "live.png") + self._break_the_images_table(database) + + with pytest.raises(DBAPIError): + store.get_metadata("live.png") + + def test_missing_row_still_raises_not_found(self, store: ImageRecordStorage) -> None: + """The genuine not-found path is untouched.""" + with pytest.raises(ImageRecordNotFoundException): + store.get("never-existed.png") + with pytest.raises(ImageRecordNotFoundException): + store.get_metadata("never-existed.png") + + +class TestDeleteIntermediatesByNames: + """delete_intermediates_by_names() deletes only rows that are still intermediates.""" + + def test_promoted_image_keeps_its_record(self, database: Database, store: ImageRecordStorage) -> None: + """An image promoted out of intermediate status after the snapshot must survive.""" + _save(store, "tmp.png", subfolder="x", is_intermediate=True) + _save(store, "promoted.png", subfolder="x", is_intermediate=True) + snapshot = [name for name, _ in _intermediates(database)] + assert set(snapshot) == {"tmp.png", "promoted.png"} + + # Simulate the race: the image stops being an intermediate between the snapshot and delete. + store.update("promoted.png", ImageRecordChanges(is_intermediate=False)) + + deleted = store.delete_intermediates_by_names(snapshot) + + assert deleted == ["tmp.png"] + # promoted.png is excluded from the returned names, so the caller never purges its files. + assert store.get("promoted.png").is_intermediate is False + with pytest.raises(ImageRecordNotFoundException): + store.get("tmp.png") + + def test_a_promotion_in_flight_is_waited_for_and_keeps_the_record( + self, database: Database, store: ImageRecordStorage + ) -> None: + """A delete that meets a promotion still in flight waits for it, and then keeps the promoted record.""" + _save(store, "tmp.png", is_intermediate=True) + _save(store, "promoted.png", is_intermediate=True) + deleted: list[str] = [] + + def promote(q: Queries) -> None: + q.images.update("promoted.png", ImageRecordChanges(is_intermediate=False)) + + errors = while_in_flight( + database, promote, lambda: deleted.extend(store.delete_intermediates_by_names(["tmp.png", "promoted.png"])) + ) + + assert errors == [] + assert deleted == ["tmp.png"] + assert store.get("promoted.png").is_intermediate is False + + def test_unknown_and_empty_names_are_no_ops(self, store: ImageRecordStorage) -> None: + _save(store, "keep.png", is_intermediate=False) + + assert store.delete_intermediates_by_names([]) == [] + # "gone.png" has no record at all and "keep.png" is not an intermediate, so neither is + # deleted or returned; keep.png must still be present afterwards. + assert store.delete_intermediates_by_names(["gone.png", "keep.png"]) == [] + assert store.get("keep.png").image_name == "keep.png" + + def test_names_spanning_several_statements( + self, database: Database, store: ImageRecordStorage, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Chunking must not lose rows: a name list spanning several chunks.""" + monkeypatch.setattr(image_queries, "IN_CHUNK", 3) + names = [f"tmp{i:02d}.png" for i in range(3 * 2 + 1)] + for name in names: + _save(store, name, is_intermediate=True) + # One image in the middle chunk is promoted and must survive. + survivor = names[4] + store.update(survivor, ImageRecordChanges(is_intermediate=False)) + + deleted = store.delete_intermediates_by_names(names[::-1]) + + assert deleted == [name for name in names[::-1] if name != survivor] + assert store.get(survivor).is_intermediate is False + assert _intermediates(database) == [] + + @pytest.mark.parametrize("operation", ["delete_intermediates_by_names", "delete_many", "get_subfolders"]) + def test_every_name_is_covered_and_no_statement_binds_more_than_sqlite_takes( + self, database: Database, store: ImageRecordStorage, operation: str + ) -> None: + """999 is the SQLITE_MAX_VARIABLE_NUMBER default on builds older than 3.32.""" + names = [f"tmp{i:05d}.png" for i in range(2 * 500 + 7)] + with database.begin(write=True) as conn: + conn.execute( + insert(images), + [ + { + "image_name": name, + "image_origin": ResourceOrigin.INTERNAL.value, + "image_category": ImageCategory.GENERAL.value, + "width": 8, + "height": 8, + "is_intermediate": True, + "image_subfolder": "sub", + } + for name in names + ], + ) + + with capture_statements(database) as statements: + result = getattr(store, operation)(names) + + widest = max(len(parameters) for _, parameters in statements) + assert 0 < widest <= 999 + if operation == "get_subfolders": + assert result == dict.fromkeys(names, "sub") + else: + assert store.get_image_names().image_names == [] + if operation == "delete_intermediates_by_names": + assert result == names + + def test_a_guard_narrows_what_is_deleted(self, store: ImageRecordStorage) -> None: + _save(store, "a.png", is_intermediate=True) + _save(store, "b.png", is_intermediate=True) + asked: list[list[str]] = [] + + def guard(q: Queries, names: Sequence[str]) -> list[str]: + asked.append(list(names)) + return [name for name in names if name != "b.png"] + + assert store.delete_intermediates_by_names(["a.png", "b.png"], guard) == ["a.png"] + assert asked == [["a.png", "b.png"]] + assert store.get("b.png").is_intermediate is True + + +class TestDeleteMany: + def test_names_spanning_several_statements_are_all_deleted( + self, database: Database, store: ImageRecordStorage, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(image_queries, "IN_CHUNK", 3) + names = [f"img{i:02d}.png" for i in range(3 * 2 + 1)] + for name in [*names, "kept.png"]: + _save(store, name) + + store.delete_many(names) + + assert store.get_image_names(starred_first=False).image_names == ["kept.png"] + + +class TestBoardCover: + def test_the_cover_is_the_newest_image_that_is_no_intermediate_a_starred_one_first( + self, database: Database, stores: Stores + ) -> None: + image_store, board_store, board_image_store = stores + board = board_store.save(board_name="Board", user_id="system") + for name, created_at, starred, intermediate in ( + ("old-starred.png", "2026-01-01 00:00:00.000", True, False), + ("newer.png", "2026-01-02 00:00:00.000", False, False), + ("newest-starred-intermediate.png", "2026-01-03 00:00:00.000", True, True), + ): + _save(image_store, name, is_intermediate=intermediate) + image_store.update(name, ImageRecordChanges(starred=starred)) + _set_created_at(database, name, created_at) + board_image_store.add_image_to_board(board_id=board.board_id, image_name=name) + + cover = image_store.get_most_recent_image_for_board(board.board_id) + + assert cover is not None and cover.image_name == "old-starred.png" + assert image_store.get_most_recent_image_for_board(board_store.save("Empty", "system").board_id) is None + + +class TestSearch: + def test_search_ignores_case_and_matches_percent_and_underscore_literally(self, store: ImageRecordStorage) -> None: + _save(store, "literal.png", metadata='{"positive_prompt": "100% Cotton_shirt"}') + # Matched too if `%` and `_` in the term were wildcards. + _save(store, "wildcard.png", metadata='{"positive_prompt": "100 cotton shirts"}') + + result = store.get_image_names(search_term="100% COTTON_") + + assert result.image_names == ["literal.png"] + assert store.get_many(limit=10, search_term="100% COTTON_").total == 1 + + +class TestSearchBeyondAscii: + def test_a_search_ignores_the_case_of_letters_beyond_ascii(self, store: ImageRecordStorage) -> None: + _save(store, "apples.png", metadata='{"prompt": "frische äpfel"}') + + assert store.get_image_names(search_term="ÄPFEL").image_names == ["apples.png"] + + +class TestOwnershipFilteringOmittedBoard: + """get_many()/get_image_names() enforce per-user isolation when board_id is omitted. + + Without this, a non-admin could enumerate every user's images (including images + on other users' private boards) simply by omitting the board_id query parameter. + """ + + def _seed_two_users(self, stores: Stores) -> str: + """user1: one image on a private board + one uncategorized. user2: one uncategorized.""" + image_store, board_store, board_image_store = stores + _save(image_store, "u1-boarded.png", user_id="user1") + _save(image_store, "u1-uncat.png", user_id="user1") + _save(image_store, "u2-uncat.png", user_id="user2") + board = board_store.save(board_name="User1 Private Board", user_id="user1") + board_image_store.add_image_to_board(board_id=board.board_id, image_name="u1-boarded.png") + return board.board_id + + def test_get_many_omitted_board_filters_by_owner(self, stores: Stores) -> None: + self._seed_two_users(stores) + image_store = stores[0] + + result = image_store.get_many(limit=10, user_id="user2", is_admin=False) + + assert {r.image_name for r in result.items} == {"u2-uncat.png"} + assert result.total == 1 + + def test_get_many_omitted_board_owner_sees_boarded_and_uncategorized(self, stores: Stores) -> None: + self._seed_two_users(stores) + image_store = stores[0] + + result = image_store.get_many(limit=10, user_id="user1", is_admin=False) + + assert {r.image_name for r in result.items} == {"u1-boarded.png", "u1-uncat.png"} + assert result.total == 2 + + def test_get_many_omitted_board_admin_sees_all(self, stores: Stores) -> None: + self._seed_two_users(stores) + image_store = stores[0] + + result = image_store.get_many(limit=10, user_id="admin", is_admin=True) + + assert {r.image_name for r in result.items} == {"u1-boarded.png", "u1-uncat.png", "u2-uncat.png"} + assert result.total == 3 + + def test_get_many_omitted_board_single_user_mode_sees_all(self, stores: Stores) -> None: + """user_id=None (single-user mode) applies no ownership filter.""" + self._seed_two_users(stores) + image_store = stores[0] + + result = image_store.get_many(limit=10, user_id=None, is_admin=False) + + assert result.total == 3 + + @pytest.mark.parametrize(("user_id", "is_admin"), [("admin", True), (None, False)]) + def test_get_many_none_board_unscoped_lists_every_owners_unboarded_images( + self, stores: Stores, user_id: Optional[str], is_admin: bool + ) -> None: + self._seed_two_users(stores) + + result = stores[0].get_many(limit=10, board_id="none", user_id=user_id, is_admin=is_admin) + + assert {r.image_name for r in result.items} == {"u1-uncat.png", "u2-uncat.png"} + + def test_get_many_none_board_still_filters_by_owner(self, stores: Stores) -> None: + """board_id="none" (uncategorized) keeps its existing per-user isolation.""" + self._seed_two_users(stores) + image_store = stores[0] + + result = image_store.get_many(limit=10, board_id="none", user_id="user1", is_admin=False) + + assert {r.image_name for r in result.items} == {"u1-uncat.png"} + + def test_get_many_explicit_board_returns_board_contents(self, stores: Stores) -> None: + """An explicit board_id lists that board's images; read access is the router's job.""" + board_id = self._seed_two_users(stores) + image_store = stores[0] + + result = image_store.get_many(limit=10, board_id=board_id, user_id="user1", is_admin=False) + + assert {r.image_name for r in result.items} == {"u1-boarded.png"} + + def test_get_image_names_omitted_board_filters_by_owner(self, stores: Stores) -> None: + self._seed_two_users(stores) + image_store = stores[0] + + result = image_store.get_image_names(user_id="user2", is_admin=False) + + assert result.image_names == ["u2-uncat.png"] + assert result.total_count == 1 + + def test_get_image_names_omitted_board_admin_sees_all(self, stores: Stores) -> None: + self._seed_two_users(stores) + image_store = stores[0] + + result = image_store.get_image_names(user_id="admin", is_admin=True) + + assert set(result.image_names) == {"u1-boarded.png", "u1-uncat.png", "u2-uncat.png"} + assert result.total_count == 3 + + def test_get_image_names_omitted_board_single_user_mode_sees_all(self, stores: Stores) -> None: + self._seed_two_users(stores) + image_store = stores[0] + + result = image_store.get_image_names(user_id=None, is_admin=False) + + assert result.total_count == 3 + + def test_get_image_names_none_board_still_filters_by_owner(self, stores: Stores) -> None: + self._seed_two_users(stores) + image_store = stores[0] + + result = image_store.get_image_names(board_id="none", user_id="user1", is_admin=False) + + assert result.image_names == ["u1-uncat.png"] + + def test_get_image_names_explicit_board_returns_board_contents(self, stores: Stores) -> None: + board_id = self._seed_two_users(stores) + image_store = stores[0] + + result = image_store.get_image_names(board_id=board_id, user_id="user1", is_admin=False) + + assert result.image_names == ["u1-boarded.png"] + + +class TestAllReadableBoardsFiltering: + def _seed_visibility_matrix(self, database: Database, stores: Stores) -> None: + image_store, board_store, board_image_store = stores + + def add_boarded_image( + image_name: str, + *, + owner: str, + visibility: BoardVisibility = BoardVisibility.Private, + archived: bool = False, + ) -> str: + _save(image_store, image_name, user_id=owner) + board = board_store.save(board_name=image_name, user_id=owner) + board_store.update( + board.board_id, + BoardChanges(board_visibility=visibility, archived=archived), + ) + board_image_store.add_image_to_board(board_id=board.board_id, image_name=image_name) + return board.board_id + + add_boarded_image("own-private.png", owner="user1") + shared_private_id = add_boarded_image("explicit-share.png", owner="user2") + add_boarded_image("shared-visibility.png", owner="user2", visibility=BoardVisibility.Shared) + add_boarded_image("public-visibility.png", owner="user2", visibility=BoardVisibility.Public) + add_boarded_image("other-private.png", owner="user2") + add_boarded_image("own-archived.png", owner="user1", archived=True) + add_boarded_image("other-archived-public.png", owner="user2", visibility=BoardVisibility.Public, archived=True) + _save(image_store, "own-uncategorized.png", user_id="user1") + _save(image_store, "other-uncategorized.png", user_id="user2") + + with database.begin(write=True) as conn: + conn.execute(insert(users).values(user_id="user1", email="user1@example.com", password_hash="unused")) + conn.execute(insert(shared_boards).values(board_id=shared_private_id, user_id="user1")) + + @pytest.mark.parametrize( + ("user_id", "is_admin", "expected"), + [ + ( + "user1", + False, + { + "own-private.png", + "explicit-share.png", + "shared-visibility.png", + "public-visibility.png", + "own-uncategorized.png", + }, + ), + ("user3", False, {"shared-visibility.png", "public-visibility.png"}), + ( + "admin", + True, + { + "own-private.png", + "explicit-share.png", + "shared-visibility.png", + "public-visibility.png", + "other-private.png", + "own-uncategorized.png", + "other-uncategorized.png", + }, + ), + ], + ) + def test_all_scope_authorization_and_counts_are_consistent( + self, + database: Database, + stores: Stores, + user_id: str, + is_admin: bool, + expected: set[str], + ) -> None: + self._seed_visibility_matrix(database, stores) + image_store = stores[0] + for name in ("own-private.png", "other-private.png"): + image_store.update(name, ImageRecordChanges(starred=True)) + + dtos = image_store.get_many(limit=100, board_id="all", user_id=user_id, is_admin=is_admin) + names = image_store.get_image_names(board_id="all", user_id=user_id, is_admin=is_admin) + + assert {image.image_name for image in dtos.items} == expected + assert set(names.image_names) == expected + assert dtos.total == names.total_count == len(expected) + assert names.starred_count == len(expected & {"own-private.png", "other-private.png"}) + + def test_all_scope_combines_with_inclusive_date_filters(self, database: Database, stores: Stores) -> None: + self._seed_visibility_matrix(database, stores) + image_store = stores[0] + _set_created_at(database, "own-private.png", "2026-06-01 00:00:00.000") + _set_created_at(database, "explicit-share.png", "2026-06-01 23:59:59.999") + _set_created_at(database, "public-visibility.png", "2026-06-02 00:00:00.000") + + dtos = image_store.get_many( + limit=100, + board_id="all", + user_id="user1", + created_from="2026-06-01", + created_to="2026-06-01", + ) + names = image_store.get_image_names( + board_id="all", + user_id="user1", + created_from="2026-06-01", + created_to="2026-06-01", + ) + + assert {image.image_name for image in dtos.items} == {"own-private.png", "explicit-share.png"} + assert set(names.image_names) == {"own-private.png", "explicit-share.png"} + assert dtos.total == names.total_count == 2 + + @pytest.mark.sqlite_only # Foreign keys are switched off for it, as only a migrated SQLite database can need. + @pytest.mark.parametrize(("user_id", "is_admin"), [("user1", False), ("admin", True)]) + def test_all_scope_excludes_dangling_board_associations( + self, + database: Database, + stores: Stores, + user_id: str, + is_admin: bool, + ) -> None: + image_store = stores[0] + _save(image_store, "dangling.png", user_id="user1") + # Production foreign keys prevent this state, but imported/legacy DBs + # may contain it. Seed it deliberately to lock down fail-closed reads. + # (A PRAGMA takes effect only outside a transaction, so it goes to the connection itself.) + database.sqlite.conn.execute("PRAGMA foreign_keys = OFF") + with sqlite_cursor(database) as cursor: + cursor.execute( + "INSERT INTO board_images (board_id, image_name) VALUES (?, ?)", ("deleted-board", "dangling.png") + ) + database.sqlite.conn.execute("PRAGMA foreign_keys = ON") + + dtos = image_store.get_many(limit=100, board_id="all", user_id=user_id, is_admin=is_admin) + names = image_store.get_image_names(board_id="all", user_id=user_id, is_admin=is_admin) + + assert dtos.items == [] + assert dtos.total == 0 + assert names.image_names == [] + assert names.total_count == names.starred_count == 0 + + +class TestCreatedAtRangeFiltering: + """get_many()/get_image_names() filter by inclusive created_from/created_to dates. + + Bounds are date-only strings compared lexicographically against the ISO text + column, which stores both space- and T-separated timestamps. + """ + + def _seed_dated(self, database: Database, store: ImageRecordStorage) -> None: + dated = { + "jan30.png": "2026-01-30 12:00:00.000", + "jan31-morning.png": "2026-01-31 08:30:00.000", + "jan31-last-second.png": "2026-01-31 23:59:59.999", + "feb01-midnight.png": "2026-02-01 00:00:00.000", + "feb15-t-sep.png": "2026-02-15T10:00:00.000", + } + for name, created_at in dated.items(): + _save(store, name) + _set_created_at(database, name, created_at) + + def test_created_from_is_inclusive(self, database: Database, store: ImageRecordStorage) -> None: + self._seed_dated(database, store) + + result = store.get_many(limit=10, created_from="2026-01-31") + + assert {r.image_name for r in result.items} == { + "jan31-morning.png", + "jan31-last-second.png", + "feb01-midnight.png", + "feb15-t-sep.png", + } + + def test_created_to_includes_end_of_day_and_excludes_next_midnight( + self, database: Database, store: ImageRecordStorage + ) -> None: + self._seed_dated(database, store) + + result = store.get_many(limit=10, created_to="2026-01-31") + + assert {r.image_name for r in result.items} == { + "jan30.png", + "jan31-morning.png", + "jan31-last-second.png", + } + + def test_created_to_handles_month_rollover(self, database: Database, store: ImageRecordStorage) -> None: + """created_to on the last day of a month must not lexicographically leak into the next month.""" + self._seed_dated(database, store) + + result = store.get_many(limit=10, created_from="2026-01-31", created_to="2026-01-31") + + assert {r.image_name for r in result.items} == {"jan31-morning.png", "jan31-last-second.png"} + + def test_range_matches_t_separated_timestamps(self, database: Database, store: ImageRecordStorage) -> None: + self._seed_dated(database, store) + + result = store.get_many(limit=10, created_from="2026-02-15", created_to="2026-02-15") + + assert {r.image_name for r in result.items} == {"feb15-t-sep.png"} + + def test_range_combines_with_search_term(self, database: Database, store: ImageRecordStorage) -> None: + self._seed_dated(database, store) + + result = store.get_many(limit=10, created_from="2026-01-31", created_to="2026-02-01", search_term="feb01") + + assert {r.image_name for r in result.items} == set() + # search_term matches metadata/created_at, not names; a created_at match works + result = store.get_many(limit=10, created_from="2026-01-31", created_to="2026-02-01", search_term="2026-02-01") + assert {r.image_name for r in result.items} == {"feb01-midnight.png"} + + def test_range_combines_with_board_filter(self, database: Database, stores: Stores) -> None: + image_store, board_store, board_image_store = stores + self._seed_dated(database, image_store) + board = board_store.save(board_name="Dated Board", user_id="user1") + board_image_store.add_image_to_board(board_id=board.board_id, image_name="jan30.png") + board_image_store.add_image_to_board(board_id=board.board_id, image_name="feb01-midnight.png") + + result = image_store.get_many(limit=10, board_id=board.board_id, created_from="2026-02-01") + + assert {r.image_name for r in result.items} == {"feb01-midnight.png"} + + def test_get_many_total_and_get_image_names_counts_are_consistent( + self, database: Database, store: ImageRecordStorage + ) -> None: + self._seed_dated(database, store) + store.update("jan31-morning.png", ImageRecordChanges(starred=True)) + + dtos = store.get_many(limit=10, created_from="2026-01-31", created_to="2026-02-01") + names = store.get_image_names(created_from="2026-01-31", created_to="2026-02-01") + + assert dtos.total == names.total_count == 3 + assert set(names.image_names) == {r.image_name for r in dtos.items} + assert names.starred_count == 1 + + def test_no_range_returns_everything(self, database: Database, store: ImageRecordStorage) -> None: + self._seed_dated(database, store) + + result = store.get_many(limit=10) + + assert result.total == 5 + + def test_names_by_date_are_the_days_images_that_are_not_intermediates( + self, database: Database, store: ImageRecordStorage + ) -> None: + self._seed_dated(database, store) + _save(store, "jan31-intermediate.png", is_intermediate=True) + _set_created_at(database, "jan31-intermediate.png", "2026-01-31 09:00:00.000") + + result = store.get_image_names_by_date("2026-01-31", starred_first=False, order_dir=SQLiteDirection.Ascending) + + assert result.image_names == ["jan31-morning.png", "jan31-last-second.png"] + # A date that is no ISO day matches nothing, as SQLite's DATE() did. + assert store.get_image_names_by_date("2026-1-31").image_names == [] + assert store.get_image_names_by_date("20260131").image_names == [] + assert store.get_image_names_by_date("2026-01-31 00:00").image_names == [] + assert store.get_image_names_by_date("").image_names == [] + + def test_the_last_day_there_is_matches_nothing_as_sqlite_did( + self, database: Database, store: ImageRecordStorage + ) -> None: + # Its next day cannot be represented; SQLite's DATE(day, '+1 day') gave NULL. + self._seed_dated(database, store) + + assert store.get_many(limit=10, created_to="9999-12-31").items == [] + assert store.get_image_names(created_to="9999-12-31").image_names == [] + + +class TestOrdering: + def test_images_of_the_same_moment_are_listed_by_name(self, database: Database, store: ImageRecordStorage) -> None: + # Saved against name order, so that only the tie-break lists them by name. + for name in ("c.png", "a.png", "b.png"): + _save(store, name) + _set_created_at(database, name, "2026-01-01 00:00:00.000") + store.update("b.png", ImageRecordChanges(starred=True)) + + def listed(starred_first: bool, order_dir: SQLiteDirection) -> list[str]: + names = store.get_image_names(starred_first=starred_first, order_dir=order_dir).image_names + page = [ + r.image_name for r in store.get_many(limit=10, starred_first=starred_first, order_dir=order_dir).items + ] + assert page == names + return names + + assert listed(False, SQLiteDirection.Ascending) == ["a.png", "b.png", "c.png"] + assert listed(False, SQLiteDirection.Descending) == ["c.png", "b.png", "a.png"] + assert listed(True, SQLiteDirection.Descending) == ["b.png", "c.png", "a.png"] + assert listed(True, SQLiteDirection.Ascending) == ["b.png", "a.png", "c.png"] + + def test_pages_walk_the_listing_and_a_page_of_none_still_counts( + self, database: Database, store: ImageRecordStorage + ) -> None: + for i, name in enumerate(("c.png", "a.png", "b.png")): + _save(store, name) + _set_created_at(database, name, f"2026-01-0{i + 1} 00:00:00.000") + names = store.get_image_names().image_names + + pages = [store.get_many(offset=offset, limit=1) for offset in range(3)] + count_only = store.get_many(limit=0) + + assert [[r.image_name for r in page.items] for page in pages] == [[name] for name in names] + assert {page.total for page in pages} == {3} + assert count_only.items == [] and count_only.total == 3 + + +@pytest.mark.sqlite_only +class TestGetImageNamesQueryPlans: + def test_default_category_query_has_no_forced_index(self, database: Database, store: ImageRecordStorage) -> None: + _save(store, "image.png", user_id="alice") + + result, statement, _ = _capture_names_plan( + database, + store, + categories=[ImageCategory.GENERAL], + is_intermediate=False, + is_admin=True, + ) + + assert result.image_names == ["image.png"] + assert "OUTER JOIN" not in statement + assert "INDEXED BY" not in statement + assert "NOT INDEXED" not in statement + + @pytest.mark.parametrize( + "kwargs", + [ + {"search_term": "does-not-match"}, + {"user_id": "alice", "is_admin": False}, + {"user_id": "alice", "is_admin": False, "starred_first": False}, + {"user_id": "alice", "is_admin": False, "order_dir": SQLiteDirection.Ascending}, + ], + ) + def test_query_shapes_do_not_force_indexes( + self, database: Database, store: ImageRecordStorage, kwargs: dict[str, Any] + ) -> None: + _save(store, "image.png", user_id="alice") + + _, statement, _ = _capture_names_plan( + database, + store, + categories=[ImageCategory.GENERAL], + is_intermediate=False, + **kwargs, + ) + + assert "OUTER JOIN" not in statement + assert "INDEXED BY" not in statement + assert "NOT INDEXED" not in statement + + def test_none_board_uses_anti_membership_filter(self, database: Database, stores: Stores) -> None: + image_store, board_store, board_image_store = stores + _save(image_store, "boarded.png", user_id="alice") + _save(image_store, "unboarded.png", user_id="alice") + board = board_store.save("Board", "alice") + board_image_store.add_image_to_board(board.board_id, "boarded.png") + + result, statement, _ = _capture_names_plan( + database, + image_store, + board_id="none", + user_id="alice", + is_admin=False, + starred_first=False, + ) + + assert result.image_names == ["unboarded.png"] + assert "NOT (EXISTS" in statement + assert "OUTER JOIN" not in statement + + @pytest.mark.parametrize("board_id", ["none", "all"]) + def test_membership_checks_search_their_indexes(self, database: Database, stores: Stores, board_id: str) -> None: + """A membership check that scanned a table would cost every listed image all memberships.""" + image_store, board_store, board_image_store = stores + _save(image_store, "boarded.png", user_id="alice") + _save(image_store, "unboarded.png", user_id="alice") + board = board_store.save("Board", "alice") + board_image_store.add_image_to_board(board.board_id, "boarded.png") + + _, _, plan = _capture_names_plan(database, image_store, board_id=board_id, user_id="alice", is_admin=False) + + assert any("board_images" in step for step in plan) + scans = [step for step in plan if step.startswith("SCAN")] + assert not [step for step in scans if any(table in step for table in ("board_images", "boards"))] + + def test_nonadmin_asset_query_has_no_forced_index(self, database: Database, store: ImageRecordStorage) -> None: + _save(store, "asset.png", user_id="alice", category=ImageCategory.CONTROL) + + result, statement, _ = _capture_names_plan( + database, + store, + categories=[ImageCategory.CONTROL], + is_intermediate=False, + user_id="alice", + is_admin=False, + ) + + assert result.image_names == ["asset.png"] + assert "INDEXED BY" not in statement + assert "NOT INDEXED" not in statement + + +class TestIterAllImageLocations: + def test_keyset_pages_keep_archived_and_intermediate_records_without_offset_skips( + self, store: ImageRecordStorage + ) -> None: + _save(store, "a.png", subfolder="nested", user_id="user-a") + _save(store, "b.png", subfolder="nested", is_intermediate=True, user_id="user-b") + _save(store, "c.png", subfolder="archived", user_id="user-a") + _save(store, "d.png", subfolder="", user_id="user-b") + + locations = iter(store.iter_all_image_locations(batch_size=2)) + assert next(locations) == ("a.png", "nested") + assert next(locations) == ("b.png", "nested") + store.delete_many(["a.png", "b.png"]) + + assert list(locations) == [("c.png", "archived"), ("d.png", "")] + + def test_iteration_rejects_nonpositive_batch_size(self, store: ImageRecordStorage) -> None: + with pytest.raises(ValueError, match="batch_size"): + list(store.iter_all_image_locations(batch_size=0)) diff --git a/tests/app/services/image_records/test_image_records_sqlite.py b/tests/app/services/image_records/test_image_records_sqlite.py deleted file mode 100644 index 8606429f617..00000000000 --- a/tests/app/services/image_records/test_image_records_sqlite.py +++ /dev/null @@ -1,790 +0,0 @@ -"""DB-backed tests for SqliteImageRecordStorage. - -Verifies that image_subfolder round-trips correctly through save(), get(), -get_many(), and get_subfolders() against a real (in-memory) SQLite database, -and that get_many()/get_image_names() enforce per-user ownership isolation. -""" - -import sqlite3 - -import pytest - -from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage -from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardVisibility -from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage -from invokeai.app.services.config.config_default import InvokeAIAppConfig -from invokeai.app.services.image_records.image_records_common import ( - ImageCategory, - ImageRecordChanges, - ImageRecordNotFoundException, - ResourceOrigin, -) -from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database - - -@pytest.fixture -def store() -> SqliteImageRecordStorage: - config = InvokeAIAppConfig(use_memory_db=True) - logger = InvokeAILogger.get_logger(config=config) - db = create_mock_sqlite_database(config, logger) - return SqliteImageRecordStorage(db=db) - - -@pytest.fixture -def stores() -> tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage]: - """Image, board, and board-image storages sharing one in-memory database.""" - config = InvokeAIAppConfig(use_memory_db=True) - logger = InvokeAILogger.get_logger(config=config) - db = create_mock_sqlite_database(config, logger) - return ( - SqliteImageRecordStorage(db=db), - SqliteBoardRecordStorage(db=db), - SqliteBoardImageRecordStorage(db=db), - ) - - -def _save( - store: SqliteImageRecordStorage, - name: str, - subfolder: str = "", - is_intermediate: bool = False, - user_id: str | None = None, - category: ImageCategory = ImageCategory.GENERAL, -) -> None: - store.save( - image_name=name, - image_origin=ResourceOrigin.INTERNAL, - image_category=category, - width=64, - height=64, - has_workflow=False, - is_intermediate=is_intermediate, - image_subfolder=subfolder, - user_id=user_id, - ) - - -def _capture_names_plan(store: SqliteImageRecordStorage, **kwargs): - statements: list[str] = [] - store._db._conn.set_trace_callback(statements.append) - try: - result = store.get_image_names(**kwargs) - finally: - store._db._conn.set_trace_callback(None) - statement = next(statement for statement in statements if "SELECT images.image_name" in statement) - details = [row[3] for row in store._db._conn.execute(f"EXPLAIN QUERY PLAN {statement}").fetchall()] - return result, statement, details - - -class TestImageSubfolderRoundTrip: - """save() -> get() preserves image_subfolder.""" - - def test_default_empty_subfolder(self, store: SqliteImageRecordStorage) -> None: - _save(store, "img_default.png") - record = store.get("img_default.png") - assert record.image_subfolder == "" - - def test_custom_subfolder(self, store: SqliteImageRecordStorage) -> None: - _save(store, "img_sub.png", subfolder="2026/04/11") - record = store.get("img_sub.png") - assert record.image_subfolder == "2026/04/11" - - def test_nested_subfolder(self, store: SqliteImageRecordStorage) -> None: - _save(store, "img_nested.png", subfolder="a/b/c/d") - record = store.get("img_nested.png") - assert record.image_subfolder == "a/b/c/d" - - -class TestGetManySubfolder: - """get_many() deserializes image_subfolder for every row.""" - - def test_get_many_returns_subfolders(self, store: SqliteImageRecordStorage) -> None: - _save(store, "flat.png", subfolder="") - _save(store, "dated.png", subfolder="2026/01") - _save(store, "hashed.png", subfolder="ab") - - result = store.get_many(limit=10, order_dir=SQLiteDirection.Ascending) - by_name = {r.image_name: r.image_subfolder for r in result.items} - - assert by_name["flat.png"] == "" - assert by_name["dated.png"] == "2026/01" - assert by_name["hashed.png"] == "ab" - - -class TestImageRecordExists: - def test_exists_returns_true_for_saved_image(self, store: SqliteImageRecordStorage) -> None: - _save(store, "exists.png") - - assert store.exists("exists.png") is True - - def test_exists_returns_false_for_missing_image(self, store: SqliteImageRecordStorage) -> None: - assert store.exists("missing.png") is False - - -def _intermediates(store: SqliteImageRecordStorage) -> list[tuple[str, str]]: - with store._db.transaction() as cursor: - cursor.execute("SELECT image_name, image_subfolder FROM images WHERE is_intermediate = TRUE;") - return [(row[0], row[1]) for row in cursor.fetchall()] - - -class TestGetSubfolders: - """get_subfolders() maps the named rows to their on-disk subfolders without touching them.""" - - def test_returns_subfolders_of_existing_rows_only(self, store: SqliteImageRecordStorage) -> None: - _save(store, "keep.png", subfolder="general", is_intermediate=False) - _save(store, "tmp1.png", subfolder="intermediate", is_intermediate=True) - - assert store.get_subfolders(["keep.png", "tmp1.png", "missing.png"]) == { - "keep.png": "general", - "tmp1.png": "intermediate", - } - assert store.get("tmp1.png").image_subfolder == "intermediate" - - def test_intermediates_are_deleted_via_delete_intermediates_by_names(self, store: SqliteImageRecordStorage) -> None: - _save(store, "tmp.png", subfolder="x", is_intermediate=True) - deleted = store.delete_intermediates_by_names([name for name, _ in _intermediates(store)]) - - assert deleted == ["tmp.png"] - with pytest.raises(ImageRecordNotFoundException): - store.get("tmp.png") - - -class TestIterAllImageLocations: - def test_keyset_pages_keep_archived_and_intermediate_records_without_offset_skips( - self, store: SqliteImageRecordStorage - ) -> None: - _save(store, "a.png", subfolder="nested", user_id="user-a") - _save(store, "b.png", subfolder="nested", is_intermediate=True, user_id="user-b") - _save(store, "c.png", subfolder="archived", user_id="user-a") - _save(store, "d.png", subfolder="", user_id="user-b") - - locations = iter(store.iter_all_image_locations(batch_size=2)) - assert next(locations) == ("a.png", "nested") - assert next(locations) == ("b.png", "nested") - store.delete_many(["a.png", "b.png"]) - - assert list(locations) == [("c.png", "archived"), ("d.png", "")] - - def test_iteration_rejects_nonpositive_batch_size(self, store: SqliteImageRecordStorage) -> None: - with pytest.raises(ValueError, match="batch_size"): - list(store.iter_all_image_locations(batch_size=0)) - - -class TestQueryFaultsAreNotNotFound: - """A failing query means the database is unavailable, not that the image is missing. - - Reporting a query fault as "not found" propagates all the way to the API, where it becomes a 404 - and tells the frontend to drop a live image from its cache. - """ - - def _break_the_images_table(self, store: SqliteImageRecordStorage) -> None: - store._db._conn.execute("ALTER TABLE images RENAME TO images_moved;") - - def test_get_raises_the_db_error_not_not_found(self, store: SqliteImageRecordStorage) -> None: - _save(store, "live.png") - self._break_the_images_table(store) - - with pytest.raises(sqlite3.Error): - store.get("live.png") - - def test_get_metadata_raises_the_db_error_not_not_found(self, store: SqliteImageRecordStorage) -> None: - _save(store, "live.png") - self._break_the_images_table(store) - - with pytest.raises(sqlite3.Error): - store.get_metadata("live.png") - - def test_missing_row_still_raises_not_found(self, store: SqliteImageRecordStorage) -> None: - """The genuine not-found path is untouched.""" - with pytest.raises(ImageRecordNotFoundException): - store.get("never-existed.png") - with pytest.raises(ImageRecordNotFoundException): - store.get_metadata("never-existed.png") - - -class TestDeleteIntermediatesByNames: - """delete_intermediates_by_names() deletes only rows that are still intermediates.""" - - def test_promoted_image_keeps_its_record(self, store: SqliteImageRecordStorage) -> None: - """An image promoted out of intermediate status after the snapshot must survive.""" - _save(store, "tmp.png", subfolder="x", is_intermediate=True) - _save(store, "promoted.png", subfolder="x", is_intermediate=True) - snapshot = [name for name, _ in _intermediates(store)] - assert set(snapshot) == {"tmp.png", "promoted.png"} - - # Simulate the race: the image stops being an intermediate between the snapshot and delete. - store.update("promoted.png", ImageRecordChanges(is_intermediate=False)) - - deleted = store.delete_intermediates_by_names(snapshot) - - assert deleted == ["tmp.png"] - # promoted.png is excluded from the returned names, so the caller never purges its files. - assert store.get("promoted.png").is_intermediate is False - with pytest.raises(ImageRecordNotFoundException): - store.get("tmp.png") - - def test_promotion_interleaved_inside_the_call_keeps_the_record(self, store: SqliteImageRecordStorage) -> None: - """The is_intermediate predicate must ride on the DELETE, not on a preceding SELECT. - - Python's legacy sqlite3 transaction control opens a transaction only before a write, so a - SELECT inside this method holds no read lock. A writer that promotes an image after that - SELECT but before the DELETE must still not lose its record. - """ - _save(store, "tmp.png", is_intermediate=True) - _save(store, "promoted.png", is_intermediate=True) - snapshot = [name for name, _ in _intermediates(store)] - - # Promote from inside the call, between the first SELECT and the DELETE. - real_execute = store._db._conn.execute - promoted = False - - def trace(statement: str) -> None: - nonlocal promoted - # The trace fires when a statement *begins*, so hooking the first SELECT would promote - # before that SELECT reads anything — indistinguishable from promoting up front. Hooking - # the DELETE puts the promotion after the SELECT has already seen the row as an - # intermediate, which is the interleaving that a SELECT-then-unconditional-DELETE - # implementation gets wrong. - if not promoted and statement.strip().upper().startswith("DELETE FROM IMAGES"): - promoted = True - real_execute("UPDATE images SET is_intermediate = 0 WHERE image_name = 'promoted.png'") - - store._db._conn.set_trace_callback(trace) - try: - deleted = store.delete_intermediates_by_names(snapshot) - finally: - store._db._conn.set_trace_callback(None) - - assert promoted, "the interleaved promotion never ran; the test proves nothing" - assert deleted == ["tmp.png"] - assert store.get("promoted.png").is_intermediate is False - - def test_unknown_and_empty_names_are_no_ops(self, store: SqliteImageRecordStorage) -> None: - _save(store, "keep.png", is_intermediate=False) - - assert store.delete_intermediates_by_names([]) == [] - # "gone.png" has no record at all and "keep.png" is not an intermediate, so neither is - # deleted or returned; keep.png must still be present afterwards. - assert store.delete_intermediates_by_names(["gone.png", "keep.png"]) == [] - assert store.get("keep.png").image_name == "keep.png" - - def test_more_names_than_sql_variable_limit(self, store: SqliteImageRecordStorage) -> None: - """Chunking must not lose rows: exercise a name list spanning several chunks.""" - chunk = SqliteImageRecordStorage._MAX_SQL_VARIABLES - names = [f"tmp{i:05d}.png" for i in range(chunk * 2 + 7)] - for name in names: - _save(store, name, is_intermediate=True) - # One image in the middle chunk is promoted and must survive. - survivor = names[chunk + 3] - store.update(survivor, ImageRecordChanges(is_intermediate=False)) - - deleted = store.delete_intermediates_by_names(names) - - assert set(deleted) == set(names) - {survivor} - assert survivor not in deleted - assert store.get(survivor).is_intermediate is False - assert _intermediates(store) == [] - - def test_chunking_stays_within_the_declared_variable_limit(self, store: SqliteImageRecordStorage) -> None: - """No statement may bind more parameters than the declared limit.""" - chunk = SqliteImageRecordStorage._MAX_SQL_VARIABLES - names = [f"tmp{i:05d}.png" for i in range(chunk * 2 + 7)] - for name in names: - _save(store, name, is_intermediate=True) - - # The trace callback reports statements with their parameters already expanded, so count the - # bound image names in each one rather than the placeholders. - widest = 0 - - def trace(statement: str) -> None: - nonlocal widest - if "images WHERE image_name IN (" in statement: - widest = max(widest, statement.count(".png")) - - store._db._conn.set_trace_callback(trace) - try: - store.delete_intermediates_by_names(names) - finally: - store._db._conn.set_trace_callback(None) - - # 999 is the SQLITE_MAX_VARIABLE_NUMBER default on builds older than 3.32. Asserting the - # literal rather than _MAX_SQL_VARIABLES keeps the test meaningful if that constant is raised. - assert 0 < widest <= 999 - - -class TestOwnershipFilteringOmittedBoard: - """get_many()/get_image_names() enforce per-user isolation when board_id is omitted. - - Without this, a non-admin could enumerate every user's images (including images - on other users' private boards) simply by omitting the board_id query parameter. - """ - - def _seed_two_users( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> str: - """user1: one image on a private board + one uncategorized. user2: one uncategorized.""" - image_store, board_store, board_image_store = stores - _save(image_store, "u1-boarded.png", user_id="user1") - _save(image_store, "u1-uncat.png", user_id="user1") - _save(image_store, "u2-uncat.png", user_id="user2") - board = board_store.save(board_name="User1 Private Board", user_id="user1") - board_image_store.add_image_to_board(board_id=board.board_id, image_name="u1-boarded.png") - return board.board_id - - def test_get_many_omitted_board_filters_by_owner( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> None: - self._seed_two_users(stores) - image_store = stores[0] - - result = image_store.get_many(limit=10, user_id="user2", is_admin=False) - - assert {r.image_name for r in result.items} == {"u2-uncat.png"} - assert result.total == 1 - - def test_get_many_omitted_board_owner_sees_boarded_and_uncategorized( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> None: - self._seed_two_users(stores) - image_store = stores[0] - - result = image_store.get_many(limit=10, user_id="user1", is_admin=False) - - assert {r.image_name for r in result.items} == {"u1-boarded.png", "u1-uncat.png"} - assert result.total == 2 - - def test_get_many_omitted_board_admin_sees_all( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> None: - self._seed_two_users(stores) - image_store = stores[0] - - result = image_store.get_many(limit=10, user_id="admin", is_admin=True) - - assert {r.image_name for r in result.items} == {"u1-boarded.png", "u1-uncat.png", "u2-uncat.png"} - assert result.total == 3 - - def test_get_many_omitted_board_single_user_mode_sees_all( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> None: - """user_id=None (single-user mode) applies no ownership filter.""" - self._seed_two_users(stores) - image_store = stores[0] - - result = image_store.get_many(limit=10, user_id=None, is_admin=False) - - assert result.total == 3 - - def test_get_many_none_board_still_filters_by_owner( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> None: - """board_id="none" (uncategorized) keeps its existing per-user isolation.""" - self._seed_two_users(stores) - image_store = stores[0] - - result = image_store.get_many(limit=10, board_id="none", user_id="user1", is_admin=False) - - assert {r.image_name for r in result.items} == {"u1-uncat.png"} - - def test_get_many_explicit_board_returns_board_contents( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> None: - """An explicit board_id lists that board's images; read access is the router's job.""" - board_id = self._seed_two_users(stores) - image_store = stores[0] - - result = image_store.get_many(limit=10, board_id=board_id, user_id="user1", is_admin=False) - - assert {r.image_name for r in result.items} == {"u1-boarded.png"} - - def test_get_image_names_omitted_board_filters_by_owner( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> None: - self._seed_two_users(stores) - image_store = stores[0] - - result = image_store.get_image_names(user_id="user2", is_admin=False) - - assert result.image_names == ["u2-uncat.png"] - assert result.total_count == 1 - - def test_get_image_names_omitted_board_admin_sees_all( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> None: - self._seed_two_users(stores) - image_store = stores[0] - - result = image_store.get_image_names(user_id="admin", is_admin=True) - - assert set(result.image_names) == {"u1-boarded.png", "u1-uncat.png", "u2-uncat.png"} - assert result.total_count == 3 - - def test_get_image_names_omitted_board_single_user_mode_sees_all( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> None: - self._seed_two_users(stores) - image_store = stores[0] - - result = image_store.get_image_names(user_id=None, is_admin=False) - - assert result.total_count == 3 - - def test_get_image_names_none_board_still_filters_by_owner( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> None: - self._seed_two_users(stores) - image_store = stores[0] - - result = image_store.get_image_names(board_id="none", user_id="user1", is_admin=False) - - assert result.image_names == ["u1-uncat.png"] - - -class TestAllReadableBoardsFiltering: - def _seed_visibility_matrix( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> None: - image_store, board_store, board_image_store = stores - - def add_boarded_image( - image_name: str, - *, - owner: str, - visibility: BoardVisibility = BoardVisibility.Private, - archived: bool = False, - ) -> str: - _save(image_store, image_name, user_id=owner) - board = board_store.save(board_name=image_name, user_id=owner) - board_store.update( - board.board_id, - BoardChanges(board_visibility=visibility, archived=archived), - ) - board_image_store.add_image_to_board(board_id=board.board_id, image_name=image_name) - return board.board_id - - add_boarded_image("own-private.png", owner="user1") - shared_private_id = add_boarded_image("explicit-share.png", owner="user2") - add_boarded_image("shared-visibility.png", owner="user2", visibility=BoardVisibility.Shared) - add_boarded_image("public-visibility.png", owner="user2", visibility=BoardVisibility.Public) - add_boarded_image("other-private.png", owner="user2") - add_boarded_image("own-archived.png", owner="user1", archived=True) - add_boarded_image("other-archived-public.png", owner="user2", visibility=BoardVisibility.Public, archived=True) - _save(image_store, "own-uncategorized.png", user_id="user1") - _save(image_store, "other-uncategorized.png", user_id="user2") - - with image_store._db.transaction() as cursor: - cursor.execute( - "INSERT OR IGNORE INTO users (user_id, email, password_hash) VALUES (?, ?, ?)", - ("user1", "user1@example.com", "unused"), - ) - cursor.execute( - "INSERT OR IGNORE INTO shared_boards (board_id, user_id) VALUES (?, ?)", - (shared_private_id, "user1"), - ) - - @pytest.mark.parametrize( - ("user_id", "is_admin", "expected"), - [ - ( - "user1", - False, - { - "own-private.png", - "explicit-share.png", - "shared-visibility.png", - "public-visibility.png", - "own-uncategorized.png", - }, - ), - ( - "admin", - True, - { - "own-private.png", - "explicit-share.png", - "shared-visibility.png", - "public-visibility.png", - "other-private.png", - "own-uncategorized.png", - "other-uncategorized.png", - }, - ), - ], - ) - def test_all_scope_authorization_and_counts_are_consistent( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - user_id: str, - is_admin: bool, - expected: set[str], - ) -> None: - self._seed_visibility_matrix(stores) - image_store = stores[0] - with image_store._db.transaction() as cursor: - cursor.executemany( - "UPDATE images SET starred = TRUE WHERE image_name = ?", - [("own-private.png",), ("other-private.png",)], - ) - - dtos = image_store.get_many(limit=100, board_id="all", user_id=user_id, is_admin=is_admin) - names = image_store.get_image_names(board_id="all", user_id=user_id, is_admin=is_admin) - - assert {image.image_name for image in dtos.items} == expected - assert set(names.image_names) == expected - assert dtos.total == names.total_count == len(expected) - assert names.starred_count == len(expected & {"own-private.png", "other-private.png"}) - - def test_all_scope_combines_with_inclusive_date_filters( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> None: - self._seed_visibility_matrix(stores) - image_store = stores[0] - _set_created_at(image_store, "own-private.png", "2026-06-01 00:00:00.000") - _set_created_at(image_store, "explicit-share.png", "2026-06-01 23:59:59.999") - _set_created_at(image_store, "public-visibility.png", "2026-06-02 00:00:00.000") - - dtos = image_store.get_many( - limit=100, - board_id="all", - user_id="user1", - created_from="2026-06-01", - created_to="2026-06-01", - ) - names = image_store.get_image_names( - board_id="all", - user_id="user1", - created_from="2026-06-01", - created_to="2026-06-01", - ) - - assert {image.image_name for image in dtos.items} == {"own-private.png", "explicit-share.png"} - assert set(names.image_names) == {"own-private.png", "explicit-share.png"} - assert dtos.total == names.total_count == 2 - - @pytest.mark.parametrize(("user_id", "is_admin"), [("user1", False), ("admin", True)]) - def test_all_scope_excludes_dangling_board_associations( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - user_id: str, - is_admin: bool, - ) -> None: - image_store = stores[0] - _save(image_store, "dangling.png", user_id="user1") - # Production foreign keys prevent this state, but imported/legacy DBs - # may contain it. Seed it deliberately to lock down fail-closed reads. - image_store._db._conn.execute("PRAGMA foreign_keys = OFF") - image_store._db._conn.execute( - "INSERT INTO board_images (board_id, image_name) VALUES (?, ?)", - ("deleted-board", "dangling.png"), - ) - image_store._db._conn.commit() - image_store._db._conn.execute("PRAGMA foreign_keys = ON") - - dtos = image_store.get_many(limit=100, board_id="all", user_id=user_id, is_admin=is_admin) - names = image_store.get_image_names(board_id="all", user_id=user_id, is_admin=is_admin) - - assert dtos.items == [] - assert dtos.total == 0 - assert names.image_names == [] - assert names.total_count == names.starred_count == 0 - - -def _set_created_at(store: SqliteImageRecordStorage, name: str, created_at: str) -> None: - """created_at is written by a SQL column default; tests override it directly.""" - with store._db.transaction() as cursor: - cursor.execute("UPDATE images SET created_at = ? WHERE image_name = ?", (created_at, name)) - - -class TestCreatedAtRangeFiltering: - """get_many()/get_image_names() filter by inclusive created_from/created_to dates. - - Bounds are date-only strings compared lexicographically against the ISO text - column, which stores both space- and T-separated timestamps. - """ - - def _seed_dated(self, store: SqliteImageRecordStorage) -> None: - _save(store, "jan30.png") - _save(store, "jan31-morning.png") - _save(store, "jan31-last-second.png") - _save(store, "feb01-midnight.png") - _save(store, "feb15-t-sep.png") - _set_created_at(store, "jan30.png", "2026-01-30 12:00:00.000") - _set_created_at(store, "jan31-morning.png", "2026-01-31 08:30:00.000") - _set_created_at(store, "jan31-last-second.png", "2026-01-31 23:59:59.999") - _set_created_at(store, "feb01-midnight.png", "2026-02-01 00:00:00.000") - _set_created_at(store, "feb15-t-sep.png", "2026-02-15T10:00:00.000") - - def test_created_from_is_inclusive(self, store: SqliteImageRecordStorage) -> None: - self._seed_dated(store) - - result = store.get_many(limit=10, created_from="2026-01-31") - - assert {r.image_name for r in result.items} == { - "jan31-morning.png", - "jan31-last-second.png", - "feb01-midnight.png", - "feb15-t-sep.png", - } - - def test_created_to_includes_end_of_day_and_excludes_next_midnight(self, store: SqliteImageRecordStorage) -> None: - self._seed_dated(store) - - result = store.get_many(limit=10, created_to="2026-01-31") - - assert {r.image_name for r in result.items} == { - "jan30.png", - "jan31-morning.png", - "jan31-last-second.png", - } - - def test_created_to_handles_month_rollover(self, store: SqliteImageRecordStorage) -> None: - """created_to on the last day of a month must not lexicographically leak into the next month.""" - self._seed_dated(store) - - result = store.get_many(limit=10, created_from="2026-01-31", created_to="2026-01-31") - - assert {r.image_name for r in result.items} == {"jan31-morning.png", "jan31-last-second.png"} - - def test_range_matches_t_separated_timestamps(self, store: SqliteImageRecordStorage) -> None: - self._seed_dated(store) - - result = store.get_many(limit=10, created_from="2026-02-15", created_to="2026-02-15") - - assert {r.image_name for r in result.items} == {"feb15-t-sep.png"} - - def test_range_combines_with_search_term(self, store: SqliteImageRecordStorage) -> None: - self._seed_dated(store) - - result = store.get_many(limit=10, created_from="2026-01-31", created_to="2026-02-01", search_term="feb01") - - assert {r.image_name for r in result.items} == set() - # search_term matches metadata/created_at, not names; a created_at match works - result = store.get_many(limit=10, created_from="2026-01-31", created_to="2026-02-01", search_term="2026-02-01") - assert {r.image_name for r in result.items} == {"feb01-midnight.png"} - - def test_range_combines_with_board_filter( - self, - stores: tuple[SqliteImageRecordStorage, SqliteBoardRecordStorage, SqliteBoardImageRecordStorage], - ) -> None: - image_store, board_store, board_image_store = stores - self._seed_dated(image_store) - board = board_store.save(board_name="Dated Board", user_id="user1") - board_image_store.add_image_to_board(board_id=board.board_id, image_name="jan30.png") - board_image_store.add_image_to_board(board_id=board.board_id, image_name="feb01-midnight.png") - - result = image_store.get_many(limit=10, board_id=board.board_id, created_from="2026-02-01") - - assert {r.image_name for r in result.items} == {"feb01-midnight.png"} - - def test_get_many_total_and_get_image_names_counts_are_consistent(self, store: SqliteImageRecordStorage) -> None: - self._seed_dated(store) - with store._db.transaction() as cursor: - cursor.execute("UPDATE images SET starred = TRUE WHERE image_name = ?", ("jan31-morning.png",)) - - dtos = store.get_many(limit=10, created_from="2026-01-31", created_to="2026-02-01") - names = store.get_image_names(created_from="2026-01-31", created_to="2026-02-01") - - assert dtos.total == names.total_count == 3 - assert set(names.image_names) == {r.image_name for r in dtos.items} - assert names.starred_count == 1 - - def test_no_range_returns_everything(self, store: SqliteImageRecordStorage) -> None: - self._seed_dated(store) - - result = store.get_many(limit=10) - - assert result.total == 5 - - -class TestGetImageNamesQueryPlans: - def test_default_category_query_has_no_forced_index(self, store: SqliteImageRecordStorage) -> None: - _save(store, "image.png", user_id="alice") - - result, statement, _ = _capture_names_plan( - store, - categories=[ImageCategory.GENERAL], - is_intermediate=False, - is_admin=True, - ) - - assert result.image_names == ["image.png"] - assert "LEFT JOIN board_images" not in statement - assert "INDEXED BY" not in statement - assert "NOT INDEXED" not in statement - - @pytest.mark.parametrize( - "kwargs", - [ - {"search_term": "does-not-match"}, - {"user_id": "alice", "is_admin": False}, - {"user_id": "alice", "is_admin": False, "starred_first": False}, - {"user_id": "alice", "is_admin": False, "order_dir": SQLiteDirection.Ascending}, - ], - ) - def test_query_shapes_do_not_force_indexes( - self, - store: SqliteImageRecordStorage, - kwargs, - ) -> None: - _save(store, "image.png", user_id="alice") - - _, statement, _ = _capture_names_plan( - store, - categories=[ImageCategory.GENERAL], - is_intermediate=False, - **kwargs, - ) - - assert "LEFT JOIN board_images" not in statement - assert "INDEXED BY" not in statement - assert "NOT INDEXED" not in statement - - def test_none_board_uses_anti_membership_filter(self, stores) -> None: - image_store, board_store, board_image_store = stores - _save(image_store, "boarded.png", user_id="alice") - _save(image_store, "unboarded.png", user_id="alice") - board = board_store.save("Board", "alice") - board_image_store.add_image_to_board(board.board_id, "boarded.png") - - result, statement, _ = _capture_names_plan( - image_store, - board_id="none", - user_id="alice", - is_admin=False, - starred_first=False, - ) - - assert result.image_names == ["unboarded.png"] - assert "NOT EXISTS" in statement - assert "LEFT JOIN board_images" not in statement - - def test_nonadmin_asset_query_has_no_forced_index(self, store: SqliteImageRecordStorage) -> None: - _save(store, "asset.png", user_id="alice", category=ImageCategory.CONTROL) - - result, statement, _ = _capture_names_plan( - store, - categories=[ImageCategory.CONTROL], - is_intermediate=False, - user_id="alice", - is_admin=False, - ) - - assert result.image_names == ["asset.png"] - assert "INDEXED BY" not in statement - assert "NOT INDEXED" not in statement diff --git a/tests/app/services/images/test_images_default.py b/tests/app/services/images/test_images_default.py index ccb3b10ec7d..1a4424d9fd4 100644 --- a/tests/app/services/images/test_images_default.py +++ b/tests/app/services/images/test_images_default.py @@ -26,13 +26,13 @@ ImageCategory, ImageRecord, ImageRecordChanges, - ImageRecordDeleteException, ImageRecordNotFoundException, ResourceOrigin, ) -from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage from invokeai.app.services.images.images_default import ImageService -from invokeai.app.services.shared.sqlite.sqlite_util import init_db +from invokeai.app.services.shared.database.errors import LockTimeoutError +from invokeai.app.services.shared.database.startup import init_database from invokeai.app.util.misc import get_iso_timestamp from invokeai.backend.util.logging import InvokeAILogger from tests.fixtures.sqlite_database import create_mock_sqlite_database @@ -78,7 +78,7 @@ def _make_record( @pytest.fixture -def real_image_service(tmp_path: Path) -> tuple[ImageService, SqliteImageRecordStorage, DiskImageFileStorage]: +def real_image_service(tmp_path: Path) -> tuple[ImageService, ImageRecordStorage, DiskImageFileStorage]: logger = InvokeAILogger.get_logger() config = InvokeAIAppConfig(use_memory_db=True, image_subfolder_strategy="flat") config._root = tmp_path @@ -86,8 +86,8 @@ def real_image_service(tmp_path: Path) -> tuple[ImageService, SqliteImageRecordS invoker = MagicMock() invoker.services.configuration.pil_compress_level = 6 storage.start(invoker) - db = init_db(config=config, logger=logger, image_files=storage) - records = SqliteImageRecordStorage(db=db) + db = init_database(config=config, logger=logger, image_files=storage) + records = ImageRecordStorage(db) invoker.services.configuration.image_subfolder_strategy = "flat" invoker.services.names.create_image_name.return_value = "uploaded.png" @@ -103,7 +103,7 @@ def real_image_service(tmp_path: Path) -> tuple[ImageService, SqliteImageRecordS def test_create_rolls_back_record_and_files_when_thumbnail_save_fails( - real_image_service: tuple[ImageService, SqliteImageRecordStorage, DiskImageFileStorage], + real_image_service: tuple[ImageService, ImageRecordStorage, DiskImageFileStorage], ) -> None: service, records, storage = real_image_service image = Image.new("RGB", (32, 32), "red") @@ -131,7 +131,7 @@ def test_create_rolls_back_record_and_files_when_thumbnail_save_fails( def test_create_accepts_large_16_bit_image( - real_image_service: tuple[ImageService, SqliteImageRecordStorage, DiskImageFileStorage], + real_image_service: tuple[ImageService, ImageRecordStorage, DiskImageFileStorage], ) -> None: service, records, storage = real_image_service image = Image.new("I;16", (1024, 1024), 32768) @@ -419,11 +419,11 @@ def fsync(fd: int) -> None: @pytest.fixture -def wired(tmp_path: Path) -> tuple[ImageService, SqliteImageRecordStorage, DiskImageFileStorage]: +def wired(tmp_path: Path) -> tuple[ImageService, ImageRecordStorage, DiskImageFileStorage]: """ImageService wired to a real record store and a real disk store — no stub decides anything.""" config = InvokeAIAppConfig(use_memory_db=True) logger = InvokeAILogger.get_logger(config=config) - records = SqliteImageRecordStorage(db=create_mock_sqlite_database(config, logger)) + records = ImageRecordStorage(create_mock_sqlite_database(config, logger)) storage = DiskImageFileStorage(tmp_path / "outputs") svc = ImageService() @@ -440,14 +440,14 @@ def wired(tmp_path: Path) -> tuple[ImageService, SqliteImageRecordStorage, DiskI @pytest.fixture def wired_with_move_service( tmp_path: Path, -) -> tuple[ImageService, SqliteImageRecordStorage, DiskImageFileStorage, ImageMoveService]: +) -> tuple[ImageService, ImageRecordStorage, DiskImageFileStorage, ImageMoveService]: """ImageService and ImageMoveService wired to one real db and disk store, as production does.""" config = InvokeAIAppConfig(use_memory_db=True, image_subfolder_strategy="flat") logger = InvokeAILogger.get_logger(config=config) db = create_mock_sqlite_database(config, logger) - records = SqliteImageRecordStorage(db=db) + records = ImageRecordStorage(db) storage = DiskImageFileStorage(tmp_path / "outputs") - moves = ImageMoveService(db=db, image_files=storage, config=config, logger=logger) + moves = ImageMoveService(db, image_files=storage, config=config, logger=logger) svc = ImageService() invoker = MagicMock() @@ -463,7 +463,7 @@ def wired_with_move_service( return svc, records, storage, moves -def _seed_record_at_subfolder(records: SqliteImageRecordStorage, name: str, image_subfolder: str) -> None: +def _seed_record_at_subfolder(records: ImageRecordStorage, name: str, image_subfolder: str) -> None: records.save( image_name=name, image_origin=ResourceOrigin.INTERNAL, @@ -476,7 +476,7 @@ def _seed_record_at_subfolder(records: SqliteImageRecordStorage, name: str, imag ) -def _seed_record(records: SqliteImageRecordStorage, name: str, is_intermediate: bool = True) -> None: +def _seed_record(records: ImageRecordStorage, name: str, is_intermediate: bool = True) -> None: records.save( image_name=name, image_origin=ResourceOrigin.INTERNAL, @@ -488,7 +488,7 @@ def _seed_record(records: SqliteImageRecordStorage, name: str, is_intermediate: ) -def _restart_file_storage(storage: DiskImageFileStorage, records: SqliteImageRecordStorage) -> DiskImageFileStorage: +def _restart_file_storage(storage: DiskImageFileStorage, records: ImageRecordStorage) -> DiskImageFileStorage: """Simulates a restart over the same output folder, running delete-journal recovery.""" invoker = MagicMock() invoker.services.image_records = records @@ -505,7 +505,7 @@ class TestDeleteTransactional: """delete() journals its intent, deletes the record, then purges — never losing files on failure.""" def test_notify_deleted_is_public_and_callbacks_observe_removed_record( - self, real_image_service: tuple[ImageService, SqliteImageRecordStorage, DiskImageFileStorage] + self, real_image_service: tuple[ImageService, ImageRecordStorage, DiskImageFileStorage] ) -> None: service, records, _storage = real_image_service _seed_record(records, "maintenance-removed.png", is_intermediate=False) @@ -557,11 +557,11 @@ def test_delete_db_failure_leaves_files_and_raises(self, disk_image_service: Ima storage = invoker.services.image_files _save_image_file(storage, "img.png") invoker.services.image_records.get.return_value = _make_record(image_name="img.png") - invoker.services.image_records.delete.side_effect = ImageRecordDeleteException() + invoker.services.image_records.delete.side_effect = LockTimeoutError("database is locked") deleted_callbacks: list[str] = [] disk_image_service.on_deleted(deleted_callbacks.append) - with pytest.raises(ImageRecordDeleteException): + with pytest.raises(LockTimeoutError): disk_image_service.delete("img.png") # Nothing was moved, so the image and its thumbnail are still exactly where they were. @@ -572,12 +572,12 @@ def test_delete_db_failure_leaves_files_and_raises(self, disk_image_service: Ima def test_delete_journal_cleanup_failure_still_raises_db_error(self, image_service: ImageService): invoker = image_service._ImageService__invoker # type: ignore - invoker.services.image_records.delete.side_effect = ImageRecordDeleteException() + invoker.services.image_records.delete.side_effect = LockTimeoutError("database is locked") invoker.services.image_files.abandon_delete.side_effect = ImageFileDeleteException("journal locked") deleted_callbacks: list[str] = [] image_service.on_deleted(deleted_callbacks.append) - with pytest.raises(ImageRecordDeleteException): + with pytest.raises(LockTimeoutError): image_service.delete("test.png") invoker.services.image_files.abandon_delete.assert_called_once_with( @@ -703,11 +703,13 @@ def test_file_purge_failure_is_logged_and_does_not_raise(self, image_service: Im def test_db_failure_raises_and_purges_nothing(self, image_service: ImageService): invoker = image_service._ImageService__invoker # type: ignore invoker.services.image_records.get_subfolders.return_value = {"tmp1.png": "", "tmp2.png": ""} - invoker.services.image_records.delete_intermediates_by_names.side_effect = ImageRecordDeleteException() + invoker.services.image_records.delete_intermediates_by_names.side_effect = LockTimeoutError( + "database is locked" + ) deleted_callbacks: list[str] = [] image_service.on_deleted(deleted_callbacks.append) - with pytest.raises(ImageRecordDeleteException): + with pytest.raises(LockTimeoutError): image_service.delete_intermediates_by_names(["tmp1.png", "tmp2.png"]) # No record was removed, so no file may be purged and the journal must be discarded. @@ -754,13 +756,13 @@ class TestDeleteIntermediatesAgainstRealRecords: and where the concurrency hazards JPPhoto reported would surface. """ - def _seed(self, records: SqliteImageRecordStorage, storage: DiskImageFileStorage, name: str) -> None: + def _seed(self, records: ImageRecordStorage, storage: DiskImageFileStorage, name: str) -> None: _seed_record(records, name) _save_image_file(storage, name) def _promote_after_snapshot( self, - records: SqliteImageRecordStorage, + records: ImageRecordStorage, monkeypatch, image_name: str, ) -> None: @@ -898,13 +900,13 @@ def competing_delete_then_fail(image_name: str) -> None: # this delete is still in flight, and only then does this one's own delete fail. real_delete(image_name) storage.delete(image_name) - raise ImageRecordDeleteException() + raise LockTimeoutError("database is locked") monkeypatch.setattr(records, "delete", competing_delete_then_fail) deleted_callbacks: list[str] = [] svc.on_deleted(deleted_callbacks.append) - with pytest.raises(ImageRecordDeleteException): + with pytest.raises(LockTimeoutError): svc.delete("img.png") assert not storage.get_path("img.png").exists() @@ -918,11 +920,11 @@ def test_a_database_failure_leaves_the_image_completely_intact(self, wired, monk _save_image_file(storage, "img.png") def failing_delete(image_name: str) -> None: - raise ImageRecordDeleteException() + raise LockTimeoutError("database is locked") monkeypatch.setattr(records, "delete", failing_delete) - with pytest.raises(ImageRecordDeleteException): + with pytest.raises(LockTimeoutError): svc.delete("img.png") assert storage.get_path("img.png").exists() @@ -1037,7 +1039,7 @@ def test_a_surviving_record_keeps_its_files(self, image_service: ImageService) - """If the record cannot be deleted the image is still referenced, so nothing may be purged.""" invoker = image_service._ImageService__invoker # type: ignore invoker.services.image_files.save.side_effect = ImageFileSaveException() - invoker.services.image_records.delete.side_effect = ImageRecordDeleteException() + invoker.services.image_records.delete.side_effect = LockTimeoutError("database is locked") with pytest.raises(ImageFileSaveException): image_service.create( @@ -1074,7 +1076,7 @@ def test_the_record_is_kept_when_the_cleanup_cannot_be_journaled(self, image_ser @pytest.mark.skipif(os.name == "nt", reason="Windows cannot open a directory for fsync") def test_a_leftover_file_is_never_orphaned_without_a_journal( self, - real_image_service: tuple[ImageService, SqliteImageRecordStorage, DiskImageFileStorage], + real_image_service: tuple[ImageService, ImageRecordStorage, DiskImageFileStorage], ) -> None: """Real storage, real records: the save leaves a file behind and the disk refuses to make the cleanup journal durable. The record must survive, and an ordinary delete once the disk @@ -1167,7 +1169,7 @@ class TestDeleteVersusSubfolderMove: def test_move_cannot_interleave_with_delete_intermediates( self, - wired_with_move_service: tuple[ImageService, SqliteImageRecordStorage, DiskImageFileStorage, ImageMoveService], + wired_with_move_service: tuple[ImageService, ImageRecordStorage, DiskImageFileStorage, ImageMoveService], monkeypatch, ) -> None: """While delete_intermediates_by_names() runs its unit, a concurrent move_all_images() must wait. @@ -1224,7 +1226,7 @@ def delete_after_pause(names: list[str], guard=None) -> list[str]: def test_move_holds_the_lock_across_its_batch_cycle( self, - wired_with_move_service: tuple[ImageService, SqliteImageRecordStorage, DiskImageFileStorage, ImageMoveService], + wired_with_move_service: tuple[ImageService, ImageRecordStorage, DiskImageFileStorage, ImageMoveService], ) -> None: """While a move unit runs, a concurrent delete must wait for it rather than race it. @@ -1256,7 +1258,7 @@ def test_move_holds_the_lock_across_its_batch_cycle( def test_failed_save_cleanup_purges_the_relocated_subfolder_too( self, - wired_with_move_service: tuple[ImageService, SqliteImageRecordStorage, DiskImageFileStorage, ImageMoveService], + wired_with_move_service: tuple[ImageService, ImageRecordStorage, DiskImageFileStorage, ImageMoveService], ) -> None: """A failed save whose record was relocated mid-save must not strand the moved files. @@ -1279,8 +1281,7 @@ def test_failed_save_cleanup_purges_the_relocated_subfolder_too( new_thumbnail.parent.mkdir(parents=True, exist_ok=True) old_image.replace(new_image) old_thumbnail.replace(new_thumbnail) - with records._db.transaction() as cursor: - cursor.execute("UPDATE images SET image_subfolder = '' WHERE image_name = 'failed.png';") + moves._queries.image_moves.repoint_image("failed.png", "old", "") assert records.get("failed.png").image_subfolder == "" # The failed save's cleanup still holds the subfolder the save captured. @@ -1296,7 +1297,7 @@ def test_failed_save_cleanup_purges_the_relocated_subfolder_too( def test_move_cannot_interleave_with_create( self, - wired_with_move_service: tuple[ImageService, SqliteImageRecordStorage, DiskImageFileStorage, ImageMoveService], + wired_with_move_service: tuple[ImageService, ImageRecordStorage, DiskImageFileStorage, ImageMoveService], monkeypatch, ) -> None: """While create() runs its save unit, a concurrent move_all_images() must wait. diff --git a/tests/app/services/intermediates/test_intermediates_records.py b/tests/app/services/intermediates/test_intermediates_records.py new file mode 100644 index 00000000000..ca2714be89c --- /dev/null +++ b/tests/app/services/intermediates/test_intermediates_records.py @@ -0,0 +1,716 @@ +"""The intermediates storage on every database backend: what protects an intermediate from a cleanup, how a +cleanup pages and checks, and the state the cleanup keeps in tables. + +The service tests (`test_intermediates_service.py`) drive the same storage through the queue and the file +storages, on SQLite only until the queue is ported. +""" + +import itertools +from collections.abc import Callable, Iterator +from datetime import datetime, timedelta, timezone +from typing import Any, Optional + +import pytest +from sqlalchemy import insert, select, update + +from invokeai.app.services.client_state_persistence.client_state_persistence_default import ( + ClientStatePersistence, +) +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage +from invokeai.app.services.intermediates import intermediates_records_default +from invokeai.app.services.intermediates.intermediates_common import BROWSER_HOLD_TTL_SECONDS +from invokeai.app.services.intermediates.intermediates_records_default import ( + IntermediatesRecords, + ReferenceOwner, +) +from invokeai.app.services.project_records.project_records_default import ProjectRecordsStorage +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.locks import DatabaseLock +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.intermediates import ( + intermediates_browser_holds, + intermediates_session_holds, + intermediates_unmeasurable, +) +from invokeai.app.services.shared.database.schema.session_queue import session_queue +from invokeai.app.services.shared.database.schema.users import users +from invokeai.app.services.shared.database.schema.videos import videos +from invokeai.app.services.shared.media_references import MediaReferences +from invokeai.app.services.video_records.video_records_default import VideoRecordStorage +from invokeai.app.services.workflow_records.workflow_records_common import ( + Workflow, + WorkflowCategory, + WorkflowMeta, + WorkflowWithoutID, +) +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage +from tests.fixtures.races import while_in_flight + +OLD = "2020-01-01 00:00:00.000" +item_ids = itertools.count(1000) + + +@pytest.fixture +def records(database: Database) -> IntermediatesRecords: + return IntermediatesRecords(database) + + +@pytest.fixture +def accounts(database: Database) -> None: + with database.begin(write=True) as conn: + conn.execute( + insert(users), + [{"user_id": user, "email": f"{user}@example.com", "password_hash": "-"} for user in ("alice", "bob")], + ) + + +def _image( + database: Database, + name: str, + *, + user_id: str = "alice", + created_at: str = OLD, + session_id: Optional[str] = None, + file_size_bytes: Optional[int] = 10, + is_intermediate: bool = True, +) -> None: + with database.begin(write=True) as conn: + conn.execute( + insert(images).values( + image_name=name, + image_origin="internal", + image_category="general", + width=8, + height=8, + is_intermediate=is_intermediate, + user_id=user_id, + created_at=created_at, + session_id=session_id, + file_size_bytes=file_size_bytes, + ) + ) + + +def _queue_item( + database: Database, + session_id: str, + *, + status: str = "pending", + session: str = "{}", + root_item_id: Optional[int] = None, +) -> int: + item_id = next(item_ids) + with database.begin(write=True) as conn: + conn.execute( + insert(session_queue).values( + item_id=item_id, + batch_id="batch", + queue_id="default", + session_id=session_id, + session=session, + status=status, + root_item_id=root_item_id, + ) + ) + return item_id + + +def _reference(database: Database, name: str, *, user_id: str = "alice", owner_id: str = "p1") -> None: + database.queries.media_references.replace( + owner_kind="project", user_id=user_id, owner_id=owner_id, references=MediaReferences(images={name}) + ) + + +def _deletable( + records: IntermediatesRecords, *, mode: str = "safe", caller: str = "admin", admin: bool = True +) -> set[str]: + batches = records.iter_deletable_batches( + "image", + user_id=None, + targets=None, + mode=mode, # type: ignore[arg-type] + is_admin=admin, + caller_user_id=caller, + recent_cutoff=None, + limit=50, + ) + return {name for batch in batches for name, _ in batch} + + +def _guarded(database: Database, records: IntermediatesRecords, names: list[str], **guard: Any) -> list[str]: + settings = {"mode": "safe", "allowed_user_ids": None, "caller_user_id": "admin", "is_admin": True, **guard} + return ImageRecordStorage(database).delete_intermediates_by_names( + names, records.make_delete_guard("image", **settings) + ) + + +def _exists(database: Database, name: str) -> bool: + with database.begin(write=False) as conn: + return conn.execute(select(images.c.image_name).where(images.c.image_name == name)).first() is not None + + +def _protect_by_queue_input(database: Database, records: IntermediatesRecords) -> None: + _queue_item(database, "s-input", session='{"node": {"image": {"image_name": "protected.png"}}}') + + +def _protect_by_producing_session(database: Database, records: IntermediatesRecords) -> None: + with database.begin(write=True) as conn: + conn.execute(update(images).where(images.c.image_name == "protected.png").values(session_id="s-producer")) + _queue_item(database, "s-producer", status="in_progress") + + +def _protect_by_completed_child_of_an_active_root(database: Database, records: IntermediatesRecords) -> None: + root = _queue_item(database, "s-root", status="waiting") + _queue_item( + database, + "s-child", + status="completed", + root_item_id=root, + session='{"out": {"image_name": "protected.png"}}', + ) + + +def _protect_as_output_of_a_completed_child_of_an_active_root( + database: Database, records: IntermediatesRecords +) -> None: + root = _queue_item(database, "s-root", status="waiting") + _queue_item(database, "s-child", status="completed", root_item_id=root) + with database.begin(write=True) as conn: + conn.execute(update(images).where(images.c.image_name == "protected.png").values(session_id="s-child")) + + +def _protect_by_browser_lease(database: Database, records: IntermediatesRecords) -> None: + records.replace_browser_hold("alice", "tab-1", ["protected.png"], []) + + +def _protect_by_cached_output(database: Database, records: IntermediatesRecords) -> None: + _queue_item(database, "s-reuser", status="in_progress") + assert records.hold_cached_media("s-reuser", MediaReferences(images={"protected.png"})) + + +def _protect_by_recency(database: Database, records: IntermediatesRecords) -> None: + now = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S.000") + with database.begin(write=True) as conn: + conn.execute(update(images).where(images.c.image_name == "protected.png").values(created_at=now)) + + +def _protect_by_reference(database: Database, records: IntermediatesRecords) -> None: + _reference(database, "protected.png") + + +@pytest.mark.parametrize( + "protect", + [ + _protect_by_queue_input, + _protect_by_producing_session, + _protect_by_completed_child_of_an_active_root, + _protect_as_output_of_a_completed_child_of_an_active_root, + _protect_by_browser_lease, + _protect_by_cached_output, + _protect_by_recency, + _protect_by_reference, + ], +) +def test_a_protected_intermediate_is_neither_paged_nor_deleted_by_a_safe_cleanup( + database: Database, records: IntermediatesRecords, protect: Callable[[Database, IntermediatesRecords], None] +) -> None: + _image(database, "protected.png") + _image(database, "safe.png") + _image(database, "durable.png", is_intermediate=False) + protect(database, records) + + assert _deletable(records) == {"safe.png"} + assert _guarded(database, records, ["protected.png", "safe.png", "durable.png"]) == ["safe.png"] + assert _exists(database, "protected.png") and _exists(database, "durable.png") + + +def test_the_summary_counts_each_classification_with_its_bytes( + database: Database, records: IntermediatesRecords +) -> None: + _image(database, "safe.png", file_size_bytes=5) + _image(database, "unmeasured.png", file_size_bytes=None) + _image(database, "referenced.png", file_size_bytes=7) + _reference(database, "referenced.png") + _image(database, "active.png") + records.replace_browser_hold("alice", "tab-1", ["active.png"], []) + _image(database, "bob.png", user_id="bob") + + rows = records.summarize(None) + alice = rows[("alice", None)]["image"] + only_bob = records.summarize(["bob"]) + + assert (alice.counts.safe, alice.counts.referenced, alice.counts.active, alice.counts.recent) == (2, 1, 1, 0) + assert (alice.safe_bytes, alice.referenced_bytes, alice.unknown_size_count) == (5, 7, 1) + assert list(only_bob) == [("bob", None)] and only_bob[("bob", None)]["image"].counts.safe == 1 + assert records.summarize([]) == {} + + +def test_windows_page_in_creation_order_with_names_breaking_ties( + database: Database, records: IntermediatesRecords +) -> None: + for name, created_at in ( + ("c.png", "2020-01-01 00:00:00.000"), + ("a.png", "2020-01-02 00:00:00.000"), + ("b.png", "2020-01-02 00:00:00.000"), + ("d.png", "2020-01-03 00:00:00.000"), + ("e.png", "2020-01-03 00:00:00.000"), + ): + _image(database, name, created_at=created_at) + + batches = list( + records.iter_deletable_batches( + "image", + user_id="alice", + targets=None, + mode="safe", + is_admin=False, + caller_user_id="alice", + recent_cutoff=None, + limit=2, + ) + ) + + assert [[name for name, _ in batch] for batch in batches] == [["c.png", "a.png"], ["b.png", "d.png"], ["e.png"]] + + +def test_a_force_cleanup_breaks_only_documents_its_caller_may_break( + database: Database, records: IntermediatesRecords +) -> None: + _image(database, "own-reference.png") + _reference(database, "own-reference.png", user_id="alice", owner_id="mine") + _image(database, "bobs-reference.png") + _reference(database, "bobs-reference.png", user_id="bob", owner_id="his") + + _image(database, "shared-reference.png") + _reference(database, "shared-reference.png", user_id="alice", owner_id="mine-too") + _reference(database, "shared-reference.png", user_id="bob", owner_id="his-too") + + caller = {"caller": "alice", "admin": False} + assert _deletable(records, mode="force", **caller) == {"own-reference.png"} + assert _deletable(records, mode="force") == {"own-reference.png", "bobs-reference.png", "shared-reference.png"} + + names = ["own-reference.png", "bobs-reference.png"] + mine = frozenset({ReferenceOwner("project", "alice", "mine")}) + # A referenced item goes only while every document naming it is acknowledged. + assert _guarded(database, records, names, mode="force", is_admin=True, acknowledged=frozenset()) == [] + assert _guarded(database, records, names, mode="force", is_admin=True, acknowledged=mine) == ["own-reference.png"] + + +def test_the_guard_re_reads_ownership(database: Database, records: IntermediatesRecords) -> None: + _image(database, "alices.png") + _image(database, "bobs.png", user_id="bob") + + kept = _guarded(database, records, ["alices.png", "bobs.png"], allowed_user_ids=frozenset({"alice"})) + + assert kept == ["alices.png"] and _exists(database, "bobs.png") + + +def test_a_cached_output_is_held_while_its_session_runs_then_for_the_grace( + database: Database, records: IntermediatesRecords +) -> None: + _image(database, "cached.png") + item = _queue_item(database, "s-running", status="in_progress") + + assert not records.hold_cached_media("s-ended", MediaReferences(images={"cached.png"})) + assert not records.hold_cached_media("s-running", MediaReferences(images={"cached.png", "gone.png"})) + assert records.hold_cached_media("s-running", MediaReferences(images={"cached.png"})) + summary = records.summarize(None)[("alice", None)]["image"].counts + assert summary.active == 1 + + with database.begin(write=True) as conn: + conn.execute(update(session_queue).where(session_queue.c.item_id == item).values(status="completed")) + assert records.summarize(None)[("alice", None)]["image"].counts.recent == 1 + + with database.begin(write=True) as conn: + conn.execute(update(intermediates_session_holds).values(released_at="2000-01-01 00:00:00.000")) + assert records.summarize(None)[("alice", None)]["image"].counts.safe == 1 + with database.begin(write=False) as conn: + assert conn.execute(select(intermediates_session_holds)).all() == [] + + +def test_a_new_process_starts_without_the_last_ones_holds_and_marks( + database: Database, records: IntermediatesRecords +) -> None: + _image(database, "cached.png") + _image(database, "leased.png") + _queue_item(database, "s-running", status="in_progress") + assert records.hold_cached_media("s-running", MediaReferences(images={"cached.png"})) + records.mark_unmeasurable("image", ["cached.png"]) + records.replace_browser_hold("alice", "tab-1", ["leased.png"], []) + + IntermediatesRecords(database).start() + + with database.begin(write=False) as conn: + assert conn.execute(select(intermediates_session_holds)).all() == [] + assert conn.execute(select(intermediates_unmeasurable)).all() == [] + # A browser lease is the tab's, not the process's: an open editor keeps its media across a restart. + assert "leased.png" not in _deletable(records) + + +def test_past_the_lease_cap_the_leases_refreshed_longest_ago_lapse( + database: Database, records: IntermediatesRecords, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(intermediates_records_default, "MAX_BROWSER_HOLD_LEASES_PER_USER", 2) + start = datetime(2030, 1, 1, tzinfo=timezone.utc) + instants: Iterator[datetime] = (start + timedelta(seconds=i) for i in itertools.count()) + monkeypatch.setattr(intermediates_records_default, "_utc_now", lambda: next(instants)) + for name in ("one.png", "two.png", "three.png"): + _image(database, name) + + records.replace_browser_hold("alice", "tab-1", ["one.png"], []) + records.replace_browser_hold("alice", "tab-2", ["two.png"], []) + records.replace_browser_hold("alice", "tab-3", ["three.png"], []) + # Refreshing tab-2 keeps it, so tab-1 lapses rather than tab-2. + records.replace_browser_hold("alice", "tab-2", ["two.png"], []) + records.replace_browser_hold("alice", "tab-4", ["three.png"], []) + + with database.begin(write=False) as conn: + leases = {row[0] for row in conn.execute(select(intermediates_browser_holds.c.lease_id))} + assert leases == {"tab-2", "tab-4"} + + +def test_a_lease_holds_only_the_accounts_own_intermediates(database: Database, records: IntermediatesRecords) -> None: + _image(database, "alices.png") + _image(database, "bobs.png", user_id="bob") + _image(database, "durable.png", is_intermediate=False) + + records.replace_browser_hold("alice", "tab-1", ["alices.png", "bobs.png", "durable.png", "missing.png"], []) + + assert _deletable(records) == {"bobs.png"} + + +class TestMeasurementQueue: + def test_the_unmeasured_come_oldest_first_and_young_or_unmeasurable_rows_wait( + self, database: Database, records: IntermediatesRecords + ) -> None: + young = (datetime.now(timezone.utc) + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S.000") + _image(database, "measured.png", file_size_bytes=100) + _image(database, "durable.png", file_size_bytes=None, is_intermediate=False) + _image(database, "unreadable.png", file_size_bytes=None) + _image(database, "young.png", file_size_bytes=None, created_at=young) + _image(database, "z-first.png", file_size_bytes=None, created_at="2020-01-01 00:00:00.000") + _image(database, "a-second.png", file_size_bytes=None, created_at="2020-01-02 00:00:00.000") + records.mark_unmeasurable("image", ["unreadable.png"]) + records.mark_unmeasurable("image", ["unreadable.png"]) + + assert records.next_unmeasured("image", 200, min_age_seconds=10) == [("z-first.png", ""), ("a-second.png", "")] + assert records.next_unmeasured("image", 1, min_age_seconds=10) == [("z-first.png", "")] + assert records.has_unmeasured_intermediates() + + records.mark_unmeasurable("image", ["z-first.png", "a-second.png", "young.png"]) + assert not records.has_unmeasured_intermediates() + + @pytest.mark.sqlite_only # Counts SQLite's virtual machine steps. + def test_late_batches_do_not_rescan_the_completed_work( + self, database: Database, records: IntermediatesRecords + ) -> None: + with database.begin(write=True) as conn: + conn.execute( + insert(images), + [ + { + "image_name": f"asset-{index:06d}", + "image_origin": "internal", + "image_category": "general", + "width": 1, + "height": 1, + "is_intermediate": True, + "created_at": OLD, + } + for index in range(100_000) + ], + ) + + def next_batch() -> tuple[list[tuple[str, str]], int]: + steps = 0 + + def count() -> int: + nonlocal steps + steps += 100 + return 0 + + with database.begin(write=False) as conn: + driver = conn.connection.driver_connection + assert driver is not None + driver.set_progress_handler(count, 100) + try: + return records.next_unmeasured("image", 200, min_age_seconds=10), steps + finally: + driver.set_progress_handler(None, 0) + + first, first_work = next_batch() + with database.begin(write=True) as conn: + conn.execute(update(images).where(images.c.image_name < "asset-099800").values(file_size_bytes=100)) + late, late_work = next_batch() + + assert first == [(f"asset-{index:06d}", "") for index in range(200)] + assert late == [(f"asset-{index:06d}", "") for index in range(99800, 100_000)] + # Deterministic enough for a generous bound, and independent of machine speed: a scan of the measured rows + # would exceed it by two orders of magnitude on the late batch. + assert 0 < first_work < 10_000 + assert late_work < 2 * first_work + + +@pytest.mark.parametrize("kind", ["image", "video"]) +def test_a_cleanup_waits_for_a_reference_in_flight_and_then_keeps_the_media( + database: Database, records: IntermediatesRecords, kind: str +) -> None: + """A reference written but not committed is invisible to a server's other transactions: without the lock, the + cleanup's check would find the media unreferenced and delete it under the save.""" + if kind == "image": + _image(database, "media") + storage: Any = ImageRecordStorage(database) + else: + with database.begin(write=True) as conn: + conn.execute( + insert(videos).values( + video_name="media", + video_origin="internal", + video_category="general", + width=8, + height=8, + is_intermediate=True, + created_at=OLD, + ) + ) + storage = VideoRecordStorage(database) + references = MediaReferences(images={"media"}) if kind == "image" else MediaReferences(videos={"media"}) + + def reference(q: Queries) -> None: + q.locks.acquire(DatabaseLock.MEDIA_PROTECTION, shared=True) + q.media_references.replace(owner_kind="project", user_id="alice", owner_id="p1", references=references) + + guard = records.make_delete_guard(kind, mode="safe", allowed_user_ids=None, caller_user_id="admin", is_admin=True) # type: ignore[arg-type] + deleted: list[str] = [] + errors = while_in_flight( + database, reference, lambda: deleted.extend(storage.delete_intermediates_by_names(["media"], guard)) + ) + + assert errors == [] and deleted == [] + + +def _workflow(name: str) -> WorkflowWithoutID: + return WorkflowWithoutID( + name=name, + author="", + description="", + version="1.0.0", + contact="", + tags="", + notes="", + exposedFields=[], + meta=WorkflowMeta(version="3.0.0", category=WorkflowCategory.User), + nodes=[{"id": "n", "data": {"image": {"image_name": "img.png"}}}], + edges=[], + ) + + +@pytest.mark.usefixtures("accounts") +@pytest.mark.parametrize( + "write", + [ + "project create", + "project update", + "workflow create", + "workflow update", + "client state", + "browser lease", + "cached output", + ], +) +def test_every_write_that_protects_media_waits_for_a_cleanup_in_flight( + database: Database, records: IntermediatesRecords, write: str +) -> None: + """Shared with the cleanup's exclusive lock: a cleanup's check and delete are never interleaved with a write + that makes the media it checked protected.""" + _image(database, "img.png") + projects = ProjectRecordsStorage(database) + workflows = WorkflowRecordsStorage(database) + project = projects.create("alice", "P", {}) + workflow = workflows.create(_workflow("W"), user_id="alice") + _queue_item(database, "s-reuser", status="in_progress") + writes: dict[str, Callable[[], object]] = { + "project create": lambda: projects.create("alice", "Q", {"imageName": "img.png"}), + "project update": lambda: projects.update( + "alice", project.project_id, project.revision, "P", {"imageName": "img.png"} + ), + "workflow create": lambda: workflows.create(_workflow("V"), user_id="alice"), + "workflow update": lambda: workflows.update(Workflow(**workflow.workflow.model_dump()), user_id="alice"), + "client state": lambda: ClientStatePersistence(database).set_by_key( + "alice", "canvas", '{"imageName": "img.png"}' + ), + "browser lease": lambda: records.replace_browser_hold("alice", "tab", ["img.png"], []), + "cached output": lambda: records.hold_cached_media("s-reuser", MediaReferences(images={"img.png"})), + } + + errors = while_in_flight(database, lambda q: q.locks.acquire(DatabaseLock.MEDIA_PROTECTION), writes[write]) + + assert errors == [] + + +@pytest.mark.parametrize("selected", [3, 33]) +def test_a_preview_of_selected_rows_counts_exactly_those_rows( + database: Database, records: IntermediatesRecords, selected: int +) -> None: + # Statements of up to 32 rows, each padded to a power of two: 3 rows fill 4 slots, 33 rows need two statements. + for owner in range(40): + _image(database, f"owner-{owner}.png", user_id=f"owner-{owner}") + + preview = records.preview_scope( + user_id=None, + targets=[(f"owner-{owner}", None) for owner in range(selected)], + mode="safe", + is_admin=True, + caller_user_id="admin", + max_acknowledged=10, + ) + + assert preview.deletable["image"].count == selected + assert preview.rows == {(f"owner-{owner}", None) for owner in range(selected)} + + +def test_windows_page_ties_within_a_second_one_row_at_a_time(database: Database, records: IntermediatesRecords) -> None: + for name, created_at in ( + ("c.png", "2020-01-02 00:00:00.123"), + ("a.png", "2020-01-02 00:00:00.500"), + ("b.png", "2020-01-02 00:00:00.500"), + ("d.png", "2020-01-02 00:00:00.999"), + ("e.png", "2020-01-02 00:00:01.000"), + ): + _image(database, name, created_at=created_at) + + batches = records.iter_deletable_batches( + "image", + user_id=None, + targets=None, + mode="safe", + is_admin=True, + caller_user_id="admin", + recent_cutoff=None, + limit=1, + ) + + assert [[name for name, _ in batch] for batch in batches] == [["c.png"], ["a.png"], ["b.png"], ["d.png"], ["e.png"]] + + +@pytest.mark.usefixtures("accounts") +def test_a_project_row_holds_the_media_of_its_project_and_the_unassigned_row_the_rest( + database: Database, records: IntermediatesRecords +) -> None: + project = ProjectRecordsStorage(database).create("alice", "P", {}) + _image(database, "in-project.png") + _image(database, "dangling.png") + _image(database, "loose.png") + with database.begin(write=True) as conn: + for name, project_id in (("in-project.png", project.project_id), ("dangling.png", "deleted-project")): + conn.execute(update(images).where(images.c.image_name == name).values(project_id=project_id)) + + def deletable(targets: list[tuple[str, Optional[str]]]) -> set[str]: + batches = records.iter_deletable_batches( + "image", + user_id=None, + targets=targets, + mode="safe", + is_admin=True, + caller_user_id="admin", + recent_cutoff=None, + limit=10, + ) + return {name for batch in batches for name, _ in batch} + + assert deletable([("alice", project.project_id)]) == {"in-project.png"} + # Media of a project that no longer exists belong to the owner's unassigned row. + assert deletable([("alice", None)]) == {"dangling.png", "loose.png"} + preview = records.preview_scope( + user_id=None, + targets=[("alice", project.project_id)], + mode="safe", + is_admin=True, + caller_user_id="admin", + max_acknowledged=10, + ) + assert preview.rows == {("alice", project.project_id)} and preview.deletable["image"].count == 1 + summary = records.summarize(["alice"]) + assert set(summary) == {("alice", project.project_id), ("alice", None)} + assert summary[("alice", None)]["image"].counts.safe == 2 + + +def test_an_expired_lease_stops_protecting_and_is_swept( + database: Database, records: IntermediatesRecords, monkeypatch: pytest.MonkeyPatch +) -> None: + now = [datetime(2030, 1, 1, tzinfo=timezone.utc)] + monkeypatch.setattr(intermediates_records_default, "_utc_now", lambda: now[0]) + _image(database, "held.png") + _image(database, "other.png") + records.replace_browser_hold("alice", "tab-1", ["held.png"], []) + assert records.summarize(None)[("alice", None)]["image"].counts.active == 1 + + now[0] += timedelta(seconds=BROWSER_HOLD_TTL_SECONDS + 1) + + assert records.summarize(None)[("alice", None)]["image"].counts.active == 0 + # The next lease of any tab sweeps the expired one away. + records.replace_browser_hold("alice", "tab-2", ["other.png"], []) + with database.begin(write=False) as conn: + assert {row[0] for row in conn.execute(select(intermediates_browser_holds.c.lease_id))} == {"tab-2"} + + +def test_a_lease_refreshed_twice_in_one_transaction_keeps_one_row_per_medium( + database: Database, records: IntermediatesRecords +) -> None: + """As when two refreshes of one lease meet on a server: both cleared the lease, and both insert its rows.""" + _image(database, "held.png") + + with database.queries.transaction() as q: + for expires_at in ("2030-01-01 00:00:00.000", "2030-01-01 00:01:00.000"): + q.intermediates.hold_for_lease( + "image", user_id="alice", lease_id="tab", names={"held.png"}, expires_at=expires_at + ) + + with database.begin(write=False) as conn: + assert conn.execute(select(intermediates_browser_holds.c.expires_at)).scalars().all() == [ + "2030-01-01 00:01:00.000" + ] + + +def test_documents_found_by_several_statements_are_counted_once_up_to_the_cap( + database: Database, records: IntermediatesRecords +) -> None: + for owner in range(40): + _image(database, f"o{owner}.png", user_id=f"owner-{owner}") + # One document naming media of two rows that fall into different statements (32 rows each). + database.queries.media_references.replace( + owner_kind="project", + user_id="owner-0", + owner_id="doc", + references=MediaReferences(images={"o0.png", "o32.png"}), + ) + + def preview(targets: Optional[list[tuple[str, Optional[str]]]]) -> Any: + return records.preview_scope( + user_id=None, targets=targets, mode="force", is_admin=True, caller_user_id="admin", max_acknowledged=1 + ) + + merged = preview([(f"owner-{owner}", None) for owner in range(33)]) + assert not merged.acknowledged_overflow + assert merged.acknowledged == {ReferenceOwner("project", "owner-0", "doc"): 2} + + database.queries.media_references.replace( + owner_kind="project", user_id="owner-1", owner_id="doc-2", references=MediaReferences(images={"o1.png"}) + ) + assert preview(None).acknowledged_overflow + + +def test_a_hold_stays_active_while_its_finished_session_belongs_to_an_active_root( + database: Database, records: IntermediatesRecords +) -> None: + """A workflow-call child completes before its outputs reach its waiting parent; the media it reused stay held.""" + _image(database, "cached.png") + root = _queue_item(database, "s-root", status="waiting") + child = _queue_item(database, "s-child", status="in_progress", root_item_id=root) + assert records.hold_cached_media("s-child", MediaReferences(images={"cached.png"})) + + with database.begin(write=True) as conn: + conn.execute(update(session_queue).where(session_queue.c.item_id == child).values(status="completed")) + + assert records.summarize(None)[("alice", None)]["image"].counts.active == 1 diff --git a/tests/app/services/intermediates/test_intermediates_service.py b/tests/app/services/intermediates/test_intermediates_service.py index cee09827163..052764f8dee 100644 --- a/tests/app/services/intermediates/test_intermediates_service.py +++ b/tests/app/services/intermediates/test_intermediates_service.py @@ -12,6 +12,7 @@ import pytest from PIL import Image +from sqlalchemy import update from invokeai.app.services.image_files.image_files_disk import DiskImageFileStorage from invokeai.app.services.image_records.image_records_common import ( @@ -36,12 +37,16 @@ IntermediatesUnavailableError, ) from invokeai.app.services.intermediates.intermediates_default import IntermediatesService -from invokeai.app.services.intermediates.intermediates_records_sqlite import IntermediatesRecordsSqlite +from invokeai.app.services.intermediates.intermediates_records_default import IntermediatesRecords from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.intermediates import intermediates_session_holds from invokeai.app.services.video_files.video_files_disk import DiskVideoFileStorage from invokeai.app.services.video_records.video_records_common import VideoRecordNotFoundException from invokeai.app.services.videos.videos_default import VideoService +from tests.fixtures.database import capture_statements, explain_query_plan +from tests.fixtures.sqlite_database import sqlite_cursor, sqlite_cursor_of ADMIN = IntermediatesCaller(user_id="admin", is_admin=True) ALICE = IntermediatesCaller(user_id="alice", is_admin=False) @@ -49,7 +54,7 @@ @pytest.fixture -def invoker(mock_invoker: Invoker, tmp_path: Path) -> Invoker: +def invoker(mock_invoker: Invoker, mock_sqlite_database: Database, tmp_path: Path) -> Invoker: """The shared mock invoker with real image/video file storage and a real queue table.""" services = mock_invoker.services services.image_files = DiskImageFileStorage(tmp_path / "outputs") @@ -57,11 +62,11 @@ def invoker(mock_invoker: Invoker, tmp_path: Path) -> Invoker: services.video_files = DiskVideoFileStorage(str(tmp_path / "outputs" / "videos")) services.videos = VideoService() services.videos.start(mock_invoker) - services.session_queue = SqliteSessionQueue(db=services.image_records._db) + services.session_queue = SessionQueue(mock_sqlite_database) services.image_moves = None # These are multi-account scenarios; the single-user case is covered explicitly below. services.configuration.multiuser = True - with services.image_records._db.transaction() as cursor: + with sqlite_cursor(mock_sqlite_database) as cursor: cursor.executemany( "INSERT INTO users (user_id, email, display_name, password_hash, is_admin, is_active)" " VALUES (?, ?, ?, 'x', ?, 1);", @@ -76,7 +81,7 @@ def invoker(mock_invoker: Invoker, tmp_path: Path) -> Invoker: @pytest.fixture def service(invoker: Invoker) -> IntermediatesService: - svc = IntermediatesService(records=IntermediatesRecordsSqlite(db=invoker.services.image_records._db)) + svc = IntermediatesService(records=IntermediatesRecords(invoker.services.session_queue._queries._database)) invoker.services.intermediates = svc svc.start(invoker) yield svc @@ -109,7 +114,7 @@ def _seed_image( project_id=project_id, ) if created_at is not None: - with records._db.transaction() as cursor: + with sqlite_cursor_of(invoker.services.session_queue) as cursor: cursor.execute("UPDATE images SET created_at = ? WHERE image_name = ?;", (created_at, name)) if with_file: invoker.services.image_files.save(image=Image.new("RGB", (8, 8)), image_name=name) @@ -131,7 +136,7 @@ def _seed_video(invoker: Invoker, name: str, *, user_id: str = "alice", is_inter is_intermediate=is_intermediate, user_id=user_id, ) - with records._db.transaction() as cursor: + with sqlite_cursor_of(invoker.services.session_queue) as cursor: cursor.execute("UPDATE videos SET created_at = '2020-01-01 00:00:00.000' WHERE video_name = ?;", (name,)) files = invoker.services.video_files path = files.get_path(name) @@ -150,7 +155,7 @@ def _enqueue_row( root_item_id: Optional[int] = None, parent_item_id: Optional[int] = None, ) -> int: - with invoker.services.image_records._db.transaction() as cursor: + with sqlite_cursor_of(invoker.services.session_queue) as cursor: cursor.execute( "INSERT INTO session_queue " "(queue_id, session, session_id, batch_id, priority, user_id, status, root_item_id, parent_item_id)" @@ -167,10 +172,11 @@ def _project(invoker: Invoker, user_id: str, name: str, data: dict) -> str: def _end_cached_media_grace(service: IntermediatesService, invoker: Invoker) -> None: """Lets the recency grace of cached media whose consuming sessions have ended run out.""" service._records.summarize(None) - with invoker.services.image_records._db.transaction() as cursor: - cursor.execute( - "UPDATE temp.intermediates_session_media SET released_at = '2000-01-01 00:00:00.000' " - "WHERE released_at IS NOT NULL;" + with invoker.services.session_queue._queries._database.begin(write=True) as conn: + conn.execute( + update(intermediates_session_holds) + .where(intermediates_session_holds.c.released_at.is_not(None)) + .values(released_at="2000-01-01 00:00:00.000") ) @@ -513,12 +519,12 @@ def test_a_row_that_was_recent_at_preview_time_stays_recent_for_the_operation( ) -> None: from datetime import datetime, timedelta, timezone - from invokeai.app.services.intermediates import intermediates_records_sqlite + from invokeai.app.services.intermediates import intermediates_records_default from invokeai.app.services.intermediates.intermediates_common import RECENT_GRACE_SECONDS start = datetime(2030, 1, 1, tzinfo=timezone.utc) readings = iter([start] + [start + timedelta(hours=1)] * 100) - monkeypatch.setattr(intermediates_records_sqlite, "_utc_now", lambda: next(readings)) + monkeypatch.setattr(intermediates_records_default, "_utc_now", lambda: next(readings)) crossing = start - timedelta(seconds=RECENT_GRACE_SECONDS - 1) _seed_image(invoker, "crossing.png", created_at=crossing.strftime("%Y-%m-%d %H:%M:%S.000")) _seed_image(invoker, "old.png") @@ -819,7 +825,7 @@ def test_a_restart_forgets_operations_and_a_fresh_run_finishes_the_work( assert failed.progress.failed_images == 1 service.stop() - restored = IntermediatesService(IntermediatesRecordsSqlite(invoker.services.image_records._db)) + restored = IntermediatesService(IntermediatesRecords(invoker.services.session_queue._queries._database)) invoker.services.intermediates = restored restored.start(invoker) monkeypatch.setattr(invoker.services.images, "delete_intermediates_by_names", real_delete) @@ -883,9 +889,11 @@ def test_inputs_added_to_an_active_session_after_it_was_scanned_are_protected( first = service.create_preview(IntermediatesPreviewRequest(mode="safe", scope=_owner("alice")), ALICE) assert first.impact.delete_images == 1 - with invoker.services.image_records._db.transaction() as cursor: + with sqlite_cursor_of(invoker.services.session_queue) as cursor: + # Every write of a session counts its revision, which tells the scan to read it again. cursor.execute( - "UPDATE session_queue SET session = ?, status = 'pending' WHERE item_id = ?;", + "UPDATE session_queue SET session = ?, session_revision = session_revision + 1, status = 'pending' " + "WHERE item_id = ?;", ('{"results": {"n": {"image": {"image_name": "child-output.png"}}}}', item_id), ) @@ -912,10 +920,10 @@ def test_a_save_cannot_slip_between_the_final_check_and_the_delete( def slow_guard_factory(kind, **kwargs): guard = real_guard_factory(kind, **kwargs) - def paused(cursor, names): + def paused(q, names): guard_entered.set() assert guard_may_finish.wait(timeout=10) - return guard(cursor, names) + return guard(q, names) return paused @@ -957,10 +965,11 @@ def save_reference() -> None: def test_active_inputs_survive_a_rolled_back_batch( invoker: Invoker, service: IntermediatesService, monkeypatch ) -> None: - """A failed batch rolls the temp table back; the next check must still see everything active.""" + """A failed batch rolls back what its guard wrote; the next check must still see everything active.""" monkeypatch.setattr(intermediates_default, "DELETE_BATCH_SIZE", 1) - _seed_image(invoker, "first.png") - _seed_image(invoker, "consumed.png") + # Windows page in creation order: "first.png" is the first batch. + _seed_image(invoker, "first.png", created_at="2020-01-01 00:00:00.000") + _seed_image(invoker, "consumed.png", created_at="2020-01-02 00:00:00.000") preview = service.create_preview(IntermediatesPreviewRequest(mode="safe", scope=_owner("alice")), ALICE) assert preview.impact.delete_images == 2 real_delete = invoker.services.image_records.delete_intermediates_by_names @@ -969,10 +978,10 @@ def test_active_inputs_survive_a_rolled_back_batch( def fail_first_batch_after_the_guard(names, guard=None): calls["n"] += 1 if calls["n"] == 1: - with invoker.services.image_records._db.transaction() as cursor: + with invoker.services.session_queue._queries._database.queries.transaction() as q: if guard is not None: - guard(cursor, names) - raise RuntimeError("simulated I/O failure after the guard rebuilt the temp table") + guard(q, names) + raise RuntimeError("simulated I/O failure after the guard ran") # Enqueued after the second batch was paged: only its final check can protect it. _enqueue_row( invoker, session_id="s-late", status="pending", session_json='{"a": {"image_name": "consumed.png"}}' @@ -1032,7 +1041,7 @@ def test_an_account_holds_media_for_a_bounded_number_of_leases(invoker: Invoker, _seed_image(invoker, f"held-{tab}.png") service.replace_browser_hold(ALICE, f"tab-{tab}", IntermediatesBrowserHoldRequest(images=[f"held-{tab}.png"])) # Refresh order decides which lease lapses; stagger it so the first tab is the stalest. - with invoker.services.image_records._db.transaction() as cursor: + with sqlite_cursor_of(invoker.services.session_queue) as cursor: cursor.execute( "UPDATE intermediates_browser_holds SET expires_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW', ?)" " WHERE user_id = 'alice' AND lease_id = ?;", @@ -1117,19 +1126,19 @@ def test_the_final_check_judges_cached_media_recency_as_of_the_preview( counted recent, and the next page would then collect the row the preview promised to keep.""" from datetime import datetime, timedelta, timezone - from invokeai.app.services.intermediates import intermediates_records_sqlite + from invokeai.app.services.intermediates import intermediates_records_default from invokeai.app.services.shared.media_references import MediaReferences monkeypatch.setattr(intermediates_default, "DELETE_BATCH_SIZE", 1) start = datetime(2030, 1, 1, tzinfo=timezone.utc) release_time = start - timedelta(minutes=29) readings = iter([release_time, release_time, start] + [start + timedelta(hours=1)] * 100) - monkeypatch.setattr(intermediates_records_sqlite, "_utc_now", lambda: next(readings)) + monkeypatch.setattr(intermediates_records_default, "_utc_now", lambda: next(readings)) _seed_image(invoker, "old.png") _seed_image(invoker, "cached.png") consumer = _enqueue_row(invoker, session_id="consumer", status="in_progress") assert service.hold_cached_media("consumer", MediaReferences(images={"cached.png"})) - with invoker.services.image_records._db.transaction() as cursor: + with sqlite_cursor_of(invoker.services.session_queue) as cursor: cursor.execute("UPDATE session_queue SET status = 'completed' WHERE item_id = ?;", (consumer,)) service._records.summarize(None) @@ -1176,49 +1185,42 @@ def test_matching_scope_resolves_filtered_rows_minus_exclusions( def test_selection_paging_seeks_each_row_without_sorting_the_scope( - invoker: Invoker, service: IntermediatesService, monkeypatch: pytest.MonkeyPatch + invoker: Invoker, service: IntermediatesService ) -> None: - """A window of an explicit-row scope is an index seek in rowid order. One OR over many rows would make - SQLite sort the whole remaining scope on every window, holding the database lock for the run.""" + """A window's page of an explicit-row scope is an index seek that starts at the keyset, in creation order. One + OR over many rows would make SQLite sort the whole remaining scope on every window, holding the database for + the run; only the page itself is classified and put in order.""" records = service._records + database = invoker.services.session_queue._queries._database targets = [("alice", _project(invoker, "alice", f"P{index}", {})) for index in range(3)] + [("alice", None)] for index in range(len(targets)): _seed_image(invoker, f"row-{index}.png", project_id=targets[index][1]) - real_scope_clauses = records._scope_clauses - paging_clauses: list[list[tuple[str, list[object]]]] = [] - - def recording_scope_clauses(**kwargs): - clauses = real_scope_clauses(**kwargs) - if kwargs.get("per_statement") == 1: - paging_clauses.append(clauses) - return clauses - - monkeypatch.setattr(records, "_scope_clauses", recording_scope_clauses) - batches = list( - records.iter_deletable_batches( - "image", - user_id=None, - targets=targets, - mode="safe", - is_admin=False, - caller_user_id="alice", - recent_cutoff=None, - limit=10, - ) - ) - assert sorted(name for batch in batches for name, _ in batch) == [f"row-{index}.png" for index in range(4)] - assert len(paging_clauses) == 1 and len(paging_clauses[0]) == len(targets) - deletable_sql, deletable_params = records._deletable_sql("image", "safe", is_admin=False, caller_user_id="alice") - with invoker.services.image_records._db.transaction() as cursor: - clock = records._prepare(cursor) - for clause, params in paging_clauses[0]: - cursor.execute( - "EXPLAIN QUERY PLAN " + records._window_sql("image", clause, deletable_sql), - [*deletable_params, *clock, *params, 0, 10], + with capture_statements(database) as statements: + batches = list( + records.iter_deletable_batches( + "image", + user_id=None, + targets=targets, + mode="safe", + is_admin=False, + caller_user_id="alice", + recent_cutoff=None, + limit=10, ) - plan = " ".join(str(row[-1]) for row in cursor.fetchall()) - assert "B-TREE FOR ORDER BY" not in plan, plan + ) + + assert sorted(name for batch in batches for name, _ in batch) == [f"row-{index}.png" for index in range(4)] + windows = [ + (statement, parameters) for statement, parameters in statements if "ORDER BY classified.created_at" in statement + ] + # A first window of each row finds it, a second finds nothing more. + assert len(windows) == 2 * len(targets) + for statement, parameters in windows: + plan = explain_query_plan(database, statement, parameters) + page = plan[: plan.index("SCAN page")] + assert any("(created_at,image_name)>(?,?)" in step for step in page), plan + assert not any("B-TREE" in step for step in page), plan @pytest.mark.parametrize("cleanup_wins_race", [False, True]) @@ -1270,7 +1272,7 @@ def queued_session(): runner.run_node(first.session.next(), first) assert first.session.is_complete() and not first.session.has_error() services.session_queue.save_queue_item_session(first.item_id, first.session) - with services.image_records._db.transaction() as cursor: + with sqlite_cursor_of(services.session_queue) as cursor: cursor.execute("UPDATE session_queue SET status = 'completed' WHERE item_id = ?", (first.item_id,)) cursor.execute("UPDATE images SET created_at = '2020-01-01 00:00:00.000' WHERE image_name = ?", (image_name,)) @@ -1309,7 +1311,7 @@ def get_then_cleanup(key): runner.run_node(second.session.next(), second) assert second.session.is_complete() and not second.session.has_error() services.session_queue.save_queue_item_session(second.item_id, second.session) - with services.image_records._db.transaction() as cursor: + with sqlite_cursor_of(services.session_queue) as cursor: cursor.execute("UPDATE session_queue SET status = 'completed' WHERE item_id = ?", (second.item_id,)) if not cleanup_wins_race: _end_cached_media_grace(service, invoker) @@ -1335,11 +1337,11 @@ def test_cached_media_holds_keep_rows_out_of_a_running_cleanup_until_all_consumi assert protected.progress.deleted_images == protected.progress.deleted_videos == 0 assert _exists(invoker, "cached.png") - with invoker.services.image_records._db.transaction() as cursor: + with sqlite_cursor_of(invoker.services.session_queue) as cursor: cursor.execute("UPDATE session_queue SET status = 'completed' WHERE item_id = ?", (first,)) preview, _ = _run(service, ALICE, _owner("alice"), mode=mode) assert preview.impact.keep_active_images == preview.impact.keep_active_videos == 1 - with invoker.services.image_records._db.transaction() as cursor: + with sqlite_cursor_of(invoker.services.session_queue) as cursor: cursor.execute("UPDATE session_queue SET status = 'canceled' WHERE item_id = ?", (second,)) preview, _ = _run(service, ALICE, _owner("alice"), mode=mode) # A cache hit reused old rows; the grace a fresh output gets now runs from the session's end. @@ -1407,10 +1409,9 @@ def test_completed_child_media_survives_until_its_root_finishes( assert queue.prune("default").deleted == 0 if reload_records: + # A new process starts without the holds of the last one. service.stop() - with invoker.services.image_records._db.transaction() as cursor: - cursor.execute("DROP TABLE IF EXISTS temp.intermediates_session_media;") - service = IntermediatesService(IntermediatesRecordsSqlite(invoker.services.image_records._db)) + service = IntermediatesService(IntermediatesRecords(invoker.services.session_queue._queries._database)) service.start(invoker) try: if not reload_records: @@ -1444,7 +1445,7 @@ def test_completed_nested_child_keeps_media_while_its_root_is_active( _seed_image(invoker, "produced.png", session_id="leaf") _seed_video(invoker, "cached.mp4") assert service.hold_cached_media("leaf", MediaReferences(videos={"cached.mp4"})) - with invoker.services.image_records._db.transaction() as cursor: + with sqlite_cursor_of(invoker.services.session_queue) as cursor: cursor.execute("UPDATE session_queue SET status = 'completed' WHERE item_id = ?", (child,)) for mode in ("safe", "force"): @@ -1452,7 +1453,7 @@ def test_completed_nested_child_keeps_media_while_its_root_is_active( assert preview.impact.keep_active_images == preview.impact.keep_active_videos == 1 assert completed.progress.deleted_images == completed.progress.deleted_videos == 0 - with invoker.services.image_records._db.transaction() as cursor: + with sqlite_cursor_of(invoker.services.session_queue) as cursor: cursor.execute("UPDATE session_queue SET status = 'completed' WHERE item_id = ?", (root,)) _end_cached_media_grace(service, invoker) _, completed = _run(service, ALICE, _owner("alice")) @@ -1542,22 +1543,21 @@ def test_legacy_client_state_protects_its_canvas_media_until_it_changes( def test_quarantined_projects_are_named_when_a_force_clear_would_break_them( invoker: Invoker, service: IntermediatesService ) -> None: - from invokeai.app.services.shared.media_references import MediaReferences, replace_media_references + from invokeai.app.services.shared.media_references import MediaReferences _seed_image(invoker, "recovered.png") - with invoker.services.image_records._db.transaction() as cursor: + with sqlite_cursor_of(invoker.services.session_queue) as cursor: cursor.execute( "CREATE TABLE IF NOT EXISTS orphaned_projects_2026_08_06 (project_id TEXT, user_id TEXT, name TEXT," " data TEXT, PRIMARY KEY (user_id, project_id));" ) cursor.execute("INSERT INTO orphaned_projects_2026_08_06 VALUES ('p1', 'alice', 'Before boards', '{}');") - replace_media_references( - cursor, - owner_kind="quarantined_project", - user_id="alice", - owner_id="p1", - references=MediaReferences(images={"recovered.png"}), - ) + invoker.services.session_queue._queries._database.queries.media_references.replace( + owner_kind="quarantined_project", + user_id="alice", + owner_id="p1", + references=MediaReferences(images={"recovered.png"}), + ) safe, _ = _run(service, ALICE, _owner("alice")) assert safe.impact.keep_referenced_images == 1 diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index f06423953e4..c76cc4ce6fa 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -202,7 +202,7 @@ def test_destination_is_hidden_from_orphan_cleanup_during_install( ) from invokeai.app.services.orphaned_models import OrphanedModelsService - orphan_service = OrphanedModelsService(config=mm2_app_config, db=mm2_installer.record_store._db) + orphan_service = OrphanedModelsService(config=mm2_app_config, store=mm2_installer.record_store) observations: list[tuple[str, bool, set[str]]] = [] real_move = shutil.move @@ -239,7 +239,7 @@ def test_managed_local_source_is_hidden_from_orphan_cleanup_during_install( source_root = mm2_app_config.models_path / f"local-source-{uuid.uuid4().hex}" shutil.copytree(diffusers_dir, source_root) - orphan_service = OrphanedModelsService(config=mm2_app_config, db=mm2_installer.record_store._db) + orphan_service = OrphanedModelsService(config=mm2_app_config, store=mm2_installer.record_store) observations: list[tuple[bool, set[str], str]] = [] real_probe = mm2_installer._probe real_move_with_retries = mm2_installer._move_with_retries @@ -285,7 +285,7 @@ def test_orphan_delete_claim_prevents_a_local_install_from_starting( source_root = mm2_app_config.models_path / f"local-source-{uuid.uuid4().hex}" shutil.copytree(diffusers_dir, source_root) - orphan_service = OrphanedModelsService(config=mm2_app_config, db=mm2_installer.record_store._db) + orphan_service = OrphanedModelsService(config=mm2_app_config, store=mm2_installer.record_store) real_rmtree = orphaned_models_service.shutil.rmtree deletion_claimed = threading.Event() release_deletion = threading.Event() @@ -410,7 +410,7 @@ def test_unregistered_destination_remains_protected_after_registration_failure( from invokeai.app.services.model_install.model_install_common import has_recovery_sentinel from invokeai.app.services.orphaned_models import OrphanedModelsService - orphan_service = OrphanedModelsService(config=mm2_app_config, db=mm2_installer.record_store._db) + orphan_service = OrphanedModelsService(config=mm2_app_config, store=mm2_installer.record_store) def fail_recording(_config: Any) -> None: raise RuntimeError("simulated model record failure") @@ -1399,15 +1399,17 @@ def test_install_registration_failure_preserves_complete_recoverable_files( store = mm2_installer.record_store assert isinstance(store, ModelRecordServiceSQL) - # Exercise a genuine SQLite write rejection without changing permissions on any real user database. - with store._db.transaction() as cursor: - cursor.execute("PRAGMA query_only = ON") + database = store._queries._database + # Exercise a genuine SQLite write rejection without changing permissions on any real user database. The + # pragma is set in read transactions: a write transaction could not even begin while it is on. + with database.begin(write=False) as conn: + conn.exec_driver_sql("PRAGMA query_only = ON") try: with pytest.raises(sqlite3.OperationalError, match="readonly database"): mm2_installer.install_path(diffusers_dir) finally: - with store._db.transaction() as cursor: - cursor.execute("PRAGMA query_only = OFF") + with database.begin(write=False) as conn: + conn.exec_driver_sql("PRAGMA query_only = OFF") assert mm2_installer.record_store.all_models() == [] new_paths = set(mm2_app_config.models_path.iterdir()) - existing_paths @@ -2585,7 +2587,7 @@ def test_paused_remote_install_stays_hidden_from_orphan_cleanup( job._multifile_job = download_job installer._download_cache[download_job.id] = job create_active_install_sentinel(tmpdir) - orphan_service = OrphanedModelsService(config=mm2_app_config, db=installer.record_store._db) + orphan_service = OrphanedModelsService(config=mm2_app_config, store=installer.record_store) try: installer._download_cancelled_callback(download_job) diff --git a/tests/app/services/model_records/test_model_records_sql.py b/tests/app/services/model_records/test_model_records_sql.py index 25b052d7727..fe6dc18bc38 100644 --- a/tests/app/services/model_records/test_model_records_sql.py +++ b/tests/app/services/model_records/test_model_records_sql.py @@ -2,13 +2,14 @@ Test the refactored model config classes. """ +import json from hashlib import sha256 from typing import Any, Optional import pytest from pydantic import ValidationError +from sqlalchemy import insert -from invokeai.app.services.config import InvokeAIAppConfig from invokeai.app.services.model_records import ( DuplicateModelException, ModelRecordOrderBy, @@ -17,7 +18,10 @@ UnknownModelException, ) from invokeai.app.services.model_records.model_records_base import ModelRecordChanges -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries.models import MAX_KEY_LENGTH, MAX_PATH_LENGTH, ModelQueries +from invokeai.app.services.shared.database.schema.models import models +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.backend.model_manager.configs.controlnet import ControlAdapterDefaultSettings from invokeai.backend.model_manager.configs.lora import LoRA_LyCORIS_SDXL_Config from invokeai.backend.model_manager.configs.main import ( @@ -41,18 +45,12 @@ SchedulerPredictionType, ) from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database +from tests.fixtures.races import when_called @pytest.fixture -def store( - datadir: Any, -) -> ModelRecordServiceSQL: - config = InvokeAIAppConfig() - config._root = datadir - logger = InvokeAILogger.get_logger(config=config) - db = create_mock_sqlite_database(config, logger) - return ModelRecordServiceSQL(db, logger) +def store(database: Database) -> ModelRecordServiceSQL: + return ModelRecordServiceSQL(database, InvokeAILogger.get_logger()) def example_ti_config(key: Optional[str] = None) -> TI_File_SD1_Config: @@ -80,7 +78,7 @@ def test_type(store: ModelRecordServiceBase): def test_raises_on_violating_uniqueness(store: ModelRecordServiceBase): - # Models have a uniqueness constraint by their name, base and type + # The key and the path are unique: config1 again repeats both, config2 its path. config1 = example_ti_config("key1") config2 = config1.model_copy(deep=True) config2.key = "key2" @@ -313,8 +311,7 @@ def test_unique_by_path(store: ModelRecordServiceBase): prediction_type=SchedulerPredictionType.Epsilon, repo_variant=ModelRepoVariant.Default, ) - # config1, config2 and config3 are compatible because they have unique paths - # of name, type and base + # config1, config2 and config3 can be installed together: their paths differ, though not their names for c in config1, config2, config3: c.key = sha256(c.path.encode("utf-8")).hexdigest() store.add_model(c) @@ -486,3 +483,202 @@ def test_model_record_changes(): changes = ModelRecordChanges.model_validate({"default_settings": {"vae": "value"}}) assert isinstance(changes.default_settings, MainModelDefaultSettings) + + +def _embedding(key: str, *, name: str = "embedding") -> TI_File_SD1_Config: + return TI_File_SD1_Config( + key=key, + source="test/source/", + source_type=ModelSourceType.Path, + path=f"/tmp/{key}.bin", + file_size=1024, + name=name, + base=BaseModelType.StableDiffusion1, + type=ModelType.TextualInversion, + format=ModelFormat.EmbeddingFile, + hash="ABC123", + ) + + +def test_a_duplicate_is_reported_by_what_it_repeats(store: ModelRecordServiceBase) -> None: + store.add_model(_embedding("key1")) + + with pytest.raises(DuplicateModelException, match="with path"): + store.add_model(_embedding("key2").model_copy(update={"path": "/tmp/key1.bin"})) + with pytest.raises(DuplicateModelException, match="with key"): + store.add_model(_embedding("key1").model_copy(update={"path": "/tmp/elsewhere.bin"})) + + +def test_models_that_tie_are_listed_by_key_and_names_ignore_case(store: ModelRecordServiceBase) -> None: + # Added against key order, so that only the tie-break lists equal names by key. By bytes "Beta" sorts before + # "alpha", so only a case-insensitive order lists it last. + for key, name in (("k4", "Beta"), ("k3", "alpha"), ("k2", "ALPHA"), ("k1", "Alpha")): + store.add_model(_embedding(key, name=name)) + + def keys(order_by: ModelRecordOrderBy, direction: SQLiteDirection) -> list[str]: + return [model.key for model in store.search_by_attr(order_by=order_by, direction=direction)] + + # The embeddings share their type, base and format, so the default order goes by their names too. + for order_by in (ModelRecordOrderBy.Name, ModelRecordOrderBy.Default): + assert keys(order_by, SQLiteDirection.Ascending) == ["k1", "k2", "k3", "k4"] + assert keys(order_by, SQLiteDirection.Descending) == ["k4", "k3", "k2", "k1"] + + +_MAIN = { + "variant": ModelVariantType.Normal, + "prediction_type": SchedulerPredictionType.Epsilon, + "repo_variant": ModelRepoVariant.Default, +} +# Models that differ on every ordered attribute: key -> (class, fields of the class, name, path, size, added, modified). +# Timestamps are written as stored, so that no test waits for the clock. +_ORDERED_MODELS: dict[str, tuple[type, dict[str, Any], str, str, int, str, str]] = { + "a": ( + Main_Diffusers_SDXL_Config, + _MAIN, + "delta", + "/m/2", + 3000, + "2000-01-03 00:00:00.000", + "2000-02-02 00:00:00.000", + ), + "b": ( + VAE_Diffusers_SD1_Config, + {"repo_variant": ModelRepoVariant.Default}, + "Charlie", + "/m/4", + 1000, + "2000-01-01 00:00:00.000", + "2000-02-04 00:00:00.000", + ), + "c": (TI_File_SD1_Config, {}, "bravo", "/m/1", 4000, "2000-01-02 00:00:00.000", "2000-02-03 00:00:00.000"), + "d": (LoRA_LyCORIS_SDXL_Config, {}, "Alpha", "/m/3", 2000, "2000-01-04 00:00:00.000", "2000-02-01 00:00:00.000"), + "e": (Main_Diffusers_SD1_Config, _MAIN, "echo", "/m/5", 5000, "2000-01-05 00:00:00.000", "2000-02-05 00:00:00.000"), +} +# Each order's listing, ascending; descending is its reverse, ties included. +_ASCENDING = { + ModelRecordOrderBy.Default: "cdeab", # by type, then base: the two main models + ModelRecordOrderBy.Type: "cdaeb", # the two main models tie, by key + ModelRecordOrderBy.Base: "bcead", + ModelRecordOrderBy.Name: "dcbae", # ignoring case; by bytes it would be "dbcae" + ModelRecordOrderBy.Format: "abecd", + ModelRecordOrderBy.Size: "bdace", + ModelRecordOrderBy.DateAdded: "bcade", + ModelRecordOrderBy.DateModified: "dacbe", + ModelRecordOrderBy.Path: "cadbe", +} + + +@pytest.fixture +def ordered_store(database: Database, store: ModelRecordServiceSQL) -> ModelRecordServiceSQL: + with database.begin(write=True) as conn: + for key, (cls, fields, name, path, size, added, modified) in _ORDERED_MODELS.items(): + config = cls( + key=key, + name=name, + path=path, + file_size=size, + hash=f"hash-{key}", + source="test/source/", + source_type=ModelSourceType.Path, + **fields, + ) + conn.execute( + insert(models).values(id=key, config=config.model_dump_json(), created_at=added, updated_at=modified) + ) + return store + + +def _listed(store: ModelRecordServiceBase, order_by: ModelRecordOrderBy, direction: SQLiteDirection) -> str: + return "".join(model.key for model in store.search_by_attr(order_by=order_by, direction=direction)) + + +def test_each_order_has_its_own_listing() -> None: + # Else a test below could not tell one order's column from another's. + assert len(set(_ASCENDING.values())) == len(_ASCENDING) == len(ModelRecordOrderBy) + + +@pytest.mark.parametrize("order_by", list(ModelRecordOrderBy)) +def test_each_order_sorts_by_its_own_column(ordered_store: ModelRecordServiceSQL, order_by: ModelRecordOrderBy) -> None: + assert _listed(ordered_store, order_by, SQLiteDirection.Ascending) == _ASCENDING[order_by] + assert _listed(ordered_store, order_by, SQLiteDirection.Descending) == _ASCENDING[order_by][::-1] + + +def test_an_update_makes_a_model_the_last_modified(ordered_store: ModelRecordServiceSQL) -> None: + ordered_store.update_model("c", ModelRecordChanges(description="edited")) + + assert _listed(ordered_store, ModelRecordOrderBy.DateModified, SQLiteDirection.Ascending) == "dabec" + assert _listed(ordered_store, ModelRecordOrderBy.DateAdded, SQLiteDirection.Ascending) == "bcade" + + +def test_the_format_filter_applies(ordered_store: ModelRecordServiceSQL) -> None: + diffusers = ordered_store.search_by_attr(model_format=ModelFormat.Diffusers) + + assert [model.key for model in diffusers] == ["e", "a", "b"] + + +def test_a_replaced_model_is_stored_and_an_unknown_one_is_not_found(ordered_store: ModelRecordServiceSQL) -> None: + replaced = ordered_store.get_model("c").model_copy(update={"name": "renamed"}) + + ordered_store.replace_model("c", replaced) + + assert ordered_store.get_model("c").name == "renamed" + with pytest.raises(ValueError): + ordered_store.replace_model("c", ordered_store.get_model("a")) + with pytest.raises(UnknownModelException): + ordered_store.replace_model("ghost", replaced.model_copy(update={"key": "ghost", "path": "/m/ghost"})) + with pytest.raises(UnknownModelException): + ordered_store.del_model("ghost") + + +def test_concurrent_updates_of_a_model_both_apply( + store: ModelRecordServiceBase, monkeypatch: pytest.MonkeyPatch +) -> None: + store.add_model(example_ti_config("key1")) + + def second() -> object: + return store.update_model("key1", ModelRecordChanges(description="second")) + + # The second update starts while the first holds the model's row: it must wait, then build on the first. + ended = when_called(monkeypatch, ModelQueries, "lock", second) + store.update_model("key1", ModelRecordChanges(name="first")) + + assert ended() == [] + updated = store.get_model("key1") + assert (updated.name, updated.description) == ("first", "second") + + +def test_model_paths_include_records_that_no_longer_validate(store: ModelRecordServiceBase, database: Database) -> None: + store.add_model(_embedding("key1")) + # A record of a model type this version does not know, as after going back to an older version. + unknown = _embedding("key2").model_dump(mode="json") | {"type": "from_the_future"} + database.queries.models.insert("key2", json.dumps(unknown)) + + assert sorted(store.get_model_paths()) == ["/tmp/key1.bin", "/tmp/key2.bin"] + assert [model.key for model in store.search_by_attr()] == ["key1"] + + +def test_the_oldest_model_of_a_file_comes_first(store: ModelRecordServiceBase, database: Database) -> None: + # The same file under two paths, the older record with the larger key: the route that recalls a model by its hash + # takes the first. + with database.begin(write=True) as conn: + for key, added in (("k1", "2000-01-02 00:00:00.000"), ("k2", "2000-01-01 00:00:00.000")): + config = _embedding(key).model_dump_json() + conn.execute(insert(models).values(id=key, config=config, created_at=added, updated_at=added)) + + assert [model.key for model in store.search_by_hash("ABC123")] == ["k2", "k1"] + + +def test_a_key_or_path_longer_than_a_record_holds_is_refused(store: ModelRecordServiceBase) -> None: + # SQLite would hold them, a server would not: every backend refuses them alike. + with pytest.raises(ValueError, match="key"): + store.add_model(_embedding("k" * (MAX_KEY_LENGTH + 1))) + too_long_path = "/" + "p" * MAX_PATH_LENGTH + with pytest.raises(ValueError, match="path"): + store.add_model(_embedding("key1").model_copy(update={"path": too_long_path})) + + longest_path = "/" + "p" * (MAX_PATH_LENGTH - 1) + store.add_model(_embedding("k" * MAX_KEY_LENGTH).model_copy(update={"path": longest_path})) + store.add_model(_embedding("key1")) + with pytest.raises(ValueError, match="path"): + store.update_model("key1", ModelRecordChanges(path=too_long_path)) + assert store.get_model("key1").path == "/tmp/key1.bin" diff --git a/tests/app/services/model_relationship_records/test_model_relationship_records.py b/tests/app/services/model_relationship_records/test_model_relationship_records.py new file mode 100644 index 00000000000..4a121d19369 --- /dev/null +++ b/tests/app/services/model_relationship_records/test_model_relationship_records.py @@ -0,0 +1,99 @@ +"""Model relationships on every database backend.""" + +import pytest + +from invokeai.app.services.model_records import ModelRecordServiceSQL +from invokeai.app.services.model_relationship_records.model_relationship_records_default import ( + ModelRelationshipRecordStorage, +) +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import ForeignKeyViolation +from invokeai.backend.model_manager.configs.textual_inversion import TI_File_SD1_Config +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelSourceType, ModelType +from invokeai.backend.util.logging import InvokeAILogger + + +@pytest.fixture +def models(database: Database) -> ModelRecordServiceSQL: + store = ModelRecordServiceSQL(database, InvokeAILogger.get_logger()) + for key in ("a", "b", "c", "d"): + store.add_model( + TI_File_SD1_Config( + key=key, + path=f"/tmp/{key}.bin", + name=key, + hash="ABC123", + file_size=1024, + source="test/source/", + source_type=ModelSourceType.Path, + base=BaseModelType.StableDiffusion1, + type=ModelType.TextualInversion, + format=ModelFormat.EmbeddingFile, + ) + ) + return store + + +@pytest.fixture +def relationships(database: Database, models: ModelRecordServiceSQL) -> ModelRelationshipRecordStorage: + return ModelRelationshipRecordStorage(database) + + +def test_a_relationship_is_seen_from_both_models_in_order(relationships: ModelRelationshipRecordStorage) -> None: + relationships.add_model_relationship("d", "c") + relationships.add_model_relationship("c", "a") + relationships.add_model_relationship("b", "c") + + # A pair is stored with its smaller key first, so c's relatives come from both sides: d after it, a and b before. + assert relationships.get_related_model_keys("c") == ["a", "b", "d"] + assert relationships.get_related_model_keys("a") == ["c"] + assert relationships.get_related_model_keys("d") == ["c"] + + +def test_relating_again_from_either_side_stores_the_relationship_once( + relationships: ModelRelationshipRecordStorage, +) -> None: + relationships.add_model_relationship("a", "c") + relationships.add_model_relationship("c", "a") + relationships.add_model_relationship("a", "c") + + # Stored once, so one removal ends it. + relationships.remove_model_relationship("c", "a") + + assert relationships.get_related_model_keys("a") == [] + + +def test_removing_a_relationship_leaves_the_others_and_again_is_harmless( + relationships: ModelRelationshipRecordStorage, +) -> None: + relationships.add_model_relationship("a", "b") + relationships.add_model_relationship("a", "c") + + relationships.remove_model_relationship("b", "a") + relationships.remove_model_relationship("b", "a") + + assert relationships.get_related_model_keys("a") == ["c"] + assert relationships.get_related_model_keys("b") == [] + + +def test_a_model_is_related_neither_to_itself_nor_to_one_that_does_not_exist( + relationships: ModelRelationshipRecordStorage, +) -> None: + with pytest.raises(ValueError): + relationships.add_model_relationship("a", "a") + # Only a relationship that exists already is skipped, not a missing model. + with pytest.raises(ForeignKeyViolation): + relationships.add_model_relationship("a", "unknown") + + +def test_deleting_a_model_ends_its_relationships( + relationships: ModelRelationshipRecordStorage, models: ModelRecordServiceSQL +) -> None: + relationships.add_model_relationship("a", "b") + relationships.add_model_relationship("b", "c") + relationships.add_model_relationship("a", "d") + + models.del_model("b") + + assert relationships.get_related_model_keys("a") == ["d"] + assert relationships.get_related_model_keys("c") == [] diff --git a/tests/app/services/model_relationship_records/test_model_relationship_records_sqlite.py b/tests/app/services/model_relationship_records/test_model_relationship_records_sqlite.py deleted file mode 100644 index 373d31a4fff..00000000000 --- a/tests/app/services/model_relationship_records/test_model_relationship_records_sqlite.py +++ /dev/null @@ -1,72 +0,0 @@ -"""Tests for SQLite model relationship storage contracts.""" - -import logging -import sqlite3 -from pathlib import Path - -import pytest - -from invokeai.app.services.config import InvokeAIAppConfig -from invokeai.app.services.model_records import ModelRecordServiceSQL -from invokeai.app.services.model_relationship_records.model_relationship_records_sqlite import ( - SqliteModelRelationshipRecordStorage, -) -from invokeai.backend.model_manager.configs.textual_inversion import TI_File_SD1_Config -from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelSourceType, ModelType -from tests.fixtures.sqlite_database import create_mock_sqlite_database - - -@pytest.fixture -def storage(tmp_path: Path) -> SqliteModelRelationshipRecordStorage: - logger = logging.getLogger(__name__) - config = InvokeAIAppConfig() - config._root = tmp_path - db = create_mock_sqlite_database(config, logger) - records = ModelRecordServiceSQL(db, logger) - for key in ("model-a", "model-b"): - records.add_model( - TI_File_SD1_Config( - key=key, - source=f"test/{key}", - source_type=ModelSourceType.Path, - path=f"/tmp/{key}.safetensors", - file_size=1, - name=key, - base=BaseModelType.StableDiffusion1, - type=ModelType.TextualInversion, - format=ModelFormat.EmbeddingFile, - hash=f"hash-{key}", - ) - ) - return SqliteModelRelationshipRecordStorage(db) - - -def test_add_relationship_is_bidirectional_and_duplicate_is_idempotent( - storage: SqliteModelRelationshipRecordStorage, -) -> None: - storage.add_model_relationship("model-b", "model-a") - - assert storage.get_related_model_keys("model-a") == ["model-b"] - assert storage.get_related_model_keys("model-b") == ["model-a"] - - storage.add_model_relationship("model-a", "model-b") - assert storage.get_related_model_keys("model-a") == ["model-b"] - assert storage.get_related_model_keys("model-b") == ["model-a"] - - -def test_add_relationship_rejects_missing_model_key(storage: SqliteModelRelationshipRecordStorage) -> None: - with pytest.raises(sqlite3.IntegrityError, match="FOREIGN KEY"): - storage.add_model_relationship("model-a", "missing-model") - - -def test_remove_missing_relationship_is_idempotent(storage: SqliteModelRelationshipRecordStorage) -> None: - storage.remove_model_relationship("model-a", "model-b") - assert storage.get_related_model_keys("model-a") == [] - - -def test_delete_model_cascades_relationship_in_migrated_schema(storage: SqliteModelRelationshipRecordStorage) -> None: - storage.add_model_relationship("model-a", "model-b") - with storage._db.transaction() as cursor: - cursor.execute("DELETE FROM models WHERE id = ?", ("model-a",)) - - assert storage.get_related_model_keys("model-b") == [] diff --git a/tests/app/services/orphaned_models/test_orphaned_models_scratch_dir.py b/tests/app/services/orphaned_models/test_orphaned_models_scratch_dir.py index afe636a452a..f15714d2925 100644 --- a/tests/app/services/orphaned_models/test_orphaned_models_scratch_dir.py +++ b/tests/app/services/orphaned_models/test_orphaned_models_scratch_dir.py @@ -8,23 +8,24 @@ be invisible to the scan by construction. """ -from logging import Logger +import json from pathlib import Path import pytest from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.model_install.model_install_common import INSTALL_RECOVERY_SENTINEL +from invokeai.app.services.model_records import ModelRecordServiceSQL from invokeai.app.services.orphaned_models import CONVERSION_SCRATCH_DIRNAME, OrphanedModelsService -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database +from invokeai.backend.model_manager.configs.textual_inversion import TI_File_SD1_Config +from invokeai.backend.model_manager.taxonomy import ModelSourceType +from invokeai.backend.util.logging import InvokeAILogger @pytest.fixture -def db() -> SqliteDatabase: - database = SqliteDatabase(db_path=None, logger=Logger("test_orphaned_models"), verbose=False) - database._conn.execute("CREATE TABLE models (id TEXT PRIMARY KEY, config TEXT NOT NULL);") - database._conn.commit() - return database +def store(database: Database) -> ModelRecordServiceSQL: + return ModelRecordServiceSQL(database, InvokeAILogger.get_logger()) @pytest.fixture @@ -34,11 +35,11 @@ def models_path(tmp_path: Path) -> Path: return path -def _service(models_path: Path, db: SqliteDatabase) -> OrphanedModelsService: +def _service(models_path: Path, store: ModelRecordServiceSQL) -> OrphanedModelsService: config = InvokeAIAppConfig() config._root = models_path.parent assert config.models_path == models_path - return OrphanedModelsService(config=config, db=db) + return OrphanedModelsService(config=config, store=store) def _write_model_file(directory: Path, name: str = "model.safetensors") -> None: @@ -46,12 +47,14 @@ def _write_model_file(directory: Path, name: str = "model.safetensors") -> None: (directory / name).write_bytes(b"not really a model") -def test_conversion_scratch_directory_is_not_reported_as_orphaned(models_path: Path, db: SqliteDatabase) -> None: +def test_conversion_scratch_directory_is_not_reported_as_orphaned( + models_path: Path, store: ModelRecordServiceSQL +) -> None: # What a conversion looks like on disk while it runs: a TemporaryDirectory under the scratch # area, holding the diffusers copy it has written so far. _write_model_file(models_path / CONVERSION_SCRATCH_DIRNAME / "tmp8f2b1c" / "sd-v1-5") - orphans = _service(models_path, db).find_orphaned_models() + orphans = _service(models_path, store).find_orphaned_models() assert orphans == [], ( "The scan reported a conversion's working directory. `DELETE /sync/orphaned` would delete " @@ -59,33 +62,94 @@ def test_conversion_scratch_directory_is_not_reported_as_orphaned(models_path: P ) -def test_a_real_orphan_is_still_reported(models_path: Path, db: SqliteDatabase) -> None: +def test_a_real_orphan_is_still_reported(models_path: Path, store: ModelRecordServiceSQL) -> None: """The control: the skip must be targeted, not a scan that has stopped finding anything.""" _write_model_file(models_path / "some-unregistered-model") - orphans = _service(models_path, db).find_orphaned_models() + orphans = _service(models_path, store).find_orphaned_models() assert [orphan.path for orphan in orphans] == ["some-unregistered-model"] -def test_conversion_scratch_directory_cannot_be_deleted(models_path: Path, db: SqliteDatabase) -> None: +def test_conversion_scratch_directory_cannot_be_deleted(models_path: Path, store: ModelRecordServiceSQL) -> None: _write_model_file(models_path / CONVERSION_SCRATCH_DIRNAME / "active-conversion" / "sd-v1-5") - result = _service(models_path, db).delete_orphaned_models([CONVERSION_SCRATCH_DIRNAME]) + result = _service(models_path, store).delete_orphaned_models([CONVERSION_SCRATCH_DIRNAME]) assert result[CONVERSION_SCRATCH_DIRNAME].startswith("error:") assert (models_path / CONVERSION_SCRATCH_DIRNAME).exists() +def test_a_registered_model_is_no_orphan_even_when_its_config_no_longer_validates( + models_path: Path, store: ModelRecordServiceSQL, database: Database +) -> None: + _write_model_file(models_path / "future-model") + # A record of a model type this version does not know, as after going back to an older version: deleting its + # files would break the model for the version that wrote it. + record = { + "key": "future", + "hash": "blake3:0", + "base": "any", + "type": "from_the_future", + "format": "checkpoint", + "name": "future model", + "source": "somewhere", + "source_type": "path", + "file_size": 18, + "path": "future-model/model.safetensors", + } + database.queries.models.insert("future", json.dumps(record)) + + assert _service(models_path, store).find_orphaned_models() == [] + + +def test_a_model_registered_by_its_absolute_path_is_no_orphan(models_path: Path, store: ModelRecordServiceSQL) -> None: + # An in-place install records the file where it is, which may be under the models root. + _write_model_file(models_path / "in-place") + store.add_model( + TI_File_SD1_Config( + path=str(models_path / "in-place" / "model.safetensors"), + name="in place", + hash="ABC123", + file_size=18, + source="test/source/", + source_type=ModelSourceType.Path, + ) + ) + + assert _service(models_path, store).find_orphaned_models() == [] + + +def test_a_directory_holding_a_registered_model_is_no_orphan_whatever_else_it_holds( + models_path: Path, store: ModelRecordServiceSQL +) -> None: + # Orphans are reported, and deleted, by their directory under the models root: reporting this one for its loose + # weights would delete the registered model with them. + _write_model_file(models_path / "family" / "model-a") + _write_model_file(models_path / "family" / "loose-weights") + store.add_model( + TI_File_SD1_Config( + path="family/model-a/model.safetensors", + name="model a", + hash="ABC123", + file_size=18, + source="test/source/", + source_type=ModelSourceType.Path, + ) + ) + + assert _service(models_path, store).find_orphaned_models() == [] + + @pytest.mark.parametrize("root_name", ["tmpinstall_recovery", "recovered-model"]) def test_install_recovery_root_is_not_reported_or_deleted( - root_name: str, models_path: Path, db: SqliteDatabase + root_name: str, models_path: Path, store: ModelRecordServiceSQL ) -> None: recovery_root = models_path / root_name _write_model_file(recovery_root) sentinel = models_path / f".{root_name}{INSTALL_RECOVERY_SENTINEL}" sentinel.write_text("preserve", encoding="utf-8") - service = _service(models_path, db) + service = _service(models_path, store) assert service.find_orphaned_models() == [] result = service.delete_orphaned_models([root_name]) @@ -95,7 +159,7 @@ def test_install_recovery_root_is_not_reported_or_deleted( assert sentinel.exists() -def test_parent_of_recovery_root_cannot_be_deleted(models_path: Path, db: SqliteDatabase) -> None: +def test_parent_of_recovery_root_cannot_be_deleted(models_path: Path, store: ModelRecordServiceSQL) -> None: parent = models_path / "model-group" _write_model_file(parent / "ordinary-orphan") recovery_root = parent / "recovered-model" @@ -103,7 +167,7 @@ def test_parent_of_recovery_root_cannot_be_deleted(models_path: Path, db: Sqlite sentinel = parent / f".{recovery_root.name}{INSTALL_RECOVERY_SENTINEL}" sentinel.write_text("preserve", encoding="utf-8") - service = _service(models_path, db) + service = _service(models_path, store) assert "model-group" not in {orphan.path for orphan in service.find_orphaned_models()} result = service.delete_orphaned_models(["model-group"]) @@ -114,10 +178,10 @@ def test_parent_of_recovery_root_cannot_be_deleted(models_path: Path, db: Sqlite assert sentinel.exists() -def test_models_root_cannot_be_deleted_as_an_orphan(models_path: Path, db: SqliteDatabase) -> None: +def test_models_root_cannot_be_deleted_as_an_orphan(models_path: Path, store: ModelRecordServiceSQL) -> None: _write_model_file(models_path / "real-model") - result = _service(models_path, db).delete_orphaned_models(["."]) + result = _service(models_path, store).delete_orphaned_models(["."]) assert result["."].startswith("error:") assert (models_path / "real-model" / "model.safetensors").exists() diff --git a/tests/app/services/project_records/test_project_records.py b/tests/app/services/project_records/test_project_records.py index aebea1ffbf7..cedab6d8780 100644 --- a/tests/app/services/project_records/test_project_records.py +++ b/tests/app/services/project_records/test_project_records.py @@ -1,11 +1,22 @@ -"""Tests for the project records service: CRUD, optimistic concurrency, boards, and user isolation.""" +"""Project records on every database backend: CRUD, optimistic concurrency, boards, and user isolation.""" + +from collections.abc import Sequence +from typing import Any, Optional import pytest +from sqlalchemy import func, insert, select, update -from invokeai.app.services.config.config_default import InvokeAIAppConfig -from invokeai.app.services.project_records import project_records_sqlite +from invokeai.app.services.board_records.board_records_common import ( + BOARD_NAME_MAX_LENGTH, + BoardChanges, + BoardRecordProjectOwnedException, + BoardVisibility, +) +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage +from invokeai.app.services.project_records import project_records_default from invokeai.app.services.project_records.project_records_common import ( ProjectBoardNotFoundError, + ProjectBoardTooLargeError, ProjectBoardUnavailableError, ProjectCanvasSchemaDowngradeError, ProjectCanvasSchemaUnsupportedError, @@ -15,37 +26,37 @@ ProjectRecordExistsError, ProjectRecordNotFoundError, ) -from invokeai.app.services.project_records.project_records_sqlite import ProjectRecordsSqlite -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.users.users_common import UserCreateRequest -from invokeai.app.services.users.users_default import UserService -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database +from invokeai.app.services.project_records.project_records_default import ProjectRecordsStorage +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import ConflictError +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.boards import BoardQueries +from invokeai.app.services.shared.database.queries.projects import ProjectQueries +from invokeai.app.services.shared.database.schema.boards import board_images, board_videos, boards, shared_boards +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.media_references import media_references +from invokeai.app.services.shared.database.schema.projects import projects +from invokeai.app.services.shared.database.schema.users import users +from invokeai.app.services.shared.database.schema.videos import videos +from tests.fixtures.races import when_called, while_in_flight SYSTEM_USER_ID = "system" +OTHER_USER_ID = "other" @pytest.fixture -def db() -> SqliteDatabase: - config = InvokeAIAppConfig(use_memory_db=True) - return create_mock_sqlite_database(config=config, logger=InvokeAILogger.get_logger()) - - -@pytest.fixture -def project_records(db: SqliteDatabase) -> ProjectRecordsSqlite: - return ProjectRecordsSqlite(db=db) +def project_records(database: Database) -> ProjectRecordsStorage: + return ProjectRecordsStorage(database) @pytest.fixture -def other_user_id(db: SqliteDatabase) -> str: - users = UserService(db=db) - user = users.create( - UserCreateRequest(email="other@example.com", display_name="Other", password="TestPass123", is_admin=False) - ) - return user.user_id +def other_user_id(database: Database) -> str: + with database.begin(write=True) as conn: + conn.execute(insert(users).values(user_id=OTHER_USER_ID, email="other@example.com", password_hash="-")) + return OTHER_USER_ID -def test_create_and_get_roundtrip(project_records: ProjectRecordsSqlite) -> None: +def test_create_and_get_roundtrip(project_records: ProjectRecordsStorage) -> None: data = {"layout": {"centerViewId": "canvas"}, "widgets": [1, 2, 3], "nested": {"a": None, "b": True}} created = project_records.create(SYSTEM_USER_ID, "My Project", data) @@ -59,24 +70,23 @@ def test_create_and_get_roundtrip(project_records: ProjectRecordsSqlite) -> None def test_oversized_documents_are_rejected_before_a_board_is_created( - project_records: ProjectRecordsSqlite, db: SqliteDatabase, monkeypatch: pytest.MonkeyPatch + project_records: ProjectRecordsStorage, database: Database, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr(project_records_sqlite, "PROJECT_DOCUMENT_MAX_BYTES", 16) + monkeypatch.setattr(project_records_default, "PROJECT_DOCUMENT_MAX_BYTES", 16) + before = _board_count(database) with pytest.raises(ProjectDocumentTooLargeError) as exc_info: project_records.create(SYSTEM_USER_ID, "Too large", {"value": "x" * 32}) assert exc_info.value.max_bytes == 16 - with db.transaction() as cursor: - cursor.execute("SELECT COUNT(*) FROM boards WHERE board_name = ?", ("Too large",)) - assert cursor.fetchone()[0] == 0 + assert _board_count(database) == before def test_oversized_updates_preserve_the_acknowledged_revision( - project_records: ProjectRecordsSqlite, monkeypatch: pytest.MonkeyPatch + project_records: ProjectRecordsStorage, monkeypatch: pytest.MonkeyPatch ) -> None: created = project_records.create(SYSTEM_USER_ID, "Project", {"value": "small"}) - monkeypatch.setattr(project_records_sqlite, "PROJECT_DOCUMENT_MAX_BYTES", 16) + monkeypatch.setattr(project_records_default, "PROJECT_DOCUMENT_MAX_BYTES", 16) with pytest.raises(ProjectDocumentTooLargeError): project_records.update( @@ -94,9 +104,9 @@ def test_oversized_updates_preserve_the_acknowledged_revision( def test_document_limit_counts_compact_utf8_bytes( - project_records: ProjectRecordsSqlite, monkeypatch: pytest.MonkeyPatch + project_records: ProjectRecordsStorage, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr(project_records_sqlite, "PROJECT_DOCUMENT_MAX_BYTES", 14) + monkeypatch.setattr(project_records_default, "PROJECT_DOCUMENT_MAX_BYTES", 14) created = project_records.create(SYSTEM_USER_ID, "Exact", {"value": "é"}) assert created.data == {"value": "é"} @@ -108,16 +118,25 @@ def test_document_limit_counts_compact_utf8_bytes( @pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf"), "\ud800"]) -def test_non_json_safe_document_values_are_rejected(project_records: ProjectRecordsSqlite, value: object) -> None: +def test_non_json_safe_document_values_are_rejected(project_records: ProjectRecordsStorage, value: object) -> None: with pytest.raises(ProjectDocumentInvalidError): project_records.create(SYSTEM_USER_ID, "Invalid", {"value": value}) +def test_a_document_is_stored_as_written(project_records: ProjectRecordsStorage) -> None: + """Text beyond the Basic Multilingual Plane and characters JSON escapes come back unchanged on every backend.""" + data = {"name": 'Kürbis 🎃 \u0000 "quoted" \\ back', "list": ["日本語", 1.5]} + + created = project_records.create(SYSTEM_USER_ID, "Unicode", data) + + assert project_records.get(SYSTEM_USER_ID, created.project_id).data == data + + def test_oversized_legacy_documents_remain_readable_and_deletable( - project_records: ProjectRecordsSqlite, monkeypatch: pytest.MonkeyPatch + project_records: ProjectRecordsStorage, monkeypatch: pytest.MonkeyPatch ) -> None: created = project_records.create(SYSTEM_USER_ID, "Legacy", {"value": "x" * 32}) - monkeypatch.setattr(project_records_sqlite, "PROJECT_DOCUMENT_MAX_BYTES", 16) + monkeypatch.setattr(project_records_default, "PROJECT_DOCUMENT_MAX_BYTES", 16) assert project_records.get(SYSTEM_USER_ID, created.project_id).data == {"value": "x" * 32} project_records.delete(SYSTEM_USER_ID, created.project_id) @@ -126,7 +145,7 @@ def test_oversized_legacy_documents_remain_readable_and_deletable( project_records.get(SYSTEM_USER_ID, created.project_id) -def test_create_with_client_id_and_duplicate_rejected(project_records: ProjectRecordsSqlite) -> None: +def test_create_with_client_id_and_duplicate_rejected(project_records: ProjectRecordsStorage) -> None: created = project_records.create(SYSTEM_USER_ID, "Imported", {"x": 1}, project_id="project-abc") assert created.project_id == "project-abc" @@ -134,7 +153,9 @@ def test_create_with_client_id_and_duplicate_rejected(project_records: ProjectRe project_records.create(SYSTEM_USER_ID, "Imported again", {"x": 2}, project_id="project-abc") -def test_same_project_id_allowed_for_different_users(project_records: ProjectRecordsSqlite, other_user_id: str) -> None: +def test_same_project_id_allowed_for_different_users( + project_records: ProjectRecordsStorage, other_user_id: str +) -> None: project_records.create(SYSTEM_USER_ID, "Mine", {"owner": "system"}, project_id="project-shared-id") other = project_records.create(other_user_id, "Theirs", {"owner": "other"}, project_id="project-shared-id") @@ -142,20 +163,42 @@ def test_same_project_id_allowed_for_different_users(project_records: ProjectRec assert project_records.get(other_user_id, other.project_id).data == {"owner": "other"} -def test_list_returns_summaries_for_own_projects_only( - project_records: ProjectRecordsSqlite, other_user_id: str +def test_list_returns_summaries_of_own_projects_oldest_first( + project_records: ProjectRecordsStorage, database: Database, other_user_id: str ) -> None: - first = project_records.create(SYSTEM_USER_ID, "First", {"n": 1}) - second = project_records.create(SYSTEM_USER_ID, "Second", {"n": 2}) + project_records.create(SYSTEM_USER_ID, "Newer", {"n": 1}, project_id="project-a") + project_records.create(SYSTEM_USER_ID, "Older", {"n": 2}, project_id="project-b") project_records.create(other_user_id, "Not mine", {"n": 3}) + with database.begin(write=True) as conn: + for project_id, created_at in ( + ("project-b", "2026-01-01 00:00:00.000"), + ("project-a", "2026-01-01 00:00:00.001"), + ): + conn.execute(update(projects).where(projects.c.project_id == project_id).values(created_at=created_at)) summaries = project_records.list(SYSTEM_USER_ID) - assert [summary.project_id for summary in summaries] == [first.project_id, second.project_id] + assert [summary.project_id for summary in summaries] == ["project-b", "project-a"] assert all(not hasattr(summary, "data") for summary in summaries) -def test_update_increments_revision(project_records: ProjectRecordsSqlite) -> None: +def test_projects_created_in_the_same_millisecond_are_listed_by_id( + project_records: ProjectRecordsStorage, database: Database +) -> None: + """Oldest first; a tie is broken by the project id, as no backend but SQLite keeps an insertion order.""" + for project_id in ("project-b", "project-c", "project-a"): + project_records.create(SYSTEM_USER_ID, project_id, {}, project_id=project_id) + with database.begin(write=True) as conn: + conn.execute(update(projects).values(created_at="2026-01-01 00:00:00.000")) + + assert [summary.project_id for summary in project_records.list(SYSTEM_USER_ID)] == [ + "project-a", + "project-b", + "project-c", + ] + + +def test_update_increments_revision(project_records: ProjectRecordsStorage) -> None: created = project_records.create(SYSTEM_USER_ID, "Project", {"v": 1}) updated = project_records.update( @@ -165,10 +208,23 @@ def test_update_increments_revision(project_records: ProjectRecordsSqlite) -> No assert updated.revision == 2 assert updated.name == "Renamed" assert updated.data == {"v": 2} + assert project_records.get(SYSTEM_USER_ID, created.project_id) == updated + + +def test_a_save_renews_when_the_project_was_updated(project_records: ProjectRecordsStorage, database: Database) -> None: + # Only the application sets it on a server. (On SQLite a trigger renews it as well.) + project_records.create(SYSTEM_USER_ID, "Project", {}, project_id="project") + with database.begin(write=True) as conn: + conn.execute(update(projects).values(updated_at="2000-01-01 00:00:00.000")) + + saved = project_records.update(SYSTEM_USER_ID, "project", expected_revision=1, name="Project", data={"v": 2}) + + assert saved.updated_at > "2000-01-01 00:00:00.000" + assert project_records.get(SYSTEM_USER_ID, "project").updated_at == saved.updated_at def test_writes_return_their_row_without_a_post_commit_read( - project_records: ProjectRecordsSqlite, monkeypatch: pytest.MonkeyPatch + project_records: ProjectRecordsStorage, monkeypatch: pytest.MonkeyPatch ) -> None: def reject_post_commit_get(*_args: object, **_kwargs: object) -> None: raise AssertionError("writes must not open a second transaction to build their response") @@ -190,7 +246,7 @@ def reject_post_commit_get(*_args: object, **_kwargs: object) -> None: assert updated.data == {"v": 2} -def test_projects_default_to_canvas_schema_v2(project_records: ProjectRecordsSqlite) -> None: +def test_projects_default_to_canvas_schema_v2(project_records: ProjectRecordsStorage) -> None: created = project_records.create(SYSTEM_USER_ID, "Project", {}) assert created.minimum_canvas_schema_version == 2 @@ -199,11 +255,9 @@ def test_projects_default_to_canvas_schema_v2(project_records: ProjectRecordsSql def test_create_refuses_a_schema_the_client_cannot_edit_without_creating_a_board( - project_records: ProjectRecordsSqlite, db: SqliteDatabase + project_records: ProjectRecordsStorage, database: Database ) -> None: - with db.transaction() as cursor: - cursor.execute("SELECT COUNT(*) FROM boards;") - before = cursor.fetchone()[0] + before = _board_count(database) with pytest.raises(ProjectCanvasSchemaUnsupportedError): project_records.create( @@ -214,12 +268,10 @@ def test_create_refuses_a_schema_the_client_cannot_edit_without_creating_a_board max_canvas_schema_version=2, ) - with db.transaction() as cursor: - cursor.execute("SELECT COUNT(*) FROM boards;") - assert cursor.fetchone()[0] == before + assert _board_count(database) == before -def test_an_incapable_client_cannot_read_or_write_a_newer_project(project_records: ProjectRecordsSqlite) -> None: +def test_an_incapable_client_cannot_read_or_write_a_newer_project(project_records: ProjectRecordsStorage) -> None: created = project_records.create( SYSTEM_USER_ID, "Future", @@ -246,7 +298,7 @@ def test_an_incapable_client_cannot_read_or_write_a_newer_project(project_record def test_update_atomically_raises_the_canvas_schema_floor_with_the_document( - project_records: ProjectRecordsSqlite, + project_records: ProjectRecordsStorage, ) -> None: created = project_records.create(SYSTEM_USER_ID, "Project", {"canvas": {"version": 2}}) @@ -267,7 +319,7 @@ def test_update_atomically_raises_the_canvas_schema_floor_with_the_document( project_records.get(SYSTEM_USER_ID, created.project_id) -def test_a_stale_save_cannot_raise_the_canvas_schema_floor(project_records: ProjectRecordsSqlite) -> None: +def test_a_stale_save_cannot_raise_the_canvas_schema_floor(project_records: ProjectRecordsStorage) -> None: created = project_records.create(SYSTEM_USER_ID, "Project", {"canvas": {"version": 2}}) project_records.update(SYSTEM_USER_ID, created.project_id, 1, "Project", {"canvas": {"version": 2}}) @@ -289,7 +341,7 @@ def test_a_stale_save_cannot_raise_the_canvas_schema_floor(project_records: Proj def test_a_capable_stale_client_gets_a_revision_conflict_before_floor_downgrade_validation( - project_records: ProjectRecordsSqlite, + project_records: ProjectRecordsStorage, ) -> None: created = project_records.create(SYSTEM_USER_ID, "Project", {"canvas": {"version": 2}}) project_records.update( @@ -319,7 +371,7 @@ def test_a_capable_stale_client_gets_a_revision_conflict_before_floor_downgrade_ assert preserved.data == {"canvas": {"version": 3}} -def test_the_canvas_schema_floor_cannot_be_lowered(project_records: ProjectRecordsSqlite) -> None: +def test_the_canvas_schema_floor_cannot_be_lowered(project_records: ProjectRecordsStorage) -> None: created = project_records.create( SYSTEM_USER_ID, "Future", @@ -345,7 +397,7 @@ def test_the_canvas_schema_floor_cannot_be_lowered(project_records: ProjectRecor assert preserved.data == {"canvas": {"version": 3}} -def test_update_with_stale_revision_raises_conflict(project_records: ProjectRecordsSqlite) -> None: +def test_update_with_stale_revision_raises_conflict(project_records: ProjectRecordsStorage) -> None: created = project_records.create(SYSTEM_USER_ID, "Project", {"v": 1}) project_records.update(SYSTEM_USER_ID, created.project_id, expected_revision=1, name="Project", data={"v": 2}) @@ -357,17 +409,17 @@ def test_update_with_stale_revision_raises_conflict(project_records: ProjectReco assert project_records.get(SYSTEM_USER_ID, created.project_id).data == {"v": 2} -def test_update_missing_project_raises_not_found(project_records: ProjectRecordsSqlite) -> None: +def test_update_missing_project_raises_not_found(project_records: ProjectRecordsStorage) -> None: with pytest.raises(ProjectRecordNotFoundError): project_records.update(SYSTEM_USER_ID, "does-not-exist", expected_revision=1, name="x", data={}) -def test_get_missing_project_raises_not_found(project_records: ProjectRecordsSqlite) -> None: +def test_get_missing_project_raises_not_found(project_records: ProjectRecordsStorage) -> None: with pytest.raises(ProjectRecordNotFoundError): project_records.get(SYSTEM_USER_ID, "does-not-exist") -def test_delete_is_idempotent(project_records: ProjectRecordsSqlite) -> None: +def test_delete_is_idempotent(project_records: ProjectRecordsStorage) -> None: created = project_records.create(SYSTEM_USER_ID, "Doomed", {}) project_records.delete(SYSTEM_USER_ID, created.project_id) @@ -377,7 +429,7 @@ def test_delete_is_idempotent(project_records: ProjectRecordsSqlite) -> None: project_records.get(SYSTEM_USER_ID, created.project_id) -def test_users_cannot_touch_each_others_projects(project_records: ProjectRecordsSqlite, other_user_id: str) -> None: +def test_users_cannot_touch_each_others_projects(project_records: ProjectRecordsStorage, other_user_id: str) -> None: created = project_records.create(SYSTEM_USER_ID, "Private", {"secret": True}) with pytest.raises(ProjectRecordNotFoundError): @@ -392,11 +444,11 @@ def test_users_cannot_touch_each_others_projects(project_records: ProjectRecords assert project_records.get(SYSTEM_USER_ID, created.project_id).data == {"secret": True} -# --- boards ----------------------------------------------------------------------------------- +# region boards def _insert_board( - db: SqliteDatabase, + database: Database, board_id: str, *, user_id: str = SYSTEM_USER_ID, @@ -404,93 +456,69 @@ def _insert_board( visibility: str = "private", archived: bool = False, ) -> str: - with db.transaction() as cursor: - cursor.execute( - "INSERT INTO boards (board_id, board_name, user_id, board_visibility, archived) VALUES (?, ?, ?, ?, ?);", - (board_id, name, user_id, visibility, archived), + with database.begin(write=True) as conn: + conn.execute( + insert(boards).values( + board_id=board_id, board_name=name, user_id=user_id, board_visibility=visibility, archived=archived + ) ) return board_id -def _board_is_archived(db: SqliteDatabase, board_id: str) -> bool: - with db.transaction() as cursor: - cursor.execute("SELECT archived FROM boards WHERE board_id = ?;", (board_id,)) - row = cursor.fetchone() - return bool(row[0]) - - -def _board_name(db: SqliteDatabase, board_id: str) -> str | None: - with db.transaction() as cursor: - cursor.execute("SELECT board_name FROM boards WHERE board_id = ?;", (board_id,)) - row = cursor.fetchone() - return None if row is None else row[0] +def _board(database: Database, board_id: str) -> Optional[tuple[str, bool]]: + """The board's name and archived flag.""" + with database.begin(write=False) as conn: + row = conn.execute(select(boards.c.board_name, boards.c.archived).where(boards.c.board_id == board_id)).first() + return None if row is None else (row[0], bool(row[1])) -def _board_count(db: SqliteDatabase) -> int: - with db.transaction() as cursor: - cursor.execute("SELECT COUNT(*) FROM boards;") - return cursor.fetchone()[0] - - -def _add_image_to_board(db: SqliteDatabase, image_name: str, board_id: str) -> None: - with db.transaction() as cursor: - cursor.execute( - "INSERT INTO images (image_name, image_origin, image_category, width, height)" - " VALUES (?, 'internal', 'general', 64, 64);", - (image_name,), - ) - cursor.execute("INSERT INTO board_images (board_id, image_name) VALUES (?, ?);", (board_id, image_name)) +def _board_count(database: Database) -> int: + with database.begin(write=False) as conn: + return conn.execute(select(func.count()).select_from(boards)).scalar_one() def test_create_makes_one_board_named_after_the_project( - project_records: ProjectRecordsSqlite, db: SqliteDatabase + project_records: ProjectRecordsStorage, database: Database ) -> None: - before = _board_count(db) + before = _board_count(database) created = project_records.create(SYSTEM_USER_ID, "My Project", {}) - assert _board_count(db) == before + 1 - assert created.board_id - assert _board_name(db, created.board_id) == "My Project" + assert _board_count(database) == before + 1 + assert _board(database, created.board_id) == ("My Project", False) assert project_records.get_board_id(SYSTEM_USER_ID, created.project_id) == created.board_id def test_create_claims_a_supplied_board_and_renames_it( - project_records: ProjectRecordsSqlite, db: SqliteDatabase + project_records: ProjectRecordsStorage, database: Database ) -> None: - _insert_board(db, "staging-board", name="Untitled") - before = _board_count(db) + _insert_board(database, "staging-board", name="Untitled") + before = _board_count(database) created = project_records.create(SYSTEM_USER_ID, "Imported", {}, board_id="staging-board") assert created.board_id == "staging-board" - assert _board_name(db, "staging-board") == "Imported" + assert _board(database, "staging-board") == ("Imported", False) # Claiming reuses the board rather than making another. - assert _board_count(db) == before + assert _board_count(database) == before -def test_claiming_a_board_un_archives_it(project_records: ProjectRecordsSqlite, db: SqliteDatabase) -> None: +def test_claiming_a_board_un_archives_it(project_records: ProjectRecordsStorage, database: Database) -> None: """A project's board follows the project, and `PATCH /boards/{id}` refuses `archived` on a claimed board — so a board claimed while archived would be invisible in every listing with no route left to bring it back.""" - _insert_board(db, "staging-board", name="Untitled", archived=True) + _insert_board(database, "staging-board", name="Untitled", archived=True) created = project_records.create(SYSTEM_USER_ID, "Imported", {}, board_id="staging-board") assert created.board_id == "staging-board" - assert _board_is_archived(db, "staging-board") is False - - -def test_a_project_creates_its_board_un_archived(project_records: ProjectRecordsSqlite, db: SqliteDatabase) -> None: - created = project_records.create(SYSTEM_USER_ID, "My Project", {}) - - assert _board_is_archived(db, created.board_id) is False + assert _board(database, "staging-board") == ("Imported", False) def test_claiming_a_missing_or_foreign_board_reports_it_as_missing( - project_records: ProjectRecordsSqlite, db: SqliteDatabase, other_user_id: str + project_records: ProjectRecordsStorage, database: Database, other_user_id: str ) -> None: - _insert_board(db, "theirs", user_id=other_user_id) + _insert_board(database, "theirs", user_id=other_user_id) with pytest.raises(ProjectBoardNotFoundError): project_records.create(SYSTEM_USER_ID, "P", {}, board_id="no-such-board") @@ -502,66 +530,66 @@ def test_claiming_a_missing_or_foreign_board_reports_it_as_missing( @pytest.mark.parametrize("visibility", ["shared", "public"]) def test_claiming_a_non_private_board_is_rejected( - project_records: ProjectRecordsSqlite, db: SqliteDatabase, visibility: str + project_records: ProjectRecordsStorage, database: Database, visibility: str ) -> None: - _insert_board(db, "visible", visibility=visibility) + _insert_board(database, "visible", visibility=visibility) with pytest.raises(ProjectBoardUnavailableError): project_records.create(SYSTEM_USER_ID, "P", {}, board_id="visible") def test_claiming_an_explicitly_shared_board_is_rejected( - project_records: ProjectRecordsSqlite, db: SqliteDatabase, other_user_id: str + project_records: ProjectRecordsStorage, database: Database, other_user_id: str ) -> None: - _insert_board(db, "lent-out") - with db.transaction() as cursor: - cursor.execute("INSERT INTO shared_boards (board_id, user_id) VALUES ('lent-out', ?);", (other_user_id,)) + _insert_board(database, "lent-out") + with database.begin(write=True) as conn: + conn.execute(insert(shared_boards).values(board_id="lent-out", user_id=other_user_id)) with pytest.raises(ProjectBoardUnavailableError): project_records.create(SYSTEM_USER_ID, "P", {}, board_id="lent-out") -def test_a_board_can_only_ever_be_claimed_once(project_records: ProjectRecordsSqlite, db: SqliteDatabase) -> None: - _insert_board(db, "contested") +def test_a_board_can_only_ever_be_claimed_once(project_records: ProjectRecordsStorage, database: Database) -> None: + _insert_board(database, "contested") project_records.create(SYSTEM_USER_ID, "First", {}, board_id="contested") with pytest.raises(ProjectBoardUnavailableError): project_records.create(SYSTEM_USER_ID, "Second", {}, board_id="contested") - assert _board_name(db, "contested") == "First" + assert _board(database, "contested") == ("First", False) -def test_a_rejected_create_leaves_no_orphan_board(project_records: ProjectRecordsSqlite, db: SqliteDatabase) -> None: +def test_a_rejected_create_leaves_no_orphan_board(project_records: ProjectRecordsStorage, database: Database) -> None: project_records.create(SYSTEM_USER_ID, "Taken", {}, project_id="project-taken") - before = _board_count(db) + before = _board_count(database) with pytest.raises(ProjectRecordExistsError): project_records.create(SYSTEM_USER_ID, "Again", {}, project_id="project-taken") # The board insert shares the failed insert's transaction, so it rolled back with it. - assert _board_count(db) == before + assert _board_count(database) == before -def test_a_rejected_claim_leaves_the_board_name_alone( - project_records: ProjectRecordsSqlite, db: SqliteDatabase +def test_a_rejected_claim_leaves_the_board_as_it_was( + project_records: ProjectRecordsStorage, database: Database ) -> None: - _insert_board(db, "staging", name="Untitled") + _insert_board(database, "staging", name="Untitled", archived=True) project_records.create(SYSTEM_USER_ID, "Taken", {}, project_id="project-taken") with pytest.raises(ProjectRecordExistsError): project_records.create(SYSTEM_USER_ID, "Renamer", {}, project_id="project-taken", board_id="staging") - assert _board_name(db, "staging") == "Untitled" + assert _board(database, "staging") == ("Untitled", True) def test_repeating_a_create_says_the_project_exists_rather_than_blaming_the_board( - project_records: ProjectRecordsSqlite, db: SqliteDatabase, other_user_id: str + project_records: ProjectRecordsStorage, database: Database, other_user_id: str ) -> None: """A client whose create response was lost re-sends it to find out whether the first one landed, and the answer decides whether it deletes the media it already uploaded. Reporting the board as unavailable would read as "somebody else took it" — the opposite of what happened, and the reading that gets the uploads thrown away.""" - _insert_board(db, "staging", name="Untitled") + _insert_board(database, "staging", name="Untitled") project_records.create(SYSTEM_USER_ID, "Imported", {}, project_id="project-import", board_id="staging") with pytest.raises(ProjectRecordExistsError): @@ -570,63 +598,285 @@ def test_repeating_a_create_says_the_project_exists_rather_than_blaming_the_boar # A different account claiming the same board is still refused: the exclusion is for the project # being created, not for board claims in general. (Foreign ownership is caught first, as a 404, # so the board is re-pointed at the other account to reach the claimed check.) - db._conn.execute("UPDATE boards SET user_id = ? WHERE board_id = 'staging';", (other_user_id,)) - db._conn.commit() + with database.begin(write=True) as conn: + conn.execute(update(boards).where(boards.c.board_id == "staging").values(user_id=other_user_id)) with pytest.raises(ProjectBoardUnavailableError): project_records.create(other_user_id, "Theirs", {}, project_id="project-import", board_id="staging") def test_rename_renames_the_board_but_only_on_a_winning_save( - project_records: ProjectRecordsSqlite, db: SqliteDatabase + project_records: ProjectRecordsStorage, database: Database ) -> None: created = project_records.create(SYSTEM_USER_ID, "Before", {}) project_records.update(SYSTEM_USER_ID, created.project_id, expected_revision=1, name="After", data={}) - assert _board_name(db, created.board_id) == "After" + assert _board(database, created.board_id) == ("After", False) with pytest.raises(ProjectRecordConflictError): project_records.update(SYSTEM_USER_ID, created.project_id, expected_revision=1, name="Loser", data={}) - assert _board_name(db, created.board_id) == "After" + assert _board(database, created.board_id) == ("After", False) + + +def test_a_long_project_name_is_cut_to_the_board_name_limit( + project_records: ProjectRecordsStorage, database: Database +) -> None: + limit = BOARD_NAME_MAX_LENGTH + + created = project_records.create(SYSTEM_USER_ID, "n" * (limit + 1), {}) + assert _board(database, created.board_id) == ("n" * limit, False) + + project_records.update(SYSTEM_USER_ID, created.project_id, expected_revision=1, name="m" * (limit + 1), data={}) + assert _board(database, created.board_id) == ("m" * limit, False) + assert project_records.get(SYSTEM_USER_ID, created.project_id).name == "m" * (limit + 1) def test_delete_removes_the_board_and_its_membership_but_keeps_the_media( - project_records: ProjectRecordsSqlite, db: SqliteDatabase + project_records: ProjectRecordsStorage, database: Database ) -> None: created = project_records.create(SYSTEM_USER_ID, "Doomed", {}) - _add_image_to_board(db, "kept.png", created.board_id) + _put_image(database, "kept.png", created.board_id) project_records.delete(SYSTEM_USER_ID, created.project_id) - assert _board_name(db, created.board_id) is None - with db.transaction() as cursor: - cursor.execute("SELECT COUNT(*) FROM board_images WHERE image_name = 'kept.png';") - assert cursor.fetchone()[0] == 0 + assert _board(database, created.board_id) is None + with database.begin(write=False) as conn: + assert conn.execute(select(board_images.c.image_name)).first() is None # The image itself survives, uncategorized. - cursor.execute("SELECT COUNT(*) FROM images WHERE image_name = 'kept.png';") - assert cursor.fetchone()[0] == 1 + assert conn.execute(select(images.c.image_name)).scalars().all() == ["kept.png"] -def test_summaries_and_records_both_carry_the_board(project_records: ProjectRecordsSqlite) -> None: +def test_summaries_and_records_both_carry_the_board(project_records: ProjectRecordsStorage) -> None: created = project_records.create(SYSTEM_USER_ID, "Listed", {}) (summary,) = [s for s in project_records.list(SYSTEM_USER_ID) if s.project_id == created.project_id] assert summary.board_id == created.board_id -def test_get_board_id_is_user_scoped(project_records: ProjectRecordsSqlite, other_user_id: str) -> None: +def test_get_board_id_is_user_scoped(project_records: ProjectRecordsStorage, other_user_id: str) -> None: created = project_records.create(SYSTEM_USER_ID, "Private", {}) with pytest.raises(ProjectRecordNotFoundError): project_records.get_board_id(other_user_id, created.project_id) -# --- board snapshot --------------------------------------------------------------------------- +# endregion + +# region writes racing each other +# +# A write locks the rows it decides on before it reads them, a project's before its board's: a write in flight is +# either seen by the write that waited for it or finds what it decided on unchanged, and no two writes deadlock. + + +def _replace_project(q: Queries, *, minimum_canvas_schema_version: int = 2) -> None: + """Deletes the project `project` and creates it again under the same id, with the board `new-board`, as another + client can.""" + project = q.projects.lock(SYSTEM_USER_ID, "project") + assert project is not None + q.boards.lock(project.board_id) + q.projects.delete(SYSTEM_USER_ID, "project") + q.boards.delete_if_unclaimed(project.board_id) + q.boards.insert(board_id="new-board", board_name="New", user_id=SYSTEM_USER_ID) + q.projects.insert( + project_id="project", + user_id=SYSTEM_USER_ID, + name="New", + data='{"new":true}', + board_id="new-board", + minimum_canvas_schema_version=minimum_canvas_schema_version, + ) + + +@pytest.mark.parametrize("in_flight", ["publish", "delete", "claim"]) +def test_a_claim_sees_the_board_change_it_waited_for( + project_records: ProjectRecordsStorage, + database: Database, + monkeypatch: pytest.MonkeyPatch, + lost_races: list[ConflictError], + in_flight: str, +) -> None: + _insert_board(database, "staging", name="Untitled") + + def claim() -> object: + return project_records.create(SYSTEM_USER_ID, "Second", {}, project_id="second", board_id="staging") + + refusal: type[Exception] + if in_flight == "publish": + ended = when_called(monkeypatch, BoardQueries, "update", claim) + BoardRecordStorage(database).update("staging", BoardChanges(board_visibility=BoardVisibility.Public)) + errors, refusal = ended(), ProjectBoardUnavailableError + elif in_flight == "delete": + # The boards service deletes a board with this one statement. + errors = while_in_flight(database, lambda q: q.boards.delete_if_unclaimed("staging"), claim) + refusal = ProjectBoardNotFoundError + else: + ended = when_called(monkeypatch, ProjectQueries, "insert", claim) + project_records.create(SYSTEM_USER_ID, "First", {}, project_id="first", board_id="staging") + errors, refusal = ended(), ProjectBoardUnavailableError + + assert [type(e) for e in errors] == [refusal] + assert lost_races == [] + with pytest.raises(ProjectRecordNotFoundError): + project_records.get(SYSTEM_USER_ID, "second") + + +@pytest.mark.parametrize("change", ["publish", "delete"]) +def test_a_board_change_waits_for_a_claim_in_flight_and_then_finds_the_board_claimed( + project_records: ProjectRecordsStorage, + database: Database, + monkeypatch: pytest.MonkeyPatch, + lost_races: list[ConflictError], + change: str, +) -> None: + _insert_board(database, "staging", name="Untitled") + board_records = BoardRecordStorage(database) + deleted: list[bool] = [] + + def change_board() -> None: + if change == "delete": + deleted.append(board_records.delete_if_unclaimed("staging")) + else: + board_records.update("staging", BoardChanges(board_visibility=BoardVisibility.Public)) + + ended = when_called(monkeypatch, ProjectQueries, "insert", change_board) + project_records.create(SYSTEM_USER_ID, "First", {}, project_id="first", board_id="staging") + errors = ended() + + if change == "delete": + assert (errors, deleted) == ([], [False]) + else: + assert [type(e) for e in errors] == [BoardRecordProjectOwnedException] + assert lost_races == [] + assert _board(database, "staging") == ("First", False) + assert project_records.get_board_id(SYSTEM_USER_ID, "first") == "staging" + + +@pytest.mark.parametrize("in_flight", ["save", "delete"]) +def test_a_save_waits_for_a_write_of_its_project_in_flight( + project_records: ProjectRecordsStorage, + database: Database, + monkeypatch: pytest.MonkeyPatch, + lost_races: list[ConflictError], + in_flight: str, +) -> None: + created = project_records.create(SYSTEM_USER_ID, "Created", {"v": 0}, project_id="project") + + def late_save() -> object: + return project_records.update(SYSTEM_USER_ID, "project", expected_revision=1, name="Late", data={"v": "late"}) + + if in_flight == "save": + ended = when_called(monkeypatch, ProjectQueries, "save", late_save) + project_records.update(SYSTEM_USER_ID, "project", expected_revision=1, name="First", data={"v": "first"}) + (error,) = ended() + assert isinstance(error, ProjectRecordConflictError) and error.current_revision == 2 + stored = project_records.get(SYSTEM_USER_ID, "project") + assert (stored.name, stored.data, stored.revision) == ("First", {"v": "first"}, 2) + assert _board(database, created.board_id) == ("First", False) + else: + ended = when_called(monkeypatch, ProjectQueries, "delete", late_save) + project_records.delete(SYSTEM_USER_ID, "project") + assert [type(e) for e in ended()] == [ProjectRecordNotFoundError] + assert lost_races == [] + + +@pytest.mark.parametrize("client_schema", [2, 3]) +def test_a_save_checks_and_writes_the_project_it_waited_for( + project_records: ProjectRecordsStorage, database: Database, lost_races: list[ConflictError], client_schema: int +) -> None: + """A save that waited for another client to replace its project under the same id checks the replacement -- + an older client must not write over a project of a newer canvas schema -- and saves into it, board and all.""" + project_records.create(SYSTEM_USER_ID, "Old", {"old": True}, project_id="project") + + errors = while_in_flight( + database, + lambda q: _replace_project(q, minimum_canvas_schema_version=3), + lambda: project_records.update( + SYSTEM_USER_ID, + "project", + expected_revision=1, + name="Saved", + data={"saved": True}, + max_canvas_schema_version=client_schema, + ), + ) + + stored = project_records.get(SYSTEM_USER_ID, "project", max_canvas_schema_version=3) + if client_schema == 2: + assert [type(e) for e in errors] == [ProjectCanvasSchemaUnsupportedError] + assert (stored.data, stored.minimum_canvas_schema_version, stored.revision) == ({"new": True}, 3, 1) + assert _board(database, "new-board") == ("New", False) + else: + assert errors == [] + assert (stored.data, stored.minimum_canvas_schema_version, stored.revision) == ({"saved": True}, 3, 2) + assert _board(database, "new-board") == ("Saved", False) + assert lost_races == [] + + +def test_deleting_a_project_deletes_the_board_it_locked( + project_records: ProjectRecordsStorage, database: Database, lost_races: list[ConflictError] +) -> None: + """A delete reads the board to delete from the project row it locks: a delete that waited for another client to + replace the project deletes the replacement's board, not the one it saw first.""" + project_records.create(SYSTEM_USER_ID, "First", {}, project_id="project") + + errors = while_in_flight(database, _replace_project, lambda: project_records.delete(SYSTEM_USER_ID, "project")) + + assert (errors, lost_races) == ([], []) + assert project_records.list(SYSTEM_USER_ID) == [] + assert _board(database, "new-board") is None + + +def test_a_resent_claim_waits_for_a_save_of_its_project_in_flight( + project_records: ProjectRecordsStorage, + database: Database, + monkeypatch: pytest.MonkeyPatch, + lost_races: list[ConflictError], +) -> None: + _insert_board(database, "staging", name="Untitled") + project_records.create(SYSTEM_USER_ID, "Imported", {}, project_id="project", board_id="staging") + + ended = when_called( + monkeypatch, + ProjectQueries, + "save", + lambda: project_records.create(SYSTEM_USER_ID, "Imported", {}, project_id="project", board_id="staging"), + ) + project_records.update(SYSTEM_USER_ID, "project", expected_revision=1, name="Saved", data={}) + + assert [type(e) for e in ended()] == [ProjectRecordExistsError] + assert lost_races == [] + assert _board(database, "staging") == ("Saved", False) + + +def test_a_board_delete_waits_for_a_delete_of_its_project_in_flight( + project_records: ProjectRecordsStorage, + database: Database, + monkeypatch: pytest.MonkeyPatch, + lost_races: list[ConflictError], +) -> None: + created = project_records.create(SYSTEM_USER_ID, "Doomed", {}, project_id="project") + deleted: list[bool] = [] + + ended = when_called( + monkeypatch, + ProjectQueries, + "delete", + lambda: deleted.append(BoardRecordStorage(database).delete_if_unclaimed(created.board_id)), + ) + project_records.delete(SYSTEM_USER_ID, "project") + + assert (ended(), deleted, lost_races) == ([], [False], []) + assert _board(database, created.board_id) is None + + +# endregion + +# region board snapshot def _put_image( - db: SqliteDatabase, + database: Database, name: str, board_id: str, *, @@ -634,17 +884,23 @@ def _put_image( starred: bool = False, is_intermediate: bool = False, ) -> None: - with db.transaction() as cursor: - cursor.execute( - "INSERT INTO images (image_name, image_origin, image_category, width, height, starred, is_intermediate)" - " VALUES (?, 'internal', ?, 64, 64, ?, ?);", - (name, category, starred, is_intermediate), + with database.begin(write=True) as conn: + conn.execute( + insert(images).values( + image_name=name, + image_origin="internal", + image_category=category, + width=64, + height=64, + starred=starred, + is_intermediate=is_intermediate, + ) ) - cursor.execute("INSERT INTO board_images (board_id, image_name) VALUES (?, ?);", (board_id, name)) + conn.execute(insert(board_images).values(board_id=board_id, image_name=name)) def _put_video( - db: SqliteDatabase, + database: Database, name: str, board_id: str, *, @@ -652,22 +908,28 @@ def _put_video( starred: bool = False, is_intermediate: bool = False, ) -> None: - with db.transaction() as cursor: - cursor.execute( - "INSERT INTO videos (video_name, video_origin, video_category, width, height, starred, is_intermediate)" - " VALUES (?, 'internal', ?, 64, 64, ?, ?);", - (name, category, starred, is_intermediate), + with database.begin(write=True) as conn: + conn.execute( + insert(videos).values( + video_name=name, + video_origin="internal", + video_category=category, + width=64, + height=64, + starred=starred, + is_intermediate=is_intermediate, + ) ) - cursor.execute("INSERT INTO board_videos (board_id, video_name) VALUES (?, ?);", (board_id, name)) + conn.execute(insert(board_videos).values(board_id=board_id, video_name=name)) def test_the_snapshot_lists_every_visible_category_of_both_kinds( - project_records: ProjectRecordsSqlite, db: SqliteDatabase + project_records: ProjectRecordsStorage, database: Database ) -> None: created = project_records.create(SYSTEM_USER_ID, "Full", {}) for category in ("general", "control", "mask", "user"): - _put_image(db, f"{category}.png", created.board_id, category=category) - _put_video(db, f"{category}.mp4", created.board_id, category=category) + _put_image(database, f"{category}.png", created.board_id, category=category) + _put_video(database, f"{category}.mp4", created.board_id, category=category) snapshot = project_records.get_board_snapshot(SYSTEM_USER_ID, created.project_id) @@ -679,27 +941,29 @@ def test_the_snapshot_lists_every_visible_category_of_both_kinds( def test_the_snapshot_excludes_what_the_gallery_does_not_show( - project_records: ProjectRecordsSqlite, db: SqliteDatabase + project_records: ProjectRecordsStorage, database: Database ) -> None: created = project_records.create(SYSTEM_USER_ID, "Filtered", {}) - _put_image(db, "shown.png", created.board_id) + _put_image(database, "shown.png", created.board_id) # `other` is the canvas's private category — in neither gallery view, so not board membership. - _put_image(db, "canvas.png", created.board_id, category="other") - _put_image(db, "scratch.png", created.board_id, is_intermediate=True) - _put_video(db, "shown.mp4", created.board_id) - _put_video(db, "canvas.mp4", created.board_id, category="other") - _put_video(db, "scratch.mp4", created.board_id, is_intermediate=True) + _put_image(database, "canvas.png", created.board_id, category="other") + _put_image(database, "scratch.png", created.board_id, is_intermediate=True) + _put_video(database, "shown.mp4", created.board_id) + _put_video(database, "canvas.mp4", created.board_id, category="other") + _put_video(database, "scratch.mp4", created.board_id, is_intermediate=True) snapshot = project_records.get_board_snapshot(SYSTEM_USER_ID, created.project_id) assert [item.name for item in snapshot.items] == ["shown.png", "shown.mp4"] -def test_the_snapshot_excludes_media_on_other_boards(project_records: ProjectRecordsSqlite, db: SqliteDatabase) -> None: +def test_the_snapshot_excludes_media_on_other_boards( + project_records: ProjectRecordsStorage, database: Database +) -> None: mine = project_records.create(SYSTEM_USER_ID, "Mine", {}) theirs = project_records.create(SYSTEM_USER_ID, "Other", {}) - _put_image(db, "mine.png", mine.board_id) - _put_image(db, "theirs.png", theirs.board_id) + _put_image(database, "mine.png", mine.board_id) + _put_image(database, "theirs.png", theirs.board_id) snapshot = project_records.get_board_snapshot(SYSTEM_USER_ID, mine.project_id) @@ -707,16 +971,19 @@ def test_the_snapshot_excludes_media_on_other_boards(project_records: ProjectRec def test_the_snapshot_carries_starring_and_is_ordered_by_kind_then_name( - project_records: ProjectRecordsSqlite, db: SqliteDatabase + project_records: ProjectRecordsStorage, database: Database ) -> None: created = project_records.create(SYSTEM_USER_ID, "Ordered", {}) - _put_image(db, "b.png", created.board_id, starred=True) - _put_image(db, "a.png", created.board_id) - _put_video(db, "a.mp4", created.board_id) + # Upper case sorts before lower case: names are compared byte for byte on every backend. + _put_image(database, "b.png", created.board_id, starred=True) + _put_image(database, "a.png", created.board_id) + _put_image(database, "C.png", created.board_id) + _put_video(database, "a.mp4", created.board_id) snapshot = project_records.get_board_snapshot(SYSTEM_USER_ID, created.project_id) assert [(item.kind, item.name, item.starred) for item in snapshot.items] == [ + ("image", "C.png", False), ("image", "a.png", False), ("image", "b.png", True), ("video", "a.mp4", False), @@ -724,26 +991,42 @@ def test_the_snapshot_carries_starring_and_is_ordered_by_kind_then_name( def test_a_same_name_image_and_video_are_separate_entries( - project_records: ProjectRecordsSqlite, db: SqliteDatabase + project_records: ProjectRecordsStorage, database: Database ) -> None: """Images and videos are separate namespaces, so one name can legitimately be both.""" created = project_records.create(SYSTEM_USER_ID, "Twins", {}) - _put_image(db, "twin", created.board_id) - _put_video(db, "twin", created.board_id) + _put_image(database, "twin", created.board_id) + _put_video(database, "twin", created.board_id) snapshot = project_records.get_board_snapshot(SYSTEM_USER_ID, created.project_id) assert [(item.kind, item.name) for item in snapshot.items] == [("image", "twin"), ("video", "twin")] -def test_an_empty_board_snapshots_to_an_empty_list(project_records: ProjectRecordsSqlite) -> None: +def test_an_empty_board_snapshots_to_an_empty_list(project_records: ProjectRecordsStorage) -> None: created = project_records.create(SYSTEM_USER_ID, "Empty", {}) assert project_records.get_board_snapshot(SYSTEM_USER_ID, created.project_id).items == [] +def test_a_board_over_the_snapshot_limit_is_refused( + project_records: ProjectRecordsStorage, database: Database, monkeypatch: pytest.MonkeyPatch +) -> None: + created = project_records.create(SYSTEM_USER_ID, "Full", {}) + for name in ("a.png", "b.png"): + _put_image(database, name, created.board_id) + _put_video(database, "c.mp4", created.board_id) + + monkeypatch.setattr(project_records_default, "PROJECT_BOARD_SNAPSHOT_MAX_ITEMS", 3) + assert len(project_records.get_board_snapshot(SYSTEM_USER_ID, created.project_id).items) == 3 + + monkeypatch.setattr(project_records_default, "PROJECT_BOARD_SNAPSHOT_MAX_ITEMS", 2) + with pytest.raises(ProjectBoardTooLargeError): + project_records.get_board_snapshot(SYSTEM_USER_ID, created.project_id) + + def test_snapshotting_someone_elses_project_is_not_found( - project_records: ProjectRecordsSqlite, other_user_id: str + project_records: ProjectRecordsStorage, other_user_id: str ) -> None: created = project_records.create(SYSTEM_USER_ID, "Private", {}) @@ -751,24 +1034,30 @@ def test_snapshotting_someone_elses_project_is_not_found( project_records.get_board_snapshot(other_user_id, created.project_id) -def _media_references(db: SqliteDatabase, project_id: str) -> set[tuple[str, str]]: - with db.transaction() as cursor: - cursor.execute( - "SELECT media_kind, media_name FROM media_references WHERE owner_kind = 'project' AND owner_id = ?;", - (project_id,), - ) - return {tuple(row) for row in cursor.fetchall()} +# endregion + +# region media references + + +def _media_references(database: Database, project_id: str) -> set[tuple[str, str]]: + with database.begin(write=False) as conn: + rows: Sequence[Sequence[Any]] = conn.execute( + select(media_references.c.media_kind, media_references.c.media_name).where( + media_references.c.owner_kind == "project", media_references.c.owner_id == project_id + ) + ).all() + return {(kind, name) for kind, name in rows} def test_saving_a_project_indexes_the_media_it_references( - db: SqliteDatabase, project_records: ProjectRecordsSqlite + database: Database, project_records: ProjectRecordsStorage ) -> None: created = project_records.create( SYSTEM_USER_ID, "Refs", {"canvas": {"stagingArea": {"pendingImages": [{"imageName": "staged.png"}]}}, "video": {"video_name": "a.mp4"}}, ) - assert _media_references(db, created.project_id) == {("image", "staged.png"), ("video", "a.mp4")} + assert _media_references(database, created.project_id) == {("image", "staged.png"), ("video", "a.mp4")} project_records.update( SYSTEM_USER_ID, @@ -777,11 +1066,11 @@ def test_saving_a_project_indexes_the_media_it_references( name="Refs", data={"layers": [{"image_name": "layer.png"}]}, ) - assert _media_references(db, created.project_id) == {("image", "layer.png")} + assert _media_references(database, created.project_id) == {("image", "layer.png")} def test_a_refused_save_leaves_the_reference_index_untouched( - db: SqliteDatabase, project_records: ProjectRecordsSqlite + database: Database, project_records: ProjectRecordsStorage ) -> None: created = project_records.create(SYSTEM_USER_ID, "Refs", {"image_name": "first.png"}) @@ -790,12 +1079,41 @@ def test_a_refused_save_leaves_the_reference_index_untouched( SYSTEM_USER_ID, created.project_id, expected_revision=99, name="Refs", data={"image_name": "second.png"} ) - assert _media_references(db, created.project_id) == {("image", "first.png")} + assert _media_references(database, created.project_id) == {("image", "first.png")} + + +def test_a_refused_create_writes_no_references(database: Database, project_records: ProjectRecordsStorage) -> None: + project_records.create(SYSTEM_USER_ID, "Refs", {"image_name": "first.png"}, project_id="project") + + with pytest.raises(ProjectRecordExistsError): + project_records.create(SYSTEM_USER_ID, "Refs", {"image_name": "second.png"}, project_id="project") + + assert _media_references(database, "project") == {("image", "first.png")} + + +def test_deleting_a_missing_project_leaves_references_under_its_id_alone( + database: Database, project_records: ProjectRecordsStorage +) -> None: + """References go only with the project they belong to: on a server, a delete that finds no project can run while + the project is created, and must not take the references that creation writes.""" + with database.begin(write=True) as conn: + conn.execute( + insert(media_references).values( + owner_kind="project", user_id=SYSTEM_USER_ID, owner_id="project", media_kind="image", media_name="a.png" + ) + ) + + project_records.delete(SYSTEM_USER_ID, "project") + + assert _media_references(database, "project") == {("image", "a.png")} -def test_deleting_a_project_drops_its_references(db: SqliteDatabase, project_records: ProjectRecordsSqlite) -> None: +def test_deleting_a_project_drops_its_references(database: Database, project_records: ProjectRecordsStorage) -> None: created = project_records.create(SYSTEM_USER_ID, "Refs", {"image_name": "first.png"}) project_records.delete(SYSTEM_USER_ID, created.project_id) - assert _media_references(db, created.project_id) == set() + assert _media_references(database, created.project_id) == set() + + +# endregion diff --git a/tests/app/services/session_processor/test_for_loop_processor_sqlite.py b/tests/app/services/session_processor/test_for_loop_processor_sqlite.py index bb5e1b2eba2..2705f76d533 100644 --- a/tests/app/services/session_processor/test_for_loop_processor_sqlite.py +++ b/tests/app/services/session_processor/test_for_loop_processor_sqlite.py @@ -21,7 +21,8 @@ DefaultSessionProcessor, DefaultSessionRunner, ) -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.execution_effects import ExecutionEffectsRecorder, ExecutionInterface from invokeai.app.services.shared.execution_state_migration import dump_execution_state from invokeai.app.services.shared.graph import ( @@ -33,6 +34,7 @@ _GenericGraphSchedulerAdapter, ) from invokeai.app.services.shared.invocation_context import InvocationContext +from tests.fixtures.sqlite_database import sqlite_cursor_of from tests.test_nodes import create_edge, create_loop_linkage @@ -268,14 +270,14 @@ def _stop_processor(processor: DefaultSessionProcessor) -> None: assert not worker.thread.is_alive() -def _insert_session(queue: SqliteSessionQueue, graph: Graph, *, versioned: bool = False) -> int: +def _insert_session(queue: SessionQueue, graph: Graph, *, versioned: bool = False) -> int: session = GraphExecutionState(graph=graph) session_json = ( json.dumps(dump_execution_state(session)) if versioned else session.model_dump_json(warnings=False, exclude_none=True) ) - with queue._db.transaction() as cursor: + with sqlite_cursor_of(queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -331,6 +333,7 @@ def reset_stats(self, graph_execution_state_id) -> None: def test_processor_sqlite_queue_nested_iterate_for_cleanup( monkeypatch: pytest.MonkeyPatch, mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, outcome: str, ) -> None: @@ -340,7 +343,7 @@ def test_processor_sqlite_queue_nested_iterate_for_cleanup( _build_test_invocation_context, ) - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() @@ -415,6 +418,7 @@ async def recording_status_handler(event) -> None: def test_processor_sqlite_nested_iterate_for_cancel_retry_reloads_fresh_stream_state( monkeypatch: pytest.MonkeyPatch, mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, ) -> None: monkeypatch.setattr( @@ -422,7 +426,7 @@ def test_processor_sqlite_nested_iterate_for_cancel_retry_reloads_fresh_stream_s _build_test_invocation_context, ) - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() @@ -463,9 +467,9 @@ def cancel_after_first_inner_return(invocation, queue_item, output) -> None: # Simulate an interrupted process after the partial state was persisted. Startup must cancel # the stale row while preserving the nested frame/stream snapshot for inspection or retry. - with queue._db.transaction() as cursor: + with sqlite_cursor_of(queue) as cursor: cursor.execute("UPDATE session_queue SET status = 'in_progress' WHERE item_id = ?", (item_id,)) - restarted_queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + restarted_queue = SessionQueue(mock_sqlite_database) restarted_queue.start(mock_invoker) reloaded_canceled_item = restarted_queue.get_queue_item(item_id) assert reloaded_canceled_item.status == "canceled" @@ -520,6 +524,7 @@ def cancel_after_first_inner_return(invocation, queue_item, output) -> None: def test_processor_sqlite_three_level_nested_for_cancel_restart_retry_isolates_execution_state( monkeypatch: pytest.MonkeyPatch, mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, ) -> None: monkeypatch.setattr( @@ -527,7 +532,7 @@ def test_processor_sqlite_three_level_nested_for_cancel_restart_retry_isolates_e _build_test_invocation_context, ) - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() @@ -587,9 +592,9 @@ def cancel_after_first_inner_return(invocation, queue_item, output) -> None: for reference in canceled_item.session.execution_refs.values() ) - with queue._db.transaction() as cursor: + with sqlite_cursor_of(queue) as cursor: cursor.execute("UPDATE session_queue SET status = 'in_progress' WHERE item_id = ?", (item_id,)) - restarted_queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + restarted_queue = SessionQueue(mock_sqlite_database) restarted_queue.start(mock_invoker) mock_invoker.services.session_queue = restarted_queue reloaded_canceled_item = restarted_queue.get_queue_item(item_id) @@ -667,6 +672,7 @@ def cancel_after_first_inner_return(invocation, queue_item, output) -> None: def test_processor_sqlite_four_level_nested_for_cancel_restart_retry_isolates_execution_state( monkeypatch: pytest.MonkeyPatch, mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, ) -> None: monkeypatch.setattr( @@ -674,7 +680,7 @@ def test_processor_sqlite_four_level_nested_for_cancel_restart_retry_isolates_ex _build_test_invocation_context, ) - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() @@ -735,9 +741,9 @@ def cancel_after_first_deepest_return(invocation, queue_item, output) -> None: for reference in canceled_item.session.execution_refs.values() ) - with queue._db.transaction() as cursor: + with sqlite_cursor_of(queue) as cursor: cursor.execute("UPDATE session_queue SET status = 'in_progress' WHERE item_id = ?", (item_id,)) - restarted_queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + restarted_queue = SessionQueue(mock_sqlite_database) restarted_queue.start(mock_invoker) mock_invoker.services.session_queue = restarted_queue reloaded_canceled_item = restarted_queue.get_queue_item(item_id) @@ -820,6 +826,7 @@ def cancel_after_first_deepest_return(invocation, queue_item, output) -> None: def test_processor_sqlite_two_sibling_nested_for_cancel_root_retry_cleans_identity( monkeypatch: pytest.MonkeyPatch, mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, ) -> None: test_timeout = 30 @@ -828,7 +835,7 @@ def test_processor_sqlite_two_sibling_nested_for_cancel_root_retry_cleans_identi _build_test_invocation_context, ) - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() @@ -885,9 +892,9 @@ def cancel_after_first_child_finalization(invocation, queue_item, output) -> Non for reference in canceled_item.session.execution_refs.values() ) - with queue._db.transaction() as cursor: + with sqlite_cursor_of(queue) as cursor: cursor.execute("UPDATE session_queue SET status = 'in_progress' WHERE item_id = ?", (item_id,)) - restarted_queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + restarted_queue = SessionQueue(mock_sqlite_database) restarted_queue.start(mock_invoker) mock_invoker.services.session_queue = restarted_queue reloaded_canceled_item = restarted_queue.get_queue_item(item_id) @@ -963,6 +970,7 @@ def cancel_after_first_child_finalization(invocation, queue_item, output) -> Non def test_processor_sqlite_two_sibling_nested_for_failure_cleans_runtime( monkeypatch: pytest.MonkeyPatch, mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, ) -> None: monkeypatch.setattr( @@ -970,7 +978,7 @@ def test_processor_sqlite_two_sibling_nested_for_failure_cleans_runtime( _build_test_invocation_context, ) - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() @@ -1043,6 +1051,7 @@ def capture_set_queue_item_session(item_id: int, session: GraphExecutionState): def test_processor_sqlite_three_level_nested_for_failure_cleans_runtime( monkeypatch: pytest.MonkeyPatch, mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, ) -> None: monkeypatch.setattr( @@ -1050,7 +1059,7 @@ def test_processor_sqlite_three_level_nested_for_failure_cleans_runtime( _build_test_invocation_context, ) - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() @@ -1119,7 +1128,7 @@ def capture_set_queue_item_session(item_id: int, session: GraphExecutionState): assert failed_item.error_message == "Refusing loop value 1" assert queue.get_current("default") is None - reloaded_queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + reloaded_queue = SessionQueue(mock_sqlite_database) reloaded_queue.start(mock_invoker) reloaded_item = reloaded_queue.get_queue_item(item_id) assert reloaded_item.status == "failed" @@ -1149,6 +1158,7 @@ def capture_set_queue_item_session(item_id: int, session: GraphExecutionState): def test_processor_sqlite_four_level_nested_for_failure_cleans_runtime( monkeypatch: pytest.MonkeyPatch, mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, ) -> None: monkeypatch.setattr( @@ -1156,7 +1166,7 @@ def test_processor_sqlite_four_level_nested_for_failure_cleans_runtime( _build_test_invocation_context, ) - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() @@ -1225,7 +1235,7 @@ def capture_set_queue_item_session(item_id: int, session: GraphExecutionState): assert failed_item.error_message == "Refusing loop value 1" assert queue.get_current("default") is None - reloaded_queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + reloaded_queue = SessionQueue(mock_sqlite_database) reloaded_queue.start(mock_invoker) reloaded_item = reloaded_queue.get_queue_item(item_id) assert reloaded_item.status == "failed" @@ -1290,6 +1300,7 @@ def capture_set_queue_item_session(item_id: int, session: GraphExecutionState): def test_processor_sqlite_queue_nested_for_cleanup( monkeypatch: pytest.MonkeyPatch, mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, outcome: str, ) -> None: @@ -1299,7 +1310,7 @@ def test_processor_sqlite_queue_nested_for_cleanup( _build_test_invocation_context, ) - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() @@ -1402,6 +1413,7 @@ async def recording_status_handler(event) -> None: def test_processor_sqlite_iterate_collect_cancel_retry_does_not_leak_stream_state( monkeypatch: pytest.MonkeyPatch, mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, ) -> None: test_timeout = 30 @@ -1410,7 +1422,7 @@ def test_processor_sqlite_iterate_collect_cancel_retry_does_not_leak_stream_stat _build_test_invocation_context, ) - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() @@ -1508,6 +1520,7 @@ def test_processor_sqlite_iterate_collect_cancel_retry_does_not_leak_stream_stat def test_processor_sqlite_direct_planner_cancel_retry_isolates_execution_state( monkeypatch: pytest.MonkeyPatch, mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, cancel_after: str, ) -> None: @@ -1517,7 +1530,7 @@ def test_processor_sqlite_direct_planner_cancel_retry_isolates_execution_state( _build_test_invocation_context, ) - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() @@ -1640,6 +1653,7 @@ def cancel_at_boundary(invocation, queue_item, output) -> None: def test_processor_sqlite_flat_for_cancel_retry_isolates_execution_state( monkeypatch: pytest.MonkeyPatch, mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, force_compatibility_scheduler: bool, ) -> None: @@ -1650,7 +1664,7 @@ def test_processor_sqlite_flat_for_cancel_retry_isolates_execution_state( if force_compatibility_scheduler: monkeypatch.setattr(GraphExecutionState, "_can_use_generic_for_scheduler", lambda self: False) - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() diff --git a/tests/app/services/session_processor/test_if_processor_sqlite.py b/tests/app/services/session_processor/test_if_processor_sqlite.py index 9751aca1847..effa5ea349e 100644 --- a/tests/app/services/session_processor/test_if_processor_sqlite.py +++ b/tests/app/services/session_processor/test_if_processor_sqlite.py @@ -17,8 +17,10 @@ DefaultSessionProcessor, DefaultSessionRunner, ) -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from tests.fixtures.sqlite_database import sqlite_cursor_of from tests.test_nodes import create_edge @@ -72,9 +74,9 @@ def _build_if_graph(*, condition: bool = True) -> Graph: return graph -def _insert_session(queue: SqliteSessionQueue, graph: Graph) -> int: +def _insert_session(queue: SessionQueue, graph: Graph) -> int: session = GraphExecutionState(graph=graph) - with queue._db.transaction() as cursor: + with sqlite_cursor_of(queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -147,6 +149,7 @@ def registered_event_bus() -> Iterator[_RecordingRegisteredEventService]: ) def test_if_cancellation_preserves_activation_and_never_resumes_unselected_branch( mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, condition: bool, selected_branch: str, @@ -154,7 +157,7 @@ def test_if_cancellation_preserves_activation_and_never_resumes_unselected_branc selected_field: str, ) -> None: test_timeout = 30 - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() @@ -214,6 +217,7 @@ def cancel_after_if(invocation, queue_item, output) -> None: ) def test_if_retry_starts_fresh_and_preserves_selected_output_without_stale_activation( mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, condition: bool, selected_branch: str, @@ -222,7 +226,7 @@ def test_if_retry_starts_fresh_and_preserves_selected_output_without_stale_activ sink_value: int, cancel_after_source: str, ) -> None: - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() @@ -287,13 +291,14 @@ def test_if_retry_starts_fresh_and_preserves_selected_output_without_stale_activ ) def test_if_early_cancellation_before_resolution_does_not_run_any_branch_continuation( mock_invoker: Invoker, + mock_sqlite_database: Database, registered_event_bus: _RecordingRegisteredEventService, condition: bool, selected_branch: str, unselected_branch: str, selected_field: str, ) -> None: - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) mock_invoker.services.events = registered_event_bus mock_invoker.services.session_queue = queue mock_invoker.services.performance_statistics = _Stats() diff --git a/tests/app/services/session_processor/test_session_processor_image_maintenance.py b/tests/app/services/session_processor/test_session_processor_image_maintenance.py index c08959280b7..a5e5be2cf16 100644 --- a/tests/app/services/session_processor/test_session_processor_image_maintenance.py +++ b/tests/app/services/session_processor/test_session_processor_image_maintenance.py @@ -8,7 +8,7 @@ from invokeai.app.services.image_files.image_files_disk import DiskImageFileStorage from invokeai.app.services.image_moves.image_moves_default import ImageMoveQueueActive, ImageMoveService from invokeai.app.services.session_processor.session_processor_default import DefaultSessionProcessor, _SessionWorker -from invokeai.app.services.shared.sqlite.sqlite_util import init_db +from invokeai.app.services.shared.database.database import Database from invokeai.backend.util.logging import InvokeAILogger @@ -32,13 +32,12 @@ def dequeue(self, device: str | None = None): return self.item -def _services(tmp_path: Path, queue: _Queue) -> tuple[DefaultSessionProcessor, ImageMoveService]: +def _services(tmp_path: Path, database: Database, queue: _Queue) -> tuple[DefaultSessionProcessor, ImageMoveService]: config = InvokeAIAppConfig(use_memory_db=True) config._root = tmp_path logger = InvokeAILogger.get_logger(config=config) image_files = DiskImageFileStorage(tmp_path / "images") - db = init_db(config=config, logger=logger, image_files=image_files) - moves = ImageMoveService(db=db, image_files=image_files, config=config, logger=logger) + moves = ImageMoveService(database, image_files=image_files, config=config, logger=logger) invoker = SimpleNamespace(services=SimpleNamespace(session_queue=queue)) moves.start(invoker) @@ -49,10 +48,12 @@ def _services(tmp_path: Path, queue: _Queue) -> tuple[DefaultSessionProcessor, I return processor, moves -def test_in_flight_queue_claim_is_visible_before_gallery_maintenance_can_start(tmp_path: Path) -> None: +def test_in_flight_queue_claim_is_visible_before_gallery_maintenance_can_start( + tmp_path: Path, database: Database +) -> None: queue = _Queue(pending=1) queue.allow_dequeue.clear() - processor, moves = _services(tmp_path, queue) + processor, moves = _services(tmp_path, database, queue) claimed: list[tuple[bool, object | None]] = [] claim_errors: list[Exception] = [] @@ -107,9 +108,9 @@ def observe_reservation_lock(): assert moves.is_maintenance_active() is False -def test_queue_work_admitted_after_reservation_waits_without_being_claimed(tmp_path: Path) -> None: +def test_queue_work_admitted_after_reservation_waits_without_being_claimed(tmp_path: Path, database: Database) -> None: queue = _Queue(pending=0) - processor, moves = _services(tmp_path, queue) + processor, moves = _services(tmp_path, database, queue) try: with moves.reserve_gallery_maintenance(): @@ -126,10 +127,10 @@ def test_queue_work_admitted_after_reservation_waits_without_being_claimed(tmp_p moves.stop() -def test_worker_stops_while_gallery_maintenance_holds_the_mutation_lock(tmp_path: Path) -> None: +def test_worker_stops_while_gallery_maintenance_holds_the_mutation_lock(tmp_path: Path, database: Database) -> None: queue = _Queue(pending=0) queue.allow_dequeue.set() - processor, moves = _services(tmp_path, queue) + processor, moves = _services(tmp_path, database, queue) maintenance_entered = threading.Event() release_maintenance = threading.Event() maintenance_errors: list[Exception] = [] diff --git a/tests/app/services/session_queue/test_for_loop_session_queue.py b/tests/app/services/session_queue/test_for_loop_session_queue.py index e4a7d86ed83..0b4e08d68f6 100644 --- a/tests/app/services/session_queue/test_for_loop_session_queue.py +++ b/tests/app/services/session_queue/test_for_loop_session_queue.py @@ -12,14 +12,16 @@ StateSetInvocation, ) from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from tests.fixtures.sqlite_database import sqlite_cursor_of from tests.test_nodes import AnyTypeTestInvocation, create_edge, create_loop_linkage @pytest.fixture -def session_queue(mock_invoker: Invoker) -> SqliteSessionQueue: - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) +def session_queue(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: + queue = SessionQueue(mock_sqlite_database) queue.start(mock_invoker) return queue @@ -88,11 +90,11 @@ def _empty_for_missing_state_graph() -> Graph: return graph -def _insert_session(queue: SqliteSessionQueue, state: GraphExecutionState) -> int: +def _insert_session(queue: SessionQueue, state: GraphExecutionState) -> int: session_id = str(uuid.uuid4()) batch_id = str(uuid.uuid4()) session_json = state.model_dump_json(warnings=False, exclude_none=True) - with queue._db.transaction() as cursor: + with sqlite_cursor_of(queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -105,7 +107,7 @@ def _insert_session(queue: SqliteSessionQueue, state: GraphExecutionState) -> in return cursor.lastrowid # type: ignore[return-value] -def test_sqlite_queue_resumes_partial_stateful_for_loop(session_queue: SqliteSessionQueue) -> None: +def test_sqlite_queue_resumes_partial_stateful_for_loop(session_queue: SessionQueue) -> None: item_id = _insert_session(session_queue, GraphExecutionState(graph=_stateful_for_graph())) queue_item = session_queue.dequeue() @@ -168,7 +170,7 @@ def test_sqlite_queue_resumes_partial_stateful_for_loop(session_queue: SqliteSes def test_sqlite_queue_round_trips_a_checkpoint_before_each_next_node( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: """Every manually persisted node boundary must survive a process restart before the next node. @@ -220,7 +222,7 @@ def test_sqlite_queue_round_trips_a_checkpoint_before_each_next_node( assert final_item.session.is_complete() -def test_sqlite_queue_round_trips_empty_for_final_output(session_queue: SqliteSessionQueue) -> None: +def test_sqlite_queue_round_trips_empty_for_final_output(session_queue: SessionQueue) -> None: item_id = _insert_session(session_queue, GraphExecutionState(graph=_empty_for_graph())) queue_item = session_queue.dequeue() @@ -241,7 +243,7 @@ def test_sqlite_queue_round_trips_empty_for_final_output(session_queue: SqliteSe def test_sqlite_queue_round_trips_missing_loop_state_value_after_empty_for( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: item_id = _insert_session(session_queue, GraphExecutionState(graph=_empty_for_missing_state_graph())) @@ -263,7 +265,7 @@ def test_sqlite_queue_round_trips_missing_loop_state_value_after_empty_for( def test_sqlite_queue_resumes_nested_for_after_first_outer_iteration( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: item_id = _insert_session(session_queue, GraphExecutionState(graph=_nested_for_graph())) diff --git a/tests/app/services/session_queue/test_session_queue_bulk_cancel_events.py b/tests/app/services/session_queue/test_session_queue_bulk_cancel_events.py index d795a6f4b37..f7ce57d4ce5 100644 --- a/tests/app/services/session_queue/test_session_queue_bulk_cancel_events.py +++ b/tests/app/services/session_queue/test_session_queue_bulk_cancel_events.py @@ -13,22 +13,24 @@ from invokeai.app.services.events.events_common import QueueItemsCanceledEvent from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from tests.fixtures.sqlite_database import sqlite_cursor_of from tests.test_nodes import PromptTestInvocation @pytest.fixture -def session_queue(mock_invoker: Invoker) -> SqliteSessionQueue: - """Create a SqliteSessionQueue backed by the mock invoker's in-memory database.""" - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) +def session_queue(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: + """Create a SessionQueue backed by the mock invoker's in-memory database.""" + db = mock_sqlite_database + queue = SessionQueue(db) queue.start(mock_invoker) return queue def _insert_queue_item( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, queue_id: str, user_id: str, batch_id: str | None = None, @@ -38,7 +40,7 @@ def _insert_queue_item( """Directly insert a minimal pending queue item for the given user and return its item_id.""" session_id = str(uuid.uuid4()) batch_id = batch_id or str(uuid.uuid4()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id, user_id) @@ -51,7 +53,7 @@ def _insert_queue_item( def _insert_dequeueable_queue_item( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, queue_id: str, user_id: str, destination: str | None = None, @@ -62,7 +64,7 @@ def _insert_dequeueable_queue_item( session = GraphExecutionState(graph=graph) session_json = session.model_dump_json(warnings=False, exclude_none=True) batch_id = str(uuid.uuid4()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id, user_id) @@ -74,9 +76,9 @@ def _insert_dequeueable_queue_item( return cursor.lastrowid -def _status_of(session_queue: SqliteSessionQueue, item_id: int) -> str: +def _status_of(session_queue: SessionQueue, item_id: int) -> str: """Reads an item's status directly; the minimal inserted rows carry no parseable session.""" - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute("SELECT status FROM session_queue WHERE item_id = ?", (item_id,)) return str(cursor.fetchone()[0]) @@ -86,7 +88,7 @@ def _canceled_events(mock_invoker: Invoker) -> list[QueueItemsCanceledEvent]: def test_cancel_all_except_current_emits_queue_items_canceled_grouped_by_owner( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker ) -> None: """An admin cancel (no user_id) affects every user's pending items; each owner's client needs its own item ids so its queue list and badge counts refetch.""" @@ -107,7 +109,7 @@ def test_cancel_all_except_current_emits_queue_items_canceled_grouped_by_owner( def test_cancel_all_except_current_scoped_to_user_only_names_their_items( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker ) -> None: """A non-admin cancel is scoped to the caller; other users' items are untouched and must not appear in the event.""" @@ -123,7 +125,7 @@ def test_cancel_all_except_current_scoped_to_user_only_names_their_items( def test_cancel_all_except_current_scoped_to_origin_prefix_leaves_other_origins( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker ) -> None: """A project-scoped cancel-all (the webv2 queue widget's default scope) must only sweep items whose origin carries that project's prefix; the prefix is a plain prefix, not a pattern.""" @@ -147,7 +149,7 @@ def test_cancel_all_except_current_scoped_to_origin_prefix_leaves_other_origins( def test_cancel_by_queue_id_scoped_to_origin_prefix_and_user( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker ) -> None: """The scoped cancel-all behind the webv2 queue widget's "Cancel All Items": pending items in the prefix are swept in one statement, other origins and other users' items stay untouched.""" @@ -164,7 +166,7 @@ def test_cancel_by_queue_id_scoped_to_origin_prefix_and_user( def test_delete_all_except_current_emits_queue_items_canceled( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker ) -> None: """Bulk delete removes rows entirely — same notification requirement as bulk cancel.""" a1 = _insert_queue_item(session_queue, "default", "user_a") @@ -178,7 +180,7 @@ def test_delete_all_except_current_emits_queue_items_canceled( assert events[0].canceled_item_ids_by_user == {"user_a": [a1], "user_b": [b1]} -def test_no_event_when_nothing_was_canceled(session_queue: SqliteSessionQueue, mock_invoker: Invoker) -> None: +def test_no_event_when_nothing_was_canceled(session_queue: SessionQueue, mock_invoker: Invoker) -> None: """An empty result must not broadcast a pointless refetch signal to every client.""" result = session_queue.cancel_all_except_current("default") @@ -186,9 +188,7 @@ def test_no_event_when_nothing_was_canceled(session_queue: SqliteSessionQueue, m assert _canceled_events(mock_invoker) == [] -def test_cancel_by_batch_ids_emits_queue_items_canceled( - session_queue: SqliteSessionQueue, mock_invoker: Invoker -) -> None: +def test_cancel_by_batch_ids_emits_queue_items_canceled(session_queue: SessionQueue, mock_invoker: Invoker) -> None: """Canceling batches bulk-updates pending rows with no per-item events; other clients need the bulk event to learn the items left the queue.""" batch_id = str(uuid.uuid4()) @@ -204,9 +204,7 @@ def test_cancel_by_batch_ids_emits_queue_items_canceled( assert events[0].canceled_item_ids_by_user == {"user_a": [a1, a2]} -def test_cancel_by_destination_emits_queue_items_canceled( - session_queue: SqliteSessionQueue, mock_invoker: Invoker -) -> None: +def test_cancel_by_destination_emits_queue_items_canceled(session_queue: SessionQueue, mock_invoker: Invoker) -> None: """Discarding a canvas staging area cancels by destination; the bulk event keeps other clients' queue badges in sync.""" a1 = _insert_queue_item(session_queue, "default", "user_a", destination="canvas:sess-1") @@ -221,9 +219,7 @@ def test_cancel_by_destination_emits_queue_items_canceled( assert events[0].canceled_item_ids_by_user == {"user_a": [a1], "user_b": [b1]} -def test_delete_by_destination_emits_queue_items_canceled( - session_queue: SqliteSessionQueue, mock_invoker: Invoker -) -> None: +def test_delete_by_destination_emits_queue_items_canceled(session_queue: SessionQueue, mock_invoker: Invoker) -> None: """Bulk delete by destination removes rows entirely - same notification requirement.""" a1 = _insert_queue_item(session_queue, "default", "user_a", destination="canvas:sess-1") _insert_queue_item(session_queue, "default", "user_a", destination="other") @@ -237,7 +233,7 @@ def test_delete_by_destination_emits_queue_items_canceled( def test_delete_by_destination_excludes_separately_canceled_current_item( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker ) -> None: """The in-progress current item matching the destination is canceled via cancel_queue_item, which emits its own per-item queue_item_status_changed - the bulk event must not signal the @@ -256,9 +252,7 @@ def test_delete_by_destination_excludes_separately_canceled_current_item( assert events[0].canceled_item_ids_by_user == {"user_a": [pending_id]} -def test_cancel_by_queue_id_emits_queue_items_canceled( - session_queue: SqliteSessionQueue, mock_invoker: Invoker -) -> None: +def test_cancel_by_queue_id_emits_queue_items_canceled(session_queue: SessionQueue, mock_invoker: Invoker) -> None: """Canceling a whole queue bulk-updates every user's non-terminal items.""" a1 = _insert_queue_item(session_queue, "default", "user_a") b1 = _insert_queue_item(session_queue, "default", "user_b") diff --git a/tests/app/services/session_queue/test_session_queue_clear.py b/tests/app/services/session_queue/test_session_queue_clear.py index b0dfab8bea7..e7dce35530b 100644 --- a/tests/app/services/session_queue/test_session_queue_clear.py +++ b/tests/app/services/session_queue/test_session_queue_clear.py @@ -6,21 +6,23 @@ from invokeai.app.services.events.events_common import QueueClearedEvent, QueueItemStatusChangedEvent from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from tests.fixtures.sqlite_database import sqlite_cursor_of from tests.test_nodes import PromptTestInvocation @pytest.fixture -def session_queue(mock_invoker: Invoker) -> SqliteSessionQueue: - """Create a SqliteSessionQueue backed by the mock invoker's in-memory database.""" - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) +def session_queue(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: + """Create a SessionQueue backed by the mock invoker's in-memory database.""" + db = mock_sqlite_database + queue = SessionQueue(db) queue.start(mock_invoker) return queue -def _insert_queue_item(session_queue: SqliteSessionQueue, queue_id: str, user_id: str, status: str = "pending") -> int: +def _insert_queue_item(session_queue: SessionQueue, queue_id: str, user_id: str, status: str = "pending") -> int: """Directly insert a minimal queue item for the given user; returns its item_id. The session payload must be a valid GraphExecutionState: canceling an in-progress item @@ -30,7 +32,7 @@ def _insert_queue_item(session_queue: SqliteSessionQueue, queue_id: str, user_id graph.add_node(PromptTestInvocation(id="prompt", prompt="test")) session = GraphExecutionState(graph=graph) batch_id = str(uuid.uuid4()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id, user_id, status) @@ -55,9 +57,9 @@ def _insert_queue_item(session_queue: SqliteSessionQueue, queue_id: str, user_id return cursor.lastrowid -def _count_items(session_queue: SqliteSessionQueue, queue_id: str, user_id: str | None = None) -> int: +def _count_items(session_queue: SessionQueue, queue_id: str, user_id: str | None = None) -> int: """Count items in the queue, optionally filtered by user_id.""" - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: if user_id is not None: cursor.execute( "SELECT COUNT(*) FROM session_queue WHERE queue_id = ? AND user_id = ?", @@ -71,7 +73,7 @@ def _count_items(session_queue: SqliteSessionQueue, queue_id: str, user_id: str return cursor.fetchone()[0] -def test_clear_with_user_id_only_deletes_own_items(session_queue: SqliteSessionQueue) -> None: +def test_clear_with_user_id_only_deletes_own_items(session_queue: SessionQueue) -> None: """Non-admin clear (user_id provided) should only remove that user's items.""" queue_id = "default" user_a = "user_a" @@ -88,7 +90,7 @@ def test_clear_with_user_id_only_deletes_own_items(session_queue: SqliteSessionQ assert _count_items(session_queue, queue_id, user_b) == 1 -def test_clear_without_user_id_deletes_all_items(session_queue: SqliteSessionQueue) -> None: +def test_clear_without_user_id_deletes_all_items(session_queue: SessionQueue) -> None: """Admin clear (no user_id) should remove all items in the queue.""" queue_id = "default" @@ -102,7 +104,7 @@ def test_clear_without_user_id_deletes_all_items(session_queue: SqliteSessionQue assert _count_items(session_queue, queue_id) == 0 -def test_clear_with_user_id_does_not_affect_other_queues(session_queue: SqliteSessionQueue) -> None: +def test_clear_with_user_id_does_not_affect_other_queues(session_queue: SessionQueue) -> None: """Clearing one queue should not affect items in another queue.""" queue_a = "queue_a" queue_b = "queue_b" @@ -118,7 +120,7 @@ def test_clear_with_user_id_does_not_affect_other_queues(session_queue: SqliteSe assert _count_items(session_queue, queue_b) == 1 -def test_clear_returns_zero_when_no_matching_items(session_queue: SqliteSessionQueue) -> None: +def test_clear_returns_zero_when_no_matching_items(session_queue: SessionQueue) -> None: """Clear should return 0 deleted when there are no items for the given user.""" queue_id = "default" @@ -130,15 +132,15 @@ def test_clear_returns_zero_when_no_matching_items(session_queue: SqliteSessionQ assert _count_items(session_queue, queue_id) == 1 -def _status_of(session_queue: SqliteSessionQueue, item_id: int) -> str | None: - with session_queue._db.transaction() as cursor: +def _status_of(session_queue: SessionQueue, item_id: int) -> str | None: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute("SELECT status FROM session_queue WHERE item_id = ?", (item_id,)) row = cursor.fetchone() return row[0] if row is not None else None def test_user_scoped_clear_cancels_only_that_users_in_progress_items( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker ) -> None: """With multiple workers, several items can be in_progress at once — one per user here. Alice's clear must cancel ALL of Alice's running items (each via its own status-changed @@ -171,7 +173,7 @@ def test_user_scoped_clear_cancels_only_that_users_in_progress_items( assert [(e.queue_id, e.user_id) for e in cleared_events] == [(queue_id, "alice")] -def test_admin_clear_cancels_all_in_progress_items(session_queue: SqliteSessionQueue, mock_invoker: Invoker) -> None: +def test_admin_clear_cancels_all_in_progress_items(session_queue: SessionQueue, mock_invoker: Invoker) -> None: """An unscoped (admin) clear cancels every running item before deleting the rows.""" queue_id = "default" alice_running = _insert_queue_item(session_queue, queue_id, "alice", status="in_progress") diff --git a/tests/app/services/session_queue/test_session_queue_dequeue.py b/tests/app/services/session_queue/test_session_queue_dequeue.py index 8d60c48d4dd..d68cc22bb8e 100644 --- a/tests/app/services/session_queue/test_session_queue_dequeue.py +++ b/tests/app/services/session_queue/test_session_queue_dequeue.py @@ -15,11 +15,8 @@ from invokeai.app.services.events.events_common import QueueItemsRetriedEvent from invokeai.app.services.invoker import Invoker from invokeai.app.services.session_queue.session_queue_common import get_session_for_queue_read -from invokeai.app.services.session_queue.session_queue_sqlite import ( - AFFINITY_MAX_LOOKAHEAD, - ROUND_ROBIN_DEQUEUE_QUERY, - SqliteSessionQueue, -) +from invokeai.app.services.session_queue.session_queue_default import AFFINITY_MAX_LOOKAHEAD, SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.execution_state_migration import ( CURRENT_EXECUTION_STATE_VERSION, dump_execution_state, @@ -33,23 +30,25 @@ GraphExecutionState, IterateInvocation, ) +from tests.fixtures.database import capture_statements, explain_query_plan +from tests.fixtures.sqlite_database import sqlite_cursor_of from tests.test_nodes import TestEventService _EMPTY_SESSION_JSON = json.dumps(to_jsonable_python(GraphExecutionState(graph=Graph()).model_dump())) @pytest.fixture -def session_queue_fifo(mock_invoker: Invoker) -> SqliteSessionQueue: +def session_queue_fifo(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: """Queue backed by a single-user (FIFO) invoker.""" # Default config has multiuser=False, so FIFO is always used. - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) + db = mock_sqlite_database + queue = SessionQueue(db) queue.start(mock_invoker) return queue @pytest.fixture -def session_queue_round_robin(mock_invoker: Invoker) -> SqliteSessionQueue: +def session_queue_round_robin(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: """Queue backed by a multiuser invoker with round_robin mode.""" mock_invoker.services.configuration = InvokeAIAppConfig( use_memory_db=True, @@ -57,8 +56,8 @@ def session_queue_round_robin(mock_invoker: Invoker) -> SqliteSessionQueue: multiuser=True, session_queue_mode="round_robin", ) - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) + db = mock_sqlite_database + queue = SessionQueue(db) queue.start(mock_invoker) return queue @@ -70,7 +69,7 @@ def event_bus(mock_invoker: Invoker) -> TestEventService: def _insert_queue_item( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, queue_id: str, user_id: str, priority: int = 0, @@ -84,7 +83,7 @@ def _insert_queue_item( """ session_id = str(uuid.uuid4()) batch_id = str(uuid.uuid4()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue (item_id, queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id, user_id) @@ -95,7 +94,7 @@ def _insert_queue_item( return cursor.lastrowid # type: ignore[return-value] -def _dequeue_user_ids(session_queue: SqliteSessionQueue, count: int) -> list[Optional[str]]: +def _dequeue_user_ids(session_queue: SessionQueue, count: int) -> list[Optional[str]]: """Dequeue `count` items and return the list of user_ids in dequeue order.""" result = [] for _ in range(count): @@ -109,7 +108,7 @@ def _dequeue_user_ids(session_queue: SqliteSessionQueue, count: int) -> list[Opt # --------------------------------------------------------------------------- -def test_fifo_single_user_order(session_queue_fifo: SqliteSessionQueue) -> None: +def test_fifo_single_user_order(session_queue_fifo: SessionQueue) -> None: """FIFO: items from a single user are dequeued in insertion order.""" queue_id = "default" _insert_queue_item(session_queue_fifo, queue_id, "user_a") @@ -120,7 +119,7 @@ def test_fifo_single_user_order(session_queue_fifo: SqliteSessionQueue) -> None: assert user_ids == ["user_a", "user_a", "user_a"] -def test_fifo_multi_user_preserves_insertion_order(session_queue_fifo: SqliteSessionQueue) -> None: +def test_fifo_multi_user_preserves_insertion_order(session_queue_fifo: SessionQueue) -> None: """FIFO: jobs from multiple users are dequeued in strict insertion order, not interleaved.""" queue_id = "default" # Insert A1, A2, B1, C1, C2, A3 – FIFO should preserve this exact order. @@ -135,7 +134,7 @@ def test_fifo_multi_user_preserves_insertion_order(session_queue_fifo: SqliteSes assert user_ids == ["user_a", "user_a", "user_b", "user_c", "user_c", "user_a"] -def test_fifo_priority_respected(session_queue_fifo: SqliteSessionQueue) -> None: +def test_fifo_priority_respected(session_queue_fifo: SessionQueue) -> None: """FIFO: higher-priority items are dequeued before lower-priority ones.""" queue_id = "default" _insert_queue_item(session_queue_fifo, queue_id, "user_a", priority=0) @@ -146,13 +145,13 @@ def test_fifo_priority_respected(session_queue_fifo: SqliteSessionQueue) -> None assert user_ids == ["user_a", "user_a"] -def test_fifo_returns_none_when_empty(session_queue_fifo: SqliteSessionQueue) -> None: +def test_fifo_returns_none_when_empty(session_queue_fifo: SessionQueue) -> None: """FIFO: dequeue returns None when the queue is empty.""" assert session_queue_fifo.dequeue() is None def test_fifo_quarantines_future_snapshot_and_dequeues_later_work( - session_queue_fifo: SqliteSessionQueue, + session_queue_fifo: SessionQueue, ) -> None: future_session = json.loads(_EMPTY_SESSION_JSON) future_session["execution_state_version"] = CURRENT_EXECUTION_STATE_VERSION + 1 @@ -168,7 +167,7 @@ def test_fifo_quarantines_future_snapshot_and_dequeues_later_work( assert dequeued is not None assert dequeued.item_id == valid_item_id - with session_queue_fifo._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_fifo) as cursor: cursor.execute( "SELECT status, error_type, error_message FROM session_queue WHERE item_id = ?", (future_item_id,), @@ -188,7 +187,7 @@ def test_fifo_quarantines_future_snapshot_and_dequeues_later_work( ], ) def test_fifo_quarantines_unreadable_snapshot_and_dequeues_later_work( - session_queue_fifo: SqliteSessionQueue, + session_queue_fifo: SessionQueue, session_json: str, ) -> None: bad_item_id = _insert_queue_item(session_queue_fifo, "default", "bad-user", session_json=session_json) @@ -198,7 +197,7 @@ def test_fifo_quarantines_unreadable_snapshot_and_dequeues_later_work( assert dequeued is not None assert dequeued.item_id == valid_item_id - with session_queue_fifo._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_fifo) as cursor: cursor.execute("SELECT status, error_message FROM session_queue WHERE item_id = ?", (bad_item_id,)) status, error_message = cursor.fetchone() assert status == "failed" @@ -206,7 +205,7 @@ def test_fifo_quarantines_unreadable_snapshot_and_dequeues_later_work( def test_fifo_quarantines_snapshot_with_invalid_if_collector_roots( - session_queue_fifo: SqliteSessionQueue, + session_queue_fifo: SessionQueue, ) -> None: nodes = [ IntegerInvocation(id="integer", value=1), @@ -262,7 +261,7 @@ def test_fifo_quarantines_snapshot_with_invalid_if_collector_roots( assert dequeued is not None assert dequeued.item_id == valid_item_id - with session_queue_fifo._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_fifo) as cursor: cursor.execute("SELECT status, error_message FROM session_queue WHERE item_id = ?", (invalid_item_id,)) status, error_message = cursor.fetchone() assert status == "failed" @@ -274,7 +273,7 @@ def test_fifo_quarantines_snapshot_with_invalid_if_collector_roots( "history-user", session_json=json.dumps(older_snapshot), ) - with session_queue_fifo._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_fifo) as cursor: cursor.execute("UPDATE session_queue SET status = 'completed' WHERE item_id = ?", (completed_item_id,)) history_item = session_queue_fifo.get_queue_item_for_api(completed_item_id) @@ -286,9 +285,9 @@ def test_fifo_quarantines_snapshot_with_invalid_if_collector_roots( def test_affinity_quarantines_unreadable_snapshot_and_dequeues_valid_work( - session_queue_round_robin: SqliteSessionQueue, + session_queue_round_robin: SessionQueue, ) -> None: - _install_fake_cache(session_queue_round_robin._SqliteSessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) + _install_fake_cache(session_queue_round_robin._SessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) cold_id = _insert_queue_item(session_queue_round_robin, "default", "user_a") future_session = json.loads(_session_with_model_key(_WARM_MODEL_KEY)) future_session["execution_state_version"] = CURRENT_EXECUTION_STATE_VERSION + 1 @@ -303,13 +302,13 @@ def test_affinity_quarantines_unreadable_snapshot_and_dequeues_valid_work( assert dequeued is not None assert dequeued.item_id == cold_id - with session_queue_round_robin._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_round_robin) as cursor: cursor.execute("SELECT status FROM session_queue WHERE item_id = ?", (future_id,)) assert cursor.fetchone()[0] == "failed" def test_unreadable_snapshot_is_safe_for_detail_list_and_retry( - session_queue_fifo: SqliteSessionQueue, + session_queue_fifo: SessionQueue, ) -> None: future_session = json.loads(_EMPTY_SESSION_JSON) future_session["execution_state_version"] = CURRENT_EXECUTION_STATE_VERSION + 1 @@ -337,7 +336,7 @@ def test_unreadable_snapshot_is_safe_for_detail_list_and_retry( def test_unreadable_terminal_snapshot_is_not_reported_as_completed( - session_queue_fifo: SqliteSessionQueue, + session_queue_fifo: SessionQueue, ) -> None: future_session = json.loads(_EMPTY_SESSION_JSON) future_session["execution_state_version"] = CURRENT_EXECUTION_STATE_VERSION + 1 @@ -347,7 +346,7 @@ def test_unreadable_terminal_snapshot_is_not_reported_as_completed( "bad-user", session_json=json.dumps(future_session), ) - with session_queue_fifo._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_fifo) as cursor: cursor.execute("UPDATE session_queue SET status = 'completed' WHERE item_id = ?", (bad_id,)) detail = session_queue_fifo.get_queue_item(bad_id) @@ -359,10 +358,10 @@ def test_unreadable_terminal_snapshot_is_not_reported_as_completed( def test_unreadable_field_values_are_safe_for_summary( - session_queue_fifo: SqliteSessionQueue, + session_queue_fifo: SessionQueue, ) -> None: item_id = _insert_queue_item(session_queue_fifo, "default", "summary-user") - with session_queue_fifo._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_fifo) as cursor: cursor.execute( "UPDATE session_queue SET field_values = ? WHERE item_id = ?", ("{not valid json", item_id), @@ -377,7 +376,7 @@ def test_unreadable_field_values_are_safe_for_summary( def test_summary_read_does_not_hydrate_runtime_session( - session_queue_fifo: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue_fifo: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: item_id = _insert_queue_item(session_queue_fifo, "default", "summary-user") @@ -395,7 +394,7 @@ def fail_runtime_hydration(_queue_item_dict: dict) -> GraphExecutionState: def test_queue_list_api_read_does_not_hydrate_runtime_session( - session_queue_fifo: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue_fifo: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: item_id = _insert_queue_item(session_queue_fifo, "default", "list-user") @@ -413,7 +412,7 @@ def fail_runtime_hydration(_queue_item_dict: dict) -> GraphExecutionState: def test_queue_list_service_read_preserves_full_runtime_hydration( - session_queue_fifo: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue_fifo: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: item_id = _insert_queue_item(session_queue_fifo, "default", "list-user") @@ -431,10 +430,10 @@ def fail_projection(_queue_item_dict: dict) -> GraphExecutionState: def test_retry_read_hydrates_runtime_once_per_queue_item( - session_queue_fifo: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue_fifo: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: item_id = _insert_queue_item(session_queue_fifo, "default", "retry-user") - with session_queue_fifo._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_fifo) as cursor: cursor.execute("UPDATE session_queue SET status = 'failed' WHERE item_id = ?", (item_id,)) from invokeai.app.services.session_queue import session_queue_common @@ -459,7 +458,7 @@ def count_runtime_hydration(queue_item_dict: dict) -> GraphExecutionState: def test_retry_read_rejects_malformed_execution_state( - session_queue_fifo: SqliteSessionQueue, + session_queue_fifo: SessionQueue, ) -> None: malformed_session = json.loads(_EMPTY_SESSION_JSON) malformed_session["executed"] = 42 @@ -469,7 +468,7 @@ def test_retry_read_rejects_malformed_execution_state( "retry-user", session_json=json.dumps(malformed_session), ) - with session_queue_fifo._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_fifo) as cursor: cursor.execute("UPDATE session_queue SET status = 'failed' WHERE item_id = ?", (item_id,)) retry_result = session_queue_fifo.retry_items_by_id("default", [item_id]) @@ -478,7 +477,7 @@ def test_retry_read_rejects_malformed_execution_state( def test_retry_read_rejects_unknown_persisted_effect( - session_queue_fifo: SqliteSessionQueue, + session_queue_fifo: SessionQueue, ) -> None: malformed_session = json.loads(_EMPTY_SESSION_JSON) malformed_session["execution_state_version"] = CURRENT_EXECUTION_STATE_VERSION @@ -489,7 +488,7 @@ def test_retry_read_rejects_unknown_persisted_effect( "retry-user", session_json=json.dumps(malformed_session), ) - with session_queue_fifo._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_fifo) as cursor: cursor.execute("UPDATE session_queue SET status = 'failed' WHERE item_id = ?", (item_id,)) retry_result = session_queue_fifo.retry_items_by_id("default", [item_id]) @@ -498,12 +497,12 @@ def test_retry_read_rejects_unknown_persisted_effect( def test_api_item_reads_use_projection_without_runtime_rehydration( - session_queue_fifo: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue_fifo: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: current_id = _insert_queue_item(session_queue_fifo, "default", "api-user") _insert_queue_item(session_queue_fifo, "default", "api-user") next_id = _insert_queue_item(session_queue_fifo, "default", "api-user", priority=10) - with session_queue_fifo._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_fifo) as cursor: cursor.execute("UPDATE session_queue SET status = 'in_progress' WHERE item_id = ?", (current_id,)) def fail_runtime_hydration(_queue_item_dict: dict) -> GraphExecutionState: @@ -534,7 +533,7 @@ def test_queue_read_projection_preserves_nested_runtime_field_exclusions() -> No def test_nested_persisted_effects_are_validated_on_queue_reads( - session_queue_fifo: SqliteSessionQueue, + session_queue_fifo: SessionQueue, ) -> None: child_snapshot = dump_execution_state(GraphExecutionState(graph=Graph())) child_snapshot["execution_effects"] = {"unknown-reference": [{"kind": "unknown-effect"}]} @@ -550,7 +549,7 @@ def test_nested_persisted_effects_are_validated_on_queue_reads( "retry-user", session_json=json.dumps(snapshot), ) - with session_queue_fifo._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_fifo) as cursor: cursor.execute("UPDATE session_queue SET status = 'failed' WHERE item_id = ?", (item_id,)) retry_result = session_queue_fifo.retry_items_by_id("default", [item_id]) @@ -560,7 +559,7 @@ def test_nested_persisted_effects_are_validated_on_queue_reads( @pytest.mark.slow @pytest.mark.parametrize("iteration_count", [100, 3000]) -def test_queue_read_benchmark(iteration_count: int, session_queue_fifo: SqliteSessionQueue, capsys) -> None: +def test_queue_read_benchmark(iteration_count: int, session_queue_fifo: SessionQueue, capsys) -> None: """Record representative queue read costs without making machine-sensitive timing claims.""" from tests.app.services.shared.test_execution_state_migration import _make_completed_iterate_state @@ -569,7 +568,7 @@ def test_queue_read_benchmark(iteration_count: int, session_queue_fifo: SqliteSe session_json = json.dumps(snapshot, default=to_jsonable_python) detail_id = _insert_queue_item(session_queue_fifo, "default", "benchmark-user", session_json=session_json) retry_id = _insert_queue_item(session_queue_fifo, "default", "benchmark-user", session_json=session_json) - with session_queue_fifo._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_fifo) as cursor: cursor.execute("UPDATE session_queue SET status = 'failed' WHERE item_id = ?", (retry_id,)) tracemalloc.start() @@ -629,9 +628,7 @@ def test_queue_read_benchmark(iteration_count: int, session_queue_fifo: SqliteSe ) -def test_unreadable_child_retries_readable_root( - session_queue_fifo: SqliteSessionQueue, event_bus: TestEventService -) -> None: +def test_unreadable_child_retries_readable_root(session_queue_fifo: SessionQueue, event_bus: TestEventService) -> None: root_id = _insert_queue_item(session_queue_fifo, "default", "workflow-user") future_session = json.loads(_EMPTY_SESSION_JSON) future_session["execution_state_version"] = CURRENT_EXECUTION_STATE_VERSION + 1 @@ -641,7 +638,7 @@ def test_unreadable_child_retries_readable_root( "workflow-user", session_json=json.dumps(future_session), ) - with session_queue_fifo._db.transaction() as cursor: + with sqlite_cursor_of(session_queue_fifo) as cursor: cursor.execute( "UPDATE session_queue SET status = 'failed' WHERE item_id = ?", (root_id,), @@ -665,7 +662,7 @@ def test_unreadable_child_retries_readable_root( # --------------------------------------------------------------------------- -def test_round_robin_interleaves_users(session_queue_round_robin: SqliteSessionQueue) -> None: +def test_round_robin_interleaves_users(session_queue_round_robin: SessionQueue) -> None: """Round-robin: jobs from multiple users are interleaved one per user per round. Queue insertion order (matching the issue example): @@ -686,7 +683,7 @@ def test_round_robin_interleaves_users(session_queue_round_robin: SqliteSessionQ assert user_ids == ["user_a", "user_b", "user_c", "user_a", "user_c", "user_a"] -def test_round_robin_single_user_behaves_like_fifo(session_queue_round_robin: SqliteSessionQueue) -> None: +def test_round_robin_single_user_behaves_like_fifo(session_queue_round_robin: SessionQueue) -> None: """Round-robin with only one user produces the same order as FIFO.""" queue_id = "default" _insert_queue_item(session_queue_round_robin, queue_id, "user_a") @@ -697,7 +694,7 @@ def test_round_robin_single_user_behaves_like_fifo(session_queue_round_robin: Sq assert user_ids == ["user_a", "user_a", "user_a"] -def test_round_robin_handles_user_joining_mid_queue(session_queue_round_robin: SqliteSessionQueue) -> None: +def test_round_robin_handles_user_joining_mid_queue(session_queue_round_robin: SessionQueue) -> None: """Round-robin: a user who joins later is correctly interleaved.""" queue_id = "default" _insert_queue_item(session_queue_round_robin, queue_id, "user_a") @@ -710,12 +707,12 @@ def test_round_robin_handles_user_joining_mid_queue(session_queue_round_robin: S assert user_ids == ["user_a", "user_b", "user_a"] -def test_round_robin_returns_none_when_empty(session_queue_round_robin: SqliteSessionQueue) -> None: +def test_round_robin_returns_none_when_empty(session_queue_round_robin: SessionQueue) -> None: """Round-robin: dequeue returns None when the queue is empty.""" assert session_queue_round_robin.dequeue() is None -def test_round_robin_priority_within_user_respected(session_queue_round_robin: SqliteSessionQueue) -> None: +def test_round_robin_priority_within_user_respected(session_queue_round_robin: SessionQueue) -> None: """Round-robin: within a single user's items, higher priority is dequeued first.""" queue_id = "default" # Insert low-priority item first, then high-priority for same user. @@ -737,13 +734,13 @@ def test_round_robin_priority_within_user_respected(session_queue_round_robin: S def _seed_completed_history( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, queue_id: str, user_id: str, count: int, ) -> None: """Insert `count` completed items (with started_at set) for a user, simulating retained history.""" - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: for i in range(count): session_id = str(uuid.uuid4()) batch_id = str(uuid.uuid4()) @@ -772,7 +769,9 @@ def _seed_completed_history( ) -def test_round_robin_dequeue_does_not_scan_full_history(session_queue_round_robin: SqliteSessionQueue) -> None: +def test_round_robin_dequeue_does_not_scan_full_history( + session_queue_round_robin: SessionQueue, mock_sqlite_database: Database +) -> None: """Round-robin dequeue cost must scale with active users, not retained queue history. Regression guard for the scaling concern: the per-user "last served" lookup must be an @@ -789,9 +788,11 @@ def test_round_robin_dequeue_does_not_scan_full_history(session_queue_round_robi _seed_completed_history(session_queue_round_robin, queue_id, u, count=500) _insert_queue_item(session_queue_round_robin, queue_id, u) - with session_queue_round_robin._db.transaction() as cursor: - plan_rows = cursor.execute("EXPLAIN QUERY PLAN " + ROUND_ROBIN_DEQUEUE_QUERY).fetchall() - details = [row["detail"] for row in plan_rows] + database = mock_sqlite_database + with capture_statements(database) as statements: + item = session_queue_round_robin.dequeue() + statement, parameters = next((sql, parameters) for sql, parameters in statements if "user_next_item" in sql) + details = explain_query_plan(database, statement, parameters) # No step may scan the session_queue base table — that is the full-history scan we are # eliminating. (CTE result scans like "SCAN uni" / "SCAN (subquery-N)" are fine; those are @@ -807,12 +808,11 @@ def test_round_robin_dequeue_does_not_scan_full_history(session_queue_round_robi # And the dequeue must still return the least-recently-served user (correctness under history). # user_a's history ends earliest only if seeded first; all three were seeded equal counts with # identical timestamps, so item_id ASC tie-breaks to the first-inserted pending item (user_a). - item = session_queue_round_robin.dequeue() assert item is not None assert item.user_id == "user_a" -def test_round_robin_ignored_in_single_user_mode(mock_invoker: Invoker) -> None: +def test_round_robin_ignored_in_single_user_mode(mock_invoker: Invoker, mock_sqlite_database: Database) -> None: """When multiuser=False, round_robin config is ignored and FIFO is used.""" mock_invoker.services.configuration = InvokeAIAppConfig( use_memory_db=True, @@ -820,8 +820,8 @@ def test_round_robin_ignored_in_single_user_mode(mock_invoker: Invoker) -> None: multiuser=False, session_queue_mode="round_robin", ) - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) + db = mock_sqlite_database + queue = SessionQueue(db) queue.start(mock_invoker) queue_id = "default" @@ -863,10 +863,10 @@ def _install_fake_cache(invoker: Invoker, device: str, resident_keys: set[str]) invoker.services.model_manager = model_manager -def test_affinity_prefers_warm_model_within_user(session_queue_round_robin: SqliteSessionQueue) -> None: +def test_affinity_prefers_warm_model_within_user(session_queue_round_robin: SessionQueue) -> None: """Among one user's equal-priority items, the item whose model is warm on the claiming device is dequeued first, ahead of an older item that would need a model load.""" - _install_fake_cache(session_queue_round_robin._SqliteSessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) + _install_fake_cache(session_queue_round_robin._SessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) queue_id = "default" cold_id = _insert_queue_item(session_queue_round_robin, queue_id, "user_a") warm_id = _insert_queue_item( @@ -879,10 +879,10 @@ def test_affinity_prefers_warm_model_within_user(session_queue_round_robin: Sqli assert second is not None and second.item_id == cold_id -def test_affinity_never_overrides_round_robin_user_choice(session_queue_round_robin: SqliteSessionQueue) -> None: +def test_affinity_never_overrides_round_robin_user_choice(session_queue_round_robin: SessionQueue) -> None: """Affinity must not steal the turn from the fairness-chosen user, even when another user's pending item has a warm model.""" - _install_fake_cache(session_queue_round_robin._SqliteSessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) + _install_fake_cache(session_queue_round_robin._SessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) queue_id = "default" a_id = _insert_queue_item(session_queue_round_robin, queue_id, "user_a") _insert_queue_item( @@ -894,9 +894,9 @@ def test_affinity_never_overrides_round_robin_user_choice(session_queue_round_ro assert first is not None and first.item_id == a_id -def test_affinity_never_overrides_priority(session_queue_round_robin: SqliteSessionQueue) -> None: +def test_affinity_never_overrides_priority(session_queue_round_robin: SessionQueue) -> None: """A warm low-priority item must not jump ahead of a cold higher-priority item.""" - _install_fake_cache(session_queue_round_robin._SqliteSessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) + _install_fake_cache(session_queue_round_robin._SessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) queue_id = "default" hi_id = _insert_queue_item(session_queue_round_robin, queue_id, "user_a", priority=10) _insert_queue_item( @@ -907,9 +907,9 @@ def test_affinity_never_overrides_priority(session_queue_round_robin: SqliteSess assert first is not None and first.item_id == hi_id -def test_affinity_noop_without_device_or_matching_cache(session_queue_round_robin: SqliteSessionQueue) -> None: +def test_affinity_noop_without_device_or_matching_cache(session_queue_round_robin: SessionQueue) -> None: """With no device, or a device with no registered cache, ordering is unchanged.""" - _install_fake_cache(session_queue_round_robin._SqliteSessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) + _install_fake_cache(session_queue_round_robin._SessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) queue_id = "default" first_id = _insert_queue_item(session_queue_round_robin, queue_id, "user_a") _insert_queue_item( @@ -921,10 +921,10 @@ def test_affinity_noop_without_device_or_matching_cache(session_queue_round_robi assert first is not None and first.item_id == first_id -def test_affinity_applies_in_fifo_mode(session_queue_fifo: SqliteSessionQueue) -> None: +def test_affinity_applies_in_fifo_mode(session_queue_fifo: SessionQueue) -> None: """Single-user installs (FIFO query, but default session_queue_mode=round_robin) also benefit: warm items are preferred within the same priority.""" - _install_fake_cache(session_queue_fifo._SqliteSessionQueue__invoker, "cuda:1", {_WARM_MODEL_KEY}) + _install_fake_cache(session_queue_fifo._SessionQueue__invoker, "cuda:1", {_WARM_MODEL_KEY}) queue_id = "default" cold_id = _insert_queue_item(session_queue_fifo, queue_id, "user_a") warm_id = _insert_queue_item( @@ -937,7 +937,7 @@ def test_affinity_applies_in_fifo_mode(session_queue_fifo: SqliteSessionQueue) - assert second is not None and second.item_id == cold_id -def test_affinity_disabled_by_explicit_fifo_mode(mock_invoker: Invoker) -> None: +def test_affinity_disabled_by_explicit_fifo_mode(mock_invoker: Invoker, mock_sqlite_database: Database) -> None: """An admin who explicitly sets session_queue_mode=FIFO is promised strict insertion order, so affinity reordering must not apply.""" mock_invoker.services.configuration = InvokeAIAppConfig( @@ -945,8 +945,8 @@ def test_affinity_disabled_by_explicit_fifo_mode(mock_invoker: Invoker) -> None: node_cache_size=0, session_queue_mode="FIFO", ) - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) + db = mock_sqlite_database + queue = SessionQueue(db) queue.start(mock_invoker) _install_fake_cache(mock_invoker, "cuda:0", {_WARM_MODEL_KEY}) @@ -960,11 +960,11 @@ def test_affinity_disabled_by_explicit_fifo_mode(mock_invoker: Invoker) -> None: assert second is not None and second.item_id == warm_id -def test_affinity_prefers_more_matching_keys(session_queue_round_robin: SqliteSessionQueue) -> None: +def test_affinity_prefers_more_matching_keys(session_queue_round_robin: SessionQueue) -> None: """An item matching two resident models outranks an older item matching one — this also exercises the multi-key score expression and its parameter binding order.""" other_key = "bbbbbbbb-5555-6666-7777-888888888888" - _install_fake_cache(session_queue_round_robin._SqliteSessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY, other_key}) + _install_fake_cache(session_queue_round_robin._SessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY, other_key}) queue_id = "default" one_match_id = _insert_queue_item( session_queue_round_robin, queue_id, "user_a", session_json=_session_with_model_key(_WARM_MODEL_KEY) @@ -982,10 +982,10 @@ def test_affinity_prefers_more_matching_keys(session_queue_round_robin: SqliteSe assert second is not None and second.item_id == one_match_id -def test_affinity_lookahead_window_bounds_deferral(session_queue_round_robin: SqliteSessionQueue) -> None: +def test_affinity_lookahead_window_bounds_deferral(session_queue_round_robin: SessionQueue) -> None: """A warm item further than AFFINITY_MAX_LOOKAHEAD past the fairness candidate must not be swapped in — this is the bound that prevents unbounded starvation of cold items.""" - _install_fake_cache(session_queue_round_robin._SqliteSessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) + _install_fake_cache(session_queue_round_robin._SessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) queue_id = "default" cold_id = _insert_queue_item(session_queue_round_robin, queue_id, "user_a", item_id=100) in_window_id = _insert_queue_item( @@ -1013,9 +1013,9 @@ def test_affinity_lookahead_window_bounds_deferral(session_queue_round_robin: Sq assert second is not None and second.item_id == cold_id -def test_affinity_does_not_scan_older_history(session_queue_round_robin: SqliteSessionQueue) -> None: +def test_affinity_does_not_scan_older_history(session_queue_round_robin: SessionQueue) -> None: """The affinity window must be bounded on both sides of the fairness candidate.""" - _install_fake_cache(session_queue_round_robin._SqliteSessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) + _install_fake_cache(session_queue_round_robin._SessionQueue__invoker, "cuda:0", {_WARM_MODEL_KEY}) old_warm_id = _insert_queue_item( session_queue_round_robin, "default", diff --git a/tests/app/services/session_queue/test_session_queue_dequeue_concurrency.py b/tests/app/services/session_queue/test_session_queue_dequeue_concurrency.py index 2d61bf0f37a..fa1f77788f6 100644 --- a/tests/app/services/session_queue/test_session_queue_dequeue_concurrency.py +++ b/tests/app/services/session_queue/test_session_queue_dequeue_concurrency.py @@ -6,26 +6,28 @@ import pytest from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from tests.fixtures.sqlite_database import sqlite_cursor_of from tests.test_nodes import PromptTestInvocation @pytest.fixture -def session_queue(mock_invoker: Invoker) -> SqliteSessionQueue: - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) +def session_queue(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: + db = mock_sqlite_database + queue = SessionQueue(db) queue.start(mock_invoker) return queue -def _insert_queue_item(session_queue: SqliteSessionQueue, user_id: str = "system") -> int: +def _insert_queue_item(session_queue: SessionQueue, user_id: str = "system") -> int: graph = Graph() graph.add_node(PromptTestInvocation(id="prompt", prompt="test")) session = GraphExecutionState(graph=graph) session_json = session.model_dump_json(warnings=False, exclude_none=True) batch_id = str(uuid.uuid4()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -39,7 +41,7 @@ def _insert_queue_item(session_queue: SqliteSessionQueue, user_id: str = "system return cursor.lastrowid -def test_concurrent_dequeue_never_claims_same_item_twice(session_queue: SqliteSessionQueue) -> None: +def test_concurrent_dequeue_never_claims_same_item_twice(session_queue: SessionQueue) -> None: item_count = 50 worker_count = 8 for _ in range(item_count): @@ -70,7 +72,7 @@ def worker() -> None: assert len(set(claimed_ids)) == item_count -def test_dequeue_records_processing_device(session_queue: SqliteSessionQueue) -> None: +def test_dequeue_records_processing_device(session_queue: SessionQueue) -> None: _insert_queue_item(session_queue) item = session_queue.dequeue(device="cuda:1") @@ -82,7 +84,7 @@ def test_dequeue_records_processing_device(session_queue: SqliteSessionQueue) -> assert completed.device == "cuda:1" -def test_dequeue_without_device_leaves_device_unset(session_queue: SqliteSessionQueue) -> None: +def test_dequeue_without_device_leaves_device_unset(session_queue: SessionQueue) -> None: _insert_queue_item(session_queue) item = session_queue.dequeue() diff --git a/tests/app/services/session_queue/test_session_queue_enqueue_idempotency.py b/tests/app/services/session_queue/test_session_queue_enqueue_idempotency.py index aa4e3182c68..c6773c3e7cb 100644 --- a/tests/app/services/session_queue/test_session_queue_enqueue_idempotency.py +++ b/tests/app/services/session_queue/test_session_queue_enqueue_idempotency.py @@ -5,26 +5,28 @@ import pytest from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue import session_queue_sqlite +from invokeai.app.services.session_queue import session_queue_default from invokeai.app.services.session_queue.session_queue_common import ( Batch, EnqueueIdempotencyConflictError, EnqueueProjectNotFoundError, EnqueueReceiptLimitError, ) -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph +from tests.fixtures.sqlite_database import sqlite_cursor from tests.test_nodes import PromptTestInvocation @pytest.fixture -def session_queue(mock_invoker: Invoker) -> SqliteSessionQueue: - with mock_invoker.services.board_records._db.transaction() as cursor: +def session_queue(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: + with sqlite_cursor(mock_sqlite_database) as cursor: cursor.executemany( "INSERT INTO users (user_id, email, password_hash) VALUES (?, ?, ?);", [("user-1", "user-1@example.com", "test"), ("user-2", "user-2@example.com", "test")], ) - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) + queue = SessionQueue(mock_sqlite_database) queue.start(mock_invoker) return queue @@ -50,7 +52,7 @@ def _batch( ) -def test_enqueue_batch_retries_return_the_original_items(session_queue: SqliteSessionQueue) -> None: +def test_enqueue_batch_retries_return_the_original_items(session_queue: SessionQueue) -> None: batch = _batch() first = asyncio.run(session_queue.enqueue_batch("default", batch, False, "user-1")) retry = asyncio.run(session_queue.enqueue_batch("default", batch, False, "user-1")) @@ -61,7 +63,7 @@ def test_enqueue_batch_retries_return_the_original_items(session_queue: SqliteSe assert len(session_queue.list_all_queue_items("default")) == 2 -def test_distinct_idempotent_enqueues_cannot_share_a_caller_batch_id(session_queue: SqliteSessionQueue) -> None: +def test_distinct_idempotent_enqueues_cannot_share_a_caller_batch_id(session_queue: SessionQueue) -> None: caller_batch_id = "caller-reused-batch" first = asyncio.run( session_queue.enqueue_batch( @@ -87,7 +89,7 @@ def test_distinct_idempotent_enqueues_cannot_share_a_caller_batch_id(session_que assert all(session_queue.get_queue_item(item_id).status == "canceled" for item_id in second.item_ids) -def test_concurrent_enqueue_retries_commit_once(session_queue: SqliteSessionQueue) -> None: +def test_concurrent_enqueue_retries_commit_once(session_queue: SessionQueue) -> None: batch = _batch() async def enqueue_twice(): @@ -103,14 +105,14 @@ async def enqueue_twice(): def test_acknowledgement_does_not_erase_identity_for_an_inflight_retry( - session_queue: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: batch = _batch() first_prepared = threading.Event() second_prepared = threading.Event() release_second = threading.Event() prepare_call_count = 0 - original = session_queue_sqlite.prepare_values_to_insert + original = session_queue_default.prepare_values_to_insert def interleave_prepare(*args: Any, **kwargs: Any): nonlocal prepare_call_count @@ -124,7 +126,7 @@ def interleave_prepare(*args: Any, **kwargs: Any): assert release_second.wait(timeout=5) return values - monkeypatch.setattr(session_queue_sqlite, "prepare_values_to_insert", interleave_prepare) + monkeypatch.setattr(session_queue_default, "prepare_values_to_insert", interleave_prepare) async def exercise_interleaving(): first_task = asyncio.create_task(session_queue.enqueue_batch("default", batch, False, "user-1")) @@ -142,7 +144,7 @@ async def exercise_interleaving(): assert len(session_queue.list_all_queue_items("default")) == 2 -def test_enqueue_batch_id_is_scoped_by_user(session_queue: SqliteSessionQueue) -> None: +def test_enqueue_batch_id_is_scoped_by_user(session_queue: SessionQueue) -> None: first = asyncio.run(session_queue.enqueue_batch("default", _batch(), False, "user-1")) second = asyncio.run(session_queue.enqueue_batch("default", _batch(), False, "user-2")) @@ -152,7 +154,7 @@ def test_enqueue_batch_id_is_scoped_by_user(session_queue: SqliteSessionQueue) - assert len(session_queue.list_all_queue_items("default")) == 4 -def test_enqueue_batch_id_is_scoped_by_queue(session_queue: SqliteSessionQueue) -> None: +def test_enqueue_batch_id_is_scoped_by_queue(session_queue: SessionQueue) -> None: first = asyncio.run(session_queue.enqueue_batch("queue-1", _batch(), False, "user-1")) second = asyncio.run(session_queue.enqueue_batch("queue-2", _batch(), False, "user-1")) @@ -161,7 +163,7 @@ def test_enqueue_batch_id_is_scoped_by_queue(session_queue: SqliteSessionQueue) assert first.batch.batch_id != second.batch.batch_id -def test_enqueue_batch_rejects_identity_reuse_for_changed_payload(session_queue: SqliteSessionQueue) -> None: +def test_enqueue_batch_rejects_identity_reuse_for_changed_payload(session_queue: SessionQueue) -> None: asyncio.run(session_queue.enqueue_batch("default", _batch(), False, "user-1")) with pytest.raises(EnqueueIdempotencyConflictError): @@ -174,7 +176,7 @@ def test_enqueue_batch_rejects_identity_reuse_for_changed_payload(session_queue: asyncio.run(session_queue.enqueue_batch("default", _batch(), True, "user-1")) -def test_enqueue_receipt_survives_queue_item_deletion(session_queue: SqliteSessionQueue) -> None: +def test_enqueue_receipt_survives_queue_item_deletion(session_queue: SessionQueue) -> None: batch = _batch() first = asyncio.run(session_queue.enqueue_batch("default", batch, False, "user-1")) session_queue.delete_queue_items_by_id(first.item_ids) @@ -187,12 +189,12 @@ def test_enqueue_receipt_survives_queue_item_deletion(session_queue: SqliteSessi def test_enqueue_receipt_does_not_expire_while_unacknowledged( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker, mock_sqlite_database: Database ) -> None: batch = _batch() first = asyncio.run(session_queue.enqueue_batch("default", batch, False, "user-1")) session_queue.delete_queue_items_by_id(first.item_ids) - with mock_invoker.services.board_records._db.transaction() as cursor: + with sqlite_cursor(mock_sqlite_database) as cursor: cursor.execute("UPDATE session_queue_enqueue_receipts SET created_at = '2000-01-01';") retry = asyncio.run(session_queue.enqueue_batch("default", batch, False, "user-1")) @@ -201,7 +203,7 @@ def test_enqueue_receipt_does_not_expire_while_unacknowledged( assert session_queue.list_all_queue_items("default") == [] -def test_acknowledged_enqueue_receipt_still_settles_delayed_retries(session_queue: SqliteSessionQueue) -> None: +def test_acknowledged_enqueue_receipt_still_settles_delayed_retries(session_queue: SessionQueue) -> None: batch = _batch() first = asyncio.run(session_queue.enqueue_batch("default", batch, False, "user-1")) session_queue.delete_queue_items_by_id(first.item_ids) @@ -213,7 +215,7 @@ def test_acknowledged_enqueue_receipt_still_settles_delayed_retries(session_queu assert session_queue.list_all_queue_items("default") == [] -def test_enqueue_receipt_lookup_is_scoped_by_user_and_queue(session_queue: SqliteSessionQueue) -> None: +def test_enqueue_receipt_lookup_is_scoped_by_user_and_queue(session_queue: SessionQueue) -> None: batch = _batch() accepted = asyncio.run(session_queue.enqueue_batch("queue-1", batch, False, "user-1")) @@ -229,10 +231,10 @@ def test_enqueue_receipt_lookup_is_scoped_by_user_and_queue(session_queue: Sqlit def test_acknowledge_enqueue_is_scoped_by_user_and_queue( - session_queue: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr( - "invokeai.app.services.session_queue.session_queue_sqlite.MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPTS_PER_OWNER", 1 + "invokeai.app.services.session_queue.session_queue_default.MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPTS_PER_OWNER", 1 ) batch = _batch() asyncio.run(session_queue.enqueue_batch("queue-1", batch, False, "user-1")) @@ -249,12 +251,12 @@ def test_acknowledge_enqueue_is_scoped_by_user_and_queue( def test_expired_acknowledged_receipts_are_collected_on_admission( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker, mock_sqlite_database: Database ) -> None: first = _batch(idempotency_key="first") asyncio.run(session_queue.enqueue_batch("default", first, False, "user-1")) session_queue.acknowledge_enqueue("default", "first", "user-1") - with mock_invoker.services.board_records._db.transaction() as cursor: + with sqlite_cursor(mock_sqlite_database) as cursor: cursor.execute( "UPDATE session_queue_enqueue_receipts SET acknowledged_at = '2000-01-01' WHERE idempotency_key = 'first';" ) @@ -265,12 +267,12 @@ def test_expired_acknowledged_receipts_are_collected_on_admission( def test_acknowledged_receipts_are_retained_for_seven_days( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker, mock_sqlite_database: Database ) -> None: first = _batch(idempotency_key="first") asyncio.run(session_queue.enqueue_batch("default", first, False, "user-1")) session_queue.acknowledge_enqueue("default", "first", "user-1") - with mock_invoker.services.board_records._db.transaction() as cursor: + with sqlite_cursor(mock_sqlite_database) as cursor: cursor.execute( """--sql UPDATE session_queue_enqueue_receipts @@ -285,10 +287,10 @@ def test_acknowledged_receipts_are_retained_for_seven_days( def test_unacknowledged_receipts_are_bounded_without_eviction( - session_queue: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr( - "invokeai.app.services.session_queue.session_queue_sqlite.MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPTS_PER_OWNER", 1 + "invokeai.app.services.session_queue.session_queue_default.MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPTS_PER_OWNER", 1 ) asyncio.run(session_queue.enqueue_batch("default", _batch(idempotency_key="first"), False, "user-1")) @@ -301,10 +303,10 @@ def test_unacknowledged_receipts_are_bounded_without_eviction( def test_unacknowledged_receipt_count_limit_is_per_user_across_queues( - session_queue: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr( - "invokeai.app.services.session_queue.session_queue_sqlite.MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPTS_PER_OWNER", 1 + "invokeai.app.services.session_queue.session_queue_default.MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPTS_PER_OWNER", 1 ) first_batch = _batch(idempotency_key="first") asyncio.run(session_queue.enqueue_batch("queue-1", first_batch, False, "user-1")) @@ -319,10 +321,10 @@ def test_unacknowledged_receipt_count_limit_is_per_user_across_queues( def test_unacknowledged_receipt_bytes_are_bounded_atomically( - session_queue: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr( - "invokeai.app.services.session_queue.session_queue_sqlite.MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPT_BYTES_PER_OWNER", + "invokeai.app.services.session_queue.session_queue_default.MAX_UNACKNOWLEDGED_ENQUEUE_RECEIPT_BYTES_PER_OWNER", 1, ) @@ -334,9 +336,9 @@ def test_unacknowledged_receipt_bytes_are_bounded_atomically( def test_total_receipt_count_includes_acknowledged_tombstones( - session_queue: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr("invokeai.app.services.session_queue.session_queue_sqlite.MAX_ENQUEUE_RECEIPTS_PER_OWNER", 1) + monkeypatch.setattr("invokeai.app.services.session_queue.session_queue_default.MAX_ENQUEUE_RECEIPTS_PER_OWNER", 1) first_batch = _batch(idempotency_key="first") asyncio.run(session_queue.enqueue_batch("default", first_batch, False, "user-1")) session_queue.acknowledge_enqueue("default", "first", "user-1") @@ -348,10 +350,10 @@ def test_total_receipt_count_includes_acknowledged_tombstones( def test_total_receipt_bytes_are_bounded_atomically( - session_queue: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr( - "invokeai.app.services.session_queue.session_queue_sqlite.MAX_ENQUEUE_RECEIPT_BYTES_PER_OWNER", 1 + "invokeai.app.services.session_queue.session_queue_default.MAX_ENQUEUE_RECEIPT_BYTES_PER_OWNER", 1 ) with pytest.raises(EnqueueReceiptLimitError): @@ -360,9 +362,7 @@ def test_total_receipt_bytes_are_bounded_atomically( assert session_queue.list_all_queue_items("default") == [] -def test_project_scoped_enqueue_requires_an_owned_project( - session_queue: SqliteSessionQueue, mock_invoker: Invoker -) -> None: +def test_project_scoped_enqueue_requires_an_owned_project(session_queue: SessionQueue, mock_invoker: Invoker) -> None: mock_invoker.services.project_records.create("system", "Owned", {}, project_id="project-1") with pytest.raises(EnqueueProjectNotFoundError): @@ -381,9 +381,7 @@ def test_project_scoped_enqueue_requires_an_owned_project( assert session_queue.list_all_queue_items("default") == [] -def test_project_deletion_wins_before_new_queue_admission( - session_queue: SqliteSessionQueue, mock_invoker: Invoker -) -> None: +def test_project_deletion_wins_before_new_queue_admission(session_queue: SessionQueue, mock_invoker: Invoker) -> None: mock_invoker.services.project_records.create("system", "Doomed", {}, project_id="project-1") mock_invoker.services.project_records.delete("system", "project-1") @@ -398,7 +396,7 @@ def test_project_deletion_wins_before_new_queue_admission( def test_existing_enqueue_receipt_settles_after_project_deletion( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker ) -> None: mock_invoker.services.project_records.create("system", "Doomed", {}, project_id="project-1") batch = _batch(project_id="project-1") @@ -411,7 +409,7 @@ def test_existing_enqueue_receipt_settles_after_project_deletion( def test_full_queue_does_not_consume_the_idempotency_key( - session_queue: SqliteSessionQueue, mock_invoker: Invoker, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, mock_invoker: Invoker, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(mock_invoker.services.configuration, "max_queue_size", 0) batch = _batch() @@ -425,7 +423,7 @@ def test_full_queue_does_not_consume_the_idempotency_key( def test_partial_acceptance_is_settled_idempotently_after_capacity_changes( - session_queue: SqliteSessionQueue, mock_invoker: Invoker, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, mock_invoker: Invoker, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(mock_invoker.services.configuration, "max_queue_size", 1) batch = _batch(runs=2) @@ -441,18 +439,18 @@ def test_partial_acceptance_is_settled_idempotently_after_capacity_changes( def test_full_queue_does_not_prepare_session_payloads( - session_queue: SqliteSessionQueue, mock_invoker: Invoker, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, mock_invoker: Invoker, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(mock_invoker.services.configuration, "max_queue_size", 1) asyncio.run(session_queue.enqueue_batch("default", _batch(idempotency_key="first", runs=1), False, "user-1")) requested_prepare_limits: list[int] = [] - original = session_queue_sqlite.prepare_values_to_insert + original = session_queue_default.prepare_values_to_insert def observe_prepare_limit(*args: Any, **kwargs: Any): requested_prepare_limits.append(kwargs["max_new_queue_items"]) return original(*args, **kwargs) - monkeypatch.setattr(session_queue_sqlite, "prepare_values_to_insert", observe_prepare_limit) + monkeypatch.setattr(session_queue_default, "prepare_values_to_insert", observe_prepare_limit) result = asyncio.run(session_queue.enqueue_batch("default", _batch(idempotency_key="second"), False, "user-1")) @@ -460,14 +458,14 @@ def observe_prepare_limit(*args: Any, **kwargs: Any): assert requested_prepare_limits == [0] -def test_enqueue_batch_rejects_identity_reuse_for_another_origin(session_queue: SqliteSessionQueue) -> None: +def test_enqueue_batch_rejects_identity_reuse_for_another_origin(session_queue: SessionQueue) -> None: asyncio.run(session_queue.enqueue_batch("default", _batch(), False, "user-1")) with pytest.raises(EnqueueIdempotencyConflictError): asyncio.run(session_queue.enqueue_batch("default", _batch(origin="webv2:project-2:item-2"), False, "user-1")) -def test_enqueued_items_carry_their_project(session_queue: SqliteSessionQueue, mock_invoker: Invoker) -> None: +def test_enqueued_items_carry_their_project(session_queue: SessionQueue, mock_invoker: Invoker) -> None: mock_invoker.services.project_records.create("system", "Owned", {}, project_id="project-1") asyncio.run( diff --git a/tests/app/services/session_queue/test_session_queue_list_pagination.py b/tests/app/services/session_queue/test_session_queue_list_pagination.py new file mode 100644 index 00000000000..f807183993e --- /dev/null +++ b/tests/app/services/session_queue/test_session_queue_list_pagination.py @@ -0,0 +1,96 @@ +"""Tests for list_queue_items() cursor pagination staying inside the requested scope.""" + +import uuid +from typing import Optional + +import pytest +from sqlalchemy import insert + +from invokeai.app.services.invoker import Invoker +from invokeai.app.services.session_queue.session_queue_common import QUEUE_ITEM_STATUS +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.session_queue import session_queue as session_queue_table +from invokeai.app.services.shared.graph import Graph, GraphExecutionState + +_SESSION_JSON = GraphExecutionState(graph=Graph()).model_dump_json() +# Far more pages than any test here inserts rows for. A cursor that stops advancing must fail the +# test by name instead of paging forever. +_MAX_PAGES = 50 + + +@pytest.fixture +def session_queue(mock_invoker: Invoker, database: Database) -> SessionQueue: + queue = SessionQueue(database) + queue.start(mock_invoker) + return queue + + +def _insert( + database: Database, + queue_id: str, + status: QUEUE_ITEM_STATUS = "pending", + destination: str = "canvas", + priority: int = 0, +) -> int: + values = { + "queue_id": queue_id, + "session": _SESSION_JSON, + "session_id": str(uuid.uuid4()), + "batch_id": str(uuid.uuid4()), + "priority": priority, + "destination": destination, + "status": status, + } + with database.begin(write=True) as conn: + return int(conn.execute(insert(session_queue_table).values(values)).inserted_primary_key[0]) + + +def _walk_pages( + session_queue: SessionQueue, + queue_id: str, + limit: int, + status: Optional[QUEUE_ITEM_STATUS] = None, + destination: Optional[str] = None, +) -> list[int]: + """Follow the cursor the way a client does: resume after the last item of the previous page.""" + item_ids: list[int] = [] + cursor: Optional[int] = None + priority = 0 + for _ in range(_MAX_PAGES): + page = session_queue.list_queue_items( + queue_id=queue_id, limit=limit, priority=priority, cursor=cursor, status=status, destination=destination + ) + item_ids.extend(item.item_id for item in page.items) + if not page.has_more: + return item_ids + cursor = page.items[-1].item_id + priority = page.items[-1].priority + pytest.fail(f"list_queue_items did not finish paging within {_MAX_PAGES} pages; collected {item_ids[:20]}") + + +def test_later_pages_stay_inside_the_requested_scope(session_queue: SessionQueue, database: Database) -> None: + in_scope: list[int] = [] + for _ in range(4): + in_scope.append(_insert(database, "default", destination="canvas")) + # Same priority and a higher item_id than every in-scope row before them, so a cursor + # predicate that escapes the filters would pull each of these onto the next page. + _insert(database, "other_queue", destination="canvas") + _insert(database, "default", status="completed", destination="canvas") + _insert(database, "default", destination="gallery") + + item_ids = _walk_pages(session_queue, "default", limit=2, status="pending", destination="canvas") + + assert item_ids == in_scope + + +def test_pages_continue_across_priority_levels(session_queue: SessionQueue, database: Database) -> None: + low = [_insert(database, "default", priority=0) for _ in range(3)] + high = [_insert(database, "default", priority=10) for _ in range(3)] + _insert(database, "other_queue", priority=10) + _insert(database, "other_queue", priority=0) + + item_ids = _walk_pages(session_queue, "default", limit=2) + + # Highest priority first, oldest first within a priority, each item exactly once. + assert item_ids == high + low diff --git a/tests/app/services/session_queue/test_session_queue_list_queue_items.py b/tests/app/services/session_queue/test_session_queue_list_queue_items.py index a197f5951a0..01055d7ff60 100644 --- a/tests/app/services/session_queue/test_session_queue_list_queue_items.py +++ b/tests/app/services/session_queue/test_session_queue_list_queue_items.py @@ -1,38 +1,43 @@ """Queue listings stay inside their filters: cursor pages of list_queue_items, and the item id listing whose limited form is the head of its full order.""" -import json import uuid +from typing import Any import pytest -from pydantic_core import to_jsonable_python +from sqlalchemy import insert from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_common import QUEUE_ITEM_STATUS +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.session_queue import session_queue as session_queue_table from invokeai.app.services.shared.graph import Graph, GraphExecutionState -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import SQLiteDirection -_SESSION_JSON = json.dumps(to_jsonable_python(GraphExecutionState(graph=Graph()).model_dump())) +_SESSION_JSON = GraphExecutionState(graph=Graph()).model_dump_json() @pytest.fixture -def session_queue(mock_invoker: Invoker) -> SqliteSessionQueue: - return SqliteSessionQueue(db=mock_invoker.services.board_records._db) +def session_queue(mock_invoker: Invoker, database: Database) -> SessionQueue: + queue = SessionQueue(database) + queue.start(mock_invoker) + return queue -def _insert(session_queue: SqliteSessionQueue, queue_id: str, status: str, destination: str, priority: int = 0) -> int: - with session_queue._db.transaction() as cursor: - cursor.execute( - """--sql - INSERT INTO session_queue (queue_id, session, session_id, batch_id, priority, status, destination) - VALUES (?, ?, ?, 'batch', ?, ?, ?); - """, - (queue_id, _SESSION_JSON, str(uuid.uuid4()), priority, status, destination), - ) - return cursor.lastrowid # type: ignore[return-value] +def _insert_row(session_queue: SessionQueue, **values: Any) -> int: + row = {"session": _SESSION_JSON, "session_id": str(uuid.uuid4()), "batch_id": "batch", **values} + with session_queue._queries._database.begin(write=True) as conn: + return int(conn.execute(insert(session_queue_table).values(row)).inserted_primary_key[0]) -def test_cursor_page_excludes_items_outside_the_filters(session_queue: SqliteSessionQueue) -> None: +def _insert( + session_queue: SessionQueue, queue_id: str, status: QUEUE_ITEM_STATUS, destination: str, priority: int = 0 +) -> int: + return _insert_row(session_queue, queue_id=queue_id, status=status, destination=destination, priority=priority) + + +def test_cursor_page_excludes_items_outside_the_filters(session_queue: SessionQueue) -> None: first = _insert(session_queue, "q1", "pending", "canvas", priority=1) second = _insert(session_queue, "q1", "pending", "canvas", priority=1) _insert(session_queue, "q2", "pending", "canvas", priority=1) @@ -63,16 +68,8 @@ def test_cursor_page_excludes_items_outside_the_filters(session_queue: SqliteSes ] -def _insert_listed(session_queue: SqliteSessionQueue, queue_id: str, created_at: str, user_id: str, origin: str) -> int: - with session_queue._db.transaction() as cursor: - cursor.execute( - """--sql - INSERT INTO session_queue (queue_id, session, session_id, batch_id, created_at, user_id, origin) - VALUES (?, ?, ?, 'batch', ?, ?, ?); - """, - (queue_id, _SESSION_JSON, str(uuid.uuid4()), created_at, user_id, origin), - ) - return cursor.lastrowid # type: ignore[return-value] +def _insert_listed(session_queue: SessionQueue, queue_id: str, created_at: str, user_id: str, origin: str) -> int: + return _insert_row(session_queue, queue_id=queue_id, created_at=created_at, user_id=user_id, origin=origin) @pytest.mark.parametrize( @@ -90,7 +87,7 @@ def _insert_listed(session_queue: SqliteSessionQueue, queue_id: str, created_at: ], ) def test_limited_item_ids_are_the_head_of_the_filtered_order( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, direction: SQLiteDirection, filters: dict[str, str], expected_rows: list[int], diff --git a/tests/app/services/session_queue/test_session_queue_multigpu_cancel.py b/tests/app/services/session_queue/test_session_queue_multigpu_cancel.py index c013d386c60..b921b4cc87e 100644 --- a/tests/app/services/session_queue/test_session_queue_multigpu_cancel.py +++ b/tests/app/services/session_queue/test_session_queue_multigpu_cancel.py @@ -12,21 +12,23 @@ from invokeai.app.services.events.events_common import QueueItemStatusChangedEvent from invokeai.app.services.invoker import Invoker from invokeai.app.services.session_queue.session_queue_common import SessionQueueItemNotFoundError -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from tests.fixtures.sqlite_database import sqlite_cursor_of from tests.test_nodes import PromptTestInvocation, TestEventService @pytest.fixture -def session_queue(mock_invoker: Invoker) -> SqliteSessionQueue: - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) +def session_queue(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: + db = mock_sqlite_database + queue = SessionQueue(db) queue.start(mock_invoker) return queue def _insert( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, batch_id: str, destination: str | None = None, user_id: str = "system", @@ -38,7 +40,7 @@ def _insert( graph.add_node(PromptTestInvocation(id="prompt", prompt="test")) session = GraphExecutionState(graph=graph) session_json = session.model_dump_json(warnings=False, exclude_none=True) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -73,7 +75,7 @@ def _canceled_event_item_ids(mock_invoker: Invoker) -> set[int]: } -def _dequeue_two_on_separate_devices(session_queue: SqliteSessionQueue) -> tuple[int, int]: +def _dequeue_two_on_separate_devices(session_queue: SessionQueue) -> tuple[int, int]: a = session_queue.dequeue(device="cuda:0") b = session_queue.dequeue(device="cuda:1") assert a is not None and b is not None @@ -81,7 +83,7 @@ def _dequeue_two_on_separate_devices(session_queue: SqliteSessionQueue) -> tuple return a.item_id, b.item_id -def test_cancel_by_batch_ids_cancels_all_in_progress(session_queue: SqliteSessionQueue, mock_invoker: Invoker): +def test_cancel_by_batch_ids_cancels_all_in_progress(session_queue: SessionQueue, mock_invoker: Invoker): batch_id = str(uuid.uuid4()) _insert(session_queue, batch_id=batch_id) _insert(session_queue, batch_id=batch_id) @@ -96,7 +98,7 @@ def test_cancel_by_batch_ids_cancels_all_in_progress(session_queue: SqliteSessio assert {id_a, id_b} <= _canceled_event_item_ids(mock_invoker) -def test_cancel_by_destination_cancels_all_in_progress(session_queue: SqliteSessionQueue, mock_invoker: Invoker): +def test_cancel_by_destination_cancels_all_in_progress(session_queue: SessionQueue, mock_invoker: Invoker): _insert(session_queue, batch_id=str(uuid.uuid4()), destination="canvas") _insert(session_queue, batch_id=str(uuid.uuid4()), destination="canvas") id_a, id_b = _dequeue_two_on_separate_devices(session_queue) @@ -109,7 +111,7 @@ def test_cancel_by_destination_cancels_all_in_progress(session_queue: SqliteSess assert {id_a, id_b} <= _canceled_event_item_ids(mock_invoker) -def test_cancel_by_queue_id_cancels_all_in_progress(session_queue: SqliteSessionQueue, mock_invoker: Invoker): +def test_cancel_by_queue_id_cancels_all_in_progress(session_queue: SessionQueue, mock_invoker: Invoker): _insert(session_queue, batch_id=str(uuid.uuid4())) _insert(session_queue, batch_id=str(uuid.uuid4())) id_a, id_b = _dequeue_two_on_separate_devices(session_queue) @@ -122,7 +124,7 @@ def test_cancel_by_queue_id_cancels_all_in_progress(session_queue: SqliteSession assert {id_a, id_b} <= _canceled_event_item_ids(mock_invoker) -def test_delete_by_destination_cancels_all_in_progress(session_queue: SqliteSessionQueue, mock_invoker: Invoker): +def test_delete_by_destination_cancels_all_in_progress(session_queue: SessionQueue, mock_invoker: Invoker): """delete_by_destination must signal every running worker (not just get_current()) before deleting their rows, or the un-canceled workers keep running and then fail to update a deleted row.""" @@ -141,14 +143,14 @@ def test_delete_by_destination_cancels_all_in_progress(session_queue: SqliteSess session_queue.get_queue_item(item_id) -def _make_workflow_chain(session_queue: SqliteSessionQueue, device: str) -> tuple[int, int, int]: +def _make_workflow_chain(session_queue: SessionQueue, device: str) -> tuple[int, int, int]: """Create a workflow-call chain mid-execution: a 'waiting' parent, a child dequeued to 'in_progress' on `device`, and a 'pending' next child. Returns (parent, running, pending) ids.""" parent_id = _insert(session_queue, batch_id=str(uuid.uuid4()), status="waiting") # Insert directly as in_progress with the device, rather than via dequeue(): dequeue picks the # globally best pending item, which in multi-chain setups may belong to another chain. running_child_id = _insert(session_queue, batch_id=str(uuid.uuid4()), parent_item_id=parent_id) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( "UPDATE session_queue SET status = 'in_progress', device = ? WHERE item_id = ?", (device, running_child_id), @@ -158,7 +160,7 @@ def _make_workflow_chain(session_queue: SqliteSessionQueue, device: str) -> tupl def test_cancel_all_except_current_preserves_every_active_workflow_chain( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker ): """With one in-progress item per GPU, 'cancel all except current' must leave EVERY active worker's workflow-call chain intact — not just the chain of one arbitrarily selected item.""" @@ -179,7 +181,7 @@ def test_cancel_all_except_current_preserves_every_active_workflow_chain( def test_delete_all_except_current_preserves_every_active_workflow_chain( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker ): parent_a, running_a, pending_a = _make_workflow_chain(session_queue, "cuda:0") parent_b, running_b, pending_b = _make_workflow_chain(session_queue, "cuda:1") @@ -196,7 +198,7 @@ def test_delete_all_except_current_preserves_every_active_workflow_chain( def test_bulk_cancel_tolerates_rows_deleted_mid_cancel( - session_queue: SqliteSessionQueue, mock_invoker: Invoker, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, mock_invoker: Invoker, monkeypatch: pytest.MonkeyPatch ): """A concurrent clear/delete can remove an in-progress row between _cancel_in_progress_matching's id SELECT and its per-item status update. The bulk cancel must complete instead of raising @@ -210,7 +212,7 @@ def test_bulk_cancel_tolerates_rows_deleted_mid_cancel( def delete_then_transition(item_id: int, status: str, **kwargs): # Simulate the concurrent deletion landing just before the per-item cancel. - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute("DELETE FROM session_queue WHERE item_id = ?", (item_id,)) return original_transition(item_id, status, **kwargs) @@ -224,7 +226,7 @@ def delete_then_transition(item_id: int, status: str, **kwargs): def test_concurrent_bulk_cancels_count_each_item_once( - session_queue: SqliteSessionQueue, mock_invoker: Invoker, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, mock_invoker: Invoker, monkeypatch: pytest.MonkeyPatch ): """Two bulk cancellations that select the same in-progress rows before either updates them must not both report those rows as newly canceled. The terminal guard is part of the cancel UPDATE @@ -275,7 +277,7 @@ def run() -> None: assert session_queue.get_queue_item(id_b).status == "canceled" -def test_cancel_by_batch_ids_respects_user_scope(session_queue: SqliteSessionQueue, mock_invoker: Invoker): +def test_cancel_by_batch_ids_respects_user_scope(session_queue: SessionQueue, mock_invoker: Invoker): """A user-scoped cancel must not cancel another user's in-progress item in the same batch.""" batch_id = str(uuid.uuid4()) _insert(session_queue, batch_id=batch_id, user_id="alice") diff --git a/tests/app/services/session_queue/test_session_queue_query_plans.py b/tests/app/services/session_queue/test_session_queue_query_plans.py index dc63031c8a3..b3ef9e9302e 100644 --- a/tests/app/services/session_queue/test_session_queue_query_plans.py +++ b/tests/app/services/session_queue/test_session_queue_query_plans.py @@ -1,94 +1,136 @@ -"""Queue refresh reads stay index-only, and point reads stay off the listing index. +"""Queue refresh reads stay index-only, and point reads and bulk changes stay off the listing index. Every session_queue column these reads use except queue_id is stored after the session blob, so a -plan that reads table rows walks each row's overflow pages. The database keeps no table statistics, -so the planner treats `queue_id = ?` as selective; reads that must cost their own rows rather than -the queue's history keep queue_id off the listing index. +plan that reads table rows walks each row's overflow pages. SQLite keeps no table statistics, so the +planner treats `queue_id = ?` as selective; statements that must cost their own rows rather than the +queue's history keep queue_id off the listing index. The plans are SQLite's: the listing index exists +only there. """ import re from collections.abc import Callable +from dataclasses import dataclass from functools import partial +from typing import Any import pytest +from sqlalchemy import insert from invokeai.app.services.invoker import Invoker from invokeai.app.services.session_queue.session_queue_common import Batch -from invokeai.app.services.session_queue.session_queue_sqlite import ( - PENDING_QUEUE_ITEM_COUNT_QUERY, - SqliteSessionQueue, -) +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.session_queue import session_queue as session_queue_table from invokeai.app.services.shared.graph import Graph -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import SQLiteDirection +from tests.fixtures.database import capture_statements, explain_query_plan from tests.test_nodes import PromptTestInvocation -LISTING_INDEX = "COVERING INDEX idx_session_queue_listing" +pytestmark = pytest.mark.sqlite_only + +LISTING_INDEX = "idx_session_queue_listing" +COVERING_LISTING_INDEX = f"COVERING INDEX {LISTING_INDEX}" +BATCH_SEEK = "idx_session_queue_batch_id (batch_id=?)" +STATUS_SEEK = "idx_session_queue_round_robin_pending (status=?" +USER_STATUS_SEEK = "idx_session_queue_round_robin_pending (status=? AND user_id=?)" +# A seek on any index that leads with the status: the live rows, never the queue's history. +ANY_STATUS_SEEK = re.compile(r"SEARCH session_queue USING (?:COVERING )?INDEX \w+ \(status=\?") +# The primary key, or the unique index on it, which SQLite may pick for a DELETE. +ID_SEEK = "rowid=?" +# A user filter SQLite may look up through an index (`+column` keeps one off every index). +INDEXED_USER_FILTER = re.compile(r"(? bool: + return self.sql.startswith(("UPDATE", "DELETE")) + + def binds(self, value: str) -> bool: + return value in self.parameters @pytest.fixture -def session_queue(mock_invoker: Invoker) -> SqliteSessionQueue: - queue = SqliteSessionQueue(db=mock_invoker.services.board_records._db) +def session_queue(mock_invoker: Invoker, database: Database) -> SessionQueue: + queue = SessionQueue(database) queue.start(mock_invoker) return queue -def _plans(session_queue: SqliteSessionQueue, read: Callable[[], object]) -> list[str]: - """The query plan of each statement `read` executes, as the service issued it.""" - conn = session_queue._db._conn - statements: list[str] = [] - conn.set_trace_callback(statements.append) - try: - read() - finally: - conn.set_trace_callback(None) - return [ - " | ".join(row["detail"] for row in conn.execute(f"EXPLAIN QUERY PLAN {statement}")) - for statement in statements - if "session_queue" in statement - ] +def _statements(database: Database, action: Callable[[], object]) -> list[Statement]: + """Each statement on the queue's rows that `action` runs (its inserts left out), with SQLite's plan for it.""" + with capture_statements(database) as captured: + action() + statements = [] + for sql, parameters in captured: + flat = " ".join(sql.split()) + if "FROM session_queue " not in flat and not flat.startswith("UPDATE session_queue "): + continue + statements.append(Statement(flat, parameters, " | ".join(explain_query_plan(database, sql, parameters)))) + assert statements, "the action ran no statement on session_queue" + return statements + + +def _plans(statements: list[Statement]) -> list[str]: + return [statement.plan for statement in statements] -def test_queue_listing_and_counts_read_only_the_listing_index(session_queue: SqliteSessionQueue) -> None: - newest_first = _plans( - session_queue, - lambda: session_queue.get_queue_item_ids("default", SQLiteDirection.Descending, origin_prefix="webv2:p:1:q:"), +def _enqueue(session_queue: SessionQueue) -> None: + """Two pending items of alice's, so that a bulk change has rows to change.""" + graph = Graph() + graph.add_node(PromptTestInvocation(id="prompt", prompt="test")) + batch = Batch(graph=graph, runs=2, batch_id="batch", origin=f"{ORIGIN_PREFIX}canvas") + session_queue._enqueue_batch("default", batch, False, "alice") + + +def test_queue_listing_and_counts_read_only_the_listing_index(session_queue: SessionQueue, database: Database) -> None: + newest_first = _statements( + database, + lambda: session_queue.get_queue_item_ids("default", SQLiteDirection.Descending, origin_prefix=ORIGIN_PREFIX), ) - assert len(newest_first) == 1 - assert LISTING_INDEX in newest_first[0] - assert "TEMP B-TREE" not in newest_first[0] + for plan in _plans(newest_first): + assert COVERING_LISTING_INDEX in plan + assert "TEMP B-TREE" not in plan # The recent window reads the index in order with no sort in between, so LIMIT can stop the walk # (test_limited_newest_first_listing_stops_after_the_window shows that it does). - recent_window = _plans( - session_queue, lambda: session_queue.get_queue_item_ids("default", SQLiteDirection.Descending, limit=50) + recent_window = _statements( + database, lambda: session_queue.get_queue_item_ids("default", SQLiteDirection.Descending, limit=50) ) - assert len(recent_window) == 1 - assert LISTING_INDEX in recent_window[0] - assert "TEMP B-TREE" not in recent_window[0] + for plan in _plans(recent_window): + assert COVERING_LISTING_INDEX in plan + assert "TEMP B-TREE" not in plan # Oldest first sorts only items that share a created_at, by item_id. - oldest_first = _plans(session_queue, lambda: session_queue.get_queue_item_ids("default", SQLiteDirection.Ascending)) - assert len(oldest_first) == 1 - assert LISTING_INDEX in oldest_first[0] + oldest_first = _statements(database, lambda: session_queue.get_queue_item_ids("default", SQLiteDirection.Ascending)) # SQLite words that partial sort differently by version; a full sort reads "FOR ORDER BY". partial_sort = re.compile(r"TEMP B-TREE FOR (LAST TERM|RIGHT PART) OF ORDER BY") - assert "TEMP B-TREE" not in partial_sort.sub("", oldest_first[0]) + for plan in _plans(oldest_first): + assert COVERING_LISTING_INDEX in plan + assert "TEMP B-TREE" not in partial_sort.sub("", plan) - status = _plans( - session_queue, + status = _statements( + database, lambda: session_queue.get_queue_status( - "default", user_id="alice", acting_user_id="alice", origin_prefix="webv2:p:1:q:" + "default", user_id="alice", acting_user_id="alice", origin_prefix=ORIGIN_PREFIX ), ) - # The global and per-user counts come from one pass over the scope's history. - counts = [plan for plan in status if "idx_session_queue_listing" in plan] + # The global and per-user counts come from one pass over the scope's history; the item in + # progress is found by its status. + counts = [plan for plan in _plans(status) if LISTING_INDEX in plan] assert len(counts) == 1, status - assert LISTING_INDEX in counts[0] + assert COVERING_LISTING_INDEX in counts[0] + assert counts[0].count("SEARCH session_queue") == 1, counts[0] -def _vm_steps(session_queue: SqliteSessionQueue, read: Callable[[], object]) -> int: - """Virtual machine steps (in units of 10) the service's statements execute during `read`.""" - conn = session_queue._db._conn +def _vm_steps(database: Database, read: Callable[[], object]) -> int: + """Virtual machine steps (in units of 10) the statements executed during `read` take.""" steps = 0 def count() -> int: @@ -96,34 +138,41 @@ def count() -> int: steps += 1 return 0 - conn.set_progress_handler(count, 10) + with database.begin(write=False) as conn: + driver = conn.connection.driver_connection + assert driver is not None + driver.set_progress_handler(count, 10) try: read() finally: - conn.set_progress_handler(None, 0) + driver.set_progress_handler(None, 0) return steps -def _append_history(session_queue: SqliteSessionQueue, start: int, count: int) -> None: - with session_queue._db.transaction() as cursor: - cursor.executemany( - """--sql - INSERT INTO session_queue (queue_id, session, session_id, batch_id, created_at) - VALUES ('default', '{}', ?, 'batch', ?); - """, - [(f"session-{index}", f"2026-01-01 00:00:00.{index:06d}") for index in range(start, start + count)], - ) +def _append_history(database: Database, start: int, count: int) -> None: + rows = [ + { + "queue_id": "default", + "session": "{}", + "session_id": f"session-{index}", + "batch_id": "batch", + "created_at": f"2026-01-01 00:00:00.{index:06d}", + } + for index in range(start, start + count) + ] + with database.begin(write=True) as conn: + conn.execute(insert(session_queue_table), rows) -def test_limited_newest_first_listing_stops_after_the_window(session_queue: SqliteSessionQueue) -> None: +def test_limited_newest_first_listing_stops_after_the_window(session_queue: SessionQueue, database: Database) -> None: def recent_window() -> object: return session_queue.get_queue_item_ids("default", SQLiteDirection.Descending, limit=50) - _append_history(session_queue, 0, 100) - small = _vm_steps(session_queue, recent_window) - _append_history(session_queue, 100, 4_900) - large = _vm_steps(session_queue, recent_window) - unlimited = _vm_steps(session_queue, lambda: session_queue.get_queue_item_ids("default")) + _append_history(database, 0, 100) + small = _vm_steps(database, recent_window) + _append_history(database, 100, 4_900) + large = _vm_steps(database, recent_window) + unlimited = _vm_steps(database, lambda: session_queue.get_queue_item_ids("default")) # Fifty times the history costs the limited read about the same; the full listing shows the # measure does see history. @@ -131,69 +180,94 @@ def recent_window() -> object: assert unlimited >= 20 * large, (large, unlimited) -def test_point_reads_do_not_scan_the_queue_through_the_listing_index(session_queue: SqliteSessionQueue) -> None: +def test_point_reads_do_not_scan_the_queue_through_the_listing_index( + session_queue: SessionQueue, database: Database +) -> None: # A handful of ids already tips an unguarded plan onto the listing index; a page of them is typical. - summaries = _plans( - session_queue, lambda: session_queue.get_queue_item_summaries_by_ids("default", list(range(1, 51))) + summaries = _statements( + database, lambda: session_queue.get_queue_item_summaries_by_ids("default", list(range(1, 51))) ) - assert len(summaries) == 1 - assert "PRIMARY KEY" in summaries[0] + for plan in _plans(summaries): + assert "PRIMARY KEY" in plan + assert LISTING_INDEX not in plan - with session_queue._db.transaction() as cursor: - plan = cursor.execute(f"EXPLAIN QUERY PLAN {PENDING_QUEUE_ITEM_COUNT_QUERY}", ("default",)).fetchall() - assert "idx_session_queue_listing" not in " | ".join(row["detail"] for row in plan) + # The admission of every enqueue counts the pending backlog. + enqueue = _statements(database, lambda: _enqueue(session_queue)) + pending_counts = [statement for statement in enqueue if "count(" in statement.sql] + assert pending_counts, enqueue + for statement in pending_counts: + assert LISTING_INDEX not in statement.plan + assert STATUS_SEEK in statement.plan -def test_batch_lookups_stay_on_the_batch_index(session_queue: SqliteSessionQueue) -> None: +def test_batch_lookups_stay_on_the_batch_index(session_queue: SessionQueue, database: Database) -> None: """A batch is a handful of rows; a user_id index would read the user's whole history.""" - graph = Graph() - graph.add_node(PromptTestInvocation(id="prompt", prompt="test")) - enqueue = _plans( - session_queue, - lambda: session_queue._enqueue_batch("default", Batch(graph=graph, runs=2), False, "alice"), - ) + enqueue = _statements(database, lambda: _enqueue(session_queue)) # The batch-id uniqueness probe and the readback of the inserted item ids. - batch_reads = [plan for plan in enqueue if "batch_id" in plan or "user_id" in plan] + batch_reads = [statement for statement in enqueue if "batch_id" in statement.sql] assert len(batch_reads) == 2, enqueue - assert all("idx_session_queue_batch_id (batch_id=?)" in plan for plan in batch_reads), enqueue + assert all(BATCH_SEEK in statement.plan for statement in batch_reads), enqueue + + owner_status = _statements(database, lambda: session_queue.get_batch_status("default", "batch", user_id="alice")) + for plan in _plans(owner_status): + assert BATCH_SEEK in plan - owner_status = _plans(session_queue, lambda: session_queue.get_batch_status("default", "batch", user_id="alice")) - assert len(owner_status) == 1 - assert "idx_session_queue_batch_id (batch_id=?)" in owner_status[0] +def _assert_writes_go_by_id(statements: list[Statement]) -> None: + writes = [statement for statement in statements if statement.is_write] + assert writes, statements + assert all(ID_SEEK in statement.plan for statement in writes), writes -def test_bulk_cancel_and_clear_find_their_rows_through_the_status_indexes(session_queue: SqliteSessionQueue) -> None: + +def test_bulk_cancel_and_clear_find_their_rows_through_the_status_indexes( + session_queue: SessionQueue, database: Database +) -> None: """The rows a bulk cancel touches are the live ones, a handful next to the queue's history.""" - status_seek = "idx_session_queue_round_robin_pending (status=?" - user_status_seek = "idx_session_queue_round_robin_pending (status=? AND user_id=?)" - - # The live-row select, the bulk UPDATE, and the in-progress lookup for the per-item cancels. - whole_queue = _plans(session_queue, lambda: session_queue.cancel_by_queue_id("default")) - assert len(whole_queue) == 3, whole_queue - assert all(status_seek in plan for plan in whole_queue), whole_queue - - # A user's filter must neither keep the statements on the listing index nor move them onto a - # user index that reads the user's whole history. - user_scoped = _plans( - session_queue, - lambda: session_queue.cancel_by_queue_id("default", user_id="alice", origin_prefix="webv2:p:1:q:"), + + def assert_rows_found_by_status(statements: list[Statement]) -> None: + for statement in statements: + assert LISTING_INDEX not in statement.plan, statements + if statement.is_write: + continue + assert ANY_STATUS_SEEK.search(statement.plan), statement + # A user's filter must not move the read onto a user index that reads the user's whole history. + if INDEXED_USER_FILTER.search(statement.sql): + assert USER_STATUS_SEEK in statement.plan, statement + _assert_writes_go_by_id(statements) + + _enqueue(session_queue) + whole_queue = _statements(database, lambda: session_queue.cancel_by_queue_id("default")) + assert_rows_found_by_status(whole_queue) + assert all(STATUS_SEEK in statement.plan for statement in whole_queue if not statement.is_write), whole_queue + + _enqueue(session_queue) + user_scoped = _statements( + database, lambda: session_queue.cancel_by_queue_id("default", user_id="alice", origin_prefix=ORIGIN_PREFIX) ) - assert len(user_scoped) == 3, user_scoped - assert all(user_status_seek in plan for plan in user_scoped), user_scoped + assert_rows_found_by_status(user_scoped) + assert all(USER_STATUS_SEEK in statement.plan for statement in user_scoped if not statement.is_write) - # After the current chain lookup: the live-row select and the bulk UPDATE or DELETE. + # Beside the live rows, these look up the chains of the items running now, by their status. for bulk in (session_queue.cancel_all_except_current, session_queue.delete_all_except_current): - plans = _plans(session_queue, partial(bulk, "default", user_id="alice")) - assert len(plans) == 4, plans - assert all(user_status_seek in plan for plan in plans[2:]), plans - - # The in-progress lookup before the delete; the delete itself visits every row in scope. - cleared = _plans(session_queue, lambda: session_queue.clear("default", user_id="alice")) - assert len(cleared) == 2, cleared - assert user_status_seek in cleared[0], cleared - - # The batch's rows come from the batch index; only the in-progress lookup seeks by status. - batch = _plans(session_queue, lambda: session_queue.cancel_by_batch_ids("default", ["batch"], user_id="alice")) - assert len(batch) == 3, batch - assert all("idx_session_queue_batch_id (batch_id=?)" in plan for plan in batch[:2]), batch - assert status_seek in batch[2], batch + _enqueue(session_queue) + assert_rows_found_by_status(_statements(database, partial(bulk, "default", user_id="alice"))) + + # The in-progress lookup seeks by status; the delete itself visits every row in scope. + _enqueue(session_queue) + cleared = _statements(database, lambda: session_queue.clear("default", user_id="alice")) + in_progress = [statement for statement in cleared if statement.binds("in_progress")] + assert in_progress, cleared + assert all(USER_STATUS_SEEK in statement.plan for statement in in_progress), cleared + assert all(LISTING_INDEX not in statement.plan for statement in cleared), cleared + _assert_writes_go_by_id(cleared) + + # The batch's live rows come from the batch index; only the in-progress lookup seeks by status. + _enqueue(session_queue) + batch = _statements(database, lambda: session_queue.cancel_by_batch_ids("default", ["batch"], user_id="alice")) + reads = [statement for statement in batch if not statement.is_write] + live_rows = [statement for statement in reads if statement.binds("pending")] + in_progress = [statement for statement in reads if statement.binds("in_progress")] + assert live_rows and in_progress, batch + assert all(BATCH_SEEK in statement.plan for statement in live_rows), batch + assert all(STATUS_SEEK in statement.plan for statement in in_progress), batch + _assert_writes_go_by_id(batch) diff --git a/tests/app/services/session_queue/test_session_queue_reads.py b/tests/app/services/session_queue/test_session_queue_reads.py new file mode 100644 index 00000000000..e1f26bd0cc6 --- /dev/null +++ b/tests/app/services/session_queue/test_session_queue_reads.py @@ -0,0 +1,433 @@ +"""The session queue's reads on every backend: what each returns, in which order, and under which filters. + +The queue's writes still run on the SQLite cursor, so these tests write their rows directly, and read through the +service with only its database: its reads need nothing else. +""" + +import json +from typing import Any, Optional + +import pytest +from pydantic_core import to_jsonable_python +from sqlalchemy import insert + +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.session_queue import session_queue, session_queue_enqueue_receipts +from invokeai.app.services.shared.database.schema.users import users +from invokeai.app.services.shared.execution_state_migration import dump_execution_state +from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from invokeai.app.services.shared.pagination import SQLiteDirection + +QUEUE = "default" + + +@pytest.fixture +def accounts(database: Database) -> None: + with database.begin(write=True) as conn: + conn.execute( + insert(users), + [ + {"user_id": "alice", "email": "alice@example.com", "display_name": "Alice", "password_hash": "-"}, + {"user_id": "bob", "email": "bob@example.com", "display_name": "Bob", "password_hash": "-"}, + ], + ) + + +@pytest.fixture +def service(database: Database) -> SessionQueue: + return SessionQueue(database) + + +def _item( + database: Database, + *, + status: str = "pending", + priority: int = 0, + user_id: str = "alice", + created_at: str = "2026-01-01 00:00:00.000", + batch_id: str = "batch", + queue_id: str = QUEUE, + origin: Optional[str] = None, + destination: Optional[str] = None, + parent_item_id: Optional[int] = None, + workflow: Optional[str] = None, +) -> int: + session = GraphExecutionState(graph=Graph()) + values: dict[str, Any] = { + "queue_id": queue_id, + "batch_id": batch_id, + "session_id": session.id, + "session": json.dumps(dump_execution_state(session), default=to_jsonable_python), + "status": status, + "priority": priority, + "user_id": user_id, + "created_at": created_at, + "origin": origin, + "destination": destination, + "parent_item_id": parent_item_id, + "workflow": workflow, + "field_values": json.dumps([{"node_path": "n", "field_name": "seed", "value": 7}]), + } + with database.begin(write=True) as conn: + return int(conn.execute(insert(session_queue).values(**values)).inserted_primary_key[0]) + + +@pytest.mark.usefixtures("accounts") +class TestItems: + def test_an_item_carries_every_column_and_its_owners_name(self, database: Database) -> None: + item_id = _item(database, workflow='{"name": "w"}', origin="canvas", destination="gallery") + + item = database.queries.session_queue.item(item_id) + + assert item is not None + assert {column.name for column in session_queue.columns} <= item.keys() + assert (item["item_id"], item["user_display_name"], item["user_email"]) == ( + item_id, + "Alice", + "alice@example.com", + ) + assert (item["origin"], item["destination"], item["workflow"]) == ("canvas", "gallery", '{"name": "w"}') + assert database.queries.session_queue.workflow(item_id) == '{"name": "w"}' + + def test_a_missing_item_reads_as_none(self, database: Database) -> None: + assert database.queries.session_queue.item(404) is None + assert database.queries.session_queue.workflow(404) is None + + def test_an_item_of_an_account_that_is_gone_keeps_its_row(self, database: Database) -> None: + item_id = _item(database, user_id="nobody") + + item = database.queries.session_queue.item(item_id) + + assert item is not None and item["user_display_name"] is None and item["user_email"] is None + + def test_the_next_pending_item_is_the_highest_priority_then_the_oldest(self, database: Database) -> None: + _item(database, priority=0, created_at="2026-01-01 00:00:00.000") + _item(database, priority=1, created_at="2026-01-04 00:00:00.000") + oldest = _item(database, priority=1, created_at="2026-01-03 00:00:00.000") + _item(database, priority=1, created_at="2026-01-03 00:00:00.000") + _item(database, priority=5, status="completed") + + item = database.queries.session_queue.item_by_status(QUEUE, "pending", None) + + assert item is not None and item["item_id"] == oldest + + def test_the_current_item_is_the_earliest_enqueued_of_several(self, database: Database) -> None: + first = _item(database, status="in_progress") + _item(database, status="in_progress") + + item = database.queries.session_queue.item_by_status(QUEUE, "in_progress", None) + + assert item is not None and item["item_id"] == first + + def test_an_origin_prefix_ignores_case_and_takes_wildcards_as_written(self, database: Database) -> None: + _item(database, origin="canvasXworkflow", priority=2) + wanted = _item(database, origin="Canvas_workflow:1", priority=1) + + item = database.queries.session_queue.item_by_status(QUEUE, "pending", "canvas_") + + assert item is not None and item["item_id"] == wanted + + def test_other_queues_are_left_out(self, database: Database) -> None: + _item(database, queue_id="other") + + assert database.queries.session_queue.item_by_status(QUEUE, "pending", None) is None + assert database.queries.session_queue.count(QUEUE) == 0 + assert database.queries.session_queue.count("other") == 1 + + +@pytest.mark.usefixtures("accounts") +class TestListings: + def test_all_items_come_highest_priority_first_then_in_enqueue_order(self, database: Database) -> None: + low = _item(database, priority=0) + high = _item(database, priority=3, destination="canvas") + middle = _item(database, priority=1) + _item(database, queue_id="other") + + listed = [item["item_id"] for item in database.queries.session_queue.all_items(QUEUE, None)] + on_canvas = [item["item_id"] for item in database.queries.session_queue.all_items(QUEUE, "canvas")] + + assert listed == [high, middle, low] + assert on_canvas == [high] + + def test_a_later_page_keeps_every_filter(self, database: Database) -> None: + # Before the port, the keyset's OR escaped the filters from the second page on: the items after the cursor + # at its own priority were listed whatever their queue, status or destination. + kept = [_item(database, priority=1, status="pending", destination="canvas") for _ in range(3)] + _item(database, priority=1, status="completed", destination="canvas") + _item(database, priority=1, status="pending", destination="gallery") + _item(database, priority=1, status="pending", destination="canvas", queue_id="other") + kept.append(_item(database, priority=0, status="pending", destination="canvas")) + + def page(after: Optional[tuple[int, int]]) -> list[int]: + rows = database.queries.session_queue.page( + QUEUE, limit=2, after=after, status="pending", destination="canvas" + ) + return [row["item_id"] for row in rows] + + first = page(None) + second = page((1, first[-1])) + + assert first + second == kept + + def test_item_ids_follow_creation_and_keep_a_batch_in_enqueue_order(self, database: Database) -> None: + # A batch's items share their creation time; both directions list them in enqueue order, as SQLite always + # returned them, so that the queue's views window a large batch the same way. + old = _item(database, created_at="2026-01-01 00:00:00.000") + batch = [_item(database, created_at="2026-01-02 00:00:00.000") for _ in range(3)] + _item(database, user_id="bob", created_at="2026-01-03 00:00:00.000") + _item(database, queue_id="other", created_at="2026-01-04 00:00:00.000") + + newest_first = database.queries.session_queue.item_ids( + QUEUE, descending=True, user_id="alice", origin_prefix=None + ) + oldest_first = database.queries.session_queue.item_ids( + QUEUE, descending=False, user_id="alice", origin_prefix=None + ) + + assert newest_first == [*batch, old] + assert oldest_first == [old, *batch] + + def test_item_ids_by_origin(self, database: Database) -> None: + wanted = _item(database, origin="workflows:1") + _item(database, origin="canvas:1") + + assert database.queries.session_queue.item_ids( + QUEUE, descending=True, user_id=None, origin_prefix="WORKFLOWS" + ) == [wanted] + + def test_summaries_are_of_the_named_items_of_the_queue(self, database: Database) -> None: + item_ids = [_item(database, destination=f"d{i}") for i in range(5)] + foreign = _item(database, queue_id="other") + + summaries = database.queries.session_queue.summaries(QUEUE, [*item_ids[1:], foreign, 404]) + + assert sorted(summary["item_id"] for summary in summaries) == item_ids[1:] + assert {summary["user_display_name"] for summary in summaries} == {"Alice"} + + +@pytest.mark.usefixtures("accounts") +class TestCounts: + def test_status_counts_are_global_with_the_accounts_own_and_the_earliest_current_item( + self, database: Database + ) -> None: + for status, user_id in [ + ("pending", "alice"), + ("pending", "bob"), + ("pending", "bob"), + ("in_progress", "bob"), + ("in_progress", "alice"), + ("completed", "alice"), + ]: + _item(database, status=status, user_id=user_id) + _item(database, status="pending", queue_id="other") + + counts, own, current = database.queries.session_queue.status_counts(QUEUE, "alice", None) + + assert counts == {"pending": 3, "in_progress": 2, "completed": 1} + assert own == {"pending": 1, "in_progress": 1, "completed": 1} + assert current is not None and current.user_id == "bob" + + def test_status_counts_by_origin_also_scope_the_accounts_own(self, database: Database) -> None: + # A project's badge: the account's items of other origins are not its own count in this one. + _item(database, status="in_progress", origin="canvas") + _item(database, status="pending", origin="project:1:q:") + _item(database, status="pending", origin="project:1:q:", user_id="bob") + _item(database, status="pending", origin="project:2:q:") + + counts, own, current = database.queries.session_queue.status_counts(QUEUE, "alice", "project:1:q:") + nobodys = database.queries.session_queue.status_counts(QUEUE, None, "project:1:q:") + + assert (counts, own, current) == ({"pending": 2}, {"pending": 1}, None) + assert nobodys.own_counts == {} + + def test_batch_counts_and_the_batchs_origin(self, database: Database) -> None: + _item(database, batch_id="b", status="pending", origin="canvas", destination="gallery") + _item(database, batch_id="b", status="failed", origin="canvas", destination="gallery", user_id="bob") + _item(database, batch_id="b", status="failed", queue_id="other") + _item(database, batch_id="other", status="pending") + + everyone = database.queries.session_queue.batch_counts(QUEUE, "b", None) + bobs = database.queries.session_queue.batch_counts(QUEUE, "b", "bob") + none = database.queries.session_queue.batch_counts(QUEUE, "missing", None) + + assert everyone == ({"pending": 1, "failed": 1}, "canvas", "gallery") + assert bobs.counts == {"failed": 1} + assert none == ({}, None, None) + + def test_destination_counts(self, database: Database) -> None: + _item(database, destination="canvas", status="pending") + _item(database, destination="canvas", status="pending", user_id="bob") + _item(database, destination="canvas", status="pending", queue_id="other") + _item(database, destination="gallery", status="pending") + + assert database.queries.session_queue.destination_counts(QUEUE, "canvas", None) == {"pending": 2} + assert database.queries.session_queue.destination_counts(QUEUE, "canvas", "bob") == {"pending": 1} + + +@pytest.mark.usefixtures("accounts") +class TestWorkflowCallChains: + def test_a_chain_is_the_ancestors_the_item_and_the_roots_descendants_breadth_first( + self, database: Database + ) -> None: + # Deep enough that breadth first and depth first differ: depth first would list a1 before b. + root = _item(database, status="waiting") + a = _item(database, parent_item_id=root) + b = _item(database, parent_item_id=root) + a1 = _item(database, parent_item_id=a) + b1 = _item(database, parent_item_id=b) + + assert database.queries.session_queue.chain_item_ids(a1) == [a, root, a1, b, b1] + assert database.queries.session_queue.descendant_ids(root) == [a, b, a1, b1] + + def test_a_chain_with_a_missing_item_or_ancestor_reads_as_none(self, database: Database) -> None: + orphan = _item(database, parent_item_id=404) + + assert database.queries.session_queue.chain_item_ids(orphan) is None + assert database.queries.session_queue.chain_item_ids(405) is None + + def test_a_cycle_ends_the_walk(self, database: Database) -> None: + # The app links no cycle; a corrupt one must not hold the transaction forever. + first = _item(database) + second = _item(database, parent_item_id=first) + with database.begin(write=True) as conn: + conn.execute(session_queue.update().where(session_queue.c.item_id == first).values(parent_item_id=second)) + + assert set(database.queries.session_queue.chain_item_ids(second) or []) == {first, second} + assert database.queries.session_queue.descendant_ids(first) == [second] + + def test_the_current_chains_are_those_of_every_item_in_progress(self, database: Database) -> None: + first_root = _item(database, status="waiting") + first = _item(database, status="in_progress", parent_item_id=first_root) + second = _item(database, status="in_progress") + _item(database, status="pending") + _item(database, status="in_progress", queue_id="other") + + assert database.queries.session_queue.current_chain_item_ids(QUEUE) == {first_root, first, second} + + def test_a_running_item_whose_parent_row_is_gone_still_protects_its_chain(self, database: Database) -> None: + # "Cancel all except current" must neither fail nor cancel a running item's children because an ancestor + # row vanished: what of the chain is left stays protected. + grandparent = _item(database, status="waiting") + parent = _item(database, status="waiting", parent_item_id=grandparent) + running = _item(database, status="in_progress", parent_item_id=parent) + child = _item(database, status="pending", parent_item_id=running) + with database.begin(write=True) as conn: + conn.execute(session_queue.delete().where(session_queue.c.item_id == grandparent)) + other = _item(database, status="in_progress") + + assert database.queries.session_queue.current_chain_item_ids(QUEUE) == {parent, running, child, other} + + def test_without_one_in_progress_the_chain_of_the_next_waiting_child_is_current(self, database: Database) -> None: + # The gap between two children of a suspended workflow call: the child that runs next keeps its chain, not + # another suspended chain, and not a pending child of a parent that no longer waits. + waiting = _item(database, status="waiting") + next_child = _item(database, status="pending", parent_item_id=waiting, priority=2) + sibling = _item(database, status="pending", parent_item_id=waiting, priority=1) + other_waiting = _item(database, status="waiting") + _item(database, status="pending", parent_item_id=other_waiting, priority=0) + other_queue_waiting = _item(database, status="waiting", queue_id="other") + _item(database, status="pending", parent_item_id=other_queue_waiting, priority=9, queue_id="other") + finished = _item(database, status="completed") + _item(database, status="pending", parent_item_id=finished, priority=3) + _item(database, status="pending", priority=4) + + assert database.queries.session_queue.current_chain_item_ids(QUEUE) == {waiting, next_child, sibling} + + def test_nothing_running_means_no_current_chain(self, database: Database) -> None: + _item(database, status="pending") + + assert database.queries.session_queue.current_chain_item_ids(QUEUE) == set() + + +@pytest.mark.usefixtures("accounts") +class TestServiceReads: + """The DTOs the service builds from these rows, on every backend.""" + + def test_an_item_reads_whole_and_for_the_api(self, database: Database, service: SessionQueue) -> None: + item_id = _item(database, origin="canvas", destination="gallery") + + item = service.get_queue_item(item_id) + for_api = service.get_queue_item_for_api(item_id) + + assert item._snapshot_readable and for_api._snapshot_readable + assert (item.item_id, item.user_display_name, item.field_values[0].value) == (item_id, "Alice", 7) # type: ignore[index] + assert (for_api.status, for_api.origin, for_api.destination) == ("pending", "canvas", "gallery") + + def test_summaries_follow_the_callers_order(self, database: Database, service: SessionQueue) -> None: + item_ids = [_item(database) for _ in range(3)] + child = _item(database, parent_item_id=item_ids[0]) + + summaries = service.get_queue_item_summaries_by_ids(QUEUE, [child, *reversed(item_ids)]) + + assert [summary.item_id for summary in summaries] == [child, *reversed(item_ids)] + assert (summaries[0].parent_item_id, summaries[0].user_email) == (item_ids[0], "alice@example.com") + + def test_batch_status_and_destination_counts(self, database: Database, service: SessionQueue) -> None: + _item(database, batch_id="b", status="pending", origin="canvas", destination="gallery") + _item(database, batch_id="b", status="completed", origin="canvas", destination="gallery") + + batch = service.get_batch_status(QUEUE, "b") + destination = service.get_counts_by_destination(QUEUE, "gallery") + + assert (batch.origin, batch.destination, batch.pending, batch.completed, batch.total) == ( + "canvas", + "gallery", + 1, + 1, + 2, + ) + assert (destination.pending, destination.completed, destination.total) == (1, 1, 2) + + def test_queue_status_hides_another_accounts_current_item(self, database: Database, service: SessionQueue) -> None: + bobs = _item(database, status="in_progress", user_id="bob") + _item(database, status="pending") + + as_alice = service.get_queue_status(QUEUE, user_id="alice") + as_bob = service.get_queue_status(QUEUE, user_id="bob") + + assert (as_alice.item_id, as_alice.user_pending, as_alice.in_progress) == (None, 1, 1) + assert (as_bob.item_id, as_bob.user_in_progress, as_bob.total) == (bobs, 1, 2) + + def test_current_next_and_ids(self, database: Database, service: SessionQueue) -> None: + running = _item(database, status="in_progress", created_at="2026-01-01 00:00:00.000") + waiting = _item(database, status="pending", created_at="2026-01-02 00:00:00.000") + + current, next_item = service.get_current(QUEUE), service.get_next_for_api(QUEUE) + ids = service.get_queue_item_ids(QUEUE, SQLiteDirection.Descending) + + assert (current and current.item_id, next_item and next_item.item_id) == (running, waiting) + assert (ids.item_ids, ids.total_count) == ([waiting, running], 2) + + def test_list_queue_items_pages_by_its_cursor(self, database: Database, service: SessionQueue) -> None: + item_ids = [_item(database, priority=1) for _ in range(3)] + + first = service.list_queue_items(QUEUE, limit=2, priority=0) + rest = service.list_queue_items(QUEUE, limit=2, priority=1, cursor=item_ids[1]) + + assert ([item.item_id for item in first.items], first.has_more) == (item_ids[:2], True) + assert ([item.item_id for item in rest.items], rest.has_more) == (item_ids[2:], False) + + +@pytest.mark.usefixtures("accounts") +def test_an_enqueue_receipt_is_found_by_queue_account_and_key(database: Database) -> None: + with database.begin(write=True) as conn: + conn.execute( + insert(session_queue_enqueue_receipts).values( + queue_id=QUEUE, + user_id="alice", + idempotency_key="key", + payload_hash="hash", + batch_id="b", + requested=3, + enqueued=2, + priority=0, + item_ids="[1,2]", + byte_size=10, + ) + ) + + receipt = database.queries.session_queue.enqueue_receipt(QUEUE, "alice", "key") + + assert receipt == ("b", 3, 2, "[1,2]") + assert database.queries.session_queue.enqueue_receipt(QUEUE, "bob", "key") is None diff --git a/tests/app/services/session_queue/test_session_queue_status_event_isolation.py b/tests/app/services/session_queue/test_session_queue_status_event_isolation.py index 5c46650c7cc..8e354fe695c 100644 --- a/tests/app/services/session_queue/test_session_queue_status_event_isolation.py +++ b/tests/app/services/session_queue/test_session_queue_status_event_isolation.py @@ -13,26 +13,28 @@ from invokeai.app.invocations.call_saved_workflow import CallSavedWorkflowInvocation from invokeai.app.services.events.events_common import QueueItemStatusChangedEvent from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from tests.fixtures.sqlite_database import sqlite_cursor_of from tests.test_nodes import PromptTestInvocation, TestEventService @pytest.fixture -def session_queue(mock_invoker: Invoker) -> SqliteSessionQueue: - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) +def session_queue(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: + db = mock_sqlite_database + queue = SessionQueue(db) queue.start(mock_invoker) return queue -def _insert_queue_item(session_queue: SqliteSessionQueue, user_id: str) -> int: +def _insert_queue_item(session_queue: SessionQueue, user_id: str) -> int: graph = Graph() graph.add_node(PromptTestInvocation(id="prompt", prompt="test")) session = GraphExecutionState(graph=graph) session_json = session.model_dump_json(warnings=False, exclude_none=True) batch_id = str(uuid.uuid4()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -46,9 +48,7 @@ def _insert_queue_item(session_queue: SqliteSessionQueue, user_id: str) -> int: return cursor.lastrowid # type: ignore[return-value] -def _insert_waiting_workflow_call_parent( - session_queue: SqliteSessionQueue, user_id: str -) -> tuple[int, GraphExecutionState]: +def _insert_waiting_workflow_call_parent(session_queue: SessionQueue, user_id: str) -> tuple[int, GraphExecutionState]: parent_graph = Graph() parent_graph.add_node(CallSavedWorkflowInvocation(id="call-node", workflow_id="workflow-a")) parent_session = GraphExecutionState(graph=parent_graph) @@ -61,7 +61,7 @@ def _insert_waiting_workflow_call_parent( parent_session.attach_waiting_workflow_call_child_session(child_session) batch_id = str(uuid.uuid4()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -94,9 +94,7 @@ def _last_status_event_for_item(event_bus: TestEventService, item_id: int) -> Qu return matches[-1] -def test_event_redacts_other_users_current_item_identifiers( - session_queue: SqliteSessionQueue, mock_invoker: Invoker -) -> None: +def test_event_redacts_other_users_current_item_identifiers(session_queue: SessionQueue, mock_invoker: Invoker) -> None: """When user A's pending item is canceled while user B's item is in_progress, the embedded queue_status in A's status-changed event must not expose B's identifiers.""" user_a = "user-a" @@ -140,7 +138,7 @@ def test_event_redacts_other_users_current_item_identifiers( def test_get_queue_status_does_not_load_full_current_session( - session_queue: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: item_id = _insert_queue_item(session_queue, user_id="user-a") in_progress = session_queue.dequeue() @@ -158,9 +156,7 @@ def fail_get_current(queue_id: str): assert status.batch_id == in_progress.batch_id -def test_event_preserves_owner_current_item_identifiers( - session_queue: SqliteSessionQueue, mock_invoker: Invoker -) -> None: +def test_event_preserves_owner_current_item_identifiers(session_queue: SessionQueue, mock_invoker: Invoker) -> None: """When the current in-progress item belongs to the same user as the changed item, the embedded queue_status must continue to expose the identifiers (no over-redaction).""" user_a = "user-a" @@ -186,7 +182,7 @@ def test_event_preserves_owner_current_item_identifiers( def test_event_redaction_uses_same_lightweight_snapshot( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, mock_invoker: Invoker, ) -> None: """The status query must use one lightweight snapshot for identifiers and redaction.""" @@ -217,7 +213,7 @@ def test_event_redaction_uses_same_lightweight_snapshot( def test_event_preserves_identifiers_when_current_item_is_the_changed_item( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker ) -> None: """The dequeue() transition makes the changed item itself the in-progress current item. queue_status must expose its identifiers since they belong to the event's owner.""" @@ -239,7 +235,7 @@ def test_event_preserves_identifiers_when_current_item_is_the_changed_item( assert a_event.queue_status.batch_id == in_progress.batch_id -def test_event_carries_owner_per_user_counts(session_queue: SqliteSessionQueue, mock_invoker: Invoker) -> None: +def test_event_carries_owner_per_user_counts(session_queue: SessionQueue, mock_invoker: Invoker) -> None: """The embedded queue_status carries the event owner's per-user counts so the owner's client can update personal UI (progress bar, spinner, favicon) from the event immediately, without waiting for a status refetch. The sanitized companion nulls these before reaching non-owners.""" @@ -269,7 +265,7 @@ def test_event_carries_owner_per_user_counts(session_queue: SqliteSessionQueue, def test_workflow_call_child_enqueue_event_redacts_other_users_current_item_identifiers( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker ) -> None: """The child enqueue path emits QueueItemStatusChangedEvent without going through _set_queue_item_status, so it must apply the same per-owner current-item redaction.""" diff --git a/tests/app/services/session_queue/test_session_queue_status_sequence.py b/tests/app/services/session_queue/test_session_queue_status_sequence.py index 0dcf8fd0f71..03caa60bf78 100644 --- a/tests/app/services/session_queue/test_session_queue_status_sequence.py +++ b/tests/app/services/session_queue/test_session_queue_status_sequence.py @@ -4,21 +4,23 @@ from invokeai.app.services.events.events_common import QueueItemStatusChangedEvent from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from tests.fixtures.sqlite_database import sqlite_cursor_of from tests.test_nodes import PromptTestInvocation, TestEventService @pytest.fixture -def session_queue(mock_invoker: Invoker) -> SqliteSessionQueue: - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) +def session_queue(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: + db = mock_sqlite_database + queue = SessionQueue(db) queue.start(mock_invoker) return queue def _insert_queue_item( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, queue_id: str = "default", destination: str | None = None, ) -> int: @@ -27,7 +29,7 @@ def _insert_queue_item( session = GraphExecutionState(graph=graph) session_json = session.model_dump_json(warnings=False, exclude_none=True) batch_id = str(uuid.uuid4()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -51,7 +53,7 @@ def _insert_queue_item( def test_status_sequence_increments_for_queue_item_lifecycle( - session_queue: SqliteSessionQueue, mock_invoker: Invoker + session_queue: SessionQueue, mock_invoker: Invoker ) -> None: item_id = _insert_queue_item(session_queue) @@ -77,7 +79,7 @@ def test_status_sequence_increments_for_queue_item_lifecycle( assert [event.status_sequence for event in status_events] == [1, 2] -def test_status_sequence_increments_for_bulk_cancel_paths(session_queue: SqliteSessionQueue) -> None: +def test_status_sequence_increments_for_bulk_cancel_paths(session_queue: SessionQueue) -> None: first_item_id = _insert_queue_item(session_queue) second_item_id = _insert_queue_item(session_queue) @@ -90,7 +92,7 @@ def test_status_sequence_increments_for_bulk_cancel_paths(session_queue: SqliteS assert session_queue.get_queue_item(second_item_id).status_sequence == 1 -def test_status_sequence_continues_after_dequeue_then_cancel(session_queue: SqliteSessionQueue) -> None: +def test_status_sequence_continues_after_dequeue_then_cancel(session_queue: SessionQueue) -> None: item_id = _insert_queue_item(session_queue) in_progress_item = session_queue.dequeue() diff --git a/tests/app/services/session_queue/test_session_queue_status_user_scoping.py b/tests/app/services/session_queue/test_session_queue_status_user_scoping.py index 69c88fe5697..19c0e767fd7 100644 --- a/tests/app/services/session_queue/test_session_queue_status_user_scoping.py +++ b/tests/app/services/session_queue/test_session_queue_status_user_scoping.py @@ -21,26 +21,28 @@ import pytest from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from tests.fixtures.sqlite_database import sqlite_cursor_of from tests.test_nodes import PromptTestInvocation @pytest.fixture -def session_queue(mock_invoker: Invoker) -> SqliteSessionQueue: - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) +def session_queue(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: + db = mock_sqlite_database + queue = SessionQueue(db) queue.start(mock_invoker) return queue -def _insert_queue_item(session_queue: SqliteSessionQueue, user_id: str | None, origin: str | None = None) -> int: +def _insert_queue_item(session_queue: SessionQueue, user_id: str | None, origin: str | None = None) -> int: graph = Graph() graph.add_node(PromptTestInvocation(id="prompt", prompt="test")) session = GraphExecutionState(graph=graph) session_json = session.model_dump_json(warnings=False, exclude_none=True) batch_id = str(uuid.uuid4()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -54,7 +56,7 @@ def _insert_queue_item(session_queue: SqliteSessionQueue, user_id: str | None, o return cursor.lastrowid # type: ignore[return-value] -def test_status_aggregate_counts_are_global_with_user_subcounts(session_queue: SqliteSessionQueue) -> None: +def test_status_aggregate_counts_are_global_with_user_subcounts(session_queue: SessionQueue) -> None: """A non-admin caller (user_id set) sees global aggregate counts plus their own subcounts.""" user_a = "user-a" user_b = "user-b" @@ -73,7 +75,7 @@ def test_status_aggregate_counts_are_global_with_user_subcounts(session_queue: S def test_status_counts_every_status_with_the_callers_share_inside_the_origin_scope( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: items = [ ("alice", "project-a:1", "pending"), @@ -89,7 +91,7 @@ def test_status_counts_every_status_with_the_callers_share_inside_the_origin_sco ] for user_id, origin, status in items: item_id = _insert_queue_item(session_queue, user_id=user_id, origin=origin) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute("UPDATE session_queue SET status = ? WHERE item_id = ?", (status, item_id)) def counts(user_id: str, origin_prefix: str | None = None) -> tuple[int, ...]: @@ -114,7 +116,7 @@ def counts(user_id: str, origin_prefix: str | None = None) -> tuple[int, ...]: assert counts("carol", "project-b:") == (1, 1, 0, 0, 0, 1, 3, 0, 0) -def test_status_admin_global_call_omits_user_subcounts(session_queue: SqliteSessionQueue) -> None: +def test_status_admin_global_call_omits_user_subcounts(session_queue: SessionQueue) -> None: """An admin/global caller (user_id=None) gets global counts and no per-user subcounts.""" _insert_queue_item(session_queue, user_id="user-a") _insert_queue_item(session_queue, user_id="user-b") @@ -127,7 +129,7 @@ def test_status_admin_global_call_omits_user_subcounts(session_queue: SqliteSess assert status.user_in_progress is None -def test_status_current_item_redacted_for_non_owner_but_counts_global(session_queue: SqliteSessionQueue) -> None: +def test_status_current_item_redacted_for_non_owner_but_counts_global(session_queue: SessionQueue) -> None: """When the in-progress item belongs to another user, its identifiers are hidden from a non-owner, but the aggregate counts (and the user's own subcounts) remain populated.""" user_a = "user-a" @@ -151,7 +153,7 @@ def test_status_current_item_redacted_for_non_owner_but_counts_global(session_qu def test_status_admin_caller_gets_user_subcounts_and_current_item_identifiers( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: """An admin caller (user_id set, is_admin=True) gets their own subcounts - so personal UI like the progress bar can distinguish the admin's own activity from other users' - while @@ -177,7 +179,7 @@ def test_status_admin_caller_gets_user_subcounts_and_current_item_identifiers( def test_status_origin_scope_preserves_admin_subcounts_and_current_item_visibility( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: """Origin scoping applies consistently to aggregate counts, personal counts, and the current-item snapshot without hiding another owner's identifiers from an admin.""" @@ -205,7 +207,7 @@ def test_status_origin_scope_preserves_admin_subcounts_and_current_item_visibili assert status.user_in_progress == 0 -def test_status_origin_scope_excludes_out_of_scope_current_item(session_queue: SqliteSessionQueue) -> None: +def test_status_origin_scope_excludes_out_of_scope_current_item(session_queue: SessionQueue) -> None: """The current-item snapshot must honour origin_prefix, not just the aggregate counts. get_queue_status reads the in-progress item with a lightweight 4-column query rather than @@ -231,7 +233,7 @@ def test_status_origin_scope_excludes_out_of_scope_current_item(session_queue: S assert status.total == 1 -def test_get_queue_item_ids_returns_all_users_ids(session_queue: SqliteSessionQueue) -> None: +def test_get_queue_item_ids_returns_all_users_ids(session_queue: SessionQueue) -> None: """get_queue_item_ids returns ids for every user so the virtualized list can show the (redacted) entries belonging to other users. Redaction happens at hydration time.""" a_item_id = _insert_queue_item(session_queue, user_id="user-a") @@ -244,11 +246,11 @@ def test_get_queue_item_ids_returns_all_users_ids(session_queue: SqliteSessionQu def test_get_queue_item_summaries_by_ids_returns_only_requested_queue_items_in_order( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: first_id = _insert_queue_item(session_queue, user_id="user-a") second_id = _insert_queue_item(session_queue, user_id="user-b") - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql UPDATE session_queue @@ -280,16 +282,16 @@ def test_get_queue_item_summaries_by_ids_returns_only_requested_queue_items_in_o assert summaries[0].status == "pending" -def test_get_queue_item_summaries_by_ids_chunks_past_the_sqlite_bind_limit( - session_queue: SqliteSessionQueue, +def test_get_queue_item_summaries_by_ids_takes_more_ids_than_sqlite_binds( + session_queue: SessionQueue, ) -> None: - """Every id becomes one bind parameter, so a single IN (...) would raise OperationalError once - the list outgrows SQLite's per-statement variable limit. The query is chunked instead.""" + """A list longer than SQLite's per-statement variable limit is answered, in the caller's order: + an IN (...) binding every id would raise OperationalError past that limit.""" real_ids = [_insert_queue_item(session_queue, user_id="user-a") for _ in range(3)] # Size the request off the limit this SQLite build actually enforces (999 on builds older than # 3.32, 32766 since), so the test stays meaningful wherever it runs. - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: bind_limit = cursor.connection.getlimit(sqlite3.SQLITE_LIMIT_VARIABLE_NUMBER) # Padding ids do not exist, which also covers chunks that match nothing at all. diff --git a/tests/app/services/session_queue/test_session_queue_workflow_call_metadata.py b/tests/app/services/session_queue/test_session_queue_workflow_call_metadata.py index a7169f96c3d..596b070d953 100644 --- a/tests/app/services/session_queue/test_session_queue_workflow_call_metadata.py +++ b/tests/app/services/session_queue/test_session_queue_workflow_call_metadata.py @@ -24,18 +24,20 @@ SessionQueueItemNotFoundError, TooManySessionsError, ) -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph, GraphExecutionState from invokeai.app.services.shared.invocation_context import ImagesInterface, InvocationContextData from invokeai.app.services.urls.urls_default import LocalUrlService from invokeai.app.services.workflow_records.workflow_records_common import WorkflowMeta, WorkflowWithoutID +from tests.fixtures.sqlite_database import sqlite_cursor_of from tests.test_nodes import TestEventService @pytest.fixture -def session_queue(mock_invoker: Invoker) -> SqliteSessionQueue: - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) +def session_queue(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: + db = mock_sqlite_database + queue = SessionQueue(db) queue.start(mock_invoker) return queue @@ -47,7 +49,7 @@ def event_bus(mock_invoker: Invoker) -> TestEventService: def _insert_queue_item( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, *, session: GraphExecutionState, status: str, @@ -63,7 +65,7 @@ def _insert_queue_item( workflow_call_depth: int | None = None, workflow: WorkflowWithoutID | None = None, ) -> int: - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -130,7 +132,7 @@ def _workflow_without_id() -> WorkflowWithoutID: def _build_waiting_workflow_call_parent( - session_queue: SqliteSessionQueue, child_count: int, project_id: str | None = None + session_queue: SessionQueue, child_count: int, project_id: str | None = None ) -> tuple[int, GraphExecutionState, list[GraphExecutionState]]: graph = Graph() graph.add_node(CallSavedWorkflowInvocation(id="call-node", workflow_id="workflow-a")) @@ -147,11 +149,11 @@ def _build_waiting_workflow_call_parent( return parent_item_id, parent_session, child_sessions -def test_get_queue_item_round_trips_workflow_call_metadata(session_queue: SqliteSessionQueue) -> None: +def test_get_queue_item_round_trips_workflow_call_metadata(session_queue: SessionQueue) -> None: session = GraphExecutionState(graph=Graph()) session_json = session.model_dump_json(warnings=False) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -205,7 +207,7 @@ def test_get_queue_item_round_trips_workflow_call_metadata(session_queue: Sqlite def test_save_queue_item_session_does_not_reload_full_queue_item( - session_queue: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: session = GraphExecutionState(graph=Graph()) item_id = _insert_queue_item(session_queue, session=session, status="pending") @@ -218,7 +220,7 @@ def fail_get_queue_item(item_id: int): session_queue.save_queue_item_session(item_id, session) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute("SELECT session FROM session_queue WHERE item_id = ?", (item_id,)) row = cursor.fetchone() assert row is not None @@ -226,7 +228,7 @@ def fail_get_queue_item(item_id: int): assert persisted.errors == {"node": "updated"} -def test_active_save_queue_item_session_does_not_overwrite_terminal_item(session_queue: SqliteSessionQueue) -> None: +def test_active_save_queue_item_session_does_not_overwrite_terminal_item(session_queue: SessionQueue) -> None: session = GraphExecutionState(graph=Graph()) item_id = _insert_queue_item(session_queue, session=session, status="pending") stale_session = session_queue.get_queue_item(item_id).session @@ -240,9 +242,7 @@ def test_active_save_queue_item_session_does_not_overwrite_terminal_item(session assert persisted.session.errors == {} -def test_failed_child_transitions_sqlite_parent_to_failed( - session_queue: SqliteSessionQueue, mock_invoker: Invoker -) -> None: +def test_failed_child_transitions_sqlite_parent_to_failed(session_queue: SessionQueue, mock_invoker: Invoker) -> None: parent_item_id, _parent_session, _child_sessions = _build_waiting_workflow_call_parent(session_queue, child_count=1) mock_invoker.services.session_queue = session_queue runner = DefaultSessionRunner() @@ -258,7 +258,7 @@ def test_failed_child_transitions_sqlite_parent_to_failed( def test_enqueue_workflow_call_children_rejects_stale_parent_session( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: parent_item_id, _parent_session, child_sessions = _build_waiting_workflow_call_parent(session_queue, child_count=2) stale_parent = session_queue.get_queue_item(parent_item_id) @@ -276,13 +276,13 @@ def test_enqueue_workflow_call_children_rejects_stale_parent_session( assert persisted_parent.status == "in_progress" assert persisted_parent.session.errors == {"sibling": "newer parent state"} assert [item.item_id for item in session_queue.list_all_queue_items("default")] == [parent_item_id] - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute("SELECT COUNT(*) FROM session_queue") assert cursor.fetchone()[0] == 1 def test_status_transition_reuses_loaded_queue_item( - session_queue: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: session = GraphExecutionState(graph=Graph()) item_id = _insert_queue_item(session_queue, session=session, status="in_progress") @@ -301,7 +301,7 @@ def fail_get_queue_item(item_id: int): def test_enqueue_workflow_call_child_inherits_workflow_for_image_metadata( - session_queue: SqliteSessionQueue, mock_invoker: Invoker, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, mock_invoker: Invoker, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: parent_graph = Graph() parent_graph.add_node(CallSavedWorkflowInvocation(id="call-node", workflow_id="workflow-a")) @@ -314,7 +314,7 @@ def test_enqueue_workflow_call_child_inherits_workflow_for_image_metadata( parent_session.begin_waiting_on_workflow_call(frame) parent_session.attach_waiting_workflow_call_child_session(child_session) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -407,7 +407,7 @@ def count_workflow_lookups(item_id: int) -> str | None: def test_enqueue_workflow_call_child_rejects_canceled_stale_parent( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: parent_graph = Graph() parent_graph.add_node(CallSavedWorkflowInvocation(id="call-node", workflow_id="workflow-a")) @@ -420,7 +420,7 @@ def test_enqueue_workflow_call_child_rejects_canceled_stale_parent( parent_session.begin_waiting_on_workflow_call(frame) parent_session.attach_waiting_workflow_call_child_session(child_session) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -451,13 +451,13 @@ def test_enqueue_workflow_call_child_rejects_canceled_stale_parent( with pytest.raises(ValueError, match="terminal parent"): session_queue.enqueue_workflow_call_child(stale_parent, child_session) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute("SELECT COUNT(*) FROM session_queue WHERE parent_item_id = ?", (parent_item_id,)) assert cursor.fetchone()[0] == 0 def test_enqueue_workflow_call_children_publishes_parent_state_before_children( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: parent_item_id, _parent_session, child_sessions = _build_waiting_workflow_call_parent( session_queue, child_count=2, project_id="project-1" @@ -488,7 +488,7 @@ def test_enqueue_workflow_call_children_publishes_parent_state_before_children( def test_concurrent_workflow_call_child_completions_preserve_both_siblings( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: parent_item_id, _parent_session, child_sessions = _build_waiting_workflow_call_parent(session_queue, child_count=2) child_queue_items = session_queue.enqueue_workflow_call_children( @@ -563,7 +563,7 @@ def prune(): assert prune() == 3 -def test_enqueue_workflow_call_child_rejects_full_pending_queue(session_queue: SqliteSessionQueue) -> None: +def test_enqueue_workflow_call_child_rejects_full_pending_queue(session_queue: SessionQueue) -> None: parent_graph = Graph() parent_graph.add_node(CallSavedWorkflowInvocation(id="call-node", workflow_id="workflow-a")) parent_session = GraphExecutionState(graph=parent_graph) @@ -578,7 +578,7 @@ def test_enqueue_workflow_call_child_rejects_full_pending_queue(session_queue: S parent_item_id = _insert_queue_item(session_queue, session=parent_session, status="waiting") _insert_queue_item(session_queue, session=GraphExecutionState(graph=Graph()), status="pending") parent_queue_item = session_queue.get_queue_item(parent_item_id) - session_queue._SqliteSessionQueue__invoker.services.configuration.max_queue_size = 1 + session_queue._SessionQueue__invoker.services.configuration.max_queue_size = 1 with pytest.raises(TooManySessionsError, match="remaining queue capacity"): session_queue.enqueue_workflow_call_child(parent_queue_item, child_session) @@ -587,7 +587,7 @@ def test_enqueue_workflow_call_child_rejects_full_pending_queue(session_queue: S def test_enqueue_workflow_call_child_preserves_workflow_through_nested_calls( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: workflow = _workflow_without_id() @@ -635,7 +635,7 @@ def test_enqueue_workflow_call_child_preserves_workflow_through_nested_calls( assert session_queue.get_queue_item_workflow_json(root_item_id) == workflow.model_dump_json() -def test_enqueue_workflow_call_child_persists_batch_field_values(session_queue: SqliteSessionQueue) -> None: +def test_enqueue_workflow_call_child_persists_batch_field_values(session_queue: SessionQueue) -> None: parent_graph = Graph() parent_graph.add_node(CallSavedWorkflowInvocation(id="call-node", workflow_id="workflow-a")) parent_session = GraphExecutionState(graph=parent_graph) @@ -647,7 +647,7 @@ def test_enqueue_workflow_call_child_persists_batch_field_values(session_queue: parent_session.begin_waiting_on_workflow_call(frame) parent_session.attach_waiting_workflow_call_child_session(child_session) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -693,7 +693,7 @@ def test_enqueue_workflow_call_child_persists_batch_field_values(session_queue: def test_suspend_and_enqueue_child_emit_waiting_then_pending_status_events( - session_queue: SqliteSessionQueue, event_bus: TestEventService + session_queue: SessionQueue, event_bus: TestEventService ) -> None: parent_graph = Graph() parent_graph.add_node(CallSavedWorkflowInvocation(id="call-node", workflow_id="workflow-a")) @@ -706,7 +706,7 @@ def test_suspend_and_enqueue_child_emit_waiting_then_pending_status_events( parent_session.begin_waiting_on_workflow_call(frame) parent_session.attach_waiting_workflow_call_child_session(child_session) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -764,11 +764,11 @@ def test_suspend_and_enqueue_child_emit_waiting_then_pending_status_events( assert child_pending_event.queue_status.item_id is None -def test_get_queue_status_counts_waiting_items(session_queue: SqliteSessionQueue) -> None: +def test_get_queue_status_counts_waiting_items(session_queue: SessionQueue) -> None: session = GraphExecutionState(graph=Graph()) session_json = session.model_dump_json(warnings=False) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -809,13 +809,13 @@ def test_get_queue_status_counts_waiting_items(session_queue: SqliteSessionQueue assert queue_status.total == 1 -def test_startup_cancellation_cancels_waiting_workflow_call_chain(session_queue: SqliteSessionQueue) -> None: +def test_startup_cancellation_cancels_waiting_workflow_call_chain(session_queue: SessionQueue) -> None: parent_session = GraphExecutionState(graph=Graph()) child_session = GraphExecutionState(graph=Graph()) sibling_session = GraphExecutionState(graph=Graph()) batch_id = str(uuid.uuid4()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -898,7 +898,7 @@ def test_startup_cancellation_cancels_waiting_workflow_call_chain(session_queue: session_queue._set_in_progress_to_canceled() assert session_queue.get_queue_item(parent_item_id).status == "canceled" - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql SELECT status @@ -912,12 +912,12 @@ def test_startup_cancellation_cancels_waiting_workflow_call_chain(session_queue: assert child_statuses == ["canceled", "canceled"] -def test_cancel_queue_item_cascades_from_waiting_parent_to_child_chain(session_queue: SqliteSessionQueue) -> None: +def test_cancel_queue_item_cascades_from_waiting_parent_to_child_chain(session_queue: SessionQueue) -> None: parent_session = GraphExecutionState(graph=Graph()) child_session = GraphExecutionState(graph=Graph()) grandchild_session = GraphExecutionState(graph=Graph()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -993,7 +993,7 @@ def test_cancel_queue_item_cascades_from_waiting_parent_to_child_chain(session_q def test_cancel_child_uses_metadata_chain_walk_and_hydrates_each_item_once( - session_queue: SqliteSessionQueue, monkeypatch: pytest.MonkeyPatch + session_queue: SessionQueue, monkeypatch: pytest.MonkeyPatch ) -> None: parent_item_id, parent_session, child_sessions = _build_waiting_workflow_call_parent(session_queue, child_count=1) child_item_id = _insert_queue_item( @@ -1024,11 +1024,11 @@ def count_hydration(raw_queue_item: dict, *, quarantine: bool): assert hydrate_calls == 2 -def test_cancel_queue_item_cascades_from_child_to_waiting_parents(session_queue: SqliteSessionQueue) -> None: +def test_cancel_queue_item_cascades_from_child_to_waiting_parents(session_queue: SessionQueue) -> None: parent_session = GraphExecutionState(graph=Graph()) child_session = GraphExecutionState(graph=Graph()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -1078,12 +1078,12 @@ def test_cancel_queue_item_cascades_from_child_to_waiting_parents(session_queue: assert session_queue.get_queue_item(parent_item_id).status == "canceled" -def test_delete_queue_item_removes_entire_workflow_call_chain(session_queue: SqliteSessionQueue) -> None: +def test_delete_queue_item_removes_entire_workflow_call_chain(session_queue: SessionQueue) -> None: parent_session = GraphExecutionState(graph=Graph()) child_session = GraphExecutionState(graph=Graph()) grandchild_session = GraphExecutionState(graph=Graph()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -1162,12 +1162,12 @@ def test_delete_queue_item_removes_entire_workflow_call_chain(session_queue: Sql def test_delete_queue_item_cancels_active_workflow_call_chain_before_deleting( - session_queue: SqliteSessionQueue, event_bus: TestEventService + session_queue: SessionQueue, event_bus: TestEventService ) -> None: parent_session = GraphExecutionState(graph=Graph()) child_session = GraphExecutionState(graph=Graph()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -1223,12 +1223,12 @@ def test_delete_queue_item_cancels_active_workflow_call_chain_before_deleting( def test_cancel_queue_item_cascade_emits_canceled_events_for_waiting_parent_and_running_child( - session_queue: SqliteSessionQueue, event_bus: TestEventService + session_queue: SessionQueue, event_bus: TestEventService ) -> None: parent_session = GraphExecutionState(graph=Graph()) child_session = GraphExecutionState(graph=Graph()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -1283,7 +1283,7 @@ def test_cancel_queue_item_cascade_emits_canceled_events_for_waiting_parent_and_ assert canceled_events[-1].queue_status.in_progress == 0 -def test_cancel_workflow_call_children_cancels_nested_descendants(session_queue: SqliteSessionQueue) -> None: +def test_cancel_workflow_call_children_cancels_nested_descendants(session_queue: SessionQueue) -> None: root_session = GraphExecutionState(graph=Graph()) waiting_child_session = GraphExecutionState(graph=Graph()) nested_child_session = GraphExecutionState(graph=Graph()) @@ -1330,7 +1330,7 @@ def test_cancel_workflow_call_children_cancels_nested_descendants(session_queue: def test_cancel_all_except_current_cancels_waiting_chains_outside_current_chain( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: current_parent_session = GraphExecutionState(graph=Graph()) current_child_session = GraphExecutionState(graph=Graph()) @@ -1370,7 +1370,7 @@ def test_cancel_all_except_current_cancels_waiting_chains_outside_current_chain( def test_delete_all_except_current_deletes_waiting_chains_outside_current_chain( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: current_parent_session = GraphExecutionState(graph=Graph()) current_child_session = GraphExecutionState(graph=Graph()) @@ -1412,7 +1412,7 @@ def test_delete_all_except_current_deletes_waiting_chains_outside_current_chain( def test_cancel_all_except_current_preserves_waiting_chain_during_pending_child_handoff( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: parent_session = GraphExecutionState(graph=Graph()) child_session = GraphExecutionState(graph=Graph()) @@ -1440,7 +1440,7 @@ def test_cancel_all_except_current_preserves_waiting_chain_during_pending_child_ def test_delete_all_except_current_preserves_waiting_chain_during_pending_child_handoff( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: parent_session = GraphExecutionState(graph=Graph()) child_session = GraphExecutionState(graph=Graph()) @@ -1468,7 +1468,7 @@ def test_delete_all_except_current_preserves_waiting_chain_during_pending_child_ session_queue.get_queue_item(unrelated_item_id) -def test_cancel_by_queue_id_cancels_current_workflow_call_descendants(session_queue: SqliteSessionQueue) -> None: +def test_cancel_by_queue_id_cancels_current_workflow_call_descendants(session_queue: SessionQueue) -> None: parent_session = GraphExecutionState(graph=Graph()) child_session = GraphExecutionState(graph=Graph()) nested_child_session = GraphExecutionState(graph=Graph()) @@ -1503,7 +1503,7 @@ def test_cancel_by_queue_id_cancels_current_workflow_call_descendants(session_qu def test_cancel_by_batch_ids_cancels_waiting_parent_and_in_progress_child( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: batch_id = str(uuid.uuid4()) parent_session = GraphExecutionState(graph=Graph()) @@ -1534,7 +1534,7 @@ def test_cancel_by_batch_ids_cancels_waiting_parent_and_in_progress_child( def test_cancel_by_destination_cancels_waiting_parent_and_in_progress_child( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: destination = "gallery" parent_session = GraphExecutionState(graph=Graph()) @@ -1564,11 +1564,11 @@ def test_cancel_by_destination_cancels_waiting_parent_and_in_progress_child( assert session_queue.get_queue_item(child_item_id).status == "canceled" -def test_retry_items_by_id_retries_root_once_for_child_chain_item(session_queue: SqliteSessionQueue) -> None: +def test_retry_items_by_id_retries_root_once_for_child_chain_item(session_queue: SessionQueue) -> None: root_session = GraphExecutionState(graph=Graph()) child_session = GraphExecutionState(graph=Graph()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -1626,13 +1626,13 @@ def test_retry_items_by_id_retries_root_once_for_child_chain_item(session_queue: def test_retry_items_by_id_emits_root_only_retry_event_for_nested_failure_chain( - session_queue: SqliteSessionQueue, event_bus: TestEventService + session_queue: SessionQueue, event_bus: TestEventService ) -> None: root_session = GraphExecutionState(graph=Graph()) child_session = GraphExecutionState(graph=Graph()) grandchild_session = GraphExecutionState(graph=Graph()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -1709,11 +1709,11 @@ def test_retry_items_by_id_emits_root_only_retry_event_for_nested_failure_chain( assert retry_events[0].retried_item_ids == [root_item_id] -def test_retry_items_by_id_respects_remaining_queue_capacity(session_queue: SqliteSessionQueue) -> None: +def test_retry_items_by_id_respects_remaining_queue_capacity(session_queue: SessionQueue) -> None: root_session = GraphExecutionState(graph=Graph()) pending_session = GraphExecutionState(graph=Graph()) - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -1750,7 +1750,7 @@ def test_retry_items_by_id_respects_remaining_queue_capacity(session_queue: Sqli ) root_item_id = cursor.lastrowid - session_queue._SqliteSessionQueue__invoker.services.configuration.max_queue_size = 1 + session_queue._SessionQueue__invoker.services.configuration.max_queue_size = 1 retry_result = session_queue.retry_items_by_id("default", [root_item_id]) assert retry_result.retried_item_ids == [] diff --git a/tests/app/services/session_queue/test_session_queue_workflow_call_retry.py b/tests/app/services/session_queue/test_session_queue_workflow_call_retry.py index 4fb0082fb7f..120287f5715 100644 --- a/tests/app/services/session_queue/test_session_queue_workflow_call_retry.py +++ b/tests/app/services/session_queue/test_session_queue_workflow_call_retry.py @@ -6,15 +6,17 @@ from invokeai.app.services.events.events_common import QueueItemsRetriedEvent from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from tests.fixtures.sqlite_database import sqlite_cursor_of from tests.test_nodes import TestEventService @pytest.fixture -def session_queue(mock_invoker: Invoker) -> SqliteSessionQueue: - db = mock_invoker.services.board_records._db - queue = SqliteSessionQueue(db=db) +def session_queue(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: + db = mock_sqlite_database + queue = SessionQueue(db) queue.start(mock_invoker) return queue @@ -26,7 +28,7 @@ def event_bus(mock_invoker: Invoker) -> TestEventService: def _insert_queue_item( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, *, session: GraphExecutionState, status: str, @@ -35,7 +37,7 @@ def _insert_queue_item( project_id: str | None = None, queue_id: str = "default", ) -> int: - with session_queue._db.transaction() as cursor: + with sqlite_cursor_of(session_queue) as cursor: cursor.execute( """--sql INSERT INTO session_queue ( @@ -85,7 +87,7 @@ def _insert_queue_item( def test_retry_items_by_id_retries_root_once_for_child_chain_item( - session_queue: SqliteSessionQueue, event_bus: TestEventService + session_queue: SessionQueue, event_bus: TestEventService ) -> None: root_session = GraphExecutionState(graph=Graph()) child_session = GraphExecutionState(graph=Graph()) @@ -119,7 +121,7 @@ def test_retry_items_by_id_retries_root_once_for_child_chain_item( def test_retry_items_by_id_emits_unique_owner_ids_for_multiple_roots( - session_queue: SqliteSessionQueue, event_bus: TestEventService + session_queue: SessionQueue, event_bus: TestEventService ) -> None: first_root_item_id = _insert_queue_item( session_queue, session=GraphExecutionState(graph=Graph()), user_id="user-1", status="failed" @@ -142,7 +144,7 @@ def test_retry_items_by_id_emits_unique_owner_ids_for_multiple_roots( def test_retried_items_inherit_the_project_of_the_root( - session_queue: SqliteSessionQueue, + session_queue: SessionQueue, ) -> None: root_item_id = _insert_queue_item( session_queue, diff --git a/tests/app/services/session_queue/test_session_queue_writes.py b/tests/app/services/session_queue/test_session_queue_writes.py new file mode 100644 index 00000000000..dadb0cd5e5a --- /dev/null +++ b/tests/app/services/session_queue/test_session_queue_writes.py @@ -0,0 +1,734 @@ +"""The session queue's writes on every backend: enqueueing, claiming, status changes, bulk cancels and deletes, sessions +and history, including what each does while another transaction is in flight. + +On MySQL and MariaDB transactions run side by side, so every check-then-write of the queue is one conditional +statement or runs under a lock; the race tests hold one transaction open and show the other waits for it or loses +cleanly. On SQLite every transaction excludes the others, so those tests pass there trivially. +""" + +import asyncio +import json +import logging +from types import SimpleNamespace +from typing import Any, Optional, cast + +import pytest +from pydantic_core import to_jsonable_python +from sqlalchemy import insert, select, update + +from invokeai.app.invocations.call_saved_workflow import CallSavedWorkflowInvocation +from invokeai.app.services.config.config_default import InvokeAIAppConfig +from invokeai.app.services.events.events_common import ( + BatchEnqueuedEvent, + QueueItemsCanceledEvent, + QueueItemStatusChangedEvent, +) +from invokeai.app.services.invoker import Invoker +from invokeai.app.services.session_queue.session_queue_common import Batch, EnqueueProjectNotFoundError +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries import session_queue as queue_queries +from invokeai.app.services.shared.database.queries.locks import DatabaseLock +from invokeai.app.services.shared.database.queries.session_queue import SessionQueueQueries +from invokeai.app.services.shared.database.schema.session_queue import session_queue +from invokeai.app.services.shared.database.schema.users import users +from invokeai.app.services.shared.execution_state_migration import dump_execution_state +from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from tests.fixtures.races import waiting, when_called, while_in_flight +from tests.test_nodes import PromptTestInvocation, TestEventService + +QUEUE = "default" + + +@pytest.fixture +def accounts(database: Database) -> None: + with database.begin(write=True) as conn: + conn.execute( + insert(users), + [ + {"user_id": user, "email": f"{user}@example.com", "display_name": user.title(), "password_hash": "-"} + for user in ("alice", "bob", "carol") + ], + ) + + +@pytest.fixture +def events() -> TestEventService: + return TestEventService() + + +def _service(database: Database, events: TestEventService, **config: Any) -> SessionQueue: + configuration = InvokeAIAppConfig(use_memory_db=True, **config) + services = SimpleNamespace( + configuration=configuration, events=events, logger=logging.getLogger("queue"), model_manager=None + ) + queue = SessionQueue(database) + queue.start(cast(Invoker, SimpleNamespace(services=services))) + return queue + + +@pytest.fixture +def queue(database: Database, events: TestEventService, accounts: None) -> SessionQueue: + return _service(database, events) + + +def _batch(prompt: str = "test", runs: int = 1, *, key: Optional[str] = None, origin: Optional[str] = None) -> Batch: + graph = Graph() + graph.add_node(PromptTestInvocation(id="prompt", prompt=prompt)) + return Batch(graph=graph, runs=runs, idempotency_key=key, origin=origin) + + +def _item( + database: Database, + *, + status: str = "pending", + user_id: str = "alice", + priority: int = 0, + origin: Optional[str] = None, + destination: Optional[str] = None, + parent_item_id: Optional[int] = None, + root_item_id: Optional[int] = None, + prompt: str = "", + created_at: str = "2026-01-01 00:00:00.000", + started_at: Optional[str] = None, +) -> int: + graph = Graph() + graph.add_node(PromptTestInvocation(id="prompt", prompt=prompt)) + session = GraphExecutionState(graph=graph) + values = { + "queue_id": QUEUE, + "batch_id": "batch", + "session_id": session.id, + "session": json.dumps(dump_execution_state(session), default=to_jsonable_python), + "status": status, + "priority": priority, + "user_id": user_id, + "origin": origin, + "destination": destination, + "parent_item_id": parent_item_id, + "root_item_id": root_item_id, + "created_at": created_at, + "started_at": started_at, + } + with database.begin(write=True) as conn: + return int(conn.execute(insert(session_queue).values(**values)).inserted_primary_key[0]) + + +def _row(database: Database, item_id: int) -> Optional[dict[str, Any]]: + with database.begin(write=False) as conn: + row = conn.execute(select(session_queue).where(session_queue.c.item_id == item_id)).first() + return row._asdict() if row is not None else None + + +def _status_events(events: TestEventService) -> list[tuple[int, str]]: + return [(event.item_id, event.status) for event in events.events if isinstance(event, QueueItemStatusChangedEvent)] + + +class TestEnqueue: + def test_an_enqueue_stamps_its_items_and_a_retry_with_its_key_returns_them( + self, database: Database, queue: SessionQueue, events: TestEventService + ) -> None: + batch = _batch(runs=2, key="k") + first = asyncio.run(queue.enqueue_batch(QUEUE, batch, False, "alice")) + again = asyncio.run(queue.enqueue_batch(QUEUE, batch, False, "alice")) + + assert (first.enqueued, len(first.item_ids), again.item_ids) == (2, 2, first.item_ids) + row = _row(database, first.item_ids[0]) + assert row is not None and row["created_at"] == row["updated_at"] and row["status"] == "pending" + assert queue.get_enqueue_receipt(QUEUE, "k", "alice") is not None + assert [type(event) for event in events.events].count(BatchEnqueuedEvent) == 1 + + def test_an_enqueue_takes_only_what_fits_and_prepends_above_the_pending( + self, database: Database, events: TestEventService, accounts: None + ) -> None: + queue = _service(database, events, max_queue_size=3) + asyncio.run(queue.enqueue_batch(QUEUE, _batch(runs=2), False, "alice")) + prepended = asyncio.run(queue.enqueue_batch(QUEUE, _batch(runs=2), True, "bob")) + + assert (prepended.enqueued, prepended.priority) == (1, 1) + assert queue.get_queue_status(QUEUE).pending == 3 + + def test_an_enqueue_for_a_project_of_another_account_is_refused(self, queue: SessionQueue) -> None: + batch = _batch().model_copy(update={"project_id": "not-alices"}) + + with pytest.raises(EnqueueProjectNotFoundError): + asyncio.run(queue.enqueue_batch(QUEUE, batch, False, "alice")) + + def test_an_enqueue_waits_for_one_admitting(self, database: Database, queue: SessionQueue) -> None: + # Two enqueues counting the pending items side by side could both fit into the last free place. + errors = while_in_flight( + database, + lambda q: q.locks.acquire(DatabaseLock.SESSION_QUEUE_ADMISSION), + lambda: asyncio.run(queue.enqueue_batch(QUEUE, _batch(), False, "alice")), + ) + + assert errors == [] + assert queue.get_queue_status(QUEUE).pending == 1 + + def test_an_enqueue_and_a_retry_share_the_media_protection(self, database: Database, queue: SessionQueue) -> None: + # The intermediates cleanup holds it exclusively between its check and its delete. + errors = while_in_flight( + database, + lambda q: q.locks.acquire(DatabaseLock.MEDIA_PROTECTION), + lambda: asyncio.run(queue.enqueue_batch(QUEUE, _batch(), False, "alice")), + ) + + assert errors == [] + + +class TestDequeue: + def test_a_claim_records_its_start_and_device_and_signals_once( + self, database: Database, queue: SessionQueue, events: TestEventService + ) -> None: + item_id = _item(database) + + claimed = queue.dequeue(device="cuda:1") + + assert claimed is not None and (claimed.item_id, claimed.status, claimed.device) == ( + item_id, + "in_progress", + "cuda:1", + ) + row = _row(database, item_id) + assert row is not None and row["started_at"] is not None and row["status_sequence"] == 1 + assert _status_events(events) == [(item_id, "in_progress")] + assert queue.dequeue() is None + + def test_two_workers_never_claim_one_item( + self, database: Database, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch + ) -> None: + # The second worker claims the item the first one chose, before the first claims it: the first then + # loses its claim and takes the next item. + first_item, second_item = _item(database), _item(database) + real = SessionQueue._apply_device_affinity + claimed_by_other: list[Any] = [] + + def other_worker_first(self: SessionQueue, candidate: Any, resident_keys: set[str]) -> Any: + if not claimed_by_other: + claimed_by_other.append(None) + claimed_by_other[0] = queue.dequeue(device="other") + return real(self, candidate, resident_keys) + + monkeypatch.setattr(SessionQueue, "_apply_device_affinity", other_worker_first) + + mine = queue.dequeue(device="mine") + + other = claimed_by_other[0] + assert other is not None and mine is not None + assert {other.item_id, mine.item_id} == {first_item, second_item} + assert (other.device, mine.device) == ("other", "mine") + + def test_round_robin_serves_the_account_served_longest_ago( + self, database: Database, events: TestEventService, accounts: None + ) -> None: + queue = _service(database, events, multiuser=True, session_queue_mode="round_robin") + _item(database, user_id="alice", status="completed", started_at="2026-01-02 00:00:00.000") + _item(database, user_id="bob", status="completed", started_at="2026-01-01 00:00:00.000") + _item(database, user_id="alice", priority=5) + _item(database, user_id="bob", priority=0) + bobs_best = _item(database, user_id="bob", priority=1) + + served = queue.dequeue() + + assert served is not None and served.item_id == bobs_best + + def test_a_worker_prefers_an_item_whose_models_it_has_loaded( + self, database: Database, events: TestEventService, accounts: None + ) -> None: + queue = _service(database, events, multiuser=True, session_queue_mode="round_robin") + cache = SimpleNamespace(cached_model_keys=lambda: ["model-key-b"]) + queue._SessionQueue__invoker.services.model_manager = SimpleNamespace( # type: ignore[attr-defined] + load=SimpleNamespace(ram_caches={"cuda:0": cache}) + ) + _item(database, prompt="model-key-a") + warm = _item(database, prompt="model-key-b") + + served = queue.dequeue(device="cuda:0") + + assert served is not None and served.item_id == warm + + +class TestStatusChanges: + def test_a_finish_stamps_its_time_and_a_finished_item_stays_finished( + self, database: Database, queue: SessionQueue, events: TestEventService + ) -> None: + item_id = _item(database, status="in_progress") + + completed = queue.complete_queue_item(item_id) + canceled = queue.cancel_queue_item(item_id) + + assert (completed.status, canceled.status) == ("completed", "completed") + row = _row(database, item_id) + assert row is not None and row["completed_at"] is not None and row["status_sequence"] == 1 + assert _status_events(events) == [(item_id, "completed")] + + def test_a_status_change_racing_another_loses_cleanly(self, database: Database, queue: SessionQueue) -> None: + # The other transaction completes the item but has not committed: the cancel waits, then finds it finished. + item_id = _item(database, status="in_progress") + outcomes: list[tuple[Any, bool]] = [] + + def cancel() -> None: + outcomes.append(queue._transition_queue_item_status(item_id, "canceled")) + + errors = while_in_flight(database, lambda q: q.session_queue.transition(item_id, "completed"), cancel) + + assert errors == [] + ((item, changed),) = outcomes + assert (item.status, changed) == ("completed", False) + + def test_a_move_to_in_progress_stamps_its_start(self, database: Database) -> None: + item_id = _item(database, status="waiting") + + exists, changed = database.queries.session_queue.transition(item_id, "in_progress", device="cpu") + + assert exists and changed is not None + assert (changed["status"], changed["device"], changed["completed_at"]) == ("in_progress", "cpu", None) + assert changed["started_at"] is not None + + def test_startup_cancels_interrupted_chains( + self, database: Database, events: TestEventService, accounts: None + ) -> None: + root = _item(database, status="waiting") + child = _item(database, status="in_progress", parent_item_id=root, root_item_id=root) + sibling = _item(database, status="pending", parent_item_id=root, root_item_id=root) + unrelated = _item(database, status="pending") + + _service(database, events) + + statuses = {item_id: (_row(database, item_id) or {})["status"] for item_id in (root, child, sibling, unrelated)} + assert statuses == {root: "canceled", child: "canceled", sibling: "canceled", unrelated: "pending"} + + +class TestBulk: + def test_cancel_by_origin_takes_its_wildcards_as_written_and_signals_running_items_one_by_one( + self, database: Database, queue: SessionQueue, events: TestEventService + ) -> None: + waiting = _item(database, origin="canvas_1", user_id="alice") + bobs = _item(database, origin="Canvas_2", user_id="bob") + running = _item(database, origin="canvas_3", status="in_progress") + _item(database, origin="canvasX", user_id="alice") + _item(database, origin="canvas_4", status="completed") + + result = queue.cancel_by_queue_id(QUEUE, origin_prefix="canvas_") + + assert result.canceled == 3 + bulk = [event for event in events.events if isinstance(event, QueueItemsCanceledEvent)] + assert [event.canceled_item_ids_by_user for event in bulk] == [{"alice": [waiting], "bob": [bobs]}] + assert _status_events(events) == [(running, "canceled")] + + def test_cancel_by_batch_cancels_only_the_batchs_items( + self, database: Database, queue: SessionQueue, events: TestEventService + ) -> None: + first = asyncio.run(queue.enqueue_batch(QUEUE, _batch(runs=2), False, "alice")) + second = asyncio.run(queue.enqueue_batch(QUEUE, _batch(runs=2), False, "alice")) + + result = queue.cancel_by_batch_ids(QUEUE, [first.batch.batch_id]) + + assert result.canceled == 2 + assert [queue.get_queue_item(item_id).status for item_id in first.item_ids] == ["canceled", "canceled"] + assert [queue.get_queue_item(item_id).status for item_id in second.item_ids] == ["pending", "pending"] + + def test_cancel_all_except_current_spares_the_running_chain(self, database: Database, queue: SessionQueue) -> None: + root = _item(database, status="waiting") + _item(database, status="in_progress", parent_item_id=root, root_item_id=root) + next_child = _item(database, status="pending", parent_item_id=root, root_item_id=root) + other = _item(database, status="pending") + + result = queue.cancel_all_except_current(QUEUE) + + assert result.canceled == 1 + assert (_row(database, other) or {})["status"] == "canceled" + assert (_row(database, next_child) or {})["status"] == "pending" + + def test_delete_by_destination_deletes_running_items_after_canceling_them( + self, database: Database, queue: SessionQueue, events: TestEventService + ) -> None: + running = _item(database, destination="canvas", status="in_progress") + waiting = _item(database, destination="canvas") + kept = _item(database, destination="gallery") + + result = queue.delete_by_destination(QUEUE, "canvas") + + assert result.deleted == 2 + assert (_row(database, running), _row(database, waiting)) == (None, None) + assert _row(database, kept) is not None + bulk = [event for event in events.events if isinstance(event, QueueItemsCanceledEvent)] + assert [event.canceled_item_ids for event in bulk] == [[waiting]] + + def test_clear_cancels_the_accounts_running_items_and_deletes_only_its_own( + self, database: Database, queue: SessionQueue, events: TestEventService + ) -> None: + alices = [_item(database, status="in_progress"), _item(database)] + bobs = _item(database, user_id="bob", status="in_progress") + + result = queue.clear(QUEUE, user_id="alice") + + assert result.deleted == 2 + assert [_row(database, item_id) for item_id in alices] == [None, None] + assert (_row(database, bobs) or {})["status"] == "in_progress" + assert _status_events(events) == [(alices[0], "canceled")] + + def test_prune_keeps_history_an_active_workflow_call_still_needs( + self, database: Database, queue: SessionQueue + ) -> None: + root = _item(database, status="waiting") + needed = _item(database, status="completed", parent_item_id=root, root_item_id=root) + old = [_item(database, status=status) for status in ("completed", "failed", "canceled")] + pending = _item(database) + + assert queue.prune(QUEUE).deleted == 3 + assert [_row(database, item_id) for item_id in old] == [None, None, None] + assert all(_row(database, item_id) is not None for item_id in (root, needed, pending)) + + def test_history_is_pruned_to_its_latest_items( + self, database: Database, events: TestEventService, accounts: None + ) -> None: + finished = [_item(database, status="completed") for _ in range(4)] + with database.begin(write=True) as conn: + for day, item_id in enumerate(finished): + conn.execute( + update(session_queue) + .where(session_queue.c.item_id == item_id) + .values(completed_at=f"2026-01-0{day + 1} 00:00:00.000") + ) + + _service(database, events, max_queue_history=2) + + assert [_row(database, item_id) is not None for item_id in finished] == [False, False, True, True] + + +class TestSessions: + def test_every_session_write_moves_the_revision(self, database: Database, queue: SessionQueue) -> None: + item_id = _item(database, status="in_progress") + session = queue.get_queue_item(item_id).session + + queue.save_queue_item_session(item_id, session) + queue.save_queue_item_session(item_id, session) + + assert (_row(database, item_id) or {})["session_revision"] == 2 + assert queue._save_queue_item_session_if_active(item_id, session) is True + queue.complete_queue_item(item_id) + assert queue._save_queue_item_session_if_active(item_id, session) is False + + def test_a_session_write_shares_the_media_protection(self, database: Database, queue: SessionQueue) -> None: + item_id = _item(database, status="in_progress") + session = queue.get_queue_item(item_id).session + + errors = while_in_flight( + database, + lambda q: q.locks.acquire(DatabaseLock.MEDIA_PROTECTION), + lambda: queue.save_queue_item_session(item_id, session), + ) + + assert errors == [] + + def test_a_parent_waits_for_its_children_only_with_the_session_it_was_read_with(self, database: Database) -> None: + item_id = _item(database, status="in_progress") + stored = (_row(database, item_id) or {})["session"] + + stale = database.queries.session_queue.wait_for_children(item_id, session="{}", expected_session="other") + waited = database.queries.session_queue.wait_for_children(item_id, session="{}", expected_session=stored) + + row = _row(database, item_id) or {} + assert (stale, waited) == (False, True) + assert (row["status"], row["session"], row["session_revision"]) == ("waiting", "{}", 1) + + def test_a_retry_enqueues_failed_roots_once( + self, database: Database, queue: SessionQueue, events: TestEventService + ) -> None: + failed = _item(database, status="failed") + _item(database, status="pending") + + result = queue.retry_items_by_id(QUEUE, [failed, failed, 404]) + + assert result.retried_item_ids == [failed] + assert queue.get_queue_status(QUEUE).pending == 2 + + +def _session_json(session: GraphExecutionState) -> str: + return json.dumps(dump_execution_state(session), default=to_jsonable_python) + + +def _waiting_parent(database: Database, children: int = 2) -> tuple[int, list[GraphExecutionState]]: + """An item in progress whose session waits on a workflow call with `children` child sessions to enqueue.""" + graph = Graph() + graph.add_node(CallSavedWorkflowInvocation(id="call-node", workflow_id="workflow-a")) + parent_session = GraphExecutionState(graph=graph) + invocation = parent_session.next() + frame = parent_session.build_workflow_call_frame(invocation.id, invocation.workflow_id) + parent_session.begin_waiting_on_workflow_call(frame) + child_sessions = [parent_session.create_child_workflow_execution_state(Graph(), frame) for _ in range(children)] + parent_session.attach_waiting_workflow_call_child_sessions(child_sessions) + parent = _item(database, status="in_progress") + with database.begin(write=True) as conn: + conn.execute( + update(session_queue).where(session_queue.c.item_id == parent).values(session=_session_json(parent_session)) + ) + return parent, child_sessions + + +def _pending_values(parent: Optional[int] = None) -> dict[str, Any]: + session = GraphExecutionState(graph=Graph()) + return { + "queue_id": QUEUE, + "batch_id": "other", + "session_id": session.id, + "session": _session_json(session), + "status": "pending", + "user_id": "bob", + "parent_item_id": parent, + "root_item_id": parent, + } + + +class TestRaces: + """What a write does while another transaction is in flight; on SQLite the other one has to wait anyway.""" + + def test_capacity_is_recounted_after_admission_waits( + self, database: Database, events: TestEventService, accounts: None + ) -> None: + # An enqueue that counted the free places before another filled them must not overfill the queue. + queue = _service(database, events, max_queue_size=2) + + def fill(q: Queries) -> None: + q.locks.acquire(DatabaseLock.SESSION_QUEUE_ADMISSION) + for _ in range(2): + q.session_queue.insert_item(_pending_values()) + + errors = while_in_flight( + database, fill, lambda: asyncio.run(queue.enqueue_batch(QUEUE, _batch(runs=2), False, "alice")) + ) + + assert errors == [] + assert queue.get_queue_status(QUEUE).pending == 2 + + @pytest.mark.parametrize("lock", [DatabaseLock.SESSION_QUEUE_ADMISSION, DatabaseLock.MEDIA_PROTECTION]) + def test_a_retry_waits_for_the_admission_and_the_cleanup( + self, database: Database, queue: SessionQueue, lock: DatabaseLock + ) -> None: + failed = _item(database, status="failed") + + errors = while_in_flight( + database, lambda q: q.locks.acquire(lock), lambda: queue.retry_items_by_id(QUEUE, [failed]) + ) + + assert errors == [] + + @pytest.mark.parametrize("lock", [DatabaseLock.SESSION_QUEUE_ADMISSION, DatabaseLock.MEDIA_PROTECTION]) + def test_a_child_enqueue_waits_for_the_admission_and_the_cleanup( + self, database: Database, queue: SessionQueue, lock: DatabaseLock + ) -> None: + parent, child_sessions = _waiting_parent(database) + + errors = while_in_flight( + database, + lambda q: q.locks.acquire(lock), + lambda: queue.enqueue_workflow_call_children( + queue.get_queue_item(parent), [(session, None) for session in child_sessions] + ), + ) + + assert errors == [] + + def test_a_child_completion_shares_the_media_protection(self, database: Database, queue: SessionQueue) -> None: + parent, child_sessions = _waiting_parent(database) + children = queue.enqueue_workflow_call_children( + queue.get_queue_item(parent), [(session, None) for session in child_sessions] + ) + queue.complete_queue_item(children[0].item_id) + + errors = while_in_flight( + database, + lambda q: q.locks.acquire(DatabaseLock.MEDIA_PROTECTION), + lambda: queue.record_workflow_call_child_completion(parent, children[0].item_id, {"r": 1}), + ) + + assert errors == [] + + def test_two_children_completing_at_once_both_record( + self, database: Database, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch + ) -> None: + # The second records into the session the first wrote, not into the one both read. + parent, child_sessions = _waiting_parent(database) + children = queue.enqueue_workflow_call_children( + queue.get_queue_item(parent), [(session, None) for session in child_sessions] + ) + for child in children: + queue.complete_queue_item(child.item_id) + results: list[Any] = [] + ended = when_called( + monkeypatch, + SessionQueueQueries, + "set_session", + lambda: results.append(queue.record_workflow_call_child_completion(parent, children[1].item_id, {"r": 2})), + ) + + results.append(queue.record_workflow_call_child_completion(parent, children[0].item_id, {"r": 1})) + + assert ended() == [] + execution = queue.get_queue_item(parent).session.waiting_workflow_call_execution + assert execution is not None + assert set(execution.completed_child_item_ids) == {children[0].item_id, children[1].item_id} + assert sum(result is not None and result.should_resume for result in results) == 1 + + def test_a_bulk_cancel_does_not_report_an_item_claimed_meanwhile( + self, database: Database, queue: SessionQueue, events: TestEventService + ) -> None: + # The claim in flight wins: the bulk cancel must neither name the item in its event nor count it twice. + claimed = _item(database) + outcomes: list[Any] = [] + + errors = while_in_flight( + database, + lambda q: q.session_queue.claim(claimed, "gpu"), + lambda: outcomes.append(queue.cancel_by_queue_id(QUEUE)), + ) + + assert errors == [] + (result,) = outcomes + bulk = [event for event in events.events if isinstance(event, QueueItemsCanceledEvent)] + assert (result.canceled, [item for event in bulk for item in event.canceled_item_ids]) == (1, []) + assert (_row(database, claimed) or {})["status"] == "canceled" + + def test_a_claim_waits_for_a_cancel_of_everything_but_the_running_chains( + self, database: Database, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch + ) -> None: + # The cancel locks its rows before it reads which chains run: a worker claiming one of them meanwhile waits, + # and then finds it canceled, instead of starting a child whose parent the cancel takes away. + _item(database, status="in_progress", user_id="bob") + parent = _item(database, status="waiting") + child = _item(database, parent_item_id=parent, root_item_id=parent) + real_current_chain = queue_queries._current_chain + claims: list[Any] = [] + + def current_chain_then_claim(conn: Any, queue_id: str) -> set[int]: + if not claims: + claims.append(waiting(lambda: database.queries.session_queue.claim(child, "gpu"))) + return real_current_chain(conn, queue_id) + + monkeypatch.setattr(queue_queries, "_current_chain", current_chain_then_claim) + + queue.cancel_all_except_current(QUEUE) + + assert claims[0]() == [] + assert [(_row(database, item_id) or {})["status"] for item_id in (parent, child)] == ["canceled", "canceled"] + + @pytest.mark.parametrize("operation", ["cancel", "delete"]) + def test_the_children_a_running_item_enqueues_meanwhile_are_spared( + self, database: Database, queue: SessionQueue, operation: str + ) -> None: + # The running item enqueues its children and waits for them while the cancel or delete is in flight: the + # chains are read after the rows are locked, so the children and the parent count as running. + parent = _item(database, status="in_progress") + children: list[int] = [] + + def enqueue_children(q: Queries) -> None: + children.extend(q.session_queue.insert_item(_pending_values(parent)) for _ in range(2)) + q.session_queue.transition(parent, "waiting") + + change = queue.cancel_all_except_current if operation == "cancel" else queue.delete_all_except_current + + errors = while_in_flight(database, enqueue_children, lambda: change(QUEUE)) + + assert errors == [] + assert [(_row(database, item_id) or {}).get("status") for item_id in children] == ["pending", "pending"] + + +class TestStamps: + def test_a_bulk_cancel_and_a_startup_cancel_stamp_completion( + self, database: Database, events: TestEventService, accounts: None + ) -> None: + queue = _service(database, events) + bulk = _item(database) + queue.cancel_by_queue_id(QUEUE) + interrupted = _item(database, status="in_progress") + + _service(database, events) + + for item_id in (bulk, interrupted): + row = _row(database, item_id) or {} + assert row["status"] == "canceled" and row["completed_at"] is not None + + def test_a_status_change_moves_updated_at(self, database: Database, queue: SessionQueue) -> None: + item_id = _item(database, status="in_progress") + with database.begin(write=True) as conn: + conn.execute( + update(session_queue) + .where(session_queue.c.item_id == item_id) + .values(updated_at="2020-01-01 00:00:00.000") + ) + + queue.complete_queue_item(item_id) + + assert (_row(database, item_id) or {})["updated_at"] > "2020-01-01 00:00:00.000" + + +class TestScopes: + def test_a_prune_of_an_account_keeps_the_others_history(self, database: Database, queue: SessionQueue) -> None: + alices = _item(database, status="completed") + bobs = _item(database, status="completed", user_id="bob") + + assert queue.prune(QUEUE, user_id="alice").deleted == 1 + assert (_row(database, alices), _row(database, bobs) is not None) == (None, True) + + def test_delete_all_except_current_keeps_history(self, database: Database, queue: SessionQueue) -> None: + done = _item(database, status="completed") + waiting_item = _item(database) + + assert queue.delete_all_except_current(QUEUE).deleted == 1 + assert (_row(database, done) is not None, _row(database, waiting_item)) == (True, None) + + def test_capacity_counts_only_pending_items( + self, database: Database, events: TestEventService, accounts: None + ) -> None: + queue = _service(database, events, max_queue_size=1) + _item(database, status="completed") + + assert asyncio.run(queue.enqueue_batch(QUEUE, _batch(), False, "alice")).enqueued == 1 + + def test_startup_keeps_finished_chain_items_and_cancels_lone_waiting_items_and_orphans( + self, database: Database, events: TestEventService, accounts: None + ) -> None: + root = _item(database, status="waiting") + done_child = _item(database, status="completed", parent_item_id=root, root_item_id=root) + next_child = _item(database, status="pending", parent_item_id=root, root_item_id=root) + orphan_parent = _item(database, status="waiting", parent_item_id=10**9) + orphan_child = _item(database, status="pending", parent_item_id=orphan_parent) + + _service(database, events) + + statuses = [ + (_row(database, item_id) or {})["status"] + for item_id in (root, done_child, next_child, orphan_parent, orphan_child) + ] + assert statuses == ["canceled", "completed", "canceled", "canceled", "canceled"] + + def test_fifo_serves_the_highest_priority_first( + self, database: Database, events: TestEventService, accounts: None + ) -> None: + queue = _service(database, events, session_queue_mode="FIFO") + _item(database, priority=0) + high = _item(database, priority=5) + + served = queue.dequeue() + + assert served is not None and served.item_id == high + + def test_a_status_event_carries_its_batch_counts( + self, database: Database, queue: SessionQueue, events: TestEventService + ) -> None: + first, _ = _item(database), _item(database) + + queue.cancel_queue_item(first) + + (event,) = [event for event in events.events if isinstance(event, QueueItemStatusChangedEvent)] + assert (event.batch_status.batch_id, event.batch_status.canceled, event.batch_status.pending) == ("batch", 1, 1) + + def test_an_enqueue_event_goes_to_its_owner( + self, database: Database, queue: SessionQueue, events: TestEventService + ) -> None: + asyncio.run(queue.enqueue_batch(QUEUE, _batch(), False, "bob")) + + (event,) = [event for event in events.events if isinstance(event, BatchEnqueuedEvent)] + assert event.user_id == "bob" diff --git a/tests/app/services/shared/database/conftest.py b/tests/app/services/shared/database/conftest.py new file mode 100644 index 00000000000..2dc71bff003 --- /dev/null +++ b/tests/app/services/shared/database/conftest.py @@ -0,0 +1,26 @@ +from collections.abc import Iterator +from pathlib import Path + +import pytest + +from invokeai.app.services.shared.database.database import Database +from invokeai.backend.util.logging import InvokeAILogger +from tests.fixtures.database_probe import ProbeQueries, create_probe_tables + + +@pytest.fixture +def probe(empty_database: Database) -> ProbeQueries: + """Queries over two probe tables, created in the empty database under test.""" + create_probe_tables(empty_database) + return ProbeQueries(empty_database) + + +@pytest.fixture +def sqlite_file_database(tmp_path: Path) -> Iterator[Database]: + """A SQLite database file with the probe tables, for tests that open it from a second connection too.""" + database = Database.open_sqlite(tmp_path / "probe.db", InvokeAILogger.get_logger("test_database")) + create_probe_tables(database) + try: + yield database + finally: + database.dispose() diff --git a/tests/app/services/shared/database/test_copy.py b/tests/app/services/shared/database/test_copy.py new file mode 100644 index 00000000000..76766e2fb7b --- /dev/null +++ b/tests/app/services/shared/database/test_copy.py @@ -0,0 +1,118 @@ +"""Copying rows between databases keeps every value, and keeps ids from being issued twice.""" + +import json +from collections.abc import Iterator + +import pytest +from sqlalchemy import delete, insert, select + +from invokeai.app.services.shared.database import copy as copy_module +from invokeai.app.services.shared.database.copy import copy_rows +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema import metadata +from invokeai.app.services.shared.database.schema.boards import board_images, boards +from invokeai.app.services.shared.database.schema.image_index import image_embeddings +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.models import models +from invokeai.app.services.shared.database.schema.session_queue import session_queue +from invokeai.app.services.shared.database.schema.users import users +from invokeai.app.services.shared.database.schema.videos import videos +from invokeai.backend.util.logging import InvokeAILogger + +MODEL_CONFIG = { + "hash": "blake3:0123", + "base": "sdxl", + "type": "main", + "path": "sdxl/model.safetensors", + "format": "checkpoint", + "name": "Ünïcödé model", + "source": "https://example.com/model", + "source_type": "url", + "file_size": 2**33, +} + + +@pytest.fixture +def source() -> Iterator[Database]: + """A SQLite database with the application schema, the copy's source as the copy tool's will be.""" + database = Database.open_sqlite(None, InvokeAILogger.get_logger("test_copy")) + with database.begin(write=True) as conn: + metadata.create_all(conn) + try: + yield database + finally: + database.dispose() + + +def _create_schema(database: Database) -> None: + with database.begin(write=True) as conn: + metadata.create_all(conn) + + +def _rows(database: Database) -> dict[str, list[tuple[object, ...]]]: + with database.begin(write=False) as conn: + return { + table.name: [tuple(row) for row in conn.execute(select(table).order_by(*table.primary_key.columns))] + for table in metadata.sorted_tables + } + + +def test_rows_keep_their_values_across_a_copy(source: Database, empty_database: Database) -> None: + image = {"image_origin": "internal", "image_category": "general", "width": 8, "height": 8} + with source.begin(write=True) as conn: + conn.execute( + insert(users).values(user_id="u1", email="zoë@example.com", password_hash="h", display_name="Zoë 🎨") + ) + document = json.dumps({"x": "y" * 100_000}) + conn.execute(insert(images).values({**image, "image_name": "a.png", "starred": True, "metadata": document})) + conn.execute( + insert(images).values({**image, "image_name": "b.png", "is_intermediate": None, "file_size_bytes": 2**40}) + ) + conn.execute(insert(boards).values(board_id="b1", board_name="Board", cover_image_name="a.png", user_id="u1")) + conn.execute(insert(board_images).values(board_id="b1", image_name="a.png")) + conn.execute( + insert(image_embeddings).values(image_name="a.png", model_id="clip", dim=2, embedding=bytes(range(256))) + ) + conn.execute( + insert(videos).values( + video_name="v.mp4", video_origin="internal", video_category="general", width=8, height=8, duration=1 / 3 + ) + ) + conn.execute(insert(models).values(id="m1", config=json.dumps(MODEL_CONFIG))) + _create_schema(empty_database) + + copied = copy_rows(source, empty_database) + + assert copied["images"] == 2 and copied["board_images"] == 1 and copied["models"] == 1 + assert _rows(empty_database) == _rows(source) + + +def test_ids_continue_after_the_highest_the_source_issued(source: Database, empty_database: Database) -> None: + item = {"batch_id": "b", "queue_id": "default", "session": "{}"} + with source.begin(write=True) as conn: + conn.execute(insert(session_queue), [{**item, "session_id": f"s{n}"} for n in (1, 2, 3)]) + # Item 3 was issued and is gone; its id must not come back. + conn.execute(delete(session_queue).where(session_queue.c.session_id == "s3")) + _create_schema(empty_database) + + copy_rows(source, empty_database, tables=[session_queue]) + with empty_database.begin(write=True) as conn: + inserted = conn.execute(insert(session_queue).values({**item, "session_id": "s4"})).inserted_primary_key + + assert inserted is not None + assert inserted.item_id == 4 + + +def test_a_table_is_copied_in_batches( + source: Database, empty_database: Database, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(copy_module, "_BATCH_ROWS", 2) + item = {"batch_id": "b", "queue_id": "default", "session": "{}"} + with source.begin(write=True) as conn: + conn.execute(insert(session_queue), [{**item, "session_id": f"s{n}"} for n in range(5)]) + _create_schema(empty_database) + + copied = copy_rows(source, empty_database, tables=[session_queue]) + + assert copied == {"session_queue": 5} + assert _rows(empty_database)["session_queue"] == _rows(source)["session_queue"] diff --git a/tests/app/services/shared/sqlite/test_sqlite_database.py b/tests/app/services/shared/database/test_database_clean.py similarity index 53% rename from tests/app/services/shared/sqlite/test_sqlite_database.py rename to tests/app/services/shared/database/test_database_clean.py index 42d792409d1..1d5731230e0 100644 --- a/tests/app/services/shared/sqlite/test_sqlite_database.py +++ b/tests/app/services/shared/database/test_database_clean.py @@ -1,4 +1,4 @@ -"""Regression test: clean() must serialize with transaction() users. +"""Regression test: clean() must serialize with transactions. Manual VACUUM may run while services and their worker threads are live; without the shared lock it can fail with "cannot VACUUM - SQL statements in progress". @@ -9,18 +9,19 @@ import pytest -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database from invokeai.backend.util.logging import InvokeAILogger +from tests.fixtures.sqlite_database import sqlite_cursor def test_clean_waits_for_in_flight_transactions(tmp_path) -> None: - db = SqliteDatabase(db_path=tmp_path / "test.db", logger=InvokeAILogger.get_logger()) + db = Database.open_sqlite(tmp_path / "test.db", InvokeAILogger.get_logger()) in_transaction = threading.Event() release = threading.Event() errors: list[Exception] = [] def hold_transaction() -> None: - with db.transaction() as cursor: + with sqlite_cursor(db) as cursor: cursor.execute("CREATE TABLE t (x INTEGER);") cursor.execute("INSERT INTO t VALUES (1);") in_transaction.set() @@ -49,18 +50,22 @@ def run_clean() -> None: assert errors == [] -def test_backup_to_captures_committed_wal_data_and_refuses_overwrite(tmp_path) -> None: +def test_a_backup_holds_committed_wal_data_and_overwrites_no_file(tmp_path) -> None: source = tmp_path / "source.db" - backup = tmp_path / "backup.db" - db = SqliteDatabase(db_path=source, logger=InvokeAILogger.get_logger()) - with db.transaction() as cursor: - cursor.execute("CREATE TABLE values_table (value TEXT NOT NULL);") - cursor.execute("INSERT INTO values_table VALUES ('committed in WAL');") - - db.backup_to(backup) - - with sqlite3.connect(backup) as connection: - assert connection.execute("SELECT value FROM values_table;").fetchall() == [("committed in WAL",)] - - with pytest.raises(FileExistsError): - db.backup_to(backup) + backup = tmp_path / "backups" / "backup.db" + db = Database.open_sqlite(source, InvokeAILogger.get_logger()) + try: + with sqlite_cursor(db) as cursor: + cursor.execute("CREATE TABLE values_table (value TEXT NOT NULL);") + cursor.execute("INSERT INTO values_table VALUES ('committed in WAL');") + + # The directory is created; the WAL's committed pages are in the copy. + db.backup(backup) + + with sqlite3.connect(backup) as connection: + assert connection.execute("SELECT value FROM values_table;").fetchall() == [("committed in WAL",)] + + with pytest.raises(FileExistsError): + db.backup(backup) + finally: + db.dispose() diff --git a/tests/app/services/shared/database/test_database_errors.py b/tests/app/services/shared/database/test_database_errors.py new file mode 100644 index 00000000000..f000154a7b9 --- /dev/null +++ b/tests/app/services/shared/database/test_database_errors.py @@ -0,0 +1,133 @@ +"""Constraint violations reach callers as backend-neutral errors, whatever the database.""" + +import pytest +from sqlalchemy.exc import DBAPIError + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import ( + CheckViolation, + DatabaseError, + ForeignKeyViolation, + NotNullViolation, + TransactionFailedError, + UniqueViolation, +) +from tests.fixtures.database_probe import ProbeQueries + + +def test_a_duplicate_unique_value_is_a_unique_violation(probe: ProbeQueries) -> None: + probe.items.add(1, "a") + + with pytest.raises(UniqueViolation): + probe.items.add(2, "a") + + +def test_a_duplicate_primary_key_is_a_unique_violation(probe: ProbeQueries) -> None: + probe.items.add(1, "a") + + with pytest.raises(UniqueViolation): + probe.items.add(1, "b") + + +def test_a_failed_check_is_a_check_violation(probe: ProbeQueries) -> None: + with pytest.raises(CheckViolation): + probe.items.add(1, "a", size=-1) + + +def test_a_missing_required_value_is_a_not_null_violation(probe: ProbeQueries) -> None: + with pytest.raises(NotNullViolation): + probe.items.add(1, None) + + +def test_foreign_keys_are_enforced_and_cascade(probe: ProbeQueries) -> None: + with pytest.raises(ForeignKeyViolation): + probe.items.link(1, item_id=99) + + probe.items.add(1, "a") + probe.items.link(1, item_id=1) + probe.items.remove(1) + + assert probe.items.link_ids() == [] + + +class DuplicateName(Exception): + pass + + +def test_a_violation_can_be_caught_inside_a_transaction_that_then_rolls_back(probe: ProbeQueries) -> None: + # The way a service turns a violation into its own error: caught where it happens, which needs the + # backend-neutral error there already, not only once the transaction has ended. + probe.items.add(1, "a") + + with pytest.raises(DuplicateName): + with probe.transaction() as q: + q.items.add(2, "b") + try: + q.items.add(3, "a") + except UniqueViolation as error: + raise DuplicateName("a") from error + + assert probe.items.names() == ["a"] + + +def test_a_transaction_takes_no_further_calls_after_a_statement_failed(probe: ProbeQueries) -> None: + probe.items.add(1, "a") + + with pytest.raises(TransactionFailedError): + with probe.transaction() as q: + with pytest.raises(UniqueViolation): + q.items.add(2, "a") + q.items.add(3, "c") + + assert probe.items.names() == ["a"] + + +def test_a_transaction_whose_failed_statement_was_swallowed_does_not_commit(probe: ProbeQueries) -> None: + # Committing what came before the failure would keep half of a unit of work -- and after a deadlock on + # MySQL, whose rollback already undid that half, only what came after it. + probe.items.add(1, "a") + + with pytest.raises(TransactionFailedError): + with probe.transaction() as q: + q.items.add(2, "b") + try: + q.items.add(3, "a") + except UniqueViolation: + pass + + assert probe.items.names() == ["a"] + + +def test_names_compare_byte_for_byte(probe: ProbeQueries) -> None: + # SQLite's default comparison: case and trailing spaces both make a different name. On a server this + # holds in a database created with the binary NO PAD collation the layer requires (as the fixture's + # is); the PAD SPACE `utf8mb4_bin` would reject the third name, a case-folding one the second. + probe.items.add(1, "Board") + probe.items.add(2, "board") + probe.items.add(3, "Board ") + + assert probe.items.names() == ["Board", "board", "Board "] + + +def test_an_error_without_a_backend_neutral_meaning_is_not_translated(empty_database: Database) -> None: + with pytest.raises(DBAPIError) as raised: + with empty_database.begin(write=False) as conn: + conn.exec_driver_sql("SELECT * FROM a_table_that_does_not_exist") + + assert not isinstance(raised.value, DatabaseError) + + +@pytest.mark.sqlite_only +def test_a_trigger_raising_an_error_is_not_a_foreign_key_violation( + probe: ProbeQueries, empty_database: Database +) -> None: + # SQLite reports both with the same error code: it enforces ON DELETE RESTRICT with a trigger program. + with empty_database.begin(write=True) as conn: + conn.exec_driver_sql( + "CREATE TRIGGER refuse_names AFTER INSERT ON probe_items BEGIN SELECT RAISE(ABORT, 'refused'); END" + ) + + with pytest.raises(DBAPIError) as raised: + probe.items.add(1, "a") + + assert not isinstance(raised.value, DatabaseError) diff --git a/tests/app/services/shared/database/test_database_startup.py b/tests/app/services/shared/database/test_database_startup.py new file mode 100644 index 00000000000..d104d2fa1a9 --- /dev/null +++ b/tests/app/services/shared/database/test_database_startup.py @@ -0,0 +1,224 @@ +"""Opening the database a config names, and what a server database is checked for before the app uses it.""" + +import logging +import time +from pathlib import Path +from typing import Optional + +import pytest +from sqlalchemy import URL, delete + +from invokeai.app.services.config.config_default import DefaultInvokeAIAppConfig, InvokeAIAppConfig +from invokeai.app.services.shared.database import database as database_module +from invokeai.app.services.shared.database import startup +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import DatabaseInUseError +from invokeai.app.services.shared.database.schema.migrator import applied_migrations +from invokeai.app.services.shared.database.session_lock import InstanceLock, SessionLock +from invokeai.app.services.shared.database.startup import ( + DatabaseSetupError, + init_database, + open_database, + open_migrated_database, + redacted_database_url, +) +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import NoImageFiles +from invokeai.backend.util.logging import InvokeAILogger +from tests.fixtures.database import external_test_db_url + +server_only = pytest.mark.skipif( + external_test_db_url() is None, reason="needs a MySQL or MariaDB server (INVOKEAI_TEST_DB_URL)" +) +LOGGER = InvokeAILogger.get_logger("test_startup") + + +def _config(root: Path, **values: object) -> InvokeAIAppConfig: + # Read from no file and no environment: the migrations clean up files under the root. + config = DefaultInvokeAIAppConfig(**values) + config._root = root + return config + + +def test_a_memory_database_and_a_url_exclude_each_other(tmp_path: Path) -> None: + config = _config(tmp_path, use_memory_db=True, db_url="mariadb+pymysql://invokeai:secret@db/invokeai") + with pytest.raises(DatabaseSetupError, match="exclude each other"): + open_database(config, LOGGER) + + +def test_a_url_is_shown_without_its_password() -> None: + assert redacted_database_url("mariadb+pymysql://invokeai:s3cret@db/invokeai") == ( + "mariadb+pymysql://invokeai:***@db/invokeai" + ) + assert redacted_database_url(None) is None + assert redacted_database_url("no url at all") == "***" + + +@pytest.mark.parametrize( + ("reported", "version"), + [("8.4.3", (8, 4)), ("10.11.6-MariaDB-1:10.11.6+maria~ubu2204", (10, 11)), ("11.4.2-MariaDB-log", (11, 4))], +) +def test_server_versions_are_read_from_what_the_server_reports(reported: str, version: tuple[int, int]) -> None: + assert startup._version_of(reported) == version + + +def test_a_tool_refuses_a_database_the_app_has_not_updated(tmp_path: Path) -> None: + config = _config(tmp_path) + database = init_database(config, LOGGER, NoImageFiles()) + with database.begin(write=True) as conn: + conn.execute( + delete(applied_migrations).where(applied_migrations.c.migration_id == "2026_10_07_drop_sqlite_triggers") + ) + database.dispose() + + with pytest.raises(DatabaseSetupError, match="start InvokeAI once"): + open_migrated_database(config, LOGGER) + + # As the app does when it starts. + init_database(config, LOGGER, NoImageFiles()).dispose() + open_migrated_database(config, LOGGER).dispose() + + +def test_a_tool_creates_no_database_where_there_is_none(tmp_path: Path) -> None: + config = _config(tmp_path) + + with pytest.raises(DatabaseSetupError, match="no database at"): + open_migrated_database(config, LOGGER) + + assert not config.db_path.exists() + + +@server_only +def test_one_process_at_a_time_holds_a_server_database(_external_test_schema: Optional[URL]) -> None: + assert _external_test_schema is not None + url = _external_test_schema.render_as_string(hide_password=False) + first = Database.open_url(url, LOGGER) + second = Database.open_url(url, LOGGER) + try: + first.hold_instance_lock() + with pytest.raises(DatabaseInUseError, match="Another InvokeAI process"): + second.hold_instance_lock() + + # The lock ends with the process that holds it, or, as here, when its database is closed. + first.dispose() + second.hold_instance_lock() + finally: + first.dispose() + second.dispose() + + +@server_only +def test_a_server_without_room_for_large_projects_or_too_old_is_refused( + empty_database: Database, monkeypatch: pytest.MonkeyPatch +) -> None: + with empty_database.begin(write=False) as conn: + packet = int(conn.exec_driver_sql("SELECT @@max_allowed_packet").scalar_one()) + + monkeypatch.setattr(startup, "MINIMUM_MAX_ALLOWED_PACKET", packet) + monkeypatch.setitem(startup.MINIMUM_SERVER_VERSIONS, empty_database.dialect_name, (8, 0)) + startup.check_server(empty_database) + + monkeypatch.setattr(startup, "MINIMUM_MAX_ALLOWED_PACKET", packet + 1) + with pytest.raises(DatabaseSetupError, match="max_allowed_packet"): + startup.check_server(empty_database) + + monkeypatch.setattr(startup, "MINIMUM_MAX_ALLOWED_PACKET", packet) + monkeypatch.setitem(startup.MINIMUM_SERVER_VERSIONS, empty_database.dialect_name, (99, 0)) + with pytest.raises(DatabaseSetupError, match="too old"): + startup.check_server(empty_database) + + +@server_only +def test_a_new_server_database_beside_a_sqlite_one_is_pointed_out( + empty_database: Database, tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + config = _config(tmp_path, db_url="mariadb+pymysql://invokeai:secret@db/invokeai") + config.db_path.parent.mkdir(parents=True) + config.db_path.write_bytes(b"") + + with caplog.at_level(logging.WARNING): + startup._warn_when_new_beside_sqlite(empty_database, config, LOGGER) + + [warning] = [record.getMessage() for record in caplog.records if record.levelno == logging.WARNING] + assert "invoke-db-copy" in warning + assert "secret" not in warning + + +def test_a_url_that_cannot_be_read_is_refused_without_repeating_its_password() -> None: + # An `@` in the password that is not percent-encoded. + url = "mariadb+pymysql://invokeai:se@cret@db.example/invokeai" + + assert redacted_database_url(url) == "***" + with pytest.raises(ValueError, match="Percent-encode") as refused: + Database.open_url(url, LOGGER) + assert "cret" not in str(refused.value) + + +def test_a_missing_driver_is_named_with_the_extra_that_installs_it(monkeypatch: pytest.MonkeyPatch) -> None: + def no_driver(_url: object) -> None: + raise ModuleNotFoundError("No module named 'pymysql'", name="pymysql") + + monkeypatch.setattr(database_module, "create_mysql_engine", no_driver) + + with pytest.raises(ValueError, match=r"invokeai\[mysql\]"): + Database.open_url("mariadb+pymysql://invokeai:secret@db.example/invokeai", LOGGER) + + +@server_only +def test_the_app_refuses_a_server_database_another_process_holds( + _external_test_schema: Optional[URL], tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + assert _external_test_schema is not None + # The test servers keep their own defaults; what is checked here is the lock. + monkeypatch.setattr(startup, "MINIMUM_MAX_ALLOWED_PACKET", 1) + url = _external_test_schema.render_as_string(hide_password=False) + holder = Database.open_url(url, LOGGER) + try: + holder.hold_instance_lock() + with pytest.raises(DatabaseInUseError, match="IS_USED_LOCK"): + init_database(_config(tmp_path, db_url=url), LOGGER, NoImageFiles()) + finally: + holder.dispose() + + +@server_only +def test_a_lock_the_server_released_is_taken_again( + _external_test_schema: Optional[URL], caplog: pytest.LogCaptureFixture +) -> None: + assert _external_test_schema is not None + url = _external_test_schema.render_as_string(hide_password=False) + database, other = Database.open_url(url, LOGGER), Database.open_url(url, LOGGER) + lock = InstanceLock(database.engine, LOGGER) + try: + # The server ends the session that holds the lock, as after an idle timeout or a network failure. + probe = SessionLock(other.engine, "instance") + probe.close() + with other.begin(write=False) as conn: + holder = conn.exec_driver_sql("SELECT IS_USED_LOCK(%s)", (probe.name,)).scalar_one() + conn.exec_driver_sql(f"KILL {int(holder)}") + + # Another process takes it meanwhile: two serve the database, which is reported. + rival = _take_when_free(other) + with caplog.at_level(logging.ERROR): + lock.check() + assert any("two processes now serve it" in record.getMessage() for record in caplog.records) + + rival.close() + lock.check() + with pytest.raises(DatabaseInUseError): + InstanceLock(other.engine, LOGGER) + finally: + lock.close() + database.dispose() + other.dispose() + + +def _take_when_free(database: Database) -> InstanceLock: + # A killed session releases its locks as the server gets to it. + deadline = time.monotonic() + 5 + while True: + try: + return InstanceLock(database.engine, LOGGER) + except DatabaseInUseError: + if time.monotonic() > deadline: + raise + time.sleep(0.05) diff --git a/tests/app/services/shared/database/test_database_transactions.py b/tests/app/services/shared/database/test_database_transactions.py new file mode 100644 index 00000000000..3d7458542e8 --- /dev/null +++ b/tests/app/services/shared/database/test_database_transactions.py @@ -0,0 +1,605 @@ +"""Transactions of the database layer: what commits together, what may nest, and how each backend isolates.""" + +import logging +import sqlite3 +import threading +from collections.abc import Callable, Iterator +from contextlib import closing, contextmanager +from pathlib import Path +from typing import Any, Optional + +import pytest +from sqlalchemy import URL, Connection, Row, select +from sqlalchemy.exc import DBAPIError, InvalidRequestError + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import ( + ConflictError, + LockTimeoutError, + NestedTransactionError, + ReadOnlyTransactionError, +) +from invokeai.app.services.shared.database.queries.base import OwnTransaction, QueryModule, mapped, read +from invokeai.backend.util.logging import InvokeAILogger +from tests.fixtures.database import external_test_db_url +from tests.fixtures.database_probe import ProbeQueries, probe_items + +server_only = pytest.mark.skipif( + external_test_db_url() is None, reason="needs a MySQL or MariaDB server (INVOKEAI_TEST_DB_URL)" +) + + +@contextmanager +def _retries(database: Database, on_retry: Optional[Callable[[], None]] = None) -> Iterator[list[str]]: + """Collects the database's retry warnings, calling `on_retry` (if given) as each is logged.""" + retries: list[str] = [] + + class Collect(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + if "Retrying a database transaction" in record.getMessage(): + retries.append(record.getMessage()) + if on_retry is not None: + on_retry() + + handler = Collect() + database.logger.addHandler(handler) + try: + yield retries + finally: + database.logger.removeHandler(handler) + + +def test_a_single_call_commits_on_its_own(probe: ProbeQueries) -> None: + probe.items.add(1, "a") + + assert probe.items.names() == ["a"] + + +def test_calls_in_one_transaction_commit_together(probe: ProbeQueries) -> None: + with probe.transaction() as q: + q.items.add(1, "a") + q.items.add(2, "b") + + assert probe.items.names() == ["a", "b"] + + +def test_a_failing_transaction_leaves_no_trace(probe: ProbeQueries) -> None: + with pytest.raises(RuntimeError): + with probe.transaction() as q: + q.items.add(1, "a") + raise RuntimeError("the unit of work fails after its first write") + + assert probe.items.names() == [] + + +class TestNesting: + def test_a_call_on_the_unbound_queries_inside_a_transaction_is_refused(self, probe: ProbeQueries) -> None: + with probe.transaction() as q: + q.items.add(1, "a") + with pytest.raises(NestedTransactionError): + probe.items.add(2, "b") + # Refusing one call does not release the transaction's claim on the thread. + with pytest.raises(NestedTransactionError): + probe.items.add(3, "c") + + # The refused calls never reached the database, so the transaction around them still commits. + assert probe.items.names() == ["a"] + + def test_a_transaction_inside_a_transaction_is_refused(self, probe: ProbeQueries) -> None: + with probe.transaction() as q: + with pytest.raises(NestedTransactionError): + with probe.transaction(): + pass + with pytest.raises(NestedTransactionError): + with q.transaction(): + pass + + def test_a_transaction_on_another_thread_is_not_nesting(self, probe: ProbeQueries) -> None: + outer_wrote = threading.Event() + errors: list[BaseException] = [] + + def write_from_another_thread() -> None: + try: + outer_wrote.wait(10) + probe.items.add(2, "b") + except BaseException as error: + errors.append(error) + + other = threading.Thread(target=write_from_another_thread) + other.start() + with probe.transaction() as q: + q.items.add(1, "a") + outer_wrote.set() + # On SQLite the other thread waits for this transaction; on a server it commits meanwhile. + other.join(0.3) + other.join(10) + + assert errors == [] + assert probe.items.names() == ["a", "b"] + + +class TestReadOnlyTransactions: + def test_a_read_only_transaction_refuses_writes(self, probe: ProbeQueries) -> None: + with probe.transaction(read_only=True) as q: + with pytest.raises(ReadOnlyTransactionError): + q.items.add(1, "a") + + assert probe.items.names() == [] + + def test_a_read_only_transaction_reads_one_snapshot(self, probe: ProbeQueries) -> None: + probe.items.add(1, "a") + first_read_done = threading.Event() + writer_done = threading.Event() + + def write_concurrently() -> None: + first_read_done.wait(10) + probe.items.add(2, "b") + writer_done.set() + + writer = threading.Thread(target=write_concurrently) + writer.start() + with probe.transaction(read_only=True) as q: + before = q.items.count() + first_read_done.set() + # A server commits the write meanwhile; SQLite makes the writer wait for this transaction. + writer_done.wait(0.3) + after = q.items.count() + writer.join(10) + + assert before == after == 1 + assert probe.items.count() == 2 + + +def test_a_mapper_runs_once_the_transaction_of_its_call_has_ended( + probe: ProbeQueries, empty_database: Database +) -> None: + # On SQLite the transaction holds the process-wide lock, which building DTOs need not hold. + probe.items.add(1, "a") + in_a_transaction: list[bool] = [] + + def note_whether_in_a_transaction(rows: list[Row[Any]]) -> list[str]: + try: + with empty_database.begin(write=False): + in_a_transaction.append(False) + except NestedTransactionError: + in_a_transaction.append(True) + return [name for (name,) in rows] + + class Names(QueryModule): + @mapped(note_whether_in_a_transaction) + @read + def names(self, conn: Connection) -> list[Row[Any]]: + return list(conn.execute(select(probe_items.c.name)).all()) + + assert Names(OwnTransaction(empty_database)).names() == ["a"] + assert in_a_transaction == [False] + + +def test_queries_of_an_ended_transaction_are_refused(probe: ProbeQueries) -> None: + with probe.transaction() as q: + pass + + with pytest.raises(RuntimeError, match="already ended"): + q.items.names() + + +@server_only +def test_a_server_database_is_left_to_its_operator_to_back_up(empty_database: Database, tmp_path: Path) -> None: + with pytest.raises(RuntimeError, match="no SQLite connection"): + empty_database.backup(tmp_path / "backup.db") + assert not (tmp_path / "backup.db").exists() + + +def test_a_disposed_database_refuses_further_work(empty_database: Database) -> None: + empty_database.dispose() + empty_database.dispose() # Repeating it is harmless. + + with pytest.raises(RuntimeError, match="closed"): + with empty_database.queries.transaction(): + pass + + +@pytest.mark.parametrize( + "url", + [ + "postgresql://user:secret@localhost/invokeai", + "sqlite:///invokeai.db", + # The right server through another driver, whose errors `translate_error` would not understand. + "mysql+mysqldb://user:secret@localhost/invokeai", + ], +) +def test_only_mysql_and_mariadb_through_pymysql_can_be_opened_by_url(url: str) -> None: + with pytest.raises(ValueError, match="Unsupported database URL scheme"): + Database.open_url(url, InvokeAILogger.get_logger("test_database")) + + +@pytest.mark.sqlite_only +class TestSqlite: + """The SQLite file is shared with other processes (the user-management commands, database tools).""" + + @pytest.fixture + def other_process(self, sqlite_file_database: Database) -> Iterator[sqlite3.Connection]: + assert sqlite_file_database.sqlite.path is not None + conn = sqlite3.connect(sqlite_file_database.sqlite.path, timeout=0) + try: + yield conn + finally: + conn.close() + + def test_a_write_transaction_holds_the_write_lock_before_it_writes( + self, sqlite_file_database: Database, other_process: sqlite3.Connection + ) -> None: + with ProbeQueries(sqlite_file_database).transaction() as q: + q.items.names() + with pytest.raises(sqlite3.OperationalError, match="locked"): + other_process.execute("BEGIN IMMEDIATE") + + def test_a_write_left_uncommitted_on_the_raw_connection_is_reported_and_kept( + self, sqlite_file_database: Database + ) -> None: + # The migrator writes through the raw connection; a write it left uncommitted would otherwise fail the next + # transaction's BEGIN obscurely, and that transaction's rollback would discard the write. + queries = ProbeQueries(sqlite_file_database) + raw = sqlite_file_database.sqlite.conn + raw.execute("INSERT INTO probe_items (id, name, size) VALUES (1, 'raw', 0)") + + with pytest.raises(RuntimeError, match="raw connection without being committed"): + queries.items.names() + + raw.commit() + assert queries.items.names() == ["raw"] + + def test_raw_cursors_read_rows_by_name_after_a_transaction(self, sqlite_file_database: Database) -> None: + # Queries read their rows without the driver's `sqlite3.Row`; the migrator's raw cursors keep it. + queries = ProbeQueries(sqlite_file_database) + queries.items.add(1, "a") + + row = sqlite_file_database.sqlite.conn.execute("SELECT name FROM probe_items").fetchone() + + assert row["name"] == "a" + + def test_a_read_transaction_leaves_other_processes_free_to_write( + self, sqlite_file_database: Database, other_process: sqlite3.Connection + ) -> None: + queries = ProbeQueries(sqlite_file_database) + with queries.transaction(read_only=True) as q: + assert q.items.names() == [] + other_process.execute("BEGIN IMMEDIATE") + other_process.execute("INSERT INTO probe_items (id, name, size) VALUES (1, 'other', 0)") + other_process.commit() + assert q.items.names() == [] + + assert queries.items.names() == ["other"] + + def test_a_backup_holds_what_only_the_write_ahead_log_holds( + self, sqlite_file_database: Database, tmp_path: Path + ) -> None: + ProbeQueries(sqlite_file_database).items.add(1, "logged") + live = sqlite_file_database.sqlite.path + assert live is not None + # Not checkpointed yet: a copy of the database file alone would miss the row. + assert live.with_name(live.name + "-wal").stat().st_size > 0 + + sqlite_file_database.backup(tmp_path / "backup.db") + + with closing(sqlite3.connect(tmp_path / "backup.db")) as copy: + assert copy.execute("SELECT name FROM probe_items").fetchall() == [("logged",)] + + def test_a_backup_never_replaces_a_file_nor_copies_a_transaction_in_flight( + self, sqlite_file_database: Database, tmp_path: Path + ) -> None: + earlier = tmp_path / "earlier.db" + earlier.write_bytes(b"an earlier backup") + with pytest.raises(FileExistsError): + sqlite_file_database.backup(earlier) + assert earlier.read_bytes() == b"an earlier backup" + + with ProbeQueries(sqlite_file_database).transaction() as q: + q.items.add(1, "uncommitted") + with pytest.raises(NestedTransactionError): + sqlite_file_database.backup(tmp_path / "during.db") + assert not (tmp_path / "during.db").exists() + + def _commit_elsewhere(self, other_process: sqlite3.Connection) -> None: + other_process.execute("BEGIN IMMEDIATE") + other_process.execute("UPDATE probe_items SET size = size + 1") + other_process.commit() + + def test_a_single_call_that_loses_a_race_is_retried( + self, sqlite_file_database: Database, other_process: sqlite3.Connection + ) -> None: + queries = ProbeQueries(sqlite_file_database) + queries.items.add(1, "a") + first_attempt = iter([True]) + + def commit_elsewhere_once() -> None: + if next(first_attempt, False): + self._commit_elsewhere(other_process) + + with _retries(sqlite_file_database) as retries: + queries.items.read_then_write(commit_elsewhere_once) + + assert len(retries) == 1 + + def test_a_unit_run_by_run_that_loses_a_race_is_retried( + self, sqlite_file_database: Database, other_process: sqlite3.Connection + ) -> None: + queries = ProbeQueries(sqlite_file_database) + queries.items.add(1, "a") + first_attempt = iter([True]) + + def commit_elsewhere_once() -> None: + if next(first_attempt, False): + self._commit_elsewhere(other_process) + + with _retries(sqlite_file_database) as retries: + queries.run(lambda q: q.items.read_then_write(commit_elsewhere_once), read_only=True) + + assert len(retries) == 1 + + def test_a_call_that_keeps_losing_races_gives_up_after_three_attempts( + self, sqlite_file_database: Database, other_process: sqlite3.Connection + ) -> None: + queries = ProbeQueries(sqlite_file_database) + queries.items.add(1, "a") + + with _retries(sqlite_file_database) as retries: + with pytest.raises(ConflictError): + queries.items.read_then_write(lambda: self._commit_elsewhere(other_process)) + + assert len(retries) == 2 + + def test_a_write_lock_held_too_long_elsewhere_fails_the_work_without_a_retry( + self, sqlite_file_database: Database, other_process: sqlite3.Connection + ) -> None: + # A retry would wait the whole busy timeout again, holding up every other thread meanwhile. + sqlite_file_database.sqlite.conn.execute("PRAGMA busy_timeout = 20;") + other_process.execute("BEGIN IMMEDIATE") + queries = ProbeQueries(sqlite_file_database) + + with _retries(sqlite_file_database) as retries: + with pytest.raises(LockTimeoutError): + queries.items.add(1, "a") + with pytest.raises(LockTimeoutError): + with queries.transaction() as q: + q.items.add(1, "a") + other_process.rollback() + + assert retries == [] + assert queries.items.names() == [] + + def test_a_failed_commit_leaves_the_database_usable( + self, sqlite_file_database: Database, other_process: sqlite3.Connection + ) -> None: + # With a rollback journal (where WAL could not engage), a COMMIT needs every reader gone. A reader + # in another process that stays past the busy timeout makes the COMMIT itself fail. + sqlite_file_database.sqlite.conn.execute("PRAGMA journal_mode = DELETE;") + sqlite_file_database.sqlite.conn.execute("PRAGMA busy_timeout = 20;") + other_process.execute("BEGIN") + other_process.execute("SELECT COUNT(*) FROM probe_items").fetchone() + queries = ProbeQueries(sqlite_file_database) + + with pytest.raises(LockTimeoutError): + with queries.transaction() as q: + q.items.add(1, "a") + other_process.rollback() + + # The failed transaction was rolled back: it neither blocks the next one nor left its write behind. + queries.items.add(2, "b") + assert queries.items.names() == ["b"] + other_process.execute("BEGIN IMMEDIATE") + other_process.rollback() + + def test_a_second_connection_through_the_engine_is_refused_and_harms_no_transaction( + self, sqlite_file_database: Database + ) -> None: + # Another connection object over the shared driver connection would run outside the lock, and closing + # it would roll back whatever was open on that connection -- here, another thread's transaction. + queries = ProbeQueries(sqlite_file_database) + attempted = threading.Event() + refused: list[BaseException] = [] + + def connect_meanwhile() -> None: + try: + sqlite_file_database.engine.connect() + except BaseException as error: + refused.append(error) + attempted.set() + + with sqlite_file_database.begin(write=True) as conn: + conn.exec_driver_sql("INSERT INTO probe_items (id, name, size) VALUES (1, 'a', 0)") + intruder = threading.Thread(target=connect_meanwhile) + intruder.start() + assert attempted.wait(10) + intruder.join(10) + + assert len(refused) == 1 + assert queries.items.names() == ["a"] + + def test_an_interrupt_during_a_statement_leaves_the_database_open( + self, sqlite_file_database: Database, monkeypatch: pytest.MonkeyPatch + ) -> None: + dialect = sqlite_file_database.engine.dialect + execute = dialect.do_execute + + def interrupted_once(*args: Any, **kwargs: Any) -> None: + monkeypatch.setattr(dialect, "do_execute", execute) + raise KeyboardInterrupt + + monkeypatch.setattr(dialect, "do_execute", interrupted_once) + queries = ProbeQueries(sqlite_file_database) + + with pytest.raises(KeyboardInterrupt): + queries.items.add(1, "a") + + # The statement's transaction rolled back, and the one connection is still there. + queries.items.add(2, "b") + assert queries.items.names() == ["b"] + + +@server_only +class TestServerBackends: + def test_a_deadlock_fails_a_transaction_with_a_conflict(self, probe: ProbeQueries) -> None: + probe.items.add(1, "a") + probe.items.add(2, "b") + first_rows_locked = threading.Barrier(2, timeout=10) + outcomes: dict[str, str] = {} + + def lock_in_order(name: str, first: int, second: int) -> None: + try: + with probe.transaction() as q: + q.items.resize(first, 1) + first_rows_locked.wait() + q.items.resize(second, 1) + outcomes[name] = "committed" + except ConflictError: + outcomes[name] = "conflict" + + workers = [ + threading.Thread(target=lock_in_order, args=("forward", 1, 2)), + threading.Thread(target=lock_in_order, args=("backward", 2, 1)), + ] + for worker in workers: + worker.start() + for worker in workers: + worker.join(30) + + assert sorted(outcomes.values()) == ["committed", "conflict"] + + def test_a_single_call_that_deadlocks_is_retried(self, probe: ProbeQueries, empty_database: Database) -> None: + probe.items.add(1, "a") + probe.items.add(2, "b") + first_rows_locked = threading.Barrier(2, timeout=10) + attempted = threading.local() + + def meet_on_the_first_attempt() -> None: + # Both calls hold their first row before either asks for its second: a deadlock. The retry of + # the call that loses it runs straight through. + if not getattr(attempted, "once", False): + attempted.once = True + first_rows_locked.wait() + + errors: list[BaseException] = [] + + def resize(first: int, second: int) -> None: + try: + probe.items.resize_both(first, second, meet_on_the_first_attempt) + except BaseException as error: + errors.append(error) + + with _retries(empty_database) as retries: + workers = [threading.Thread(target=resize, args=(1, 2)), threading.Thread(target=resize, args=(2, 1))] + for worker in workers: + worker.start() + for worker in workers: + worker.join(30) + + assert errors == [] + assert len(retries) == 1 + + def test_a_unit_run_by_run_that_deadlocks_is_retried(self, probe: ProbeQueries, empty_database: Database) -> None: + probe.items.add(1, "a") + probe.items.add(2, "b") + first_rows_locked = threading.Barrier(2, timeout=10) + attempted = threading.local() + + def resize_in_one_unit(first: int, second: int) -> Callable[[ProbeQueries], None]: + def work(q: ProbeQueries) -> None: + q.items.resize(first, 1) + # Both units hold their first row before either asks for its second: a deadlock. The retry of + # the unit that loses it runs straight through. + if not getattr(attempted, "once", False): + attempted.once = True + first_rows_locked.wait() + q.items.resize(second, 1) + + return work + + errors: list[BaseException] = [] + + def resize(first: int, second: int) -> None: + try: + probe.run(resize_in_one_unit(first, second)) + except BaseException as error: + errors.append(error) + + with _retries(empty_database) as retries: + workers = [threading.Thread(target=resize, args=(1, 2)), threading.Thread(target=resize, args=(2, 1))] + for worker in workers: + worker.start() + for worker in workers: + worker.join(30) + + assert errors == [] + assert len(retries) == 1 + + def test_a_write_transaction_reads_what_others_committed_meanwhile(self, probe: ProbeQueries) -> None: + # Writers run at READ COMMITTED: they take no gap locks, and a guard they evaluate sees current rows. + # (Readers keep one snapshot; see test_a_read_only_transaction_reads_one_snapshot.) + probe.items.add(1, "a") + + with probe.transaction() as q: + before = q.items.count() + other = threading.Thread(target=probe.items.add, args=(2, "b")) + other.start() + other.join(10) + after = q.items.count() + + assert (before, after) == (1, 2) + + @pytest.mark.uses_database + def test_a_url_naming_the_other_server_flavor_is_refused(self, _external_test_schema: URL) -> None: + # Tables are created with the collations and column definitions of the flavor the URL names. + other = "mysql+pymysql" if _external_test_schema.get_backend_name() == "mariadb" else "mariadb+pymysql" + database = Database.open_url( + _external_test_schema.set(drivername=other).render_as_string(hide_password=False), + InvokeAILogger.get_logger("test_database"), + ) + try: + with pytest.raises((ValueError, InvalidRequestError), match="MariaDB"): + with database.begin(write=False): + pass + finally: + database.dispose() + + def test_literals_compare_byte_for_byte(self, empty_database: Database) -> None: + # The connection's collation, which comparisons of parameters and literals use: case and trailing + # spaces matter, as in SQLite. + with empty_database.begin(write=False) as conn: + assert tuple(conn.exec_driver_sql("SELECT 'A' = 'a', 'a ' = 'a'").one()) == (0, 0) + + @pytest.fixture + def lax_session_database(self, probe: ProbeQueries, _external_test_schema: URL) -> Iterator[Database]: + """The schema under test, through connections that start with an empty sql_mode. + + Both servers default to strict modes, so only a session that starts lax shows that the layer + pins its own mode instead of inheriting the server's. + """ + lax = _external_test_schema.update_query_dict({"init_command": "SET SESSION sql_mode = ''"}) + database = Database.open_url(lax.render_as_string(hide_password=False), InvokeAILogger.get_logger("lax")) + try: + yield database + finally: + database.dispose() + + def test_a_value_too_long_for_its_column_is_rejected(self, lax_session_database: Database) -> None: + queries = ProbeQueries(lax_session_database) + + with pytest.raises(DBAPIError) as raised: + queries.items.add(1, "x" * 40) + + assert raised.value.orig is not None + assert raised.value.orig.args[0] == 1406 # ER_DATA_TOO_LONG, instead of silent truncation + assert queries.items.names() == [] + + def test_an_ambiguous_group_by_is_rejected(self, lax_session_database: Database) -> None: + queries = ProbeQueries(lax_session_database) + queries.items.add(1, "a", size=1) + queries.items.add(2, "b", size=1) + + with pytest.raises(DBAPIError) as raised: + with lax_session_database.begin(write=False) as conn: + conn.exec_driver_sql("SELECT name, COUNT(*) FROM probe_items GROUP BY size").all() + + assert raised.value.orig is not None + assert raised.value.orig.args[0] == 1055 # ER_WRONG_FIELD_WITH_GROUP diff --git a/tests/app/services/shared/database/test_dialect.py b/tests/app/services/shared/database/test_dialect.py new file mode 100644 index 00000000000..fe59280a7d2 --- /dev/null +++ b/tests/app/services/shared/database/test_dialect.py @@ -0,0 +1,145 @@ +"""Constructs of `dialect.py` beyond what the query modules exercise on every backend.""" + +import pytest +from sqlalchemy import Column, Index, Integer, MetaData, String, Table, bindparam, insert, select +from sqlalchemy.dialects import mysql, postgresql, sqlite + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.dialect import ( + CaseInsensitiveLike, + InBoundSet, + bound_set, + fixed_limit, + insert_ignore, + upsert, +) +from invokeai.app.services.shared.database.schema.client_state import client_state +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.session_queue import session_queue +from invokeai.app.services.shared.database.schema.users import users + + +def test_an_upsert_refuses_a_table_with_a_unique_key_besides_its_primary_key() -> None: + # MySQL would update the row whichever unique key matched; the other backends only on the primary key. + with pytest.raises(ValueError, match="unique key besides its primary key"): + upsert("sqlite", users, update=["display_name"]) + + +def test_an_insert_ignore_refuses_a_table_with_a_unique_key_besides_its_primary_key() -> None: + # MySQL would skip the row whichever unique key matched; the other backends only on the primary key. + with pytest.raises(ValueError, match="unique key besides its primary key"): + insert_ignore("sqlite", users) + + +def _keyed_table(*unique_index_columns: str) -> Table: + table = Table("keyed", MetaData(), Column("name", String(10), primary_key=True), Column("value", Integer)) + if unique_index_columns: + Index("keyed_unique", *(table.c[name] for name in unique_index_columns), unique=True) + return table + + +def test_a_unique_index_of_the_primary_key_columns_is_that_same_key() -> None: + # Some migrated SQLite tables index their primary key once more, uniquely (images.image_name). + table = _keyed_table("name") + + sql = str(insert_ignore("sqlite", table).compile(dialect=sqlite.dialect())) + upsert("sqlite", table, update=["value"]) + + assert "ON CONFLICT (name) DO NOTHING" in sql + + +@pytest.mark.parametrize("columns", [("value",), ("name", "value")]) +def test_a_unique_index_of_other_columns_is_refused(columns: tuple[str, ...]) -> None: + # A key on more columns than the primary key's is another key too: MySQL resolves a conflict on either. + with pytest.raises(ValueError, match="unique key besides its primary key"): + insert_ignore("sqlite", _keyed_table(*columns)) + + +def test_an_upsert_refuses_a_table_without_a_primary_key() -> None: + table = Table("no_key", MetaData(), Column("value", Integer)) + + with pytest.raises(ValueError, match="no primary key"): + upsert("sqlite", table, update=["value"]) + + +@pytest.mark.parametrize("dialect", [sqlite.dialect(), mysql.dialect()]) +def test_a_case_insensitive_like_is_a_condition_as_it_stands(dialect: object) -> None: + # Compared with 1, a LIKE keeps SQLite from searching an index with its collation by range. + statement = select(client_state.c.key).where(CaseInsensitiveLike(client_state.c.key, bindparam("pattern"))) + + sql = str(statement.compile(dialect=dialect)) # type: ignore[arg-type] + + assert "LIKE" in sql and not sql.rstrip().endswith(("= 1", "= true")) + + +@pytest.mark.parametrize("dialect", [sqlite.dialect(), mysql.dialect()]) +def test_a_fixed_limit_is_written_into_the_statement(dialect: object) -> None: + # Bound, it would cost SQLite 10-20 µs per execution. + compiled = select(users.c.email).limit(fixed_limit(1)).compile(dialect=dialect) # type: ignore[arg-type] + + assert "LIMIT 1" in str(compiled) + assert 1 not in compiled.params.values() + + +def test_an_upsert_compiles_for_postgresql() -> None: + statement = upsert("postgresql", client_state, update=["value", "updated_at"]) + + sql = str(statement.compile(dialect=postgresql.dialect())) + + assert "ON CONFLICT (user_id, key) DO UPDATE SET value = excluded.value, updated_at = excluded.updated_at" in sql + + +def test_an_insert_ignore_compiles_for_postgresql() -> None: + statement = insert_ignore("postgresql", client_state) + + sql = str(statement.compile(dialect=postgresql.dialect())) + + assert "ON CONFLICT (user_id, key) DO NOTHING" in sql + + +def test_a_bound_set_matches_exactly_the_values_it_holds(database: Database) -> None: + quoted = 'it\'s "quoted" \\ here.png' + stored = ["a.png", "A.png", quoted, "ünïcode.png", "x" * 251 + ".png", "other.png"] + with database.begin(write=True) as conn: + conn.execute( + insert(images), + [ + {"image_name": name, "image_origin": "internal", "image_category": "general", "width": 1, "height": 1} + for name in stored + ], + ) + statement = select(images.c.image_name).where(InBoundSet(images.c.image_name, bindparam("names"))) + + def matching(values: list[str]) -> set[str]: + with database.begin(write=False) as conn: + return set(conn.execute(statement, {"names": bound_set(values)}).scalars()) + + # Compared as stored: case, quotes, backslashes and letters beyond ASCII, up to the longest name. + wanted = ["a.png", quoted, "ünïcode.png", "x" * 251 + ".png", "missing.png"] + assert matching(wanted) == set(wanted) - {"missing.png"} + assert matching([]) == set() + + +def test_a_bound_set_of_an_integer_column_matches_its_integers(database: Database) -> None: + # Beyond 32 bits, so that reading the set as INT would miss it. + large = 2**40 + with database.begin(write=True) as conn: + for item_id in (large, large + 1, large + 2): + conn.execute( + insert(session_queue).values( + item_id=item_id, queue_id="q", batch_id="b", session_id=f"s{item_id}", session="{}" + ) + ) + statement = select(session_queue.c.item_id).where(InBoundSet(session_queue.c.item_id, bindparam("ids"))) + + with database.begin(write=False) as conn: + found = set(conn.execute(statement, {"ids": bound_set([large, large + 2, 7])}).scalars()) + + assert found == {large, large + 2} + + +def test_a_bound_set_of_an_integer_column_is_read_as_integers_on_every_server() -> None: + statement = select(session_queue.c.item_id).where(InBoundSet(session_queue.c.item_id, bindparam("ids"))) + + assert "value BIGINT PATH '$'" in str(statement.compile(dialect=mysql.dialect())) + assert "CAST(value AS BIGINT)" in str(statement.compile(dialect=postgresql.dialect())) diff --git a/tests/app/services/shared/database/test_import_weight.py b/tests/app/services/shared/database/test_import_weight.py new file mode 100644 index 00000000000..5332c7368b8 --- /dev/null +++ b/tests/app/services/shared/database/test_import_weight.py @@ -0,0 +1,18 @@ +"""The database layer stays light to import: the user CLIs, migrations and the copy tool load it without the app.""" + +import subprocess +import sys + + +def test_the_layer_does_not_load_torch_or_the_invocations() -> None: + # A query module that imported invocation fields once made every database-only program load torch (~9 s). + probe = ( + "import sys\n" + "import invokeai.app.services.shared.database.database\n" + "import invokeai.app.services.shared.database.startup\n" + "print(sorted(m for m in ('torch', 'invokeai.app.invocations') if m in sys.modules))\n" + ) + + loaded = subprocess.run([sys.executable, "-c", probe], capture_output=True, text=True, check=True).stdout + + assert loaded.strip() == "[]" diff --git a/tests/app/services/shared/database/test_locks.py b/tests/app/services/shared/database/test_locks.py new file mode 100644 index 00000000000..ba2a34d5cff --- /dev/null +++ b/tests/app/services/shared/database/test_locks.py @@ -0,0 +1,47 @@ +"""Database locks: named locks (`queries/locks.py`) and the locks of single rows. What they serialise is tested with +the work they guard.""" + +import pytest +from sqlalchemy import delete + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries.locks import DatabaseLock +from invokeai.app.services.shared.database.schema.locks import db_locks + + +def test_every_lock_has_its_row_in_a_new_database(database: Database) -> None: + with database.queries.transaction() as q: + q.locks.acquire(*DatabaseLock) + + +def test_a_lock_without_its_row_is_refused_rather_than_locking_nothing(database: Database) -> None: + with database.begin(write=True) as conn: + conn.execute(delete(db_locks)) + + with pytest.raises(RuntimeError, match="admin_accounts"): + with database.queries.transaction() as q: + q.locks.acquire(DatabaseLock.ADMIN_ACCOUNTS) + + +def test_a_lock_outside_a_transaction_is_refused(database: Database) -> None: + # It would be released as soon as it was taken. + with pytest.raises(RuntimeError, match="first call of a transaction"): + database.queries.locks.acquire(DatabaseLock.ADMIN_ACCOUNTS) + + +def test_a_lock_after_other_work_in_the_transaction_is_refused(database: Database) -> None: + # What the transaction read before would not be guarded. + with pytest.raises(RuntimeError, match="first call of a transaction"): + with database.queries.transaction() as q: + q.users.count_active_admins() + q.locks.acquire(DatabaseLock.ADMIN_ACCOUNTS) + + +@pytest.mark.parametrize("row", ["board", "project"]) +def test_a_row_lock_outside_a_transaction_is_refused(database: Database, row: str) -> None: + # It would be released as soon as it was taken. + with pytest.raises(RuntimeError, match="only a transaction"): + if row == "board": + database.queries.boards.lock("board") + else: + database.queries.projects.lock("system", "project") diff --git a/tests/app/services/shared/database/test_schema.py b/tests/app/services/shared/database/test_schema.py new file mode 100644 index 00000000000..f7f3e3c4596 --- /dev/null +++ b/tests/app/services/shared/database/test_schema.py @@ -0,0 +1,352 @@ +"""The application schema on the backend under test: what its constraints, generated columns and defaults do. + +`test_schema_parity.py` shows that the metadata is the migrated SQLite schema; these show that the tables it +creates behave alike on every backend. +""" + +import json +import re +from typing import Any + +import pytest +from sqlalchemy import ( + ForeignKeyConstraint, + Table, + UniqueConstraint, + create_mock_engine, + delete, + func, + insert, + select, + update, +) + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.engines import MARIADB_BINARY_COLLATION, MYSQL_BINARY_COLLATION +from invokeai.app.services.shared.database.errors import ( + CheckViolation, + ForeignKeyViolation, + NotNullViolation, + UniqueViolation, +) +from invokeai.app.services.shared.database.schema import metadata +from invokeai.app.services.shared.database.schema.boards import board_images, boards +from invokeai.app.services.shared.database.schema.fonts import fonts +from invokeai.app.services.shared.database.schema.image_index import image_embeddings, image_index_vocab_terms +from invokeai.app.services.shared.database.schema.image_moves import ( + image_subfolder_move_items, + image_subfolder_move_jobs, +) +from invokeai.app.services.shared.database.schema.images import images +from invokeai.app.services.shared.database.schema.models import models +from invokeai.app.services.shared.database.schema.projects import projects +from invokeai.app.services.shared.database.schema.session_queue import session_queue +from invokeai.app.services.shared.database.schema.users import users +from invokeai.app.services.shared.database.schema.videos import videos +from invokeai.app.services.shared.database.schema.wildcards import wildcards +from invokeai.app.services.shared.database.schema.workflows import workflow_library +from invokeai.app.services.shared.database.types import Key, LongText, NoCaseKey +from tests.fixtures.database import external_test_db_url + +server_only = pytest.mark.skipif( + external_test_db_url() is None, reason="needs a MySQL or MariaDB server (INVOKEAI_TEST_DB_URL)" +) + +MODEL_CONFIG: dict[str, Any] = { + "hash": "blake3:0123456789abcdef", + "base": "sdxl", + "type": "main", + "path": "sdxl/main/model.safetensors", + "format": "checkpoint", + "name": "Model", + "description": None, + "source": "https://example.com/model.safetensors", + "source_type": "url", + "file_size": 6_938_078_334, + "trigger_phrases": ["a photo"], +} + + +def _create(database: Database, *tables: Table) -> None: + with database.begin(write=True) as conn: + metadata.create_all(conn, tables=list(tables)) + + +def _image(name: str) -> dict[str, Any]: + return {"image_name": name, "image_origin": "internal", "image_category": "general", "width": 8, "height": 8} + + +@pytest.mark.parametrize("url", ["mysql+pymysql://", "mariadb+pymysql://", "postgresql://"]) +def test_the_schema_compiles_for_each_server_backend(url: str) -> None: + # Without a server or a driver: every type and construct compiles for the backend, PostgreSQL (not yet + # supported) included, and an index that differs per backend is created once. + statements: list[str] = [] + + def compile_statement(statement: Any, *args: Any, **kwargs: Any) -> None: + statements.append(str(statement.compile(dialect=engine.dialect))) + + engine = create_mock_engine(url, compile_statement) + metadata.create_all(engine, checkfirst=False) + + creates_index = re.compile(r"\s*CREATE (?:UNIQUE )?INDEX (\S+)") + indexes = [match.group(1) for match in map(creates_index.match, statements) if match] + assert len(statements) == len(metadata.tables) + len(indexes) + assert sorted(set(indexes)) == sorted(indexes) + if not url.startswith("postgresql"): + # Keys of up to 3072 bytes need the DYNAMIC row format, whatever a server defaults to. + assert all("ROW_FORMAT=DYNAMIC" in s for s in statements if not creates_index.match(s)) + + +def test_text_columns_are_bounded_exactly_where_a_server_indexes_them_whole() -> None: + for table in metadata.sorted_tables: + indexed_whole = {column.name for column in table.primary_key.columns} + for constraint in table.constraints: + if isinstance(constraint, (UniqueConstraint, ForeignKeyConstraint)): + indexed_whole.update(column.name for column in constraint.columns) + for index in table.indexes: + if index._ddl_if is not None and index._ddl_if.dialect == "sqlite": + continue + prefix = index.dialect_options["mysql"]["length"] or {} + indexed_whole.update( + column.name for column in index.columns if not (isinstance(prefix, int) or column.name in prefix) + ) + for column in table.columns: + if isinstance(column.type, (Key, NoCaseKey)): + assert column.name in indexed_whole, f"{table.name}.{column.name}: a Key no server index needs" + elif isinstance(column.type, LongText): + assert column.name not in indexed_whole, f"{table.name}.{column.name}: LongText indexed whole" + + +@server_only +def test_server_tables_compare_text_byte_for_byte_whatever_the_database_default(empty_database: Database) -> None: + binary = MARIADB_BINARY_COLLATION if empty_database.dialect_name == "mariadb" else MYSQL_BINARY_COLLATION + # A server's default collation (and so a new database's) folds case. + with empty_database.begin(write=True) as conn: + conn.exec_driver_sql("ALTER DATABASE COLLATE utf8mb4_general_ci") + try: + with empty_database.begin(write=True) as conn: + metadata.create_all(conn) + finally: + with empty_database.begin(write=True) as conn: + conn.exec_driver_sql(f"ALTER DATABASE COLLATE {binary}") + + with empty_database.begin(write=False) as conn: + tables = conn.exec_driver_sql( + "SELECT engine, table_collation FROM information_schema.tables WHERE table_schema = DATABASE()" + ).all() + assert len(tables) == len(metadata.tables) + assert set(tables) == {("InnoDB", binary)} + + +def test_generated_columns_hold_members_of_the_json_document(empty_database: Database) -> None: + _create(empty_database, models, workflow_library) + workflow = {"name": "Flow", "description": "", "meta": {"category": "user"}, "tags": "a, b"} + + with empty_database.begin(write=True) as conn: + conn.execute(insert(models).values(id="m1", config=json.dumps(MODEL_CONFIG))) + conn.execute(insert(workflow_library).values(workflow_id="w1", workflow=json.dumps(workflow))) + + with empty_database.begin(write=False) as conn: + model = conn.execute(select(models)).one()._mapping + flow = conn.execute(select(workflow_library)).one()._mapping + + for member in ("hash", "base", "type", "path", "format", "name", "source", "source_type", "file_size"): + assert model[member] == MODEL_CONFIG[member], member + # A JSON null is NULL, not the text 'null'. + assert model["description"] is None + assert (flow["name"], flow["description"], flow["category"], flow["tags"]) == ("Flow", "", "user", "a, b") + + +@pytest.mark.parametrize("member", ["hash", "file_size"]) +@pytest.mark.parametrize("missing", ["null", "absent"]) +def test_a_model_config_without_a_required_member_is_refused( + empty_database: Database, member: str, missing: str +) -> None: + _create(empty_database, models) + config = {key: value for key, value in MODEL_CONFIG.items() if key != member} + if missing == "null": + config[member] = None + + # NOT NULL on the generated column; on MariaDB, which has none for generated columns, a CHECK. + expected = CheckViolation if empty_database.dialect_name == "mariadb" else NotNullViolation + with pytest.raises(expected): + with empty_database.begin(write=True) as conn: + conn.execute(insert(models).values(id="m1", config=json.dumps(config))) + + +def test_foreign_keys_cascade_set_null_or_restrict_as_declared(empty_database: Database) -> None: + _create(empty_database, users, images, boards, board_images, projects) + with empty_database.begin(write=True) as conn: + conn.execute(insert(users).values(user_id="u1", email="u1@example.com", password_hash="hash")) + conn.execute(insert(images).values(_image("i1.png"))) + conn.execute(insert(boards).values(board_id="b1", board_name="Board", cover_image_name="i1.png")) + conn.execute(insert(boards).values(board_id="b2", board_name="Project board")) + conn.execute(insert(board_images).values(board_id="b1", image_name="i1.png")) + conn.execute(insert(projects).values(project_id="p1", user_id="u1", name="Project", data="{}", board_id="b2")) + + with empty_database.begin(write=True) as conn: + conn.execute(delete(images).where(images.c.image_name == "i1.png")) + with empty_database.begin(write=False) as conn: + assert conn.execute(select(boards.c.cover_image_name).where(boards.c.board_id == "b1")).scalar_one() is None + assert conn.execute(select(func.count()).select_from(board_images)).scalar_one() == 0 + + # A project's board cannot be deleted while the project exists. + with pytest.raises(ForeignKeyViolation): + with empty_database.begin(write=True) as conn: + conn.execute(delete(boards).where(boards.c.board_id == "b2")) + + with empty_database.begin(write=True) as conn: + conn.execute(delete(users).where(users.c.user_id == "u1")) + with empty_database.begin(write=False) as conn: + assert conn.execute(select(func.count()).select_from(projects)).scalar_one() == 0 + + +def test_a_move_job_with_items_cannot_be_deleted(empty_database: Database) -> None: + # The items' foreign key has no ON DELETE action: the job's rows are the audit trail of its moves. + _create(empty_database, images, image_subfolder_move_jobs, image_subfolder_move_items) + with empty_database.begin(write=True) as conn: + conn.execute(insert(images).values(_image("i1.png"))) + job = conn.execute(insert(image_subfolder_move_jobs).values(state="planned")).inserted_primary_key + assert job is not None + conn.execute( + insert(image_subfolder_move_items).values( + job_id=job.id, image_name="i1.png", old_subfolder="", new_subfolder="a", state="planned" + ) + ) + + with pytest.raises(ForeignKeyViolation): + with empty_database.begin(write=True) as conn: + conn.execute(delete(image_subfolder_move_jobs)) + + +def test_vocabulary_terms_are_unique_whatever_their_case(empty_database: Database) -> None: + _create(empty_database, image_index_vocab_terms) + with empty_database.begin(write=True) as conn: + conn.execute(insert(image_index_vocab_terms).values(term="Cat")) + # Trailing spaces still make a different term, as in SQLite. + conn.execute(insert(image_index_vocab_terms).values(term="Cat ")) + + with pytest.raises(UniqueViolation): + with empty_database.begin(write=True) as conn: + conn.execute(insert(image_index_vocab_terms).values(term="cat")) + + +def test_a_directory_font_path_is_unique_and_uploaded_fonts_have_none(empty_database: Database) -> None: + # A partial unique index on SQLite, a full one on a server: the same constraint, because only directory + # fonts have a source path. + _create(empty_database, users, fonts) + font = {"filename": "f.ttf", "family": "F", "label": "F", "style": "normal", "weight": 400, "byte_size": 1} + with empty_database.begin(write=True) as conn: + conn.execute(insert(users).values(user_id="u1", email="u1@example.com", password_hash="hash")) + for n in (1, 2): + uploaded = {"id": f"u{n}", "owner_id": "u1", "scope": "private", "source": "uploaded"} + conn.execute( + insert(fonts).values({**font, **uploaded, "storage_path": f"u{n}.ttf", "content_hash": f"{n}" * 64}) + ) + directory = {"scope": "shared", "source": "directory", "source_path": "a/f.ttf"} + conn.execute(insert(fonts).values({**font, **directory, "id": "d1", "content_hash": "3" * 64})) + + with pytest.raises(UniqueViolation): + with empty_database.begin(write=True) as conn: + conn.execute(insert(fonts).values({**font, **directory, "id": "d2", "content_hash": "4" * 64})) + + +def test_other_names_differ_by_case(empty_database: Database) -> None: + _create(empty_database, users, wildcards) + with empty_database.begin(write=True) as conn: + conn.execute(insert(users).values(user_id="u1", email="u1@example.com", password_hash="hash")) + conn.execute(insert(wildcards).values(id="w1", name="Animals", user_id="u1")) + conn.execute(insert(wildcards).values(id="w2", name="animals", user_id="u1")) + + with empty_database.begin(write=False) as conn: + assert conn.execute(select(func.count()).select_from(wildcards)).scalar_one() == 2 + + +def test_queue_item_ids_are_never_reused(empty_database: Database) -> None: + _create(empty_database, session_queue) + item = {"batch_id": "b", "queue_id": "default", "session": "{}"} + with empty_database.begin(write=True) as conn: + conn.execute(insert(session_queue).values({**item, "session_id": "s1"})) + conn.execute(insert(session_queue).values({**item, "session_id": "s2"})) + conn.execute(delete(session_queue).where(session_queue.c.session_id == "s2")) + third = conn.execute(insert(session_queue).values({**item, "session_id": "s3"})).inserted_primary_key + + assert third is not None + assert third.item_id == 3 + + +def test_omitted_columns_take_their_defaults(empty_database: Database) -> None: + _create(empty_database, images, boards) + with empty_database.begin(write=True) as conn: + conn.execute(insert(boards).values(board_id="b1", board_name="Board")) + conn.execute(insert(images).values(_image("i1.png"))) + + with empty_database.begin(write=False) as conn: + board = conn.execute(select(boards)).one()._mapping + image = conn.execute(select(images.c.image_subfolder, images.c.is_intermediate, images.c.user_id)).one() + + assert (board["archived"], board["user_id"], board["is_public"], board["board_visibility"]) == ( + False, + "system", + False, + "private", + ) + for timestamp in (board["created_at"], board["updated_at"]): + assert re.fullmatch(r"\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3}", timestamp) + assert tuple(image) == ("", False, "system") + + +def test_an_update_sets_updated_at(empty_database: Database) -> None: + _create(empty_database, images, boards) + long_ago = "2020-01-01 00:00:00.000" + with empty_database.begin(write=True) as conn: + conn.execute(insert(boards).values(board_id="b1", board_name="Board", created_at=long_ago, updated_at=long_ago)) + conn.execute(update(boards).values(board_name="Renamed")) + + with empty_database.begin(write=False) as conn: + created_at, updated_at = conn.execute(select(boards.c.created_at, boards.c.updated_at)).one() + + assert created_at == long_ago + assert isinstance(updated_at, str) and updated_at > long_ago + + +def test_every_updated_at_column_is_set_by_updates() -> None: + # No trigger keeps `updated_at` current on any backend: an `update()` of its table must set it. (Fonts set + # theirs where a directory font changes; the quarantine table of old projects keeps theirs as found.) + set_elsewhere = {"fonts", "orphaned_projects_2026_08_06"} + without = [ + name + for name, table in metadata.tables.items() + if "updated_at" in table.c and name not in set_elsewhere and table.c.updated_at.onupdate is None + ] + assert without == [] + + +def test_columns_hold_long_and_large_values(empty_database: Database) -> None: + _create(empty_database, models, images, image_embeddings, videos) + # A model name is indexed by its first characters on a server; a path, which is unique, is limited to + # what a server's index holds whole. + config = MODEL_CONFIG | {"name": "n" * 1000, "path": "p" * 768} + embedding = bytes(range(256)) * 400 # beyond the 64 KiB of a server's BLOB + data = json.dumps({"layers": "x" * 100_000}) # beyond the 64 KiB of a server's TEXT + + with empty_database.begin(write=True) as conn: + conn.execute(insert(models).values(id="m1", config=json.dumps(config))) + conn.execute(insert(images).values({**_image("i1.png"), "metadata": data, "file_size_bytes": 2**40})) + conn.execute(insert(image_embeddings).values(image_name="i1.png", model_id="clip", dim=1, embedding=embedding)) + conn.execute( + insert(videos).values( + video_name="v1.mp4", + video_origin="internal", + video_category="general", + width=8, + height=8, + duration=1 / 3, + ) + ) + + with empty_database.begin(write=False) as conn: + assert tuple(conn.execute(select(models.c.name, models.c.path)).one()) == (config["name"], config["path"]) + assert tuple(conn.execute(select(images.c.metadata, images.c.file_size_bytes)).one()) == (data, 2**40) + assert conn.execute(select(image_embeddings.c.embedding)).scalar_one() == embedding + assert conn.execute(select(videos.c.duration)).scalar_one() == 1 / 3 diff --git a/tests/app/services/shared/database/test_schema_parity.py b/tests/app/services/shared/database/test_schema_parity.py new file mode 100644 index 00000000000..c5e1c812f48 --- /dev/null +++ b/tests/app/services/shared/database/test_schema_parity.py @@ -0,0 +1,304 @@ +"""The schema metadata describes the schema the SQLite migrations build, exactly. + +Server databases are created from the metadata, and queries are compiled against it, so a difference here +would be a difference between backends: a missing index, a cascade that does not happen, a default that is +not there. Compared on SQLite, where both can be built: the migration chain on one in-memory database, +`metadata.create_all()` on another. +""" + +import re +import sqlite3 +from collections import defaultdict +from collections.abc import Iterator +from typing import Any, Optional + +import pytest +from sqlalchemy import create_engine + +from invokeai.app.services.config.config_default import DefaultInvokeAIAppConfig +from invokeai.app.services.shared.database.schema import metadata +from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2026_08_06_add_project_boards import ( + AddProjectBoardsMigrationCallback, +) +from invokeai.backend.util.logging import InvokeAILogger +from tests.fixtures.sqlite_database import create_mock_sqlite_database + + +@pytest.fixture(scope="module") +def migrated(tmp_path_factory: pytest.TempPathFactory) -> Iterator[sqlite3.Connection]: + # Migrations clean up legacy files under the root; it must be this temporary one, whatever the environment says. + config = DefaultInvokeAIAppConfig(use_memory_db=True) + config._root = tmp_path_factory.mktemp("root") + db = create_mock_sqlite_database(config, InvokeAILogger.get_logger("test_schema_parity")) + conn = db.sqlite.conn + # Created by the project boards migration only when it has projects to rescue, which a new database has not. + AddProjectBoardsMigrationCallback(InvokeAILogger.get_logger("test_schema_parity"))._create_quarantine_table( + conn.cursor() + ) + conn.commit() + yield conn + db.dispose() + + +@pytest.fixture(scope="module") +def created() -> Iterator[sqlite3.Connection]: + conn = sqlite3.connect(":memory:") + engine = create_engine("sqlite://", creator=lambda: conn) + metadata.create_all(engine) + yield conn + engine.dispose() + conn.close() + + +def test_the_metadata_creates_the_migrated_schema(created: sqlite3.Connection, migrated: sqlite3.Connection) -> None: + created_schema = _schema(created) + migrated_schema = _schema(migrated) + + differences = [ + f"{table}.{aspect}:\n metadata {created_schema.get(table, {}).get(aspect)!r}\n" + f" migrated {migrated_schema.get(table, {}).get(aspect)!r}" + for table in sorted(created_schema.keys() | migrated_schema.keys()) + for aspect in ("table", "columns", "checks", "foreign_keys", "indexes") + if created_schema.get(table, {}).get(aspect) != migrated_schema.get(table, {}).get(aspect) + ] + assert differences == [] + + +def test_neither_schema_has_triggers(created: sqlite3.Connection, migrated: sqlite3.Connection) -> None: + # The application sets what the triggers of the earlier migrations set; 2026_10_07_drop_sqlite_triggers drops them. + assert _names(migrated, "trigger") == set() + assert _names(created, "trigger") == set() + assert _names(created, "view") == _names(migrated, "view") == set() + + +def _names(conn: sqlite3.Connection, kind: str) -> set[str]: + return {name for (name,) in conn.execute("SELECT name FROM sqlite_master WHERE type = ?", (kind,))} + + +def _schema(conn: sqlite3.Connection) -> dict[str, dict[str, Any]]: + """What SQLite reports about each table, and what only its declaration shows. + + Its declaration shows the generation expressions, CHECKs, conflict clauses, deferred constraints and + AUTOINCREMENT. + """ + schema: dict[str, dict[str, Any]] = {} + for name, sql in conn.execute( + "SELECT name, sql FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%'" + ).fetchall(): + definitions = _definitions(sql) + code = _code(sql) + generated = { + _column_name(definition): _expression_after(definition, "GENERATED ALWAYS AS") + for definition in definitions + if re.search(r"\bGENERATED\s+ALWAYS\s+AS\b", definition, re.IGNORECASE) + } + kind, column_count, without_rowid, strict = conn.execute( + "SELECT type, ncol, wr, strict FROM pragma_table_list WHERE schema = 'main' AND name = ?", (name,) + ).fetchone() + schema[name] = { + "table": { + "kind": kind, + "columns": column_count, + "without_rowid": without_rowid, + "strict": strict, + "autoincrement": re.search(r"\bAUTOINCREMENT\b", code, re.IGNORECASE) is not None, + "conflict_clauses": sorted( + clause.upper() for clause in re.findall(r"\bON\s+CONFLICT\s+(\w+)", code, re.IGNORECASE) + ), + "deferred": sorted( + " ".join(clause.upper().split()) + for clause in re.findall( + r"\b(?:NOT\s+)?DEFERRABLE(?:\s+INITIALLY\s+\w+)?", code, flags=re.IGNORECASE + ) + ), + }, + "columns": [ + { + "name": column, + "type": declared_type, + "not_null": not_null, + "default": _default(default), + "primary_key_position": pk, + "hidden": hidden, + "collation": _collation(conn, name, column), + "generated": generated.get(column), + } + # Read whole first: `_collation` changes the schema, which an open statement over it would block. + for _, column, declared_type, not_null, default, pk, hidden in conn.execute( + f"PRAGMA table_xinfo({_quote(name)})" + ).fetchall() + ], + "checks": sorted(check for definition in definitions for check in _expressions_after(definition, "CHECK")), + "foreign_keys": _foreign_keys(conn, name), + "indexes": sorted(_indexes(conn, name), key=repr), + } + return schema + + +def _foreign_keys(conn: sqlite3.Connection, table: str) -> list[tuple[Any, ...]]: + """Each foreign key whole: a composite key's columns in order, with its actions.""" + keys: dict[int, list[Any]] = defaultdict(list) + for key_id, _, target, source, referenced, on_update, on_delete, match in conn.execute( + f"PRAGMA foreign_key_list({_quote(table)})" + ).fetchall(): + keys[key_id].append((target, source, referenced, on_update, on_delete, match)) + return sorted( + ( + rows[0][0], + tuple(row[1] for row in rows), + tuple(row[2] for row in rows), + *rows[0][3:], + ) + for rows in keys.values() + ) + + +def _indexes(conn: sqlite3.Connection, table: str) -> Iterator[tuple[Any, ...]]: + for _, index, unique, origin, partial in conn.execute(f"PRAGMA index_list({_quote(table)})").fetchall(): + keys = tuple( + (column, descending, collation) + for _, _, column, descending, collation, key in conn.execute( + f"PRAGMA index_xinfo({_quote(index)})" + ).fetchall() + if key + ) + sql: Optional[str] = conn.execute("SELECT sql FROM sqlite_master WHERE name = ?", (index,)).fetchone()[0] + # SQLite reports an expression key without its expression, so the declared key list stands in for it. + expressions = _normalized(_balanced(sql, sql.index("("))) if sql and any(k[0] is None for k in keys) else None + predicate: Optional[str] = None + if partial and sql: + predicate = _normalized(re.split(r"\bWHERE\b", sql, maxsplit=1, flags=re.IGNORECASE)[1]) + # Indexes SQLite creates for a PRIMARY KEY or UNIQUE constraint are named by constraint position. + yield (index if origin == "c" else None, unique, origin, keys, expressions, predicate) + + +def _collation(conn: sqlite3.Connection, table: str, column: str) -> str: + """The column's collation, as SQLite resolves it for an index on the column alone.""" + conn.execute(f"CREATE INDEX _parity_probe ON {_quote(table)} ({_quote(column)})") + try: + ((collation,),) = conn.execute( + "SELECT coll FROM pragma_index_xinfo('_parity_probe') WHERE seqno = 0" + ).fetchall() + return str(collation) + finally: + conn.execute("DROP INDEX _parity_probe") + + +def _default(default: Optional[str]) -> Optional[str]: + # SQLite reads a double-quoted default that names no column as a string literal: "user" is 'user'. + if default is not None and re.fullmatch(r'"[^"]*"', default): + return f"'{default[1:-1]}'" + return default + + +def _definitions(create_table: str) -> list[str]: + """The column definitions and table constraints of a CREATE TABLE statement, without comments.""" + body = _balanced(create_table, create_table.index("(")) + definitions: list[str] = [] + current: list[str] = [] + depth = 0 + for token in _tokens(body): + if token.startswith("--"): + token = " " + elif token == "(": + depth += 1 + elif token == ")": + depth -= 1 + elif token == "," and depth == 0: + definitions.append(_normalized("".join(current))) + current = [] + continue + current.append(token) + definitions.append(_normalized("".join(current))) + return [definition for definition in definitions if definition] + + +def _column_name(definition: str) -> str: + name = definition.split()[0] + return name[1:-1] if name[0] in '"`[' else name + + +def _expression_after(definition: str, keyword: str) -> str: + (expression,) = _expressions_after(definition, keyword) + return expression + + +def _expressions_after(definition: str, keyword: str) -> list[str]: + pattern = r"\b" + r"\s+".join(keyword.split()) + r"\s*\(" + return [ + _normalized(_balanced(definition, match.end() - 1)) + for match in re.finditer(pattern, _code(definition), re.IGNORECASE) + ] + + +def _balanced(sql: str, open_at: int) -> str: + """The text inside the parentheses that open at `open_at`.""" + depth = 0 + position = open_at + for token in _tokens(sql[open_at:]): + if token == "(": + depth += 1 + elif token == ")": + depth -= 1 + if depth == 0: + return sql[open_at + 1 : position] + position += len(token) + raise ValueError(f"Unbalanced parentheses in {sql!r}") + + +def _tokens(sql: str) -> Iterator[str]: + """Quoted strings and identifiers whole, comments whole, every other character on its own.""" + position = 0 + while position < len(sql): + char = sql[position] + if sql.startswith("--", position): + end = sql.find("\n", position) + end = len(sql) if end == -1 else end + yield sql[position:end] + position = end + continue + if char in "'\"`": + end = position + 1 + while end < len(sql): + if sql[end] == char: + if sql[end + 1 : end + 2] == char: + end += 2 + continue + break + end += 1 + yield sql[position : end + 1] + position = end + 1 + continue + yield char + position += 1 + + +def _code(sql: str) -> str: + """The statement with comments as spaces and each string literal's text blanked, at the same positions. + + For finding keywords: only outside comments and literals are they keywords. + """ + return "".join( + " " * len(token) + if token.startswith("--") + else ("'" + " " * (len(token) - 2) + "'" if token.startswith("'") else token) + for token in _tokens(sql) + ) + + +def _normalized(sql: str) -> str: + """Comments dropped and whitespace collapsed, none just inside parentheses; string literals as written.""" + parts: list[str] = [] + for token in _tokens(sql): + if token.isspace() or token.startswith("--"): + if parts and parts[-1] not in (" ", "("): + parts.append(" ") + continue + if token == ")" and parts and parts[-1] == " ": + parts.pop() + parts.append(token) + return "".join(parts).strip() + + +def _quote(identifier: str) -> str: + return '"' + identifier.replace('"', '""') + '"' diff --git a/tests/app/services/shared/sqlite/test_sqlite_synchronous.py b/tests/app/services/shared/database/test_sqlite_synchronous.py similarity index 82% rename from tests/app/services/shared/sqlite/test_sqlite_synchronous.py rename to tests/app/services/shared/database/test_sqlite_synchronous.py index d51116dee54..04526f6065c 100644 --- a/tests/app/services/shared/sqlite/test_sqlite_synchronous.py +++ b/tests/app/services/shared/database/test_sqlite_synchronous.py @@ -14,20 +14,20 @@ from pydantic import ValidationError from invokeai.app.services.config.config_default import InvokeAIAppConfig -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.shared.sqlite.sqlite_util import init_db +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.startup import init_database from invokeai.backend.util.logging import InvokeAILogger # What `PRAGMA synchronous` reports back, as documented by SQLite. PRAGMA_VALUES = {"off": 0, "normal": 1, "full": 2, "extra": 3} -def _synchronous_of(db: SqliteDatabase) -> int: - return int(db._conn.execute("PRAGMA synchronous;").fetchone()[0]) +def _synchronous_of(db: Database) -> int: + return int(db.sqlite.conn.execute("PRAGMA synchronous;").fetchone()[0]) -def _journal_mode_of(db: SqliteDatabase) -> str: - return str(db._conn.execute("PRAGMA journal_mode;").fetchone()[0]).lower() +def _journal_mode_of(db: Database) -> str: + return str(db.sqlite.conn.execute("PRAGMA journal_mode;").fetchone()[0]).lower() def _warnings_about_durability(caplog: pytest.LogCaptureFixture) -> list[str]: @@ -42,7 +42,7 @@ def sqlite_cannot_use_wal(monkeypatch: pytest.MonkeyPatch) -> None: SQLite refuses to move a transient database into WAL mode, so the connection opened here answers `PRAGMA journal_mode = WAL;` with a rollback-journal mode instead -- a real refusal from real - SQLite. The interception is at the `sqlite3.connect` boundary; everything `SqliteDatabase` does + SQLite. The interception is at the `sqlite3.connect` boundary; everything `Database` does after it is the production path. """ real_connect = sqlite3.connect @@ -71,7 +71,7 @@ def test_other_sqlite_values_are_rejected(self, value): class TestTheConstructorValidatesTheSetting: - """`SqliteDatabase` is a public constructor -- `user_management.py` and tests call it directly -- + """`Database` is a public constructor -- `user_management.py` and tests call it directly -- and the value ends up interpolated into a PRAGMA statement, which takes no bind parameters. The `Literal` annotation is checked by a type checker, not at runtime, so the constructor checks it.""" @@ -80,7 +80,7 @@ def test_a_value_outside_the_literal_never_reaches_sql(self, tmp_path, value): db_path = tmp_path / "rejected.db" with pytest.raises(ValueError, match="synchronous"): - SqliteDatabase(db_path=db_path, logger=InvokeAILogger.get_logger(), synchronous=value) + Database.open_sqlite(db_path, InvokeAILogger.get_logger(), synchronous=value) # Rejected before anything was opened, so there is nothing for the statement to have run against. assert not db_path.exists() @@ -88,14 +88,14 @@ def test_a_value_outside_the_literal_never_reaches_sql(self, tmp_path, value): class TestThePragmaIsApplied: def test_the_default_setting_uses_full(self, tmp_path): - db = SqliteDatabase(db_path=tmp_path / "default.db", logger=InvokeAILogger.get_logger()) + db = Database.open_sqlite(tmp_path / "default.db", InvokeAILogger.get_logger()) assert _synchronous_of(db) == PRAGMA_VALUES["full"] @pytest.mark.parametrize("setting", ["full", "normal"]) def test_each_setting_reaches_the_connection(self, tmp_path, setting): - db = SqliteDatabase( - db_path=tmp_path / f"{setting}.db", - logger=InvokeAILogger.get_logger(), + db = Database.open_sqlite( + tmp_path / f"{setting}.db", + InvokeAILogger.get_logger(), synchronous=setting, ) assert _synchronous_of(db) == PRAGMA_VALUES[setting] @@ -104,7 +104,7 @@ def test_an_in_memory_database_keeps_normal_despite_having_no_journal(self): """An in-memory database reports journal mode `memory`, never `wal`, but it has no durability to trade away in the first place -- so the WAL requirement below must not reach it. The migrator and several tests run on this path.""" - db = SqliteDatabase(db_path=None, logger=InvokeAILogger.get_logger(), synchronous="normal") + db = Database.open_sqlite(None, InvokeAILogger.get_logger(), synchronous="normal") assert _journal_mode_of(db) != "wal" assert _synchronous_of(db) == PRAGMA_VALUES["normal"] @@ -113,7 +113,7 @@ def test_wal_is_still_on(self, tmp_path): """`normal` is only safe against corruption *because* of WAL. If journal mode ever stopped being WAL for an ordinary database file, this setting would become a different trade than the one documented.""" - db = SqliteDatabase(db_path=tmp_path / "wal.db", logger=InvokeAILogger.get_logger(), synchronous="normal") + db = Database.open_sqlite(tmp_path / "wal.db", InvokeAILogger.get_logger(), synchronous="normal") assert _journal_mode_of(db) == "wal" def test_normal_is_refused_when_wal_did_not_engage(self, tmp_path, caplog, sqlite_cannot_use_wal): @@ -124,7 +124,7 @@ def test_normal_is_refused_when_wal_did_not_engage(self, tmp_path, caplog, sqlit logger = InvokeAILogger.get_logger() with caplog.at_level(logging.WARNING): - db = SqliteDatabase(db_path=tmp_path / "no-wal.db", logger=logger, synchronous="normal") + db = Database.open_sqlite(tmp_path / "no-wal.db", logger, synchronous="normal") journal_mode = _journal_mode_of(db) assert journal_mode != "wal" @@ -136,7 +136,7 @@ def test_full_is_applied_without_a_warning_when_wal_did_not_engage(self, tmp_pat logger = InvokeAILogger.get_logger() with caplog.at_level(logging.WARNING): - db = SqliteDatabase(db_path=tmp_path / "no-wal-full.db", logger=logger) + db = Database.open_sqlite(tmp_path / "no-wal-full.db", logger) assert _synchronous_of(db) == PRAGMA_VALUES["full"] assert _warnings_about_durability(caplog) == [] @@ -149,15 +149,15 @@ class TestTheSettingReachesTheDatabase: boots, every other test passes, and the setting simply does nothing. This happened once during development, which is why it is pinned rather than assumed. - Driven through the real `init_db` against an in-memory database, migrations and all, so what is - asserted is the connection the app ends up with. Stubbing `SqliteDatabase` here and checking the + Driven through the real `init_database` against an in-memory database, migrations and all, so what is + asserted is the connection the app ends up with. Stubbing `Database` here and checking the keyword arrived would pin the argument's *name* instead: an implementation that accepts it and then applies it conditionally would still pass, which is the regression worth catching. """ @pytest.mark.parametrize("setting", ["full", "normal"]) def test_init_db_carries_the_configured_value_to_the_connection(self, setting): - db = init_db( + db = init_database( config=InvokeAIAppConfig(use_memory_db=True, db_synchronous=setting), logger=InvokeAILogger.get_logger(), image_files=Mock(), diff --git a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_07_30_repair_projects_table.py b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_07_30_repair_projects_table.py index 55ad9b2868b..19677c70baa 100644 --- a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_07_30_repair_projects_table.py +++ b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_07_30_repair_projects_table.py @@ -3,14 +3,14 @@ from pathlib import Path from unittest.mock import MagicMock -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.sqlite_migrator.migration_loader import MigrationBuildContext, build_migrations from invokeai.app.services.shared.sqlite_migrator.migrations.migration_34 import build_migration_34 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2026_07_30_repair_projects_table import ( RepairProjectsTableMigrationCallback, build_migration, ) -from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import Migrator def _table_exists(cursor: sqlite3.Cursor, name: str) -> bool: @@ -66,10 +66,10 @@ def test_repairs_a_database_that_came_from_an_upstream_build(tmp_path: Path) -> context = MigrationBuildContext(app_config=MagicMock(), logger=logger, image_files=MagicMock()) all_migrations = build_migrations(context) - db = SqliteDatabase(db_path=tmp_path / "upstream.db", logger=logger, verbose=False) + db = Database.open_sqlite(tmp_path / "upstream.db", logger) # 1. Bring the database to legacy version 32 using only the migrations at or below 32. - to_v32 = SqliteMigrator(db=db) + to_v32 = Migrator(db) for migration in all_migrations: if migration.to_version is not None and migration.to_version <= 32: to_v32.register_migration(migration) @@ -77,23 +77,22 @@ def test_repairs_a_database_that_came_from_an_upstream_build(tmp_path: Path) -> # 2. Simulate UPSTREAM's migration_33 having run. Its body is what this fork carries as # migration_34, so run that and stamp the legacy version/id upstream would have recorded. - cursor = db._conn.cursor() + cursor = db.sqlite.conn.cursor() build_migration_34().callback(cursor) cursor.execute("INSERT INTO migrations (version) VALUES (33);") cursor.execute("INSERT INTO applied_migrations (migration_id, legacy_version) VALUES ('migration_33', 33);") - db._conn.commit() + db.sqlite.conn.commit() assert not _table_exists(cursor, "projects") # 3. Open it with the full fork migrator, as the app does on startup. - full = SqliteMigrator(db=db) + full = Migrator(db) for migration in all_migrations: full.register_migration(migration) full.run_migrations() - # Without the repair migration the fork's migration_33 stays skipped and this is still missing. + # Without the repair migration the fork's migration_33 stays skipped and this is still missing. (The trigger the + # repair creates is gone again: a later migration drops every trigger.) assert _table_exists(cursor, "projects") - cursor.execute("SELECT name FROM sqlite_master WHERE type='trigger' AND name='tg_projects_updated_at';") - assert cursor.fetchone() is not None def test_migration_has_a_dated_id_so_it_cannot_collide_with_an_upstream_number() -> None: diff --git a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_08_06_add_project_boards.py b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_08_06_add_project_boards.py index 0c18ce3bdc8..c5a2d067e66 100644 --- a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_08_06_add_project_boards.py +++ b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_08_06_add_project_boards.py @@ -8,13 +8,13 @@ import pytest -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.sqlite_migrator.migration_loader import MigrationBuildContext, build_migrations from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2026_08_06_add_project_boards import ( AddProjectBoardsMigrationCallback, build_migration, ) -from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import Migrator MIGRATION_ID = "2026_08_06_add_project_boards" @@ -531,29 +531,29 @@ def test_it_runs_once_through_the_real_migrator_and_is_not_reapplied(tmp_path: P context = MigrationBuildContext(app_config=MagicMock(), logger=logger, image_files=MagicMock()) all_migrations = build_migrations(context) - db = SqliteDatabase(db_path=tmp_path / "projects.db", logger=logger, verbose=False) + db = Database.open_sqlite(tmp_path / "projects.db", logger) - before = SqliteMigrator(db=db) + before = Migrator(db) for migration in all_migrations: if migration.id == MIGRATION_ID: break before.register_migration(migration) before.run_migrations() - cursor = db._conn.cursor() + cursor = db.sqlite.conn.cursor() cursor.execute("INSERT INTO boards (board_id, board_name, user_id) VALUES ('b1', 'Old name', 'system');") cursor.execute( "INSERT INTO projects (project_id, user_id, name, data) VALUES ('p1', 'system', 'Kept', ?);", (_gallery_document("b1"),), ) - db._conn.commit() + db.sqlite.conn.commit() - full = SqliteMigrator(db=db) + full = Migrator(db) for migration in all_migrations: full.register_migration(migration) assert full.run_migrations() is True - # SqliteDatabase sets row_factory = sqlite3.Row, which does not compare equal to a tuple. + # Database sets row_factory = sqlite3.Row, which does not compare equal to a tuple. cursor.execute("SELECT name, board_id FROM projects WHERE project_id = 'p1';") assert tuple(cursor.fetchone()) == ("Kept", "b1") cursor.execute("SELECT board_name FROM boards WHERE board_id = 'b1';") @@ -562,7 +562,7 @@ def test_it_runs_once_through_the_real_migrator_and_is_not_reapplied(tmp_path: P assert cursor.fetchall() == [] # Migration tracking, not idempotent DDL, is what stops a second run rebuilding the table again. - again = SqliteMigrator(db=db) + again = Migrator(db) for migration in all_migrations: again.register_migration(migration) assert again.run_migrations() is False @@ -576,21 +576,20 @@ def test_it_runs_through_the_real_migrator_on_a_database_with_no_projects(tmp_pa context = MigrationBuildContext(app_config=MagicMock(), logger=logger, image_files=MagicMock()) all_migrations = build_migrations(context) - db = SqliteDatabase(db_path=tmp_path / "projects.db", logger=logger, verbose=False) - migrator = SqliteMigrator(db=db) + db = Database.open_sqlite(tmp_path / "projects.db", logger) + migrator = Migrator(db) for migration in all_migrations: migrator.register_migration(migration) assert migrator.run_migrations() is True - cursor = db._conn.cursor() + cursor = db.sqlite.conn.cursor() columns = {row[1] for row in cursor.execute("PRAGMA table_info(projects);").fetchall()} assert "board_id" in columns cursor.execute("SELECT COUNT(*) FROM projects;") assert cursor.fetchone()[0] == 0 - # The scratch table must not outlive the rebuild, and the trigger must be back. + # The scratch table must not outlive the rebuild. (The trigger it puts back, a later migration drops with the + # others; `test_the_rebuilt_table_keeps_its_key_columns_ordering_and_trigger` covers this migration alone.) cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='projects_with_boards';") assert cursor.fetchone() is None - cursor.execute("SELECT name FROM sqlite_master WHERE type='trigger' AND name='tg_projects_updated_at';") - assert cursor.fetchone() is not None cursor.execute("PRAGMA foreign_key_check;") assert cursor.fetchall() == [] diff --git a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_09_22_index_unmeasured_intermediates.py b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_09_22_index_unmeasured_intermediates.py deleted file mode 100644 index c0d247910de..00000000000 --- a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_09_22_index_unmeasured_intermediates.py +++ /dev/null @@ -1,98 +0,0 @@ -from collections.abc import Iterator -from logging import getLogger - -import pytest - -from invokeai.app.services.intermediates.intermediates_records_sqlite import IntermediatesRecordsSqlite, MediaKind -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2026_09_22_add_intermediates_management import ( - build_migration as build_intermediates_migration, -) -from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2026_09_22_index_unmeasured_intermediates import ( - build_migration, -) - - -@pytest.fixture -def db() -> Iterator[SqliteDatabase]: - database = SqliteDatabase(db_path=None, logger=getLogger(__name__)) - with database.transaction() as cursor: - cursor.execute("CREATE TABLE session_queue (item_id INTEGER PRIMARY KEY);") - for table, kind in (("images", "image"), ("videos", "video")): - cursor.execute( - f"""CREATE TABLE {table} ( - {kind}_name TEXT PRIMARY KEY, - {kind}_subfolder TEXT NOT NULL DEFAULT '', - user_id TEXT NOT NULL DEFAULT 'system', - is_intermediate BOOLEAN NOT NULL DEFAULT TRUE, - created_at TEXT NOT NULL DEFAULT '2020-01-01 00:00:00.000' - );""" - ) - # Existing installations have already applied this migration before receiving the index. - build_intermediates_migration(getLogger(__name__)).callback(cursor) - build_migration().callback(cursor) - yield database - database._conn.close() - - -@pytest.mark.parametrize("kind", ["image", "video"]) -def test_late_measurement_batches_do_not_rescan_completed_work(db: SqliteDatabase, kind: MediaKind) -> None: - table = f"{kind}s" - with db.transaction() as cursor: - cursor.executemany( - f"INSERT INTO {table} ({kind}_name) VALUES (?);", - [(f"asset-{index}",) for index in range(100_000)], - ) - records = IntermediatesRecordsSqlite(db) - - def next_batch() -> tuple[list[tuple[str, str]], int]: - instructions = 0 - - def count_instructions() -> int: - nonlocal instructions - instructions += 100 - return 0 - - db._conn.set_progress_handler(count_instructions, 100) - try: - return records.next_unmeasured(kind, 200, min_age_seconds=10), instructions - finally: - db._conn.set_progress_handler(None, 0) - - first, first_work = next_batch() - assert first == [(f"asset-{index}", "") for index in range(200)] - with db.transaction() as cursor: - cursor.execute(f"UPDATE {table} SET file_size_bytes = 100 WHERE rowid <= 99800;") - late, late_work = next_batch() - assert late == [(f"asset-{index}", "") for index in range(99800, 100_000)] - # SQLite VM work is deterministic enough for a generous growth bound and independent of - # machine speed. A table scan exceeds this by two orders of magnitude on the late batch. - assert 0 < first_work < 10_000 - assert late_work < 2 * first_work - - -@pytest.mark.parametrize("kind", ["image", "video"]) -def test_backfill_preserves_young_and_unmeasurable_rows(db: SqliteDatabase, kind: MediaKind) -> None: - table = f"{kind}s" - with db.transaction() as cursor: - cursor.executemany( - f"INSERT INTO {table} ({kind}_name, is_intermediate, file_size_bytes, created_at) VALUES (?, ?, ?, ?);", - [ - ("durable", False, None, "2020-01-01"), - ("measured", True, 100, "2020-01-01"), - ("unreadable", True, None, "2020-01-01"), - ("young", True, None, "2999-01-01"), - ("z-first", True, None, "2020-01-01"), - ("a-second", True, None, "2020-01-01"), - ], - ) - # The index migration can be replayed without rewriting measurement state. - build_migration().callback(cursor) - records = IntermediatesRecordsSqlite(db) - records.mark_unmeasurable(kind, ["unreadable"]) - - assert records.next_unmeasured(kind, 200, min_age_seconds=10) == [("z-first", ""), ("a-second", "")] - with db.transaction() as cursor: - cursor.execute(f"UPDATE {table} SET file_size_bytes = 100 WHERE {kind}_name IN ('z-first', 'a-second');") - cursor.execute(f"UPDATE {table} SET created_at = '2020-01-01' WHERE {kind}_name = 'young';") - assert records.next_unmeasured(kind, 200, min_age_seconds=10) == [("young", "")] diff --git a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_09_23_index_client_state_references.py b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_09_23_index_client_state_references.py index ededffd88a3..83e111b008a 100644 --- a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_09_23_index_client_state_references.py +++ b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_09_23_index_client_state_references.py @@ -4,7 +4,7 @@ import sqlite3 from logging import getLogger -from invokeai.app.services.shared.media_references import create_media_references_table +from invokeai.app.services.shared.sqlite_migrator.migrations._media_references_v1 import create_media_references_table from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2026_09_23_index_client_state_references import ( build_migration, ) diff --git a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_09_30_store_embeddings_fp16.py b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_09_30_store_embeddings_fp16.py index b2084287764..bc3f2e53bb6 100644 --- a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_09_30_store_embeddings_fp16.py +++ b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_09_30_store_embeddings_fp16.py @@ -6,7 +6,7 @@ import pytest from invokeai.app.services.image_index.image_index_common import blob_to_embedding, embedding_to_blob -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.sqlite_migrator.migration_loader import ( MigrationBuildContext, build_migrations, @@ -16,7 +16,7 @@ build_migration, ) from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration, MigrationError -from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import Migrator MIGRATION_ID = "2026_09_30_store_embeddings_fp16" DEPENDENCY_ID = "2026_09_12_add_video_embeddings" @@ -72,14 +72,14 @@ def _create_schema(connection: sqlite3.Connection, *, image_encoding: bool = Fal connection.commit() -def _make_db(tmp_path: Path) -> SqliteDatabase: - db = SqliteDatabase(db_path=tmp_path / "embeddings.db", logger=logging.getLogger(__name__)) - _create_schema(db._conn) +def _make_db(tmp_path: Path) -> Database: + db = Database.open_sqlite(tmp_path / "embeddings.db", logging.getLogger(__name__)) + _create_schema(db.sqlite.conn) return db -def _migrator(db: SqliteDatabase) -> SqliteMigrator: - migrator = SqliteMigrator(db) +def _migrator(db: Database) -> Migrator: + migrator = Migrator(db) migrator.register_migration(Migration(id=DEPENDENCY_ID, callback=lambda cursor: None)) migrator.register_migration(build_migration(logger=logging.getLogger(__name__))) return migrator @@ -134,7 +134,7 @@ def test_migration_is_discovered_with_stable_id_and_dependency(tmp_path: Path) - def test_actual_migrator_converts_both_tables_and_preserves_metadata_and_relations(tmp_path: Path) -> None: db = _make_db(tmp_path) - connection = db._conn + connection = db.sqlite.conn image_vector = _vector(1) video_vector = _vector(2) image_blob = image_vector.astype(np.float32).tobytes() @@ -186,13 +186,13 @@ def test_actual_migrator_converts_both_tables_and_preserves_metadata_and_relatio connection.execute("DELETE FROM videos WHERE video_name = 'video.mp4';") assert tuple(connection.execute("SELECT COUNT(*) FROM image_embeddings;").fetchone()) == (0,) assert tuple(connection.execute("SELECT COUNT(*) FROM video_embeddings;").fetchone()) == (0,) - db._conn.close() + db.dispose() def test_empty_and_mixed_rows_are_safe_and_callback_is_idempotent(tmp_path: Path) -> None: - db = SqliteDatabase(db_path=tmp_path / "mixed.db", logger=logging.getLogger(__name__)) - _create_schema(db._conn, image_encoding=True) - connection = db._conn + db = Database.open_sqlite(tmp_path / "mixed.db", logging.getLogger(__name__)) + _create_schema(db.sqlite.conn, image_encoding=True) + connection = db.sqlite.conn preserved_blob = embedding_to_blob(_vector(4)) _insert(connection, "image_embeddings", "already-fp16.png", "model-a", preserved_blob, encoding="float16") legacy_vector = _vector(5) @@ -220,13 +220,13 @@ def test_empty_and_mixed_rows_are_safe_and_callback_is_idempotent(tmp_path: Path assert after_repeat[0][2] == "float16" assert after_repeat[1][2] == "float16" assert tuple(connection.execute("SELECT COUNT(*) FROM video_embeddings;").fetchone()) == (0,) - db._conn.rollback() - db._conn.close() + db.sqlite.conn.rollback() + db.dispose() def test_keyset_batches_convert_more_than_two_batches_across_models(tmp_path: Path) -> None: db = _make_db(tmp_path) - connection = db._conn + connection = db.sqlite.conn row_count = 514 for index in range(257): vector = _vector(index + 20) @@ -261,12 +261,12 @@ def test_keyset_batches_convert_more_than_two_batches_across_models(tmp_path: Pa assert len(selects) >= 3 assert all("limit 256" in statement for statement in selects) assert all("offset" not in statement for statement in selects) - db._conn.close() + db.dispose() def test_migration_logs_progress_and_row_totals(caplog: pytest.LogCaptureFixture, tmp_path: Path) -> None: db = _make_db(tmp_path) - connection = db._conn + connection = db.sqlite.conn vector_blob = _vector(30).astype(np.float32).tobytes() names = [(f"image-{index:05d}.png",) for index in range(10_001)] connection.executemany("INSERT INTO images (image_name) VALUES (?);", names) @@ -287,12 +287,12 @@ def test_migration_logs_progress_and_row_totals(caplog: pytest.LogCaptureFixture "elapsed=" in message for message in messages ) - db._conn.close() + db.dispose() def test_late_bad_video_rolls_back_schema_and_prior_updates_then_retries(tmp_path: Path) -> None: db = _make_db(tmp_path) - connection = db._conn + connection = db.sqlite.conn image_blob = _vector(6).astype(np.float32).tobytes() video_blob = _vector(7).astype(np.float32).tobytes() for index in range(257): @@ -338,4 +338,4 @@ def test_late_bad_video_rolls_back_schema_and_prior_updates_then_retries(tmp_pat assert connection.execute( "SELECT migration_id FROM applied_migrations WHERE migration_id = ?;", (MIGRATION_ID,) ).fetchone() - db._conn.close() + db.dispose() diff --git a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_10_04_add_db_locks.py b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_10_04_add_db_locks.py new file mode 100644 index 00000000000..e19e0449efc --- /dev/null +++ b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_10_04_add_db_locks.py @@ -0,0 +1,25 @@ +from pathlib import Path + +from sqlalchemy import delete, select + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.locks import db_locks +from invokeai.app.services.shared.database.schema.migrator import applied_migrations +from tests.fixtures.database import migrate_to_newest + + +def test_a_run_that_stopped_after_creating_the_table_adds_the_lock_rows(database: Database, tmp_path: Path) -> None: + # On a server each DDL statement commits as it runs: a run that failed after creating the table leaves it + # there, empty, and the migration unrecorded. + with database.begin(write=True) as conn: + conn.execute(delete(db_locks)) + conn.execute(delete(applied_migrations).where(applied_migrations.c.migration_id == "2026_10_04_add_db_locks")) + + migrate_to_newest(database, tmp_path) + + with database.begin(write=False) as conn: + assert list(conn.execute(select(db_locks.c.name)).scalars()) == ["admin_accounts"] + recorded = select(applied_migrations.c.migration_id).where( + applied_migrations.c.migration_id == "2026_10_04_add_db_locks" + ) + assert conn.execute(recorded).first() is not None diff --git a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_10_04_add_session_queue_listing_index.py b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_10_04_add_session_queue_listing_index.py index 8b8a42cbe08..b7d80b3164e 100644 --- a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_10_04_add_session_queue_listing_index.py +++ b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_10_04_add_session_queue_listing_index.py @@ -1,22 +1,31 @@ -"""An existing queue gains the listing index without any queue read changing what it returns.""" +"""An existing queue gains the listing index without any queue read changing what it returns. + +The index exists only on SQLite; on a server the migration does nothing. +""" import uuid -from logging import Logger +from collections.abc import Iterator, Sequence +from pathlib import Path from typing import Any -from unittest.mock import MagicMock - -from invokeai.app.services.config.config_default import InvokeAIAppConfig -from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from unittest import mock + +import pytest +from sqlalchemy import delete, insert, inspect + +from invokeai.app.services.config.config_default import DefaultInvokeAIAppConfig +from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.migrator import applied_migrations +from invokeai.app.services.shared.database.schema.session_queue import session_queue as session_queue_table +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.services.shared.sqlite_migrator.migration_loader import MigrationBuildContext, build_migrations -from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2026_10_04_add_session_queue_listing_index import ( - build_migration, -) -from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration -from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import MigrationBase +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import Migrator +from invokeai.backend.util.logging import InvokeAILogger MIGRATION_ID = "2026_10_04_add_session_queue_listing_index" +INDEX = "idx_session_queue_listing" PREFIX = "webv2:p:project-1:q:" # (created_at, user_id, origin, status), inserted in this order. One enqueue writes its rows within the @@ -32,18 +41,41 @@ ] -def _run(db: SqliteDatabase, migrations: list[Migration]) -> bool: - migrator = SqliteMigrator(db=db) +@pytest.fixture +def migrations(tmp_path: Path) -> list[MigrationBase]: + # Migrations clean up legacy files under the root. No settings from the environment or a config file may + # point them anywhere else. + config = DefaultInvokeAIAppConfig() + config._root = tmp_path + context = MigrationBuildContext( + app_config=config, + logger=InvokeAILogger.get_logger("test_listing_index_migration"), + image_files=mock.Mock(spec=ImageFileStorageBase), + ) + return build_migrations(context) + + +@pytest.fixture +def sqlite_database() -> Iterator[Database]: + database = Database.open_sqlite(None, InvokeAILogger.get_logger("test_listing_index_migration")) + try: + yield database + finally: + database.dispose() + + +def _run(database: Database, migrations: Sequence[MigrationBase]) -> bool: + migrator = Migrator(database) for migration in migrations: migrator.register_migration(migration) return migrator.run_migrations() -def _previous_state(migrations: list[Migration]) -> list[Migration]: +def _previous_state(migrations: list[MigrationBase]) -> list[MigrationBase]: """Every migration except this one and those that depend on it, directly or not.""" by_id = {migration.id: migration for migration in migrations} - def depends_on_this(migration: Migration) -> bool: + def depends_on_this(migration: MigrationBase) -> bool: current: str | None = migration.id while current is not None: if current == MIGRATION_ID: @@ -54,12 +86,12 @@ def depends_on_this(migration: Migration) -> bool: return [migration for migration in migrations if not depends_on_this(migration)] -def _has_index(db: SqliteDatabase) -> bool: - row = db._conn.execute("SELECT 1 FROM sqlite_master WHERE name = 'idx_session_queue_listing';").fetchone() - return row is not None +def _has_index(database: Database) -> bool: + with database.begin(write=False) as conn: + return any(index["name"] == INDEX for index in inspect(conn).get_indexes("session_queue")) -def _reads(queue: SqliteSessionQueue) -> dict[str, Any]: +def _reads(queue: SessionQueue) -> dict[str, Any]: def ids(direction: SQLiteDirection, **filters: str) -> list[int]: return queue.get_queue_item_ids("default", direction, **filters).item_ids @@ -79,25 +111,30 @@ def ids(direction: SQLiteDirection, **filters: str) -> list[int]: } -def test_indexes_an_existing_queue_without_changing_its_reads() -> None: - logger = Logger("test") - context = MigrationBuildContext( - app_config=InvokeAIAppConfig(use_memory_db=True, node_cache_size=0), logger=logger, image_files=MagicMock() - ) - migrations = build_migrations(context) - db = SqliteDatabase(db_path=None, logger=logger) - assert _run(db, _previous_state(migrations)) - assert not _has_index(db) - - with db.transaction() as cursor: - cursor.executemany( - """--sql - INSERT INTO session_queue (queue_id, session, session_id, batch_id, created_at, user_id, origin, status) - VALUES ('default', '{}', ?, 'batch', ?, ?, ?, ?); - """, - [(str(uuid.uuid4()), *row) for row in ROWS], +def test_indexes_an_existing_queue_without_changing_its_reads( + sqlite_database: Database, migrations: list[MigrationBase] +) -> None: + assert _run(sqlite_database, _previous_state(migrations)) + assert not _has_index(sqlite_database) + + with sqlite_database.begin(write=True) as conn: + conn.execute( + insert(session_queue_table), + [ + { + "queue_id": "default", + "session": "{}", + "session_id": str(uuid.uuid4()), + "batch_id": "batch", + "created_at": created_at, + "user_id": user_id, + "origin": origin, + "status": status, + } + for created_at, user_id, origin, status in ROWS + ], ) - queue = SqliteSessionQueue(db=db) + queue = SessionQueue(sqlite_database) before = _reads(queue) # Newest first; rows enqueued together keep their enqueue order in either direction. @@ -110,10 +147,15 @@ def test_indexes_an_existing_queue_without_changing_its_reads() -> None: # In the requested order; an unknown id is skipped. assert before["summaries"] == [(3, "completed", f"{PREFIX}c"), (1, "completed", f"{PREFIX}a")] - assert _run(db, migrations) - assert _has_index(db) + assert _run(sqlite_database, migrations) + assert _has_index(sqlite_database) assert _reads(queue) == before - assert not _run(db, migrations) - build_migration().callback(db._conn.cursor()) + assert not _run(sqlite_database, migrations) + # Run again with its record lost, as a run that failed after creating the index leaves it: the migration finds + # the index and keeps it. + with sqlite_database.begin(write=True) as conn: + conn.execute(delete(applied_migrations).where(applied_migrations.c.migration_id == MIGRATION_ID)) + assert _run(sqlite_database, migrations) + assert _has_index(sqlite_database) assert _reads(queue) == before diff --git a/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_10_06_add_intermediates_state.py b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_10_06_add_intermediates_state.py new file mode 100644 index 00000000000..6e74db876b9 --- /dev/null +++ b/tests/app/services/shared/sqlite_migrator/migrations/test_migration_2026_10_06_add_intermediates_state.py @@ -0,0 +1,29 @@ +from pathlib import Path + +from sqlalchemy import delete, inspect, select, text + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.locks import db_locks +from invokeai.app.services.shared.database.schema.migrator import applied_migrations +from tests.fixtures.database import migrate_to_newest + +MIGRATION = "2026_10_06_add_intermediates_state" + + +def test_a_run_that_stopped_halfway_completes_what_it_left(database: Database, tmp_path: Path) -> None: + # On a server each DDL statement commits as it runs: a run that failed after its first steps leaves them done, + # the rest undone and the migration unrecorded. + with database.begin(write=True) as conn: + conn.execute(delete(db_locks).where(db_locks.c.name == "media_protection")) + conn.execute(delete(applied_migrations).where(applied_migrations.c.migration_id == MIGRATION)) + on_table = "" if conn.dialect.name == "sqlite" else " ON images" + conn.execute(text(f"DROP INDEX idx_images_intermediates_owner{on_table}")) + + migrate_to_newest(database, tmp_path) + + with database.begin(write=False) as conn: + assert conn.execute(select(db_locks.c.name).where(db_locks.c.name == "media_protection")).first() is not None + indexes = {index["name"]: index["column_names"] for index in inspect(conn).get_indexes("images")} + assert "idx_images_intermediates_owner" in indexes + recorded = select(applied_migrations.c.migration_id).where(applied_migrations.c.migration_id == MIGRATION) + assert conn.execute(recorded).first() is not None diff --git a/tests/app/services/shared/sqlite_migrator/server_schema_at_cutover/mariadb.sql b/tests/app/services/shared/sqlite_migrator/server_schema_at_cutover/mariadb.sql new file mode 100644 index 00000000000..ab00f931108 --- /dev/null +++ b/tests/app/services/shared/sqlite_migrator/server_schema_at_cutover/mariadb.sql @@ -0,0 +1,529 @@ +-- The schema of a mariadb database at the portable-migration cutover (2026_10_01_add_anima_variant): every table +-- as 10.11.19-MariaDB-ubu2204 rendered it (SHOW CREATE TABLE) after creating it from the schema metadata of that time. +-- Frozen: test_portable_migrations.py migrates a database created from it and compares the result with +-- one created from today's metadata. Do not edit or regenerate it. + +CREATE TABLE `app_settings` ( + `key` varchar(255) NOT NULL, + `value` longtext NOT NULL, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + PRIMARY KEY (`key`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `applied_migrations` ( + `migration_id` varchar(255) NOT NULL, + `legacy_version` bigint(20) DEFAULT NULL, + `migrated_at` varchar(32) NOT NULL, + PRIMARY KEY (`migration_id`), + UNIQUE KEY `uq_applied_migrations_legacy_version` (`legacy_version`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `image_index_vocab_terms` ( + `term` varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_uca1400_nopad_as_ci NOT NULL, + `created_at` varchar(32) NOT NULL, + PRIMARY KEY (`term`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `image_subfolder_move_jobs` ( + `id` bigint(20) NOT NULL AUTO_INCREMENT, + `state` longtext NOT NULL, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `error_message` longtext DEFAULT NULL, + PRIMARY KEY (`id`), + CONSTRAINT `ck_image_subfolder_move_jobs_state` CHECK (`state` in ('planned','moving','moved','committed','error')) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `images` ( + `image_name` varchar(255) NOT NULL, + `image_origin` varchar(32) NOT NULL, + `image_category` varchar(32) NOT NULL, + `width` bigint(20) NOT NULL, + `height` bigint(20) NOT NULL, + `session_id` longtext DEFAULT NULL, + `node_id` longtext DEFAULT NULL, + `metadata` longtext DEFAULT NULL, + `is_intermediate` tinyint(1) DEFAULT 0, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `deleted_at` varchar(32) DEFAULT NULL, + `starred` tinyint(1) DEFAULT 0, + `has_workflow` tinyint(1) DEFAULT 0, + `user_id` varchar(64) DEFAULT 'system', + `image_subfolder` longtext NOT NULL DEFAULT '', + `project_id` varchar(255) DEFAULT NULL, + `file_size_bytes` bigint(20) DEFAULT NULL, + PRIMARY KEY (`image_name`), + KEY `idx_images_created_at` (`created_at`), + KEY `idx_images_unmeasured_intermediates` (`is_intermediate`,`file_size_bytes`,`created_at`), + KEY `idx_images_image_origin` (`image_origin`), + KEY `idx_images_user_id` (`user_id`), + KEY `idx_images_image_category` (`image_category`), + KEY `idx_images_intermediate_scope` (`is_intermediate`,`user_id`,`project_id`), + KEY `idx_images_starred` (`starred`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `intermediates_browser_holds` ( + `user_id` varchar(64) NOT NULL, + `lease_id` varchar(255) NOT NULL, + `media_kind` varchar(32) NOT NULL, + `media_name` varchar(255) NOT NULL, + `expires_at` varchar(32) NOT NULL, + PRIMARY KEY (`user_id`,`lease_id`,`media_kind`,`media_name`), + KEY `idx_intermediates_browser_holds_expires_at` (`expires_at`), + KEY `idx_intermediates_browser_holds_media` (`media_kind`,`media_name`,`expires_at`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `media_references` ( + `owner_kind` varchar(32) NOT NULL, + `user_id` varchar(64) NOT NULL, + `owner_id` varchar(255) NOT NULL, + `media_kind` varchar(32) NOT NULL, + `media_name` varchar(255) NOT NULL, + PRIMARY KEY (`owner_kind`,`user_id`,`owner_id`,`media_kind`,`media_name`), + KEY `idx_media_references_media` (`media_kind`,`media_name`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `migrations` ( + `version` bigint(20) NOT NULL, + `migrated_at` varchar(32) NOT NULL, + PRIMARY KEY (`version`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `models` ( + `id` varchar(255) NOT NULL, + `hash` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`config`,'$.hash')) = 'NULL' then NULL else json_unquote(json_extract(`config`,'$.hash')) end) VIRTUAL, + `base` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`config`,'$.base')) = 'NULL' then NULL else json_unquote(json_extract(`config`,'$.base')) end) VIRTUAL, + `type` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`config`,'$.type')) = 'NULL' then NULL else json_unquote(json_extract(`config`,'$.type')) end) VIRTUAL, + `path` varchar(768) GENERATED ALWAYS AS (case when json_type(json_extract(`config`,'$.path')) = 'NULL' then NULL else json_unquote(json_extract(`config`,'$.path')) end) VIRTUAL, + `format` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`config`,'$.format')) = 'NULL' then NULL else json_unquote(json_extract(`config`,'$.format')) end) VIRTUAL, + `name` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`config`,'$.name')) = 'NULL' then NULL else json_unquote(json_extract(`config`,'$.name')) end) VIRTUAL, + `description` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`config`,'$.description')) = 'NULL' then NULL else json_unquote(json_extract(`config`,'$.description')) end) VIRTUAL, + `source` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`config`,'$.source')) = 'NULL' then NULL else json_unquote(json_extract(`config`,'$.source')) end) VIRTUAL, + `source_type` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`config`,'$.source_type')) = 'NULL' then NULL else json_unquote(json_extract(`config`,'$.source_type')) end) VIRTUAL, + `source_api_response` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`config`,'$.source_api_response')) = 'NULL' then NULL else json_unquote(json_extract(`config`,'$.source_api_response')) end) VIRTUAL, + `trigger_phrases` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`config`,'$.trigger_phrases')) = 'NULL' then NULL else json_unquote(json_extract(`config`,'$.trigger_phrases')) end) VIRTUAL, + `file_size` bigint(20) GENERATED ALWAYS AS (cast(case when json_type(json_extract(`config`,'$.file_size')) = 'NULL' then NULL else json_unquote(json_extract(`config`,'$.file_size')) end as signed)) VIRTUAL, + `config` longtext NOT NULL, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY `uq_models_path` (`path`), + KEY `type_index` (`type`(255)), + KEY `base_index` (`base`(255)), + KEY `name_index` (`name`(255)), + CONSTRAINT `ck_models_hash_not_null` CHECK (`hash` is not null), + CONSTRAINT `ck_models_base_not_null` CHECK (`base` is not null), + CONSTRAINT `ck_models_type_not_null` CHECK (`type` is not null), + CONSTRAINT `ck_models_path_not_null` CHECK (`path` is not null), + CONSTRAINT `ck_models_format_not_null` CHECK (`format` is not null), + CONSTRAINT `ck_models_name_not_null` CHECK (`name` is not null), + CONSTRAINT `ck_models_source_not_null` CHECK (`source` is not null), + CONSTRAINT `ck_models_source_type_not_null` CHECK (`source_type` is not null), + CONSTRAINT `ck_models_file_size_not_null` CHECK (`file_size` is not null) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `orphaned_projects_2026_08_06` ( + `project_id` varchar(255) NOT NULL, + `user_id` varchar(64) NOT NULL, + `name` longtext NOT NULL, + `data` longtext NOT NULL, + `revision` bigint(20) NOT NULL, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `quarantined_at` varchar(32) NOT NULL, + PRIMARY KEY (`user_id`,`project_id`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `session_queue` ( + `item_id` bigint(20) NOT NULL AUTO_INCREMENT, + `batch_id` varchar(255) NOT NULL, + `queue_id` longtext NOT NULL, + `session_id` varchar(255) NOT NULL, + `field_values` longtext DEFAULT NULL, + `session` longtext NOT NULL, + `status` varchar(32) NOT NULL DEFAULT 'pending', + `priority` bigint(20) NOT NULL DEFAULT 0, + `error_traceback` longtext DEFAULT NULL, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `started_at` varchar(32) DEFAULT NULL, + `completed_at` varchar(32) DEFAULT NULL, + `workflow` longtext DEFAULT NULL, + `error_type` longtext DEFAULT NULL, + `error_message` longtext DEFAULT NULL, + `origin` longtext DEFAULT NULL, + `destination` longtext DEFAULT NULL, + `retried_from_item_id` bigint(20) DEFAULT NULL, + `user_id` varchar(64) DEFAULT 'system', + `status_sequence` bigint(20) DEFAULT 0, + `device` longtext DEFAULT NULL, + `workflow_call_id` varchar(255) DEFAULT NULL, + `parent_item_id` bigint(20) DEFAULT NULL, + `parent_session_id` varchar(255) DEFAULT NULL, + `root_item_id` bigint(20) DEFAULT NULL, + `workflow_call_depth` bigint(20) DEFAULT NULL, + `project_id` longtext DEFAULT NULL, + `session_revision` bigint(20) NOT NULL DEFAULT 0, + PRIMARY KEY (`item_id`), + UNIQUE KEY `uq_session_queue_session_id` (`session_id`), + KEY `idx_session_queue_parent_session_id` (`parent_session_id`), + KEY `idx_session_queue_batch_id` (`batch_id`), + KEY `idx_session_queue_workflow_call_depth` (`workflow_call_depth`), + KEY `idx_session_queue_root_item_id` (`root_item_id`), + KEY `idx_session_queue_parent_item_id` (`parent_item_id`), + KEY `idx_session_queue_workflow_call_id` (`workflow_call_id`), + KEY `idx_session_queue_round_robin_pending` (`status`,`user_id`,`priority` DESC,`item_id`), + KEY `idx_session_queue_created_priority` (`priority`), + KEY `idx_session_queue_user_started_at` (`user_id`,`started_at`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `style_presets` ( + `id` varchar(255) NOT NULL, + `name` longtext NOT NULL, + `preset_data` longtext NOT NULL, + `type` longtext NOT NULL DEFAULT 'user', + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `user_id` varchar(64) DEFAULT 'system', + `is_public` tinyint(1) NOT NULL DEFAULT 0, + PRIMARY KEY (`id`), + KEY `idx_style_presets_user_id` (`user_id`), + KEY `idx_style_presets_is_public` (`is_public`), + KEY `idx_style_presets_name` (`name`(255)) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `system_prompts` ( + `id` varchar(255) NOT NULL, + `name` longtext NOT NULL, + `content` longtext NOT NULL, + `user_id` varchar(64) NOT NULL DEFAULT 'system', + `is_public` tinyint(1) NOT NULL DEFAULT 0, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `max_tokens` bigint(20) DEFAULT NULL, + PRIMARY KEY (`id`), + KEY `idx_system_prompts_user_id` (`user_id`), + KEY `idx_system_prompts_name` (`name`(255)) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `users` ( + `user_id` varchar(64) NOT NULL, + `email` varchar(255) NOT NULL, + `display_name` longtext DEFAULT NULL, + `password_hash` longtext NOT NULL, + `is_admin` tinyint(1) NOT NULL DEFAULT 0, + `is_active` tinyint(1) NOT NULL DEFAULT 1, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `last_login_at` varchar(32) DEFAULT NULL, + `token_epoch` bigint(20) NOT NULL DEFAULT 0, + PRIMARY KEY (`user_id`), + UNIQUE KEY `uq_users_email` (`email`), + KEY `idx_users_is_admin` (`is_admin`), + KEY `idx_users_is_active` (`is_active`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `videos` ( + `video_name` varchar(255) NOT NULL, + `video_origin` varchar(32) NOT NULL, + `video_category` varchar(32) NOT NULL, + `width` bigint(20) NOT NULL, + `height` bigint(20) NOT NULL, + `duration` double NOT NULL DEFAULT 0, + `fps` double DEFAULT NULL, + `session_id` longtext DEFAULT NULL, + `node_id` longtext DEFAULT NULL, + `metadata` longtext DEFAULT NULL, + `is_intermediate` tinyint(1) DEFAULT 0, + `starred` tinyint(1) DEFAULT 0, + `has_workflow` tinyint(1) DEFAULT 0, + `user_id` varchar(64) NOT NULL DEFAULT 'system', + `video_subfolder` longtext NOT NULL DEFAULT '', + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `deleted_at` varchar(32) DEFAULT NULL, + `project_id` varchar(255) DEFAULT NULL, + `file_size_bytes` bigint(20) DEFAULT NULL, + PRIMARY KEY (`video_name`), + KEY `idx_videos_video_category` (`video_category`), + KEY `idx_videos_video_origin` (`video_origin`), + KEY `idx_videos_intermediate_scope` (`is_intermediate`,`user_id`,`project_id`), + KEY `idx_videos_starred` (`starred`), + KEY `idx_videos_unmeasured_intermediates` (`is_intermediate`,`file_size_bytes`,`created_at`), + KEY `idx_videos_user_id` (`user_id`), + KEY `idx_videos_created_at` (`created_at`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `workflow_library` ( + `workflow_id` varchar(255) NOT NULL, + `workflow` longtext NOT NULL, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `category` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`workflow`,'$.meta.category')) = 'NULL' then NULL else json_unquote(json_extract(`workflow`,'$.meta.category')) end) VIRTUAL, + `name` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`workflow`,'$.name')) = 'NULL' then NULL else json_unquote(json_extract(`workflow`,'$.name')) end) VIRTUAL, + `description` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`workflow`,'$.description')) = 'NULL' then NULL else json_unquote(json_extract(`workflow`,'$.description')) end) VIRTUAL, + `tags` longtext GENERATED ALWAYS AS (case when json_type(json_extract(`workflow`,'$.tags')) = 'NULL' then NULL else json_unquote(json_extract(`workflow`,'$.tags')) end) VIRTUAL, + `opened_at` varchar(32) DEFAULT NULL, + `user_id` varchar(64) DEFAULT 'system', + `is_public` tinyint(1) NOT NULL DEFAULT 0, + `last_run_at` varchar(32) DEFAULT NULL, + `revision` bigint(20) NOT NULL DEFAULT 1, + PRIMARY KEY (`workflow_id`), + KEY `idx_workflow_library_is_public` (`is_public`), + KEY `idx_workflow_library_category` (`category`(255)), + KEY `idx_workflow_library_description` (`description`(255)), + KEY `idx_workflow_library_opened_at` (`opened_at`), + KEY `idx_workflow_library_created_at` (`created_at`), + KEY `idx_workflow_library_updated_at` (`updated_at`), + KEY `idx_workflow_library_name` (`name`(255)), + KEY `idx_workflow_library_user_id` (`user_id`), + CONSTRAINT `ck_workflow_library_category_not_null` CHECK (`category` is not null), + CONSTRAINT `ck_workflow_library_name_not_null` CHECK (`name` is not null), + CONSTRAINT `ck_workflow_library_description_not_null` CHECK (`description` is not null) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `boards` ( + `board_id` varchar(255) NOT NULL, + `board_name` longtext NOT NULL, + `cover_image_name` varchar(255) DEFAULT NULL, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `deleted_at` varchar(32) DEFAULT NULL, + `archived` tinyint(1) DEFAULT 0, + `user_id` varchar(64) DEFAULT 'system', + `is_public` tinyint(1) NOT NULL DEFAULT 0, + `board_visibility` varchar(32) NOT NULL DEFAULT 'private', + PRIMARY KEY (`board_id`), + KEY `fk_boards_cover_image_name_images` (`cover_image_name`), + KEY `idx_boards_board_visibility` (`board_visibility`), + KEY `idx_boards_created_at` (`created_at`), + KEY `idx_boards_user_id` (`user_id`), + KEY `idx_boards_is_public` (`is_public`), + CONSTRAINT `fk_boards_cover_image_name_images` FOREIGN KEY (`cover_image_name`) REFERENCES `images` (`image_name`) ON DELETE SET NULL +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `client_state` ( + `user_id` varchar(64) NOT NULL, + `key` varchar(255) NOT NULL, + `value` longtext NOT NULL, + `updated_at` varchar(32) NOT NULL, + PRIMARY KEY (`user_id`,`key`), + CONSTRAINT `fk_client_state_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `fonts` ( + `id` varchar(255) NOT NULL, + `owner_id` varchar(64) DEFAULT NULL, + `scope` varchar(32) NOT NULL, + `source` varchar(32) NOT NULL, + `filename` longtext NOT NULL, + `storage_path` longtext DEFAULT NULL, + `source_path` varchar(768) DEFAULT NULL, + `family` longtext NOT NULL, + `label` longtext NOT NULL, + `style` longtext NOT NULL, + `weight` bigint(20) NOT NULL, + `content_hash` varchar(64) NOT NULL, + `byte_size` bigint(20) NOT NULL, + `axes_json` longtext NOT NULL DEFAULT '[]', + `instances_json` longtext NOT NULL DEFAULT '[]', + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY `idx_fonts_directory_source_path` (`source_path`), + KEY `fk_fonts_owner_id_users` (`owner_id`), + KEY `idx_fonts_uploaded_hash` (`source`,`scope`,`owner_id`,`content_hash`), + CONSTRAINT `fk_fonts_owner_id_users` FOREIGN KEY (`owner_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE, + CONSTRAINT `ck_fonts_scope` CHECK (`scope` in ('private','shared')), + CONSTRAINT `ck_fonts_source` CHECK (`source` in ('uploaded','directory')), + CONSTRAINT `ck_fonts_weight` CHECK (`weight` >= 1 and `weight` <= 1000), + CONSTRAINT `ck_fonts_content_hash` CHECK (octet_length(`content_hash`) = 64), + CONSTRAINT `ck_fonts_byte_size` CHECK (`byte_size` > 0), + CONSTRAINT `ck_fonts_source_fields` CHECK (`source` = 'directory' and `owner_id` is null and `scope` = 'shared' and `storage_path` is null and `source_path` is not null or `source` = 'uploaded' and `storage_path` is not null and `source_path` is null and (`scope` = 'private' and `owner_id` is not null or `scope` = 'shared' and `owner_id` is null)) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `image_embeddings` ( + `image_name` varchar(255) NOT NULL, + `model_id` varchar(255) NOT NULL, + `dim` bigint(20) NOT NULL, + `embedding` longblob NOT NULL, + `created_at` varchar(32) NOT NULL, + `encoding` longtext NOT NULL DEFAULT 'float32', + PRIMARY KEY (`image_name`,`model_id`), + KEY `idx_image_embeddings_model_id` (`model_id`), + CONSTRAINT `fk_image_embeddings_image_name_images` FOREIGN KEY (`image_name`) REFERENCES `images` (`image_name`) ON DELETE CASCADE, + CONSTRAINT `ck_image_embeddings_encoding` CHECK (`encoding` in ('float32','float16')) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `image_projections` ( + `user_id` varchar(64) NOT NULL, + `model_id` varchar(255) NOT NULL, + `scope_hash` longtext NOT NULL, + `params` longtext NOT NULL, + `point_count` bigint(20) NOT NULL, + `image_names` longtext NOT NULL, + `coords` longblob NOT NULL, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `item_kinds` longtext DEFAULT NULL, + PRIMARY KEY (`user_id`,`model_id`), + CONSTRAINT `fk_image_projections_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `image_subfolder_move_items` ( + `job_id` bigint(20) NOT NULL, + `image_name` varchar(255) NOT NULL, + `old_subfolder` longtext NOT NULL, + `new_subfolder` longtext NOT NULL, + `is_intermediate` tinyint(1) NOT NULL DEFAULT 0, + `old_path` longtext DEFAULT NULL, + `new_path` longtext DEFAULT NULL, + `old_thumbnail_path` longtext DEFAULT NULL, + `new_thumbnail_path` longtext DEFAULT NULL, + `state` varchar(32) NOT NULL, + `error_message` longtext DEFAULT NULL, + PRIMARY KEY (`job_id`,`image_name`), + KEY `idx_image_subfolder_move_items_image_name` (`image_name`), + KEY `idx_image_subfolder_move_items_job_state` (`job_id`,`state`), + CONSTRAINT `fk_image_subfolder_move_items_image_name_images` FOREIGN KEY (`image_name`) REFERENCES `images` (`image_name`) ON DELETE CASCADE, + CONSTRAINT `fk_image_subfolder_move_items_job_id_image_subfolder_move_jobs` FOREIGN KEY (`job_id`) REFERENCES `image_subfolder_move_jobs` (`id`), + CONSTRAINT `ck_image_subfolder_move_items_state` CHECK (`state` in ('planned','moved','committed','error')) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `model_relationships` ( + `model_key_1` varchar(255) NOT NULL, + `model_key_2` varchar(255) NOT NULL, + `created_at` varchar(32) NOT NULL, + PRIMARY KEY (`model_key_1`,`model_key_2`), + KEY `keyx_model_relationships_model_key_2` (`model_key_2`), + CONSTRAINT `fk_model_relationships_model_key_1_models` FOREIGN KEY (`model_key_1`) REFERENCES `models` (`id`) ON DELETE CASCADE, + CONSTRAINT `fk_model_relationships_model_key_2_models` FOREIGN KEY (`model_key_2`) REFERENCES `models` (`id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `session_queue_enqueue_receipts` ( + `queue_id` varchar(255) NOT NULL, + `user_id` varchar(64) NOT NULL, + `idempotency_key` varchar(255) NOT NULL, + `payload_hash` longtext NOT NULL, + `batch_id` longtext NOT NULL, + `requested` bigint(20) NOT NULL, + `enqueued` bigint(20) NOT NULL, + `priority` bigint(20) NOT NULL, + `item_ids` longtext NOT NULL, + `byte_size` bigint(20) NOT NULL, + `created_at` varchar(32) NOT NULL, + `acknowledged_at` varchar(32) DEFAULT NULL, + PRIMARY KEY (`queue_id`,`user_id`,`idempotency_key`), + KEY `idx_session_queue_enqueue_receipts_owner_ack` (`user_id`,`acknowledged_at`,`byte_size`), + CONSTRAINT `fk_session_queue_enqueue_receipts_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `user_invitations` ( + `invitation_id` varchar(255) NOT NULL, + `email` varchar(255) NOT NULL, + `invited_by` varchar(64) NOT NULL, + `invitation_code` varchar(255) NOT NULL, + `is_admin` tinyint(1) NOT NULL DEFAULT 0, + `expires_at` varchar(32) NOT NULL, + `used_at` varchar(32) DEFAULT NULL, + `created_at` varchar(32) NOT NULL, + PRIMARY KEY (`invitation_id`), + UNIQUE KEY `uq_user_invitations_invitation_code` (`invitation_code`), + KEY `fk_user_invitations_invited_by_users` (`invited_by`), + KEY `idx_user_invitations_expires_at` (`expires_at`), + KEY `idx_user_invitations_email` (`email`), + CONSTRAINT `fk_user_invitations_invited_by_users` FOREIGN KEY (`invited_by`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `user_sessions` ( + `session_id` varchar(255) NOT NULL, + `user_id` varchar(64) NOT NULL, + `token_hash` varchar(255) NOT NULL, + `expires_at` varchar(32) NOT NULL, + `created_at` varchar(32) NOT NULL, + `last_activity_at` varchar(32) NOT NULL, + PRIMARY KEY (`session_id`), + KEY `idx_user_sessions_token_hash` (`token_hash`), + KEY `idx_user_sessions_expires_at` (`expires_at`), + KEY `idx_user_sessions_user_id` (`user_id`), + CONSTRAINT `fk_user_sessions_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `video_embeddings` ( + `video_name` varchar(255) NOT NULL, + `model_id` varchar(255) NOT NULL, + `dim` bigint(20) NOT NULL, + `embedding` longblob NOT NULL, + `created_at` varchar(32) NOT NULL, + `encoding` longtext NOT NULL DEFAULT 'float32', + PRIMARY KEY (`video_name`,`model_id`), + KEY `idx_video_embeddings_model_id` (`model_id`), + CONSTRAINT `fk_video_embeddings_video_name_videos` FOREIGN KEY (`video_name`) REFERENCES `videos` (`video_name`) ON DELETE CASCADE, + CONSTRAINT `ck_video_embeddings_encoding` CHECK (`encoding` in ('float32','float16')) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `wildcards` ( + `id` varchar(255) NOT NULL, + `name` varchar(255) NOT NULL, + `values_json` longtext NOT NULL DEFAULT '[]', + `user_id` varchar(64) NOT NULL, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY `idx_wildcards_user_id_name` (`user_id`,`name`), + CONSTRAINT `fk_wildcards_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `board_images` ( + `board_id` varchar(255) NOT NULL, + `image_name` varchar(255) NOT NULL, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `deleted_at` varchar(32) DEFAULT NULL, + PRIMARY KEY (`image_name`), + KEY `idx_board_images_board_id_created_at` (`board_id`,`created_at`), + CONSTRAINT `fk_board_images_board_id_boards` FOREIGN KEY (`board_id`) REFERENCES `boards` (`board_id`) ON DELETE CASCADE, + CONSTRAINT `fk_board_images_image_name_images` FOREIGN KEY (`image_name`) REFERENCES `images` (`image_name`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `board_videos` ( + `board_id` varchar(255) NOT NULL, + `video_name` varchar(255) NOT NULL, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `deleted_at` varchar(32) DEFAULT NULL, + PRIMARY KEY (`video_name`), + KEY `idx_board_videos_board_id_created_at` (`board_id`,`created_at`), + CONSTRAINT `fk_board_videos_board_id_boards` FOREIGN KEY (`board_id`) REFERENCES `boards` (`board_id`) ON DELETE CASCADE, + CONSTRAINT `fk_board_videos_video_name_videos` FOREIGN KEY (`video_name`) REFERENCES `videos` (`video_name`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `projects` ( + `project_id` varchar(255) NOT NULL, + `user_id` varchar(64) NOT NULL, + `name` longtext NOT NULL, + `data` longtext NOT NULL, + `board_id` varchar(255) NOT NULL, + `revision` bigint(20) NOT NULL DEFAULT 1, + `created_at` varchar(32) NOT NULL, + `updated_at` varchar(32) NOT NULL, + `minimum_canvas_schema_version` bigint(20) NOT NULL DEFAULT 2, + PRIMARY KEY (`user_id`,`project_id`), + UNIQUE KEY `uq_projects_board_id` (`board_id`), + CONSTRAINT `fk_projects_board_id_boards` FOREIGN KEY (`board_id`) REFERENCES `boards` (`board_id`), + CONSTRAINT `fk_projects_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE, + CONSTRAINT `ck_projects_minimum_canvas_schema_version` CHECK (`minimum_canvas_schema_version` >= 1) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `shared_boards` ( + `board_id` varchar(255) NOT NULL, + `user_id` varchar(64) NOT NULL, + `can_edit` tinyint(1) NOT NULL DEFAULT 0, + `shared_at` varchar(32) NOT NULL, + PRIMARY KEY (`board_id`,`user_id`), + KEY `idx_shared_boards_user_id` (`user_id`), + CONSTRAINT `fk_shared_boards_board_id_boards` FOREIGN KEY (`board_id`) REFERENCES `boards` (`board_id`) ON DELETE CASCADE, + CONSTRAINT `fk_shared_boards_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_nopad_bin ROW_FORMAT=DYNAMIC; diff --git a/tests/app/services/shared/sqlite_migrator/server_schema_at_cutover/mysql.sql b/tests/app/services/shared/sqlite_migrator/server_schema_at_cutover/mysql.sql new file mode 100644 index 00000000000..31616972bab --- /dev/null +++ b/tests/app/services/shared/sqlite_migrator/server_schema_at_cutover/mysql.sql @@ -0,0 +1,517 @@ +-- The schema of a mysql database at the portable-migration cutover (2026_10_01_add_anima_variant): every table +-- as 8.4.11 rendered it (SHOW CREATE TABLE) after creating it from the schema metadata of that time. +-- Frozen: test_portable_migrations.py migrates a database created from it and compares the result with +-- one created from today's metadata. Do not edit or regenerate it. + +CREATE TABLE `app_settings` ( + `key` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `value` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`key`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `applied_migrations` ( + `migration_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `legacy_version` bigint DEFAULT NULL, + `migrated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`migration_id`), + UNIQUE KEY `uq_applied_migrations_legacy_version` (`legacy_version`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `image_index_vocab_terms` ( + `term` varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_as_ci NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`term`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `image_subfolder_move_jobs` ( + `id` bigint NOT NULL AUTO_INCREMENT, + `state` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `error_message` longtext COLLATE utf8mb4_0900_bin, + PRIMARY KEY (`id`), + CONSTRAINT `ck_image_subfolder_move_jobs_state` CHECK ((`state` in (_utf8mb4'planned',_utf8mb4'moving',_utf8mb4'moved',_utf8mb4'committed',_utf8mb4'error'))) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `images` ( + `image_name` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `image_origin` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `image_category` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `width` bigint NOT NULL, + `height` bigint NOT NULL, + `session_id` longtext COLLATE utf8mb4_0900_bin, + `node_id` longtext COLLATE utf8mb4_0900_bin, + `metadata` longtext COLLATE utf8mb4_0900_bin, + `is_intermediate` tinyint(1) DEFAULT '0', + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `deleted_at` varchar(32) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `starred` tinyint(1) DEFAULT '0', + `has_workflow` tinyint(1) DEFAULT '0', + `user_id` varchar(64) COLLATE utf8mb4_0900_bin DEFAULT (_utf8mb4'system'), + `image_subfolder` longtext COLLATE utf8mb4_0900_bin NOT NULL DEFAULT (_utf8mb4''), + `project_id` varchar(255) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `file_size_bytes` bigint DEFAULT NULL, + PRIMARY KEY (`image_name`), + KEY `idx_images_created_at` (`created_at`), + KEY `idx_images_unmeasured_intermediates` (`is_intermediate`,`file_size_bytes`,`created_at`), + KEY `idx_images_image_origin` (`image_origin`), + KEY `idx_images_user_id` (`user_id`), + KEY `idx_images_image_category` (`image_category`), + KEY `idx_images_intermediate_scope` (`is_intermediate`,`user_id`,`project_id`), + KEY `idx_images_starred` (`starred`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `intermediates_browser_holds` ( + `user_id` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL, + `lease_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `media_kind` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `media_name` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `expires_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`user_id`,`lease_id`,`media_kind`,`media_name`), + KEY `idx_intermediates_browser_holds_expires_at` (`expires_at`), + KEY `idx_intermediates_browser_holds_media` (`media_kind`,`media_name`,`expires_at`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `media_references` ( + `owner_kind` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `user_id` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL, + `owner_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `media_kind` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `media_name` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`owner_kind`,`user_id`,`owner_id`,`media_kind`,`media_name`), + KEY `idx_media_references_media` (`media_kind`,`media_name`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `migrations` ( + `version` bigint NOT NULL, + `migrated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`version`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `models` ( + `id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `hash` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`config`,_utf8mb4'$.hash')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`config`,_utf8mb4'$.hash')) end)) VIRTUAL NOT NULL, + `base` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`config`,_utf8mb4'$.base')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`config`,_utf8mb4'$.base')) end)) VIRTUAL NOT NULL, + `type` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`config`,_utf8mb4'$.type')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`config`,_utf8mb4'$.type')) end)) VIRTUAL NOT NULL, + `path` varchar(768) COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`config`,_utf8mb4'$.path')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`config`,_utf8mb4'$.path')) end)) VIRTUAL NOT NULL, + `format` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`config`,_utf8mb4'$.format')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`config`,_utf8mb4'$.format')) end)) VIRTUAL NOT NULL, + `name` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`config`,_utf8mb4'$.name')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`config`,_utf8mb4'$.name')) end)) VIRTUAL NOT NULL, + `description` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`config`,_utf8mb4'$.description')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`config`,_utf8mb4'$.description')) end)) VIRTUAL, + `source` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`config`,_utf8mb4'$.source')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`config`,_utf8mb4'$.source')) end)) VIRTUAL NOT NULL, + `source_type` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`config`,_utf8mb4'$.source_type')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`config`,_utf8mb4'$.source_type')) end)) VIRTUAL NOT NULL, + `source_api_response` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`config`,_utf8mb4'$.source_api_response')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`config`,_utf8mb4'$.source_api_response')) end)) VIRTUAL, + `trigger_phrases` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`config`,_utf8mb4'$.trigger_phrases')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`config`,_utf8mb4'$.trigger_phrases')) end)) VIRTUAL, + `file_size` bigint GENERATED ALWAYS AS (cast((case when (json_type(json_extract(`config`,_utf8mb4'$.file_size')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`config`,_utf8mb4'$.file_size')) end) as signed)) VIRTUAL NOT NULL, + `config` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY `uq_models_path` (`path`), + KEY `type_index` (`type`(255)), + KEY `base_index` (`base`(255)), + KEY `name_index` (`name`(255)) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `orphaned_projects_2026_08_06` ( + `project_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `user_id` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL, + `name` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `data` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `revision` bigint NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `quarantined_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`user_id`,`project_id`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `session_queue` ( + `item_id` bigint NOT NULL AUTO_INCREMENT, + `batch_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `queue_id` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `session_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `field_values` longtext COLLATE utf8mb4_0900_bin, + `session` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `status` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL DEFAULT (_utf8mb4'pending'), + `priority` bigint NOT NULL DEFAULT '0', + `error_traceback` longtext COLLATE utf8mb4_0900_bin, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `started_at` varchar(32) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `completed_at` varchar(32) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `workflow` longtext COLLATE utf8mb4_0900_bin, + `error_type` longtext COLLATE utf8mb4_0900_bin, + `error_message` longtext COLLATE utf8mb4_0900_bin, + `origin` longtext COLLATE utf8mb4_0900_bin, + `destination` longtext COLLATE utf8mb4_0900_bin, + `retried_from_item_id` bigint DEFAULT NULL, + `user_id` varchar(64) COLLATE utf8mb4_0900_bin DEFAULT (_utf8mb4'system'), + `status_sequence` bigint DEFAULT '0', + `device` longtext COLLATE utf8mb4_0900_bin, + `workflow_call_id` varchar(255) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `parent_item_id` bigint DEFAULT NULL, + `parent_session_id` varchar(255) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `root_item_id` bigint DEFAULT NULL, + `workflow_call_depth` bigint DEFAULT NULL, + `project_id` longtext COLLATE utf8mb4_0900_bin, + `session_revision` bigint NOT NULL DEFAULT '0', + PRIMARY KEY (`item_id`), + UNIQUE KEY `uq_session_queue_session_id` (`session_id`), + KEY `idx_session_queue_parent_session_id` (`parent_session_id`), + KEY `idx_session_queue_batch_id` (`batch_id`), + KEY `idx_session_queue_workflow_call_depth` (`workflow_call_depth`), + KEY `idx_session_queue_root_item_id` (`root_item_id`), + KEY `idx_session_queue_parent_item_id` (`parent_item_id`), + KEY `idx_session_queue_workflow_call_id` (`workflow_call_id`), + KEY `idx_session_queue_round_robin_pending` (`status`,`user_id`,`priority` DESC,`item_id`), + KEY `idx_session_queue_created_priority` (`priority`), + KEY `idx_session_queue_user_started_at` (`user_id`,`started_at`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `style_presets` ( + `id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `name` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `preset_data` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `type` longtext COLLATE utf8mb4_0900_bin NOT NULL DEFAULT (_utf8mb4'user'), + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `user_id` varchar(64) COLLATE utf8mb4_0900_bin DEFAULT (_utf8mb4'system'), + `is_public` tinyint(1) NOT NULL DEFAULT '0', + PRIMARY KEY (`id`), + KEY `idx_style_presets_user_id` (`user_id`), + KEY `idx_style_presets_is_public` (`is_public`), + KEY `idx_style_presets_name` (`name`(255)) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `system_prompts` ( + `id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `name` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `content` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `user_id` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL DEFAULT (_utf8mb4'system'), + `is_public` tinyint(1) NOT NULL DEFAULT '0', + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `max_tokens` bigint DEFAULT NULL, + PRIMARY KEY (`id`), + KEY `idx_system_prompts_user_id` (`user_id`), + KEY `idx_system_prompts_name` (`name`(255)) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `users` ( + `user_id` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL, + `email` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `display_name` longtext COLLATE utf8mb4_0900_bin, + `password_hash` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `is_admin` tinyint(1) NOT NULL DEFAULT '0', + `is_active` tinyint(1) NOT NULL DEFAULT '1', + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `last_login_at` varchar(32) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `token_epoch` bigint NOT NULL DEFAULT '0', + PRIMARY KEY (`user_id`), + UNIQUE KEY `uq_users_email` (`email`), + KEY `idx_users_is_admin` (`is_admin`), + KEY `idx_users_is_active` (`is_active`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `videos` ( + `video_name` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `video_origin` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `video_category` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `width` bigint NOT NULL, + `height` bigint NOT NULL, + `duration` double NOT NULL DEFAULT (0.0), + `fps` double DEFAULT NULL, + `session_id` longtext COLLATE utf8mb4_0900_bin, + `node_id` longtext COLLATE utf8mb4_0900_bin, + `metadata` longtext COLLATE utf8mb4_0900_bin, + `is_intermediate` tinyint(1) DEFAULT '0', + `starred` tinyint(1) DEFAULT '0', + `has_workflow` tinyint(1) DEFAULT '0', + `user_id` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL DEFAULT (_utf8mb4'system'), + `video_subfolder` longtext COLLATE utf8mb4_0900_bin NOT NULL DEFAULT (_utf8mb4''), + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `deleted_at` varchar(32) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `project_id` varchar(255) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `file_size_bytes` bigint DEFAULT NULL, + PRIMARY KEY (`video_name`), + KEY `idx_videos_video_category` (`video_category`), + KEY `idx_videos_video_origin` (`video_origin`), + KEY `idx_videos_intermediate_scope` (`is_intermediate`,`user_id`,`project_id`), + KEY `idx_videos_starred` (`starred`), + KEY `idx_videos_unmeasured_intermediates` (`is_intermediate`,`file_size_bytes`,`created_at`), + KEY `idx_videos_user_id` (`user_id`), + KEY `idx_videos_created_at` (`created_at`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `workflow_library` ( + `workflow_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `workflow` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `category` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`workflow`,_utf8mb4'$.meta.category')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`workflow`,_utf8mb4'$.meta.category')) end)) VIRTUAL NOT NULL, + `name` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`workflow`,_utf8mb4'$.name')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`workflow`,_utf8mb4'$.name')) end)) VIRTUAL NOT NULL, + `description` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`workflow`,_utf8mb4'$.description')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`workflow`,_utf8mb4'$.description')) end)) VIRTUAL NOT NULL, + `tags` longtext COLLATE utf8mb4_0900_bin GENERATED ALWAYS AS ((case when (json_type(json_extract(`workflow`,_utf8mb4'$.tags')) = _utf8mb4'NULL') then NULL else json_unquote(json_extract(`workflow`,_utf8mb4'$.tags')) end)) VIRTUAL, + `opened_at` varchar(32) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `user_id` varchar(64) COLLATE utf8mb4_0900_bin DEFAULT (_utf8mb4'system'), + `is_public` tinyint(1) NOT NULL DEFAULT '0', + `last_run_at` varchar(32) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `revision` bigint NOT NULL DEFAULT '1', + PRIMARY KEY (`workflow_id`), + KEY `idx_workflow_library_is_public` (`is_public`), + KEY `idx_workflow_library_opened_at` (`opened_at`), + KEY `idx_workflow_library_created_at` (`created_at`), + KEY `idx_workflow_library_updated_at` (`updated_at`), + KEY `idx_workflow_library_user_id` (`user_id`), + KEY `idx_workflow_library_category` (`category`(255)), + KEY `idx_workflow_library_description` (`description`(255)), + KEY `idx_workflow_library_name` (`name`(255)) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `boards` ( + `board_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `board_name` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `cover_image_name` varchar(255) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `deleted_at` varchar(32) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `archived` tinyint(1) DEFAULT '0', + `user_id` varchar(64) COLLATE utf8mb4_0900_bin DEFAULT (_utf8mb4'system'), + `is_public` tinyint(1) NOT NULL DEFAULT '0', + `board_visibility` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL DEFAULT (_utf8mb4'private'), + PRIMARY KEY (`board_id`), + KEY `fk_boards_cover_image_name_images` (`cover_image_name`), + KEY `idx_boards_board_visibility` (`board_visibility`), + KEY `idx_boards_created_at` (`created_at`), + KEY `idx_boards_user_id` (`user_id`), + KEY `idx_boards_is_public` (`is_public`), + CONSTRAINT `fk_boards_cover_image_name_images` FOREIGN KEY (`cover_image_name`) REFERENCES `images` (`image_name`) ON DELETE SET NULL +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `client_state` ( + `user_id` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL, + `key` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `value` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`user_id`,`key`), + CONSTRAINT `fk_client_state_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `fonts` ( + `id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `owner_id` varchar(64) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `scope` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `source` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `filename` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `storage_path` longtext COLLATE utf8mb4_0900_bin, + `source_path` varchar(768) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `family` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `label` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `style` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `weight` bigint NOT NULL, + `content_hash` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL, + `byte_size` bigint NOT NULL, + `axes_json` longtext COLLATE utf8mb4_0900_bin NOT NULL DEFAULT (_utf8mb4'[]'), + `instances_json` longtext COLLATE utf8mb4_0900_bin NOT NULL DEFAULT (_utf8mb4'[]'), + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY `idx_fonts_directory_source_path` (`source_path`), + KEY `fk_fonts_owner_id_users` (`owner_id`), + KEY `idx_fonts_uploaded_hash` (`source`,`scope`,`owner_id`,`content_hash`), + CONSTRAINT `fk_fonts_owner_id_users` FOREIGN KEY (`owner_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE, + CONSTRAINT `ck_fonts_byte_size` CHECK ((`byte_size` > 0)), + CONSTRAINT `ck_fonts_content_hash` CHECK ((length(`content_hash`) = 64)), + CONSTRAINT `ck_fonts_scope` CHECK ((`scope` in (_utf8mb4'private',_utf8mb4'shared'))), + CONSTRAINT `ck_fonts_source` CHECK ((`source` in (_utf8mb4'uploaded',_utf8mb4'directory'))), + CONSTRAINT `ck_fonts_source_fields` CHECK ((((`source` = _utf8mb4'directory') and (`owner_id` is null) and (`scope` = _utf8mb4'shared') and (`storage_path` is null) and (`source_path` is not null)) or ((`source` = _utf8mb4'uploaded') and (`storage_path` is not null) and (`source_path` is null) and (((`scope` = _utf8mb4'private') and (`owner_id` is not null)) or ((`scope` = _utf8mb4'shared') and (`owner_id` is null)))))), + CONSTRAINT `ck_fonts_weight` CHECK (((`weight` >= 1) and (`weight` <= 1000))) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `image_embeddings` ( + `image_name` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `model_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `dim` bigint NOT NULL, + `embedding` longblob NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `encoding` longtext COLLATE utf8mb4_0900_bin NOT NULL DEFAULT (_utf8mb4'float32'), + PRIMARY KEY (`image_name`,`model_id`), + KEY `idx_image_embeddings_model_id` (`model_id`), + CONSTRAINT `fk_image_embeddings_image_name_images` FOREIGN KEY (`image_name`) REFERENCES `images` (`image_name`) ON DELETE CASCADE, + CONSTRAINT `ck_image_embeddings_encoding` CHECK ((`encoding` in (_utf8mb4'float32',_utf8mb4'float16'))) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `image_projections` ( + `user_id` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL, + `model_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `scope_hash` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `params` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `point_count` bigint NOT NULL, + `image_names` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `coords` longblob NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `item_kinds` longtext COLLATE utf8mb4_0900_bin, + PRIMARY KEY (`user_id`,`model_id`), + CONSTRAINT `fk_image_projections_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `image_subfolder_move_items` ( + `job_id` bigint NOT NULL, + `image_name` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `old_subfolder` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `new_subfolder` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `is_intermediate` tinyint(1) NOT NULL DEFAULT '0', + `old_path` longtext COLLATE utf8mb4_0900_bin, + `new_path` longtext COLLATE utf8mb4_0900_bin, + `old_thumbnail_path` longtext COLLATE utf8mb4_0900_bin, + `new_thumbnail_path` longtext COLLATE utf8mb4_0900_bin, + `state` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `error_message` longtext COLLATE utf8mb4_0900_bin, + PRIMARY KEY (`job_id`,`image_name`), + KEY `idx_image_subfolder_move_items_image_name` (`image_name`), + KEY `idx_image_subfolder_move_items_job_state` (`job_id`,`state`), + CONSTRAINT `fk_image_subfolder_move_items_image_name_images` FOREIGN KEY (`image_name`) REFERENCES `images` (`image_name`) ON DELETE CASCADE, + CONSTRAINT `fk_image_subfolder_move_items_job_id_image_subfolder_move_jobs` FOREIGN KEY (`job_id`) REFERENCES `image_subfolder_move_jobs` (`id`), + CONSTRAINT `ck_image_subfolder_move_items_state` CHECK ((`state` in (_utf8mb4'planned',_utf8mb4'moved',_utf8mb4'committed',_utf8mb4'error'))) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `model_relationships` ( + `model_key_1` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `model_key_2` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`model_key_1`,`model_key_2`), + KEY `keyx_model_relationships_model_key_2` (`model_key_2`), + CONSTRAINT `fk_model_relationships_model_key_1_models` FOREIGN KEY (`model_key_1`) REFERENCES `models` (`id`) ON DELETE CASCADE, + CONSTRAINT `fk_model_relationships_model_key_2_models` FOREIGN KEY (`model_key_2`) REFERENCES `models` (`id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `session_queue_enqueue_receipts` ( + `queue_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `user_id` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL, + `idempotency_key` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `payload_hash` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `batch_id` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `requested` bigint NOT NULL, + `enqueued` bigint NOT NULL, + `priority` bigint NOT NULL, + `item_ids` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `byte_size` bigint NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `acknowledged_at` varchar(32) COLLATE utf8mb4_0900_bin DEFAULT NULL, + PRIMARY KEY (`queue_id`,`user_id`,`idempotency_key`), + KEY `idx_session_queue_enqueue_receipts_owner_ack` (`user_id`,`acknowledged_at`,`byte_size`), + CONSTRAINT `fk_session_queue_enqueue_receipts_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `user_invitations` ( + `invitation_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `email` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `invited_by` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL, + `invitation_code` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `is_admin` tinyint(1) NOT NULL DEFAULT '0', + `expires_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `used_at` varchar(32) COLLATE utf8mb4_0900_bin DEFAULT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`invitation_id`), + UNIQUE KEY `uq_user_invitations_invitation_code` (`invitation_code`), + KEY `fk_user_invitations_invited_by_users` (`invited_by`), + KEY `idx_user_invitations_expires_at` (`expires_at`), + KEY `idx_user_invitations_email` (`email`), + CONSTRAINT `fk_user_invitations_invited_by_users` FOREIGN KEY (`invited_by`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `user_sessions` ( + `session_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `user_id` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL, + `token_hash` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `expires_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `last_activity_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`session_id`), + KEY `idx_user_sessions_token_hash` (`token_hash`), + KEY `idx_user_sessions_expires_at` (`expires_at`), + KEY `idx_user_sessions_user_id` (`user_id`), + CONSTRAINT `fk_user_sessions_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `video_embeddings` ( + `video_name` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `model_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `dim` bigint NOT NULL, + `embedding` longblob NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `encoding` longtext COLLATE utf8mb4_0900_bin NOT NULL DEFAULT (_utf8mb4'float32'), + PRIMARY KEY (`video_name`,`model_id`), + KEY `idx_video_embeddings_model_id` (`model_id`), + CONSTRAINT `fk_video_embeddings_video_name_videos` FOREIGN KEY (`video_name`) REFERENCES `videos` (`video_name`) ON DELETE CASCADE, + CONSTRAINT `ck_video_embeddings_encoding` CHECK ((`encoding` in (_utf8mb4'float32',_utf8mb4'float16'))) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `wildcards` ( + `id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `name` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `values_json` longtext COLLATE utf8mb4_0900_bin NOT NULL DEFAULT (_utf8mb4'[]'), + `user_id` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`id`), + UNIQUE KEY `idx_wildcards_user_id_name` (`user_id`,`name`), + CONSTRAINT `fk_wildcards_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `board_images` ( + `board_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `image_name` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `deleted_at` varchar(32) COLLATE utf8mb4_0900_bin DEFAULT NULL, + PRIMARY KEY (`image_name`), + KEY `idx_board_images_board_id_created_at` (`board_id`,`created_at`), + CONSTRAINT `fk_board_images_board_id_boards` FOREIGN KEY (`board_id`) REFERENCES `boards` (`board_id`) ON DELETE CASCADE, + CONSTRAINT `fk_board_images_image_name_images` FOREIGN KEY (`image_name`) REFERENCES `images` (`image_name`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `board_videos` ( + `board_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `video_name` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `deleted_at` varchar(32) COLLATE utf8mb4_0900_bin DEFAULT NULL, + PRIMARY KEY (`video_name`), + KEY `idx_board_videos_board_id_created_at` (`board_id`,`created_at`), + CONSTRAINT `fk_board_videos_board_id_boards` FOREIGN KEY (`board_id`) REFERENCES `boards` (`board_id`) ON DELETE CASCADE, + CONSTRAINT `fk_board_videos_video_name_videos` FOREIGN KEY (`video_name`) REFERENCES `videos` (`video_name`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `projects` ( + `project_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `user_id` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL, + `name` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `data` longtext COLLATE utf8mb4_0900_bin NOT NULL, + `board_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `revision` bigint NOT NULL DEFAULT '1', + `created_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `updated_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + `minimum_canvas_schema_version` bigint NOT NULL DEFAULT '2', + PRIMARY KEY (`user_id`,`project_id`), + UNIQUE KEY `uq_projects_board_id` (`board_id`), + CONSTRAINT `fk_projects_board_id_boards` FOREIGN KEY (`board_id`) REFERENCES `boards` (`board_id`) ON DELETE RESTRICT, + CONSTRAINT `fk_projects_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE, + CONSTRAINT `ck_projects_minimum_canvas_schema_version` CHECK ((`minimum_canvas_schema_version` >= 1)) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; + +CREATE TABLE `shared_boards` ( + `board_id` varchar(255) COLLATE utf8mb4_0900_bin NOT NULL, + `user_id` varchar(64) COLLATE utf8mb4_0900_bin NOT NULL, + `can_edit` tinyint(1) NOT NULL DEFAULT '0', + `shared_at` varchar(32) COLLATE utf8mb4_0900_bin NOT NULL, + PRIMARY KEY (`board_id`,`user_id`), + KEY `idx_shared_boards_user_id` (`user_id`), + CONSTRAINT `fk_shared_boards_board_id_boards` FOREIGN KEY (`board_id`) REFERENCES `boards` (`board_id`) ON DELETE CASCADE, + CONSTRAINT `fk_shared_boards_user_id_users` FOREIGN KEY (`user_id`) REFERENCES `users` (`user_id`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin ROW_FORMAT=DYNAMIC; diff --git a/tests/app/services/shared/sqlite_migrator/test_portable_migrations.py b/tests/app/services/shared/sqlite_migrator/test_portable_migrations.py new file mode 100644 index 00000000000..e1bfd243083 --- /dev/null +++ b/tests/app/services/shared/sqlite_migrator/test_portable_migrations.py @@ -0,0 +1,462 @@ +"""Migrations on every database backend: portable migrations, and how a new database reaches the newest schema.""" + +import re +import threading +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import Any, Optional +from unittest import mock + +import pytest +from pydantic import ValidationError +from sqlalchemy import ( + Column, + ForeignKey, + Integer, + MetaData, + Table, + column, + delete, + func, + insert, + inspect, + select, + table, +) + +from invokeai.app.services.config.config_default import DefaultInvokeAIAppConfig +from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase +from invokeai.app.services.shared.database.copy import copy_rows +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.engines import MARIADB_BINARY_COLLATION, MYSQL_BINARY_COLLATION +from invokeai.app.services.shared.database.errors import ForeignKeyViolation +from invokeai.app.services.shared.database.schema import metadata +from invokeai.app.services.shared.database.schema.app_settings import app_settings +from invokeai.app.services.shared.database.schema.boards import board_images +from invokeai.app.services.shared.database.schema.migrator import applied_migrations +from invokeai.app.services.shared.database.schema.users import users +from invokeai.app.services.shared.database.types import Key +from invokeai.app.services.shared.sqlite_migrator import sqlite_migrator_impl +from invokeai.app.services.shared.sqlite_migrator.migration_loader import ( + PORTABLE_CUTOVER, + MigrationBuildContext, + build_migrations, +) +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import ( + Migration, + MigrationBase, + MigrationError, + PortableMigration, + PortableMigrationContext, +) +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import Migrator +from invokeai.backend.util.logging import InvokeAILogger +from tests.fixtures.database import external_test_db_url + +server_only = pytest.mark.skipif( + external_test_db_url() is None, reason="needs a MySQL or MariaDB server (INVOKEAI_TEST_DB_URL)" +) + +# The migrations that run on SQLite only: every one up to the cutover. Their number grows only by a migration +# written before the cutover and merged after it (2026_09_30_store_embeddings_fp16), never by a new one. +SQLITE_ONLY_MIGRATIONS = 64 +QUARANTINE_TABLE = "orphaned_projects_2026_08_06" +CUTOVER_SCHEMAS = Path(__file__).parent / "server_schema_at_cutover" + +probe = table("portable_probe", column("id")) +parent = table("probe_parent", column("id"), column("name")) +child = table("probe_child", column("id"), column("parent_id")) + + +@pytest.fixture +def application_migrations(tmp_path: Path) -> list[MigrationBase]: + # Migrations clean up legacy files under the root. No settings from the environment or a config file may + # point them anywhere else. + config = DefaultInvokeAIAppConfig() + config._root = tmp_path + context = MigrationBuildContext( + app_config=config, + logger=InvokeAILogger.get_logger("test_portable_migrations"), + image_files=mock.Mock(spec=ImageFileStorageBase), + ) + return build_migrations(context) + + +def _migrator(database: Database, migrations: Sequence[MigrationBase], *more: MigrationBase) -> Migrator: + migrator = Migrator(database) + for migration in [*migrations, *more]: + migrator.register_migration(migration) + return migrator + + +def _portable(callback: Callable[[PortableMigrationContext], None], name: str = "probe") -> PortableMigration: + return PortableMigration(id=f"2099_01_01_{name}", depends_on=PORTABLE_CUTOVER, callback=callback) + + +def _applied(database: Database) -> set[str]: + with database.begin(write=False) as conn: + return set(conn.execute(select(applied_migrations.c.migration_id)).scalars()) + + +def _has_table(database: Database, name: str) -> bool: + with database.begin(write=False) as conn: + return inspect(conn).has_table(name) + + +def _rows(database: Database, of: Any) -> list[tuple[Any, ...]]: + with database.begin(write=False) as conn: + return [tuple(row) for row in conn.execute(select(of).order_by(of.c.id))] + + +def _create_probe_table(context: PortableMigrationContext) -> None: + # Idempotent, as a portable migration must be: on a server, a failed run may have created the table already. + if not inspect(context.conn).has_table("portable_probe"): + context.create_table("portable_probe", Column("id", Integer, primary_key=True, autoincrement=False)) + + +def _create_parent_and_child(context: PortableMigrationContext) -> None: + if not inspect(context.conn).has_table("probe_parent"): + context.create_table( + "probe_parent", + Column("id", Integer, primary_key=True, autoincrement=False), + Column("name", Key(), nullable=False), + ) + context.create_table( + "probe_child", + Column("id", Integer, primary_key=True, autoincrement=False), + Column("parent_id", Integer, ForeignKey("probe_parent.id", ondelete="CASCADE"), nullable=False), + ) + + +def test_a_portable_migration_has_no_legacy_version() -> None: + with pytest.raises(ValidationError, match="no legacy version"): + PortableMigration(from_version=0, to_version=1, callback=_create_probe_table) + + +def test_only_the_migrations_up_to_the_cutover_run_on_sqlite_only( + application_migrations: list[MigrationBase], +) -> None: + sqlite_only = [migration.id for migration in application_migrations if isinstance(migration, Migration)] + + # A SQLite-only migration added now, of whatever date, would be refused by every existing server database. + assert len(sqlite_only) == SQLITE_ONLY_MIGRATIONS + assert sqlite_only[-1] == PORTABLE_CUTOVER + + +def test_a_new_database_reaches_the_newest_schema_with_the_seeded_rows( + empty_database: Database, application_migrations: list[MigrationBase] +) -> None: + assert _migrator(empty_database, application_migrations).run_migrations() + + # The rows the migration chain seeds into a SQLite database. (On SQLite this compares the chain with itself; + # on a server, the bootstrap with the chain.) + reference = Database.open_sqlite(None, InvokeAILogger.get_logger("test_portable_migrations")) + try: + _migrator(reference, application_migrations).run_migrations() + with reference.begin(write=False) as conn: + expected_rows = { + name: conn.execute(select(func.count()).select_from(each)).scalar_one() + for name, each in metadata.tables.items() + if name != QUARANTINE_TABLE + } + finally: + reference.dispose() + + with empty_database.begin(write=False) as conn: + tables = set(inspect(conn).get_table_names()) + rows = { + name: conn.execute(select(func.count()).select_from(metadata.tables[name])).scalar_one() + for name in expected_rows + } + secret: str = conn.execute(select(app_settings.c.value).where(app_settings.c.key == "jwt_secret")).scalar_one() + user_ids: list[str] = list(conn.execute(select(users.c.user_id)).scalars()) + # The quarantine table is created on a server; on SQLite only where a migration had projects to keep. + expected_tables = set(metadata.tables) - ({QUARANTINE_TABLE} if empty_database.dialect_name == "sqlite" else set()) + assert tables == expected_tables + assert rows == expected_rows + assert _applied(empty_database) == {migration.id for migration in application_migrations} + assert user_ids == ["system"] + assert len(secret) == 64 + + # Up to date: a second run has nothing to do. + assert not _migrator(empty_database, application_migrations).run_migrations() + + +def test_a_portable_migration_runs_once_and_is_recorded( + empty_database: Database, application_migrations: list[MigrationBase] +) -> None: + def add_probe(context: PortableMigrationContext) -> None: + _create_probe_table(context) + context.conn.execute(insert(probe).values(id=1)) + + migration = _portable(add_probe) + + assert _migrator(empty_database, application_migrations, migration).run_migrations() + # Not run again: its insert would fail on the existing row. + assert not _migrator(empty_database, application_migrations, migration).run_migrations() + + assert _rows(empty_database, probe) == [(1,)] + assert migration.id in _applied(empty_database) + + +def test_a_failed_portable_migration_is_not_recorded_and_runs_again( + empty_database: Database, application_migrations: list[MigrationBase] +) -> None: + def fail_halfway(context: PortableMigrationContext) -> None: + _create_probe_table(context) + raise RuntimeError("interrupted") + + failing = _portable(fail_halfway) + with pytest.raises(MigrationError, match="interrupted"): + _migrator(empty_database, application_migrations, failing).run_migrations() + + assert failing.id not in _applied(empty_database) + # SQLite rolls the DDL back; a server commits each DDL statement as it runs, which is why a portable + # migration checks what is there before it changes it. + assert _has_table(empty_database, "portable_probe") is (empty_database.dialect_name != "sqlite") + + fixed = _portable(_create_probe_table) + assert _migrator(empty_database, application_migrations, fixed).run_migrations() + assert fixed.id in _applied(empty_database) + assert _has_table(empty_database, "portable_probe") + + +def test_rebuilding_a_table_keeps_the_rows_that_reference_it( + empty_database: Database, application_migrations: list[MigrationBase] +) -> None: + def create_rows(context: PortableMigrationContext) -> None: + _create_parent_and_child(context) + context.conn.execute(insert(parent).values(id=1, name="parent")) + context.conn.execute(insert(child).values(id=1, parent_id=1)) + + def rebuild_parent(context: PortableMigrationContext) -> None: + # On SQLite, batch mode copies the table and drops the original, which would cascade to the children. + with context.op.batch_alter_table("probe_parent") as batch: + batch.alter_column("name", existing_type=Key(), nullable=True) + + created = _portable(create_rows, "create_rows") + rebuilt = PortableMigration(id="2099_01_02_rebuild", depends_on=created.id, callback=rebuild_parent) + assert _migrator(empty_database, application_migrations, created, rebuilt).run_migrations() + + assert _rows(empty_database, child) == [(1, 1)] + assert rebuilt.id in _applied(empty_database) + + # Foreign keys are on again: deleting the parent deletes its child, by cascade. + with empty_database.begin(write=True) as conn: + conn.execute(delete(parent).where(parent.c.id == 1)) + assert _rows(empty_database, child) == [] + + +def test_a_portable_migration_that_orphans_rows_fails( + empty_database: Database, application_migrations: list[MigrationBase] +) -> None: + def orphan(context: PortableMigrationContext) -> None: + _create_parent_and_child(context) + context.conn.execute(insert(child).values(id=1, parent_id=99)) + + migration = _portable(orphan) + # A server refuses the row at once; SQLite, which runs portable migrations with foreign keys off, when the + # migrator checks them before committing. + with pytest.raises(MigrationError): + _migrator(empty_database, application_migrations, migration).run_migrations() + + assert migration.id not in _applied(empty_database) + if _has_table(empty_database, "probe_child"): + assert _rows(empty_database, child) == [] + # Foreign keys are on again after the failure, too. + with pytest.raises(ForeignKeyViolation): + with empty_database.begin(write=True) as conn: + conn.execute(insert(board_images).values(board_id="no board", image_name="no image.png")) + + +@pytest.mark.sqlite_only +def test_rows_orphaned_before_a_portable_migration_do_not_fail_it( + empty_database: Database, application_migrations: list[MigrationBase] +) -> None: + # Tools that deleted with foreign keys off left orphans in real databases; a migration did not make them. + _migrator(empty_database, application_migrations).run_migrations() + raw = empty_database.sqlite.conn + raw.execute("PRAGMA foreign_keys = OFF") + raw.execute("INSERT INTO board_images (board_id, image_name) VALUES ('deleted board', 'deleted.png')") + raw.commit() + raw.execute("PRAGMA foreign_keys = ON") + + migration = _portable(_create_probe_table) + assert _migrator(empty_database, application_migrations, migration).run_migrations() + + assert migration.id in _applied(empty_database) + + +def test_a_database_from_a_newer_version_is_refused( + empty_database: Database, application_migrations: list[MigrationBase] +) -> None: + _migrator(empty_database, application_migrations).run_migrations() + with empty_database.begin(write=True) as conn: + conn.execute(insert(applied_migrations).values(migration_id="2999_01_01_from_a_newer_version")) + + with pytest.raises(MigrationError, match="unknown applied migration IDs"): + _migrator(empty_database, application_migrations).run_migrations() + + +@server_only +def test_a_sqlite_only_migration_is_refused_on_a_server( + empty_database: Database, application_migrations: list[MigrationBase] +) -> None: + def cursor_callback(cursor: object) -> None: + pass + + sqlite_only = Migration(id="2099_01_01_sqlite_only", depends_on=PORTABLE_CUTOVER, callback=cursor_callback) + + with pytest.raises(MigrationError, match="runs on SQLite only"): + _migrator(empty_database, application_migrations, sqlite_only).run_migrations() + + +@server_only +def test_a_server_database_with_other_tables_is_refused( + empty_database: Database, application_migrations: list[MigrationBase] +) -> None: + with empty_database.begin(write=True) as conn: + Table("orders", MetaData(), Column("id", Integer, primary_key=True)).create(conn) + + with pytest.raises(MigrationError, match="not an InvokeAI database"): + _migrator(empty_database, application_migrations).run_migrations() + + +@server_only +def test_an_interrupted_creation_is_refused_rather_than_taken_for_complete( + empty_database: Database, application_migrations: list[MigrationBase], monkeypatch: pytest.MonkeyPatch +) -> None: + def copy_until_users(source: Database, target: Database, *, tables: Optional[Sequence[Table]] = None) -> None: + # Copies, in the order the copy goes, what precedes the accounts, then stops as a crash would. + wanted = set(tables if tables is not None else metadata.sorted_tables) + before_users: list[Table] = [] + for each in metadata.sorted_tables: + if each.name == users.name: + break + if each in wanted: + before_users.append(each) + copy_rows(source, target, tables=before_users) + raise RuntimeError("interrupted") + + monkeypatch.setattr(sqlite_migrator_impl, "copy_rows", copy_until_users) + with pytest.raises(RuntimeError, match="interrupted"): + _migrator(empty_database, application_migrations).run_migrations() + monkeypatch.undo() + + with pytest.raises(MigrationError, match="creation was interrupted"): + _migrator(empty_database, application_migrations).run_migrations() + + +@server_only +def test_the_creation_reads_no_settings_from_the_environment( + empty_database: Database, + application_migrations: list[MigrationBase], + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + # A legacy core model the reference chain's clean-up deletes from the models directory it is given. + models_dir = tmp_path / "elsewhere" / "models" + legacy_model = models_dir / "core" / "upscaling" / "realesrgan" / "RealESRGAN_x4plus.pth" + legacy_model.parent.mkdir(parents=True) + legacy_model.write_bytes(b"weights") + monkeypatch.setenv("INVOKEAI_MODELS_DIR", str(models_dir)) + + assert _migrator(empty_database, application_migrations).run_migrations() + + assert legacy_model.exists() + + +@server_only +def test_one_process_at_a_time_creates_or_migrates_a_server_database( + empty_database: Database, application_migrations: list[MigrationBase], monkeypatch: pytest.MonkeyPatch +) -> None: + creating = threading.Event() + release = threading.Event() + outcome: list[object] = [] + + def paused_copy_rows(source: Database, target: Database, *, tables: Optional[Sequence[Table]] = None) -> Any: + creating.set() + release.wait(60) + return copy_rows(source, target, tables=tables) + + def create() -> None: + try: + outcome.append(_migrator(empty_database, application_migrations).run_migrations()) + except Exception as error: + outcome.append(error) + + monkeypatch.setattr(sqlite_migrator_impl, "copy_rows", paused_copy_rows) + first = threading.Thread(target=create) + first.start() + try: + # The first process is creating the database: its tables exist, its rows not yet. + while not creating.wait(0.1): + assert first.is_alive(), outcome + monkeypatch.setattr(sqlite_migrator_impl, "MIGRATION_LOCK_TIMEOUT_SECONDS", 1) + with pytest.raises(MigrationError, match="Another process"): + _migrator(empty_database, application_migrations).run_migrations() + finally: + release.set() + first.join(60) + + assert outcome == [True] + + +@server_only +def test_portable_migrations_take_a_server_from_the_cutover_to_the_newest_schema( + empty_database: Database, application_migrations: list[MigrationBase] +) -> None: + # A server database as it was created at the cutover, which the migrations since must bring to what the + # metadata creates today. Its default collation folds case, as a stock server's does, so that a table a + # migration creates without the schema's table options shows. + binary = MARIADB_BINARY_COLLATION if empty_database.dialect_name == "mariadb" else MYSQL_BINARY_COLLATION + statements = _statements(CUTOVER_SCHEMAS / f"{empty_database.dialect_name}.sql") + with empty_database.begin(write=True) as conn: + conn.exec_driver_sql("ALTER DATABASE COLLATE utf8mb4_general_ci") + try: + with empty_database.begin(write=True) as conn: + for statement in statements: + conn.exec_driver_sql(statement) + conn.execute( + insert(applied_migrations), + [ + {"migration_id": migration.id, "legacy_version": migration.to_version} + for migration in application_migrations + if isinstance(migration, Migration) + ], + ) + _migrator(empty_database, application_migrations).run_migrations() + migrated = _server_schema(empty_database) + finally: + with empty_database.begin(write=True) as conn: + conn.exec_driver_sql(f"ALTER DATABASE COLLATE {binary}") + + with empty_database.begin(write=True) as conn: + metadata.drop_all(conn) + metadata.create_all(conn) + + assert migrated == _server_schema(empty_database) + + +def _statements(path: Path) -> list[str]: + text = "\n".join(line for line in path.read_text(encoding="utf-8").splitlines() if not line.startswith("--")) + return [statement.strip() for statement in text.split(";\n") if statement.strip()] + + +def _server_schema(database: Database) -> dict[str, tuple[list[str], list[str], str]]: + """Each table as the server renders it: columns in order, keys and constraints in any order, and options.""" + schema: dict[str, tuple[list[str], list[str], str]] = {} + with database.begin(write=False) as conn: + for name in inspect(conn).get_table_names(): + ddl: str = conn.exec_driver_sql(f"SHOW CREATE TABLE `{name}`").one()[1] + # MySQL names the character set of a column whose collation was given explicitly, as it is in the + # frozen statements; the collation names it either way. + ddl = re.sub(r"CHARACTER SET \w+ COLLATE", "COLLATE", ddl) + lines = [line.strip().rstrip(",") for line in ddl.splitlines()] + body = lines[1:-1] + schema[name] = ( + [line for line in body if line.startswith("`")], + sorted(line for line in body if not line.startswith("`")), + re.sub(r" AUTO_INCREMENT=\d+", "", lines[-1]), + ) + return schema diff --git a/tests/app/services/shared/test_invocation_context_boards.py b/tests/app/services/shared/test_invocation_context_boards.py index f771bf2db61..ab42e736a04 100644 --- a/tests/app/services/shared/test_invocation_context_boards.py +++ b/tests/app/services/shared/test_invocation_context_boards.py @@ -5,7 +5,7 @@ from invokeai.app.services.board_records.board_records_common import BoardRecordOrderBy from invokeai.app.services.boards.boards_base import BoardServiceABC from invokeai.app.services.shared.invocation_context import BoardsInterface -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import SQLiteDirection def _make_interface(multiuser: bool, user: MagicMock | None) -> tuple[BoardsInterface, MagicMock]: diff --git a/tests/app/services/shared/test_media_references.py b/tests/app/services/shared/test_media_references.py index 43ca3968626..3be6b8bc829 100644 --- a/tests/app/services/shared/test_media_references.py +++ b/tests/app/services/shared/test_media_references.py @@ -1,14 +1,13 @@ -"""The media reference extractor and index writer shared by project, workflow and migration code.""" +"""The media reference extractor, and the index writer frozen with the migrations that backfill the index.""" import sqlite3 import pytest -from invokeai.app.services.shared.media_references import ( - MediaReferences, +from invokeai.app.services.shared.media_references import MediaReferences, extract_media_references +from invokeai.app.services.shared.sqlite_migrator.migrations._media_references_v1 import ( create_media_references_table, delete_media_references, - extract_media_references, replace_media_references, ) diff --git a/tests/app/services/test_image_category_views.py b/tests/app/services/test_image_category_views.py index 8d049f85efd..0125547c097 100644 --- a/tests/app/services/test_image_category_views.py +++ b/tests/app/services/test_image_category_views.py @@ -55,21 +55,3 @@ def test_images_view_excludes_canvas_owned_images(mock_services: InvocationServi result = mock_services.image_records.get_many(offset=0, limit=10, categories=IMAGE_CATEGORIES) assert [record.image_name for record in result.items] == ["generated.png"] - - -def test_date_board_counts_ignore_canvas_owned_images(mock_services: InvocationServices) -> None: - """`get_image_dates` derives its counts from the category constants. - - It previously classified every non-`general` image as an asset, which would - keep counting canvas-owned images after they left the assets view — leaving - the date board's badge disagreeing with the images it actually lists. - """ - _save(mock_services, "painted.png", ImageCategory.OTHER) - _save(mock_services, "uploaded.png", ImageCategory.USER) - _save(mock_services, "generated.png", ImageCategory.GENERAL) - - dates = mock_services.image_records.get_image_dates() - - assert len(dates) == 1 - assert dates[0].image_count == 1 - assert dates[0].asset_count == 1 diff --git a/tests/app/services/test_sql_injection_protection.py b/tests/app/services/test_sql_injection_protection.py index bf5867f2179..d593c8e13df 100644 --- a/tests/app/services/test_sql_injection_protection.py +++ b/tests/app/services/test_sql_injection_protection.py @@ -4,17 +4,17 @@ BoardRecordNotFoundException, BoardRecordOrderBy, ) -from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage from invokeai.app.services.config.config_default import InvokeAIAppConfig -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.backend.util.logging import InvokeAILogger from tests.fixtures.sqlite_database import create_mock_sqlite_database -def _create_board_storage() -> SqliteBoardRecordStorage: +def _create_board_storage() -> BoardRecordStorage: config = InvokeAIAppConfig(use_memory_db=True) db = create_mock_sqlite_database(config=config, logger=InvokeAILogger.get_logger()) - return SqliteBoardRecordStorage(db=db) + return BoardRecordStorage(db) def test_sql_injection_payload_in_board_name_is_stored_as_plain_text() -> None: diff --git a/tests/app/services/test_style_preset_records.py b/tests/app/services/test_style_preset_records.py new file mode 100644 index 00000000000..4ffc5acdb41 --- /dev/null +++ b/tests/app/services/test_style_preset_records.py @@ -0,0 +1,174 @@ +"""Style preset records on every database backend: which presets an account sees, and the bundled ones kept in step +with the file that ships them.""" + +import json +from pathlib import Path +from unittest import mock + +import pytest + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries.style_presets import StylePresetQueries +from invokeai.app.services.style_preset_records import style_preset_records_default +from invokeai.app.services.style_preset_records.style_preset_records_common import ( + PresetData, + PresetType, + StylePresetChanges, + StylePresetNotFoundError, + StylePresetWithoutId, +) +from invokeai.app.services.style_preset_records.style_preset_records_default import ( + SYSTEM_USER_ID, + StylePresetRecordsStorage, +) + +ALICE = "alice" +BOB = "bob" + + +@pytest.fixture +def presets(database: Database) -> StylePresetRecordsStorage: + return StylePresetRecordsStorage(database) + + +def _preset(name: str, *, type: PresetType = PresetType.User, is_public: bool = False) -> StylePresetWithoutId: + return StylePresetWithoutId( + name=name, + preset_data=PresetData(positive_prompt=f"{name}, {{prompt}}", negative_prompt="blurry"), + type=type, + is_public=is_public, + ) + + +def _bundled_names() -> list[str]: + path = Path(style_preset_records_default.__file__).parent / "default_style_presets.json" + return sorted(preset["name"] for preset in json.loads(path.read_text(encoding="utf-8"))) + + +def test_an_account_sees_its_own_the_shared_and_the_bundled_presets(presets: StylePresetRecordsStorage) -> None: + own = presets.create(_preset("alice private"), ALICE) + shared = presets.create(_preset("bob shared", is_public=True), BOB) + hidden = presets.create(_preset("bob private"), BOB) + bundled = presets.create(_preset("bundled", type=PresetType.Default), SYSTEM_USER_ID) + + def visible(user_id: str | None, is_admin: bool = False) -> set[str]: + return {preset.id for preset in presets.get_many(user_id=user_id, is_admin=is_admin)} + + assert visible(ALICE) == {own.id, shared.id, bundled.id} + assert visible(BOB) == {shared.id, hidden.id, bundled.id} + # Without an account only what everyone may see is listed. + assert visible(None) == {shared.id, bundled.id} + assert visible(ALICE, is_admin=True) == {own.id, shared.id, hidden.id, bundled.id} + + +def test_a_type_filter_narrows_what_the_account_sees(presets: StylePresetRecordsStorage) -> None: + own = presets.create(_preset("mine"), ALICE) + theirs = presets.create(_preset("theirs"), BOB) + bundled = presets.create(_preset("bundled", type=PresetType.Default), SYSTEM_USER_ID) + + assert [preset.id for preset in presets.get_many(type=PresetType.User, user_id=ALICE)] == [own.id] + assert [preset.id for preset in presets.get_many(type=PresetType.Default, user_id=ALICE)] == [bundled.id] + # The export: every account's own presets, and no bundled one. + exported = presets.get_many(type=PresetType.User, user_id=ALICE, is_admin=True) + assert {preset.id for preset in exported} == {own.id, theirs.id} + + +def test_presets_are_listed_by_name_ignoring_case_and_equal_names_by_id( + database: Database, presets: StylePresetRecordsStorage +) -> None: + # Stored against id order, so that only the tie-break lists equal names by id. + named = (("p1", "beta"), ("p4", "alpha"), ("p3", "ALPHA"), ("p2", "Alpha")) + database.queries.style_presets.insert({preset_id: _preset(name) for preset_id, name in named}, ALICE) + + assert [preset.id for preset in presets.get_many(user_id=ALICE)] == ["p2", "p3", "p4", "p1"] + + +def test_a_created_preset_is_stored_as_returned(presets: StylePresetRecordsStorage) -> None: + created = presets.create(_preset("shared", is_public=True), ALICE) + + assert (created.name, created.type, created.user_id, created.is_public) == ("shared", PresetType.User, ALICE, True) + assert presets.get(created.id) == created + + +def test_an_update_sets_only_the_fields_it_names_of_only_that_preset(presets: StylePresetRecordsStorage) -> None: + created = presets.create(_preset("before"), ALICE) + bystander = presets.create(_preset("bystander", is_public=True), BOB) + data = PresetData(positive_prompt="after, {prompt}", negative_prompt="") + + renamed = presets.update(created.id, StylePresetChanges(name="after", type=None)) + shared = presets.update(created.id, StylePresetChanges(preset_data=data, is_public=True, type=None)) + unshared = presets.update(created.id, StylePresetChanges(is_public=False, type=None)) + + assert (renamed.name, renamed.preset_data, renamed.is_public) == ("after", created.preset_data, False) + assert (shared.name, shared.preset_data, shared.is_public) == ("after", data, True) + assert (unshared.name, unshared.preset_data, unshared.is_public) == ("after", data, False) + assert presets.get(created.id) == unshared + # Made private again, it is hidden from the other accounts again. + assert created.id not in {preset.id for preset in presets.get_many(user_id=BOB)} + assert presets.get(bystander.id) == bystander + + +def test_deleting_removes_only_that_preset_and_deleting_it_again_is_harmless( + presets: StylePresetRecordsStorage, +) -> None: + created = presets.create(_preset("short-lived"), ALICE) + bystander = presets.create(_preset("bystander"), BOB) + + presets.delete(created.id) + presets.delete(created.id) + + with pytest.raises(StylePresetNotFoundError): + presets.get(created.id) + with pytest.raises(StylePresetNotFoundError): + presets.update(created.id, StylePresetChanges(name="revived", type=None)) + assert presets.get(bystander.id) == bystander + + +def test_an_import_stores_every_preset_for_the_account(presets: StylePresetRecordsStorage) -> None: + presets.create_many([_preset("one"), _preset("two")], ALICE) + # An import of an empty file stores nothing rather than failing. + presets.create_many([], ALICE) + + imported = presets.get_many(user_id=ALICE) + + assert [(preset.name, preset.user_id, preset.type) for preset in imported] == [ + ("one", ALICE, PresetType.User), + ("two", ALICE, PresetType.User), + ] + assert len({preset.id for preset in imported}) == 2 + + +def test_starting_replaces_the_bundled_presets_and_keeps_the_accounts_own(presets: StylePresetRecordsStorage) -> None: + own = presets.create(_preset("mine"), ALICE) + # In single-user mode every preset belongs to the system user, so only the type tells the bundled ones apart. + single_user = presets.create(_preset("single-user preset"), SYSTEM_USER_ID) + retired = presets.create(_preset("retired", type=PresetType.Default), SYSTEM_USER_ID) + + presets.start(mock.Mock()) + # A restart stores them again in place of the first copies, not beside them. + presets.start(mock.Mock()) + + bundled = presets.get_many(type=PresetType.Default, is_admin=True) + assert sorted(preset.name for preset in bundled) == _bundled_names() + assert {preset.user_id for preset in bundled} == {SYSTEM_USER_ID} + assert presets.get(own.id) == own + assert presets.get(single_user.id) == single_user + with pytest.raises(StylePresetNotFoundError): + presets.get(retired.id) + + +def test_a_failed_sync_keeps_the_bundled_presets_it_had( + presets: StylePresetRecordsStorage, monkeypatch: pytest.MonkeyPatch +) -> None: + presets.start(mock.Mock()) + + def failing_insert(self: StylePresetQueries, bundled: object, user_id: object) -> None: + raise RuntimeError("the database went away") + + monkeypatch.setattr(StylePresetQueries, "insert", failing_insert) + with pytest.raises(RuntimeError): + presets.start(mock.Mock()) + + assert ( + sorted(preset.name for preset in presets.get_many(type=PresetType.Default, is_admin=True)) == _bundled_names() + ) diff --git a/tests/app/services/test_system_prompt_records.py b/tests/app/services/test_system_prompt_records.py index 3bc54eaa59b..443ed36ae66 100644 --- a/tests/app/services/test_system_prompt_records.py +++ b/tests/app/services/test_system_prompt_records.py @@ -10,28 +10,22 @@ import pytest from pydantic import ValidationError -from invokeai.app.services.config.config_default import InvokeAIAppConfig +from invokeai.app.services.shared.database.database import Database from invokeai.app.services.system_prompt_records.system_prompt_records_common import ( EXPAND_PROMPT_MAX_TOKENS_MAX, SystemPromptChanges, SystemPromptNotFoundError, SystemPromptWithoutId, ) -from invokeai.app.services.system_prompt_records.system_prompt_records_sqlite import ( - SqliteSystemPromptRecordsStorage, -) -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database +from invokeai.app.services.system_prompt_records.system_prompt_records_default import SystemPromptRecordsStorage -def _storage() -> SqliteSystemPromptRecordsStorage: - config = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0) - db = create_mock_sqlite_database(config, InvokeAILogger.get_logger()) - return SqliteSystemPromptRecordsStorage(db=db) +@pytest.fixture +def svc(database: Database) -> SystemPromptRecordsStorage: + return SystemPromptRecordsStorage(database) -def test_seeded_defaults_visible_to_every_user() -> None: - svc = _storage() +def test_seeded_defaults_visible_to_every_user(svc: SystemPromptRecordsStorage) -> None: seeded = svc.get_many(user_id=None) assert len(seeded) >= 1 assert all(p.user_id == "system" and p.is_public for p in seeded) @@ -42,8 +36,7 @@ def test_seeded_defaults_visible_to_every_user() -> None: assert seeded_ids.issubset({p.id for p in for_alice}) -def test_private_prompt_hidden_from_other_user_visible_to_owner_and_admin() -> None: - svc = _storage() +def test_private_prompt_hidden_from_other_user_visible_to_owner_and_admin(svc: SystemPromptRecordsStorage) -> None: alice = svc.create( SystemPromptWithoutId(name="alice secret", content="x"), user_id="alice", @@ -63,8 +56,7 @@ def test_private_prompt_hidden_from_other_user_visible_to_owner_and_admin() -> N assert alice.id in {p.id for p in admin_view} -def test_public_prompt_visible_to_everyone() -> None: - svc = _storage() +def test_public_prompt_visible_to_everyone(svc: SystemPromptRecordsStorage) -> None: bob_pub = svc.create( SystemPromptWithoutId(name="bob shared", content="x"), user_id="bob", @@ -74,8 +66,7 @@ def test_public_prompt_visible_to_everyone() -> None: assert bob_pub.id in {p.id for p in alice_view} -def test_update_with_non_owner_user_id_raises_and_does_not_mutate() -> None: - svc = _storage() +def test_update_with_non_owner_user_id_raises_and_does_not_mutate(svc: SystemPromptRecordsStorage) -> None: alice = svc.create( SystemPromptWithoutId(name="alice secret", content="original"), user_id="alice", @@ -92,8 +83,7 @@ def test_update_with_non_owner_user_id_raises_and_does_not_mutate() -> None: assert after.content == "original" -def test_delete_with_non_owner_user_id_raises_and_does_not_delete() -> None: - svc = _storage() +def test_delete_with_non_owner_user_id_raises_and_does_not_delete(svc: SystemPromptRecordsStorage) -> None: alice = svc.create( SystemPromptWithoutId(name="alice secret", content="x"), user_id="alice", @@ -109,10 +99,9 @@ def test_delete_with_non_owner_user_id_raises_and_does_not_delete() -> None: assert svc.get(alice.id).id == alice.id -def test_delete_of_unknown_id_raises() -> None: +def test_delete_of_unknown_id_raises(svc: SystemPromptRecordsStorage) -> None: # Unscoped (single-user / admin) deletes must not report success for an id that does not # exist -- the router turns this into a 404, matching GET. - svc = _storage() try: svc.delete("does-not-exist", user_id=None) raise AssertionError("expected SystemPromptNotFoundError") @@ -120,8 +109,7 @@ def test_delete_of_unknown_id_raises() -> None: pass -def test_deleting_the_same_row_twice_raises_the_second_time() -> None: - svc = _storage() +def test_deleting_the_same_row_twice_raises_the_second_time(svc: SystemPromptRecordsStorage) -> None: created = svc.create(SystemPromptWithoutId(name="temp", content="x"), user_id="alice", is_public=False) svc.delete(created.id, user_id=None) try: @@ -131,8 +119,7 @@ def test_deleting_the_same_row_twice_raises_the_second_time() -> None: pass -def test_admin_can_delete_any_row() -> None: - svc = _storage() +def test_admin_can_delete_any_row(svc: SystemPromptRecordsStorage) -> None: alice = svc.create( SystemPromptWithoutId(name="alice secret", content="x"), user_id="alice", @@ -146,8 +133,7 @@ def test_admin_can_delete_any_row() -> None: pass -def test_owner_can_flip_is_public() -> None: - svc = _storage() +def test_owner_can_flip_is_public(svc: SystemPromptRecordsStorage) -> None: alice = svc.create( SystemPromptWithoutId(name="alice", content="x"), user_id="alice", @@ -160,16 +146,43 @@ def test_owner_can_flip_is_public() -> None: bob_view = svc.get_many(user_id="bob") assert alice.id in {p.id for p in bob_view} + # And private again, hidden again + unshared = svc.update(alice.id, SystemPromptChanges(is_public=False), user_id="alice") + assert unshared.is_public is False + assert alice.id not in {p.id for p in svc.get_many(user_id="bob")} + + +def test_a_created_prompt_is_stored_as_returned(svc: SystemPromptRecordsStorage) -> None: + created = svc.create(SystemPromptWithoutId(name="shared", content="x", max_tokens=700), "alice", is_public=True) + + assert (created.user_id, created.is_public, created.max_tokens) == ("alice", True, 700) + assert created.updated_at == created.created_at + assert svc.get(created.id) == created + + +def test_updates_and_deletes_touch_only_the_prompt_they_name(svc: SystemPromptRecordsStorage) -> None: + alice = svc.create(SystemPromptWithoutId(name="alice", content="x"), user_id="alice") + alices_other = svc.create(SystemPromptWithoutId(name="alice too", content="x"), user_id="alice") + bobs = svc.create(SystemPromptWithoutId(name="bob", content="x"), user_id="bob", is_public=True) + + svc.update(alice.id, SystemPromptChanges(content="by its owner"), user_id="alice") + svc.update(alice.id, SystemPromptChanges(name="by an admin"), user_id=None) + assert svc.get(alices_other.id) == alices_other + assert svc.get(bobs.id) == bobs + + svc.delete(alice.id, user_id="alice") + assert svc.get(alices_other.id) == alices_other + svc.delete(alices_other.id, user_id=None) + assert svc.get(bobs.id) == bobs -def test_created_prompt_defaults_to_no_cap() -> None: + +def test_created_prompt_defaults_to_no_cap(svc: SystemPromptRecordsStorage) -> None: # NULL is what tells the client to fall back to the endpoint default. - svc = _storage() created = svc.create(SystemPromptWithoutId(name="alice", content="x"), user_id="alice") assert created.max_tokens is None -def test_create_and_update_round_trip_a_cap() -> None: - svc = _storage() +def test_create_and_update_round_trip_a_cap(svc: SystemPromptRecordsStorage) -> None: created = svc.create(SystemPromptWithoutId(name="alice", content="x", max_tokens=500), user_id="alice") assert created.max_tokens == 500 @@ -177,8 +190,7 @@ def test_create_and_update_round_trip_a_cap() -> None: assert raised.max_tokens == 900 -def test_update_that_omits_max_tokens_leaves_the_cap_alone() -> None: - svc = _storage() +def test_update_that_omits_max_tokens_leaves_the_cap_alone(svc: SystemPromptRecordsStorage) -> None: created = svc.create(SystemPromptWithoutId(name="alice", content="x", max_tokens=500), user_id="alice") renamed = svc.update(created.id, SystemPromptChanges(name="renamed"), user_id="alice") @@ -186,10 +198,9 @@ def test_update_that_omits_max_tokens_leaves_the_cap_alone() -> None: assert renamed.max_tokens == 500 -def test_explicit_null_clears_the_cap_back_to_the_default() -> None: +def test_explicit_null_clears_the_cap_back_to_the_default(svc: SystemPromptRecordsStorage) -> None: # The one field where null is a value rather than "no change" -- without this a cap could be # set but never removed. - svc = _storage() created = svc.create(SystemPromptWithoutId(name="alice", content="x", max_tokens=500), user_id="alice") cleared = svc.update(created.id, SystemPromptChanges(max_tokens=None), user_id="alice") @@ -197,8 +208,7 @@ def test_explicit_null_clears_the_cap_back_to_the_default() -> None: assert cleared.max_tokens is None -def test_non_owner_cannot_change_the_cap() -> None: - svc = _storage() +def test_non_owner_cannot_change_the_cap(svc: SystemPromptRecordsStorage) -> None: alice = svc.create(SystemPromptWithoutId(name="alice", content="x", max_tokens=500), user_id="alice") try: @@ -215,3 +225,27 @@ def test_out_of_range_caps_are_rejected_by_the_model() -> None: SystemPromptWithoutId(name="alice", content="x", max_tokens=out_of_range) with pytest.raises(ValidationError): SystemPromptChanges(max_tokens=out_of_range) + + +def test_an_update_without_changes_still_needs_a_prompt_the_account_may_change(svc: SystemPromptRecordsStorage) -> None: + # An empty PATCH changes nothing, but must not succeed where one with changes would not. + alice = svc.create(SystemPromptWithoutId(name="alice", content="x"), user_id="alice") + + assert svc.update(alice.id, SystemPromptChanges(), user_id="alice") == alice + assert svc.update(alice.id, SystemPromptChanges(), user_id=None) == alice + with pytest.raises(SystemPromptNotFoundError): + svc.update(alice.id, SystemPromptChanges(), user_id="bob") + with pytest.raises(SystemPromptNotFoundError): + svc.update("does-not-exist", SystemPromptChanges(), user_id=None) + + +def test_prompts_are_listed_by_name_ignoring_case_and_equal_names_by_id( + database: Database, svc: SystemPromptRecordsStorage +) -> None: + # Stored against id order, so that only the tie-break lists equal names by id. + for prompt_id, name in (("p1", "Beta"), ("p3", "alpha"), ("p2", "Alpha")): + database.queries.system_prompts.insert(prompt_id, SystemPromptWithoutId(name=name, content="x"), "alice", False) + + for user_id in ("alice", None): + listed = [prompt.id for prompt in svc.get_many(user_id=user_id) if prompt.user_id == "alice"] + assert listed == ["p2", "p3", "p1"] diff --git a/tests/app/services/test_wildcard_records.py b/tests/app/services/test_wildcard_records.py index ea594b4e0f7..d759f00b711 100644 --- a/tests/app/services/test_wildcard_records.py +++ b/tests/app/services/test_wildcard_records.py @@ -1,11 +1,10 @@ -import logging -import sqlite3 - import pytest from dynamicprompts.generators import CombinatorialPromptGenerator +from sqlalchemy import delete, insert -from invokeai.app.services.config.config_default import InvokeAIAppConfig -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import ForeignKeyViolation +from invokeai.app.services.shared.database.schema.users import users from invokeai.app.services.wildcard_records.wildcard_records_common import ( WildcardChanges, WildcardNameConflictError, @@ -13,31 +12,27 @@ WildcardWithoutId, build_wildcard_manager, ) -from invokeai.app.services.wildcard_records.wildcard_records_sqlite import SqliteWildcardRecordsStorage +from invokeai.app.services.wildcard_records.wildcard_records_default import WildcardRecordsStorage from invokeai.app.util.dynamicprompts import expand_dynamic_prompt -from tests.fixtures.sqlite_database import create_mock_sqlite_database @pytest.fixture -def db() -> SqliteDatabase: - database = create_mock_sqlite_database(InvokeAIAppConfig(use_memory_db=True), logging.getLogger(__name__)) +def store(database: Database) -> WildcardRecordsStorage: # Wildcards cascade from their owner, so the owners have to exist. Inserted # directly rather than through UserService: these tests are about storage, and # the only columns the foreign key needs are the ones without a default. - with database.transaction() as cursor: - cursor.executemany( - "INSERT INTO users (user_id, email, password_hash) VALUES (?, ?, ?);", - [(user_id, f"{user_id}@example.com", "unused") for user_id in ("user-1", "user-2")], + with database.begin(write=True) as conn: + conn.execute( + insert(users), + [ + {"user_id": user_id, "email": f"{user_id}@example.com", "password_hash": "unused"} + for user_id in ("user-1", "user-2") + ], ) - return database - - -@pytest.fixture -def store(db: SqliteDatabase) -> SqliteWildcardRecordsStorage: - return SqliteWildcardRecordsStorage(db=db) + return WildcardRecordsStorage(database) -def test_create_and_list_is_scoped_to_the_owner(store: SqliteWildcardRecordsStorage) -> None: +def test_create_and_list_is_scoped_to_the_owner(store: WildcardRecordsStorage) -> None: store.create(WildcardWithoutId(name="colors", values=["red", "green"]), user_id="user-1") store.create(WildcardWithoutId(name="animals", values=["cat"]), user_id="user-2") @@ -45,14 +40,14 @@ def test_create_and_list_is_scoped_to_the_owner(store: SqliteWildcardRecordsStor assert [w.name for w in store.get_many("user-2")] == ["animals"] -def test_get_many_is_ordered_by_name(store: SqliteWildcardRecordsStorage) -> None: +def test_get_many_is_ordered_by_name(store: WildcardRecordsStorage) -> None: for name in ["zebra", "apple", "mango"]: store.create(WildcardWithoutId(name=name, values=["x"]), user_id="user-1") assert [w.name for w in store.get_many("user-1")] == ["apple", "mango", "zebra"] -def test_a_name_is_unique_per_user_but_not_across_users(store: SqliteWildcardRecordsStorage) -> None: +def test_a_name_is_unique_per_user_but_not_across_users(store: WildcardRecordsStorage) -> None: store.create(WildcardWithoutId(name="colors", values=["red"]), user_id="user-1") store.create(WildcardWithoutId(name="colors", values=["blue"]), user_id="user-2") @@ -60,16 +55,25 @@ def test_a_name_is_unique_per_user_but_not_across_users(store: SqliteWildcardRec store.create(WildcardWithoutId(name="colors", values=["green"]), user_id="user-1") -def test_update_changes_name_and_values(store: SqliteWildcardRecordsStorage) -> None: +def test_a_created_wildcard_is_stored_as_returned(store: WildcardRecordsStorage) -> None: + created = store.create(WildcardWithoutId(name="colors", values=[" red ", "green"]), user_id="user-1") + + assert (created.name, created.values, created.user_id) == ("colors", ["red", "green"], "user-1") + assert store.get(created.id) == created + + +def test_update_changes_name_and_values_of_only_that_wildcard(store: WildcardRecordsStorage) -> None: created = store.create(WildcardWithoutId(name="colors", values=["red"]), user_id="user-1") + bystander = store.create(WildcardWithoutId(name="animals", values=["cat"]), user_id="user-2") updated = store.update(created.id, WildcardChanges(name="colours", values=["red", "green"])) assert updated.name == "colours" assert updated.values == ["red", "green"] + assert store.get(bystander.id) == bystander -def test_update_to_a_taken_name_conflicts(store: SqliteWildcardRecordsStorage) -> None: +def test_update_to_a_taken_name_conflicts(store: WildcardRecordsStorage) -> None: store.create(WildcardWithoutId(name="colors", values=["red"]), user_id="user-1") other = store.create(WildcardWithoutId(name="animals", values=["cat"]), user_id="user-1") @@ -77,36 +81,36 @@ def test_update_to_a_taken_name_conflicts(store: SqliteWildcardRecordsStorage) - store.update(other.id, WildcardChanges(name="colors")) -def test_delete_removes_the_record(store: SqliteWildcardRecordsStorage) -> None: +def test_delete_removes_the_record(store: WildcardRecordsStorage) -> None: created = store.create(WildcardWithoutId(name="colors", values=["red"]), user_id="user-1") + bystander = store.create(WildcardWithoutId(name="colors", values=["red"]), user_id="user-2") store.delete(created.id) assert store.get_many("user-1") == [] + assert store.get_many("user-2") == [bystander] with pytest.raises(WildcardNotFoundError): store.get(created.id) -def test_deleting_the_owner_takes_their_wildcards_with_it( - db: SqliteDatabase, store: SqliteWildcardRecordsStorage -) -> None: +def test_deleting_the_owner_takes_their_wildcards_with_it(database: Database, store: WildcardRecordsStorage) -> None: # `UserService.delete` is a bare DELETE FROM users; without the cascade these # rows would survive their owner, unreachable but forever. kept = store.create(WildcardWithoutId(name="colors", values=["red"]), user_id="user-2") orphaned = store.create(WildcardWithoutId(name="colors", values=["red"]), user_id="user-1") - with db.transaction() as cursor: - cursor.execute("DELETE FROM users WHERE user_id = ?;", ("user-1",)) + with database.begin(write=True) as conn: + conn.execute(delete(users).where(users.c.user_id == "user-1")) with pytest.raises(WildcardNotFoundError): store.get(orphaned.id) assert [w.id for w in store.get_many("user-2")] == [kept.id] -def test_an_unknown_owner_is_not_reported_as_a_name_conflict(store: SqliteWildcardRecordsStorage) -> None: +def test_an_unknown_owner_is_not_reported_as_a_name_conflict(store: WildcardRecordsStorage) -> None: # The name is free; it is the owner that does not exist. Mapping this to a # conflict would send the caller renaming a wildcard that is fine. - with pytest.raises(sqlite3.IntegrityError): + with pytest.raises(ForeignKeyViolation): store.create(WildcardWithoutId(name="colors", values=["red"]), user_id="ghost") @@ -140,7 +144,7 @@ def test_values_are_trimmed_and_blanks_dropped() -> None: assert WildcardWithoutId(name="colors", values=[" red ", "", " ", "green"]).values == ["red", "green"] -def test_build_wildcard_manager_resolves_a_users_wildcards(store: SqliteWildcardRecordsStorage) -> None: +def test_build_wildcard_manager_resolves_a_users_wildcards(store: WildcardRecordsStorage) -> None: store.create(WildcardWithoutId(name="colors", values=["red", "green", "blue"]), user_id="user-1") manager = build_wildcard_manager(store.get_many("user-1")) @@ -152,7 +156,7 @@ def test_build_wildcard_manager_resolves_a_users_wildcards(store: SqliteWildcard ] -def test_build_wildcard_manager_omits_valueless_wildcards(store: SqliteWildcardRecordsStorage) -> None: +def test_build_wildcard_manager_omits_valueless_wildcards(store: WildcardRecordsStorage) -> None: # A name that is "known" but expands to nothing would silently drop the prompt, so it must stay # unresolvable and be reported as missing instead. store.create(WildcardWithoutId(name="empty", values=[]), user_id="user-1") @@ -166,7 +170,7 @@ def test_build_wildcard_manager_omits_valueless_wildcards(store: SqliteWildcardR @pytest.mark.timeout(20) def test_a_resolvable_wildcard_in_a_variant_generates_instead_of_hanging( - store: SqliteWildcardRecordsStorage, + store: WildcardRecordsStorage, ) -> None: store.create(WildcardWithoutId(name="colors", values=["red", "green"]), user_id="user-1") manager = build_wildcard_manager(store.get_many("user-1")) @@ -177,7 +181,7 @@ def test_a_resolvable_wildcard_in_a_variant_generates_instead_of_hanging( assert sorted(expanded.prompts) == ["green", "red", "x"] -def test_one_users_wildcards_do_not_resolve_for_another(store: SqliteWildcardRecordsStorage) -> None: +def test_one_users_wildcards_do_not_resolve_for_another(store: WildcardRecordsStorage) -> None: store.create(WildcardWithoutId(name="colors", values=["red", "green"]), user_id="user-1") other = expand_dynamic_prompt( diff --git a/tests/app/services/users/test_last_admin_invariant.py b/tests/app/services/users/test_last_admin_invariant.py index 10ecda6e269..65bfc882b70 100644 --- a/tests/app/services/users/test_last_admin_invariant.py +++ b/tests/app/services/users/test_last_admin_invariant.py @@ -18,13 +18,19 @@ """ import threading -from logging import Logger +from collections.abc import Callable +from typing import Any import pytest +from sqlalchemy import update from invokeai.app.services.auth.password_utils import hash_password -from invokeai.app.services.shared.media_references import create_media_references_table -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.locks import DatabaseLock +from invokeai.app.services.shared.database.queries.users import UserQueries +from invokeai.app.services.shared.database.schema.users import users as users_table +from invokeai.app.services.users import users_default from invokeai.app.services.users.users_common import ( SYSTEM_USER_ID, LastAdministratorError, @@ -38,31 +44,8 @@ @pytest.fixture -def db() -> SqliteDatabase: - db = SqliteDatabase(db_path=None, logger=Logger("test_last_admin"), verbose=False) - db._conn.execute(""" - CREATE TABLE users ( - user_id TEXT NOT NULL PRIMARY KEY, - email TEXT NOT NULL UNIQUE, - display_name TEXT, - password_hash TEXT NOT NULL, - is_admin BOOLEAN NOT NULL DEFAULT FALSE, - is_active BOOLEAN NOT NULL DEFAULT TRUE, - 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')), - last_login_at DATETIME, - token_epoch INTEGER NOT NULL DEFAULT 0 - ); - """) - # Deleting an account also drops the media references its documents held. - create_media_references_table(db._conn.cursor()) - db._conn.commit() - return db - - -@pytest.fixture -def users(db: SqliteDatabase) -> UserService: - return UserService(db) +def users(database: Database) -> UserService: + return UserService(database) def _make(users: UserService, email: str, *, is_admin: bool) -> str: @@ -242,23 +225,160 @@ def demote(user_id: str) -> None: assert sum(isinstance(e, LastAdministratorError) for e in errors) == 1 -# endregion +def _while_another_transaction_holds_the_lock( + database: Database, change: Callable[[], None], holders_work: Callable[[Queries], None] +) -> list[BaseException]: + """Runs `change` while another transaction holds `ADMIN_ACCOUNTS`; that transaction then does `holders_work` + and commits, and `change` goes on. Returns what `change` raised.""" + locked = threading.Event() + proceed = threading.Event() -# region the system account + def hold() -> None: + with database.queries.transaction() as q: + q.locks.acquire(DatabaseLock.ADMIN_ACCOUNTS) + locked.set() + proceed.wait(timeout=30) + holders_work(q) + + errors: list[BaseException] = [] + + def run_change() -> None: + try: + change() + except BaseException as e: # noqa: BLE001 - asserted on by the caller + errors.append(e) + + holder = threading.Thread(target=hold) + holder.start() + try: + assert locked.wait(timeout=10), "the lock was not taken" + changing = threading.Thread(target=run_change) + changing.start() + changing.join(timeout=0.5) + assert changing.is_alive(), "the change did not wait for the lock" + finally: + proceed.set() + holder.join(timeout=10) + changing.join(timeout=10) + assert not changing.is_alive() + return errors + + +@pytest.fixture +def quick_hashing(monkeypatch: pytest.MonkeyPatch) -> None: + """Hashing takes a large part of a second on its own; the tests below time how long a change waits.""" + monkeypatch.setattr(users_default, "hash_password", lambda password: f"hash of {password}") + + +@pytest.mark.parametrize("change", ["delete", "demote", "deactivate"]) +def test_a_waiting_change_counts_the_admins_the_lock_holder_left( + database: Database, users: UserService, change: str +) -> None: + """On a server transactions run side by side: what decides is that each change takes the lock before it + counts, so it counts what the transaction holding the lock committed. (On SQLite the held transaction blocks + every other one anyway.)""" + first = _make(users, "a1@test.com", is_admin=True) + second = _make(users, "a2@test.com", is_admin=True) + + def revoke_second() -> None: + if change == "delete": + users.delete(second) + else: + revoke = UserUpdateRequest(is_admin=False) if change == "demote" else UserUpdateRequest(is_active=False) + users.update(second, revoke, strict_password_checking=False) + + errors = _while_another_transaction_holds_the_lock( + database, revoke_second, lambda q: q.users.update(first, is_admin=False) + ) + + assert [type(e) for e in errors] == [LastAdministratorError] + assert users.count_admins() == 1 -def _seed_system_user(db: SqliteDatabase) -> None: - """The row migration_27 creates: active, non-admin, and with an empty password hash.""" - db._conn.execute( - """ - INSERT INTO users (user_id, email, display_name, password_hash, is_admin, is_active) - VALUES ('system', 'system@system.invokeai', 'System', '', FALSE, TRUE); - """ +@pytest.mark.usefixtures("quick_hashing") +def test_a_waiting_first_admin_setup_sees_the_admin_the_lock_holder_created( + database: Database, users: UserService +) -> None: + def setup() -> None: + users.create_admin( + UserCreateRequest(email="second@test.com", password=PASSWORD), strict_password_checking=False + ) + + errors = _while_another_transaction_holds_the_lock( + database, + setup, + lambda q: q.users.insert( + user_id="first", email="first@test.com", display_name=None, password_hash="hash", is_admin=True + ), ) - db._conn.commit() + assert [str(e) for e in errors] == ["Admin user already exists"] + assert users.count_admins() == 1 + + +@pytest.mark.usefixtures("quick_hashing") +@pytest.mark.parametrize("change", ["create_admin_account", "promote", "reactivate"]) +def test_first_admin_setup_waits_for_an_admin_being_added( + database: Database, users: UserService, monkeypatch: pytest.MonkeyPatch, change: str +) -> None: + """`invoke-useradd --admin` or `invoke-usermod` can add an administrator while the unauthenticated setup is + open: the setup must not count before that change has committed.""" + if change == "promote": + target = _make(users, "user@test.com", is_admin=False) + elif change == "reactivate": + target = _make(users, "admin@test.com", is_admin=True) + with database.queries.transaction() as q: + q.users.update(target, is_active=False) + written = threading.Event() + proceed = threading.Event() + patched = "insert" if change == "create_admin_account" else "update" + real = getattr(UserQueries, patched) + + def write_then_wait(self: UserQueries, *args: Any, **kwargs: Any) -> None: + real(self, *args, **kwargs) + written.set() + assert proceed.wait(timeout=30) + + monkeypatch.setattr(UserQueries, patched, write_then_wait) + + def add_admin() -> None: + if change == "create_admin_account": + users.create(UserCreateRequest(email="added@test.com", password=PASSWORD, is_admin=True), False) + else: + added = UserUpdateRequest(is_admin=True) if change == "promote" else UserUpdateRequest(is_active=True) + users.update(target, added, strict_password_checking=False) + + adding = threading.Thread(target=add_admin) + adding.start() + errors: list[BaseException] = [] + + def setup() -> None: + try: + users.create_admin(UserCreateRequest(email="setup@test.com", password=PASSWORD), False) + except BaseException as e: # noqa: BLE001 - asserted on below + errors.append(e) + + try: + assert written.wait(timeout=10), "the admin was not added" + setting_up = threading.Thread(target=setup) + setting_up.start() + setting_up.join(timeout=0.5) + assert setting_up.is_alive(), "the setup did not wait for the lock" + finally: + proceed.set() + adding.join(timeout=10) + setting_up.join(timeout=10) + + assert [str(e) for e in errors] == ["Admin user already exists"] + assert users.count_admins() == 1 -def test_the_system_user_cannot_be_promoted(db: SqliteDatabase, users: UserService) -> None: + +# endregion + +# region the system account + + +def test_the_system_user_cannot_be_promoted(users: UserService) -> None: """`count_admins()` counts admin rows, but the invariant that matters is "an admin who can log in". The system row is active and can never authenticate — it has no password — so promoting it would inflate the count with an unusable administrator, which is enough @@ -270,7 +390,6 @@ def test_the_system_user_cannot_be_promoted(db: SqliteDatabase, users: UserServi leaving the instance with no usable administration and no authenticated way back. """ - _seed_system_user(db) admin = _make(users, "admin@test.com", is_admin=True) with pytest.raises(SystemUserProtectedError): @@ -283,10 +402,9 @@ def test_the_system_user_cannot_be_promoted(db: SqliteDatabase, users: UserServi users.update(admin, UserUpdateRequest(is_admin=False), strict_password_checking=False) -def test_the_system_user_cannot_be_given_a_password(db: SqliteDatabase, users: UserService) -> None: +def test_the_system_user_cannot_be_given_a_password(users: UserService) -> None: """The other end of the same hole: a password turns the owner of every pre-multiuser board, image, and workflow into a login account.""" - _seed_system_user(db) with pytest.raises(SystemUserProtectedError): users.update(SYSTEM_USER_ID, UserUpdateRequest(password=PASSWORD), strict_password_checking=False) @@ -294,11 +412,10 @@ def test_the_system_user_cannot_be_given_a_password(db: SqliteDatabase, users: U assert users.authenticate("system@system.invokeai", PASSWORD) is None -def test_the_system_user_cannot_be_deleted_or_deactivated(db: SqliteDatabase, users: UserService) -> None: +def test_the_system_user_cannot_be_deleted_or_deactivated(users: UserService) -> None: """The routes already refuse both, but `invoke-userdel` / `invoke-usermod` construct this service directly and never reach a route — the same reason the last-admin guard lives here.""" - _seed_system_user(db) with pytest.raises(SystemUserProtectedError): users.delete(SYSTEM_USER_ID) @@ -309,16 +426,15 @@ def test_the_system_user_cannot_be_deleted_or_deactivated(db: SqliteDatabase, us assert system is not None and system.is_active is True -def test_renaming_the_system_user_is_allowed(db: SqliteDatabase, users: UserService) -> None: +def test_renaming_the_system_user_is_allowed(users: UserService) -> None: """Not a blanket lock on the row — only the changes that would make it dangerous.""" - _seed_system_user(db) updated = users.update(SYSTEM_USER_ID, UserUpdateRequest(display_name="Renamed"), strict_password_checking=False) assert updated.display_name == "Renamed" -def test_a_system_row_carrying_a_password_still_cannot_log_in(db: SqliteDatabase, users: UserService) -> None: +def test_a_system_row_carrying_a_password_still_cannot_log_in(database: Database, users: UserService) -> None: """The guard above only stops a password being set *from now on*. An instance that set one through the old `PATCH /auth/users/system` hole still carries @@ -327,27 +443,25 @@ def test_a_system_row_carrying_a_password_still_cannot_log_in(db: SqliteDatabase migration clears it, but migrations run once and cannot reach a row damaged afterwards by direct SQL, so `authenticate` refuses the account outright whatever the row holds. """ - _seed_system_user(db) - db._conn.execute( - "UPDATE users SET password_hash = ? WHERE user_id = 'system'", - (hash_password(PASSWORD),), - ) - db._conn.commit() + with database.begin(write=True) as conn: + conn.execute( + update(users_table) + .where(users_table.c.user_id == SYSTEM_USER_ID) + .values(password_hash=hash_password(PASSWORD)) + ) assert users.authenticate("system@system.invokeai", PASSWORD) is None -def test_refusing_the_system_account_does_not_block_other_logins(db: SqliteDatabase, users: UserService) -> None: +def test_refusing_the_system_account_does_not_block_other_logins(users: UserService) -> None: """The refusal is keyed on the user id, not on anything a real account shares.""" - _seed_system_user(db) _make(users, "real@test.com", is_admin=False) assert users.authenticate("real@test.com", PASSWORD) is not None -def test_the_system_error_is_a_value_error(db: SqliteDatabase, users: UserService) -> None: +def test_the_system_error_is_a_value_error(users: UserService) -> None: """Same reason as the last-admin error: existing route and CLI handlers catch ValueError.""" - _seed_system_user(db) with pytest.raises(ValueError): users.delete(SYSTEM_USER_ID) diff --git a/tests/app/services/users/test_user_service.py b/tests/app/services/users/test_user_service.py index 11b58390a18..12fc7e9018b 100644 --- a/tests/app/services/users/test_user_service.py +++ b/tests/app/services/users/test_user_service.py @@ -1,52 +1,29 @@ """Tests for user service.""" import threading -from logging import Logger import pytest +from pydantic import ValidationError +from sqlalchemy import insert, select, update -from invokeai.app.services.shared.media_references import create_media_references_table -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries.base import IN_CHUNK +from invokeai.app.services.shared.database.schema.media_references import media_references +from invokeai.app.services.shared.database.schema.users import users as users_table from invokeai.app.services.users import users_default -from invokeai.app.services.users.users_common import UserCreateRequest, UserUpdateRequest -from invokeai.app.services.users.users_default import USER_LOOKUP_CHUNK_SIZE, UserService +from invokeai.app.services.users.users_common import ( + MAX_EMAIL_LENGTH, + SYSTEM_USER_ID, + UserCreateRequest, + UserUpdateRequest, +) +from invokeai.app.services.users.users_default import UserService @pytest.fixture -def logger() -> Logger: - """Create a logger for testing.""" - return Logger("test_user_service") - - -@pytest.fixture -def db(logger: Logger) -> SqliteDatabase: - """Create an in-memory database for testing.""" - db = SqliteDatabase(db_path=None, logger=logger, verbose=False) - # Create users table manually for testing - db._conn.execute(""" - CREATE TABLE users ( - user_id TEXT NOT NULL PRIMARY KEY, - email TEXT NOT NULL UNIQUE, - display_name TEXT, - password_hash TEXT NOT NULL, - is_admin BOOLEAN NOT NULL DEFAULT FALSE, - is_active BOOLEAN NOT NULL DEFAULT TRUE, - 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')), - last_login_at DATETIME, - token_epoch INTEGER NOT NULL DEFAULT 0 - ); - """) - # Deleting an account also drops the media references its documents held. - create_media_references_table(db._conn.cursor()) - db._conn.commit() - return db - - -@pytest.fixture -def user_service(db: SqliteDatabase) -> UserService: +def user_service(database: Database) -> UserService: """Create a user service for testing.""" - return UserService(db) + return UserService(database) def test_create_user(user_service: UserService): @@ -182,19 +159,47 @@ def test_delete_user(user_service: UserService): assert retrieved_user is None -def test_delete_user_drops_references_only_of_documents_that_cascade(user_service: UserService, db: SqliteDatabase): - user = user_service.create( - UserCreateRequest(email="test@example.com", display_name="T", password="TestPassword123") - ) - db._conn.executemany( - "INSERT INTO media_references VALUES (?, ?, 'owner', 'image', 'a.png');", - [(kind, user.user_id) for kind in ("project", "client_state", "workflow", "quarantined_project")], - ) +def test_delete_user_drops_references_only_of_documents_that_cascade(user_service: UserService, database: Database): + leaving = user_service.create(UserCreateRequest(email="leaving@example.com", password="TestPassword123")).user_id + staying = user_service.create(UserCreateRequest(email="staying@example.com", password="TestPassword123")).user_id + kinds = ("project", "client_state", "workflow", "quarantined_project") + with database.begin(write=True) as conn: + conn.execute( + insert(media_references), + [ + {"owner_kind": kind, "user_id": user, "owner_id": "owner", "media_kind": "image", "media_name": "a.png"} + for kind in kinds + for user in (leaving, staying) + ], + ) - user_service.delete(user.user_id) + user_service.delete(leaving) + + with database.begin(write=False) as conn: + remaining = { + (user_id, kind) + for user_id, kind in conn.execute(select(media_references.c.user_id, media_references.c.owner_kind)) + } + assert remaining == {(leaving, "quarantined_project"), (leaving, "workflow")} | {(staying, kind) for kind in kinds} - rows = db._conn.execute("SELECT owner_kind FROM media_references ORDER BY owner_kind;").fetchall() - assert [row[0] for row in rows] == ["quarantined_project", "workflow"] + +def test_an_account_created_meanwhile_with_the_same_email_is_refused( + user_service: UserService, database: Database, monkeypatch: pytest.MonkeyPatch +): + """The email is checked before the password is hashed, which takes a while; the unique key decides.""" + real_hash_password = users_default.hash_password + + def hash_while_someone_else_signs_up(password: str) -> str: + with database.begin(write=True) as conn: + conn.execute(insert(users_table).values(user_id="first", email="taken@example.com", password_hash="hash")) + return real_hash_password(password) + + monkeypatch.setattr(users_default, "hash_password", hash_while_someone_else_signs_up) + + with pytest.raises(ValueError, match="Failed to create user"): + user_service.create(UserCreateRequest(email="taken@example.com", password="TestPassword123")) + + assert [user.user_id for user in user_service.list_users() if user.email == "taken@example.com"] == ["first"] def test_authenticate_valid_credentials(user_service: UserService): @@ -213,6 +218,19 @@ def test_authenticate_valid_credentials(user_service: UserService): assert authenticated_user.last_login_at is not None +def test_a_login_is_recorded_on_that_account_only(user_service: UserService): + user_service.create(UserCreateRequest(email="first@example.com", password="TestPassword123")) + second = user_service.create(UserCreateRequest(email="second@example.com", password="TestPassword123")) + + logged_in = user_service.authenticate("first@example.com", "TestPassword123") + + assert logged_in is not None + first = user_service.get(logged_in.user_id) + assert first is not None and first.last_login_at == logged_in.last_login_at + unchanged = user_service.get(second.user_id) + assert unchanged is not None and unchanged.last_login_at is None + + def test_authenticate_invalid_password(user_service: UserService): """Test authenticating with invalid password.""" user_data = UserCreateRequest( @@ -343,13 +361,33 @@ def test_list_users(user_service: UserService): ) user_service.create(user_data) + # Every database starts with the system account. users = user_service.list_users() - assert len(users) == 5 + assert SYSTEM_USER_ID in {user.user_id for user in users} + assert len(users) == 6 limited_users = user_service.list_users(limit=2) assert len(limited_users) == 2 +def test_list_users_lists_newest_first(user_service: UserService, database: Database): + created = { + user_service.create(UserCreateRequest(email=f"user{day}@example.com", password="TestPassword123")).user_id: day + for day in (2, 3, 1) + } + with database.begin(write=True) as conn: + for user_id, day in created.items(): + conn.execute( + update(users_table) + .where(users_table.c.user_id == user_id) + .values(created_at=f"2099-01-0{day} 00:00:00.000") + ) + + listed = [user.email for user in user_service.list_users(limit=3)] + + assert listed == ["user3@example.com", "user2@example.com", "user1@example.com"] + + def test_get_many_returns_users_keyed_by_id(user_service: UserService): """Batch lookup: dedups input, keys by user_id, and omits unknown ids.""" created = [ @@ -398,8 +436,17 @@ def test_get_many_chunks_beyond_sqlite_parameter_limit(user_service: UserService user = user_service.create( UserCreateRequest(email="chunked@example.com", display_name="Chunked", password="TestPassword123") ) - ids = [f"missing-{index}" for index in range(USER_LOOKUP_CHUNK_SIZE * 2 + 5)] + [user.user_id] + ids = [f"missing-{index}" for index in range(IN_CHUNK * 2 + 5)] + [user.user_id] users = user_service.get_many(ids) assert set(users) == {user.user_id} + + +def test_an_address_at_a_special_use_domain_is_accepted_up_to_the_longest_address() -> None: + domain = "@studio.local" + longest = "a" * (MAX_EMAIL_LENGTH - len(domain)) + domain + + assert UserCreateRequest(email=longest, password="TestPassword123").email == longest + with pytest.raises(ValidationError, match=f"at most {MAX_EMAIL_LENGTH} characters"): + UserCreateRequest(email="a" + longest, password="TestPassword123") diff --git a/tests/app/services/video_files/test_video_files_disk.py b/tests/app/services/video_files/test_video_files_disk.py index de941b295df..8f4afdb98fe 100644 --- a/tests/app/services/video_files/test_video_files_disk.py +++ b/tests/app/services/video_files/test_video_files_disk.py @@ -206,7 +206,7 @@ def test_start_keeps_staged_delete_when_the_record_cannot_be_read(storage: DiskV The recovery decides between purging the staged files and restoring them by asking whether the record is still there, so it needs that read to distinguish "gone" from "could not - look". `SqliteVideoRecordStorage.get` used to translate every sqlite3.Error into + look". `VideoRecordStorage.get` used to translate every sqlite3.Error into VideoRecordNotFoundException, which made an unreadable database delete the user's video files outright. Now the staged copy survives for a later attempt. """ diff --git a/tests/app/services/video_records/test_video_records.py b/tests/app/services/video_records/test_video_records.py new file mode 100644 index 00000000000..f0b4565e2b6 --- /dev/null +++ b/tests/app/services/video_records/test_video_records.py @@ -0,0 +1,688 @@ +"""Video records on every database backend. + +The video queries are written apart from the image ones, so they are tested apart: listings and their filters, +per-user isolation, fields, intermediate deletes and the `media_origin` marker. + +Per-user isolation covers JPPhoto's code-review finding (PR #9163): when ``board_id`` was omitted from /v1/videos/ +and /v1/videos/names, a non-admin caller saw every user's videos. +""" + +import json +from collections.abc import Sequence +from typing import Any + +import pytest +from sqlalchemy import insert, update +from sqlalchemy.exc import DBAPIError + +from invokeai.app.invocations.fields import MetadataFieldValidator +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage +from invokeai.app.services.board_video_records.board_video_records_default import BoardVideoRecordStorage +from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.schema.videos import videos +from invokeai.app.services.shared.pagination import SQLiteDirection +from invokeai.app.services.users.users_common import UserCreateRequest +from invokeai.app.services.users.users_default import UserService +from invokeai.app.services.video_records.video_records_common import VideoRecordChanges, VideoRecordNotFoundException +from invokeai.app.services.video_records.video_records_default import VideoRecordStorage +from tests.fixtures.database import capture_statements +from tests.fixtures.races import while_in_flight + + +@pytest.fixture +def store(database: Database) -> VideoRecordStorage: + return VideoRecordStorage(database) + + +def _drop_the_videos_table(database: Database) -> None: + with database.begin(write=True) as conn: + conn.exec_driver_sql("DROP TABLE videos") + + +def _save(store: VideoRecordStorage, name: str, user_id: str, is_intermediate: bool = False) -> None: + store.save( + video_name=name, + video_origin=ResourceOrigin.INTERNAL, + video_category=ImageCategory.GENERAL, + width=64, + height=64, + duration=1.0, + fps=8.0, + has_workflow=False, + is_intermediate=is_intermediate, + user_id=user_id, + ) + + +@pytest.fixture +def seeded_store(store: VideoRecordStorage) -> VideoRecordStorage: + # Two videos per user; all without board association (the bug occurred when board_id + # was omitted from the query). + _save(store, "alice_1.mp4", user_id="alice") + _save(store, "alice_2.mp4", user_id="alice") + _save(store, "bob_1.mp4", user_id="bob") + _save(store, "bob_2.mp4", user_id="bob") + return store + + +class TestGetManyOmittedBoardIdMultiuser: + """get_many() with board_id=None must filter by user_id for non-admin callers.""" + + def test_non_admin_only_sees_own_videos(self, seeded_store: VideoRecordStorage) -> None: + result = seeded_store.get_many(user_id="alice", is_admin=False) + names = {v.video_name for v in result.items} + assert names == {"alice_1.mp4", "alice_2.mp4"} + assert result.total == 2 + + def test_admin_sees_every_users_videos(self, seeded_store: VideoRecordStorage) -> None: + result = seeded_store.get_many(user_id="alice", is_admin=True) + names = {v.video_name for v in result.items} + assert names == {"alice_1.mp4", "alice_2.mp4", "bob_1.mp4", "bob_2.mp4"} + + def test_no_user_id_returns_all(self, seeded_store: VideoRecordStorage) -> None: + # No user_id means the caller is bypassing user filtering entirely (e.g. internal calls). + result = seeded_store.get_many(user_id=None, is_admin=False) + names = {v.video_name for v in result.items} + assert names == {"alice_1.mp4", "alice_2.mp4", "bob_1.mp4", "bob_2.mp4"} + + +class TestGetVideoNamesOmittedBoardIdMultiuser: + """get_video_names() with board_id=None must filter by user_id for non-admin callers.""" + + def test_non_admin_only_sees_own_videos(self, seeded_store: VideoRecordStorage) -> None: + result = seeded_store.get_video_names(user_id="alice", is_admin=False) + assert set(result.video_names) == {"alice_1.mp4", "alice_2.mp4"} + assert result.total_count == 2 + + def test_admin_sees_every_users_videos(self, seeded_store: VideoRecordStorage) -> None: + result = seeded_store.get_video_names(user_id="alice", is_admin=True) + assert set(result.video_names) == {"alice_1.mp4", "alice_2.mp4", "bob_1.mp4", "bob_2.mp4"} + + def test_explicit_none_board_still_isolates(self, seeded_store: VideoRecordStorage) -> None: + # The "none" sentinel (uncategorized) must also apply the user filter — this was the + # only path that was correct *before* the fix; the test guards against accidental + # regression there too. + result = seeded_store.get_video_names(board_id="none", user_id="alice", is_admin=False) + assert set(result.video_names) == {"alice_1.mp4", "alice_2.mp4"} + + +class TestDeterministicOrdering: + @staticmethod + def _set_same_timestamp(database: Database) -> None: + with database.begin(write=True) as conn: + conn.execute(update(videos).values(created_at="2026-07-20 00:00:00.000")) + + def test_same_timestamp_pages_have_stable_mirrored_order( + self, database: Database, seeded_store: VideoRecordStorage + ) -> None: + self._set_same_timestamp(database) + + ascending = [ + seeded_store.get_many( + offset=offset, + limit=1, + starred_first=False, + order_dir=SQLiteDirection.Ascending, + user_id="alice", + ) + .items[0] + .video_name + for offset in range(2) + ] + descending = [ + seeded_store.get_many( + offset=offset, + limit=1, + starred_first=False, + order_dir=SQLiteDirection.Descending, + user_id="alice", + ) + .items[0] + .video_name + for offset in range(2) + ] + + assert ascending == ["alice_1.mp4", "alice_2.mp4"] + assert descending == list(reversed(ascending)) + + def test_same_timestamp_name_list_has_stable_mirrored_order( + self, database: Database, seeded_store: VideoRecordStorage + ) -> None: + self._set_same_timestamp(database) + + ascending = seeded_store.get_video_names( + starred_first=False, + order_dir=SQLiteDirection.Ascending, + user_id="alice", + ).video_names + descending = seeded_store.get_video_names( + starred_first=False, + order_dir=SQLiteDirection.Descending, + user_id="alice", + ).video_names + + assert ascending == ["alice_1.mp4", "alice_2.mp4"] + assert descending == list(reversed(ascending)) + + def test_board_cover_uses_name_as_same_timestamp_tie_breaker( + self, database: Database, store: VideoRecordStorage + ) -> None: + board = BoardRecordStorage(database).save("Board", "system") + _save(store, "a.mp4", user_id="system") + _save(store, "b.mp4", user_id="system") + self._set_same_timestamp(database) + for name in ("a.mp4", "b.mp4"): + BoardVideoRecordStorage(database).add_video_to_board(board.board_id, name) + + assert store.get_most_recent_video_for_board(board.board_id).video_name == "b.mp4" + + +class TestUserDeletionLifecycle: + """Documents the intended videos↔users lifecycle (JPPhoto PR #9163 July-10 follow-up). + + ``videos.user_id`` deliberately has no FK to ``users`` — exactly like ``images``, + ``boards`` and ``workflows``, whose user_id columns (migration_27) are index-only. + Deleting a user therefore leaves their videos in place instead of cascading a row + delete that would strand the files on disk; the orphaned records stay visible to + administrators (and only to administrators), who can clean them up or reassign them. + These tests pin that platform-wide behavior for videos so any future change to the + user-deletion story is a deliberate decision rather than an accident. + """ + + def test_videos_survive_owner_deletion_and_remain_admin_only(self, database: Database) -> None: + users = UserService(database) + store = VideoRecordStorage(database) + + owner = users.create( + UserCreateRequest( + email="doomed@example.com", + display_name="Doomed User", + password="TestPassword123", + is_admin=False, + ) + ) + _save(store, "doomed.mp4", user_id=owner.user_id) + + users.delete(owner.user_id) + assert users.get(owner.user_id) is None + + # The record survives, still attributed to the deleted owner... + assert store.get_user_id("doomed.mp4") == owner.user_id + # ...is visible to admins for cleanup... + admin_view = store.get_many(user_id="some-admin", is_admin=True) + assert "doomed.mp4" in {v.video_name for v in admin_view.items} + # ...and no regular user inherits it. + other_view = store.get_many(user_id="bystander", is_admin=False) + assert "doomed.mp4" not in {v.video_name for v in other_view.items} + + +@pytest.mark.sqlite_only # A dropped table would stay dropped for every later test on a shared server schema. +def test_get_propagates_a_storage_error_instead_of_reporting_the_row_missing( + database: Database, store: VideoRecordStorage +): + """An unreadable database must not be indistinguishable from a deleted video. + + `get` used to translate every sqlite3.Error into VideoRecordNotFoundException, which made + that exception mean "the row is absent, OR the read failed". Two callers act destructively + on it: `_assert_video_read_access` answers 404 for a positive not-found and the clients drop + their reference to the video on one, and the staged-delete recovery reads it as proof the + delete committed and purges the staged files. + """ + _save(store, "video-1.mp4", "user-1") + # A real storage failure rather than a patched one: the SELECT below cannot run at all, the + # same shape a locked or corrupt database presents. The row's absence is not what is being + # reported, and the caller must be able to tell. + _drop_the_videos_table(database) + + with pytest.raises(DBAPIError): + store.get("video-1.mp4") + + +def test_get_still_reports_a_positively_absent_row_as_missing(store: VideoRecordStorage): + """The narrowing stays exactly that narrow: absence is still absence.""" + with pytest.raises(VideoRecordNotFoundException): + store.get("never-existed.mp4") + + +def test_exists_reports_a_row_get_cannot_deserialize(database: Database, store: VideoRecordStorage): + """Presence, not readability. `get` would raise on an enum value this version does not know + — a row written by a newer one — and the refusal path reads that as absence, which would + report a live video gone.""" + _save(store, "video-1.mp4", "user-1") + with database.begin(write=True) as conn: + conn.execute( + update(videos).where(videos.c.video_name == "video-1.mp4").values(video_category="from_the_future") + ) + + with pytest.raises(ValueError): + store.get("video-1.mp4") + assert store.exists("video-1.mp4") is True + + +@pytest.mark.sqlite_only # See above. +def test_exists_propagates_a_storage_error(database: Database, store: VideoRecordStorage): + """ "Could not look" is not "not there" — the caller answers 404 on a False.""" + _drop_the_videos_table(database) + + with pytest.raises(DBAPIError): + store.exists("video-1.mp4") + + +@pytest.mark.sqlite_only # See above. +def test_get_metadata_propagates_a_storage_error(database: Database, store: VideoRecordStorage): + """As for `get`: a metadata read that fails is not a video that is gone.""" + _save(store, "video-1.mp4", "user-1") + _drop_the_videos_table(database) + + with pytest.raises(DBAPIError): + store.get_metadata("video-1.mp4") + + +def _save_with_metadata(store: VideoRecordStorage, name: str, metadata: str | None) -> None: + store.save( + video_name=name, + video_origin=ResourceOrigin.EXTERNAL, + video_category=ImageCategory.USER, + width=640, + height=360, + duration=1.0, + fps=24.0, + has_workflow=False, + is_intermediate=False, + metadata=metadata, + user_id="alice", + ) + + +def test_media_origin_is_projected_out_of_the_metadata_blob(store: VideoRecordStorage) -> None: + """The record carries `media_origin` so clients need no separate /metadata request. + + An audio upload the ingest converter wrapped into a waveform video is marked + `audio_upload`; the frontend starts such a reference on its soundtrack alone, and a + wrong value there costs a video reference slot and roughly doubles the packed sequence. + """ + _save_with_metadata(store, "wrapped.mp4", '{"media_origin": "audio_upload", "note": "kept"}') + + assert store.get("wrapped.mp4").media_origin == "audio_upload" + # The rest of the blob stays behind the /metadata route rather than riding every row. + assert store.get_metadata("wrapped.mp4") is not None + + +@pytest.mark.parametrize( + "metadata", + [ + pytest.param(None, id="no metadata at all"), + pytest.param("{}", id="metadata without the key"), + pytest.param('{"note": "kept"}', id="metadata with other keys"), + ], +) +def test_media_origin_is_none_when_unmarked(store: VideoRecordStorage, metadata: str | None) -> None: + """`json_extract` over a NULL or key-less blob yields NULL rather than raising.""" + _save_with_metadata(store, "plain.mp4", metadata) + + assert store.get("plain.mp4").media_origin is None + + +def test_media_origin_survives_a_listing(store: VideoRecordStorage) -> None: + """`get_many` selects the same columns, so a listed row carries the marker too.""" + _save_with_metadata(store, "wrapped.mp4", '{"media_origin": "audio_upload"}') + _save_with_metadata(store, "plain.mp4", None) + + listed = store.get_many(offset=0, limit=10, order_dir=SQLiteDirection.Descending, user_id="alice", is_admin=False) + origins = {record.video_name: record.media_origin for record in listed.items} + + assert origins == {"wrapped.mp4": "audio_upload", "plain.mp4": None} + + +@pytest.mark.parametrize( + "raw", + [ + pytest.param("1", id="a JSON number"), + pytest.param("true", id="a JSON boolean"), + pytest.param("1.5", id="a JSON float"), + pytest.param('""', id="an empty string, which the frontend also reads as no marker"), + # `json_extract` hands an object or array back as its SERIALIZED TEXT, so these + # arrive as `str` and an isinstance check alone would propagate them. + pytest.param('{"a": "b"}', id="a JSON object, returned as its text"), + pytest.param("[1, 2]", id="a JSON array, returned as its text"), + pytest.param('"has a space"', id="a string that is not marker-shaped"), + ], +) +def test_a_non_string_marker_does_not_break_the_record(store: VideoRecordStorage, raw: str) -> None: + """A client may store any JSON value under `media_origin`; the row must still deserialize. + + `MetadataField` validates only that the blob is an object, so `json_extract` can hand + back an int or a float. Passing one to the `Optional[str]` field raised while building + the record -- which took out the video's DTO *and* every gallery listing containing it, + so one odd upload broke the gallery. + """ + _save_with_metadata(store, "odd.mp4", '{"media_origin": ' + raw + "}") + + assert store.get("odd.mp4").media_origin is None + + +def test_metadata_that_is_not_an_object_leaves_the_marker_unset(store: VideoRecordStorage) -> None: + """`json_extract` over a non-object blob yields NULL rather than raising.""" + _save_with_metadata(store, "array.mp4", "[1, 2, 3]") + + assert store.get("array.mp4").media_origin is None + + +def test_a_malformed_metadata_blob_does_not_fail_the_query(store: VideoRecordStorage) -> None: + """`json_extract` RAISES on unparseable text, and this expression runs on every listed row. + + Unguarded, one malformed blob would fail the whole page — `get_many` deserializes rows in + a comprehension — rather than the one video. The column is plain TEXT with no CHECK, so + only convention keeps such a row out; the `json_valid` guard makes that not matter. + """ + _save_with_metadata(store, "bad.mp4", "not json at all") + _save_with_metadata(store, "good.mp4", '{"media_origin": "audio_upload"}') + + listed = store.get_many(offset=0, limit=10, order_dir=SQLiteDirection.Descending, user_id="alice", is_admin=False) + + assert {record.video_name: record.media_origin for record in listed.items} == { + "bad.mp4": None, + "good.mp4": "audio_upload", + } + + +def test_an_overlong_marker_is_dropped_rather_than_echoed_on_every_row(store: VideoRecordStorage) -> None: + """Upload metadata is client-supplied and unbounded, and this key now rides every row. + + Without a cap, one upload carrying a huge marker is echoed back on every gallery page + that includes the video. + """ + _save_with_metadata(store, "huge.mp4", json.dumps({"media_origin": "A" * 5000})) + + assert store.get("huge.mp4").media_origin is None + # A marker of a plausible length still passes. + _save_with_metadata(store, "fine.mp4", json.dumps({"media_origin": "audio_upload"})) + assert store.get("fine.mp4").media_origin == "audio_upload" + + +def test_only_videos_still_intermediate_are_deleted_and_reported_in_the_given_order(store: VideoRecordStorage) -> None: + for name in ("a.mp4", "b.mp4", "c.mp4"): + _save(store, name, "alice", is_intermediate=True) + store.update("b.mp4", VideoRecordChanges(is_intermediate=False)) + + deleted = store.delete_intermediates_by_names(["c.mp4", "b.mp4", "gone.mp4", "a.mp4"]) + + assert deleted == ["c.mp4", "a.mp4"] + assert [store.exists(name) for name in ("a.mp4", "b.mp4", "c.mp4")] == [False, True, False] + + +def test_a_promotion_in_flight_is_waited_for_and_keeps_the_record( + database: Database, store: VideoRecordStorage +) -> None: + """A delete that meets a promotion still in flight waits for it, and then keeps the promoted record.""" + _save(store, "tmp.mp4", "alice", is_intermediate=True) + _save(store, "promoted.mp4", "alice", is_intermediate=True) + deleted: list[str] = [] + + def promote(q: Queries) -> None: + q.videos.update("promoted.mp4", VideoRecordChanges(is_intermediate=False)) + + errors = while_in_flight( + database, promote, lambda: deleted.extend(store.delete_intermediates_by_names(["tmp.mp4", "promoted.mp4"])) + ) + + assert errors == [] + assert deleted == ["tmp.mp4"] + assert store.get("promoted.mp4").is_intermediate is False + + +def test_a_guard_narrows_what_is_deleted(store: VideoRecordStorage) -> None: + _save(store, "a.mp4", "alice", is_intermediate=True) + _save(store, "b.mp4", "alice", is_intermediate=True) + asked: list[list[str]] = [] + + def guard(q: Queries, names: Sequence[str]) -> list[str]: + asked.append(list(names)) + return [name for name in names if name != "b.mp4"] + + assert store.delete_intermediates_by_names(["a.mp4", "b.mp4"], guard) == ["a.mp4"] + assert asked == [["a.mp4", "b.mp4"]] + assert store.get("b.mp4").is_intermediate is True + + +class TestListingFilters: + """Each filter of a listing, on pages and names alike, for an administrator's view.""" + + @pytest.fixture + def board_id(self, database: Database, store: VideoRecordStorage) -> str: + board = BoardRecordStorage(database).save("Board", "alice") + # name, category, origin, metadata, starred, intermediate, on the board, created at + rows = [ + ( + "boarded.mp4", + ImageCategory.GENERAL, + ResourceOrigin.INTERNAL, + '{"prompt": "a 100% Cat_Photo"}', + True, + False, + True, + "01", + ), + ( + "user.mp4", + ImageCategory.USER, + ResourceOrigin.INTERNAL, + '{"prompt": "a 100x catXphoto"}', + False, + False, + False, + "02", + ), + ("external.mp4", ImageCategory.GENERAL, ResourceOrigin.EXTERNAL, None, False, False, False, "03"), + ("intermediate.mp4", ImageCategory.GENERAL, ResourceOrigin.INTERNAL, None, True, True, True, "04"), + ] + for name, category, origin, metadata, starred, intermediate, on_board, day in rows: + store.save( + video_name=name, + video_origin=origin, + video_category=category, + width=8, + height=8, + duration=1.0, + fps=8.0, + has_workflow=False, + is_intermediate=intermediate, + starred=starred, + metadata=metadata, + user_id="alice", + ) + with database.begin(write=True) as conn: + conn.execute( + update(videos).where(videos.c.video_name == name).values(created_at=f"2026-01-{day} 00:00:00.000") + ) + if on_board: + BoardVideoRecordStorage(database).add_video_to_board(board.board_id, name) + return board.board_id + + @pytest.mark.parametrize( + ("filters", "expected"), + [ + pytest.param({"board_id": "board"}, {"boarded.mp4", "intermediate.mp4"}, id="one board"), + pytest.param({"board_id": "none"}, {"user.mp4", "external.mp4"}, id="no board"), + pytest.param({"categories": [ImageCategory.USER]}, {"user.mp4"}, id="categories"), + pytest.param({"video_origin": ResourceOrigin.EXTERNAL}, {"external.mp4"}, id="origin"), + pytest.param({"is_intermediate": False}, {"boarded.mp4", "user.mp4", "external.mp4"}, id="intermediate"), + # Literal % and _, and either case: the unescaped pattern would also match "100x catXphoto". + pytest.param({"search_term": "100% cat_"}, {"boarded.mp4"}, id="search"), + ], + ) + def test_a_filter_narrows_pages_and_names_alike( + self, store: VideoRecordStorage, board_id: str, filters: dict[str, Any], expected: set[str] + ) -> None: + if filters.get("board_id") == "board": + filters = {**filters, "board_id": board_id} + + page = store.get_many(limit=10, is_admin=True, **filters) + names = store.get_video_names(is_admin=True, **filters) + + assert {record.video_name for record in page.items} == expected + assert page.total == names.total_count == len(expected) + assert names.starred_count == len(expected & {"boarded.mp4", "intermediate.mp4"}) + assert [record.video_name for record in page.items] == names.video_names + + def test_starred_videos_come_first_and_are_counted(self, store: VideoRecordStorage, board_id: str) -> None: + names = store.get_video_names(starred_first=True, order_dir=SQLiteDirection.Descending, is_admin=True) + unstarred_first = store.get_video_names(starred_first=False, order_dir=SQLiteDirection.Ascending, is_admin=True) + + assert names.video_names == ["intermediate.mp4", "boarded.mp4", "external.mp4", "user.mp4"] + assert names.starred_count == 2 + assert unstarred_first.video_names == ["boarded.mp4", "user.mp4", "external.mp4", "intermediate.mp4"] + + def test_the_cover_is_the_newest_video_that_is_no_intermediate( + self, store: VideoRecordStorage, board_id: str + ) -> None: + cover = store.get_most_recent_video_for_board(board_id) + + assert cover is not None and cover.video_name == "boarded.mp4" + + +class TestFields: + def test_save_stores_every_field_and_update_changes_every_field(self, store: VideoRecordStorage) -> None: + store.save( + video_name="full.mp4", + video_origin=ResourceOrigin.EXTERNAL, + video_category=ImageCategory.CONTROL, + width=640, + height=480, + duration=2.5, + fps=30.0, + has_workflow=True, + is_intermediate=True, + starred=True, + session_id="session-1", + node_id="node-1", + metadata='{"seed": 1}', + user_id="user-1", + video_subfolder="sub/dir", + project_id="project-1", + ) + + saved = store.get("full.mp4") + store.update( + "full.mp4", + VideoRecordChanges( + video_category=ImageCategory.USER, session_id="session-2", is_intermediate=False, starred=False + ), + ) + updated = store.get("full.mp4") + + assert ( + saved.video_origin, + saved.video_category, + saved.width, + saved.height, + saved.duration, + saved.fps, + saved.has_workflow, + saved.is_intermediate, + saved.starred, + saved.session_id, + saved.node_id, + saved.video_subfolder, + saved.project_id, + ) == ( + ResourceOrigin.EXTERNAL, + ImageCategory.CONTROL, + 640, + 480, + 2.5, + 30.0, + True, + True, + True, + "session-1", + "node-1", + "sub/dir", + "project-1", + ) + assert store.get_user_id("full.mp4") == "user-1" + assert store.get_metadata("full.mp4") == MetadataFieldValidator.validate_json('{"seed": 1}') + assert (updated.video_category, updated.session_id, updated.is_intermediate, updated.starred) == ( + ImageCategory.USER, + "session-2", + False, + False, + ) + assert updated.node_id == "node-1" and updated.has_workflow is True + + def test_saving_a_taken_name_keeps_the_record_and_returns_its_time( + self, database: Database, store: VideoRecordStorage + ) -> None: + _save(store, "taken.mp4", "alice") + with database.begin(write=True) as conn: + conn.execute( + update(videos).where(videos.c.video_name == "taken.mp4").values(created_at="2026-01-01 00:00:00.000") + ) + + created_at = store.save( + video_name="taken.mp4", + video_origin=ResourceOrigin.EXTERNAL, + video_category=ImageCategory.USER, + width=1, + height=1, + duration=9.0, + fps=None, + has_workflow=True, + user_id="bob", + ) + + record = store.get("taken.mp4") + assert created_at.isoformat(" ", "milliseconds") == "2026-01-01 00:00:00.000" + assert (record.video_origin, record.width, store.get_user_id("taken.mp4")) == ( + ResourceOrigin.INTERNAL, + 64, + "alice", + ) + + +def test_a_size_backfill_fills_only_sizes_not_known_yet(store: VideoRecordStorage) -> None: + _save(store, "unknown.mp4", "alice") + _save(store, "measured.mp4", "alice") + store.set_file_size_bytes("measured.mp4", 100) + + store.set_file_sizes_bytes({"unknown.mp4": 7, "measured.mp4": 9}) + + assert [store.get(name).file_size_bytes for name in ("unknown.mp4", "measured.mp4")] == [7, 100] + + +@pytest.mark.parametrize("operation", ["delete_intermediates_by_names", "delete_many", "get_subfolders"]) +def test_every_name_is_covered_and_no_statement_binds_more_than_sqlite_takes( + database: Database, store: VideoRecordStorage, operation: str +) -> None: + """999 is the SQLITE_MAX_VARIABLE_NUMBER default on builds older than 3.32.""" + names = [f"tmp{i:05d}.mp4" for i in range(2 * 500 + 7)] + with database.begin(write=True) as conn: + conn.execute( + insert(videos), + [ + { + "video_name": name, + "video_origin": ResourceOrigin.INTERNAL.value, + "video_category": ImageCategory.GENERAL.value, + "width": 8, + "height": 8, + "is_intermediate": True, + "video_subfolder": "sub", + } + for name in names + ], + ) + + with capture_statements(database) as statements: + result = getattr(store, operation)(names) + + widest = max(len(parameters) for _, parameters in statements) + assert 0 < widest <= 999 + if operation == "get_subfolders": + assert result == dict.fromkeys(names, "sub") + else: + assert store.get_video_names().video_names == [] + if operation == "delete_intermediates_by_names": + assert result == names diff --git a/tests/app/services/video_records/test_video_records_sqlite.py b/tests/app/services/video_records/test_video_records_sqlite.py deleted file mode 100644 index a7fef798f70..00000000000 --- a/tests/app/services/video_records/test_video_records_sqlite.py +++ /dev/null @@ -1,383 +0,0 @@ -"""Regression tests for SqliteVideoRecordStorage multiuser isolation. - -Covers JPPhoto's code-review finding (PR #9163): when ``board_id`` was omitted -from /v1/videos/ and /v1/videos/names, the SQL builder applied no user filter -and a non-admin caller saw every user's videos. The fix added an -``elif user_id is not None and not is_admin`` branch; these tests pin the -behaviour so the regression cannot reappear. -""" - -import json -import sqlite3 - -import pytest - -from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage -from invokeai.app.services.config.config_default import InvokeAIAppConfig -from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.users.users_common import UserCreateRequest -from invokeai.app.services.users.users_default import UserService -from invokeai.app.services.video_records.video_records_common import VideoRecordNotFoundException -from invokeai.app.services.video_records.video_records_sqlite import SqliteVideoRecordStorage -from invokeai.backend.util.logging import InvokeAILogger -from tests.fixtures.sqlite_database import create_mock_sqlite_database - - -@pytest.fixture -def store() -> SqliteVideoRecordStorage: - config = InvokeAIAppConfig(use_memory_db=True) - logger = InvokeAILogger.get_logger(config=config) - db = create_mock_sqlite_database(config, logger) - return SqliteVideoRecordStorage(db=db) - - -def _save(store: SqliteVideoRecordStorage, name: str, user_id: str) -> None: - store.save( - video_name=name, - video_origin=ResourceOrigin.INTERNAL, - video_category=ImageCategory.GENERAL, - width=64, - height=64, - duration=1.0, - fps=8.0, - has_workflow=False, - is_intermediate=False, - user_id=user_id, - ) - - -@pytest.fixture -def seeded_store(store: SqliteVideoRecordStorage) -> SqliteVideoRecordStorage: - # Two videos per user; all without board association (the bug occurred when board_id - # was omitted from the query). - _save(store, "alice_1.mp4", user_id="alice") - _save(store, "alice_2.mp4", user_id="alice") - _save(store, "bob_1.mp4", user_id="bob") - _save(store, "bob_2.mp4", user_id="bob") - return store - - -class TestGetManyOmittedBoardIdMultiuser: - """get_many() with board_id=None must filter by user_id for non-admin callers.""" - - def test_non_admin_only_sees_own_videos(self, seeded_store: SqliteVideoRecordStorage) -> None: - result = seeded_store.get_many(user_id="alice", is_admin=False) - names = {v.video_name for v in result.items} - assert names == {"alice_1.mp4", "alice_2.mp4"} - assert result.total == 2 - - def test_admin_sees_every_users_videos(self, seeded_store: SqliteVideoRecordStorage) -> None: - result = seeded_store.get_many(user_id="alice", is_admin=True) - names = {v.video_name for v in result.items} - assert names == {"alice_1.mp4", "alice_2.mp4", "bob_1.mp4", "bob_2.mp4"} - - def test_no_user_id_returns_all(self, seeded_store: SqliteVideoRecordStorage) -> None: - # No user_id means the caller is bypassing user filtering entirely (e.g. internal calls). - result = seeded_store.get_many(user_id=None, is_admin=False) - names = {v.video_name for v in result.items} - assert names == {"alice_1.mp4", "alice_2.mp4", "bob_1.mp4", "bob_2.mp4"} - - -class TestGetVideoNamesOmittedBoardIdMultiuser: - """get_video_names() with board_id=None must filter by user_id for non-admin callers.""" - - def test_non_admin_only_sees_own_videos(self, seeded_store: SqliteVideoRecordStorage) -> None: - result = seeded_store.get_video_names(user_id="alice", is_admin=False) - assert set(result.video_names) == {"alice_1.mp4", "alice_2.mp4"} - assert result.total_count == 2 - - def test_admin_sees_every_users_videos(self, seeded_store: SqliteVideoRecordStorage) -> None: - result = seeded_store.get_video_names(user_id="alice", is_admin=True) - assert set(result.video_names) == {"alice_1.mp4", "alice_2.mp4", "bob_1.mp4", "bob_2.mp4"} - - def test_explicit_none_board_still_isolates(self, seeded_store: SqliteVideoRecordStorage) -> None: - # The "none" sentinel (uncategorized) must also apply the user filter — this was the - # only path that was correct *before* the fix; the test guards against accidental - # regression there too. - result = seeded_store.get_video_names(board_id="none", user_id="alice", is_admin=False) - assert set(result.video_names) == {"alice_1.mp4", "alice_2.mp4"} - - -class TestDeterministicOrdering: - @staticmethod - def _set_same_timestamp(store: SqliteVideoRecordStorage) -> None: - with store._db.transaction() as cursor: - cursor.execute("UPDATE videos SET created_at = '2026-07-20 00:00:00.000'") - - def test_same_timestamp_pages_have_stable_mirrored_order(self, seeded_store: SqliteVideoRecordStorage) -> None: - self._set_same_timestamp(seeded_store) - - ascending = [ - seeded_store.get_many( - offset=offset, - limit=1, - starred_first=False, - order_dir=SQLiteDirection.Ascending, - user_id="alice", - ) - .items[0] - .video_name - for offset in range(2) - ] - descending = [ - seeded_store.get_many( - offset=offset, - limit=1, - starred_first=False, - order_dir=SQLiteDirection.Descending, - user_id="alice", - ) - .items[0] - .video_name - for offset in range(2) - ] - - assert ascending == ["alice_1.mp4", "alice_2.mp4"] - assert descending == list(reversed(ascending)) - - def test_same_timestamp_name_list_has_stable_mirrored_order(self, seeded_store: SqliteVideoRecordStorage) -> None: - self._set_same_timestamp(seeded_store) - - ascending = seeded_store.get_video_names( - starred_first=False, - order_dir=SQLiteDirection.Ascending, - user_id="alice", - ).video_names - descending = seeded_store.get_video_names( - starred_first=False, - order_dir=SQLiteDirection.Descending, - user_id="alice", - ).video_names - - assert ascending == ["alice_1.mp4", "alice_2.mp4"] - assert descending == list(reversed(ascending)) - - def test_board_cover_uses_name_as_same_timestamp_tie_breaker(self, store: SqliteVideoRecordStorage) -> None: - board = SqliteBoardRecordStorage(store._db).save("Board", "system") - _save(store, "a.mp4", user_id="system") - _save(store, "b.mp4", user_id="system") - self._set_same_timestamp(store) - with store._db.transaction() as cursor: - cursor.executemany( - "INSERT INTO board_videos (board_id, video_name) VALUES (?, ?)", - [(board.board_id, "a.mp4"), (board.board_id, "b.mp4")], - ) - - assert store.get_most_recent_video_for_board(board.board_id).video_name == "b.mp4" - - -class TestUserDeletionLifecycle: - """Documents the intended videos↔users lifecycle (JPPhoto PR #9163 July-10 follow-up). - - ``videos.user_id`` deliberately has no FK to ``users`` — exactly like ``images``, - ``boards`` and ``workflows``, whose user_id columns (migration_27) are index-only. - Deleting a user therefore leaves their videos in place instead of cascading a row - delete that would strand the files on disk; the orphaned records stay visible to - administrators (and only to administrators), who can clean them up or reassign them. - These tests pin that platform-wide behavior for videos so any future change to the - user-deletion story is a deliberate decision rather than an accident. - """ - - @pytest.fixture - def migrated_db(self) -> SqliteDatabase: - config = InvokeAIAppConfig(use_memory_db=True) - logger = InvokeAILogger.get_logger(config=config) - return create_mock_sqlite_database(config, logger) - - def test_videos_survive_owner_deletion_and_remain_admin_only(self, migrated_db: SqliteDatabase) -> None: - users = UserService(migrated_db) - store = SqliteVideoRecordStorage(db=migrated_db) - - owner = users.create( - UserCreateRequest( - email="doomed@example.com", - display_name="Doomed User", - password="TestPassword123", - is_admin=False, - ) - ) - _save(store, "doomed.mp4", user_id=owner.user_id) - - users.delete(owner.user_id) - assert users.get(owner.user_id) is None - - # The record survives, still attributed to the deleted owner... - assert store.get_user_id("doomed.mp4") == owner.user_id - # ...is visible to admins for cleanup... - admin_view = store.get_many(user_id="some-admin", is_admin=True) - assert "doomed.mp4" in {v.video_name for v in admin_view.items} - # ...and no regular user inherits it. - other_view = store.get_many(user_id="bystander", is_admin=False) - assert "doomed.mp4" not in {v.video_name for v in other_view.items} - - -def test_get_propagates_a_storage_error_instead_of_reporting_the_row_missing(store: SqliteVideoRecordStorage): - """An unreadable database must not be indistinguishable from a deleted video. - - `get` used to translate every sqlite3.Error into VideoRecordNotFoundException, which made - that exception mean "the row is absent, OR the read failed". Two callers act destructively - on it: `_assert_video_read_access` answers 404 for a positive not-found and the clients drop - their reference to the video on one, and the staged-delete recovery reads it as proof the - delete committed and purges the staged files. - """ - _save(store, "video-1.mp4", "user-1") - # A real storage failure rather than a patched one: the SELECT below cannot run at all, the - # same shape a locked or corrupt database presents. The row's absence is not what is being - # reported, and the caller must be able to tell. - with store._db.transaction() as cursor: - cursor.execute("DROP TABLE videos;") - - with pytest.raises(sqlite3.OperationalError): - store.get("video-1.mp4") - - -def test_get_still_reports_a_positively_absent_row_as_missing(store: SqliteVideoRecordStorage): - """The narrowing stays exactly that narrow: absence is still absence.""" - with pytest.raises(VideoRecordNotFoundException): - store.get("never-existed.mp4") - - -def test_exists_reports_a_row_get_cannot_deserialize(store: SqliteVideoRecordStorage): - """Presence, not readability. `get` would raise on an enum value this version does not know - — a row written by a newer one — and the refusal path reads that as absence, which would - report a live video gone.""" - _save(store, "video-1.mp4", "user-1") - with store._db.transaction() as cursor: - cursor.execute("UPDATE videos SET video_category = 'from_the_future' WHERE video_name = ?;", ("video-1.mp4",)) - - with pytest.raises(ValueError): - store.get("video-1.mp4") - assert store.exists("video-1.mp4") is True - - -def test_exists_propagates_a_storage_error(store: SqliteVideoRecordStorage): - """ "Could not look" is not "not there" — the caller answers 404 on a False.""" - with store._db.transaction() as cursor: - cursor.execute("DROP TABLE videos;") - - with pytest.raises(sqlite3.OperationalError): - store.exists("video-1.mp4") - - -def _save_with_metadata(store: SqliteVideoRecordStorage, name: str, metadata: str | None) -> None: - store.save( - video_name=name, - video_origin=ResourceOrigin.EXTERNAL, - video_category=ImageCategory.USER, - width=640, - height=360, - duration=1.0, - fps=24.0, - has_workflow=False, - is_intermediate=False, - metadata=metadata, - user_id="alice", - ) - - -def test_media_origin_is_projected_out_of_the_metadata_blob(store: SqliteVideoRecordStorage) -> None: - """The record carries `media_origin` so clients need no separate /metadata request. - - An audio upload the ingest converter wrapped into a waveform video is marked - `audio_upload`; the frontend starts such a reference on its soundtrack alone, and a - wrong value there costs a video reference slot and roughly doubles the packed sequence. - """ - _save_with_metadata(store, "wrapped.mp4", '{"media_origin": "audio_upload", "note": "kept"}') - - assert store.get("wrapped.mp4").media_origin == "audio_upload" - # The rest of the blob stays behind the /metadata route rather than riding every row. - assert store.get_metadata("wrapped.mp4") is not None - - -@pytest.mark.parametrize( - "metadata", - [ - pytest.param(None, id="no metadata at all"), - pytest.param("{}", id="metadata without the key"), - pytest.param('{"note": "kept"}', id="metadata with other keys"), - ], -) -def test_media_origin_is_none_when_unmarked(store: SqliteVideoRecordStorage, metadata: str | None) -> None: - """`json_extract` over a NULL or key-less blob yields NULL rather than raising.""" - _save_with_metadata(store, "plain.mp4", metadata) - - assert store.get("plain.mp4").media_origin is None - - -def test_media_origin_survives_a_listing(store: SqliteVideoRecordStorage) -> None: - """`get_many` selects the same columns, so a listed row carries the marker too.""" - _save_with_metadata(store, "wrapped.mp4", '{"media_origin": "audio_upload"}') - _save_with_metadata(store, "plain.mp4", None) - - listed = store.get_many(offset=0, limit=10, order_dir=SQLiteDirection.Descending, user_id="alice", is_admin=False) - origins = {record.video_name: record.media_origin for record in listed.items} - - assert origins == {"wrapped.mp4": "audio_upload", "plain.mp4": None} - - -@pytest.mark.parametrize( - "raw", - [ - pytest.param("1", id="a JSON number"), - pytest.param("true", id="a JSON boolean"), - pytest.param("1.5", id="a JSON float"), - pytest.param('""', id="an empty string, which the frontend also reads as no marker"), - # `json_extract` hands an object or array back as its SERIALIZED TEXT, so these - # arrive as `str` and an isinstance check alone would propagate them. - pytest.param('{"a": "b"}', id="a JSON object, returned as its text"), - pytest.param("[1, 2]", id="a JSON array, returned as its text"), - pytest.param('"has a space"', id="a string that is not marker-shaped"), - ], -) -def test_a_non_string_marker_does_not_break_the_record(store: SqliteVideoRecordStorage, raw: str) -> None: - """A client may store any JSON value under `media_origin`; the row must still deserialize. - - `MetadataField` validates only that the blob is an object, so `json_extract` can hand - back an int or a float. Passing one to the `Optional[str]` field raised while building - the record -- which took out the video's DTO *and* every gallery listing containing it, - so one odd upload broke the gallery. - """ - _save_with_metadata(store, "odd.mp4", '{"media_origin": ' + raw + "}") - - assert store.get("odd.mp4").media_origin is None - - -def test_metadata_that_is_not_an_object_leaves_the_marker_unset(store: SqliteVideoRecordStorage) -> None: - """`json_extract` over a non-object blob yields NULL rather than raising.""" - _save_with_metadata(store, "array.mp4", "[1, 2, 3]") - - assert store.get("array.mp4").media_origin is None - - -def test_a_malformed_metadata_blob_does_not_fail_the_query(store: SqliteVideoRecordStorage) -> None: - """`json_extract` RAISES on unparseable text, and this expression runs on every listed row. - - Unguarded, one malformed blob would fail the whole page — `get_many` deserializes rows in - a comprehension — rather than the one video. The column is plain TEXT with no CHECK, so - only convention keeps such a row out; the `json_valid` guard makes that not matter. - """ - _save_with_metadata(store, "bad.mp4", "not json at all") - _save_with_metadata(store, "good.mp4", '{"media_origin": "audio_upload"}') - - listed = store.get_many(offset=0, limit=10, order_dir=SQLiteDirection.Descending, user_id="alice", is_admin=False) - - assert {record.video_name: record.media_origin for record in listed.items} == { - "bad.mp4": None, - "good.mp4": "audio_upload", - } - - -def test_an_overlong_marker_is_dropped_rather_than_echoed_on_every_row(store: SqliteVideoRecordStorage) -> None: - """Upload metadata is client-supplied and unbounded, and this key now rides every row. - - Without a cap, one upload carrying a huge marker is echoed back on every gallery page - that includes the video. - """ - _save_with_metadata(store, "huge.mp4", json.dumps({"media_origin": "A" * 5000})) - - assert store.get("huge.mp4").media_origin is None - # A marker of a plausible length still passes. - _save_with_metadata(store, "fine.mp4", json.dumps({"media_origin": "audio_upload"})) - assert store.get("fine.mp4").media_origin == "audio_upload" diff --git a/tests/app/services/workflow_records/conftest.py b/tests/app/services/workflow_records/conftest.py new file mode 100644 index 00000000000..0abc56858e5 --- /dev/null +++ b/tests/app/services/workflow_records/conftest.py @@ -0,0 +1,17 @@ +from unittest import mock + +import pytest + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage +from invokeai.backend.util.logging import InvokeAILogger + + +@pytest.fixture +def workflow_records(database: Database) -> WorkflowRecordsStorage: + """The storage on the backend under test, started as the app starts it: the bundled workflows are stored.""" + records = WorkflowRecordsStorage(database) + invoker = mock.Mock() + invoker.services.logger = InvokeAILogger.get_logger() + records.start(invoker) + return records diff --git a/tests/app/services/workflow_records/test_workflow_last_run_at.py b/tests/app/services/workflow_records/test_workflow_last_run_at.py index 1492a45e460..3f65bd72df5 100644 --- a/tests/app/services/workflow_records/test_workflow_last_run_at.py +++ b/tests/app/services/workflow_records/test_workflow_last_run_at.py @@ -5,22 +5,14 @@ the frontend show "Your last run · 2 days ago" on library cards. """ -import pytest - -from invokeai.app.services.invoker import Invoker -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.services.workflow_records.workflow_records_common import ( WorkflowCategory, WorkflowMeta, WorkflowRecordOrderBy, WorkflowWithoutID, ) -from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage - - -@pytest.fixture -def workflow_records_service(mock_invoker: Invoker) -> SqliteWorkflowRecordsStorage: - return mock_invoker.services.workflow_records +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage def create_minimal_user_workflow() -> WorkflowWithoutID: @@ -40,17 +32,17 @@ def create_minimal_user_workflow() -> WorkflowWithoutID: ) -def test_update_last_run_at_sets_timestamp(workflow_records_service: SqliteWorkflowRecordsStorage) -> None: +def test_update_last_run_at_sets_timestamp(workflow_records: WorkflowRecordsStorage) -> None: workflow = create_minimal_user_workflow() - created = workflow_records_service.create(workflow=workflow) + created = workflow_records.create(workflow=workflow) assert created.last_run_at is None - workflow_records_service.update_last_run_at(created.workflow_id) + workflow_records.update_last_run_at(created.workflow_id) - fetched = workflow_records_service.get(created.workflow_id) + fetched = workflow_records.get(created.workflow_id) assert fetched.last_run_at is not None - listed = workflow_records_service.get_many( + listed = workflow_records.get_many( order_by=WorkflowRecordOrderBy.CreatedAt, direction=SQLiteDirection.Descending, categories=None, @@ -61,24 +53,24 @@ def test_update_last_run_at_sets_timestamp(workflow_records_service: SqliteWorkf assert row.last_run_at is not None -def test_update_last_run_at_scoped_to_owning_user(workflow_records_service: SqliteWorkflowRecordsStorage) -> None: +def test_update_last_run_at_scoped_to_owning_user(workflow_records: WorkflowRecordsStorage) -> None: """Mirrors update_opened_at's ownership scoping: passing a mismatched user_id is a no-op, and the row updates only once the correct owner's user_id is supplied.""" workflow = create_minimal_user_workflow() - created = workflow_records_service.create(workflow=workflow, user_id="user-a") + created = workflow_records.create(workflow=workflow, user_id="user-a") - workflow_records_service.update_last_run_at(created.workflow_id, user_id="user-b") - assert workflow_records_service.get(created.workflow_id).last_run_at is None + workflow_records.update_last_run_at(created.workflow_id, user_id="user-b") + assert workflow_records.get(created.workflow_id).last_run_at is None - workflow_records_service.update_last_run_at(created.workflow_id, user_id="user-a") - assert workflow_records_service.get(created.workflow_id).last_run_at is not None + workflow_records.update_last_run_at(created.workflow_id, user_id="user-a") + assert workflow_records.get(created.workflow_id).last_run_at is not None def test_update_last_run_at_missing_workflow_is_a_no_op( - workflow_records_service: SqliteWorkflowRecordsStorage, + workflow_records: WorkflowRecordsStorage, ) -> None: """The service method mirrors update_opened_at: it has no existence check and silently does nothing for an unknown workflow_id. The router turns a missing workflow into a 404 by calling `get()` first (see tests/app/routers/test_multiuser_authorization.py::TestWorkflowMutationAuth ::test_update_last_run_at_missing_workflow_404s).""" - workflow_records_service.update_last_run_at("does-not-exist") + workflow_records.update_last_run_at("does-not-exist") diff --git a/tests/app/services/workflow_records/test_workflow_library_queries.py b/tests/app/services/workflow_records/test_workflow_library_queries.py new file mode 100644 index 00000000000..83839d42dc4 --- /dev/null +++ b/tests/app/services/workflow_records/test_workflow_library_queries.py @@ -0,0 +1,318 @@ +"""Listing, searching and counting the workflow library, on every database backend.""" + +from pathlib import Path +from typing import Any, Optional + +import pytest +from sqlalchemy import update + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import workflows as workflow_queries +from invokeai.app.services.shared.database.schema.workflows import workflow_library +from invokeai.app.services.shared.pagination import SQLiteDirection +from invokeai.app.services.workflow_records import workflow_records_default +from invokeai.app.services.workflow_records.workflow_records_common import ( + Workflow, + WorkflowCategory, + WorkflowMeta, + WorkflowRecordOrderBy, + WorkflowWithoutID, +) +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage + +USER = "user-1" +OTHER = "user-2" + + +def _workflow(name: str, *, description: str = "", tags: str = "") -> WorkflowWithoutID: + return WorkflowWithoutID( + name=name, + author="", + description=description, + version="1.0.0", + contact="", + tags=tags, + notes="", + exposedFields=[], + meta=WorkflowMeta(version="3.0.0", category=WorkflowCategory.User), + nodes=[], + edges=[], + ) + + +def _set(database: Database, workflow_id: str, **values: Any) -> None: + with database.begin(write=True) as conn: + conn.execute(update(workflow_library).where(workflow_library.c.workflow_id == workflow_id).values(**values)) + + +def _names( + records: WorkflowRecordsStorage, + *, + query: Optional[str] = None, + tags: Optional[list[str]] = None, + user_id: Optional[str] = None, + is_public: Optional[bool] = None, + has_been_opened: Optional[bool] = None, +) -> list[str]: + listed = records.get_many( + order_by=WorkflowRecordOrderBy.Name, + direction=SQLiteDirection.Ascending, + categories=[WorkflowCategory.User], + query=query, + tags=tags, + user_id=user_id, + is_public=is_public, + has_been_opened=has_been_opened, + ) + return [item.name for item in listed.items] + + +def test_a_search_looks_in_the_name_description_and_tags_ignoring_case( + workflow_records: WorkflowRecordsStorage, +) -> None: + workflow_records.create(_workflow("Hires FIX"), user_id=USER) + workflow_records.create(_workflow("Faces", description="Fixes faces"), user_id=USER) + workflow_records.create(_workflow("Tagged", tags="fixer"), user_id=USER) + workflow_records.create(_workflow("Unrelated"), user_id=USER) + + assert _names(workflow_records, query="fix") == ["Faces", "Hires FIX", "Tagged"] + assert _names(workflow_records, query=" FIX ") == ["Faces", "Hires FIX", "Tagged"] + + +def test_search_texts_and_tags_match_wildcard_characters_only_as_themselves( + workflow_records: WorkflowRecordsStorage, +) -> None: + for name in ("50% done", "500 done", "my_flow", "my-flow", "back\\slash", "backslash"): + workflow_records.create(_workflow(name, tags=name), user_id=USER) + + assert _names(workflow_records, query="50%") == ["50% done"] + assert _names(workflow_records, query="my_flow") == ["my_flow"] + assert _names(workflow_records, query="back\\slash") == ["back\\slash"] + assert _names(workflow_records, tags=["my_flow"]) == ["my_flow"] + + +def test_tag_counts_agree_with_the_tag_filter(workflow_records: WorkflowRecordsStorage) -> None: + for name, tags in (("A", "Beta"), ("B", "50% off"), ("C", "500 off"), ("D", "my_flow"), ("E", "my-flow")): + workflow_records.create(_workflow(name, tags=tags), user_id=USER) + + for tag in ("beta", "50%", "my_flow"): + assert len(_names(workflow_records, tags=[tag])) == 1 + assert workflow_records.counts_by_tag([tag], categories=[WorkflowCategory.User]) == {tag: 1} + + +def test_a_tag_filter_keeps_the_workflows_with_any_of_its_tags(workflow_records: WorkflowRecordsStorage) -> None: + workflow_records.create(_workflow("A", tags="alpha, common"), user_id=USER) + workflow_records.create(_workflow("B", tags="Beta"), user_id=USER) + workflow_records.create(_workflow("C", tags="gamma"), user_id=USER) + + assert _names(workflow_records, tags=["alpha", "beta"]) == ["A", "B"] + # Three tags take four slots of a statement: the spare one matches nothing. + assert _names(workflow_records, tags=["alpha", "beta", "missing"]) == ["A", "B"] + assert _names(workflow_records, tags=["common"]) == ["A"] + + +def test_more_tags_than_a_statement_has_slots_are_listed_and_counted_like_few( + workflow_records: WorkflowRecordsStorage, +) -> None: + # Zero-padded, so that no tag contains another. + workflow_records.create(_workflow("Late", tags="t099"), user_id=USER) + workflow_records.create(_workflow("Early", tags="t000"), user_id=USER) + tags = [f"t{i:03d}" for i in range(100)] + + assert _names(workflow_records, tags=tags) == ["Early", "Late"] + counts = workflow_records.counts_by_tag(tags, categories=[WorkflowCategory.User]) + assert (counts["t000"], counts["t050"], counts["t099"]) == (1, 0, 1) + + +def test_an_account_sees_its_own_and_the_bundled_workflows(workflow_records: WorkflowRecordsStorage) -> None: + mine = workflow_records.create(_workflow("Mine"), user_id=USER) + theirs = workflow_records.create(_workflow("Theirs"), user_id=OTHER, is_public=True) + + listed = workflow_records.get_many( + order_by=WorkflowRecordOrderBy.CreatedAt, + direction=SQLiteDirection.Descending, + categories=None, + user_id=USER, + ).items + + ids = {item.workflow_id for item in listed} + assert mine.workflow_id in ids and theirs.workflow_id not in ids + assert {item.category for item in listed} == {WorkflowCategory.User, WorkflowCategory.Default} + assert _names(workflow_records, is_public=True) == ["Theirs"] + assert _names(workflow_records, is_public=False) == ["Mine"] + + +def test_listings_tell_opened_workflows_apart(workflow_records: WorkflowRecordsStorage) -> None: + opened = workflow_records.create(_workflow("Opened"), user_id=USER) + workflow_records.create(_workflow("Never opened"), user_id=USER) + # Opening by another account changes nothing: it is scoped to the owner. + workflow_records.update_opened_at(opened.workflow_id, user_id=OTHER) + assert _names(workflow_records, has_been_opened=True) == [] + + workflow_records.update_opened_at(opened.workflow_id, user_id=USER) + + assert _names(workflow_records, has_been_opened=True) == ["Opened"] + assert _names(workflow_records, has_been_opened=False) == ["Never opened"] + + +@pytest.mark.parametrize("direction", [SQLiteDirection.Ascending, SQLiteDirection.Descending]) +def test_pages_of_equal_names_keep_one_order( + workflow_records: WorkflowRecordsStorage, direction: SQLiteDirection +) -> None: + created = [workflow_records.create(_workflow("Same"), user_id=USER).workflow_id for _ in range(5)] + + pages = [ + workflow_records.get_many( + order_by=WorkflowRecordOrderBy.Name, + direction=direction, + categories=[WorkflowCategory.User], + page=page, + per_page=2, + ) + for page in range(3) + ] + + listed = [item.workflow_id for page in pages for item in page.items] + assert listed == sorted(created, reverse=direction == SQLiteDirection.Descending) + assert [(page.total, page.pages, page.per_page) for page in pages] == [(5, 3, 2)] * 3 + + +@pytest.mark.parametrize("order_by", list(WorkflowRecordOrderBy)) +def test_each_order_key_sorts_by_its_column( + workflow_records: WorkflowRecordsStorage, database: Database, order_by: WorkflowRecordOrderBy +) -> None: + # Each column puts the three workflows in a different order. + rows: dict[str, dict[str, Any]] = { + "c": { + "created_at": "2026-01-01 00:00:00.000", + "updated_at": "2026-01-03 00:00:00.000", + "opened_at": "2026-01-02 00:00:00.000", + "is_public": True, + }, + "a": { + "created_at": "2026-01-02 00:00:00.000", + "updated_at": "2026-01-01 00:00:00.000", + "opened_at": "2026-01-03 00:00:00.000", + "is_public": False, + }, + "b": { + "created_at": "2026-01-03 00:00:00.000", + "updated_at": "2026-01-02 00:00:00.000", + "opened_at": "2026-01-01 00:00:00.000", + "is_public": False, + }, + } + ids = {name: workflow_records.create(_workflow(name), user_id=USER).workflow_id for name in rows} + for name, values in rows.items(): + _set(database, ids[name], **values) + + listed = workflow_records.get_many(order_by, SQLiteDirection.Ascending, [WorkflowCategory.User]).items + + def key(name: str) -> tuple[Any, str]: + return (name if order_by == WorkflowRecordOrderBy.Name else rows[name][order_by.value]), ids[name] + + assert [item.name for item in listed] == sorted(rows, key=key) + + +def test_tags_are_counted_one_by_one_under_the_filters(workflow_records: WorkflowRecordsStorage) -> None: + workflow_records.create(_workflow("A", tags="alpha, beta"), user_id=USER) + workflow_records.create(_workflow("B", tags="beta"), user_id=USER, is_public=True) + workflow_records.create(_workflow("C", tags="alpha"), user_id=OTHER) + + counts = workflow_records.counts_by_tag(["alpha", "beta", "missing"], categories=[WorkflowCategory.User]) + assert counts == {"alpha": 2, "beta": 2, "missing": 0} + assert workflow_records.counts_by_tag(["alpha", "beta"], categories=[WorkflowCategory.User], user_id=USER) == { + "alpha": 1, + "beta": 2, + } + assert workflow_records.counts_by_tag(["beta"], categories=[WorkflowCategory.User], is_public=True) == {"beta": 1} + assert workflow_records.counts_by_tag([]) == {} + + +def test_counts_and_tags_follow_the_visibility_and_opened_filters(workflow_records: WorkflowRecordsStorage) -> None: + shared = workflow_records.create(_workflow("S", tags="alpha, shared-only"), user_id=USER, is_public=True) + workflow_records.create(_workflow("P", tags="alpha"), user_id=USER) + workflow_records.update_opened_at(shared.workflow_id) + + assert workflow_records.counts_by_category([WorkflowCategory.User], is_public=True) == {"user": 1} + user = [WorkflowCategory.User] + assert workflow_records.counts_by_tag(["alpha"], categories=user, has_been_opened=True) == {"alpha": 1} + assert workflow_records.get_all_tags(categories=[WorkflowCategory.User], is_public=False) == ["alpha"] + + +def test_categories_are_counted_one_by_one_under_the_filters(workflow_records: WorkflowRecordsStorage) -> None: + bundled = len(list((Path(workflow_records_default.__file__).parent / "default_workflows").glob("*.json"))) + workflow_records.create(_workflow("Mine"), user_id=USER) + workflow_records.create(_workflow("Theirs"), user_id=OTHER) + + assert workflow_records.counts_by_category([WorkflowCategory.User, WorkflowCategory.Default]) == { + "user": 2, + "default": bundled, + } + assert workflow_records.counts_by_category([WorkflowCategory.User], user_id=USER) == {"user": 1} + assert workflow_records.counts_by_category([WorkflowCategory.User], has_been_opened=True) == {"user": 0} + assert workflow_records.counts_by_category([]) == {} + + +def test_all_tags_are_split_trimmed_and_sorted(workflow_records: WorkflowRecordsStorage) -> None: + workflow_records.create(_workflow("A", tags="beta, alpha"), user_id=USER) + workflow_records.create(_workflow("B", tags=" alpha ,gamma,"), user_id=USER) + workflow_records.create(_workflow("C"), user_id=USER) + workflow_records.create(_workflow("D", tags="theirs"), user_id=OTHER) + + assert workflow_records.get_all_tags(categories=[WorkflowCategory.User], user_id=USER) == ["alpha", "beta", "gamma"] + + +@pytest.mark.parametrize("write", ["save", "open", "run", "share"]) +def test_every_write_renews_when_the_workflow_was_updated( + workflow_records: WorkflowRecordsStorage, database: Database, write: str +) -> None: + created = workflow_records.create(_workflow("Written"), user_id=USER) + _set(database, created.workflow_id, updated_at="2000-01-01 00:00:00.000") + + if write == "save": + workflow_records.update(Workflow(**_workflow("Written again").model_dump(), id=created.workflow_id)) + elif write == "open": + workflow_records.update_opened_at(created.workflow_id, user_id=USER) + elif write == "run": + workflow_records.update_last_run_at(created.workflow_id, user_id=USER) + else: + workflow_records.update_is_public(created.workflow_id, True, user_id=USER) + + assert str(workflow_records.get(created.workflow_id).updated_at) > "2000-01-01 00:00:00.000" + + +def test_a_flag_passed_as_a_number_filters_as_the_flag(workflow_records: WorkflowRecordsStorage) -> None: + """Statements are cached by their conditions' shape, in which 1 equals True: a number must not stand for a + statement without the condition.""" + workflow_queries._list.cache_clear() + workflow_records.create(_workflow("Shared"), user_id=USER, is_public=True) + workflow_records.create(_workflow("Private"), user_id=USER) + + assert _names(workflow_records, is_public=1) == ["Shared"] # type: ignore[arg-type] + assert _names(workflow_records, is_public=True) == ["Shared"] + assert _names(workflow_records, is_public=0) == ["Private"] # type: ignore[arg-type] + assert _names(workflow_records, has_been_opened=0) == ["Private", "Shared"] # type: ignore[arg-type] + + +def test_statement_shapes_stay_few_whatever_a_client_sends( + workflow_records: WorkflowRecordsStorage, database: Database +) -> None: + """SQLAlchemy keeps every statement it compiled: how many tags or categories a client sends must not decide how + many there are.""" + builders = (workflow_queries._list, workflow_queries._count, workflow_queries._tag_counts) + for builder in builders: + builder.cache_clear() + compiled = database.engine._compiled_cache + assert compiled is not None + compiled_before = len(compiled) + + for count in range(1, 101): + tags = [f"tag {i}" for i in range(count)] + categories = [WorkflowCategory.User] * count + workflow_records.get_many(WorkflowRecordOrderBy.Name, SQLiteDirection.Ascending, categories, tags=tags) + workflow_records.counts_by_tag(tags, categories=categories) + + # Tag slots 1, 2, 4 ... 32: six shapes each. A listing for more tags is compiled for its call alone. + assert [builder.cache_info().currsize for builder in builders] == [6, 6, 6] + assert len(compiled) - compiled_before <= 3 * 6 + 5 diff --git a/tests/app/services/workflow_records/test_workflow_media_references.py b/tests/app/services/workflow_records/test_workflow_media_references.py index 1dd5d8a81ae..efa2d646003 100644 --- a/tests/app/services/workflow_records/test_workflow_media_references.py +++ b/tests/app/services/workflow_records/test_workflow_media_references.py @@ -1,8 +1,13 @@ """Library workflows keep the media reference index current, like project documents do.""" +from collections.abc import Sequence +from typing import Any + import pytest +from sqlalchemy import select -from invokeai.app.services.invoker import Invoker +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.media_references import media_references from invokeai.app.services.workflow_records.workflow_records_common import ( Workflow, WorkflowAccessDeniedError, @@ -10,12 +15,7 @@ WorkflowMeta, WorkflowWithoutID, ) -from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage - - -@pytest.fixture -def workflow_records(mock_invoker: Invoker) -> SqliteWorkflowRecordsStorage: - return mock_invoker.services.workflow_records +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage def _workflow(image_name: str) -> WorkflowWithoutID: @@ -34,27 +34,32 @@ def _workflow(image_name: str) -> WorkflowWithoutID: ) -def _references(records: SqliteWorkflowRecordsStorage, workflow_id: str) -> set[tuple[str, str, str]]: - with records._db.transaction() as cursor: - cursor.execute( - "SELECT user_id, media_kind, media_name FROM media_references WHERE owner_kind = 'workflow' AND owner_id = ?;", - (workflow_id,), - ) - return {tuple(row) for row in cursor.fetchall()} +def _references(database: Database, workflow_id: str) -> set[tuple[str, str, str]]: + with database.begin(write=False) as conn: + rows: Sequence[Sequence[Any]] = conn.execute( + select(media_references.c.user_id, media_references.c.media_kind, media_references.c.media_name).where( + media_references.c.owner_kind == "workflow", media_references.c.owner_id == workflow_id + ) + ).all() + return {(user_id, kind, name) for user_id, kind, name in rows} -def test_create_update_and_delete_keep_the_index_current(workflow_records: SqliteWorkflowRecordsStorage) -> None: +def test_create_update_and_delete_keep_the_index_current( + database: Database, workflow_records: WorkflowRecordsStorage +) -> None: created = workflow_records.create(_workflow("input.png"), user_id="user-1") - assert _references(workflow_records, created.workflow_id) == {("user-1", "image", "input.png")} + assert _references(database, created.workflow_id) == {("user-1", "image", "input.png")} workflow_records.update(Workflow(**_workflow("replaced.png").model_dump(), id=created.workflow_id)) - assert _references(workflow_records, created.workflow_id) == {("user-1", "image", "replaced.png")} + assert _references(database, created.workflow_id) == {("user-1", "image", "replaced.png")} workflow_records.delete(created.workflow_id, user_id="user-1") - assert _references(workflow_records, created.workflow_id) == set() + assert _references(database, created.workflow_id) == set() -def test_an_update_refused_by_ownership_leaves_the_index_alone(workflow_records: SqliteWorkflowRecordsStorage) -> None: +def test_an_update_refused_by_ownership_leaves_the_index_alone( + database: Database, workflow_records: WorkflowRecordsStorage +) -> None: created = workflow_records.create(_workflow("input.png"), user_id="user-1") with pytest.raises(WorkflowAccessDeniedError): @@ -62,17 +67,27 @@ def test_an_update_refused_by_ownership_leaves_the_index_alone(workflow_records: Workflow(**_workflow("stolen.png").model_dump(), id=created.workflow_id), user_id="user-2" ) - assert _references(workflow_records, created.workflow_id) == {("user-1", "image", "input.png")} + assert _references(database, created.workflow_id) == {("user-1", "image", "input.png")} def test_visibility_change_keeps_the_reference_index_of_the_document_it_writes( - workflow_records: SqliteWorkflowRecordsStorage, + database: Database, workflow_records: WorkflowRecordsStorage ) -> None: created = workflow_records.create(_workflow("old.png"), user_id="user-1") shared = workflow_records.update_is_public(created.workflow_id, True, user_id="user-1") assert shared.is_public is True - assert _references(workflow_records, created.workflow_id) == {("user-1", "image", "old.png")} + assert _references(database, created.workflow_id) == {("user-1", "image", "old.png")} workflow_records.update(Workflow(**_workflow("new.png").model_dump(), id=created.workflow_id), user_id="user-1") - assert _references(workflow_records, created.workflow_id) == {("user-1", "image", "new.png")} + assert _references(database, created.workflow_id) == {("user-1", "image", "new.png")} + + +def test_a_delete_naming_no_account_drops_the_owners_references( + database: Database, workflow_records: WorkflowRecordsStorage +) -> None: + created = workflow_records.create(_workflow("input.png"), user_id="user-1") + + workflow_records.delete(created.workflow_id) # as an administrator deletes it + + assert _references(database, created.workflow_id) == set() diff --git a/tests/app/services/workflow_records/test_workflow_races.py b/tests/app/services/workflow_records/test_workflow_races.py new file mode 100644 index 00000000000..aec17d76c58 --- /dev/null +++ b/tests/app/services/workflow_records/test_workflow_races.py @@ -0,0 +1,187 @@ +"""Writes of one workflow racing each other. + +A write locks the workflow's row before it reads what it decides on, so a write in flight is either seen by the +write that waited for it or finds the row unchanged; see `tests/fixtures/races.py`. +""" + +import uuid + +import pytest +from sqlalchemy import select + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import ConflictError +from invokeai.app.services.shared.database.queries.workflows import WorkflowQueries +from invokeai.app.services.shared.database.schema.media_references import media_references +from invokeai.app.services.workflow_records.workflow_records_common import ( + Workflow, + WorkflowCategory, + WorkflowIdConflictError, + WorkflowMeta, + WorkflowNotFoundError, + WorkflowRecordDTO, + WorkflowRevisionConflictError, + WorkflowWithoutID, +) +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage +from tests.fixtures.races import when_called + +USER = "user-1" + + +def _workflow(name: str, image_name: str) -> WorkflowWithoutID: + return WorkflowWithoutID( + name=name, + author="", + description="", + version="1.0.0", + contact="", + tags="", + notes="", + exposedFields=[], + meta=WorkflowMeta(version="3.0.0", category=WorkflowCategory.User), + nodes=[{"id": "n1", "data": {"inputs": {"image": {"value": {"image_name": image_name}}}}}], + edges=[], + ) + + +def _edit(created: WorkflowRecordDTO, name: str, image_name: str) -> Workflow: + return Workflow(**_workflow(name, image_name).model_dump(), id=created.workflow_id) + + +def _references(database: Database, workflow_id: str) -> set[str]: + with database.begin(write=False) as conn: + return set( + conn.execute( + select(media_references.c.media_name).where( + media_references.c.owner_kind == "workflow", media_references.c.owner_id == workflow_id + ) + ).scalars() + ) + + +@pytest.mark.parametrize("expects_a_revision", [True, False]) +def test_a_save_waits_for_a_save_in_flight( + workflow_records: WorkflowRecordsStorage, + database: Database, + monkeypatch: pytest.MonkeyPatch, + lost_races: list[ConflictError], + expects_a_revision: bool, +) -> None: + created = workflow_records.create(_workflow("Created", "created.png"), user_id=USER) + + def late_save() -> object: + revision = 1 if expects_a_revision else None + edit = _edit(created, "Late", "late.png") + return workflow_records.update(edit, user_id=USER, expected_revision=revision) + + ended = when_called(monkeypatch, WorkflowQueries, "save", late_save) + workflow_records.update(_edit(created, "First", "first.png"), user_id=USER, expected_revision=1) + errors = ended() + + stored = workflow_records.get(created.workflow_id) + if expects_a_revision: + (error,) = errors + assert isinstance(error, WorkflowRevisionConflictError) and error.current_revision == 2 + assert (stored.workflow.name, stored.revision) == ("First", 2) + assert _references(database, created.workflow_id) == {"first.png"} + else: + # A caller that names no revision writes after the save it waited for. + assert errors == [] + assert (stored.workflow.name, stored.revision) == ("Late", 3) + assert _references(database, created.workflow_id) == {"late.png"} + assert lost_races == [] + + +def test_a_visibility_change_rewrites_the_document_of_the_save_it_waited_for( + workflow_records: WorkflowRecordsStorage, + database: Database, + monkeypatch: pytest.MonkeyPatch, + lost_races: list[ConflictError], +) -> None: + created = workflow_records.create(_workflow("Created", "created.png"), user_id=USER) + + ended = when_called( + monkeypatch, + WorkflowQueries, + "save", + lambda: workflow_records.update_is_public(created.workflow_id, True, user_id=USER), + ) + workflow_records.update(_edit(created, "Edited", "edited.png"), user_id=USER) + + assert ended() == [] + stored = workflow_records.get(created.workflow_id) + assert (stored.workflow.name, stored.workflow.tags, stored.is_public, stored.revision) == ( + "Edited", + "shared", + True, + 2, + ) + assert _references(database, created.workflow_id) == {"edited.png"} + assert lost_races == [] + + +@pytest.mark.parametrize("in_flight", ["save", "delete"]) +def test_a_save_and_a_delete_of_one_workflow_wait_for_each_other( + workflow_records: WorkflowRecordsStorage, + database: Database, + monkeypatch: pytest.MonkeyPatch, + lost_races: list[ConflictError], + in_flight: str, +) -> None: + created = workflow_records.create(_workflow("Created", "created.png"), user_id=USER) + + if in_flight == "save": + ended = when_called( + monkeypatch, WorkflowQueries, "save", lambda: workflow_records.delete(created.workflow_id, user_id=USER) + ) + workflow_records.update(_edit(created, "Edited", "edited.png"), user_id=USER) + assert ended() == [] + else: + ended = when_called( + monkeypatch, + WorkflowQueries, + "delete", + lambda: workflow_records.update(_edit(created, "Edited", "edited.png"), user_id=USER), + ) + workflow_records.delete(created.workflow_id, user_id=USER) + assert [type(e) for e in ended()] == [WorkflowNotFoundError] + + with pytest.raises(WorkflowNotFoundError): + workflow_records.get(created.workflow_id) + assert _references(database, created.workflow_id) == set() + assert lost_races == [] + + +@pytest.mark.parametrize("same_content", [True, False]) +def test_a_creation_sent_twice_at_once_makes_one_workflow( + workflow_records: WorkflowRecordsStorage, + database: Database, + monkeypatch: pytest.MonkeyPatch, + lost_races: list[ConflictError], + same_content: bool, +) -> None: + """On a server the second creation looks for the reserved id while the first has not committed, and only its + insert finds the id taken; it then answers as if it had seen the first creation's record.""" + reserved = str(uuid.uuid4()) + first = _workflow("Created", "first.png") + second = first if same_content else _workflow("Created", "other.png") + results: list[WorkflowRecordDTO] = [] + + ended = when_called( + monkeypatch, + WorkflowQueries, + "insert", + lambda: results.append(workflow_records.create(second, user_id=USER, workflow_id=reserved)), + ) + created = workflow_records.create(first, user_id=USER, workflow_id=reserved) + errors = ended() + + if same_content: + assert errors == [] + assert [(result.workflow_id, result.revision) for result in results] == [(reserved, 1)] + else: + assert [type(e) for e in errors] == [WorkflowIdConflictError] + assert workflow_records.get(reserved) == created + assert _references(database, reserved) == {"first.png"} + assert lost_races == [] diff --git a/tests/app/services/workflow_records/test_workflow_revisions.py b/tests/app/services/workflow_records/test_workflow_revisions.py index 36f15ce0b87..d2bb382921b 100644 --- a/tests/app/services/workflow_records/test_workflow_revisions.py +++ b/tests/app/services/workflow_records/test_workflow_revisions.py @@ -3,9 +3,12 @@ import uuid import pytest +from sqlalchemy import insert, select, update -from invokeai.app.services.invoker import Invoker -from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.media_references import media_references +from invokeai.app.services.shared.database.schema.workflows import workflow_library +from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.services.workflow_records.workflow_records_common import ( Workflow, WorkflowAccessDeniedError, @@ -18,12 +21,7 @@ WorkflowRevisionConflictError, WorkflowWithoutID, ) -from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage - - -@pytest.fixture -def records(mock_invoker: Invoker) -> SqliteWorkflowRecordsStorage: - return mock_invoker.services.workflow_records +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage def _workflow(name: str = "Template", image_name: str | None = None) -> WorkflowWithoutID: @@ -47,17 +45,19 @@ def _with_id(workflow: WorkflowWithoutID, workflow_id: str) -> Workflow: return Workflow(**workflow.model_dump(), id=workflow_id) -def _references(records: SqliteWorkflowRecordsStorage, workflow_id: str) -> set[str]: - with records._db.transaction() as cursor: - cursor.execute( - "SELECT media_name FROM media_references WHERE owner_kind = 'workflow' AND owner_id = ?;", - (workflow_id,), +def _references(database: Database, workflow_id: str) -> set[str]: + with database.begin(write=False) as conn: + return set( + conn.execute( + select(media_references.c.media_name).where( + media_references.c.owner_kind == "workflow", media_references.c.owner_id == workflow_id + ) + ).scalars() ) - return {row[0] for row in cursor.fetchall()} -def _bundled_id(records: SqliteWorkflowRecordsStorage) -> str: - bundled = records.get_many( +def _bundled_id(workflow_records: WorkflowRecordsStorage) -> str: + bundled = workflow_records.get_many( order_by=WorkflowRecordOrderBy.Name, direction=SQLiteDirection.Ascending, categories=[WorkflowCategory.Default], @@ -71,19 +71,19 @@ def _bundled_id(records: SqliteWorkflowRecordsStorage) -> str: return bundled[0].workflow_id -def test_every_content_write_advances_the_revision(records: SqliteWorkflowRecordsStorage) -> None: - created = records.create(_workflow(), user_id="user-1") +def test_every_content_write_advances_the_revision(workflow_records: WorkflowRecordsStorage) -> None: + created = workflow_records.create(_workflow(), user_id="user-1") assert created.revision == 1 - first = records.update(_with_id(_workflow("Edited"), created.workflow_id), user_id="user-1") + first = workflow_records.update(_with_id(_workflow("Edited"), created.workflow_id), user_id="user-1") assert first.revision == 2 # Legacy callers that carry no expected revision still write, and still advance it. - second = records.update(_with_id(_workflow("Edited again"), created.workflow_id)) + second = workflow_records.update(_with_id(_workflow("Edited again"), created.workflow_id)) assert second.revision == 3 - assert records.get(created.workflow_id).workflow.name == "Edited again" + assert workflow_records.get(created.workflow_id).workflow.name == "Edited again" - listed = records.get_many( + listed = workflow_records.get_many( order_by=WorkflowRecordOrderBy.CreatedAt, direction=SQLiteDirection.Descending, categories=[WorkflowCategory.User], @@ -96,12 +96,14 @@ def test_every_content_write_advances_the_revision(records: SqliteWorkflowRecord assert next(item.revision for item in listed if item.workflow_id == created.workflow_id) == 3 -def test_a_stale_expected_revision_is_refused_and_writes_nothing(records: SqliteWorkflowRecordsStorage) -> None: - created = records.create(_workflow(image_name="v1.png"), user_id="user-1") - records.update(_with_id(_workflow(image_name="v2.png"), created.workflow_id), user_id="user-1") +def test_a_stale_expected_revision_is_refused_and_writes_nothing( + workflow_records: WorkflowRecordsStorage, database: Database +) -> None: + created = workflow_records.create(_workflow(image_name="v1.png"), user_id="user-1") + workflow_records.update(_with_id(_workflow(image_name="v2.png"), created.workflow_id), user_id="user-1") with pytest.raises(WorkflowRevisionConflictError) as conflict: - records.update( + workflow_records.update( _with_id(_workflow(image_name="stale.png"), created.workflow_id), user_id="user-1", expected_revision=1, @@ -109,98 +111,107 @@ def test_a_stale_expected_revision_is_refused_and_writes_nothing(records: Sqlite assert conflict.value.current_revision == 2 assert conflict.value.expected_revision == 1 - current = records.get(created.workflow_id) + current = workflow_records.get(created.workflow_id) assert current.revision == 2 - assert _references(records, created.workflow_id) == {"v2.png"} + assert _references(database, created.workflow_id) == {"v2.png"} - accepted = records.update( + accepted = workflow_records.update( _with_id(_workflow(image_name="v3.png"), created.workflow_id), user_id="user-1", expected_revision=2 ) assert accepted.revision == 3 - assert _references(records, created.workflow_id) == {"v3.png"} + assert _references(database, created.workflow_id) == {"v3.png"} -def test_bookkeeping_writes_leave_the_revision_alone(records: SqliteWorkflowRecordsStorage) -> None: - created = records.create(_workflow(image_name="in.png"), user_id="user-1") +def test_bookkeeping_writes_leave_the_revision_alone( + workflow_records: WorkflowRecordsStorage, database: Database +) -> None: + created = workflow_records.create(_workflow(image_name="in.png"), user_id="user-1") - records.update_opened_at(created.workflow_id, user_id="user-1") - records.update_last_run_at(created.workflow_id, user_id="user-1") - shared = records.update_is_public(created.workflow_id, True, user_id="user-1") + workflow_records.update_opened_at(created.workflow_id, user_id="user-1") + workflow_records.update_last_run_at(created.workflow_id, user_id="user-1") + shared = workflow_records.update_is_public(created.workflow_id, True, user_id="user-1") assert shared.revision == 1 assert shared.is_public is True assert "shared" in shared.workflow.tags - assert _references(records, created.workflow_id) == {"in.png"} + assert _references(database, created.workflow_id) == {"in.png"} # An editor holding revision 1 is therefore not asked to resolve a conflict it cannot see. - assert records.update(_with_id(_workflow(), created.workflow_id), expected_revision=1).revision == 2 + assert workflow_records.update(_with_id(_workflow(), created.workflow_id), expected_revision=1).revision == 2 -def test_ownership_is_enforced_atomically_with_the_write(records: SqliteWorkflowRecordsStorage) -> None: - created = records.create(_workflow(image_name="mine.png"), user_id="user-1") +def test_ownership_is_enforced_atomically_with_the_write( + workflow_records: WorkflowRecordsStorage, database: Database +) -> None: + created = workflow_records.create(_workflow(image_name="mine.png"), user_id="user-1") with pytest.raises(WorkflowAccessDeniedError): - records.update(_with_id(_workflow(image_name="stolen.png"), created.workflow_id), user_id="user-2") + workflow_records.update(_with_id(_workflow(image_name="stolen.png"), created.workflow_id), user_id="user-2") with pytest.raises(WorkflowAccessDeniedError): - records.delete(created.workflow_id, user_id="user-2") + workflow_records.delete(created.workflow_id, user_id="user-2") - current = records.get(created.workflow_id) + current = workflow_records.get(created.workflow_id) assert current.revision == 1 - assert _references(records, created.workflow_id) == {"mine.png"} + assert _references(database, created.workflow_id) == {"mine.png"} with pytest.raises(WorkflowNotFoundError): - records.update(_with_id(_workflow(), "missing"), user_id="user-1") + workflow_records.update(_with_id(_workflow(), "missing"), user_id="user-1") -def test_bundled_workflows_are_immutable_whatever_the_request_claims(records: SqliteWorkflowRecordsStorage) -> None: - bundled_id = _bundled_id(records) - before = records.get(bundled_id) +def test_bundled_workflows_are_immutable_whatever_the_request_claims(workflow_records: WorkflowRecordsStorage) -> None: + bundled_id = _bundled_id(workflow_records) + before = workflow_records.get(bundled_id) # The body says "user"; the stored record decides. with pytest.raises(WorkflowImmutableError): - records.update(_with_id(_workflow("Hijacked"), bundled_id)) + workflow_records.update(_with_id(_workflow("Hijacked"), bundled_id)) with pytest.raises(WorkflowImmutableError): - records.update(_with_id(_workflow("Hijacked"), bundled_id), user_id=before.user_id, expected_revision=1) + workflow_records.update( + _with_id(_workflow("Hijacked"), bundled_id), user_id=before.user_id, expected_revision=1 + ) with pytest.raises(WorkflowImmutableError): - records.delete(bundled_id) + workflow_records.delete(bundled_id) - after = records.get(bundled_id) + after = workflow_records.get(bundled_id) assert after.workflow == before.workflow assert after.revision == before.revision def test_bundled_sync_advances_the_revision_only_when_content_changes( - records: SqliteWorkflowRecordsStorage, + workflow_records: WorkflowRecordsStorage, + database: Database, ) -> None: - bundled_id = _bundled_id(records) - assert records.get(bundled_id).revision == 1 + bundled_id = _bundled_id(workflow_records) + assert workflow_records.get(bundled_id).revision == 1 # A second start with unchanged bundles is a no-op for every record. - records._sync_default_workflows() - assert records.get(bundled_id).revision == 1 - - with records._db.transaction() as cursor: - cursor.execute( - "UPDATE workflow_library SET workflow = json_set(workflow, '$.notes', 'stale copy') WHERE workflow_id = ?;", - (bundled_id,), + workflow_records._sync_default_workflows() + assert workflow_records.get(bundled_id).revision == 1 + + stale = workflow_records.get(bundled_id).workflow.model_copy(update={"notes": "stale copy"}) + with database.begin(write=True) as conn: + conn.execute( + update(workflow_library) + .where(workflow_library.c.workflow_id == bundled_id) + .values(workflow=stale.model_dump_json()) ) - records._sync_default_workflows() - resynced = records.get(bundled_id) + workflow_records._sync_default_workflows() + resynced = workflow_records.get(bundled_id) assert resynced.revision == 2 assert resynced.workflow.notes != "stale copy" -def test_a_reserved_id_makes_creation_retry_safe(records: SqliteWorkflowRecordsStorage) -> None: +def test_a_reserved_id_makes_creation_retry_safe(workflow_records: WorkflowRecordsStorage, database: Database) -> None: reserved = str(uuid.uuid4()) workflow = _workflow(image_name="first.png") - created = records.create(workflow, user_id="user-1", workflow_id=reserved) - retried = records.create(workflow, user_id="user-1", workflow_id=reserved) + created = workflow_records.create(workflow, user_id="user-1", workflow_id=reserved) + retried = workflow_records.create(workflow, user_id="user-1", workflow_id=reserved) assert created.workflow_id == reserved assert retried.workflow_id == reserved assert retried.revision == 1 - assert _references(records, reserved) == {"first.png"} - listed = records.get_many( + assert _references(database, reserved) == {"first.png"} + listed = workflow_records.get_many( order_by=WorkflowRecordOrderBy.CreatedAt, direction=SQLiteDirection.Descending, categories=[WorkflowCategory.User], @@ -213,24 +224,92 @@ def test_a_reserved_id_makes_creation_retry_safe(records: SqliteWorkflowRecordsS assert [item.workflow_id for item in listed.items].count(reserved) == 1 -def test_a_reserved_id_that_names_something_else_is_a_conflict(records: SqliteWorkflowRecordsStorage) -> None: +def test_a_reserved_id_that_names_something_else_is_a_conflict( + workflow_records: WorkflowRecordsStorage, database: Database +) -> None: reserved = str(uuid.uuid4()) - records.create(_workflow(image_name="first.png"), user_id="user-1", workflow_id=reserved) + workflow_records.create(_workflow(image_name="first.png"), user_id="user-1", workflow_id=reserved) with pytest.raises(WorkflowIdConflictError): - records.create(_workflow(image_name="other.png"), user_id="user-1", workflow_id=reserved) + workflow_records.create(_workflow(image_name="other.png"), user_id="user-1", workflow_id=reserved) with pytest.raises(WorkflowIdConflictError): - records.create(_workflow(image_name="first.png"), user_id="user-2", workflow_id=reserved) + workflow_records.create(_workflow(image_name="first.png"), user_id="user-2", workflow_id=reserved) with pytest.raises(ValueError): - records.create(_workflow(), user_id="user-1", workflow_id="not-a-uuid") + workflow_records.create(_workflow(), user_id="user-1", workflow_id="not-a-uuid") - assert records.get(reserved).user_id == "user-1" - assert _references(records, reserved) == {"first.png"} + assert workflow_records.get(reserved).user_id == "user-1" + assert _references(database, reserved) == {"first.png"} -def test_creation_without_a_reserved_id_keeps_generating_ids(records: SqliteWorkflowRecordsStorage) -> None: - first = records.create(_workflow(), user_id="user-1") - second = records.create(_workflow(), user_id="user-1") +def test_creation_without_a_reserved_id_keeps_generating_ids(workflow_records: WorkflowRecordsStorage) -> None: + first = workflow_records.create(_workflow(), user_id="user-1") + second = workflow_records.create(_workflow(), user_id="user-1") assert first.workflow_id != second.workflow_id assert uuid.UUID(first.workflow_id) + + +def test_unsharing_removes_the_shared_tag_sharing_added(workflow_records: WorkflowRecordsStorage) -> None: + created = workflow_records.create(_workflow().model_copy(update={"tags": "mine"}), user_id="user-1") + + workflow_records.update_is_public(created.workflow_id, True, user_id="user-1") + unshared = workflow_records.update_is_public(created.workflow_id, False, user_id="user-1") + + assert (unshared.is_public, unshared.workflow.tags) == (False, "mine") + assert workflow_records.get(created.workflow_id) == unshared + + +def test_a_visibility_change_leaves_others_and_bundled_workflows_alone( + workflow_records: WorkflowRecordsStorage, +) -> None: + created = workflow_records.create(_workflow(), user_id="user-1") + bundled_id = _bundled_id(workflow_records) + + refused = workflow_records.update_is_public(created.workflow_id, True, user_id="user-2") + bundled = workflow_records.update_is_public(bundled_id, True) + + assert (refused.is_public, workflow_records.get(created.workflow_id).is_public) == (False, False) + assert (bundled.is_public, workflow_records.get(bundled_id).is_public) == (False, False) + with pytest.raises(WorkflowNotFoundError): + workflow_records.update_is_public("missing", True) + + +def test_a_start_removes_bundled_workflows_no_longer_shipped( + workflow_records: WorkflowRecordsStorage, database: Database +) -> None: + retired = workflow_records.get(_bundled_id(workflow_records)).workflow.model_copy(update={"id": "default_retired"}) + with database.begin(write=True) as conn: + conn.execute(insert(workflow_library).values(workflow_id=retired.id, workflow=retired.model_dump_json())) + + workflow_records._sync_default_workflows() + + with pytest.raises(WorkflowNotFoundError): + workflow_records.get(retired.id) + + +def test_writes_into_the_default_category_are_refused(workflow_records: WorkflowRecordsStorage) -> None: + created = workflow_records.create(_workflow(), user_id="user-1") + bundled_kind = _workflow().model_copy( + update={"meta": WorkflowMeta(version="3.0.0", category=WorkflowCategory.Default)} + ) + + with pytest.raises(ValueError): + workflow_records.create(bundled_kind, user_id="user-1") + with pytest.raises(ValueError): + workflow_records.update(_with_id(bundled_kind, created.workflow_id), user_id="user-1") + + +def test_a_workflow_without_an_owner_belongs_to_the_system_account( + workflow_records: WorkflowRecordsStorage, database: Database +) -> None: + created = workflow_records.create(_workflow(image_name="in.png"), user_id="system") + with database.begin(write=True) as conn: + conn.execute( + update(workflow_library).where(workflow_library.c.workflow_id == created.workflow_id).values(user_id=None) + ) + + with pytest.raises(WorkflowAccessDeniedError): + workflow_records.delete(created.workflow_id, user_id="user-1") + workflow_records.delete(created.workflow_id, user_id="system") + + assert _references(database, created.workflow_id) == set() diff --git a/tests/app/services/workflow_thumbnails/test_workflow_thumbnails_disk.py b/tests/app/services/workflow_thumbnails/test_workflow_thumbnails_disk.py new file mode 100644 index 00000000000..8b802db3572 --- /dev/null +++ b/tests/app/services/workflow_thumbnails/test_workflow_thumbnails_disk.py @@ -0,0 +1,28 @@ +"""Where a workflow's thumbnail is: bundled workflows ship theirs, the accounts' own are stored.""" + +from pathlib import Path +from unittest import mock + +from invokeai.app.services.workflow_records.workflow_records_common import WorkflowCategory +from invokeai.app.services.workflow_thumbnails.workflow_thumbnails_disk import WorkflowThumbnailFileStorageDisk + + +def test_a_caller_that_knows_the_category_spares_reading_the_workflow(tmp_path: Path) -> None: + thumbnails = WorkflowThumbnailFileStorageDisk(tmp_path) + invoker = mock.Mock() + invoker.services.workflow_records.get.side_effect = AssertionError("the workflow was read") + thumbnails.start(invoker) + + assert thumbnails.get_path("mine", category=WorkflowCategory.User) == tmp_path / "mine.webp" + assert thumbnails.get_path("default_x", category=WorkflowCategory.Default).name == "default_x.png" + assert thumbnails.get_path("default_x", category=WorkflowCategory.Default).parent != tmp_path + + +def test_without_a_category_the_workflow_decides(tmp_path: Path) -> None: + thumbnails = WorkflowThumbnailFileStorageDisk(tmp_path) + invoker = mock.Mock() + invoker.services.workflow_records.get.return_value.workflow.meta.category = WorkflowCategory.Default + thumbnails.start(invoker) + + assert thumbnails.get_path("default_x").name == "default_x.png" + invoker.services.workflow_records.get.assert_called_once_with("default_x") diff --git a/tests/app/test_api_dependencies.py b/tests/app/test_api_dependencies.py index e3380cf6c94..186b4659eaf 100644 --- a/tests/app/test_api_dependencies.py +++ b/tests/app/test_api_dependencies.py @@ -5,7 +5,7 @@ from invokeai.app.api import dependencies from invokeai.app.api.dependencies import ApiDependencies from invokeai.app.services.config.config_default import InvokeAIAppConfig -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database def test_initialize_does_not_vacuum_database(monkeypatch, tmp_path) -> None: @@ -18,17 +18,17 @@ def test_initialize_does_not_vacuum_database(monkeypatch, tmp_path) -> None: MagicMock(return_value=MagicMock()), ) - original_init_db = dependencies.init_db - initialized_databases: list[SqliteDatabase] = [] + original_init_database = dependencies.init_database + initialized_databases: list[Database] = [] - def initialize_database(*args, **kwargs) -> SqliteDatabase: - database = original_init_db(*args, **kwargs) + def initialize_database(*args, **kwargs) -> Database: + database = original_init_database(*args, **kwargs) initialized_databases.append(database) return database - monkeypatch.setattr(dependencies, "init_db", initialize_database) - clean_calls: list[SqliteDatabase] = [] - monkeypatch.setattr(SqliteDatabase, "clean", lambda database: clean_calls.append(database)) + monkeypatch.setattr(dependencies, "init_database", initialize_database) + clean_calls: list[Database] = [] + monkeypatch.setattr(Database, "clean", lambda database: clean_calls.append(database)) config = InvokeAIAppConfig( use_memory_db=True, @@ -54,4 +54,7 @@ def initialize_database(*args, **kwargs) -> SqliteDatabase: delattr(ApiDependencies, "invoker") else: ApiDependencies.invoker = original_invoker + for database in initialized_databases: + database.dispose() + ApiDependencies.database = None loop.close() diff --git a/tests/app/test_database_access_guard.py b/tests/app/test_database_access_guard.py new file mode 100644 index 00000000000..59e113d4fc3 --- /dev/null +++ b/tests/app/test_database_access_guard.py @@ -0,0 +1,155 @@ +"""Every SQL statement of the application lives in the database layer. + +Everything else reaches the database through `Database.queries`, so a backend port, a schema change or a +query fix happens in one place (`invokeai/app/services/shared/database/`). This test fails when code +outside the database layer imports a database driver, SQLAlchemy or Alembic, executes a statement, +carries SQL text, or reaches into a database object's internals. +""" + +import ast +import re +from dataclasses import dataclass +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +SCANNED_ROOTS = ("invokeai", "scripts") +DATABASE_LAYER = ( + "invokeai/app/services/shared/database/", + # The 63 migrations up to the portable cutover are SQLite DDL and stay as written. + "invokeai/app/services/shared/sqlite_migrator/", +) + +# Code these rules flag that touches no application database: path -> why. Each entry must still be flagged. +NOT_THE_APPLICATION_DATABASE = { + "invokeai/app/services/gallery_maintenance/gallery_maintenance_default.py": ( + "keeps its scan inventory in a SQLite file of its own, in a temporary directory, not in the app's database" + ), + "invokeai/app/api/routers/gallery_maintenance.py": "`service.execute()` runs a gallery maintenance operation", + "invokeai/app/services/image_files/image_files_disk.py": ( + "walks the output folders with their pending directories in a SQLite file of its own, in a temporary directory" + ), +} + +_DRIVER_MODULES = ("sqlite3", "sqlalchemy", "alembic", "pymysql") +_EXECUTE_METHODS = {"execute", "executemany", "executescript", "exec_driver_sql"} +# Internals of the database layer that only it may import. +_LAYER_PACKAGE = "invokeai.app.services.shared.database" +_LAYER_INTERNAL_MODULES = ("engines", "schema", "dialect") +_LAYER_INTERNALS = tuple(f"{_LAYER_PACKAGE}.{module}" for module in _LAYER_INTERNAL_MODULES) +_SQL_STATEMENT = re.compile( + r"^\s*(SELECT|INSERT|UPDATE|DELETE|CREATE|ALTER|DROP|PRAGMA|WITH|REPLACE|VACUUM|BEGIN)\b[\s\S]*" + r"\b(FROM|INTO|SET|TABLE|INDEX|TRIGGER|AS|IMMEDIATE|TRANSACTION|WHERE|VALUES)\b" +) + + +@dataclass(frozen=True) +class Violation: + path: str + line: int + rule: str + detail: str + + def __str__(self) -> str: + return f"{self.path}:{self.line}: {self.rule}: {self.detail}" + + +def _docstring_nodes(tree: ast.AST) -> set[int]: + docstrings: set[int] = set() + for node in ast.walk(tree): + if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)): + body = node.body + if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant): + docstrings.add(id(body[0].value)) + return docstrings + + +def _is_driver_module(name: str) -> bool: + return any(name == module or name.startswith(f"{module}.") for module in _DRIVER_MODULES) + + +def _looks_like_sql(text: str) -> bool: + return "--sql" in text or _SQL_STATEMENT.match(text) is not None + + +def scan(path: Path, relative: str) -> list[Violation]: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=relative) + docstrings = _docstring_nodes(tree) + f_string_parts = {id(part) for node in ast.walk(tree) if isinstance(node, ast.JoinedStr) for part in node.values} + violations: list[Violation] = [] + + def add(node: ast.AST, rule: str, detail: str) -> None: + violations.append(Violation(relative, getattr(node, "lineno", 0), rule, detail)) + + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if _is_driver_module(alias.name): + add(node, "imports a database driver", alias.name) + elif alias.name.startswith(_LAYER_INTERNALS): + add(node, "imports database layer internals", alias.name) + elif isinstance(node, ast.ImportFrom) and node.module is not None and node.level == 0: + if _is_driver_module(node.module): + add(node, "imports a database driver", node.module) + elif node.module.startswith(_LAYER_INTERNALS) or ( + node.module == _LAYER_PACKAGE and any(alias.name in _LAYER_INTERNAL_MODULES for alias in node.names) + ): + add(node, "imports database layer internals", node.module) + elif isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + if node.func.attr in _EXECUTE_METHODS: + add(node, "executes SQL", f".{node.func.attr}()") + elif isinstance(node, ast.Attribute): + owner_is_self = isinstance(node.value, ast.Name) and node.value.id == "self" + if node.attr in ("_db", "_conn") and not owner_is_self: + add(node, "reaches into a database object", f".{node.attr}") + elif node.attr == "engine" and isinstance(node.ctx, ast.Load) and not owner_is_self: + add(node, "uses the SQLAlchemy engine", ".engine") + elif isinstance(node, ast.JoinedStr): + # An f-string's constant parts are checked together, with the interpolations as placeholders. + text = "".join( + part.value if isinstance(part, ast.Constant) and isinstance(part.value, str) else "{}" + for part in node.values + ) + if _looks_like_sql(text): + add(node, "carries SQL text", " ".join(text.split())[:60]) + elif ( + isinstance(node, ast.Constant) + and isinstance(node.value, str) + and id(node) not in docstrings + and id(node) not in f_string_parts + and _looks_like_sql(node.value) + ): + add(node, "carries SQL text", " ".join(node.value.split())[:60]) + return violations + + +def _scanned_files() -> list[tuple[Path, str]]: + files: list[tuple[Path, str]] = [] + for root in SCANNED_ROOTS: + for path in sorted((REPO_ROOT / root).rglob("*.py")): + relative = path.relative_to(REPO_ROOT).as_posix() + if "node_modules/" in relative or "/__pycache__/" in relative: + continue + if relative.startswith(DATABASE_LAYER): + continue + files.append((path, relative)) + return files + + +def test_only_the_database_layer_touches_the_database() -> None: + offending: list[str] = [] + for path, relative in _scanned_files(): + if relative in NOT_THE_APPLICATION_DATABASE: + continue + violations = scan(path, relative) + if violations: + offending.append(f"{relative}: {len(violations)} violations") + offending.extend(f" {violation}" for violation in violations) + assert not offending, "Database access outside the database layer; add a query module instead:\n" + "\n".join( + offending + ) + + +def test_every_exemption_is_still_needed() -> None: + scanned = {relative: path for path, relative in _scanned_files()} + stale = [relative for relative in NOT_THE_APPLICATION_DATABASE if not scan(scanned[relative], relative)] + assert not stale, f"Remove these from NOT_THE_APPLICATION_DATABASE, the rules no longer flag them: {stale}" diff --git a/tests/app/util/test_database_copy.py b/tests/app/util/test_database_copy.py new file mode 100644 index 00000000000..047fe2c70c0 --- /dev/null +++ b/tests/app/util/test_database_copy.py @@ -0,0 +1,423 @@ +"""`invoke-db-copy`: an install's SQLite database into a new server database, checked before and verified after.""" + +import json +from pathlib import Path +from typing import Optional +from unittest.mock import Mock + +import pytest +from sqlalchemy import URL, delete, func, insert, select, update + +from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage +from invokeai.app.services.config.config_default import load_config_from_root +from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase +from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage +from invokeai.app.services.shared.database import startup +from invokeai.app.services.shared.database.copy import ( + Problem, + copy_database, + copy_records, + count_normalized, + delete_orphans, + find_orphans, + find_oversized, + verify_copy, +) +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.schema.boards import board_images, boards +from invokeai.app.services.shared.database.schema.image_index import image_index_vocab_terms +from invokeai.app.services.shared.database.schema.models import models +from invokeai.app.services.shared.database.schema.session_queue import session_queue +from invokeai.app.services.shared.database.startup import init_database, open_migrated_database +from invokeai.app.services.users.users_common import UserCreateRequest +from invokeai.app.services.users.users_default import UserService +from invokeai.app.util import database_copy +from invokeai.app.util.database_copy import copy, copy_to_sqlite +from invokeai.backend.util.logging import InvokeAILogger +from tests.fixtures.database import external_test_db_url + +server_only = pytest.mark.skipif( + external_test_db_url() is None, reason="needs a MySQL or MariaDB server (INVOKEAI_TEST_DB_URL)" +) +LOGGER = InvokeAILogger.get_logger("test_database_copy") +MODEL_CONFIG = { + "hash": "blake3:0123456789abcdef", + "base": "sdxl", + "type": "main", + "path": "sdxl/main/model.safetensors", + "format": "checkpoint", + "name": "Model", + "source": "https://example.com/model.safetensors", + "source_type": "url", + # Stored as JSON wrote it, which a server's integer column refuses. + "file_size": 6_938_078_334.5, +} + + +def _install(root: Path, *, orphans: bool) -> None: + """An install whose SQLite database holds an account, a board with an image, a model, two vocabulary terms a + server treats as one, and (with `orphans`) an image membership of a board that is gone and a board whose cover + image is gone.""" + (root / "invokeai.yaml").write_text('schema_version: "4.0.3"\n') + database = init_database(load_config_from_root(root), LOGGER, Mock(spec=ImageFileStorageBase)) + try: + user = UserService(database).create(UserCreateRequest(email="alice@test.com", password="AlicePass123")) + board = BoardRecordStorage(database).save("Holiday", user.user_id) + ImageRecordStorage(database).save( + image_name="beach.png", + image_origin=ResourceOrigin.INTERNAL, + image_category=ImageCategory.GENERAL, + width=64, + height=64, + has_workflow=False, + user_id=user.user_id, + ) + BoardImageRecordStorage(database).add_image_to_board(board.board_id, "beach.png") + with database.begin(write=True) as conn: + conn.execute(insert(models).values(id="model-1", config=json.dumps(MODEL_CONFIG))) + conn.execute(insert(image_index_vocab_terms), [{"term": "Äpfel"}, {"term": "äpfel"}, {"term": "beach"}]) + if orphans: + raw = database.sqlite.conn + raw.execute("PRAGMA foreign_keys = OFF") + raw.execute("DELETE FROM board_images") + raw.execute("INSERT INTO board_images (board_id, image_name) VALUES ('gone', 'beach.png')") + raw.execute("UPDATE boards SET cover_image_name = 'gone.png'") + raw.commit() + raw.execute("PRAGMA foreign_keys = ON") + finally: + database.dispose() + + +def _open(root: Path) -> Database: + return open_migrated_database(load_config_from_root(root), LOGGER) + + +def test_orphans_are_found_and_left_out_or_cleared(tmp_path: Path) -> None: + _install(tmp_path, orphans=True) + database = _open(tmp_path) + try: + assert [(problem.table, problem.count) for problem in find_orphans(database)] == [ + ("board_images", 1), + ("boards", 1), + ] + assert [(problem.table, problem.count) for problem in count_normalized(database)] == [("models", 1)] + + assert delete_orphans(database) == [ + Problem("board_images", 1, "left out"), + Problem("boards", 1, "with a reference to a missing row, which was cleared"), + ] + assert find_orphans(database) == [] + # The board stays, without the cover the database would have cleared itself. + with database.begin(write=False) as conn: + assert conn.execute(select(boards.c.board_name, boards.c.cover_image_name)).all() == [("Holiday", None)] + finally: + database.dispose() + + +def test_values_a_server_cannot_store_are_found(tmp_path: Path) -> None: + _install(tmp_path, orphans=False) + database = _open(tmp_path) + try: + assert find_oversized(database, max_allowed_packet=64 * 1024 * 1024) == [] + + with database.begin(write=True) as conn: + conn.execute(insert(image_index_vocab_terms).values(term="t" * 256)) + # The server computes `path` from the config, into a column of 768 characters. + config = {**MODEL_CONFIG, "path": "p" * 769} + conn.execute(insert(models).values(id="m" * 256, config=json.dumps(config))) + found = find_oversized(database, max_allowed_packet=64 * 1024 * 1024) + assert sorted(found) == [ + Problem("image_index_vocab_terms", 1, "with a term longer than 255 characters"), + Problem("models", 1, "with a path longer than 768 characters"), + Problem("models", 1, "with an id longer than 255 characters"), + ] + + # A row is sent in one statement: it must fit half a packet, which leaves room for escaping. + assert Problem("models", 1, "larger than half the server's max_allowed_packet") in find_oversized( + database, max_allowed_packet=1600 + ) + finally: + database.dispose() + + +def test_a_copy_is_verified_row_by_row(tmp_path: Path) -> None: + _install(tmp_path, orphans=False) + source = _open(tmp_path) + target = Database.open_sqlite(tmp_path / "target.db", LOGGER) + try: + copy_database(source, target) + assert verify_copy(source, target) == [] + copy_records(source, target) + assert verify_copy(source, target, records=True) == [] + + with target.begin(write=True) as conn: + conn.execute(update(boards).values(board_name="Renamed")) + conn.execute(delete(board_images)) + assert [mismatch.split(":")[0] for mismatch in verify_copy(source, target)] == ["boards", "board_images"] + finally: + source.dispose() + target.dispose() + + +def test_a_merged_table_holds_one_of_each_set_of_equal_rows_and_nothing_else(tmp_path: Path) -> None: + _install(tmp_path, orphans=False) + source = _open(tmp_path) + target = Database.open_sqlite(tmp_path / "target.db", LOGGER) + try: + with source.begin(write=True) as conn: + # Equal to "beach" on a server, whose collations ignore a soft hyphen. + conn.execute(insert(image_index_vocab_terms).values(term="bea\u00adch")) + copy_database(source, target) + # As a server merges them. + with target.begin(write=True) as conn: + conn.execute( + delete(image_index_vocab_terms).where(image_index_vocab_terms.c.term.in_(["äpfel", "bea\u00adch"])) + ) + assert verify_copy(source, target) == [] + + with target.begin(write=True) as conn: + conn.execute(delete(image_index_vocab_terms).where(image_index_vocab_terms.c.term == "beach")) + assert verify_copy(source, target) == ["image_index_vocab_terms: rows or contents differ"] + + with target.begin(write=True) as conn: + conn.execute(insert(image_index_vocab_terms).values(term="beach", created_at="2000-01-01 00:00:00.000")) + assert verify_copy(source, target) == ["image_index_vocab_terms: rows or contents differ"] + finally: + source.dispose() + target.dispose() + + +def test_without_a_target_the_copy_asks_for_one(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + _install(tmp_path, orphans=False) + + assert copy(None, root=tmp_path, check_only=True, skip_orphans=False) == 1 + assert "Name the target" in capsys.readouterr().out + + +def test_the_copy_reports_a_target_it_cannot_use(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + _install(tmp_path, orphans=False) + + assert copy("postgresql://u:p@localhost/db", root=tmp_path, check_only=True, skip_orphans=False) == 1 + assert "Cannot copy" in capsys.readouterr().out + + +@server_only +def test_a_copy_holds_every_row_of_its_source( + tmp_path: Path, + empty_database: Database, + _external_test_schema: Optional[URL], + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + assert _external_test_schema is not None + # The test servers keep their own defaults; what is checked here is the copy. + monkeypatch.setattr(startup, "MINIMUM_MAX_ALLOWED_PACKET", 1) + _install(tmp_path, orphans=True) + url = _external_test_schema.render_as_string(hide_password=False) + # The target is the database the install's config names. + (tmp_path / "invokeai.yaml").write_text(f'schema_version: "4.0.3"\ndb_url: "{url}"\n') + + assert copy(None, root=tmp_path, check_only=False, skip_orphans=False) == 1 + assert "--orphans skip" in capsys.readouterr().out + + assert copy(None, root=tmp_path, check_only=True, skip_orphans=True) == 0 + assert "Nothing was copied" in capsys.readouterr().out + + assert copy(None, root=tmp_path, check_only=False, skip_orphans=True) == 0 + output = capsys.readouterr().out + assert "board_images: 1 row left out" in output + assert "every table matches its source" in output + assert "image_index_vocab_terms: 1 row the target treats as equal to others were merged" in output + + # The install's own database keeps what the copy left out. + source = open_migrated_database(load_config_from_root(tmp_path).model_copy(update={"db_url": None}), LOGGER) + try: + assert [(problem.table, problem.count) for problem in find_orphans(source)] == [ + ("board_images", 1), + ("boards", 1), + ] + finally: + source.dispose() + + # The app takes the copy as it takes a database it created. + target = open_migrated_database(load_config_from_root(tmp_path), LOGGER) + try: + [user] = [user for user in UserService(target).list_users() if user.email == "alice@test.com"] + [board] = BoardRecordStorage(target).get_all(user.user_id, False, "board_name", "ASC") # type: ignore[arg-type] + assert (board.board_name, board.cover_image_name) == ("Holiday", None) + assert ImageRecordStorage(target).exists("beach.png") + with target.begin(write=False) as conn: + assert conn.execute(select(models.c.file_size)).scalar_one() == round(MODEL_CONFIG["file_size"]) + assert conn.execute(select(func.count()).select_from(board_images)).scalar_one() == 0 + finally: + target.dispose() + + # A database is copied once: into a new, empty one. + assert copy(None, root=tmp_path, check_only=False, skip_orphans=True) == 1 + assert "not empty" in capsys.readouterr().out + + +@server_only +def test_a_failed_copy_leaves_a_target_the_app_refuses_and_says_how_to_start_again( + tmp_path: Path, + empty_database: Database, + _external_test_schema: Optional[URL], + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + assert _external_test_schema is not None + monkeypatch.setattr(startup, "MINIMUM_MAX_ALLOWED_PACKET", 1) + _install(tmp_path, orphans=False) + url = _external_test_schema.render_as_string(hide_password=False) + + def lost(source: Database, target: Database) -> None: + raise ConnectionError("the connection to the server was lost") + + monkeypatch.setattr(database_copy, "copy_records", lost) + + assert copy(url, root=tmp_path, check_only=False, skip_orphans=False) == 1 + assert "Drop the target database" in capsys.readouterr().out + config = load_config_from_root(tmp_path).model_copy(update={"db_url": url}) + with pytest.raises(Exception, match="no record of the migrations"): + open_migrated_database(config, LOGGER).dispose() + + +def test_a_copy_to_sqlite_asks_for_its_source(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + _install(tmp_path, orphans=False) + + assert copy_to_sqlite(tmp_path / "back.db", source_url=None, root=tmp_path, check_only=False) == 1 + assert "Name the source" in capsys.readouterr().out + + +def test_a_copy_to_sqlite_overwrites_no_file(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + _install(tmp_path, orphans=False) + existing = tmp_path / "databases" / "invokeai.db" + before = existing.read_bytes() + + url = "mariadb+pymysql://invokeai:secret@127.0.0.1:1/invokeai" + assert copy_to_sqlite(existing, source_url=url, root=tmp_path, check_only=False) == 1 + assert "overwrites no database" in capsys.readouterr().out + assert existing.read_bytes() == before + + +def _moved_to_a_server(tmp_path: Path, url: str, monkeypatch: pytest.MonkeyPatch) -> None: + """An install whose database `invoke-db-copy` moved to the server, which its config now names.""" + monkeypatch.setattr(startup, "MINIMUM_MAX_ALLOWED_PACKET", 1) + _install(tmp_path, orphans=False) + (tmp_path / "invokeai.yaml").write_text(f'schema_version: "4.0.3"\ndb_url: "{url}"\n') + assert copy(None, root=tmp_path, check_only=False, skip_orphans=False) == 0 + + +@server_only +def test_a_server_database_goes_back_to_sqlite( + tmp_path: Path, + empty_database: Database, + _external_test_schema: Optional[URL], + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + assert _external_test_schema is not None + url = _external_test_schema.render_as_string(hide_password=False) + _moved_to_a_server(tmp_path, url, monkeypatch) + + # Work done on the server after the move, which the old SQLite file does not hold. + server = open_migrated_database(load_config_from_root(tmp_path), LOGGER) + try: + with server.engine.connect() as conn: + # As MySQL's background statistics update does: it then serves the id counter of that moment from a + # cache, until the cache expires a day later. + conn.exec_driver_sql("ANALYZE TABLE session_queue").all() + conn.commit() + [user] = [user for user in UserService(server).list_users() if user.email == "alice@test.com"] + BoardRecordStorage(server).save("Made on the server", user.user_id) + with server.begin(write=True) as conn: + item = {"batch_id": "b", "queue_id": "default", "session": "{}"} + conn.execute(insert(session_queue), [{**item, "session_id": f"s{n}"} for n in range(3)]) + newest = conn.execute(select(func.max(session_queue.c.item_id))).scalar_one() + conn.execute(delete(session_queue).where(session_queue.c.item_id == newest)) + finally: + server.dispose() + capsys.readouterr() + + back = tmp_path / "back" / "invokeai.db" + back.parent.mkdir() + assert copy_to_sqlite(back, source_url=None, root=tmp_path, check_only=False) == 0 + assert "every table matches its source" in capsys.readouterr().out + assert sorted(path.name for path in back.parent.iterdir()) == ["invokeai.db"] + + # The app takes the file as an up-to-date database, holding what the server held. + config = load_config_from_root(tmp_path).model_copy(update={"db_url": None, "db_dir": back.parent}) + copied = open_migrated_database(config, LOGGER) + try: + boards_held = BoardRecordStorage(copied).get_all(user.user_id, False, "board_name", "ASC") # type: ignore[arg-type] + assert sorted(board.board_name for board in boards_held) == ["Holiday", "Made on the server"] + with copied.begin(write=True) as conn: + conn.execute(insert(session_queue).values(batch_id="b", queue_id="default", session="{}", session_id="new")) + # The id the server issued and deleted is not issued again. + assert conn.execute(select(func.max(session_queue.c.item_id))).scalar_one() > newest + finally: + copied.dispose() + + +@server_only +def test_a_copy_to_sqlite_waits_for_no_running_app_and_keeps_no_failed_copy( + tmp_path: Path, + empty_database: Database, + _external_test_schema: Optional[URL], + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + assert _external_test_schema is not None + url = _external_test_schema.render_as_string(hide_password=False) + _moved_to_a_server(tmp_path, url, monkeypatch) + back = tmp_path / "back.db" + + running = Database.open_url(url, LOGGER) + try: + running.hold_instance_lock() + assert copy_to_sqlite(back, source_url=None, root=tmp_path, check_only=False) == 1 + assert "Another InvokeAI process" in capsys.readouterr().out + finally: + running.dispose() + + assert copy_to_sqlite(back, source_url=None, root=tmp_path, check_only=True) == 0 + assert "Nothing was copied" in capsys.readouterr().out + assert not back.exists() + + def lost(source: Database, target: Database) -> None: + raise ConnectionError("the connection to the server was lost") + + monkeypatch.setattr(database_copy, "copy_records", lost) + assert copy_to_sqlite(back, source_url=None, root=tmp_path, check_only=False) == 1 + assert "not kept" in capsys.readouterr().out + assert not back.exists() + assert not [path for path in tmp_path.iterdir() if path.name.startswith("invoke-db-copy-")] + + # A file made at the target path while copying is not overwritten. + def made_meanwhile(source: Database, target: Database) -> dict[str, int]: + back.write_bytes(b"made meanwhile") + return copy_records(source, target) + + monkeypatch.setattr(database_copy, "copy_records", made_meanwhile) + assert copy_to_sqlite(back, source_url=None, root=tmp_path, check_only=False) == 1 + assert "overwrites no file" in capsys.readouterr().out + assert back.read_bytes() == b"made meanwhile" + assert not [path for path in tmp_path.iterdir() if path.name.startswith("invoke-db-copy-")] + + +@pytest.mark.parametrize( + "arguments", + [ + ["--to-sqlite", "back.db", "--target", "mysql+pymysql://u:p@db/invokeai"], + ["--to-sqlite", "back.db", "--orphans", "skip"], + ["--source", "mysql+pymysql://u:p@db/invokeai"], + ], +) +def test_arguments_of_the_other_direction_are_refused(monkeypatch: pytest.MonkeyPatch, arguments: list[str]) -> None: + monkeypatch.setattr("sys.argv", ["invoke-db-copy", *arguments]) + + with pytest.raises(SystemExit) as refused: + database_copy.main() + assert refused.value.code == 2 diff --git a/tests/app/util/test_user_management.py b/tests/app/util/test_user_management.py index 8d8c1215353..1e8e6ad4ebf 100644 --- a/tests/app/util/test_user_management.py +++ b/tests/app/util/test_user_management.py @@ -2,6 +2,7 @@ import logging from pathlib import Path +from types import SimpleNamespace from unittest.mock import Mock import pytest @@ -10,7 +11,7 @@ from invokeai.app.services.fonts.fonts_common import FontScope from invokeai.app.services.fonts.fonts_default import FontService from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase -from invokeai.app.services.shared.sqlite.sqlite_util import init_db +from invokeai.app.services.shared.database.startup import init_database from invokeai.app.services.users.users_common import UserCreateRequest from invokeai.app.services.users.users_default import UserService from invokeai.app.util import user_management @@ -26,9 +27,9 @@ def deletion_fixture(tmp_path: Path): ) config._root = tmp_path logger = logging.getLogger("user-management-font-tests") - db = init_db(config, logger, Mock(spec=ImageFileStorageBase)) + db = init_database(config, logger, Mock(spec=ImageFileStorageBase)) fonts = FontService( - db=db, + db, fonts_dir=config.fonts_path, storage_dir=config.fonts_storage_path, logger=logger, @@ -66,7 +67,7 @@ def test_user_delete_commands_clean_private_fonts_and_keep_shared_fonts( assert shared_path.is_file() assert pending_path.is_file() - monkeypatch.setattr("invokeai.app.services.config.get_config", lambda: config) + monkeypatch.setattr(user_management, "_cli_config", lambda: config) if interactive: answers = iter([email, "yes"]) monkeypatch.setattr("builtins.input", lambda _prompt: next(answers)) @@ -79,3 +80,11 @@ def test_user_delete_commands_clean_private_fonts_and_keep_shared_fonts( assert shared_path.is_file() assert pending_path.is_file() assert fonts.get_accessible(user_id="system", font_id=shared_id).id == shared_id + + +def test_the_commands_use_the_database_the_install_configures(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + # `get_config` outside the app reads no `invokeai.yaml`: a database configured there would be missed. + (tmp_path / "invokeai.yaml").write_text('schema_version: "4.0.3"\ndb_dir: elsewhere\n') + monkeypatch.setattr("invokeai.app.services.config.get_config", lambda: SimpleNamespace(root_path=tmp_path)) + + assert user_management._cli_config().db_path == (tmp_path / "elsewhere" / "invokeai.db").resolve() diff --git a/tests/backend/image_util/test_vendor_mutable_defaults.py b/tests/backend/image_util/test_vendor_mutable_defaults.py index 8cd74cbd6c3..3daba8864db 100644 --- a/tests/backend/image_util/test_vendor_mutable_defaults.py +++ b/tests/backend/image_util/test_vendor_mutable_defaults.py @@ -1,10 +1,4 @@ -"""Tests for the mutable default argument fix in imwatermark/vendor.py -and the bare except fix in sqlite_database.py.""" - -from logging import Logger -from unittest import mock - -import pytest +"""Tests for the mutable default argument fix in imwatermark/vendor.py.""" from invokeai.backend.image_util.imwatermark.vendor import EmbedMaxDct, WatermarkEncoder @@ -61,46 +55,3 @@ def test_explicit_args_still_work(self): assert e._wmLen == 4 assert e._scales == sc assert e._block == 8 - - -class TestTransactionExceptException: - """The transaction() context manager used to have a bare `except:`. - After the fix it uses `except Exception:`, so BaseException subclasses - like KeyboardInterrupt and SystemExit should propagate instead of - being silently caught and rolled back.""" - - @staticmethod - def _make_db(): - """Create a minimal SqliteDatabase-like object with transaction().""" - # Import here so the test stays focused; we just need the real class. - from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase - - logger = mock.MagicMock(spec=Logger) - db = SqliteDatabase(db_path=None, logger=logger, verbose=False) - return db - - def test_regular_exception_rolls_back(self): - db = self._make_db() - - # create a table first in a successful transaction - with db.transaction() as cursor: - cursor.execute("CREATE TABLE t (id INTEGER)") - - # now try to insert and fail — the insert should be rolled back - with pytest.raises(ValueError): - with db.transaction() as cursor: - cursor.execute("INSERT INTO t VALUES (42)") - raise ValueError("boom") - - # the row should not exist after rollback - with db.transaction() as cursor: - cursor.execute("SELECT * FROM t") - assert cursor.fetchone() is None - - def test_keyboard_interrupt_propagates(self): - with pytest.raises(KeyboardInterrupt): - raise KeyboardInterrupt() - - def test_system_exit_propagates(self): - with pytest.raises(SystemExit): - raise SystemExit(1) diff --git a/tests/conftest.py b/tests/conftest.py index 3d555e61a2b..badeaca0e79 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,38 +11,62 @@ import pytest -from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage -from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage -from invokeai.app.services.board_video_records.board_video_records_sqlite import SqliteBoardVideoRecordStorage +from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage +from invokeai.app.services.board_video_records.board_video_records_default import BoardVideoRecordStorage from invokeai.app.services.boards.boards_default import BoardService from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService -from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ClientStatePersistenceSqlite +from invokeai.app.services.client_state_persistence.client_state_persistence_default import ClientStatePersistence from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.external_generation.external_generation_default import ExternalGenerationService -from invokeai.app.services.gallery.gallery_default import SqliteGalleryService +from invokeai.app.services.gallery.gallery_default import GalleryService from invokeai.app.services.image_index.image_index_default import ImageIndexService -from invokeai.app.services.image_index.image_index_records_sqlite import ImageIndexRecordsSqlite -from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage +from invokeai.app.services.image_index.image_index_records_default import ImageIndexRecords +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage from invokeai.app.services.images.images_default import ImageService from invokeai.app.services.intermediates.intermediates_default import IntermediatesService -from invokeai.app.services.intermediates.intermediates_records_sqlite import IntermediatesRecordsSqlite +from invokeai.app.services.intermediates.intermediates_records_default import IntermediatesRecords from invokeai.app.services.invocation_cache.invocation_cache_memory import MemoryInvocationCache from invokeai.app.services.invocation_services import InvocationServices from invokeai.app.services.invocation_stats.invocation_stats_default import InvocationStatsService from invokeai.app.services.invoker import Invoker -from invokeai.app.services.project_records.project_records_sqlite import ProjectRecordsSqlite -from invokeai.app.services.system_prompt_records.system_prompt_records_sqlite import ( - SqliteSystemPromptRecordsStorage, +from invokeai.app.services.project_records.project_records_default import ProjectRecordsStorage +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.system_prompt_records.system_prompt_records_default import ( + SystemPromptRecordsStorage, ) from invokeai.app.services.users.users_default import UserService -from invokeai.app.services.video_records.video_records_sqlite import SqliteVideoRecordStorage -from invokeai.app.services.wildcard_records.wildcard_records_sqlite import SqliteWildcardRecordsStorage -from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage +from invokeai.app.services.video_records.video_records_default import VideoRecordStorage +from invokeai.app.services.wildcard_records.wildcard_records_default import WildcardRecordsStorage +from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage from invokeai.backend.util.logging import InvokeAILogger from tests.backend.model_manager.model_manager_fixtures import * # noqa: F403 +from tests.fixtures.database import ( # noqa: F401 + _external_application_schema, + _external_test_schema, + _migrated_sqlite, + database, + empty_database, + external_test_db_url, +) +from tests.fixtures.races import lost_races # noqa: F401 from tests.fixtures.sqlite_database import create_mock_sqlite_database # noqa: F401 from tests.test_nodes import TestEventService +# Fixtures that put a test on the backend under test, which `-m uses_database` selects. +_DATABASE_FIXTURES = {"empty_database", "database"} + + +@pytest.hookimpl(tryfirst=True) +def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: + # tryfirst: the marks must exist before `-m` deselects by them. + external = external_test_db_url() is not None + for item in items: + if _DATABASE_FIXTURES & set(getattr(item, "fixturenames", ())): + item.add_marker(pytest.mark.uses_database) + if external and item.get_closest_marker("sqlite_only") is not None: + item.add_marker(pytest.mark.skip(reason="needs SQLite, but INVOKEAI_TEST_DB_URL names another backend")) + @pytest.fixture(autouse=True) def _clear_deferred_empty_cache(): @@ -56,26 +80,33 @@ def _clear_deferred_empty_cache(): @pytest.fixture -def mock_services() -> InvocationServices: +def mock_sqlite_database() -> Database: + """The in-memory database behind `mock_services`, for tests that need SQL of their own or a service that + `mock_services` leaves out (the session queue).""" + return create_mock_sqlite_database(InvokeAIAppConfig(use_memory_db=True), InvokeAILogger.get_logger()) + + +@pytest.fixture +def mock_services(mock_sqlite_database: Database) -> InvocationServices: # Image indexing is on by default, but `model_manager` below is None: starting # the indexer against these stub services would fail while resolving the # embedding model. Tests that exercise the index enable it themselves. configuration = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0, image_index_enabled=False) logger = InvokeAILogger.get_logger() - db = create_mock_sqlite_database(configuration, logger) + db = mock_sqlite_database # NOTE: none of these are actually called by the test invocations return InvocationServices( - board_image_records=SqliteBoardImageRecordStorage(db=db), + board_image_records=BoardImageRecordStorage(db), board_images=None, # type: ignore - board_records=SqliteBoardRecordStorage(db=db), + board_records=BoardRecordStorage(db), boards=BoardService(), bulk_download=BulkDownloadService(), configuration=configuration, database=db, events=TestEventService(), image_files=None, # type: ignore - image_records=SqliteImageRecordStorage(db=db), + image_records=ImageRecordStorage(db), images=ImageService(), invocation_cache=MemoryInvocationCache(max_cache_size=0), logger=logging, # type: ignore @@ -88,29 +119,29 @@ def mock_services() -> InvocationServices: session_processor=None, # type: ignore session_queue=None, # type: ignore urls=None, # type: ignore - workflow_records=SqliteWorkflowRecordsStorage(db=db), + workflow_records=WorkflowRecordsStorage(db), tensors=None, # type: ignore conditioning=None, # type: ignore style_preset_records=None, # type: ignore style_preset_image_files=None, # type: ignore - system_prompt_records=SqliteSystemPromptRecordsStorage(db=db), + system_prompt_records=SystemPromptRecordsStorage(db), workflow_thumbnails=None, # type: ignore model_relationship_records=None, # type: ignore model_relationships=None, # type: ignore - client_state_persistence=ClientStatePersistenceSqlite(db=db), - project_records=ProjectRecordsSqlite(db=db), + client_state_persistence=ClientStatePersistence(db), + project_records=ProjectRecordsStorage(db), users=UserService(db), - wildcard_records=SqliteWildcardRecordsStorage(db=db), + wildcard_records=WildcardRecordsStorage(db), videos=None, # type: ignore video_files=None, # type: ignore - video_records=SqliteVideoRecordStorage(db=db), - board_video_records=SqliteBoardVideoRecordStorage(db=db), + video_records=VideoRecordStorage(db), + board_video_records=BoardVideoRecordStorage(db), # Real SQLite-backed gallery service: the virtual-boards router reads dates and # per-date item names through it, and MagicMock cannot exercise the filter SQL. - gallery=SqliteGalleryService(db=db), - image_index_records=ImageIndexRecordsSqlite(db=db), + gallery=GalleryService(db), + image_index_records=ImageIndexRecords(db), image_index=ImageIndexService(), - intermediates=IntermediatesService(records=IntermediatesRecordsSqlite(db=db), logger=logger), + intermediates=IntermediatesService(records=IntermediatesRecords(db), logger=logger), ) diff --git a/tests/fixtures/database.py b/tests/fixtures/database.py new file mode 100644 index 00000000000..9bfa90a59b5 --- /dev/null +++ b/tests/fixtures/database.py @@ -0,0 +1,252 @@ +"""Databases on the backend under test. + +Tests run against SQLite unless INVOKEAI_TEST_DB_URL names a MySQL or MariaDB server, e.g. +`mysql+pymysql://root:secret@127.0.0.1:3306`. Each xdist worker of each test run then gets schemas of +its own on that server, so neither workers nor concurrent runs (an IDE and a terminal, two worktrees) see +each other's rows. The account needs the privilege to create and drop databases. +""" + +import os +import re +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional +from unittest import mock + +import pytest +from sqlalchemy import URL, Engine, create_engine, delete, event, insert, make_url, select + +from invokeai.app.services.config.config_default import DefaultInvokeAIAppConfig +from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.engines import ( + MARIADB_BINARY_COLLATION, + MYSQL_BINARY_COLLATION, + create_mysql_engine, +) +from invokeai.app.services.shared.database.schema import metadata +from invokeai.app.services.shared.sqlite_migrator.migration_loader import MigrationBuildContext, build_migrations +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import Migrator +from invokeai.backend.util.logging import InvokeAILogger + +TEST_DB_URL_ENV = "INVOKEAI_TEST_DB_URL" +_SCHEMA_PREFIX = "invokeai_test_" +_GENERATED_SCHEMA = re.compile(r"invokeai_test_(?:[0-9a-f]{8}_)?(?:main|gw[0-9]+)(?:_app)?") + + +def external_test_db_url() -> Optional[str]: + """The server named by INVOKEAI_TEST_DB_URL, or None when tests run against SQLite.""" + return os.environ.get(TEST_DB_URL_ENV) or None + + +def migrate_to_newest(database: Database, root: Path) -> None: + """Runs the application's migrations: a new database is then at the newest schema, with its seeded rows. + + The migrations clean up legacy files under the root, so it must be a test's own directory; the config + reads nothing from the environment that could point them elsewhere. + """ + config = DefaultInvokeAIAppConfig() + config._root = root + migrator = Migrator(database) + context = MigrationBuildContext( + app_config=config, + logger=InvokeAILogger.get_logger("test_database"), + image_files=mock.Mock(spec=ImageFileStorageBase), + ) + for migration in build_migrations(context): + migrator.register_migration(migration) + migrator.run_migrations() + + +@contextmanager +def _server_schema(base: URL, suffix: str) -> Iterator[URL]: + """Creates a schema of this run and worker on the server, and drops it again.""" + run = os.environ.get("PYTEST_XDIST_TESTRUNUID", uuid.uuid4().hex)[:8] + schema = f"{_SCHEMA_PREFIX}{run}_{os.environ.get('PYTEST_XDIST_WORKER', 'main')}{suffix}" + # `URL.set(database=None)` keeps the component, so the server-level URL is built from parts. + server = create_engine( + URL.create(base.drivername, base.username, base.password, base.host, base.port, query=base.query) + ) + # Held for the whole session: it tells other runs this schema is in use. A schema whose lock nobody + # holds was left behind by a run that crashed, and is dropped. + owner = server.connect() + try: + owner.exec_driver_sql("SELECT GET_LOCK(%s, 0)", (schema,)) + _drop_abandoned_schemas(server) + with server.begin() as conn: + is_mariadb = "mariadb" in str(conn.exec_driver_sql("SELECT VERSION()").scalar_one()).lower() + # The binary NO PAD collation the layer requires, for the tables tests create. + collation = MARIADB_BINARY_COLLATION if is_mariadb else MYSQL_BINARY_COLLATION + conn.exec_driver_sql(f"CREATE DATABASE `{schema}` CHARACTER SET utf8mb4 COLLATE {collation}") + yield base.set(database=schema) + finally: + with server.begin() as conn: + conn.exec_driver_sql(f"DROP DATABASE IF EXISTS `{schema}`") + owner.close() + server.dispose() + + +def _drop_abandoned_schemas(server: Engine) -> None: + with server.connect() as conn: + schemas: list[str] = list( + conn.exec_driver_sql( + "SELECT schema_name FROM information_schema.schemata WHERE schema_name LIKE %s", + (_SCHEMA_PREFIX.replace("_", "\\_") + "%",), + ).scalars() + ) + for schema in schemas: + # Only names this fixture generates (and its older fixed ones), never someone's own database. + if not _GENERATED_SCHEMA.fullmatch(schema): + continue + if conn.exec_driver_sql("SELECT GET_LOCK(%s, 0)", (schema,)).scalar_one() == 1: + conn.exec_driver_sql(f"DROP DATABASE IF EXISTS `{schema}`") + conn.exec_driver_sql("SELECT RELEASE_LOCK(%s)", (schema,)) + conn.commit() + + +def _open(url: URL) -> Database: + return Database.open_url(url.render_as_string(hide_password=False), InvokeAILogger.get_logger("test_database")) + + +@pytest.fixture(scope="session") +def _external_test_schema() -> Iterator[Optional[URL]]: + url = external_test_db_url() + if url is None: + yield None + return + with _server_schema(make_url(url), "") as schema: + yield schema + + +@dataclass(frozen=True) +class _ApplicationSchema: + engine: Engine + # The rows of a new database, per table: the generated columns left out, as the database computes them. + rows: dict[str, list[dict[str, Any]]] + + +@pytest.fixture(scope="session") +def _external_application_schema(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Optional[_ApplicationSchema]]: + url = external_test_db_url() + if url is None: + yield None + return + with _server_schema(make_url(url), "_app") as schema: + engine = create_mysql_engine(schema) + try: + database = Database(engine, InvokeAILogger.get_logger("test_database")) + migrate_to_newest(database, tmp_path_factory.mktemp("migrations")) + with database.begin(write=False) as conn: + rows = { + table.name: [ + row._asdict() + for row in conn.execute( + select(*[column for column in table.columns if column.computed is None]) + ) + ] + for table in metadata.sorted_tables + } + yield _ApplicationSchema(engine=engine, rows={name: each for name, each in rows.items() if each}) + finally: + engine.dispose() + + +@pytest.fixture(scope="session") +def _migrated_sqlite(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Database]: + """An in-memory SQLite database migrated once per run, which the `database` fixture copies for each test.""" + database = Database.open_sqlite(None, InvokeAILogger.get_logger("test_database")) + try: + migrate_to_newest(database, tmp_path_factory.mktemp("migrations")) + yield database + finally: + database.dispose() + + +@pytest.fixture +def empty_database(_external_test_schema: Optional[URL]) -> Iterator[Database]: + """A database without tables: SQLite in memory, or this worker's schema on the external server.""" + if _external_test_schema is None: + database = Database.open_sqlite(None, InvokeAILogger.get_logger("test_database")) + else: + database = _open(_external_test_schema) + try: + if _external_test_schema is not None: + # Also before the test: a teardown that failed must not hand its tables to the next test. + _drop_all_tables(database) + yield database + finally: + try: + if _external_test_schema is not None: + _drop_all_tables(database) + finally: + database.dispose() + + +@pytest.fixture +def database( + _external_application_schema: Optional[_ApplicationSchema], request: pytest.FixtureRequest +) -> Iterator[Database]: + """A database at the newest schema, holding the rows a new install starts with (the system account, the + JWT secret, the lock rows): a copy in memory of a SQLite database migrated once per run, or this worker's + schema on the external server, migrated once per run and reset to those rows before each test. + + On a server the ids a table generates go on counting from test to test: assert on the ids a test got back, + not on literal ones. + """ + logger = InvokeAILogger.get_logger("test_database") + if _external_application_schema is None: + migrated: Database = request.getfixturevalue("_migrated_sqlite") + database = Database.open_sqlite(None, logger) + try: + migrated.sqlite.conn.backup(database.sqlite.conn) + yield database + finally: + database.dispose() + return + # The engine, and its pool of connections, serves every test of the run. + database = Database(_external_application_schema.engine, logger) + with database.begin(write=True) as conn: + for table in reversed(metadata.sorted_tables): + conn.execute(delete(table)) + for table in metadata.sorted_tables: + if rows := _external_application_schema.rows.get(table.name): + conn.execute(insert(table), rows) + yield database + + +def _drop_all_tables(database: Database) -> None: + with database.engine.connect() as conn: + tables: list[str] = list( + conn.exec_driver_sql( + "SELECT table_name FROM information_schema.tables WHERE table_schema = DATABASE()" + ).scalars() + ) + conn.exec_driver_sql("SET FOREIGN_KEY_CHECKS = 0") + for table in tables: + conn.exec_driver_sql(f"DROP TABLE `{table}`") + conn.exec_driver_sql("SET FOREIGN_KEY_CHECKS = 1") + conn.commit() + + +@contextmanager +def capture_statements(database: Database) -> Iterator[list[tuple[str, Any]]]: + """The statements the database runs inside the block, each with its parameters as the driver gets them.""" + captured: list[tuple[str, Any]] = [] + + def record(conn: Any, cursor: Any, statement: str, parameters: Any, context: Any, executemany: bool) -> None: + captured.append((statement, parameters)) + + event.listen(database.engine, "before_cursor_execute", record) + try: + yield captured + finally: + event.remove(database.engine, "before_cursor_execute", record) + + +def explain_query_plan(database: Database, statement: str, parameters: Any) -> list[str]: + """The details of SQLite's plan for a captured statement.""" + with database.begin(write=False) as conn: + return [row[3] for row in conn.exec_driver_sql(f"EXPLAIN QUERY PLAN {statement}", parameters).all()] diff --git a/tests/fixtures/database_probe.py b/tests/fixtures/database_probe.py new file mode 100644 index 00000000000..5f5d0486aff --- /dev/null +++ b/tests/fixtures/database_probe.py @@ -0,0 +1,104 @@ +"""Probe tables and a probe query module for exercising the database layer without the application schema.""" + +from collections.abc import Callable +from typing import Optional + +from sqlalchemy import ( + CheckConstraint, + Column, + Connection, + ForeignKey, + Integer, + MetaData, + String, + Table, + delete, + func, + insert, + select, + update, +) + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.queries import Queries +from invokeai.app.services.shared.database.queries.base import QueryModule, QueryScope, read, write + +probe_metadata = MetaData() + +probe_items = Table( + "probe_items", + probe_metadata, + Column("id", Integer, primary_key=True, autoincrement=False), + Column("name", String(32), nullable=False, unique=True), + Column("size", Integer, nullable=False), + # Table-level: MariaDB rejects a named CHECK inside a column definition. + CheckConstraint("size >= 0", name="ck_probe_items_size"), + mysql_engine="InnoDB", +) + +probe_links = Table( + "probe_links", + probe_metadata, + Column("id", Integer, primary_key=True, autoincrement=False), + Column("item_id", Integer, ForeignKey("probe_items.id", ondelete="CASCADE"), nullable=False), + mysql_engine="InnoDB", +) + + +class ProbeItems(QueryModule): + @write + def add(self, conn: Connection, item_id: int, name: Optional[str], size: int = 0) -> None: + conn.execute(insert(probe_items).values(id=item_id, name=name, size=size)) + + @write + def link(self, conn: Connection, link_id: int, item_id: int) -> None: + conn.execute(insert(probe_links).values(id=link_id, item_id=item_id)) + + @write + def resize(self, conn: Connection, item_id: int, size: int) -> None: + conn.execute(update(probe_items).where(probe_items.c.id == item_id).values(size=size)) + + @write + def resize_both(self, conn: Connection, first: int, second: int, between: Callable[[], None]) -> None: + """Updates two rows in the given order, calling `between` in between (to interleave two writers).""" + conn.execute(update(probe_items).where(probe_items.c.id == first).values(size=1)) + between() + conn.execute(update(probe_items).where(probe_items.c.id == second).values(size=1)) + + @read + def read_then_write(self, conn: Connection, between: Callable[[], None]) -> None: + """Reads, calls `between`, then writes, all in a read transaction -- which no real query method does. + + On SQLite a read transaction holds a snapshot; if another connection commits in `between`, the + write cannot proceed on that stale snapshot and fails at once: a lost race, as a deadlock is. + """ + conn.execute(select(func.count()).select_from(probe_items)).scalar_one() + between() + conn.execute(update(probe_items).values(size=0)) + + @write + def remove(self, conn: Connection, item_id: int) -> None: + conn.execute(delete(probe_items).where(probe_items.c.id == item_id)) + + @read + def names(self, conn: Connection) -> list[str]: + return list(conn.execute(select(probe_items.c.name).order_by(probe_items.c.id)).scalars()) + + @read + def link_ids(self, conn: Connection) -> list[int]: + return list(conn.execute(select(probe_links.c.id).order_by(probe_links.c.id)).scalars()) + + @read + def count(self, conn: Connection) -> int: + return conn.execute(select(func.count()).select_from(probe_items)).scalar_one() + + +class ProbeQueries(Queries): + def __init__(self, database: Database, scope: Optional[QueryScope] = None) -> None: + super().__init__(database, scope) + self.items = ProbeItems(self._scope) + + +def create_probe_tables(database: Database) -> None: + with database.begin(write=True) as conn: + probe_metadata.create_all(conn) diff --git a/tests/fixtures/races.py b/tests/fixtures/races.py new file mode 100644 index 00000000000..e29bac2f63d --- /dev/null +++ b/tests/fixtures/races.py @@ -0,0 +1,122 @@ +"""Transactions racing each other, for tests of what a write does while another one is in flight. + +On MySQL and MariaDB transactions run side by side: a test holds one transaction open at a chosen point and starts +another write, which must wait for it (asserted after half a second) and then sees what it committed. On SQLite a +transaction in flight blocks every other one, so these tests pass there trivially. +""" + +import threading +from collections.abc import Callable + +import pytest + +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.errors import ConflictError +from invokeai.app.services.shared.database.queries import Queries + + +@pytest.fixture +def lost_races(monkeypatch: pytest.MonkeyPatch) -> list[ConflictError]: + """Every deadlock a retried transaction of the test loses, including those it then wins on another attempt. (A + transaction in flight is not retried: it raises what it loses.)""" + lost: list[ConflictError] = [] + retry_conflicts = Database.retry_conflicts + + def recording(self: Database, transaction: Callable[[], object]) -> object: + def recorded() -> object: + try: + return transaction() + except ConflictError as error: + lost.append(error) + raise + + return retry_conflicts(self, recorded) + + monkeypatch.setattr(Database, "retry_conflicts", recording) + return lost + + +def waiting(change: Callable[[], object]) -> Callable[[], list[BaseException]]: + """Starts `change` on another thread and asserts that half a second later it still waits, for a lock the + caller's transaction holds. Returns a function that waits for `change` to end and returns what it raised.""" + errors: list[BaseException] = [] + + def run() -> None: + try: + change() + except BaseException as e: # noqa: BLE001 - returned to the caller + errors.append(e) + + thread = threading.Thread(target=run) + thread.start() + thread.join(timeout=0.5) + assert thread.is_alive(), f"the change did not wait for the transaction in flight; it raised {errors!r}" + + def ended() -> list[BaseException]: + thread.join(timeout=30) + assert not thread.is_alive(), "the change still waits" + return errors + + return ended + + +def while_in_flight( + database: Database, work: Callable[[Queries], object], change: Callable[[], object] +) -> list[BaseException]: + """Runs `change` while another transaction has done `work` and has not committed; that transaction then commits, + and `change` goes on. Returns what `change` raised.""" + done = threading.Event() + proceed = threading.Event() + failures: list[BaseException] = [] + + def in_flight() -> None: + try: + with database.queries.transaction() as q: + work(q) + done.set() + proceed.wait(timeout=30) + except BaseException as e: # noqa: BLE001 - raised again below + failures.append(e) + done.set() + + holder = threading.Thread(target=in_flight) + holder.start() + try: + assert done.wait(timeout=10), "the transaction in flight did not get to its work" + if failures: + raise failures[0] + ended = waiting(change) + finally: + proceed.set() + holder.join(timeout=10) + if failures: + raise failures[0] + return ended() + + +def when_called( + monkeypatch: pytest.MonkeyPatch, cls: type, method: str, change: Callable[[], object] +) -> Callable[[], list[BaseException]]: + """Makes the first call of `cls.method` start `change` when it returns, still in its transaction, and go on once + `change` waits for that transaction. Call it in a unit of `queries.run()`; the returned function, called after the + unit, returns what `change` raised.""" + real = getattr(cls, method) + first = True + started: list[Callable[[], list[BaseException]]] = [] + + def then_change(self: object, *args: object, **kwargs: object) -> object: + nonlocal first + result = real(self, *args, **kwargs) + if first: + # Cleared before the change starts: a change that calls the method itself must not start another one. + first = False + started.append(waiting(change)) + return result + + monkeypatch.setattr(cls, method, then_change) + + def ended() -> list[BaseException]: + assert started, f"{cls.__name__}.{method} was not called" + return started[0]() + + return ended diff --git a/tests/fixtures/sqlite_database.py b/tests/fixtures/sqlite_database.py index 7b2dd029a95..99f27fee25e 100644 --- a/tests/fixtures/sqlite_database.py +++ b/tests/fixtures/sqlite_database.py @@ -1,13 +1,39 @@ +import sqlite3 +from collections.abc import Iterator +from contextlib import contextmanager from logging import Logger +from typing import TYPE_CHECKING from unittest import mock from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase -from invokeai.app.services.shared.sqlite.sqlite_util import init_db +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.database.startup import init_database +if TYPE_CHECKING: + from invokeai.app.services.session_queue.session_queue_default import SessionQueue -def create_mock_sqlite_database(config: InvokeAIAppConfig, logger: Logger) -> SqliteDatabase: - image_files = mock.Mock(spec=ImageFileStorageBase) - db = init_db(config=config, logger=logger, image_files=image_files) - return db + +def create_mock_sqlite_database(config: InvokeAIAppConfig, logger: Logger) -> Database: + """A SQLite database at the newest schema, as the app opens it, with image files that are not there.""" + return init_database(config=config, logger=logger, image_files=mock.Mock(spec=ImageFileStorageBase)) + + +@contextmanager +def sqlite_cursor(database: Database) -> Iterator[sqlite3.Cursor]: + """A raw cursor in a write transaction of a SQLite database, for tests that set up or inspect rows with SQL of + their own; its rows are `sqlite3.Row`s. The transaction commits when the block exits normally.""" + with database.begin(write=True) as conn: + cursor = conn.connection.driver_connection.cursor() + cursor.row_factory = sqlite3.Row + try: + yield cursor + finally: + cursor.close() + + +@contextmanager +def sqlite_cursor_of(queue: "SessionQueue") -> Iterator[sqlite3.Cursor]: + """`sqlite_cursor` on the database a session queue uses.""" + with sqlite_cursor(queue._queries._database) as cursor: + yield cursor diff --git a/tests/frontend/install/test_import_images.py b/tests/frontend/install/test_import_images.py new file mode 100644 index 00000000000..8ef205ff88f --- /dev/null +++ b/tests/frontend/install/test_import_images.py @@ -0,0 +1,67 @@ +"""The database work of the image import script, on a database file of its own.""" + +import datetime +import sqlite3 +from contextlib import closing +from pathlib import Path + +from invokeai.app.services.board_image_records.board_image_records_default import BoardImageRecordStorage +from invokeai.app.services.board_records.board_records_common import BoardChanges +from invokeai.app.services.board_records.board_records_default import BoardRecordStorage +from invokeai.app.services.image_records.image_records_default import ImageRecordStorage +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.users.users_common import UserCreateRequest +from invokeai.app.services.users.users_default import UserService +from invokeai.backend.util.logging import InvokeAILogger +from invokeai.frontend.install.import_images import DatabaseMapper +from tests.fixtures.database import migrate_to_newest + + +def test_imports_images_onto_boards_found_or_created_by_name(tmp_path: Path) -> None: + path = tmp_path / "databases" / "invokeai.db" + database = Database.open_sqlite(path, InvokeAILogger.get_logger("test")) + migrate_to_newest(database, tmp_path) + database.dispose() + mapper = DatabaseMapper(str(path), str(tmp_path / "databases" / "backup")) + mapper.connect() + try: + assert mapper.get_board_names() == [] + board_id = mapper.get_board_id_with_create("Import") + assert mapper.get_board_id_with_create("IMPORT") == board_id + assert mapper.get_board_names() == ["Import"] + + # Any board of the name serves, whoever owns it, archived or not. + assert mapper.database is not None + other = UserService(mapper.database).create( + UserCreateRequest(email="other@test.com", password="OtherPass123", display_name="other") + ) + boards = BoardRecordStorage(mapper.database) + theirs = boards.save("Theirs", other.user_id).board_id + archived = boards.save("Archived", "system").board_id + boards.update(archived, BoardChanges(archived=True)) + assert mapper.get_board_id_with_create("THEIRS") == theirs + assert mapper.get_board_id_with_create("archived") == archived + + assert not mapper.does_image_exist("old.png") + modified = datetime.datetime(2001, 2, 3, 4, 5, 6, 789000) + mapper.add_new_image_to_database("old.png", 64, 32, '{"seed": 1}', modified) + mapper.add_image_to_board("old.png", board_id) + assert mapper.does_image_exist("old.png") + + mapper.backup("20010203T040506Z") + finally: + mapper.disconnect() + + database = Database.open_sqlite(path, InvokeAILogger.get_logger("test")) + try: + record = ImageRecordStorage(database).get("old.png") + # The image keeps the time its file was made, so it sorts among the images of that time. + assert str(record.created_at).startswith("2001-02-03 04:05:06.789") + assert (record.width, record.height) == (64, 32) + assert ImageRecordStorage(database).get_metadata("old.png") is not None + assert BoardImageRecordStorage(database).get_board_for_image("old.png") == board_id + finally: + database.dispose() + backup = tmp_path / "databases" / "backup" / "backup-20010203T040506Z-invokeai.db" + with closing(sqlite3.connect(backup)) as copy: + assert copy.execute("SELECT image_name FROM images").fetchall() == [("old.png",)] diff --git a/tests/test_config.py b/tests/test_config.py index e372baf6e33..43fb76ab101 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -12,6 +12,7 @@ ensure_fonts_dir, get_config, load_and_migrate_config, + load_config_from_root, ) from invokeai.app.services.shared.graph import Graph from invokeai.frontend.cli.arg_parser import InvokeAIArgs @@ -101,6 +102,30 @@ def test_db_synchronous_defaults_to_full_and_loads_from_yaml(tmp_path: Path, pat assert load_and_migrate_config(temp_config_file).db_synchronous == "normal" +def test_load_config_from_root_reads_that_roots_file_and_writes_nothing(tmp_path: Path, patch_rootdir: None) -> None: + root = tmp_path / "install" + root.mkdir() + (root / "invokeai.yaml").write_text( + 'schema_version: "4.0.3"\ndb_dir: elsewhere\noutputs_dir: pictures\ndb_synchronous: normal\n' + ) + + config = load_config_from_root(root) + + assert config.db_path == (root / "elsewhere" / "invokeai.db").resolve() + assert config.outputs_path == (root / "pictures").resolve() + assert config.db_synchronous == "normal" + assert [path.name for path in root.iterdir()] == ["invokeai.yaml"] + + +def test_load_config_from_root_lets_the_environment_win_over_the_file( + tmp_path: Path, patch_rootdir: None, monkeypatch: pytest.MonkeyPatch +) -> None: + (tmp_path / "invokeai.yaml").write_text('schema_version: "4.0.3"\ndb_synchronous: normal\n') + monkeypatch.setenv("INVOKEAI_DB_SYNCHRONOUS", "full") + + assert load_config_from_root(tmp_path).db_synchronous == "full" + + def test_read_config_from_file(tmp_path: Path, patch_rootdir: None): """Test reading configuration from a file.""" temp_config_file = tmp_path / "temp_invokeai.yaml" diff --git a/tests/test_sqlite_migrator.py b/tests/test_sqlite_migrator.py index a737689ce0c..48f40c4a384 100644 --- a/tests/test_sqlite_migrator.py +++ b/tests/test_sqlite_migrator.py @@ -1,5 +1,6 @@ import sqlite3 from contextlib import closing +from datetime import datetime from logging import Logger from pathlib import Path from tempfile import TemporaryDirectory @@ -7,7 +8,8 @@ import pytest from pydantic import ValidationError -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.sqlite_migrator import sqlite_migrator_impl from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import ( MigrateCallback, Migration, @@ -16,7 +18,7 @@ MigrationVersionError, ) from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import ( - SqliteMigrator, + Migrator, ) @@ -36,9 +38,9 @@ def memory_db_cursor(memory_db_conn: sqlite3.Connection) -> sqlite3.Cursor: @pytest.fixture -def migrator(logger: Logger) -> SqliteMigrator: - db = SqliteDatabase(db_path=None, logger=logger, verbose=False) - return SqliteMigrator(db=db) +def migrator(logger: Logger) -> Migrator: + db = Database.open_sqlite(None, logger) + return Migrator(db) @pytest.fixture @@ -134,15 +136,13 @@ def test_explicit_migration_id_and_dependency_are_preserved(no_op_migrate_callba assert migration.to_version is None -def test_migration_set_add_migration(migrator: SqliteMigrator, migration_no_op: Migration) -> None: +def test_migration_set_add_migration(migrator: Migrator, migration_no_op: Migration) -> None: migration = migration_no_op migrator._migration_set.register(migration) assert migration in migrator._migration_set._migrations -def test_migration_set_may_not_register_dupes( - migrator: SqliteMigrator, no_op_migrate_callback: MigrateCallback -) -> None: +def test_migration_set_may_not_register_dupes(migrator: Migrator, no_op_migrate_callback: MigrateCallback) -> None: migrate_0_to_1_a = Migration(from_version=0, to_version=1, callback=no_op_migrate_callback) migrate_0_to_1_b = Migration(from_version=0, to_version=1, callback=no_op_migrate_callback) migrator._migration_set.register(migrate_0_to_1_a) @@ -275,28 +275,28 @@ def test_migration_runs(memory_db_cursor: sqlite3.Cursor, migrate_callback_creat assert memory_db_cursor.fetchone() is not None -def test_migrator_registers_migration(migrator: SqliteMigrator, migration_no_op: Migration) -> None: +def test_migrator_registers_migration(migrator: Migrator, migration_no_op: Migration) -> None: migration = migration_no_op migrator.register_migration(migration) assert migration in migrator._migration_set._migrations -def test_migrator_creates_migrations_table(migrator: SqliteMigrator) -> None: - cursor = migrator._db._conn.cursor() +def test_migrator_creates_migrations_table(migrator: Migrator) -> None: + cursor = migrator._database.sqlite.conn.cursor() migrator._create_migrations_table(cursor) cursor.execute("SELECT * FROM sqlite_master WHERE type='table' AND name='migrations';") assert cursor.fetchone() is not None -def test_migrator_creates_applied_migrations_table(migrator: SqliteMigrator) -> None: - cursor = migrator._db._conn.cursor() +def test_migrator_creates_applied_migrations_table(migrator: Migrator) -> None: + cursor = migrator._database.sqlite.conn.cursor() migrator._create_applied_migrations_table(cursor) cursor.execute("SELECT * FROM sqlite_master WHERE type='table' AND name='applied_migrations';") assert cursor.fetchone() is not None -def test_migrator_migration_sets_version(migrator: SqliteMigrator, migration_no_op: Migration) -> None: - cursor = migrator._db._conn.cursor() +def test_migrator_migration_sets_version(migrator: Migrator, migration_no_op: Migration) -> None: + cursor = migrator._database.sqlite.conn.cursor() migrator._create_migrations_table(cursor) migrator.register_migration(migration_no_op) migrator.run_migrations() @@ -304,8 +304,8 @@ def test_migrator_migration_sets_version(migrator: SqliteMigrator, migration_no_ assert cursor.fetchone()[0] == 1 -def test_migrator_gets_current_version(migrator: SqliteMigrator, migration_no_op: Migration) -> None: - cursor = migrator._db._conn.cursor() +def test_migrator_gets_current_version(migrator: Migrator, migration_no_op: Migration) -> None: + cursor = migrator._database.sqlite.conn.cursor() assert migrator._get_current_version(cursor) == 0 migrator._create_migrations_table(cursor) assert migrator._get_current_version(cursor) == 0 @@ -314,8 +314,8 @@ def test_migrator_gets_current_version(migrator: SqliteMigrator, migration_no_op assert migrator._get_current_version(cursor) == 1 -def test_migrator_runs_single_migration(migrator: SqliteMigrator, migration_create_test_table: Migration) -> None: - cursor = migrator._db._conn.cursor() +def test_migrator_runs_single_migration(migrator: Migrator, migration_create_test_table: Migration) -> None: + cursor = migrator._database.sqlite.conn.cursor() migrator._create_migrations_table(cursor) migrator._run_migration(migration_create_test_table) assert migrator._get_current_version(cursor) == 1 @@ -323,8 +323,8 @@ def test_migrator_runs_single_migration(migrator: SqliteMigrator, migration_crea assert cursor.fetchone() is not None -def test_migrator_runs_all_migrations_in_memory(migrator: SqliteMigrator) -> None: - cursor = migrator._db._conn.cursor() +def test_migrator_runs_all_migrations_in_memory(migrator: Migrator) -> None: + cursor = migrator._database.sqlite.conn.cursor() migrations = [Migration(from_version=i, to_version=i + 1, callback=create_migrate(i)) for i in range(0, 3)] for migration in migrations: migrator.register_migration(migration) @@ -333,9 +333,9 @@ def test_migrator_runs_all_migrations_in_memory(migrator: SqliteMigrator) -> Non def test_migrator_bootstraps_applied_migrations_from_legacy_versions( - migrator: SqliteMigrator, no_op_migrate_callback: MigrateCallback + migrator: Migrator, no_op_migrate_callback: MigrateCallback ) -> None: - cursor = migrator._db._conn.cursor() + cursor = migrator._database.sqlite.conn.cursor() migrator._create_migrations_table(cursor) cursor.execute("INSERT INTO migrations (version) VALUES (1);") cursor.execute("INSERT INTO migrations (version) VALUES (2);") @@ -356,8 +356,8 @@ def test_migrator_backs_up_file_db_before_metadata_only_bootstrap( ) -> None: with TemporaryDirectory() as tempdir: original_db_path = Path(tempdir) / "invokeai.db" - db = SqliteDatabase(db_path=original_db_path, logger=logger, verbose=False) - cursor = db._conn.cursor() + db = Database.open_sqlite(original_db_path, logger) + cursor = db.sqlite.conn.cursor() cursor.execute("CREATE TABLE test (id INTEGER PRIMARY KEY);") cursor.execute( """--sql @@ -369,14 +369,14 @@ def test_migrator_backs_up_file_db_before_metadata_only_bootstrap( ) cursor.execute("INSERT INTO migrations (version) VALUES (0);") cursor.execute("INSERT INTO migrations (version) VALUES (1);") - db._conn.commit() + db.sqlite.conn.commit() - migrator = SqliteMigrator(db=db) + migrator = Migrator(db) migrator.register_migration(Migration(from_version=0, to_version=1, callback=no_op_migrate_callback)) assert migrator.run_migrations() is False - db._conn.close() + db.sqlite.conn.close() assert migrator._backup_path is not None with closing(sqlite3.connect(migrator._backup_path)) as backup_db_conn: backup_db_cursor = backup_db_conn.cursor() @@ -386,8 +386,8 @@ def test_migrator_backs_up_file_db_before_metadata_only_bootstrap( assert backup_db_cursor.fetchone() is None -def test_migrator_runs_branching_graph_migrations(migrator: SqliteMigrator) -> None: - cursor = migrator._db._conn.cursor() +def test_migrator_runs_branching_graph_migrations(migrator: Migrator) -> None: + cursor = migrator._database.sqlite.conn.cursor() executed: list[str] = [] def create_migration(migration_id: str, depends_on: str | None) -> Migration: @@ -413,9 +413,9 @@ def migrate(cursor: sqlite3.Cursor) -> None: def test_migrator_rejects_unknown_applied_migration( - migrator: SqliteMigrator, no_op_migrate_callback: MigrateCallback + migrator: Migrator, no_op_migrate_callback: MigrateCallback ) -> None: - cursor = migrator._db._conn.cursor() + cursor = migrator._database.sqlite.conn.cursor() migrator._create_migrations_table(cursor) migrator._create_applied_migrations_table(cursor) cursor.execute("INSERT INTO applied_migrations (migration_id) VALUES ('future_migration');") @@ -427,9 +427,9 @@ def test_migrator_rejects_unknown_applied_migration( def test_migrator_rejects_unknown_applied_migration_before_creating_legacy_table( - migrator: SqliteMigrator, no_op_migrate_callback: MigrateCallback + migrator: Migrator, no_op_migrate_callback: MigrateCallback ) -> None: - cursor = migrator._db._conn.cursor() + cursor = migrator._database.sqlite.conn.cursor() migrator._create_applied_migrations_table(cursor) cursor.execute("INSERT INTO applied_migrations (migration_id) VALUES ('future_migration');") cursor.connection.commit() @@ -443,9 +443,9 @@ def test_migrator_rejects_unknown_applied_migration_before_creating_legacy_table def test_migrator_rejects_inconsistent_applied_legacy_version( - migrator: SqliteMigrator, no_op_migrate_callback: MigrateCallback + migrator: Migrator, no_op_migrate_callback: MigrateCallback ) -> None: - cursor = migrator._db._conn.cursor() + cursor = migrator._database.sqlite.conn.cursor() migrator._create_migrations_table(cursor) migrator._create_applied_migrations_table(cursor) cursor.execute("INSERT INTO migrations (version) VALUES (1);") @@ -468,9 +468,9 @@ def migration_callback(cursor: sqlite3.Cursor) -> None: def test_migrator_rejects_applied_legacy_migration_missing_legacy_row( - migrator: SqliteMigrator, no_op_migrate_callback: MigrateCallback + migrator: Migrator, no_op_migrate_callback: MigrateCallback ) -> None: - cursor = migrator._db._conn.cursor() + cursor = migrator._database.sqlite.conn.cursor() migrator._create_migrations_table(cursor) migrator._create_applied_migrations_table(cursor) cursor.execute("INSERT INTO migrations (version) VALUES (1);") @@ -494,10 +494,8 @@ def graph_migration_callback(cursor: sqlite3.Cursor) -> None: assert callback_ran is False -def test_migrator_rejects_unknown_legacy_version( - migrator: SqliteMigrator, no_op_migrate_callback: MigrateCallback -) -> None: - cursor = migrator._db._conn.cursor() +def test_migrator_rejects_unknown_legacy_version(migrator: Migrator, no_op_migrate_callback: MigrateCallback) -> None: + cursor = migrator._database.sqlite.conn.cursor() migrator._create_migrations_table(cursor) cursor.execute("INSERT INTO migrations (version) VALUES (1);") cursor.execute("INSERT INTO migrations (version) VALUES (2);") @@ -509,9 +507,9 @@ def test_migrator_rejects_unknown_legacy_version( def test_migrator_rejects_unknown_legacy_version_before_creating_applied_table( - migrator: SqliteMigrator, no_op_migrate_callback: MigrateCallback + migrator: Migrator, no_op_migrate_callback: MigrateCallback ) -> None: - cursor = migrator._db._conn.cursor() + cursor = migrator._database.sqlite.conn.cursor() migrator._create_migrations_table(cursor) cursor.execute("INSERT INTO migrations (version) VALUES (1);") cursor.execute("INSERT INTO migrations (version) VALUES (2);") @@ -528,35 +526,35 @@ def test_migrator_rejects_unknown_legacy_version_before_creating_applied_table( def test_migrator_runs_all_migrations_file(logger: Logger) -> None: with TemporaryDirectory() as tempdir: original_db_path = Path(tempdir) / "invokeai.db" - db = SqliteDatabase(db_path=original_db_path, logger=logger, verbose=False) - migrator = SqliteMigrator(db=db) + db = Database.open_sqlite(original_db_path, logger) + migrator = Migrator(db) migrations = [Migration(from_version=i, to_version=i + 1, callback=create_migrate(i)) for i in range(0, 3)] for migration in migrations: migrator.register_migration(migration) migrator.run_migrations() with closing(sqlite3.connect(original_db_path)) as original_db_conn: original_db_cursor = original_db_conn.cursor() - assert SqliteMigrator._get_current_version(original_db_cursor) == 3 + assert Migrator._get_current_version(original_db_cursor) == 3 # Must manually close else we get an error on Windows - db._conn.close() + db.sqlite.conn.close() def test_migrator_backs_up_db(logger: Logger) -> None: with TemporaryDirectory() as tempdir: original_db_path = Path(tempdir) / "invokeai.db" - db = SqliteDatabase(db_path=original_db_path, logger=logger, verbose=False) + db = Database.open_sqlite(original_db_path, logger) # Write some data to the db to test for successful backup - temp_cursor = db._conn.cursor() + temp_cursor = db.sqlite.conn.cursor() temp_cursor.execute("CREATE TABLE test (id INTEGER PRIMARY KEY);") - db._conn.commit() + db.sqlite.conn.commit() # Set up the migrator - migrator = SqliteMigrator(db=db) + migrator = Migrator(db) migrations = [Migration(from_version=i, to_version=i + 1, callback=create_migrate(i)) for i in range(0, 3)] for migration in migrations: migrator.register_migration(migration) migrator.run_migrations() # Must manually close else we get an error on Windows - db._conn.close() + db.sqlite.conn.close() assert original_db_path.exists() # We should have a backup file when we migrated a file db assert migrator._backup_path @@ -568,9 +566,9 @@ def test_migrator_backs_up_db(logger: Logger) -> None: def test_migrator_makes_no_changes_on_failed_migration( - migrator: SqliteMigrator, migration_no_op: Migration, failing_migrate_callback: MigrateCallback + migrator: Migrator, migration_no_op: Migration, failing_migrate_callback: MigrateCallback ) -> None: - cursor = migrator._db._conn.cursor() + cursor = migrator._database.sqlite.conn.cursor() migrator.register_migration(migration_no_op) migrator.run_migrations() assert migrator._get_current_version(cursor) == 1 @@ -583,7 +581,7 @@ def test_migrator_makes_no_changes_on_failed_migration( def test_migrator_rolls_back_ddl_a_failed_migration_already_issued( - migrator: SqliteMigrator, migration_no_op: Migration + migrator: Migrator, migration_no_op: Migration ) -> None: """A migration that only ever issues DDL must still be all-or-nothing. @@ -600,7 +598,7 @@ def migrate(cursor: sqlite3.Cursor, **kwargs) -> None: cursor.execute("CREATE TABLE ddl_only_replacement (id INTEGER PRIMARY KEY);") raise Exception("Bad migration") - cursor = migrator._db._conn.cursor() + cursor = migrator._database.sqlite.conn.cursor() migrator.register_migration(migration_no_op) migrator.run_migrations() migrator.register_migration(Migration(from_version=1, to_version=2, callback=migrate)) @@ -616,7 +614,7 @@ def migrate(cursor: sqlite3.Cursor, **kwargs) -> None: def test_migrator_rolls_back_a_failed_migration_that_mixed_ddl_and_dml( - migrator: SqliteMigrator, migration_no_op: Migration + migrator: Migrator, migration_no_op: Migration ) -> None: """The same guarantee where the DDL precedes the first row written.""" @@ -625,7 +623,7 @@ def migrate(cursor: sqlite3.Cursor, **kwargs) -> None: cursor.execute("INSERT INTO mixed (id) VALUES (1);") raise Exception("Bad migration") - cursor = migrator._db._conn.cursor() + cursor = migrator._database.sqlite.conn.cursor() migrator.register_migration(migration_no_op) migrator.run_migrations() migrator.register_migration(Migration(from_version=1, to_version=2, callback=migrate)) @@ -637,8 +635,8 @@ def migrate(cursor: sqlite3.Cursor, **kwargs) -> None: assert cursor.fetchone() is None -def test_idempotent_migrations(migrator: SqliteMigrator, migration_create_test_table: Migration) -> None: - cursor = migrator._db._conn.cursor() +def test_idempotent_migrations(migrator: Migrator, migration_create_test_table: Migration) -> None: + cursor = migrator._database.sqlite.conn.cursor() migrator.register_migration(migration_create_test_table) migrator.run_migrations() # not throwing is sufficient @@ -650,20 +648,20 @@ def test_migration_27_creates_users_table(logger: Logger) -> None: """Test that migration 27 creates the users table and related tables.""" from invokeai.app.services.shared.sqlite_migrator.migrations.migration_27 import Migration27Callback - db = SqliteDatabase(db_path=None, logger=logger, verbose=False) - cursor = db._conn.cursor() + db = Database.open_sqlite(None, logger) + cursor = db.sqlite.conn.cursor() # Create minimal tables that migration 27 expects to exist cursor.execute("CREATE TABLE IF NOT EXISTS boards (board_id TEXT PRIMARY KEY);") cursor.execute("CREATE TABLE IF NOT EXISTS images (image_name TEXT PRIMARY KEY);") cursor.execute("CREATE TABLE IF NOT EXISTS workflows (workflow_id TEXT PRIMARY KEY);") cursor.execute("CREATE TABLE IF NOT EXISTS session_queue (item_id INTEGER PRIMARY KEY);") - db._conn.commit() + db.sqlite.conn.commit() # Run migration callback directly (not through migrator to avoid chain validation) migration_callback = Migration27Callback() migration_callback(cursor) - db._conn.commit() + db.sqlite.conn.commit() # Verify users table exists cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='users';") @@ -720,7 +718,7 @@ def test_migration_27_creates_users_table(logger: Logger) -> None: assert jwt_row is not None assert len(jwt_row[0]) == 64 # 32 bytes = 64 hex characters - db._conn.close() + db.sqlite.conn.close() def test_migration_27_with_existing_client_state_data(logger: Logger) -> None: @@ -730,8 +728,8 @@ def test_migration_27_with_existing_client_state_data(logger: Logger) -> None: from invokeai.app.services.shared.sqlite_migrator.migrations.migration_21 import Migration21Callback from invokeai.app.services.shared.sqlite_migrator.migrations.migration_27 import Migration27Callback - db = SqliteDatabase(db_path=None, logger=logger, verbose=False) - cursor = db._conn.cursor() + db = Database.open_sqlite(None, logger) + cursor = db.sqlite.conn.cursor() # Run migration 21 to create old-style client_state with data column Migration21Callback()(cursor) @@ -740,7 +738,7 @@ def test_migration_27_with_existing_client_state_data(logger: Logger) -> None: "INSERT INTO client_state (id, data) VALUES (1, ?);", (json.dumps({"galleryView": "images", "lastBoardId": "board123"}),), ) - db._conn.commit() + db.sqlite.conn.commit() # Run migration 27 pre-reqs cursor.execute("CREATE TABLE IF NOT EXISTS boards (board_id TEXT PRIMARY KEY);") @@ -748,11 +746,11 @@ def test_migration_27_with_existing_client_state_data(logger: Logger) -> None: cursor.execute("CREATE TABLE IF NOT EXISTS workflows (workflow_id TEXT PRIMARY KEY);") 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() + db.sqlite.conn.commit() # Run migration 27 Migration27Callback()(cursor) - db._conn.commit() + db.sqlite.conn.commit() # Verify new client_state schema cursor.execute("PRAGMA table_info(client_state);") @@ -770,15 +768,15 @@ def test_migration_27_with_existing_client_state_data(logger: Logger) -> None: assert ("system", "galleryView", "images") in rows assert ("system", "lastBoardId", "board123") in rows - db._conn.close() + db.sqlite.conn.close() 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 - db = SqliteDatabase(db_path=None, logger=logger, verbose=False) - cursor = db._conn.cursor() + db = Database.open_sqlite(None, logger) + cursor = db.sqlite.conn.cursor() # Create old client_state WITHOUT data column (simulating an older migration 21) cursor.execute( @@ -789,7 +787,7 @@ def test_migration_27_without_client_state_data_column(logger: Logger) -> None: ); """ ) - db._conn.commit() + db.sqlite.conn.commit() # Run migration 27 pre-reqs cursor.execute("CREATE TABLE IF NOT EXISTS boards (board_id TEXT PRIMARY KEY);") @@ -797,11 +795,11 @@ def test_migration_27_without_client_state_data_column(logger: Logger) -> None: cursor.execute("CREATE TABLE IF NOT EXISTS workflows (workflow_id TEXT PRIMARY KEY);") 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() + db.sqlite.conn.commit() # Run migration 27 - should not raise even without data column Migration27Callback()(cursor) - db._conn.commit() + db.sqlite.conn.commit() # Verify new client_state schema cursor.execute("PRAGMA table_info(client_state);") @@ -815,4 +813,30 @@ def test_migration_27_without_client_state_data_column(logger: Logger) -> None: cursor.execute("SELECT COUNT(*) FROM client_state;") assert cursor.fetchone()[0] == 0 - db._conn.close() + db.sqlite.conn.close() + + +def test_a_backup_of_the_same_second_is_kept_and_the_next_one_numbered( + tmp_path: Path, logger: Logger, monkeypatch: pytest.MonkeyPatch +) -> None: + # Migrating a database again right away, as on a restart after a failed migration, must not fail on, or + # replace, the backup the first run made. + class SameSecond(datetime): + @classmethod + def now(cls, tz: object = None) -> "SameSecond": + return cls(2026, 10, 7, 12, 0, 0) + + monkeypatch.setattr(sqlite_migrator_impl, "datetime", SameSecond) + db = Database.open_sqlite(tmp_path / "invokeai.db", logger) + db.sqlite.conn.execute("CREATE TABLE test (id INTEGER PRIMARY KEY);") + db.sqlite.conn.commit() + + first = Migrator(db) + first._backup_db() + second = Migrator(db) + second._backup_db() + + assert first._backup_path == tmp_path / "invokeai_backup_20261007-120000.db" + assert second._backup_path == tmp_path / "invokeai_backup_20261007-120000-1.db" + assert first._backup_path.is_file() and second._backup_path.is_file() + db.dispose() diff --git a/uv.lock b/uv.lock index 11102acdf0d..87e4a1256c7 100644 --- a/uv.lock +++ b/uv.lock @@ -79,6 +79,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a8/db/253133d7e7cb40d3af384bb2f5c0b4a2b7fdcffbc95c688cc67a20a3c103/accelerate-1.14.0-py3-none-any.whl", hash = "sha256:e94390c2863b873be18f623f9df48a0d8fe5eff13ea7f1a00092b0a7904888c6", size = 389246, upload-time = "2026-06-11T13:45:50.477Z" }, ] +[[package]] +name = "alembic" +version = "1.20.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mako" }, + { name = "sqlalchemy" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ed/aa/02910bdb8e2f1444f6654d5b296cd827d126f82209050ee7b1000f92ac4b/alembic-1.20.0.tar.gz", hash = "sha256:db505480647bc60386c5369402f4a57a506b7539c9e9ef5e270d45cbbe4939bf", size = 2093272, upload-time = "2026-09-11T19:09:11.126Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/27/78a89b55b0904d222183164e079b4ca56208e94eff1d35ad1f1ad5be9b06/alembic-1.20.0-py3-none-any.whl", hash = "sha256:77eb101048d95f982c0353e9233404889dcd7a6fc244c107836c0e2fc9cf7d9d", size = 268719, upload-time = "2026-09-11T19:09:12.88Z" }, +] + [[package]] name = "annotated-doc" version = "0.0.4" @@ -1082,6 +1096,7 @@ name = "invokeai" source = { editable = "." } dependencies = [ { name = "accelerate" }, + { name = "alembic" }, { name = "bcrypt" }, { name = "bitsandbytes", marker = "sys_platform == 'linux' or sys_platform == 'win32' or (extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-cuda') or (extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-rocm') or (extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-xpu') or (extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-rocm') or (extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-xpu') or (extra == 'extra-8-invokeai-rocm' and extra == 'extra-8-invokeai-xpu')" }, { name = "blake3" }, @@ -1124,6 +1139,7 @@ dependencies = [ { name = "semver" }, { name = "sentencepiece" }, { name = "spandrel" }, + { name = "sqlalchemy" }, { name = "torch", version = "2.7.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine != 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and sys_platform != 'win32') or (sys_platform == 'linux' and extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-cuda') or (sys_platform == 'linux' and extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-rocm') or (sys_platform == 'linux' and extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-xpu') or (sys_platform == 'linux' and extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-rocm') or (sys_platform == 'linux' and extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-xpu') or (sys_platform == 'linux' and extra == 'extra-8-invokeai-rocm' and extra == 'extra-8-invokeai-xpu') or (sys_platform == 'win32' and extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-cuda') or (sys_platform == 'win32' and extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-rocm') or (sys_platform == 'win32' and extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-xpu') or (sys_platform == 'win32' and extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-rocm') or (sys_platform == 'win32' and extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-xpu') or (sys_platform == 'win32' and extra == 'extra-8-invokeai-rocm' and extra == 'extra-8-invokeai-xpu')" }, { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-invokeai-cpu' and extra != 'extra-8-invokeai-cuda' and extra != 'extra-8-invokeai-rocm' and extra != 'extra-8-invokeai-xpu') or (sys_platform == 'win32' and extra == 'extra-8-invokeai-rocm') or (sys_platform == 'win32' and extra != 'extra-8-invokeai-cpu' and extra != 'extra-8-invokeai-cuda' and extra != 'extra-8-invokeai-xpu') or (extra == 'extra-8-invokeai-rocm' and extra == 'extra-8-invokeai-xpu') or (extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-cuda') or (extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-rocm') or (extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-xpu') or (extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-rocm') or (extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-xpu')" }, { name = "torch", version = "2.13.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-invokeai-cpu') or (platform_machine != 'x86_64' and extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-cuda') or (platform_machine != 'x86_64' and extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-rocm') or (platform_machine != 'x86_64' and extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-xpu') or (platform_machine != 'x86_64' and extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-rocm') or (platform_machine != 'x86_64' and extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-xpu') or (platform_machine != 'x86_64' and extra == 'extra-8-invokeai-rocm' and extra == 'extra-8-invokeai-xpu') or (sys_platform == 'win32' and extra == 'extra-8-invokeai-cpu') or (sys_platform != 'linux' and extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-cuda') or (sys_platform != 'linux' and extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-rocm') or (sys_platform != 'linux' and extra == 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-xpu') or (sys_platform != 'linux' and extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-rocm') or (sys_platform != 'linux' and extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-xpu') or (sys_platform != 'linux' and extra == 'extra-8-invokeai-rocm' and extra == 'extra-8-invokeai-xpu') or (extra != 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-rocm') or (extra != 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-cuda' and extra == 'extra-8-invokeai-xpu') or (extra != 'extra-8-invokeai-cpu' and extra == 'extra-8-invokeai-rocm' and extra == 'extra-8-invokeai-xpu')" }, @@ -1166,6 +1182,9 @@ dist = [ { name = "pipdeptree" }, { name = "twine" }, ] +mysql = [ + { name = "pymysql" }, +] onnx = [ { name = "onnxruntime" }, ] @@ -1212,6 +1231,7 @@ xpu = [ [package.metadata] requires-dist = [ { name = "accelerate" }, + { name = "alembic", specifier = ">=1.20,<2" }, { name = "bcrypt", specifier = "<4.0.0" }, { name = "bitsandbytes", marker = "sys_platform != 'darwin'" }, { name = "blake3" }, @@ -1257,6 +1277,7 @@ requires-dist = [ { name = "pydantic" }, { name = "pydantic-settings" }, { name = "pyjwt", specifier = ">=2.15.0" }, + { name = "pymysql", marker = "extra == 'mysql'", specifier = ">=1.2.3,<2" }, { name = "pypatchmatch" }, { name = "pytest", marker = "extra == 'test'", specifier = ">6.0.0" }, { name = "pytest-cov", marker = "extra == 'test'" }, @@ -1276,6 +1297,7 @@ requires-dist = [ { name = "sentencepiece", specifier = "==0.2.2" }, { name = "snakeviz", marker = "extra == 'dev'" }, { name = "spandrel" }, + { name = "sqlalchemy", specifier = ">=2.1,<2.2" }, { name = "torch", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')", specifier = ">=2.7.0,<3.0" }, { name = "torch", marker = "sys_platform == 'darwin'", specifier = ">=2.7.0,<2.8.0" }, { name = "torch", marker = "sys_platform == 'darwin' and extra == 'cpu'", specifier = "==2.7.1" }, @@ -1313,7 +1335,7 @@ requires-dist = [ { name = "xformers", marker = "(platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'xformers') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32' and extra == 'xformers')", specifier = ">=0.0.28.post1" }, { name = "xformers", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'xformers') or (sys_platform == 'win32' and extra == 'xformers')", specifier = ">=0.0.28.post1", index = "https://download.pytorch.org/whl/cu130" }, ] -provides-extras = ["xformers", "cpu", "cuda", "rocm", "xpu", "onnx", "onnx-cuda", "onnx-directml", "dist", "dev", "test"] +provides-extras = ["xformers", "cpu", "cuda", "rocm", "xpu", "mysql", "onnx", "onnx-cuda", "onnx-directml", "dist", "dev", "test"] [[package]] name = "jaraco-classes" @@ -1557,6 +1579,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7b/f0/92f2d609d6642b5f30cb50a885d2bf1483301c69d5786286500d15651ef2/lsprotocol-2025.0.0-py3-none-any.whl", hash = "sha256:f9d78f25221f2a60eaa4a96d3b4ffae011b107537facee61d3da3313880995c7", size = 76250, upload-time = "2025-06-17T21:30:19.455Z" }, ] +[[package]] +name = "mako" +version = "1.4.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5a/09/e07c4b5579a79f4b16f8d4f29f6c54514ac787c4ad506b8c4f28a0e6b0bf/mako-1.4.3.tar.gz", hash = "sha256:cd6537fe88d5fec315c55c2f8529bc4ce7a9a352ad7db3eeaa6a66e2dd4ec37a", size = 412799, upload-time = "2026-09-22T20:54:31.509Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6d/a0/053d6af3e8f871e0073b4a36732d9e65be77a72e5434c31b94f6af78a6bb/mako-1.4.3-py3-none-any.whl", hash = "sha256:723296007c870bfd6b3f0c3230dba7198096e5269297ebf5e4eff9e7ffa39d4f", size = 80164, upload-time = "2026-09-22T20:54:33.128Z" }, +] + [[package]] name = "markdown-it-py" version = "4.2.0" @@ -2535,6 +2569,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/50/ca/44de4e75f8aadc457f0634be3b542815078ded46dca30efb960edeecad6e/pyjwt-2.15.1-py3-none-any.whl", hash = "sha256:42d59d631f7768a1028a64c7ff581a9bf7519804daf91fc5b6c56e30eec5e193", size = 33860, upload-time = "2026-09-28T18:40:41.429Z" }, ] +[[package]] +name = "pymysql" +version = "1.2.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b1/d4/c15b459e25a23767d2f4065ef40968920320f04e302889574310c21c96a3/pymysql-1.2.3.tar.gz", hash = "sha256:d5b288529782e536ae171866df3ca9dc4f6cbfb3cc2f18e6f837fbb90dbc262b", size = 50629, upload-time = "2026-09-17T12:22:49.146Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a9/4b/0a906d8184f011ff8dbd4722743783867589b33269d2c5fff238d636fdcb/pymysql-1.2.3-py3-none-any.whl", hash = "sha256:14f1c68e2ed859243ae5ca41ffbe677027fc46bc136a9f0be8a4e928e5e7415a", size = 46740, upload-time = "2026-09-17T12:22:47.826Z" }, +] + [[package]] name = "pynndescent" version = "0.6.0" @@ -3125,6 +3168,26 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/74/31/411ea965835534c43d4b98d451968354876e0e867ea1fd42669e4cca0732/spandrel-0.4.2-py3-none-any.whl", hash = "sha256:6c93e3ecbeb0e548fd2df45a605472b34c1614287c56b51bb33cdef7ae5235b5", size = 320811, upload-time = "2026-02-21T01:52:25.015Z" }, ] +[[package]] +name = "sqlalchemy" +version = "2.1.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1f/44/311bac6b6ef81e4dfd0287d04900108b1f5c00c9761dd3c0a2b7b9d0f86b/sqlalchemy-2.1.4.tar.gz", hash = "sha256:7bd7ad604487daa7eab8716471c29a7185f17b5287ce73bb7bc79fea050d8cfd", size = 10544216, upload-time = "2026-10-07T17:33:59.116Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/5e/cb5b078e007340661b010fa8bd31ce27468f88e09b35266544df4e0c52ca/sqlalchemy-2.1.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f953be9ba26039a24a5205c65d33518b608ce6f4f0f4e9b9c14eaf42a10dfc52", size = 2466469, upload-time = "2026-10-07T18:17:24.049Z" }, + { url = "https://files.pythonhosted.org/packages/b1/98/44e2fdc5bc053dae559bf4f4eb7967ceecbad162299ecfc8de2edc3fcbe7/sqlalchemy-2.1.4-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ac64fce94c5b389062d2e3806db5dc780447591e0dfd5ead218c884f0703f2e", size = 4668213, upload-time = "2026-10-07T18:37:42.294Z" }, + { url = "https://files.pythonhosted.org/packages/08/25/ed2262f964687b06f10c2c98b2dc9c9ed211f7cc11702879969a9ac217e4/sqlalchemy-2.1.4-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3e5045fb6aadbb0f978ab9b9d8822f7b7a97d2281814e7d13d791155664eace3", size = 4720857, upload-time = "2026-10-07T18:24:46.842Z" }, + { url = "https://files.pythonhosted.org/packages/d9/e4/33413f0fafbcf3b332320aac2c1e40f3b4f17e56359a9474cb10de4bee8b/sqlalchemy-2.1.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:71040390ef01c85e9d26e5c83cb0c5942dcc8725c49186430af160ce2f54234d", size = 4591496, upload-time = "2026-10-07T18:37:44.433Z" }, + { url = "https://files.pythonhosted.org/packages/01/e3/168a0f93efd6ec40f59645a7e45ab08918e0bc8ecf07656e4ca09acdcc30/sqlalchemy-2.1.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a577e2127e52b0fe2bc54c73abb375a20ffe6f59fbc5568ccafc233f5bfcf8ef", size = 4667997, upload-time = "2026-10-07T18:24:48.72Z" }, + { url = "https://files.pythonhosted.org/packages/54/79/0a852ef65864acd8d577d7aa6f67146167382bd6faee7a7586b9e6e28275/sqlalchemy-2.1.4-cp312-cp312-win32.whl", hash = "sha256:6c79e0c824d51c586757ecd342160bbdede9010df04bb71b9bbfffd5c7b6ee29", size = 2376664, upload-time = "2026-10-07T18:25:00.637Z" }, + { url = "https://files.pythonhosted.org/packages/27/b9/a5934263bb1d712f743289ca224ab3b87e3570ac157802291e37ab85d365/sqlalchemy-2.1.4-cp312-cp312-win_amd64.whl", hash = "sha256:dffa69d2f3ba1933c1c1882dbef8fb3231b33eb19263e8b8c5cea24995071f06", size = 2429357, upload-time = "2026-10-07T18:25:02.565Z" }, + { url = "https://files.pythonhosted.org/packages/a5/fa/a2323d81384ff214aa189057b7455b63623e66f28208b982e86c3cb042f5/sqlalchemy-2.1.4-cp312-cp312-win_arm64.whl", hash = "sha256:e30524ae24e31d83e1b5f734862882c442f4158e3566f2c5f5e9bd3c659bb517", size = 2388756, upload-time = "2026-10-07T18:22:36.025Z" }, + { url = "https://files.pythonhosted.org/packages/f7/62/dbf11a262f6fbb41390cab2d8e47a30ec0961018b68201607b599dd489f5/sqlalchemy-2.1.4-py3-none-any.whl", hash = "sha256:0b96edcc2cd60fe1e35f67a46f4eb076e57297841b9eae949ac5f196593f00a7", size = 2054935, upload-time = "2026-10-07T18:01:16.403Z" }, +] + [[package]] name = "starlette" version = "1.7.0"