codekingpro/portable-devtools

1import multiprocessing
2import multiprocessing.context
3import sys
4import time
5from multiprocessing.synchronize import Event
6
7import chromadb
8from chromadb.api.client import Client
9from chromadb.api.models.Collection import Collection
10from chromadb.cli import cli
11from chromadb.cli.cli import build_cli_args
12from chromadb.config import Settings, System
13from chromadb.db.base import get_sql
14from chromadb.db.impl.sqlite import SqliteDB
15from pypika import Table
16import numpy as np
17
18from chromadb.test.property import invariants
19
20
21def wait_for_server(
22        host: str, port: int,
23    max_retries: int = 5, initial_delay: float = 1.0
24) -> bool:
25    """Wait for server to be ready using exponential backoff.
26    Args:
27        client: ChromaDB client instance
28        max_retries: Maximum number of retry attempts
29        initial_delay: Initial delay in seconds before first retry
30    Returns:
31        bool: True if server is ready, False if max retries exceeded
32    """
33    delay = initial_delay
34    for attempt in range(max_retries):
35        try:
36            client = chromadb.HttpClient(host=host, port=port)
37            heartbeat = client.heartbeat()
38            if heartbeat > 0:
39                return True
40        except Exception:
41            print("Heartbeat failed, trying again...")
42            pass
43
44        if attempt < max_retries - 1:
45            time.sleep(delay)
46            delay *= 2
47
48    return False
49
50def start_app(args: list[str]) -> None:
51    sys.argv = args
52    cli.app()
53
54def test_app() -> None:
55    kwargs = {"path": "chroma_test_data", "port": 8001}
56    args = ["chroma", "run"]
57    args.extend(build_cli_args(**kwargs))
58    print(args)
59    server_process = multiprocessing.Process(target=start_app, args=(args,))
60    server_process.start()
61    time.sleep(5)
62
63    assert wait_for_server(host="localhost", port=8001), "Server failed to start within maximum retry attempts"
64
65    server_process.terminate()
66    server_process.join()
67
68
69def test_vacuum(sqlite_persistent: System) -> None:
70    system = sqlite_persistent
71    sqlite = system.instance(SqliteDB)
72
73    # This is True because it's a fresh system, so let's set it to False to test that the vacuum command enables it
74    config = sqlite.config
75    config.set_parameter("automatically_purge", False)
76    sqlite.set_config(config)
77
78    # Add some data
79    client = Client.from_system(system)
80    collection1 = client.create_collection("collection1")
81    collection2 = client.create_collection("collection2")
82
83    def add_records(collection: Collection, num: int) -> None:
84        ids = [str(i) for i in range(num)]
85        embeddings = np.random.rand(num, 2)
86        collection.add(ids=ids, embeddings=embeddings)
87
88    add_records(collection1, 100)
89    add_records(collection2, 2_000)
90
91    # Maintenance log should be empty
92    with sqlite.tx() as cur:
93        t = Table("maintenance_log")
94        q = sqlite.querybuilder().from_(t).select("*")
95        sql, params = get_sql(q)
96        cur.execute(sql, params)
97        assert cur.fetchall() == []
98
99    sys.argv = ["chroma", "vacuum", "--path", system.settings.persist_directory, "--force"]
100    cli.app()
101
102    # Maintenance log should have a vacuum entry
103    with sqlite.tx() as cur:
104        t = Table("maintenance_log")
105        q = sqlite.querybuilder().from_(t).select("*")
106        sql, params = get_sql(q)
107        cur.execute(sql, params)
108        rows = cur.fetchall()
109        assert len(rows) == 1
110        assert rows[0][2] == "vacuum"
111
112    # Automatic pruning should have been enabled
113    if hasattr(sqlite, "config"):
114        del (
115            sqlite.config
116        )  # the CLI will end up starting a new instance of sqlite, so we need to force-refresh the cached config here
117    assert sqlite.config.get_parameter("automatically_purge").value
118
119    # Log should be clean
120    invariants.log_size_below_max(system, [collection1, collection2], True)
121
122
123def simulate_transactional_write(
124    settings: Settings, ready_event: Event, shutdown_event: Event
125) -> None:
126    system = System(settings=settings)
127    system.start()
128    sqlite = system.instance(SqliteDB)
129
130    with sqlite.tx() as cur:
131        cur.execute("INSERT INTO tenants DEFAULT VALUES")
132        ready_event.set()
133        shutdown_event.wait()
134
135    system.stop()
136
137
138def test_vacuum_errors_if_locked(sqlite_persistent: System, capfd) -> None:
139    """Vacuum command should fail with details if there is a long-lived lock on the database."""
140    ctx = multiprocessing.get_context("spawn")
141    ready_event = ctx.Event()
142    shutdown_event = ctx.Event()
143    process = ctx.Process(
144        target=simulate_transactional_write,
145        args=(sqlite_persistent.settings, ready_event, shutdown_event),
146    )
147    process.start()
148    ready_event.wait()
149
150    try:
151        sys.argv = ["chroma", "vacuum", "--path", sqlite_persistent.settings.persist_directory, "--force", "--timeout", "10"]
152        cli.app()
153        captured = capfd.readouterr()
154        assert "Failed to vacuum Chroma" in captured.err.strip()
155    finally:
156        shutdown_event.set()
157        process.join()
158