Compare commits
54 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| ad3ff57260 | |||
| c9f1a5dddc | |||
| 2042b345aa | |||
| b72e8ed5e7 | |||
| e79fe339aa | |||
| 96ae76e70c | |||
| ed2c02bcf9 | |||
| 2f3031cfc7 | |||
| 459adc077a | |||
| 19f18ef480 | |||
| 6148b28cba | |||
| 012632111e | |||
| b9a59a8649 | |||
| f8a8c9fdff | |||
| d37b444dd5 | |||
| 02ad9a29a7 | |||
| 24c504518d | |||
| 92f07ef3d7 | |||
| 22e010c554 | |||
| 8ebc1bb9e1 | |||
| 3fc21dcd99 | |||
| ad5e0e8a72 | |||
| c384df20ce | |||
| 4e944f3d05 | |||
| e0455daa2e | |||
| f3f1afe9af | |||
| 94dc9cfc3f | |||
| a9c0ced1ca | |||
| c8626a9afd | |||
| 55a1f66e57 | |||
| 5b5f35821f | |||
| e8161e6417 | |||
| 1a10c32b12 | |||
| cb8cb87d9f | |||
| 96117d54b9 | |||
| 53f58e45c3 | |||
| 34d708be7e | |||
| 418e8dc052 | |||
| b8567ebe08 | |||
| 827bfa6370 | |||
| ceba64e34a | |||
| 8646d58f06 | |||
| 8ce6ff11c3 | |||
| 5c7aaebaba | |||
| be7d086c0b | |||
| 91a7b69f1d | |||
| 4e078b35d4 | |||
| d8a8259c73 | |||
| 9e7ce25b45 | |||
| b634ee50d1 | |||
| 5a47c516fd | |||
| d34b8673e1 | |||
| 943302f789 | |||
| e23a951dbe |
@@ -3,6 +3,9 @@ node_modules/
|
|||||||
dist_ts/
|
dist_ts/
|
||||||
dist_*/
|
dist_*/
|
||||||
|
|
||||||
|
# rust build artifacts
|
||||||
|
rust/target/
|
||||||
|
|
||||||
# config
|
# config
|
||||||
.nogit/
|
.nogit/
|
||||||
|
|
||||||
@@ -10,5 +13,8 @@ dist_*/
|
|||||||
package-lock.json
|
package-lock.json
|
||||||
yarn.lock
|
yarn.lock
|
||||||
|
|
||||||
|
# generated bundle (rebuilt on every build, embeds version)
|
||||||
|
ts_debugserver/bundled.ts
|
||||||
|
|
||||||
# playwright
|
# playwright
|
||||||
.playwright-mcp/
|
.playwright-mcp/
|
||||||
|
|||||||
+18
-4
@@ -5,19 +5,19 @@
|
|||||||
"githost": "code.foss.global",
|
"githost": "code.foss.global",
|
||||||
"gitscope": "push.rocks",
|
"gitscope": "push.rocks",
|
||||||
"gitrepo": "smartdb",
|
"gitrepo": "smartdb",
|
||||||
"description": "A pure TypeScript MongoDB wire-protocol-compatible database server with pluggable storage, indexing, transactions, and zero external binary dependencies.",
|
"description": "A MongoDB-compatible embedded database server with wire protocol support, backed by a high-performance Rust engine.",
|
||||||
"npmPackagename": "@push.rocks/smartdb",
|
"npmPackagename": "@push.rocks/smartdb",
|
||||||
"license": "MIT",
|
"license": "MIT",
|
||||||
"projectDomain": "push.rocks",
|
"projectDomain": "push.rocks",
|
||||||
"keywords": [
|
"keywords": [
|
||||||
"mongodb",
|
"mongodb-compatible",
|
||||||
"wire protocol",
|
"wire protocol",
|
||||||
"typescript database",
|
"embedded database",
|
||||||
"in-memory database",
|
"in-memory database",
|
||||||
"testing",
|
"testing",
|
||||||
"local database",
|
"local database",
|
||||||
"database server",
|
"database server",
|
||||||
"typescript"
|
"rust"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
"release": {
|
"release": {
|
||||||
@@ -32,6 +32,20 @@
|
|||||||
"@git.zone/tsdoc": {
|
"@git.zone/tsdoc": {
|
||||||
"legal": "\n## License and Legal Information\n\nThis repository contains open-source code that is licensed under the MIT License. A copy of the MIT License can be found in the [license](license) file within this repository. \n\n**Please note:** The MIT License does not grant permission to use the trade names, trademarks, service marks, or product names of the project, except as required for reasonable and customary use in describing the origin of the work and reproducing the content of the NOTICE file.\n\n### Trademarks\n\nThis project is owned and maintained by Task Venture Capital GmbH. The names and logos associated with Task Venture Capital GmbH and any related products or services are trademarks of Task Venture Capital GmbH and are not included within the scope of the MIT license granted herein. Use of these trademarks must comply with Task Venture Capital GmbH's Trademark Guidelines, and any usage must be approved in writing by Task Venture Capital GmbH.\n\n### Company Information\n\nTask Venture Capital GmbH \nRegistered at District court Bremen HRB 35230 HB, Germany\n\nFor any legal inquiries or if you require further information, please contact us via email at hello@task.vc.\n\nBy using this repository, you acknowledge that you have read this section, agree to comply with its terms, and understand that the licensing of the code does not imply endorsement by Task Venture Capital GmbH of any derivative works.\n"
|
"legal": "\n## License and Legal Information\n\nThis repository contains open-source code that is licensed under the MIT License. A copy of the MIT License can be found in the [license](license) file within this repository. \n\n**Please note:** The MIT License does not grant permission to use the trade names, trademarks, service marks, or product names of the project, except as required for reasonable and customary use in describing the origin of the work and reproducing the content of the NOTICE file.\n\n### Trademarks\n\nThis project is owned and maintained by Task Venture Capital GmbH. The names and logos associated with Task Venture Capital GmbH and any related products or services are trademarks of Task Venture Capital GmbH and are not included within the scope of the MIT license granted herein. Use of these trademarks must comply with Task Venture Capital GmbH's Trademark Guidelines, and any usage must be approved in writing by Task Venture Capital GmbH.\n\n### Company Information\n\nTask Venture Capital GmbH \nRegistered at District court Bremen HRB 35230 HB, Germany\n\nFor any legal inquiries or if you require further information, please contact us via email at hello@task.vc.\n\nBy using this repository, you acknowledge that you have read this section, agree to comply with its terms, and understand that the licensing of the code does not imply endorsement by Task Venture Capital GmbH of any derivative works.\n"
|
||||||
},
|
},
|
||||||
|
"@git.zone/tsbundle": {
|
||||||
|
"bundles": [
|
||||||
|
{
|
||||||
|
"from": "./ts_debugui/index.ts",
|
||||||
|
"to": "./ts_debugserver/bundled.ts",
|
||||||
|
"outputMode": "base64ts",
|
||||||
|
"bundler": "esbuild",
|
||||||
|
"includeFiles": ["./html/index.html"]
|
||||||
|
}
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"@git.zone/tsrust": {
|
||||||
|
"targets": ["linux_amd64", "linux_arm64"]
|
||||||
|
},
|
||||||
"@ship.zone/szci": {
|
"@ship.zone/szci": {
|
||||||
"npmGlobalTools": []
|
"npmGlobalTools": []
|
||||||
}
|
}
|
||||||
|
|||||||
+157
@@ -1,5 +1,162 @@
|
|||||||
# Changelog
|
# Changelog
|
||||||
|
|
||||||
|
## 2026-05-02 - 2.9.0 - feat(server)
|
||||||
|
add tenant management, health checks, and database export/import APIs
|
||||||
|
|
||||||
|
- adds TypeScript and Rust management commands for creating, listing, deleting, and rotating isolated database tenants
|
||||||
|
- introduces health reporting with storage, auth, database, collection, and uptime information
|
||||||
|
- supports exporting and importing single-database snapshots and increases IPC payload size for larger transfers
|
||||||
|
- adds integration coverage for tenant isolation, password rotation, persistence across restart, and database restore flows
|
||||||
|
|
||||||
|
## 2026-04-29 - 2.8.0 - feat(transactions)
|
||||||
|
add single-node transaction support with session-aware reads, commits, aborts, and transaction metrics
|
||||||
|
|
||||||
|
- Buffer insert, update, delete, find, count, distinct, and findAndModify operations inside driver sessions and apply them on commit with write-conflict checks
|
||||||
|
- Return MongoDB-compatible NoSuchTransaction and WriteConflict errors for transaction lifecycle failures
|
||||||
|
- Expose authenticated users in connectionStatus and add session, transaction, auth, and oplog data to serverStatus and management metrics
|
||||||
|
- Document transaction support and extend bridge metrics typings and integration tests accordingly
|
||||||
|
|
||||||
|
## 2026-04-29 - 2.7.1 - fix(repo)
|
||||||
|
no changes to commit
|
||||||
|
|
||||||
|
|
||||||
|
## 2026-04-14 - 2.7.0 - feat(update)
|
||||||
|
add aggregation pipeline updates and enforce immutable _id handling
|
||||||
|
|
||||||
|
- support aggregation pipeline syntax in update and findOneAndUpdate operations, including upserts
|
||||||
|
- add $unset stage support for aggregation-based document transformations
|
||||||
|
- return an ImmutableField error when updates attempt to change _id and preserve _id when omitted from replacements
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.6.2 - fix(readme)
|
||||||
|
align architecture diagram formatting in the documentation
|
||||||
|
|
||||||
|
- Adjusts spacing and box alignment in the README architecture diagram for clearer presentation.
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.6.1 - fix(readme)
|
||||||
|
correct ASCII diagram spacing in architecture overview
|
||||||
|
|
||||||
|
- Adjusts alignment in the README architecture diagram for clearer visual formatting.
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.6.0 - feat(readme)
|
||||||
|
document index enforcement, storage reliability, and data integrity validation features
|
||||||
|
|
||||||
|
- Add documentation for engine-level unique index enforcement and duplicate key behavior
|
||||||
|
- Describe storage engine reliability features including WAL, CRC32 checks, compaction, hint file staleness detection, and stale socket cleanup
|
||||||
|
- Add usage documentation for the offline data integrity validation CLI
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.5.9 - fix(rustdb-storage)
|
||||||
|
run collection compaction during file storage initialization after crashes
|
||||||
|
|
||||||
|
- Triggers compaction for all loaded collections before starting the periodic background compaction task.
|
||||||
|
- Helps clean up dead weight left from before a crash during startup.
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.5.8 - fix(rustdb-storage)
|
||||||
|
detect stale hint files using data file size metadata and add restart persistence regression tests
|
||||||
|
|
||||||
|
- Store the current data.rdb size in hint file headers and validate it on load to rebuild KeyDir when hints are stale or written in the old format.
|
||||||
|
- Persist updated hint metadata after compaction and shutdown to avoid missing appended tombstones after restart.
|
||||||
|
- Add validation reporting for stale hint files based on recorded versus actual data file size.
|
||||||
|
- Add regression tests covering delete persistence across restarts, missing hint recovery, stale socket cleanup, and unique index enforcement persistence.
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.5.7 - fix(repo)
|
||||||
|
no changes to commit
|
||||||
|
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.5.6 - fix(repo)
|
||||||
|
no changes to commit
|
||||||
|
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.5.5 - fix(repo)
|
||||||
|
no changes to commit
|
||||||
|
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.5.4 - fix(package)
|
||||||
|
bump package version to 2.5.3
|
||||||
|
|
||||||
|
- Updates the package metadata version by one patch release.
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.5.3 - fix(rustdb-commands)
|
||||||
|
restore persisted index initialization before writes to enforce unique constraints after restart
|
||||||
|
|
||||||
|
- load stored index specifications from storage when creating command context index engines
|
||||||
|
- rebuild index data from existing documents so custom indexes are active before insert, update, and upsert operations
|
||||||
|
- add @push.rocks/smartdata as a runtime dependency
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.5.2 - fix(rustdb-indexes)
|
||||||
|
persist created indexes and restore them on server startup
|
||||||
|
|
||||||
|
- Save index specifications to storage when indexes are created.
|
||||||
|
- Remove persisted index metadata when indexes are dropped by name, key spec, or wildcard.
|
||||||
|
- Rebuild in-memory index engines from stored definitions and existing documents during startup.
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.5.1 - fix(docs)
|
||||||
|
update project documentation
|
||||||
|
|
||||||
|
- Modifies a single documentation-related file with a minimal text change.
|
||||||
|
- No source code, API, or package metadata changes are indicated in the diff summary.
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.5.0 - feat(storage)
|
||||||
|
add offline data validation and strengthen storage/index integrity checks
|
||||||
|
|
||||||
|
- adds a `--validate-data <PATH>` CLI mode to run offline integrity checks on storage directories
|
||||||
|
- introduces storage validation reporting for headers, checksums, duplicate ids, tombstones, and stale or orphaned hint entries
|
||||||
|
- pre-checks unique index constraints before insert, update, upsert, and findAndModify writes to prevent duplicate-key violations before storage changes
|
||||||
|
- validates hint files against data files during collection load and rebuilds indexes from data when hints are stale
|
||||||
|
- ensures new data files always receive a SMARTDB header and persists fresh hint files after successful compaction
|
||||||
|
- cleans up stale local Unix socket files before starting the TypeScript local server
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.4.1 - fix(package)
|
||||||
|
update package metadata
|
||||||
|
|
||||||
|
- Adjusts package manifest content with a minimal one-line change.
|
||||||
|
|
||||||
|
## 2026-04-05 - 2.4.0 - feat(rustdb)
|
||||||
|
add restore and periodic persistence support for in-memory storage
|
||||||
|
|
||||||
|
- Restore previously persisted state during startup when a persist path is configured.
|
||||||
|
- Spawn a background task to periodically persist in-memory data using the configured interval.
|
||||||
|
- Warn when running purely in-memory without durable persistence configured.
|
||||||
|
|
||||||
|
## 2026-04-04 - 2.3.1 - fix(package)
|
||||||
|
update package metadata
|
||||||
|
|
||||||
|
- Adjusts a single package-level metadata entry in the project configuration.
|
||||||
|
|
||||||
|
## 2026-04-04 - 2.3.0 - feat(test)
|
||||||
|
add integration coverage for file storage, compaction, migration, and LocalSmartDb workflows
|
||||||
|
|
||||||
|
- adds end-to-end tests for file-backed storage creation, CRUD operations, bulk updates, persistence, and index file generation
|
||||||
|
- adds compaction stress tests covering repeated updates, tombstones, file shrinking behavior, and restart integrity
|
||||||
|
- adds migration tests for automatic v0 JSON layout detection, v1 conversion, restart persistence, and post-migration writes
|
||||||
|
- adds LocalSmartDb lifecycle and unix socket tests, including restart persistence, custom socket paths, and database isolation
|
||||||
|
|
||||||
|
## 2026-04-04 - 2.2.0 - feat(storage)
|
||||||
|
add Bitcask storage migration, binary WAL, and data compaction support
|
||||||
|
|
||||||
|
- add TypeScript storage migration from legacy JSON collections to the v1 Bitcask binary format before starting the Rust engine
|
||||||
|
- replace the legacy JSON WAL with a binary write-ahead log plus shared binary record and KeyDir infrastructure in rustdb-storage
|
||||||
|
- introduce data file compaction with dead-record reclamation and tests, and add the bson dependency for BSON serialization during migration
|
||||||
|
|
||||||
|
## 2026-04-02 - 2.1.1 - fix(package)
|
||||||
|
update package metadata
|
||||||
|
|
||||||
|
- Adjusts a single package metadata entry in package.json.
|
||||||
|
|
||||||
|
## 2026-04-02 - 2.1.0 - feat(smartdb)
|
||||||
|
add operation log APIs, point-in-time revert support, and a web-based debug dashboard
|
||||||
|
|
||||||
|
- records insert, update, and delete operations with before/after document snapshots in the Rust oplog
|
||||||
|
- adds management and TypeScript APIs for metrics, oplog queries, collection browsing, document browsing, and revert-to-sequence operations
|
||||||
|
- introduces new debugserver and debugui package exports with bundled browser assets served through typedserver
|
||||||
|
|
||||||
|
## 2026-03-26 - 2.0.0 - BREAKING CHANGE(core)
|
||||||
|
replace the TypeScript database engine with a Rust-backed embedded server and bridge
|
||||||
|
|
||||||
|
- adds a Rust workspace implementing wire protocol handling, commands, storage, indexing, aggregation, sessions, and transactions
|
||||||
|
- switches the TypeScript package to lifecycle orchestration via RustDbBridge and @push.rocks/smartrust
|
||||||
|
- removes previously exported TypeScript internals such as query, update, index, transaction, session, WAL, checksum, and router utilities
|
||||||
|
- updates build/test tooling and package metadata to compile and ship Rust binaries
|
||||||
|
|
||||||
## 2026-03-26 - 1.0.1 - fix(repo)
|
## 2026-03-26 - 1.0.1 - fix(repo)
|
||||||
no changes to commit
|
no changes to commit
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,13 @@
|
|||||||
|
<!DOCTYPE html>
|
||||||
|
<html lang="en">
|
||||||
|
<head>
|
||||||
|
<meta charset="UTF-8" />
|
||||||
|
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
||||||
|
<title>SmartDB Debug</title>
|
||||||
|
<style>body { margin: 0; background: #09090b; }</style>
|
||||||
|
<script type="module" src="/bundle.js"></script>
|
||||||
|
</head>
|
||||||
|
<body>
|
||||||
|
<smartdb-debugui></smartdb-debugui>
|
||||||
|
</body>
|
||||||
|
</html>
|
||||||
@@ -1,4 +1,6 @@
|
|||||||
Copyright (c) 2021 Task Venture Capital GmbH (hello@task.vc)
|
MIT License
|
||||||
|
|
||||||
|
Copyright (c) 2024 Task Venture Capital GmbH
|
||||||
|
|
||||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||||
of this software and associated documentation files (the "Software"), to deal
|
of this software and associated documentation files (the "Software"), to deal
|
||||||
|
|||||||
+16
-13
@@ -1,34 +1,37 @@
|
|||||||
{
|
{
|
||||||
"name": "@push.rocks/smartdb",
|
"name": "@push.rocks/smartdb",
|
||||||
"version": "1.0.1",
|
"version": "2.9.0",
|
||||||
"private": false,
|
"private": false,
|
||||||
"description": "A pure TypeScript MongoDB wire-protocol-compatible database server with pluggable storage, indexing, transactions, and zero external binary dependencies.",
|
"description": "A MongoDB-compatible embedded database server with wire protocol support, backed by a high-performance Rust engine.",
|
||||||
"exports": {
|
"exports": {
|
||||||
".": "./dist_ts/index.js"
|
".": "./dist_ts/index.js",
|
||||||
|
"./debugui": "./dist_ts_debugui/index.js",
|
||||||
|
"./debugserver": "./dist_ts_debugserver/index.js"
|
||||||
},
|
},
|
||||||
"type": "module",
|
"type": "module",
|
||||||
"author": "Task Venture Capital GmbH",
|
"author": "Task Venture Capital GmbH",
|
||||||
"license": "MIT",
|
"license": "MIT",
|
||||||
"scripts": {
|
"scripts": {
|
||||||
|
"test:before": "(tsrust)",
|
||||||
"test": "(tstest test/. --verbose --logfile --timeout 60)",
|
"test": "(tstest test/. --verbose --logfile --timeout 60)",
|
||||||
"build": "(tsbuild tsfolders)",
|
"build": "(tsbundle) && (tsbuild tsfolders) && (tsrust)",
|
||||||
"buildDocs": "tsdoc"
|
"buildDocs": "tsdoc"
|
||||||
},
|
},
|
||||||
"devDependencies": {
|
"devDependencies": {
|
||||||
"@git.zone/tsbuild": "^4.4.0",
|
"@git.zone/tsbuild": "^4.4.0",
|
||||||
"@git.zone/tsbundle": "^2.10.0",
|
"@git.zone/tsbundle": "^2.10.0",
|
||||||
"@git.zone/tsrun": "^2.0.2",
|
"@git.zone/tsrun": "^2.0.2",
|
||||||
|
"@git.zone/tsrust": "^1.3.2",
|
||||||
"@git.zone/tstest": "^3.6.1",
|
"@git.zone/tstest": "^3.6.1",
|
||||||
"@types/node": "^25.5.0",
|
"@types/node": "^25.5.0",
|
||||||
"mongodb": "^7.1.1"
|
"mongodb": "^7.1.1"
|
||||||
},
|
},
|
||||||
"dependencies": {
|
"dependencies": {
|
||||||
"@push.rocks/smartfs": "^1.5.0",
|
"@api.global/typedserver": "^8.0.0",
|
||||||
"@push.rocks/smartpath": "^6.0.0",
|
"@design.estate/dees-element": "^2.0.0",
|
||||||
"@push.rocks/smartpromise": "^4.2.3",
|
"@push.rocks/smartdata": "7.1.5",
|
||||||
"@push.rocks/smartrx": "^3.0.10",
|
"@push.rocks/smartrust": "^1.3.2",
|
||||||
"bson": "^7.2.0",
|
"bson": "^7.2.0"
|
||||||
"mingo": "^7.2.0"
|
|
||||||
},
|
},
|
||||||
"browserslist": [
|
"browserslist": [
|
||||||
"last 1 chrome versions"
|
"last 1 chrome versions"
|
||||||
@@ -46,14 +49,14 @@
|
|||||||
"readme.md"
|
"readme.md"
|
||||||
],
|
],
|
||||||
"keywords": [
|
"keywords": [
|
||||||
"mongodb",
|
"mongodb-compatible",
|
||||||
"wire protocol",
|
"wire protocol",
|
||||||
"typescript database",
|
"embedded database",
|
||||||
"in-memory database",
|
"in-memory database",
|
||||||
"testing",
|
"testing",
|
||||||
"local database",
|
"local database",
|
||||||
"database server",
|
"database server",
|
||||||
"typescript"
|
"rust"
|
||||||
],
|
],
|
||||||
"homepage": "https://code.foss.global/push.rocks/smartdb#readme",
|
"homepage": "https://code.foss.global/push.rocks/smartdb#readme",
|
||||||
"repository": {
|
"repository": {
|
||||||
|
|||||||
Generated
+1369
-19
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,6 @@
|
|||||||
# @push.rocks/smartdb
|
# @push.rocks/smartdb
|
||||||
|
|
||||||
A pure TypeScript MongoDB wire-protocol-compatible database server. Zero binary dependencies, instant startup, pluggable storage — use the official MongoDB driver and it just works. ⚡
|
A MongoDB-compatible embedded database server powered by Rust 🦀⚡ — use the official `mongodb` driver and it just works. No binary downloads, instant startup, zero config. Features a built-in **operation log** with **point-in-time revert** and a web-based **debug dashboard**.
|
||||||
|
|
||||||
## Install
|
## Install
|
||||||
|
|
||||||
@@ -14,24 +14,74 @@ npm install @push.rocks/smartdb
|
|||||||
|
|
||||||
For reporting bugs, issues, or security vulnerabilities, please visit [community.foss.global/](https://community.foss.global/). This is the central community hub for all issue reporting. Developers who sign and comply with our contribution agreement and go through identification can also get a [code.foss.global/](https://code.foss.global/) account to submit Pull Requests directly.
|
For reporting bugs, issues, or security vulnerabilities, please visit [community.foss.global/](https://community.foss.global/). This is the central community hub for all issue reporting. Developers who sign and comply with our contribution agreement and go through identification can also get a [code.foss.global/](https://code.foss.global/) account to submit Pull Requests directly.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
## What It Does
|
## What It Does
|
||||||
|
|
||||||
`@push.rocks/smartdb` is a **real database server** written entirely in TypeScript that speaks the MongoDB binary wire protocol. Connect with the official `mongodb` Node.js driver — no mocks, no stubs, no MongoDB binary required.
|
`@push.rocks/smartdb` is a **real database server** that speaks the wire protocol used by MongoDB drivers. The core engine is written in Rust for high performance, with a thin TypeScript orchestration layer. Connect with the standard `mongodb` Node.js driver — no mocks, no stubs, no external binaries required.
|
||||||
|
|
||||||
### Why SmartDB?
|
### Why SmartDB?
|
||||||
|
|
||||||
| | SmartDB | Real MongoDB |
|
| | SmartDB | External DB Server |
|
||||||
|---|---|---|
|
|---|---|---|
|
||||||
| **Startup time** | ~5ms | ~2-5s |
|
| **Startup time** | ~30ms | ~2-5s |
|
||||||
| **Binary download** | None | ~200MB |
|
| **Binary download** | Bundled (~7MB) | ~200MB+ |
|
||||||
| **Node.js only** | ✅ | ❌ |
|
| **Install** | `pnpm add` | System package / Docker |
|
||||||
| **Persistence** | Memory or file-based | Full disk engine |
|
| **Persistence** | Memory or file-based | Full disk engine |
|
||||||
| **Perfect for** | Unit tests, CI/CD, prototyping, local dev | Production |
|
| **Debug UI** | Built-in 🖥️ | External tooling |
|
||||||
|
| **Point-in-time revert** | Built-in ⏪ | Requires oplog tailing |
|
||||||
|
| **Perfect for** | Unit tests, CI/CD, prototyping, local dev, embedded | Production at scale |
|
||||||
|
|
||||||
### Two Ways to Use It
|
### Three Ways to Use It
|
||||||
|
|
||||||
|
- 🎯 **`LocalSmartDb`** — Zero-config convenience. Give it a folder path, get a persistent database over a Unix socket. Done.
|
||||||
- 🏗️ **`SmartdbServer`** — Full control. Configure port, host, storage backend, Unix sockets. Great for test fixtures or custom setups.
|
- 🏗️ **`SmartdbServer`** — Full control. Configure port, host, storage backend, Unix sockets. Great for test fixtures or custom setups.
|
||||||
- 🎯 **`LocalSmartDb`** — Zero-config convenience. Give it a folder path, get a persistent MongoDB-compatible database over a Unix socket. Done.
|
- 🖥️ **`SmartdbDebugServer`** — Launch a web dashboard to visually browse collections, inspect the operation log, and revert to any point in time.
|
||||||
|
|
||||||
|
### Architecture: TypeScript + Rust 🦀
|
||||||
|
|
||||||
|
SmartDB uses a **sidecar binary** pattern — TypeScript handles lifecycle, Rust handles all database operations:
|
||||||
|
|
||||||
|
```
|
||||||
|
┌──────────────────────────────────────────────────────────────┐
|
||||||
|
│ Your Application │
|
||||||
|
│ (TypeScript / Node.js) │
|
||||||
|
│ ┌──────────────────┐ ┌───────────────────────────┐ │
|
||||||
|
│ │ SmartdbServer │─────▶│ RustDbBridge (IPC) │ │
|
||||||
|
│ │ or LocalSmartDb │ │ @push.rocks/smartrust │ │
|
||||||
|
│ └──────────────────┘ └───────────┬───────────────┘ │
|
||||||
|
└────────────────────────────────────────┼─────────────────────┘
|
||||||
|
│ spawn + JSON IPC
|
||||||
|
▼
|
||||||
|
┌──────────────────────────────────────────────────────────────┐
|
||||||
|
│ rustdb binary │
|
||||||
|
│ │
|
||||||
|
│ ┌──────────────┐ ┌──────────────┐ ┌───────────────┐ │
|
||||||
|
│ │ Wire Protocol│→ │Command Router│→ │ Handlers │ │
|
||||||
|
│ │ (OP_MSG) │ │ (40+ cmds) │ │ Find,Insert.. │ │
|
||||||
|
│ └──────────────┘ └──────────────┘ └───────┬───────┘ │
|
||||||
|
│ │ │
|
||||||
|
│ ┌─────────┐ ┌────────┐ ┌───────────┐ ┌──────┴──────┐ │
|
||||||
|
│ │ Query │ │ Update │ │Aggregation│ │ Index │ │
|
||||||
|
│ │ Matcher │ │ Engine │ │ Engine │ │ Engine │ │
|
||||||
|
│ └─────────┘ └────────┘ └───────────┘ └─────────────┘ │
|
||||||
|
│ │
|
||||||
|
│ ┌──────────────────┐ ┌──────────────────┐ ┌──────────┐ │
|
||||||
|
│ │ MemoryStorage │ │ FileStorage │ │ OpLog │ │
|
||||||
|
│ └──────────────────┘ └──────────────────┘ └──────────┘ │
|
||||||
|
└──────────────────────────────────────────────────────────────┘
|
||||||
|
▲
|
||||||
|
│ TCP / Unix Socket (wire protocol)
|
||||||
|
│
|
||||||
|
┌─────────────┴────────────────────────────────────────────────┐
|
||||||
|
│ MongoClient (mongodb npm driver) │
|
||||||
|
│ Connects directly to Rust binary │
|
||||||
|
└──────────────────────────────────────────────────────────────┘
|
||||||
|
```
|
||||||
|
|
||||||
|
The TypeScript layer handles **lifecycle only** (start/stop/configure via IPC). All database operations flow directly from the `MongoClient` to the Rust binary over TCP or Unix sockets — **zero per-query IPC overhead**.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
## Quick Start
|
## Quick Start
|
||||||
|
|
||||||
@@ -47,11 +97,11 @@ import { MongoClient } from 'mongodb';
|
|||||||
const db = new LocalSmartDb({ folderPath: './my-data' });
|
const db = new LocalSmartDb({ folderPath: './my-data' });
|
||||||
const { connectionUri } = await db.start();
|
const { connectionUri } = await db.start();
|
||||||
|
|
||||||
// Connect with the standard MongoDB driver
|
// Connect with the standard driver
|
||||||
const client = new MongoClient(connectionUri, { directConnection: true });
|
const client = new MongoClient(connectionUri, { directConnection: true });
|
||||||
await client.connect();
|
await client.connect();
|
||||||
|
|
||||||
// Use exactly like MongoDB
|
// Use it like any wire-protocol-compatible database
|
||||||
const users = client.db('myapp').collection('users');
|
const users = client.db('myapp').collection('users');
|
||||||
await users.insertOne({ name: 'Alice', email: 'alice@example.com' });
|
await users.insertOne({ name: 'Alice', email: 'alice@example.com' });
|
||||||
const user = await users.findOne({ name: 'Alice' });
|
const user = await users.findOne({ name: 'Alice' });
|
||||||
@@ -83,13 +133,94 @@ await client.close();
|
|||||||
await server.stop();
|
await server.stop();
|
||||||
```
|
```
|
||||||
|
|
||||||
|
### Option 3: Debug Server (Visual Dashboard) 🖥️
|
||||||
|
|
||||||
|
Launch a web-based dashboard to inspect your database in real time:
|
||||||
|
|
||||||
|
```typescript
|
||||||
|
import { SmartdbServer } from '@push.rocks/smartdb';
|
||||||
|
import { SmartdbDebugServer } from '@push.rocks/smartdb/debugserver';
|
||||||
|
|
||||||
|
const server = new SmartdbServer({ storage: 'memory' });
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
const debugServer = new SmartdbDebugServer(server, { port: 4000 });
|
||||||
|
await debugServer.start();
|
||||||
|
// Open http://localhost:4000 in your browser 🚀
|
||||||
|
```
|
||||||
|
|
||||||
|
The debug dashboard gives you:
|
||||||
|
- 📊 **Dashboard** — server status, uptime, database/collection counts, operation breakdown
|
||||||
|
- 📁 **Collection Browser** — browse databases, collections, and documents interactively
|
||||||
|
- 📝 **OpLog Timeline** — every insert, update, and delete with expandable field-level diffs
|
||||||
|
- ⏪ **Point-in-Time Revert** — select any oplog sequence, preview what will be undone, and execute
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 📝 Operation Log & Point-in-Time Revert
|
||||||
|
|
||||||
|
Every write operation (insert, update, delete) is automatically recorded in an in-memory **operation log (OpLog)** with full before/after document snapshots. The OpLog lives in RAM and resets on restart — it covers the current session only. This enables:
|
||||||
|
|
||||||
|
- **Change tracking** — see exactly what changed, when, and in which collection
|
||||||
|
- **Field-level diffs** — compare previous and new document states
|
||||||
|
- **Point-in-time revert** — undo operations back to any sequence number
|
||||||
|
- **Dry-run preview** — see what would be reverted before executing
|
||||||
|
|
||||||
|
### Programmatic OpLog API
|
||||||
|
|
||||||
|
```typescript
|
||||||
|
import { SmartdbServer } from '@push.rocks/smartdb';
|
||||||
|
|
||||||
|
const server = new SmartdbServer({ port: 27017 });
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
// ... perform some CRUD operations via MongoClient ...
|
||||||
|
|
||||||
|
// Get oplog entries
|
||||||
|
const oplog = await server.getOpLog({ limit: 50 });
|
||||||
|
console.log(oplog.entries);
|
||||||
|
// [{ seq: 1, op: 'insert', db: 'myapp', collection: 'users', document: {...}, previousDocument: null }, ...]
|
||||||
|
|
||||||
|
// Get aggregate stats
|
||||||
|
const stats = await server.getOpLogStats();
|
||||||
|
console.log(stats);
|
||||||
|
// { currentSeq: 42, totalEntries: 42, entriesByOp: { insert: 20, update: 15, delete: 7 } }
|
||||||
|
|
||||||
|
// Preview a revert (dry run)
|
||||||
|
const preview = await server.revertToSeq(30, true);
|
||||||
|
console.log(`Would undo ${preview.reverted} operations`);
|
||||||
|
|
||||||
|
// Execute the revert — undoes all operations after seq 30
|
||||||
|
const result = await server.revertToSeq(30, false);
|
||||||
|
console.log(`Reverted ${result.reverted} operations`);
|
||||||
|
|
||||||
|
// Browse collections programmatically
|
||||||
|
const collections = await server.getCollections();
|
||||||
|
const docs = await server.getDocuments('myapp', 'users', 50, 0);
|
||||||
|
```
|
||||||
|
|
||||||
|
### OpLog Entry Structure
|
||||||
|
|
||||||
|
Each entry contains:
|
||||||
|
|
||||||
|
| Field | Type | Description |
|
||||||
|
|---|---|---|
|
||||||
|
| `seq` | `number` | Monotonically increasing sequence number |
|
||||||
|
| `timestampMs` | `number` | Unix timestamp in milliseconds |
|
||||||
|
| `op` | `'insert' \| 'update' \| 'delete'` | Operation type |
|
||||||
|
| `db` | `string` | Database name |
|
||||||
|
| `collection` | `string` | Collection name |
|
||||||
|
| `documentId` | `string` | Document `_id` as hex string |
|
||||||
|
| `document` | `object \| null` | New document state (null for deletes) |
|
||||||
|
| `previousDocument` | `object \| null` | Previous document state (null for inserts) |
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## API Reference
|
## API Reference
|
||||||
|
|
||||||
### SmartdbServer
|
### SmartdbServer
|
||||||
|
|
||||||
The core server class. Speaks MongoDB wire protocol over TCP or Unix sockets.
|
The core server class. Manages the Rust database engine and exposes connection details.
|
||||||
|
|
||||||
#### Constructor Options (`ISmartdbServerOptions`)
|
#### Constructor Options (`ISmartdbServerOptions`)
|
||||||
|
|
||||||
@@ -117,21 +248,83 @@ const server = new SmartdbServer({
|
|||||||
persistPath: './data/snapshot.json',
|
persistPath: './data/snapshot.json',
|
||||||
persistIntervalMs: 30000, // Save every 30s
|
persistIntervalMs: 30000, // Save every 30s
|
||||||
});
|
});
|
||||||
|
|
||||||
|
// TLS transport for TCP mode
|
||||||
|
const tlsServer = new SmartdbServer({
|
||||||
|
port: 27017,
|
||||||
|
tls: {
|
||||||
|
enabled: true,
|
||||||
|
certPath: './certs/server.pem',
|
||||||
|
keyPath: './certs/server.key',
|
||||||
|
// caPath: './certs/client-ca.pem',
|
||||||
|
// requireClientCert: true, // Enables mTLS client certificate checks
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
// SCRAM-SHA-256 authentication
|
||||||
|
const secureServer = new SmartdbServer({
|
||||||
|
port: 27017,
|
||||||
|
auth: {
|
||||||
|
enabled: true,
|
||||||
|
usersPath: './data/smartdb-users.json', // Optional: persists derived SCRAM credentials
|
||||||
|
users: [
|
||||||
|
{
|
||||||
|
username: 'root',
|
||||||
|
password: 'change-me',
|
||||||
|
database: 'admin',
|
||||||
|
roles: ['root'],
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
});
|
||||||
|
```
|
||||||
|
|
||||||
|
When `auth.enabled` is true, protected commands require successful SCRAM-SHA-256 authentication through the official MongoDB driver:
|
||||||
|
|
||||||
|
```typescript
|
||||||
|
const client = new MongoClient('mongodb://root:change-me@127.0.0.1:27017/admin?authSource=admin', {
|
||||||
|
directConnection: true,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
```
|
||||||
|
|
||||||
|
TLS is available for TCP listeners. `getConnectionUri()` includes `?tls=true` when TLS is enabled; pass the trusted CA to the MongoDB driver with `tlsCAFile`, `ca`, or `secureContext`.
|
||||||
|
|
||||||
|
Authentication verifies SCRAM credentials, denies unauthenticated commands, and enforces command-level built-in roles for supported operations. `connectionStatus` reports the authenticated users and roles for the current socket.
|
||||||
|
|
||||||
|
Supported built-in role names are `root`, `read`, `readWrite`, `dbAdmin`, `userAdmin`, `clusterMonitor`, plus `readAnyDatabase`, `readWriteAnyDatabase`, `dbAdminAnyDatabase`, and `userAdminAnyDatabase`. When `usersPath` is set, SmartDB persists SCRAM credential material atomically and does not store plaintext passwords.
|
||||||
|
|
||||||
|
Single-node transactions are supported through official MongoDB driver sessions. Writes with `startTransaction` and `autocommit: false` are buffered per logical session, reads inside the transaction see the buffered overlay, `commitTransaction` applies the write set with conflict checks, and `abortTransaction` discards it.
|
||||||
|
|
||||||
|
Basic user management commands are available for authenticated users with `root` or `userAdmin` privileges:
|
||||||
|
|
||||||
|
```typescript
|
||||||
|
await client.db('admin').command({
|
||||||
|
createUser: 'reader',
|
||||||
|
pwd: 'readpass',
|
||||||
|
roles: [{ role: 'read', db: 'myapp' }],
|
||||||
|
});
|
||||||
|
|
||||||
|
await client.db('admin').command({ usersInfo: 'reader' });
|
||||||
```
|
```
|
||||||
|
|
||||||
#### Methods & Properties
|
#### Methods & Properties
|
||||||
|
|
||||||
| Method / Property | Type | Description |
|
| Method / Property | Type | Description |
|
||||||
|---|---|---|
|
|---|---|---|
|
||||||
| `start()` | `Promise<void>` | Start the server |
|
| `start()` | `Promise<void>` | Start the server (spawns Rust binary) |
|
||||||
| `stop()` | `Promise<void>` | Stop the server and clean up |
|
| `stop()` | `Promise<void>` | Stop the server and clean up |
|
||||||
| `getConnectionUri()` | `string` | Get the MongoDB connection URI |
|
| `getConnectionUri()` | `string` | Get the `mongodb://` connection URI |
|
||||||
| `running` | `boolean` | Whether the server is currently running |
|
| `running` | `boolean` | Whether the server is currently running |
|
||||||
| `port` | `number` | Bound port (TCP mode) |
|
| `port` | `number` | Configured port (TCP mode) |
|
||||||
| `host` | `string` | Bound host (TCP mode) |
|
| `host` | `string` | Configured host (TCP mode) |
|
||||||
| `socketPath` | `string` | Socket path (socket mode) |
|
| `socketPath` | `string \| undefined` | Socket path (socket mode) |
|
||||||
| `getUptime()` | `number` | Seconds since start |
|
| `getMetrics()` | `Promise<ISmartDbMetrics>` | Server metrics (db/collection counts, sessions, transactions, auth, uptime) |
|
||||||
| `getConnectionCount()` | `number` | Active client connections |
|
| `getOpLog(params?)` | `Promise<IOpLogResult>` | Query oplog entries with optional filters |
|
||||||
|
| `getOpLogStats()` | `Promise<IOpLogStats>` | Aggregate oplog statistics |
|
||||||
|
| `revertToSeq(seq, dryRun?)` | `Promise<IRevertResult>` | Revert to a specific oplog sequence |
|
||||||
|
| `getCollections(db?)` | `Promise<ICollectionInfo[]>` | List all collections with counts |
|
||||||
|
| `getDocuments(db, coll, limit?, skip?)` | `Promise<IDocumentsResult>` | Browse documents with pagination |
|
||||||
|
|
||||||
### LocalSmartDb
|
### LocalSmartDb
|
||||||
|
|
||||||
@@ -155,24 +348,45 @@ const db = new LocalSmartDb({
|
|||||||
| `start()` | `Promise<ILocalSmartDbConnectionInfo>` | Start and return connection info |
|
| `start()` | `Promise<ILocalSmartDbConnectionInfo>` | Start and return connection info |
|
||||||
| `stop()` | `Promise<void>` | Stop the server |
|
| `stop()` | `Promise<void>` | Stop the server |
|
||||||
| `getConnectionInfo()` | `ILocalSmartDbConnectionInfo` | Get current connection info |
|
| `getConnectionInfo()` | `ILocalSmartDbConnectionInfo` | Get current connection info |
|
||||||
| `getConnectionUri()` | `string` | Get the MongoDB URI |
|
| `getConnectionUri()` | `string` | Get the connection URI |
|
||||||
| `getServer()` | `SmartdbServer` | Access the underlying server |
|
| `getServer()` | `SmartdbServer` | Access the underlying server |
|
||||||
| `running` | `boolean` | Whether the server is running |
|
| `running` | `boolean` | Whether the server is running |
|
||||||
|
|
||||||
#### Connection Info (`ILocalSmartDbConnectionInfo`)
|
### SmartdbDebugServer
|
||||||
|
|
||||||
|
Web-based debug dashboard served via `@api.global/typedserver`. Import from the `debugserver` subpath:
|
||||||
|
|
||||||
```typescript
|
```typescript
|
||||||
interface ILocalSmartDbConnectionInfo {
|
import { SmartdbDebugServer } from '@push.rocks/smartdb/debugserver';
|
||||||
socketPath: string; // e.g., /tmp/smartdb-abc123.sock
|
|
||||||
connectionUri: string; // e.g., mongodb://%2Ftmp%2Fsmartdb-abc123.sock
|
const debugServer = new SmartdbDebugServer(server, { port: 4000 });
|
||||||
}
|
await debugServer.start();
|
||||||
|
// Dashboard at http://localhost:4000
|
||||||
|
|
||||||
|
await debugServer.stop();
|
||||||
|
```
|
||||||
|
|
||||||
|
The UI is bundled as base64-encoded content (via `@git.zone/tsbundle`) and served from memory — no static file directory needed.
|
||||||
|
|
||||||
|
### SmartdbDebugUi (Web Component)
|
||||||
|
|
||||||
|
For embedding the debug UI directly into your own web application, import the `<smartdb-debugui>` web component:
|
||||||
|
|
||||||
|
```typescript
|
||||||
|
import { SmartdbDebugUi } from '@push.rocks/smartdb/debugui';
|
||||||
|
|
||||||
|
// In your HTML/lit template:
|
||||||
|
// <smartdb-debugui .server=${mySmartdbServer}></smartdb-debugui>
|
||||||
|
//
|
||||||
|
// Or in HTTP mode (when served by SmartdbDebugServer):
|
||||||
|
// <smartdb-debugui apiBaseUrl=""></smartdb-debugui>
|
||||||
```
|
```
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Supported MongoDB Operations
|
## Supported Operations
|
||||||
|
|
||||||
SmartDB supports the core MongoDB operations through the wire protocol. Use the standard `mongodb` driver — these all work:
|
SmartDB supports the core operations through the wire protocol. Use the standard `mongodb` driver — these all work:
|
||||||
|
|
||||||
### CRUD
|
### CRUD
|
||||||
|
|
||||||
@@ -258,7 +472,7 @@ const results = await collection.aggregate([
|
|||||||
]).toArray();
|
]).toArray();
|
||||||
```
|
```
|
||||||
|
|
||||||
**Supported stages:** `$match`, `$project`, `$group`, `$sort`, `$limit`, `$skip`, `$unwind`, `$lookup`, `$addFields`, `$count`, `$facet`, `$replaceRoot`, `$set`, `$unset`
|
**Supported stages:** `$match`, `$project`, `$group`, `$sort`, `$limit`, `$skip`, `$unwind`, `$lookup`, `$addFields`, `$count`, `$facet`, `$replaceRoot`, `$set`, `$unionWith`, `$out`, `$merge`
|
||||||
|
|
||||||
**Group accumulators:** `$sum`, `$avg`, `$min`, `$max`, `$first`, `$last`, `$push`, `$addToSet`, `$count`
|
**Group accumulators:** `$sum`, `$avg`, `$min`, `$max`, `$first`, `$last`, `$push`, `$addToSet`, `$count`
|
||||||
|
|
||||||
@@ -273,6 +487,8 @@ await collection.dropIndex('email_1');
|
|||||||
await collection.dropIndexes(); // drop all except _id
|
await collection.dropIndexes(); // drop all except _id
|
||||||
```
|
```
|
||||||
|
|
||||||
|
> 🛡️ **Unique indexes are enforced at the engine level.** Duplicate values are rejected with a `DuplicateKey` error (code 11000) *before* the document is written to disk — on `insertOne`, `updateOne`, `findAndModify`, and upserts. Index definitions are persisted to `indexes.json` and automatically restored on restart.
|
||||||
|
|
||||||
### Database & Admin
|
### Database & Admin
|
||||||
|
|
||||||
```typescript
|
```typescript
|
||||||
@@ -316,148 +532,69 @@ const names = await collection.distinct('name');
|
|||||||
| **CRUD** | `find`, `insert`, `update`, `delete`, `findAndModify`, `getMore`, `killCursors` |
|
| **CRUD** | `find`, `insert`, `update`, `delete`, `findAndModify`, `getMore`, `killCursors` |
|
||||||
| **Aggregation** | `aggregate`, `count`, `distinct` |
|
| **Aggregation** | `aggregate`, `count`, `distinct` |
|
||||||
| **Indexes** | `createIndexes`, `dropIndexes`, `listIndexes` |
|
| **Indexes** | `createIndexes`, `dropIndexes`, `listIndexes` |
|
||||||
| **Transactions** | `startTransaction`, `commitTransaction`, `abortTransaction` |
|
| **Sessions** | `startSession`, `endSessions` |
|
||||||
| **Sessions** | `startSession`, `endSessions`, `refreshSessions` |
|
| **Transactions** | `startTransaction`, `commitTransaction`, `abortTransaction` through driver sessions |
|
||||||
| **Admin** | `ping`, `listDatabases`, `listCollections`, `drop`, `dropDatabase`, `create`, `serverStatus`, `buildInfo`, `dbStats`, `collStats`, `connectionStatus`, `currentOp`, `collMod`, `renameCollection` |
|
| **Admin** | `ping`, `listDatabases`, `listCollections`, `drop`, `dropDatabase`, `create`, `serverStatus`, `buildInfo`, `dbStats`, `collStats`, `connectionStatus`, `currentOp`, `renameCollection` |
|
||||||
|
|
||||||
Compatible with MongoDB wire protocol versions 0–21 (MongoDB 3.6 through 7.0 drivers).
|
Compatible with wire protocol versions 0–21 (driver versions 3.6 through 7.0).
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Architecture
|
## Rust Crate Architecture 🦀
|
||||||
|
|
||||||
```
|
The Rust engine is organized as a Cargo workspace with 9 focused crates:
|
||||||
┌─────────────────────────────────────────────────────────────┐
|
|
||||||
│ Official MongoDB Driver │
|
|
||||||
│ (mongodb npm) │
|
|
||||||
└─────────────────────────┬───────────────────────────────────┘
|
|
||||||
│ TCP / Unix Socket + OP_MSG / BSON
|
|
||||||
▼
|
|
||||||
┌─────────────────────────────────────────────────────────────┐
|
|
||||||
│ SmartdbServer │
|
|
||||||
│ ┌──────────────┐ ┌──────────────┐ ┌──────────────────┐ │
|
|
||||||
│ │ WireProtocol │→ │CommandRouter │→ │ Handlers │ │
|
|
||||||
│ │ (OP_MSG) │ │ │ │ (Find, Insert..) │ │
|
|
||||||
│ └──────────────┘ └──────────────┘ └──────────────────┘ │
|
|
||||||
└─────────────────────────┬───────────────────────────────────┘
|
|
||||||
│
|
|
||||||
▼
|
|
||||||
┌─────────────────────────────────────────────────────────────┐
|
|
||||||
│ Engines │
|
|
||||||
│ ┌─────────┐ ┌────────┐ ┌───────────┐ ┌───────┐ ┌───────┐ │
|
|
||||||
│ │ Query │ │ Update │ │Aggregation│ │ Index │ │Session│ │
|
|
||||||
│ │ Planner │ │ Engine │ │ Engine │ │Engine │ │Engine │ │
|
|
||||||
│ └─────────┘ └────────┘ └───────────┘ └───────┘ └───────┘ │
|
|
||||||
│ ┌──────────────────────┐ │
|
|
||||||
│ │ Transaction Engine │ │
|
|
||||||
│ └──────────────────────┘ │
|
|
||||||
└─────────────────────────┬───────────────────────────────────┘
|
|
||||||
│
|
|
||||||
▼
|
|
||||||
┌─────────────────────────────────────────────────────────────┐
|
|
||||||
│ Storage Layer │
|
|
||||||
│ ┌──────────────────┐ ┌──────────────────┐ ┌──────────┐ │
|
|
||||||
│ │ MemoryStorage │ │ FileStorage │ │ WAL │ │
|
|
||||||
│ │ │ │ (+ Checksums) │ │ │ │
|
|
||||||
│ └──────────────────┘ └──────────────────┘ └──────────┘ │
|
|
||||||
└─────────────────────────────────────────────────────────────┘
|
|
||||||
```
|
|
||||||
|
|
||||||
### Key Components
|
| Crate | Purpose |
|
||||||
|
|
||||||
| Component | What It Does |
|
|
||||||
|---|---|
|
|---|---|
|
||||||
| **WireProtocol** | Parses/encodes MongoDB OP_MSG binary frames |
|
| `rustdb` | Binary entry point: TCP/Unix listener, management IPC, CLI |
|
||||||
| **CommandRouter** | Routes parsed commands to the right handler |
|
| `rustdb-config` | Server configuration types (serde, camelCase JSON) |
|
||||||
| **QueryPlanner** | Picks COLLSCAN vs IXSCAN based on available indexes |
|
| `rustdb-wire` | Wire protocol parser/encoder (OP_MSG, OP_QUERY, OP_REPLY) |
|
||||||
| **QueryEngine** | Filter matching powered by [mingo](https://github.com/kofrasa/mingo) |
|
| `rustdb-query` | Query matcher, update engine, aggregation, sort, projection |
|
||||||
| **UpdateEngine** | Processes `$set`, `$inc`, `$push`, and all update operators |
|
| `rustdb-storage` | Storage backends (memory, file), OpLog with point-in-time replay |
|
||||||
| **AggregationEngine** | Runs aggregation pipelines via mingo |
|
| `rustdb-index` | B-tree/hash indexes, query planner (IXSCAN/COLLSCAN) |
|
||||||
| **IndexEngine** | B-tree (range) and hash (equality) indexes |
|
| `rustdb-txn` | Transaction + session management with snapshot isolation |
|
||||||
| **TransactionEngine** | ACID transactions with snapshot isolation |
|
| `rustdb-auth` | SCRAM-SHA-256 credential handling, user metadata persistence, RBAC checks |
|
||||||
| **SessionEngine** | Client session tracking with automatic timeouts |
|
| `rustdb-commands` | 40+ command handlers wiring everything together |
|
||||||
| **WAL** | Write-ahead logging with CRC32 checksums for crash recovery |
|
|
||||||
|
Cross-compiled for `linux_amd64` and `linux_arm64` via [@git.zone/tsrust](https://www.npmjs.com/package/@git.zone/tsrust).
|
||||||
|
|
||||||
|
### Storage Engine Reliability 🔒
|
||||||
|
|
||||||
|
The Bitcask-style file storage engine includes several reliability features:
|
||||||
|
|
||||||
|
- **Write-ahead log (WAL)** — every write is logged before being applied, with crash recovery on restart
|
||||||
|
- **CRC32 checksums** — every record is integrity-checked on read
|
||||||
|
- **Automatic compaction** — dead records are reclaimed when they exceed 50% of file size, runs on startup and after every write
|
||||||
|
- **Hint file staleness detection** — the hint file records the data file size at write time; if data.rdb changed since (e.g. crash after a delete), the engine falls back to a full scan to ensure tombstones are not lost
|
||||||
|
- **Torn-tail repair** — startup scans `data.rdb` to the last valid record, truncates invalid trailing bytes, and preserves all verified records after interrupted writes
|
||||||
|
- **Stale socket cleanup** — orphaned `/tmp/smartdb-*.sock` files from crashed instances are automatically cleaned up on startup
|
||||||
|
|
||||||
|
### Data Integrity CLI 🔍
|
||||||
|
|
||||||
|
The Rust binary includes an offline integrity checker:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Check all collections in a data directory
|
||||||
|
./dist_rust/rustdb_linux_amd64 --validate-data /path/to/data
|
||||||
|
|
||||||
|
# Output:
|
||||||
|
# === SmartDB Data Integrity Report ===
|
||||||
|
#
|
||||||
|
# Database: mydb
|
||||||
|
# Collection: users
|
||||||
|
# Header: OK
|
||||||
|
# Records: 1,234 (1,200 live, 34 tombstones)
|
||||||
|
# Data size: 2.1 MB
|
||||||
|
# Duplicates: 0
|
||||||
|
# CRC errors: 0
|
||||||
|
# Hint file: OK
|
||||||
|
```
|
||||||
|
|
||||||
|
Checks file headers, record CRC32 checksums, duplicate `_id` entries, and hint file consistency. Exit code 1 if any errors are found.
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Advanced Usage
|
## Testing Example
|
||||||
|
|
||||||
### Storage Adapters
|
|
||||||
|
|
||||||
```typescript
|
|
||||||
import { SmartdbServer } from '@push.rocks/smartdb';
|
|
||||||
|
|
||||||
// In-memory (default) — fast, data lost on stop
|
|
||||||
const server = new SmartdbServer({ storage: 'memory' });
|
|
||||||
|
|
||||||
// In-memory with periodic persistence
|
|
||||||
const server = new SmartdbServer({
|
|
||||||
storage: 'memory',
|
|
||||||
persistPath: './data/snapshot.json',
|
|
||||||
persistIntervalMs: 30000,
|
|
||||||
});
|
|
||||||
|
|
||||||
// File-based — persistent storage with CRC32 checksums
|
|
||||||
const server = new SmartdbServer({
|
|
||||||
storage: 'file',
|
|
||||||
storagePath: './data/smartdb',
|
|
||||||
});
|
|
||||||
```
|
|
||||||
|
|
||||||
### Query Planner (Debugging)
|
|
||||||
|
|
||||||
```typescript
|
|
||||||
import { QueryPlanner, IndexEngine, MemoryStorageAdapter } from '@push.rocks/smartdb';
|
|
||||||
|
|
||||||
const storage = new MemoryStorageAdapter();
|
|
||||||
await storage.initialize();
|
|
||||||
const indexEngine = new IndexEngine('mydb', 'mycoll', storage);
|
|
||||||
const planner = new QueryPlanner(indexEngine);
|
|
||||||
|
|
||||||
const plan = await planner.plan({ age: { $gte: 18 } });
|
|
||||||
console.log(plan);
|
|
||||||
// { type: 'IXSCAN_RANGE', indexName: 'age_1', selectivity: 0.3, usesRange: true, ... }
|
|
||||||
|
|
||||||
const explain = await planner.explain({ age: 18 });
|
|
||||||
// Returns winning plan, rejected plans, and detailed analysis
|
|
||||||
```
|
|
||||||
|
|
||||||
### Data Integrity Checksums
|
|
||||||
|
|
||||||
```typescript
|
|
||||||
import { calculateDocumentChecksum, addChecksum, verifyChecksum, removeChecksum } from '@push.rocks/smartdb';
|
|
||||||
|
|
||||||
const doc = { name: 'Alice', age: 30 };
|
|
||||||
|
|
||||||
const protected = addChecksum(doc); // Adds _checksum field
|
|
||||||
const valid = verifyChecksum(protected); // true
|
|
||||||
protected.age = 31; // Tamper!
|
|
||||||
const still = verifyChecksum(protected); // false
|
|
||||||
const clean = removeChecksum(protected); // Removes _checksum
|
|
||||||
```
|
|
||||||
|
|
||||||
### Write-Ahead Logging
|
|
||||||
|
|
||||||
```typescript
|
|
||||||
import { WAL } from '@push.rocks/smartdb';
|
|
||||||
|
|
||||||
const wal = new WAL('./data/wal.log', { checkpointInterval: 100 });
|
|
||||||
await wal.initialize();
|
|
||||||
|
|
||||||
// Entries include: LSN, timestamp, operation, BSON data, CRC32 checksum
|
|
||||||
const lsn = await wal.logInsert('mydb', 'users', doc);
|
|
||||||
const entries = wal.getEntriesAfter(lastCheckpoint);
|
|
||||||
const recovered = wal.recoverDocument(entry);
|
|
||||||
|
|
||||||
await wal.checkpoint();
|
|
||||||
await wal.close();
|
|
||||||
```
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Testing Examples
|
|
||||||
|
|
||||||
### Unit Tests with @git.zone/tstest
|
|
||||||
|
|
||||||
```typescript
|
```typescript
|
||||||
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
@@ -481,6 +618,12 @@ tap.test('should insert and find', async () => {
|
|||||||
expect(item?.price).toEqual(9.99);
|
expect(item?.price).toEqual(9.99);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
tap.test('should track changes in oplog', async () => {
|
||||||
|
const oplog = await server.getOpLog();
|
||||||
|
expect(oplog.entries.length).toBeGreaterThan(0);
|
||||||
|
expect(oplog.entries[0].op).toEqual('insert');
|
||||||
|
});
|
||||||
|
|
||||||
tap.test('teardown', async () => {
|
tap.test('teardown', async () => {
|
||||||
await client.close();
|
await client.close();
|
||||||
await server.stop();
|
await server.stop();
|
||||||
@@ -489,27 +632,11 @@ tap.test('teardown', async () => {
|
|||||||
export default tap.start();
|
export default tap.start();
|
||||||
```
|
```
|
||||||
|
|
||||||
### With LocalSmartDb (Persistent Tests)
|
---
|
||||||
|
|
||||||
```typescript
|
|
||||||
import { LocalSmartDb } from '@push.rocks/smartdb';
|
|
||||||
import { MongoClient } from 'mongodb';
|
|
||||||
|
|
||||||
const db = new LocalSmartDb({ folderPath: './test-data' });
|
|
||||||
const { connectionUri } = await db.start();
|
|
||||||
|
|
||||||
const client = new MongoClient(connectionUri, { directConnection: true });
|
|
||||||
await client.connect();
|
|
||||||
|
|
||||||
// Tests here — data persists between test runs!
|
|
||||||
|
|
||||||
await client.close();
|
|
||||||
await db.stop();
|
|
||||||
```
|
|
||||||
|
|
||||||
## License and Legal Information
|
## License and Legal Information
|
||||||
|
|
||||||
This repository contains open-source code licensed under the MIT License. A copy of the license can be found in the [LICENSE](./LICENSE) file.
|
This repository contains open-source code licensed under the MIT License. A copy of the license can be found in the [license](./license) file.
|
||||||
|
|
||||||
**Please note:** The MIT License does not grant permission to use the trade names, trademarks, service marks, or product names of the project, except as required for reasonable and customary use in describing the origin of the work and reproducing the content of the NOTICE file.
|
**Please note:** The MIT License does not grant permission to use the trade names, trademarks, service marks, or product names of the project, except as required for reasonable and customary use in describing the origin of the work and reproducing the content of the NOTICE file.
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,2 @@
|
|||||||
|
[target.aarch64-unknown-linux-gnu]
|
||||||
|
linker = "aarch64-linux-gnu-gcc"
|
||||||
Generated
+1756
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,93 @@
|
|||||||
|
[workspace]
|
||||||
|
resolver = "2"
|
||||||
|
members = [
|
||||||
|
"crates/rustdb",
|
||||||
|
"crates/rustdb-config",
|
||||||
|
"crates/rustdb-wire",
|
||||||
|
"crates/rustdb-query",
|
||||||
|
"crates/rustdb-storage",
|
||||||
|
"crates/rustdb-index",
|
||||||
|
"crates/rustdb-txn",
|
||||||
|
"crates/rustdb-auth",
|
||||||
|
"crates/rustdb-commands",
|
||||||
|
]
|
||||||
|
|
||||||
|
[workspace.package]
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2021"
|
||||||
|
license = "MIT"
|
||||||
|
authors = ["Lossless GmbH <hello@lossless.com>"]
|
||||||
|
|
||||||
|
[workspace.dependencies]
|
||||||
|
# Async runtime
|
||||||
|
tokio = { version = "1", features = ["full"] }
|
||||||
|
|
||||||
|
# Serialization
|
||||||
|
serde = { version = "1", features = ["derive"] }
|
||||||
|
serde_json = "1"
|
||||||
|
|
||||||
|
# BSON serialization (bson crate)
|
||||||
|
bson = "2"
|
||||||
|
|
||||||
|
# Binary buffer manipulation
|
||||||
|
bytes = "1"
|
||||||
|
|
||||||
|
# CLI
|
||||||
|
clap = { version = "4", features = ["derive"] }
|
||||||
|
|
||||||
|
# Structured logging
|
||||||
|
tracing = "0.1"
|
||||||
|
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||||
|
|
||||||
|
# Error handling
|
||||||
|
thiserror = "2"
|
||||||
|
anyhow = "1"
|
||||||
|
|
||||||
|
# Lock-free atomics
|
||||||
|
arc-swap = "1"
|
||||||
|
|
||||||
|
# Concurrent maps
|
||||||
|
dashmap = "6"
|
||||||
|
|
||||||
|
# Cancellation / utility
|
||||||
|
tokio-util = { version = "0.7", features = ["codec"] }
|
||||||
|
|
||||||
|
# TLS transport
|
||||||
|
tokio-rustls = { version = "0.26", default-features = false, features = ["ring", "tls12"] }
|
||||||
|
rustls-pemfile = "2"
|
||||||
|
|
||||||
|
# mimalloc allocator
|
||||||
|
mimalloc = "0.1"
|
||||||
|
|
||||||
|
# CRC32 checksums
|
||||||
|
crc32fast = "1"
|
||||||
|
|
||||||
|
# Regex for $regex operator
|
||||||
|
regex = "1"
|
||||||
|
|
||||||
|
# Auth crypto
|
||||||
|
base64 = "0.22"
|
||||||
|
hmac = "0.12"
|
||||||
|
pbkdf2 = { version = "0.12", features = ["hmac"] }
|
||||||
|
rand = "0.8"
|
||||||
|
sha2 = "0.10"
|
||||||
|
subtle = "2"
|
||||||
|
|
||||||
|
# UUID for sessions
|
||||||
|
uuid = { version = "1", features = ["v4", "serde"] }
|
||||||
|
|
||||||
|
# Async traits
|
||||||
|
async-trait = "0.1"
|
||||||
|
|
||||||
|
# Test utilities
|
||||||
|
tempfile = "3"
|
||||||
|
|
||||||
|
# Internal crates
|
||||||
|
rustdb-config = { path = "crates/rustdb-config" }
|
||||||
|
rustdb-wire = { path = "crates/rustdb-wire" }
|
||||||
|
rustdb-query = { path = "crates/rustdb-query" }
|
||||||
|
rustdb-storage = { path = "crates/rustdb-storage" }
|
||||||
|
rustdb-index = { path = "crates/rustdb-index" }
|
||||||
|
rustdb-txn = { path = "crates/rustdb-txn" }
|
||||||
|
rustdb-auth = { path = "crates/rustdb-auth" }
|
||||||
|
rustdb-commands = { path = "crates/rustdb-commands" }
|
||||||
@@ -0,0 +1,20 @@
|
|||||||
|
[package]
|
||||||
|
name = "rustdb-auth"
|
||||||
|
version.workspace = true
|
||||||
|
edition.workspace = true
|
||||||
|
license.workspace = true
|
||||||
|
authors.workspace = true
|
||||||
|
description = "Authentication primitives for RustDb"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
base64 = { workspace = true }
|
||||||
|
bson = { workspace = true }
|
||||||
|
hmac = { workspace = true }
|
||||||
|
pbkdf2 = { workspace = true }
|
||||||
|
rand = { workspace = true }
|
||||||
|
rustdb-config = { workspace = true }
|
||||||
|
serde = { workspace = true }
|
||||||
|
serde_json = { workspace = true }
|
||||||
|
sha2 = { workspace = true }
|
||||||
|
subtle = { workspace = true }
|
||||||
|
thiserror = { workspace = true }
|
||||||
@@ -0,0 +1,593 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
|
use std::io::Write;
|
||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
use std::sync::RwLock;
|
||||||
|
|
||||||
|
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
|
||||||
|
use hmac::{Hmac, Mac};
|
||||||
|
use pbkdf2::pbkdf2_hmac;
|
||||||
|
use rand::{rngs::OsRng, RngCore};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use sha2::{Digest, Sha256};
|
||||||
|
use subtle::ConstantTimeEq;
|
||||||
|
|
||||||
|
use rustdb_config::{AuthOptions, AuthUserOptions};
|
||||||
|
|
||||||
|
type HmacSha256 = Hmac<Sha256>;
|
||||||
|
|
||||||
|
const SCRAM_SHA_256: &str = "SCRAM-SHA-256";
|
||||||
|
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum AuthError {
|
||||||
|
#[error("authentication is disabled")]
|
||||||
|
Disabled,
|
||||||
|
#[error("unsupported authentication mechanism: {0}")]
|
||||||
|
UnsupportedMechanism(String),
|
||||||
|
#[error("invalid SCRAM payload: {0}")]
|
||||||
|
InvalidPayload(String),
|
||||||
|
#[error("authentication failed")]
|
||||||
|
AuthenticationFailed,
|
||||||
|
#[error("unknown SASL conversation")]
|
||||||
|
UnknownConversation,
|
||||||
|
#[error("user already exists: {0}")]
|
||||||
|
UserAlreadyExists(String),
|
||||||
|
#[error("user not found: {0}")]
|
||||||
|
UserNotFound(String),
|
||||||
|
#[error("auth metadata persistence failed: {0}")]
|
||||||
|
Persistence(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub enum AuthAction {
|
||||||
|
Read,
|
||||||
|
Write,
|
||||||
|
DbAdmin,
|
||||||
|
UserAdmin,
|
||||||
|
ClusterMonitor,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct AuthenticatedUser {
|
||||||
|
pub username: String,
|
||||||
|
pub database: String,
|
||||||
|
pub roles: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
struct ScramCredential {
|
||||||
|
salt: Vec<u8>,
|
||||||
|
iterations: u32,
|
||||||
|
stored_key: Vec<u8>,
|
||||||
|
server_key: Vec<u8>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
struct AuthUser {
|
||||||
|
username: String,
|
||||||
|
database: String,
|
||||||
|
roles: Vec<String>,
|
||||||
|
scram_sha256: ScramCredential,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||||
|
struct PersistedAuthState {
|
||||||
|
users: Vec<AuthUser>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct ScramConversation {
|
||||||
|
user: AuthenticatedUser,
|
||||||
|
client_first_bare: String,
|
||||||
|
server_first: String,
|
||||||
|
nonce: String,
|
||||||
|
stored_key: Vec<u8>,
|
||||||
|
server_key: Vec<u8>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct ScramStartResult {
|
||||||
|
pub payload: Vec<u8>,
|
||||||
|
pub conversation: ScramConversation,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct ScramContinueResult {
|
||||||
|
pub payload: Vec<u8>,
|
||||||
|
pub user: AuthenticatedUser,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct AuthEngine {
|
||||||
|
enabled: bool,
|
||||||
|
users: RwLock<HashMap<String, AuthUser>>,
|
||||||
|
users_path: Option<PathBuf>,
|
||||||
|
scram_iterations: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AuthEngine {
|
||||||
|
pub fn from_options(options: &AuthOptions) -> Result<Self, AuthError> {
|
||||||
|
let users_path = options.users_path.as_ref().map(PathBuf::from);
|
||||||
|
let mut users = if let Some(ref path) = users_path {
|
||||||
|
load_users(path)?
|
||||||
|
} else {
|
||||||
|
HashMap::new()
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut changed = false;
|
||||||
|
for user_options in &options.users {
|
||||||
|
let key = user_key(&user_options.database, &user_options.username);
|
||||||
|
if !users.contains_key(&key) {
|
||||||
|
let user = AuthUser::from_options(user_options, options.scram_iterations);
|
||||||
|
users.insert(key, user);
|
||||||
|
changed = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if changed {
|
||||||
|
if let Some(ref path) = users_path {
|
||||||
|
persist_users(path, &users)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(Self {
|
||||||
|
enabled: options.enabled,
|
||||||
|
users: RwLock::new(users),
|
||||||
|
users_path,
|
||||||
|
scram_iterations: options.scram_iterations,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn disabled() -> Self {
|
||||||
|
Self {
|
||||||
|
enabled: false,
|
||||||
|
users: RwLock::new(HashMap::new()),
|
||||||
|
users_path: None,
|
||||||
|
scram_iterations: 15000,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn enabled(&self) -> bool {
|
||||||
|
self.enabled
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn user_count(&self) -> usize {
|
||||||
|
self.users
|
||||||
|
.read()
|
||||||
|
.unwrap_or_else(|poisoned| poisoned.into_inner())
|
||||||
|
.len()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn supported_mechanisms(&self, namespace_user: &str) -> Vec<String> {
|
||||||
|
let Some((database, username)) = namespace_user.split_once('.') else {
|
||||||
|
return Vec::new();
|
||||||
|
};
|
||||||
|
let users = self.users.read().unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||||
|
if users.contains_key(&user_key(database, username)) {
|
||||||
|
vec![SCRAM_SHA_256.to_string()]
|
||||||
|
} else {
|
||||||
|
Vec::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn is_authorized(
|
||||||
|
&self,
|
||||||
|
authenticated_users: &[AuthenticatedUser],
|
||||||
|
target_db: &str,
|
||||||
|
action: AuthAction,
|
||||||
|
) -> bool {
|
||||||
|
authenticated_users
|
||||||
|
.iter()
|
||||||
|
.any(|user| user.roles.iter().any(|role| role_allows(role, user, target_db, action)))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn create_user(
|
||||||
|
&self,
|
||||||
|
database: &str,
|
||||||
|
username: &str,
|
||||||
|
password: &str,
|
||||||
|
roles: Vec<String>,
|
||||||
|
) -> Result<(), AuthError> {
|
||||||
|
let key = user_key(database, username);
|
||||||
|
let mut users = self.users.write().unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||||
|
if users.contains_key(&key) {
|
||||||
|
return Err(AuthError::UserAlreadyExists(format!("{database}.{username}")));
|
||||||
|
}
|
||||||
|
let options = AuthUserOptions {
|
||||||
|
username: username.to_string(),
|
||||||
|
password: password.to_string(),
|
||||||
|
database: database.to_string(),
|
||||||
|
roles,
|
||||||
|
};
|
||||||
|
users.insert(key, AuthUser::from_options(&options, self.scram_iterations));
|
||||||
|
self.persist_locked(&users)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn drop_user(&self, database: &str, username: &str) -> Result<(), AuthError> {
|
||||||
|
let key = user_key(database, username);
|
||||||
|
let mut users = self.users.write().unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||||
|
if users.remove(&key).is_none() {
|
||||||
|
return Err(AuthError::UserNotFound(format!("{database}.{username}")));
|
||||||
|
}
|
||||||
|
self.persist_locked(&users)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn update_user(
|
||||||
|
&self,
|
||||||
|
database: &str,
|
||||||
|
username: &str,
|
||||||
|
password: Option<&str>,
|
||||||
|
roles: Option<Vec<String>>,
|
||||||
|
) -> Result<(), AuthError> {
|
||||||
|
let key = user_key(database, username);
|
||||||
|
let mut users = self.users.write().unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||||
|
let user = users
|
||||||
|
.get_mut(&key)
|
||||||
|
.ok_or_else(|| AuthError::UserNotFound(format!("{database}.{username}")))?;
|
||||||
|
if let Some(new_roles) = roles {
|
||||||
|
user.roles = new_roles;
|
||||||
|
}
|
||||||
|
if let Some(new_password) = password {
|
||||||
|
let options = AuthUserOptions {
|
||||||
|
username: username.to_string(),
|
||||||
|
password: new_password.to_string(),
|
||||||
|
database: database.to_string(),
|
||||||
|
roles: user.roles.clone(),
|
||||||
|
};
|
||||||
|
user.scram_sha256 = AuthUser::from_options(&options, self.scram_iterations).scram_sha256;
|
||||||
|
}
|
||||||
|
self.persist_locked(&users)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn grant_roles(
|
||||||
|
&self,
|
||||||
|
database: &str,
|
||||||
|
username: &str,
|
||||||
|
roles: Vec<String>,
|
||||||
|
) -> Result<(), AuthError> {
|
||||||
|
let key = user_key(database, username);
|
||||||
|
let mut users = self.users.write().unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||||
|
let user = users
|
||||||
|
.get_mut(&key)
|
||||||
|
.ok_or_else(|| AuthError::UserNotFound(format!("{database}.{username}")))?;
|
||||||
|
for role in roles {
|
||||||
|
if !user.roles.contains(&role) {
|
||||||
|
user.roles.push(role);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
self.persist_locked(&users)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn revoke_roles(
|
||||||
|
&self,
|
||||||
|
database: &str,
|
||||||
|
username: &str,
|
||||||
|
roles: Vec<String>,
|
||||||
|
) -> Result<(), AuthError> {
|
||||||
|
let key = user_key(database, username);
|
||||||
|
let mut users = self.users.write().unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||||
|
let user = users
|
||||||
|
.get_mut(&key)
|
||||||
|
.ok_or_else(|| AuthError::UserNotFound(format!("{database}.{username}")))?;
|
||||||
|
user.roles.retain(|role| !roles.contains(role));
|
||||||
|
self.persist_locked(&users)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn users_info(&self, database: &str, username: Option<&str>) -> Vec<AuthenticatedUser> {
|
||||||
|
let users = self.users.read().unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||||
|
users
|
||||||
|
.values()
|
||||||
|
.filter(|user| user.database == database)
|
||||||
|
.filter(|user| username.map(|name| user.username == name).unwrap_or(true))
|
||||||
|
.map(AuthUser::to_authenticated_user)
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn list_users(&self) -> Vec<AuthenticatedUser> {
|
||||||
|
let users = self.users.read().unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||||
|
let mut result: Vec<AuthenticatedUser> = users
|
||||||
|
.values()
|
||||||
|
.map(AuthUser::to_authenticated_user)
|
||||||
|
.collect();
|
||||||
|
result.sort_by(|a, b| a.database.cmp(&b.database).then(a.username.cmp(&b.username)));
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn drop_users_for_database(&self, database: &str) -> Result<usize, AuthError> {
|
||||||
|
let mut users = self.users.write().unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||||
|
let before = users.len();
|
||||||
|
users.retain(|_, user| user.database != database);
|
||||||
|
let dropped = before.saturating_sub(users.len());
|
||||||
|
if dropped > 0 {
|
||||||
|
self.persist_locked(&users)?;
|
||||||
|
}
|
||||||
|
Ok(dropped)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn start_scram_sha256(
|
||||||
|
&self,
|
||||||
|
database: &str,
|
||||||
|
payload: &[u8],
|
||||||
|
) -> Result<ScramStartResult, AuthError> {
|
||||||
|
if !self.enabled {
|
||||||
|
return Err(AuthError::Disabled);
|
||||||
|
}
|
||||||
|
|
||||||
|
let message = std::str::from_utf8(payload)
|
||||||
|
.map_err(|_| AuthError::InvalidPayload("payload is not valid UTF-8".to_string()))?;
|
||||||
|
let client_first_bare = message
|
||||||
|
.strip_prefix("n,,")
|
||||||
|
.ok_or_else(|| AuthError::InvalidPayload("expected SCRAM gs2 header 'n,,'".to_string()))?;
|
||||||
|
let attrs = parse_scram_attrs(client_first_bare);
|
||||||
|
let raw_username = attrs
|
||||||
|
.get("n")
|
||||||
|
.ok_or_else(|| AuthError::InvalidPayload("missing username".to_string()))?;
|
||||||
|
let username = decode_scram_name(raw_username);
|
||||||
|
let client_nonce = attrs
|
||||||
|
.get("r")
|
||||||
|
.ok_or_else(|| AuthError::InvalidPayload("missing client nonce".to_string()))?;
|
||||||
|
|
||||||
|
let users = self.users.read().unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||||
|
let user = users
|
||||||
|
.get(&user_key(database, &username))
|
||||||
|
.ok_or(AuthError::AuthenticationFailed)?;
|
||||||
|
|
||||||
|
let nonce = format!("{}{}", client_nonce, secure_base64(18));
|
||||||
|
let server_first = format!(
|
||||||
|
"r={},s={},i={}",
|
||||||
|
nonce,
|
||||||
|
BASE64_STANDARD.encode(&user.scram_sha256.salt),
|
||||||
|
user.scram_sha256.iterations,
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(ScramStartResult {
|
||||||
|
payload: server_first.as_bytes().to_vec(),
|
||||||
|
conversation: ScramConversation {
|
||||||
|
user: user.to_authenticated_user(),
|
||||||
|
client_first_bare: client_first_bare.to_string(),
|
||||||
|
server_first: server_first.clone(),
|
||||||
|
nonce,
|
||||||
|
stored_key: user.scram_sha256.stored_key.clone(),
|
||||||
|
server_key: user.scram_sha256.server_key.clone(),
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn continue_scram_sha256(
|
||||||
|
&self,
|
||||||
|
conversation: ScramConversation,
|
||||||
|
payload: &[u8],
|
||||||
|
) -> Result<ScramContinueResult, AuthError> {
|
||||||
|
let message = std::str::from_utf8(payload)
|
||||||
|
.map_err(|_| AuthError::InvalidPayload("payload is not valid UTF-8".to_string()))?;
|
||||||
|
let proof_marker = ",p=";
|
||||||
|
let proof_pos = message
|
||||||
|
.rfind(proof_marker)
|
||||||
|
.ok_or_else(|| AuthError::InvalidPayload("missing client proof".to_string()))?;
|
||||||
|
let client_final_without_proof = &message[..proof_pos];
|
||||||
|
let proof_b64 = &message[proof_pos + proof_marker.len()..];
|
||||||
|
let attrs = parse_scram_attrs(client_final_without_proof);
|
||||||
|
let nonce = attrs
|
||||||
|
.get("r")
|
||||||
|
.ok_or_else(|| AuthError::InvalidPayload("missing nonce".to_string()))?;
|
||||||
|
if nonce != &conversation.nonce {
|
||||||
|
return Err(AuthError::AuthenticationFailed);
|
||||||
|
}
|
||||||
|
|
||||||
|
let client_proof = BASE64_STANDARD
|
||||||
|
.decode(proof_b64.as_bytes())
|
||||||
|
.map_err(|_| AuthError::InvalidPayload("invalid client proof encoding".to_string()))?;
|
||||||
|
if client_proof.len() != 32 || conversation.stored_key.len() != 32 {
|
||||||
|
return Err(AuthError::AuthenticationFailed);
|
||||||
|
}
|
||||||
|
|
||||||
|
let auth_message = format!(
|
||||||
|
"{},{},{}",
|
||||||
|
conversation.client_first_bare,
|
||||||
|
conversation.server_first,
|
||||||
|
client_final_without_proof,
|
||||||
|
);
|
||||||
|
let client_signature = hmac_sha256(&conversation.stored_key, auth_message.as_bytes());
|
||||||
|
let client_key: Vec<u8> = client_proof
|
||||||
|
.iter()
|
||||||
|
.zip(client_signature.iter())
|
||||||
|
.map(|(proof_byte, signature_byte)| proof_byte ^ signature_byte)
|
||||||
|
.collect();
|
||||||
|
let computed_stored_key = Sha256::digest(&client_key).to_vec();
|
||||||
|
|
||||||
|
if computed_stored_key.ct_eq(&conversation.stored_key).unwrap_u8() != 1 {
|
||||||
|
return Err(AuthError::AuthenticationFailed);
|
||||||
|
}
|
||||||
|
|
||||||
|
let server_signature = hmac_sha256(&conversation.server_key, auth_message.as_bytes());
|
||||||
|
let server_final = format!("v={}", BASE64_STANDARD.encode(server_signature));
|
||||||
|
|
||||||
|
Ok(ScramContinueResult {
|
||||||
|
payload: server_final.as_bytes().to_vec(),
|
||||||
|
user: conversation.user,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn persist_locked(&self, users: &HashMap<String, AuthUser>) -> Result<(), AuthError> {
|
||||||
|
if let Some(ref path) = self.users_path {
|
||||||
|
persist_users(path, users)?;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for AuthEngine {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::disabled()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AuthUser {
|
||||||
|
fn from_options(options: &AuthUserOptions, iterations: u32) -> Self {
|
||||||
|
let salt = secure_random(24);
|
||||||
|
let salted_password = salted_password(options.password.as_bytes(), &salt, iterations);
|
||||||
|
let client_key = hmac_sha256(&salted_password, b"Client Key");
|
||||||
|
let stored_key = Sha256::digest(&client_key).to_vec();
|
||||||
|
let server_key = hmac_sha256(&salted_password, b"Server Key");
|
||||||
|
|
||||||
|
Self {
|
||||||
|
username: options.username.clone(),
|
||||||
|
database: options.database.clone(),
|
||||||
|
roles: options.roles.clone(),
|
||||||
|
scram_sha256: ScramCredential {
|
||||||
|
salt,
|
||||||
|
iterations,
|
||||||
|
stored_key,
|
||||||
|
server_key,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn to_authenticated_user(&self) -> AuthenticatedUser {
|
||||||
|
AuthenticatedUser {
|
||||||
|
username: self.username.clone(),
|
||||||
|
database: self.database.clone(),
|
||||||
|
roles: self.roles.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn role_allows(role: &str, user: &AuthenticatedUser, target_db: &str, action: AuthAction) -> bool {
|
||||||
|
let (role_db, role_name) = role.split_once('.').unwrap_or(("", role));
|
||||||
|
if role_name == "root" {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
|
||||||
|
let any_database = role_name.ends_with("AnyDatabase");
|
||||||
|
let scoped_db = if role_db.is_empty() { &user.database } else { role_db };
|
||||||
|
if !any_database && scoped_db != target_db {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
match role_name {
|
||||||
|
"read" | "readAnyDatabase" => action == AuthAction::Read,
|
||||||
|
"readWrite" | "readWriteAnyDatabase" => {
|
||||||
|
matches!(action, AuthAction::Read | AuthAction::Write)
|
||||||
|
}
|
||||||
|
"dbAdmin" | "dbAdminAnyDatabase" => action == AuthAction::DbAdmin,
|
||||||
|
"userAdmin" | "userAdminAnyDatabase" => action == AuthAction::UserAdmin,
|
||||||
|
"clusterMonitor" => action == AuthAction::ClusterMonitor,
|
||||||
|
_ => false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn load_users(path: &Path) -> Result<HashMap<String, AuthUser>, AuthError> {
|
||||||
|
if !path.exists() {
|
||||||
|
return Ok(HashMap::new());
|
||||||
|
}
|
||||||
|
let data = std::fs::read_to_string(path).map_err(|e| AuthError::Persistence(e.to_string()))?;
|
||||||
|
let persisted: PersistedAuthState = serde_json::from_str(&data)
|
||||||
|
.map_err(|e| AuthError::Persistence(format!("failed to parse users file: {e}")))?;
|
||||||
|
Ok(persisted
|
||||||
|
.users
|
||||||
|
.into_iter()
|
||||||
|
.map(|user| (user_key(&user.database, &user.username), user))
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn persist_users(path: &Path, users: &HashMap<String, AuthUser>) -> Result<(), AuthError> {
|
||||||
|
if let Some(parent) = path.parent() {
|
||||||
|
std::fs::create_dir_all(parent).map_err(|e| AuthError::Persistence(e.to_string()))?;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut user_list: Vec<AuthUser> = users.values().cloned().collect();
|
||||||
|
user_list.sort_by(|a, b| a.database.cmp(&b.database).then(a.username.cmp(&b.username)));
|
||||||
|
let payload = serde_json::to_vec_pretty(&PersistedAuthState { users: user_list })
|
||||||
|
.map_err(|e| AuthError::Persistence(e.to_string()))?;
|
||||||
|
|
||||||
|
let tmp_path = path.with_extension("tmp");
|
||||||
|
{
|
||||||
|
let mut file = std::fs::File::create(&tmp_path)
|
||||||
|
.map_err(|e| AuthError::Persistence(e.to_string()))?;
|
||||||
|
file.write_all(&payload)
|
||||||
|
.map_err(|e| AuthError::Persistence(e.to_string()))?;
|
||||||
|
file.sync_all()
|
||||||
|
.map_err(|e| AuthError::Persistence(e.to_string()))?;
|
||||||
|
}
|
||||||
|
std::fs::rename(&tmp_path, path).map_err(|e| AuthError::Persistence(e.to_string()))?;
|
||||||
|
if let Some(parent) = path.parent() {
|
||||||
|
if let Ok(dir) = std::fs::File::open(parent) {
|
||||||
|
let _ = dir.sync_all();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn user_key(database: &str, username: &str) -> String {
|
||||||
|
format!("{}\0{}", database, username)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn salted_password(password: &[u8], salt: &[u8], iterations: u32) -> Vec<u8> {
|
||||||
|
let mut output = [0u8; 32];
|
||||||
|
pbkdf2_hmac::<Sha256>(password, salt, iterations, &mut output);
|
||||||
|
output.to_vec()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn hmac_sha256(key: &[u8], message: &[u8]) -> Vec<u8> {
|
||||||
|
let mut mac = HmacSha256::new_from_slice(key).expect("HMAC-SHA256 accepts keys of any size");
|
||||||
|
mac.update(message);
|
||||||
|
mac.finalize().into_bytes().to_vec()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn secure_random(len: usize) -> Vec<u8> {
|
||||||
|
let mut bytes = vec![0u8; len];
|
||||||
|
OsRng.fill_bytes(&mut bytes);
|
||||||
|
bytes
|
||||||
|
}
|
||||||
|
|
||||||
|
fn secure_base64(len: usize) -> String {
|
||||||
|
BASE64_STANDARD.encode(secure_random(len))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_scram_attrs(input: &str) -> HashMap<String, String> {
|
||||||
|
let mut result = HashMap::new();
|
||||||
|
for part in input.split(',') {
|
||||||
|
if let Some((key, value)) = part.split_once('=') {
|
||||||
|
result.insert(key.to_string(), value.to_string());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
fn decode_scram_name(input: &str) -> String {
|
||||||
|
input.replace("=2C", ",").replace("=3D", "=")
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn mechanism_lookup_returns_scram_sha256() {
|
||||||
|
let options = AuthOptions {
|
||||||
|
enabled: true,
|
||||||
|
users: vec![AuthUserOptions {
|
||||||
|
username: "root".to_string(),
|
||||||
|
password: "secret".to_string(),
|
||||||
|
database: "admin".to_string(),
|
||||||
|
roles: vec!["root".to_string()],
|
||||||
|
}],
|
||||||
|
users_path: None,
|
||||||
|
scram_iterations: 4096,
|
||||||
|
};
|
||||||
|
let engine = AuthEngine::from_options(&options).unwrap();
|
||||||
|
assert_eq!(engine.supported_mechanisms("admin.root"), vec![SCRAM_SHA_256.to_string()]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn read_write_role_allows_read_and_write_only_on_own_db() {
|
||||||
|
let user = AuthenticatedUser {
|
||||||
|
username: "app".to_string(),
|
||||||
|
database: "appdb".to_string(),
|
||||||
|
roles: vec!["readWrite".to_string()],
|
||||||
|
};
|
||||||
|
assert!(role_allows("readWrite", &user, "appdb", AuthAction::Read));
|
||||||
|
assert!(role_allows("readWrite", &user, "appdb", AuthAction::Write));
|
||||||
|
assert!(!role_allows("readWrite", &user, "other", AuthAction::Read));
|
||||||
|
assert!(!role_allows("readWrite", &user, "appdb", AuthAction::DbAdmin));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,25 @@
|
|||||||
|
[package]
|
||||||
|
name = "rustdb-commands"
|
||||||
|
version.workspace = true
|
||||||
|
edition.workspace = true
|
||||||
|
license.workspace = true
|
||||||
|
authors.workspace = true
|
||||||
|
description = "MongoDB-compatible command routing and handlers for RustDb"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
bson = { workspace = true }
|
||||||
|
serde = { workspace = true }
|
||||||
|
serde_json = { workspace = true }
|
||||||
|
dashmap = { workspace = true }
|
||||||
|
tokio = { workspace = true }
|
||||||
|
thiserror = { workspace = true }
|
||||||
|
tracing = { workspace = true }
|
||||||
|
uuid = { workspace = true }
|
||||||
|
async-trait = { workspace = true }
|
||||||
|
rustdb-config = { workspace = true }
|
||||||
|
rustdb-wire = { workspace = true }
|
||||||
|
rustdb-query = { workspace = true }
|
||||||
|
rustdb-storage = { workspace = true }
|
||||||
|
rustdb-index = { workspace = true }
|
||||||
|
rustdb-txn = { workspace = true }
|
||||||
|
rustdb-auth = { workspace = true }
|
||||||
@@ -0,0 +1,138 @@
|
|||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use bson::{Bson, Document};
|
||||||
|
use dashmap::DashMap;
|
||||||
|
use rustdb_auth::{AuthEngine, AuthenticatedUser, ScramConversation};
|
||||||
|
use rustdb_index::{IndexEngine, IndexOptions};
|
||||||
|
use rustdb_storage::{OpLog, StorageAdapter};
|
||||||
|
use rustdb_txn::{SessionEngine, TransactionEngine};
|
||||||
|
|
||||||
|
/// Shared command execution context, passed to all handlers.
|
||||||
|
pub struct CommandContext {
|
||||||
|
/// The storage backend.
|
||||||
|
pub storage: Arc<dyn StorageAdapter>,
|
||||||
|
/// Index engines per namespace: "db.collection" -> IndexEngine.
|
||||||
|
pub indexes: Arc<DashMap<String, IndexEngine>>,
|
||||||
|
/// Transaction engine for multi-document transactions.
|
||||||
|
pub transactions: Arc<TransactionEngine>,
|
||||||
|
/// Session engine for logical sessions.
|
||||||
|
pub sessions: Arc<SessionEngine>,
|
||||||
|
/// Active cursors for getMore / killCursors.
|
||||||
|
pub cursors: Arc<DashMap<i64, CursorState>>,
|
||||||
|
/// Server start time (for uptime reporting).
|
||||||
|
pub start_time: std::time::Instant,
|
||||||
|
/// Operation log for point-in-time replay.
|
||||||
|
pub oplog: Arc<OpLog>,
|
||||||
|
/// Authentication engine and user store.
|
||||||
|
pub auth: Arc<AuthEngine>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CommandContext {
|
||||||
|
/// Get or lazily initialize an IndexEngine for a namespace.
|
||||||
|
///
|
||||||
|
/// If no IndexEngine exists yet for this namespace, loads persisted index
|
||||||
|
/// specs from `indexes.json` via the storage adapter, creates the engine
|
||||||
|
/// with those specs, and rebuilds index data from existing documents.
|
||||||
|
/// This ensures unique indexes are enforced even on the very first write
|
||||||
|
/// after a restart.
|
||||||
|
pub async fn get_or_init_index_engine(&self, db: &str, coll: &str) -> dashmap::mapref::one::RefMut<'_, String, IndexEngine> {
|
||||||
|
let ns_key = format!("{}.{}", db, coll);
|
||||||
|
|
||||||
|
// Fast path: engine already exists.
|
||||||
|
if self.indexes.contains_key(&ns_key) {
|
||||||
|
return self.indexes.entry(ns_key).or_insert_with(IndexEngine::new);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Slow path: load from persisted specs.
|
||||||
|
let mut engine = IndexEngine::new();
|
||||||
|
let mut has_custom = false;
|
||||||
|
|
||||||
|
if let Ok(specs) = self.storage.get_indexes(db, coll).await {
|
||||||
|
for spec in &specs {
|
||||||
|
let name = spec.get_str("name").unwrap_or("").to_string();
|
||||||
|
if name == "_id_" || name.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let key = match spec.get("key") {
|
||||||
|
Some(Bson::Document(k)) => k.clone(),
|
||||||
|
_ => continue,
|
||||||
|
};
|
||||||
|
let unique = matches!(spec.get("unique"), Some(Bson::Boolean(true)));
|
||||||
|
let sparse = matches!(spec.get("sparse"), Some(Bson::Boolean(true)));
|
||||||
|
let expire_after_seconds = match spec.get("expireAfterSeconds") {
|
||||||
|
Some(Bson::Int32(n)) => Some(*n as u64),
|
||||||
|
Some(Bson::Int64(n)) => Some(*n as u64),
|
||||||
|
_ => None,
|
||||||
|
};
|
||||||
|
let options = IndexOptions {
|
||||||
|
name: Some(name),
|
||||||
|
unique,
|
||||||
|
sparse,
|
||||||
|
expire_after_seconds,
|
||||||
|
};
|
||||||
|
let _ = engine.create_index(key, options);
|
||||||
|
has_custom = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if has_custom {
|
||||||
|
// Rebuild index data from existing documents.
|
||||||
|
if let Ok(docs) = self.storage.find_all(db, coll).await {
|
||||||
|
if !docs.is_empty() {
|
||||||
|
engine.rebuild_from_documents(&docs);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
self.indexes.entry(ns_key).or_insert(engine)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Per-client connection state. Authentication is socket-scoped in MongoDB.
|
||||||
|
pub struct ConnectionState {
|
||||||
|
pub authenticated_users: Vec<AuthenticatedUser>,
|
||||||
|
pub sasl_conversations: std::collections::HashMap<i32, ScramConversation>,
|
||||||
|
next_conversation_id: i32,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ConnectionState {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
authenticated_users: Vec::new(),
|
||||||
|
sasl_conversations: std::collections::HashMap::new(),
|
||||||
|
next_conversation_id: 1,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn is_authenticated(&self) -> bool {
|
||||||
|
!self.authenticated_users.is_empty()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn next_conversation_id(&mut self) -> i32 {
|
||||||
|
let id = self.next_conversation_id;
|
||||||
|
self.next_conversation_id += 1;
|
||||||
|
id
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn authenticate(&mut self, user: AuthenticatedUser) {
|
||||||
|
self.authenticated_users.push(user);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for ConnectionState {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// State of an open cursor from a find or aggregate command.
|
||||||
|
pub struct CursorState {
|
||||||
|
/// Documents remaining to be returned.
|
||||||
|
pub documents: Vec<Document>,
|
||||||
|
/// Current read position within `documents`.
|
||||||
|
pub position: usize,
|
||||||
|
/// Database the cursor belongs to.
|
||||||
|
pub database: String,
|
||||||
|
/// Collection the cursor belongs to.
|
||||||
|
pub collection: String,
|
||||||
|
}
|
||||||
@@ -0,0 +1,108 @@
|
|||||||
|
use thiserror::Error;
|
||||||
|
|
||||||
|
/// Errors that can occur during command processing.
|
||||||
|
#[derive(Debug, Error)]
|
||||||
|
pub enum CommandError {
|
||||||
|
#[error("command not implemented: {0}")]
|
||||||
|
NotImplemented(String),
|
||||||
|
|
||||||
|
#[error("invalid argument: {0}")]
|
||||||
|
InvalidArgument(String),
|
||||||
|
|
||||||
|
#[error("storage error: {0}")]
|
||||||
|
StorageError(String),
|
||||||
|
|
||||||
|
#[error("index error: {0}")]
|
||||||
|
IndexError(String),
|
||||||
|
|
||||||
|
#[error("transaction error: {0}")]
|
||||||
|
TransactionError(String),
|
||||||
|
|
||||||
|
#[error("no such transaction: {0}")]
|
||||||
|
NoSuchTransaction(String),
|
||||||
|
|
||||||
|
#[error("write conflict: {0}")]
|
||||||
|
WriteConflict(String),
|
||||||
|
|
||||||
|
#[error("namespace not found: {0}")]
|
||||||
|
NamespaceNotFound(String),
|
||||||
|
|
||||||
|
#[error("namespace already exists: {0}")]
|
||||||
|
NamespaceExists(String),
|
||||||
|
|
||||||
|
#[error("duplicate key: {0}")]
|
||||||
|
DuplicateKey(String),
|
||||||
|
|
||||||
|
#[error("immutable field: {0}")]
|
||||||
|
ImmutableField(String),
|
||||||
|
|
||||||
|
#[error("unauthorized: {0}")]
|
||||||
|
Unauthorized(String),
|
||||||
|
|
||||||
|
#[error("authentication failed")]
|
||||||
|
AuthenticationFailed,
|
||||||
|
|
||||||
|
#[error("illegal operation: {0}")]
|
||||||
|
IllegalOperation(String),
|
||||||
|
|
||||||
|
#[error("internal error: {0}")]
|
||||||
|
InternalError(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CommandError {
|
||||||
|
/// Convert a CommandError to a BSON error response document.
|
||||||
|
pub fn to_error_doc(&self) -> bson::Document {
|
||||||
|
let (code, code_name) = match self {
|
||||||
|
CommandError::NotImplemented(_) => (59, "CommandNotFound"),
|
||||||
|
CommandError::InvalidArgument(_) => (14, "TypeMismatch"),
|
||||||
|
CommandError::StorageError(_) => (1, "InternalError"),
|
||||||
|
CommandError::IndexError(_) => (27, "IndexNotFound"),
|
||||||
|
CommandError::TransactionError(_) => (112, "WriteConflict"),
|
||||||
|
CommandError::NoSuchTransaction(_) => (251, "NoSuchTransaction"),
|
||||||
|
CommandError::WriteConflict(_) => (112, "WriteConflict"),
|
||||||
|
CommandError::NamespaceNotFound(_) => (26, "NamespaceNotFound"),
|
||||||
|
CommandError::NamespaceExists(_) => (48, "NamespaceExists"),
|
||||||
|
CommandError::DuplicateKey(_) => (11000, "DuplicateKey"),
|
||||||
|
CommandError::ImmutableField(_) => (66, "ImmutableField"),
|
||||||
|
CommandError::Unauthorized(_) => (13, "Unauthorized"),
|
||||||
|
CommandError::AuthenticationFailed => (18, "AuthenticationFailed"),
|
||||||
|
CommandError::IllegalOperation(_) => (20, "IllegalOperation"),
|
||||||
|
CommandError::InternalError(_) => (1, "InternalError"),
|
||||||
|
};
|
||||||
|
|
||||||
|
bson::doc! {
|
||||||
|
"ok": 0,
|
||||||
|
"errmsg": self.to_string(),
|
||||||
|
"code": code,
|
||||||
|
"codeName": code_name,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl From<rustdb_storage::StorageError> for CommandError {
|
||||||
|
fn from(e: rustdb_storage::StorageError) -> Self {
|
||||||
|
CommandError::StorageError(e.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl From<rustdb_txn::TransactionError> for CommandError {
|
||||||
|
fn from(e: rustdb_txn::TransactionError) -> Self {
|
||||||
|
match e {
|
||||||
|
rustdb_txn::TransactionError::NotFound(message) => {
|
||||||
|
CommandError::NoSuchTransaction(message)
|
||||||
|
}
|
||||||
|
rustdb_txn::TransactionError::WriteConflict(message) => {
|
||||||
|
CommandError::WriteConflict(message)
|
||||||
|
}
|
||||||
|
other => CommandError::TransactionError(other.to_string()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl From<rustdb_index::IndexError> for CommandError {
|
||||||
|
fn from(e: rustdb_index::IndexError) -> Self {
|
||||||
|
CommandError::IndexError(e.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub type CommandResult<T> = Result<T, CommandError>;
|
||||||
@@ -0,0 +1,875 @@
|
|||||||
|
use bson::{doc, Bson, Document};
|
||||||
|
use rustdb_index::IndexEngine;
|
||||||
|
use tracing::debug;
|
||||||
|
|
||||||
|
use crate::context::{CommandContext, ConnectionState, CursorState};
|
||||||
|
use crate::error::{CommandError, CommandResult};
|
||||||
|
use crate::transactions;
|
||||||
|
|
||||||
|
/// Handle various admin / diagnostic / session / auth commands.
|
||||||
|
pub async fn handle(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
command_name: &str,
|
||||||
|
connection: &ConnectionState,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
match command_name {
|
||||||
|
"ping" => Ok(doc! { "ok": 1.0 }),
|
||||||
|
|
||||||
|
"buildInfo" | "buildinfo" => Ok(doc! {
|
||||||
|
"version": "7.0.0",
|
||||||
|
"gitVersion": "unknown",
|
||||||
|
"modules": [],
|
||||||
|
"sysInfo": "rustdb",
|
||||||
|
"versionArray": [7_i32, 0_i32, 0_i32, 0_i32],
|
||||||
|
"ok": 1.0,
|
||||||
|
}),
|
||||||
|
|
||||||
|
"serverStatus" => handle_server_status(ctx),
|
||||||
|
|
||||||
|
"hostInfo" => Ok(doc! {
|
||||||
|
"system": {
|
||||||
|
"hostname": "localhost",
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
}),
|
||||||
|
|
||||||
|
"whatsmyuri" => Ok(doc! {
|
||||||
|
"you": "127.0.0.1:0",
|
||||||
|
"ok": 1.0,
|
||||||
|
}),
|
||||||
|
|
||||||
|
"getLog" => {
|
||||||
|
let _log_type = cmd.get_str("getLog").unwrap_or("global");
|
||||||
|
Ok(doc! {
|
||||||
|
"totalLinesWritten": 0_i32,
|
||||||
|
"log": [],
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
"replSetGetStatus" => {
|
||||||
|
// Not a replica set.
|
||||||
|
Ok(doc! {
|
||||||
|
"ok": 0.0,
|
||||||
|
"errmsg": "not running with --replSet",
|
||||||
|
"code": 76_i32,
|
||||||
|
"codeName": "NoReplicationEnabled",
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
"getCmdLineOpts" => Ok(doc! {
|
||||||
|
"argv": ["rustdb"],
|
||||||
|
"parsed": {},
|
||||||
|
"ok": 1.0,
|
||||||
|
}),
|
||||||
|
|
||||||
|
"getParameter" => Ok(doc! {
|
||||||
|
"ok": 1.0,
|
||||||
|
}),
|
||||||
|
|
||||||
|
"getFreeMonitoringStatus" | "setFreeMonitoring" => Ok(doc! {
|
||||||
|
"state": "disabled",
|
||||||
|
"ok": 1.0,
|
||||||
|
}),
|
||||||
|
|
||||||
|
"getShardMap" | "shardingState" => Ok(doc! {
|
||||||
|
"enabled": false,
|
||||||
|
"ok": 1.0,
|
||||||
|
}),
|
||||||
|
|
||||||
|
"atlasVersion" => Ok(doc! {
|
||||||
|
"ok": 0.0,
|
||||||
|
"errmsg": "not supported",
|
||||||
|
"code": 59_i32,
|
||||||
|
"codeName": "CommandNotFound",
|
||||||
|
}),
|
||||||
|
|
||||||
|
"connectionStatus" => Ok(handle_connection_status(connection)),
|
||||||
|
|
||||||
|
"createUser" => handle_create_user(cmd, db, ctx).await,
|
||||||
|
|
||||||
|
"updateUser" => handle_update_user(cmd, db, ctx).await,
|
||||||
|
|
||||||
|
"dropUser" => handle_drop_user(cmd, db, ctx).await,
|
||||||
|
|
||||||
|
"usersInfo" => handle_users_info(cmd, db, ctx).await,
|
||||||
|
|
||||||
|
"grantRolesToUser" => handle_grant_roles_to_user(cmd, db, ctx).await,
|
||||||
|
|
||||||
|
"revokeRolesFromUser" => handle_revoke_roles_from_user(cmd, db, ctx).await,
|
||||||
|
|
||||||
|
"listDatabases" => handle_list_databases(cmd, ctx).await,
|
||||||
|
|
||||||
|
"listCollections" => handle_list_collections(cmd, db, ctx).await,
|
||||||
|
|
||||||
|
"create" => handle_create(cmd, db, ctx).await,
|
||||||
|
|
||||||
|
"drop" => handle_drop(cmd, db, ctx).await,
|
||||||
|
|
||||||
|
"dropDatabase" => handle_drop_database(db, ctx).await,
|
||||||
|
|
||||||
|
"renameCollection" => handle_rename_collection(cmd, ctx).await,
|
||||||
|
|
||||||
|
"collStats" | "validate" => handle_coll_stats(cmd, db, ctx, command_name).await,
|
||||||
|
|
||||||
|
"dbStats" => handle_db_stats(db, ctx).await,
|
||||||
|
|
||||||
|
"explain" => Ok(doc! {
|
||||||
|
"queryPlanner": {},
|
||||||
|
"ok": 1.0,
|
||||||
|
}),
|
||||||
|
|
||||||
|
"startSession" => {
|
||||||
|
let session_id = uuid::Uuid::new_v4().to_string();
|
||||||
|
ctx.sessions.get_or_create_session(&session_id);
|
||||||
|
Ok(doc! {
|
||||||
|
"id": { "id": &session_id },
|
||||||
|
"timeoutMinutes": 30_i32,
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
"endSessions" | "killSessions" => {
|
||||||
|
// Attempt to end listed sessions.
|
||||||
|
if let Ok(sessions) = cmd
|
||||||
|
.get_array("endSessions")
|
||||||
|
.or_else(|_| cmd.get_array("killSessions"))
|
||||||
|
{
|
||||||
|
for s in sessions {
|
||||||
|
if let Some(sid) = rustdb_txn::SessionEngine::extract_session_id(s) {
|
||||||
|
ctx.sessions.end_session(&sid);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(doc! { "ok": 1.0 })
|
||||||
|
}
|
||||||
|
|
||||||
|
"commitTransaction" => transactions::commit_transaction_command(cmd, ctx).await,
|
||||||
|
|
||||||
|
"abortTransaction" => transactions::abort_transaction_command(cmd, ctx),
|
||||||
|
|
||||||
|
// Auth stubs - accept silently.
|
||||||
|
"saslStart" => Ok(doc! {
|
||||||
|
"conversationId": 1_i32,
|
||||||
|
"done": true,
|
||||||
|
"payload": bson::Binary { subtype: bson::spec::BinarySubtype::Generic, bytes: vec![] },
|
||||||
|
"ok": 1.0,
|
||||||
|
}),
|
||||||
|
|
||||||
|
"saslContinue" => Ok(doc! {
|
||||||
|
"conversationId": 1_i32,
|
||||||
|
"done": true,
|
||||||
|
"payload": bson::Binary { subtype: bson::spec::BinarySubtype::Generic, bytes: vec![] },
|
||||||
|
"ok": 1.0,
|
||||||
|
}),
|
||||||
|
|
||||||
|
"authenticate" | "logout" => Ok(doc! { "ok": 1.0 }),
|
||||||
|
|
||||||
|
"currentOp" => Ok(doc! {
|
||||||
|
"inprog": [],
|
||||||
|
"ok": 1.0,
|
||||||
|
}),
|
||||||
|
|
||||||
|
"killOp" | "top" | "profile" | "compact" | "reIndex"
|
||||||
|
| "fsync" | "connPoolSync" => Ok(doc! { "ok": 1.0 }),
|
||||||
|
|
||||||
|
other => {
|
||||||
|
// Catch-all for any admin command we missed.
|
||||||
|
Ok(doc! {
|
||||||
|
"ok": 1.0,
|
||||||
|
"note": format!("stub response for command: {}", other),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn handle_server_status(ctx: &CommandContext) -> CommandResult<Document> {
|
||||||
|
let oplog_stats = ctx.oplog.stats();
|
||||||
|
Ok(doc! {
|
||||||
|
"host": "localhost",
|
||||||
|
"version": "7.0.0",
|
||||||
|
"process": "rustdb",
|
||||||
|
"uptime": ctx.start_time.elapsed().as_secs() as i64,
|
||||||
|
"connections": {
|
||||||
|
"current": 0_i32,
|
||||||
|
"available": i32::MAX,
|
||||||
|
},
|
||||||
|
"logicalSessionRecordCache": {
|
||||||
|
"activeSessionsCount": ctx.sessions.len() as i64,
|
||||||
|
},
|
||||||
|
"transactions": {
|
||||||
|
"currentActive": ctx.transactions.len() as i64,
|
||||||
|
},
|
||||||
|
"oplog": {
|
||||||
|
"currentSeq": oplog_stats.current_seq as i64,
|
||||||
|
"totalEntries": oplog_stats.total_entries as i64,
|
||||||
|
"oldestSeq": oplog_stats.oldest_seq as i64,
|
||||||
|
"entriesByOp": {
|
||||||
|
"insert": oplog_stats.inserts as i64,
|
||||||
|
"update": oplog_stats.updates as i64,
|
||||||
|
"delete": oplog_stats.deletes as i64,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"security": {
|
||||||
|
"authentication": ctx.auth.enabled(),
|
||||||
|
"users": ctx.auth.user_count() as i64,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn handle_connection_status(connection: &ConnectionState) -> Document {
|
||||||
|
let authenticated_users: Vec<Bson> = connection
|
||||||
|
.authenticated_users
|
||||||
|
.iter()
|
||||||
|
.map(|user| {
|
||||||
|
Bson::Document(doc! {
|
||||||
|
"user": user.username.clone(),
|
||||||
|
"db": user.database.clone(),
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let authenticated_roles: Vec<Bson> = connection
|
||||||
|
.authenticated_users
|
||||||
|
.iter()
|
||||||
|
.flat_map(|user| {
|
||||||
|
user.roles
|
||||||
|
.iter()
|
||||||
|
.map(|role| Bson::Document(role_to_document(&user.database, role)))
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
doc! {
|
||||||
|
"authInfo": {
|
||||||
|
"authenticatedUsers": authenticated_users,
|
||||||
|
"authenticatedUserRoles": authenticated_roles,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn handle_create_user(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let username = cmd
|
||||||
|
.get_str("createUser")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'createUser' field".into()))?;
|
||||||
|
let password = cmd
|
||||||
|
.get_str("pwd")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'pwd' field".into()))?;
|
||||||
|
let roles = parse_roles(cmd, db, "roles")?;
|
||||||
|
ctx.auth
|
||||||
|
.create_user(db, username, password, roles)
|
||||||
|
.map_err(auth_error_to_command_error)?;
|
||||||
|
Ok(doc! { "ok": 1.0 })
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn handle_update_user(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let username = cmd
|
||||||
|
.get_str("updateUser")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'updateUser' field".into()))?;
|
||||||
|
let password = cmd.get_str("pwd").ok();
|
||||||
|
let roles = if cmd.contains_key("roles") {
|
||||||
|
Some(parse_roles(cmd, db, "roles")?)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
ctx.auth
|
||||||
|
.update_user(db, username, password, roles)
|
||||||
|
.map_err(auth_error_to_command_error)?;
|
||||||
|
Ok(doc! { "ok": 1.0 })
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn handle_drop_user(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let username = cmd
|
||||||
|
.get_str("dropUser")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'dropUser' field".into()))?;
|
||||||
|
ctx.auth
|
||||||
|
.drop_user(db, username)
|
||||||
|
.map_err(auth_error_to_command_error)?;
|
||||||
|
Ok(doc! { "ok": 1.0 })
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn handle_users_info(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let username = match cmd.get("usersInfo") {
|
||||||
|
Some(Bson::String(name)) => Some(name.as_str()),
|
||||||
|
Some(Bson::Document(user_doc)) => user_doc.get_str("user").ok(),
|
||||||
|
_ => None,
|
||||||
|
};
|
||||||
|
let users = ctx.auth.users_info(db, username);
|
||||||
|
let user_docs: Vec<Bson> = users
|
||||||
|
.into_iter()
|
||||||
|
.map(|user| {
|
||||||
|
let roles: Vec<Bson> = user
|
||||||
|
.roles
|
||||||
|
.iter()
|
||||||
|
.map(|role| Bson::Document(role_to_document(&user.database, role)))
|
||||||
|
.collect();
|
||||||
|
Bson::Document(doc! {
|
||||||
|
"user": user.username,
|
||||||
|
"db": user.database,
|
||||||
|
"roles": roles,
|
||||||
|
"mechanisms": ["SCRAM-SHA-256"],
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
Ok(doc! { "users": user_docs, "ok": 1.0 })
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn handle_grant_roles_to_user(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let username = cmd
|
||||||
|
.get_str("grantRolesToUser")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'grantRolesToUser' field".into()))?;
|
||||||
|
let roles = parse_roles(cmd, db, "roles")?;
|
||||||
|
ctx.auth
|
||||||
|
.grant_roles(db, username, roles)
|
||||||
|
.map_err(auth_error_to_command_error)?;
|
||||||
|
Ok(doc! { "ok": 1.0 })
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn handle_revoke_roles_from_user(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let username = cmd
|
||||||
|
.get_str("revokeRolesFromUser")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'revokeRolesFromUser' field".into()))?;
|
||||||
|
let roles = parse_roles(cmd, db, "roles")?;
|
||||||
|
ctx.auth
|
||||||
|
.revoke_roles(db, username, roles)
|
||||||
|
.map_err(auth_error_to_command_error)?;
|
||||||
|
Ok(doc! { "ok": 1.0 })
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_roles(cmd: &Document, db: &str, key: &str) -> CommandResult<Vec<String>> {
|
||||||
|
let role_values = cmd
|
||||||
|
.get_array(key)
|
||||||
|
.map_err(|_| CommandError::InvalidArgument(format!("missing '{key}' array")))?;
|
||||||
|
let mut roles = Vec::with_capacity(role_values.len());
|
||||||
|
for role_value in role_values {
|
||||||
|
match role_value {
|
||||||
|
Bson::String(role) => roles.push(role.clone()),
|
||||||
|
Bson::Document(role_doc) => {
|
||||||
|
let role = role_doc
|
||||||
|
.get_str("role")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("role document missing 'role'".into()))?;
|
||||||
|
let role_db = role_doc.get_str("db").unwrap_or(db);
|
||||||
|
if role_db == db {
|
||||||
|
roles.push(role.to_string());
|
||||||
|
} else {
|
||||||
|
roles.push(format!("{role_db}.{role}"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => return Err(CommandError::InvalidArgument("roles must be strings or documents".into())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(roles)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn role_to_document(default_db: &str, role: &str) -> Document {
|
||||||
|
if let Some((role_db, role_name)) = role.split_once('.') {
|
||||||
|
doc! { "role": role_name, "db": role_db }
|
||||||
|
} else {
|
||||||
|
doc! { "role": role, "db": default_db }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn auth_error_to_command_error(error: rustdb_auth::AuthError) -> CommandError {
|
||||||
|
match error {
|
||||||
|
rustdb_auth::AuthError::UserAlreadyExists(message) => CommandError::DuplicateKey(message),
|
||||||
|
rustdb_auth::AuthError::UserNotFound(message) => CommandError::NamespaceNotFound(message),
|
||||||
|
rustdb_auth::AuthError::Persistence(message) => CommandError::InternalError(message),
|
||||||
|
rustdb_auth::AuthError::AuthenticationFailed => CommandError::AuthenticationFailed,
|
||||||
|
rustdb_auth::AuthError::InvalidPayload(message) => CommandError::InvalidArgument(message),
|
||||||
|
rustdb_auth::AuthError::UnsupportedMechanism(message) => CommandError::InvalidArgument(message),
|
||||||
|
rustdb_auth::AuthError::Disabled => CommandError::Unauthorized("authentication is disabled".into()),
|
||||||
|
rustdb_auth::AuthError::UnknownConversation => {
|
||||||
|
CommandError::InvalidArgument("unknown SASL conversation".into())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle `listDatabases` command.
|
||||||
|
async fn handle_list_databases(
|
||||||
|
cmd: &Document,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let dbs = ctx.storage.list_databases().await?;
|
||||||
|
|
||||||
|
let name_only = match cmd.get("nameOnly") {
|
||||||
|
Some(Bson::Boolean(true)) => true,
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let filter = match cmd.get("filter") {
|
||||||
|
Some(Bson::Document(d)) => Some(d.clone()),
|
||||||
|
_ => None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut db_docs: Vec<Bson> = Vec::new();
|
||||||
|
let mut total_size: i64 = 0;
|
||||||
|
|
||||||
|
for db_name in &dbs {
|
||||||
|
let mut db_info = doc! { "name": db_name.as_str() };
|
||||||
|
|
||||||
|
if !name_only {
|
||||||
|
// Estimate size by counting documents across collections.
|
||||||
|
let mut db_size: i64 = 0;
|
||||||
|
if let Ok(collections) = ctx.storage.list_collections(db_name).await {
|
||||||
|
for coll in &collections {
|
||||||
|
if let Ok(count) = ctx.storage.count(db_name, coll).await {
|
||||||
|
// Rough estimate: 200 bytes per document.
|
||||||
|
db_size += count as i64 * 200;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
db_info.insert("sizeOnDisk", db_size);
|
||||||
|
db_info.insert("empty", db_size == 0);
|
||||||
|
total_size += db_size;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Apply filter if specified.
|
||||||
|
if let Some(ref f) = filter {
|
||||||
|
if !rustdb_query::QueryMatcher::matches(&db_info, f) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
db_docs.push(Bson::Document(db_info));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut response = doc! {
|
||||||
|
"databases": db_docs,
|
||||||
|
"ok": 1.0,
|
||||||
|
};
|
||||||
|
|
||||||
|
if !name_only {
|
||||||
|
response.insert("totalSize", total_size);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(response)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle `listCollections` command.
|
||||||
|
async fn handle_list_collections(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let collections = ctx.storage.list_collections(db).await?;
|
||||||
|
|
||||||
|
let filter = match cmd.get("filter") {
|
||||||
|
Some(Bson::Document(d)) => Some(d.clone()),
|
||||||
|
_ => None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let name_only = match cmd.get("nameOnly") {
|
||||||
|
Some(Bson::Boolean(true)) => true,
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let batch_size = cmd
|
||||||
|
.get_document("cursor")
|
||||||
|
.ok()
|
||||||
|
.and_then(|c| {
|
||||||
|
c.get_i32("batchSize")
|
||||||
|
.ok()
|
||||||
|
.map(|v| v as usize)
|
||||||
|
.or_else(|| c.get_i64("batchSize").ok().map(|v| v as usize))
|
||||||
|
})
|
||||||
|
.unwrap_or(usize::MAX);
|
||||||
|
|
||||||
|
let ns = format!("{}.$cmd.listCollections", db);
|
||||||
|
|
||||||
|
let mut coll_docs: Vec<Document> = Vec::new();
|
||||||
|
|
||||||
|
for coll_name in &collections {
|
||||||
|
let info_doc = if name_only {
|
||||||
|
doc! {
|
||||||
|
"name": coll_name.as_str(),
|
||||||
|
"type": "collection",
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
doc! {
|
||||||
|
"name": coll_name.as_str(),
|
||||||
|
"type": "collection",
|
||||||
|
"options": {},
|
||||||
|
"info": {
|
||||||
|
"readOnly": false,
|
||||||
|
},
|
||||||
|
"idIndex": {
|
||||||
|
"v": 2_i32,
|
||||||
|
"key": { "_id": 1_i32 },
|
||||||
|
"name": "_id_",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Apply filter if specified.
|
||||||
|
if let Some(ref f) = filter {
|
||||||
|
if !rustdb_query::QueryMatcher::matches(&info_doc, f) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
coll_docs.push(info_doc);
|
||||||
|
}
|
||||||
|
|
||||||
|
if coll_docs.len() <= batch_size {
|
||||||
|
let first_batch: Vec<Bson> = coll_docs.into_iter().map(Bson::Document).collect();
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"cursor": {
|
||||||
|
"id": 0_i64,
|
||||||
|
"ns": &ns,
|
||||||
|
"firstBatch": first_batch,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
let first_batch: Vec<Bson> = coll_docs[..batch_size]
|
||||||
|
.iter()
|
||||||
|
.cloned()
|
||||||
|
.map(Bson::Document)
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let remaining: Vec<Document> = coll_docs[batch_size..].to_vec();
|
||||||
|
let cursor_id = generate_cursor_id();
|
||||||
|
|
||||||
|
ctx.cursors.insert(
|
||||||
|
cursor_id,
|
||||||
|
CursorState {
|
||||||
|
documents: remaining,
|
||||||
|
position: 0,
|
||||||
|
database: db.to_string(),
|
||||||
|
collection: String::new(),
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"cursor": {
|
||||||
|
"id": cursor_id,
|
||||||
|
"ns": &ns,
|
||||||
|
"firstBatch": first_batch,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle `create` command.
|
||||||
|
async fn handle_create(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let coll = cmd
|
||||||
|
.get_str("create")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'create' field".into()))?;
|
||||||
|
|
||||||
|
debug!(db = db, collection = coll, "create command");
|
||||||
|
|
||||||
|
// Create database (ignore AlreadyExists).
|
||||||
|
if let Err(e) = ctx.storage.create_database(db).await {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
|
||||||
|
return Err(CommandError::StorageError(msg));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create collection.
|
||||||
|
if let Err(e) = ctx.storage.create_collection(db, coll).await {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if msg.contains("AlreadyExists") || msg.contains("already exists") {
|
||||||
|
return Err(CommandError::NamespaceExists(format!("{}.{}", db, coll)));
|
||||||
|
}
|
||||||
|
return Err(CommandError::StorageError(msg));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Initialize index engine for the new collection.
|
||||||
|
let ns_key = format!("{}.{}", db, coll);
|
||||||
|
ctx.indexes
|
||||||
|
.entry(ns_key)
|
||||||
|
.or_insert_with(IndexEngine::new);
|
||||||
|
|
||||||
|
Ok(doc! { "ok": 1.0 })
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle `drop` command.
|
||||||
|
async fn handle_drop(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let coll = cmd
|
||||||
|
.get_str("drop")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'drop' field".into()))?;
|
||||||
|
|
||||||
|
let ns_key = format!("{}.{}", db, coll);
|
||||||
|
|
||||||
|
debug!(db = db, collection = coll, "drop command");
|
||||||
|
|
||||||
|
// Check if collection exists.
|
||||||
|
match ctx.storage.collection_exists(db, coll).await {
|
||||||
|
Ok(false) => {
|
||||||
|
return Err(CommandError::NamespaceNotFound(format!(
|
||||||
|
"ns not found: {}",
|
||||||
|
ns_key
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
Err(_) => {}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Drop from storage.
|
||||||
|
ctx.storage.drop_collection(db, coll).await?;
|
||||||
|
|
||||||
|
// Remove from indexes.
|
||||||
|
ctx.indexes.remove(&ns_key);
|
||||||
|
|
||||||
|
// Count of indexes that were on this collection (at least _id_).
|
||||||
|
Ok(doc! {
|
||||||
|
"ns": &ns_key,
|
||||||
|
"nIndexesWas": 1_i32,
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle `dropDatabase` command.
|
||||||
|
async fn handle_drop_database(
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
debug!(db = db, "dropDatabase command");
|
||||||
|
|
||||||
|
// Remove all index entries for this database.
|
||||||
|
let prefix = format!("{}.", db);
|
||||||
|
let keys_to_remove: Vec<String> = ctx
|
||||||
|
.indexes
|
||||||
|
.iter()
|
||||||
|
.filter(|entry| entry.key().starts_with(&prefix))
|
||||||
|
.map(|entry| entry.key().clone())
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
for key in keys_to_remove {
|
||||||
|
ctx.indexes.remove(&key);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Drop from storage.
|
||||||
|
ctx.storage.drop_database(db).await?;
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"dropped": db,
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle `renameCollection` command.
|
||||||
|
async fn handle_rename_collection(
|
||||||
|
cmd: &Document,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let source_ns = cmd
|
||||||
|
.get_str("renameCollection")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'renameCollection' field".into()))?;
|
||||||
|
|
||||||
|
let target_ns = cmd
|
||||||
|
.get_str("to")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'to' field".into()))?;
|
||||||
|
|
||||||
|
let drop_target = match cmd.get("dropTarget") {
|
||||||
|
Some(Bson::Boolean(b)) => *b,
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Parse "db.collection" format.
|
||||||
|
let (source_db, source_coll) = parse_namespace(source_ns)?;
|
||||||
|
let (target_db, target_coll) = parse_namespace(target_ns)?;
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
source = source_ns,
|
||||||
|
target = target_ns,
|
||||||
|
drop_target = drop_target,
|
||||||
|
"renameCollection command"
|
||||||
|
);
|
||||||
|
|
||||||
|
// If cross-database rename, that's more complex. For now, support same-db rename.
|
||||||
|
if source_db != target_db {
|
||||||
|
return Err(CommandError::InvalidArgument(
|
||||||
|
"cross-database renameCollection not yet supported".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
// If dropTarget, drop the target collection first.
|
||||||
|
if drop_target {
|
||||||
|
let _ = ctx.storage.drop_collection(target_db, target_coll).await;
|
||||||
|
let target_ns_key = format!("{}.{}", target_db, target_coll);
|
||||||
|
ctx.indexes.remove(&target_ns_key);
|
||||||
|
} else {
|
||||||
|
// Check if target already exists.
|
||||||
|
if let Ok(true) = ctx.storage.collection_exists(target_db, target_coll).await {
|
||||||
|
return Err(CommandError::NamespaceExists(target_ns.to_string()));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Rename in storage.
|
||||||
|
ctx.storage
|
||||||
|
.rename_collection(source_db, source_coll, target_coll)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
// Update index engine: move from old namespace to new.
|
||||||
|
let source_ns_key = format!("{}.{}", source_db, source_coll);
|
||||||
|
let target_ns_key = format!("{}.{}", target_db, target_coll);
|
||||||
|
|
||||||
|
if let Some((_, engine)) = ctx.indexes.remove(&source_ns_key) {
|
||||||
|
ctx.indexes.insert(target_ns_key, engine);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(doc! { "ok": 1.0 })
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle `collStats` command.
|
||||||
|
async fn handle_coll_stats(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
command_name: &str,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let coll = cmd
|
||||||
|
.get_str(command_name)
|
||||||
|
.unwrap_or("unknown");
|
||||||
|
|
||||||
|
let ns = format!("{}.{}", db, coll);
|
||||||
|
|
||||||
|
let count = ctx
|
||||||
|
.storage
|
||||||
|
.count(db, coll)
|
||||||
|
.await
|
||||||
|
.unwrap_or(0);
|
||||||
|
|
||||||
|
let n_indexes = match ctx.indexes.get(&ns) {
|
||||||
|
Some(engine) => engine.list_indexes().len() as i32,
|
||||||
|
None => 1_i32,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Rough size estimate.
|
||||||
|
let data_size = count as i64 * 200;
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"ns": &ns,
|
||||||
|
"count": count as i64,
|
||||||
|
"size": data_size,
|
||||||
|
"avgObjSize": if count > 0 { 200_i64 } else { 0_i64 },
|
||||||
|
"storageSize": data_size,
|
||||||
|
"nindexes": n_indexes,
|
||||||
|
"totalIndexSize": 0_i64,
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle `dbStats` command.
|
||||||
|
async fn handle_db_stats(
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let collections = ctx
|
||||||
|
.storage
|
||||||
|
.list_collections(db)
|
||||||
|
.await
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
let num_collections = collections.len() as i32;
|
||||||
|
let mut total_objects: i64 = 0;
|
||||||
|
let mut total_indexes: i32 = 0;
|
||||||
|
|
||||||
|
for coll in &collections {
|
||||||
|
if let Ok(count) = ctx.storage.count(db, coll).await {
|
||||||
|
total_objects += count as i64;
|
||||||
|
}
|
||||||
|
|
||||||
|
let ns_key = format!("{}.{}", db, coll);
|
||||||
|
if let Some(engine) = ctx.indexes.get(&ns_key) {
|
||||||
|
total_indexes += engine.list_indexes().len() as i32;
|
||||||
|
} else {
|
||||||
|
total_indexes += 1; // At least _id_.
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let data_size = total_objects * 200;
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"db": db,
|
||||||
|
"collections": num_collections,
|
||||||
|
"objects": total_objects,
|
||||||
|
"avgObjSize": if total_objects > 0 { 200_i64 } else { 0_i64 },
|
||||||
|
"dataSize": data_size,
|
||||||
|
"storageSize": data_size,
|
||||||
|
"indexes": total_indexes,
|
||||||
|
"indexSize": 0_i64,
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parse a namespace string "db.collection" into (db, collection).
|
||||||
|
fn parse_namespace(ns: &str) -> CommandResult<(&str, &str)> {
|
||||||
|
let dot_pos = ns.find('.').ok_or_else(|| {
|
||||||
|
CommandError::InvalidArgument(format!(
|
||||||
|
"invalid namespace '{}': expected 'db.collection' format",
|
||||||
|
ns
|
||||||
|
))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
let db = &ns[..dot_pos];
|
||||||
|
let coll = &ns[dot_pos + 1..];
|
||||||
|
|
||||||
|
if db.is_empty() || coll.is_empty() {
|
||||||
|
return Err(CommandError::InvalidArgument(format!(
|
||||||
|
"invalid namespace '{}': db and collection must not be empty",
|
||||||
|
ns
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok((db, coll))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Generate a pseudo-random cursor ID.
|
||||||
|
fn generate_cursor_id() -> i64 {
|
||||||
|
use std::collections::hash_map::RandomState;
|
||||||
|
use std::hash::{BuildHasher, Hasher};
|
||||||
|
let s = RandomState::new();
|
||||||
|
let mut hasher = s.build_hasher();
|
||||||
|
hasher.write_u64(
|
||||||
|
std::time::SystemTime::now()
|
||||||
|
.duration_since(std::time::UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_nanos() as u64,
|
||||||
|
);
|
||||||
|
let id = hasher.finish() as i64;
|
||||||
|
if id == 0 {
|
||||||
|
1
|
||||||
|
} else {
|
||||||
|
id.abs()
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,310 @@
|
|||||||
|
use bson::{doc, Bson, Document};
|
||||||
|
use rustdb_query::AggregationEngine;
|
||||||
|
use rustdb_query::error::QueryError;
|
||||||
|
use tracing::debug;
|
||||||
|
|
||||||
|
use crate::context::{CommandContext, CursorState};
|
||||||
|
use crate::error::{CommandError, CommandResult};
|
||||||
|
|
||||||
|
/// A CollectionResolver that reads from the storage adapter.
|
||||||
|
struct StorageResolver<'a> {
|
||||||
|
storage: &'a dyn rustdb_storage::StorageAdapter,
|
||||||
|
/// We use a tokio runtime handle to call async methods synchronously,
|
||||||
|
/// since the CollectionResolver trait is synchronous.
|
||||||
|
handle: tokio::runtime::Handle,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'a> rustdb_query::aggregation::CollectionResolver for StorageResolver<'a> {
|
||||||
|
fn resolve(&self, db: &str, coll: &str) -> Result<Vec<Document>, QueryError> {
|
||||||
|
self.handle
|
||||||
|
.block_on(async { self.storage.find_all(db, coll).await })
|
||||||
|
.map_err(|e| QueryError::AggregationError(format!("Failed to resolve {}.{}: {}", db, coll, e)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle the `aggregate` command.
|
||||||
|
pub async fn handle(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
// The aggregate field can be a string (collection name) or integer 1 (db-level).
|
||||||
|
let (coll, is_db_level) = match cmd.get("aggregate") {
|
||||||
|
Some(Bson::String(s)) => (s.as_str().to_string(), false),
|
||||||
|
Some(Bson::Int32(1)) => (String::new(), true),
|
||||||
|
Some(Bson::Int64(1)) => (String::new(), true),
|
||||||
|
_ => {
|
||||||
|
return Err(CommandError::InvalidArgument(
|
||||||
|
"missing or invalid 'aggregate' field".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let pipeline_bson = cmd
|
||||||
|
.get_array("pipeline")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'pipeline' array".into()))?;
|
||||||
|
|
||||||
|
// Convert pipeline to Vec<Document>.
|
||||||
|
let mut pipeline: Vec<Document> = Vec::with_capacity(pipeline_bson.len());
|
||||||
|
for stage in pipeline_bson {
|
||||||
|
match stage {
|
||||||
|
Bson::Document(d) => pipeline.push(d.clone()),
|
||||||
|
_ => {
|
||||||
|
return Err(CommandError::InvalidArgument(
|
||||||
|
"pipeline stage must be a document".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check for $out and $merge as the last stage (handle after pipeline execution).
|
||||||
|
let out_stage = if let Some(last) = pipeline.last() {
|
||||||
|
if last.contains_key("$out") || last.contains_key("$merge") {
|
||||||
|
Some(pipeline.pop().unwrap())
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
|
let batch_size = cmd
|
||||||
|
.get_document("cursor")
|
||||||
|
.ok()
|
||||||
|
.and_then(|c| {
|
||||||
|
c.get_i32("batchSize")
|
||||||
|
.ok()
|
||||||
|
.map(|v| v as usize)
|
||||||
|
.or_else(|| c.get_i64("batchSize").ok().map(|v| v as usize))
|
||||||
|
})
|
||||||
|
.unwrap_or(101);
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
db = db,
|
||||||
|
collection = %coll,
|
||||||
|
stages = pipeline.len(),
|
||||||
|
"aggregate command"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Load source documents.
|
||||||
|
let source_docs = if is_db_level {
|
||||||
|
// Database-level aggregate: start with empty set (useful for $currentOp, etc.)
|
||||||
|
Vec::new()
|
||||||
|
} else {
|
||||||
|
ctx.storage.find_all(db, &coll).await?
|
||||||
|
};
|
||||||
|
|
||||||
|
// Create a resolver for $lookup and similar stages.
|
||||||
|
let handle = tokio::runtime::Handle::current();
|
||||||
|
let resolver = StorageResolver {
|
||||||
|
storage: ctx.storage.as_ref(),
|
||||||
|
handle,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Run the aggregation pipeline.
|
||||||
|
let result_docs = AggregationEngine::aggregate(
|
||||||
|
source_docs,
|
||||||
|
&pipeline,
|
||||||
|
Some(&resolver),
|
||||||
|
db,
|
||||||
|
)
|
||||||
|
.map_err(|e| CommandError::InternalError(e.to_string()))?;
|
||||||
|
|
||||||
|
// Handle $out stage: write results to target collection.
|
||||||
|
if let Some(out) = out_stage {
|
||||||
|
if let Some(out_spec) = out.get("$out") {
|
||||||
|
handle_out_stage(db, out_spec, &result_docs, ctx).await?;
|
||||||
|
} else if let Some(merge_spec) = out.get("$merge") {
|
||||||
|
handle_merge_stage(db, merge_spec, &result_docs, ctx).await?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build cursor response.
|
||||||
|
let ns = if is_db_level {
|
||||||
|
format!("{}.$cmd.aggregate", db)
|
||||||
|
} else {
|
||||||
|
format!("{}.{}", db, coll)
|
||||||
|
};
|
||||||
|
|
||||||
|
if result_docs.len() <= batch_size {
|
||||||
|
// All results fit in first batch.
|
||||||
|
let first_batch: Vec<Bson> = result_docs
|
||||||
|
.into_iter()
|
||||||
|
.map(Bson::Document)
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"cursor": {
|
||||||
|
"firstBatch": first_batch,
|
||||||
|
"id": 0_i64,
|
||||||
|
"ns": &ns,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
// Need to create a cursor for remaining results.
|
||||||
|
let first_batch: Vec<Bson> = result_docs[..batch_size]
|
||||||
|
.iter()
|
||||||
|
.cloned()
|
||||||
|
.map(Bson::Document)
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let remaining: Vec<Document> = result_docs[batch_size..].to_vec();
|
||||||
|
let cursor_id = generate_cursor_id();
|
||||||
|
|
||||||
|
ctx.cursors.insert(
|
||||||
|
cursor_id,
|
||||||
|
CursorState {
|
||||||
|
documents: remaining,
|
||||||
|
position: 0,
|
||||||
|
database: db.to_string(),
|
||||||
|
collection: coll.to_string(),
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"cursor": {
|
||||||
|
"firstBatch": first_batch,
|
||||||
|
"id": cursor_id,
|
||||||
|
"ns": &ns,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle $out stage: drop and replace target collection with pipeline results.
|
||||||
|
async fn handle_out_stage(
|
||||||
|
db: &str,
|
||||||
|
out_spec: &Bson,
|
||||||
|
docs: &[Document],
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<()> {
|
||||||
|
let (target_db, target_coll) = match out_spec {
|
||||||
|
Bson::String(coll_name) => (db.to_string(), coll_name.clone()),
|
||||||
|
Bson::Document(d) => {
|
||||||
|
let tdb = d.get_str("db").unwrap_or(db).to_string();
|
||||||
|
let tcoll = d
|
||||||
|
.get_str("coll")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("$out requires 'coll'".into()))?
|
||||||
|
.to_string();
|
||||||
|
(tdb, tcoll)
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
return Err(CommandError::InvalidArgument(
|
||||||
|
"$out requires a string or document".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Drop existing target collection (ignore errors).
|
||||||
|
let _ = ctx.storage.drop_collection(&target_db, &target_coll).await;
|
||||||
|
|
||||||
|
// Create target collection.
|
||||||
|
let _ = ctx.storage.create_database(&target_db).await;
|
||||||
|
let _ = ctx.storage.create_collection(&target_db, &target_coll).await;
|
||||||
|
|
||||||
|
// Insert all result documents.
|
||||||
|
for doc in docs {
|
||||||
|
let _ = ctx
|
||||||
|
.storage
|
||||||
|
.insert_one(&target_db, &target_coll, doc.clone())
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle $merge stage: merge pipeline results into target collection.
|
||||||
|
async fn handle_merge_stage(
|
||||||
|
db: &str,
|
||||||
|
merge_spec: &Bson,
|
||||||
|
docs: &[Document],
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<()> {
|
||||||
|
let (target_db, target_coll) = match merge_spec {
|
||||||
|
Bson::String(coll_name) => (db.to_string(), coll_name.clone()),
|
||||||
|
Bson::Document(d) => {
|
||||||
|
let into_val = d.get("into");
|
||||||
|
match into_val {
|
||||||
|
Some(Bson::String(s)) => (db.to_string(), s.clone()),
|
||||||
|
Some(Bson::Document(into_doc)) => {
|
||||||
|
let tdb = into_doc.get_str("db").unwrap_or(db).to_string();
|
||||||
|
let tcoll = into_doc
|
||||||
|
.get_str("coll")
|
||||||
|
.map_err(|_| {
|
||||||
|
CommandError::InvalidArgument("$merge.into requires 'coll'".into())
|
||||||
|
})?
|
||||||
|
.to_string();
|
||||||
|
(tdb, tcoll)
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
return Err(CommandError::InvalidArgument(
|
||||||
|
"$merge requires 'into' field".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
return Err(CommandError::InvalidArgument(
|
||||||
|
"$merge requires a string or document".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Ensure target collection exists.
|
||||||
|
let _ = ctx.storage.create_database(&target_db).await;
|
||||||
|
let _ = ctx
|
||||||
|
.storage
|
||||||
|
.create_collection(&target_db, &target_coll)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
// Simple merge: upsert by _id.
|
||||||
|
for doc in docs {
|
||||||
|
let id_str = match doc.get("_id") {
|
||||||
|
Some(Bson::ObjectId(oid)) => oid.to_hex(),
|
||||||
|
Some(Bson::String(s)) => s.clone(),
|
||||||
|
Some(other) => format!("{}", other),
|
||||||
|
None => {
|
||||||
|
// No _id, just insert.
|
||||||
|
let _ = ctx
|
||||||
|
.storage
|
||||||
|
.insert_one(&target_db, &target_coll, doc.clone())
|
||||||
|
.await;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Try update first, insert if it fails.
|
||||||
|
match ctx
|
||||||
|
.storage
|
||||||
|
.update_by_id(&target_db, &target_coll, &id_str, doc.clone())
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(_) => {
|
||||||
|
let _ = ctx
|
||||||
|
.storage
|
||||||
|
.insert_one(&target_db, &target_coll, doc.clone())
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Generate a pseudo-random cursor ID.
|
||||||
|
fn generate_cursor_id() -> i64 {
|
||||||
|
use std::collections::hash_map::RandomState;
|
||||||
|
use std::hash::{BuildHasher, Hasher};
|
||||||
|
let s = RandomState::new();
|
||||||
|
let mut hasher = s.build_hasher();
|
||||||
|
hasher.write_u64(std::time::SystemTime::now()
|
||||||
|
.duration_since(std::time::UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_nanos() as u64);
|
||||||
|
let id = hasher.finish() as i64;
|
||||||
|
// Ensure positive and non-zero.
|
||||||
|
if id == 0 { 1 } else { id.abs() }
|
||||||
|
}
|
||||||
@@ -0,0 +1,87 @@
|
|||||||
|
use bson::{doc, Binary, Bson, Document};
|
||||||
|
|
||||||
|
use crate::context::{CommandContext, ConnectionState};
|
||||||
|
use crate::error::{CommandError, CommandResult};
|
||||||
|
|
||||||
|
pub async fn handle_sasl_start(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
connection: &mut ConnectionState,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let mechanism = cmd
|
||||||
|
.get_str("mechanism")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing SASL mechanism".into()))?;
|
||||||
|
if mechanism != "SCRAM-SHA-256" {
|
||||||
|
return Err(CommandError::InvalidArgument(format!(
|
||||||
|
"unsupported SASL mechanism: {mechanism}"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
let payload = payload_bytes(cmd)?;
|
||||||
|
let result = ctx
|
||||||
|
.auth
|
||||||
|
.start_scram_sha256(db, &payload)
|
||||||
|
.map_err(map_auth_error)?;
|
||||||
|
let conversation_id = connection.next_conversation_id();
|
||||||
|
connection
|
||||||
|
.sasl_conversations
|
||||||
|
.insert(conversation_id, result.conversation);
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"conversationId": conversation_id,
|
||||||
|
"done": false,
|
||||||
|
"payload": Binary { subtype: bson::spec::BinarySubtype::Generic, bytes: result.payload },
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn handle_sasl_continue(
|
||||||
|
cmd: &Document,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
connection: &mut ConnectionState,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let conversation_id = cmd
|
||||||
|
.get_i32("conversationId")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing SASL conversationId".into()))?;
|
||||||
|
let payload = payload_bytes(cmd)?;
|
||||||
|
let conversation = connection
|
||||||
|
.sasl_conversations
|
||||||
|
.remove(&conversation_id)
|
||||||
|
.ok_or_else(|| CommandError::InvalidArgument("unknown SASL conversation".into()))?;
|
||||||
|
let result = ctx
|
||||||
|
.auth
|
||||||
|
.continue_scram_sha256(conversation, &payload)
|
||||||
|
.map_err(map_auth_error)?;
|
||||||
|
connection.authenticate(result.user);
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"conversationId": conversation_id,
|
||||||
|
"done": true,
|
||||||
|
"payload": Binary { subtype: bson::spec::BinarySubtype::Generic, bytes: result.payload },
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn payload_bytes(cmd: &Document) -> CommandResult<Vec<u8>> {
|
||||||
|
match cmd.get("payload") {
|
||||||
|
Some(Bson::Binary(binary)) => Ok(binary.bytes.clone()),
|
||||||
|
Some(Bson::String(value)) => Ok(value.as_bytes().to_vec()),
|
||||||
|
_ => Err(CommandError::InvalidArgument("missing SASL payload".into())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn map_auth_error(error: rustdb_auth::AuthError) -> CommandError {
|
||||||
|
match error {
|
||||||
|
rustdb_auth::AuthError::InvalidPayload(message) => CommandError::InvalidArgument(message),
|
||||||
|
rustdb_auth::AuthError::UnsupportedMechanism(message) => CommandError::InvalidArgument(message),
|
||||||
|
rustdb_auth::AuthError::Disabled => CommandError::Unauthorized("authentication is disabled".into()),
|
||||||
|
rustdb_auth::AuthError::UnknownConversation => {
|
||||||
|
CommandError::InvalidArgument("unknown SASL conversation".into())
|
||||||
|
}
|
||||||
|
rustdb_auth::AuthError::AuthenticationFailed => CommandError::AuthenticationFailed,
|
||||||
|
rustdb_auth::AuthError::UserAlreadyExists(message) => CommandError::DuplicateKey(message),
|
||||||
|
rustdb_auth::AuthError::UserNotFound(message) => CommandError::NamespaceNotFound(message),
|
||||||
|
rustdb_auth::AuthError::Persistence(message) => CommandError::InternalError(message),
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,226 @@
|
|||||||
|
use std::collections::HashSet;
|
||||||
|
|
||||||
|
use bson::{doc, Bson, Document};
|
||||||
|
use rustdb_query::QueryMatcher;
|
||||||
|
use rustdb_storage::OpType;
|
||||||
|
use tracing::debug;
|
||||||
|
|
||||||
|
use crate::context::CommandContext;
|
||||||
|
use crate::error::{CommandError, CommandResult};
|
||||||
|
use crate::transactions;
|
||||||
|
|
||||||
|
/// Handle the `delete` command.
|
||||||
|
pub async fn handle(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let coll = cmd
|
||||||
|
.get_str("delete")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'delete' field".into()))?;
|
||||||
|
|
||||||
|
let deletes = cmd
|
||||||
|
.get_array("deletes")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'deletes' array".into()))?;
|
||||||
|
|
||||||
|
// Ordered flag (default true).
|
||||||
|
let ordered = match cmd.get("ordered") {
|
||||||
|
Some(Bson::Boolean(b)) => *b,
|
||||||
|
_ => true,
|
||||||
|
};
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
db = db,
|
||||||
|
collection = coll,
|
||||||
|
count = deletes.len(),
|
||||||
|
"delete command"
|
||||||
|
);
|
||||||
|
|
||||||
|
let ns_key = format!("{}.{}", db, coll);
|
||||||
|
let txn_id = transactions::active_transaction_id(ctx, cmd);
|
||||||
|
let mut total_deleted: i32 = 0;
|
||||||
|
let mut write_errors: Vec<Document> = Vec::new();
|
||||||
|
|
||||||
|
for (idx, del_spec) in deletes.iter().enumerate() {
|
||||||
|
let del_doc = match del_spec {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 14_i32,
|
||||||
|
"codeName": "TypeMismatch",
|
||||||
|
"errmsg": "delete spec must be a document",
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Extract filter (q) and limit.
|
||||||
|
let filter = match del_doc.get_document("q") {
|
||||||
|
Ok(f) => f.clone(),
|
||||||
|
Err(_) => Document::new(), // empty filter matches everything
|
||||||
|
};
|
||||||
|
|
||||||
|
let limit = match del_doc.get("limit") {
|
||||||
|
Some(Bson::Int32(n)) => *n,
|
||||||
|
Some(Bson::Int64(n)) => *n as i32,
|
||||||
|
Some(Bson::Double(n)) => *n as i32,
|
||||||
|
_ => 0, // default: delete all matches
|
||||||
|
};
|
||||||
|
|
||||||
|
match delete_matching(db, coll, &ns_key, &filter, limit, ctx, txn_id.as_deref()).await {
|
||||||
|
Ok(count) => {
|
||||||
|
total_deleted += count;
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 1_i32,
|
||||||
|
"codeName": "InternalError",
|
||||||
|
"errmsg": e.to_string(),
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build response.
|
||||||
|
let mut response = doc! {
|
||||||
|
"n": total_deleted,
|
||||||
|
"ok": 1.0,
|
||||||
|
};
|
||||||
|
|
||||||
|
if !write_errors.is_empty() {
|
||||||
|
response.insert(
|
||||||
|
"writeErrors",
|
||||||
|
write_errors
|
||||||
|
.into_iter()
|
||||||
|
.map(Bson::Document)
|
||||||
|
.collect::<Vec<_>>(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(response)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Find and delete documents matching a filter, returning the number deleted.
|
||||||
|
async fn delete_matching(
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
ns_key: &str,
|
||||||
|
filter: &Document,
|
||||||
|
limit: i32,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
txn_id: Option<&str>,
|
||||||
|
) -> Result<i32, CommandError> {
|
||||||
|
if let Some(txn_id) = txn_id {
|
||||||
|
let docs = transactions::load_transaction_docs(ctx, txn_id, db, coll).await?;
|
||||||
|
let matched = QueryMatcher::filter(&docs, filter);
|
||||||
|
let to_delete: &[Document] = if limit == 1 && !matched.is_empty() {
|
||||||
|
&matched[..1]
|
||||||
|
} else {
|
||||||
|
&matched
|
||||||
|
};
|
||||||
|
|
||||||
|
for doc in to_delete {
|
||||||
|
transactions::record_delete(ctx, txn_id, db, coll, doc.clone()).await?;
|
||||||
|
}
|
||||||
|
|
||||||
|
return Ok(to_delete.len() as i32);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if the collection exists; if not, nothing to delete.
|
||||||
|
match ctx.storage.collection_exists(db, coll).await {
|
||||||
|
Ok(false) => return Ok(0),
|
||||||
|
Err(_) => return Ok(0),
|
||||||
|
Ok(true) => {}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try to use index to narrow candidates.
|
||||||
|
let candidate_ids: Option<HashSet<String>> = {
|
||||||
|
if let Some(engine) = ctx.indexes.get(ns_key) {
|
||||||
|
engine.find_candidate_ids(filter)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Load candidate documents.
|
||||||
|
let docs = if let Some(ids) = candidate_ids {
|
||||||
|
if ids.is_empty() {
|
||||||
|
return Ok(0);
|
||||||
|
}
|
||||||
|
ctx.storage
|
||||||
|
.find_by_ids(db, coll, ids)
|
||||||
|
.await
|
||||||
|
.map_err(|e| CommandError::StorageError(e.to_string()))?
|
||||||
|
} else {
|
||||||
|
ctx.storage
|
||||||
|
.find_all(db, coll)
|
||||||
|
.await
|
||||||
|
.map_err(|e| CommandError::StorageError(e.to_string()))?
|
||||||
|
};
|
||||||
|
|
||||||
|
// Apply filter to get matched documents.
|
||||||
|
let matched = QueryMatcher::filter(&docs, filter);
|
||||||
|
|
||||||
|
// Apply limit: 0 means delete all, 1 means delete only the first match.
|
||||||
|
let to_delete: &[Document] = if limit == 1 && !matched.is_empty() {
|
||||||
|
&matched[..1]
|
||||||
|
} else {
|
||||||
|
&matched
|
||||||
|
};
|
||||||
|
|
||||||
|
if to_delete.is_empty() {
|
||||||
|
return Ok(0);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut deleted_count: i32 = 0;
|
||||||
|
|
||||||
|
for doc in to_delete {
|
||||||
|
// Extract the _id as a hex string for storage deletion.
|
||||||
|
let id_str = extract_id_string(doc)?;
|
||||||
|
|
||||||
|
ctx.storage
|
||||||
|
.delete_by_id(db, coll, &id_str)
|
||||||
|
.await
|
||||||
|
.map_err(|e| CommandError::StorageError(e.to_string()))?;
|
||||||
|
|
||||||
|
// Record in oplog.
|
||||||
|
ctx.oplog.append(
|
||||||
|
OpType::Delete,
|
||||||
|
db,
|
||||||
|
coll,
|
||||||
|
&id_str,
|
||||||
|
None,
|
||||||
|
Some(doc.clone()),
|
||||||
|
);
|
||||||
|
|
||||||
|
// Update index engine.
|
||||||
|
if let Some(mut engine) = ctx.indexes.get_mut(ns_key) {
|
||||||
|
engine.on_delete(doc);
|
||||||
|
}
|
||||||
|
|
||||||
|
deleted_count += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(deleted_count)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extract the `_id` field from a document as a hex string suitable for the
|
||||||
|
/// storage adapter.
|
||||||
|
fn extract_id_string(doc: &Document) -> Result<String, CommandError> {
|
||||||
|
match doc.get("_id") {
|
||||||
|
Some(Bson::ObjectId(oid)) => Ok(oid.to_hex()),
|
||||||
|
Some(Bson::String(s)) => Ok(s.clone()),
|
||||||
|
Some(other) => Ok(format!("{}", other)),
|
||||||
|
None => Err(CommandError::InvalidArgument(
|
||||||
|
"document missing _id field".into(),
|
||||||
|
)),
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,409 @@
|
|||||||
|
use std::sync::atomic::{AtomicI64, Ordering};
|
||||||
|
|
||||||
|
use bson::{doc, Bson, Document};
|
||||||
|
use tracing::debug;
|
||||||
|
|
||||||
|
use rustdb_query::{QueryMatcher, sort_documents, apply_projection, distinct_values};
|
||||||
|
|
||||||
|
use crate::context::{CommandContext, CursorState};
|
||||||
|
use crate::error::{CommandError, CommandResult};
|
||||||
|
use crate::transactions;
|
||||||
|
|
||||||
|
/// Atomic counter for generating unique cursor IDs.
|
||||||
|
static CURSOR_ID_COUNTER: AtomicI64 = AtomicI64::new(1);
|
||||||
|
|
||||||
|
/// Generate a new unique, positive cursor ID.
|
||||||
|
fn next_cursor_id() -> i64 {
|
||||||
|
CURSOR_ID_COUNTER.fetch_add(1, Ordering::Relaxed)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Helpers to defensively extract values from BSON command documents
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
fn get_str<'a>(doc: &'a Document, key: &str) -> Option<&'a str> {
|
||||||
|
match doc.get(key)? {
|
||||||
|
Bson::String(s) => Some(s.as_str()),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_i32(doc: &Document, key: &str) -> Option<i32> {
|
||||||
|
match doc.get(key)? {
|
||||||
|
Bson::Int32(v) => Some(*v),
|
||||||
|
Bson::Int64(v) => Some(*v as i32),
|
||||||
|
Bson::Double(v) => Some(*v as i32),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_i64(doc: &Document, key: &str) -> Option<i64> {
|
||||||
|
match doc.get(key)? {
|
||||||
|
Bson::Int64(v) => Some(*v),
|
||||||
|
Bson::Int32(v) => Some(*v as i64),
|
||||||
|
Bson::Double(v) => Some(*v as i64),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_bool(doc: &Document, key: &str) -> Option<bool> {
|
||||||
|
match doc.get(key)? {
|
||||||
|
Bson::Boolean(v) => Some(*v),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_document<'a>(doc: &'a Document, key: &str) -> Option<&'a Document> {
|
||||||
|
match doc.get(key)? {
|
||||||
|
Bson::Document(d) => Some(d),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// find
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Handle the `find` command.
|
||||||
|
pub async fn handle(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let coll = get_str(cmd, "find").unwrap_or("unknown");
|
||||||
|
let ns = format!("{}.{}", db, coll);
|
||||||
|
|
||||||
|
// Extract optional parameters.
|
||||||
|
let filter = get_document(cmd, "filter").cloned().unwrap_or_default();
|
||||||
|
let sort_spec = get_document(cmd, "sort").cloned();
|
||||||
|
let projection = get_document(cmd, "projection").cloned();
|
||||||
|
let skip = get_i64(cmd, "skip").unwrap_or(0).max(0) as usize;
|
||||||
|
let limit = get_i64(cmd, "limit").unwrap_or(0).max(0) as usize;
|
||||||
|
let batch_size = get_i32(cmd, "batchSize").unwrap_or(101).max(0) as usize;
|
||||||
|
let single_batch = get_bool(cmd, "singleBatch").unwrap_or(false);
|
||||||
|
let txn_id = transactions::active_transaction_id(ctx, cmd);
|
||||||
|
|
||||||
|
// If the collection does not exist, return an empty cursor.
|
||||||
|
let exists = if txn_id.is_some() {
|
||||||
|
true
|
||||||
|
} else {
|
||||||
|
ctx.storage.collection_exists(db, coll).await?
|
||||||
|
};
|
||||||
|
if !exists {
|
||||||
|
return Ok(doc! {
|
||||||
|
"cursor": {
|
||||||
|
"firstBatch": [],
|
||||||
|
"id": 0_i64,
|
||||||
|
"ns": &ns,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try index-accelerated lookup.
|
||||||
|
let index_key = format!("{}.{}", db, coll);
|
||||||
|
let docs = if let Some(ref txn_id) = txn_id {
|
||||||
|
transactions::load_transaction_docs(ctx, txn_id, db, coll).await?
|
||||||
|
} else if let Some(idx_ref) = ctx.indexes.get(&index_key) {
|
||||||
|
if let Some(candidate_ids) = idx_ref.find_candidate_ids(&filter) {
|
||||||
|
debug!(
|
||||||
|
ns = %ns,
|
||||||
|
candidates = candidate_ids.len(),
|
||||||
|
"using index acceleration"
|
||||||
|
);
|
||||||
|
ctx.storage.find_by_ids(db, coll, candidate_ids).await?
|
||||||
|
} else {
|
||||||
|
ctx.storage.find_all(db, coll).await?
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
ctx.storage.find_all(db, coll).await?
|
||||||
|
};
|
||||||
|
|
||||||
|
// Apply filter.
|
||||||
|
let mut docs = QueryMatcher::filter(&docs, &filter);
|
||||||
|
|
||||||
|
// Apply sort.
|
||||||
|
if let Some(ref sort) = sort_spec {
|
||||||
|
sort_documents(&mut docs, sort);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Apply skip.
|
||||||
|
if skip > 0 {
|
||||||
|
if skip >= docs.len() {
|
||||||
|
docs = Vec::new();
|
||||||
|
} else {
|
||||||
|
docs = docs.split_off(skip);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Apply limit.
|
||||||
|
if limit > 0 && docs.len() > limit {
|
||||||
|
docs.truncate(limit);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Apply projection.
|
||||||
|
if let Some(ref proj) = projection {
|
||||||
|
docs = docs.iter().map(|d| apply_projection(d, proj)).collect();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Determine first batch.
|
||||||
|
if docs.len() <= batch_size || single_batch {
|
||||||
|
// Everything fits in a single batch.
|
||||||
|
let batch: Vec<Bson> = docs.into_iter().map(Bson::Document).collect();
|
||||||
|
Ok(doc! {
|
||||||
|
"cursor": {
|
||||||
|
"firstBatch": batch,
|
||||||
|
"id": 0_i64,
|
||||||
|
"ns": &ns,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
// Split into first batch and remainder, store cursor.
|
||||||
|
let remaining = docs.split_off(batch_size);
|
||||||
|
let first_batch: Vec<Bson> = docs.into_iter().map(Bson::Document).collect();
|
||||||
|
|
||||||
|
let cursor_id = next_cursor_id();
|
||||||
|
ctx.cursors.insert(cursor_id, CursorState {
|
||||||
|
documents: remaining,
|
||||||
|
position: 0,
|
||||||
|
database: db.to_string(),
|
||||||
|
collection: coll.to_string(),
|
||||||
|
});
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"cursor": {
|
||||||
|
"firstBatch": first_batch,
|
||||||
|
"id": cursor_id,
|
||||||
|
"ns": &ns,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// getMore
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Handle the `getMore` command.
|
||||||
|
pub async fn handle_get_more(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
// Defensively extract cursor id.
|
||||||
|
let cursor_id = get_i64(cmd, "getMore").ok_or_else(|| {
|
||||||
|
CommandError::InvalidArgument("getMore requires a cursor id".into())
|
||||||
|
})?;
|
||||||
|
|
||||||
|
let coll = get_str(cmd, "collection").unwrap_or("unknown");
|
||||||
|
let ns = format!("{}.{}", db, coll);
|
||||||
|
|
||||||
|
let batch_size = get_i64(cmd, "batchSize")
|
||||||
|
.or_else(|| get_i32(cmd, "batchSize").map(|v| v as i64))
|
||||||
|
.unwrap_or(101)
|
||||||
|
.max(0) as usize;
|
||||||
|
|
||||||
|
// Look up the cursor.
|
||||||
|
let mut cursor_entry = ctx.cursors.get_mut(&cursor_id).ok_or_else(|| {
|
||||||
|
CommandError::InvalidArgument(format!("cursor id {} not found", cursor_id))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
let cursor = cursor_entry.value_mut();
|
||||||
|
let start = cursor.position;
|
||||||
|
let end = (start + batch_size).min(cursor.documents.len());
|
||||||
|
let batch: Vec<Bson> = cursor.documents[start..end]
|
||||||
|
.iter()
|
||||||
|
.cloned()
|
||||||
|
.map(Bson::Document)
|
||||||
|
.collect();
|
||||||
|
cursor.position = end;
|
||||||
|
|
||||||
|
let exhausted = cursor.position >= cursor.documents.len();
|
||||||
|
|
||||||
|
// Must drop the mutable reference before removing.
|
||||||
|
drop(cursor_entry);
|
||||||
|
|
||||||
|
if exhausted {
|
||||||
|
ctx.cursors.remove(&cursor_id);
|
||||||
|
Ok(doc! {
|
||||||
|
"cursor": {
|
||||||
|
"nextBatch": batch,
|
||||||
|
"id": 0_i64,
|
||||||
|
"ns": &ns,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
Ok(doc! {
|
||||||
|
"cursor": {
|
||||||
|
"nextBatch": batch,
|
||||||
|
"id": cursor_id,
|
||||||
|
"ns": &ns,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// killCursors
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Handle the `killCursors` command.
|
||||||
|
pub async fn handle_kill_cursors(
|
||||||
|
cmd: &Document,
|
||||||
|
_db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let cursor_ids = match cmd.get("cursors") {
|
||||||
|
Some(Bson::Array(arr)) => arr,
|
||||||
|
_ => {
|
||||||
|
return Ok(doc! {
|
||||||
|
"cursorsKilled": [],
|
||||||
|
"cursorsNotFound": [],
|
||||||
|
"cursorsAlive": [],
|
||||||
|
"cursorsUnknown": [],
|
||||||
|
"ok": 1.0,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut killed: Vec<Bson> = Vec::new();
|
||||||
|
let mut not_found: Vec<Bson> = Vec::new();
|
||||||
|
|
||||||
|
for id_bson in cursor_ids {
|
||||||
|
let id = match id_bson {
|
||||||
|
Bson::Int64(v) => *v,
|
||||||
|
Bson::Int32(v) => *v as i64,
|
||||||
|
_ => continue,
|
||||||
|
};
|
||||||
|
if ctx.cursors.remove(&id).is_some() {
|
||||||
|
killed.push(Bson::Int64(id));
|
||||||
|
} else {
|
||||||
|
not_found.push(Bson::Int64(id));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"cursorsKilled": killed,
|
||||||
|
"cursorsNotFound": not_found,
|
||||||
|
"cursorsAlive": [],
|
||||||
|
"cursorsUnknown": [],
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// count
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Handle the `count` command.
|
||||||
|
pub async fn handle_count(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let coll = get_str(cmd, "count").unwrap_or("unknown");
|
||||||
|
let txn_id = transactions::active_transaction_id(ctx, cmd);
|
||||||
|
|
||||||
|
// Check collection existence.
|
||||||
|
let exists = if txn_id.is_some() {
|
||||||
|
true
|
||||||
|
} else {
|
||||||
|
ctx.storage.collection_exists(db, coll).await?
|
||||||
|
};
|
||||||
|
if !exists {
|
||||||
|
return Ok(doc! { "n": 0_i64, "ok": 1.0 });
|
||||||
|
}
|
||||||
|
|
||||||
|
let query = get_document(cmd, "query").cloned().unwrap_or_default();
|
||||||
|
let skip = get_i64(cmd, "skip").unwrap_or(0).max(0) as usize;
|
||||||
|
let limit = get_i64(cmd, "limit").unwrap_or(0).max(0) as usize;
|
||||||
|
|
||||||
|
if let Some(ref txn_id) = txn_id {
|
||||||
|
let docs = transactions::load_transaction_docs(ctx, txn_id, db, coll).await?;
|
||||||
|
let filtered = if query.is_empty() {
|
||||||
|
docs
|
||||||
|
} else {
|
||||||
|
QueryMatcher::filter(&docs, &query)
|
||||||
|
};
|
||||||
|
let mut n = filtered.len().saturating_sub(skip);
|
||||||
|
if limit > 0 {
|
||||||
|
n = n.min(limit);
|
||||||
|
}
|
||||||
|
return Ok(doc! {
|
||||||
|
"n": n as i64,
|
||||||
|
"ok": 1.0,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let count: u64 = if query.is_empty() && skip == 0 && limit == 0 {
|
||||||
|
// Fast path: use storage-level count.
|
||||||
|
ctx.storage.count(db, coll).await?
|
||||||
|
} else if query.is_empty() {
|
||||||
|
// No filter but skip/limit apply.
|
||||||
|
let total = ctx.storage.count(db, coll).await? as usize;
|
||||||
|
let after_skip = total.saturating_sub(skip);
|
||||||
|
let result = if limit > 0 { after_skip.min(limit) } else { after_skip };
|
||||||
|
result as u64
|
||||||
|
} else {
|
||||||
|
// Need to load and filter.
|
||||||
|
let docs = ctx.storage.find_all(db, coll).await?;
|
||||||
|
let filtered = QueryMatcher::filter(&docs, &query);
|
||||||
|
let mut n = filtered.len();
|
||||||
|
// Apply skip.
|
||||||
|
n = n.saturating_sub(skip);
|
||||||
|
// Apply limit.
|
||||||
|
if limit > 0 {
|
||||||
|
n = n.min(limit);
|
||||||
|
}
|
||||||
|
n as u64
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"n": count as i64,
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// distinct
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Handle the `distinct` command.
|
||||||
|
pub async fn handle_distinct(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let coll = get_str(cmd, "distinct").unwrap_or("unknown");
|
||||||
|
let key = get_str(cmd, "key").ok_or_else(|| {
|
||||||
|
CommandError::InvalidArgument("distinct requires a 'key' field".into())
|
||||||
|
})?;
|
||||||
|
let txn_id = transactions::active_transaction_id(ctx, cmd);
|
||||||
|
|
||||||
|
// Check collection existence.
|
||||||
|
let exists = if txn_id.is_some() {
|
||||||
|
true
|
||||||
|
} else {
|
||||||
|
ctx.storage.collection_exists(db, coll).await?
|
||||||
|
};
|
||||||
|
if !exists {
|
||||||
|
return Ok(doc! { "values": [], "ok": 1.0 });
|
||||||
|
}
|
||||||
|
|
||||||
|
let query = get_document(cmd, "query").cloned();
|
||||||
|
let docs = if let Some(txn_id) = txn_id {
|
||||||
|
transactions::load_transaction_docs(ctx, &txn_id, db, coll).await?
|
||||||
|
} else {
|
||||||
|
ctx.storage.find_all(db, coll).await?
|
||||||
|
};
|
||||||
|
let values = distinct_values(&docs, key, query.as_ref());
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"values": values,
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
@@ -0,0 +1,43 @@
|
|||||||
|
use bson::{doc, Bson, Document};
|
||||||
|
|
||||||
|
use crate::context::CommandContext;
|
||||||
|
use crate::error::CommandResult;
|
||||||
|
|
||||||
|
/// Handle `hello`, `ismaster`, and `isMaster` commands.
|
||||||
|
///
|
||||||
|
/// Returns server capabilities matching wire protocol expectations.
|
||||||
|
pub async fn handle(
|
||||||
|
cmd: &Document,
|
||||||
|
_db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let mut response = doc! {
|
||||||
|
"ismaster": true,
|
||||||
|
"helloOk": true,
|
||||||
|
"isWritablePrimary": true,
|
||||||
|
"maxBsonObjectSize": 16_777_216_i32,
|
||||||
|
"maxMessageSizeBytes": 48_000_000_i32,
|
||||||
|
"maxWriteBatchSize": 100_000_i32,
|
||||||
|
"localTime": bson::DateTime::now(),
|
||||||
|
"logicalSessionTimeoutMinutes": 30_i32,
|
||||||
|
"connectionId": 1_i32,
|
||||||
|
"minWireVersion": 0_i32,
|
||||||
|
"maxWireVersion": 21_i32,
|
||||||
|
"readOnly": false,
|
||||||
|
"ok": 1.0,
|
||||||
|
};
|
||||||
|
|
||||||
|
if ctx.auth.enabled() {
|
||||||
|
if let Ok(namespace_user) = cmd.get_str("saslSupportedMechs") {
|
||||||
|
let mechanisms: Vec<Bson> = ctx
|
||||||
|
.auth
|
||||||
|
.supported_mechanisms(namespace_user)
|
||||||
|
.into_iter()
|
||||||
|
.map(Bson::String)
|
||||||
|
.collect();
|
||||||
|
response.insert("saslSupportedMechs", Bson::Array(mechanisms));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(response)
|
||||||
|
}
|
||||||
@@ -0,0 +1,380 @@
|
|||||||
|
use bson::{doc, Bson, Document};
|
||||||
|
use rustdb_index::{IndexEngine, IndexOptions};
|
||||||
|
use tracing::debug;
|
||||||
|
|
||||||
|
use crate::context::CommandContext;
|
||||||
|
use crate::error::{CommandError, CommandResult};
|
||||||
|
|
||||||
|
/// Handle `createIndexes`, `dropIndexes`, and `listIndexes` commands.
|
||||||
|
pub async fn handle(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
command_name: &str,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
match command_name {
|
||||||
|
"createIndexes" => handle_create_indexes(cmd, db, ctx).await,
|
||||||
|
"dropIndexes" => handle_drop_indexes(cmd, db, ctx).await,
|
||||||
|
"listIndexes" => handle_list_indexes(cmd, db, ctx).await,
|
||||||
|
_ => Ok(doc! { "ok": 1.0 }),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle the `createIndexes` command.
|
||||||
|
async fn handle_create_indexes(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let coll = cmd
|
||||||
|
.get_str("createIndexes")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'createIndexes' field".into()))?;
|
||||||
|
|
||||||
|
let indexes = cmd
|
||||||
|
.get_array("indexes")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'indexes' array".into()))?;
|
||||||
|
|
||||||
|
let ns_key = format!("{}.{}", db, coll);
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
db = db,
|
||||||
|
collection = coll,
|
||||||
|
count = indexes.len(),
|
||||||
|
"createIndexes command"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Auto-create collection if needed.
|
||||||
|
let created_automatically = ensure_collection_exists(db, coll, ctx).await?;
|
||||||
|
|
||||||
|
// Get the number of indexes before creating new ones.
|
||||||
|
let num_before = {
|
||||||
|
let engine = ctx
|
||||||
|
.indexes
|
||||||
|
.entry(ns_key.clone())
|
||||||
|
.or_insert_with(IndexEngine::new);
|
||||||
|
engine.list_indexes().len() as i32
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut created_count = 0_i32;
|
||||||
|
|
||||||
|
for index_bson in indexes {
|
||||||
|
let index_spec = match index_bson {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => {
|
||||||
|
return Err(CommandError::InvalidArgument(
|
||||||
|
"index spec must be a document".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let key = match index_spec.get("key") {
|
||||||
|
Some(Bson::Document(k)) => k.clone(),
|
||||||
|
_ => {
|
||||||
|
return Err(CommandError::InvalidArgument(
|
||||||
|
"index spec must have a 'key' document".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let name = index_spec.get_str("name").ok().map(|s| s.to_string());
|
||||||
|
|
||||||
|
let unique = match index_spec.get("unique") {
|
||||||
|
Some(Bson::Boolean(b)) => *b,
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let sparse = match index_spec.get("sparse") {
|
||||||
|
Some(Bson::Boolean(b)) => *b,
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let expire_after_seconds = match index_spec.get("expireAfterSeconds") {
|
||||||
|
Some(Bson::Int32(n)) => Some(*n as u64),
|
||||||
|
Some(Bson::Int64(n)) => Some(*n as u64),
|
||||||
|
_ => None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let options = IndexOptions {
|
||||||
|
name,
|
||||||
|
unique,
|
||||||
|
sparse,
|
||||||
|
expire_after_seconds,
|
||||||
|
};
|
||||||
|
|
||||||
|
let options_for_persist = IndexOptions {
|
||||||
|
name: options.name.clone(),
|
||||||
|
unique: options.unique,
|
||||||
|
sparse: options.sparse,
|
||||||
|
expire_after_seconds: options.expire_after_seconds,
|
||||||
|
};
|
||||||
|
let key_for_persist = key.clone();
|
||||||
|
|
||||||
|
// Create the index in-memory.
|
||||||
|
let mut engine = ctx
|
||||||
|
.indexes
|
||||||
|
.entry(ns_key.clone())
|
||||||
|
.or_insert_with(IndexEngine::new);
|
||||||
|
|
||||||
|
match engine.create_index(key, options) {
|
||||||
|
Ok(index_name) => {
|
||||||
|
debug!(index_name = %index_name, "Created index");
|
||||||
|
|
||||||
|
// Persist index spec to disk.
|
||||||
|
let mut spec = doc! { "key": key_for_persist };
|
||||||
|
if options_for_persist.unique {
|
||||||
|
spec.insert("unique", true);
|
||||||
|
}
|
||||||
|
if options_for_persist.sparse {
|
||||||
|
spec.insert("sparse", true);
|
||||||
|
}
|
||||||
|
if let Some(ttl) = options_for_persist.expire_after_seconds {
|
||||||
|
spec.insert("expireAfterSeconds", ttl as i64);
|
||||||
|
}
|
||||||
|
if let Err(e) = ctx.storage.save_index(db, coll, &index_name, spec).await {
|
||||||
|
tracing::warn!(index = %index_name, error = %e, "failed to persist index spec");
|
||||||
|
}
|
||||||
|
|
||||||
|
created_count += 1;
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
return Err(CommandError::IndexError(e.to_string()));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// If we created indexes on an existing collection, rebuild from documents.
|
||||||
|
if created_count > 0 && !created_automatically {
|
||||||
|
// Load all documents and rebuild indexes.
|
||||||
|
if let Ok(all_docs) = ctx.storage.find_all(db, coll).await {
|
||||||
|
if !all_docs.is_empty() {
|
||||||
|
let mut engine = ctx
|
||||||
|
.indexes
|
||||||
|
.entry(ns_key.clone())
|
||||||
|
.or_insert_with(IndexEngine::new);
|
||||||
|
engine.rebuild_from_documents(&all_docs);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let num_after = {
|
||||||
|
let engine = ctx
|
||||||
|
.indexes
|
||||||
|
.entry(ns_key.clone())
|
||||||
|
.or_insert_with(IndexEngine::new);
|
||||||
|
engine.list_indexes().len() as i32
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"createdCollectionAutomatically": created_automatically,
|
||||||
|
"numIndexesBefore": num_before,
|
||||||
|
"numIndexesAfter": num_after,
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle the `dropIndexes` command.
|
||||||
|
async fn handle_drop_indexes(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let coll = cmd
|
||||||
|
.get_str("dropIndexes")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'dropIndexes' field".into()))?;
|
||||||
|
|
||||||
|
let ns_key = format!("{}.{}", db, coll);
|
||||||
|
|
||||||
|
// Get current index count.
|
||||||
|
let n_indexes_was = {
|
||||||
|
match ctx.indexes.get(&ns_key) {
|
||||||
|
Some(engine) => engine.list_indexes().len() as i32,
|
||||||
|
None => 1_i32, // At minimum the _id_ index.
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let index_spec = cmd.get("index");
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
db = db,
|
||||||
|
collection = coll,
|
||||||
|
index_spec = ?index_spec,
|
||||||
|
"dropIndexes command"
|
||||||
|
);
|
||||||
|
|
||||||
|
match index_spec {
|
||||||
|
Some(Bson::String(name)) if name == "*" => {
|
||||||
|
// Drop all indexes except _id_.
|
||||||
|
// Collect names to drop from storage first.
|
||||||
|
let names_to_drop: Vec<String> = if let Some(engine) = ctx.indexes.get(&ns_key) {
|
||||||
|
engine.list_indexes().iter()
|
||||||
|
.filter(|info| info.name != "_id_")
|
||||||
|
.map(|info| info.name.clone())
|
||||||
|
.collect()
|
||||||
|
} else {
|
||||||
|
Vec::new()
|
||||||
|
};
|
||||||
|
if let Some(mut engine) = ctx.indexes.get_mut(&ns_key) {
|
||||||
|
engine.drop_all_indexes();
|
||||||
|
}
|
||||||
|
for idx_name in &names_to_drop {
|
||||||
|
let _ = ctx.storage.drop_index(db, coll, idx_name).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(Bson::String(name)) => {
|
||||||
|
// Drop by name.
|
||||||
|
if let Some(mut engine) = ctx.indexes.get_mut(&ns_key) {
|
||||||
|
engine.drop_index(name).map_err(|e| {
|
||||||
|
CommandError::IndexError(e.to_string())
|
||||||
|
})?;
|
||||||
|
} else {
|
||||||
|
return Err(CommandError::IndexError(format!(
|
||||||
|
"index not found: {}",
|
||||||
|
name
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
let _ = ctx.storage.drop_index(db, coll, name).await;
|
||||||
|
}
|
||||||
|
Some(Bson::Document(key_spec)) => {
|
||||||
|
// Drop by key spec: find the index with matching key.
|
||||||
|
if let Some(mut engine) = ctx.indexes.get_mut(&ns_key) {
|
||||||
|
let index_name = engine
|
||||||
|
.list_indexes()
|
||||||
|
.iter()
|
||||||
|
.find(|info| info.key == *key_spec)
|
||||||
|
.map(|info| info.name.clone());
|
||||||
|
|
||||||
|
if let Some(name) = index_name {
|
||||||
|
engine.drop_index(&name).map_err(|e| {
|
||||||
|
CommandError::IndexError(e.to_string())
|
||||||
|
})?;
|
||||||
|
let _ = ctx.storage.drop_index(db, coll, &name).await;
|
||||||
|
} else {
|
||||||
|
return Err(CommandError::IndexError(
|
||||||
|
"index not found with specified key".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
return Err(CommandError::IndexError(
|
||||||
|
"no indexes found for collection".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
return Err(CommandError::InvalidArgument(
|
||||||
|
"dropIndexes requires 'index' field (string, document, or \"*\")".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"nIndexesWas": n_indexes_was,
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle the `listIndexes` command.
|
||||||
|
async fn handle_list_indexes(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let coll = cmd
|
||||||
|
.get_str("listIndexes")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'listIndexes' field".into()))?;
|
||||||
|
|
||||||
|
let ns_key = format!("{}.{}", db, coll);
|
||||||
|
let ns = format!("{}.{}", db, coll);
|
||||||
|
|
||||||
|
// Check if collection exists.
|
||||||
|
match ctx.storage.collection_exists(db, coll).await {
|
||||||
|
Ok(false) => {
|
||||||
|
return Err(CommandError::NamespaceNotFound(format!(
|
||||||
|
"ns not found: {}",
|
||||||
|
ns
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
Err(_) => {
|
||||||
|
// If we can't check, try to proceed anyway.
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
|
||||||
|
let indexes = match ctx.indexes.get(&ns_key) {
|
||||||
|
Some(engine) => engine.list_indexes(),
|
||||||
|
None => {
|
||||||
|
// Return at least the default _id_ index.
|
||||||
|
let engine = IndexEngine::new();
|
||||||
|
engine.list_indexes()
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let first_batch: Vec<Bson> = indexes
|
||||||
|
.into_iter()
|
||||||
|
.map(|info| {
|
||||||
|
let mut doc = doc! {
|
||||||
|
"v": info.v,
|
||||||
|
"key": info.key,
|
||||||
|
"name": info.name,
|
||||||
|
};
|
||||||
|
if info.unique {
|
||||||
|
doc.insert("unique", true);
|
||||||
|
}
|
||||||
|
if info.sparse {
|
||||||
|
doc.insert("sparse", true);
|
||||||
|
}
|
||||||
|
if let Some(ttl) = info.expire_after_seconds {
|
||||||
|
doc.insert("expireAfterSeconds", ttl as i64);
|
||||||
|
}
|
||||||
|
Bson::Document(doc)
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"cursor": {
|
||||||
|
"id": 0_i64,
|
||||||
|
"ns": &ns,
|
||||||
|
"firstBatch": first_batch,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Ensure the target database and collection exist. Returns true if the collection
|
||||||
|
/// was newly created (i.e., `createdCollectionAutomatically`).
|
||||||
|
async fn ensure_collection_exists(
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<bool> {
|
||||||
|
// Create database (ignore AlreadyExists).
|
||||||
|
if let Err(e) = ctx.storage.create_database(db).await {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
|
||||||
|
return Err(CommandError::StorageError(msg));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if collection exists.
|
||||||
|
match ctx.storage.collection_exists(db, coll).await {
|
||||||
|
Ok(true) => Ok(false),
|
||||||
|
Ok(false) => {
|
||||||
|
if let Err(e) = ctx.storage.create_collection(db, coll).await {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
|
||||||
|
return Err(CommandError::StorageError(msg));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(true)
|
||||||
|
}
|
||||||
|
Err(_) => {
|
||||||
|
// Try creating anyway.
|
||||||
|
if let Err(e) = ctx.storage.create_collection(db, coll).await {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
|
||||||
|
return Err(CommandError::StorageError(msg));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(true)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,240 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
use bson::{doc, oid::ObjectId, Bson, Document};
|
||||||
|
use rustdb_storage::OpType;
|
||||||
|
use tracing::debug;
|
||||||
|
|
||||||
|
use crate::context::CommandContext;
|
||||||
|
use crate::error::{CommandError, CommandResult};
|
||||||
|
use crate::transactions;
|
||||||
|
|
||||||
|
/// Handle the `insert` command.
|
||||||
|
pub async fn handle(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
document_sequences: Option<&HashMap<String, Vec<Document>>>,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let coll = cmd
|
||||||
|
.get_str("insert")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'insert' field".into()))?;
|
||||||
|
|
||||||
|
// Determine whether writes are ordered (default: true).
|
||||||
|
let ordered = match cmd.get("ordered") {
|
||||||
|
Some(Bson::Boolean(b)) => *b,
|
||||||
|
_ => true,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Collect documents from either the command body or OP_MSG document sequences.
|
||||||
|
let docs: Vec<Document> = if let Some(seqs) = document_sequences {
|
||||||
|
if let Some(seq_docs) = seqs.get("documents") {
|
||||||
|
seq_docs.clone()
|
||||||
|
} else {
|
||||||
|
extract_docs_from_array(cmd)?
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
extract_docs_from_array(cmd)?
|
||||||
|
};
|
||||||
|
|
||||||
|
if docs.is_empty() {
|
||||||
|
return Err(CommandError::InvalidArgument(
|
||||||
|
"no documents to insert".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
db = db,
|
||||||
|
collection = coll,
|
||||||
|
count = docs.len(),
|
||||||
|
"insert command"
|
||||||
|
);
|
||||||
|
|
||||||
|
let txn_id = transactions::active_transaction_id(ctx, cmd);
|
||||||
|
|
||||||
|
// Auto-create database and collection if they don't exist. Transactional
|
||||||
|
// writes defer collection creation until commit so abort remains clean.
|
||||||
|
if txn_id.is_none() {
|
||||||
|
ensure_collection_exists(db, coll, ctx).await?;
|
||||||
|
}
|
||||||
|
|
||||||
|
let ns_key = format!("{}.{}", db, coll);
|
||||||
|
let mut inserted_count: i32 = 0;
|
||||||
|
let mut write_errors: Vec<Document> = Vec::new();
|
||||||
|
|
||||||
|
// Ensure the IndexEngine is loaded (with persisted specs from indexes.json).
|
||||||
|
// This must happen BEFORE any writes, so unique constraints are enforced
|
||||||
|
// even on the first write after a restart.
|
||||||
|
drop(ctx.get_or_init_index_engine(db, coll).await);
|
||||||
|
|
||||||
|
for (idx, mut doc) in docs.into_iter().enumerate() {
|
||||||
|
// Auto-generate _id if not present.
|
||||||
|
if !doc.contains_key("_id") {
|
||||||
|
doc.insert("_id", ObjectId::new());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Pre-check unique index constraints BEFORE storage write.
|
||||||
|
// The engine is guaranteed to exist from the get_or_init call above.
|
||||||
|
if let Some(engine) = ctx.indexes.get(&ns_key) {
|
||||||
|
if let Err(e) = engine.check_unique_constraints(&doc) {
|
||||||
|
let err_msg = e.to_string();
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 11000_i32,
|
||||||
|
"codeName": "DuplicateKey",
|
||||||
|
"errmsg": &err_msg,
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(ref txn_id) = txn_id {
|
||||||
|
match transactions::record_insert(ctx, txn_id, db, coll, doc.clone()).await {
|
||||||
|
Ok(_) => inserted_count += 1,
|
||||||
|
Err(e) => {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 11000_i32,
|
||||||
|
"codeName": "DuplicateKey",
|
||||||
|
"errmsg": e.to_string(),
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Attempt storage insert.
|
||||||
|
match ctx.storage.insert_one(db, coll, doc.clone()).await {
|
||||||
|
Ok(id_str) => {
|
||||||
|
// Record in oplog.
|
||||||
|
ctx.oplog.append(
|
||||||
|
OpType::Insert,
|
||||||
|
db,
|
||||||
|
coll,
|
||||||
|
&id_str,
|
||||||
|
Some(doc.clone()),
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
|
||||||
|
// Update index engine (already initialized above).
|
||||||
|
if let Some(mut engine) = ctx.indexes.get_mut(&ns_key) {
|
||||||
|
if let Err(e) = engine.on_insert(&doc) {
|
||||||
|
tracing::error!(
|
||||||
|
namespace = %ns_key,
|
||||||
|
error = %e,
|
||||||
|
"index update failed after successful insert"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
inserted_count += 1;
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
let err_msg = e.to_string();
|
||||||
|
let (code, code_name) = if err_msg.contains("AlreadyExists")
|
||||||
|
|| err_msg.contains("duplicate")
|
||||||
|
{
|
||||||
|
(11000_i32, "DuplicateKey")
|
||||||
|
} else {
|
||||||
|
(1_i32, "InternalError")
|
||||||
|
};
|
||||||
|
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": code,
|
||||||
|
"codeName": code_name,
|
||||||
|
"errmsg": &err_msg,
|
||||||
|
});
|
||||||
|
|
||||||
|
if ordered {
|
||||||
|
// Stop on first error when ordered.
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build response document.
|
||||||
|
let mut response = doc! {
|
||||||
|
"n": inserted_count,
|
||||||
|
"ok": 1.0,
|
||||||
|
};
|
||||||
|
|
||||||
|
if !write_errors.is_empty() {
|
||||||
|
response.insert(
|
||||||
|
"writeErrors",
|
||||||
|
write_errors
|
||||||
|
.into_iter()
|
||||||
|
.map(Bson::Document)
|
||||||
|
.collect::<Vec<_>>(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(response)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extract documents from the `documents` array field in the command BSON.
|
||||||
|
fn extract_docs_from_array(cmd: &Document) -> CommandResult<Vec<Document>> {
|
||||||
|
match cmd.get_array("documents") {
|
||||||
|
Ok(arr) => {
|
||||||
|
let mut docs = Vec::with_capacity(arr.len());
|
||||||
|
for item in arr {
|
||||||
|
match item {
|
||||||
|
Bson::Document(d) => docs.push(d.clone()),
|
||||||
|
_ => {
|
||||||
|
return Err(CommandError::InvalidArgument(
|
||||||
|
"documents array contains non-document element".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(docs)
|
||||||
|
}
|
||||||
|
Err(_) => Ok(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Ensure the target database and collection exist, creating them if needed.
|
||||||
|
async fn ensure_collection_exists(
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<()> {
|
||||||
|
// Create database (no-op if it already exists in most backends).
|
||||||
|
if let Err(e) = ctx.storage.create_database(db).await {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
|
||||||
|
return Err(CommandError::StorageError(msg));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create collection if it doesn't exist.
|
||||||
|
match ctx.storage.collection_exists(db, coll).await {
|
||||||
|
Ok(true) => {}
|
||||||
|
Ok(false) => {
|
||||||
|
if let Err(e) = ctx.storage.create_collection(db, coll).await {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
|
||||||
|
return Err(CommandError::StorageError(msg));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
// Database might not exist yet; try creating collection anyway.
|
||||||
|
if let Err(e2) = ctx.storage.create_collection(db, coll).await {
|
||||||
|
let msg = e2.to_string();
|
||||||
|
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
|
||||||
|
return Err(CommandError::StorageError(format!(
|
||||||
|
"collection_exists failed: {e}; create_collection failed: {msg}"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
@@ -0,0 +1,9 @@
|
|||||||
|
pub mod admin_handler;
|
||||||
|
pub mod aggregate_handler;
|
||||||
|
pub mod auth_handler;
|
||||||
|
pub mod delete_handler;
|
||||||
|
pub mod find_handler;
|
||||||
|
pub mod hello_handler;
|
||||||
|
pub mod index_handler;
|
||||||
|
pub mod insert_handler;
|
||||||
|
pub mod update_handler;
|
||||||
@@ -0,0 +1,947 @@
|
|||||||
|
use std::collections::HashSet;
|
||||||
|
|
||||||
|
use bson::{doc, oid::ObjectId, Bson, Document};
|
||||||
|
use rustdb_query::{QueryMatcher, UpdateEngine, sort_documents, apply_projection};
|
||||||
|
use rustdb_storage::OpType;
|
||||||
|
use tracing::debug;
|
||||||
|
|
||||||
|
use crate::context::CommandContext;
|
||||||
|
use crate::error::{CommandError, CommandResult};
|
||||||
|
use crate::transactions;
|
||||||
|
|
||||||
|
/// Handle `update` and `findAndModify` commands.
|
||||||
|
pub async fn handle(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
command_name: &str,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
match command_name {
|
||||||
|
"findAndModify" | "findandmodify" => handle_find_and_modify(cmd, db, ctx).await,
|
||||||
|
_ => handle_update(cmd, db, ctx).await,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
enum TUpdateSpec {
|
||||||
|
Document(Document),
|
||||||
|
Pipeline(Vec<Document>),
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle the `update` command.
|
||||||
|
async fn handle_update(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let coll = cmd
|
||||||
|
.get_str("update")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'update' field".into()))?;
|
||||||
|
|
||||||
|
let updates = cmd
|
||||||
|
.get_array("updates")
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'updates' array".into()))?;
|
||||||
|
|
||||||
|
let ordered = match cmd.get("ordered") {
|
||||||
|
Some(Bson::Boolean(b)) => *b,
|
||||||
|
_ => true,
|
||||||
|
};
|
||||||
|
|
||||||
|
debug!(db = db, collection = coll, count = updates.len(), "update command");
|
||||||
|
|
||||||
|
let txn_id = transactions::active_transaction_id(ctx, cmd);
|
||||||
|
|
||||||
|
// Transactional writes defer namespace creation until commit.
|
||||||
|
if txn_id.is_none() {
|
||||||
|
ensure_collection_exists(db, coll, ctx).await?;
|
||||||
|
}
|
||||||
|
|
||||||
|
let ns_key = format!("{}.{}", db, coll);
|
||||||
|
|
||||||
|
// Ensure the IndexEngine is loaded with persisted specs from indexes.json.
|
||||||
|
drop(ctx.get_or_init_index_engine(db, coll).await);
|
||||||
|
|
||||||
|
let mut total_n: i32 = 0;
|
||||||
|
let mut total_n_modified: i32 = 0;
|
||||||
|
let mut upserted_list: Vec<Document> = Vec::new();
|
||||||
|
let mut write_errors: Vec<Document> = Vec::new();
|
||||||
|
|
||||||
|
for (idx, update_bson) in updates.iter().enumerate() {
|
||||||
|
let update_spec = match update_bson {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 14_i32,
|
||||||
|
"codeName": "TypeMismatch",
|
||||||
|
"errmsg": "update spec must be a document",
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let filter = match update_spec.get("q") {
|
||||||
|
Some(Bson::Document(d)) => d.clone(),
|
||||||
|
_ => Document::new(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let update = match update_spec.get("u") {
|
||||||
|
Some(update_value) => match parse_update_spec(update_value) {
|
||||||
|
Ok(parsed) => parsed,
|
||||||
|
Err(err) => {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 14_i32,
|
||||||
|
"codeName": "TypeMismatch",
|
||||||
|
"errmsg": err,
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
},
|
||||||
|
None => {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 14_i32,
|
||||||
|
"codeName": "TypeMismatch",
|
||||||
|
"errmsg": "missing or invalid 'u' field in update spec",
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let multi = match update_spec.get("multi") {
|
||||||
|
Some(Bson::Boolean(b)) => *b,
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let upsert = match update_spec.get("upsert") {
|
||||||
|
Some(Bson::Boolean(b)) => *b,
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let array_filters: Option<Vec<Document>> =
|
||||||
|
update_spec.get_array("arrayFilters").ok().map(|arr| {
|
||||||
|
arr.iter()
|
||||||
|
.filter_map(|v| {
|
||||||
|
if let Bson::Document(d) = v {
|
||||||
|
Some(d.clone())
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
});
|
||||||
|
|
||||||
|
// Load all documents and filter.
|
||||||
|
let all_docs = load_filtered_docs(db, coll, &filter, &ns_key, ctx, txn_id.as_deref()).await?;
|
||||||
|
|
||||||
|
if all_docs.is_empty() && upsert {
|
||||||
|
// Upsert: create a new document.
|
||||||
|
let new_doc = build_upsert_doc(&filter);
|
||||||
|
|
||||||
|
// Apply update operators or replacement.
|
||||||
|
match apply_update_spec(&new_doc, &update, array_filters.as_deref()) {
|
||||||
|
Ok(mut updated) => {
|
||||||
|
apply_set_on_insert_if_present(&update, &mut updated);
|
||||||
|
|
||||||
|
// Ensure _id exists.
|
||||||
|
let new_id = ensure_document_id(&mut updated);
|
||||||
|
|
||||||
|
// Pre-check unique index constraints before upsert insert.
|
||||||
|
if let Some(engine) = ctx.indexes.get(&ns_key) {
|
||||||
|
if let Err(e) = engine.check_unique_constraints(&updated) {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 11000_i32,
|
||||||
|
"codeName": "DuplicateKey",
|
||||||
|
"errmsg": e.to_string(),
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(ref txn_id) = txn_id {
|
||||||
|
match transactions::record_insert(ctx, txn_id, db, coll, updated.clone()).await {
|
||||||
|
Ok(_) => {
|
||||||
|
total_n += 1;
|
||||||
|
upserted_list.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"_id": new_id,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 1_i32,
|
||||||
|
"codeName": "InternalError",
|
||||||
|
"errmsg": e.to_string(),
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Insert the new document.
|
||||||
|
match ctx.storage.insert_one(db, coll, updated.clone()).await {
|
||||||
|
Ok(id_str) => {
|
||||||
|
// Record upsert in oplog as an insert.
|
||||||
|
ctx.oplog.append(
|
||||||
|
OpType::Insert,
|
||||||
|
db,
|
||||||
|
coll,
|
||||||
|
&id_str,
|
||||||
|
Some(updated.clone()),
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
|
||||||
|
// Update index (engine already initialized above).
|
||||||
|
if let Some(mut engine) = ctx.indexes.get_mut(&ns_key) {
|
||||||
|
if let Err(e) = engine.on_insert(&updated) {
|
||||||
|
tracing::error!(namespace = %ns_key, error = %e, "index update failed after upsert insert");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
total_n += 1;
|
||||||
|
upserted_list.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"_id": new_id,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 1_i32,
|
||||||
|
"codeName": "InternalError",
|
||||||
|
"errmsg": e.to_string(),
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 14_i32,
|
||||||
|
"codeName": "TypeMismatch",
|
||||||
|
"errmsg": e.to_string(),
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// Update matched documents.
|
||||||
|
let docs_to_update = if multi {
|
||||||
|
all_docs
|
||||||
|
} else {
|
||||||
|
all_docs.into_iter().take(1).collect()
|
||||||
|
};
|
||||||
|
|
||||||
|
for matched_doc in &docs_to_update {
|
||||||
|
match apply_update_spec(matched_doc, &update, array_filters.as_deref()) {
|
||||||
|
Ok(mut updated_doc) => {
|
||||||
|
if let Err(e) = ensure_immutable_id(matched_doc, &mut updated_doc) {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 66_i32,
|
||||||
|
"codeName": "ImmutableField",
|
||||||
|
"errmsg": e.to_string(),
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Pre-check unique index constraints before storage write.
|
||||||
|
if let Some(engine) = ctx.indexes.get(&ns_key) {
|
||||||
|
if let Err(e) = engine.check_unique_constraints_for_update(matched_doc, &updated_doc) {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 11000_i32,
|
||||||
|
"codeName": "DuplicateKey",
|
||||||
|
"errmsg": e.to_string(),
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let id_str = extract_id_string(matched_doc);
|
||||||
|
if let Some(ref txn_id) = txn_id {
|
||||||
|
match transactions::record_update(
|
||||||
|
ctx,
|
||||||
|
txn_id,
|
||||||
|
db,
|
||||||
|
coll,
|
||||||
|
matched_doc.clone(),
|
||||||
|
updated_doc.clone(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(_) => {
|
||||||
|
total_n += 1;
|
||||||
|
if matched_doc != &updated_doc {
|
||||||
|
total_n_modified += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 1_i32,
|
||||||
|
"codeName": "InternalError",
|
||||||
|
"errmsg": e.to_string(),
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
match ctx
|
||||||
|
.storage
|
||||||
|
.update_by_id(db, coll, &id_str, updated_doc.clone())
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(()) => {
|
||||||
|
// Record in oplog.
|
||||||
|
ctx.oplog.append(
|
||||||
|
OpType::Update,
|
||||||
|
db,
|
||||||
|
coll,
|
||||||
|
&id_str,
|
||||||
|
Some(updated_doc.clone()),
|
||||||
|
Some(matched_doc.clone()),
|
||||||
|
);
|
||||||
|
|
||||||
|
// Update index.
|
||||||
|
if let Some(mut engine) = ctx.indexes.get_mut(&ns_key) {
|
||||||
|
if let Err(e) = engine.on_update(matched_doc, &updated_doc) {
|
||||||
|
tracing::error!(namespace = %ns_key, error = %e, "index update failed after update");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
total_n += 1;
|
||||||
|
// Check if the document actually changed.
|
||||||
|
if matched_doc != &updated_doc {
|
||||||
|
total_n_modified += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 1_i32,
|
||||||
|
"codeName": "InternalError",
|
||||||
|
"errmsg": e.to_string(),
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
write_errors.push(doc! {
|
||||||
|
"index": idx as i32,
|
||||||
|
"code": 14_i32,
|
||||||
|
"codeName": "TypeMismatch",
|
||||||
|
"errmsg": e.to_string(),
|
||||||
|
});
|
||||||
|
if ordered {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build response.
|
||||||
|
let mut response = doc! {
|
||||||
|
"n": total_n,
|
||||||
|
"nModified": total_n_modified,
|
||||||
|
"ok": 1.0,
|
||||||
|
};
|
||||||
|
|
||||||
|
if !upserted_list.is_empty() {
|
||||||
|
response.insert(
|
||||||
|
"upserted",
|
||||||
|
upserted_list
|
||||||
|
.into_iter()
|
||||||
|
.map(Bson::Document)
|
||||||
|
.collect::<Vec<_>>(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
if !write_errors.is_empty() {
|
||||||
|
response.insert(
|
||||||
|
"writeErrors",
|
||||||
|
write_errors
|
||||||
|
.into_iter()
|
||||||
|
.map(Bson::Document)
|
||||||
|
.collect::<Vec<_>>(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(response)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle the `findAndModify` command.
|
||||||
|
async fn handle_find_and_modify(
|
||||||
|
cmd: &Document,
|
||||||
|
db: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let coll = cmd
|
||||||
|
.get_str("findAndModify")
|
||||||
|
.or_else(|_| cmd.get_str("findandmodify"))
|
||||||
|
.map_err(|_| CommandError::InvalidArgument("missing 'findAndModify' field".into()))?;
|
||||||
|
|
||||||
|
let query = match cmd.get("query") {
|
||||||
|
Some(Bson::Document(d)) => d.clone(),
|
||||||
|
_ => Document::new(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let sort = match cmd.get("sort") {
|
||||||
|
Some(Bson::Document(d)) => Some(d.clone()),
|
||||||
|
_ => None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let update_doc = match cmd.get("update") {
|
||||||
|
Some(update_value) => Some(
|
||||||
|
parse_update_spec(update_value)
|
||||||
|
.map_err(CommandError::InvalidArgument)?
|
||||||
|
),
|
||||||
|
None => None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let remove = match cmd.get("remove") {
|
||||||
|
Some(Bson::Boolean(b)) => *b,
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let return_new = match cmd.get("new") {
|
||||||
|
Some(Bson::Boolean(b)) => *b,
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let upsert = match cmd.get("upsert") {
|
||||||
|
Some(Bson::Boolean(b)) => *b,
|
||||||
|
_ => false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let fields = match cmd.get("fields") {
|
||||||
|
Some(Bson::Document(d)) => Some(d.clone()),
|
||||||
|
_ => None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let array_filters: Option<Vec<Document>> =
|
||||||
|
cmd.get_array("arrayFilters").ok().map(|arr| {
|
||||||
|
arr.iter()
|
||||||
|
.filter_map(|v| {
|
||||||
|
if let Bson::Document(d) = v {
|
||||||
|
Some(d.clone())
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
});
|
||||||
|
|
||||||
|
let txn_id = transactions::active_transaction_id(ctx, cmd);
|
||||||
|
|
||||||
|
// Transactional writes defer namespace creation until commit.
|
||||||
|
if txn_id.is_none() {
|
||||||
|
ensure_collection_exists(db, coll, ctx).await?;
|
||||||
|
}
|
||||||
|
|
||||||
|
let ns_key = format!("{}.{}", db, coll);
|
||||||
|
|
||||||
|
// Ensure the IndexEngine is loaded with persisted specs.
|
||||||
|
drop(ctx.get_or_init_index_engine(db, coll).await);
|
||||||
|
|
||||||
|
// Load and filter documents.
|
||||||
|
let mut matched = load_filtered_docs(db, coll, &query, &ns_key, ctx, txn_id.as_deref()).await?;
|
||||||
|
|
||||||
|
// Sort if specified.
|
||||||
|
if let Some(ref sort_spec) = sort {
|
||||||
|
sort_documents(&mut matched, sort_spec);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Take the first match.
|
||||||
|
let target = matched.into_iter().next();
|
||||||
|
|
||||||
|
if remove {
|
||||||
|
// Remove operation.
|
||||||
|
if let Some(ref doc) = target {
|
||||||
|
let id_str = extract_id_string(doc);
|
||||||
|
if let Some(ref txn_id) = txn_id {
|
||||||
|
transactions::record_delete(ctx, txn_id, db, coll, doc.clone()).await?;
|
||||||
|
|
||||||
|
let value = apply_fields_projection(doc, &fields);
|
||||||
|
|
||||||
|
return Ok(doc! {
|
||||||
|
"value": value,
|
||||||
|
"lastErrorObject": {
|
||||||
|
"n": 1_i32,
|
||||||
|
"updatedExisting": false,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
ctx.storage.delete_by_id(db, coll, &id_str).await?;
|
||||||
|
|
||||||
|
// Record in oplog.
|
||||||
|
ctx.oplog.append(
|
||||||
|
OpType::Delete,
|
||||||
|
db,
|
||||||
|
coll,
|
||||||
|
&id_str,
|
||||||
|
None,
|
||||||
|
Some(doc.clone()),
|
||||||
|
);
|
||||||
|
|
||||||
|
// Update index.
|
||||||
|
if let Some(mut engine) = ctx.indexes.get_mut(&ns_key) {
|
||||||
|
engine.on_delete(doc);
|
||||||
|
}
|
||||||
|
|
||||||
|
let value = apply_fields_projection(doc, &fields);
|
||||||
|
|
||||||
|
return Ok(doc! {
|
||||||
|
"value": value,
|
||||||
|
"lastErrorObject": {
|
||||||
|
"n": 1_i32,
|
||||||
|
"updatedExisting": false,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
});
|
||||||
|
} else {
|
||||||
|
return Ok(doc! {
|
||||||
|
"value": Bson::Null,
|
||||||
|
"lastErrorObject": {
|
||||||
|
"n": 0_i32,
|
||||||
|
"updatedExisting": false,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Update operation.
|
||||||
|
let update = match update_doc {
|
||||||
|
Some(u) => u,
|
||||||
|
None => {
|
||||||
|
return Ok(doc! {
|
||||||
|
"value": Bson::Null,
|
||||||
|
"lastErrorObject": {
|
||||||
|
"n": 0_i32,
|
||||||
|
"updatedExisting": false,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Some(original_doc) = target {
|
||||||
|
// Update the matched document.
|
||||||
|
let mut updated_doc = apply_update_spec(
|
||||||
|
&original_doc,
|
||||||
|
&update,
|
||||||
|
array_filters.as_deref(),
|
||||||
|
)
|
||||||
|
.map_err(CommandError::InternalError)?;
|
||||||
|
|
||||||
|
ensure_immutable_id(&original_doc, &mut updated_doc)?;
|
||||||
|
|
||||||
|
// Pre-check unique index constraints before storage write.
|
||||||
|
if let Some(engine) = ctx.indexes.get(&ns_key) {
|
||||||
|
if let Err(e) = engine.check_unique_constraints_for_update(&original_doc, &updated_doc) {
|
||||||
|
return Err(CommandError::StorageError(e.to_string()));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let id_str = extract_id_string(&original_doc);
|
||||||
|
if let Some(ref txn_id) = txn_id {
|
||||||
|
transactions::record_update(
|
||||||
|
ctx,
|
||||||
|
txn_id,
|
||||||
|
db,
|
||||||
|
coll,
|
||||||
|
original_doc.clone(),
|
||||||
|
updated_doc.clone(),
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let return_doc = if return_new {
|
||||||
|
&updated_doc
|
||||||
|
} else {
|
||||||
|
&original_doc
|
||||||
|
};
|
||||||
|
|
||||||
|
let value = apply_fields_projection(return_doc, &fields);
|
||||||
|
|
||||||
|
return Ok(doc! {
|
||||||
|
"value": value,
|
||||||
|
"lastErrorObject": {
|
||||||
|
"n": 1_i32,
|
||||||
|
"updatedExisting": true,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
ctx.storage
|
||||||
|
.update_by_id(db, coll, &id_str, updated_doc.clone())
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
// Record in oplog.
|
||||||
|
ctx.oplog.append(
|
||||||
|
OpType::Update,
|
||||||
|
db,
|
||||||
|
coll,
|
||||||
|
&id_str,
|
||||||
|
Some(updated_doc.clone()),
|
||||||
|
Some(original_doc.clone()),
|
||||||
|
);
|
||||||
|
|
||||||
|
// Update index.
|
||||||
|
if let Some(mut engine) = ctx.indexes.get_mut(&ns_key) {
|
||||||
|
if let Err(e) = engine.on_update(&original_doc, &updated_doc) {
|
||||||
|
tracing::error!(namespace = %ns_key, error = %e, "index update failed after findAndModify update");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let return_doc = if return_new {
|
||||||
|
&updated_doc
|
||||||
|
} else {
|
||||||
|
&original_doc
|
||||||
|
};
|
||||||
|
|
||||||
|
let value = apply_fields_projection(return_doc, &fields);
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"value": value,
|
||||||
|
"lastErrorObject": {
|
||||||
|
"n": 1_i32,
|
||||||
|
"updatedExisting": true,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
} else if upsert {
|
||||||
|
// Upsert: create a new document.
|
||||||
|
let new_doc = build_upsert_doc(&query);
|
||||||
|
|
||||||
|
let mut updated_doc = apply_update_spec(
|
||||||
|
&new_doc,
|
||||||
|
&update,
|
||||||
|
array_filters.as_deref(),
|
||||||
|
)
|
||||||
|
.map_err(CommandError::InternalError)?;
|
||||||
|
|
||||||
|
apply_set_on_insert_if_present(&update, &mut updated_doc);
|
||||||
|
|
||||||
|
// Ensure _id.
|
||||||
|
let upserted_id = ensure_document_id(&mut updated_doc);
|
||||||
|
|
||||||
|
// Pre-check unique index constraints before upsert insert.
|
||||||
|
if let Some(engine) = ctx.indexes.get(&ns_key) {
|
||||||
|
if let Err(e) = engine.check_unique_constraints(&updated_doc) {
|
||||||
|
return Err(CommandError::StorageError(e.to_string()));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(ref txn_id) = txn_id {
|
||||||
|
transactions::record_insert(ctx, txn_id, db, coll, updated_doc.clone()).await?;
|
||||||
|
|
||||||
|
let value = if return_new {
|
||||||
|
apply_fields_projection(&updated_doc, &fields)
|
||||||
|
} else {
|
||||||
|
Bson::Null
|
||||||
|
};
|
||||||
|
|
||||||
|
return Ok(doc! {
|
||||||
|
"value": value,
|
||||||
|
"lastErrorObject": {
|
||||||
|
"n": 1_i32,
|
||||||
|
"updatedExisting": false,
|
||||||
|
"upserted": upserted_id,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let inserted_id_str = ctx.storage
|
||||||
|
.insert_one(db, coll, updated_doc.clone())
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
// Record upsert in oplog as an insert.
|
||||||
|
ctx.oplog.append(
|
||||||
|
OpType::Insert,
|
||||||
|
db,
|
||||||
|
coll,
|
||||||
|
&inserted_id_str,
|
||||||
|
Some(updated_doc.clone()),
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
|
||||||
|
// Update index.
|
||||||
|
{
|
||||||
|
if let Some(mut engine) = ctx.indexes.get_mut(&ns_key) {
|
||||||
|
if let Err(e) = engine.on_insert(&updated_doc) {
|
||||||
|
tracing::error!(namespace = %ns_key, error = %e, "index update failed after findAndModify upsert");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let value = if return_new {
|
||||||
|
apply_fields_projection(&updated_doc, &fields)
|
||||||
|
} else {
|
||||||
|
Bson::Null
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(doc! {
|
||||||
|
"value": value,
|
||||||
|
"lastErrorObject": {
|
||||||
|
"n": 1_i32,
|
||||||
|
"updatedExisting": false,
|
||||||
|
"upserted": upserted_id,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
Ok(doc! {
|
||||||
|
"value": Bson::Null,
|
||||||
|
"lastErrorObject": {
|
||||||
|
"n": 0_i32,
|
||||||
|
"updatedExisting": false,
|
||||||
|
},
|
||||||
|
"ok": 1.0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- Helpers ----
|
||||||
|
|
||||||
|
/// Load documents from storage, optionally using index for candidate narrowing, then filter.
|
||||||
|
async fn load_filtered_docs(
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
filter: &Document,
|
||||||
|
ns_key: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
txn_id: Option<&str>,
|
||||||
|
) -> CommandResult<Vec<Document>> {
|
||||||
|
if let Some(txn_id) = txn_id {
|
||||||
|
let docs = transactions::load_transaction_docs(ctx, txn_id, db, coll).await?;
|
||||||
|
return if filter.is_empty() {
|
||||||
|
Ok(docs)
|
||||||
|
} else {
|
||||||
|
Ok(QueryMatcher::filter(&docs, filter))
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try to use index to narrow candidates.
|
||||||
|
let candidate_ids: Option<HashSet<String>> = ctx
|
||||||
|
.indexes
|
||||||
|
.get(ns_key)
|
||||||
|
.and_then(|engine| engine.find_candidate_ids(filter));
|
||||||
|
|
||||||
|
let docs = if let Some(ids) = candidate_ids {
|
||||||
|
if ids.is_empty() {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
ctx.storage.find_by_ids(db, coll, ids).await?
|
||||||
|
} else {
|
||||||
|
ctx.storage.find_all(db, coll).await?
|
||||||
|
};
|
||||||
|
|
||||||
|
// Apply filter.
|
||||||
|
if filter.is_empty() {
|
||||||
|
Ok(docs)
|
||||||
|
} else {
|
||||||
|
Ok(QueryMatcher::filter(&docs, filter))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a base document for an upsert from the filter's equality conditions.
|
||||||
|
fn build_upsert_doc(filter: &Document) -> Document {
|
||||||
|
let mut doc = Document::new();
|
||||||
|
for (key, value) in filter {
|
||||||
|
if key.starts_with('$') {
|
||||||
|
// Skip top-level operators like $and, $or.
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
match value {
|
||||||
|
Bson::Document(d) if d.keys().any(|k| k.starts_with('$')) => {
|
||||||
|
// If the value has operators (e.g., $gt), extract $eq if present.
|
||||||
|
if let Some(eq_val) = d.get("$eq") {
|
||||||
|
doc.insert(key.clone(), eq_val.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
doc.insert(key.clone(), value.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
doc
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_update_spec(update_value: &Bson) -> Result<TUpdateSpec, String> {
|
||||||
|
match update_value {
|
||||||
|
Bson::Document(d) => Ok(TUpdateSpec::Document(d.clone())),
|
||||||
|
Bson::Array(stages) => {
|
||||||
|
if stages.is_empty() {
|
||||||
|
return Err("aggregation pipeline update cannot be empty".into());
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut pipeline = Vec::with_capacity(stages.len());
|
||||||
|
for stage in stages {
|
||||||
|
match stage {
|
||||||
|
Bson::Document(d) => pipeline.push(d.clone()),
|
||||||
|
_ => {
|
||||||
|
return Err(
|
||||||
|
"aggregation pipeline update stages must be documents".into(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(TUpdateSpec::Pipeline(pipeline))
|
||||||
|
}
|
||||||
|
_ => Err("missing or invalid 'u' field in update spec".into()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_update_spec(
|
||||||
|
doc: &Document,
|
||||||
|
update: &TUpdateSpec,
|
||||||
|
array_filters: Option<&[Document]>,
|
||||||
|
) -> Result<Document, String> {
|
||||||
|
match update {
|
||||||
|
TUpdateSpec::Document(update_doc) => UpdateEngine::apply_update(doc, update_doc, array_filters)
|
||||||
|
.map_err(|e| e.to_string()),
|
||||||
|
TUpdateSpec::Pipeline(pipeline) => {
|
||||||
|
if array_filters.is_some_and(|filters| !filters.is_empty()) {
|
||||||
|
return Err(
|
||||||
|
"arrayFilters are not supported with aggregation pipeline updates"
|
||||||
|
.into(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
UpdateEngine::apply_pipeline_update(doc, pipeline).map_err(|e| e.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_set_on_insert_if_present(update: &TUpdateSpec, doc: &mut Document) {
|
||||||
|
if let TUpdateSpec::Document(update_doc) = update {
|
||||||
|
if let Some(Bson::Document(soi)) = update_doc.get("$setOnInsert") {
|
||||||
|
UpdateEngine::apply_set_on_insert(doc, soi);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn ensure_document_id(doc: &mut Document) -> Bson {
|
||||||
|
if let Some(id) = doc.get("_id") {
|
||||||
|
id.clone()
|
||||||
|
} else {
|
||||||
|
let oid = ObjectId::new();
|
||||||
|
doc.insert("_id", oid);
|
||||||
|
Bson::ObjectId(oid)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn ensure_immutable_id(original_doc: &Document, updated_doc: &mut Document) -> CommandResult<()> {
|
||||||
|
if let Some(original_id) = original_doc.get("_id") {
|
||||||
|
match updated_doc.get("_id") {
|
||||||
|
Some(updated_id) if updated_id == original_id => Ok(()),
|
||||||
|
Some(_) => Err(CommandError::ImmutableField(
|
||||||
|
"cannot modify immutable field '_id'".into(),
|
||||||
|
)),
|
||||||
|
None => {
|
||||||
|
updated_doc.insert("_id", original_id.clone());
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extract _id as a string for storage operations.
|
||||||
|
fn extract_id_string(doc: &Document) -> String {
|
||||||
|
match doc.get("_id") {
|
||||||
|
Some(Bson::ObjectId(oid)) => oid.to_hex(),
|
||||||
|
Some(Bson::String(s)) => s.clone(),
|
||||||
|
Some(other) => format!("{}", other),
|
||||||
|
None => String::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Apply fields projection if specified, returning Bson.
|
||||||
|
fn apply_fields_projection(doc: &Document, fields: &Option<Document>) -> Bson {
|
||||||
|
match fields {
|
||||||
|
Some(proj) if !proj.is_empty() => Bson::Document(apply_projection(doc, proj)),
|
||||||
|
_ => Bson::Document(doc.clone()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Ensure the target database and collection exist, creating them if needed.
|
||||||
|
async fn ensure_collection_exists(
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<()> {
|
||||||
|
if let Err(e) = ctx.storage.create_database(db).await {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
|
||||||
|
return Err(CommandError::StorageError(msg));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
match ctx.storage.collection_exists(db, coll).await {
|
||||||
|
Ok(true) => {}
|
||||||
|
Ok(false) => {
|
||||||
|
if let Err(e) = ctx.storage.create_collection(db, coll).await {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
|
||||||
|
return Err(CommandError::StorageError(msg));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
if let Err(e2) = ctx.storage.create_collection(db, coll).await {
|
||||||
|
let msg = e2.to_string();
|
||||||
|
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
|
||||||
|
return Err(CommandError::StorageError(format!(
|
||||||
|
"collection_exists failed: {e}; create_collection failed: {msg}"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
@@ -0,0 +1,9 @@
|
|||||||
|
mod context;
|
||||||
|
pub mod error;
|
||||||
|
pub mod handlers;
|
||||||
|
pub mod transactions;
|
||||||
|
mod router;
|
||||||
|
|
||||||
|
pub use context::{CommandContext, ConnectionState, CursorState};
|
||||||
|
pub use error::{CommandError, CommandResult};
|
||||||
|
pub use router::CommandRouter;
|
||||||
@@ -0,0 +1,210 @@
|
|||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use bson::{Bson, Document};
|
||||||
|
use tracing::{debug, warn};
|
||||||
|
|
||||||
|
use rustdb_wire::ParsedCommand;
|
||||||
|
use rustdb_auth::AuthAction;
|
||||||
|
|
||||||
|
use crate::context::{CommandContext, ConnectionState};
|
||||||
|
use crate::error::CommandError;
|
||||||
|
use crate::{handlers, transactions};
|
||||||
|
|
||||||
|
/// Routes parsed wire protocol commands to the appropriate handler.
|
||||||
|
pub struct CommandRouter {
|
||||||
|
ctx: Arc<CommandContext>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CommandRouter {
|
||||||
|
/// Create a new command router with the given context.
|
||||||
|
pub fn new(ctx: Arc<CommandContext>) -> Self {
|
||||||
|
Self { ctx }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Route a parsed command to the appropriate handler, returning a BSON response document.
|
||||||
|
pub async fn route(&self, cmd: &ParsedCommand, connection: &mut ConnectionState) -> Document {
|
||||||
|
let db = &cmd.database;
|
||||||
|
let command_name = cmd.command_name.as_str();
|
||||||
|
|
||||||
|
debug!(command = %command_name, database = %db, "routing command");
|
||||||
|
|
||||||
|
if self.ctx.auth.enabled()
|
||||||
|
&& !connection.is_authenticated()
|
||||||
|
&& !allows_unauthenticated(command_name)
|
||||||
|
{
|
||||||
|
return CommandError::Unauthorized(format!(
|
||||||
|
"command '{}' requires authentication",
|
||||||
|
command_name,
|
||||||
|
))
|
||||||
|
.to_error_doc();
|
||||||
|
}
|
||||||
|
|
||||||
|
if self.ctx.auth.enabled() && connection.is_authenticated() {
|
||||||
|
if let Some(action) = required_action(command_name, &cmd.command) {
|
||||||
|
if !self
|
||||||
|
.ctx
|
||||||
|
.auth
|
||||||
|
.is_authorized(&connection.authenticated_users, db, action)
|
||||||
|
{
|
||||||
|
return CommandError::Unauthorized(format!(
|
||||||
|
"command '{}' is not authorized for database '{}'",
|
||||||
|
command_name, db,
|
||||||
|
))
|
||||||
|
.to_error_doc();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Err(e) = transactions::prepare_transaction_for_command(
|
||||||
|
&self.ctx,
|
||||||
|
&cmd.command,
|
||||||
|
command_name,
|
||||||
|
) {
|
||||||
|
return e.to_error_doc();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Extract session id if present, and touch the session.
|
||||||
|
if let Some(lsid) = cmd.command.get("lsid") {
|
||||||
|
if let Some(session_id) = rustdb_txn::SessionEngine::extract_session_id(lsid) {
|
||||||
|
self.ctx.sessions.get_or_create_session(&session_id);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let result = match command_name {
|
||||||
|
// -- handshake / monitoring --
|
||||||
|
"hello" | "ismaster" | "isMaster" => {
|
||||||
|
handlers::hello_handler::handle(&cmd.command, db, &self.ctx).await
|
||||||
|
}
|
||||||
|
|
||||||
|
// -- authentication --
|
||||||
|
"saslStart" => {
|
||||||
|
handlers::auth_handler::handle_sasl_start(&cmd.command, db, &self.ctx, connection).await
|
||||||
|
}
|
||||||
|
"saslContinue" => {
|
||||||
|
handlers::auth_handler::handle_sasl_continue(&cmd.command, &self.ctx, connection).await
|
||||||
|
}
|
||||||
|
|
||||||
|
// -- query commands --
|
||||||
|
"find" => {
|
||||||
|
handlers::find_handler::handle(&cmd.command, db, &self.ctx).await
|
||||||
|
}
|
||||||
|
"getMore" => {
|
||||||
|
handlers::find_handler::handle_get_more(&cmd.command, db, &self.ctx).await
|
||||||
|
}
|
||||||
|
"killCursors" => {
|
||||||
|
handlers::find_handler::handle_kill_cursors(&cmd.command, db, &self.ctx).await
|
||||||
|
}
|
||||||
|
"count" => {
|
||||||
|
handlers::find_handler::handle_count(&cmd.command, db, &self.ctx).await
|
||||||
|
}
|
||||||
|
"distinct" => {
|
||||||
|
handlers::find_handler::handle_distinct(&cmd.command, db, &self.ctx).await
|
||||||
|
}
|
||||||
|
|
||||||
|
// -- write commands --
|
||||||
|
"insert" => {
|
||||||
|
handlers::insert_handler::handle(&cmd.command, db, &self.ctx, cmd.document_sequences.as_ref()).await
|
||||||
|
}
|
||||||
|
"update" | "findAndModify" => {
|
||||||
|
handlers::update_handler::handle(&cmd.command, db, &self.ctx, command_name).await
|
||||||
|
}
|
||||||
|
"delete" => {
|
||||||
|
handlers::delete_handler::handle(&cmd.command, db, &self.ctx).await
|
||||||
|
}
|
||||||
|
|
||||||
|
// -- aggregation --
|
||||||
|
"aggregate" => {
|
||||||
|
handlers::aggregate_handler::handle(&cmd.command, db, &self.ctx).await
|
||||||
|
}
|
||||||
|
|
||||||
|
// -- index management --
|
||||||
|
"createIndexes" | "dropIndexes" | "listIndexes" => {
|
||||||
|
handlers::index_handler::handle(&cmd.command, db, &self.ctx, command_name).await
|
||||||
|
}
|
||||||
|
|
||||||
|
// -- admin commands --
|
||||||
|
"ping" | "buildInfo" | "buildinfo" | "serverStatus" | "hostInfo"
|
||||||
|
| "whatsmyuri" | "getLog" | "replSetGetStatus" | "getCmdLineOpts"
|
||||||
|
| "getParameter" | "getFreeMonitoringStatus" | "setFreeMonitoring"
|
||||||
|
| "getShardMap" | "shardingState" | "atlasVersion"
|
||||||
|
| "connectionStatus" | "listDatabases" | "listCollections"
|
||||||
|
| "create" | "drop" | "dropDatabase" | "renameCollection"
|
||||||
|
| "dbStats" | "collStats" | "validate" | "explain"
|
||||||
|
| "startSession" | "endSessions" | "killSessions"
|
||||||
|
| "commitTransaction" | "abortTransaction"
|
||||||
|
| "authenticate" | "logout"
|
||||||
|
| "createUser" | "updateUser" | "dropUser" | "usersInfo"
|
||||||
|
| "grantRolesToUser" | "revokeRolesFromUser"
|
||||||
|
| "currentOp" | "killOp" | "top" | "profile"
|
||||||
|
| "compact" | "reIndex" | "fsync" | "connPoolSync" => {
|
||||||
|
handlers::admin_handler::handle(&cmd.command, db, &self.ctx, command_name, connection).await
|
||||||
|
}
|
||||||
|
|
||||||
|
// -- unknown command --
|
||||||
|
other => {
|
||||||
|
warn!(command = %other, "unknown command");
|
||||||
|
Err(CommandError::NotImplemented(other.to_string()))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Ok(doc) => doc,
|
||||||
|
Err(e) => e.to_error_doc(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn allows_unauthenticated(command_name: &str) -> bool {
|
||||||
|
matches!(
|
||||||
|
command_name,
|
||||||
|
"hello" | "ismaster" | "isMaster" | "saslStart" | "saslContinue" | "getnonce"
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn required_action(command_name: &str, command: &Document) -> Option<AuthAction> {
|
||||||
|
match command_name {
|
||||||
|
"hello" | "ismaster" | "isMaster" | "saslStart" | "saslContinue" | "getnonce" => None,
|
||||||
|
"ping" | "buildInfo" | "buildinfo" | "hostInfo" | "whatsmyuri" | "getLog"
|
||||||
|
| "getCmdLineOpts" | "getParameter" | "getFreeMonitoringStatus" | "setFreeMonitoring"
|
||||||
|
| "getShardMap" | "shardingState" | "atlasVersion" | "connectionStatus"
|
||||||
|
| "startSession" | "endSessions" | "killSessions" | "authenticate" | "logout" => None,
|
||||||
|
|
||||||
|
"find" | "getMore" | "killCursors" | "count" | "distinct" | "listIndexes"
|
||||||
|
| "listCollections" | "collStats" | "dbStats" | "validate" | "explain" => {
|
||||||
|
Some(AuthAction::Read)
|
||||||
|
}
|
||||||
|
|
||||||
|
"aggregate" => Some(if aggregate_writes(command) {
|
||||||
|
AuthAction::Write
|
||||||
|
} else {
|
||||||
|
AuthAction::Read
|
||||||
|
}),
|
||||||
|
|
||||||
|
"insert" | "update" | "findAndModify" | "delete" | "commitTransaction"
|
||||||
|
| "abortTransaction" => Some(AuthAction::Write),
|
||||||
|
|
||||||
|
"createIndexes" | "dropIndexes" | "create" | "drop" | "dropDatabase"
|
||||||
|
| "renameCollection" | "compact" | "reIndex" | "fsync" | "profile" => {
|
||||||
|
Some(AuthAction::DbAdmin)
|
||||||
|
}
|
||||||
|
|
||||||
|
"createUser" | "updateUser" | "dropUser" | "usersInfo" | "grantRolesToUser"
|
||||||
|
| "revokeRolesFromUser" => Some(AuthAction::UserAdmin),
|
||||||
|
|
||||||
|
"serverStatus" | "listDatabases" | "currentOp" | "killOp" | "top" => {
|
||||||
|
Some(AuthAction::ClusterMonitor)
|
||||||
|
}
|
||||||
|
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn aggregate_writes(command: &Document) -> bool {
|
||||||
|
let Ok(pipeline) = command.get_array("pipeline") else {
|
||||||
|
return false;
|
||||||
|
};
|
||||||
|
pipeline.last().and_then(|stage| match stage {
|
||||||
|
Bson::Document(doc) => Some(doc.contains_key("$out") || doc.contains_key("$merge")),
|
||||||
|
_ => None,
|
||||||
|
}).unwrap_or(false)
|
||||||
|
}
|
||||||
@@ -0,0 +1,367 @@
|
|||||||
|
use bson::{doc, Bson, Document};
|
||||||
|
use rustdb_storage::OpType;
|
||||||
|
use rustdb_txn::{TransactionState, WriteEntry, WriteOp};
|
||||||
|
|
||||||
|
use crate::context::CommandContext;
|
||||||
|
use crate::error::{CommandError, CommandResult};
|
||||||
|
|
||||||
|
pub fn command_starts_transaction(cmd: &Document) -> bool {
|
||||||
|
matches!(cmd.get("startTransaction"), Some(Bson::Boolean(true)))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn command_uses_transaction(cmd: &Document) -> bool {
|
||||||
|
command_starts_transaction(cmd) || matches!(cmd.get("autocommit"), Some(Bson::Boolean(false)))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn active_transaction_id(ctx: &CommandContext, cmd: &Document) -> Option<String> {
|
||||||
|
if !command_uses_transaction(cmd) {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
let session_id = cmd
|
||||||
|
.get("lsid")
|
||||||
|
.and_then(rustdb_txn::SessionEngine::extract_session_id)?;
|
||||||
|
ctx.sessions.get_transaction_id(&session_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn prepare_transaction_for_command(
|
||||||
|
ctx: &CommandContext,
|
||||||
|
cmd: &Document,
|
||||||
|
command_name: &str,
|
||||||
|
) -> CommandResult<()> {
|
||||||
|
if matches!(command_name, "commitTransaction" | "abortTransaction") {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
let starts_transaction = command_starts_transaction(cmd);
|
||||||
|
let uses_transaction = command_uses_transaction(cmd);
|
||||||
|
if !uses_transaction {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
let session_id = session_id_from_command(cmd)?;
|
||||||
|
require_txn_number(cmd)?;
|
||||||
|
ctx.sessions.get_or_create_session(&session_id);
|
||||||
|
|
||||||
|
if starts_transaction {
|
||||||
|
let txn_id = ctx.transactions.start_transaction(&session_id)?;
|
||||||
|
ctx.sessions.start_transaction(&session_id, &txn_id)?;
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
if ctx.sessions.get_transaction_id(&session_id).is_none() {
|
||||||
|
return Err(CommandError::NoSuchTransaction(format!(
|
||||||
|
"session {session_id} has no active transaction"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn load_transaction_docs(
|
||||||
|
ctx: &CommandContext,
|
||||||
|
txn_id: &str,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
) -> CommandResult<Vec<Document>> {
|
||||||
|
let ns = namespace(db, coll);
|
||||||
|
if !ctx.transactions.has_snapshot(txn_id, &ns) {
|
||||||
|
let docs = match ctx.storage.collection_exists(db, coll).await {
|
||||||
|
Ok(true) => ctx.storage.find_all(db, coll).await?,
|
||||||
|
Ok(false) => Vec::new(),
|
||||||
|
Err(_) => Vec::new(),
|
||||||
|
};
|
||||||
|
ctx.transactions.set_snapshot(txn_id, &ns, docs);
|
||||||
|
}
|
||||||
|
|
||||||
|
ctx.transactions
|
||||||
|
.get_snapshot(txn_id, &ns)
|
||||||
|
.ok_or_else(|| CommandError::NoSuchTransaction(txn_id.to_string()))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn record_insert(
|
||||||
|
ctx: &CommandContext,
|
||||||
|
txn_id: &str,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
doc: Document,
|
||||||
|
) -> CommandResult<String> {
|
||||||
|
let id = document_id_string(&doc)?;
|
||||||
|
let docs = load_transaction_docs(ctx, txn_id, db, coll).await?;
|
||||||
|
if docs.iter().any(|existing| document_id_string(existing).ok().as_deref() == Some(id.as_str())) {
|
||||||
|
return Err(CommandError::DuplicateKey(format!(
|
||||||
|
"duplicate _id '{}' in transaction",
|
||||||
|
id
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
ctx.transactions.record_write(
|
||||||
|
txn_id,
|
||||||
|
&namespace(db, coll),
|
||||||
|
&id,
|
||||||
|
WriteOp::Insert,
|
||||||
|
Some(doc),
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
Ok(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn record_update(
|
||||||
|
ctx: &CommandContext,
|
||||||
|
txn_id: &str,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
original: Document,
|
||||||
|
updated: Document,
|
||||||
|
) -> CommandResult<String> {
|
||||||
|
let id = document_id_string(&original)?;
|
||||||
|
ctx.transactions.record_write(
|
||||||
|
txn_id,
|
||||||
|
&namespace(db, coll),
|
||||||
|
&id,
|
||||||
|
WriteOp::Update,
|
||||||
|
Some(updated),
|
||||||
|
Some(original),
|
||||||
|
);
|
||||||
|
Ok(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn record_delete(
|
||||||
|
ctx: &CommandContext,
|
||||||
|
txn_id: &str,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
original: Document,
|
||||||
|
) -> CommandResult<String> {
|
||||||
|
let id = document_id_string(&original)?;
|
||||||
|
ctx.transactions.record_write(
|
||||||
|
txn_id,
|
||||||
|
&namespace(db, coll),
|
||||||
|
&id,
|
||||||
|
WriteOp::Delete,
|
||||||
|
None,
|
||||||
|
Some(original),
|
||||||
|
);
|
||||||
|
Ok(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn commit_transaction_command(
|
||||||
|
cmd: &Document,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<Document> {
|
||||||
|
let session_id = session_id_from_command(cmd)?;
|
||||||
|
let txn_id = ctx
|
||||||
|
.sessions
|
||||||
|
.get_transaction_id(&session_id)
|
||||||
|
.ok_or_else(|| CommandError::NoSuchTransaction(format!(
|
||||||
|
"session {session_id} has no active transaction"
|
||||||
|
)))?;
|
||||||
|
let state = ctx.transactions.take_transaction(&txn_id)?;
|
||||||
|
|
||||||
|
preflight_transaction(&state, ctx).await?;
|
||||||
|
apply_transaction(state, ctx).await?;
|
||||||
|
ctx.sessions.end_transaction(&session_id);
|
||||||
|
|
||||||
|
Ok(doc! { "ok": 1.0 })
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn abort_transaction_command(cmd: &Document, ctx: &CommandContext) -> CommandResult<Document> {
|
||||||
|
let session_id = session_id_from_command(cmd)?;
|
||||||
|
let txn_id = ctx
|
||||||
|
.sessions
|
||||||
|
.get_transaction_id(&session_id)
|
||||||
|
.ok_or_else(|| CommandError::NoSuchTransaction(format!(
|
||||||
|
"session {session_id} has no active transaction"
|
||||||
|
)))?;
|
||||||
|
ctx.transactions.abort_transaction(&txn_id)?;
|
||||||
|
ctx.sessions.end_transaction(&session_id);
|
||||||
|
Ok(doc! { "ok": 1.0 })
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn document_id_string(doc: &Document) -> CommandResult<String> {
|
||||||
|
match doc.get("_id") {
|
||||||
|
Some(Bson::ObjectId(oid)) => Ok(oid.to_hex()),
|
||||||
|
Some(Bson::String(s)) => Ok(s.clone()),
|
||||||
|
Some(other) => Ok(format!("{}", other)),
|
||||||
|
None => Err(CommandError::InvalidArgument("document missing _id field".into())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn session_id_from_command(cmd: &Document) -> CommandResult<String> {
|
||||||
|
cmd.get("lsid")
|
||||||
|
.and_then(rustdb_txn::SessionEngine::extract_session_id)
|
||||||
|
.ok_or_else(|| CommandError::InvalidArgument("transaction command requires lsid".into()))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn require_txn_number(cmd: &Document) -> CommandResult<()> {
|
||||||
|
match cmd.get("txnNumber") {
|
||||||
|
Some(Bson::Int64(_)) | Some(Bson::Int32(_)) => Ok(()),
|
||||||
|
_ => Err(CommandError::InvalidArgument(
|
||||||
|
"transaction command requires txnNumber".into(),
|
||||||
|
)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn namespace(db: &str, coll: &str) -> String {
|
||||||
|
format!("{db}.{coll}")
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn preflight_transaction(state: &TransactionState, ctx: &CommandContext) -> CommandResult<()> {
|
||||||
|
for (ns, writes) in &state.write_set {
|
||||||
|
let (db, coll) = split_namespace(ns)?;
|
||||||
|
drop(ctx.get_or_init_index_engine(db, coll).await);
|
||||||
|
|
||||||
|
for (doc_id, entry) in writes {
|
||||||
|
let current = current_doc(ctx, db, coll, doc_id).await?;
|
||||||
|
match entry.op {
|
||||||
|
WriteOp::Insert => {
|
||||||
|
if current.is_some() {
|
||||||
|
return Err(CommandError::DuplicateKey(format!(
|
||||||
|
"duplicate _id '{}' on transaction commit",
|
||||||
|
doc_id
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
if let Some(ref doc) = entry.doc {
|
||||||
|
if let Some(engine) = ctx.indexes.get(ns) {
|
||||||
|
engine.check_unique_constraints(doc)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
WriteOp::Update => {
|
||||||
|
assert_unchanged(doc_id, current.as_ref(), entry.original_doc.as_ref())?;
|
||||||
|
if let (Some(current_doc), Some(updated_doc)) = (current.as_ref(), entry.doc.as_ref()) {
|
||||||
|
if let Some(engine) = ctx.indexes.get(ns) {
|
||||||
|
engine.check_unique_constraints_for_update(current_doc, updated_doc)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
WriteOp::Delete => {
|
||||||
|
assert_unchanged(doc_id, current.as_ref(), entry.original_doc.as_ref())?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn apply_transaction(state: TransactionState, ctx: &CommandContext) -> CommandResult<()> {
|
||||||
|
let mut namespaces: Vec<_> = state.write_set.into_iter().collect();
|
||||||
|
namespaces.sort_by(|a, b| a.0.cmp(&b.0));
|
||||||
|
|
||||||
|
for (ns, writes) in namespaces {
|
||||||
|
let (db, coll) = split_namespace(&ns)?;
|
||||||
|
ensure_collection_exists(db, coll, ctx).await?;
|
||||||
|
drop(ctx.get_or_init_index_engine(db, coll).await);
|
||||||
|
|
||||||
|
let mut writes: Vec<(String, WriteEntry)> = writes.into_iter().collect();
|
||||||
|
writes.sort_by(|a, b| a.0.cmp(&b.0));
|
||||||
|
|
||||||
|
for (doc_id, entry) in writes {
|
||||||
|
match entry.op {
|
||||||
|
WriteOp::Insert => {
|
||||||
|
let Some(doc) = entry.doc else { continue; };
|
||||||
|
let inserted_id = ctx.storage.insert_one(db, coll, doc.clone()).await?;
|
||||||
|
ctx.oplog.append(OpType::Insert, db, coll, &inserted_id, Some(doc.clone()), None);
|
||||||
|
if let Some(mut engine) = ctx.indexes.get_mut(&ns) {
|
||||||
|
engine.on_insert(&doc)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
WriteOp::Update => {
|
||||||
|
let Some(doc) = entry.doc else { continue; };
|
||||||
|
ctx.storage.update_by_id(db, coll, &doc_id, doc.clone()).await?;
|
||||||
|
ctx.oplog.append(
|
||||||
|
OpType::Update,
|
||||||
|
db,
|
||||||
|
coll,
|
||||||
|
&doc_id,
|
||||||
|
Some(doc.clone()),
|
||||||
|
entry.original_doc.clone(),
|
||||||
|
);
|
||||||
|
if let (Some(mut engine), Some(ref original)) =
|
||||||
|
(ctx.indexes.get_mut(&ns), entry.original_doc.as_ref())
|
||||||
|
{
|
||||||
|
engine.on_update(original, &doc)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
WriteOp::Delete => {
|
||||||
|
ctx.storage.delete_by_id(db, coll, &doc_id).await?;
|
||||||
|
ctx.oplog.append(
|
||||||
|
OpType::Delete,
|
||||||
|
db,
|
||||||
|
coll,
|
||||||
|
&doc_id,
|
||||||
|
None,
|
||||||
|
entry.original_doc.clone(),
|
||||||
|
);
|
||||||
|
if let (Some(mut engine), Some(ref original)) =
|
||||||
|
(ctx.indexes.get_mut(&ns), entry.original_doc.as_ref())
|
||||||
|
{
|
||||||
|
engine.on_delete(original);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn current_doc(
|
||||||
|
ctx: &CommandContext,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
doc_id: &str,
|
||||||
|
) -> CommandResult<Option<Document>> {
|
||||||
|
match ctx.storage.collection_exists(db, coll).await {
|
||||||
|
Ok(true) => Ok(ctx.storage.find_by_id(db, coll, doc_id).await?),
|
||||||
|
Ok(false) => Ok(None),
|
||||||
|
Err(_) => Ok(None),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn assert_unchanged(
|
||||||
|
doc_id: &str,
|
||||||
|
current: Option<&Document>,
|
||||||
|
original: Option<&Document>,
|
||||||
|
) -> CommandResult<()> {
|
||||||
|
if current == original {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
Err(CommandError::WriteConflict(format!(
|
||||||
|
"document '{}' changed during transaction",
|
||||||
|
doc_id
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn ensure_collection_exists(
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
ctx: &CommandContext,
|
||||||
|
) -> CommandResult<()> {
|
||||||
|
if let Err(e) = ctx.storage.create_database(db).await {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
|
||||||
|
return Err(CommandError::StorageError(msg));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
match ctx.storage.collection_exists(db, coll).await {
|
||||||
|
Ok(true) => Ok(()),
|
||||||
|
Ok(false) | Err(_) => {
|
||||||
|
if let Err(e) = ctx.storage.create_collection(db, coll).await {
|
||||||
|
let msg = e.to_string();
|
||||||
|
if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
|
||||||
|
return Err(CommandError::StorageError(msg));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn split_namespace(ns: &str) -> CommandResult<(&str, &str)> {
|
||||||
|
ns.split_once('.')
|
||||||
|
.ok_or_else(|| CommandError::InvalidArgument(format!("invalid namespace '{ns}'")))
|
||||||
|
}
|
||||||
@@ -0,0 +1,12 @@
|
|||||||
|
[package]
|
||||||
|
name = "rustdb-config"
|
||||||
|
version.workspace = true
|
||||||
|
edition.workspace = true
|
||||||
|
license.workspace = true
|
||||||
|
authors.workspace = true
|
||||||
|
description = "Configuration types for RustDb, compatible with SmartDB JSON schema"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
serde = { workspace = true }
|
||||||
|
serde_json = { workspace = true }
|
||||||
|
thiserror = { workspace = true }
|
||||||
@@ -0,0 +1,342 @@
|
|||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
/// Storage backend type.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||||
|
#[serde(rename_all = "lowercase")]
|
||||||
|
pub enum StorageType {
|
||||||
|
Memory,
|
||||||
|
File,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for StorageType {
|
||||||
|
fn default() -> Self {
|
||||||
|
StorageType::Memory
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Top-level configuration for RustDb server.
|
||||||
|
/// Field names use camelCase to match the TypeScript SmartdbServer options.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
#[serde(rename_all = "camelCase")]
|
||||||
|
pub struct RustDbOptions {
|
||||||
|
/// TCP port to listen on (default: 27017)
|
||||||
|
#[serde(default = "default_port")]
|
||||||
|
pub port: u16,
|
||||||
|
|
||||||
|
/// Host/IP to bind to (default: "127.0.0.1")
|
||||||
|
#[serde(default = "default_host")]
|
||||||
|
pub host: String,
|
||||||
|
|
||||||
|
/// Unix socket path (overrides TCP if set)
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub socket_path: Option<String>,
|
||||||
|
|
||||||
|
/// Storage backend type
|
||||||
|
#[serde(default)]
|
||||||
|
pub storage: StorageType,
|
||||||
|
|
||||||
|
/// Base path for file storage (required when storage = "file")
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub storage_path: Option<String>,
|
||||||
|
|
||||||
|
/// Path for periodic persistence of in-memory data
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub persist_path: Option<String>,
|
||||||
|
|
||||||
|
/// Interval in ms for periodic persistence (default: 60000)
|
||||||
|
#[serde(default = "default_persist_interval")]
|
||||||
|
pub persist_interval_ms: u64,
|
||||||
|
|
||||||
|
/// Authentication configuration.
|
||||||
|
#[serde(default)]
|
||||||
|
pub auth: AuthOptions,
|
||||||
|
|
||||||
|
/// TLS transport configuration for TCP listeners.
|
||||||
|
#[serde(default)]
|
||||||
|
pub tls: TlsOptions,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Authentication configuration for the embedded server.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
#[serde(rename_all = "camelCase")]
|
||||||
|
pub struct AuthOptions {
|
||||||
|
/// Whether clients must authenticate before issuing protected commands.
|
||||||
|
#[serde(default)]
|
||||||
|
pub enabled: bool,
|
||||||
|
|
||||||
|
/// Bootstrap users loaded at startup. Passwords are converted into SCRAM credentials in memory.
|
||||||
|
#[serde(default)]
|
||||||
|
pub users: Vec<AuthUserOptions>,
|
||||||
|
|
||||||
|
/// Optional path for persisted SCRAM user metadata. Stores derived credentials, never plaintext passwords.
|
||||||
|
#[serde(default)]
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub users_path: Option<String>,
|
||||||
|
|
||||||
|
/// SCRAM iteration count used for bootstrap credentials.
|
||||||
|
#[serde(default = "default_scram_iterations")]
|
||||||
|
pub scram_iterations: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for AuthOptions {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
enabled: false,
|
||||||
|
users: Vec::new(),
|
||||||
|
users_path: None,
|
||||||
|
scram_iterations: default_scram_iterations(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// TLS transport configuration for the embedded server.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
#[serde(rename_all = "camelCase")]
|
||||||
|
pub struct TlsOptions {
|
||||||
|
/// Whether TCP client connections must use TLS.
|
||||||
|
#[serde(default)]
|
||||||
|
pub enabled: bool,
|
||||||
|
|
||||||
|
/// PEM-encoded server certificate chain.
|
||||||
|
#[serde(default)]
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub cert_path: Option<String>,
|
||||||
|
|
||||||
|
/// PEM-encoded server private key.
|
||||||
|
#[serde(default)]
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub key_path: Option<String>,
|
||||||
|
|
||||||
|
/// PEM-encoded client CA roots for mTLS verification.
|
||||||
|
#[serde(default)]
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub ca_path: Option<String>,
|
||||||
|
|
||||||
|
/// Require clients to present a certificate signed by caPath.
|
||||||
|
#[serde(default)]
|
||||||
|
pub require_client_cert: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for TlsOptions {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
enabled: false,
|
||||||
|
cert_path: None,
|
||||||
|
key_path: None,
|
||||||
|
ca_path: None,
|
||||||
|
require_client_cert: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A bootstrap user for SCRAM authentication.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
#[serde(rename_all = "camelCase")]
|
||||||
|
pub struct AuthUserOptions {
|
||||||
|
pub username: String,
|
||||||
|
pub password: String,
|
||||||
|
#[serde(default = "default_auth_database")]
|
||||||
|
pub database: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub roles: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_port() -> u16 {
|
||||||
|
27017
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_host() -> String {
|
||||||
|
"127.0.0.1".to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_persist_interval() -> u64 {
|
||||||
|
60000
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_scram_iterations() -> u32 {
|
||||||
|
15000
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_auth_database() -> String {
|
||||||
|
"admin".to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for RustDbOptions {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
port: default_port(),
|
||||||
|
host: default_host(),
|
||||||
|
socket_path: None,
|
||||||
|
storage: StorageType::default(),
|
||||||
|
storage_path: None,
|
||||||
|
persist_path: None,
|
||||||
|
persist_interval_ms: default_persist_interval(),
|
||||||
|
auth: AuthOptions::default(),
|
||||||
|
tls: TlsOptions::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RustDbOptions {
|
||||||
|
/// Load options from a JSON config file.
|
||||||
|
pub fn from_file(path: &str) -> Result<Self, ConfigError> {
|
||||||
|
let content = std::fs::read_to_string(path)
|
||||||
|
.map_err(|e| ConfigError::IoError(e.to_string()))?;
|
||||||
|
let options: Self = serde_json::from_str(&content)
|
||||||
|
.map_err(|e| ConfigError::ParseError(e.to_string()))?;
|
||||||
|
options.validate()?;
|
||||||
|
Ok(options)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Validate the configuration.
|
||||||
|
pub fn validate(&self) -> Result<(), ConfigError> {
|
||||||
|
if self.storage == StorageType::File && self.storage_path.is_none() {
|
||||||
|
return Err(ConfigError::ValidationError(
|
||||||
|
"storagePath is required when storage is 'file'".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if self.auth.enabled {
|
||||||
|
if self.auth.users.is_empty() && self.auth.users_path.is_none() {
|
||||||
|
return Err(ConfigError::ValidationError(
|
||||||
|
"auth.users or auth.usersPath must be set when auth.enabled is true".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if self.auth.scram_iterations < 4096 {
|
||||||
|
return Err(ConfigError::ValidationError(
|
||||||
|
"auth.scramIterations must be at least 4096".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
for user in &self.auth.users {
|
||||||
|
if user.username.is_empty() {
|
||||||
|
return Err(ConfigError::ValidationError(
|
||||||
|
"auth.users[].username must not be empty".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if user.password.is_empty() {
|
||||||
|
return Err(ConfigError::ValidationError(
|
||||||
|
format!("auth user '{}' must have a non-empty password", user.username),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if user.database.is_empty() {
|
||||||
|
return Err(ConfigError::ValidationError(
|
||||||
|
format!("auth user '{}' must have a non-empty database", user.username),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if self.tls.enabled {
|
||||||
|
if self.socket_path.is_some() {
|
||||||
|
return Err(ConfigError::ValidationError(
|
||||||
|
"tls.enabled is only supported for TCP listeners".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if self.tls.cert_path.as_deref().unwrap_or_default().is_empty() {
|
||||||
|
return Err(ConfigError::ValidationError(
|
||||||
|
"tls.certPath is required when tls.enabled is true".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if self.tls.key_path.as_deref().unwrap_or_default().is_empty() {
|
||||||
|
return Err(ConfigError::ValidationError(
|
||||||
|
"tls.keyPath is required when tls.enabled is true".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if self.tls.require_client_cert
|
||||||
|
&& self.tls.ca_path.as_deref().unwrap_or_default().is_empty()
|
||||||
|
{
|
||||||
|
return Err(ConfigError::ValidationError(
|
||||||
|
"tls.caPath is required when tls.requireClientCert is true".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the connection URI for this server configuration.
|
||||||
|
pub fn connection_uri(&self) -> String {
|
||||||
|
if let Some(ref socket_path) = self.socket_path {
|
||||||
|
let encoded = urlencoding(socket_path);
|
||||||
|
format!("mongodb://{}", encoded)
|
||||||
|
} else {
|
||||||
|
let base = format!("mongodb://{}:{}", self.host, self.port);
|
||||||
|
if self.tls.enabled {
|
||||||
|
format!("{}/?tls=true", base)
|
||||||
|
} else {
|
||||||
|
base
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Simple URL encoding for socket paths (encode / as %2F, etc.)
|
||||||
|
fn urlencoding(s: &str) -> String {
|
||||||
|
s.chars()
|
||||||
|
.map(|c| match c {
|
||||||
|
'/' => "%2F".to_string(),
|
||||||
|
':' => "%3A".to_string(),
|
||||||
|
' ' => "%20".to_string(),
|
||||||
|
_ => c.to_string(),
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Configuration errors.
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum ConfigError {
|
||||||
|
#[error("IO error: {0}")]
|
||||||
|
IoError(String),
|
||||||
|
#[error("Parse error: {0}")]
|
||||||
|
ParseError(String),
|
||||||
|
#[error("Validation error: {0}")]
|
||||||
|
ValidationError(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_default_options() {
|
||||||
|
let opts = RustDbOptions::default();
|
||||||
|
assert_eq!(opts.port, 27017);
|
||||||
|
assert_eq!(opts.host, "127.0.0.1");
|
||||||
|
assert!(opts.socket_path.is_none());
|
||||||
|
assert_eq!(opts.storage, StorageType::Memory);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_deserialize_from_json() {
|
||||||
|
let json = r#"{"port": 27018, "storage": "file", "storagePath": "./data"}"#;
|
||||||
|
let opts: RustDbOptions = serde_json::from_str(json).unwrap();
|
||||||
|
assert_eq!(opts.port, 27018);
|
||||||
|
assert_eq!(opts.storage, StorageType::File);
|
||||||
|
assert_eq!(opts.storage_path, Some("./data".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_connection_uri_tcp() {
|
||||||
|
let opts = RustDbOptions::default();
|
||||||
|
assert_eq!(opts.connection_uri(), "mongodb://127.0.0.1:27017");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_connection_uri_socket() {
|
||||||
|
let opts = RustDbOptions {
|
||||||
|
socket_path: Some("/tmp/smartdb-test.sock".to_string()),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
assert_eq!(
|
||||||
|
opts.connection_uri(),
|
||||||
|
"mongodb://%2Ftmp%2Fsmartdb-test.sock"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_validation_file_storage_requires_path() {
|
||||||
|
let opts = RustDbOptions {
|
||||||
|
storage: StorageType::File,
|
||||||
|
storage_path: None,
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
assert!(opts.validate().is_err());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,15 @@
|
|||||||
|
[package]
|
||||||
|
name = "rustdb-index"
|
||||||
|
version.workspace = true
|
||||||
|
edition.workspace = true
|
||||||
|
license.workspace = true
|
||||||
|
authors.workspace = true
|
||||||
|
description = "MongoDB-compatible B-tree and hash index engine with query planner for RustDb"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
bson = { workspace = true }
|
||||||
|
serde = { workspace = true }
|
||||||
|
serde_json = { workspace = true }
|
||||||
|
thiserror = { workspace = true }
|
||||||
|
tracing = { workspace = true }
|
||||||
|
rustdb-query = { workspace = true }
|
||||||
@@ -0,0 +1,740 @@
|
|||||||
|
use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
|
||||||
|
|
||||||
|
use bson::{Bson, Document};
|
||||||
|
use tracing::{debug, trace};
|
||||||
|
|
||||||
|
use rustdb_query::get_nested_value;
|
||||||
|
|
||||||
|
use crate::error::IndexError;
|
||||||
|
|
||||||
|
/// Options for creating an index.
|
||||||
|
#[derive(Debug, Clone, Default)]
|
||||||
|
pub struct IndexOptions {
|
||||||
|
/// Custom name for the index. Auto-generated if None.
|
||||||
|
pub name: Option<String>,
|
||||||
|
/// Whether the index enforces unique values.
|
||||||
|
pub unique: bool,
|
||||||
|
/// Whether the index skips documents missing the indexed field.
|
||||||
|
pub sparse: bool,
|
||||||
|
/// TTL in seconds (for date fields). None means no expiry.
|
||||||
|
pub expire_after_seconds: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Metadata about an existing index.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct IndexInfo {
|
||||||
|
/// Index version (always 2).
|
||||||
|
pub v: i32,
|
||||||
|
/// The key specification document (e.g. {"name": 1}).
|
||||||
|
pub key: Document,
|
||||||
|
/// The index name.
|
||||||
|
pub name: String,
|
||||||
|
/// Whether the index enforces uniqueness.
|
||||||
|
pub unique: bool,
|
||||||
|
/// Whether the index is sparse.
|
||||||
|
pub sparse: bool,
|
||||||
|
/// TTL expiry in seconds, if set.
|
||||||
|
pub expire_after_seconds: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Internal data for a single index.
|
||||||
|
struct IndexData {
|
||||||
|
/// The key specification (field -> direction).
|
||||||
|
key: Document,
|
||||||
|
/// The index name.
|
||||||
|
name: String,
|
||||||
|
/// Whether uniqueness is enforced.
|
||||||
|
unique: bool,
|
||||||
|
/// Whether the index is sparse.
|
||||||
|
sparse: bool,
|
||||||
|
/// TTL in seconds.
|
||||||
|
expire_after_seconds: Option<u64>,
|
||||||
|
/// B-tree for range queries: serialized key bytes -> set of document _id hex strings.
|
||||||
|
btree: BTreeMap<Vec<u8>, BTreeSet<String>>,
|
||||||
|
/// Hash map for equality lookups: serialized key bytes -> set of document _id hex strings.
|
||||||
|
hash: HashMap<Vec<u8>, HashSet<String>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl IndexData {
|
||||||
|
fn new(key: Document, name: String, unique: bool, sparse: bool, expire_after_seconds: Option<u64>) -> Self {
|
||||||
|
Self {
|
||||||
|
key,
|
||||||
|
name,
|
||||||
|
unique,
|
||||||
|
sparse,
|
||||||
|
expire_after_seconds,
|
||||||
|
btree: BTreeMap::new(),
|
||||||
|
hash: HashMap::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn to_info(&self) -> IndexInfo {
|
||||||
|
IndexInfo {
|
||||||
|
v: 2,
|
||||||
|
key: self.key.clone(),
|
||||||
|
name: self.name.clone(),
|
||||||
|
unique: self.unique,
|
||||||
|
sparse: self.sparse,
|
||||||
|
expire_after_seconds: self.expire_after_seconds,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Manages indexes for a single collection.
|
||||||
|
pub struct IndexEngine {
|
||||||
|
/// All indexes keyed by name.
|
||||||
|
indexes: HashMap<String, IndexData>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl IndexEngine {
|
||||||
|
/// Create a new IndexEngine with the default `_id_` index.
|
||||||
|
pub fn new() -> Self {
|
||||||
|
let mut indexes = HashMap::new();
|
||||||
|
let id_key = bson::doc! { "_id": 1 };
|
||||||
|
let id_index = IndexData::new(id_key, "_id_".to_string(), true, false, None);
|
||||||
|
indexes.insert("_id_".to_string(), id_index);
|
||||||
|
Self { indexes }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create a new index. Returns the index name.
|
||||||
|
pub fn create_index(&mut self, key: Document, options: IndexOptions) -> Result<String, IndexError> {
|
||||||
|
if key.is_empty() {
|
||||||
|
return Err(IndexError::InvalidIndex("Index key must have at least one field".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let name = options.name.unwrap_or_else(|| Self::generate_index_name(&key));
|
||||||
|
|
||||||
|
if self.indexes.contains_key(&name) {
|
||||||
|
debug!(index_name = %name, "Index already exists, returning existing");
|
||||||
|
return Ok(name);
|
||||||
|
}
|
||||||
|
|
||||||
|
debug!(index_name = %name, unique = options.unique, sparse = options.sparse, "Creating index");
|
||||||
|
|
||||||
|
let index_data = IndexData::new(
|
||||||
|
key,
|
||||||
|
name.clone(),
|
||||||
|
options.unique,
|
||||||
|
options.sparse,
|
||||||
|
options.expire_after_seconds,
|
||||||
|
);
|
||||||
|
self.indexes.insert(name.clone(), index_data);
|
||||||
|
|
||||||
|
Ok(name)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Drop an index by name. Returns true if the index existed.
|
||||||
|
/// Cannot drop the `_id_` index.
|
||||||
|
pub fn drop_index(&mut self, name: &str) -> Result<bool, IndexError> {
|
||||||
|
if name == "_id_" {
|
||||||
|
return Err(IndexError::ProtectedIndex("_id_".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let existed = self.indexes.remove(name).is_some();
|
||||||
|
if existed {
|
||||||
|
debug!(index_name = %name, "Dropped index");
|
||||||
|
}
|
||||||
|
Ok(existed)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Drop all indexes except `_id_`.
|
||||||
|
pub fn drop_all_indexes(&mut self) {
|
||||||
|
self.indexes.retain(|name, _| name == "_id_");
|
||||||
|
debug!("Dropped all non-_id indexes");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all indexes.
|
||||||
|
pub fn list_indexes(&self) -> Vec<IndexInfo> {
|
||||||
|
self.indexes.values().map(|idx| idx.to_info()).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check whether an index with the given name exists.
|
||||||
|
pub fn index_exists(&self, name: &str) -> bool {
|
||||||
|
self.indexes.contains_key(name)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check unique constraints for a document without modifying the index.
|
||||||
|
/// Returns Ok(()) if no conflict, Err(DuplicateKey) if a unique constraint
|
||||||
|
/// would be violated. This is a read-only check (immutable &self).
|
||||||
|
pub fn check_unique_constraints(&self, doc: &Document) -> Result<(), IndexError> {
|
||||||
|
for idx in self.indexes.values() {
|
||||||
|
if idx.unique {
|
||||||
|
let key_bytes = Self::extract_key_bytes(doc, &idx.key, idx.sparse);
|
||||||
|
if let Some(ref kb) = key_bytes {
|
||||||
|
if let Some(existing_ids) = idx.hash.get(kb) {
|
||||||
|
if !existing_ids.is_empty() {
|
||||||
|
return Err(IndexError::DuplicateKey {
|
||||||
|
index: idx.name.clone(),
|
||||||
|
key: format!("{:?}", kb),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check unique constraints for an update, excluding the document being updated.
|
||||||
|
/// Returns Ok(()) if no conflict. This is a read-only check (immutable &self).
|
||||||
|
pub fn check_unique_constraints_for_update(
|
||||||
|
&self,
|
||||||
|
old_doc: &Document,
|
||||||
|
new_doc: &Document,
|
||||||
|
) -> Result<(), IndexError> {
|
||||||
|
let doc_id = Self::extract_id(old_doc);
|
||||||
|
for idx in self.indexes.values() {
|
||||||
|
if idx.unique {
|
||||||
|
let new_key_bytes = Self::extract_key_bytes(new_doc, &idx.key, idx.sparse);
|
||||||
|
if let Some(ref kb) = new_key_bytes {
|
||||||
|
if let Some(existing_ids) = idx.hash.get(kb) {
|
||||||
|
let has_conflict = existing_ids.iter().any(|id| *id != doc_id);
|
||||||
|
if has_conflict {
|
||||||
|
return Err(IndexError::DuplicateKey {
|
||||||
|
index: idx.name.clone(),
|
||||||
|
key: format!("{:?}", kb),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Notify the engine that a document has been inserted.
|
||||||
|
/// Checks unique constraints and updates all index structures.
|
||||||
|
pub fn on_insert(&mut self, doc: &Document) -> Result<(), IndexError> {
|
||||||
|
let doc_id = Self::extract_id(doc);
|
||||||
|
|
||||||
|
// First pass: check unique constraints
|
||||||
|
for idx in self.indexes.values() {
|
||||||
|
if idx.unique {
|
||||||
|
let key_bytes = Self::extract_key_bytes(doc, &idx.key, idx.sparse);
|
||||||
|
if let Some(ref kb) = key_bytes {
|
||||||
|
if let Some(existing_ids) = idx.hash.get(kb) {
|
||||||
|
if !existing_ids.is_empty() {
|
||||||
|
return Err(IndexError::DuplicateKey {
|
||||||
|
index: idx.name.clone(),
|
||||||
|
key: format!("{:?}", kb),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Second pass: insert into all indexes
|
||||||
|
for idx in self.indexes.values_mut() {
|
||||||
|
let key_bytes = Self::extract_key_bytes(doc, &idx.key, idx.sparse);
|
||||||
|
if let Some(kb) = key_bytes {
|
||||||
|
idx.btree.entry(kb.clone()).or_default().insert(doc_id.clone());
|
||||||
|
idx.hash.entry(kb).or_default().insert(doc_id.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
trace!(doc_id = %doc_id, "Indexed document on insert");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Notify the engine that a document has been updated.
|
||||||
|
pub fn on_update(&mut self, old_doc: &Document, new_doc: &Document) -> Result<(), IndexError> {
|
||||||
|
let doc_id = Self::extract_id(old_doc);
|
||||||
|
|
||||||
|
// Check unique constraints for the new document (excluding the document itself)
|
||||||
|
for idx in self.indexes.values() {
|
||||||
|
if idx.unique {
|
||||||
|
let new_key_bytes = Self::extract_key_bytes(new_doc, &idx.key, idx.sparse);
|
||||||
|
if let Some(ref kb) = new_key_bytes {
|
||||||
|
if let Some(existing_ids) = idx.hash.get(kb) {
|
||||||
|
// If there are existing entries that aren't this document, it's a conflict
|
||||||
|
let other_ids: HashSet<_> = existing_ids.iter()
|
||||||
|
.filter(|id| **id != doc_id)
|
||||||
|
.collect();
|
||||||
|
if !other_ids.is_empty() {
|
||||||
|
return Err(IndexError::DuplicateKey {
|
||||||
|
index: idx.name.clone(),
|
||||||
|
key: format!("{:?}", kb),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Remove old entries and insert new ones
|
||||||
|
for idx in self.indexes.values_mut() {
|
||||||
|
let old_key_bytes = Self::extract_key_bytes(old_doc, &idx.key, idx.sparse);
|
||||||
|
if let Some(ref kb) = old_key_bytes {
|
||||||
|
if let Some(set) = idx.btree.get_mut(kb) {
|
||||||
|
set.remove(&doc_id);
|
||||||
|
if set.is_empty() {
|
||||||
|
idx.btree.remove(kb);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if let Some(set) = idx.hash.get_mut(kb) {
|
||||||
|
set.remove(&doc_id);
|
||||||
|
if set.is_empty() {
|
||||||
|
idx.hash.remove(kb);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let new_key_bytes = Self::extract_key_bytes(new_doc, &idx.key, idx.sparse);
|
||||||
|
if let Some(kb) = new_key_bytes {
|
||||||
|
idx.btree.entry(kb.clone()).or_default().insert(doc_id.clone());
|
||||||
|
idx.hash.entry(kb).or_default().insert(doc_id.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
trace!(doc_id = %doc_id, "Re-indexed document on update");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Notify the engine that a document has been deleted.
|
||||||
|
pub fn on_delete(&mut self, doc: &Document) {
|
||||||
|
let doc_id = Self::extract_id(doc);
|
||||||
|
|
||||||
|
for idx in self.indexes.values_mut() {
|
||||||
|
let key_bytes = Self::extract_key_bytes(doc, &idx.key, idx.sparse);
|
||||||
|
if let Some(ref kb) = key_bytes {
|
||||||
|
if let Some(set) = idx.btree.get_mut(kb) {
|
||||||
|
set.remove(&doc_id);
|
||||||
|
if set.is_empty() {
|
||||||
|
idx.btree.remove(kb);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if let Some(set) = idx.hash.get_mut(kb) {
|
||||||
|
set.remove(&doc_id);
|
||||||
|
if set.is_empty() {
|
||||||
|
idx.hash.remove(kb);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
trace!(doc_id = %doc_id, "Removed document from indexes");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Attempt to find candidate document IDs using indexes for the given filter.
|
||||||
|
/// Returns `None` if no suitable index is found (meaning a COLLSCAN is needed).
|
||||||
|
/// Returns `Some(set)` with candidate IDs that should be checked against the full filter.
|
||||||
|
pub fn find_candidate_ids(&self, filter: &Document) -> Option<HashSet<String>> {
|
||||||
|
if filter.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try each index to see which can serve this query
|
||||||
|
let mut best_candidates: Option<HashSet<String>> = None;
|
||||||
|
let mut best_score: f64 = 0.0;
|
||||||
|
|
||||||
|
for idx in self.indexes.values() {
|
||||||
|
if let Some((candidates, score)) = self.try_index_lookup(idx, filter) {
|
||||||
|
if score > best_score {
|
||||||
|
best_score = score;
|
||||||
|
best_candidates = Some(candidates);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
best_candidates
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Rebuild all indexes from a full set of documents.
|
||||||
|
pub fn rebuild_from_documents(&mut self, docs: &[Document]) {
|
||||||
|
// Clear all index data
|
||||||
|
for idx in self.indexes.values_mut() {
|
||||||
|
idx.btree.clear();
|
||||||
|
idx.hash.clear();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Re-index all documents
|
||||||
|
for doc in docs {
|
||||||
|
let doc_id = Self::extract_id(doc);
|
||||||
|
for idx in self.indexes.values_mut() {
|
||||||
|
let key_bytes = Self::extract_key_bytes(doc, &idx.key, idx.sparse);
|
||||||
|
if let Some(kb) = key_bytes {
|
||||||
|
idx.btree.entry(kb.clone()).or_default().insert(doc_id.clone());
|
||||||
|
idx.hash.entry(kb).or_default().insert(doc_id.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
debug!(num_docs = docs.len(), num_indexes = self.indexes.len(), "Rebuilt all indexes");
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- Internal helpers ----
|
||||||
|
|
||||||
|
/// Try to use an index for the given filter. Returns candidate IDs and a score.
|
||||||
|
fn try_index_lookup(&self, idx: &IndexData, filter: &Document) -> Option<(HashSet<String>, f64)> {
|
||||||
|
let index_fields: Vec<String> = idx.key.keys().map(|k| k.to_string()).collect();
|
||||||
|
|
||||||
|
// Check if the filter uses fields covered by this index
|
||||||
|
let mut matched_any = false;
|
||||||
|
let mut result_set: Option<HashSet<String>> = None;
|
||||||
|
let mut total_score: f64 = 0.0;
|
||||||
|
|
||||||
|
for field in &index_fields {
|
||||||
|
if let Some(condition) = filter.get(field) {
|
||||||
|
matched_any = true;
|
||||||
|
|
||||||
|
let (candidates, score) = self.lookup_field(idx, field, condition);
|
||||||
|
total_score += score;
|
||||||
|
|
||||||
|
// Add unique bonus
|
||||||
|
if idx.unique {
|
||||||
|
total_score += 0.5;
|
||||||
|
}
|
||||||
|
|
||||||
|
result_set = Some(match result_set {
|
||||||
|
Some(existing) => existing.intersection(&candidates).cloned().collect(),
|
||||||
|
None => candidates,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if !matched_any {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
result_set.map(|rs| (rs, total_score))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Look up candidates for a single field condition in an index.
|
||||||
|
fn lookup_field(&self, idx: &IndexData, field: &str, condition: &Bson) -> (HashSet<String>, f64) {
|
||||||
|
match condition {
|
||||||
|
// Equality match
|
||||||
|
Bson::Document(cond_doc) if Self::has_operators(cond_doc) => {
|
||||||
|
self.lookup_operator(idx, field, cond_doc)
|
||||||
|
}
|
||||||
|
// Direct equality
|
||||||
|
_ => {
|
||||||
|
let key_bytes = Self::bson_to_key_bytes(condition);
|
||||||
|
let candidates = idx.hash
|
||||||
|
.get(&key_bytes)
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or_default();
|
||||||
|
(candidates, 2.0) // equality score
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle operator-based lookups ($eq, $in, $gt, $lt, etc.).
|
||||||
|
fn lookup_operator(&self, idx: &IndexData, field: &str, operators: &Document) -> (HashSet<String>, f64) {
|
||||||
|
let mut candidates = HashSet::new();
|
||||||
|
let mut score: f64 = 0.0;
|
||||||
|
let mut has_range = false;
|
||||||
|
|
||||||
|
for (op, value) in operators {
|
||||||
|
match op.as_str() {
|
||||||
|
"$eq" => {
|
||||||
|
let key_bytes = Self::bson_to_key_bytes(value);
|
||||||
|
if let Some(ids) = idx.hash.get(&key_bytes) {
|
||||||
|
candidates = if candidates.is_empty() {
|
||||||
|
ids.clone()
|
||||||
|
} else {
|
||||||
|
candidates.intersection(ids).cloned().collect()
|
||||||
|
};
|
||||||
|
}
|
||||||
|
score += 2.0;
|
||||||
|
}
|
||||||
|
"$in" => {
|
||||||
|
if let Bson::Array(arr) = value {
|
||||||
|
let mut in_candidates = HashSet::new();
|
||||||
|
for v in arr {
|
||||||
|
let key_bytes = Self::bson_to_key_bytes(v);
|
||||||
|
if let Some(ids) = idx.hash.get(&key_bytes) {
|
||||||
|
in_candidates.extend(ids.iter().cloned());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
candidates = if candidates.is_empty() {
|
||||||
|
in_candidates
|
||||||
|
} else {
|
||||||
|
candidates.intersection(&in_candidates).cloned().collect()
|
||||||
|
};
|
||||||
|
score += 1.5;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"$gt" | "$gte" | "$lt" | "$lte" => {
|
||||||
|
let range_candidates = self.range_scan(idx, field, op.as_str(), value);
|
||||||
|
candidates = if candidates.is_empty() && !has_range {
|
||||||
|
range_candidates
|
||||||
|
} else {
|
||||||
|
candidates.intersection(&range_candidates).cloned().collect()
|
||||||
|
};
|
||||||
|
has_range = true;
|
||||||
|
score += 1.0;
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
// Operators like $ne, $nin, $exists, $regex are not efficiently indexable
|
||||||
|
// Return all indexed IDs for this index
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// If we only had non-indexable operators, return empty with 0 score
|
||||||
|
if score == 0.0 {
|
||||||
|
return (HashSet::new(), 0.0);
|
||||||
|
}
|
||||||
|
|
||||||
|
(candidates, score)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Perform a range scan on the B-tree index.
|
||||||
|
fn range_scan(&self, idx: &IndexData, _field: &str, op: &str, bound: &Bson) -> HashSet<String> {
|
||||||
|
let bound_bytes = Self::bson_to_key_bytes(bound);
|
||||||
|
let mut result = HashSet::new();
|
||||||
|
|
||||||
|
match op {
|
||||||
|
"$gt" => {
|
||||||
|
use std::ops::Bound;
|
||||||
|
for (_key, ids) in idx.btree.range((Bound::Excluded(bound_bytes), Bound::Unbounded)) {
|
||||||
|
result.extend(ids.iter().cloned());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"$gte" => {
|
||||||
|
for (_key, ids) in idx.btree.range(bound_bytes..) {
|
||||||
|
result.extend(ids.iter().cloned());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"$lt" => {
|
||||||
|
for (_key, ids) in idx.btree.range(..bound_bytes) {
|
||||||
|
result.extend(ids.iter().cloned());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"$lte" => {
|
||||||
|
for (_key, ids) in idx.btree.range(..=bound_bytes) {
|
||||||
|
result.extend(ids.iter().cloned());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Generate an index name from the key spec (e.g. {"name": 1, "age": -1} -> "name_1_age_-1").
|
||||||
|
fn generate_index_name(key: &Document) -> String {
|
||||||
|
key.iter()
|
||||||
|
.map(|(field, dir)| {
|
||||||
|
let dir_val = match dir {
|
||||||
|
Bson::Int32(n) => n.to_string(),
|
||||||
|
Bson::Int64(n) => n.to_string(),
|
||||||
|
Bson::String(s) => s.clone(),
|
||||||
|
_ => "1".to_string(),
|
||||||
|
};
|
||||||
|
format!("{}_{}", field, dir_val)
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join("_")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extract the `_id` field from a document as a hex string.
|
||||||
|
fn extract_id(doc: &Document) -> String {
|
||||||
|
match doc.get("_id") {
|
||||||
|
Some(Bson::ObjectId(oid)) => oid.to_hex(),
|
||||||
|
Some(Bson::String(s)) => s.clone(),
|
||||||
|
Some(other) => format!("{}", other),
|
||||||
|
None => String::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extract the index key bytes from a document for a given key specification.
|
||||||
|
/// Returns `None` if the document should be skipped (sparse index with missing fields).
|
||||||
|
fn extract_key_bytes(doc: &Document, key_spec: &Document, sparse: bool) -> Option<Vec<u8>> {
|
||||||
|
let fields: Vec<(&str, &Bson)> = key_spec.iter().map(|(k, v)| (k.as_str(), v)).collect();
|
||||||
|
|
||||||
|
if fields.len() == 1 {
|
||||||
|
// Single-field index
|
||||||
|
let field = fields[0].0;
|
||||||
|
let value = Self::resolve_field_value(doc, field);
|
||||||
|
if sparse && value.is_none() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let val = value.unwrap_or(Bson::Null);
|
||||||
|
Some(Self::bson_to_key_bytes(&val))
|
||||||
|
} else {
|
||||||
|
// Compound index: concatenate field values
|
||||||
|
let mut all_null = true;
|
||||||
|
let mut compound_bytes = Vec::new();
|
||||||
|
for (field, _dir) in &fields {
|
||||||
|
let value = Self::resolve_field_value(doc, field);
|
||||||
|
if value.is_some() {
|
||||||
|
all_null = false;
|
||||||
|
}
|
||||||
|
let val = value.unwrap_or(Bson::Null);
|
||||||
|
let field_bytes = Self::bson_to_key_bytes(&val);
|
||||||
|
// Length-prefix each field for unambiguous concatenation
|
||||||
|
compound_bytes.extend_from_slice(&(field_bytes.len() as u32).to_be_bytes());
|
||||||
|
compound_bytes.extend_from_slice(&field_bytes);
|
||||||
|
}
|
||||||
|
|
||||||
|
if sparse && all_null {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
Some(compound_bytes)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolve a field value from a document, supporting dot notation.
|
||||||
|
fn resolve_field_value(doc: &Document, field: &str) -> Option<Bson> {
|
||||||
|
if field.contains('.') {
|
||||||
|
get_nested_value(doc, field)
|
||||||
|
} else {
|
||||||
|
doc.get(field).cloned()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Serialize a BSON value to bytes for use as an index key.
|
||||||
|
fn bson_to_key_bytes(value: &Bson) -> Vec<u8> {
|
||||||
|
// Use BSON raw serialization for consistent byte representation.
|
||||||
|
// We wrap in a document since raw BSON requires a top-level document.
|
||||||
|
let wrapper = bson::doc! { "k": value.clone() };
|
||||||
|
let raw = bson::to_vec(&wrapper).unwrap_or_default();
|
||||||
|
raw
|
||||||
|
}
|
||||||
|
|
||||||
|
fn has_operators(doc: &Document) -> bool {
|
||||||
|
doc.keys().any(|k| k.starts_with('$'))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for IndexEngine {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use bson::oid::ObjectId;
|
||||||
|
|
||||||
|
fn make_doc(name: &str, age: i32) -> Document {
|
||||||
|
bson::doc! {
|
||||||
|
"_id": ObjectId::new(),
|
||||||
|
"name": name,
|
||||||
|
"age": age,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_default_id_index() {
|
||||||
|
let engine = IndexEngine::new();
|
||||||
|
assert!(engine.index_exists("_id_"));
|
||||||
|
assert_eq!(engine.list_indexes().len(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_create_and_drop_index() {
|
||||||
|
let mut engine = IndexEngine::new();
|
||||||
|
let name = engine.create_index(
|
||||||
|
bson::doc! { "name": 1 },
|
||||||
|
IndexOptions::default(),
|
||||||
|
).unwrap();
|
||||||
|
assert_eq!(name, "name_1");
|
||||||
|
assert!(engine.index_exists("name_1"));
|
||||||
|
|
||||||
|
assert!(engine.drop_index("name_1").unwrap());
|
||||||
|
assert!(!engine.index_exists("name_1"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_cannot_drop_id_index() {
|
||||||
|
let mut engine = IndexEngine::new();
|
||||||
|
let result = engine.drop_index("_id_");
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_unique_constraint() {
|
||||||
|
let mut engine = IndexEngine::new();
|
||||||
|
engine.create_index(
|
||||||
|
bson::doc! { "email": 1 },
|
||||||
|
IndexOptions { unique: true, ..Default::default() },
|
||||||
|
).unwrap();
|
||||||
|
|
||||||
|
let doc1 = bson::doc! { "_id": ObjectId::new(), "email": "a@b.com" };
|
||||||
|
let doc2 = bson::doc! { "_id": ObjectId::new(), "email": "a@b.com" };
|
||||||
|
|
||||||
|
engine.on_insert(&doc1).unwrap();
|
||||||
|
let result = engine.on_insert(&doc2);
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_find_candidates_equality() {
|
||||||
|
let mut engine = IndexEngine::new();
|
||||||
|
engine.create_index(
|
||||||
|
bson::doc! { "name": 1 },
|
||||||
|
IndexOptions::default(),
|
||||||
|
).unwrap();
|
||||||
|
|
||||||
|
let doc1 = make_doc("Alice", 30);
|
||||||
|
let doc2 = make_doc("Bob", 25);
|
||||||
|
let doc3 = make_doc("Alice", 35);
|
||||||
|
|
||||||
|
engine.on_insert(&doc1).unwrap();
|
||||||
|
engine.on_insert(&doc2).unwrap();
|
||||||
|
engine.on_insert(&doc3).unwrap();
|
||||||
|
|
||||||
|
let filter = bson::doc! { "name": "Alice" };
|
||||||
|
let candidates = engine.find_candidate_ids(&filter);
|
||||||
|
assert!(candidates.is_some());
|
||||||
|
assert_eq!(candidates.unwrap().len(), 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_on_delete() {
|
||||||
|
let mut engine = IndexEngine::new();
|
||||||
|
engine.create_index(
|
||||||
|
bson::doc! { "name": 1 },
|
||||||
|
IndexOptions::default(),
|
||||||
|
).unwrap();
|
||||||
|
|
||||||
|
let doc = make_doc("Alice", 30);
|
||||||
|
engine.on_insert(&doc).unwrap();
|
||||||
|
|
||||||
|
let filter = bson::doc! { "name": "Alice" };
|
||||||
|
assert!(engine.find_candidate_ids(&filter).is_some());
|
||||||
|
|
||||||
|
engine.on_delete(&doc);
|
||||||
|
let candidates = engine.find_candidate_ids(&filter);
|
||||||
|
assert!(candidates.is_some());
|
||||||
|
assert!(candidates.unwrap().is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_rebuild_from_documents() {
|
||||||
|
let mut engine = IndexEngine::new();
|
||||||
|
engine.create_index(
|
||||||
|
bson::doc! { "name": 1 },
|
||||||
|
IndexOptions::default(),
|
||||||
|
).unwrap();
|
||||||
|
|
||||||
|
let docs = vec![
|
||||||
|
make_doc("Alice", 30),
|
||||||
|
make_doc("Bob", 25),
|
||||||
|
];
|
||||||
|
|
||||||
|
engine.rebuild_from_documents(&docs);
|
||||||
|
|
||||||
|
let filter = bson::doc! { "name": "Alice" };
|
||||||
|
let candidates = engine.find_candidate_ids(&filter);
|
||||||
|
assert!(candidates.is_some());
|
||||||
|
assert_eq!(candidates.unwrap().len(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_drop_all_indexes() {
|
||||||
|
let mut engine = IndexEngine::new();
|
||||||
|
engine.create_index(bson::doc! { "a": 1 }, IndexOptions::default()).unwrap();
|
||||||
|
engine.create_index(bson::doc! { "b": 1 }, IndexOptions::default()).unwrap();
|
||||||
|
assert_eq!(engine.list_indexes().len(), 3);
|
||||||
|
|
||||||
|
engine.drop_all_indexes();
|
||||||
|
assert_eq!(engine.list_indexes().len(), 1);
|
||||||
|
assert!(engine.index_exists("_id_"));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,15 @@
|
|||||||
|
/// Errors from index operations.
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum IndexError {
|
||||||
|
#[error("Duplicate key error: index '{index}' has duplicate value for key {key}")]
|
||||||
|
DuplicateKey { index: String, key: String },
|
||||||
|
|
||||||
|
#[error("Index not found: {0}")]
|
||||||
|
IndexNotFound(String),
|
||||||
|
|
||||||
|
#[error("Invalid index specification: {0}")]
|
||||||
|
InvalidIndex(String),
|
||||||
|
|
||||||
|
#[error("Cannot drop protected index: {0}")]
|
||||||
|
ProtectedIndex(String),
|
||||||
|
}
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
mod engine;
|
||||||
|
mod planner;
|
||||||
|
pub mod error;
|
||||||
|
|
||||||
|
pub use engine::{IndexEngine, IndexInfo, IndexOptions};
|
||||||
|
pub use planner::{QueryPlan, QueryPlanner};
|
||||||
|
pub use error::IndexError;
|
||||||
@@ -0,0 +1,239 @@
|
|||||||
|
use std::collections::HashSet;
|
||||||
|
|
||||||
|
use bson::{Bson, Document};
|
||||||
|
use tracing::debug;
|
||||||
|
|
||||||
|
use crate::engine::IndexEngine;
|
||||||
|
|
||||||
|
/// The execution plan for a query.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub enum QueryPlan {
|
||||||
|
/// Full collection scan - no suitable index found.
|
||||||
|
CollScan,
|
||||||
|
/// Index scan with exact/equality matches.
|
||||||
|
IxScan {
|
||||||
|
/// Name of the index used.
|
||||||
|
index_name: String,
|
||||||
|
/// Candidate document IDs from the index.
|
||||||
|
candidate_ids: HashSet<String>,
|
||||||
|
},
|
||||||
|
/// Index scan with range-based matches.
|
||||||
|
IxScanRange {
|
||||||
|
/// Name of the index used.
|
||||||
|
index_name: String,
|
||||||
|
/// Candidate document IDs from the range scan.
|
||||||
|
candidate_ids: HashSet<String>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Plans query execution by selecting the best available index.
|
||||||
|
pub struct QueryPlanner;
|
||||||
|
|
||||||
|
impl QueryPlanner {
|
||||||
|
/// Analyze a filter and the available indexes to produce a query plan.
|
||||||
|
pub fn plan(filter: &Document, engine: &IndexEngine) -> QueryPlan {
|
||||||
|
if filter.is_empty() {
|
||||||
|
debug!("Empty filter -> CollScan");
|
||||||
|
return QueryPlan::CollScan;
|
||||||
|
}
|
||||||
|
|
||||||
|
let indexes = engine.list_indexes();
|
||||||
|
let mut best_plan: Option<QueryPlan> = None;
|
||||||
|
let mut best_score: f64 = 0.0;
|
||||||
|
|
||||||
|
for idx_info in &indexes {
|
||||||
|
let index_fields: Vec<String> = idx_info.key.keys().map(|k| k.to_string()).collect();
|
||||||
|
|
||||||
|
let mut matched = false;
|
||||||
|
let mut score: f64 = 0.0;
|
||||||
|
let mut is_range = false;
|
||||||
|
|
||||||
|
for field in &index_fields {
|
||||||
|
if let Some(condition) = filter.get(field) {
|
||||||
|
matched = true;
|
||||||
|
let field_score = Self::score_condition(condition);
|
||||||
|
score += field_score;
|
||||||
|
|
||||||
|
if Self::is_range_condition(condition) {
|
||||||
|
is_range = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if !matched {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Unique index bonus
|
||||||
|
if idx_info.unique {
|
||||||
|
score += 0.5;
|
||||||
|
}
|
||||||
|
|
||||||
|
if score > best_score {
|
||||||
|
best_score = score;
|
||||||
|
|
||||||
|
// Try to get candidates from the engine
|
||||||
|
// We build a sub-filter with only the fields this index covers
|
||||||
|
let mut sub_filter = Document::new();
|
||||||
|
for field in &index_fields {
|
||||||
|
if let Some(val) = filter.get(field) {
|
||||||
|
sub_filter.insert(field.clone(), val.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(candidates) = engine.find_candidate_ids(&sub_filter) {
|
||||||
|
if is_range {
|
||||||
|
best_plan = Some(QueryPlan::IxScanRange {
|
||||||
|
index_name: idx_info.name.clone(),
|
||||||
|
candidate_ids: candidates,
|
||||||
|
});
|
||||||
|
} else {
|
||||||
|
best_plan = Some(QueryPlan::IxScan {
|
||||||
|
index_name: idx_info.name.clone(),
|
||||||
|
candidate_ids: candidates,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
match best_plan {
|
||||||
|
Some(plan) => {
|
||||||
|
debug!(score = best_score, "Selected index plan");
|
||||||
|
plan
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
debug!("No suitable index found -> CollScan");
|
||||||
|
QueryPlan::CollScan
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Score a filter condition for index selectivity.
|
||||||
|
/// Higher scores indicate more selective (better) index usage.
|
||||||
|
fn score_condition(condition: &Bson) -> f64 {
|
||||||
|
match condition {
|
||||||
|
Bson::Document(doc) if Self::has_operators(doc) => {
|
||||||
|
let mut score: f64 = 0.0;
|
||||||
|
for (op, _) in doc {
|
||||||
|
score += match op.as_str() {
|
||||||
|
"$eq" => 2.0,
|
||||||
|
"$in" => 1.5,
|
||||||
|
"$gt" | "$gte" | "$lt" | "$lte" => 1.0,
|
||||||
|
_ => 0.0,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
score
|
||||||
|
}
|
||||||
|
// Direct equality
|
||||||
|
_ => 2.0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if a condition involves range operators.
|
||||||
|
fn is_range_condition(condition: &Bson) -> bool {
|
||||||
|
match condition {
|
||||||
|
Bson::Document(doc) => {
|
||||||
|
doc.keys().any(|k| matches!(k.as_str(), "$gt" | "$gte" | "$lt" | "$lte"))
|
||||||
|
}
|
||||||
|
_ => false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn has_operators(doc: &Document) -> bool {
|
||||||
|
doc.keys().any(|k| k.starts_with('$'))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::engine::IndexOptions;
|
||||||
|
use bson::oid::ObjectId;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_empty_filter_collscan() {
|
||||||
|
let engine = IndexEngine::new();
|
||||||
|
let plan = QueryPlanner::plan(&bson::doc! {}, &engine);
|
||||||
|
assert!(matches!(plan, QueryPlan::CollScan));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_id_equality_ixscan() {
|
||||||
|
let mut engine = IndexEngine::new();
|
||||||
|
let oid = ObjectId::new();
|
||||||
|
let doc = bson::doc! { "_id": oid.clone(), "name": "Alice" };
|
||||||
|
engine.on_insert(&doc).unwrap();
|
||||||
|
|
||||||
|
let filter = bson::doc! { "_id": oid };
|
||||||
|
let plan = QueryPlanner::plan(&filter, &engine);
|
||||||
|
assert!(matches!(plan, QueryPlan::IxScan { .. }));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_indexed_field_ixscan() {
|
||||||
|
let mut engine = IndexEngine::new();
|
||||||
|
engine.create_index(
|
||||||
|
bson::doc! { "status": 1 },
|
||||||
|
IndexOptions::default(),
|
||||||
|
).unwrap();
|
||||||
|
|
||||||
|
let doc = bson::doc! { "_id": ObjectId::new(), "status": "active" };
|
||||||
|
engine.on_insert(&doc).unwrap();
|
||||||
|
|
||||||
|
let filter = bson::doc! { "status": "active" };
|
||||||
|
let plan = QueryPlanner::plan(&filter, &engine);
|
||||||
|
assert!(matches!(plan, QueryPlan::IxScan { .. }));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_unindexed_field_collscan() {
|
||||||
|
let engine = IndexEngine::new();
|
||||||
|
let filter = bson::doc! { "unindexed_field": "value" };
|
||||||
|
let plan = QueryPlanner::plan(&filter, &engine);
|
||||||
|
assert!(matches!(plan, QueryPlan::CollScan));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_range_query_ixscan_range() {
|
||||||
|
let mut engine = IndexEngine::new();
|
||||||
|
engine.create_index(
|
||||||
|
bson::doc! { "age": 1 },
|
||||||
|
IndexOptions::default(),
|
||||||
|
).unwrap();
|
||||||
|
|
||||||
|
let doc = bson::doc! { "_id": ObjectId::new(), "age": 30 };
|
||||||
|
engine.on_insert(&doc).unwrap();
|
||||||
|
|
||||||
|
let filter = bson::doc! { "age": { "$gte": 25, "$lt": 35 } };
|
||||||
|
let plan = QueryPlanner::plan(&filter, &engine);
|
||||||
|
assert!(matches!(plan, QueryPlan::IxScanRange { .. }));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_unique_index_preferred() {
|
||||||
|
let mut engine = IndexEngine::new();
|
||||||
|
engine.create_index(
|
||||||
|
bson::doc! { "email": 1 },
|
||||||
|
IndexOptions { unique: true, ..Default::default() },
|
||||||
|
).unwrap();
|
||||||
|
engine.create_index(
|
||||||
|
bson::doc! { "email": 1, "name": 1 },
|
||||||
|
IndexOptions { name: Some("email_name".to_string()), ..Default::default() },
|
||||||
|
).unwrap();
|
||||||
|
|
||||||
|
let doc = bson::doc! { "_id": ObjectId::new(), "email": "a@b.com", "name": "Alice" };
|
||||||
|
engine.on_insert(&doc).unwrap();
|
||||||
|
|
||||||
|
let filter = bson::doc! { "email": "a@b.com" };
|
||||||
|
let plan = QueryPlanner::plan(&filter, &engine);
|
||||||
|
|
||||||
|
// The unique index on email should be preferred (higher score)
|
||||||
|
match plan {
|
||||||
|
QueryPlan::IxScan { index_name, .. } => {
|
||||||
|
assert_eq!(index_name, "email_1");
|
||||||
|
}
|
||||||
|
_ => panic!("Expected IxScan"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,15 @@
|
|||||||
|
[package]
|
||||||
|
name = "rustdb-query"
|
||||||
|
version.workspace = true
|
||||||
|
edition.workspace = true
|
||||||
|
license.workspace = true
|
||||||
|
authors.workspace = true
|
||||||
|
description = "MongoDB-compatible query matching, update operators, aggregation, sort, and projection engine"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
bson = { workspace = true }
|
||||||
|
serde = { workspace = true }
|
||||||
|
serde_json = { workspace = true }
|
||||||
|
regex = { workspace = true }
|
||||||
|
thiserror = { workspace = true }
|
||||||
|
tracing = { workspace = true }
|
||||||
@@ -0,0 +1,719 @@
|
|||||||
|
use bson::{Bson, Document};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
use crate::error::QueryError;
|
||||||
|
use crate::field_path::{get_nested_value, remove_nested_value};
|
||||||
|
use crate::matcher::QueryMatcher;
|
||||||
|
use crate::projection::apply_projection;
|
||||||
|
use crate::sort::sort_documents;
|
||||||
|
|
||||||
|
/// Aggregation pipeline engine.
|
||||||
|
pub struct AggregationEngine;
|
||||||
|
|
||||||
|
/// Trait for resolving cross-collection data (for $lookup, $graphLookup, etc.).
|
||||||
|
pub trait CollectionResolver {
|
||||||
|
fn resolve(&self, db: &str, coll: &str) -> Result<Vec<Document>, QueryError>;
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AggregationEngine {
|
||||||
|
/// Execute an aggregation pipeline on a set of documents.
|
||||||
|
pub fn aggregate(
|
||||||
|
docs: Vec<Document>,
|
||||||
|
pipeline: &[Document],
|
||||||
|
resolver: Option<&dyn CollectionResolver>,
|
||||||
|
db: &str,
|
||||||
|
) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let mut current = docs;
|
||||||
|
|
||||||
|
for stage in pipeline {
|
||||||
|
let (stage_name, stage_spec) = stage
|
||||||
|
.iter()
|
||||||
|
.next()
|
||||||
|
.ok_or_else(|| QueryError::AggregationError("Empty pipeline stage".into()))?;
|
||||||
|
|
||||||
|
current = match stage_name.as_str() {
|
||||||
|
"$match" => Self::stage_match(current, stage_spec)?,
|
||||||
|
"$project" => Self::stage_project(current, stage_spec)?,
|
||||||
|
"$sort" => Self::stage_sort(current, stage_spec)?,
|
||||||
|
"$limit" => Self::stage_limit(current, stage_spec)?,
|
||||||
|
"$skip" => Self::stage_skip(current, stage_spec)?,
|
||||||
|
"$group" => Self::stage_group(current, stage_spec)?,
|
||||||
|
"$unwind" => Self::stage_unwind(current, stage_spec)?,
|
||||||
|
"$count" => Self::stage_count(current, stage_spec)?,
|
||||||
|
"$addFields" | "$set" => Self::stage_add_fields(current, stage_spec)?,
|
||||||
|
"$replaceRoot" | "$replaceWith" => Self::stage_replace_root(current, stage_spec)?,
|
||||||
|
"$unset" => Self::stage_unset(current, stage_spec)?,
|
||||||
|
"$lookup" => Self::stage_lookup(current, stage_spec, resolver, db)?,
|
||||||
|
"$facet" => Self::stage_facet(current, stage_spec, resolver, db)?,
|
||||||
|
"$unionWith" => Self::stage_union_with(current, stage_spec, resolver, db)?,
|
||||||
|
other => {
|
||||||
|
return Err(QueryError::AggregationError(format!(
|
||||||
|
"Unsupported aggregation stage: {}",
|
||||||
|
other
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(current)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_match(docs: Vec<Document>, spec: &Bson) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let filter = match spec {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::AggregationError(
|
||||||
|
"$match requires a document".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
Ok(QueryMatcher::filter(&docs, filter))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_project(docs: Vec<Document>, spec: &Bson) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let projection = match spec {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::AggregationError(
|
||||||
|
"$project requires a document".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
Ok(docs
|
||||||
|
.into_iter()
|
||||||
|
.map(|doc| apply_projection(&doc, projection))
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_sort(mut docs: Vec<Document>, spec: &Bson) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let sort_spec = match spec {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::AggregationError(
|
||||||
|
"$sort requires a document".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
sort_documents(&mut docs, sort_spec);
|
||||||
|
Ok(docs)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_limit(docs: Vec<Document>, spec: &Bson) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let n = bson_to_usize(spec)
|
||||||
|
.ok_or_else(|| QueryError::AggregationError("$limit requires a number".into()))?;
|
||||||
|
Ok(docs.into_iter().take(n).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_skip(docs: Vec<Document>, spec: &Bson) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let n = bson_to_usize(spec)
|
||||||
|
.ok_or_else(|| QueryError::AggregationError("$skip requires a number".into()))?;
|
||||||
|
Ok(docs.into_iter().skip(n).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_group(docs: Vec<Document>, spec: &Bson) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let group_spec = match spec {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::AggregationError(
|
||||||
|
"$group requires a document".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let id_expr = group_spec.get("_id").cloned().unwrap_or(Bson::Null);
|
||||||
|
|
||||||
|
// Group documents by _id
|
||||||
|
let mut groups: HashMap<String, (Bson, Vec<Document>)> = HashMap::new();
|
||||||
|
|
||||||
|
for doc in &docs {
|
||||||
|
let group_key = resolve_expression(&id_expr, doc);
|
||||||
|
let key_str = format!("{:?}", group_key);
|
||||||
|
groups
|
||||||
|
.entry(key_str)
|
||||||
|
.or_insert_with(|| (group_key.clone(), Vec::new()))
|
||||||
|
.1
|
||||||
|
.push(doc.clone());
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut result = Vec::new();
|
||||||
|
|
||||||
|
for (_key_str, (group_id, group_docs)) in groups {
|
||||||
|
let mut output = bson::doc! { "_id": group_id };
|
||||||
|
|
||||||
|
for (field, accumulator) in group_spec {
|
||||||
|
if field == "_id" {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let acc_doc = match accumulator {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
let (acc_op, acc_expr) = acc_doc.iter().next().unwrap();
|
||||||
|
|
||||||
|
let value = match acc_op.as_str() {
|
||||||
|
"$sum" => accumulate_sum(&group_docs, acc_expr),
|
||||||
|
"$avg" => accumulate_avg(&group_docs, acc_expr),
|
||||||
|
"$min" => accumulate_min(&group_docs, acc_expr),
|
||||||
|
"$max" => accumulate_max(&group_docs, acc_expr),
|
||||||
|
"$first" => accumulate_first(&group_docs, acc_expr),
|
||||||
|
"$last" => accumulate_last(&group_docs, acc_expr),
|
||||||
|
"$push" => accumulate_push(&group_docs, acc_expr),
|
||||||
|
"$addToSet" => accumulate_add_to_set(&group_docs, acc_expr),
|
||||||
|
"$count" => Bson::Int64(group_docs.len() as i64),
|
||||||
|
_ => Bson::Null,
|
||||||
|
};
|
||||||
|
|
||||||
|
output.insert(field.clone(), value);
|
||||||
|
}
|
||||||
|
|
||||||
|
result.push(output);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(result)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_unwind(docs: Vec<Document>, spec: &Bson) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let (path, preserve_null) = match spec {
|
||||||
|
Bson::String(s) => (s.trim_start_matches('$').to_string(), false),
|
||||||
|
Bson::Document(d) => {
|
||||||
|
let path = d
|
||||||
|
.get_str("path")
|
||||||
|
.map(|s| s.trim_start_matches('$').to_string())
|
||||||
|
.map_err(|_| QueryError::AggregationError("$unwind requires 'path'".into()))?;
|
||||||
|
let preserve = d.get_bool("preserveNullAndEmptyArrays").unwrap_or(false);
|
||||||
|
(path, preserve)
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::AggregationError(
|
||||||
|
"$unwind requires a string or document".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut result = Vec::new();
|
||||||
|
|
||||||
|
for doc in docs {
|
||||||
|
let value = doc.get(&path).cloned();
|
||||||
|
|
||||||
|
match value {
|
||||||
|
Some(Bson::Array(arr)) => {
|
||||||
|
if arr.is_empty() && preserve_null {
|
||||||
|
let mut new_doc = doc.clone();
|
||||||
|
new_doc.remove(&path);
|
||||||
|
result.push(new_doc);
|
||||||
|
} else {
|
||||||
|
for elem in arr {
|
||||||
|
let mut new_doc = doc.clone();
|
||||||
|
new_doc.insert(path.clone(), elem);
|
||||||
|
result.push(new_doc);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(Bson::Null) | None => {
|
||||||
|
if preserve_null {
|
||||||
|
result.push(doc);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(val) => {
|
||||||
|
// Non-array: keep as-is
|
||||||
|
let mut new_doc = doc;
|
||||||
|
new_doc.insert(path.clone(), val);
|
||||||
|
result.push(new_doc);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(result)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_count(docs: Vec<Document>, spec: &Bson) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let field = match spec {
|
||||||
|
Bson::String(s) => s.clone(),
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::AggregationError(
|
||||||
|
"$count requires a string".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
Ok(vec![bson::doc! { field: docs.len() as i64 }])
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_add_fields(docs: Vec<Document>, spec: &Bson) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let fields = match spec {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::AggregationError(
|
||||||
|
"$addFields requires a document".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(docs
|
||||||
|
.into_iter()
|
||||||
|
.map(|mut doc| {
|
||||||
|
for (key, expr) in fields {
|
||||||
|
let value = resolve_expression(expr, &doc);
|
||||||
|
doc.insert(key.clone(), value);
|
||||||
|
}
|
||||||
|
doc
|
||||||
|
})
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_replace_root(docs: Vec<Document>, spec: &Bson) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let new_root_expr = match spec {
|
||||||
|
Bson::Document(d) => d
|
||||||
|
.get("newRoot")
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or(Bson::Document(d.clone())),
|
||||||
|
Bson::String(s) => Bson::String(s.clone()),
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::AggregationError(
|
||||||
|
"$replaceRoot requires a document".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut result = Vec::new();
|
||||||
|
for doc in docs {
|
||||||
|
let new_root = resolve_expression(&new_root_expr, &doc);
|
||||||
|
if let Bson::Document(d) = new_root {
|
||||||
|
result.push(d);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(result)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_unset(docs: Vec<Document>, spec: &Bson) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let fields: Vec<String> = match spec {
|
||||||
|
Bson::String(s) => vec![s.clone()],
|
||||||
|
Bson::Array(arr) => arr
|
||||||
|
.iter()
|
||||||
|
.map(|value| match value {
|
||||||
|
Bson::String(field) => Ok(field.clone()),
|
||||||
|
_ => Err(QueryError::AggregationError(
|
||||||
|
"$unset array entries must be strings".into(),
|
||||||
|
)),
|
||||||
|
})
|
||||||
|
.collect::<Result<Vec<_>, _>>()?,
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::AggregationError(
|
||||||
|
"$unset requires a string or array of strings".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(docs
|
||||||
|
.into_iter()
|
||||||
|
.map(|mut doc| {
|
||||||
|
for field in &fields {
|
||||||
|
if field.contains('.') {
|
||||||
|
remove_nested_value(&mut doc, field);
|
||||||
|
} else {
|
||||||
|
doc.remove(field);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
doc
|
||||||
|
})
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_lookup(
|
||||||
|
docs: Vec<Document>,
|
||||||
|
spec: &Bson,
|
||||||
|
resolver: Option<&dyn CollectionResolver>,
|
||||||
|
db: &str,
|
||||||
|
) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let lookup = match spec {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::AggregationError(
|
||||||
|
"$lookup requires a document".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let from = lookup
|
||||||
|
.get_str("from")
|
||||||
|
.map_err(|_| QueryError::AggregationError("$lookup requires 'from'".into()))?;
|
||||||
|
let local_field = lookup
|
||||||
|
.get_str("localField")
|
||||||
|
.map_err(|_| QueryError::AggregationError("$lookup requires 'localField'".into()))?;
|
||||||
|
let foreign_field = lookup
|
||||||
|
.get_str("foreignField")
|
||||||
|
.map_err(|_| QueryError::AggregationError("$lookup requires 'foreignField'".into()))?;
|
||||||
|
let as_field = lookup
|
||||||
|
.get_str("as")
|
||||||
|
.map_err(|_| QueryError::AggregationError("$lookup requires 'as'".into()))?;
|
||||||
|
|
||||||
|
let resolver = resolver.ok_or_else(|| {
|
||||||
|
QueryError::AggregationError("$lookup requires a collection resolver".into())
|
||||||
|
})?;
|
||||||
|
let foreign_docs = resolver.resolve(db, from)?;
|
||||||
|
|
||||||
|
Ok(docs
|
||||||
|
.into_iter()
|
||||||
|
.map(|mut doc| {
|
||||||
|
let local_val = get_nested_value(&doc, local_field);
|
||||||
|
let matches: Vec<Bson> = foreign_docs
|
||||||
|
.iter()
|
||||||
|
.filter(|fd| {
|
||||||
|
let foreign_val = get_nested_value(fd, foreign_field);
|
||||||
|
match (&local_val, &foreign_val) {
|
||||||
|
(Some(a), Some(b)) => bson_loose_eq(a, b),
|
||||||
|
_ => false,
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.map(|fd| Bson::Document(fd.clone()))
|
||||||
|
.collect();
|
||||||
|
doc.insert(as_field.to_string(), Bson::Array(matches));
|
||||||
|
doc
|
||||||
|
})
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_facet(
|
||||||
|
docs: Vec<Document>,
|
||||||
|
spec: &Bson,
|
||||||
|
resolver: Option<&dyn CollectionResolver>,
|
||||||
|
db: &str,
|
||||||
|
) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let facets = match spec {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::AggregationError(
|
||||||
|
"$facet requires a document".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut result = Document::new();
|
||||||
|
|
||||||
|
for (facet_name, pipeline_bson) in facets {
|
||||||
|
let pipeline = match pipeline_bson {
|
||||||
|
Bson::Array(arr) => {
|
||||||
|
let mut stages = Vec::new();
|
||||||
|
for stage in arr {
|
||||||
|
if let Bson::Document(d) = stage {
|
||||||
|
stages.push(d.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
stages
|
||||||
|
}
|
||||||
|
_ => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
let facet_result = Self::aggregate(docs.clone(), &pipeline, resolver, db)?;
|
||||||
|
result.insert(
|
||||||
|
facet_name.clone(),
|
||||||
|
Bson::Array(facet_result.into_iter().map(Bson::Document).collect()),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(vec![result])
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stage_union_with(
|
||||||
|
mut docs: Vec<Document>,
|
||||||
|
spec: &Bson,
|
||||||
|
resolver: Option<&dyn CollectionResolver>,
|
||||||
|
db: &str,
|
||||||
|
) -> Result<Vec<Document>, QueryError> {
|
||||||
|
let (coll, pipeline) = match spec {
|
||||||
|
Bson::String(s) => (s.as_str(), None),
|
||||||
|
Bson::Document(d) => {
|
||||||
|
let coll = d.get_str("coll").map_err(|_| {
|
||||||
|
QueryError::AggregationError("$unionWith requires 'coll'".into())
|
||||||
|
})?;
|
||||||
|
let pipeline = d.get_array("pipeline").ok().map(|arr| {
|
||||||
|
arr.iter()
|
||||||
|
.filter_map(|s| {
|
||||||
|
if let Bson::Document(d) = s {
|
||||||
|
Some(d.clone())
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect::<Vec<Document>>()
|
||||||
|
});
|
||||||
|
(coll, pipeline)
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::AggregationError(
|
||||||
|
"$unionWith requires a string or document".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let resolver = resolver.ok_or_else(|| {
|
||||||
|
QueryError::AggregationError("$unionWith requires a collection resolver".into())
|
||||||
|
})?;
|
||||||
|
let mut other_docs = resolver.resolve(db, coll)?;
|
||||||
|
|
||||||
|
if let Some(p) = pipeline {
|
||||||
|
other_docs = Self::aggregate(other_docs, &p, Some(resolver), db)?;
|
||||||
|
}
|
||||||
|
|
||||||
|
docs.extend(other_docs);
|
||||||
|
Ok(docs)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Helper functions ---
|
||||||
|
|
||||||
|
fn resolve_expression(expr: &Bson, doc: &Document) -> Bson {
|
||||||
|
match expr {
|
||||||
|
Bson::String(s) if s.starts_with('$') => {
|
||||||
|
let field = &s[1..];
|
||||||
|
get_nested_value(doc, field).unwrap_or(Bson::Null)
|
||||||
|
}
|
||||||
|
_ => expr.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bson_to_usize(v: &Bson) -> Option<usize> {
|
||||||
|
match v {
|
||||||
|
Bson::Int32(n) => Some(*n as usize),
|
||||||
|
Bson::Int64(n) => Some(*n as usize),
|
||||||
|
Bson::Double(n) => Some(*n as usize),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bson_to_f64(v: &Bson) -> Option<f64> {
|
||||||
|
match v {
|
||||||
|
Bson::Int32(n) => Some(*n as f64),
|
||||||
|
Bson::Int64(n) => Some(*n as f64),
|
||||||
|
Bson::Double(n) => Some(*n),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bson_loose_eq(a: &Bson, b: &Bson) -> bool {
|
||||||
|
match (a, b) {
|
||||||
|
(Bson::Int32(x), Bson::Int64(y)) => (*x as i64) == *y,
|
||||||
|
(Bson::Int64(x), Bson::Int32(y)) => *x == (*y as i64),
|
||||||
|
(Bson::Int32(x), Bson::Double(y)) => (*x as f64) == *y,
|
||||||
|
(Bson::Double(x), Bson::Int32(y)) => *x == (*y as f64),
|
||||||
|
_ => a == b,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Accumulators ---
|
||||||
|
|
||||||
|
fn accumulate_sum(docs: &[Document], expr: &Bson) -> Bson {
|
||||||
|
match expr {
|
||||||
|
Bson::Int32(n) => Bson::Int64(*n as i64 * docs.len() as i64),
|
||||||
|
Bson::Int64(n) => Bson::Int64(*n * docs.len() as i64),
|
||||||
|
Bson::String(s) if s.starts_with('$') => {
|
||||||
|
let field = &s[1..];
|
||||||
|
let mut sum = 0.0f64;
|
||||||
|
let mut is_int = true;
|
||||||
|
let mut int_sum = 0i64;
|
||||||
|
for doc in docs {
|
||||||
|
if let Some(val) = get_nested_value(doc, field) {
|
||||||
|
if let Some(n) = bson_to_f64(&val) {
|
||||||
|
sum += n;
|
||||||
|
if is_int {
|
||||||
|
match &val {
|
||||||
|
Bson::Int32(i) => int_sum += *i as i64,
|
||||||
|
Bson::Int64(i) => int_sum += i,
|
||||||
|
_ => is_int = false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if is_int {
|
||||||
|
Bson::Int64(int_sum)
|
||||||
|
} else {
|
||||||
|
Bson::Double(sum)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => Bson::Int32(0),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn accumulate_avg(docs: &[Document], expr: &Bson) -> Bson {
|
||||||
|
if docs.is_empty() {
|
||||||
|
return Bson::Null;
|
||||||
|
}
|
||||||
|
let field = match expr {
|
||||||
|
Bson::String(s) if s.starts_with('$') => &s[1..],
|
||||||
|
_ => return Bson::Null,
|
||||||
|
};
|
||||||
|
let mut sum = 0.0f64;
|
||||||
|
let mut count = 0usize;
|
||||||
|
for doc in docs {
|
||||||
|
if let Some(val) = get_nested_value(doc, field) {
|
||||||
|
if let Some(n) = bson_to_f64(&val) {
|
||||||
|
sum += n;
|
||||||
|
count += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if count == 0 {
|
||||||
|
Bson::Null
|
||||||
|
} else {
|
||||||
|
Bson::Double(sum / count as f64)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn accumulate_min(docs: &[Document], expr: &Bson) -> Bson {
|
||||||
|
let field = match expr {
|
||||||
|
Bson::String(s) if s.starts_with('$') => &s[1..],
|
||||||
|
_ => return Bson::Null,
|
||||||
|
};
|
||||||
|
let mut min: Option<Bson> = None;
|
||||||
|
for doc in docs {
|
||||||
|
if let Some(val) = get_nested_value(doc, field) {
|
||||||
|
min = Some(match min {
|
||||||
|
None => val,
|
||||||
|
Some(current) => {
|
||||||
|
if let (Some(cv), Some(vv)) = (bson_to_f64(¤t), bson_to_f64(&val)) {
|
||||||
|
if vv < cv {
|
||||||
|
val
|
||||||
|
} else {
|
||||||
|
current
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
current
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
min.unwrap_or(Bson::Null)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn accumulate_max(docs: &[Document], expr: &Bson) -> Bson {
|
||||||
|
let field = match expr {
|
||||||
|
Bson::String(s) if s.starts_with('$') => &s[1..],
|
||||||
|
_ => return Bson::Null,
|
||||||
|
};
|
||||||
|
let mut max: Option<Bson> = None;
|
||||||
|
for doc in docs {
|
||||||
|
if let Some(val) = get_nested_value(doc, field) {
|
||||||
|
max = Some(match max {
|
||||||
|
None => val,
|
||||||
|
Some(current) => {
|
||||||
|
if let (Some(cv), Some(vv)) = (bson_to_f64(¤t), bson_to_f64(&val)) {
|
||||||
|
if vv > cv {
|
||||||
|
val
|
||||||
|
} else {
|
||||||
|
current
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
current
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
max.unwrap_or(Bson::Null)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn accumulate_first(docs: &[Document], expr: &Bson) -> Bson {
|
||||||
|
let field = match expr {
|
||||||
|
Bson::String(s) if s.starts_with('$') => &s[1..],
|
||||||
|
_ => return Bson::Null,
|
||||||
|
};
|
||||||
|
docs.first()
|
||||||
|
.and_then(|doc| get_nested_value(doc, field))
|
||||||
|
.unwrap_or(Bson::Null)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn accumulate_last(docs: &[Document], expr: &Bson) -> Bson {
|
||||||
|
let field = match expr {
|
||||||
|
Bson::String(s) if s.starts_with('$') => &s[1..],
|
||||||
|
_ => return Bson::Null,
|
||||||
|
};
|
||||||
|
docs.last()
|
||||||
|
.and_then(|doc| get_nested_value(doc, field))
|
||||||
|
.unwrap_or(Bson::Null)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn accumulate_push(docs: &[Document], expr: &Bson) -> Bson {
|
||||||
|
let field = match expr {
|
||||||
|
Bson::String(s) if s.starts_with('$') => &s[1..],
|
||||||
|
_ => return Bson::Array(vec![]),
|
||||||
|
};
|
||||||
|
let values: Vec<Bson> = docs
|
||||||
|
.iter()
|
||||||
|
.filter_map(|doc| get_nested_value(doc, field))
|
||||||
|
.collect();
|
||||||
|
Bson::Array(values)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn accumulate_add_to_set(docs: &[Document], expr: &Bson) -> Bson {
|
||||||
|
let field = match expr {
|
||||||
|
Bson::String(s) if s.starts_with('$') => &s[1..],
|
||||||
|
_ => return Bson::Array(vec![]),
|
||||||
|
};
|
||||||
|
let mut seen = std::collections::HashSet::new();
|
||||||
|
let mut values = Vec::new();
|
||||||
|
for doc in docs {
|
||||||
|
if let Some(val) = get_nested_value(doc, field) {
|
||||||
|
let key = format!("{:?}", val);
|
||||||
|
if seen.insert(key) {
|
||||||
|
values.push(val);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Bson::Array(values)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_match_stage() {
|
||||||
|
let docs = vec![
|
||||||
|
bson::doc! { "x": 1 },
|
||||||
|
bson::doc! { "x": 2 },
|
||||||
|
bson::doc! { "x": 3 },
|
||||||
|
];
|
||||||
|
let pipeline = vec![bson::doc! { "$match": { "x": { "$gt": 1 } } }];
|
||||||
|
let result = AggregationEngine::aggregate(docs, &pipeline, None, "test").unwrap();
|
||||||
|
assert_eq!(result.len(), 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_group_stage() {
|
||||||
|
let docs = vec![
|
||||||
|
bson::doc! { "category": "a", "value": 10 },
|
||||||
|
bson::doc! { "category": "b", "value": 20 },
|
||||||
|
bson::doc! { "category": "a", "value": 30 },
|
||||||
|
];
|
||||||
|
let pipeline = vec![bson::doc! {
|
||||||
|
"$group": {
|
||||||
|
"_id": "$category",
|
||||||
|
"total": { "$sum": "$value" }
|
||||||
|
}
|
||||||
|
}];
|
||||||
|
let result = AggregationEngine::aggregate(docs, &pipeline, None, "test").unwrap();
|
||||||
|
assert_eq!(result.len(), 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_sort_limit_skip() {
|
||||||
|
let docs = vec![
|
||||||
|
bson::doc! { "x": 3 },
|
||||||
|
bson::doc! { "x": 1 },
|
||||||
|
bson::doc! { "x": 2 },
|
||||||
|
bson::doc! { "x": 4 },
|
||||||
|
];
|
||||||
|
let pipeline = vec![
|
||||||
|
bson::doc! { "$sort": { "x": 1 } },
|
||||||
|
bson::doc! { "$skip": 1_i64 },
|
||||||
|
bson::doc! { "$limit": 2_i64 },
|
||||||
|
];
|
||||||
|
let result = AggregationEngine::aggregate(docs, &pipeline, None, "test").unwrap();
|
||||||
|
assert_eq!(result.len(), 2);
|
||||||
|
assert_eq!(result[0].get_i32("x").unwrap(), 2);
|
||||||
|
assert_eq!(result[1].get_i32("x").unwrap(), 3);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,80 @@
|
|||||||
|
use bson::{Bson, Document};
|
||||||
|
use std::collections::HashSet;
|
||||||
|
|
||||||
|
use crate::field_path::get_nested_value;
|
||||||
|
use crate::matcher::QueryMatcher;
|
||||||
|
|
||||||
|
/// Get distinct values for a field across documents, with optional filter.
|
||||||
|
/// Handles array flattening (each array element counted separately).
|
||||||
|
pub fn distinct_values(
|
||||||
|
docs: &[Document],
|
||||||
|
field: &str,
|
||||||
|
filter: Option<&Document>,
|
||||||
|
) -> Vec<Bson> {
|
||||||
|
let filtered: Vec<&Document> = if let Some(f) = filter {
|
||||||
|
docs.iter().filter(|d| QueryMatcher::matches(d, f)).collect()
|
||||||
|
} else {
|
||||||
|
docs.iter().collect()
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut seen = HashSet::new();
|
||||||
|
let mut result = Vec::new();
|
||||||
|
|
||||||
|
for doc in &filtered {
|
||||||
|
let value = if field.contains('.') {
|
||||||
|
get_nested_value(doc, field)
|
||||||
|
} else {
|
||||||
|
doc.get(field).cloned()
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Some(val) = value {
|
||||||
|
collect_distinct_values(&val, &mut seen, &mut result);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
fn collect_distinct_values(value: &Bson, seen: &mut HashSet<String>, result: &mut Vec<Bson>) {
|
||||||
|
match value {
|
||||||
|
Bson::Array(arr) => {
|
||||||
|
// Flatten: each array element is a separate value
|
||||||
|
for elem in arr {
|
||||||
|
collect_distinct_values(elem, seen, result);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
let key = format!("{:?}", value);
|
||||||
|
if seen.insert(key) {
|
||||||
|
result.push(value.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_distinct_simple() {
|
||||||
|
let docs = vec![
|
||||||
|
bson::doc! { "x": 1 },
|
||||||
|
bson::doc! { "x": 2 },
|
||||||
|
bson::doc! { "x": 1 },
|
||||||
|
bson::doc! { "x": 3 },
|
||||||
|
];
|
||||||
|
let result = distinct_values(&docs, "x", None);
|
||||||
|
assert_eq!(result.len(), 3);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_distinct_array_flattening() {
|
||||||
|
let docs = vec![
|
||||||
|
bson::doc! { "tags": ["a", "b"] },
|
||||||
|
bson::doc! { "tags": ["b", "c"] },
|
||||||
|
];
|
||||||
|
let result = distinct_values(&docs, "tags", None);
|
||||||
|
assert_eq!(result.len(), 3); // a, b, c
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,18 @@
|
|||||||
|
/// Errors from query operations.
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum QueryError {
|
||||||
|
#[error("Invalid query operator: {0}")]
|
||||||
|
InvalidOperator(String),
|
||||||
|
|
||||||
|
#[error("Type mismatch: {0}")]
|
||||||
|
TypeMismatch(String),
|
||||||
|
|
||||||
|
#[error("Invalid update: {0}")]
|
||||||
|
InvalidUpdate(String),
|
||||||
|
|
||||||
|
#[error("Aggregation error: {0}")]
|
||||||
|
AggregationError(String),
|
||||||
|
|
||||||
|
#[error("Invalid regex: {0}")]
|
||||||
|
InvalidRegex(String),
|
||||||
|
}
|
||||||
@@ -0,0 +1,115 @@
|
|||||||
|
use bson::{Bson, Document};
|
||||||
|
|
||||||
|
/// Get a nested value from a document using dot-notation path (e.g., "a.b.c").
|
||||||
|
/// Handles both nested documents and array traversal.
|
||||||
|
pub fn get_nested_value(doc: &Document, path: &str) -> Option<Bson> {
|
||||||
|
let parts: Vec<&str> = path.split('.').collect();
|
||||||
|
get_nested_recursive(&Bson::Document(doc.clone()), &parts)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_nested_recursive(value: &Bson, parts: &[&str]) -> Option<Bson> {
|
||||||
|
if parts.is_empty() {
|
||||||
|
return Some(value.clone());
|
||||||
|
}
|
||||||
|
|
||||||
|
let key = parts[0];
|
||||||
|
let rest = &parts[1..];
|
||||||
|
|
||||||
|
match value {
|
||||||
|
Bson::Document(doc) => {
|
||||||
|
let child = doc.get(key)?;
|
||||||
|
get_nested_recursive(child, rest)
|
||||||
|
}
|
||||||
|
Bson::Array(arr) => {
|
||||||
|
// Try numeric index first
|
||||||
|
if let Ok(idx) = key.parse::<usize>() {
|
||||||
|
if let Some(elem) = arr.get(idx) {
|
||||||
|
return get_nested_recursive(elem, rest);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Otherwise, collect from all elements
|
||||||
|
let results: Vec<Bson> = arr
|
||||||
|
.iter()
|
||||||
|
.filter_map(|elem| get_nested_recursive(elem, parts))
|
||||||
|
.collect();
|
||||||
|
if results.is_empty() {
|
||||||
|
None
|
||||||
|
} else if results.len() == 1 {
|
||||||
|
Some(results.into_iter().next().unwrap())
|
||||||
|
} else {
|
||||||
|
Some(Bson::Array(results))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Set a nested value in a document using dot-notation path.
|
||||||
|
pub fn set_nested_value(doc: &mut Document, path: &str, value: Bson) {
|
||||||
|
let parts: Vec<&str> = path.split('.').collect();
|
||||||
|
set_nested_recursive(doc, &parts, value);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn set_nested_recursive(doc: &mut Document, parts: &[&str], value: Bson) {
|
||||||
|
if parts.len() == 1 {
|
||||||
|
doc.insert(parts[0].to_string(), value);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let key = parts[0];
|
||||||
|
let rest = &parts[1..];
|
||||||
|
|
||||||
|
// Get or create nested document
|
||||||
|
if !doc.contains_key(key) {
|
||||||
|
doc.insert(key.to_string(), Bson::Document(Document::new()));
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(Bson::Document(ref mut nested)) = doc.get_mut(key) {
|
||||||
|
set_nested_recursive(nested, rest, value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Remove a nested value from a document using dot-notation path.
|
||||||
|
pub fn remove_nested_value(doc: &mut Document, path: &str) -> Option<Bson> {
|
||||||
|
let parts: Vec<&str> = path.split('.').collect();
|
||||||
|
remove_nested_recursive(doc, &parts)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn remove_nested_recursive(doc: &mut Document, parts: &[&str]) -> Option<Bson> {
|
||||||
|
if parts.len() == 1 {
|
||||||
|
return doc.remove(parts[0]);
|
||||||
|
}
|
||||||
|
|
||||||
|
let key = parts[0];
|
||||||
|
let rest = &parts[1..];
|
||||||
|
|
||||||
|
if let Some(Bson::Document(ref mut nested)) = doc.get_mut(key) {
|
||||||
|
remove_nested_recursive(nested, rest)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_get_nested_simple() {
|
||||||
|
let doc = bson::doc! { "a": { "b": { "c": 42 } } };
|
||||||
|
assert_eq!(get_nested_value(&doc, "a.b.c"), Some(Bson::Int32(42)));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_get_nested_missing() {
|
||||||
|
let doc = bson::doc! { "a": { "b": 1 } };
|
||||||
|
assert_eq!(get_nested_value(&doc, "a.c"), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_set_nested() {
|
||||||
|
let mut doc = bson::doc! {};
|
||||||
|
set_nested_value(&mut doc, "a.b.c", Bson::Int32(42));
|
||||||
|
assert_eq!(get_nested_value(&doc, "a.b.c"), Some(Bson::Int32(42)));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,16 @@
|
|||||||
|
mod matcher;
|
||||||
|
mod update;
|
||||||
|
mod sort;
|
||||||
|
mod projection;
|
||||||
|
mod distinct;
|
||||||
|
pub mod aggregation;
|
||||||
|
mod field_path;
|
||||||
|
pub mod error;
|
||||||
|
|
||||||
|
pub use matcher::QueryMatcher;
|
||||||
|
pub use update::UpdateEngine;
|
||||||
|
pub use sort::sort_documents;
|
||||||
|
pub use projection::apply_projection;
|
||||||
|
pub use distinct::distinct_values;
|
||||||
|
pub use aggregation::AggregationEngine;
|
||||||
|
pub use field_path::{get_nested_value, set_nested_value};
|
||||||
@@ -0,0 +1,574 @@
|
|||||||
|
use bson::{Bson, Document};
|
||||||
|
use regex::Regex;
|
||||||
|
|
||||||
|
use crate::field_path::get_nested_value;
|
||||||
|
|
||||||
|
/// Query matching engine.
|
||||||
|
/// Evaluates filter documents against BSON documents.
|
||||||
|
pub struct QueryMatcher;
|
||||||
|
|
||||||
|
impl QueryMatcher {
|
||||||
|
/// Test whether a single document matches a filter.
|
||||||
|
pub fn matches(doc: &Document, filter: &Document) -> bool {
|
||||||
|
Self::matches_filter(doc, filter)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Filter a slice of documents, returning those that match.
|
||||||
|
pub fn filter(docs: &[Document], filter: &Document) -> Vec<Document> {
|
||||||
|
if filter.is_empty() {
|
||||||
|
return docs.to_vec();
|
||||||
|
}
|
||||||
|
docs.iter()
|
||||||
|
.filter(|doc| Self::matches_filter(doc, filter))
|
||||||
|
.cloned()
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Find the first document matching a filter.
|
||||||
|
pub fn find_one(docs: &[Document], filter: &Document) -> Option<Document> {
|
||||||
|
docs.iter()
|
||||||
|
.find(|doc| Self::matches_filter(doc, filter))
|
||||||
|
.cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn matches_filter(doc: &Document, filter: &Document) -> bool {
|
||||||
|
for (key, value) in filter {
|
||||||
|
if !Self::matches_condition(doc, key, value) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
fn matches_condition(doc: &Document, key: &str, condition: &Bson) -> bool {
|
||||||
|
match key {
|
||||||
|
"$and" => Self::match_logical_and(doc, condition),
|
||||||
|
"$or" => Self::match_logical_or(doc, condition),
|
||||||
|
"$nor" => Self::match_logical_nor(doc, condition),
|
||||||
|
"$not" => Self::match_logical_not(doc, condition),
|
||||||
|
"$expr" => {
|
||||||
|
// Basic $expr support - just return true for now
|
||||||
|
true
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
// Field condition
|
||||||
|
match condition {
|
||||||
|
Bson::Document(cond_doc) if Self::has_operators(cond_doc) => {
|
||||||
|
Self::match_field_operators(doc, key, cond_doc)
|
||||||
|
}
|
||||||
|
// Implicit equality
|
||||||
|
_ => Self::match_equality(doc, key, condition),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn has_operators(doc: &Document) -> bool {
|
||||||
|
doc.keys().any(|k| k.starts_with('$'))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Public accessor for has_operators (used by update engine).
|
||||||
|
pub fn has_operators_pub(doc: &Document) -> bool {
|
||||||
|
Self::has_operators(doc)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Public accessor for bson_compare (used by update engine).
|
||||||
|
pub fn bson_compare_pub(a: &Bson, b: &Bson) -> Option<std::cmp::Ordering> {
|
||||||
|
Self::bson_compare(a, b)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn match_equality(doc: &Document, field: &str, expected: &Bson) -> bool {
|
||||||
|
let actual = Self::resolve_field(doc, field);
|
||||||
|
match actual {
|
||||||
|
Some(val) => Self::bson_equals(&val, expected),
|
||||||
|
None => matches!(expected, Bson::Null),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn match_field_operators(doc: &Document, field: &str, operators: &Document) -> bool {
|
||||||
|
let actual = Self::resolve_field(doc, field);
|
||||||
|
|
||||||
|
for (op, op_value) in operators {
|
||||||
|
let result = match op.as_str() {
|
||||||
|
"$eq" => Self::op_eq(&actual, op_value),
|
||||||
|
"$ne" => Self::op_ne(&actual, op_value),
|
||||||
|
"$gt" => Self::op_cmp(&actual, op_value, CmpOp::Gt),
|
||||||
|
"$gte" => Self::op_cmp(&actual, op_value, CmpOp::Gte),
|
||||||
|
"$lt" => Self::op_cmp(&actual, op_value, CmpOp::Lt),
|
||||||
|
"$lte" => Self::op_cmp(&actual, op_value, CmpOp::Lte),
|
||||||
|
"$in" => Self::op_in(&actual, op_value),
|
||||||
|
"$nin" => Self::op_nin(&actual, op_value),
|
||||||
|
"$exists" => Self::op_exists(&actual, op_value),
|
||||||
|
"$type" => Self::op_type(&actual, op_value),
|
||||||
|
"$regex" => Self::op_regex(&actual, op_value, operators.get("$options")),
|
||||||
|
"$not" => Self::op_not(doc, field, op_value),
|
||||||
|
"$elemMatch" => Self::op_elem_match(&actual, op_value),
|
||||||
|
"$size" => Self::op_size(&actual, op_value),
|
||||||
|
"$all" => Self::op_all(&actual, op_value),
|
||||||
|
"$mod" => Self::op_mod(&actual, op_value),
|
||||||
|
"$options" => continue, // handled by $regex
|
||||||
|
_ => true, // unknown operator, skip
|
||||||
|
};
|
||||||
|
if !result {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
fn resolve_field(doc: &Document, field: &str) -> Option<Bson> {
|
||||||
|
if field.contains('.') {
|
||||||
|
get_nested_value(doc, field)
|
||||||
|
} else {
|
||||||
|
doc.get(field).cloned()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bson_equals(a: &Bson, b: &Bson) -> bool {
|
||||||
|
match (a, b) {
|
||||||
|
(Bson::Int32(x), Bson::Int64(y)) => (*x as i64) == *y,
|
||||||
|
(Bson::Int64(x), Bson::Int32(y)) => *x == (*y as i64),
|
||||||
|
(Bson::Int32(x), Bson::Double(y)) => (*x as f64) == *y,
|
||||||
|
(Bson::Double(x), Bson::Int32(y)) => *x == (*y as f64),
|
||||||
|
(Bson::Int64(x), Bson::Double(y)) => (*x as f64) == *y,
|
||||||
|
(Bson::Double(x), Bson::Int64(y)) => *x == (*y as f64),
|
||||||
|
// For arrays, check if any element matches (implicit $elemMatch)
|
||||||
|
(Bson::Array(arr), _) if !matches!(b, Bson::Array(_)) => {
|
||||||
|
arr.iter().any(|elem| Self::bson_equals(elem, b))
|
||||||
|
}
|
||||||
|
_ => a == b,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bson_compare(a: &Bson, b: &Bson) -> Option<std::cmp::Ordering> {
|
||||||
|
use std::cmp::Ordering;
|
||||||
|
|
||||||
|
match (a, b) {
|
||||||
|
// Numeric comparisons (cross-type)
|
||||||
|
(Bson::Int32(x), Bson::Int32(y)) => Some(x.cmp(y)),
|
||||||
|
(Bson::Int64(x), Bson::Int64(y)) => Some(x.cmp(y)),
|
||||||
|
(Bson::Double(x), Bson::Double(y)) => x.partial_cmp(y),
|
||||||
|
(Bson::Int32(x), Bson::Int64(y)) => Some((*x as i64).cmp(y)),
|
||||||
|
(Bson::Int64(x), Bson::Int32(y)) => Some(x.cmp(&(*y as i64))),
|
||||||
|
(Bson::Int32(x), Bson::Double(y)) => (*x as f64).partial_cmp(y),
|
||||||
|
(Bson::Double(x), Bson::Int32(y)) => x.partial_cmp(&(*y as f64)),
|
||||||
|
(Bson::Int64(x), Bson::Double(y)) => (*x as f64).partial_cmp(y),
|
||||||
|
(Bson::Double(x), Bson::Int64(y)) => x.partial_cmp(&(*y as f64)),
|
||||||
|
|
||||||
|
// String comparisons
|
||||||
|
(Bson::String(x), Bson::String(y)) => Some(x.cmp(y)),
|
||||||
|
|
||||||
|
// DateTime comparisons
|
||||||
|
(Bson::DateTime(x), Bson::DateTime(y)) => Some(x.cmp(y)),
|
||||||
|
|
||||||
|
// Boolean comparisons
|
||||||
|
(Bson::Boolean(x), Bson::Boolean(y)) => Some(x.cmp(y)),
|
||||||
|
|
||||||
|
// ObjectId comparisons
|
||||||
|
(Bson::ObjectId(x), Bson::ObjectId(y)) => Some(x.cmp(y)),
|
||||||
|
|
||||||
|
// Null comparisons
|
||||||
|
(Bson::Null, Bson::Null) => Some(Ordering::Equal),
|
||||||
|
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Operator implementations ---
|
||||||
|
|
||||||
|
fn op_eq(actual: &Option<Bson>, expected: &Bson) -> bool {
|
||||||
|
match actual {
|
||||||
|
Some(val) => Self::bson_equals(val, expected),
|
||||||
|
None => matches!(expected, Bson::Null),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn op_ne(actual: &Option<Bson>, expected: &Bson) -> bool {
|
||||||
|
!Self::op_eq(actual, expected)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn op_cmp(actual: &Option<Bson>, expected: &Bson, op: CmpOp) -> bool {
|
||||||
|
let val = match actual {
|
||||||
|
Some(v) => v,
|
||||||
|
None => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
// For arrays, check if any element satisfies the comparison
|
||||||
|
if let Bson::Array(arr) = val {
|
||||||
|
return arr.iter().any(|elem| {
|
||||||
|
if let Some(ord) = Self::bson_compare(elem, expected) {
|
||||||
|
op.check(ord)
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(ord) = Self::bson_compare(val, expected) {
|
||||||
|
op.check(ord)
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn op_in(actual: &Option<Bson>, values: &Bson) -> bool {
|
||||||
|
let arr = match values {
|
||||||
|
Bson::Array(a) => a,
|
||||||
|
_ => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
match actual {
|
||||||
|
Some(val) => {
|
||||||
|
// For array values, check if any element is in the list
|
||||||
|
if let Bson::Array(actual_arr) = val {
|
||||||
|
actual_arr.iter().any(|elem| {
|
||||||
|
arr.iter().any(|v| Self::bson_equals(elem, v))
|
||||||
|
}) || arr.iter().any(|v| Self::bson_equals(val, v))
|
||||||
|
} else {
|
||||||
|
arr.iter().any(|v| Self::bson_equals(val, v))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None => arr.iter().any(|v| matches!(v, Bson::Null)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn op_nin(actual: &Option<Bson>, values: &Bson) -> bool {
|
||||||
|
!Self::op_in(actual, values)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn op_exists(actual: &Option<Bson>, expected: &Bson) -> bool {
|
||||||
|
let should_exist = match expected {
|
||||||
|
Bson::Boolean(b) => *b,
|
||||||
|
Bson::Int32(n) => *n != 0,
|
||||||
|
Bson::Int64(n) => *n != 0,
|
||||||
|
_ => true,
|
||||||
|
};
|
||||||
|
actual.is_some() == should_exist
|
||||||
|
}
|
||||||
|
|
||||||
|
fn op_type(actual: &Option<Bson>, expected: &Bson) -> bool {
|
||||||
|
let val = match actual {
|
||||||
|
Some(v) => v,
|
||||||
|
None => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let type_num = match expected {
|
||||||
|
Bson::Int32(n) => *n,
|
||||||
|
Bson::String(s) => match s.as_str() {
|
||||||
|
"double" => 1,
|
||||||
|
"string" => 2,
|
||||||
|
"object" => 3,
|
||||||
|
"array" => 4,
|
||||||
|
"binData" => 5,
|
||||||
|
"objectId" => 7,
|
||||||
|
"bool" => 8,
|
||||||
|
"date" => 9,
|
||||||
|
"null" => 10,
|
||||||
|
"regex" => 11,
|
||||||
|
"int" => 16,
|
||||||
|
"long" => 18,
|
||||||
|
"decimal" => 19,
|
||||||
|
"number" => -1, // special: any numeric type
|
||||||
|
_ => return false,
|
||||||
|
},
|
||||||
|
_ => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
if type_num == -1 {
|
||||||
|
return matches!(val, Bson::Int32(_) | Bson::Int64(_) | Bson::Double(_));
|
||||||
|
}
|
||||||
|
|
||||||
|
let actual_type = match val {
|
||||||
|
Bson::Double(_) => 1,
|
||||||
|
Bson::String(_) => 2,
|
||||||
|
Bson::Document(_) => 3,
|
||||||
|
Bson::Array(_) => 4,
|
||||||
|
Bson::Binary(_) => 5,
|
||||||
|
Bson::ObjectId(_) => 7,
|
||||||
|
Bson::Boolean(_) => 8,
|
||||||
|
Bson::DateTime(_) => 9,
|
||||||
|
Bson::Null => 10,
|
||||||
|
Bson::RegularExpression(_) => 11,
|
||||||
|
Bson::Int32(_) => 16,
|
||||||
|
Bson::Int64(_) => 18,
|
||||||
|
Bson::Decimal128(_) => 19,
|
||||||
|
_ => 0,
|
||||||
|
};
|
||||||
|
|
||||||
|
actual_type == type_num
|
||||||
|
}
|
||||||
|
|
||||||
|
fn op_regex(actual: &Option<Bson>, pattern: &Bson, options: Option<&Bson>) -> bool {
|
||||||
|
let val = match actual {
|
||||||
|
Some(Bson::String(s)) => s.as_str(),
|
||||||
|
_ => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let pattern_str = match pattern {
|
||||||
|
Bson::String(s) => s.as_str(),
|
||||||
|
Bson::RegularExpression(re) => re.pattern.as_str(),
|
||||||
|
_ => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let opts = match options {
|
||||||
|
Some(Bson::String(s)) => s.as_str(),
|
||||||
|
_ => match pattern {
|
||||||
|
Bson::RegularExpression(re) => re.options.as_str(),
|
||||||
|
_ => "",
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut regex_pattern = String::new();
|
||||||
|
if opts.contains('i') {
|
||||||
|
regex_pattern.push_str("(?i)");
|
||||||
|
}
|
||||||
|
if opts.contains('m') {
|
||||||
|
regex_pattern.push_str("(?m)");
|
||||||
|
}
|
||||||
|
if opts.contains('s') {
|
||||||
|
regex_pattern.push_str("(?s)");
|
||||||
|
}
|
||||||
|
regex_pattern.push_str(pattern_str);
|
||||||
|
|
||||||
|
match Regex::new(®ex_pattern) {
|
||||||
|
Ok(re) => re.is_match(val),
|
||||||
|
Err(_) => false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn op_not(doc: &Document, field: &str, condition: &Bson) -> bool {
|
||||||
|
match condition {
|
||||||
|
Bson::Document(cond_doc) => !Self::match_field_operators(doc, field, cond_doc),
|
||||||
|
_ => true,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn op_elem_match(actual: &Option<Bson>, condition: &Bson) -> bool {
|
||||||
|
let arr = match actual {
|
||||||
|
Some(Bson::Array(a)) => a,
|
||||||
|
_ => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let cond_doc = match condition {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
arr.iter().any(|elem| {
|
||||||
|
if let Bson::Document(elem_doc) = elem {
|
||||||
|
Self::matches_filter(elem_doc, cond_doc)
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn op_size(actual: &Option<Bson>, expected: &Bson) -> bool {
|
||||||
|
let arr = match actual {
|
||||||
|
Some(Bson::Array(a)) => a,
|
||||||
|
_ => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let expected_size = match expected {
|
||||||
|
Bson::Int32(n) => *n as usize,
|
||||||
|
Bson::Int64(n) => *n as usize,
|
||||||
|
_ => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
arr.len() == expected_size
|
||||||
|
}
|
||||||
|
|
||||||
|
fn op_all(actual: &Option<Bson>, expected: &Bson) -> bool {
|
||||||
|
let arr = match actual {
|
||||||
|
Some(Bson::Array(a)) => a,
|
||||||
|
_ => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let expected_arr = match expected {
|
||||||
|
Bson::Array(a) => a,
|
||||||
|
_ => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
expected_arr.iter().all(|expected_val| {
|
||||||
|
arr.iter().any(|elem| Self::bson_equals(elem, expected_val))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn op_mod(actual: &Option<Bson>, expected: &Bson) -> bool {
|
||||||
|
let val = match actual {
|
||||||
|
Some(v) => match v {
|
||||||
|
Bson::Int32(n) => *n as i64,
|
||||||
|
Bson::Int64(n) => *n,
|
||||||
|
Bson::Double(n) => *n as i64,
|
||||||
|
_ => return false,
|
||||||
|
},
|
||||||
|
None => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let arr = match expected {
|
||||||
|
Bson::Array(a) if a.len() == 2 => a,
|
||||||
|
_ => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let divisor = match &arr[0] {
|
||||||
|
Bson::Int32(n) => *n as i64,
|
||||||
|
Bson::Int64(n) => *n,
|
||||||
|
_ => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let remainder = match &arr[1] {
|
||||||
|
Bson::Int32(n) => *n as i64,
|
||||||
|
Bson::Int64(n) => *n,
|
||||||
|
_ => return false,
|
||||||
|
};
|
||||||
|
|
||||||
|
if divisor == 0 {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
val % divisor == remainder
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Logical operators ---
|
||||||
|
|
||||||
|
fn match_logical_and(doc: &Document, conditions: &Bson) -> bool {
|
||||||
|
match conditions {
|
||||||
|
Bson::Array(arr) => arr.iter().all(|cond| {
|
||||||
|
if let Bson::Document(cond_doc) = cond {
|
||||||
|
Self::matches_filter(doc, cond_doc)
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
_ => false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn match_logical_or(doc: &Document, conditions: &Bson) -> bool {
|
||||||
|
match conditions {
|
||||||
|
Bson::Array(arr) => arr.iter().any(|cond| {
|
||||||
|
if let Bson::Document(cond_doc) = cond {
|
||||||
|
Self::matches_filter(doc, cond_doc)
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
_ => false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn match_logical_nor(doc: &Document, conditions: &Bson) -> bool {
|
||||||
|
!Self::match_logical_or(doc, conditions)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn match_logical_not(doc: &Document, condition: &Bson) -> bool {
|
||||||
|
match condition {
|
||||||
|
Bson::Document(cond_doc) => !Self::matches_filter(doc, cond_doc),
|
||||||
|
_ => true,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy)]
|
||||||
|
enum CmpOp {
|
||||||
|
Gt,
|
||||||
|
Gte,
|
||||||
|
Lt,
|
||||||
|
Lte,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CmpOp {
|
||||||
|
fn check(self, ord: std::cmp::Ordering) -> bool {
|
||||||
|
use std::cmp::Ordering;
|
||||||
|
match self {
|
||||||
|
CmpOp::Gt => ord == Ordering::Greater,
|
||||||
|
CmpOp::Gte => ord == Ordering::Greater || ord == Ordering::Equal,
|
||||||
|
CmpOp::Lt => ord == Ordering::Less,
|
||||||
|
CmpOp::Lte => ord == Ordering::Less || ord == Ordering::Equal,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_simple_equality() {
|
||||||
|
let doc = bson::doc! { "name": "Alice", "age": 30 };
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! { "name": "Alice" }));
|
||||||
|
assert!(!QueryMatcher::matches(&doc, &bson::doc! { "name": "Bob" }));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_comparison_operators() {
|
||||||
|
let doc = bson::doc! { "age": 30 };
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! { "age": { "$gt": 25 } }));
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! { "age": { "$gte": 30 } }));
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! { "age": { "$lt": 35 } }));
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! { "age": { "$lte": 30 } }));
|
||||||
|
assert!(!QueryMatcher::matches(&doc, &bson::doc! { "age": { "$gt": 30 } }));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_in_operator() {
|
||||||
|
let doc = bson::doc! { "status": "active" };
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! { "status": { "$in": ["active", "pending"] } }));
|
||||||
|
assert!(!QueryMatcher::matches(&doc, &bson::doc! { "status": { "$in": ["closed"] } }));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_exists_operator() {
|
||||||
|
let doc = bson::doc! { "name": "Alice" };
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! { "name": { "$exists": true } }));
|
||||||
|
assert!(!QueryMatcher::matches(&doc, &bson::doc! { "age": { "$exists": true } }));
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! { "age": { "$exists": false } }));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_logical_or() {
|
||||||
|
let doc = bson::doc! { "age": 30 };
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! {
|
||||||
|
"$or": [{ "age": 30 }, { "age": 40 }]
|
||||||
|
}));
|
||||||
|
assert!(!QueryMatcher::matches(&doc, &bson::doc! {
|
||||||
|
"$or": [{ "age": 20 }, { "age": 40 }]
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_logical_and() {
|
||||||
|
let doc = bson::doc! { "age": 30, "name": "Alice" };
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! {
|
||||||
|
"$and": [{ "age": 30 }, { "name": "Alice" }]
|
||||||
|
}));
|
||||||
|
assert!(!QueryMatcher::matches(&doc, &bson::doc! {
|
||||||
|
"$and": [{ "age": 30 }, { "name": "Bob" }]
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_dot_notation() {
|
||||||
|
let doc = bson::doc! { "address": { "city": "NYC" } };
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! { "address.city": "NYC" }));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_ne_operator() {
|
||||||
|
let doc = bson::doc! { "status": "active" };
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! { "status": { "$ne": "closed" } }));
|
||||||
|
assert!(!QueryMatcher::matches(&doc, &bson::doc! { "status": { "$ne": "active" } }));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_cross_type_numeric_equality() {
|
||||||
|
let doc = bson::doc! { "count": 5_i32 };
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! { "count": 5_i64 }));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_empty_filter_matches_all() {
|
||||||
|
let doc = bson::doc! { "x": 1 };
|
||||||
|
assert!(QueryMatcher::matches(&doc, &bson::doc! {}));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,168 @@
|
|||||||
|
use bson::{Bson, Document};
|
||||||
|
|
||||||
|
use crate::field_path::get_nested_value;
|
||||||
|
|
||||||
|
/// Apply a projection to a document.
|
||||||
|
/// Inclusion mode: only specified fields + _id.
|
||||||
|
/// Exclusion mode: all fields except specified ones.
|
||||||
|
/// _id can be explicitly excluded in either mode.
|
||||||
|
pub fn apply_projection(doc: &Document, projection: &Document) -> Document {
|
||||||
|
if projection.is_empty() {
|
||||||
|
return doc.clone();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Determine mode: inclusion or exclusion
|
||||||
|
let mut has_inclusion = false;
|
||||||
|
let mut id_explicitly_set = false;
|
||||||
|
|
||||||
|
for (key, value) in projection {
|
||||||
|
if key == "_id" {
|
||||||
|
id_explicitly_set = true;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
match value {
|
||||||
|
Bson::Int32(0) | Bson::Int64(0) | Bson::Boolean(false) => {}
|
||||||
|
_ => has_inclusion = true,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if has_inclusion {
|
||||||
|
apply_inclusion(doc, projection, id_explicitly_set)
|
||||||
|
} else {
|
||||||
|
apply_exclusion(doc, projection)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_inclusion(doc: &Document, projection: &Document, id_explicitly_set: bool) -> Document {
|
||||||
|
let mut result = Document::new();
|
||||||
|
|
||||||
|
// Include _id by default unless explicitly excluded
|
||||||
|
let include_id = if id_explicitly_set {
|
||||||
|
is_truthy(projection.get("_id"))
|
||||||
|
} else {
|
||||||
|
true
|
||||||
|
};
|
||||||
|
|
||||||
|
if include_id {
|
||||||
|
if let Some(id) = doc.get("_id") {
|
||||||
|
result.insert("_id", id.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for (key, value) in projection {
|
||||||
|
if key == "_id" {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if !is_truthy(Some(value)) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if key.contains('.') {
|
||||||
|
if let Some(val) = get_nested_value(doc, key) {
|
||||||
|
// Rebuild nested structure
|
||||||
|
set_nested_in_result(&mut result, key, val);
|
||||||
|
}
|
||||||
|
} else if let Some(val) = doc.get(key) {
|
||||||
|
result.insert(key.clone(), val.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_exclusion(doc: &Document, projection: &Document) -> Document {
|
||||||
|
let mut result = doc.clone();
|
||||||
|
|
||||||
|
for (key, value) in projection {
|
||||||
|
if !is_truthy(Some(value)) {
|
||||||
|
if key.contains('.') {
|
||||||
|
// Remove nested field
|
||||||
|
remove_nested_from_result(&mut result, key);
|
||||||
|
} else {
|
||||||
|
result.remove(key);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_truthy(value: Option<&Bson>) -> bool {
|
||||||
|
match value {
|
||||||
|
None => false,
|
||||||
|
Some(Bson::Int32(0)) | Some(Bson::Int64(0)) | Some(Bson::Boolean(false)) => false,
|
||||||
|
_ => true,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn set_nested_in_result(doc: &mut Document, path: &str, value: Bson) {
|
||||||
|
let parts: Vec<&str> = path.split('.').collect();
|
||||||
|
set_nested_recursive(doc, &parts, value);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn set_nested_recursive(doc: &mut Document, parts: &[&str], value: Bson) {
|
||||||
|
if parts.len() == 1 {
|
||||||
|
doc.insert(parts[0].to_string(), value);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let key = parts[0];
|
||||||
|
if !doc.contains_key(key) {
|
||||||
|
doc.insert(key.to_string(), Bson::Document(Document::new()));
|
||||||
|
}
|
||||||
|
if let Some(Bson::Document(ref mut nested)) = doc.get_mut(key) {
|
||||||
|
set_nested_recursive(nested, &parts[1..], value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn remove_nested_from_result(doc: &mut Document, path: &str) {
|
||||||
|
let parts: Vec<&str> = path.split('.').collect();
|
||||||
|
remove_nested_recursive(doc, &parts);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn remove_nested_recursive(doc: &mut Document, parts: &[&str]) {
|
||||||
|
if parts.len() == 1 {
|
||||||
|
doc.remove(parts[0]);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let key = parts[0];
|
||||||
|
if let Some(Bson::Document(ref mut nested)) = doc.get_mut(key) {
|
||||||
|
remove_nested_recursive(nested, &parts[1..]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_inclusion_projection() {
|
||||||
|
let doc = bson::doc! { "_id": 1, "name": "Alice", "age": 30, "email": "a@b.c" };
|
||||||
|
let proj = bson::doc! { "name": 1, "age": 1 };
|
||||||
|
let result = apply_projection(&doc, &proj);
|
||||||
|
assert!(result.contains_key("_id"));
|
||||||
|
assert!(result.contains_key("name"));
|
||||||
|
assert!(result.contains_key("age"));
|
||||||
|
assert!(!result.contains_key("email"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_exclusion_projection() {
|
||||||
|
let doc = bson::doc! { "_id": 1, "name": "Alice", "age": 30 };
|
||||||
|
let proj = bson::doc! { "age": 0 };
|
||||||
|
let result = apply_projection(&doc, &proj);
|
||||||
|
assert!(result.contains_key("_id"));
|
||||||
|
assert!(result.contains_key("name"));
|
||||||
|
assert!(!result.contains_key("age"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_exclude_id() {
|
||||||
|
let doc = bson::doc! { "_id": 1, "name": "Alice" };
|
||||||
|
let proj = bson::doc! { "name": 1, "_id": 0 };
|
||||||
|
let result = apply_projection(&doc, &proj);
|
||||||
|
assert!(!result.contains_key("_id"));
|
||||||
|
assert!(result.contains_key("name"));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,137 @@
|
|||||||
|
use bson::{Bson, Document};
|
||||||
|
|
||||||
|
use crate::field_path::get_nested_value;
|
||||||
|
|
||||||
|
/// Sort documents according to a sort specification.
|
||||||
|
/// Sort spec: `{ field1: 1, field2: -1 }` where 1 = ascending, -1 = descending.
|
||||||
|
pub fn sort_documents(docs: &mut [Document], sort_spec: &Document) {
|
||||||
|
if sort_spec.is_empty() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
docs.sort_by(|a, b| {
|
||||||
|
for (field, direction) in sort_spec {
|
||||||
|
let ascending = match direction {
|
||||||
|
Bson::Int32(n) => *n > 0,
|
||||||
|
Bson::Int64(n) => *n > 0,
|
||||||
|
Bson::String(s) => !s.eq_ignore_ascii_case("desc") && !s.eq_ignore_ascii_case("descending"),
|
||||||
|
_ => true,
|
||||||
|
};
|
||||||
|
|
||||||
|
let a_val = get_value(a, field);
|
||||||
|
let b_val = get_value(b, field);
|
||||||
|
|
||||||
|
let ord = compare_bson_values(&a_val, &b_val);
|
||||||
|
let ord = if ascending { ord } else { ord.reverse() };
|
||||||
|
|
||||||
|
if ord != std::cmp::Ordering::Equal {
|
||||||
|
return ord;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
std::cmp::Ordering::Equal
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_value(doc: &Document, field: &str) -> Option<Bson> {
|
||||||
|
if field.contains('.') {
|
||||||
|
get_nested_value(doc, field)
|
||||||
|
} else {
|
||||||
|
doc.get(field).cloned()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Compare two BSON values for sorting purposes.
|
||||||
|
/// BSON type sort order: null < numbers < strings < objects < arrays < binData < ObjectId < bool < date
|
||||||
|
fn compare_bson_values(a: &Option<Bson>, b: &Option<Bson>) -> std::cmp::Ordering {
|
||||||
|
use std::cmp::Ordering;
|
||||||
|
|
||||||
|
match (a, b) {
|
||||||
|
(None, None) => Ordering::Equal,
|
||||||
|
(None, Some(Bson::Null)) => Ordering::Equal,
|
||||||
|
(Some(Bson::Null), None) => Ordering::Equal,
|
||||||
|
(None, Some(_)) => Ordering::Less,
|
||||||
|
(Some(_), None) => Ordering::Greater,
|
||||||
|
(Some(Bson::Null), Some(Bson::Null)) => Ordering::Equal,
|
||||||
|
(Some(Bson::Null), Some(_)) => Ordering::Less,
|
||||||
|
(Some(_), Some(Bson::Null)) => Ordering::Greater,
|
||||||
|
(Some(av), Some(bv)) => compare_typed(av, bv),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn compare_typed(a: &Bson, b: &Bson) -> std::cmp::Ordering {
|
||||||
|
use std::cmp::Ordering;
|
||||||
|
|
||||||
|
// Cross-type numeric comparison
|
||||||
|
let a_num = to_f64(a);
|
||||||
|
let b_num = to_f64(b);
|
||||||
|
if let (Some(an), Some(bn)) = (a_num, b_num) {
|
||||||
|
return an.partial_cmp(&bn).unwrap_or(Ordering::Equal);
|
||||||
|
}
|
||||||
|
|
||||||
|
match (a, b) {
|
||||||
|
(Bson::String(x), Bson::String(y)) => x.cmp(y),
|
||||||
|
(Bson::Boolean(x), Bson::Boolean(y)) => x.cmp(y),
|
||||||
|
(Bson::DateTime(x), Bson::DateTime(y)) => x.cmp(y),
|
||||||
|
(Bson::ObjectId(x), Bson::ObjectId(y)) => x.cmp(y),
|
||||||
|
_ => {
|
||||||
|
let ta = type_order(a);
|
||||||
|
let tb = type_order(b);
|
||||||
|
ta.cmp(&tb)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn to_f64(v: &Bson) -> Option<f64> {
|
||||||
|
match v {
|
||||||
|
Bson::Int32(n) => Some(*n as f64),
|
||||||
|
Bson::Int64(n) => Some(*n as f64),
|
||||||
|
Bson::Double(n) => Some(*n),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn type_order(v: &Bson) -> u8 {
|
||||||
|
match v {
|
||||||
|
Bson::Null => 0,
|
||||||
|
Bson::Int32(_) | Bson::Int64(_) | Bson::Double(_) | Bson::Decimal128(_) => 1,
|
||||||
|
Bson::String(_) => 2,
|
||||||
|
Bson::Document(_) => 3,
|
||||||
|
Bson::Array(_) => 4,
|
||||||
|
Bson::Binary(_) => 5,
|
||||||
|
Bson::ObjectId(_) => 7,
|
||||||
|
Bson::Boolean(_) => 8,
|
||||||
|
Bson::DateTime(_) => 9,
|
||||||
|
_ => 10,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_sort_ascending() {
|
||||||
|
let mut docs = vec![
|
||||||
|
bson::doc! { "x": 3 },
|
||||||
|
bson::doc! { "x": 1 },
|
||||||
|
bson::doc! { "x": 2 },
|
||||||
|
];
|
||||||
|
sort_documents(&mut docs, &bson::doc! { "x": 1 });
|
||||||
|
assert_eq!(docs[0].get_i32("x").unwrap(), 1);
|
||||||
|
assert_eq!(docs[1].get_i32("x").unwrap(), 2);
|
||||||
|
assert_eq!(docs[2].get_i32("x").unwrap(), 3);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_sort_descending() {
|
||||||
|
let mut docs = vec![
|
||||||
|
bson::doc! { "x": 1 },
|
||||||
|
bson::doc! { "x": 3 },
|
||||||
|
bson::doc! { "x": 2 },
|
||||||
|
];
|
||||||
|
sort_documents(&mut docs, &bson::doc! { "x": -1 });
|
||||||
|
assert_eq!(docs[0].get_i32("x").unwrap(), 3);
|
||||||
|
assert_eq!(docs[1].get_i32("x").unwrap(), 2);
|
||||||
|
assert_eq!(docs[2].get_i32("x").unwrap(), 1);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,661 @@
|
|||||||
|
use bson::{doc, Bson, Document};
|
||||||
|
|
||||||
|
use crate::aggregation::AggregationEngine;
|
||||||
|
use crate::error::QueryError;
|
||||||
|
use crate::field_path::{get_nested_value, remove_nested_value, set_nested_value};
|
||||||
|
use crate::matcher::QueryMatcher;
|
||||||
|
|
||||||
|
/// Update engine — applies update operators to documents.
|
||||||
|
pub struct UpdateEngine;
|
||||||
|
|
||||||
|
impl UpdateEngine {
|
||||||
|
/// Apply an update specification to a document.
|
||||||
|
/// Returns the updated document.
|
||||||
|
pub fn apply_update(
|
||||||
|
doc: &Document,
|
||||||
|
update: &Document,
|
||||||
|
_array_filters: Option<&[Document]>,
|
||||||
|
) -> Result<Document, QueryError> {
|
||||||
|
// Check if this is a replacement (no $ operators)
|
||||||
|
if !update.keys().any(|k| k.starts_with('$')) {
|
||||||
|
return Self::apply_replacement(doc, update);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut result = doc.clone();
|
||||||
|
|
||||||
|
for (op, value) in update {
|
||||||
|
let fields = match value {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
match op.as_str() {
|
||||||
|
"$set" => Self::apply_set(&mut result, fields)?,
|
||||||
|
"$unset" => Self::apply_unset(&mut result, fields)?,
|
||||||
|
"$inc" => Self::apply_inc(&mut result, fields)?,
|
||||||
|
"$mul" => Self::apply_mul(&mut result, fields)?,
|
||||||
|
"$min" => Self::apply_min(&mut result, fields)?,
|
||||||
|
"$max" => Self::apply_max(&mut result, fields)?,
|
||||||
|
"$rename" => Self::apply_rename(&mut result, fields)?,
|
||||||
|
"$currentDate" => Self::apply_current_date(&mut result, fields)?,
|
||||||
|
"$setOnInsert" => {} // handled separately during upsert
|
||||||
|
"$push" => Self::apply_push(&mut result, fields)?,
|
||||||
|
"$pop" => Self::apply_pop(&mut result, fields)?,
|
||||||
|
"$pull" => Self::apply_pull(&mut result, fields)?,
|
||||||
|
"$pullAll" => Self::apply_pull_all(&mut result, fields)?,
|
||||||
|
"$addToSet" => Self::apply_add_to_set(&mut result, fields)?,
|
||||||
|
"$bit" => Self::apply_bit(&mut result, fields)?,
|
||||||
|
other => {
|
||||||
|
return Err(QueryError::InvalidUpdate(format!(
|
||||||
|
"Unknown update operator: {}",
|
||||||
|
other
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(result)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Apply an aggregation pipeline update specification to a document.
|
||||||
|
pub fn apply_pipeline_update(
|
||||||
|
doc: &Document,
|
||||||
|
pipeline: &[Document],
|
||||||
|
) -> Result<Document, QueryError> {
|
||||||
|
if pipeline.is_empty() {
|
||||||
|
return Err(QueryError::InvalidUpdate(
|
||||||
|
"aggregation pipeline update cannot be empty".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
for stage in pipeline {
|
||||||
|
let (stage_name, _) = stage.iter().next().ok_or_else(|| {
|
||||||
|
QueryError::InvalidUpdate(
|
||||||
|
"aggregation pipeline update stages must not be empty".into(),
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
|
||||||
|
if !matches!(
|
||||||
|
stage_name.as_str(),
|
||||||
|
"$addFields" | "$set" | "$project" | "$unset" | "$replaceRoot" | "$replaceWith"
|
||||||
|
) {
|
||||||
|
return Err(QueryError::InvalidUpdate(format!(
|
||||||
|
"Unsupported aggregation pipeline update stage: {}",
|
||||||
|
stage_name
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut results = AggregationEngine::aggregate(vec![doc.clone()], pipeline, None, "")
|
||||||
|
.map_err(|e| QueryError::InvalidUpdate(e.to_string()))?;
|
||||||
|
|
||||||
|
match results.len() {
|
||||||
|
1 => Ok(results.remove(0)),
|
||||||
|
_ => Err(QueryError::InvalidUpdate(
|
||||||
|
"aggregation pipeline update must produce exactly one document".into(),
|
||||||
|
)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Apply $setOnInsert fields (used during upsert only).
|
||||||
|
pub fn apply_set_on_insert(doc: &mut Document, fields: &Document) {
|
||||||
|
for (key, value) in fields {
|
||||||
|
if key.contains('.') {
|
||||||
|
set_nested_value(doc, key, value.clone());
|
||||||
|
} else {
|
||||||
|
doc.insert(key.clone(), value.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Deep clone a BSON document.
|
||||||
|
pub fn deep_clone(doc: &Document) -> Document {
|
||||||
|
doc.clone()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_replacement(doc: &Document, replacement: &Document) -> Result<Document, QueryError> {
|
||||||
|
let mut result = replacement.clone();
|
||||||
|
// Preserve _id
|
||||||
|
if let Some(id) = doc.get("_id") {
|
||||||
|
result.insert("_id", id.clone());
|
||||||
|
}
|
||||||
|
Ok(result)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_set(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
for (key, value) in fields {
|
||||||
|
if key.contains('.') {
|
||||||
|
set_nested_value(doc, key, value.clone());
|
||||||
|
} else {
|
||||||
|
doc.insert(key.clone(), value.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_unset(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
for (key, _) in fields {
|
||||||
|
if key.contains('.') {
|
||||||
|
remove_nested_value(doc, key);
|
||||||
|
} else {
|
||||||
|
doc.remove(key);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_inc(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
for (key, inc_value) in fields {
|
||||||
|
let current = if key.contains('.') {
|
||||||
|
get_nested_value(doc, key)
|
||||||
|
} else {
|
||||||
|
doc.get(key).cloned()
|
||||||
|
};
|
||||||
|
|
||||||
|
let new_value = match (¤t, inc_value) {
|
||||||
|
(Some(Bson::Int32(a)), Bson::Int32(b)) => Bson::Int32(a + b),
|
||||||
|
(Some(Bson::Int64(a)), Bson::Int64(b)) => Bson::Int64(a + b),
|
||||||
|
(Some(Bson::Int32(a)), Bson::Int64(b)) => Bson::Int64(*a as i64 + b),
|
||||||
|
(Some(Bson::Int64(a)), Bson::Int32(b)) => Bson::Int64(a + *b as i64),
|
||||||
|
(Some(Bson::Double(a)), Bson::Double(b)) => Bson::Double(a + b),
|
||||||
|
(Some(Bson::Int32(a)), Bson::Double(b)) => Bson::Double(*a as f64 + b),
|
||||||
|
(Some(Bson::Double(a)), Bson::Int32(b)) => Bson::Double(a + *b as f64),
|
||||||
|
(Some(Bson::Int64(a)), Bson::Double(b)) => Bson::Double(*a as f64 + b),
|
||||||
|
(Some(Bson::Double(a)), Bson::Int64(b)) => Bson::Double(a + *b as f64),
|
||||||
|
(None, v) => v.clone(), // treat missing as 0
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::TypeMismatch(format!(
|
||||||
|
"Cannot apply $inc to non-numeric field: {}",
|
||||||
|
key
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
if key.contains('.') {
|
||||||
|
set_nested_value(doc, key, new_value);
|
||||||
|
} else {
|
||||||
|
doc.insert(key.clone(), new_value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_mul(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
for (key, mul_value) in fields {
|
||||||
|
let current = if key.contains('.') {
|
||||||
|
get_nested_value(doc, key)
|
||||||
|
} else {
|
||||||
|
doc.get(key).cloned()
|
||||||
|
};
|
||||||
|
|
||||||
|
let new_value = match (¤t, mul_value) {
|
||||||
|
(Some(Bson::Int32(a)), Bson::Int32(b)) => Bson::Int32(a * b),
|
||||||
|
(Some(Bson::Int64(a)), Bson::Int64(b)) => Bson::Int64(a * b),
|
||||||
|
(Some(Bson::Int32(a)), Bson::Int64(b)) => Bson::Int64(*a as i64 * b),
|
||||||
|
(Some(Bson::Int64(a)), Bson::Int32(b)) => Bson::Int64(a * *b as i64),
|
||||||
|
(Some(Bson::Double(a)), Bson::Double(b)) => Bson::Double(a * b),
|
||||||
|
(Some(Bson::Int32(a)), Bson::Double(b)) => Bson::Double(*a as f64 * b),
|
||||||
|
(Some(Bson::Double(a)), Bson::Int32(b)) => Bson::Double(a * *b as f64),
|
||||||
|
(None, _) => Bson::Int32(0), // missing field * anything = 0
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::TypeMismatch(format!(
|
||||||
|
"Cannot apply $mul to non-numeric field: {}",
|
||||||
|
key
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
if key.contains('.') {
|
||||||
|
set_nested_value(doc, key, new_value);
|
||||||
|
} else {
|
||||||
|
doc.insert(key.clone(), new_value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_min(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
for (key, min_value) in fields {
|
||||||
|
let current = if key.contains('.') {
|
||||||
|
get_nested_value(doc, key)
|
||||||
|
} else {
|
||||||
|
doc.get(key).cloned()
|
||||||
|
};
|
||||||
|
|
||||||
|
let should_update = match ¤t {
|
||||||
|
None => true,
|
||||||
|
Some(cur) => {
|
||||||
|
if let Some(ord) = QueryMatcher::bson_compare_pub(min_value, cur) {
|
||||||
|
ord == std::cmp::Ordering::Less
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
if should_update {
|
||||||
|
if key.contains('.') {
|
||||||
|
set_nested_value(doc, key, min_value.clone());
|
||||||
|
} else {
|
||||||
|
doc.insert(key.clone(), min_value.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_max(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
for (key, max_value) in fields {
|
||||||
|
let current = if key.contains('.') {
|
||||||
|
get_nested_value(doc, key)
|
||||||
|
} else {
|
||||||
|
doc.get(key).cloned()
|
||||||
|
};
|
||||||
|
|
||||||
|
let should_update = match ¤t {
|
||||||
|
None => true,
|
||||||
|
Some(cur) => {
|
||||||
|
if let Some(ord) = QueryMatcher::bson_compare_pub(max_value, cur) {
|
||||||
|
ord == std::cmp::Ordering::Greater
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
if should_update {
|
||||||
|
if key.contains('.') {
|
||||||
|
set_nested_value(doc, key, max_value.clone());
|
||||||
|
} else {
|
||||||
|
doc.insert(key.clone(), max_value.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_rename(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
for (old_name, new_name_bson) in fields {
|
||||||
|
let new_name = match new_name_bson {
|
||||||
|
Bson::String(s) => s.clone(),
|
||||||
|
_ => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Some(value) = doc.remove(old_name) {
|
||||||
|
doc.insert(new_name, value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_current_date(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
let now = bson::DateTime::now();
|
||||||
|
for (key, spec) in fields {
|
||||||
|
let value = match spec {
|
||||||
|
Bson::Boolean(true) => Bson::DateTime(now),
|
||||||
|
Bson::Document(d) => match d.get_str("$type").unwrap_or("date") {
|
||||||
|
"date" => Bson::DateTime(now),
|
||||||
|
"timestamp" => Bson::Timestamp(bson::Timestamp {
|
||||||
|
time: (now.timestamp_millis() / 1000) as u32,
|
||||||
|
increment: 0,
|
||||||
|
}),
|
||||||
|
_ => Bson::DateTime(now),
|
||||||
|
},
|
||||||
|
_ => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
if key.contains('.') {
|
||||||
|
set_nested_value(doc, key, value);
|
||||||
|
} else {
|
||||||
|
doc.insert(key.clone(), value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_push(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
for (key, value) in fields {
|
||||||
|
let arr = Self::get_or_create_array(doc, key);
|
||||||
|
|
||||||
|
match value {
|
||||||
|
Bson::Document(d) if d.contains_key("$each") => {
|
||||||
|
let each = match d.get("$each") {
|
||||||
|
Some(Bson::Array(a)) => a.clone(),
|
||||||
|
_ => {
|
||||||
|
return Err(QueryError::InvalidUpdate("$each must be an array".into()))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let position = d.get("$position").and_then(|v| match v {
|
||||||
|
Bson::Int32(n) => Some(*n as usize),
|
||||||
|
Bson::Int64(n) => Some(*n as usize),
|
||||||
|
_ => None,
|
||||||
|
});
|
||||||
|
|
||||||
|
if let Some(pos) = position {
|
||||||
|
let pos = pos.min(arr.len());
|
||||||
|
for (i, item) in each.into_iter().enumerate() {
|
||||||
|
arr.insert(pos + i, item);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
arr.extend(each);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Apply $sort if present
|
||||||
|
if let Some(sort_spec) = d.get("$sort") {
|
||||||
|
Self::sort_array(arr, sort_spec);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Apply $slice if present
|
||||||
|
if let Some(slice) = d.get("$slice") {
|
||||||
|
Self::slice_array(arr, slice);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
arr.push(value.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_pop(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
for (key, direction) in fields {
|
||||||
|
if let Some(Bson::Array(arr)) = doc.get_mut(key) {
|
||||||
|
if arr.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
match direction {
|
||||||
|
Bson::Int32(-1) | Bson::Int64(-1) => {
|
||||||
|
arr.remove(0);
|
||||||
|
}
|
||||||
|
Bson::Int32(1) | Bson::Int64(1) => {
|
||||||
|
arr.pop();
|
||||||
|
}
|
||||||
|
Bson::Double(f) if *f == 1.0 => {
|
||||||
|
arr.pop();
|
||||||
|
}
|
||||||
|
Bson::Double(f) if *f == -1.0 => {
|
||||||
|
arr.remove(0);
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
arr.pop();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_pull(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
for (key, condition) in fields {
|
||||||
|
if let Some(Bson::Array(arr)) = doc.get_mut(key) {
|
||||||
|
match condition {
|
||||||
|
Bson::Document(cond_doc) if QueryMatcher::has_operators_pub(cond_doc) => {
|
||||||
|
arr.retain(|elem| {
|
||||||
|
if let Bson::Document(elem_doc) = elem {
|
||||||
|
!QueryMatcher::matches(elem_doc, cond_doc)
|
||||||
|
} else {
|
||||||
|
// For primitive matching with operators
|
||||||
|
let wrapper = doc! { "v": elem.clone() };
|
||||||
|
let cond_wrapper = doc! { "v": condition.clone() };
|
||||||
|
!QueryMatcher::matches(&wrapper, &cond_wrapper)
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
arr.retain(|elem| elem != condition);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_pull_all(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
for (key, values) in fields {
|
||||||
|
if let (Some(Bson::Array(arr)), Bson::Array(to_remove)) = (doc.get_mut(key), values) {
|
||||||
|
arr.retain(|elem| !to_remove.contains(elem));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_add_to_set(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
for (key, value) in fields {
|
||||||
|
let arr = Self::get_or_create_array(doc, key);
|
||||||
|
|
||||||
|
match value {
|
||||||
|
Bson::Document(d) if d.contains_key("$each") => {
|
||||||
|
if let Some(Bson::Array(each)) = d.get("$each") {
|
||||||
|
for item in each {
|
||||||
|
if !arr.contains(item) {
|
||||||
|
arr.push(item.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
if !arr.contains(value) {
|
||||||
|
arr.push(value.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_bit(doc: &mut Document, fields: &Document) -> Result<(), QueryError> {
|
||||||
|
for (key, ops) in fields {
|
||||||
|
let ops_doc = match ops {
|
||||||
|
Bson::Document(d) => d,
|
||||||
|
_ => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
let current = doc.get(key).cloned().unwrap_or(Bson::Int32(0));
|
||||||
|
let mut val = match ¤t {
|
||||||
|
Bson::Int32(n) => *n as i64,
|
||||||
|
Bson::Int64(n) => *n,
|
||||||
|
_ => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
for (bit_op, operand) in ops_doc {
|
||||||
|
let operand_val = match operand {
|
||||||
|
Bson::Int32(n) => *n as i64,
|
||||||
|
Bson::Int64(n) => *n,
|
||||||
|
_ => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
match bit_op.as_str() {
|
||||||
|
"and" => val &= operand_val,
|
||||||
|
"or" => val |= operand_val,
|
||||||
|
"xor" => val ^= operand_val,
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let new_value = match ¤t {
|
||||||
|
Bson::Int32(_) => Bson::Int32(val as i32),
|
||||||
|
_ => Bson::Int64(val),
|
||||||
|
};
|
||||||
|
doc.insert(key.clone(), new_value);
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Helpers ---
|
||||||
|
|
||||||
|
fn get_or_create_array<'a>(doc: &'a mut Document, key: &str) -> &'a mut Vec<Bson> {
|
||||||
|
// Ensure an array exists at this key
|
||||||
|
let needs_init = match doc.get(key) {
|
||||||
|
Some(Bson::Array(_)) => false,
|
||||||
|
_ => true,
|
||||||
|
};
|
||||||
|
if needs_init {
|
||||||
|
doc.insert(key.to_string(), Bson::Array(Vec::new()));
|
||||||
|
}
|
||||||
|
match doc.get_mut(key).unwrap() {
|
||||||
|
Bson::Array(arr) => arr,
|
||||||
|
_ => unreachable!(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn sort_array(arr: &mut Vec<Bson>, sort_spec: &Bson) {
|
||||||
|
match sort_spec {
|
||||||
|
Bson::Int32(dir) => {
|
||||||
|
let ascending = *dir > 0;
|
||||||
|
arr.sort_by(|a, b| {
|
||||||
|
let ord = partial_cmp_bson(a, b);
|
||||||
|
if ascending {
|
||||||
|
ord
|
||||||
|
} else {
|
||||||
|
ord.reverse()
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Bson::Document(spec) => {
|
||||||
|
arr.sort_by(|a, b| {
|
||||||
|
for (field, dir) in spec {
|
||||||
|
let ascending = match dir {
|
||||||
|
Bson::Int32(n) => *n > 0,
|
||||||
|
_ => true,
|
||||||
|
};
|
||||||
|
let a_val = if let Bson::Document(d) = a {
|
||||||
|
d.get(field)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
let b_val = if let Bson::Document(d) = b {
|
||||||
|
d.get(field)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
let ord = match (a_val, b_val) {
|
||||||
|
(Some(av), Some(bv)) => partial_cmp_bson(av, bv),
|
||||||
|
(Some(_), None) => std::cmp::Ordering::Greater,
|
||||||
|
(None, Some(_)) => std::cmp::Ordering::Less,
|
||||||
|
(None, None) => std::cmp::Ordering::Equal,
|
||||||
|
};
|
||||||
|
let ord = if ascending { ord } else { ord.reverse() };
|
||||||
|
if ord != std::cmp::Ordering::Equal {
|
||||||
|
return ord;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
std::cmp::Ordering::Equal
|
||||||
|
});
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn slice_array(arr: &mut Vec<Bson>, slice: &Bson) {
|
||||||
|
let n = match slice {
|
||||||
|
Bson::Int32(n) => *n as i64,
|
||||||
|
Bson::Int64(n) => *n,
|
||||||
|
_ => return,
|
||||||
|
};
|
||||||
|
|
||||||
|
if n >= 0 {
|
||||||
|
arr.truncate(n as usize);
|
||||||
|
} else {
|
||||||
|
let keep = (-n) as usize;
|
||||||
|
if keep < arr.len() {
|
||||||
|
let start = arr.len() - keep;
|
||||||
|
*arr = arr[start..].to_vec();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn partial_cmp_bson(a: &Bson, b: &Bson) -> std::cmp::Ordering {
|
||||||
|
use std::cmp::Ordering;
|
||||||
|
match (a, b) {
|
||||||
|
(Bson::Int32(x), Bson::Int32(y)) => x.cmp(y),
|
||||||
|
(Bson::Int64(x), Bson::Int64(y)) => x.cmp(y),
|
||||||
|
(Bson::Double(x), Bson::Double(y)) => x.partial_cmp(y).unwrap_or(Ordering::Equal),
|
||||||
|
(Bson::String(x), Bson::String(y)) => x.cmp(y),
|
||||||
|
(Bson::Boolean(x), Bson::Boolean(y)) => x.cmp(y),
|
||||||
|
_ => Ordering::Equal,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_set() {
|
||||||
|
let doc = doc! { "_id": 1, "name": "Alice" };
|
||||||
|
let update = doc! { "$set": { "name": "Bob", "age": 30 } };
|
||||||
|
let result = UpdateEngine::apply_update(&doc, &update, None).unwrap();
|
||||||
|
assert_eq!(result.get_str("name").unwrap(), "Bob");
|
||||||
|
assert_eq!(result.get_i32("age").unwrap(), 30);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_inc() {
|
||||||
|
let doc = doc! { "_id": 1, "count": 5 };
|
||||||
|
let update = doc! { "$inc": { "count": 3 } };
|
||||||
|
let result = UpdateEngine::apply_update(&doc, &update, None).unwrap();
|
||||||
|
assert_eq!(result.get_i32("count").unwrap(), 8);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_unset() {
|
||||||
|
let doc = doc! { "_id": 1, "name": "Alice", "age": 30 };
|
||||||
|
let update = doc! { "$unset": { "age": "" } };
|
||||||
|
let result = UpdateEngine::apply_update(&doc, &update, None).unwrap();
|
||||||
|
assert!(result.get("age").is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_replacement() {
|
||||||
|
let doc = doc! { "_id": 1, "name": "Alice", "age": 30 };
|
||||||
|
let update = doc! { "name": "Bob" };
|
||||||
|
let result = UpdateEngine::apply_update(&doc, &update, None).unwrap();
|
||||||
|
assert_eq!(result.get_i32("_id").unwrap(), 1); // preserved
|
||||||
|
assert_eq!(result.get_str("name").unwrap(), "Bob");
|
||||||
|
assert!(result.get("age").is_none()); // removed
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_push() {
|
||||||
|
let doc = doc! { "_id": 1, "tags": ["a"] };
|
||||||
|
let update = doc! { "$push": { "tags": "b" } };
|
||||||
|
let result = UpdateEngine::apply_update(&doc, &update, None).unwrap();
|
||||||
|
let tags = result.get_array("tags").unwrap();
|
||||||
|
assert_eq!(tags.len(), 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_add_to_set() {
|
||||||
|
let doc = doc! { "_id": 1, "tags": ["a", "b"] };
|
||||||
|
let update = doc! { "$addToSet": { "tags": "a" } };
|
||||||
|
let result = UpdateEngine::apply_update(&doc, &update, None).unwrap();
|
||||||
|
let tags = result.get_array("tags").unwrap();
|
||||||
|
assert_eq!(tags.len(), 2); // no duplicate
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_pipeline_update() {
|
||||||
|
let doc = doc! { "_id": 1, "name": "Alice", "age": 30, "legacy": true };
|
||||||
|
let pipeline = vec![
|
||||||
|
doc! { "$set": { "displayName": "$name", "status": "updated" } },
|
||||||
|
doc! { "$unset": ["legacy"] },
|
||||||
|
];
|
||||||
|
|
||||||
|
let result = UpdateEngine::apply_pipeline_update(&doc, &pipeline).unwrap();
|
||||||
|
assert_eq!(result.get_str("displayName").unwrap(), "Alice");
|
||||||
|
assert_eq!(result.get_str("status").unwrap(), "updated");
|
||||||
|
assert!(result.get("legacy").is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_pipeline_update_rejects_unsupported_stage() {
|
||||||
|
let doc = doc! { "_id": 1, "name": "Alice" };
|
||||||
|
let pipeline = vec![doc! { "$match": { "name": "Alice" } }];
|
||||||
|
|
||||||
|
let result = UpdateEngine::apply_pipeline_update(&doc, &pipeline);
|
||||||
|
assert!(matches!(result, Err(QueryError::InvalidUpdate(_))));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,22 @@
|
|||||||
|
[package]
|
||||||
|
name = "rustdb-storage"
|
||||||
|
version.workspace = true
|
||||||
|
edition.workspace = true
|
||||||
|
license.workspace = true
|
||||||
|
authors.workspace = true
|
||||||
|
description = "Storage adapters (memory, file) with WAL and OpLog for RustDb"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
bson = { workspace = true }
|
||||||
|
serde = { workspace = true }
|
||||||
|
serde_json = { workspace = true }
|
||||||
|
dashmap = { workspace = true }
|
||||||
|
tokio = { workspace = true }
|
||||||
|
thiserror = { workspace = true }
|
||||||
|
tracing = { workspace = true }
|
||||||
|
crc32fast = { workspace = true }
|
||||||
|
uuid = { workspace = true }
|
||||||
|
async-trait = { workspace = true }
|
||||||
|
|
||||||
|
[dev-dependencies]
|
||||||
|
tempfile = { workspace = true }
|
||||||
@@ -0,0 +1,185 @@
|
|||||||
|
use std::collections::HashSet;
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use bson::Document;
|
||||||
|
|
||||||
|
use crate::error::StorageResult;
|
||||||
|
|
||||||
|
/// Core storage adapter trait that all backends must implement.
|
||||||
|
#[async_trait]
|
||||||
|
pub trait StorageAdapter: Send + Sync {
|
||||||
|
// ---- lifecycle ----
|
||||||
|
|
||||||
|
/// Initialize the storage backend (create directories, open files, etc.).
|
||||||
|
async fn initialize(&self) -> StorageResult<()>;
|
||||||
|
|
||||||
|
/// Gracefully shut down the storage backend.
|
||||||
|
async fn close(&self) -> StorageResult<()>;
|
||||||
|
|
||||||
|
// ---- database operations ----
|
||||||
|
|
||||||
|
/// List all database names.
|
||||||
|
async fn list_databases(&self) -> StorageResult<Vec<String>>;
|
||||||
|
|
||||||
|
/// Create a new database.
|
||||||
|
async fn create_database(&self, db: &str) -> StorageResult<()>;
|
||||||
|
|
||||||
|
/// Drop a database and all its collections.
|
||||||
|
async fn drop_database(&self, db: &str) -> StorageResult<()>;
|
||||||
|
|
||||||
|
/// Check whether a database exists.
|
||||||
|
async fn database_exists(&self, db: &str) -> StorageResult<bool>;
|
||||||
|
|
||||||
|
// ---- collection operations ----
|
||||||
|
|
||||||
|
/// List all collection names in a database.
|
||||||
|
async fn list_collections(&self, db: &str) -> StorageResult<Vec<String>>;
|
||||||
|
|
||||||
|
/// Create a new collection inside a database.
|
||||||
|
async fn create_collection(&self, db: &str, coll: &str) -> StorageResult<()>;
|
||||||
|
|
||||||
|
/// Drop a collection.
|
||||||
|
async fn drop_collection(&self, db: &str, coll: &str) -> StorageResult<()>;
|
||||||
|
|
||||||
|
/// Check whether a collection exists.
|
||||||
|
async fn collection_exists(&self, db: &str, coll: &str) -> StorageResult<bool>;
|
||||||
|
|
||||||
|
/// Rename a collection within the same database.
|
||||||
|
async fn rename_collection(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
old_name: &str,
|
||||||
|
new_name: &str,
|
||||||
|
) -> StorageResult<()>;
|
||||||
|
|
||||||
|
// ---- document write operations ----
|
||||||
|
|
||||||
|
/// Insert a single document. Returns the `_id` as hex string.
|
||||||
|
async fn insert_one(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
doc: Document,
|
||||||
|
) -> StorageResult<String>;
|
||||||
|
|
||||||
|
/// Insert many documents. Returns the `_id` hex strings.
|
||||||
|
async fn insert_many(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
docs: Vec<Document>,
|
||||||
|
) -> StorageResult<Vec<String>>;
|
||||||
|
|
||||||
|
/// Replace a document by its `_id` hex string.
|
||||||
|
async fn update_by_id(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
id: &str,
|
||||||
|
doc: Document,
|
||||||
|
) -> StorageResult<()>;
|
||||||
|
|
||||||
|
/// Delete a single document by `_id` hex string.
|
||||||
|
async fn delete_by_id(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
id: &str,
|
||||||
|
) -> StorageResult<()>;
|
||||||
|
|
||||||
|
/// Delete multiple documents by `_id` hex strings.
|
||||||
|
async fn delete_by_ids(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
ids: &[String],
|
||||||
|
) -> StorageResult<()>;
|
||||||
|
|
||||||
|
// ---- document read operations ----
|
||||||
|
|
||||||
|
/// Return all documents in a collection.
|
||||||
|
async fn find_all(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
) -> StorageResult<Vec<Document>>;
|
||||||
|
|
||||||
|
/// Return documents whose `_id` hex is in the given set.
|
||||||
|
async fn find_by_ids(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
ids: HashSet<String>,
|
||||||
|
) -> StorageResult<Vec<Document>>;
|
||||||
|
|
||||||
|
/// Return a single document by `_id` hex.
|
||||||
|
async fn find_by_id(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
id: &str,
|
||||||
|
) -> StorageResult<Option<Document>>;
|
||||||
|
|
||||||
|
/// Count documents in a collection.
|
||||||
|
async fn count(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
) -> StorageResult<u64>;
|
||||||
|
|
||||||
|
// ---- index operations ----
|
||||||
|
|
||||||
|
/// Persist an index specification for a collection.
|
||||||
|
async fn save_index(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
name: &str,
|
||||||
|
spec: Document,
|
||||||
|
) -> StorageResult<()>;
|
||||||
|
|
||||||
|
/// Return all saved index specs for a collection.
|
||||||
|
async fn get_indexes(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
) -> StorageResult<Vec<Document>>;
|
||||||
|
|
||||||
|
/// Drop a named index.
|
||||||
|
async fn drop_index(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
name: &str,
|
||||||
|
) -> StorageResult<()>;
|
||||||
|
|
||||||
|
// ---- snapshot / conflict detection ----
|
||||||
|
|
||||||
|
/// Create a logical snapshot timestamp for a collection. Returns a timestamp (ms).
|
||||||
|
async fn create_snapshot(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
) -> StorageResult<i64>;
|
||||||
|
|
||||||
|
/// Check if any of the given document ids have been modified after `snapshot_time`.
|
||||||
|
async fn has_conflicts(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
ids: &HashSet<String>,
|
||||||
|
snapshot_time: i64,
|
||||||
|
) -> StorageResult<bool>;
|
||||||
|
|
||||||
|
// ---- optional persistence (for in-memory backends) ----
|
||||||
|
|
||||||
|
/// Persist current state to durable storage. Default: no-op.
|
||||||
|
async fn persist(&self) -> StorageResult<()> {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Restore state from durable storage. Default: no-op.
|
||||||
|
async fn restore(&self) -> StorageResult<()> {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,499 @@
|
|||||||
|
//! Binary Write-Ahead Log for crash recovery.
|
||||||
|
//!
|
||||||
|
//! # Protocol
|
||||||
|
//!
|
||||||
|
//! Every mutation follows this sequence:
|
||||||
|
//! 1. Append WAL record → fsync
|
||||||
|
//! 2. Perform the actual data write
|
||||||
|
//! 3. Append WAL commit marker → fsync
|
||||||
|
//!
|
||||||
|
//! On recovery, uncommitted entries (those without a matching commit marker)
|
||||||
|
//! are replayed or verified.
|
||||||
|
//!
|
||||||
|
//! # Record format
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! ┌──────────┬──────────┬──────────┬──────────┬──────────┬──────────┬────────────┐
|
||||||
|
//! │ magic │ seq │ op │ key_len │ val_len │ crc32 │ payload │
|
||||||
|
//! │ u16 LE │ u64 LE │ u8 │ u32 LE │ u32 LE │ u32 LE │ [key][val] │
|
||||||
|
//! │ 0xWA01 │ │ │ │ │ │ │
|
||||||
|
//! └──────────┴──────────┴──────────┴──────────┴──────────┴──────────┴────────────┘
|
||||||
|
//! ```
|
||||||
|
//!
|
||||||
|
//! # Commit marker
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! ┌──────────┬──────────┬──────────┐
|
||||||
|
//! │ magic │ seq │ crc32 │
|
||||||
|
//! │ u16 LE │ u64 LE │ u32 LE │
|
||||||
|
//! │ 0xCA01 │ │ │
|
||||||
|
//! └──────────┴──────────┴──────────┘
|
||||||
|
//! ```
|
||||||
|
|
||||||
|
use std::io::{self, BufReader, Read, Write};
|
||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
|
||||||
|
use crate::error::{StorageError, StorageResult};
|
||||||
|
use crate::record::{FileHeader, FileType, FILE_HEADER_SIZE};
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Constants
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
const WAL_RECORD_MAGIC: u16 = 0xAA01;
|
||||||
|
const WAL_COMMIT_MAGIC: u16 = 0xCC01;
|
||||||
|
|
||||||
|
/// WAL record header: magic(2) + seq(8) + op(1) + key_len(4) + val_len(4) + crc(4) = 23
|
||||||
|
const WAL_RECORD_HEADER: usize = 23;
|
||||||
|
|
||||||
|
/// Commit marker size: magic(2) + seq(8) + crc(4) = 14
|
||||||
|
const WAL_COMMIT_SIZE: usize = 14;
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// WAL operation type
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
#[repr(u8)]
|
||||||
|
pub enum WalOpType {
|
||||||
|
Insert = 1,
|
||||||
|
Update = 2,
|
||||||
|
Delete = 3,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl WalOpType {
|
||||||
|
fn from_u8(v: u8) -> StorageResult<Self> {
|
||||||
|
match v {
|
||||||
|
1 => Ok(WalOpType::Insert),
|
||||||
|
2 => Ok(WalOpType::Update),
|
||||||
|
3 => Ok(WalOpType::Delete),
|
||||||
|
_ => Err(StorageError::WalError(format!("unknown WAL op: {v}"))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// WAL entry (parsed from file)
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct WalEntry {
|
||||||
|
pub seq: u64,
|
||||||
|
pub op: WalOpType,
|
||||||
|
pub key: Vec<u8>,
|
||||||
|
pub value: Vec<u8>,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Internal: what we read from the WAL file
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
enum WalItem {
|
||||||
|
Record(WalEntry),
|
||||||
|
Commit(u64), // seq that was committed
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// BinaryWal
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Binary write-ahead log backed by a single file.
|
||||||
|
pub struct BinaryWal {
|
||||||
|
path: PathBuf,
|
||||||
|
next_seq: AtomicU64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl BinaryWal {
|
||||||
|
/// Create a new WAL. Does not touch the filesystem until `initialize()`.
|
||||||
|
pub fn new(path: PathBuf) -> Self {
|
||||||
|
Self {
|
||||||
|
path,
|
||||||
|
next_seq: AtomicU64::new(1),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Initialize: create parent dirs, recover sequence counter from existing file.
|
||||||
|
pub fn initialize(&self) -> StorageResult<()> {
|
||||||
|
if let Some(parent) = self.path.parent() {
|
||||||
|
std::fs::create_dir_all(parent)?;
|
||||||
|
}
|
||||||
|
|
||||||
|
if self.path.exists() {
|
||||||
|
// Scan to find highest seq
|
||||||
|
let items = self.read_all_items()?;
|
||||||
|
let max_seq = items
|
||||||
|
.iter()
|
||||||
|
.map(|item| match item {
|
||||||
|
WalItem::Record(e) => e.seq,
|
||||||
|
WalItem::Commit(s) => *s,
|
||||||
|
})
|
||||||
|
.max()
|
||||||
|
.unwrap_or(0);
|
||||||
|
self.next_seq.store(max_seq + 1, Ordering::SeqCst);
|
||||||
|
} else {
|
||||||
|
// Create the file with a header
|
||||||
|
let mut f = std::fs::File::create(&self.path)?;
|
||||||
|
let hdr = FileHeader::new(FileType::Wal);
|
||||||
|
f.write_all(&hdr.encode())?;
|
||||||
|
f.flush()?;
|
||||||
|
f.sync_all()?;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append a WAL record. Returns the sequence number. Fsyncs.
|
||||||
|
pub fn append(
|
||||||
|
&self,
|
||||||
|
op: WalOpType,
|
||||||
|
key: &[u8],
|
||||||
|
value: &[u8],
|
||||||
|
) -> StorageResult<u64> {
|
||||||
|
let seq = self.next_seq.fetch_add(1, Ordering::SeqCst);
|
||||||
|
let key_len = key.len() as u32;
|
||||||
|
let val_len = value.len() as u32;
|
||||||
|
|
||||||
|
// Build header bytes (without CRC)
|
||||||
|
let mut hdr = Vec::with_capacity(WAL_RECORD_HEADER);
|
||||||
|
hdr.extend_from_slice(&WAL_RECORD_MAGIC.to_le_bytes());
|
||||||
|
hdr.extend_from_slice(&seq.to_le_bytes());
|
||||||
|
hdr.push(op as u8);
|
||||||
|
hdr.extend_from_slice(&key_len.to_le_bytes());
|
||||||
|
hdr.extend_from_slice(&val_len.to_le_bytes());
|
||||||
|
// CRC placeholder
|
||||||
|
hdr.extend_from_slice(&0u32.to_le_bytes());
|
||||||
|
|
||||||
|
// Compute CRC over header (without crc field) + payload
|
||||||
|
let mut hasher = crc32fast::Hasher::new();
|
||||||
|
hasher.update(&hdr[0..19]); // magic + seq + op + key_len + val_len
|
||||||
|
hasher.update(key);
|
||||||
|
hasher.update(value);
|
||||||
|
let crc = hasher.finalize();
|
||||||
|
hdr[19..23].copy_from_slice(&crc.to_le_bytes());
|
||||||
|
|
||||||
|
// Append to file
|
||||||
|
let mut f = std::fs::OpenOptions::new()
|
||||||
|
.create(true)
|
||||||
|
.append(true)
|
||||||
|
.open(&self.path)?;
|
||||||
|
f.write_all(&hdr)?;
|
||||||
|
f.write_all(key)?;
|
||||||
|
f.write_all(value)?;
|
||||||
|
f.sync_all()?;
|
||||||
|
|
||||||
|
Ok(seq)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append a commit marker for the given sequence. Fsyncs.
|
||||||
|
pub fn append_commit(&self, seq: u64) -> StorageResult<()> {
|
||||||
|
let mut buf = Vec::with_capacity(WAL_COMMIT_SIZE);
|
||||||
|
buf.extend_from_slice(&WAL_COMMIT_MAGIC.to_le_bytes());
|
||||||
|
buf.extend_from_slice(&seq.to_le_bytes());
|
||||||
|
|
||||||
|
// CRC over magic + seq
|
||||||
|
let mut hasher = crc32fast::Hasher::new();
|
||||||
|
hasher.update(&buf[0..10]);
|
||||||
|
let crc = hasher.finalize();
|
||||||
|
buf.extend_from_slice(&crc.to_le_bytes());
|
||||||
|
|
||||||
|
let mut f = std::fs::OpenOptions::new()
|
||||||
|
.create(true)
|
||||||
|
.append(true)
|
||||||
|
.open(&self.path)?;
|
||||||
|
f.write_all(&buf)?;
|
||||||
|
f.sync_all()?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Recover: return all WAL entries that were NOT committed.
|
||||||
|
pub fn recover(&self) -> StorageResult<Vec<WalEntry>> {
|
||||||
|
let items = self.read_all_items()?;
|
||||||
|
|
||||||
|
// Collect committed seq numbers
|
||||||
|
let committed: std::collections::HashSet<u64> = items
|
||||||
|
.iter()
|
||||||
|
.filter_map(|item| {
|
||||||
|
if let WalItem::Commit(s) = item {
|
||||||
|
Some(*s)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// Return records without a commit marker
|
||||||
|
let uncommitted: Vec<WalEntry> = items
|
||||||
|
.into_iter()
|
||||||
|
.filter_map(|item| {
|
||||||
|
if let WalItem::Record(entry) = item {
|
||||||
|
if !committed.contains(&entry.seq) {
|
||||||
|
return Some(entry);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
Ok(uncommitted)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Truncate the WAL: rewrite with just the file header (clears all entries).
|
||||||
|
pub fn truncate(&self) -> StorageResult<()> {
|
||||||
|
let mut f = std::fs::File::create(&self.path)?;
|
||||||
|
let hdr = FileHeader::new(FileType::Wal);
|
||||||
|
f.write_all(&hdr.encode())?;
|
||||||
|
f.flush()?;
|
||||||
|
f.sync_all()?;
|
||||||
|
// Don't reset next_seq — it should keep incrementing
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Path to the WAL file.
|
||||||
|
pub fn path(&self) -> &Path {
|
||||||
|
&self.path
|
||||||
|
}
|
||||||
|
|
||||||
|
// -----------------------------------------------------------------------
|
||||||
|
// Internal: read all items from the WAL file
|
||||||
|
// -----------------------------------------------------------------------
|
||||||
|
|
||||||
|
fn read_all_items(&self) -> StorageResult<Vec<WalItem>> {
|
||||||
|
if !self.path.exists() {
|
||||||
|
return Ok(vec![]);
|
||||||
|
}
|
||||||
|
|
||||||
|
let file = std::fs::File::open(&self.path)?;
|
||||||
|
let mut reader = BufReader::new(file);
|
||||||
|
|
||||||
|
// Skip file header (if present)
|
||||||
|
let file_len = std::fs::metadata(&self.path)?.len();
|
||||||
|
if file_len >= FILE_HEADER_SIZE as u64 {
|
||||||
|
let mut hdr_buf = [0u8; FILE_HEADER_SIZE];
|
||||||
|
reader.read_exact(&mut hdr_buf)?;
|
||||||
|
// Validate but don't fail hard — allow reading even slightly off headers
|
||||||
|
let _ = FileHeader::decode(&hdr_buf);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut items = Vec::new();
|
||||||
|
|
||||||
|
loop {
|
||||||
|
// Peek at the magic to determine if this is a record or commit marker
|
||||||
|
let mut magic_buf = [0u8; 2];
|
||||||
|
match reader.read_exact(&mut magic_buf) {
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break,
|
||||||
|
Err(e) => return Err(e.into()),
|
||||||
|
}
|
||||||
|
let magic = u16::from_le_bytes(magic_buf);
|
||||||
|
|
||||||
|
match magic {
|
||||||
|
WAL_RECORD_MAGIC => {
|
||||||
|
// Read rest of header: seq(8) + op(1) + key_len(4) + val_len(4) + crc(4) = 21
|
||||||
|
let mut rest = [0u8; 21];
|
||||||
|
match reader.read_exact(&mut rest) {
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break,
|
||||||
|
Err(e) => return Err(e.into()),
|
||||||
|
}
|
||||||
|
|
||||||
|
let seq = u64::from_le_bytes(rest[0..8].try_into().unwrap());
|
||||||
|
let op = WalOpType::from_u8(rest[8])?;
|
||||||
|
let key_len = u32::from_le_bytes(rest[9..13].try_into().unwrap()) as usize;
|
||||||
|
let val_len = u32::from_le_bytes(rest[13..17].try_into().unwrap()) as usize;
|
||||||
|
let stored_crc = u32::from_le_bytes(rest[17..21].try_into().unwrap());
|
||||||
|
|
||||||
|
let mut payload = vec![0u8; key_len + val_len];
|
||||||
|
match reader.read_exact(&mut payload) {
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break,
|
||||||
|
Err(e) => return Err(e.into()),
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify CRC
|
||||||
|
let mut hasher = crc32fast::Hasher::new();
|
||||||
|
hasher.update(&magic_buf);
|
||||||
|
hasher.update(&rest[0..17]); // seq + op + key_len + val_len
|
||||||
|
hasher.update(&payload);
|
||||||
|
let computed = hasher.finalize();
|
||||||
|
|
||||||
|
if computed != stored_crc {
|
||||||
|
// Corrupt WAL entry — skip it (best-effort recovery)
|
||||||
|
tracing::warn!(
|
||||||
|
seq,
|
||||||
|
"skipping corrupt WAL record: CRC mismatch (expected 0x{stored_crc:08X}, got 0x{computed:08X})"
|
||||||
|
);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let key = payload[..key_len].to_vec();
|
||||||
|
let value = payload[key_len..].to_vec();
|
||||||
|
items.push(WalItem::Record(WalEntry {
|
||||||
|
seq,
|
||||||
|
op,
|
||||||
|
key,
|
||||||
|
value,
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
WAL_COMMIT_MAGIC => {
|
||||||
|
// Read rest: seq(8) + crc(4) = 12
|
||||||
|
let mut rest = [0u8; 12];
|
||||||
|
match reader.read_exact(&mut rest) {
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break,
|
||||||
|
Err(e) => return Err(e.into()),
|
||||||
|
}
|
||||||
|
|
||||||
|
let seq = u64::from_le_bytes(rest[0..8].try_into().unwrap());
|
||||||
|
let stored_crc = u32::from_le_bytes(rest[8..12].try_into().unwrap());
|
||||||
|
|
||||||
|
let mut hasher = crc32fast::Hasher::new();
|
||||||
|
hasher.update(&magic_buf);
|
||||||
|
hasher.update(&rest[0..8]);
|
||||||
|
let computed = hasher.finalize();
|
||||||
|
|
||||||
|
if computed != stored_crc {
|
||||||
|
tracing::warn!(
|
||||||
|
seq,
|
||||||
|
"skipping corrupt WAL commit marker: CRC mismatch"
|
||||||
|
);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
items.push(WalItem::Commit(seq));
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
// Unknown magic — file is corrupt past this point
|
||||||
|
tracing::warn!("unknown WAL magic 0x{magic:04X}, stopping scan");
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(items)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Tests
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn make_wal(dir: &tempfile::TempDir) -> BinaryWal {
|
||||||
|
let path = dir.path().join("test.wal");
|
||||||
|
let wal = BinaryWal::new(path);
|
||||||
|
wal.initialize().unwrap();
|
||||||
|
wal
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn append_and_commit() {
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let wal = make_wal(&dir);
|
||||||
|
|
||||||
|
let seq = wal
|
||||||
|
.append(WalOpType::Insert, b"key1", b"value1")
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(seq, 1);
|
||||||
|
|
||||||
|
wal.append_commit(seq).unwrap();
|
||||||
|
|
||||||
|
// All committed — recover should return empty
|
||||||
|
let uncommitted = wal.recover().unwrap();
|
||||||
|
assert!(uncommitted.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn uncommitted_entries_recovered() {
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let wal = make_wal(&dir);
|
||||||
|
|
||||||
|
let s1 = wal
|
||||||
|
.append(WalOpType::Insert, b"k1", b"v1")
|
||||||
|
.unwrap();
|
||||||
|
wal.append_commit(s1).unwrap();
|
||||||
|
|
||||||
|
// s2 is NOT committed
|
||||||
|
let s2 = wal
|
||||||
|
.append(WalOpType::Update, b"k2", b"v2")
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let uncommitted = wal.recover().unwrap();
|
||||||
|
assert_eq!(uncommitted.len(), 1);
|
||||||
|
assert_eq!(uncommitted[0].seq, s2);
|
||||||
|
assert_eq!(uncommitted[0].op, WalOpType::Update);
|
||||||
|
assert_eq!(uncommitted[0].key, b"k2");
|
||||||
|
assert_eq!(uncommitted[0].value, b"v2");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn truncate_clears_wal() {
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let wal = make_wal(&dir);
|
||||||
|
|
||||||
|
wal.append(WalOpType::Insert, b"k", b"v").unwrap();
|
||||||
|
wal.truncate().unwrap();
|
||||||
|
|
||||||
|
let uncommitted = wal.recover().unwrap();
|
||||||
|
assert!(uncommitted.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn multiple_operations() {
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let wal = make_wal(&dir);
|
||||||
|
|
||||||
|
let s1 = wal.append(WalOpType::Insert, b"a", b"1").unwrap();
|
||||||
|
let s2 = wal.append(WalOpType::Update, b"b", b"2").unwrap();
|
||||||
|
let s3 = wal.append(WalOpType::Delete, b"c", b"").unwrap();
|
||||||
|
|
||||||
|
// Commit only s1 and s3
|
||||||
|
wal.append_commit(s1).unwrap();
|
||||||
|
wal.append_commit(s3).unwrap();
|
||||||
|
|
||||||
|
let uncommitted = wal.recover().unwrap();
|
||||||
|
assert_eq!(uncommitted.len(), 1);
|
||||||
|
assert_eq!(uncommitted[0].seq, s2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn sequence_numbers_persist_across_reinit() {
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let path = dir.path().join("persist.wal");
|
||||||
|
|
||||||
|
{
|
||||||
|
let wal = BinaryWal::new(path.clone());
|
||||||
|
wal.initialize().unwrap();
|
||||||
|
let s1 = wal.append(WalOpType::Insert, b"k", b"v").unwrap();
|
||||||
|
assert_eq!(s1, 1);
|
||||||
|
wal.append_commit(s1).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Re-open — seq should continue from 2+ (since max committed was 1)
|
||||||
|
{
|
||||||
|
let wal = BinaryWal::new(path);
|
||||||
|
wal.initialize().unwrap();
|
||||||
|
let s2 = wal.append(WalOpType::Insert, b"k2", b"v2").unwrap();
|
||||||
|
assert!(s2 >= 2, "seq should continue: got {s2}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn delete_has_empty_value() {
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let wal = make_wal(&dir);
|
||||||
|
|
||||||
|
let seq = wal.append(WalOpType::Delete, b"key", b"").unwrap();
|
||||||
|
|
||||||
|
let uncommitted = wal.recover().unwrap();
|
||||||
|
assert_eq!(uncommitted.len(), 1);
|
||||||
|
assert_eq!(uncommitted[0].seq, seq);
|
||||||
|
assert_eq!(uncommitted[0].op, WalOpType::Delete);
|
||||||
|
assert!(uncommitted[0].value.is_empty());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,270 @@
|
|||||||
|
//! Compaction for the Bitcask-style storage engine.
|
||||||
|
//!
|
||||||
|
//! Over time, the data file accumulates dead records (superseded by updates,
|
||||||
|
//! tombstones from deletes). Compaction rewrites the data file with only live
|
||||||
|
//! records, reclaiming disk space.
|
||||||
|
//!
|
||||||
|
//! The process is:
|
||||||
|
//! 1. Create a new `data.rdb.compact` file with a fresh file header.
|
||||||
|
//! 2. Iterate all live entries from the KeyDir.
|
||||||
|
//! 3. Read each live document from the old data file, write to the new file.
|
||||||
|
//! 4. Atomically rename `data.rdb.compact` → `data.rdb`.
|
||||||
|
//! 5. Update KeyDir entries with new offsets.
|
||||||
|
//! 6. Reset dead_bytes counter.
|
||||||
|
|
||||||
|
use std::io::{Seek, SeekFrom, Write};
|
||||||
|
use std::path::Path;
|
||||||
|
use std::sync::atomic::Ordering;
|
||||||
|
|
||||||
|
use tracing::info;
|
||||||
|
|
||||||
|
use crate::error::StorageResult;
|
||||||
|
use crate::keydir::{KeyDir, KeyDirEntry};
|
||||||
|
use crate::record::{DataRecord, FileHeader, FileType, FILE_HEADER_SIZE};
|
||||||
|
|
||||||
|
/// Result of a compaction operation.
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct CompactionResult {
|
||||||
|
/// Number of live records written.
|
||||||
|
pub records_written: u64,
|
||||||
|
/// Bytes reclaimed (old file size - new file size).
|
||||||
|
pub bytes_reclaimed: u64,
|
||||||
|
/// New data file size.
|
||||||
|
pub new_file_size: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Compact a collection's data file.
|
||||||
|
///
|
||||||
|
/// This function:
|
||||||
|
/// - Reads all live documents (entries present in the KeyDir) from the old data file
|
||||||
|
/// - Writes them sequentially to a new file
|
||||||
|
/// - Atomically renames the new file over the old one
|
||||||
|
/// - Updates all KeyDir entries with their new offsets
|
||||||
|
///
|
||||||
|
/// The caller must hold the collection's write lock during this operation.
|
||||||
|
pub fn compact_data_file(
|
||||||
|
data_path: &Path,
|
||||||
|
keydir: &KeyDir,
|
||||||
|
dead_bytes: &std::sync::atomic::AtomicU64,
|
||||||
|
data_file_size: &std::sync::atomic::AtomicU64,
|
||||||
|
) -> StorageResult<CompactionResult> {
|
||||||
|
let compact_path = data_path.with_extension("rdb.compact");
|
||||||
|
|
||||||
|
let old_file_size = std::fs::metadata(data_path)
|
||||||
|
.map(|m| m.len())
|
||||||
|
.unwrap_or(0);
|
||||||
|
|
||||||
|
// Collect all live entries with their keys
|
||||||
|
let mut live_entries: Vec<(String, KeyDirEntry)> = Vec::with_capacity(keydir.len() as usize);
|
||||||
|
keydir.for_each(|key, entry| {
|
||||||
|
live_entries.push((key.to_string(), *entry));
|
||||||
|
});
|
||||||
|
|
||||||
|
// Sort by offset for sequential reads (cache-friendly)
|
||||||
|
live_entries.sort_by_key(|(_, e)| e.offset);
|
||||||
|
|
||||||
|
// Create compact file with header
|
||||||
|
let mut compact_file = std::fs::File::create(&compact_path)?;
|
||||||
|
let hdr = FileHeader::new(FileType::Data);
|
||||||
|
compact_file.write_all(&hdr.encode())?;
|
||||||
|
|
||||||
|
let mut current_offset = FILE_HEADER_SIZE as u64;
|
||||||
|
let mut new_entries: Vec<(String, KeyDirEntry)> = Vec::with_capacity(live_entries.len());
|
||||||
|
let mut old_data_file = std::fs::File::open(data_path)?;
|
||||||
|
|
||||||
|
for (key, entry) in &live_entries {
|
||||||
|
// Read the record from the old file
|
||||||
|
old_data_file.seek(SeekFrom::Start(entry.offset))?;
|
||||||
|
let (record, _disk_size) = DataRecord::decode_from(&mut old_data_file)?
|
||||||
|
.ok_or_else(|| {
|
||||||
|
crate::error::StorageError::CorruptRecord(format!(
|
||||||
|
"compaction: unexpected EOF reading doc '{key}' at offset {}",
|
||||||
|
entry.offset
|
||||||
|
))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
// Write to compact file
|
||||||
|
let encoded = record.encode();
|
||||||
|
let new_disk_size = encoded.len() as u32;
|
||||||
|
compact_file.write_all(&encoded)?;
|
||||||
|
|
||||||
|
new_entries.push((
|
||||||
|
key.clone(),
|
||||||
|
KeyDirEntry {
|
||||||
|
offset: current_offset,
|
||||||
|
record_len: new_disk_size,
|
||||||
|
value_len: entry.value_len,
|
||||||
|
timestamp: entry.timestamp,
|
||||||
|
},
|
||||||
|
));
|
||||||
|
|
||||||
|
current_offset += new_disk_size as u64;
|
||||||
|
}
|
||||||
|
|
||||||
|
compact_file.sync_all()?;
|
||||||
|
drop(compact_file);
|
||||||
|
drop(old_data_file);
|
||||||
|
|
||||||
|
// Atomic rename
|
||||||
|
std::fs::rename(&compact_path, data_path)?;
|
||||||
|
|
||||||
|
// Update KeyDir with new offsets
|
||||||
|
for (key, new_entry) in new_entries {
|
||||||
|
keydir.insert(key, new_entry);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reset counters
|
||||||
|
dead_bytes.store(0, Ordering::Relaxed);
|
||||||
|
data_file_size.store(current_offset, Ordering::Relaxed);
|
||||||
|
|
||||||
|
let bytes_reclaimed = old_file_size.saturating_sub(current_offset);
|
||||||
|
|
||||||
|
info!(
|
||||||
|
records = live_entries.len(),
|
||||||
|
old_size = old_file_size,
|
||||||
|
new_size = current_offset,
|
||||||
|
reclaimed = bytes_reclaimed,
|
||||||
|
"compaction complete"
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(CompactionResult {
|
||||||
|
records_written: live_entries.len() as u64,
|
||||||
|
bytes_reclaimed,
|
||||||
|
new_file_size: current_offset,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if compaction is warranted for a collection.
|
||||||
|
/// Returns true if dead bytes exceed 50% of live data.
|
||||||
|
pub fn should_compact(dead_bytes: u64, data_file_size: u64) -> bool {
|
||||||
|
if data_file_size <= FILE_HEADER_SIZE as u64 {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
let useful_bytes = data_file_size - FILE_HEADER_SIZE as u64;
|
||||||
|
// Trigger when dead > 50% of total useful data
|
||||||
|
dead_bytes > useful_bytes / 2
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Tests
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::keydir::KeyDir;
|
||||||
|
use crate::record::{now_ms, DataRecord, FileHeader, FileType};
|
||||||
|
use std::io::Write;
|
||||||
|
use std::sync::atomic::AtomicU64;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn compact_removes_dead_records() {
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let data_path = dir.path().join("data.rdb");
|
||||||
|
|
||||||
|
// Write a data file: insert A, update A (new version), insert B
|
||||||
|
let mut f = std::fs::File::create(&data_path).unwrap();
|
||||||
|
let hdr = FileHeader::new(FileType::Data);
|
||||||
|
f.write_all(&hdr.encode()).unwrap();
|
||||||
|
|
||||||
|
let ts = now_ms();
|
||||||
|
|
||||||
|
// Record 1: A v1 (will be superseded)
|
||||||
|
let r1 = DataRecord {
|
||||||
|
timestamp: ts,
|
||||||
|
key: b"aaa".to_vec(),
|
||||||
|
value: b"old_value".to_vec(),
|
||||||
|
};
|
||||||
|
let r1_enc = r1.encode();
|
||||||
|
let r1_offset = FILE_HEADER_SIZE as u64;
|
||||||
|
let r1_size = r1_enc.len();
|
||||||
|
f.write_all(&r1_enc).unwrap();
|
||||||
|
|
||||||
|
// Record 2: A v2 (current)
|
||||||
|
let r2 = DataRecord {
|
||||||
|
timestamp: ts + 1,
|
||||||
|
key: b"aaa".to_vec(),
|
||||||
|
value: b"new_value".to_vec(),
|
||||||
|
};
|
||||||
|
let r2_enc = r2.encode();
|
||||||
|
let r2_offset = r1_offset + r1_size as u64;
|
||||||
|
let r2_size = r2_enc.len();
|
||||||
|
f.write_all(&r2_enc).unwrap();
|
||||||
|
|
||||||
|
// Record 3: B (live)
|
||||||
|
let r3 = DataRecord {
|
||||||
|
timestamp: ts + 2,
|
||||||
|
key: b"bbb".to_vec(),
|
||||||
|
value: b"bbb_value".to_vec(),
|
||||||
|
};
|
||||||
|
let r3_enc = r3.encode();
|
||||||
|
let r3_offset = r2_offset + r2_size as u64;
|
||||||
|
f.write_all(&r3_enc).unwrap();
|
||||||
|
f.sync_all().unwrap();
|
||||||
|
drop(f);
|
||||||
|
|
||||||
|
let total_size = std::fs::metadata(&data_path).unwrap().len();
|
||||||
|
|
||||||
|
// Build KeyDir — only points to latest versions
|
||||||
|
let keydir = KeyDir::new();
|
||||||
|
keydir.insert(
|
||||||
|
"aaa".into(),
|
||||||
|
KeyDirEntry {
|
||||||
|
offset: r2_offset,
|
||||||
|
record_len: r2_size as u32,
|
||||||
|
value_len: r2.value.len() as u32,
|
||||||
|
timestamp: ts + 1,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
keydir.insert(
|
||||||
|
"bbb".into(),
|
||||||
|
KeyDirEntry {
|
||||||
|
offset: r3_offset,
|
||||||
|
record_len: r3.encode().len() as u32,
|
||||||
|
value_len: r3.value.len() as u32,
|
||||||
|
timestamp: ts + 2,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
let dead_bytes_counter = AtomicU64::new(r1_size as u64);
|
||||||
|
let data_file_size_counter = AtomicU64::new(total_size);
|
||||||
|
|
||||||
|
let result = compact_data_file(
|
||||||
|
&data_path,
|
||||||
|
&keydir,
|
||||||
|
&dead_bytes_counter,
|
||||||
|
&data_file_size_counter,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(result.records_written, 2);
|
||||||
|
assert!(result.bytes_reclaimed > 0);
|
||||||
|
assert!(result.new_file_size < total_size);
|
||||||
|
|
||||||
|
// Verify dead_bytes was reset
|
||||||
|
assert_eq!(dead_bytes_counter.load(Ordering::Relaxed), 0);
|
||||||
|
|
||||||
|
// Verify KeyDir was updated with new offsets
|
||||||
|
let a_entry = keydir.get("aaa").unwrap();
|
||||||
|
assert_eq!(a_entry.offset, FILE_HEADER_SIZE as u64); // first record after header
|
||||||
|
assert_eq!(a_entry.value_len, b"new_value".len() as u32);
|
||||||
|
|
||||||
|
let b_entry = keydir.get("bbb").unwrap();
|
||||||
|
assert!(b_entry.offset > a_entry.offset);
|
||||||
|
|
||||||
|
// Verify the compacted file can be used to rebuild KeyDir
|
||||||
|
let (rebuilt, dead, _stats) = KeyDir::build_from_data_file(&data_path).unwrap();
|
||||||
|
assert_eq!(rebuilt.len(), 2);
|
||||||
|
assert_eq!(dead, 0); // no dead records in compacted file
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn should_compact_thresholds() {
|
||||||
|
// Under threshold
|
||||||
|
assert!(!should_compact(10, 100 + FILE_HEADER_SIZE as u64));
|
||||||
|
// Over threshold (dead > 50% of useful)
|
||||||
|
assert!(should_compact(60, 100 + FILE_HEADER_SIZE as u64));
|
||||||
|
// Empty file
|
||||||
|
assert!(!should_compact(0, FILE_HEADER_SIZE as u64));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,49 @@
|
|||||||
|
use thiserror::Error;
|
||||||
|
|
||||||
|
/// Errors that can occur in storage operations.
|
||||||
|
#[derive(Debug, Error)]
|
||||||
|
pub enum StorageError {
|
||||||
|
#[error("not found: {0}")]
|
||||||
|
NotFound(String),
|
||||||
|
|
||||||
|
#[error("already exists: {0}")]
|
||||||
|
AlreadyExists(String),
|
||||||
|
|
||||||
|
#[error("I/O error: {0}")]
|
||||||
|
IoError(#[from] std::io::Error),
|
||||||
|
|
||||||
|
#[error("serialization error: {0}")]
|
||||||
|
SerializationError(String),
|
||||||
|
|
||||||
|
#[error("conflict detected: {0}")]
|
||||||
|
ConflictError(String),
|
||||||
|
|
||||||
|
#[error("corrupt record: {0}")]
|
||||||
|
CorruptRecord(String),
|
||||||
|
|
||||||
|
#[error("checksum mismatch: expected 0x{expected:08X}, got 0x{actual:08X}")]
|
||||||
|
ChecksumMismatch { expected: u32, actual: u32 },
|
||||||
|
|
||||||
|
#[error("WAL error: {0}")]
|
||||||
|
WalError(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
impl From<serde_json::Error> for StorageError {
|
||||||
|
fn from(e: serde_json::Error) -> Self {
|
||||||
|
StorageError::SerializationError(e.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl From<bson::de::Error> for StorageError {
|
||||||
|
fn from(e: bson::de::Error) -> Self {
|
||||||
|
StorageError::SerializationError(e.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl From<bson::ser::Error> for StorageError {
|
||||||
|
fn from(e: bson::ser::Error) -> Self {
|
||||||
|
StorageError::SerializationError(e.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub type StorageResult<T> = Result<T, StorageError>;
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,610 @@
|
|||||||
|
//! KeyDir — in-memory document location index for the Bitcask storage engine.
|
||||||
|
//!
|
||||||
|
//! Maps document `_id` (hex string) to its location in the append-only data file.
|
||||||
|
//! Backed by `DashMap` for lock-free concurrent reads and fine-grained write locking.
|
||||||
|
//!
|
||||||
|
//! The KeyDir can be rebuilt from a data file scan, or loaded quickly from a
|
||||||
|
//! persisted hint file for fast restart.
|
||||||
|
|
||||||
|
use std::io::{self, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
|
||||||
|
use std::path::Path;
|
||||||
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
|
||||||
|
use dashmap::DashMap;
|
||||||
|
|
||||||
|
use crate::error::{StorageError, StorageResult};
|
||||||
|
use crate::record::{
|
||||||
|
DataRecord, FileHeader, FileType, FILE_HEADER_SIZE, FORMAT_VERSION,
|
||||||
|
};
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// KeyDirEntry
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Location of a single document in the data file.
|
||||||
|
#[derive(Debug, Clone, Copy)]
|
||||||
|
pub struct KeyDirEntry {
|
||||||
|
/// Byte offset of the record in `data.rdb`.
|
||||||
|
pub offset: u64,
|
||||||
|
/// Total record size on disk (header + payload).
|
||||||
|
pub record_len: u32,
|
||||||
|
/// BSON value length. 0 means tombstone (used during compaction accounting).
|
||||||
|
pub value_len: u32,
|
||||||
|
/// Timestamp (epoch ms) from the record. Used for conflict detection.
|
||||||
|
pub timestamp: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// BuildStats — statistics from building KeyDir from a data file scan
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Statistics collected while building a KeyDir from a data file scan.
|
||||||
|
#[derive(Debug, Clone, Default)]
|
||||||
|
pub struct BuildStats {
|
||||||
|
/// Total records scanned (live + tombstones + superseded).
|
||||||
|
pub total_records_scanned: u64,
|
||||||
|
/// Number of live documents in the final KeyDir.
|
||||||
|
pub live_documents: u64,
|
||||||
|
/// Number of tombstone records encountered.
|
||||||
|
pub tombstones: u64,
|
||||||
|
/// Number of records superseded by a later write for the same key.
|
||||||
|
pub superseded_records: u64,
|
||||||
|
/// Byte offset immediately after the last valid record.
|
||||||
|
pub valid_data_end: u64,
|
||||||
|
/// Number of invalid tail bytes after the last valid record.
|
||||||
|
pub invalid_tail_bytes: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// KeyDir
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// In-memory index mapping document ID → data file location.
|
||||||
|
pub struct KeyDir {
|
||||||
|
map: DashMap<String, KeyDirEntry>,
|
||||||
|
/// Running count of live documents.
|
||||||
|
doc_count: AtomicU64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl KeyDir {
|
||||||
|
/// Create an empty KeyDir.
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
map: DashMap::new(),
|
||||||
|
doc_count: AtomicU64::new(0),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Insert or update an entry. Returns the previous entry if one existed.
|
||||||
|
pub fn insert(&self, key: String, entry: KeyDirEntry) -> Option<KeyDirEntry> {
|
||||||
|
let prev = self.map.insert(key, entry);
|
||||||
|
if prev.is_none() {
|
||||||
|
self.doc_count.fetch_add(1, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
prev
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Look up an entry by key.
|
||||||
|
pub fn get(&self, key: &str) -> Option<KeyDirEntry> {
|
||||||
|
self.map.get(key).map(|r| *r.value())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Remove an entry. Returns the removed entry if it existed.
|
||||||
|
pub fn remove(&self, key: &str) -> Option<KeyDirEntry> {
|
||||||
|
let removed = self.map.remove(key).map(|(_, v)| v);
|
||||||
|
if removed.is_some() {
|
||||||
|
self.doc_count.fetch_sub(1, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
removed
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Number of live documents.
|
||||||
|
pub fn len(&self) -> u64 {
|
||||||
|
self.doc_count.load(Ordering::Relaxed)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Whether the index is empty.
|
||||||
|
pub fn is_empty(&self) -> bool {
|
||||||
|
self.len() == 0
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if a key exists.
|
||||||
|
pub fn contains(&self, key: &str) -> bool {
|
||||||
|
self.map.contains_key(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Iterate over all entries. The closure receives (key, entry).
|
||||||
|
pub fn for_each(&self, mut f: impl FnMut(&str, &KeyDirEntry)) {
|
||||||
|
for entry in self.map.iter() {
|
||||||
|
f(entry.key(), entry.value());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Collect all keys.
|
||||||
|
pub fn keys(&self) -> Vec<String> {
|
||||||
|
self.map.iter().map(|e| e.key().clone()).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Clear all entries.
|
||||||
|
pub fn clear(&self) {
|
||||||
|
self.map.clear();
|
||||||
|
self.doc_count.store(0, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
|
||||||
|
// -----------------------------------------------------------------------
|
||||||
|
// Build from data file
|
||||||
|
// -----------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Rebuild the KeyDir by scanning an entire data file.
|
||||||
|
/// The file must start with a valid `FileHeader`.
|
||||||
|
/// Returns `(keydir, dead_bytes, stats)` where `dead_bytes` is the total size of
|
||||||
|
/// stale records (superseded by later writes or tombstoned).
|
||||||
|
pub fn build_from_data_file(path: &Path) -> StorageResult<(Self, u64, BuildStats)> {
|
||||||
|
let file = std::fs::File::open(path)?;
|
||||||
|
let file_len = file.metadata()?.len();
|
||||||
|
let mut reader = BufReader::new(file);
|
||||||
|
|
||||||
|
// Read and validate file header
|
||||||
|
let mut hdr_buf = [0u8; FILE_HEADER_SIZE];
|
||||||
|
reader.read_exact(&mut hdr_buf)?;
|
||||||
|
let hdr = FileHeader::decode(&hdr_buf)?;
|
||||||
|
if hdr.file_type != FileType::Data {
|
||||||
|
return Err(StorageError::CorruptRecord(format!(
|
||||||
|
"expected data file (type 1), got type {:?}",
|
||||||
|
hdr.file_type
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
let keydir = KeyDir::new();
|
||||||
|
let mut dead_bytes: u64 = 0;
|
||||||
|
let mut stats = BuildStats {
|
||||||
|
valid_data_end: FILE_HEADER_SIZE as u64,
|
||||||
|
..BuildStats::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
loop {
|
||||||
|
let record_offset = stats.valid_data_end;
|
||||||
|
let (record, disk_size) = match DataRecord::decode_from(&mut reader) {
|
||||||
|
Ok(Some((record, disk_size))) => (record, disk_size),
|
||||||
|
Ok(None) => {
|
||||||
|
if file_len > record_offset {
|
||||||
|
stats.invalid_tail_bytes = file_len - record_offset;
|
||||||
|
}
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
Err(StorageError::IoError(e)) if e.kind() == io::ErrorKind::UnexpectedEof => {
|
||||||
|
stats.invalid_tail_bytes = file_len.saturating_sub(record_offset);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
Err(StorageError::ChecksumMismatch { expected, actual }) => {
|
||||||
|
tracing::warn!(
|
||||||
|
path = %path.display(),
|
||||||
|
offset = record_offset,
|
||||||
|
"stopping data file scan at checksum mismatch: expected 0x{expected:08X}, got 0x{actual:08X}"
|
||||||
|
);
|
||||||
|
stats.invalid_tail_bytes = file_len.saturating_sub(record_offset);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
Err(StorageError::CorruptRecord(message)) => {
|
||||||
|
tracing::warn!(
|
||||||
|
path = %path.display(),
|
||||||
|
offset = record_offset,
|
||||||
|
"stopping data file scan at corrupt record: {message}"
|
||||||
|
);
|
||||||
|
stats.invalid_tail_bytes = file_len.saturating_sub(record_offset);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
Err(e) => return Err(e),
|
||||||
|
};
|
||||||
|
|
||||||
|
stats.valid_data_end += disk_size as u64;
|
||||||
|
let is_tombstone = record.is_tombstone();
|
||||||
|
let disk_size = disk_size as u32;
|
||||||
|
let value_len = record.value.len() as u32;
|
||||||
|
let timestamp = record.timestamp;
|
||||||
|
let key = String::from_utf8(record.key)
|
||||||
|
.map_err(|e| StorageError::CorruptRecord(format!("invalid UTF-8 key: {e}")))?;
|
||||||
|
|
||||||
|
stats.total_records_scanned += 1;
|
||||||
|
|
||||||
|
if is_tombstone {
|
||||||
|
stats.tombstones += 1;
|
||||||
|
// Remove from index; the tombstone itself is dead weight
|
||||||
|
if let Some(prev) = keydir.remove(&key) {
|
||||||
|
dead_bytes += prev.record_len as u64;
|
||||||
|
}
|
||||||
|
dead_bytes += disk_size as u64;
|
||||||
|
} else {
|
||||||
|
let entry = KeyDirEntry {
|
||||||
|
offset: record_offset,
|
||||||
|
record_len: disk_size,
|
||||||
|
value_len,
|
||||||
|
timestamp,
|
||||||
|
};
|
||||||
|
if let Some(prev) = keydir.insert(key, entry) {
|
||||||
|
// Previous version of same key is now dead
|
||||||
|
dead_bytes += prev.record_len as u64;
|
||||||
|
stats.superseded_records += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
stats.live_documents = keydir.len();
|
||||||
|
Ok((keydir, dead_bytes, stats))
|
||||||
|
}
|
||||||
|
|
||||||
|
// -----------------------------------------------------------------------
|
||||||
|
// Hint file persistence (for fast startup)
|
||||||
|
// -----------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Persist the KeyDir to a hint file for fast restart.
|
||||||
|
///
|
||||||
|
/// `data_file_size` is the current size of data.rdb — stored in the hint header
|
||||||
|
/// so that on next load we can detect if data.rdb changed (stale hint).
|
||||||
|
///
|
||||||
|
/// Hint file format (after the 64-byte file header):
|
||||||
|
/// For each entry: [key_len:u32 LE][key bytes][offset:u64 LE][record_len:u32 LE][value_len:u32 LE][timestamp:u64 LE]
|
||||||
|
pub fn persist_to_hint_file(&self, path: &Path, data_file_size: u64) -> StorageResult<()> {
|
||||||
|
let file = std::fs::File::create(path)?;
|
||||||
|
let mut writer = BufWriter::new(file);
|
||||||
|
|
||||||
|
// Write file header with data_file_size for staleness detection
|
||||||
|
let hdr = FileHeader::new_hint(data_file_size);
|
||||||
|
writer.write_all(&hdr.encode())?;
|
||||||
|
|
||||||
|
// Write entries
|
||||||
|
for entry in self.map.iter() {
|
||||||
|
let key_bytes = entry.key().as_bytes();
|
||||||
|
let key_len = key_bytes.len() as u32;
|
||||||
|
writer.write_all(&key_len.to_le_bytes())?;
|
||||||
|
writer.write_all(key_bytes)?;
|
||||||
|
writer.write_all(&entry.value().offset.to_le_bytes())?;
|
||||||
|
writer.write_all(&entry.value().record_len.to_le_bytes())?;
|
||||||
|
writer.write_all(&entry.value().value_len.to_le_bytes())?;
|
||||||
|
writer.write_all(&entry.value().timestamp.to_le_bytes())?;
|
||||||
|
}
|
||||||
|
|
||||||
|
writer.flush()?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Load a KeyDir from a hint file. Returns None if the file doesn't exist.
|
||||||
|
/// Returns `(keydir, stored_data_file_size)` where `stored_data_file_size` is the
|
||||||
|
/// data.rdb size recorded when the hint was written (0 = old format, unknown).
|
||||||
|
pub fn load_from_hint_file(path: &Path) -> StorageResult<Option<(Self, u64)>> {
|
||||||
|
if !path.exists() {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
|
||||||
|
let file = std::fs::File::open(path)?;
|
||||||
|
let mut reader = BufReader::new(file);
|
||||||
|
|
||||||
|
// Read and validate header
|
||||||
|
let mut hdr_buf = [0u8; FILE_HEADER_SIZE];
|
||||||
|
match reader.read_exact(&mut hdr_buf) {
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
|
||||||
|
Err(e) => return Err(e.into()),
|
||||||
|
}
|
||||||
|
let hdr = FileHeader::decode(&hdr_buf)?;
|
||||||
|
if hdr.file_type != FileType::Hint {
|
||||||
|
return Err(StorageError::CorruptRecord(format!(
|
||||||
|
"expected hint file (type 3), got type {:?}",
|
||||||
|
hdr.file_type
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
if hdr.version > FORMAT_VERSION {
|
||||||
|
return Err(StorageError::CorruptRecord(format!(
|
||||||
|
"hint file version {} is newer than supported {}",
|
||||||
|
hdr.version, FORMAT_VERSION
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
let stored_data_file_size = hdr.data_file_size;
|
||||||
|
let keydir = KeyDir::new();
|
||||||
|
|
||||||
|
loop {
|
||||||
|
// Read key_len
|
||||||
|
let mut key_len_buf = [0u8; 4];
|
||||||
|
match reader.read_exact(&mut key_len_buf) {
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break,
|
||||||
|
Err(e) => return Err(e.into()),
|
||||||
|
}
|
||||||
|
let key_len = u32::from_le_bytes(key_len_buf) as usize;
|
||||||
|
|
||||||
|
// Read key
|
||||||
|
let mut key_buf = vec![0u8; key_len];
|
||||||
|
reader.read_exact(&mut key_buf)?;
|
||||||
|
let key = String::from_utf8(key_buf)
|
||||||
|
.map_err(|e| StorageError::CorruptRecord(format!("invalid UTF-8 key: {e}")))?;
|
||||||
|
|
||||||
|
// Read entry fields
|
||||||
|
let mut fields = [0u8; 8 + 4 + 4 + 8]; // offset + record_len + value_len + timestamp = 24
|
||||||
|
reader.read_exact(&mut fields)?;
|
||||||
|
|
||||||
|
let offset = u64::from_le_bytes(fields[0..8].try_into().unwrap());
|
||||||
|
let record_len = u32::from_le_bytes(fields[8..12].try_into().unwrap());
|
||||||
|
let value_len = u32::from_le_bytes(fields[12..16].try_into().unwrap());
|
||||||
|
let timestamp = u64::from_le_bytes(fields[16..24].try_into().unwrap());
|
||||||
|
|
||||||
|
keydir.insert(
|
||||||
|
key,
|
||||||
|
KeyDirEntry {
|
||||||
|
offset,
|
||||||
|
record_len,
|
||||||
|
value_len,
|
||||||
|
timestamp,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(Some((keydir, stored_data_file_size)))
|
||||||
|
}
|
||||||
|
|
||||||
|
// -----------------------------------------------------------------------
|
||||||
|
// Hint file validation
|
||||||
|
// -----------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Validate this KeyDir (loaded from a hint file) against the actual data file.
|
||||||
|
/// Returns `Ok(true)` if the hint appears consistent, `Ok(false)` if a rebuild
|
||||||
|
/// from the data file is recommended.
|
||||||
|
///
|
||||||
|
/// Checks:
|
||||||
|
/// 1. All entry offsets + record_len fit within the data file size.
|
||||||
|
/// 2. All entry offsets are >= FILE_HEADER_SIZE.
|
||||||
|
/// 3. A random sample of entries is spot-checked by reading the record at
|
||||||
|
/// the offset and verifying the key matches.
|
||||||
|
pub fn validate_against_data_file(&self, data_path: &Path, sample_size: usize) -> StorageResult<bool> {
|
||||||
|
let file_size = std::fs::metadata(data_path)
|
||||||
|
.map(|m| m.len())
|
||||||
|
.unwrap_or(0);
|
||||||
|
|
||||||
|
if file_size < FILE_HEADER_SIZE as u64 {
|
||||||
|
// Data file is too small to even contain a header
|
||||||
|
return Ok(self.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Pass 1: bounds check all entries
|
||||||
|
let mut all_keys: Vec<(String, KeyDirEntry)> = Vec::with_capacity(self.len() as usize);
|
||||||
|
let mut bounds_ok = true;
|
||||||
|
self.for_each(|key, entry| {
|
||||||
|
if entry.offset < FILE_HEADER_SIZE as u64
|
||||||
|
|| entry.offset + entry.record_len as u64 > file_size
|
||||||
|
{
|
||||||
|
bounds_ok = false;
|
||||||
|
}
|
||||||
|
all_keys.push((key.to_string(), *entry));
|
||||||
|
});
|
||||||
|
|
||||||
|
if !bounds_ok {
|
||||||
|
return Ok(false);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Pass 2: spot-check a sample of entries by reading records from data.rdb
|
||||||
|
if all_keys.is_empty() {
|
||||||
|
return Ok(true);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sort by offset for sequential I/O, take first `sample_size` entries
|
||||||
|
all_keys.sort_by_key(|(_, e)| e.offset);
|
||||||
|
let step = if all_keys.len() <= sample_size {
|
||||||
|
1
|
||||||
|
} else {
|
||||||
|
all_keys.len() / sample_size
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut file = std::fs::File::open(data_path)?;
|
||||||
|
let mut checked = 0usize;
|
||||||
|
for (i, (expected_key, entry)) in all_keys.iter().enumerate() {
|
||||||
|
if checked >= sample_size {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
if i % step != 0 {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
// Seek to the entry's offset and try to decode the record
|
||||||
|
file.seek(SeekFrom::Start(entry.offset))?;
|
||||||
|
match DataRecord::decode_from(&mut file) {
|
||||||
|
Ok(Some((record, _disk_size))) => {
|
||||||
|
let record_key = String::from_utf8_lossy(&record.key);
|
||||||
|
if record_key != *expected_key {
|
||||||
|
return Ok(false);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(None) | Err(_) => {
|
||||||
|
return Ok(false);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
checked += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(true)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for KeyDir {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Tests
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::record::DataRecord;
|
||||||
|
use std::io::Write;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn basic_insert_get_remove() {
|
||||||
|
let kd = KeyDir::new();
|
||||||
|
assert!(kd.is_empty());
|
||||||
|
|
||||||
|
let entry = KeyDirEntry {
|
||||||
|
offset: 100,
|
||||||
|
record_len: 50,
|
||||||
|
value_len: 30,
|
||||||
|
timestamp: 1700000000000,
|
||||||
|
};
|
||||||
|
|
||||||
|
assert!(kd.insert("abc".into(), entry).is_none());
|
||||||
|
assert_eq!(kd.len(), 1);
|
||||||
|
assert!(kd.contains("abc"));
|
||||||
|
|
||||||
|
let got = kd.get("abc").unwrap();
|
||||||
|
assert_eq!(got.offset, 100);
|
||||||
|
assert_eq!(got.value_len, 30);
|
||||||
|
|
||||||
|
let removed = kd.remove("abc").unwrap();
|
||||||
|
assert_eq!(removed.offset, 100);
|
||||||
|
assert_eq!(kd.len(), 0);
|
||||||
|
assert!(!kd.contains("abc"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn insert_overwrites_returns_previous() {
|
||||||
|
let kd = KeyDir::new();
|
||||||
|
let e1 = KeyDirEntry {
|
||||||
|
offset: 100,
|
||||||
|
record_len: 50,
|
||||||
|
value_len: 30,
|
||||||
|
timestamp: 1,
|
||||||
|
};
|
||||||
|
let e2 = KeyDirEntry {
|
||||||
|
offset: 200,
|
||||||
|
record_len: 60,
|
||||||
|
value_len: 40,
|
||||||
|
timestamp: 2,
|
||||||
|
};
|
||||||
|
|
||||||
|
kd.insert("k".into(), e1);
|
||||||
|
assert_eq!(kd.len(), 1);
|
||||||
|
|
||||||
|
let prev = kd.insert("k".into(), e2).unwrap();
|
||||||
|
assert_eq!(prev.offset, 100);
|
||||||
|
// Count stays at 1 (overwrite, not new)
|
||||||
|
assert_eq!(kd.len(), 1);
|
||||||
|
assert_eq!(kd.get("k").unwrap().offset, 200);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn build_from_data_file() {
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let data_path = dir.path().join("data.rdb");
|
||||||
|
|
||||||
|
// Write a data file with 3 records: insert A, insert B, delete A
|
||||||
|
{
|
||||||
|
let mut f = std::fs::File::create(&data_path).unwrap();
|
||||||
|
let hdr = FileHeader::new(FileType::Data);
|
||||||
|
f.write_all(&hdr.encode()).unwrap();
|
||||||
|
|
||||||
|
let r1 = DataRecord {
|
||||||
|
timestamp: 1,
|
||||||
|
key: b"aaa".to_vec(),
|
||||||
|
value: b"val_a".to_vec(),
|
||||||
|
};
|
||||||
|
let r2 = DataRecord {
|
||||||
|
timestamp: 2,
|
||||||
|
key: b"bbb".to_vec(),
|
||||||
|
value: b"val_b".to_vec(),
|
||||||
|
};
|
||||||
|
let r3 = DataRecord {
|
||||||
|
timestamp: 3,
|
||||||
|
key: b"aaa".to_vec(),
|
||||||
|
value: vec![], // tombstone
|
||||||
|
};
|
||||||
|
f.write_all(&r1.encode()).unwrap();
|
||||||
|
f.write_all(&r2.encode()).unwrap();
|
||||||
|
f.write_all(&r3.encode()).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
let (kd, dead_bytes, stats) = KeyDir::build_from_data_file(&data_path).unwrap();
|
||||||
|
|
||||||
|
// Only B should be live
|
||||||
|
assert_eq!(kd.len(), 1);
|
||||||
|
assert!(kd.contains("bbb"));
|
||||||
|
assert!(!kd.contains("aaa"));
|
||||||
|
|
||||||
|
// Dead bytes: r1 (aaa live, then superseded by tombstone) + r3 (tombstone itself)
|
||||||
|
assert!(dead_bytes > 0);
|
||||||
|
|
||||||
|
// Stats
|
||||||
|
assert_eq!(stats.total_records_scanned, 3);
|
||||||
|
assert_eq!(stats.live_documents, 1);
|
||||||
|
assert_eq!(stats.tombstones, 1);
|
||||||
|
assert_eq!(stats.superseded_records, 0); // aaa was removed by tombstone, not superseded
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn hint_file_roundtrip() {
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let hint_path = dir.path().join("keydir.hint");
|
||||||
|
|
||||||
|
let kd = KeyDir::new();
|
||||||
|
kd.insert(
|
||||||
|
"doc1".into(),
|
||||||
|
KeyDirEntry {
|
||||||
|
offset: 64,
|
||||||
|
record_len: 100,
|
||||||
|
value_len: 80,
|
||||||
|
timestamp: 1000,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
kd.insert(
|
||||||
|
"doc2".into(),
|
||||||
|
KeyDirEntry {
|
||||||
|
offset: 164,
|
||||||
|
record_len: 200,
|
||||||
|
value_len: 150,
|
||||||
|
timestamp: 2000,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
kd.persist_to_hint_file(&hint_path, 12345).unwrap();
|
||||||
|
let (loaded, stored_size) = KeyDir::load_from_hint_file(&hint_path).unwrap().unwrap();
|
||||||
|
|
||||||
|
assert_eq!(stored_size, 12345);
|
||||||
|
assert_eq!(loaded.len(), 2);
|
||||||
|
let e1 = loaded.get("doc1").unwrap();
|
||||||
|
assert_eq!(e1.offset, 64);
|
||||||
|
assert_eq!(e1.record_len, 100);
|
||||||
|
assert_eq!(e1.value_len, 80);
|
||||||
|
assert_eq!(e1.timestamp, 1000);
|
||||||
|
|
||||||
|
let e2 = loaded.get("doc2").unwrap();
|
||||||
|
assert_eq!(e2.offset, 164);
|
||||||
|
assert_eq!(e2.timestamp, 2000);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn hint_file_nonexistent_returns_none() {
|
||||||
|
let result = KeyDir::load_from_hint_file(Path::new("/tmp/nonexistent_hint_file.hint"));
|
||||||
|
assert!(result.unwrap().is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn for_each_and_keys() {
|
||||||
|
let kd = KeyDir::new();
|
||||||
|
let e = KeyDirEntry {
|
||||||
|
offset: 0,
|
||||||
|
record_len: 10,
|
||||||
|
value_len: 5,
|
||||||
|
timestamp: 1,
|
||||||
|
};
|
||||||
|
kd.insert("x".into(), e);
|
||||||
|
kd.insert("y".into(), e);
|
||||||
|
|
||||||
|
let mut collected = Vec::new();
|
||||||
|
kd.for_each(|k, _| collected.push(k.to_string()));
|
||||||
|
collected.sort();
|
||||||
|
assert_eq!(collected, vec!["x", "y"]);
|
||||||
|
|
||||||
|
let mut keys = kd.keys();
|
||||||
|
keys.sort();
|
||||||
|
assert_eq!(keys, vec!["x", "y"]);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,32 @@
|
|||||||
|
//! `rustdb-storage` -- Storage adapters for RustDb.
|
||||||
|
//!
|
||||||
|
//! Provides the [`StorageAdapter`] trait and two concrete implementations:
|
||||||
|
//! - [`MemoryStorageAdapter`] -- fast in-memory store backed by `DashMap`
|
||||||
|
//! - [`FileStorageAdapter`] -- Bitcask-style append-only log with crash recovery
|
||||||
|
//!
|
||||||
|
//! Also includes an [`OpLog`] for operation logging, a [`BinaryWal`] for
|
||||||
|
//! write-ahead logging, and [`compaction`] for dead record reclamation.
|
||||||
|
|
||||||
|
pub mod adapter;
|
||||||
|
pub mod binary_wal;
|
||||||
|
pub mod compaction;
|
||||||
|
pub mod error;
|
||||||
|
pub mod file;
|
||||||
|
pub mod keydir;
|
||||||
|
pub mod memory;
|
||||||
|
pub mod oplog;
|
||||||
|
pub mod record;
|
||||||
|
pub mod validate;
|
||||||
|
|
||||||
|
pub use adapter::StorageAdapter;
|
||||||
|
pub use binary_wal::{BinaryWal, WalEntry, WalOpType};
|
||||||
|
pub use compaction::{compact_data_file, should_compact, CompactionResult};
|
||||||
|
pub use error::{StorageError, StorageResult};
|
||||||
|
pub use file::FileStorageAdapter;
|
||||||
|
pub use keydir::{BuildStats, KeyDir, KeyDirEntry};
|
||||||
|
pub use memory::MemoryStorageAdapter;
|
||||||
|
pub use oplog::{OpLog, OpLogEntry, OpLogStats, OpType};
|
||||||
|
pub use record::{
|
||||||
|
DataRecord, FileHeader, FileType, RecordScanner, FILE_HEADER_SIZE, FILE_MAGIC, FORMAT_VERSION,
|
||||||
|
RECORD_HEADER_SIZE, RECORD_MAGIC,
|
||||||
|
};
|
||||||
@@ -0,0 +1,613 @@
|
|||||||
|
use std::collections::HashSet;
|
||||||
|
use std::path::PathBuf;
|
||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use bson::{doc, oid::ObjectId, Document};
|
||||||
|
use dashmap::DashMap;
|
||||||
|
use tracing::{debug, warn};
|
||||||
|
|
||||||
|
use crate::adapter::StorageAdapter;
|
||||||
|
use crate::error::{StorageError, StorageResult};
|
||||||
|
|
||||||
|
/// Per-document timestamp tracking for conflict detection.
|
||||||
|
type TimestampMap = DashMap<String, i64>;
|
||||||
|
|
||||||
|
/// db -> coll -> id_hex -> Document
|
||||||
|
type DataStore = DashMap<String, DashMap<String, DashMap<String, Document>>>;
|
||||||
|
|
||||||
|
/// db -> coll -> Vec<index spec Document>
|
||||||
|
type IndexStore = DashMap<String, DashMap<String, Vec<Document>>>;
|
||||||
|
|
||||||
|
/// db -> coll -> id_hex -> last_modified_ms
|
||||||
|
type ModificationStore = DashMap<String, DashMap<String, TimestampMap>>;
|
||||||
|
|
||||||
|
fn now_ms() -> i64 {
|
||||||
|
SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap()
|
||||||
|
.as_millis() as i64
|
||||||
|
}
|
||||||
|
|
||||||
|
/// In-memory storage adapter backed by `DashMap`.
|
||||||
|
///
|
||||||
|
/// Optionally persists to a JSON file at a configured path.
|
||||||
|
pub struct MemoryStorageAdapter {
|
||||||
|
data: DataStore,
|
||||||
|
indexes: IndexStore,
|
||||||
|
modifications: ModificationStore,
|
||||||
|
persist_path: Option<PathBuf>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl MemoryStorageAdapter {
|
||||||
|
/// Create a new purely in-memory adapter.
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
data: DashMap::new(),
|
||||||
|
indexes: DashMap::new(),
|
||||||
|
modifications: DashMap::new(),
|
||||||
|
persist_path: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create a new adapter that will persist state to the given JSON file.
|
||||||
|
pub fn with_persist_path(path: PathBuf) -> Self {
|
||||||
|
Self {
|
||||||
|
data: DashMap::new(),
|
||||||
|
indexes: DashMap::new(),
|
||||||
|
modifications: DashMap::new(),
|
||||||
|
persist_path: Some(path),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get or create the database entry in the data store.
|
||||||
|
fn ensure_db(&self, db: &str) {
|
||||||
|
self.data.entry(db.to_string()).or_insert_with(DashMap::new);
|
||||||
|
self.indexes
|
||||||
|
.entry(db.to_string())
|
||||||
|
.or_insert_with(DashMap::new);
|
||||||
|
self.modifications
|
||||||
|
.entry(db.to_string())
|
||||||
|
.or_insert_with(DashMap::new);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn extract_id(doc: &Document) -> StorageResult<String> {
|
||||||
|
match doc.get("_id") {
|
||||||
|
Some(bson::Bson::ObjectId(oid)) => Ok(oid.to_hex()),
|
||||||
|
_ => Err(StorageError::NotFound("document missing _id".into())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn record_modification(&self, db: &str, coll: &str, id: &str) {
|
||||||
|
if let Some(db_mods) = self.modifications.get(db) {
|
||||||
|
if let Some(coll_mods) = db_mods.get(coll) {
|
||||||
|
coll_mods.insert(id.to_string(), now_ms());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl StorageAdapter for MemoryStorageAdapter {
|
||||||
|
async fn initialize(&self) -> StorageResult<()> {
|
||||||
|
debug!("MemoryStorageAdapter initialized");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn close(&self) -> StorageResult<()> {
|
||||||
|
// Persist if configured.
|
||||||
|
self.persist().await?;
|
||||||
|
debug!("MemoryStorageAdapter closed");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- database ----
|
||||||
|
|
||||||
|
async fn list_databases(&self) -> StorageResult<Vec<String>> {
|
||||||
|
Ok(self.data.iter().map(|e| e.key().clone()).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn create_database(&self, db: &str) -> StorageResult<()> {
|
||||||
|
if self.data.contains_key(db) {
|
||||||
|
return Err(StorageError::AlreadyExists(format!("database '{db}'")));
|
||||||
|
}
|
||||||
|
self.ensure_db(db);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn drop_database(&self, db: &str) -> StorageResult<()> {
|
||||||
|
self.data.remove(db);
|
||||||
|
self.indexes.remove(db);
|
||||||
|
self.modifications.remove(db);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn database_exists(&self, db: &str) -> StorageResult<bool> {
|
||||||
|
Ok(self.data.contains_key(db))
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- collection ----
|
||||||
|
|
||||||
|
async fn list_collections(&self, db: &str) -> StorageResult<Vec<String>> {
|
||||||
|
let db_ref = self
|
||||||
|
.data
|
||||||
|
.get(db)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("database '{db}'")))?;
|
||||||
|
Ok(db_ref.iter().map(|e| e.key().clone()).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn create_collection(&self, db: &str, coll: &str) -> StorageResult<()> {
|
||||||
|
self.ensure_db(db);
|
||||||
|
let db_ref = self.data.get(db).unwrap();
|
||||||
|
if db_ref.contains_key(coll) {
|
||||||
|
return Err(StorageError::AlreadyExists(format!(
|
||||||
|
"collection '{db}.{coll}'"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
db_ref.insert(coll.to_string(), DashMap::new());
|
||||||
|
drop(db_ref);
|
||||||
|
|
||||||
|
// Create modification tracker for this collection.
|
||||||
|
if let Some(db_mods) = self.modifications.get(db) {
|
||||||
|
db_mods.insert(coll.to_string(), DashMap::new());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Auto-create _id index spec.
|
||||||
|
let idx_spec = doc! { "name": "_id_", "key": { "_id": 1 } };
|
||||||
|
if let Some(db_idx) = self.indexes.get(db) {
|
||||||
|
db_idx.insert(coll.to_string(), vec![idx_spec]);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn drop_collection(&self, db: &str, coll: &str) -> StorageResult<()> {
|
||||||
|
if let Some(db_ref) = self.data.get(db) {
|
||||||
|
db_ref.remove(coll);
|
||||||
|
}
|
||||||
|
if let Some(db_idx) = self.indexes.get(db) {
|
||||||
|
db_idx.remove(coll);
|
||||||
|
}
|
||||||
|
if let Some(db_mods) = self.modifications.get(db) {
|
||||||
|
db_mods.remove(coll);
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn collection_exists(&self, db: &str, coll: &str) -> StorageResult<bool> {
|
||||||
|
Ok(self
|
||||||
|
.data
|
||||||
|
.get(db)
|
||||||
|
.map(|db_ref| db_ref.contains_key(coll))
|
||||||
|
.unwrap_or(false))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn rename_collection(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
old_name: &str,
|
||||||
|
new_name: &str,
|
||||||
|
) -> StorageResult<()> {
|
||||||
|
let db_ref = self
|
||||||
|
.data
|
||||||
|
.get(db)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("database '{db}'")))?;
|
||||||
|
if db_ref.contains_key(new_name) {
|
||||||
|
return Err(StorageError::AlreadyExists(format!(
|
||||||
|
"collection '{db}.{new_name}'"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
let (_, coll_data) = db_ref
|
||||||
|
.remove(old_name)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("collection '{db}.{old_name}'")))?;
|
||||||
|
db_ref.insert(new_name.to_string(), coll_data);
|
||||||
|
drop(db_ref);
|
||||||
|
|
||||||
|
// Rename in indexes.
|
||||||
|
if let Some(db_idx) = self.indexes.get(db) {
|
||||||
|
if let Some((_, idx_data)) = db_idx.remove(old_name) {
|
||||||
|
db_idx.insert(new_name.to_string(), idx_data);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Rename in modifications.
|
||||||
|
if let Some(db_mods) = self.modifications.get(db) {
|
||||||
|
if let Some((_, mod_data)) = db_mods.remove(old_name) {
|
||||||
|
db_mods.insert(new_name.to_string(), mod_data);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- document writes ----
|
||||||
|
|
||||||
|
async fn insert_one(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
mut doc: Document,
|
||||||
|
) -> StorageResult<String> {
|
||||||
|
// Ensure _id exists.
|
||||||
|
if !doc.contains_key("_id") {
|
||||||
|
doc.insert("_id", ObjectId::new());
|
||||||
|
}
|
||||||
|
let id = Self::extract_id(&doc)?;
|
||||||
|
|
||||||
|
let db_ref = self
|
||||||
|
.data
|
||||||
|
.get(db)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("database '{db}'")))?;
|
||||||
|
let coll_ref = db_ref
|
||||||
|
.get(coll)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("collection '{db}.{coll}'")))?;
|
||||||
|
|
||||||
|
if coll_ref.contains_key(&id) {
|
||||||
|
return Err(StorageError::AlreadyExists(format!("document '{id}'")));
|
||||||
|
}
|
||||||
|
coll_ref.insert(id.clone(), doc);
|
||||||
|
drop(coll_ref);
|
||||||
|
drop(db_ref);
|
||||||
|
|
||||||
|
self.record_modification(db, coll, &id);
|
||||||
|
Ok(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn insert_many(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
docs: Vec<Document>,
|
||||||
|
) -> StorageResult<Vec<String>> {
|
||||||
|
let mut ids = Vec::with_capacity(docs.len());
|
||||||
|
for doc in docs {
|
||||||
|
let id = self.insert_one(db, coll, doc).await?;
|
||||||
|
ids.push(id);
|
||||||
|
}
|
||||||
|
Ok(ids)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn update_by_id(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
id: &str,
|
||||||
|
doc: Document,
|
||||||
|
) -> StorageResult<()> {
|
||||||
|
let db_ref = self
|
||||||
|
.data
|
||||||
|
.get(db)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("database '{db}'")))?;
|
||||||
|
let coll_ref = db_ref
|
||||||
|
.get(coll)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("collection '{db}.{coll}'")))?;
|
||||||
|
|
||||||
|
if !coll_ref.contains_key(id) {
|
||||||
|
return Err(StorageError::NotFound(format!("document '{id}'")));
|
||||||
|
}
|
||||||
|
coll_ref.insert(id.to_string(), doc);
|
||||||
|
drop(coll_ref);
|
||||||
|
drop(db_ref);
|
||||||
|
|
||||||
|
self.record_modification(db, coll, id);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_by_id(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
id: &str,
|
||||||
|
) -> StorageResult<()> {
|
||||||
|
let db_ref = self
|
||||||
|
.data
|
||||||
|
.get(db)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("database '{db}'")))?;
|
||||||
|
let coll_ref = db_ref
|
||||||
|
.get(coll)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("collection '{db}.{coll}'")))?;
|
||||||
|
|
||||||
|
coll_ref
|
||||||
|
.remove(id)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("document '{id}'")))?;
|
||||||
|
drop(coll_ref);
|
||||||
|
drop(db_ref);
|
||||||
|
|
||||||
|
self.record_modification(db, coll, id);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_by_ids(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
ids: &[String],
|
||||||
|
) -> StorageResult<()> {
|
||||||
|
for id in ids {
|
||||||
|
self.delete_by_id(db, coll, id).await?;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- document reads ----
|
||||||
|
|
||||||
|
async fn find_all(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
) -> StorageResult<Vec<Document>> {
|
||||||
|
let db_ref = self
|
||||||
|
.data
|
||||||
|
.get(db)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("database '{db}'")))?;
|
||||||
|
let coll_ref = db_ref
|
||||||
|
.get(coll)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("collection '{db}.{coll}'")))?;
|
||||||
|
|
||||||
|
Ok(coll_ref.iter().map(|e| e.value().clone()).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn find_by_ids(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
ids: HashSet<String>,
|
||||||
|
) -> StorageResult<Vec<Document>> {
|
||||||
|
let db_ref = self
|
||||||
|
.data
|
||||||
|
.get(db)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("database '{db}'")))?;
|
||||||
|
let coll_ref = db_ref
|
||||||
|
.get(coll)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("collection '{db}.{coll}'")))?;
|
||||||
|
|
||||||
|
let mut results = Vec::with_capacity(ids.len());
|
||||||
|
for id in &ids {
|
||||||
|
if let Some(doc) = coll_ref.get(id) {
|
||||||
|
results.push(doc.value().clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(results)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn find_by_id(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
id: &str,
|
||||||
|
) -> StorageResult<Option<Document>> {
|
||||||
|
let db_ref = self
|
||||||
|
.data
|
||||||
|
.get(db)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("database '{db}'")))?;
|
||||||
|
let coll_ref = db_ref
|
||||||
|
.get(coll)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("collection '{db}.{coll}'")))?;
|
||||||
|
|
||||||
|
Ok(coll_ref.get(id).map(|e| e.value().clone()))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn count(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
) -> StorageResult<u64> {
|
||||||
|
let db_ref = self
|
||||||
|
.data
|
||||||
|
.get(db)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("database '{db}'")))?;
|
||||||
|
let coll_ref = db_ref
|
||||||
|
.get(coll)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("collection '{db}.{coll}'")))?;
|
||||||
|
|
||||||
|
Ok(coll_ref.len() as u64)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- indexes ----
|
||||||
|
|
||||||
|
async fn save_index(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
name: &str,
|
||||||
|
spec: Document,
|
||||||
|
) -> StorageResult<()> {
|
||||||
|
let db_idx = self
|
||||||
|
.indexes
|
||||||
|
.get(db)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("database '{db}'")))?;
|
||||||
|
|
||||||
|
let mut specs = db_idx
|
||||||
|
.get_mut(coll)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("collection '{db}.{coll}'")))?;
|
||||||
|
|
||||||
|
// Remove existing index with same name, then add.
|
||||||
|
specs.retain(|s| s.get_str("name").unwrap_or("") != name);
|
||||||
|
let mut full_spec = spec;
|
||||||
|
full_spec.insert("name", name);
|
||||||
|
specs.push(full_spec);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_indexes(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
) -> StorageResult<Vec<Document>> {
|
||||||
|
let db_idx = self
|
||||||
|
.indexes
|
||||||
|
.get(db)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("database '{db}'")))?;
|
||||||
|
|
||||||
|
let specs = db_idx
|
||||||
|
.get(coll)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("collection '{db}.{coll}'")))?;
|
||||||
|
|
||||||
|
Ok(specs.clone())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn drop_index(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
name: &str,
|
||||||
|
) -> StorageResult<()> {
|
||||||
|
let db_idx = self
|
||||||
|
.indexes
|
||||||
|
.get(db)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("database '{db}'")))?;
|
||||||
|
|
||||||
|
let mut specs = db_idx
|
||||||
|
.get_mut(coll)
|
||||||
|
.ok_or_else(|| StorageError::NotFound(format!("collection '{db}.{coll}'")))?;
|
||||||
|
|
||||||
|
let before = specs.len();
|
||||||
|
specs.retain(|s| s.get_str("name").unwrap_or("") != name);
|
||||||
|
if specs.len() == before {
|
||||||
|
return Err(StorageError::NotFound(format!("index '{name}'")));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- snapshot / conflict detection ----
|
||||||
|
|
||||||
|
async fn create_snapshot(
|
||||||
|
&self,
|
||||||
|
_db: &str,
|
||||||
|
_coll: &str,
|
||||||
|
) -> StorageResult<i64> {
|
||||||
|
Ok(now_ms())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn has_conflicts(
|
||||||
|
&self,
|
||||||
|
db: &str,
|
||||||
|
coll: &str,
|
||||||
|
ids: &HashSet<String>,
|
||||||
|
snapshot_time: i64,
|
||||||
|
) -> StorageResult<bool> {
|
||||||
|
if let Some(db_mods) = self.modifications.get(db) {
|
||||||
|
if let Some(coll_mods) = db_mods.get(coll) {
|
||||||
|
for id in ids {
|
||||||
|
if let Some(ts) = coll_mods.get(id) {
|
||||||
|
if *ts.value() > snapshot_time {
|
||||||
|
return Ok(true);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(false)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- persistence ----
|
||||||
|
|
||||||
|
async fn persist(&self) -> StorageResult<()> {
|
||||||
|
let path = match &self.persist_path {
|
||||||
|
Some(p) => p,
|
||||||
|
None => return Ok(()),
|
||||||
|
};
|
||||||
|
|
||||||
|
// Serialize the entire data store to JSON.
|
||||||
|
let mut db_map = serde_json::Map::new();
|
||||||
|
for db_entry in self.data.iter() {
|
||||||
|
let db_name = db_entry.key().clone();
|
||||||
|
let mut coll_map = serde_json::Map::new();
|
||||||
|
for coll_entry in db_entry.value().iter() {
|
||||||
|
let coll_name = coll_entry.key().clone();
|
||||||
|
let mut docs_map = serde_json::Map::new();
|
||||||
|
for doc_entry in coll_entry.value().iter() {
|
||||||
|
let id = doc_entry.key().clone();
|
||||||
|
// Convert bson::Document -> serde_json::Value via bson's
|
||||||
|
// built-in extended-JSON serialization.
|
||||||
|
let json_val: serde_json::Value =
|
||||||
|
bson::to_bson(doc_entry.value())
|
||||||
|
.map_err(|e| StorageError::SerializationError(e.to_string()))
|
||||||
|
.and_then(|b| {
|
||||||
|
serde_json::to_value(&b)
|
||||||
|
.map_err(|e| StorageError::SerializationError(e.to_string()))
|
||||||
|
})?;
|
||||||
|
docs_map.insert(id, json_val);
|
||||||
|
}
|
||||||
|
coll_map.insert(coll_name, serde_json::Value::Object(docs_map));
|
||||||
|
}
|
||||||
|
db_map.insert(db_name, serde_json::Value::Object(coll_map));
|
||||||
|
}
|
||||||
|
|
||||||
|
let json = serde_json::to_string_pretty(&serde_json::Value::Object(db_map))?;
|
||||||
|
if let Some(parent) = path.parent() {
|
||||||
|
tokio::fs::create_dir_all(parent).await?;
|
||||||
|
}
|
||||||
|
tokio::fs::write(path, json).await?;
|
||||||
|
debug!("MemoryStorageAdapter persisted to {:?}", path);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn restore(&self) -> StorageResult<()> {
|
||||||
|
let path = match &self.persist_path {
|
||||||
|
Some(p) => p,
|
||||||
|
None => return Ok(()),
|
||||||
|
};
|
||||||
|
|
||||||
|
if !path.exists() {
|
||||||
|
warn!("persist file not found at {:?}, skipping restore", path);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
let json = tokio::fs::read_to_string(path).await?;
|
||||||
|
let root: serde_json::Value = serde_json::from_str(&json)?;
|
||||||
|
|
||||||
|
let root_obj = root
|
||||||
|
.as_object()
|
||||||
|
.ok_or_else(|| StorageError::SerializationError("expected object".into()))?;
|
||||||
|
|
||||||
|
self.data.clear();
|
||||||
|
self.indexes.clear();
|
||||||
|
self.modifications.clear();
|
||||||
|
|
||||||
|
for (db_name, colls_val) in root_obj {
|
||||||
|
self.ensure_db(db_name);
|
||||||
|
let db_ref = self.data.get(db_name).unwrap();
|
||||||
|
let colls = colls_val
|
||||||
|
.as_object()
|
||||||
|
.ok_or_else(|| StorageError::SerializationError("expected object".into()))?;
|
||||||
|
|
||||||
|
for (coll_name, docs_val) in colls {
|
||||||
|
let coll_map: DashMap<String, Document> = DashMap::new();
|
||||||
|
let docs = docs_val
|
||||||
|
.as_object()
|
||||||
|
.ok_or_else(|| StorageError::SerializationError("expected object".into()))?;
|
||||||
|
|
||||||
|
for (id, doc_val) in docs {
|
||||||
|
let bson_val: bson::Bson = serde_json::from_value(doc_val.clone())
|
||||||
|
.map_err(|e| StorageError::SerializationError(e.to_string()))?;
|
||||||
|
let doc = bson_val
|
||||||
|
.as_document()
|
||||||
|
.ok_or_else(|| {
|
||||||
|
StorageError::SerializationError("expected document".into())
|
||||||
|
})?
|
||||||
|
.clone();
|
||||||
|
coll_map.insert(id.clone(), doc);
|
||||||
|
}
|
||||||
|
db_ref.insert(coll_name.clone(), coll_map);
|
||||||
|
|
||||||
|
// Restore modification tracker and default _id index.
|
||||||
|
if let Some(db_mods) = self.modifications.get(db_name) {
|
||||||
|
db_mods.insert(coll_name.clone(), DashMap::new());
|
||||||
|
}
|
||||||
|
if let Some(db_idx) = self.indexes.get(db_name) {
|
||||||
|
let idx_spec = doc! { "name": "_id_", "key": { "_id": 1 } };
|
||||||
|
db_idx.insert(coll_name.clone(), vec![idx_spec]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
debug!("MemoryStorageAdapter restored from {:?}", path);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for MemoryStorageAdapter {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,203 @@
|
|||||||
|
//! Operation log (OpLog) for tracking mutations.
|
||||||
|
//!
|
||||||
|
//! The OpLog records every write operation so that changes can be replayed,
|
||||||
|
//! replicated, or used for change-stream style notifications.
|
||||||
|
//! Each entry stores both the new and previous document state, enabling
|
||||||
|
//! point-in-time replay and revert.
|
||||||
|
|
||||||
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
use bson::Document;
|
||||||
|
use dashmap::DashMap;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
/// The type of operation recorded in the oplog.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||||
|
pub enum OpType {
|
||||||
|
Insert,
|
||||||
|
Update,
|
||||||
|
Delete,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A single oplog entry.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct OpLogEntry {
|
||||||
|
/// Monotonically increasing sequence number.
|
||||||
|
pub seq: u64,
|
||||||
|
/// Timestamp in milliseconds since UNIX epoch.
|
||||||
|
pub timestamp_ms: i64,
|
||||||
|
/// Operation type.
|
||||||
|
pub op: OpType,
|
||||||
|
/// Database name.
|
||||||
|
pub db: String,
|
||||||
|
/// Collection name.
|
||||||
|
pub collection: String,
|
||||||
|
/// Document id (hex string).
|
||||||
|
pub document_id: String,
|
||||||
|
/// The new document snapshot (for insert/update; None for delete).
|
||||||
|
pub document: Option<Document>,
|
||||||
|
/// The previous document snapshot (for update/delete; None for insert).
|
||||||
|
pub previous_document: Option<Document>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Aggregate statistics about the oplog.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct OpLogStats {
|
||||||
|
pub current_seq: u64,
|
||||||
|
pub total_entries: usize,
|
||||||
|
pub oldest_seq: u64,
|
||||||
|
pub inserts: usize,
|
||||||
|
pub updates: usize,
|
||||||
|
pub deletes: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// In-memory operation log.
|
||||||
|
pub struct OpLog {
|
||||||
|
/// All entries keyed by sequence number.
|
||||||
|
entries: DashMap<u64, OpLogEntry>,
|
||||||
|
/// Next sequence number.
|
||||||
|
next_seq: AtomicU64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl OpLog {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
entries: DashMap::new(),
|
||||||
|
next_seq: AtomicU64::new(1),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append an operation to the log and return its sequence number.
|
||||||
|
pub fn append(
|
||||||
|
&self,
|
||||||
|
op: OpType,
|
||||||
|
db: &str,
|
||||||
|
collection: &str,
|
||||||
|
document_id: &str,
|
||||||
|
document: Option<Document>,
|
||||||
|
previous_document: Option<Document>,
|
||||||
|
) -> u64 {
|
||||||
|
let seq = self.next_seq.fetch_add(1, Ordering::SeqCst);
|
||||||
|
let entry = OpLogEntry {
|
||||||
|
seq,
|
||||||
|
timestamp_ms: SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap()
|
||||||
|
.as_millis() as i64,
|
||||||
|
op,
|
||||||
|
db: db.to_string(),
|
||||||
|
collection: collection.to_string(),
|
||||||
|
document_id: document_id.to_string(),
|
||||||
|
document,
|
||||||
|
previous_document,
|
||||||
|
};
|
||||||
|
self.entries.insert(seq, entry);
|
||||||
|
seq
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get a single entry by sequence number.
|
||||||
|
pub fn get_entry(&self, seq: u64) -> Option<OpLogEntry> {
|
||||||
|
self.entries.get(&seq).map(|e| e.value().clone())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get all entries with sequence number >= `since`.
|
||||||
|
pub fn entries_since(&self, since: u64) -> Vec<OpLogEntry> {
|
||||||
|
let mut result: Vec<_> = self
|
||||||
|
.entries
|
||||||
|
.iter()
|
||||||
|
.filter(|e| *e.key() >= since)
|
||||||
|
.map(|e| e.value().clone())
|
||||||
|
.collect();
|
||||||
|
result.sort_by_key(|e| e.seq);
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get entries in range [from_seq, to_seq] inclusive, sorted by seq.
|
||||||
|
pub fn entries_range(&self, from_seq: u64, to_seq: u64) -> Vec<OpLogEntry> {
|
||||||
|
let mut result: Vec<_> = self
|
||||||
|
.entries
|
||||||
|
.iter()
|
||||||
|
.filter(|e| {
|
||||||
|
let k = *e.key();
|
||||||
|
k >= from_seq && k <= to_seq
|
||||||
|
})
|
||||||
|
.map(|e| e.value().clone())
|
||||||
|
.collect();
|
||||||
|
result.sort_by_key(|e| e.seq);
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Remove all entries with seq > `after_seq` and reset the next_seq counter.
|
||||||
|
pub fn truncate_after(&self, after_seq: u64) {
|
||||||
|
let keys_to_remove: Vec<u64> = self
|
||||||
|
.entries
|
||||||
|
.iter()
|
||||||
|
.filter(|e| *e.key() > after_seq)
|
||||||
|
.map(|e| *e.key())
|
||||||
|
.collect();
|
||||||
|
for key in keys_to_remove {
|
||||||
|
self.entries.remove(&key);
|
||||||
|
}
|
||||||
|
self.next_seq.store(after_seq + 1, Ordering::SeqCst);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the current (latest) sequence number. Returns 0 if empty.
|
||||||
|
pub fn current_seq(&self) -> u64 {
|
||||||
|
self.next_seq.load(Ordering::SeqCst).saturating_sub(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get aggregate statistics.
|
||||||
|
pub fn stats(&self) -> OpLogStats {
|
||||||
|
let mut inserts = 0usize;
|
||||||
|
let mut updates = 0usize;
|
||||||
|
let mut deletes = 0usize;
|
||||||
|
let mut oldest_seq = u64::MAX;
|
||||||
|
|
||||||
|
for entry in self.entries.iter() {
|
||||||
|
match entry.value().op {
|
||||||
|
OpType::Insert => inserts += 1,
|
||||||
|
OpType::Update => updates += 1,
|
||||||
|
OpType::Delete => deletes += 1,
|
||||||
|
}
|
||||||
|
if entry.value().seq < oldest_seq {
|
||||||
|
oldest_seq = entry.value().seq;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if oldest_seq == u64::MAX {
|
||||||
|
oldest_seq = 0;
|
||||||
|
}
|
||||||
|
|
||||||
|
OpLogStats {
|
||||||
|
current_seq: self.current_seq(),
|
||||||
|
total_entries: self.entries.len(),
|
||||||
|
oldest_seq,
|
||||||
|
inserts,
|
||||||
|
updates,
|
||||||
|
deletes,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Clear all entries.
|
||||||
|
pub fn clear(&self) {
|
||||||
|
self.entries.clear();
|
||||||
|
self.next_seq.store(1, Ordering::SeqCst);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Number of entries in the log.
|
||||||
|
pub fn len(&self) -> usize {
|
||||||
|
self.entries.len()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Whether the log is empty.
|
||||||
|
pub fn is_empty(&self) -> bool {
|
||||||
|
self.entries.is_empty()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for OpLog {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,472 @@
|
|||||||
|
//! Binary data record format for the Bitcask-style storage engine.
|
||||||
|
//!
|
||||||
|
//! # File Version Header (64 bytes, at offset 0 of every .rdb / .hint file)
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! ┌──────────────┬──────────┬──────────┬──────────┬──────────┬───────────────┐
|
||||||
|
//! │ magic │ version │ file_type│ flags │ created │ reserved │
|
||||||
|
//! │ 8 bytes │ u16 LE │ u8 │ u32 LE │ u64 LE │ 41 bytes │
|
||||||
|
//! │ "SMARTDB\0" │ │ │ │ epoch_ms │ (zeros) │
|
||||||
|
//! └──────────────┴──────────┴──────────┴──────────┴──────────┴───────────────┘
|
||||||
|
//! ```
|
||||||
|
//!
|
||||||
|
//! # Data Record (appended after the header)
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! ┌──────────┬──────────┬──────────┬──────────┬──────────┬──────────────────┐
|
||||||
|
//! │ magic │ timestamp│ key_len │ val_len │ crc32 │ payload │
|
||||||
|
//! │ u16 LE │ u64 LE │ u32 LE │ u32 LE │ u32 LE │ [key][value] │
|
||||||
|
//! │ 0xDB01 │ epoch_ms │ │ 0=delete │ │ │
|
||||||
|
//! └──────────┴──────────┴──────────┴──────────┴──────────┴──────────────────┘
|
||||||
|
//! ```
|
||||||
|
|
||||||
|
use std::io::{self, Read};
|
||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
use crate::error::{StorageError, StorageResult};
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Constants
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// File-level magic: b"SMARTDB\0"
|
||||||
|
pub const FILE_MAGIC: &[u8; 8] = b"SMARTDB\0";
|
||||||
|
|
||||||
|
/// Current storage format version.
|
||||||
|
pub const FORMAT_VERSION: u16 = 1;
|
||||||
|
|
||||||
|
/// File version header size.
|
||||||
|
pub const FILE_HEADER_SIZE: usize = 64;
|
||||||
|
|
||||||
|
/// Per-record magic.
|
||||||
|
pub const RECORD_MAGIC: u16 = 0xDB01;
|
||||||
|
|
||||||
|
/// Per-record header size (before payload).
|
||||||
|
pub const RECORD_HEADER_SIZE: usize = 2 + 8 + 4 + 4 + 4; // 22 bytes
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// File type tag stored in the version header
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
#[repr(u8)]
|
||||||
|
pub enum FileType {
|
||||||
|
Data = 1,
|
||||||
|
Wal = 2,
|
||||||
|
Hint = 3,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl FileType {
|
||||||
|
pub fn from_u8(v: u8) -> StorageResult<Self> {
|
||||||
|
match v {
|
||||||
|
1 => Ok(FileType::Data),
|
||||||
|
2 => Ok(FileType::Wal),
|
||||||
|
3 => Ok(FileType::Hint),
|
||||||
|
_ => Err(StorageError::CorruptRecord(format!(
|
||||||
|
"unknown file type tag: {v}"
|
||||||
|
))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// File Version Header
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct FileHeader {
|
||||||
|
pub version: u16,
|
||||||
|
pub file_type: FileType,
|
||||||
|
pub flags: u32,
|
||||||
|
pub created_ms: u64,
|
||||||
|
/// For hint files: the data.rdb file size at the time the hint was written.
|
||||||
|
/// Used to detect stale hints after ungraceful shutdown. 0 = unknown (old format).
|
||||||
|
pub data_file_size: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl FileHeader {
|
||||||
|
/// Create a new header for the current format version.
|
||||||
|
pub fn new(file_type: FileType) -> Self {
|
||||||
|
Self {
|
||||||
|
version: FORMAT_VERSION,
|
||||||
|
file_type,
|
||||||
|
flags: 0,
|
||||||
|
created_ms: now_ms(),
|
||||||
|
data_file_size: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create a new hint header that records the data file size.
|
||||||
|
pub fn new_hint(data_file_size: u64) -> Self {
|
||||||
|
Self {
|
||||||
|
version: FORMAT_VERSION,
|
||||||
|
file_type: FileType::Hint,
|
||||||
|
flags: 0,
|
||||||
|
created_ms: now_ms(),
|
||||||
|
data_file_size,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Encode the header to a 64-byte buffer.
|
||||||
|
pub fn encode(&self) -> [u8; FILE_HEADER_SIZE] {
|
||||||
|
let mut buf = [0u8; FILE_HEADER_SIZE];
|
||||||
|
buf[0..8].copy_from_slice(FILE_MAGIC);
|
||||||
|
buf[8..10].copy_from_slice(&self.version.to_le_bytes());
|
||||||
|
buf[10] = self.file_type as u8;
|
||||||
|
buf[11..15].copy_from_slice(&self.flags.to_le_bytes());
|
||||||
|
buf[15..23].copy_from_slice(&self.created_ms.to_le_bytes());
|
||||||
|
buf[23..31].copy_from_slice(&self.data_file_size.to_le_bytes());
|
||||||
|
// bytes 31..64 are reserved (zeros)
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Decode a 64-byte header. Validates magic and version.
|
||||||
|
pub fn decode(buf: &[u8; FILE_HEADER_SIZE]) -> StorageResult<Self> {
|
||||||
|
if &buf[0..8] != FILE_MAGIC {
|
||||||
|
return Err(StorageError::CorruptRecord(
|
||||||
|
"invalid file magic — not a SmartDB file".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let version = u16::from_le_bytes([buf[8], buf[9]]);
|
||||||
|
if version > FORMAT_VERSION {
|
||||||
|
return Err(StorageError::CorruptRecord(format!(
|
||||||
|
"file format version {version} is newer than supported version {FORMAT_VERSION} — please upgrade"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
if version == 0 {
|
||||||
|
return Err(StorageError::CorruptRecord(
|
||||||
|
"file format version 0 is invalid".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let file_type = FileType::from_u8(buf[10])?;
|
||||||
|
let flags = u32::from_le_bytes([buf[11], buf[12], buf[13], buf[14]]);
|
||||||
|
let created_ms = u64::from_le_bytes([
|
||||||
|
buf[15], buf[16], buf[17], buf[18], buf[19], buf[20], buf[21], buf[22],
|
||||||
|
]);
|
||||||
|
let data_file_size = u64::from_le_bytes([
|
||||||
|
buf[23], buf[24], buf[25], buf[26], buf[27], buf[28], buf[29], buf[30],
|
||||||
|
]);
|
||||||
|
Ok(Self {
|
||||||
|
version,
|
||||||
|
file_type,
|
||||||
|
flags,
|
||||||
|
created_ms,
|
||||||
|
data_file_size,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Data Record
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// A single data record (live document or tombstone).
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct DataRecord {
|
||||||
|
pub timestamp: u64,
|
||||||
|
pub key: Vec<u8>,
|
||||||
|
/// BSON value bytes. Empty for tombstones.
|
||||||
|
pub value: Vec<u8>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl DataRecord {
|
||||||
|
/// Whether this record is a tombstone (delete marker).
|
||||||
|
pub fn is_tombstone(&self) -> bool {
|
||||||
|
self.value.is_empty()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Total size on disk (header + payload).
|
||||||
|
pub fn disk_size(&self) -> usize {
|
||||||
|
RECORD_HEADER_SIZE + self.key.len() + self.value.len()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Encode to bytes. CRC32 covers magic + timestamp + key_len + val_len + payload.
|
||||||
|
pub fn encode(&self) -> Vec<u8> {
|
||||||
|
let key_len = self.key.len() as u32;
|
||||||
|
let val_len = self.value.len() as u32;
|
||||||
|
let total = RECORD_HEADER_SIZE + self.key.len() + self.value.len();
|
||||||
|
let mut buf = Vec::with_capacity(total);
|
||||||
|
|
||||||
|
// Write fields WITHOUT crc first to compute checksum.
|
||||||
|
buf.extend_from_slice(&RECORD_MAGIC.to_le_bytes()); // 2
|
||||||
|
buf.extend_from_slice(&self.timestamp.to_le_bytes()); // 8
|
||||||
|
buf.extend_from_slice(&key_len.to_le_bytes()); // 4
|
||||||
|
buf.extend_from_slice(&val_len.to_le_bytes()); // 4
|
||||||
|
// placeholder for crc32 — we'll fill it after computing
|
||||||
|
buf.extend_from_slice(&0u32.to_le_bytes()); // 4
|
||||||
|
buf.extend_from_slice(&self.key); // key_len
|
||||||
|
buf.extend_from_slice(&self.value); // val_len
|
||||||
|
|
||||||
|
// CRC covers everything except the crc32 field itself:
|
||||||
|
// bytes [0..18] (magic+ts+key_len+val_len) + bytes [22..] (payload)
|
||||||
|
let mut hasher = crc32fast::Hasher::new();
|
||||||
|
hasher.update(&buf[0..18]);
|
||||||
|
hasher.update(&buf[22..]);
|
||||||
|
let crc = hasher.finalize();
|
||||||
|
buf[18..22].copy_from_slice(&crc.to_le_bytes());
|
||||||
|
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Decode a record from a reader. Returns the record and its total disk size.
|
||||||
|
/// On EOF at the very start (no bytes to read), returns Ok(None).
|
||||||
|
pub fn decode_from<R: Read>(reader: &mut R) -> StorageResult<Option<(Self, usize)>> {
|
||||||
|
// Read header
|
||||||
|
let mut hdr = [0u8; RECORD_HEADER_SIZE];
|
||||||
|
match reader.read_exact(&mut hdr) {
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
|
||||||
|
Err(e) => return Err(e.into()),
|
||||||
|
}
|
||||||
|
|
||||||
|
let magic = u16::from_le_bytes([hdr[0], hdr[1]]);
|
||||||
|
if magic != RECORD_MAGIC {
|
||||||
|
return Err(StorageError::CorruptRecord(format!(
|
||||||
|
"invalid record magic: 0x{magic:04X}, expected 0x{RECORD_MAGIC:04X}"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
let timestamp = u64::from_le_bytes(hdr[2..10].try_into().unwrap());
|
||||||
|
let key_len = u32::from_le_bytes(hdr[10..14].try_into().unwrap()) as usize;
|
||||||
|
let val_len = u32::from_le_bytes(hdr[14..18].try_into().unwrap()) as usize;
|
||||||
|
let stored_crc = u32::from_le_bytes(hdr[18..22].try_into().unwrap());
|
||||||
|
|
||||||
|
// Read payload
|
||||||
|
let payload_len = key_len + val_len;
|
||||||
|
let mut payload = vec![0u8; payload_len];
|
||||||
|
reader.read_exact(&mut payload)?;
|
||||||
|
|
||||||
|
// Verify CRC: covers header bytes [0..18] + payload
|
||||||
|
let mut hasher = crc32fast::Hasher::new();
|
||||||
|
hasher.update(&hdr[0..18]);
|
||||||
|
hasher.update(&payload);
|
||||||
|
let computed_crc = hasher.finalize();
|
||||||
|
if computed_crc != stored_crc {
|
||||||
|
return Err(StorageError::ChecksumMismatch {
|
||||||
|
expected: stored_crc,
|
||||||
|
actual: computed_crc,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let key = payload[..key_len].to_vec();
|
||||||
|
let value = payload[key_len..].to_vec();
|
||||||
|
let disk_size = RECORD_HEADER_SIZE + payload_len;
|
||||||
|
|
||||||
|
Ok(Some((
|
||||||
|
DataRecord {
|
||||||
|
timestamp,
|
||||||
|
key,
|
||||||
|
value,
|
||||||
|
},
|
||||||
|
disk_size,
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Record Scanner — iterate records from a byte slice or reader
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Scans records sequentially from a reader, yielding (offset, record) pairs.
|
||||||
|
/// Starts reading from the current reader position. The `base_offset` parameter
|
||||||
|
/// indicates the byte offset in the file where reading begins (typically
|
||||||
|
/// `FILE_HEADER_SIZE` for a data file).
|
||||||
|
pub struct RecordScanner<R> {
|
||||||
|
reader: R,
|
||||||
|
offset: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<R: Read> RecordScanner<R> {
|
||||||
|
pub fn new(reader: R, base_offset: u64) -> Self {
|
||||||
|
Self {
|
||||||
|
reader,
|
||||||
|
offset: base_offset,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<R: Read> Iterator for RecordScanner<R> {
|
||||||
|
/// (file_offset, record) or an error. Iteration stops on EOF or error.
|
||||||
|
type Item = StorageResult<(u64, DataRecord)>;
|
||||||
|
|
||||||
|
fn next(&mut self) -> Option<Self::Item> {
|
||||||
|
match DataRecord::decode_from(&mut self.reader) {
|
||||||
|
Ok(Some((record, disk_size))) => {
|
||||||
|
let offset = self.offset;
|
||||||
|
self.offset += disk_size as u64;
|
||||||
|
Some(Ok((offset, record)))
|
||||||
|
}
|
||||||
|
Ok(None) => None, // clean EOF
|
||||||
|
Err(e) => Some(Err(e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Helpers
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Current time in milliseconds since UNIX epoch.
|
||||||
|
pub fn now_ms() -> u64 {
|
||||||
|
SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap()
|
||||||
|
.as_millis() as u64
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Tests
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn file_header_roundtrip() {
|
||||||
|
let hdr = FileHeader::new(FileType::Data);
|
||||||
|
let buf = hdr.encode();
|
||||||
|
assert_eq!(buf.len(), FILE_HEADER_SIZE);
|
||||||
|
|
||||||
|
let decoded = FileHeader::decode(&buf).unwrap();
|
||||||
|
assert_eq!(decoded.version, FORMAT_VERSION);
|
||||||
|
assert_eq!(decoded.file_type, FileType::Data);
|
||||||
|
assert_eq!(decoded.flags, 0);
|
||||||
|
assert_eq!(decoded.created_ms, hdr.created_ms);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn file_header_rejects_bad_magic() {
|
||||||
|
let mut buf = [0u8; FILE_HEADER_SIZE];
|
||||||
|
buf[0..8].copy_from_slice(b"BADMAGIC");
|
||||||
|
assert!(FileHeader::decode(&buf).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn file_header_rejects_future_version() {
|
||||||
|
let mut hdr = FileHeader::new(FileType::Data);
|
||||||
|
hdr.version = FORMAT_VERSION + 1;
|
||||||
|
let buf = hdr.encode();
|
||||||
|
// Manually patch the version in the buffer
|
||||||
|
let mut buf2 = buf;
|
||||||
|
buf2[8..10].copy_from_slice(&(FORMAT_VERSION + 1).to_le_bytes());
|
||||||
|
assert!(FileHeader::decode(&buf2).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn record_roundtrip_live() {
|
||||||
|
let rec = DataRecord {
|
||||||
|
timestamp: 1700000000000,
|
||||||
|
key: b"abc123".to_vec(),
|
||||||
|
value: b"\x10\x00\x00\x00\x02hi\x00\x03\x00\x00\x00ok\x00\x00".to_vec(),
|
||||||
|
};
|
||||||
|
let encoded = rec.encode();
|
||||||
|
assert_eq!(encoded.len(), rec.disk_size());
|
||||||
|
|
||||||
|
let mut cursor = std::io::Cursor::new(&encoded);
|
||||||
|
let (decoded, size) = DataRecord::decode_from(&mut cursor).unwrap().unwrap();
|
||||||
|
assert_eq!(size, encoded.len());
|
||||||
|
assert_eq!(decoded.timestamp, rec.timestamp);
|
||||||
|
assert_eq!(decoded.key, rec.key);
|
||||||
|
assert_eq!(decoded.value, rec.value);
|
||||||
|
assert!(!decoded.is_tombstone());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn record_roundtrip_tombstone() {
|
||||||
|
let rec = DataRecord {
|
||||||
|
timestamp: 1700000000000,
|
||||||
|
key: b"def456".to_vec(),
|
||||||
|
value: vec![],
|
||||||
|
};
|
||||||
|
assert!(rec.is_tombstone());
|
||||||
|
let encoded = rec.encode();
|
||||||
|
|
||||||
|
let mut cursor = std::io::Cursor::new(&encoded);
|
||||||
|
let (decoded, _) = DataRecord::decode_from(&mut cursor).unwrap().unwrap();
|
||||||
|
assert!(decoded.is_tombstone());
|
||||||
|
assert_eq!(decoded.key, b"def456");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn record_detects_corruption() {
|
||||||
|
let rec = DataRecord {
|
||||||
|
timestamp: 42,
|
||||||
|
key: b"key".to_vec(),
|
||||||
|
value: b"value".to_vec(),
|
||||||
|
};
|
||||||
|
let mut encoded = rec.encode();
|
||||||
|
// Flip a bit in the payload
|
||||||
|
let last = encoded.len() - 1;
|
||||||
|
encoded[last] ^= 0xFF;
|
||||||
|
|
||||||
|
let mut cursor = std::io::Cursor::new(&encoded);
|
||||||
|
let result = DataRecord::decode_from(&mut cursor);
|
||||||
|
assert!(matches!(result, Err(StorageError::ChecksumMismatch { .. })));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn record_detects_bad_magic() {
|
||||||
|
let rec = DataRecord {
|
||||||
|
timestamp: 42,
|
||||||
|
key: b"key".to_vec(),
|
||||||
|
value: b"value".to_vec(),
|
||||||
|
};
|
||||||
|
let mut encoded = rec.encode();
|
||||||
|
encoded[0] = 0xFF;
|
||||||
|
encoded[1] = 0xFF;
|
||||||
|
|
||||||
|
let mut cursor = std::io::Cursor::new(&encoded);
|
||||||
|
let result = DataRecord::decode_from(&mut cursor);
|
||||||
|
assert!(matches!(result, Err(StorageError::CorruptRecord(_))));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn eof_returns_none() {
|
||||||
|
let empty: &[u8] = &[];
|
||||||
|
let mut cursor = std::io::Cursor::new(empty);
|
||||||
|
let result = DataRecord::decode_from(&mut cursor).unwrap();
|
||||||
|
assert!(result.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn scanner_iterates_multiple_records() {
|
||||||
|
let records = vec![
|
||||||
|
DataRecord {
|
||||||
|
timestamp: 1,
|
||||||
|
key: b"a".to_vec(),
|
||||||
|
value: b"v1".to_vec(),
|
||||||
|
},
|
||||||
|
DataRecord {
|
||||||
|
timestamp: 2,
|
||||||
|
key: b"b".to_vec(),
|
||||||
|
value: b"v2".to_vec(),
|
||||||
|
},
|
||||||
|
DataRecord {
|
||||||
|
timestamp: 3,
|
||||||
|
key: b"c".to_vec(),
|
||||||
|
value: vec![],
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
for r in &records {
|
||||||
|
buf.extend_from_slice(&r.encode());
|
||||||
|
}
|
||||||
|
|
||||||
|
let scanner = RecordScanner::new(std::io::Cursor::new(&buf), 0);
|
||||||
|
let results: Vec<_> = scanner.collect::<Result<Vec<_>, _>>().unwrap();
|
||||||
|
assert_eq!(results.len(), 3);
|
||||||
|
assert_eq!(results[0].1.key, b"a");
|
||||||
|
assert_eq!(results[1].1.key, b"b");
|
||||||
|
assert!(results[2].1.is_tombstone());
|
||||||
|
|
||||||
|
// Verify offsets are correct
|
||||||
|
assert_eq!(results[0].0, 0);
|
||||||
|
assert_eq!(results[1].0, records[0].disk_size() as u64);
|
||||||
|
assert_eq!(
|
||||||
|
results[2].0,
|
||||||
|
(records[0].disk_size() + records[1].disk_size()) as u64
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,330 @@
|
|||||||
|
//! Data integrity validation for RustDb storage directories.
|
||||||
|
//!
|
||||||
|
//! Provides offline validation of data files without starting the server.
|
||||||
|
//! Checks header magic, record CRC32 checksums, duplicate IDs, and
|
||||||
|
//! keydir.hint consistency.
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::io::{BufReader, Read};
|
||||||
|
use std::path::Path;
|
||||||
|
|
||||||
|
use crate::error::{StorageError, StorageResult};
|
||||||
|
use crate::keydir::KeyDir;
|
||||||
|
use crate::record::{FileHeader, FileType, RecordScanner, FILE_HEADER_SIZE};
|
||||||
|
|
||||||
|
/// Result of validating an entire data directory.
|
||||||
|
pub struct ValidationReport {
|
||||||
|
pub collections: Vec<CollectionReport>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Result of validating a single collection.
|
||||||
|
pub struct CollectionReport {
|
||||||
|
pub db: String,
|
||||||
|
pub collection: String,
|
||||||
|
pub header_valid: bool,
|
||||||
|
pub total_records: u64,
|
||||||
|
pub live_documents: u64,
|
||||||
|
pub tombstones: u64,
|
||||||
|
pub duplicate_ids: Vec<String>,
|
||||||
|
pub checksum_errors: u64,
|
||||||
|
pub decode_errors: u64,
|
||||||
|
pub data_file_size: u64,
|
||||||
|
pub hint_file_exists: bool,
|
||||||
|
pub orphaned_hint_entries: u64,
|
||||||
|
pub errors: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ValidationReport {
|
||||||
|
/// Whether any errors were found across all collections.
|
||||||
|
pub fn has_errors(&self) -> bool {
|
||||||
|
self.collections.iter().any(|c| {
|
||||||
|
!c.header_valid
|
||||||
|
|| !c.duplicate_ids.is_empty()
|
||||||
|
|| c.checksum_errors > 0
|
||||||
|
|| c.decode_errors > 0
|
||||||
|
|| c.orphaned_hint_entries > 0
|
||||||
|
|| !c.errors.is_empty()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Print a human-readable summary to stdout.
|
||||||
|
pub fn print_summary(&self) {
|
||||||
|
println!("=== SmartDB Data Integrity Report ===");
|
||||||
|
println!();
|
||||||
|
|
||||||
|
let mut total_errors = 0u64;
|
||||||
|
|
||||||
|
for report in &self.collections {
|
||||||
|
println!("Database: {}", report.db);
|
||||||
|
println!(" Collection: {}", report.collection);
|
||||||
|
println!(
|
||||||
|
" Header: {}",
|
||||||
|
if report.header_valid { "OK" } else { "INVALID" }
|
||||||
|
);
|
||||||
|
println!(
|
||||||
|
" Records: {} ({} live, {} tombstones)",
|
||||||
|
report.total_records, report.live_documents, report.tombstones
|
||||||
|
);
|
||||||
|
println!(" Data size: {} bytes", report.data_file_size);
|
||||||
|
|
||||||
|
if report.duplicate_ids.is_empty() {
|
||||||
|
println!(" Duplicates: 0");
|
||||||
|
} else {
|
||||||
|
let ids_preview: Vec<&str> = report.duplicate_ids.iter().take(5).map(|s| s.as_str()).collect();
|
||||||
|
let suffix = if report.duplicate_ids.len() > 5 {
|
||||||
|
format!(", ... and {} more", report.duplicate_ids.len() - 5)
|
||||||
|
} else {
|
||||||
|
String::new()
|
||||||
|
};
|
||||||
|
println!(
|
||||||
|
" Duplicates: {} (ids: {}{})",
|
||||||
|
report.duplicate_ids.len(),
|
||||||
|
ids_preview.join(", "),
|
||||||
|
suffix
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
if report.checksum_errors > 0 {
|
||||||
|
println!(" CRC errors: {}", report.checksum_errors);
|
||||||
|
} else {
|
||||||
|
println!(" CRC errors: 0");
|
||||||
|
}
|
||||||
|
|
||||||
|
if report.decode_errors > 0 {
|
||||||
|
println!(" Decode errors: {}", report.decode_errors);
|
||||||
|
}
|
||||||
|
|
||||||
|
if report.hint_file_exists {
|
||||||
|
if report.orphaned_hint_entries > 0 {
|
||||||
|
println!(
|
||||||
|
" Hint file: STALE ({} orphaned entries)",
|
||||||
|
report.orphaned_hint_entries
|
||||||
|
);
|
||||||
|
} else {
|
||||||
|
println!(" Hint file: OK");
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
println!(" Hint file: absent");
|
||||||
|
}
|
||||||
|
|
||||||
|
for err in &report.errors {
|
||||||
|
println!(" ERROR: {}", err);
|
||||||
|
}
|
||||||
|
|
||||||
|
println!();
|
||||||
|
|
||||||
|
if !report.header_valid { total_errors += 1; }
|
||||||
|
total_errors += report.duplicate_ids.len() as u64;
|
||||||
|
total_errors += report.checksum_errors;
|
||||||
|
total_errors += report.decode_errors;
|
||||||
|
total_errors += report.orphaned_hint_entries;
|
||||||
|
total_errors += report.errors.len() as u64;
|
||||||
|
}
|
||||||
|
|
||||||
|
println!(
|
||||||
|
"Summary: {} collection(s) checked, {} error(s) found.",
|
||||||
|
self.collections.len(),
|
||||||
|
total_errors
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Validate all collections in a data directory.
|
||||||
|
///
|
||||||
|
/// The directory structure is expected to be:
|
||||||
|
/// ```text
|
||||||
|
/// {base_path}/{db}/{collection}/data.rdb
|
||||||
|
/// ```
|
||||||
|
pub fn validate_data_directory(base_path: &str) -> StorageResult<ValidationReport> {
|
||||||
|
let base = Path::new(base_path);
|
||||||
|
if !base.exists() {
|
||||||
|
return Err(StorageError::IoError(std::io::Error::new(
|
||||||
|
std::io::ErrorKind::NotFound,
|
||||||
|
format!("data directory not found: {base_path}"),
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut collections = Vec::new();
|
||||||
|
|
||||||
|
// Iterate database directories
|
||||||
|
let entries = std::fs::read_dir(base)?;
|
||||||
|
for entry in entries {
|
||||||
|
let entry = entry?;
|
||||||
|
if !entry.file_type()?.is_dir() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let db_name = match entry.file_name().to_str() {
|
||||||
|
Some(s) => s.to_string(),
|
||||||
|
None => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Iterate collection directories
|
||||||
|
let db_entries = std::fs::read_dir(entry.path())?;
|
||||||
|
for coll_entry in db_entries {
|
||||||
|
let coll_entry = coll_entry?;
|
||||||
|
if !coll_entry.file_type()?.is_dir() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let coll_name = match coll_entry.file_name().to_str() {
|
||||||
|
Some(s) => s.to_string(),
|
||||||
|
None => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
let data_path = coll_entry.path().join("data.rdb");
|
||||||
|
if !data_path.exists() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let report = validate_collection(&db_name, &coll_name, &coll_entry.path());
|
||||||
|
collections.push(report);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sort for deterministic output
|
||||||
|
collections.sort_by(|a, b| (&a.db, &a.collection).cmp(&(&b.db, &b.collection)));
|
||||||
|
|
||||||
|
Ok(ValidationReport { collections })
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Validate a single collection directory.
|
||||||
|
fn validate_collection(db: &str, coll: &str, coll_dir: &Path) -> CollectionReport {
|
||||||
|
let data_path = coll_dir.join("data.rdb");
|
||||||
|
let hint_path = coll_dir.join("keydir.hint");
|
||||||
|
|
||||||
|
let mut report = CollectionReport {
|
||||||
|
db: db.to_string(),
|
||||||
|
collection: coll.to_string(),
|
||||||
|
header_valid: false,
|
||||||
|
total_records: 0,
|
||||||
|
live_documents: 0,
|
||||||
|
tombstones: 0,
|
||||||
|
duplicate_ids: Vec::new(),
|
||||||
|
checksum_errors: 0,
|
||||||
|
decode_errors: 0,
|
||||||
|
data_file_size: 0,
|
||||||
|
hint_file_exists: hint_path.exists(),
|
||||||
|
orphaned_hint_entries: 0,
|
||||||
|
errors: Vec::new(),
|
||||||
|
};
|
||||||
|
|
||||||
|
// Get file size
|
||||||
|
match std::fs::metadata(&data_path) {
|
||||||
|
Ok(m) => report.data_file_size = m.len(),
|
||||||
|
Err(e) => {
|
||||||
|
report.errors.push(format!("cannot stat data.rdb: {e}"));
|
||||||
|
return report;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Open and validate header
|
||||||
|
let file = match std::fs::File::open(&data_path) {
|
||||||
|
Ok(f) => f,
|
||||||
|
Err(e) => {
|
||||||
|
report.errors.push(format!("cannot open data.rdb: {e}"));
|
||||||
|
return report;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let mut reader = BufReader::new(file);
|
||||||
|
|
||||||
|
let mut hdr_buf = [0u8; FILE_HEADER_SIZE];
|
||||||
|
if let Err(e) = reader.read_exact(&mut hdr_buf) {
|
||||||
|
report.errors.push(format!("cannot read header: {e}"));
|
||||||
|
return report;
|
||||||
|
}
|
||||||
|
|
||||||
|
match FileHeader::decode(&hdr_buf) {
|
||||||
|
Ok(hdr) => {
|
||||||
|
if hdr.file_type != FileType::Data {
|
||||||
|
report.errors.push(format!(
|
||||||
|
"wrong file type: expected Data, got {:?}",
|
||||||
|
hdr.file_type
|
||||||
|
));
|
||||||
|
} else {
|
||||||
|
report.header_valid = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
report.errors.push(format!("invalid header: {e}"));
|
||||||
|
return report;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Scan all records
|
||||||
|
let mut id_counts: HashMap<String, u64> = HashMap::new();
|
||||||
|
let mut live_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
|
||||||
|
let scanner = RecordScanner::new(reader, FILE_HEADER_SIZE as u64);
|
||||||
|
|
||||||
|
for result in scanner {
|
||||||
|
match result {
|
||||||
|
Ok((_offset, record)) => {
|
||||||
|
report.total_records += 1;
|
||||||
|
let key = String::from_utf8_lossy(&record.key).to_string();
|
||||||
|
|
||||||
|
if record.is_tombstone() {
|
||||||
|
report.tombstones += 1;
|
||||||
|
live_ids.remove(&key);
|
||||||
|
} else {
|
||||||
|
*id_counts.entry(key.clone()).or_insert(0) += 1;
|
||||||
|
live_ids.insert(key);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
let err_str = e.to_string();
|
||||||
|
if err_str.contains("checksum") || err_str.contains("Checksum") {
|
||||||
|
report.checksum_errors += 1;
|
||||||
|
} else {
|
||||||
|
report.decode_errors += 1;
|
||||||
|
}
|
||||||
|
// Cannot continue scanning after a decode error — the stream position is lost
|
||||||
|
report.errors.push(format!("record decode error: {e}"));
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
report.live_documents = live_ids.len() as u64;
|
||||||
|
|
||||||
|
// Find duplicates (keys that appeared more than once as live inserts)
|
||||||
|
for (id, count) in &id_counts {
|
||||||
|
if *count > 1 {
|
||||||
|
report.duplicate_ids.push(id.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
report.duplicate_ids.sort();
|
||||||
|
|
||||||
|
// Validate hint file if present
|
||||||
|
if hint_path.exists() {
|
||||||
|
match KeyDir::load_from_hint_file(&hint_path) {
|
||||||
|
Ok(Some((hint_kd, stored_size))) => {
|
||||||
|
if stored_size > 0 && stored_size != report.data_file_size {
|
||||||
|
report.errors.push(format!(
|
||||||
|
"hint file is stale: recorded data size {} but actual is {}",
|
||||||
|
stored_size, report.data_file_size
|
||||||
|
));
|
||||||
|
}
|
||||||
|
// Check for orphaned entries: keys in hint but not live in data
|
||||||
|
hint_kd.for_each(|key, _entry| {
|
||||||
|
if !live_ids.contains(key) {
|
||||||
|
report.orphaned_hint_entries += 1;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// Also check if hint references offsets beyond file size
|
||||||
|
hint_kd.for_each(|_key, entry| {
|
||||||
|
if entry.offset + entry.record_len as u64 > report.data_file_size {
|
||||||
|
report.orphaned_hint_entries += 1;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Ok(None) => {
|
||||||
|
// File existed but was empty or unreadable
|
||||||
|
report.errors.push("hint file exists but is empty".into());
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
report.errors.push(format!("hint file decode error: {e}"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
report
|
||||||
|
}
|
||||||
@@ -0,0 +1,17 @@
|
|||||||
|
[package]
|
||||||
|
name = "rustdb-txn"
|
||||||
|
version.workspace = true
|
||||||
|
edition.workspace = true
|
||||||
|
license.workspace = true
|
||||||
|
authors.workspace = true
|
||||||
|
description = "MongoDB-compatible transaction and session management with snapshot isolation for RustDb"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
bson = { workspace = true }
|
||||||
|
dashmap = { workspace = true }
|
||||||
|
tokio = { workspace = true }
|
||||||
|
thiserror = { workspace = true }
|
||||||
|
tracing = { workspace = true }
|
||||||
|
uuid = { workspace = true }
|
||||||
|
rustdb-storage = { workspace = true }
|
||||||
|
async-trait = { workspace = true }
|
||||||
@@ -0,0 +1,35 @@
|
|||||||
|
use thiserror::Error;
|
||||||
|
|
||||||
|
/// Errors that can occur during transaction or session operations.
|
||||||
|
#[derive(Debug, Error)]
|
||||||
|
pub enum TransactionError {
|
||||||
|
#[error("not found: {0}")]
|
||||||
|
NotFound(String),
|
||||||
|
|
||||||
|
#[error("transaction already active for session: {0}")]
|
||||||
|
AlreadyActive(String),
|
||||||
|
|
||||||
|
#[error("write conflict detected (code 112): {0}")]
|
||||||
|
WriteConflict(String),
|
||||||
|
|
||||||
|
#[error("session expired: {0}")]
|
||||||
|
SessionExpired(String),
|
||||||
|
|
||||||
|
#[error("invalid transaction state: {0}")]
|
||||||
|
InvalidState(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TransactionError {
|
||||||
|
/// Returns the error code.
|
||||||
|
pub fn code(&self) -> i32 {
|
||||||
|
match self {
|
||||||
|
TransactionError::NotFound(_) => 251,
|
||||||
|
TransactionError::AlreadyActive(_) => 256,
|
||||||
|
TransactionError::WriteConflict(_) => 112,
|
||||||
|
TransactionError::SessionExpired(_) => 6100,
|
||||||
|
TransactionError::InvalidState(_) => 263,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub type TransactionResult<T> = Result<T, TransactionError>;
|
||||||
@@ -0,0 +1,9 @@
|
|||||||
|
pub mod error;
|
||||||
|
mod session;
|
||||||
|
mod transaction;
|
||||||
|
|
||||||
|
pub use error::{TransactionError, TransactionResult};
|
||||||
|
pub use session::{Session, SessionEngine};
|
||||||
|
pub use transaction::{
|
||||||
|
TransactionEngine, TransactionState, TransactionStatus, WriteEntry, WriteOp,
|
||||||
|
};
|
||||||
@@ -0,0 +1,215 @@
|
|||||||
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
|
use bson::Bson;
|
||||||
|
use dashmap::DashMap;
|
||||||
|
use tracing::{debug, warn};
|
||||||
|
|
||||||
|
use crate::error::{TransactionError, TransactionResult};
|
||||||
|
|
||||||
|
/// Represents a logical session.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct Session {
|
||||||
|
pub id: String,
|
||||||
|
pub created_at: Instant,
|
||||||
|
pub last_activity_at: Instant,
|
||||||
|
pub txn_id: Option<String>,
|
||||||
|
pub in_transaction: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Engine that manages logical sessions with timeout and cleanup.
|
||||||
|
pub struct SessionEngine {
|
||||||
|
sessions: DashMap<String, Session>,
|
||||||
|
timeout: Duration,
|
||||||
|
_cleanup_interval: Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SessionEngine {
|
||||||
|
/// Create a new session engine.
|
||||||
|
///
|
||||||
|
/// * `timeout_ms` - Session timeout in milliseconds (default: 30 minutes = 1_800_000).
|
||||||
|
/// * `cleanup_interval_ms` - How often to run the cleanup task in milliseconds (default: 60_000).
|
||||||
|
pub fn new(timeout_ms: u64, cleanup_interval_ms: u64) -> Self {
|
||||||
|
Self {
|
||||||
|
sessions: DashMap::new(),
|
||||||
|
timeout: Duration::from_millis(timeout_ms),
|
||||||
|
_cleanup_interval: Duration::from_millis(cleanup_interval_ms),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get an existing session or create a new one. Returns the session id.
|
||||||
|
pub fn get_or_create_session(&self, id: &str) -> String {
|
||||||
|
if let Some(mut session) = self.sessions.get_mut(id) {
|
||||||
|
session.last_activity_at = Instant::now();
|
||||||
|
return session.id.clone();
|
||||||
|
}
|
||||||
|
|
||||||
|
let now = Instant::now();
|
||||||
|
let session = Session {
|
||||||
|
id: id.to_string(),
|
||||||
|
created_at: now,
|
||||||
|
last_activity_at: now,
|
||||||
|
txn_id: None,
|
||||||
|
in_transaction: false,
|
||||||
|
};
|
||||||
|
self.sessions.insert(id.to_string(), session);
|
||||||
|
debug!(session_id = %id, "created new session");
|
||||||
|
id.to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Update the last activity timestamp for a session.
|
||||||
|
pub fn touch_session(&self, id: &str) {
|
||||||
|
if let Some(mut session) = self.sessions.get_mut(id) {
|
||||||
|
session.last_activity_at = Instant::now();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// End a session. If a transaction is active, it will be marked for abort.
|
||||||
|
pub fn end_session(&self, id: &str) {
|
||||||
|
if let Some((_, session)) = self.sessions.remove(id) {
|
||||||
|
if session.in_transaction {
|
||||||
|
warn!(
|
||||||
|
session_id = %id,
|
||||||
|
txn_id = ?session.txn_id,
|
||||||
|
"ending session with active transaction, transaction should be aborted"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
debug!(session_id = %id, "session ended");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Associate a transaction with a session.
|
||||||
|
pub fn start_transaction(&self, session_id: &str, txn_id: &str) -> TransactionResult<()> {
|
||||||
|
let mut session = self
|
||||||
|
.sessions
|
||||||
|
.get_mut(session_id)
|
||||||
|
.ok_or_else(|| TransactionError::NotFound(format!("session {}", session_id)))?;
|
||||||
|
|
||||||
|
if session.in_transaction {
|
||||||
|
return Err(TransactionError::AlreadyActive(session_id.to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
|
session.txn_id = Some(txn_id.to_string());
|
||||||
|
session.in_transaction = true;
|
||||||
|
session.last_activity_at = Instant::now();
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Disassociate the transaction from a session (after commit or abort).
|
||||||
|
pub fn end_transaction(&self, session_id: &str) {
|
||||||
|
if let Some(mut session) = self.sessions.get_mut(session_id) {
|
||||||
|
session.txn_id = None;
|
||||||
|
session.in_transaction = false;
|
||||||
|
session.last_activity_at = Instant::now();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check whether a session is currently in a transaction.
|
||||||
|
pub fn is_in_transaction(&self, session_id: &str) -> bool {
|
||||||
|
self.sessions
|
||||||
|
.get(session_id)
|
||||||
|
.map(|s| s.in_transaction)
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the active transaction id for a session, if any.
|
||||||
|
pub fn get_transaction_id(&self, session_id: &str) -> Option<String> {
|
||||||
|
self.sessions
|
||||||
|
.get(session_id)
|
||||||
|
.and_then(|s| s.txn_id.clone())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extract a session id from a BSON `lsid` value.
|
||||||
|
///
|
||||||
|
/// Handles the following formats:
|
||||||
|
/// - `{ "id": UUID }` (standard driver format)
|
||||||
|
/// - `{ "id": "string" }` (string shorthand)
|
||||||
|
/// - `{ "id": Binary(base64) }` (binary UUID)
|
||||||
|
pub fn extract_session_id(lsid: &Bson) -> Option<String> {
|
||||||
|
match lsid {
|
||||||
|
Bson::Document(doc) => {
|
||||||
|
if let Some(id_val) = doc.get("id") {
|
||||||
|
match id_val {
|
||||||
|
Bson::Binary(bin) => {
|
||||||
|
// UUID stored as Binary subtype 4.
|
||||||
|
let bytes = &bin.bytes;
|
||||||
|
if bytes.len() == 16 {
|
||||||
|
let uuid = uuid::Uuid::from_slice(bytes).ok()?;
|
||||||
|
Some(uuid.to_string())
|
||||||
|
} else {
|
||||||
|
// Fall back to base64 representation.
|
||||||
|
Some(base64_encode(bytes))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Bson::String(s) => Some(s.clone()),
|
||||||
|
_ => Some(format!("{}", id_val)),
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Bson::String(s) => Some(s.clone()),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Clean up expired sessions. Returns the number of sessions removed.
|
||||||
|
pub fn cleanup_expired(&self) -> usize {
|
||||||
|
let now = Instant::now();
|
||||||
|
let timeout = self.timeout;
|
||||||
|
let expired: Vec<String> = self
|
||||||
|
.sessions
|
||||||
|
.iter()
|
||||||
|
.filter(|entry| now.duration_since(entry.last_activity_at) > timeout)
|
||||||
|
.map(|entry| entry.id.clone())
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let count = expired.len();
|
||||||
|
for id in &expired {
|
||||||
|
debug!(session_id = %id, "cleaning up expired session");
|
||||||
|
self.sessions.remove(id);
|
||||||
|
}
|
||||||
|
count
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Number of currently tracked logical sessions.
|
||||||
|
pub fn len(&self) -> usize {
|
||||||
|
self.sessions.len()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Whether there are no tracked logical sessions.
|
||||||
|
pub fn is_empty(&self) -> bool {
|
||||||
|
self.sessions.is_empty()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for SessionEngine {
|
||||||
|
fn default() -> Self {
|
||||||
|
// 30 minutes timeout, 60 seconds cleanup interval.
|
||||||
|
Self::new(1_800_000, 60_000)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Simple base64 encoding for binary data (no external dependency needed).
|
||||||
|
fn base64_encode(data: &[u8]) -> String {
|
||||||
|
const CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
|
||||||
|
let mut result = String::with_capacity((data.len() + 2) / 3 * 4);
|
||||||
|
for chunk in data.chunks(3) {
|
||||||
|
let b0 = chunk[0] as u32;
|
||||||
|
let b1 = if chunk.len() > 1 { chunk[1] as u32 } else { 0 };
|
||||||
|
let b2 = if chunk.len() > 2 { chunk[2] as u32 } else { 0 };
|
||||||
|
let triple = (b0 << 16) | (b1 << 8) | b2;
|
||||||
|
result.push(CHARS[((triple >> 18) & 0x3F) as usize] as char);
|
||||||
|
result.push(CHARS[((triple >> 12) & 0x3F) as usize] as char);
|
||||||
|
if chunk.len() > 1 {
|
||||||
|
result.push(CHARS[((triple >> 6) & 0x3F) as usize] as char);
|
||||||
|
} else {
|
||||||
|
result.push('=');
|
||||||
|
}
|
||||||
|
if chunk.len() > 2 {
|
||||||
|
result.push(CHARS[(triple & 0x3F) as usize] as char);
|
||||||
|
} else {
|
||||||
|
result.push('=');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
result
|
||||||
|
}
|
||||||
@@ -0,0 +1,372 @@
|
|||||||
|
use std::collections::{HashMap, HashSet};
|
||||||
|
|
||||||
|
use bson::Document;
|
||||||
|
use dashmap::DashMap;
|
||||||
|
use tracing::{debug, warn};
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
use rustdb_storage::StorageAdapter;
|
||||||
|
|
||||||
|
use crate::error::{TransactionError, TransactionResult};
|
||||||
|
|
||||||
|
/// The status of a transaction.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub enum TransactionStatus {
|
||||||
|
Active,
|
||||||
|
Committed,
|
||||||
|
Aborted,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Describes a write operation within a transaction.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub enum WriteOp {
|
||||||
|
Insert,
|
||||||
|
Update,
|
||||||
|
Delete,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A single write entry recorded within a transaction.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct WriteEntry {
|
||||||
|
pub op: WriteOp,
|
||||||
|
pub doc: Option<Document>,
|
||||||
|
pub original_doc: Option<Document>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Full state of an in-flight transaction.
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct TransactionState {
|
||||||
|
pub id: String,
|
||||||
|
pub session_id: String,
|
||||||
|
pub status: TransactionStatus,
|
||||||
|
/// Tracks which documents were read: namespace -> set of doc ids.
|
||||||
|
pub read_set: HashMap<String, HashSet<String>>,
|
||||||
|
/// Tracks writes: namespace -> (doc_id -> WriteEntry).
|
||||||
|
pub write_set: HashMap<String, HashMap<String, WriteEntry>>,
|
||||||
|
/// Snapshot of collections at transaction start: namespace -> documents.
|
||||||
|
pub snapshots: HashMap<String, Vec<Document>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Engine that manages transaction lifecycle and conflict detection.
|
||||||
|
pub struct TransactionEngine {
|
||||||
|
transactions: DashMap<String, TransactionState>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TransactionEngine {
|
||||||
|
/// Create a new transaction engine.
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
transactions: DashMap::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Start a new transaction for the given session.
|
||||||
|
/// Returns a unique transaction id (UUID v4).
|
||||||
|
pub fn start_transaction(&self, session_id: &str) -> TransactionResult<String> {
|
||||||
|
let txn_id = Uuid::new_v4().to_string();
|
||||||
|
debug!(txn_id = %txn_id, session_id = %session_id, "starting transaction");
|
||||||
|
|
||||||
|
let state = TransactionState {
|
||||||
|
id: txn_id.clone(),
|
||||||
|
session_id: session_id.to_string(),
|
||||||
|
status: TransactionStatus::Active,
|
||||||
|
read_set: HashMap::new(),
|
||||||
|
write_set: HashMap::new(),
|
||||||
|
snapshots: HashMap::new(),
|
||||||
|
};
|
||||||
|
|
||||||
|
self.transactions.insert(txn_id.clone(), state);
|
||||||
|
Ok(txn_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Commit a transaction: check for conflicts, then apply buffered writes
|
||||||
|
/// to the underlying storage adapter.
|
||||||
|
pub async fn commit_transaction(
|
||||||
|
&self,
|
||||||
|
txn_id: &str,
|
||||||
|
storage: &dyn StorageAdapter,
|
||||||
|
) -> TransactionResult<()> {
|
||||||
|
// Remove the transaction so we own it exclusively.
|
||||||
|
let mut state = self
|
||||||
|
.transactions
|
||||||
|
.remove(txn_id)
|
||||||
|
.map(|(_, s)| s)
|
||||||
|
.ok_or_else(|| TransactionError::NotFound(txn_id.to_string()))?;
|
||||||
|
|
||||||
|
if state.status != TransactionStatus::Active {
|
||||||
|
return Err(TransactionError::InvalidState(format!(
|
||||||
|
"transaction {} is {:?}, cannot commit",
|
||||||
|
txn_id, state.status
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Conflict detection: check if any documents in the read set have
|
||||||
|
// been modified since the snapshot was taken.
|
||||||
|
// (Simplified: we skip real snapshot timestamps for now.)
|
||||||
|
|
||||||
|
// Apply buffered writes to storage.
|
||||||
|
for (ns, writes) in &state.write_set {
|
||||||
|
let parts: Vec<&str> = ns.splitn(2, '.').collect();
|
||||||
|
if parts.len() != 2 {
|
||||||
|
warn!(namespace = %ns, "invalid namespace format, skipping");
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let (db, coll) = (parts[0], parts[1]);
|
||||||
|
|
||||||
|
for (doc_id, entry) in writes {
|
||||||
|
match entry.op {
|
||||||
|
WriteOp::Insert => {
|
||||||
|
if let Some(ref doc) = entry.doc {
|
||||||
|
let _ = storage.insert_one(db, coll, doc.clone()).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
WriteOp::Update => {
|
||||||
|
if let Some(ref doc) = entry.doc {
|
||||||
|
let _ = storage.update_by_id(db, coll, doc_id, doc.clone()).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
WriteOp::Delete => {
|
||||||
|
let _ = storage.delete_by_id(db, coll, doc_id).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
state.status = TransactionStatus::Committed;
|
||||||
|
debug!(txn_id = %txn_id, "transaction committed");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Remove an active transaction and return its buffered state for an
|
||||||
|
/// external committer that needs to update secondary indexes and oplogs.
|
||||||
|
pub fn take_transaction(&self, txn_id: &str) -> TransactionResult<TransactionState> {
|
||||||
|
let state = self
|
||||||
|
.transactions
|
||||||
|
.remove(txn_id)
|
||||||
|
.map(|(_, s)| s)
|
||||||
|
.ok_or_else(|| TransactionError::NotFound(txn_id.to_string()))?;
|
||||||
|
|
||||||
|
if state.status != TransactionStatus::Active {
|
||||||
|
return Err(TransactionError::InvalidState(format!(
|
||||||
|
"transaction {} is {:?}, cannot commit",
|
||||||
|
txn_id, state.status
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(state)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Abort a transaction, discarding all buffered writes.
|
||||||
|
pub fn abort_transaction(&self, txn_id: &str) -> TransactionResult<()> {
|
||||||
|
let mut state = self
|
||||||
|
.transactions
|
||||||
|
.get_mut(txn_id)
|
||||||
|
.ok_or_else(|| TransactionError::NotFound(txn_id.to_string()))?;
|
||||||
|
|
||||||
|
if state.status != TransactionStatus::Active {
|
||||||
|
return Err(TransactionError::InvalidState(format!(
|
||||||
|
"transaction {} is {:?}, cannot abort",
|
||||||
|
txn_id, state.status
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
state.status = TransactionStatus::Aborted;
|
||||||
|
debug!(txn_id = %txn_id, "transaction aborted");
|
||||||
|
|
||||||
|
// Drop the mutable ref before removing.
|
||||||
|
drop(state);
|
||||||
|
self.transactions.remove(txn_id);
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check whether a transaction is currently active.
|
||||||
|
pub fn is_active(&self, txn_id: &str) -> bool {
|
||||||
|
self.transactions
|
||||||
|
.get(txn_id)
|
||||||
|
.map(|s| s.status == TransactionStatus::Active)
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Record a document read within a transaction (for conflict detection).
|
||||||
|
pub fn record_read(&self, txn_id: &str, ns: &str, doc_id: &str) {
|
||||||
|
if let Some(mut state) = self.transactions.get_mut(txn_id) {
|
||||||
|
state
|
||||||
|
.read_set
|
||||||
|
.entry(ns.to_string())
|
||||||
|
.or_default()
|
||||||
|
.insert(doc_id.to_string());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Record a document write within a transaction (buffered until commit).
|
||||||
|
pub fn record_write(
|
||||||
|
&self,
|
||||||
|
txn_id: &str,
|
||||||
|
ns: &str,
|
||||||
|
doc_id: &str,
|
||||||
|
op: WriteOp,
|
||||||
|
doc: Option<Document>,
|
||||||
|
original: Option<Document>,
|
||||||
|
) {
|
||||||
|
if let Some(mut state) = self.transactions.get_mut(txn_id) {
|
||||||
|
let writes = state.write_set.entry(ns.to_string()).or_default();
|
||||||
|
if let Some(existing) = writes.remove(doc_id) {
|
||||||
|
if let Some(merged) = merge_write_entry(existing, op, doc, original) {
|
||||||
|
writes.insert(doc_id.to_string(), merged);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
writes.insert(
|
||||||
|
doc_id.to_string(),
|
||||||
|
WriteEntry {
|
||||||
|
op,
|
||||||
|
doc,
|
||||||
|
original_doc: original,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Return true if the transaction already has a base snapshot for a namespace.
|
||||||
|
pub fn has_snapshot(&self, txn_id: &str, ns: &str) -> bool {
|
||||||
|
self.transactions
|
||||||
|
.get(txn_id)
|
||||||
|
.map(|state| state.snapshots.contains_key(ns))
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get a snapshot of documents for a namespace within a transaction,
|
||||||
|
/// applying the write overlay (inserts, updates, deletes) on top.
|
||||||
|
pub fn get_snapshot(&self, txn_id: &str, ns: &str) -> Option<Vec<Document>> {
|
||||||
|
let state = self.transactions.get(txn_id)?;
|
||||||
|
|
||||||
|
// Start with the base snapshot.
|
||||||
|
let mut docs: Vec<Document> = state
|
||||||
|
.snapshots
|
||||||
|
.get(ns)
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
// Apply write overlay.
|
||||||
|
if let Some(writes) = state.write_set.get(ns) {
|
||||||
|
// Collect ids to delete.
|
||||||
|
let delete_ids: HashSet<&String> = writes
|
||||||
|
.iter()
|
||||||
|
.filter(|(_, e)| e.op == WriteOp::Delete)
|
||||||
|
.map(|(id, _)| id)
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// Remove deleted docs.
|
||||||
|
docs.retain(|d| {
|
||||||
|
if let Some(id) = d.get_object_id("_id").ok().map(|oid| oid.to_hex()) {
|
||||||
|
!delete_ids.contains(&id)
|
||||||
|
} else {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// Apply updates.
|
||||||
|
for (doc_id, entry) in writes {
|
||||||
|
if entry.op == WriteOp::Update {
|
||||||
|
if let Some(ref new_doc) = entry.doc {
|
||||||
|
// Replace existing doc with updated version.
|
||||||
|
let hex_id = doc_id.clone();
|
||||||
|
if let Some(pos) = docs.iter().position(|d| {
|
||||||
|
d.get_object_id("_id")
|
||||||
|
.ok()
|
||||||
|
.map(|oid| oid.to_hex()) == Some(hex_id.clone())
|
||||||
|
}) {
|
||||||
|
docs[pos] = new_doc.clone();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Apply inserts.
|
||||||
|
for (_doc_id, entry) in writes {
|
||||||
|
if entry.op == WriteOp::Insert {
|
||||||
|
if let Some(ref doc) = entry.doc {
|
||||||
|
docs.push(doc.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Some(docs)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Store a base snapshot for a namespace within a transaction.
|
||||||
|
pub fn set_snapshot(&self, txn_id: &str, ns: &str, docs: Vec<Document>) {
|
||||||
|
if let Some(mut state) = self.transactions.get_mut(txn_id) {
|
||||||
|
state.snapshots.insert(ns.to_string(), docs);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Number of currently active transactions.
|
||||||
|
pub fn len(&self) -> usize {
|
||||||
|
self.transactions.len()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Whether there are no active transactions.
|
||||||
|
pub fn is_empty(&self) -> bool {
|
||||||
|
self.transactions.is_empty()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn merge_write_entry(
|
||||||
|
existing: WriteEntry,
|
||||||
|
next_op: WriteOp,
|
||||||
|
next_doc: Option<Document>,
|
||||||
|
next_original: Option<Document>,
|
||||||
|
) -> Option<WriteEntry> {
|
||||||
|
match (existing.op, next_op) {
|
||||||
|
(WriteOp::Insert, WriteOp::Update) => Some(WriteEntry {
|
||||||
|
op: WriteOp::Insert,
|
||||||
|
doc: next_doc,
|
||||||
|
original_doc: None,
|
||||||
|
}),
|
||||||
|
(WriteOp::Insert, WriteOp::Delete) => None,
|
||||||
|
(WriteOp::Insert, WriteOp::Insert) => Some(WriteEntry {
|
||||||
|
op: WriteOp::Insert,
|
||||||
|
doc: next_doc,
|
||||||
|
original_doc: None,
|
||||||
|
}),
|
||||||
|
(WriteOp::Update, WriteOp::Update) => Some(WriteEntry {
|
||||||
|
op: WriteOp::Update,
|
||||||
|
doc: next_doc,
|
||||||
|
original_doc: existing.original_doc,
|
||||||
|
}),
|
||||||
|
(WriteOp::Update, WriteOp::Delete) => Some(WriteEntry {
|
||||||
|
op: WriteOp::Delete,
|
||||||
|
doc: None,
|
||||||
|
original_doc: existing.original_doc,
|
||||||
|
}),
|
||||||
|
(WriteOp::Update, WriteOp::Insert) => Some(WriteEntry {
|
||||||
|
op: WriteOp::Update,
|
||||||
|
doc: next_doc,
|
||||||
|
original_doc: existing.original_doc,
|
||||||
|
}),
|
||||||
|
(WriteOp::Delete, WriteOp::Insert) => Some(WriteEntry {
|
||||||
|
op: if existing.original_doc.is_some() {
|
||||||
|
WriteOp::Update
|
||||||
|
} else {
|
||||||
|
WriteOp::Insert
|
||||||
|
},
|
||||||
|
doc: next_doc,
|
||||||
|
original_doc: existing.original_doc,
|
||||||
|
}),
|
||||||
|
(WriteOp::Delete, WriteOp::Update) => Some(WriteEntry {
|
||||||
|
op: WriteOp::Update,
|
||||||
|
doc: next_doc,
|
||||||
|
original_doc: existing.original_doc.or(next_original),
|
||||||
|
}),
|
||||||
|
(WriteOp::Delete, WriteOp::Delete) => Some(existing),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for TransactionEngine {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,18 @@
|
|||||||
|
[package]
|
||||||
|
name = "rustdb-wire"
|
||||||
|
version.workspace = true
|
||||||
|
edition.workspace = true
|
||||||
|
license.workspace = true
|
||||||
|
authors.workspace = true
|
||||||
|
description = "MongoDB-compatible wire protocol parser and encoder for RustDb"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
bson = { workspace = true }
|
||||||
|
bytes = { workspace = true }
|
||||||
|
tokio-util = { workspace = true }
|
||||||
|
thiserror = { workspace = true }
|
||||||
|
tracing = { workspace = true }
|
||||||
|
crc32fast = { workspace = true }
|
||||||
|
|
||||||
|
[dev-dependencies]
|
||||||
|
tokio = { workspace = true }
|
||||||
@@ -0,0 +1,49 @@
|
|||||||
|
use bytes::{Buf, BytesMut};
|
||||||
|
use tokio_util::codec::{Decoder, Encoder};
|
||||||
|
|
||||||
|
use crate::error::WireError;
|
||||||
|
use crate::parser::{parse_message, ParsedCommand};
|
||||||
|
|
||||||
|
/// Tokio codec for framing wire protocol messages on a TCP stream.
|
||||||
|
///
|
||||||
|
/// The wire protocol is naturally length-prefixed:
|
||||||
|
/// the first 4 bytes of each message contain the total message length.
|
||||||
|
pub struct WireCodec;
|
||||||
|
|
||||||
|
impl Decoder for WireCodec {
|
||||||
|
type Item = ParsedCommand;
|
||||||
|
type Error = WireError;
|
||||||
|
|
||||||
|
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
|
||||||
|
if src.len() < 4 {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Peek at message length
|
||||||
|
let msg_len = i32::from_le_bytes([src[0], src[1], src[2], src[3]]) as usize;
|
||||||
|
|
||||||
|
if src.len() < msg_len {
|
||||||
|
// Reserve space for the rest of the message
|
||||||
|
src.reserve(msg_len - src.len());
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
|
||||||
|
match parse_message(src)? {
|
||||||
|
Some((cmd, bytes_consumed)) => {
|
||||||
|
src.advance(bytes_consumed);
|
||||||
|
Ok(Some(cmd))
|
||||||
|
}
|
||||||
|
None => Ok(None),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Encoder for raw byte responses (already serialized by the command handlers).
|
||||||
|
impl Encoder<Vec<u8>> for WireCodec {
|
||||||
|
type Error = WireError;
|
||||||
|
|
||||||
|
fn encode(&mut self, item: Vec<u8>, dst: &mut BytesMut) -> Result<(), Self::Error> {
|
||||||
|
dst.extend_from_slice(&item);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,142 @@
|
|||||||
|
use bson::Document;
|
||||||
|
|
||||||
|
use crate::opcodes::*;
|
||||||
|
|
||||||
|
/// Encode an OP_MSG response.
|
||||||
|
pub fn encode_op_msg_response(
|
||||||
|
response_to: i32,
|
||||||
|
response: &Document,
|
||||||
|
request_id: i32,
|
||||||
|
) -> Vec<u8> {
|
||||||
|
let body_bson = bson::to_vec(response).expect("failed to serialize BSON response");
|
||||||
|
|
||||||
|
// Header (16) + flagBits (4) + section type (1) + body BSON
|
||||||
|
let message_length = 16 + 4 + 1 + body_bson.len();
|
||||||
|
|
||||||
|
let mut buf = Vec::with_capacity(message_length);
|
||||||
|
|
||||||
|
// Header
|
||||||
|
buf.extend_from_slice(&(message_length as i32).to_le_bytes());
|
||||||
|
buf.extend_from_slice(&request_id.to_le_bytes());
|
||||||
|
buf.extend_from_slice(&response_to.to_le_bytes());
|
||||||
|
buf.extend_from_slice(&OP_MSG.to_le_bytes());
|
||||||
|
|
||||||
|
// Flag bits (0 = no flags)
|
||||||
|
buf.extend_from_slice(&0u32.to_le_bytes());
|
||||||
|
|
||||||
|
// Section type 0 (body)
|
||||||
|
buf.push(SECTION_BODY);
|
||||||
|
|
||||||
|
// Body BSON
|
||||||
|
buf.extend_from_slice(&body_bson);
|
||||||
|
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Encode an OP_REPLY response (legacy, for OP_QUERY responses).
|
||||||
|
pub fn encode_op_reply_response(
|
||||||
|
response_to: i32,
|
||||||
|
documents: &[Document],
|
||||||
|
request_id: i32,
|
||||||
|
cursor_id: i64,
|
||||||
|
) -> Vec<u8> {
|
||||||
|
let doc_buffers: Vec<Vec<u8>> = documents
|
||||||
|
.iter()
|
||||||
|
.map(|doc| bson::to_vec(doc).expect("failed to serialize BSON document"))
|
||||||
|
.collect();
|
||||||
|
let total_docs_size: usize = doc_buffers.iter().map(|b| b.len()).sum();
|
||||||
|
|
||||||
|
// Header (16) + responseFlags (4) + cursorID (8) + startingFrom (4) + numberReturned (4) + docs
|
||||||
|
let message_length = 16 + 4 + 8 + 4 + 4 + total_docs_size;
|
||||||
|
|
||||||
|
let mut buf = Vec::with_capacity(message_length);
|
||||||
|
|
||||||
|
// Header
|
||||||
|
buf.extend_from_slice(&(message_length as i32).to_le_bytes());
|
||||||
|
buf.extend_from_slice(&request_id.to_le_bytes());
|
||||||
|
buf.extend_from_slice(&response_to.to_le_bytes());
|
||||||
|
buf.extend_from_slice(&OP_REPLY.to_le_bytes());
|
||||||
|
|
||||||
|
// OP_REPLY fields
|
||||||
|
buf.extend_from_slice(&0i32.to_le_bytes()); // responseFlags
|
||||||
|
buf.extend_from_slice(&cursor_id.to_le_bytes()); // cursorID
|
||||||
|
buf.extend_from_slice(&0i32.to_le_bytes()); // startingFrom
|
||||||
|
buf.extend_from_slice(&(documents.len() as i32).to_le_bytes()); // numberReturned
|
||||||
|
|
||||||
|
// Documents
|
||||||
|
for doc_buf in &doc_buffers {
|
||||||
|
buf.extend_from_slice(doc_buf);
|
||||||
|
}
|
||||||
|
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Encode an error response as OP_MSG.
|
||||||
|
pub fn encode_error_response(
|
||||||
|
response_to: i32,
|
||||||
|
error_code: i32,
|
||||||
|
error_message: &str,
|
||||||
|
request_id: i32,
|
||||||
|
) -> Vec<u8> {
|
||||||
|
let response = bson::doc! {
|
||||||
|
"ok": 0,
|
||||||
|
"errmsg": error_message,
|
||||||
|
"code": error_code,
|
||||||
|
"codeName": error_code_name(error_code),
|
||||||
|
};
|
||||||
|
encode_op_msg_response(response_to, &response, request_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Map error codes to their code names.
|
||||||
|
pub fn error_code_name(code: i32) -> &'static str {
|
||||||
|
match code {
|
||||||
|
0 => "OK",
|
||||||
|
1 => "InternalError",
|
||||||
|
2 => "BadValue",
|
||||||
|
13 => "Unauthorized",
|
||||||
|
26 => "NamespaceNotFound",
|
||||||
|
27 => "IndexNotFound",
|
||||||
|
48 => "NamespaceExists",
|
||||||
|
59 => "CommandNotFound",
|
||||||
|
66 => "ImmutableField",
|
||||||
|
73 => "InvalidNamespace",
|
||||||
|
85 => "IndexOptionsConflict",
|
||||||
|
112 => "WriteConflict",
|
||||||
|
121 => "DocumentValidationFailure",
|
||||||
|
211 => "KeyNotFound",
|
||||||
|
251 => "NoSuchTransaction",
|
||||||
|
11000 => "DuplicateKey",
|
||||||
|
11001 => "DuplicateKeyValue",
|
||||||
|
_ => "UnknownError",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_encode_op_msg_roundtrip() {
|
||||||
|
let doc = bson::doc! { "ok": 1 };
|
||||||
|
let encoded = encode_op_msg_response(1, &doc, 2);
|
||||||
|
|
||||||
|
// Verify header
|
||||||
|
let msg_len = i32::from_le_bytes([encoded[0], encoded[1], encoded[2], encoded[3]]);
|
||||||
|
assert_eq!(msg_len as usize, encoded.len());
|
||||||
|
|
||||||
|
let op_code = i32::from_le_bytes([encoded[12], encoded[13], encoded[14], encoded[15]]);
|
||||||
|
assert_eq!(op_code, OP_MSG);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_encode_op_reply() {
|
||||||
|
let docs = vec![bson::doc! { "ok": 1 }];
|
||||||
|
let encoded = encode_op_reply_response(1, &docs, 2, 0);
|
||||||
|
|
||||||
|
let msg_len = i32::from_le_bytes([encoded[0], encoded[1], encoded[2], encoded[3]]);
|
||||||
|
assert_eq!(msg_len as usize, encoded.len());
|
||||||
|
|
||||||
|
let op_code = i32::from_le_bytes([encoded[12], encoded[13], encoded[14], encoded[15]]);
|
||||||
|
assert_eq!(op_code, OP_REPLY);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,27 @@
|
|||||||
|
/// Errors from wire protocol parsing/encoding.
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum WireError {
|
||||||
|
#[error("Incomplete message: need {needed} bytes, have {have}")]
|
||||||
|
Incomplete { needed: usize, have: usize },
|
||||||
|
|
||||||
|
#[error("Unsupported opCode: {0}")]
|
||||||
|
UnsupportedOpCode(i32),
|
||||||
|
|
||||||
|
#[error("Missing command body section in OP_MSG")]
|
||||||
|
MissingBody,
|
||||||
|
|
||||||
|
#[error("Unknown section type: {0}")]
|
||||||
|
UnknownSectionType(u8),
|
||||||
|
|
||||||
|
#[error("BSON deserialization error: {0}")]
|
||||||
|
BsonError(#[from] bson::de::Error),
|
||||||
|
|
||||||
|
#[error("BSON serialization error: {0}")]
|
||||||
|
BsonSerError(#[from] bson::ser::Error),
|
||||||
|
|
||||||
|
#[error("IO error: {0}")]
|
||||||
|
IoError(#[from] std::io::Error),
|
||||||
|
|
||||||
|
#[error("Checksum mismatch: expected {expected}, got {actual}")]
|
||||||
|
ChecksumMismatch { expected: u32, actual: u32 },
|
||||||
|
}
|
||||||
@@ -0,0 +1,11 @@
|
|||||||
|
mod codec;
|
||||||
|
mod error;
|
||||||
|
mod opcodes;
|
||||||
|
mod parser;
|
||||||
|
mod encoder;
|
||||||
|
|
||||||
|
pub use codec::WireCodec;
|
||||||
|
pub use error::WireError;
|
||||||
|
pub use opcodes::*;
|
||||||
|
pub use parser::*;
|
||||||
|
pub use encoder::*;
|
||||||
@@ -0,0 +1,19 @@
|
|||||||
|
/// Wire protocol op codes
|
||||||
|
pub const OP_REPLY: i32 = 1;
|
||||||
|
pub const OP_UPDATE: i32 = 2001;
|
||||||
|
pub const OP_INSERT: i32 = 2002;
|
||||||
|
pub const OP_QUERY: i32 = 2004;
|
||||||
|
pub const OP_GET_MORE: i32 = 2005;
|
||||||
|
pub const OP_DELETE: i32 = 2006;
|
||||||
|
pub const OP_KILL_CURSORS: i32 = 2007;
|
||||||
|
pub const OP_COMPRESSED: i32 = 2012;
|
||||||
|
pub const OP_MSG: i32 = 2013;
|
||||||
|
|
||||||
|
/// OP_MSG section types
|
||||||
|
pub const SECTION_BODY: u8 = 0;
|
||||||
|
pub const SECTION_DOCUMENT_SEQUENCE: u8 = 1;
|
||||||
|
|
||||||
|
/// OP_MSG flag bits
|
||||||
|
pub const MSG_FLAG_CHECKSUM_PRESENT: u32 = 1 << 0;
|
||||||
|
pub const MSG_FLAG_MORE_TO_COME: u32 = 1 << 1;
|
||||||
|
pub const MSG_FLAG_EXHAUST_ALLOWED: u32 = 1 << 16;
|
||||||
@@ -0,0 +1,236 @@
|
|||||||
|
use bson::Document;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
use crate::error::WireError;
|
||||||
|
use crate::opcodes::*;
|
||||||
|
|
||||||
|
/// Parsed wire protocol message header (16 bytes).
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct MessageHeader {
|
||||||
|
pub message_length: i32,
|
||||||
|
pub request_id: i32,
|
||||||
|
pub response_to: i32,
|
||||||
|
pub op_code: i32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A parsed OP_MSG section.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub enum OpMsgSection {
|
||||||
|
/// Section type 0: single BSON document body.
|
||||||
|
Body(Document),
|
||||||
|
/// Section type 1: named document sequence for bulk operations.
|
||||||
|
DocumentSequence {
|
||||||
|
identifier: String,
|
||||||
|
documents: Vec<Document>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A fully parsed command extracted from any message type.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct ParsedCommand {
|
||||||
|
pub command_name: String,
|
||||||
|
pub command: Document,
|
||||||
|
pub database: String,
|
||||||
|
pub request_id: i32,
|
||||||
|
pub op_code: i32,
|
||||||
|
/// Document sequences from OP_MSG section type 1 (e.g., "documents" for insert).
|
||||||
|
pub document_sequences: Option<HashMap<String, Vec<Document>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parse a message header from a byte slice (must be >= 16 bytes).
|
||||||
|
pub fn parse_header(buf: &[u8]) -> MessageHeader {
|
||||||
|
MessageHeader {
|
||||||
|
message_length: i32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]),
|
||||||
|
request_id: i32::from_le_bytes([buf[4], buf[5], buf[6], buf[7]]),
|
||||||
|
response_to: i32::from_le_bytes([buf[8], buf[9], buf[10], buf[11]]),
|
||||||
|
op_code: i32::from_le_bytes([buf[12], buf[13], buf[14], buf[15]]),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parse a complete message from a buffer.
|
||||||
|
/// Returns the parsed command and bytes consumed, or None if not enough data.
|
||||||
|
pub fn parse_message(buf: &[u8]) -> Result<Option<(ParsedCommand, usize)>, WireError> {
|
||||||
|
if buf.len() < 16 {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
|
||||||
|
let header = parse_header(buf);
|
||||||
|
let msg_len = header.message_length as usize;
|
||||||
|
|
||||||
|
if buf.len() < msg_len {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
|
||||||
|
let message_buf = &buf[..msg_len];
|
||||||
|
|
||||||
|
match header.op_code {
|
||||||
|
OP_MSG => parse_op_msg(message_buf, &header).map(|cmd| Some((cmd, msg_len))),
|
||||||
|
OP_QUERY => parse_op_query(message_buf, &header).map(|cmd| Some((cmd, msg_len))),
|
||||||
|
other => Err(WireError::UnsupportedOpCode(other)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parse an OP_MSG message.
|
||||||
|
fn parse_op_msg(buf: &[u8], header: &MessageHeader) -> Result<ParsedCommand, WireError> {
|
||||||
|
let mut offset = 16; // skip header
|
||||||
|
|
||||||
|
let flag_bits = u32::from_le_bytes([buf[offset], buf[offset + 1], buf[offset + 2], buf[offset + 3]]);
|
||||||
|
offset += 4;
|
||||||
|
|
||||||
|
let mut body: Option<Document> = None;
|
||||||
|
let mut document_sequences: HashMap<String, Vec<Document>> = HashMap::new();
|
||||||
|
|
||||||
|
// Parse sections until end (or checksum)
|
||||||
|
let message_end = if flag_bits & MSG_FLAG_CHECKSUM_PRESENT != 0 {
|
||||||
|
header.message_length as usize - 4
|
||||||
|
} else {
|
||||||
|
header.message_length as usize
|
||||||
|
};
|
||||||
|
|
||||||
|
while offset < message_end {
|
||||||
|
let section_type = buf[offset];
|
||||||
|
offset += 1;
|
||||||
|
|
||||||
|
match section_type {
|
||||||
|
SECTION_BODY => {
|
||||||
|
let doc_size = i32::from_le_bytes([
|
||||||
|
buf[offset], buf[offset + 1], buf[offset + 2], buf[offset + 3],
|
||||||
|
]) as usize;
|
||||||
|
let doc = bson::from_slice(&buf[offset..offset + doc_size])?;
|
||||||
|
body = Some(doc);
|
||||||
|
offset += doc_size;
|
||||||
|
}
|
||||||
|
SECTION_DOCUMENT_SEQUENCE => {
|
||||||
|
let section_size = i32::from_le_bytes([
|
||||||
|
buf[offset], buf[offset + 1], buf[offset + 2], buf[offset + 3],
|
||||||
|
]) as usize;
|
||||||
|
let section_end = offset + section_size;
|
||||||
|
offset += 4;
|
||||||
|
|
||||||
|
// Read identifier (C string, null-terminated)
|
||||||
|
let id_start = offset;
|
||||||
|
while offset < section_end && buf[offset] != 0 {
|
||||||
|
offset += 1;
|
||||||
|
}
|
||||||
|
let identifier = std::str::from_utf8(&buf[id_start..offset])
|
||||||
|
.unwrap_or("")
|
||||||
|
.to_string();
|
||||||
|
offset += 1; // skip null terminator
|
||||||
|
|
||||||
|
// Read documents
|
||||||
|
let mut documents = Vec::new();
|
||||||
|
while offset < section_end {
|
||||||
|
let doc_size = i32::from_le_bytes([
|
||||||
|
buf[offset], buf[offset + 1], buf[offset + 2], buf[offset + 3],
|
||||||
|
]) as usize;
|
||||||
|
let doc = bson::from_slice(&buf[offset..offset + doc_size])?;
|
||||||
|
documents.push(doc);
|
||||||
|
offset += doc_size;
|
||||||
|
}
|
||||||
|
|
||||||
|
document_sequences.insert(identifier, documents);
|
||||||
|
}
|
||||||
|
other => return Err(WireError::UnknownSectionType(other)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let command = body.ok_or(WireError::MissingBody)?;
|
||||||
|
let command_name = command
|
||||||
|
.keys()
|
||||||
|
.next()
|
||||||
|
.map(|s| s.to_string())
|
||||||
|
.unwrap_or_default();
|
||||||
|
let database = command
|
||||||
|
.get_str("$db")
|
||||||
|
.unwrap_or("admin")
|
||||||
|
.to_string();
|
||||||
|
|
||||||
|
Ok(ParsedCommand {
|
||||||
|
command_name,
|
||||||
|
command,
|
||||||
|
database,
|
||||||
|
request_id: header.request_id,
|
||||||
|
op_code: header.op_code,
|
||||||
|
document_sequences: if document_sequences.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(document_sequences)
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parse an OP_QUERY message (legacy, used for initial driver handshake).
|
||||||
|
fn parse_op_query(buf: &[u8], header: &MessageHeader) -> Result<ParsedCommand, WireError> {
|
||||||
|
let mut offset = 16; // skip header
|
||||||
|
|
||||||
|
let _flags = i32::from_le_bytes([buf[offset], buf[offset + 1], buf[offset + 2], buf[offset + 3]]);
|
||||||
|
offset += 4;
|
||||||
|
|
||||||
|
// Read full collection name (C string)
|
||||||
|
let name_start = offset;
|
||||||
|
while offset < buf.len() && buf[offset] != 0 {
|
||||||
|
offset += 1;
|
||||||
|
}
|
||||||
|
let full_collection_name = std::str::from_utf8(&buf[name_start..offset])
|
||||||
|
.unwrap_or("")
|
||||||
|
.to_string();
|
||||||
|
offset += 1; // skip null terminator
|
||||||
|
|
||||||
|
let _number_to_skip = i32::from_le_bytes([buf[offset], buf[offset + 1], buf[offset + 2], buf[offset + 3]]);
|
||||||
|
offset += 4;
|
||||||
|
|
||||||
|
let _number_to_return = i32::from_le_bytes([buf[offset], buf[offset + 1], buf[offset + 2], buf[offset + 3]]);
|
||||||
|
offset += 4;
|
||||||
|
|
||||||
|
// Read query document
|
||||||
|
let doc_size = i32::from_le_bytes([buf[offset], buf[offset + 1], buf[offset + 2], buf[offset + 3]]) as usize;
|
||||||
|
let query: Document = bson::from_slice(&buf[offset..offset + doc_size])?;
|
||||||
|
|
||||||
|
// Extract database from collection name (format: "dbname.$cmd")
|
||||||
|
let parts: Vec<&str> = full_collection_name.splitn(2, '.').collect();
|
||||||
|
let database = parts.first().unwrap_or(&"admin").to_string();
|
||||||
|
|
||||||
|
let mut command_name = query
|
||||||
|
.keys()
|
||||||
|
.next()
|
||||||
|
.map(|s| s.to_string())
|
||||||
|
.unwrap_or_else(|| "find".to_string());
|
||||||
|
|
||||||
|
// Map legacy isMaster/ismaster to hello
|
||||||
|
if parts.get(1) == Some(&"$cmd") {
|
||||||
|
if command_name == "isMaster" || command_name == "ismaster" {
|
||||||
|
command_name = "hello".to_string();
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
command_name = "find".to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(ParsedCommand {
|
||||||
|
command_name,
|
||||||
|
command: query,
|
||||||
|
database,
|
||||||
|
request_id: header.request_id,
|
||||||
|
op_code: header.op_code,
|
||||||
|
document_sequences: None,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_header() {
|
||||||
|
let mut buf = [0u8; 16];
|
||||||
|
buf[0..4].copy_from_slice(&100i32.to_le_bytes()); // messageLength
|
||||||
|
buf[4..8].copy_from_slice(&42i32.to_le_bytes()); // requestID
|
||||||
|
buf[8..12].copy_from_slice(&0i32.to_le_bytes()); // responseTo
|
||||||
|
buf[12..16].copy_from_slice(&OP_MSG.to_le_bytes()); // opCode
|
||||||
|
|
||||||
|
let header = parse_header(&buf);
|
||||||
|
assert_eq!(header.message_length, 100);
|
||||||
|
assert_eq!(header.request_id, 42);
|
||||||
|
assert_eq!(header.response_to, 0);
|
||||||
|
assert_eq!(header.op_code, OP_MSG);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,41 @@
|
|||||||
|
[package]
|
||||||
|
name = "rustdb"
|
||||||
|
version.workspace = true
|
||||||
|
edition.workspace = true
|
||||||
|
license.workspace = true
|
||||||
|
authors.workspace = true
|
||||||
|
description = "MongoDB-compatible embedded database server with wire protocol support"
|
||||||
|
|
||||||
|
[[bin]]
|
||||||
|
name = "rustdb"
|
||||||
|
path = "src/main.rs"
|
||||||
|
|
||||||
|
[lib]
|
||||||
|
name = "rustdb"
|
||||||
|
path = "src/lib.rs"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
rustdb-config = { workspace = true }
|
||||||
|
rustdb-wire = { workspace = true }
|
||||||
|
rustdb-query = { workspace = true }
|
||||||
|
rustdb-storage = { workspace = true }
|
||||||
|
rustdb-index = { workspace = true }
|
||||||
|
rustdb-txn = { workspace = true }
|
||||||
|
rustdb-auth = { workspace = true }
|
||||||
|
rustdb-commands = { workspace = true }
|
||||||
|
tokio = { workspace = true }
|
||||||
|
tokio-util = { workspace = true }
|
||||||
|
tokio-rustls = { workspace = true }
|
||||||
|
rustls-pemfile = { workspace = true }
|
||||||
|
tracing = { workspace = true }
|
||||||
|
tracing-subscriber = { workspace = true }
|
||||||
|
clap = { workspace = true }
|
||||||
|
anyhow = { workspace = true }
|
||||||
|
arc-swap = { workspace = true }
|
||||||
|
serde = { workspace = true }
|
||||||
|
serde_json = { workspace = true }
|
||||||
|
bson = { workspace = true }
|
||||||
|
bytes = { workspace = true }
|
||||||
|
dashmap = { workspace = true }
|
||||||
|
mimalloc = { workspace = true }
|
||||||
|
futures-util = { version = "0.3", features = ["sink"] }
|
||||||
@@ -0,0 +1,425 @@
|
|||||||
|
pub mod management;
|
||||||
|
|
||||||
|
use std::fs::File;
|
||||||
|
use std::io::BufReader;
|
||||||
|
use std::path::PathBuf;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use anyhow::{Context, Result};
|
||||||
|
use dashmap::DashMap;
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
#[cfg(unix)]
|
||||||
|
use tokio::net::UnixListener;
|
||||||
|
use tokio_util::codec::Framed;
|
||||||
|
use tokio_util::sync::CancellationToken;
|
||||||
|
|
||||||
|
use rustdb_config::{RustDbOptions, StorageType, TlsOptions};
|
||||||
|
use rustdb_wire::{WireCodec, OP_QUERY};
|
||||||
|
use rustdb_wire::{encode_op_msg_response, encode_op_reply_response};
|
||||||
|
use rustdb_storage::{StorageAdapter, MemoryStorageAdapter, FileStorageAdapter, OpLog};
|
||||||
|
use rustdb_index::{IndexEngine, IndexOptions};
|
||||||
|
use rustdb_txn::{TransactionEngine, SessionEngine};
|
||||||
|
use rustdb_auth::AuthEngine;
|
||||||
|
use rustdb_commands::{CommandRouter, CommandContext, ConnectionState};
|
||||||
|
use tokio_rustls::rustls::{RootCertStore, ServerConfig};
|
||||||
|
use tokio_rustls::rustls::server::WebPkiClientVerifier;
|
||||||
|
use tokio_rustls::TlsAcceptor;
|
||||||
|
|
||||||
|
/// The main RustDb server.
|
||||||
|
pub struct RustDb {
|
||||||
|
options: RustDbOptions,
|
||||||
|
ctx: Arc<CommandContext>,
|
||||||
|
router: Arc<CommandRouter>,
|
||||||
|
cancel_token: CancellationToken,
|
||||||
|
listener_handle: Option<tokio::task::JoinHandle<()>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RustDb {
|
||||||
|
/// Create a new RustDb server with the given options.
|
||||||
|
pub async fn new(options: RustDbOptions) -> Result<Self> {
|
||||||
|
// Create storage adapter
|
||||||
|
let storage: Arc<dyn StorageAdapter> = match options.storage {
|
||||||
|
StorageType::Memory => {
|
||||||
|
let adapter = if let Some(ref pp) = options.persist_path {
|
||||||
|
tracing::info!("MemoryStorageAdapter with periodic persistence to {}", pp);
|
||||||
|
MemoryStorageAdapter::with_persist_path(PathBuf::from(pp))
|
||||||
|
} else {
|
||||||
|
tracing::warn!(
|
||||||
|
"SmartDB is using in-memory storage — data will NOT survive a restart. \
|
||||||
|
Set storage to 'file' for durable persistence."
|
||||||
|
);
|
||||||
|
MemoryStorageAdapter::new()
|
||||||
|
};
|
||||||
|
Arc::new(adapter)
|
||||||
|
}
|
||||||
|
StorageType::File => {
|
||||||
|
let path = options
|
||||||
|
.storage_path
|
||||||
|
.clone()
|
||||||
|
.unwrap_or_else(|| "./data".to_string());
|
||||||
|
let adapter = FileStorageAdapter::new(&path);
|
||||||
|
Arc::new(adapter)
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Initialize storage
|
||||||
|
storage.initialize().await?;
|
||||||
|
|
||||||
|
// Restore any previously persisted state (no-op for file storage and
|
||||||
|
// memory storage without a persist_path).
|
||||||
|
storage.restore().await?;
|
||||||
|
|
||||||
|
// Spawn periodic persistence task for memory storage with persist_path.
|
||||||
|
if options.storage == StorageType::Memory && options.persist_path.is_some() {
|
||||||
|
let persist_storage = storage.clone();
|
||||||
|
let interval_ms = options.persist_interval_ms;
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let mut interval = tokio::time::interval(Duration::from_millis(interval_ms));
|
||||||
|
interval.tick().await; // skip the immediate first tick
|
||||||
|
loop {
|
||||||
|
interval.tick().await;
|
||||||
|
if let Err(e) = persist_storage.persist().await {
|
||||||
|
tracing::error!("Periodic persist failed: {}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let indexes: Arc<DashMap<String, IndexEngine>> = Arc::new(DashMap::new());
|
||||||
|
|
||||||
|
// Restore persisted indexes from storage.
|
||||||
|
if let Ok(databases) = storage.list_databases().await {
|
||||||
|
for db_name in &databases {
|
||||||
|
if let Ok(collections) = storage.list_collections(db_name).await {
|
||||||
|
for coll_name in &collections {
|
||||||
|
if let Ok(specs) = storage.get_indexes(db_name, coll_name).await {
|
||||||
|
let has_custom = specs.iter().any(|s| {
|
||||||
|
s.get_str("name").unwrap_or("_id_") != "_id_"
|
||||||
|
});
|
||||||
|
if !has_custom {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let ns_key = format!("{}.{}", db_name, coll_name);
|
||||||
|
let mut engine = IndexEngine::new();
|
||||||
|
|
||||||
|
for spec in &specs {
|
||||||
|
let name = spec.get_str("name").unwrap_or("").to_string();
|
||||||
|
if name == "_id_" {
|
||||||
|
continue; // already created by IndexEngine::new()
|
||||||
|
}
|
||||||
|
let key = match spec.get("key") {
|
||||||
|
Some(bson::Bson::Document(k)) => k.clone(),
|
||||||
|
_ => continue,
|
||||||
|
};
|
||||||
|
let unique = matches!(spec.get("unique"), Some(bson::Bson::Boolean(true)));
|
||||||
|
let sparse = matches!(spec.get("sparse"), Some(bson::Bson::Boolean(true)));
|
||||||
|
let expire_after_seconds = match spec.get("expireAfterSeconds") {
|
||||||
|
Some(bson::Bson::Int32(n)) => Some(*n as u64),
|
||||||
|
Some(bson::Bson::Int64(n)) => Some(*n as u64),
|
||||||
|
_ => None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let options = IndexOptions {
|
||||||
|
name: Some(name.clone()),
|
||||||
|
unique,
|
||||||
|
sparse,
|
||||||
|
expire_after_seconds,
|
||||||
|
};
|
||||||
|
if let Err(e) = engine.create_index(key, options) {
|
||||||
|
tracing::warn!(
|
||||||
|
namespace = %ns_key,
|
||||||
|
index = %name,
|
||||||
|
error = %e,
|
||||||
|
"failed to restore index"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Rebuild index data from existing documents.
|
||||||
|
if let Ok(docs) = storage.find_all(db_name, coll_name).await {
|
||||||
|
if !docs.is_empty() {
|
||||||
|
engine.rebuild_from_documents(&docs);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
namespace = %ns_key,
|
||||||
|
indexes = engine.list_indexes().len(),
|
||||||
|
"restored indexes"
|
||||||
|
);
|
||||||
|
indexes.insert(ns_key, engine);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let auth = Arc::new(AuthEngine::from_options(&options.auth)?);
|
||||||
|
|
||||||
|
let ctx = Arc::new(CommandContext {
|
||||||
|
storage,
|
||||||
|
indexes,
|
||||||
|
transactions: Arc::new(TransactionEngine::new()),
|
||||||
|
sessions: Arc::new(SessionEngine::new(30 * 60 * 1000, 60 * 1000)),
|
||||||
|
cursors: Arc::new(DashMap::new()),
|
||||||
|
start_time: std::time::Instant::now(),
|
||||||
|
oplog: Arc::new(OpLog::new()),
|
||||||
|
auth,
|
||||||
|
});
|
||||||
|
|
||||||
|
let router = Arc::new(CommandRouter::new(ctx.clone()));
|
||||||
|
|
||||||
|
Ok(Self {
|
||||||
|
options,
|
||||||
|
ctx,
|
||||||
|
router,
|
||||||
|
cancel_token: CancellationToken::new(),
|
||||||
|
listener_handle: None,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Start listening for connections.
|
||||||
|
pub async fn start(&mut self) -> Result<()> {
|
||||||
|
let cancel = self.cancel_token.clone();
|
||||||
|
let router = self.router.clone();
|
||||||
|
|
||||||
|
if let Some(ref socket_path) = self.options.socket_path {
|
||||||
|
#[cfg(unix)]
|
||||||
|
{
|
||||||
|
// Remove stale socket file
|
||||||
|
let _ = tokio::fs::remove_file(socket_path).await;
|
||||||
|
|
||||||
|
let listener = UnixListener::bind(socket_path)?;
|
||||||
|
let socket_path_clone = socket_path.clone();
|
||||||
|
tracing::info!("RustDb listening on unix:{}", socket_path_clone);
|
||||||
|
|
||||||
|
let handle = tokio::spawn(async move {
|
||||||
|
loop {
|
||||||
|
tokio::select! {
|
||||||
|
_ = cancel.cancelled() => break,
|
||||||
|
result = listener.accept() => {
|
||||||
|
match result {
|
||||||
|
Ok((stream, _addr)) => {
|
||||||
|
let router = router.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
handle_connection(stream, router).await;
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!("Accept error: {}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
self.listener_handle = Some(handle);
|
||||||
|
}
|
||||||
|
#[cfg(not(unix))]
|
||||||
|
{
|
||||||
|
anyhow::bail!("Unix sockets are not supported on this platform");
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
let addr = format!("{}:{}", self.options.host, self.options.port);
|
||||||
|
let listener = TcpListener::bind(&addr).await?;
|
||||||
|
let tls_acceptor = if self.options.tls.enabled {
|
||||||
|
Some(build_tls_acceptor(&self.options.tls)?)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
tracing::info!(tls = self.options.tls.enabled, "RustDb listening on {}", addr);
|
||||||
|
|
||||||
|
let handle = tokio::spawn(async move {
|
||||||
|
loop {
|
||||||
|
tokio::select! {
|
||||||
|
_ = cancel.cancelled() => break,
|
||||||
|
result = listener.accept() => {
|
||||||
|
match result {
|
||||||
|
Ok((stream, _addr)) => {
|
||||||
|
let _ = stream.set_nodelay(true);
|
||||||
|
let router = router.clone();
|
||||||
|
match tls_acceptor.clone() {
|
||||||
|
Some(acceptor) => {
|
||||||
|
tokio::spawn(async move {
|
||||||
|
match acceptor.accept(stream).await {
|
||||||
|
Ok(tls_stream) => handle_connection(tls_stream, router).await,
|
||||||
|
Err(e) => tracing::debug!("TLS handshake failed: {}", e),
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
tokio::spawn(async move {
|
||||||
|
handle_connection(stream, router).await;
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!("Accept error: {}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
self.listener_handle = Some(handle);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Stop the server.
|
||||||
|
pub async fn stop(&mut self) -> Result<()> {
|
||||||
|
self.cancel_token.cancel();
|
||||||
|
|
||||||
|
if let Some(handle) = self.listener_handle.take() {
|
||||||
|
handle.abort();
|
||||||
|
let _ = handle.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close storage (persists if configured)
|
||||||
|
self.ctx.storage.close().await?;
|
||||||
|
|
||||||
|
// Clean up Unix socket file
|
||||||
|
if let Some(ref socket_path) = self.options.socket_path {
|
||||||
|
let _ = tokio::fs::remove_file(socket_path).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the connection URI.
|
||||||
|
pub fn connection_uri(&self) -> String {
|
||||||
|
self.options.connection_uri()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get a reference to the shared command context (for management IPC access to oplog, storage, etc.).
|
||||||
|
pub fn ctx(&self) -> &Arc<CommandContext> {
|
||||||
|
&self.ctx
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get the server options used for this instance.
|
||||||
|
pub fn options(&self) -> &RustDbOptions {
|
||||||
|
&self.options
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build_tls_acceptor(options: &TlsOptions) -> Result<TlsAcceptor> {
|
||||||
|
let cert_path = options
|
||||||
|
.cert_path
|
||||||
|
.as_deref()
|
||||||
|
.context("tls.certPath is required when tls.enabled is true")?;
|
||||||
|
let key_path = options
|
||||||
|
.key_path
|
||||||
|
.as_deref()
|
||||||
|
.context("tls.keyPath is required when tls.enabled is true")?;
|
||||||
|
|
||||||
|
let certs = load_certs(cert_path)?;
|
||||||
|
let key = load_private_key(key_path)?;
|
||||||
|
|
||||||
|
let config = if options.require_client_cert {
|
||||||
|
let ca_path = options
|
||||||
|
.ca_path
|
||||||
|
.as_deref()
|
||||||
|
.context("tls.caPath is required when tls.requireClientCert is true")?;
|
||||||
|
let roots = load_root_store(ca_path)?;
|
||||||
|
let verifier = WebPkiClientVerifier::builder(Arc::new(roots))
|
||||||
|
.build()
|
||||||
|
.context("failed to build TLS client certificate verifier")?;
|
||||||
|
ServerConfig::builder()
|
||||||
|
.with_client_cert_verifier(verifier)
|
||||||
|
.with_single_cert(certs, key)
|
||||||
|
.context("failed to build TLS server configuration")?
|
||||||
|
} else {
|
||||||
|
ServerConfig::builder()
|
||||||
|
.with_no_client_auth()
|
||||||
|
.with_single_cert(certs, key)
|
||||||
|
.context("failed to build TLS server configuration")?
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(TlsAcceptor::from(Arc::new(config)))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn load_certs(path: &str) -> Result<Vec<tokio_rustls::rustls::pki_types::CertificateDer<'static>>> {
|
||||||
|
let file = File::open(path).with_context(|| format!("failed to open TLS certificate file '{}'", path))?;
|
||||||
|
let mut reader = BufReader::new(file);
|
||||||
|
let certs = rustls_pemfile::certs(&mut reader)
|
||||||
|
.collect::<std::result::Result<Vec<_>, _>>()
|
||||||
|
.with_context(|| format!("failed to parse TLS certificate file '{}'", path))?;
|
||||||
|
|
||||||
|
if certs.is_empty() {
|
||||||
|
anyhow::bail!("TLS certificate file '{}' did not contain any certificates", path);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(certs)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn load_private_key(path: &str) -> Result<tokio_rustls::rustls::pki_types::PrivateKeyDer<'static>> {
|
||||||
|
let file = File::open(path).with_context(|| format!("failed to open TLS private key file '{}'", path))?;
|
||||||
|
let mut reader = BufReader::new(file);
|
||||||
|
rustls_pemfile::private_key(&mut reader)
|
||||||
|
.with_context(|| format!("failed to parse TLS private key file '{}'", path))?
|
||||||
|
.with_context(|| format!("TLS private key file '{}' did not contain a private key", path))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn load_root_store(path: &str) -> Result<RootCertStore> {
|
||||||
|
let mut roots = RootCertStore::empty();
|
||||||
|
for cert in load_certs(path)? {
|
||||||
|
roots
|
||||||
|
.add(cert)
|
||||||
|
.with_context(|| format!("failed to add TLS client CA certificate from '{}'", path))?;
|
||||||
|
}
|
||||||
|
|
||||||
|
if roots.is_empty() {
|
||||||
|
anyhow::bail!("TLS client CA file '{}' did not contain usable certificates", path);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(roots)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Handle a single client connection using the wire protocol codec.
|
||||||
|
async fn handle_connection<S>(stream: S, router: Arc<CommandRouter>)
|
||||||
|
where
|
||||||
|
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
|
||||||
|
{
|
||||||
|
use futures_util::{SinkExt, StreamExt};
|
||||||
|
|
||||||
|
let mut framed = Framed::new(stream, WireCodec);
|
||||||
|
let mut connection = ConnectionState::new();
|
||||||
|
|
||||||
|
while let Some(result) = framed.next().await {
|
||||||
|
match result {
|
||||||
|
Ok(parsed_cmd) => {
|
||||||
|
let request_id = parsed_cmd.request_id;
|
||||||
|
let op_code = parsed_cmd.op_code;
|
||||||
|
|
||||||
|
let response_doc = router.route(&parsed_cmd, &mut connection).await;
|
||||||
|
|
||||||
|
let response_id = next_request_id();
|
||||||
|
|
||||||
|
let response_bytes = if op_code == OP_QUERY {
|
||||||
|
encode_op_reply_response(request_id, &[response_doc], response_id, 0)
|
||||||
|
} else {
|
||||||
|
encode_op_msg_response(request_id, &response_doc, response_id)
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Err(e) = framed.send(response_bytes).await {
|
||||||
|
tracing::debug!("Failed to send response: {}", e);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::debug!("Wire protocol error: {}", e);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn next_request_id() -> i32 {
|
||||||
|
use std::sync::atomic::{AtomicI32, Ordering};
|
||||||
|
static COUNTER: AtomicI32 = AtomicI32::new(1);
|
||||||
|
COUNTER.fetch_add(1, Ordering::Relaxed)
|
||||||
|
}
|
||||||
@@ -0,0 +1,101 @@
|
|||||||
|
#[global_allocator]
|
||||||
|
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
|
||||||
|
|
||||||
|
use clap::Parser;
|
||||||
|
use tracing_subscriber::EnvFilter;
|
||||||
|
use anyhow::Result;
|
||||||
|
|
||||||
|
use rustdb::RustDb;
|
||||||
|
use rustdb::management;
|
||||||
|
use rustdb_config::RustDbOptions;
|
||||||
|
|
||||||
|
/// RustDb - MongoDB-compatible embedded database server
|
||||||
|
#[derive(Parser, Debug)]
|
||||||
|
#[command(name = "rustdb", version, about)]
|
||||||
|
struct Cli {
|
||||||
|
/// Path to JSON configuration file
|
||||||
|
#[arg(short, long, default_value = "config.json")]
|
||||||
|
config: String,
|
||||||
|
|
||||||
|
/// Log level (trace, debug, info, warn, error)
|
||||||
|
#[arg(short, long, default_value = "info")]
|
||||||
|
log_level: String,
|
||||||
|
|
||||||
|
/// Validate configuration without starting
|
||||||
|
#[arg(long)]
|
||||||
|
validate: bool,
|
||||||
|
|
||||||
|
/// Validate data integrity of a storage directory (offline check)
|
||||||
|
#[arg(long, value_name = "PATH")]
|
||||||
|
validate_data: Option<String>,
|
||||||
|
|
||||||
|
/// Run in management mode (JSON-over-stdin IPC for TypeScript wrapper)
|
||||||
|
#[arg(long)]
|
||||||
|
management: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() -> Result<()> {
|
||||||
|
let cli = Cli::parse();
|
||||||
|
|
||||||
|
// Initialize tracing - write to stderr so stdout is reserved for management IPC
|
||||||
|
tracing_subscriber::fmt()
|
||||||
|
.with_writer(std::io::stderr)
|
||||||
|
.with_env_filter(
|
||||||
|
EnvFilter::try_from_default_env()
|
||||||
|
.unwrap_or_else(|_| EnvFilter::new(&cli.log_level)),
|
||||||
|
)
|
||||||
|
.init();
|
||||||
|
|
||||||
|
// Management mode: JSON IPC over stdin/stdout
|
||||||
|
if cli.management {
|
||||||
|
tracing::info!("RustDb starting in management mode...");
|
||||||
|
return management::management_loop().await;
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::info!("RustDb starting...");
|
||||||
|
|
||||||
|
// Load configuration
|
||||||
|
let options = RustDbOptions::from_file(&cli.config)
|
||||||
|
.map_err(|e| anyhow::anyhow!("Failed to load config '{}': {}", cli.config, e))?;
|
||||||
|
|
||||||
|
// Validate-only mode (config)
|
||||||
|
if cli.validate {
|
||||||
|
match options.validate() {
|
||||||
|
Ok(()) => {
|
||||||
|
tracing::info!("Configuration is valid");
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!("Validation error: {}", e);
|
||||||
|
anyhow::bail!("Configuration validation failed: {}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Validate data integrity mode
|
||||||
|
if let Some(ref data_path) = cli.validate_data {
|
||||||
|
tracing::info!("Validating data integrity at {}", data_path);
|
||||||
|
let report = rustdb_storage::validate::validate_data_directory(data_path)
|
||||||
|
.map_err(|e| anyhow::anyhow!("Validation failed: {}", e))?;
|
||||||
|
report.print_summary();
|
||||||
|
if report.has_errors() {
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create and start server
|
||||||
|
let mut db = RustDb::new(options).await?;
|
||||||
|
db.start().await?;
|
||||||
|
|
||||||
|
// Wait for shutdown signal
|
||||||
|
tracing::info!("RustDb is running. Press Ctrl+C to stop.");
|
||||||
|
tokio::signal::ctrl_c().await?;
|
||||||
|
|
||||||
|
tracing::info!("Shutdown signal received");
|
||||||
|
db.stop().await?;
|
||||||
|
|
||||||
|
tracing::info!("RustDb shutdown complete");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,178 @@
|
|||||||
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
|
import * as smartdb from '../ts/index.js';
|
||||||
|
import { MongoClient } from 'mongodb';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as os from 'os';
|
||||||
|
import * as path from 'path';
|
||||||
|
|
||||||
|
let server: smartdb.SmartdbServer;
|
||||||
|
let authedClient: MongoClient;
|
||||||
|
let openClient: MongoClient;
|
||||||
|
let readerClient: MongoClient;
|
||||||
|
let tmpDir: string;
|
||||||
|
let usersPath: string;
|
||||||
|
|
||||||
|
function makeTmpDir(): string {
|
||||||
|
return fs.mkdtempSync(path.join(os.tmpdir(), 'smartdb-auth-test-'));
|
||||||
|
}
|
||||||
|
|
||||||
|
function cleanTmpDir(dir: string): void {
|
||||||
|
if (fs.existsSync(dir)) {
|
||||||
|
fs.rmSync(dir, { recursive: true, force: true });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
tap.test('auth: should start server with SCRAM-SHA-256 auth enabled', async () => {
|
||||||
|
tmpDir = makeTmpDir();
|
||||||
|
usersPath = path.join(tmpDir, 'users.json');
|
||||||
|
server = new smartdb.SmartdbServer({
|
||||||
|
port: 27118,
|
||||||
|
auth: {
|
||||||
|
enabled: true,
|
||||||
|
usersPath,
|
||||||
|
scramIterations: 4096,
|
||||||
|
users: [
|
||||||
|
{
|
||||||
|
username: 'root',
|
||||||
|
password: 'secret',
|
||||||
|
database: 'admin',
|
||||||
|
roles: ['root'],
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
});
|
||||||
|
await server.start();
|
||||||
|
expect(server.running).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('auth: should reject protected commands before authentication', async () => {
|
||||||
|
openClient = new MongoClient('mongodb://127.0.0.1:27118', {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await openClient.connect();
|
||||||
|
|
||||||
|
let threw = false;
|
||||||
|
try {
|
||||||
|
await openClient.db('admin').command({ ping: 1 });
|
||||||
|
} catch (err: any) {
|
||||||
|
threw = true;
|
||||||
|
expect(err.code).toEqual(13);
|
||||||
|
}
|
||||||
|
expect(threw).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('auth: should reject invalid credentials', async () => {
|
||||||
|
const badClient = new MongoClient('mongodb://root:wrong@127.0.0.1:27118/admin?authSource=admin', {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
|
||||||
|
let threw = false;
|
||||||
|
try {
|
||||||
|
await badClient.connect();
|
||||||
|
await badClient.db('admin').command({ ping: 1 });
|
||||||
|
} catch {
|
||||||
|
threw = true;
|
||||||
|
} finally {
|
||||||
|
await badClient.close().catch(() => undefined);
|
||||||
|
}
|
||||||
|
expect(threw).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('auth: should authenticate valid credentials', async () => {
|
||||||
|
authedClient = new MongoClient('mongodb://root:secret@127.0.0.1:27118/admin?authSource=admin', {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await authedClient.connect();
|
||||||
|
const result = await authedClient.db('admin').command({ ping: 1 });
|
||||||
|
expect(result.ok).toEqual(1);
|
||||||
|
|
||||||
|
const status = await authedClient.db('admin').command({ connectionStatus: 1 });
|
||||||
|
expect(status.ok).toEqual(1);
|
||||||
|
expect(status.authInfo.authenticatedUsers[0]).toEqual({ user: 'root', db: 'admin' });
|
||||||
|
expect(status.authInfo.authenticatedUserRoles[0]).toEqual({ role: 'root', db: 'admin' });
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('auth: should allow CRUD after authentication', async () => {
|
||||||
|
const coll = authedClient.db('securedb').collection('notes');
|
||||||
|
const inserted = await coll.insertOne({ title: 'enterprise auth' });
|
||||||
|
expect(inserted.acknowledged).toBeTrue();
|
||||||
|
|
||||||
|
const doc = await coll.findOne({ _id: inserted.insertedId });
|
||||||
|
expect(doc).toBeTruthy();
|
||||||
|
expect(doc!.title).toEqual('enterprise auth');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('auth: root should create a read-only user', async () => {
|
||||||
|
const result = await authedClient.db('admin').command({
|
||||||
|
createUser: 'reader',
|
||||||
|
pwd: 'readpass',
|
||||||
|
roles: [{ role: 'read', db: 'securedb' }],
|
||||||
|
});
|
||||||
|
expect(result.ok).toEqual(1);
|
||||||
|
|
||||||
|
const usersInfo = await authedClient.db('admin').command({ usersInfo: 'reader' });
|
||||||
|
expect(usersInfo.ok).toEqual(1);
|
||||||
|
expect(usersInfo.users.length).toEqual(1);
|
||||||
|
expect(usersInfo.users[0].user).toEqual('reader');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('auth: read-only user should read but not write', async () => {
|
||||||
|
readerClient = new MongoClient('mongodb://reader:readpass@127.0.0.1:27118/admin?authSource=admin', {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await readerClient.connect();
|
||||||
|
|
||||||
|
const doc = await readerClient.db('securedb').collection('notes').findOne({ title: 'enterprise auth' });
|
||||||
|
expect(doc).toBeTruthy();
|
||||||
|
|
||||||
|
let threw = false;
|
||||||
|
try {
|
||||||
|
await readerClient.db('securedb').collection('notes').insertOne({ title: 'denied write' });
|
||||||
|
} catch (err: any) {
|
||||||
|
threw = true;
|
||||||
|
expect(err.code).toEqual(13);
|
||||||
|
}
|
||||||
|
expect(threw).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('auth: persisted users should survive server restart', async () => {
|
||||||
|
await readerClient.close();
|
||||||
|
await authedClient.close();
|
||||||
|
await server.stop();
|
||||||
|
|
||||||
|
// Simulates a crash after writing the temporary auth metadata file but before rename.
|
||||||
|
fs.writeFileSync(path.join(tmpDir, 'users.tmp'), '{ invalid json');
|
||||||
|
|
||||||
|
server = new smartdb.SmartdbServer({
|
||||||
|
port: 27118,
|
||||||
|
auth: {
|
||||||
|
enabled: true,
|
||||||
|
usersPath,
|
||||||
|
users: [],
|
||||||
|
scramIterations: 4096,
|
||||||
|
},
|
||||||
|
});
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
readerClient = new MongoClient('mongodb://reader:readpass@127.0.0.1:27118/admin?authSource=admin', {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await readerClient.connect();
|
||||||
|
const result = await readerClient.db('admin').command({ ping: 1 });
|
||||||
|
expect(result.ok).toEqual(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('auth: cleanup', async () => {
|
||||||
|
await openClient.close();
|
||||||
|
await readerClient.close();
|
||||||
|
await server.stop();
|
||||||
|
expect(server.running).toBeFalse();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
export default tap.start();
|
||||||
@@ -0,0 +1,256 @@
|
|||||||
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
|
import * as smartdb from '../ts/index.js';
|
||||||
|
import { MongoClient, Db } from 'mongodb';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as path from 'path';
|
||||||
|
import * as os from 'os';
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Helpers
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
let tmpDir: string;
|
||||||
|
let server: smartdb.SmartdbServer;
|
||||||
|
let client: MongoClient;
|
||||||
|
let db: Db;
|
||||||
|
|
||||||
|
function makeTmpDir(): string {
|
||||||
|
return fs.mkdtempSync(path.join(os.tmpdir(), 'smartdb-compact-test-'));
|
||||||
|
}
|
||||||
|
|
||||||
|
function cleanTmpDir(dir: string): void {
|
||||||
|
if (fs.existsSync(dir)) {
|
||||||
|
fs.rmSync(dir, { recursive: true, force: true });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function getDataFileSize(storagePath: string, dbName: string, collName: string): number {
|
||||||
|
const dataPath = path.join(storagePath, dbName, collName, 'data.rdb');
|
||||||
|
if (!fs.existsSync(dataPath)) return 0;
|
||||||
|
return fs.statSync(dataPath).size;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Compaction: Setup
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('compaction: start server with file storage', async () => {
|
||||||
|
tmpDir = makeTmpDir();
|
||||||
|
server = new smartdb.SmartdbServer({
|
||||||
|
socketPath: path.join(os.tmpdir(), `smartdb-compact-${Date.now()}-${Math.random().toString(36).slice(2)}.sock`),
|
||||||
|
storage: 'file',
|
||||||
|
storagePath: tmpDir,
|
||||||
|
});
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
client = new MongoClient(server.getConnectionUri(), {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('compactdb');
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Compaction: Updates grow the data file
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('compaction: repeated updates grow the data file', async () => {
|
||||||
|
const coll = db.collection('growing');
|
||||||
|
|
||||||
|
// Insert a document
|
||||||
|
await coll.insertOne({ key: 'target', counter: 0, payload: 'x'.repeat(200) });
|
||||||
|
|
||||||
|
const sizeAfterInsert = getDataFileSize(tmpDir, 'compactdb', 'growing');
|
||||||
|
expect(sizeAfterInsert).toBeGreaterThan(0);
|
||||||
|
|
||||||
|
// Update the same document 50 times — each update appends a new record
|
||||||
|
for (let i = 1; i <= 50; i++) {
|
||||||
|
await coll.updateOne(
|
||||||
|
{ key: 'target' },
|
||||||
|
{ $set: { counter: i, payload: 'y'.repeat(200) } }
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
const sizeAfterUpdates = getDataFileSize(tmpDir, 'compactdb', 'growing');
|
||||||
|
// Compaction may have run during updates, so we can't assert the file is
|
||||||
|
// much larger. What matters is the data is correct.
|
||||||
|
|
||||||
|
// The collection still has just 1 document
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(1);
|
||||||
|
|
||||||
|
const doc = await coll.findOne({ key: 'target' });
|
||||||
|
expect(doc!.counter).toEqual(50);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Compaction: Deletes create tombstones
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('compaction: insert-then-delete creates dead space', async () => {
|
||||||
|
const coll = db.collection('tombstones');
|
||||||
|
|
||||||
|
// Insert 100 documents
|
||||||
|
const docs = [];
|
||||||
|
for (let i = 0; i < 100; i++) {
|
||||||
|
docs.push({ idx: i, data: 'delete-me-' + 'z'.repeat(100) });
|
||||||
|
}
|
||||||
|
await coll.insertMany(docs);
|
||||||
|
|
||||||
|
const sizeAfterInsert = getDataFileSize(tmpDir, 'compactdb', 'tombstones');
|
||||||
|
|
||||||
|
// Delete all 100
|
||||||
|
await coll.deleteMany({});
|
||||||
|
|
||||||
|
const sizeAfterDelete = getDataFileSize(tmpDir, 'compactdb', 'tombstones');
|
||||||
|
// File may have been compacted during deletes (dead > 50% threshold),
|
||||||
|
// but the operation itself should succeed regardless of file size.
|
||||||
|
// After deleting all docs, the file might be very small (just header + compacted).
|
||||||
|
|
||||||
|
// But count is 0
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(0);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Compaction: Data integrity after compaction trigger
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('compaction: data file shrinks after heavy updates trigger compaction', async () => {
|
||||||
|
const coll = db.collection('shrinktest');
|
||||||
|
|
||||||
|
// Insert 10 documents with large payloads
|
||||||
|
const docs = [];
|
||||||
|
for (let i = 0; i < 10; i++) {
|
||||||
|
docs.push({ idx: i, data: 'a'.repeat(500) });
|
||||||
|
}
|
||||||
|
await coll.insertMany(docs);
|
||||||
|
|
||||||
|
const sizeAfterInsert = getDataFileSize(tmpDir, 'compactdb', 'shrinktest');
|
||||||
|
|
||||||
|
// Update each document 20 times (creates 200 dead records vs 10 live)
|
||||||
|
// This should trigger compaction (dead > 50% threshold)
|
||||||
|
for (let round = 0; round < 20; round++) {
|
||||||
|
for (let i = 0; i < 10; i++) {
|
||||||
|
await coll.updateOne(
|
||||||
|
{ idx: i },
|
||||||
|
{ $set: { data: `round-${round}-` + 'b'.repeat(500) } }
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// After compaction, file should be smaller than the pre-compaction peak
|
||||||
|
// (We can't measure the peak exactly, but the final size should be reasonable)
|
||||||
|
const sizeAfterCompaction = getDataFileSize(tmpDir, 'compactdb', 'shrinktest');
|
||||||
|
|
||||||
|
// The file should not be 20x the insert size since compaction should have run
|
||||||
|
// With 10 live records of ~530 bytes each, the file should be roughly that
|
||||||
|
// plus header overhead. Without compaction it would be 210 * ~530 bytes.
|
||||||
|
const maxExpectedSize = sizeAfterInsert * 5; // generous upper bound
|
||||||
|
expect(sizeAfterCompaction).toBeLessThanOrEqual(maxExpectedSize);
|
||||||
|
|
||||||
|
// All documents should still be readable and correct
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(10);
|
||||||
|
|
||||||
|
for (let i = 0; i < 10; i++) {
|
||||||
|
const doc = await coll.findOne({ idx: i });
|
||||||
|
expect(doc).toBeTruthy();
|
||||||
|
expect(doc!.data.startsWith('round-19-')).toBeTrue();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Compaction: Persistence after compaction + restart
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('compaction: data survives compaction + restart', async () => {
|
||||||
|
await client.close();
|
||||||
|
await server.stop();
|
||||||
|
|
||||||
|
server = new smartdb.SmartdbServer({
|
||||||
|
socketPath: path.join(os.tmpdir(), `smartdb-compact-${Date.now()}-${Math.random().toString(36).slice(2)}.sock`),
|
||||||
|
storage: 'file',
|
||||||
|
storagePath: tmpDir,
|
||||||
|
});
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
client = new MongoClient(server.getConnectionUri(), {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('compactdb');
|
||||||
|
|
||||||
|
// Verify shrinktest data
|
||||||
|
const coll = db.collection('shrinktest');
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(10);
|
||||||
|
|
||||||
|
for (let i = 0; i < 10; i++) {
|
||||||
|
const doc = await coll.findOne({ idx: i });
|
||||||
|
expect(doc).toBeTruthy();
|
||||||
|
expect(doc!.data.startsWith('round-19-')).toBeTrue();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify growing collection
|
||||||
|
const growing = db.collection('growing');
|
||||||
|
const growDoc = await growing.findOne({ key: 'target' });
|
||||||
|
expect(growDoc).toBeTruthy();
|
||||||
|
expect(growDoc!.counter).toEqual(50);
|
||||||
|
|
||||||
|
// Verify tombstones collection is empty
|
||||||
|
const tombCount = await db.collection('tombstones').countDocuments();
|
||||||
|
expect(tombCount).toEqual(0);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Compaction: Mixed operations stress test
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('compaction: mixed insert-update-delete stress test', async () => {
|
||||||
|
const coll = db.collection('stress');
|
||||||
|
|
||||||
|
// Phase 1: Insert 200 documents
|
||||||
|
const batch = [];
|
||||||
|
for (let i = 0; i < 200; i++) {
|
||||||
|
batch.push({ idx: i, value: `initial-${i}`, alive: true });
|
||||||
|
}
|
||||||
|
await coll.insertMany(batch);
|
||||||
|
|
||||||
|
// Phase 2: Update every even-indexed document
|
||||||
|
for (let i = 0; i < 200; i += 2) {
|
||||||
|
await coll.updateOne({ idx: i }, { $set: { value: `updated-${i}` } });
|
||||||
|
}
|
||||||
|
|
||||||
|
// Phase 3: Delete every document where idx % 3 === 0
|
||||||
|
await coll.deleteMany({ idx: { $in: Array.from({ length: 67 }, (_, k) => k * 3) } });
|
||||||
|
|
||||||
|
// Verify: documents where idx % 3 !== 0 should remain
|
||||||
|
const remaining = await coll.find({}).toArray();
|
||||||
|
for (const doc of remaining) {
|
||||||
|
expect(doc.idx % 3).not.toEqual(0);
|
||||||
|
if (doc.idx % 2 === 0) {
|
||||||
|
expect(doc.value).toEqual(`updated-${doc.idx}`);
|
||||||
|
} else {
|
||||||
|
expect(doc.value).toEqual(`initial-${doc.idx}`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Count should be 200 - 67 = 133
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(133);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Cleanup
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('compaction: cleanup', async () => {
|
||||||
|
await client.close();
|
||||||
|
await server.stop();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
export default tap.start();
|
||||||
@@ -0,0 +1,91 @@
|
|||||||
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
|
import * as smartdb from '../ts/index.js';
|
||||||
|
import { MongoClient, Db } from 'mongodb';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as os from 'os';
|
||||||
|
import * as path from 'path';
|
||||||
|
|
||||||
|
let tmpDir: string;
|
||||||
|
let localDb: smartdb.LocalSmartDb;
|
||||||
|
let client: MongoClient;
|
||||||
|
let db: Db;
|
||||||
|
let dataPath: string;
|
||||||
|
let corruptedSize: number;
|
||||||
|
|
||||||
|
function makeTmpDir(): string {
|
||||||
|
return fs.mkdtempSync(path.join(os.tmpdir(), 'smartdb-crash-test-'));
|
||||||
|
}
|
||||||
|
|
||||||
|
function cleanTmpDir(dir: string): void {
|
||||||
|
if (fs.existsSync(dir)) {
|
||||||
|
fs.rmSync(dir, { recursive: true, force: true });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async function startAndConnect(): Promise<void> {
|
||||||
|
localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
client = new MongoClient(info.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('crashtest');
|
||||||
|
}
|
||||||
|
|
||||||
|
tap.test('crash-recovery: create baseline data', async () => {
|
||||||
|
tmpDir = makeTmpDir();
|
||||||
|
await startAndConnect();
|
||||||
|
|
||||||
|
await db.collection('docs').insertMany([
|
||||||
|
{ key: 'a', value: 1 },
|
||||||
|
{ key: 'b', value: 2 },
|
||||||
|
{ key: 'c', value: 3 },
|
||||||
|
]);
|
||||||
|
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
|
||||||
|
dataPath = path.join(tmpDir, 'crashtest', 'docs', 'data.rdb');
|
||||||
|
expect(fs.existsSync(dataPath)).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('crash-recovery: append a torn final record', async () => {
|
||||||
|
const data = fs.readFileSync(dataPath);
|
||||||
|
const partialRecord = data.subarray(64, 94);
|
||||||
|
expect(partialRecord.length).toEqual(30);
|
||||||
|
|
||||||
|
fs.appendFileSync(dataPath, partialRecord);
|
||||||
|
corruptedSize = fs.statSync(dataPath).size;
|
||||||
|
expect(corruptedSize).toEqual(data.length + partialRecord.length);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('crash-recovery: restart truncates invalid tail and preserves valid records', async () => {
|
||||||
|
await startAndConnect();
|
||||||
|
|
||||||
|
const repairedSize = fs.statSync(dataPath).size;
|
||||||
|
expect(repairedSize < corruptedSize).toBeTrue();
|
||||||
|
|
||||||
|
const docs = await db.collection('docs').find({}).sort({ key: 1 }).toArray();
|
||||||
|
expect(docs.map(doc => doc.key)).toEqual(['a', 'b', 'c']);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('crash-recovery: future writes remain durable after tail repair', async () => {
|
||||||
|
await db.collection('docs').insertOne({ key: 'd', value: 4 });
|
||||||
|
expect(await db.collection('docs').countDocuments()).toEqual(4);
|
||||||
|
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
|
||||||
|
await startAndConnect();
|
||||||
|
const docs = await db.collection('docs').find({}).sort({ key: 1 }).toArray();
|
||||||
|
expect(docs.map(doc => doc.key)).toEqual(['a', 'b', 'c', 'd']);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('crash-recovery: cleanup', async () => {
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
export default tap.start();
|
||||||
@@ -0,0 +1,191 @@
|
|||||||
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
|
import * as smartdb from '../ts/index.js';
|
||||||
|
import { MongoClient, Db } from 'mongodb';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as path from 'path';
|
||||||
|
import * as os from 'os';
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Test: Deletes persist across restart (tombstone + hint staleness detection)
|
||||||
|
// Covers: append_tombstone to data.rdb, hint file data_file_size tracking,
|
||||||
|
// stale hint detection on restart
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
let tmpDir: string;
|
||||||
|
let localDb: smartdb.LocalSmartDb;
|
||||||
|
let client: MongoClient;
|
||||||
|
let db: Db;
|
||||||
|
|
||||||
|
function makeTmpDir(): string {
|
||||||
|
return fs.mkdtempSync(path.join(os.tmpdir(), 'smartdb-delete-test-'));
|
||||||
|
}
|
||||||
|
|
||||||
|
function cleanTmpDir(dir: string): void {
|
||||||
|
if (fs.existsSync(dir)) {
|
||||||
|
fs.rmSync(dir, { recursive: true, force: true });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Setup
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('setup: start local db and insert documents', async () => {
|
||||||
|
tmpDir = makeTmpDir();
|
||||||
|
localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
client = new MongoClient(info.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('deletetest');
|
||||||
|
|
||||||
|
const coll = db.collection('items');
|
||||||
|
await coll.insertMany([
|
||||||
|
{ name: 'keep-1', value: 100 },
|
||||||
|
{ name: 'keep-2', value: 200 },
|
||||||
|
{ name: 'delete-me', value: 999 },
|
||||||
|
{ name: 'keep-3', value: 300 },
|
||||||
|
]);
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(4);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Delete and verify
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('delete-persistence: delete a document', async () => {
|
||||||
|
const coll = db.collection('items');
|
||||||
|
const result = await coll.deleteOne({ name: 'delete-me' });
|
||||||
|
expect(result.deletedCount).toEqual(1);
|
||||||
|
|
||||||
|
const remaining = await coll.countDocuments();
|
||||||
|
expect(remaining).toEqual(3);
|
||||||
|
|
||||||
|
const deleted = await coll.findOne({ name: 'delete-me' });
|
||||||
|
expect(deleted).toBeNull();
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Graceful restart: delete survives
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('delete-persistence: graceful stop and restart', async () => {
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop(); // graceful — writes hint file
|
||||||
|
|
||||||
|
localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
client = new MongoClient(info.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('deletetest');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('delete-persistence: deleted doc stays deleted after graceful restart', async () => {
|
||||||
|
const coll = db.collection('items');
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(3);
|
||||||
|
|
||||||
|
const deleted = await coll.findOne({ name: 'delete-me' });
|
||||||
|
expect(deleted).toBeNull();
|
||||||
|
|
||||||
|
// The remaining docs are intact
|
||||||
|
const keep1 = await coll.findOne({ name: 'keep-1' });
|
||||||
|
expect(keep1).toBeTruthy();
|
||||||
|
expect(keep1!.value).toEqual(100);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Simulate ungraceful restart: delete after hint write, then restart
|
||||||
|
// The hint file data_file_size check should detect the stale hint
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('delete-persistence: insert and delete more docs, then restart', async () => {
|
||||||
|
const coll = db.collection('items');
|
||||||
|
|
||||||
|
// Insert a new doc
|
||||||
|
await coll.insertOne({ name: 'temporary', value: 777 });
|
||||||
|
expect(await coll.countDocuments()).toEqual(4);
|
||||||
|
|
||||||
|
// Delete it
|
||||||
|
await coll.deleteOne({ name: 'temporary' });
|
||||||
|
expect(await coll.countDocuments()).toEqual(3);
|
||||||
|
|
||||||
|
const gone = await coll.findOne({ name: 'temporary' });
|
||||||
|
expect(gone).toBeNull();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('delete-persistence: stop and restart again', async () => {
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
|
||||||
|
localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
client = new MongoClient(info.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('deletetest');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('delete-persistence: all deletes survived second restart', async () => {
|
||||||
|
const coll = db.collection('items');
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(3);
|
||||||
|
|
||||||
|
// Both deletes are permanent
|
||||||
|
expect(await coll.findOne({ name: 'delete-me' })).toBeNull();
|
||||||
|
expect(await coll.findOne({ name: 'temporary' })).toBeNull();
|
||||||
|
|
||||||
|
// Survivors intact
|
||||||
|
const names = (await coll.find({}).toArray()).map(d => d.name).sort();
|
||||||
|
expect(names).toEqual(['keep-1', 'keep-2', 'keep-3']);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Delete all docs and verify empty after restart
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('delete-persistence: delete all remaining docs', async () => {
|
||||||
|
const coll = db.collection('items');
|
||||||
|
await coll.deleteMany({});
|
||||||
|
expect(await coll.countDocuments()).toEqual(0);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('delete-persistence: restart with empty collection', async () => {
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
|
||||||
|
localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
client = new MongoClient(info.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('deletetest');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('delete-persistence: collection is empty after restart', async () => {
|
||||||
|
const coll = db.collection('items');
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(0);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Cleanup
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('delete-persistence: cleanup', async () => {
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
export default tap.start();
|
||||||
@@ -0,0 +1,394 @@
|
|||||||
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
|
import * as smartdb from '../ts/index.js';
|
||||||
|
import { MongoClient, Db } from 'mongodb';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as path from 'path';
|
||||||
|
import * as os from 'os';
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Helpers
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
let tmpDir: string;
|
||||||
|
let server: smartdb.SmartdbServer;
|
||||||
|
let client: MongoClient;
|
||||||
|
let db: Db;
|
||||||
|
|
||||||
|
function makeTmpDir(): string {
|
||||||
|
return fs.mkdtempSync(path.join(os.tmpdir(), 'smartdb-test-'));
|
||||||
|
}
|
||||||
|
|
||||||
|
function cleanTmpDir(dir: string): void {
|
||||||
|
if (fs.existsSync(dir)) {
|
||||||
|
fs.rmSync(dir, { recursive: true, force: true });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// File Storage: Startup
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('file-storage: should start server with file storage', async () => {
|
||||||
|
tmpDir = makeTmpDir();
|
||||||
|
server = new smartdb.SmartdbServer({
|
||||||
|
port: 27118,
|
||||||
|
storage: 'file',
|
||||||
|
storagePath: tmpDir,
|
||||||
|
});
|
||||||
|
await server.start();
|
||||||
|
expect(server.running).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: should connect MongoClient', async () => {
|
||||||
|
client = new MongoClient('mongodb://127.0.0.1:27118', {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('filetest');
|
||||||
|
expect(db).toBeTruthy();
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// File Storage: Data files are created on disk
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('file-storage: inserting creates data files on disk', async () => {
|
||||||
|
const coll = db.collection('diskcheck');
|
||||||
|
await coll.insertOne({ name: 'disk-test', value: 42 });
|
||||||
|
|
||||||
|
// The storage directory should now contain a database directory
|
||||||
|
const dbDir = path.join(tmpDir, 'filetest');
|
||||||
|
expect(fs.existsSync(dbDir)).toBeTrue();
|
||||||
|
|
||||||
|
// Collection directory with data.rdb should exist
|
||||||
|
const collDir = path.join(dbDir, 'diskcheck');
|
||||||
|
expect(fs.existsSync(collDir)).toBeTrue();
|
||||||
|
|
||||||
|
const dataFile = path.join(collDir, 'data.rdb');
|
||||||
|
expect(fs.existsSync(dataFile)).toBeTrue();
|
||||||
|
|
||||||
|
// data.rdb should have the SMARTDB magic header
|
||||||
|
const header = Buffer.alloc(8);
|
||||||
|
const fd = fs.openSync(dataFile, 'r');
|
||||||
|
fs.readSync(fd, header, 0, 8, 0);
|
||||||
|
fs.closeSync(fd);
|
||||||
|
expect(header.toString('ascii')).toEqual('SMARTDB\0');
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// File Storage: Full CRUD cycle
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('file-storage: insertOne returns valid id', async () => {
|
||||||
|
const coll = db.collection('crud');
|
||||||
|
const result = await coll.insertOne({ name: 'Alice', age: 30 });
|
||||||
|
expect(result.acknowledged).toBeTrue();
|
||||||
|
expect(result.insertedId).toBeTruthy();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: insertMany returns all ids', async () => {
|
||||||
|
const coll = db.collection('crud');
|
||||||
|
const result = await coll.insertMany([
|
||||||
|
{ name: 'Bob', age: 25 },
|
||||||
|
{ name: 'Charlie', age: 35 },
|
||||||
|
{ name: 'Diana', age: 28 },
|
||||||
|
{ name: 'Eve', age: 32 },
|
||||||
|
]);
|
||||||
|
expect(result.insertedCount).toEqual(4);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: findOne retrieves correct document', async () => {
|
||||||
|
const coll = db.collection('crud');
|
||||||
|
const doc = await coll.findOne({ name: 'Alice' });
|
||||||
|
expect(doc).toBeTruthy();
|
||||||
|
expect(doc!.name).toEqual('Alice');
|
||||||
|
expect(doc!.age).toEqual(30);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: find with filter returns correct subset', async () => {
|
||||||
|
const coll = db.collection('crud');
|
||||||
|
const docs = await coll.find({ age: { $gte: 30 } }).toArray();
|
||||||
|
expect(docs.length).toEqual(3); // Alice(30), Charlie(35), Eve(32)
|
||||||
|
expect(docs.every(d => d.age >= 30)).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: updateOne modifies document', async () => {
|
||||||
|
const coll = db.collection('crud');
|
||||||
|
const result = await coll.updateOne(
|
||||||
|
{ name: 'Alice' },
|
||||||
|
{ $set: { age: 31, updated: true } }
|
||||||
|
);
|
||||||
|
expect(result.modifiedCount).toEqual(1);
|
||||||
|
|
||||||
|
const doc = await coll.findOne({ name: 'Alice' });
|
||||||
|
expect(doc!.age).toEqual(31);
|
||||||
|
expect(doc!.updated).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: deleteOne removes document', async () => {
|
||||||
|
const coll = db.collection('crud');
|
||||||
|
const result = await coll.deleteOne({ name: 'Eve' });
|
||||||
|
expect(result.deletedCount).toEqual(1);
|
||||||
|
|
||||||
|
const doc = await coll.findOne({ name: 'Eve' });
|
||||||
|
expect(doc).toBeNull();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: count reflects current state', async () => {
|
||||||
|
const coll = db.collection('crud');
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(4); // 5 inserted - 1 deleted = 4
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// File Storage: Persistence across server restart
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('file-storage: stop server for restart test', async () => {
|
||||||
|
await client.close();
|
||||||
|
await server.stop();
|
||||||
|
expect(server.running).toBeFalse();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: restart server with same data path', async () => {
|
||||||
|
server = new smartdb.SmartdbServer({
|
||||||
|
port: 27118,
|
||||||
|
storage: 'file',
|
||||||
|
storagePath: tmpDir,
|
||||||
|
});
|
||||||
|
await server.start();
|
||||||
|
expect(server.running).toBeTrue();
|
||||||
|
|
||||||
|
client = new MongoClient('mongodb://127.0.0.1:27118', {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('filetest');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: data persists after restart', async () => {
|
||||||
|
const coll = db.collection('crud');
|
||||||
|
|
||||||
|
// Alice should still be there with updated age
|
||||||
|
const alice = await coll.findOne({ name: 'Alice' });
|
||||||
|
expect(alice).toBeTruthy();
|
||||||
|
expect(alice!.age).toEqual(31);
|
||||||
|
expect(alice!.updated).toBeTrue();
|
||||||
|
|
||||||
|
// Bob, Charlie, Diana should be there
|
||||||
|
const bob = await coll.findOne({ name: 'Bob' });
|
||||||
|
expect(bob).toBeTruthy();
|
||||||
|
expect(bob!.age).toEqual(25);
|
||||||
|
|
||||||
|
const charlie = await coll.findOne({ name: 'Charlie' });
|
||||||
|
expect(charlie).toBeTruthy();
|
||||||
|
|
||||||
|
const diana = await coll.findOne({ name: 'Diana' });
|
||||||
|
expect(diana).toBeTruthy();
|
||||||
|
|
||||||
|
// Eve should still be deleted
|
||||||
|
const eve = await coll.findOne({ name: 'Eve' });
|
||||||
|
expect(eve).toBeNull();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: count is correct after restart', async () => {
|
||||||
|
const coll = db.collection('crud');
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(4);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: can write new data after restart', async () => {
|
||||||
|
const coll = db.collection('crud');
|
||||||
|
const result = await coll.insertOne({ name: 'Frank', age: 45 });
|
||||||
|
expect(result.acknowledged).toBeTrue();
|
||||||
|
|
||||||
|
const doc = await coll.findOne({ name: 'Frank' });
|
||||||
|
expect(doc).toBeTruthy();
|
||||||
|
expect(doc!.age).toEqual(45);
|
||||||
|
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(5);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// File Storage: Multiple collections in same database
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('file-storage: multiple collections are independent', async () => {
|
||||||
|
const products = db.collection('products');
|
||||||
|
const orders = db.collection('orders');
|
||||||
|
|
||||||
|
await products.insertMany([
|
||||||
|
{ sku: 'A001', name: 'Widget', price: 9.99 },
|
||||||
|
{ sku: 'A002', name: 'Gadget', price: 19.99 },
|
||||||
|
]);
|
||||||
|
|
||||||
|
await orders.insertMany([
|
||||||
|
{ orderId: 1, sku: 'A001', qty: 3 },
|
||||||
|
{ orderId: 2, sku: 'A002', qty: 1 },
|
||||||
|
{ orderId: 3, sku: 'A001', qty: 2 },
|
||||||
|
]);
|
||||||
|
|
||||||
|
const productCount = await products.countDocuments();
|
||||||
|
const orderCount = await orders.countDocuments();
|
||||||
|
expect(productCount).toEqual(2);
|
||||||
|
expect(orderCount).toEqual(3);
|
||||||
|
|
||||||
|
// Deleting from one collection doesn't affect the other
|
||||||
|
await products.deleteOne({ sku: 'A001' });
|
||||||
|
expect(await products.countDocuments()).toEqual(1);
|
||||||
|
expect(await orders.countDocuments()).toEqual(3);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// File Storage: Multiple databases
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('file-storage: multiple databases are independent', async () => {
|
||||||
|
const db2 = client.db('filetest2');
|
||||||
|
const coll2 = db2.collection('items');
|
||||||
|
|
||||||
|
await coll2.insertOne({ name: 'cross-db-test', source: 'db2' });
|
||||||
|
|
||||||
|
// db2 has 1 doc
|
||||||
|
const count2 = await coll2.countDocuments();
|
||||||
|
expect(count2).toEqual(1);
|
||||||
|
|
||||||
|
// original db is unaffected
|
||||||
|
const crudCount = await db.collection('crud').countDocuments();
|
||||||
|
expect(crudCount).toEqual(5);
|
||||||
|
|
||||||
|
await db2.dropDatabase();
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// File Storage: Large batch insert and retrieval
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('file-storage: bulk insert 1000 documents', async () => {
|
||||||
|
const coll = db.collection('bulk');
|
||||||
|
const docs = [];
|
||||||
|
for (let i = 0; i < 1000; i++) {
|
||||||
|
docs.push({ index: i, data: `value-${i}`, timestamp: Date.now() });
|
||||||
|
}
|
||||||
|
const result = await coll.insertMany(docs);
|
||||||
|
expect(result.insertedCount).toEqual(1000);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: find all 1000 documents', async () => {
|
||||||
|
const coll = db.collection('bulk');
|
||||||
|
const docs = await coll.find({}).toArray();
|
||||||
|
expect(docs.length).toEqual(1000);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: range query on 1000 documents', async () => {
|
||||||
|
const coll = db.collection('bulk');
|
||||||
|
const docs = await coll.find({ index: { $gte: 500, $lt: 600 } }).toArray();
|
||||||
|
expect(docs.length).toEqual(100);
|
||||||
|
expect(docs.every(d => d.index >= 500 && d.index < 600)).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: sorted retrieval with limit', async () => {
|
||||||
|
const coll = db.collection('bulk');
|
||||||
|
const docs = await coll.find({}).sort({ index: -1 }).limit(10).toArray();
|
||||||
|
expect(docs.length).toEqual(10);
|
||||||
|
expect(docs[0].index).toEqual(999);
|
||||||
|
expect(docs[9].index).toEqual(990);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// File Storage: Update many and verify persistence
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('file-storage: updateMany on bulk collection', async () => {
|
||||||
|
const coll = db.collection('bulk');
|
||||||
|
const result = await coll.updateMany(
|
||||||
|
{ index: { $lt: 100 } },
|
||||||
|
{ $set: { batch: 'first-hundred' } }
|
||||||
|
);
|
||||||
|
expect(result.modifiedCount).toEqual(100);
|
||||||
|
|
||||||
|
const updated = await coll.find({ batch: 'first-hundred' }).toArray();
|
||||||
|
expect(updated.length).toEqual(100);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// File Storage: Delete many and verify
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('file-storage: deleteMany removes correct documents', async () => {
|
||||||
|
const coll = db.collection('bulk');
|
||||||
|
const result = await coll.deleteMany({ index: { $gte: 900 } });
|
||||||
|
expect(result.deletedCount).toEqual(100);
|
||||||
|
|
||||||
|
const remaining = await coll.countDocuments();
|
||||||
|
expect(remaining).toEqual(900);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// File Storage: Persistence of bulk data across restart
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('file-storage: stop server for bulk restart test', async () => {
|
||||||
|
await client.close();
|
||||||
|
await server.stop();
|
||||||
|
expect(server.running).toBeFalse();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('file-storage: restart and verify bulk data', async () => {
|
||||||
|
server = new smartdb.SmartdbServer({
|
||||||
|
port: 27118,
|
||||||
|
storage: 'file',
|
||||||
|
storagePath: tmpDir,
|
||||||
|
});
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
client = new MongoClient('mongodb://127.0.0.1:27118', {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('filetest');
|
||||||
|
|
||||||
|
const coll = db.collection('bulk');
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(900);
|
||||||
|
|
||||||
|
// Verify the updateMany persisted
|
||||||
|
const firstHundred = await coll.find({ batch: 'first-hundred' }).toArray();
|
||||||
|
expect(firstHundred.length).toEqual(100);
|
||||||
|
|
||||||
|
// Verify deleted docs are gone
|
||||||
|
const over900 = await coll.find({ index: { $gte: 900 } }).toArray();
|
||||||
|
expect(over900.length).toEqual(0);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// File Storage: Index persistence
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('file-storage: default indexes.json exists on disk', async () => {
|
||||||
|
// The indexes.json is created when the collection is first created,
|
||||||
|
// containing the default _id_ index spec.
|
||||||
|
const indexFile = path.join(tmpDir, 'filetest', 'crud', 'indexes.json');
|
||||||
|
expect(fs.existsSync(indexFile)).toBeTrue();
|
||||||
|
|
||||||
|
const indexData = JSON.parse(fs.readFileSync(indexFile, 'utf-8'));
|
||||||
|
const names = indexData.map((i: any) => i.name);
|
||||||
|
expect(names).toContain('_id_');
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Cleanup
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('file-storage: cleanup', async () => {
|
||||||
|
await client.close();
|
||||||
|
await server.stop();
|
||||||
|
expect(server.running).toBeFalse();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
export default tap.start();
|
||||||
@@ -0,0 +1,126 @@
|
|||||||
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
|
import * as smartdb from '../ts/index.js';
|
||||||
|
import { MongoClient, Db } from 'mongodb';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as path from 'path';
|
||||||
|
import * as os from 'os';
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Test: Missing data.rdb header recovery + startup logging
|
||||||
|
// Covers: ensure_data_header, BuildStats, info-level startup logging
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
let tmpDir: string;
|
||||||
|
let localDb: smartdb.LocalSmartDb;
|
||||||
|
let client: MongoClient;
|
||||||
|
let db: Db;
|
||||||
|
|
||||||
|
function makeTmpDir(): string {
|
||||||
|
return fs.mkdtempSync(path.join(os.tmpdir(), 'smartdb-header-test-'));
|
||||||
|
}
|
||||||
|
|
||||||
|
function cleanTmpDir(dir: string): void {
|
||||||
|
if (fs.existsSync(dir)) {
|
||||||
|
fs.rmSync(dir, { recursive: true, force: true });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Setup: create data, then corrupt it
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('setup: start, insert data, stop', async () => {
|
||||||
|
tmpDir = makeTmpDir();
|
||||||
|
localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
client = new MongoClient(info.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('headertest');
|
||||||
|
|
||||||
|
const coll = db.collection('docs');
|
||||||
|
await coll.insertMany([
|
||||||
|
{ key: 'a', val: 1 },
|
||||||
|
{ key: 'b', val: 2 },
|
||||||
|
{ key: 'c', val: 3 },
|
||||||
|
]);
|
||||||
|
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Delete hint file and restart: should rebuild from data.rdb scan
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('header-recovery: delete hint file and restart', async () => {
|
||||||
|
// Find and delete hint files
|
||||||
|
const dbDir = path.join(tmpDir, 'headertest', 'docs');
|
||||||
|
const hintPath = path.join(dbDir, 'keydir.hint');
|
||||||
|
if (fs.existsSync(hintPath)) {
|
||||||
|
fs.unlinkSync(hintPath);
|
||||||
|
}
|
||||||
|
|
||||||
|
localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
client = new MongoClient(info.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('headertest');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('header-recovery: data intact after hint deletion', async () => {
|
||||||
|
const coll = db.collection('docs');
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(3);
|
||||||
|
|
||||||
|
const a = await coll.findOne({ key: 'a' });
|
||||||
|
expect(a!.val).toEqual(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Write new data after restart, stop, restart again
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('header-recovery: write after hint-less restart', async () => {
|
||||||
|
const coll = db.collection('docs');
|
||||||
|
await coll.insertOne({ key: 'd', val: 4 });
|
||||||
|
expect(await coll.countDocuments()).toEqual(4);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('header-recovery: restart and verify all data', async () => {
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
|
||||||
|
localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
client = new MongoClient(info.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('headertest');
|
||||||
|
|
||||||
|
const coll = db.collection('docs');
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(4);
|
||||||
|
|
||||||
|
const keys = (await coll.find({}).toArray()).map(d => d.key).sort();
|
||||||
|
expect(keys).toEqual(['a', 'b', 'c', 'd']);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Cleanup
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('header-recovery: cleanup', async () => {
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
export default tap.start();
|
||||||
@@ -0,0 +1,235 @@
|
|||||||
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
|
import * as smartdb from '../ts/index.js';
|
||||||
|
import { MongoClient, Db } from 'mongodb';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as path from 'path';
|
||||||
|
import * as os from 'os';
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Helpers
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
let tmpDir: string;
|
||||||
|
let localDb: smartdb.LocalSmartDb;
|
||||||
|
let client: MongoClient;
|
||||||
|
let db: Db;
|
||||||
|
|
||||||
|
function makeTmpDir(): string {
|
||||||
|
return fs.mkdtempSync(path.join(os.tmpdir(), 'smartdb-local-test-'));
|
||||||
|
}
|
||||||
|
|
||||||
|
function cleanTmpDir(dir: string): void {
|
||||||
|
if (fs.existsSync(dir)) {
|
||||||
|
fs.rmSync(dir, { recursive: true, force: true });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// LocalSmartDb: Lifecycle
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('localsmartdb: should start with just a folder path', async () => {
|
||||||
|
tmpDir = makeTmpDir();
|
||||||
|
localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
|
||||||
|
expect(localDb.running).toBeTrue();
|
||||||
|
expect(info.socketPath).toBeTruthy();
|
||||||
|
expect(info.connectionUri).toBeTruthy();
|
||||||
|
expect(info.connectionUri.startsWith('mongodb://')).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('localsmartdb: should connect via returned connectionUri', async () => {
|
||||||
|
const info = localDb.getConnectionInfo();
|
||||||
|
client = new MongoClient(info.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('localtest');
|
||||||
|
expect(db).toBeTruthy();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('localsmartdb: should reject double start', async () => {
|
||||||
|
let threw = false;
|
||||||
|
try {
|
||||||
|
await localDb.start();
|
||||||
|
} catch {
|
||||||
|
threw = true;
|
||||||
|
}
|
||||||
|
expect(threw).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// LocalSmartDb: CRUD via Unix socket
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('localsmartdb: insert and find documents', async () => {
|
||||||
|
const coll = db.collection('notes');
|
||||||
|
await coll.insertMany([
|
||||||
|
{ title: 'Note 1', body: 'First note', priority: 1 },
|
||||||
|
{ title: 'Note 2', body: 'Second note', priority: 2 },
|
||||||
|
{ title: 'Note 3', body: 'Third note', priority: 3 },
|
||||||
|
]);
|
||||||
|
|
||||||
|
const all = await coll.find({}).toArray();
|
||||||
|
expect(all.length).toEqual(3);
|
||||||
|
|
||||||
|
const high = await coll.findOne({ priority: 3 });
|
||||||
|
expect(high).toBeTruthy();
|
||||||
|
expect(high!.title).toEqual('Note 3');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('localsmartdb: update and verify', async () => {
|
||||||
|
const coll = db.collection('notes');
|
||||||
|
await coll.updateOne(
|
||||||
|
{ title: 'Note 2' },
|
||||||
|
{ $set: { body: 'Updated second note', edited: true } }
|
||||||
|
);
|
||||||
|
|
||||||
|
const doc = await coll.findOne({ title: 'Note 2' });
|
||||||
|
expect(doc!.body).toEqual('Updated second note');
|
||||||
|
expect(doc!.edited).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('localsmartdb: delete and verify', async () => {
|
||||||
|
const coll = db.collection('notes');
|
||||||
|
await coll.deleteOne({ title: 'Note 1' });
|
||||||
|
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(2);
|
||||||
|
|
||||||
|
const deleted = await coll.findOne({ title: 'Note 1' });
|
||||||
|
expect(deleted).toBeNull();
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// LocalSmartDb: Persistence across restart
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('localsmartdb: stop for restart', async () => {
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
expect(localDb.running).toBeFalse();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('localsmartdb: restart with same folder', async () => {
|
||||||
|
localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
expect(localDb.running).toBeTrue();
|
||||||
|
|
||||||
|
client = new MongoClient(info.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('localtest');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('localsmartdb: data persists after restart', async () => {
|
||||||
|
const coll = db.collection('notes');
|
||||||
|
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(2); // 3 inserted - 1 deleted
|
||||||
|
|
||||||
|
const note2 = await coll.findOne({ title: 'Note 2' });
|
||||||
|
expect(note2!.body).toEqual('Updated second note');
|
||||||
|
expect(note2!.edited).toBeTrue();
|
||||||
|
|
||||||
|
const note3 = await coll.findOne({ title: 'Note 3' });
|
||||||
|
expect(note3!.priority).toEqual(3);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// LocalSmartDb: Custom socket path
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('localsmartdb: works with custom socket path', async () => {
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
|
||||||
|
const customSocket = path.join(os.tmpdir(), `smartdb-custom-${Date.now()}.sock`);
|
||||||
|
const tmpDir2 = makeTmpDir();
|
||||||
|
const localDb2 = new smartdb.LocalSmartDb({
|
||||||
|
folderPath: tmpDir2,
|
||||||
|
socketPath: customSocket,
|
||||||
|
});
|
||||||
|
|
||||||
|
const info = await localDb2.start();
|
||||||
|
expect(info.socketPath).toEqual(customSocket);
|
||||||
|
|
||||||
|
const client2 = new MongoClient(info.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client2.connect();
|
||||||
|
const testDb = client2.db('customsock');
|
||||||
|
await testDb.collection('test').insertOne({ x: 1 });
|
||||||
|
const doc = await testDb.collection('test').findOne({ x: 1 });
|
||||||
|
expect(doc).toBeTruthy();
|
||||||
|
|
||||||
|
await client2.close();
|
||||||
|
await localDb2.stop();
|
||||||
|
cleanTmpDir(tmpDir2);
|
||||||
|
|
||||||
|
// Reconnect original for remaining tests
|
||||||
|
localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const origInfo = await localDb.start();
|
||||||
|
client = new MongoClient(origInfo.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('localtest');
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// LocalSmartDb: getConnectionUri and getServer helpers
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('localsmartdb: getConnectionUri returns valid uri', async () => {
|
||||||
|
const uri = localDb.getConnectionUri();
|
||||||
|
expect(uri.startsWith('mongodb://')).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('localsmartdb: getServer returns the SmartdbServer', async () => {
|
||||||
|
const srv = localDb.getServer();
|
||||||
|
expect(srv).toBeTruthy();
|
||||||
|
expect(srv.running).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// LocalSmartDb: Data isolation between databases
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('localsmartdb: databases are isolated', async () => {
|
||||||
|
const dbA = client.db('isoA');
|
||||||
|
const dbB = client.db('isoB');
|
||||||
|
|
||||||
|
await dbA.collection('shared').insertOne({ source: 'A', val: 1 });
|
||||||
|
await dbB.collection('shared').insertOne({ source: 'B', val: 2 });
|
||||||
|
|
||||||
|
const docsA = await dbA.collection('shared').find({}).toArray();
|
||||||
|
const docsB = await dbB.collection('shared').find({}).toArray();
|
||||||
|
|
||||||
|
expect(docsA.length).toEqual(1);
|
||||||
|
expect(docsA[0].source).toEqual('A');
|
||||||
|
expect(docsB.length).toEqual(1);
|
||||||
|
expect(docsB[0].source).toEqual('B');
|
||||||
|
|
||||||
|
await dbA.dropDatabase();
|
||||||
|
await dbB.dropDatabase();
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Cleanup
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('localsmartdb: cleanup', async () => {
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
expect(localDb.running).toBeFalse();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
export default tap.start();
|
||||||
@@ -0,0 +1,269 @@
|
|||||||
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
|
import * as smartdb from '../ts/index.js';
|
||||||
|
import { MongoClient, Db } from 'mongodb';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as path from 'path';
|
||||||
|
import * as os from 'os';
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Helpers
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
let tmpDir: string;
|
||||||
|
|
||||||
|
function makeTmpDir(): string {
|
||||||
|
return fs.mkdtempSync(path.join(os.tmpdir(), 'smartdb-migration-test-'));
|
||||||
|
}
|
||||||
|
|
||||||
|
function cleanTmpDir(dir: string): void {
|
||||||
|
if (fs.existsSync(dir)) {
|
||||||
|
fs.rmSync(dir, { recursive: true, force: true });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Create a v0 (legacy JSON) storage layout:
|
||||||
|
* {base}/{db}/{coll}.json
|
||||||
|
* {base}/{db}/{coll}.indexes.json
|
||||||
|
*/
|
||||||
|
function createV0Layout(basePath: string, dbName: string, collName: string, docs: any[]): void {
|
||||||
|
const dbDir = path.join(basePath, dbName);
|
||||||
|
fs.mkdirSync(dbDir, { recursive: true });
|
||||||
|
|
||||||
|
// Convert docs to the extended JSON format that the old Rust engine wrote:
|
||||||
|
// ObjectId is stored as { "$oid": "hex" }
|
||||||
|
const jsonDocs = docs.map(doc => {
|
||||||
|
const clone = { ...doc };
|
||||||
|
if (!clone._id) {
|
||||||
|
// Generate a fake ObjectId-like hex string
|
||||||
|
const hex = [...Array(24)].map(() => Math.floor(Math.random() * 16).toString(16)).join('');
|
||||||
|
clone._id = { '$oid': hex };
|
||||||
|
}
|
||||||
|
return clone;
|
||||||
|
});
|
||||||
|
|
||||||
|
const collPath = path.join(dbDir, `${collName}.json`);
|
||||||
|
fs.writeFileSync(collPath, JSON.stringify(jsonDocs, null, 2));
|
||||||
|
|
||||||
|
const indexPath = path.join(dbDir, `${collName}.indexes.json`);
|
||||||
|
fs.writeFileSync(indexPath, JSON.stringify([
|
||||||
|
{ name: '_id_', key: { _id: 1 } },
|
||||||
|
], null, 2));
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Migration: v0 → v1 basic
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('migration: detects v0 format and migrates on startup', async () => {
|
||||||
|
tmpDir = makeTmpDir();
|
||||||
|
|
||||||
|
// Create v0 layout with test data
|
||||||
|
createV0Layout(tmpDir, 'mydb', 'users', [
|
||||||
|
{ name: 'Alice', age: 30, email: 'alice@test.com' },
|
||||||
|
{ name: 'Bob', age: 25, email: 'bob@test.com' },
|
||||||
|
{ name: 'Charlie', age: 35, email: 'charlie@test.com' },
|
||||||
|
]);
|
||||||
|
|
||||||
|
createV0Layout(tmpDir, 'mydb', 'products', [
|
||||||
|
{ sku: 'W001', name: 'Widget', price: 9.99 },
|
||||||
|
{ sku: 'G001', name: 'Gadget', price: 19.99 },
|
||||||
|
]);
|
||||||
|
|
||||||
|
// Verify v0 files exist
|
||||||
|
expect(fs.existsSync(path.join(tmpDir, 'mydb', 'users.json'))).toBeTrue();
|
||||||
|
expect(fs.existsSync(path.join(tmpDir, 'mydb', 'products.json'))).toBeTrue();
|
||||||
|
|
||||||
|
// Start server — migration should run automatically
|
||||||
|
const server = new smartdb.SmartdbServer({
|
||||||
|
socketPath: path.join(os.tmpdir(), `smartdb-mig-${Date.now()}-${Math.random().toString(36).slice(2)}.sock`),
|
||||||
|
storage: 'file',
|
||||||
|
storagePath: tmpDir,
|
||||||
|
});
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
// v1 directories should now exist
|
||||||
|
expect(fs.existsSync(path.join(tmpDir, 'mydb', 'users', 'data.rdb'))).toBeTrue();
|
||||||
|
expect(fs.existsSync(path.join(tmpDir, 'mydb', 'products', 'data.rdb'))).toBeTrue();
|
||||||
|
|
||||||
|
// v0 files should still exist (not deleted)
|
||||||
|
expect(fs.existsSync(path.join(tmpDir, 'mydb', 'users.json'))).toBeTrue();
|
||||||
|
expect(fs.existsSync(path.join(tmpDir, 'mydb', 'products.json'))).toBeTrue();
|
||||||
|
|
||||||
|
// Connect and verify data is accessible
|
||||||
|
const client = new MongoClient(server.getConnectionUri(), {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
const db = client.db('mydb');
|
||||||
|
|
||||||
|
// Users collection
|
||||||
|
const users = await db.collection('users').find({}).toArray();
|
||||||
|
expect(users.length).toEqual(3);
|
||||||
|
const alice = users.find(u => u.name === 'Alice');
|
||||||
|
expect(alice).toBeTruthy();
|
||||||
|
expect(alice!.age).toEqual(30);
|
||||||
|
expect(alice!.email).toEqual('alice@test.com');
|
||||||
|
|
||||||
|
// Products collection
|
||||||
|
const products = await db.collection('products').find({}).toArray();
|
||||||
|
expect(products.length).toEqual(2);
|
||||||
|
const widget = products.find(p => p.sku === 'W001');
|
||||||
|
expect(widget).toBeTruthy();
|
||||||
|
expect(widget!.price).toEqual(9.99);
|
||||||
|
|
||||||
|
await client.close();
|
||||||
|
await server.stop();
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Migration: migrated data survives another restart
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('migration: migrated data persists across restart', async () => {
|
||||||
|
const server = new smartdb.SmartdbServer({
|
||||||
|
socketPath: path.join(os.tmpdir(), `smartdb-mig-${Date.now()}-${Math.random().toString(36).slice(2)}.sock`),
|
||||||
|
storage: 'file',
|
||||||
|
storagePath: tmpDir,
|
||||||
|
});
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
const client = new MongoClient(server.getConnectionUri(), {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
const db = client.db('mydb');
|
||||||
|
|
||||||
|
const users = await db.collection('users').find({}).toArray();
|
||||||
|
expect(users.length).toEqual(3);
|
||||||
|
|
||||||
|
const products = await db.collection('products').find({}).toArray();
|
||||||
|
expect(products.length).toEqual(2);
|
||||||
|
|
||||||
|
await client.close();
|
||||||
|
await server.stop();
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Migration: can write new data after migration
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('migration: new writes work after migration', async () => {
|
||||||
|
const server = new smartdb.SmartdbServer({
|
||||||
|
socketPath: path.join(os.tmpdir(), `smartdb-mig-${Date.now()}-${Math.random().toString(36).slice(2)}.sock`),
|
||||||
|
storage: 'file',
|
||||||
|
storagePath: tmpDir,
|
||||||
|
});
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
const client = new MongoClient(server.getConnectionUri(), {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
const db = client.db('mydb');
|
||||||
|
|
||||||
|
// Insert new documents
|
||||||
|
await db.collection('users').insertOne({ name: 'Diana', age: 28 });
|
||||||
|
const count = await db.collection('users').countDocuments();
|
||||||
|
expect(count).toEqual(4);
|
||||||
|
|
||||||
|
// Update existing migrated document
|
||||||
|
await db.collection('users').updateOne(
|
||||||
|
{ name: 'Alice' },
|
||||||
|
{ $set: { age: 31 } }
|
||||||
|
);
|
||||||
|
const alice = await db.collection('users').findOne({ name: 'Alice' });
|
||||||
|
expect(alice!.age).toEqual(31);
|
||||||
|
|
||||||
|
// Delete a migrated document
|
||||||
|
await db.collection('products').deleteOne({ sku: 'G001' });
|
||||||
|
const prodCount = await db.collection('products').countDocuments();
|
||||||
|
expect(prodCount).toEqual(1);
|
||||||
|
|
||||||
|
await client.close();
|
||||||
|
await server.stop();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Migration: skips already-migrated data
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('migration: no-op for v1 format', async () => {
|
||||||
|
tmpDir = makeTmpDir();
|
||||||
|
|
||||||
|
// Start fresh to create v1 layout
|
||||||
|
const server = new smartdb.SmartdbServer({
|
||||||
|
socketPath: path.join(os.tmpdir(), `smartdb-mig-${Date.now()}-${Math.random().toString(36).slice(2)}.sock`),
|
||||||
|
storage: 'file',
|
||||||
|
storagePath: tmpDir,
|
||||||
|
});
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
const client = new MongoClient(server.getConnectionUri(), {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
const db = client.db('v1test');
|
||||||
|
await db.collection('items').insertOne({ x: 1 });
|
||||||
|
await client.close();
|
||||||
|
await server.stop();
|
||||||
|
|
||||||
|
// Restart — migration should detect v1 and skip
|
||||||
|
const server2 = new smartdb.SmartdbServer({
|
||||||
|
socketPath: path.join(os.tmpdir(), `smartdb-mig-${Date.now()}-${Math.random().toString(36).slice(2)}.sock`),
|
||||||
|
storage: 'file',
|
||||||
|
storagePath: tmpDir,
|
||||||
|
});
|
||||||
|
await server2.start();
|
||||||
|
|
||||||
|
const client2 = new MongoClient(server2.getConnectionUri(), {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client2.connect();
|
||||||
|
const db2 = client2.db('v1test');
|
||||||
|
const doc = await db2.collection('items').findOne({ x: 1 });
|
||||||
|
expect(doc).toBeTruthy();
|
||||||
|
|
||||||
|
await client2.close();
|
||||||
|
await server2.stop();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Migration: empty storage is handled gracefully
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('migration: empty storage directory works', async () => {
|
||||||
|
tmpDir = makeTmpDir();
|
||||||
|
|
||||||
|
const server = new smartdb.SmartdbServer({
|
||||||
|
socketPath: path.join(os.tmpdir(), `smartdb-mig-${Date.now()}-${Math.random().toString(36).slice(2)}.sock`),
|
||||||
|
storage: 'file',
|
||||||
|
storagePath: tmpDir,
|
||||||
|
});
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
const client = new MongoClient(server.getConnectionUri(), {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
|
||||||
|
// Should work fine with empty storage
|
||||||
|
const db = client.db('emptytest');
|
||||||
|
await db.collection('first').insertOne({ hello: 'world' });
|
||||||
|
const doc = await db.collection('first').findOne({ hello: 'world' });
|
||||||
|
expect(doc).toBeTruthy();
|
||||||
|
|
||||||
|
await client.close();
|
||||||
|
await server.stop();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
export default tap.start();
|
||||||
@@ -1,232 +0,0 @@
|
|||||||
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
|
||||||
import * as smartdb from '../ts/index.js';
|
|
||||||
|
|
||||||
const {
|
|
||||||
calculateCRC32,
|
|
||||||
calculateCRC32Buffer,
|
|
||||||
calculateDocumentChecksum,
|
|
||||||
addChecksum,
|
|
||||||
verifyChecksum,
|
|
||||||
removeChecksum,
|
|
||||||
} = smartdb;
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// CRC32 String Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('checksum: calculateCRC32 should return consistent value for same input', async () => {
|
|
||||||
const result1 = calculateCRC32('hello world');
|
|
||||||
const result2 = calculateCRC32('hello world');
|
|
||||||
expect(result1).toEqual(result2);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: calculateCRC32 should return different values for different inputs', async () => {
|
|
||||||
const result1 = calculateCRC32('hello');
|
|
||||||
const result2 = calculateCRC32('world');
|
|
||||||
expect(result1).not.toEqual(result2);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: calculateCRC32 should return a 32-bit unsigned integer', async () => {
|
|
||||||
const result = calculateCRC32('test string');
|
|
||||||
expect(result).toBeGreaterThanOrEqual(0);
|
|
||||||
expect(result).toBeLessThanOrEqual(0xFFFFFFFF);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: calculateCRC32 should handle empty string', async () => {
|
|
||||||
const result = calculateCRC32('');
|
|
||||||
expect(typeof result).toEqual('number');
|
|
||||||
expect(result).toBeGreaterThanOrEqual(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: calculateCRC32 should handle special characters', async () => {
|
|
||||||
const result = calculateCRC32('hello\nworld\t!"#$%&\'()');
|
|
||||||
expect(typeof result).toEqual('number');
|
|
||||||
expect(result).toBeGreaterThanOrEqual(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: calculateCRC32 should handle unicode characters', async () => {
|
|
||||||
const result = calculateCRC32('hello 世界 🌍');
|
|
||||||
expect(typeof result).toEqual('number');
|
|
||||||
expect(result).toBeGreaterThanOrEqual(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// CRC32 Buffer Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('checksum: calculateCRC32Buffer should return consistent value for same input', async () => {
|
|
||||||
const buffer = Buffer.from('hello world');
|
|
||||||
const result1 = calculateCRC32Buffer(buffer);
|
|
||||||
const result2 = calculateCRC32Buffer(buffer);
|
|
||||||
expect(result1).toEqual(result2);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: calculateCRC32Buffer should return different values for different inputs', async () => {
|
|
||||||
const buffer1 = Buffer.from('hello');
|
|
||||||
const buffer2 = Buffer.from('world');
|
|
||||||
const result1 = calculateCRC32Buffer(buffer1);
|
|
||||||
const result2 = calculateCRC32Buffer(buffer2);
|
|
||||||
expect(result1).not.toEqual(result2);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: calculateCRC32Buffer should handle empty buffer', async () => {
|
|
||||||
const result = calculateCRC32Buffer(Buffer.from(''));
|
|
||||||
expect(typeof result).toEqual('number');
|
|
||||||
expect(result).toBeGreaterThanOrEqual(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: calculateCRC32Buffer should handle binary data', async () => {
|
|
||||||
const buffer = Buffer.from([0x00, 0xFF, 0x7F, 0x80, 0x01]);
|
|
||||||
const result = calculateCRC32Buffer(buffer);
|
|
||||||
expect(typeof result).toEqual('number');
|
|
||||||
expect(result).toBeGreaterThanOrEqual(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Document Checksum Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('checksum: calculateDocumentChecksum should return consistent value', async () => {
|
|
||||||
const doc = { name: 'John', age: 30 };
|
|
||||||
const result1 = calculateDocumentChecksum(doc);
|
|
||||||
const result2 = calculateDocumentChecksum(doc);
|
|
||||||
expect(result1).toEqual(result2);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: calculateDocumentChecksum should exclude _checksum field', async () => {
|
|
||||||
const doc1 = { name: 'John', age: 30 };
|
|
||||||
const doc2 = { name: 'John', age: 30, _checksum: 12345 };
|
|
||||||
const result1 = calculateDocumentChecksum(doc1);
|
|
||||||
const result2 = calculateDocumentChecksum(doc2);
|
|
||||||
expect(result1).toEqual(result2);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: calculateDocumentChecksum should handle empty document', async () => {
|
|
||||||
const result = calculateDocumentChecksum({});
|
|
||||||
expect(typeof result).toEqual('number');
|
|
||||||
expect(result).toBeGreaterThanOrEqual(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: calculateDocumentChecksum should handle nested objects', async () => {
|
|
||||||
const doc = {
|
|
||||||
name: 'John',
|
|
||||||
address: {
|
|
||||||
street: '123 Main St',
|
|
||||||
city: 'Springfield',
|
|
||||||
zip: {
|
|
||||||
code: '12345',
|
|
||||||
plus4: '6789',
|
|
||||||
},
|
|
||||||
},
|
|
||||||
};
|
|
||||||
const result = calculateDocumentChecksum(doc);
|
|
||||||
expect(typeof result).toEqual('number');
|
|
||||||
expect(result).toBeGreaterThanOrEqual(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: calculateDocumentChecksum should handle arrays', async () => {
|
|
||||||
const doc = {
|
|
||||||
name: 'John',
|
|
||||||
tags: ['developer', 'tester', 'admin'],
|
|
||||||
scores: [95, 87, 92],
|
|
||||||
};
|
|
||||||
const result = calculateDocumentChecksum(doc);
|
|
||||||
expect(typeof result).toEqual('number');
|
|
||||||
expect(result).toBeGreaterThanOrEqual(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Add/Verify/Remove Checksum Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('checksum: addChecksum should add _checksum field to document', async () => {
|
|
||||||
const doc = { name: 'John', age: 30 };
|
|
||||||
const docWithChecksum = addChecksum(doc);
|
|
||||||
|
|
||||||
expect('_checksum' in docWithChecksum).toBeTrue();
|
|
||||||
expect(typeof docWithChecksum._checksum).toEqual('number');
|
|
||||||
expect(docWithChecksum.name).toEqual('John');
|
|
||||||
expect(docWithChecksum.age).toEqual(30);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: addChecksum should not modify the original document', async () => {
|
|
||||||
const doc = { name: 'John', age: 30 };
|
|
||||||
addChecksum(doc);
|
|
||||||
expect('_checksum' in doc).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: verifyChecksum should return true for valid checksum', async () => {
|
|
||||||
const doc = { name: 'John', age: 30 };
|
|
||||||
const docWithChecksum = addChecksum(doc);
|
|
||||||
const isValid = verifyChecksum(docWithChecksum);
|
|
||||||
expect(isValid).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: verifyChecksum should return false for tampered document', async () => {
|
|
||||||
const doc = { name: 'John', age: 30 };
|
|
||||||
const docWithChecksum = addChecksum(doc);
|
|
||||||
|
|
||||||
// Tamper with the document
|
|
||||||
docWithChecksum.age = 31;
|
|
||||||
|
|
||||||
const isValid = verifyChecksum(docWithChecksum);
|
|
||||||
expect(isValid).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: verifyChecksum should return false for wrong checksum', async () => {
|
|
||||||
const doc = { name: 'John', age: 30, _checksum: 12345 };
|
|
||||||
const isValid = verifyChecksum(doc);
|
|
||||||
expect(isValid).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: verifyChecksum should return true for document without checksum', async () => {
|
|
||||||
const doc = { name: 'John', age: 30 };
|
|
||||||
const isValid = verifyChecksum(doc);
|
|
||||||
expect(isValid).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: removeChecksum should remove _checksum field', async () => {
|
|
||||||
const doc = { name: 'John', age: 30 };
|
|
||||||
const docWithChecksum = addChecksum(doc);
|
|
||||||
const docWithoutChecksum = removeChecksum(docWithChecksum);
|
|
||||||
|
|
||||||
expect('_checksum' in docWithoutChecksum).toBeFalse();
|
|
||||||
expect(docWithoutChecksum.name).toEqual('John');
|
|
||||||
expect(docWithoutChecksum.age).toEqual(30);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('checksum: removeChecksum should handle document without checksum', async () => {
|
|
||||||
const doc = { name: 'John', age: 30 };
|
|
||||||
const result = removeChecksum(doc);
|
|
||||||
|
|
||||||
expect('_checksum' in result).toBeFalse();
|
|
||||||
expect(result.name).toEqual('John');
|
|
||||||
expect(result.age).toEqual(30);
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Round-trip Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('checksum: full round-trip - add, verify, remove', async () => {
|
|
||||||
const original = { name: 'Test', value: 42, nested: { a: 1, b: 2 } };
|
|
||||||
|
|
||||||
// Add checksum
|
|
||||||
const withChecksum = addChecksum(original);
|
|
||||||
expect('_checksum' in withChecksum).toBeTrue();
|
|
||||||
|
|
||||||
// Verify checksum
|
|
||||||
expect(verifyChecksum(withChecksum)).toBeTrue();
|
|
||||||
|
|
||||||
// Remove checksum
|
|
||||||
const restored = removeChecksum(withChecksum);
|
|
||||||
expect('_checksum' in restored).toBeFalse();
|
|
||||||
|
|
||||||
// Original data should be intact
|
|
||||||
expect(restored.name).toEqual('Test');
|
|
||||||
expect(restored.value).toEqual(42);
|
|
||||||
expect(restored.nested.a).toEqual(1);
|
|
||||||
expect(restored.nested.b).toEqual(2);
|
|
||||||
});
|
|
||||||
|
|
||||||
export default tap.start();
|
|
||||||
@@ -1,417 +0,0 @@
|
|||||||
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
|
||||||
import * as smartdb from '../ts/index.js';
|
|
||||||
|
|
||||||
const { IndexEngine, MemoryStorageAdapter, ObjectId } = smartdb;
|
|
||||||
|
|
||||||
let storage: InstanceType<typeof MemoryStorageAdapter>;
|
|
||||||
let indexEngine: InstanceType<typeof IndexEngine>;
|
|
||||||
|
|
||||||
const TEST_DB = 'testdb';
|
|
||||||
const TEST_COLL = 'indextest';
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Setup
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('indexengine: should create IndexEngine instance', async () => {
|
|
||||||
storage = new MemoryStorageAdapter();
|
|
||||||
await storage.initialize();
|
|
||||||
await storage.createCollection(TEST_DB, TEST_COLL);
|
|
||||||
|
|
||||||
indexEngine = new IndexEngine(TEST_DB, TEST_COLL, storage);
|
|
||||||
expect(indexEngine).toBeTruthy();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Index Creation Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('indexengine: createIndex should create single-field index', async () => {
|
|
||||||
const indexName = await indexEngine.createIndex({ name: 1 });
|
|
||||||
|
|
||||||
expect(indexName).toEqual('name_1');
|
|
||||||
|
|
||||||
const exists = await indexEngine.indexExists('name_1');
|
|
||||||
expect(exists).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: createIndex should create compound index', async () => {
|
|
||||||
const indexName = await indexEngine.createIndex({ city: 1, state: -1 });
|
|
||||||
|
|
||||||
expect(indexName).toEqual('city_1_state_-1');
|
|
||||||
|
|
||||||
const exists = await indexEngine.indexExists('city_1_state_-1');
|
|
||||||
expect(exists).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: createIndex should use custom name if provided', async () => {
|
|
||||||
const indexName = await indexEngine.createIndex({ email: 1 }, { name: 'custom_email_index' });
|
|
||||||
|
|
||||||
expect(indexName).toEqual('custom_email_index');
|
|
||||||
|
|
||||||
const exists = await indexEngine.indexExists('custom_email_index');
|
|
||||||
expect(exists).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: createIndex should handle unique option', async () => {
|
|
||||||
const indexName = await indexEngine.createIndex({ uniqueField: 1 }, { unique: true });
|
|
||||||
|
|
||||||
expect(indexName).toEqual('uniqueField_1');
|
|
||||||
|
|
||||||
const indexes = await indexEngine.listIndexes();
|
|
||||||
const uniqueIndex = indexes.find(i => i.name === 'uniqueField_1');
|
|
||||||
expect(uniqueIndex!.unique).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: createIndex should handle sparse option', async () => {
|
|
||||||
const indexName = await indexEngine.createIndex({ sparseField: 1 }, { sparse: true });
|
|
||||||
|
|
||||||
expect(indexName).toEqual('sparseField_1');
|
|
||||||
|
|
||||||
const indexes = await indexEngine.listIndexes();
|
|
||||||
const sparseIndex = indexes.find(i => i.name === 'sparseField_1');
|
|
||||||
expect(sparseIndex!.sparse).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: createIndex should return existing index name if already exists', async () => {
|
|
||||||
const indexName1 = await indexEngine.createIndex({ existingField: 1 }, { name: 'existing_idx' });
|
|
||||||
const indexName2 = await indexEngine.createIndex({ existingField: 1 }, { name: 'existing_idx' });
|
|
||||||
|
|
||||||
expect(indexName1).toEqual(indexName2);
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Index Listing and Existence Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('indexengine: listIndexes should return all indexes', async () => {
|
|
||||||
const indexes = await indexEngine.listIndexes();
|
|
||||||
|
|
||||||
expect(indexes.length).toBeGreaterThanOrEqual(5); // _id_ + created indexes
|
|
||||||
expect(indexes.some(i => i.name === '_id_')).toBeTrue();
|
|
||||||
expect(indexes.some(i => i.name === 'name_1')).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: indexExists should return true for existing index', async () => {
|
|
||||||
const exists = await indexEngine.indexExists('name_1');
|
|
||||||
expect(exists).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: indexExists should return false for non-existent index', async () => {
|
|
||||||
const exists = await indexEngine.indexExists('nonexistent_index');
|
|
||||||
expect(exists).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Document Operations and Index Updates
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('indexengine: should insert documents for index testing', async () => {
|
|
||||||
// Create a fresh index engine for document operations
|
|
||||||
await storage.dropCollection(TEST_DB, TEST_COLL);
|
|
||||||
await storage.createCollection(TEST_DB, TEST_COLL);
|
|
||||||
|
|
||||||
indexEngine = new IndexEngine(TEST_DB, TEST_COLL, storage);
|
|
||||||
|
|
||||||
// Create indexes first
|
|
||||||
await indexEngine.createIndex({ age: 1 });
|
|
||||||
await indexEngine.createIndex({ category: 1 });
|
|
||||||
|
|
||||||
// Insert test documents
|
|
||||||
const docs = [
|
|
||||||
{ _id: new ObjectId(), name: 'Alice', age: 25, category: 'A' },
|
|
||||||
{ _id: new ObjectId(), name: 'Bob', age: 30, category: 'B' },
|
|
||||||
{ _id: new ObjectId(), name: 'Charlie', age: 35, category: 'A' },
|
|
||||||
{ _id: new ObjectId(), name: 'Diana', age: 28, category: 'C' },
|
|
||||||
{ _id: new ObjectId(), name: 'Eve', age: 30, category: 'B' },
|
|
||||||
];
|
|
||||||
|
|
||||||
for (const doc of docs) {
|
|
||||||
const stored = await storage.insertOne(TEST_DB, TEST_COLL, doc);
|
|
||||||
await indexEngine.onInsert(stored);
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: onInsert should update indexes', async () => {
|
|
||||||
const newDoc = {
|
|
||||||
_id: new ObjectId(),
|
|
||||||
name: 'Frank',
|
|
||||||
age: 40,
|
|
||||||
category: 'D',
|
|
||||||
};
|
|
||||||
|
|
||||||
const stored = await storage.insertOne(TEST_DB, TEST_COLL, newDoc);
|
|
||||||
await indexEngine.onInsert(stored);
|
|
||||||
|
|
||||||
// Find by the indexed field
|
|
||||||
const candidates = await indexEngine.findCandidateIds({ age: 40 });
|
|
||||||
expect(candidates).toBeTruthy();
|
|
||||||
expect(candidates!.size).toEqual(1);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: onUpdate should update indexes correctly', async () => {
|
|
||||||
// Get an existing document
|
|
||||||
const docs = await storage.findAll(TEST_DB, TEST_COLL);
|
|
||||||
const oldDoc = docs.find(d => d.name === 'Alice')!;
|
|
||||||
|
|
||||||
// Update the document
|
|
||||||
const newDoc = { ...oldDoc, age: 26 };
|
|
||||||
await storage.updateById(TEST_DB, TEST_COLL, oldDoc._id, newDoc);
|
|
||||||
await indexEngine.onUpdate(oldDoc, newDoc);
|
|
||||||
|
|
||||||
// Old value should not be in index
|
|
||||||
const oldCandidates = await indexEngine.findCandidateIds({ age: 25 });
|
|
||||||
expect(oldCandidates).toBeTruthy();
|
|
||||||
expect(oldCandidates!.has(oldDoc._id.toHexString())).toBeFalse();
|
|
||||||
|
|
||||||
// New value should be in index
|
|
||||||
const newCandidates = await indexEngine.findCandidateIds({ age: 26 });
|
|
||||||
expect(newCandidates).toBeTruthy();
|
|
||||||
expect(newCandidates!.has(oldDoc._id.toHexString())).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: onDelete should remove from indexes', async () => {
|
|
||||||
const docs = await storage.findAll(TEST_DB, TEST_COLL);
|
|
||||||
const docToDelete = docs.find(d => d.name === 'Frank')!;
|
|
||||||
|
|
||||||
await storage.deleteById(TEST_DB, TEST_COLL, docToDelete._id);
|
|
||||||
await indexEngine.onDelete(docToDelete);
|
|
||||||
|
|
||||||
const candidates = await indexEngine.findCandidateIds({ age: 40 });
|
|
||||||
expect(candidates).toBeTruthy();
|
|
||||||
expect(candidates!.has(docToDelete._id.toHexString())).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// findCandidateIds Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('indexengine: findCandidateIds with equality filter', async () => {
|
|
||||||
const candidates = await indexEngine.findCandidateIds({ age: 30 });
|
|
||||||
|
|
||||||
expect(candidates).toBeTruthy();
|
|
||||||
expect(candidates!.size).toEqual(2); // Bob and Eve both have age 30
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: findCandidateIds with $in filter', async () => {
|
|
||||||
const candidates = await indexEngine.findCandidateIds({ age: { $in: [28, 30] } });
|
|
||||||
|
|
||||||
expect(candidates).toBeTruthy();
|
|
||||||
expect(candidates!.size).toEqual(3); // Diana (28), Bob (30), Eve (30)
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: findCandidateIds with no matching index', async () => {
|
|
||||||
const candidates = await indexEngine.findCandidateIds({ nonIndexedField: 'value' });
|
|
||||||
|
|
||||||
// Should return null when no index can be used
|
|
||||||
expect(candidates).toBeNull();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: findCandidateIds with empty filter', async () => {
|
|
||||||
const candidates = await indexEngine.findCandidateIds({});
|
|
||||||
|
|
||||||
// Empty filter = no index can be used
|
|
||||||
expect(candidates).toBeNull();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Range Query Tests (B-Tree)
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('indexengine: findCandidateIds with $gt', async () => {
|
|
||||||
const candidates = await indexEngine.findCandidateIds({ age: { $gt: 30 } });
|
|
||||||
|
|
||||||
expect(candidates).toBeTruthy();
|
|
||||||
// Charlie (35) is > 30
|
|
||||||
expect(candidates!.size).toBeGreaterThanOrEqual(1);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: findCandidateIds with $lt', async () => {
|
|
||||||
const candidates = await indexEngine.findCandidateIds({ age: { $lt: 28 } });
|
|
||||||
|
|
||||||
expect(candidates).toBeTruthy();
|
|
||||||
// Alice (26) is < 28
|
|
||||||
expect(candidates!.size).toBeGreaterThanOrEqual(1);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: findCandidateIds with $gte', async () => {
|
|
||||||
const candidates = await indexEngine.findCandidateIds({ age: { $gte: 30 } });
|
|
||||||
|
|
||||||
expect(candidates).toBeTruthy();
|
|
||||||
// Bob (30), Eve (30), Charlie (35)
|
|
||||||
expect(candidates!.size).toBeGreaterThanOrEqual(3);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: findCandidateIds with $lte', async () => {
|
|
||||||
const candidates = await indexEngine.findCandidateIds({ age: { $lte: 28 } });
|
|
||||||
|
|
||||||
expect(candidates).toBeTruthy();
|
|
||||||
// Alice (26), Diana (28)
|
|
||||||
expect(candidates!.size).toBeGreaterThanOrEqual(2);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: findCandidateIds with range $gt and $lt', async () => {
|
|
||||||
const candidates = await indexEngine.findCandidateIds({ age: { $gt: 26, $lt: 35 } });
|
|
||||||
|
|
||||||
expect(candidates).toBeTruthy();
|
|
||||||
// Diana (28), Bob (30), Eve (30) are between 26 and 35 exclusive
|
|
||||||
expect(candidates!.size).toBeGreaterThanOrEqual(3);
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Index Selection Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('indexengine: selectIndex should return best index for equality', async () => {
|
|
||||||
const result = indexEngine.selectIndex({ age: 30 });
|
|
||||||
|
|
||||||
expect(result).toBeTruthy();
|
|
||||||
expect(result!.name).toEqual('age_1');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: selectIndex should return best index for range query', async () => {
|
|
||||||
const result = indexEngine.selectIndex({ age: { $gt: 25 } });
|
|
||||||
|
|
||||||
expect(result).toBeTruthy();
|
|
||||||
expect(result!.name).toEqual('age_1');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: selectIndex should return null for no matching filter', async () => {
|
|
||||||
const result = indexEngine.selectIndex({ nonIndexedField: 'value' });
|
|
||||||
|
|
||||||
expect(result).toBeNull();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: selectIndex should return null for empty filter', async () => {
|
|
||||||
const result = indexEngine.selectIndex({});
|
|
||||||
|
|
||||||
expect(result).toBeNull();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: selectIndex should prefer more specific indexes', async () => {
|
|
||||||
// Create a compound index
|
|
||||||
await indexEngine.createIndex({ age: 1, category: 1 }, { name: 'age_category_compound' });
|
|
||||||
|
|
||||||
// Query that matches compound index
|
|
||||||
const result = indexEngine.selectIndex({ age: 30, category: 'B' });
|
|
||||||
|
|
||||||
expect(result).toBeTruthy();
|
|
||||||
// Should prefer the compound index since it covers more fields
|
|
||||||
expect(result!.name).toEqual('age_category_compound');
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Drop Index Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('indexengine: dropIndex should remove the index', async () => {
|
|
||||||
await indexEngine.createIndex({ dropTest: 1 }, { name: 'drop_test_idx' });
|
|
||||||
expect(await indexEngine.indexExists('drop_test_idx')).toBeTrue();
|
|
||||||
|
|
||||||
await indexEngine.dropIndex('drop_test_idx');
|
|
||||||
expect(await indexEngine.indexExists('drop_test_idx')).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: dropIndex should throw for _id index', async () => {
|
|
||||||
let threw = false;
|
|
||||||
try {
|
|
||||||
await indexEngine.dropIndex('_id_');
|
|
||||||
} catch (e) {
|
|
||||||
threw = true;
|
|
||||||
}
|
|
||||||
expect(threw).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: dropIndex should throw for non-existent index', async () => {
|
|
||||||
let threw = false;
|
|
||||||
try {
|
|
||||||
await indexEngine.dropIndex('nonexistent_index');
|
|
||||||
} catch (e) {
|
|
||||||
threw = true;
|
|
||||||
}
|
|
||||||
expect(threw).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('indexengine: dropAllIndexes should remove all indexes except _id', async () => {
|
|
||||||
// Create some indexes to drop
|
|
||||||
await indexEngine.createIndex({ toDrop1: 1 });
|
|
||||||
await indexEngine.createIndex({ toDrop2: 1 });
|
|
||||||
|
|
||||||
await indexEngine.dropAllIndexes();
|
|
||||||
|
|
||||||
const indexes = await indexEngine.listIndexes();
|
|
||||||
expect(indexes.length).toEqual(1);
|
|
||||||
expect(indexes[0].name).toEqual('_id_');
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Unique Index Constraint Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('indexengine: unique index should prevent duplicate inserts', async () => {
|
|
||||||
// Create fresh collection
|
|
||||||
await storage.dropCollection(TEST_DB, 'uniquetest');
|
|
||||||
await storage.createCollection(TEST_DB, 'uniquetest');
|
|
||||||
|
|
||||||
const uniqueIndexEngine = new IndexEngine(TEST_DB, 'uniquetest', storage);
|
|
||||||
await uniqueIndexEngine.createIndex({ email: 1 }, { unique: true });
|
|
||||||
|
|
||||||
// Insert first document
|
|
||||||
const doc1 = { _id: new ObjectId(), email: 'test@example.com', name: 'Test' };
|
|
||||||
const stored1 = await storage.insertOne(TEST_DB, 'uniquetest', doc1);
|
|
||||||
await uniqueIndexEngine.onInsert(stored1);
|
|
||||||
|
|
||||||
// Try to insert duplicate
|
|
||||||
const doc2 = { _id: new ObjectId(), email: 'test@example.com', name: 'Test2' };
|
|
||||||
const stored2 = await storage.insertOne(TEST_DB, 'uniquetest', doc2);
|
|
||||||
|
|
||||||
let threw = false;
|
|
||||||
try {
|
|
||||||
await uniqueIndexEngine.onInsert(stored2);
|
|
||||||
} catch (e: any) {
|
|
||||||
threw = true;
|
|
||||||
expect(e.message).toContain('duplicate key');
|
|
||||||
}
|
|
||||||
|
|
||||||
expect(threw).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Sparse Index Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('indexengine: sparse index should not include documents without the field', async () => {
|
|
||||||
// Create fresh collection
|
|
||||||
await storage.dropCollection(TEST_DB, 'sparsetest');
|
|
||||||
await storage.createCollection(TEST_DB, 'sparsetest');
|
|
||||||
|
|
||||||
const sparseIndexEngine = new IndexEngine(TEST_DB, 'sparsetest', storage);
|
|
||||||
await sparseIndexEngine.createIndex({ optionalField: 1 }, { sparse: true });
|
|
||||||
|
|
||||||
// Insert doc with the field
|
|
||||||
const doc1 = { _id: new ObjectId(), optionalField: 'hasValue', name: 'HasField' };
|
|
||||||
const stored1 = await storage.insertOne(TEST_DB, 'sparsetest', doc1);
|
|
||||||
await sparseIndexEngine.onInsert(stored1);
|
|
||||||
|
|
||||||
// Insert doc without the field
|
|
||||||
const doc2 = { _id: new ObjectId(), name: 'NoField' };
|
|
||||||
const stored2 = await storage.insertOne(TEST_DB, 'sparsetest', doc2);
|
|
||||||
await sparseIndexEngine.onInsert(stored2);
|
|
||||||
|
|
||||||
// Search for documents with the field
|
|
||||||
const candidates = await sparseIndexEngine.findCandidateIds({ optionalField: 'hasValue' });
|
|
||||||
expect(candidates).toBeTruthy();
|
|
||||||
expect(candidates!.size).toEqual(1);
|
|
||||||
expect(candidates!.has(stored1._id.toHexString())).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Cleanup
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('indexengine: cleanup', async () => {
|
|
||||||
await storage.close();
|
|
||||||
expect(true).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
export default tap.start();
|
|
||||||
@@ -1,273 +0,0 @@
|
|||||||
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
|
||||||
import * as smartdb from '../ts/index.js';
|
|
||||||
|
|
||||||
const { QueryPlanner, IndexEngine, MemoryStorageAdapter, ObjectId } = smartdb;
|
|
||||||
|
|
||||||
let storage: InstanceType<typeof MemoryStorageAdapter>;
|
|
||||||
let indexEngine: InstanceType<typeof IndexEngine>;
|
|
||||||
let queryPlanner: InstanceType<typeof QueryPlanner>;
|
|
||||||
|
|
||||||
const TEST_DB = 'testdb';
|
|
||||||
const TEST_COLL = 'testcoll';
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Setup
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('queryplanner: should create QueryPlanner instance', async () => {
|
|
||||||
storage = new MemoryStorageAdapter();
|
|
||||||
await storage.initialize();
|
|
||||||
await storage.createCollection(TEST_DB, TEST_COLL);
|
|
||||||
|
|
||||||
indexEngine = new IndexEngine(TEST_DB, TEST_COLL, storage);
|
|
||||||
queryPlanner = new QueryPlanner(indexEngine);
|
|
||||||
|
|
||||||
expect(queryPlanner).toBeTruthy();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: should insert test documents', async () => {
|
|
||||||
// Insert test documents
|
|
||||||
const docs = [
|
|
||||||
{ _id: new ObjectId(), name: 'Alice', age: 25, city: 'NYC', category: 'A' },
|
|
||||||
{ _id: new ObjectId(), name: 'Bob', age: 30, city: 'LA', category: 'B' },
|
|
||||||
{ _id: new ObjectId(), name: 'Charlie', age: 35, city: 'NYC', category: 'A' },
|
|
||||||
{ _id: new ObjectId(), name: 'Diana', age: 28, city: 'Chicago', category: 'C' },
|
|
||||||
{ _id: new ObjectId(), name: 'Eve', age: 32, city: 'LA', category: 'B' },
|
|
||||||
];
|
|
||||||
|
|
||||||
for (const doc of docs) {
|
|
||||||
await storage.insertOne(TEST_DB, TEST_COLL, doc);
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Basic Plan Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('queryplanner: empty filter should result in COLLSCAN', async () => {
|
|
||||||
const plan = await queryPlanner.plan({});
|
|
||||||
|
|
||||||
expect(plan.type).toEqual('COLLSCAN');
|
|
||||||
expect(plan.indexCovering).toBeFalse();
|
|
||||||
expect(plan.selectivity).toEqual(1.0);
|
|
||||||
expect(plan.explanation).toContain('No filter');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: null filter should result in COLLSCAN', async () => {
|
|
||||||
const plan = await queryPlanner.plan(null as any);
|
|
||||||
|
|
||||||
expect(plan.type).toEqual('COLLSCAN');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: filter with no matching index should result in COLLSCAN', async () => {
|
|
||||||
const plan = await queryPlanner.plan({ nonExistentField: 'value' });
|
|
||||||
|
|
||||||
expect(plan.type).toEqual('COLLSCAN');
|
|
||||||
expect(plan.explanation).toContain('No suitable index');
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Index Scan Tests (with indexes)
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('queryplanner: should create test indexes', async () => {
|
|
||||||
await indexEngine.createIndex({ age: 1 }, { name: 'age_1' });
|
|
||||||
await indexEngine.createIndex({ name: 1 }, { name: 'name_1' });
|
|
||||||
await indexEngine.createIndex({ city: 1, category: 1 }, { name: 'city_category_1' });
|
|
||||||
|
|
||||||
const indexes = await indexEngine.listIndexes();
|
|
||||||
expect(indexes.length).toBeGreaterThanOrEqual(4); // _id_ + 3 created
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: simple equality filter should use IXSCAN', async () => {
|
|
||||||
const plan = await queryPlanner.plan({ age: 30 });
|
|
||||||
|
|
||||||
expect(plan.type).toEqual('IXSCAN');
|
|
||||||
expect(plan.indexName).toEqual('age_1');
|
|
||||||
expect(plan.indexFieldsUsed).toContain('age');
|
|
||||||
expect(plan.usesRange).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: $eq operator should use IXSCAN', async () => {
|
|
||||||
const plan = await queryPlanner.plan({ name: { $eq: 'Alice' } });
|
|
||||||
|
|
||||||
expect(plan.type).toEqual('IXSCAN');
|
|
||||||
expect(plan.indexName).toEqual('name_1');
|
|
||||||
expect(plan.indexFieldsUsed).toContain('name');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: range filter ($gt) should use IXSCAN_RANGE', async () => {
|
|
||||||
const plan = await queryPlanner.plan({ age: { $gt: 25 } });
|
|
||||||
|
|
||||||
expect(plan.type).toEqual('IXSCAN_RANGE');
|
|
||||||
expect(plan.indexName).toEqual('age_1');
|
|
||||||
expect(plan.usesRange).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: range filter ($lt) should use IXSCAN_RANGE', async () => {
|
|
||||||
const plan = await queryPlanner.plan({ age: { $lt: 35 } });
|
|
||||||
|
|
||||||
expect(plan.type).toEqual('IXSCAN_RANGE');
|
|
||||||
expect(plan.usesRange).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: range filter ($gte, $lte) should use IXSCAN_RANGE', async () => {
|
|
||||||
const plan = await queryPlanner.plan({ age: { $gte: 25, $lte: 35 } });
|
|
||||||
|
|
||||||
expect(plan.type).toEqual('IXSCAN_RANGE');
|
|
||||||
expect(plan.usesRange).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: $in operator should use IXSCAN', async () => {
|
|
||||||
const plan = await queryPlanner.plan({ age: { $in: [25, 30, 35] } });
|
|
||||||
|
|
||||||
expect(plan.type).toEqual('IXSCAN');
|
|
||||||
expect(plan.indexName).toEqual('age_1');
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Compound Index Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('queryplanner: compound index - first field equality should use index', async () => {
|
|
||||||
const plan = await queryPlanner.plan({ city: 'NYC' });
|
|
||||||
|
|
||||||
expect(plan.type).toEqual('IXSCAN');
|
|
||||||
expect(plan.indexName).toEqual('city_category_1');
|
|
||||||
expect(plan.indexFieldsUsed).toContain('city');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: compound index - both fields should use full index', async () => {
|
|
||||||
const plan = await queryPlanner.plan({ city: 'NYC', category: 'A' });
|
|
||||||
|
|
||||||
expect(plan.type).toEqual('IXSCAN');
|
|
||||||
expect(plan.indexName).toEqual('city_category_1');
|
|
||||||
expect(plan.indexFieldsUsed).toContain('city');
|
|
||||||
expect(plan.indexFieldsUsed).toContain('category');
|
|
||||||
expect(plan.indexFieldsUsed.length).toEqual(2);
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Selectivity Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('queryplanner: equality query should have low selectivity', async () => {
|
|
||||||
const plan = await queryPlanner.plan({ age: 30 });
|
|
||||||
|
|
||||||
expect(plan.selectivity).toBeLessThan(0.1);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: range query should have moderate selectivity', async () => {
|
|
||||||
const plan = await queryPlanner.plan({ age: { $gt: 25 } });
|
|
||||||
|
|
||||||
expect(plan.selectivity).toBeGreaterThan(0);
|
|
||||||
expect(plan.selectivity).toBeLessThan(1);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: $in query selectivity depends on array size', async () => {
|
|
||||||
const smallInPlan = await queryPlanner.plan({ age: { $in: [25] } });
|
|
||||||
const largeInPlan = await queryPlanner.plan({ age: { $in: [25, 26, 27, 28, 29, 30] } });
|
|
||||||
|
|
||||||
// Larger $in should have higher selectivity (less selective = more documents)
|
|
||||||
expect(largeInPlan.selectivity).toBeGreaterThanOrEqual(smallInPlan.selectivity);
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Index Covering Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('queryplanner: query covering all filter fields should be index covering', async () => {
|
|
||||||
const plan = await queryPlanner.plan({ age: 30 });
|
|
||||||
|
|
||||||
// All filter fields are covered by the index
|
|
||||||
expect(plan.indexCovering).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: query with residual filter should not be index covering', async () => {
|
|
||||||
const plan = await queryPlanner.plan({ city: 'NYC', name: 'Alice' });
|
|
||||||
|
|
||||||
// 'name' is not in the compound index city_category, so it's residual
|
|
||||||
expect(plan.indexCovering).toBeFalse();
|
|
||||||
expect(plan.residualFilter).toBeTruthy();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Explain Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('queryplanner: explain should return detailed plan info', async () => {
|
|
||||||
const explanation = await queryPlanner.explain({ age: 30 });
|
|
||||||
|
|
||||||
expect(explanation.queryPlanner).toBeTruthy();
|
|
||||||
expect(explanation.queryPlanner.plannerVersion).toEqual(1);
|
|
||||||
expect(explanation.queryPlanner.winningPlan).toBeTruthy();
|
|
||||||
expect(explanation.queryPlanner.rejectedPlans).toBeArray();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: explain should include winning and rejected plans', async () => {
|
|
||||||
const explanation = await queryPlanner.explain({ age: 30 });
|
|
||||||
|
|
||||||
expect(explanation.queryPlanner.winningPlan.type).toBeTruthy();
|
|
||||||
expect(explanation.queryPlanner.rejectedPlans.length).toBeGreaterThan(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: explain winning plan should be the best plan', async () => {
|
|
||||||
const explanation = await queryPlanner.explain({ age: 30 });
|
|
||||||
|
|
||||||
// Winning plan should use an index, not collection scan (if index exists)
|
|
||||||
expect(explanation.queryPlanner.winningPlan.type).toEqual('IXSCAN');
|
|
||||||
|
|
||||||
// There should be a COLLSCAN in rejected plans
|
|
||||||
const hasCOLLSCAN = explanation.queryPlanner.rejectedPlans.some(p => p.type === 'COLLSCAN');
|
|
||||||
expect(hasCOLLSCAN).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// $and Operator Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('queryplanner: $and conditions should be analyzed', async () => {
|
|
||||||
const plan = await queryPlanner.plan({
|
|
||||||
$and: [
|
|
||||||
{ age: { $gte: 25 } },
|
|
||||||
{ age: { $lte: 35 } },
|
|
||||||
],
|
|
||||||
});
|
|
||||||
|
|
||||||
expect(plan.type).toEqual('IXSCAN_RANGE');
|
|
||||||
expect(plan.indexName).toEqual('age_1');
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Edge Cases
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('queryplanner: should handle complex nested operators', async () => {
|
|
||||||
const plan = await queryPlanner.plan({
|
|
||||||
age: { $gte: 20, $lte: 40 },
|
|
||||||
city: 'NYC',
|
|
||||||
});
|
|
||||||
|
|
||||||
expect(plan).toBeTruthy();
|
|
||||||
expect(plan.type).not.toBeUndefined();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('queryplanner: should handle $exists operator', async () => {
|
|
||||||
await indexEngine.createIndex({ email: 1 }, { name: 'email_1', sparse: true });
|
|
||||||
|
|
||||||
const plan = await queryPlanner.plan({ email: { $exists: true } });
|
|
||||||
|
|
||||||
// $exists can use sparse indexes
|
|
||||||
expect(plan).toBeTruthy();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Cleanup
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('queryplanner: cleanup', async () => {
|
|
||||||
await storage.close();
|
|
||||||
expect(true).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
export default tap.start();
|
|
||||||
@@ -1,361 +0,0 @@
|
|||||||
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
|
||||||
import * as smartdb from '../ts/index.js';
|
|
||||||
|
|
||||||
const { SessionEngine } = smartdb;
|
|
||||||
|
|
||||||
let sessionEngine: InstanceType<typeof SessionEngine>;
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Setup
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('session: should create SessionEngine instance', async () => {
|
|
||||||
sessionEngine = new SessionEngine({
|
|
||||||
sessionTimeoutMs: 1000, // 1 second for testing
|
|
||||||
cleanupIntervalMs: 10000, // 10 seconds to avoid cleanup during tests
|
|
||||||
});
|
|
||||||
expect(sessionEngine).toBeTruthy();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Session Lifecycle Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('session: startSession should create session with auto-generated ID', async () => {
|
|
||||||
const session = sessionEngine.startSession();
|
|
||||||
|
|
||||||
expect(session).toBeTruthy();
|
|
||||||
expect(session.id).toBeTruthy();
|
|
||||||
expect(session.id.length).toBeGreaterThanOrEqual(32); // UUID hex string (32 or 36 with hyphens)
|
|
||||||
expect(session.createdAt).toBeGreaterThan(0);
|
|
||||||
expect(session.lastActivityAt).toBeGreaterThan(0);
|
|
||||||
expect(session.inTransaction).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: startSession should create session with specified ID', async () => {
|
|
||||||
const customId = 'custom-session-id-12345';
|
|
||||||
const session = sessionEngine.startSession(customId);
|
|
||||||
|
|
||||||
expect(session.id).toEqual(customId);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: startSession should create session with metadata', async () => {
|
|
||||||
const metadata = { client: 'test-client', version: '1.0' };
|
|
||||||
const session = sessionEngine.startSession(undefined, metadata);
|
|
||||||
|
|
||||||
expect(session.metadata).toBeTruthy();
|
|
||||||
expect(session.metadata!.client).toEqual('test-client');
|
|
||||||
expect(session.metadata!.version).toEqual('1.0');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: getSession should return session by ID', async () => {
|
|
||||||
const created = sessionEngine.startSession('get-session-test');
|
|
||||||
const retrieved = sessionEngine.getSession('get-session-test');
|
|
||||||
|
|
||||||
expect(retrieved).toBeTruthy();
|
|
||||||
expect(retrieved!.id).toEqual('get-session-test');
|
|
||||||
expect(retrieved!.id).toEqual(created.id);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: getSession should return undefined for non-existent session', async () => {
|
|
||||||
const session = sessionEngine.getSession('non-existent-session-id');
|
|
||||||
expect(session).toBeUndefined();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: touchSession should update lastActivityAt', async () => {
|
|
||||||
const session = sessionEngine.startSession('touch-test-session');
|
|
||||||
const originalLastActivity = session.lastActivityAt;
|
|
||||||
|
|
||||||
// Wait a bit to ensure time difference
|
|
||||||
await new Promise(resolve => setTimeout(resolve, 10));
|
|
||||||
|
|
||||||
const touched = sessionEngine.touchSession('touch-test-session');
|
|
||||||
expect(touched).toBeTrue();
|
|
||||||
|
|
||||||
const updated = sessionEngine.getSession('touch-test-session');
|
|
||||||
expect(updated!.lastActivityAt).toBeGreaterThanOrEqual(originalLastActivity);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: touchSession should return false for non-existent session', async () => {
|
|
||||||
const touched = sessionEngine.touchSession('non-existent-touch-session');
|
|
||||||
expect(touched).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: endSession should remove the session', async () => {
|
|
||||||
sessionEngine.startSession('end-session-test');
|
|
||||||
expect(sessionEngine.getSession('end-session-test')).toBeTruthy();
|
|
||||||
|
|
||||||
const ended = await sessionEngine.endSession('end-session-test');
|
|
||||||
expect(ended).toBeTrue();
|
|
||||||
|
|
||||||
expect(sessionEngine.getSession('end-session-test')).toBeUndefined();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: endSession should return false for non-existent session', async () => {
|
|
||||||
const ended = await sessionEngine.endSession('non-existent-end-session');
|
|
||||||
expect(ended).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Session Expiry Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('session: isSessionExpired should return false for fresh session', async () => {
|
|
||||||
const session = sessionEngine.startSession('fresh-session');
|
|
||||||
const isExpired = sessionEngine.isSessionExpired(session);
|
|
||||||
expect(isExpired).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: isSessionExpired should return true for old session', async () => {
|
|
||||||
// Create a session with old lastActivityAt
|
|
||||||
const session = sessionEngine.startSession('old-session');
|
|
||||||
// Manually set lastActivityAt to old value (sessionTimeoutMs is 1000ms)
|
|
||||||
(session as any).lastActivityAt = Date.now() - 2000;
|
|
||||||
|
|
||||||
const isExpired = sessionEngine.isSessionExpired(session);
|
|
||||||
expect(isExpired).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: getSession should return undefined for expired session', async () => {
|
|
||||||
const session = sessionEngine.startSession('expiring-session');
|
|
||||||
// Manually expire the session
|
|
||||||
(session as any).lastActivityAt = Date.now() - 2000;
|
|
||||||
|
|
||||||
const retrieved = sessionEngine.getSession('expiring-session');
|
|
||||||
expect(retrieved).toBeUndefined();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Transaction Integration Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('session: startTransaction should mark session as in transaction', async () => {
|
|
||||||
sessionEngine.startSession('txn-session-1');
|
|
||||||
const started = sessionEngine.startTransaction('txn-session-1', 'txn-id-1', 1);
|
|
||||||
|
|
||||||
expect(started).toBeTrue();
|
|
||||||
|
|
||||||
const session = sessionEngine.getSession('txn-session-1');
|
|
||||||
expect(session!.inTransaction).toBeTrue();
|
|
||||||
expect(session!.txnId).toEqual('txn-id-1');
|
|
||||||
expect(session!.txnNumber).toEqual(1);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: startTransaction should return false for non-existent session', async () => {
|
|
||||||
const started = sessionEngine.startTransaction('non-existent-txn-session', 'txn-id');
|
|
||||||
expect(started).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: endTransaction should clear transaction state', async () => {
|
|
||||||
sessionEngine.startSession('txn-session-2');
|
|
||||||
sessionEngine.startTransaction('txn-session-2', 'txn-id-2');
|
|
||||||
|
|
||||||
const ended = sessionEngine.endTransaction('txn-session-2');
|
|
||||||
expect(ended).toBeTrue();
|
|
||||||
|
|
||||||
const session = sessionEngine.getSession('txn-session-2');
|
|
||||||
expect(session!.inTransaction).toBeFalse();
|
|
||||||
expect(session!.txnId).toBeUndefined();
|
|
||||||
expect(session!.txnNumber).toBeUndefined();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: endTransaction should return false for non-existent session', async () => {
|
|
||||||
const ended = sessionEngine.endTransaction('non-existent-end-txn-session');
|
|
||||||
expect(ended).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: getTransactionId should return transaction ID', async () => {
|
|
||||||
sessionEngine.startSession('txn-id-session');
|
|
||||||
sessionEngine.startTransaction('txn-id-session', 'my-txn-id');
|
|
||||||
|
|
||||||
const txnId = sessionEngine.getTransactionId('txn-id-session');
|
|
||||||
expect(txnId).toEqual('my-txn-id');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: getTransactionId should return undefined for session without transaction', async () => {
|
|
||||||
sessionEngine.startSession('no-txn-session');
|
|
||||||
const txnId = sessionEngine.getTransactionId('no-txn-session');
|
|
||||||
expect(txnId).toBeUndefined();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: getTransactionId should return undefined for non-existent session', async () => {
|
|
||||||
const txnId = sessionEngine.getTransactionId('non-existent-txn-id-session');
|
|
||||||
expect(txnId).toBeUndefined();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: isInTransaction should return correct state', async () => {
|
|
||||||
sessionEngine.startSession('in-txn-check-session');
|
|
||||||
|
|
||||||
expect(sessionEngine.isInTransaction('in-txn-check-session')).toBeFalse();
|
|
||||||
|
|
||||||
sessionEngine.startTransaction('in-txn-check-session', 'txn-check');
|
|
||||||
expect(sessionEngine.isInTransaction('in-txn-check-session')).toBeTrue();
|
|
||||||
|
|
||||||
sessionEngine.endTransaction('in-txn-check-session');
|
|
||||||
expect(sessionEngine.isInTransaction('in-txn-check-session')).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: isInTransaction should return false for non-existent session', async () => {
|
|
||||||
expect(sessionEngine.isInTransaction('non-existent-in-txn-session')).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Session Listing Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('session: listSessions should return all active sessions', async () => {
|
|
||||||
// Close and recreate to have a clean slate
|
|
||||||
sessionEngine.close();
|
|
||||||
sessionEngine = new SessionEngine({
|
|
||||||
sessionTimeoutMs: 10000,
|
|
||||||
cleanupIntervalMs: 60000,
|
|
||||||
});
|
|
||||||
|
|
||||||
sessionEngine.startSession('list-session-1');
|
|
||||||
sessionEngine.startSession('list-session-2');
|
|
||||||
sessionEngine.startSession('list-session-3');
|
|
||||||
|
|
||||||
const sessions = sessionEngine.listSessions();
|
|
||||||
expect(sessions.length).toEqual(3);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: listSessions should not include expired sessions', async () => {
|
|
||||||
const session = sessionEngine.startSession('expired-list-session');
|
|
||||||
// Expire the session
|
|
||||||
(session as any).lastActivityAt = Date.now() - 20000;
|
|
||||||
|
|
||||||
const sessions = sessionEngine.listSessions();
|
|
||||||
const found = sessions.find(s => s.id === 'expired-list-session');
|
|
||||||
expect(found).toBeUndefined();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: getSessionCount should return correct count', async () => {
|
|
||||||
const count = sessionEngine.getSessionCount();
|
|
||||||
expect(count).toBeGreaterThanOrEqual(3); // We created 3 sessions above
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: getSessionsWithTransactions should filter correctly', async () => {
|
|
||||||
// Clean slate
|
|
||||||
sessionEngine.close();
|
|
||||||
sessionEngine = new SessionEngine({
|
|
||||||
sessionTimeoutMs: 10000,
|
|
||||||
cleanupIntervalMs: 60000,
|
|
||||||
});
|
|
||||||
|
|
||||||
sessionEngine.startSession('no-txn-1');
|
|
||||||
sessionEngine.startSession('no-txn-2');
|
|
||||||
sessionEngine.startSession('with-txn-1');
|
|
||||||
sessionEngine.startSession('with-txn-2');
|
|
||||||
|
|
||||||
sessionEngine.startTransaction('with-txn-1', 'txn-a');
|
|
||||||
sessionEngine.startTransaction('with-txn-2', 'txn-b');
|
|
||||||
|
|
||||||
const txnSessions = sessionEngine.getSessionsWithTransactions();
|
|
||||||
expect(txnSessions.length).toEqual(2);
|
|
||||||
expect(txnSessions.every(s => s.inTransaction)).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// getOrCreateSession Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('session: getOrCreateSession should create if missing', async () => {
|
|
||||||
const session = sessionEngine.getOrCreateSession('get-or-create-new');
|
|
||||||
expect(session).toBeTruthy();
|
|
||||||
expect(session.id).toEqual('get-or-create-new');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: getOrCreateSession should return existing session', async () => {
|
|
||||||
const created = sessionEngine.startSession('get-or-create-existing');
|
|
||||||
const retrieved = sessionEngine.getOrCreateSession('get-or-create-existing');
|
|
||||||
|
|
||||||
expect(retrieved.id).toEqual(created.id);
|
|
||||||
expect(retrieved.createdAt).toEqual(created.createdAt);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: getOrCreateSession should touch existing session', async () => {
|
|
||||||
const session = sessionEngine.startSession('get-or-create-touch');
|
|
||||||
const originalLastActivity = session.lastActivityAt;
|
|
||||||
|
|
||||||
await new Promise(resolve => setTimeout(resolve, 10));
|
|
||||||
|
|
||||||
sessionEngine.getOrCreateSession('get-or-create-touch');
|
|
||||||
const updated = sessionEngine.getSession('get-or-create-touch');
|
|
||||||
|
|
||||||
expect(updated!.lastActivityAt).toBeGreaterThanOrEqual(originalLastActivity);
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// extractSessionId Static Method Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('session: extractSessionId should handle UUID object', async () => {
|
|
||||||
const { ObjectId } = smartdb;
|
|
||||||
const uuid = new smartdb.plugins.bson.UUID();
|
|
||||||
const lsid = { id: uuid };
|
|
||||||
|
|
||||||
const extracted = SessionEngine.extractSessionId(lsid);
|
|
||||||
expect(extracted).toEqual(uuid.toHexString());
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: extractSessionId should handle string ID', async () => {
|
|
||||||
const lsid = { id: 'string-session-id' };
|
|
||||||
|
|
||||||
const extracted = SessionEngine.extractSessionId(lsid);
|
|
||||||
expect(extracted).toEqual('string-session-id');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: extractSessionId should handle binary format', async () => {
|
|
||||||
const binaryData = Buffer.from('test-binary-uuid', 'utf8').toString('base64');
|
|
||||||
const lsid = { id: { $binary: { base64: binaryData } } };
|
|
||||||
|
|
||||||
const extracted = SessionEngine.extractSessionId(lsid);
|
|
||||||
expect(extracted).toBeTruthy();
|
|
||||||
expect(typeof extracted).toEqual('string');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: extractSessionId should return undefined for null/undefined', async () => {
|
|
||||||
expect(SessionEngine.extractSessionId(null)).toBeUndefined();
|
|
||||||
expect(SessionEngine.extractSessionId(undefined)).toBeUndefined();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: extractSessionId should return undefined for empty object', async () => {
|
|
||||||
expect(SessionEngine.extractSessionId({})).toBeUndefined();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// refreshSession Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('session: refreshSession should update lastActivityAt', async () => {
|
|
||||||
const session = sessionEngine.startSession('refresh-session-test');
|
|
||||||
const originalLastActivity = session.lastActivityAt;
|
|
||||||
|
|
||||||
await new Promise(resolve => setTimeout(resolve, 10));
|
|
||||||
|
|
||||||
const refreshed = sessionEngine.refreshSession('refresh-session-test');
|
|
||||||
expect(refreshed).toBeTrue();
|
|
||||||
|
|
||||||
const updated = sessionEngine.getSession('refresh-session-test');
|
|
||||||
expect(updated!.lastActivityAt).toBeGreaterThanOrEqual(originalLastActivity);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('session: refreshSession should return false for non-existent session', async () => {
|
|
||||||
const refreshed = sessionEngine.refreshSession('non-existent-refresh-session');
|
|
||||||
expect(refreshed).toBeFalse();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Cleanup
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('session: close should clear all sessions', async () => {
|
|
||||||
sessionEngine.startSession('close-test-session');
|
|
||||||
expect(sessionEngine.getSessionCount()).toBeGreaterThan(0);
|
|
||||||
|
|
||||||
sessionEngine.close();
|
|
||||||
|
|
||||||
expect(sessionEngine.getSessionCount()).toEqual(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
export default tap.start();
|
|
||||||
+83
-1
@@ -1,6 +1,6 @@
|
|||||||
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
import * as smartdb from '../ts/index.js';
|
import * as smartdb from '../ts/index.js';
|
||||||
import { MongoClient, Db, Collection } from 'mongodb';
|
import { MongoClient, Db, Collection, ObjectId } from 'mongodb';
|
||||||
|
|
||||||
let server: smartdb.SmartdbServer;
|
let server: smartdb.SmartdbServer;
|
||||||
let client: MongoClient;
|
let client: MongoClient;
|
||||||
@@ -252,6 +252,71 @@ tap.test('smartdb: update - upsert creates new document', async () => {
|
|||||||
expect(inserted!.email).toEqual('new@example.com');
|
expect(inserted!.email).toEqual('new@example.com');
|
||||||
});
|
});
|
||||||
|
|
||||||
|
tap.test('smartdb: update - aggregation pipeline updateOne', async () => {
|
||||||
|
const collection = db.collection('users');
|
||||||
|
await collection.insertOne({ name: 'PipelineUser', source: 'alpha', legacy: true, visits: 2 });
|
||||||
|
|
||||||
|
const result = await collection.updateOne(
|
||||||
|
{ name: 'PipelineUser' },
|
||||||
|
[
|
||||||
|
{ $set: { sourceCopy: '$source', pipelineStatus: 'updated' } },
|
||||||
|
{ $unset: ['legacy'] },
|
||||||
|
]
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(result.matchedCount).toEqual(1);
|
||||||
|
expect(result.modifiedCount).toEqual(1);
|
||||||
|
|
||||||
|
const updated = await collection.findOne({ name: 'PipelineUser' });
|
||||||
|
expect(updated).toBeTruthy();
|
||||||
|
expect(updated!.sourceCopy).toEqual('alpha');
|
||||||
|
expect(updated!.pipelineStatus).toEqual('updated');
|
||||||
|
expect(updated!.legacy).toBeUndefined();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('smartdb: update - aggregation pipeline upsert', async () => {
|
||||||
|
const collection = db.collection('users');
|
||||||
|
const result = await collection.updateOne(
|
||||||
|
{ name: 'PipelineUpsert' },
|
||||||
|
[
|
||||||
|
{ $set: { email: 'pipeline@example.com', status: 'new', mirroredName: '$name' } },
|
||||||
|
],
|
||||||
|
{ upsert: true }
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(result.upsertedCount).toEqual(1);
|
||||||
|
|
||||||
|
const inserted = await collection.findOne({ name: 'PipelineUpsert' });
|
||||||
|
expect(inserted).toBeTruthy();
|
||||||
|
expect(inserted!.email).toEqual('pipeline@example.com');
|
||||||
|
expect(inserted!.status).toEqual('new');
|
||||||
|
expect(inserted!.mirroredName).toEqual('PipelineUpsert');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('smartdb: update - cannot modify immutable _id through pipeline', async () => {
|
||||||
|
const collection = db.collection('users');
|
||||||
|
const inserted = await collection.insertOne({ name: 'ImmutableIdUser' });
|
||||||
|
|
||||||
|
let threw = false;
|
||||||
|
try {
|
||||||
|
await collection.updateOne(
|
||||||
|
{ _id: inserted.insertedId },
|
||||||
|
[
|
||||||
|
{ $set: { _id: new ObjectId() } },
|
||||||
|
]
|
||||||
|
);
|
||||||
|
} catch (err: any) {
|
||||||
|
threw = true;
|
||||||
|
expect(err.code).toEqual(66);
|
||||||
|
}
|
||||||
|
|
||||||
|
expect(threw).toBeTrue();
|
||||||
|
|
||||||
|
const persisted = await collection.findOne({ _id: inserted.insertedId });
|
||||||
|
expect(persisted).toBeTruthy();
|
||||||
|
expect(persisted!.name).toEqual('ImmutableIdUser');
|
||||||
|
});
|
||||||
|
|
||||||
// ============================================================================
|
// ============================================================================
|
||||||
// Cursor Tests
|
// Cursor Tests
|
||||||
// ============================================================================
|
// ============================================================================
|
||||||
@@ -306,6 +371,23 @@ tap.test('smartdb: findOneAndUpdate - returns updated document', async () => {
|
|||||||
expect(result!.status).toEqual('active');
|
expect(result!.status).toEqual('active');
|
||||||
});
|
});
|
||||||
|
|
||||||
|
tap.test('smartdb: findOneAndUpdate - supports aggregation pipeline updates', async () => {
|
||||||
|
const collection = db.collection('users');
|
||||||
|
await collection.insertOne({ name: 'PipelineFindAndModify', sourceName: 'Finder' });
|
||||||
|
|
||||||
|
const result = await collection.findOneAndUpdate(
|
||||||
|
{ name: 'PipelineFindAndModify' },
|
||||||
|
[
|
||||||
|
{ $set: { displayName: '$sourceName', mode: 'pipeline' } },
|
||||||
|
],
|
||||||
|
{ returnDocument: 'after' }
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(result).toBeTruthy();
|
||||||
|
expect(result!.displayName).toEqual('Finder');
|
||||||
|
expect(result!.mode).toEqual('pipeline');
|
||||||
|
});
|
||||||
|
|
||||||
tap.test('smartdb: findOneAndDelete - returns deleted document', async () => {
|
tap.test('smartdb: findOneAndDelete - returns deleted document', async () => {
|
||||||
const collection = db.collection('users');
|
const collection = db.collection('users');
|
||||||
|
|
||||||
|
|||||||
@@ -1,411 +0,0 @@
|
|||||||
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
|
||||||
import * as smartdb from '../ts/index.js';
|
|
||||||
import * as path from 'path';
|
|
||||||
import * as fs from 'fs/promises';
|
|
||||||
|
|
||||||
const { WAL, ObjectId } = smartdb;
|
|
||||||
|
|
||||||
let wal: InstanceType<typeof WAL>;
|
|
||||||
const TEST_WAL_PATH = '/tmp/smartdb-test-wal/test.wal';
|
|
||||||
|
|
||||||
// Helper to clean up test files
|
|
||||||
async function cleanupTestFiles() {
|
|
||||||
try {
|
|
||||||
await fs.rm('/tmp/smartdb-test-wal', { recursive: true, force: true });
|
|
||||||
} catch {
|
|
||||||
// Ignore if doesn't exist
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Setup
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: cleanup before tests', async () => {
|
|
||||||
await cleanupTestFiles();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: should create WAL instance', async () => {
|
|
||||||
wal = new WAL(TEST_WAL_PATH, { checkpointInterval: 100 });
|
|
||||||
expect(wal).toBeTruthy();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: should initialize WAL', async () => {
|
|
||||||
const result = await wal.initialize();
|
|
||||||
expect(result).toBeTruthy();
|
|
||||||
expect(result.recoveredEntries).toBeArray();
|
|
||||||
expect(result.recoveredEntries.length).toEqual(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// LSN Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: getCurrentLsn should return 0 initially', async () => {
|
|
||||||
const lsn = wal.getCurrentLsn();
|
|
||||||
expect(lsn).toEqual(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: LSN should increment after logging', async () => {
|
|
||||||
const doc = { _id: new ObjectId(), name: 'Test' };
|
|
||||||
const lsn = await wal.logInsert('testdb', 'testcoll', doc as any);
|
|
||||||
|
|
||||||
expect(lsn).toEqual(1);
|
|
||||||
expect(wal.getCurrentLsn()).toEqual(1);
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Insert Logging Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: logInsert should create entry with correct structure', async () => {
|
|
||||||
const doc = { _id: new ObjectId(), name: 'InsertTest', value: 42 };
|
|
||||||
const lsn = await wal.logInsert('testdb', 'insertcoll', doc as any);
|
|
||||||
|
|
||||||
expect(lsn).toBeGreaterThan(0);
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
expect(entry).toBeTruthy();
|
|
||||||
expect(entry!.operation).toEqual('insert');
|
|
||||||
expect(entry!.dbName).toEqual('testdb');
|
|
||||||
expect(entry!.collName).toEqual('insertcoll');
|
|
||||||
expect(entry!.documentId).toEqual(doc._id.toHexString());
|
|
||||||
expect(entry!.data).toBeTruthy();
|
|
||||||
expect(entry!.timestamp).toBeGreaterThan(0);
|
|
||||||
expect(entry!.checksum).toBeGreaterThan(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: logInsert with transaction ID', async () => {
|
|
||||||
const doc = { _id: new ObjectId(), name: 'TxnInsertTest' };
|
|
||||||
const lsn = await wal.logInsert('testdb', 'insertcoll', doc as any, 'txn-123');
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
expect(entry!.txnId).toEqual('txn-123');
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Update Logging Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: logUpdate should store old and new document', async () => {
|
|
||||||
const oldDoc = { _id: new ObjectId(), name: 'OldName', value: 1 };
|
|
||||||
const newDoc = { ...oldDoc, name: 'NewName', value: 2 };
|
|
||||||
|
|
||||||
const lsn = await wal.logUpdate('testdb', 'updatecoll', oldDoc as any, newDoc as any);
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
expect(entry).toBeTruthy();
|
|
||||||
expect(entry!.operation).toEqual('update');
|
|
||||||
expect(entry!.data).toBeTruthy();
|
|
||||||
expect(entry!.previousData).toBeTruthy();
|
|
||||||
expect(entry!.documentId).toEqual(oldDoc._id.toHexString());
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Delete Logging Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: logDelete should record deleted document', async () => {
|
|
||||||
const doc = { _id: new ObjectId(), name: 'ToDelete' };
|
|
||||||
|
|
||||||
const lsn = await wal.logDelete('testdb', 'deletecoll', doc as any);
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
expect(entry).toBeTruthy();
|
|
||||||
expect(entry!.operation).toEqual('delete');
|
|
||||||
expect(entry!.previousData).toBeTruthy();
|
|
||||||
expect(entry!.data).toBeUndefined();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Transaction Logging Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: logBeginTransaction should create begin entry', async () => {
|
|
||||||
const lsn = await wal.logBeginTransaction('txn-begin-test');
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
expect(entry).toBeTruthy();
|
|
||||||
expect(entry!.operation).toEqual('begin');
|
|
||||||
expect(entry!.txnId).toEqual('txn-begin-test');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: logCommitTransaction should create commit entry', async () => {
|
|
||||||
const lsn = await wal.logCommitTransaction('txn-commit-test');
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
expect(entry).toBeTruthy();
|
|
||||||
expect(entry!.operation).toEqual('commit');
|
|
||||||
expect(entry!.txnId).toEqual('txn-commit-test');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: logAbortTransaction should create abort entry', async () => {
|
|
||||||
const lsn = await wal.logAbortTransaction('txn-abort-test');
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
expect(entry).toBeTruthy();
|
|
||||||
expect(entry!.operation).toEqual('abort');
|
|
||||||
expect(entry!.txnId).toEqual('txn-abort-test');
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// getTransactionEntries Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: getTransactionEntries should return entries for transaction', async () => {
|
|
||||||
// Log a complete transaction
|
|
||||||
const txnId = 'txn-entries-test';
|
|
||||||
await wal.logBeginTransaction(txnId);
|
|
||||||
|
|
||||||
const doc1 = { _id: new ObjectId(), name: 'TxnDoc1' };
|
|
||||||
await wal.logInsert('testdb', 'txncoll', doc1 as any, txnId);
|
|
||||||
|
|
||||||
const doc2 = { _id: new ObjectId(), name: 'TxnDoc2' };
|
|
||||||
await wal.logInsert('testdb', 'txncoll', doc2 as any, txnId);
|
|
||||||
|
|
||||||
await wal.logCommitTransaction(txnId);
|
|
||||||
|
|
||||||
const entries = wal.getTransactionEntries(txnId);
|
|
||||||
|
|
||||||
expect(entries.length).toEqual(4); // begin + 2 inserts + commit
|
|
||||||
expect(entries.every(e => e.txnId === txnId)).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: getTransactionEntries should return empty for unknown transaction', async () => {
|
|
||||||
const entries = wal.getTransactionEntries('unknown-txn-id');
|
|
||||||
expect(entries.length).toEqual(0);
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// getEntriesAfter Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: getEntriesAfter should filter by LSN', async () => {
|
|
||||||
const currentLsn = wal.getCurrentLsn();
|
|
||||||
|
|
||||||
// Add more entries
|
|
||||||
const doc = { _id: new ObjectId(), name: 'AfterTest' };
|
|
||||||
await wal.logInsert('testdb', 'aftercoll', doc as any);
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(currentLsn);
|
|
||||||
expect(entries.length).toEqual(1);
|
|
||||||
expect(entries[0].lsn).toBeGreaterThan(currentLsn);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: getEntriesAfter with LSN 0 should return all entries', async () => {
|
|
||||||
const entries = wal.getEntriesAfter(0);
|
|
||||||
expect(entries.length).toBeGreaterThan(0);
|
|
||||||
expect(entries.length).toEqual(wal.getCurrentLsn());
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Checkpoint Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: checkpoint should create checkpoint entry', async () => {
|
|
||||||
const lsn = await wal.checkpoint();
|
|
||||||
|
|
||||||
expect(lsn).toBeGreaterThan(0);
|
|
||||||
|
|
||||||
// After checkpoint, getEntriesAfter(checkpoint) should be limited
|
|
||||||
const entries = wal.getEntriesAfter(0);
|
|
||||||
expect(entries.some(e => e.operation === 'checkpoint')).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Document Recovery Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: recoverDocument should deserialize document from entry', async () => {
|
|
||||||
const originalDoc = { _id: new ObjectId(), name: 'RecoverTest', nested: { a: 1, b: 2 } };
|
|
||||||
const lsn = await wal.logInsert('testdb', 'recovercoll', originalDoc as any);
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
const recovered = wal.recoverDocument(entry!);
|
|
||||||
|
|
||||||
expect(recovered).toBeTruthy();
|
|
||||||
expect(recovered!.name).toEqual('RecoverTest');
|
|
||||||
expect(recovered!.nested.a).toEqual(1);
|
|
||||||
expect(recovered!.nested.b).toEqual(2);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: recoverDocument should return null for entry without data', async () => {
|
|
||||||
const lsn = await wal.logBeginTransaction('recover-no-data');
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
const recovered = wal.recoverDocument(entry!);
|
|
||||||
expect(recovered).toBeNull();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: recoverPreviousDocument should deserialize previous state', async () => {
|
|
||||||
const oldDoc = { _id: new ObjectId(), name: 'Old', value: 100 };
|
|
||||||
const newDoc = { ...oldDoc, name: 'New', value: 200 };
|
|
||||||
|
|
||||||
const lsn = await wal.logUpdate('testdb', 'recovercoll', oldDoc as any, newDoc as any);
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
const previous = wal.recoverPreviousDocument(entry!);
|
|
||||||
|
|
||||||
expect(previous).toBeTruthy();
|
|
||||||
expect(previous!.name).toEqual('Old');
|
|
||||||
expect(previous!.value).toEqual(100);
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: recoverPreviousDocument should return null for insert entry', async () => {
|
|
||||||
const doc = { _id: new ObjectId(), name: 'NoPrevious' };
|
|
||||||
const lsn = await wal.logInsert('testdb', 'recovercoll', doc as any);
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
const previous = wal.recoverPreviousDocument(entry!);
|
|
||||||
expect(previous).toBeNull();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// WAL Persistence and Recovery Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: should persist and recover entries', async () => {
|
|
||||||
// Close current WAL
|
|
||||||
await wal.close();
|
|
||||||
|
|
||||||
// Create new WAL instance and initialize (should recover)
|
|
||||||
const wal2 = new WAL(TEST_WAL_PATH, { checkpointInterval: 100 });
|
|
||||||
const result = await wal2.initialize();
|
|
||||||
|
|
||||||
// Should have recovered entries
|
|
||||||
expect(result.recoveredEntries).toBeArray();
|
|
||||||
// After checkpoint, there might not be many recoverable entries
|
|
||||||
// but getCurrentLsn should be preserved or reset
|
|
||||||
|
|
||||||
await wal2.close();
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Entry Checksum Tests
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: entries should have valid checksums', async () => {
|
|
||||||
wal = new WAL(TEST_WAL_PATH + '.checksum', { checkpointInterval: 100 });
|
|
||||||
await wal.initialize();
|
|
||||||
|
|
||||||
const doc = { _id: new ObjectId(), name: 'ChecksumTest' };
|
|
||||||
const lsn = await wal.logInsert('testdb', 'checksumcoll', doc as any);
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
expect(entry!.checksum).toBeGreaterThan(0);
|
|
||||||
expect(typeof entry!.checksum).toEqual('number');
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Edge Cases
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: should handle special characters in document', async () => {
|
|
||||||
const doc = {
|
|
||||||
_id: new ObjectId(),
|
|
||||||
name: 'Test\nWith\tSpecial\r\nChars',
|
|
||||||
emoji: '🎉',
|
|
||||||
unicode: '日本語',
|
|
||||||
};
|
|
||||||
|
|
||||||
const lsn = await wal.logInsert('testdb', 'specialcoll', doc as any);
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
const recovered = wal.recoverDocument(entry!);
|
|
||||||
expect(recovered!.name).toEqual('Test\nWith\tSpecial\r\nChars');
|
|
||||||
expect(recovered!.emoji).toEqual('🎉');
|
|
||||||
expect(recovered!.unicode).toEqual('日本語');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: should handle binary data in documents', async () => {
|
|
||||||
const doc = {
|
|
||||||
_id: new ObjectId(),
|
|
||||||
binaryField: Buffer.from([0x00, 0xFF, 0x7F, 0x80]),
|
|
||||||
};
|
|
||||||
|
|
||||||
const lsn = await wal.logInsert('testdb', 'binarycoll', doc as any);
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
const recovered = wal.recoverDocument(entry!);
|
|
||||||
expect(recovered).toBeTruthy();
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: should handle nested documents', async () => {
|
|
||||||
const doc = {
|
|
||||||
_id: new ObjectId(),
|
|
||||||
level1: {
|
|
||||||
level2: {
|
|
||||||
level3: {
|
|
||||||
value: 'deep',
|
|
||||||
},
|
|
||||||
},
|
|
||||||
},
|
|
||||||
};
|
|
||||||
|
|
||||||
const lsn = await wal.logInsert('testdb', 'nestedcoll', doc as any);
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
const recovered = wal.recoverDocument(entry!);
|
|
||||||
expect(recovered!.level1.level2.level3.value).toEqual('deep');
|
|
||||||
});
|
|
||||||
|
|
||||||
tap.test('wal: should handle arrays in documents', async () => {
|
|
||||||
const doc = {
|
|
||||||
_id: new ObjectId(),
|
|
||||||
tags: ['a', 'b', 'c'],
|
|
||||||
numbers: [1, 2, 3],
|
|
||||||
mixed: [1, 'two', { three: 3 }],
|
|
||||||
};
|
|
||||||
|
|
||||||
const lsn = await wal.logInsert('testdb', 'arraycoll', doc as any);
|
|
||||||
|
|
||||||
const entries = wal.getEntriesAfter(lsn - 1);
|
|
||||||
const entry = entries.find(e => e.lsn === lsn);
|
|
||||||
|
|
||||||
const recovered = wal.recoverDocument(entry!);
|
|
||||||
expect(recovered!.tags).toEqual(['a', 'b', 'c']);
|
|
||||||
expect(recovered!.numbers).toEqual([1, 2, 3]);
|
|
||||||
expect(recovered!.mixed[2].three).toEqual(3);
|
|
||||||
});
|
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Cleanup
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
tap.test('wal: cleanup', async () => {
|
|
||||||
await wal.close();
|
|
||||||
await cleanupTestFiles();
|
|
||||||
expect(true).toBeTrue();
|
|
||||||
});
|
|
||||||
|
|
||||||
export default tap.start();
|
|
||||||
@@ -0,0 +1,82 @@
|
|||||||
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
|
import * as smartdb from '../ts/index.js';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as net from 'net';
|
||||||
|
import * as path from 'path';
|
||||||
|
import * as os from 'os';
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Test: Stale socket cleanup on startup
|
||||||
|
// Covers: LocalSmartDb.cleanStaleSockets(), isSocketAlive()
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
function makeTmpDir(): string {
|
||||||
|
return fs.mkdtempSync(path.join(os.tmpdir(), 'smartdb-socket-test-'));
|
||||||
|
}
|
||||||
|
|
||||||
|
function cleanTmpDir(dir: string): void {
|
||||||
|
if (fs.existsSync(dir)) {
|
||||||
|
fs.rmSync(dir, { recursive: true, force: true });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Stale socket cleanup: active sockets are preserved
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('stale-sockets: does not remove active sockets', async () => {
|
||||||
|
const tmpDir = makeTmpDir();
|
||||||
|
const activeSocketPath = path.join(os.tmpdir(), `smartdb-active-${Date.now()}.sock`);
|
||||||
|
|
||||||
|
// Create an active socket (server still listening)
|
||||||
|
const activeServer = net.createServer();
|
||||||
|
await new Promise<void>((resolve) => activeServer.listen(activeSocketPath, resolve));
|
||||||
|
|
||||||
|
expect(fs.existsSync(activeSocketPath)).toBeTrue();
|
||||||
|
|
||||||
|
// Start LocalSmartDb — should NOT remove the active socket
|
||||||
|
const localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
await localDb.start();
|
||||||
|
|
||||||
|
expect(fs.existsSync(activeSocketPath)).toBeTrue();
|
||||||
|
|
||||||
|
// Cleanup
|
||||||
|
await localDb.stop();
|
||||||
|
await new Promise<void>((resolve) => activeServer.close(() => resolve()));
|
||||||
|
try { fs.unlinkSync(activeSocketPath); } catch {}
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Stale socket cleanup: startup works with no stale sockets
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('stale-sockets: startup works cleanly with no stale sockets', async () => {
|
||||||
|
const tmpDir = makeTmpDir();
|
||||||
|
const localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
expect(localDb.running).toBeTrue();
|
||||||
|
expect(info.socketPath).toBeTruthy();
|
||||||
|
await localDb.stop();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Stale socket cleanup: the socket file for the current instance is cleaned on stop
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('stale-sockets: own socket file is removed on stop', async () => {
|
||||||
|
const tmpDir = makeTmpDir();
|
||||||
|
const localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
|
||||||
|
expect(fs.existsSync(info.socketPath)).toBeTrue();
|
||||||
|
|
||||||
|
await localDb.stop();
|
||||||
|
|
||||||
|
// Socket file should be gone after graceful stop
|
||||||
|
expect(fs.existsSync(info.socketPath)).toBeFalse();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
export default tap.start();
|
||||||
@@ -0,0 +1,232 @@
|
|||||||
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
|
import * as smartdb from '../ts/index.js';
|
||||||
|
import { MongoClient } from 'mongodb';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as os from 'os';
|
||||||
|
import * as path from 'path';
|
||||||
|
|
||||||
|
let server: smartdb.SmartdbServer;
|
||||||
|
let tmpDir: string;
|
||||||
|
let storagePath: string;
|
||||||
|
let usersPath: string;
|
||||||
|
const port = 27129;
|
||||||
|
const openedClients: MongoClient[] = [];
|
||||||
|
|
||||||
|
let tenantA: smartdb.ISmartDbDatabaseTenantDescriptor;
|
||||||
|
let tenantB: smartdb.ISmartDbDatabaseTenantDescriptor;
|
||||||
|
let exportedTenantA: smartdb.ISmartDbDatabaseExport;
|
||||||
|
|
||||||
|
function makeTmpDir(): string {
|
||||||
|
return fs.mkdtempSync(path.join(os.tmpdir(), 'smartdb-tenants-test-'));
|
||||||
|
}
|
||||||
|
|
||||||
|
function cleanTmpDir(dir: string): void {
|
||||||
|
if (fs.existsSync(dir)) {
|
||||||
|
fs.rmSync(dir, { recursive: true, force: true });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async function connect(uri: string): Promise<MongoClient> {
|
||||||
|
const client = new MongoClient(uri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
openedClients.push(client);
|
||||||
|
return client;
|
||||||
|
}
|
||||||
|
|
||||||
|
async function expectConnectionToFail(uri: string): Promise<void> {
|
||||||
|
const client = new MongoClient(uri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
let threw = false;
|
||||||
|
try {
|
||||||
|
await client.connect();
|
||||||
|
await client.db('tenant_a').command({ ping: 1 });
|
||||||
|
} catch {
|
||||||
|
threw = true;
|
||||||
|
} finally {
|
||||||
|
await client.close().catch(() => undefined);
|
||||||
|
}
|
||||||
|
expect(threw).toBeTrue();
|
||||||
|
}
|
||||||
|
|
||||||
|
async function closeOpenedClients(): Promise<void> {
|
||||||
|
while (openedClients.length > 0) {
|
||||||
|
const client = openedClients.pop();
|
||||||
|
await client?.close().catch(() => undefined);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function createServer(): smartdb.SmartdbServer {
|
||||||
|
return new smartdb.SmartdbServer({
|
||||||
|
port,
|
||||||
|
storage: 'file',
|
||||||
|
storagePath,
|
||||||
|
auth: {
|
||||||
|
enabled: true,
|
||||||
|
usersPath,
|
||||||
|
scramIterations: 4096,
|
||||||
|
users: [
|
||||||
|
{
|
||||||
|
username: 'root',
|
||||||
|
password: 'secret',
|
||||||
|
database: 'admin',
|
||||||
|
roles: ['root'],
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
tap.test('tenants: should start durable authenticated service', async () => {
|
||||||
|
tmpDir = makeTmpDir();
|
||||||
|
storagePath = path.join(tmpDir, 'data');
|
||||||
|
usersPath = path.join(tmpDir, 'users.json');
|
||||||
|
server = createServer();
|
||||||
|
await server.start();
|
||||||
|
expect(server.running).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('tenants: should create isolated database tenants', async () => {
|
||||||
|
tenantA = await server.createDatabaseTenant({
|
||||||
|
databaseName: 'tenant_a',
|
||||||
|
username: 'tenant_a_user',
|
||||||
|
password: 'tenant-a-pass-1',
|
||||||
|
});
|
||||||
|
tenantB = await server.createDatabaseTenant({
|
||||||
|
databaseName: 'tenant_b',
|
||||||
|
username: 'tenant_b_user',
|
||||||
|
password: 'tenant-b-pass-1',
|
||||||
|
});
|
||||||
|
|
||||||
|
expect(tenantA.databaseName).toEqual('tenant_a');
|
||||||
|
expect(tenantA.authSource).toEqual('tenant_a');
|
||||||
|
expect(tenantA.roles.includes('readWrite')).toBeTrue();
|
||||||
|
expect(tenantA.roles.includes('dbAdmin')).toBeTrue();
|
||||||
|
expect(typeof tenantA.mongodbUri).toEqual('string');
|
||||||
|
|
||||||
|
const tenants = await server.listDatabaseTenants();
|
||||||
|
expect(tenants.some((tenant) => tenant.databaseName === 'tenant_a')).toBeTrue();
|
||||||
|
expect(tenants.some((tenant) => tenant.databaseName === 'tenant_b')).toBeTrue();
|
||||||
|
|
||||||
|
const descriptor = await server.getDatabaseTenantDescriptor({
|
||||||
|
databaseName: 'tenant_a',
|
||||||
|
username: 'tenant_a_user',
|
||||||
|
});
|
||||||
|
expect(descriptor.username).toEqual('tenant_a_user');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('tenants: should work with official MongoDB driver and enforce auth isolation', async () => {
|
||||||
|
const clientA = await connect(tenantA.mongodbUri!);
|
||||||
|
const clientB = await connect(tenantB.mongodbUri!);
|
||||||
|
|
||||||
|
const ping = await clientA.db('tenant_a').command({ ping: 1 });
|
||||||
|
expect(ping.ok).toEqual(1);
|
||||||
|
|
||||||
|
await clientA.db('tenant_a').collection('notes').insertOne({ title: 'tenant a note' });
|
||||||
|
await clientA.db('tenant_a').collection('notes').createIndex({ title: 1 });
|
||||||
|
await clientB.db('tenant_b').collection('notes').insertOne({ title: 'tenant b note' });
|
||||||
|
|
||||||
|
let threw = false;
|
||||||
|
try {
|
||||||
|
await clientA.db('tenant_b').collection('notes').findOne({ title: 'tenant b note' });
|
||||||
|
} catch (err: any) {
|
||||||
|
threw = true;
|
||||||
|
expect(err.code).toEqual(13);
|
||||||
|
}
|
||||||
|
expect(threw).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('tenants: should expose health and metrics for readiness checks', async () => {
|
||||||
|
const health = await server.getHealth();
|
||||||
|
expect(health.running).toBeTrue();
|
||||||
|
expect(health.storagePath).toEqual(storagePath);
|
||||||
|
expect(health.authEnabled).toBeTrue();
|
||||||
|
expect(health.databaseCount >= 2).toBeTrue();
|
||||||
|
expect(health.collectionCount >= 2).toBeTrue();
|
||||||
|
|
||||||
|
const metrics = await server.getMetrics();
|
||||||
|
expect(metrics.authEnabled).toBeTrue();
|
||||||
|
expect(metrics.databases >= 2).toBeTrue();
|
||||||
|
expect(metrics.collections >= 2).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('tenants: should rotate password without restart', async () => {
|
||||||
|
const oldUri = tenantA.mongodbUri!;
|
||||||
|
await closeOpenedClients();
|
||||||
|
|
||||||
|
tenantA = await server.rotateDatabaseTenantPassword({
|
||||||
|
username: 'tenant_a_user',
|
||||||
|
password: 'tenant-a-pass-2',
|
||||||
|
});
|
||||||
|
expect(typeof tenantA.mongodbUri).toEqual('string');
|
||||||
|
|
||||||
|
await expectConnectionToFail(oldUri);
|
||||||
|
const rotatedClient = await connect(tenantA.mongodbUri!);
|
||||||
|
const doc = await rotatedClient.db('tenant_a').collection('notes').findOne({ title: 'tenant a note' });
|
||||||
|
expect(doc).toBeTruthy();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('tenants: should persist runtime users and file-backed data across restart', async () => {
|
||||||
|
await closeOpenedClients();
|
||||||
|
await server.stop();
|
||||||
|
|
||||||
|
server = createServer();
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
const clientA = await connect(tenantA.mongodbUri!);
|
||||||
|
const clientB = await connect(tenantB.mongodbUri!);
|
||||||
|
const docA = await clientA.db('tenant_a').collection('notes').findOne({ title: 'tenant a note' });
|
||||||
|
const docB = await clientB.db('tenant_b').collection('notes').findOne({ title: 'tenant b note' });
|
||||||
|
expect(docA).toBeTruthy();
|
||||||
|
expect(docB).toBeTruthy();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('tenants: should export and restore one database without unrelated tenants', async () => {
|
||||||
|
exportedTenantA = await server.exportDatabase({ databaseName: 'tenant_a' });
|
||||||
|
expect(exportedTenantA.databaseName).toEqual('tenant_a');
|
||||||
|
expect(exportedTenantA.collections.length).toEqual(1);
|
||||||
|
expect(JSON.stringify(exportedTenantA).includes('tenant b note')).toBeFalse();
|
||||||
|
|
||||||
|
await closeOpenedClients();
|
||||||
|
const deleteResult = await server.deleteDatabaseTenant({
|
||||||
|
databaseName: 'tenant_a',
|
||||||
|
username: 'tenant_a_user',
|
||||||
|
});
|
||||||
|
expect(deleteResult.databaseDropped).toBeTrue();
|
||||||
|
expect(deleteResult.deletedUsers).toEqual(1);
|
||||||
|
|
||||||
|
await expectConnectionToFail(tenantA.mongodbUri!);
|
||||||
|
|
||||||
|
const importResult = await server.importDatabase({
|
||||||
|
databaseName: 'tenant_a',
|
||||||
|
source: exportedTenantA,
|
||||||
|
});
|
||||||
|
expect(importResult.databaseName).toEqual('tenant_a');
|
||||||
|
expect(importResult.documents).toEqual(1);
|
||||||
|
|
||||||
|
tenantA = await server.createDatabaseTenant({
|
||||||
|
databaseName: 'tenant_a',
|
||||||
|
username: 'tenant_a_user',
|
||||||
|
password: 'tenant-a-pass-3',
|
||||||
|
});
|
||||||
|
const restoredClient = await connect(tenantA.mongodbUri!);
|
||||||
|
const restoredDoc = await restoredClient.db('tenant_a').collection('notes').findOne({ title: 'tenant a note' });
|
||||||
|
expect(restoredDoc).toBeTruthy();
|
||||||
|
|
||||||
|
const clientB = await connect(tenantB.mongodbUri!);
|
||||||
|
const unrelatedDoc = await clientB.db('tenant_b').collection('notes').findOne({ title: 'tenant b note' });
|
||||||
|
expect(unrelatedDoc).toBeTruthy();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('tenants: cleanup', async () => {
|
||||||
|
await closeOpenedClients();
|
||||||
|
await server.stop();
|
||||||
|
expect(server.running).toBeFalse();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
export default tap.start();
|
||||||
@@ -0,0 +1,171 @@
|
|||||||
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
|
import * as smartdb from '../ts/index.js';
|
||||||
|
import { MongoClient } from 'mongodb';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as net from 'net';
|
||||||
|
import * as os from 'os';
|
||||||
|
import * as path from 'path';
|
||||||
|
|
||||||
|
// Static test-only CA and server certificate. The private key is intentionally
|
||||||
|
// non-secret test fixture material and must not be reused outside tests.
|
||||||
|
const CA_PEM = `-----BEGIN CERTIFICATE-----
|
||||||
|
MIIDFTCCAf2gAwIBAgIUXQlk6FLuWELDKLw9KXi0UIYmU50wDQYJKoZIhvcNAQEL
|
||||||
|
BQAwGjEYMBYGA1UEAwwPU21hcnREQiBUZXN0IENBMB4XDTI2MDQyOTIxMjYxNFoX
|
||||||
|
DTM2MDQyNjIxMjYxNFowGjEYMBYGA1UEAwwPU21hcnREQiBUZXN0IENBMIIBIjAN
|
||||||
|
BgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEApnRgZvodreKEKkSodwgDe2JKsA3N
|
||||||
|
GC4c7dmqmOBRQst0OYRoW0kjHnzCVHoGlMTAnjJWXRayPeJCroSA0WhEZIjgHAjW
|
||||||
|
FuWIr+MUYdCG7czdbDEqZYGsrBDUwv+ydgsDNhLKtbfVfcJckdmFp+TT+Po3sf8o
|
||||||
|
u5AfOlcjhM22reBLhZJ2FfM2IbqygRbBxNvU3tH5E1kgu2CpYieXQsmqBwkOPM0S
|
||||||
|
fgkCjlqFeeqV7Jjdq1P6srIItzg6n8/5KGBTxc7VB11WxVAZMIxnOtwpOCpSjbiy
|
||||||
|
jymBLKvyZxklWGpG9HT6RzUTdp0WpwnO7FlbYqD97jrbwA7PfhbJVUkTeQIDAQAB
|
||||||
|
o1MwUTAdBgNVHQ4EFgQUaqFWiFvibBYpJjluNW4XlocmqOQwHwYDVR0jBBgwFoAU
|
||||||
|
aqFWiFvibBYpJjluNW4XlocmqOQwDwYDVR0TAQH/BAUwAwEB/zANBgkqhkiG9w0B
|
||||||
|
AQsFAAOCAQEAdbmRCxeHwfq6Mw0BRXWYM81xrzDMDBwLkIyaVkBJXCEX4Ybj8QHv
|
||||||
|
tplNqgQae1Hr1qYyNzkivDI/hPnvv/wDsAnT8Wz0/udPpcASTXC03xhRtFXwBSGq
|
||||||
|
2GtLa53cZHJLoGu1S2ntM6Xo3gropXSx/+LIfefsQvqRO/5WxRrEE10OiFr19rA7
|
||||||
|
md0nD6zXdwrMRghu6ACuxX6Ext6QJbTL4r1UGbHg2a9UbdBjcb8sfFPLyEjiLpBK
|
||||||
|
DYvRjddKOwbOpFPoLwmed59Pa6bcqT9NnkRHL+aXUm3M3HfVhNKae7JJShUmCzdx
|
||||||
|
rbKNJQAUp/mMHnBOSxYS7aqgwBKCiKtP4A==
|
||||||
|
-----END CERTIFICATE-----
|
||||||
|
`;
|
||||||
|
|
||||||
|
const SERVER_CERT_PEM = `-----BEGIN CERTIFICATE-----
|
||||||
|
MIIDPTCCAiWgAwIBAgIUMfuX4VHvVJ8Vo6o1U2+f7MHU7dowDQYJKoZIhvcNAQEL
|
||||||
|
BQAwGjEYMBYGA1UEAwwPU21hcnREQiBUZXN0IENBMB4XDTI2MDQyOTIxMjYxNFoX
|
||||||
|
DTM2MDQyNjIxMjYxNFowFDESMBAGA1UEAwwJbG9jYWxob3N0MIIBIjANBgkqhkiG
|
||||||
|
9w0BAQEFAAOCAQ8AMIIBCgKCAQEA5eFz1q4juQsEE7cPN5eFrLvRJW/zOMGBmiet
|
||||||
|
VTQSqVZ/3j3NBWsgxK2xQnNbEXGMlTEE11ih0cCQacc/JnbuvwOt3QX8X6oy4pmb
|
||||||
|
LMGQJEk2FgdpP6OtGqqYbt/fT7QBY39nt6z/RzxYZI7t5g/nkHnlzmzD+ila6k9b
|
||||||
|
TzBSfSmtHHKW/c6az/Dh/xe50zDgrzlBA7e5zoleKqRJFRZlDnDoLyx0EOUbbTbQ
|
||||||
|
vipMynP5bq8l6Fc0N9DAWmXvV4o2x0ZQjfEx5LTvbxNkVWtv8w9w4t4vAZqXwrXd
|
||||||
|
5OZETMWdy7ezxL0E9Snwc6sSfatlVenD/8P5hWJ/C0vCiw21RwIDAQABo4GAMH4w
|
||||||
|
GgYDVR0RBBMwEYIJbG9jYWxob3N0hwR/AAABMAsGA1UdDwQEAwIFoDATBgNVHSUE
|
||||||
|
DDAKBggrBgEFBQcDATAdBgNVHQ4EFgQUK2nSXereMZek6gxLweY1AVt9OaswHwYD
|
||||||
|
VR0jBBgwFoAUaqFWiFvibBYpJjluNW4XlocmqOQwDQYJKoZIhvcNAQELBQADggEB
|
||||||
|
AAkC6suxamn+OEmJLMqgaGCvEtFbob5pMijYC32vJNPev+bUHMOB4Oo0FyO59sX3
|
||||||
|
zfLLwk7jagbWJi37T714aSjyJwUHd4XA7McSabP4+1hOOL0NqfiE4yRnxPhlvf3E
|
||||||
|
9otoStAAJ86067DwIs5id7jYm+qrxn6bL+P1h+P1tYxnPOoD0v1cHVbtUNV2tH2E
|
||||||
|
eBhdtTbF+NHrj+oXFGI3jiI7qcwpJ9DFUo/w0sC0POY0T5aWl4ptSXVgEc7nkE91
|
||||||
|
bbPOPyoMjjZ4WhKAW5UzfOafB0bO7+4E0GHcAkBJmS4V8g5qt56nftr+d58R/odY
|
||||||
|
0hQjpoIwzl9RCEW0h8xkqMQ=
|
||||||
|
-----END CERTIFICATE-----
|
||||||
|
`;
|
||||||
|
|
||||||
|
const SERVER_KEY_PEM = `-----BEGIN PRIVATE KEY-----
|
||||||
|
MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQDl4XPWriO5CwQT
|
||||||
|
tw83l4Wsu9Elb/M4wYGaJ61VNBKpVn/ePc0FayDErbFCc1sRcYyVMQTXWKHRwJBp
|
||||||
|
xz8mdu6/A63dBfxfqjLimZsswZAkSTYWB2k/o60aqphu399PtAFjf2e3rP9HPFhk
|
||||||
|
ju3mD+eQeeXObMP6KVrqT1tPMFJ9Ka0ccpb9zprP8OH/F7nTMOCvOUEDt7nOiV4q
|
||||||
|
pEkVFmUOcOgvLHQQ5RttNtC+KkzKc/luryXoVzQ30MBaZe9XijbHRlCN8THktO9v
|
||||||
|
E2RVa2/zD3Di3i8BmpfCtd3k5kRMxZ3Lt7PEvQT1KfBzqxJ9q2VV6cP/w/mFYn8L
|
||||||
|
S8KLDbVHAgMBAAECggEAAInWJR8US1cow8kOepFayUxJUZ6hAbWGUa+dGtF757Sh
|
||||||
|
qQoZBFW7ZmqHu0Gc6X4MF79dJQn6mwyp6e2DCtqFdaITEqz0ad7yrpAwilrLtSIM
|
||||||
|
w+FxkCoYejMDF2Nj2QJxbGO8gPQhRu/vvxCMoxjPcImwjZq4nMnjAiB8dMOGte9V
|
||||||
|
av/RoWUOFXqeiJHqAXiE372I4BupwYhGrSUQyuVj3SugDRbzvPepTQNRxaBJQPgy
|
||||||
|
4ZtZ8FjJdPFvlyxv6fmLFULHwPNcS6PLWPuwpj7oEQzG4/Q9ojYj4EPdpoOW7qoH
|
||||||
|
h1Y6ag1vk5A/m9DjvMhIDzmUJmq8mlldxqbCBpH0+QKBgQD3Eh7F0ZXdLQe/aG5t
|
||||||
|
ul9hTv68NZa5M0JzJinB6WjXl2s0bUgIvAE9ZmfUYHs8AMvTu4YwJqsrpMuzFOT9
|
||||||
|
Ct5wBSyFbPzVOt9MYE1Gipxx8RfEMSq7Sp0MjarX3h0Va8ry83NWzrN1CvyP8BQq
|
||||||
|
CuXo/IislCDgPg0uXhLD/7GsWQKBgQDuMEptldCKtpW6CdLdYih6xh0j1mdGU4Kb
|
||||||
|
7mTzo3OU3nDnGXGhqvJt/xpksPl7GPRHYQ1dqRzvLKHDtTJqhkedZBnE6A94LkVl
|
||||||
|
uNJnR8v4PkR9nKKg0uK2ug9VcfSiXUpl2yyYiDc123WjHdwH2U6BV3smb/7KwEvv
|
||||||
|
FWaP7PO6nwKBgAE2w5PxPa1ChWE5YCGF4uYVf0bpdH4gdFkgfOAJB4zXn504VDxG
|
||||||
|
wDLPB/+RIcnfryCxMS2XYwvp2V5d4eokXYdrXxagvHVHvsUfTAHmuHIO3zEFlNIq
|
||||||
|
wa7IG2jIHJh4WRzseUqZ5WPT0/3ZDiBOwWZtpzZB3A99/o6Vw73WycaxAoGAHTeR
|
||||||
|
OaYB4bIJ5bskwYEz4/N/SZEYM/k0cTop6fTnzaAHi2GEncchW7rKGwXWZHIoLMVL
|
||||||
|
5WxEH1aDNUV5vLVh/X1058FrfFt4qcSlEoQtEfNZZWscS8vygWWLUfjbgDsfUCU1
|
||||||
|
cDRtSU71PCACiHfweE8pzQo539b8uYQPg6IWN5MCgYA6z/kvGiBB9xFBUAJPsj+w
|
||||||
|
XW/UGbn7svZaCob+N5RA9Rs/0idv/bO2nAauZyHG/nn6HXII6U5pmRyVqWKhI22q
|
||||||
|
K3J0LCP42Zb6/eYzQPbP1jWHCMaL2QJQGsl4NMZixlnNJV0aG/5CButqzSC/cMbG
|
||||||
|
DX0n+YqqWmCgHWU2csnlAA==
|
||||||
|
-----END PRIVATE KEY-----
|
||||||
|
`;
|
||||||
|
|
||||||
|
let server: smartdb.SmartdbServer;
|
||||||
|
let client: MongoClient;
|
||||||
|
let tmpDir: string;
|
||||||
|
let caPath: string;
|
||||||
|
let certPath: string;
|
||||||
|
let keyPath: string;
|
||||||
|
let port: number;
|
||||||
|
|
||||||
|
function makeTmpDir(): string {
|
||||||
|
return fs.mkdtempSync(path.join(os.tmpdir(), 'smartdb-tls-test-'));
|
||||||
|
}
|
||||||
|
|
||||||
|
function cleanTmpDir(dir: string): void {
|
||||||
|
if (fs.existsSync(dir)) {
|
||||||
|
fs.rmSync(dir, { recursive: true, force: true });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async function getFreePort(): Promise<number> {
|
||||||
|
return await new Promise((resolve, reject) => {
|
||||||
|
const probe = net.createServer();
|
||||||
|
probe.once('error', reject);
|
||||||
|
probe.listen(0, '127.0.0.1', () => {
|
||||||
|
const address = probe.address();
|
||||||
|
if (!address || typeof address === 'string') {
|
||||||
|
probe.close(() => reject(new Error('Failed to allocate TCP port')));
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
probe.close(() => resolve(address.port));
|
||||||
|
});
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
tap.test('tls: should start server with TLS enabled', async () => {
|
||||||
|
tmpDir = makeTmpDir();
|
||||||
|
port = await getFreePort();
|
||||||
|
caPath = path.join(tmpDir, 'ca.pem');
|
||||||
|
certPath = path.join(tmpDir, 'server.pem');
|
||||||
|
keyPath = path.join(tmpDir, 'server.key');
|
||||||
|
|
||||||
|
fs.writeFileSync(caPath, CA_PEM);
|
||||||
|
fs.writeFileSync(certPath, SERVER_CERT_PEM);
|
||||||
|
fs.writeFileSync(keyPath, SERVER_KEY_PEM, { mode: 0o600 });
|
||||||
|
|
||||||
|
server = new smartdb.SmartdbServer({
|
||||||
|
port,
|
||||||
|
tls: {
|
||||||
|
enabled: true,
|
||||||
|
certPath,
|
||||||
|
keyPath,
|
||||||
|
},
|
||||||
|
});
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
expect(server.running).toBeTrue();
|
||||||
|
expect(server.getConnectionUri()).toEqual(`mongodb://127.0.0.1:${port}/?tls=true`);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('tls: should connect with official MongoClient and CA validation', async () => {
|
||||||
|
client = new MongoClient(server.getConnectionUri(), {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
tlsCAFile: caPath,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
|
||||||
|
const ping = await client.db('admin').command({ ping: 1 });
|
||||||
|
expect(ping.ok).toEqual(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('tls: should support CRUD over encrypted transport', async () => {
|
||||||
|
const collection = client.db('tlsdb').collection('notes');
|
||||||
|
const inserted = await collection.insertOne({ title: 'encrypted transport' });
|
||||||
|
expect(inserted.acknowledged).toBeTrue();
|
||||||
|
|
||||||
|
const doc = await collection.findOne({ _id: inserted.insertedId });
|
||||||
|
expect(doc).toBeTruthy();
|
||||||
|
expect(doc!.title).toEqual('encrypted transport');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('tls: cleanup', async () => {
|
||||||
|
await client.close();
|
||||||
|
await server.stop();
|
||||||
|
expect(server.running).toBeFalse();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
export default tap.start();
|
||||||
@@ -0,0 +1,160 @@
|
|||||||
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
|
import * as smartdb from '../ts/index.js';
|
||||||
|
import { MongoClient } from 'mongodb';
|
||||||
|
import * as net from 'net';
|
||||||
|
|
||||||
|
let server: smartdb.SmartdbServer;
|
||||||
|
let client: MongoClient;
|
||||||
|
let port: number;
|
||||||
|
|
||||||
|
async function getFreePort(): Promise<number> {
|
||||||
|
return await new Promise((resolve, reject) => {
|
||||||
|
const probe = net.createServer();
|
||||||
|
probe.once('error', reject);
|
||||||
|
probe.listen(0, '127.0.0.1', () => {
|
||||||
|
const address = probe.address();
|
||||||
|
if (!address || typeof address === 'string') {
|
||||||
|
probe.close(() => reject(new Error('Failed to allocate TCP port')));
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
probe.close(() => resolve(address.port));
|
||||||
|
});
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
tap.test('transactions: should start server and connect', async () => {
|
||||||
|
port = await getFreePort();
|
||||||
|
server = new smartdb.SmartdbServer({ port });
|
||||||
|
await server.start();
|
||||||
|
|
||||||
|
client = new MongoClient(`mongodb://127.0.0.1:${port}`, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
expect(server.running).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('transactions: should still support explicit sessions', async () => {
|
||||||
|
const result = await client.db('admin').command({ startSession: 1 });
|
||||||
|
expect(result.ok).toEqual(1);
|
||||||
|
expect(result.id).toBeTruthy();
|
||||||
|
|
||||||
|
const end = await client.db('admin').command({ endSessions: [result.id] });
|
||||||
|
expect(end.ok).toEqual(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('transactions: should reject transaction-scoped writes without txnNumber before mutation', async () => {
|
||||||
|
const db = client.db('txntest');
|
||||||
|
const coll = db.collection('docs');
|
||||||
|
await coll.insertOne({ key: 'outside', value: 1 });
|
||||||
|
|
||||||
|
let threw = false;
|
||||||
|
try {
|
||||||
|
await db.command({
|
||||||
|
insert: 'docs',
|
||||||
|
documents: [{ key: 'inside-raw', value: 2 }],
|
||||||
|
startTransaction: true,
|
||||||
|
autocommit: false,
|
||||||
|
});
|
||||||
|
} catch (err: any) {
|
||||||
|
threw = true;
|
||||||
|
expect(err.code).toEqual(14);
|
||||||
|
expect(err.codeName).toEqual('TypeMismatch');
|
||||||
|
}
|
||||||
|
expect(threw).toBeTrue();
|
||||||
|
|
||||||
|
expect(await coll.countDocuments({ key: 'inside-raw' })).toEqual(0);
|
||||||
|
expect(await coll.countDocuments({ key: 'outside' })).toEqual(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('transactions: official driver transaction should commit buffered writes', async () => {
|
||||||
|
const coll = client.db('txntest').collection('driverdocs');
|
||||||
|
await coll.insertOne({ key: 'outside-driver', value: 0 });
|
||||||
|
const session = client.startSession();
|
||||||
|
|
||||||
|
try {
|
||||||
|
session.startTransaction();
|
||||||
|
await coll.insertOne({ key: 'inside-driver', value: 1 }, { session });
|
||||||
|
const inTxn = await coll.findOne({ key: 'inside-driver' }, { session });
|
||||||
|
expect(inTxn).toBeTruthy();
|
||||||
|
expect(await coll.countDocuments({ key: 'inside-driver' })).toEqual(0);
|
||||||
|
await session.commitTransaction();
|
||||||
|
} finally {
|
||||||
|
await session.endSession();
|
||||||
|
}
|
||||||
|
|
||||||
|
expect(await coll.countDocuments({ key: 'inside-driver' })).toEqual(1);
|
||||||
|
expect(await coll.countDocuments({ key: 'outside-driver' })).toEqual(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('transactions: abort should discard buffered writes', async () => {
|
||||||
|
const coll = client.db('txntest').collection('abortdocs');
|
||||||
|
const session = client.startSession();
|
||||||
|
|
||||||
|
try {
|
||||||
|
session.startTransaction();
|
||||||
|
await coll.insertOne({ key: 'abort-me', value: 1 }, { session });
|
||||||
|
expect(await coll.findOne({ key: 'abort-me' }, { session })).toBeTruthy();
|
||||||
|
await session.abortTransaction();
|
||||||
|
} finally {
|
||||||
|
await session.endSession();
|
||||||
|
}
|
||||||
|
|
||||||
|
expect(await coll.findOne({ key: 'abort-me' })).toBeNull();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('transactions: update and delete should commit atomically', async () => {
|
||||||
|
const coll = client.db('txntest').collection('mutations');
|
||||||
|
await coll.insertMany([
|
||||||
|
{ key: 'update-me', value: 1 },
|
||||||
|
{ key: 'delete-me', value: 2 },
|
||||||
|
]);
|
||||||
|
const session = client.startSession();
|
||||||
|
|
||||||
|
try {
|
||||||
|
session.startTransaction();
|
||||||
|
await coll.updateOne({ key: 'update-me' }, { $set: { value: 10 } }, { session });
|
||||||
|
await coll.deleteOne({ key: 'delete-me' }, { session });
|
||||||
|
expect((await coll.findOne({ key: 'update-me' }, { session }))!.value).toEqual(10);
|
||||||
|
expect(await coll.findOne({ key: 'delete-me' }, { session })).toBeNull();
|
||||||
|
expect((await coll.findOne({ key: 'update-me' }))!.value).toEqual(1);
|
||||||
|
expect(await coll.findOne({ key: 'delete-me' })).toBeTruthy();
|
||||||
|
await session.commitTransaction();
|
||||||
|
} finally {
|
||||||
|
await session.endSession();
|
||||||
|
}
|
||||||
|
|
||||||
|
expect((await coll.findOne({ key: 'update-me' }))!.value).toEqual(10);
|
||||||
|
expect(await coll.findOne({ key: 'delete-me' })).toBeNull();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('transactions: commit and abort without active transaction should be explicit errors', async () => {
|
||||||
|
for (const command of [{ commitTransaction: 1 }, { abortTransaction: 1 }]) {
|
||||||
|
let threw = false;
|
||||||
|
try {
|
||||||
|
await client.db('admin').command(command);
|
||||||
|
} catch (err: any) {
|
||||||
|
threw = true;
|
||||||
|
expect(err.code).toEqual(251);
|
||||||
|
expect(err.codeName).toEqual('NoSuchTransaction');
|
||||||
|
}
|
||||||
|
expect(threw).toBeTrue();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('transactions: serverStatus should expose transaction and oplog metrics', async () => {
|
||||||
|
const status = await client.db('admin').command({ serverStatus: 1 });
|
||||||
|
expect(status.ok).toEqual(1);
|
||||||
|
expect(status.transactions.currentActive).toEqual(0);
|
||||||
|
expect(status.logicalSessionRecordCache.activeSessionsCount).toBeGreaterThanOrEqual(0);
|
||||||
|
expect(status.oplog.totalEntries).toBeGreaterThan(0);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('transactions: cleanup', async () => {
|
||||||
|
await client.close();
|
||||||
|
await server.stop();
|
||||||
|
expect(server.running).toBeFalse();
|
||||||
|
});
|
||||||
|
|
||||||
|
export default tap.start();
|
||||||
@@ -0,0 +1,180 @@
|
|||||||
|
import { expect, tap } from '@git.zone/tstest/tapbundle';
|
||||||
|
import * as smartdb from '../ts/index.js';
|
||||||
|
import { MongoClient, Db } from 'mongodb';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as path from 'path';
|
||||||
|
import * as os from 'os';
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Test: Unique index enforcement via wire protocol
|
||||||
|
// Covers: unique index pre-check, createIndexes persistence, index restoration
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
let tmpDir: string;
|
||||||
|
let localDb: smartdb.LocalSmartDb;
|
||||||
|
let client: MongoClient;
|
||||||
|
let db: Db;
|
||||||
|
|
||||||
|
function makeTmpDir(): string {
|
||||||
|
return fs.mkdtempSync(path.join(os.tmpdir(), 'smartdb-unique-test-'));
|
||||||
|
}
|
||||||
|
|
||||||
|
function cleanTmpDir(dir: string): void {
|
||||||
|
if (fs.existsSync(dir)) {
|
||||||
|
fs.rmSync(dir, { recursive: true, force: true });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Setup
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('setup: start local db', async () => {
|
||||||
|
tmpDir = makeTmpDir();
|
||||||
|
localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
client = new MongoClient(info.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('uniquetest');
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Unique index enforcement on insert
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('unique-index: createIndex with unique: true', async () => {
|
||||||
|
const coll = db.collection('users');
|
||||||
|
await coll.insertOne({ email: 'alice@example.com', name: 'Alice' });
|
||||||
|
const indexName = await coll.createIndex({ email: 1 }, { unique: true });
|
||||||
|
expect(indexName).toBeTruthy();
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('unique-index: reject duplicate on insertOne', async () => {
|
||||||
|
const coll = db.collection('users');
|
||||||
|
let threw = false;
|
||||||
|
try {
|
||||||
|
await coll.insertOne({ email: 'alice@example.com', name: 'Alice2' });
|
||||||
|
} catch (err: any) {
|
||||||
|
threw = true;
|
||||||
|
expect(err.code).toEqual(11000);
|
||||||
|
}
|
||||||
|
expect(threw).toBeTrue();
|
||||||
|
|
||||||
|
// Verify only 1 document exists
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('unique-index: allow insert with different unique value', async () => {
|
||||||
|
const coll = db.collection('users');
|
||||||
|
await coll.insertOne({ email: 'bob@example.com', name: 'Bob' });
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(2);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Unique index enforcement on update
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('unique-index: reject duplicate on updateOne that changes unique field', async () => {
|
||||||
|
const coll = db.collection('users');
|
||||||
|
let threw = false;
|
||||||
|
try {
|
||||||
|
await coll.updateOne(
|
||||||
|
{ email: 'bob@example.com' },
|
||||||
|
{ $set: { email: 'alice@example.com' } }
|
||||||
|
);
|
||||||
|
} catch (err: any) {
|
||||||
|
threw = true;
|
||||||
|
expect(err.code).toEqual(11000);
|
||||||
|
}
|
||||||
|
expect(threw).toBeTrue();
|
||||||
|
|
||||||
|
// Bob's email should be unchanged
|
||||||
|
const bob = await coll.findOne({ name: 'Bob' });
|
||||||
|
expect(bob!.email).toEqual('bob@example.com');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('unique-index: allow update that keeps same unique value', async () => {
|
||||||
|
const coll = db.collection('users');
|
||||||
|
await coll.updateOne(
|
||||||
|
{ email: 'bob@example.com' },
|
||||||
|
{ $set: { name: 'Robert' } }
|
||||||
|
);
|
||||||
|
const bob = await coll.findOne({ email: 'bob@example.com' });
|
||||||
|
expect(bob!.name).toEqual('Robert');
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Unique index enforcement on upsert
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('unique-index: reject duplicate on upsert insert', async () => {
|
||||||
|
const coll = db.collection('users');
|
||||||
|
let threw = false;
|
||||||
|
try {
|
||||||
|
await coll.updateOne(
|
||||||
|
{ email: 'new@example.com' },
|
||||||
|
{ $set: { email: 'alice@example.com', name: 'Imposter' } },
|
||||||
|
{ upsert: true }
|
||||||
|
);
|
||||||
|
} catch (err: any) {
|
||||||
|
threw = true;
|
||||||
|
}
|
||||||
|
expect(threw).toBeTrue();
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Unique index survives restart (persistence + restoration)
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('unique-index: stop and restart', async () => {
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
|
||||||
|
localDb = new smartdb.LocalSmartDb({ folderPath: tmpDir });
|
||||||
|
const info = await localDb.start();
|
||||||
|
client = new MongoClient(info.connectionUri, {
|
||||||
|
directConnection: true,
|
||||||
|
serverSelectionTimeoutMS: 5000,
|
||||||
|
});
|
||||||
|
await client.connect();
|
||||||
|
db = client.db('uniquetest');
|
||||||
|
});
|
||||||
|
|
||||||
|
tap.test('unique-index: enforcement persists after restart', async () => {
|
||||||
|
const coll = db.collection('users');
|
||||||
|
|
||||||
|
// Data should still be there
|
||||||
|
const count = await coll.countDocuments();
|
||||||
|
expect(count).toEqual(2);
|
||||||
|
|
||||||
|
// Unique constraint should still be enforced without calling createIndex again
|
||||||
|
let threw = false;
|
||||||
|
try {
|
||||||
|
await coll.insertOne({ email: 'alice@example.com', name: 'Alice3' });
|
||||||
|
} catch (err: any) {
|
||||||
|
threw = true;
|
||||||
|
expect(err.code).toEqual(11000);
|
||||||
|
}
|
||||||
|
expect(threw).toBeTrue();
|
||||||
|
|
||||||
|
// Count unchanged
|
||||||
|
const countAfter = await coll.countDocuments();
|
||||||
|
expect(countAfter).toEqual(2);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Cleanup
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
tap.test('unique-index: cleanup', async () => {
|
||||||
|
await client.close();
|
||||||
|
await localDb.stop();
|
||||||
|
cleanTmpDir(tmpDir);
|
||||||
|
});
|
||||||
|
|
||||||
|
export default tap.start();
|
||||||
@@ -3,6 +3,6 @@
|
|||||||
*/
|
*/
|
||||||
export const commitinfo = {
|
export const commitinfo = {
|
||||||
name: '@push.rocks/smartdb',
|
name: '@push.rocks/smartdb',
|
||||||
version: '1.0.1',
|
version: '2.9.0',
|
||||||
description: 'A pure TypeScript MongoDB wire-protocol-compatible database server with pluggable storage, indexing, transactions, and zero external binary dependencies.'
|
description: 'A MongoDB-compatible embedded database server with wire protocol support, backed by a high-performance Rust engine.'
|
||||||
}
|
}
|
||||||
|
|||||||
+24
@@ -7,5 +7,29 @@ export * from './ts_smartdb/index.js';
|
|||||||
export { LocalSmartDb } from './ts_local/index.js';
|
export { LocalSmartDb } from './ts_local/index.js';
|
||||||
export type { ILocalSmartDbOptions, ILocalSmartDbConnectionInfo } from './ts_local/index.js';
|
export type { ILocalSmartDbOptions, ILocalSmartDbConnectionInfo } from './ts_local/index.js';
|
||||||
|
|
||||||
|
// Export migration
|
||||||
|
export { StorageMigrator } from './ts_migration/index.js';
|
||||||
|
|
||||||
// Export commitinfo
|
// Export commitinfo
|
||||||
export { commitinfo };
|
export { commitinfo };
|
||||||
|
|
||||||
|
// Re-export oplog / debug types for convenience
|
||||||
|
export type {
|
||||||
|
IOpLogEntry,
|
||||||
|
IOpLogResult,
|
||||||
|
IOpLogStats,
|
||||||
|
IRevertResult,
|
||||||
|
ICollectionInfo,
|
||||||
|
IDocumentsResult,
|
||||||
|
ISmartDbMetrics,
|
||||||
|
ISmartDbHealth,
|
||||||
|
ISmartDbDatabaseTenantInput,
|
||||||
|
ISmartDbDeleteDatabaseTenantInput,
|
||||||
|
ISmartDbRotateDatabaseTenantPasswordInput,
|
||||||
|
ISmartDbDatabaseTenantDescriptor,
|
||||||
|
ISmartDbDeleteDatabaseTenantResult,
|
||||||
|
ISmartDbDatabaseExportCollection,
|
||||||
|
ISmartDbDatabaseExport,
|
||||||
|
ISmartDbImportDatabaseInput,
|
||||||
|
ISmartDbImportDatabaseResult,
|
||||||
|
} from './ts_smartdb/index.js';
|
||||||
|
|||||||
@@ -1,8 +1,10 @@
|
|||||||
import * as plugins from './plugins.js';
|
|
||||||
import * as crypto from 'crypto';
|
import * as crypto from 'crypto';
|
||||||
|
import * as fs from 'fs/promises';
|
||||||
|
import * as net from 'net';
|
||||||
import * as path from 'path';
|
import * as path from 'path';
|
||||||
import * as os from 'os';
|
import * as os from 'os';
|
||||||
import { SmartdbServer } from '../ts_smartdb/index.js';
|
import { SmartdbServer } from '../ts_smartdb/index.js';
|
||||||
|
import { StorageMigrator } from '../ts_migration/index.js';
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Connection information returned by LocalSmartDb.start()
|
* Connection information returned by LocalSmartDb.start()
|
||||||
@@ -10,7 +12,7 @@ import { SmartdbServer } from '../ts_smartdb/index.js';
|
|||||||
export interface ILocalSmartDbConnectionInfo {
|
export interface ILocalSmartDbConnectionInfo {
|
||||||
/** The Unix socket file path */
|
/** The Unix socket file path */
|
||||||
socketPath: string;
|
socketPath: string;
|
||||||
/** MongoDB connection URI ready for MongoClient */
|
/** Connection URI (mongodb:// scheme) ready for MongoClient */
|
||||||
connectionUri: string;
|
connectionUri: string;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -25,8 +27,8 @@ export interface ILocalSmartDbOptions {
|
|||||||
* LocalSmartDb - Lightweight local MongoDB-compatible database using Unix sockets
|
* LocalSmartDb - Lightweight local MongoDB-compatible database using Unix sockets
|
||||||
*
|
*
|
||||||
* This class wraps SmartdbServer and provides a simple interface for
|
* This class wraps SmartdbServer and provides a simple interface for
|
||||||
* starting a local file-based MongoDB-compatible server. Returns connection
|
* starting a local file-based database server. Returns connection
|
||||||
* info that you can use with your own MongoDB driver instance.
|
* info that you can use with any compatible driver instance.
|
||||||
*
|
*
|
||||||
* @example
|
* @example
|
||||||
* ```typescript
|
* ```typescript
|
||||||
@@ -36,11 +38,11 @@ export interface ILocalSmartDbOptions {
|
|||||||
* const db = new LocalSmartDb({ folderPath: './data' });
|
* const db = new LocalSmartDb({ folderPath: './data' });
|
||||||
* const { connectionUri } = await db.start();
|
* const { connectionUri } = await db.start();
|
||||||
*
|
*
|
||||||
* // Connect with your own MongoDB client
|
* // Connect with the driver
|
||||||
* const client = new MongoClient(connectionUri, { directConnection: true });
|
* const client = new MongoClient(connectionUri, { directConnection: true });
|
||||||
* await client.connect();
|
* await client.connect();
|
||||||
*
|
*
|
||||||
* // Use the MongoDB client
|
* // Use the client
|
||||||
* const collection = client.db('mydb').collection('users');
|
* const collection = client.db('mydb').collection('users');
|
||||||
* await collection.insertOne({ name: 'Alice' });
|
* await collection.insertOne({ name: 'Alice' });
|
||||||
*
|
*
|
||||||
@@ -66,6 +68,55 @@ export class LocalSmartDb {
|
|||||||
return path.join(os.tmpdir(), `smartdb-${randomId}.sock`);
|
return path.join(os.tmpdir(), `smartdb-${randomId}.sock`);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Check if a Unix socket is alive by attempting to connect.
|
||||||
|
*/
|
||||||
|
private static isSocketAlive(socketPath: string): Promise<boolean> {
|
||||||
|
return new Promise((resolve) => {
|
||||||
|
const client = net.createConnection({ path: socketPath }, () => {
|
||||||
|
client.destroy();
|
||||||
|
resolve(true);
|
||||||
|
});
|
||||||
|
client.on('error', () => {
|
||||||
|
resolve(false);
|
||||||
|
});
|
||||||
|
client.setTimeout(500, () => {
|
||||||
|
client.destroy();
|
||||||
|
resolve(false);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Remove stale smartdb-*.sock files from /tmp.
|
||||||
|
* A socket is considered stale if connecting to it fails.
|
||||||
|
*/
|
||||||
|
private static async cleanStaleSockets(): Promise<void> {
|
||||||
|
const tmpDir = os.tmpdir();
|
||||||
|
let entries: string[];
|
||||||
|
try {
|
||||||
|
entries = await fs.readdir(tmpDir);
|
||||||
|
} catch {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
const socketFiles = entries.filter(
|
||||||
|
(f) => f.startsWith('smartdb-') && f.endsWith('.sock')
|
||||||
|
);
|
||||||
|
for (const name of socketFiles) {
|
||||||
|
const fullPath = path.join(tmpDir, name);
|
||||||
|
try {
|
||||||
|
const stat = await fs.stat(fullPath);
|
||||||
|
if (!stat.isSocket()) continue;
|
||||||
|
const alive = await LocalSmartDb.isSocketAlive(fullPath);
|
||||||
|
if (!alive) {
|
||||||
|
await fs.unlink(fullPath);
|
||||||
|
}
|
||||||
|
} catch {
|
||||||
|
// File may have been removed already; ignore
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Start the local SmartDB server and return connection info
|
* Start the local SmartDB server and return connection info
|
||||||
*/
|
*/
|
||||||
@@ -74,6 +125,13 @@ export class LocalSmartDb {
|
|||||||
throw new Error('LocalSmartDb is already running');
|
throw new Error('LocalSmartDb is already running');
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Clean up stale sockets from previous crashed instances
|
||||||
|
await LocalSmartDb.cleanStaleSockets();
|
||||||
|
|
||||||
|
// Run storage migration before starting the Rust engine
|
||||||
|
const migrator = new StorageMigrator(this.options.folderPath);
|
||||||
|
await migrator.run();
|
||||||
|
|
||||||
// Use provided socket path or generate one
|
// Use provided socket path or generate one
|
||||||
this.generatedSocketPath = this.options.socketPath ?? this.generateSocketPath();
|
this.generatedSocketPath = this.options.socketPath ?? this.generateSocketPath();
|
||||||
|
|
||||||
|
|||||||
@@ -1,3 +1 @@
|
|||||||
import * as smartpromise from '@push.rocks/smartpromise';
|
// Local module plugins - currently no external dependencies needed
|
||||||
|
|
||||||
export { smartpromise };
|
|
||||||
|
|||||||
@@ -0,0 +1,93 @@
|
|||||||
|
import * as fs from 'fs';
|
||||||
|
import * as path from 'path';
|
||||||
|
import { migrateV0ToV1 } from './migrators/v0_to_v1.js';
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Detected storage format version.
|
||||||
|
* - v0: Legacy JSON format ({db}/{coll}.json files)
|
||||||
|
* - v1: Bitcask binary format ({db}/{coll}/data.rdb directories)
|
||||||
|
*/
|
||||||
|
type TStorageVersion = 0 | 1;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* StorageMigrator — runs before the Rust engine starts.
|
||||||
|
*
|
||||||
|
* Detects the current storage format version and runs the appropriate
|
||||||
|
* migration chain. The Rust engine only knows the current format (v1).
|
||||||
|
*
|
||||||
|
* Migration is safe: original files are never modified or deleted.
|
||||||
|
* On success, a console hint is printed about which old files can be removed.
|
||||||
|
*/
|
||||||
|
export class StorageMigrator {
|
||||||
|
private storagePath: string;
|
||||||
|
|
||||||
|
constructor(storagePath: string) {
|
||||||
|
this.storagePath = storagePath;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Run any needed migrations. Safe to call even if storage is already current.
|
||||||
|
*/
|
||||||
|
async run(): Promise<void> {
|
||||||
|
if (!fs.existsSync(this.storagePath)) {
|
||||||
|
return; // No data yet — nothing to migrate
|
||||||
|
}
|
||||||
|
|
||||||
|
const version = this.detectVersion();
|
||||||
|
|
||||||
|
if (version === 1) {
|
||||||
|
return; // Already current
|
||||||
|
}
|
||||||
|
|
||||||
|
if (version === 0) {
|
||||||
|
console.log(`[smartdb] Detected v0 (JSON) storage format at ${this.storagePath}`);
|
||||||
|
console.log(`[smartdb] Running migration v0 → v1 (Bitcask binary format)...`);
|
||||||
|
|
||||||
|
const deletableFiles = await migrateV0ToV1(this.storagePath);
|
||||||
|
|
||||||
|
if (deletableFiles.length > 0) {
|
||||||
|
console.log(`[smartdb] Migration v0 → v1 complete.`);
|
||||||
|
console.log(`[smartdb] The following old files can be safely deleted:`);
|
||||||
|
for (const f of deletableFiles) {
|
||||||
|
console.log(`[smartdb] ${f}`);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
console.log(`[smartdb] Migration v0 → v1 complete. No old files to clean up.`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Detect the storage format version by inspecting the directory structure.
|
||||||
|
*
|
||||||
|
* v0: {db}/{coll}.json files exist
|
||||||
|
* v1: {db}/{coll}/data.rdb directories exist
|
||||||
|
*/
|
||||||
|
private detectVersion(): TStorageVersion {
|
||||||
|
const entries = fs.readdirSync(this.storagePath, { withFileTypes: true });
|
||||||
|
|
||||||
|
for (const entry of entries) {
|
||||||
|
if (!entry.isDirectory()) continue;
|
||||||
|
|
||||||
|
const dbDir = path.join(this.storagePath, entry.name);
|
||||||
|
const dbEntries = fs.readdirSync(dbDir, { withFileTypes: true });
|
||||||
|
|
||||||
|
for (const dbEntry of dbEntries) {
|
||||||
|
// v1: subdirectory with data.rdb
|
||||||
|
if (dbEntry.isDirectory()) {
|
||||||
|
const dataRdb = path.join(dbDir, dbEntry.name, 'data.rdb');
|
||||||
|
if (fs.existsSync(dataRdb)) {
|
||||||
|
return 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// v0: .json file (not .indexes.json)
|
||||||
|
if (dbEntry.isFile() && dbEntry.name.endsWith('.json') && !dbEntry.name.endsWith('.indexes.json')) {
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Empty or unrecognized — treat as v1 (fresh start)
|
||||||
|
return 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1 @@
|
|||||||
|
export { StorageMigrator } from './classes.storagemigrator.js';
|
||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user