Skip to content

← Parallel Python with Process Pools step 5 of 24

Hard End-to-End

Give Each Worker Its Own Database Connection

Your app opens a database connection at startup and keeps it in a global. Under fork, each worker gets a copy, so the parent and the worker share one socket. Under spawn or forkserver, the global is None in every worker. The fix is to send each worker a recipe, a small picklable description such as the connection string. An initializer then builds the worker's own connection from it:

recipe = ("pg://db",)
with ProcessPoolExecutor(initializer=worker_init, initargs=(recipe,)) as pool:
    results = list(pool.map(worker_handle, rows))

The grader cannot start processes, so solve simulates the pool. It runs each worker's lifetime one after another. The starter's Conn is a fake connection. It holds a lock, so pickle rejects it, and it counts the rows it serves.

What to write

Recipe is a one-element tuple holding the connection string, (dsn,).

  • worker_init(recipe) builds a new Conn from the recipe, stores it in _CONN, and adds one to _OPENED.
  • worker_handle(row) returns the current connection's query(row). If there is no connection, it raises RuntimeError("worker_handle called before worker_init").

solve(*, dsn, rows, workers) runs these steps:

  1. Set _CONN to None and _OPENED to 0.
  2. If workers <= 0, raise ValueError("workers must be positive").
  3. Give worker w the rows at indexes w, w + workers, w + 2 * workers, and so on. Run the workers in order from 0.
  4. Skip a worker with no rows. For each other worker, call worker_init once, handle its rows in order, then set _CONN back to None. Reset _CONN even when a row raises, and let the error reach the caller.
  5. Return (results, opened). results is a tuple in the original row order. opened is how many connections worker_init built.

Track each row's index, so two equal rows stay separate. Don't change rows. Empty rows with a positive workers returns ((), 0).

Example

Conn.query returns f"{dsn}#{serial}:{row.upper()}", where serial is 1 for a connection's first row, 2 for its second, and so on. With dsn="db", rows a, b, c, d and two workers:

(("db#1:A", "db#1:B", "db#2:C", "db#2:D"), 2)

Worker 0 serves a and c. Worker 1 serves b and d. Each worker reuses its own connection, so c is the second row on worker 0's connection.

Pass mypy --strict

Annotate the global as _CONN: Conn | None. In worker_handle, copy _CONN into a local variable, raise if it is None, then call query on the local. mypy then knows the local is a Conn.

Going further

A real pool decides which worker gets which row, so a real run won't match this split or these serial numbers. Close a real connection with its own close() method. Setting the global to None is enough only for this fake.