Fix custom-profile-field reads/writes on the Rust backend

Custom profile fields use Postgres JSON types that the shim couldn't handle, so
the ProfileFieldRestServlet returned HTTP 500 on the Rust backend. Under trial's
parallel runner these 500s also polluted their worker and cascaded into
unrelated failures, which is what made runs look non-deterministically flaky —
psycopg2 was unaffected.

Three gaps, fixed here (json/jsonb decode was added separately):

  - Bind a JSON document (as text) to a `json` / `jsonb` parameter: `?::jsonb`
    types the parameter as jsonb, so `ToSql` now encodes `Text` for `json`
    (raw text) and `jsonb`/`jsonpath` (a one-byte version header then the
    text). `set_profile_field` accordingly passes the canonical JSON as text
    rather than a psycopg2 `Json` wrapper (which the shim can't bind, and which
    coupled the storage layer to psycopg2).

  - Bind a `jsonpath` parameter: `get_profile_field`'s
    `JSONB_PATH_EXISTS(fields, ?)` types the parameter as `jsonpath`; it now
    encodes the same way as `jsonb`.

  - `JSON_BUILD_OBJECT(?, ?::jsonb)` left the key parameter's type
    indeterminate in a prepared statement (SQLSTATE 42P18) — psycopg2 sends
    untyped parameters and infers at execute, but the shim prepares. Cast the
    key to `?::text` so its type is explicit.

Fixes tests.rest.client.test_profile on the Rust backend (37/37, stable under
-j4 across repeated runs); psycopg2 and sqlite are unaffected.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d
This commit is contained in:
Erik Johnston
2026-07-06 10:14:31 +00:00
co-authored by Claude Opus 4.8
parent b1492c1993
commit 62f113f0a3
2 changed files with 39 additions and 7 deletions
+33 -1
View File
@@ -271,6 +271,22 @@ impl ToSql for PgValue {
tid_to_sql(v, buf)?;
Ok(IsNull::No)
}
(PgValue::Text(v), &Type::JSON) => {
// A JSON document passed as text and bound to a `json` column /
// `?::json` cast (e.g. a custom profile field value). `json`'s
// wire format is just the raw text.
buf.extend_from_slice(v.as_bytes());
Ok(IsNull::No)
}
(PgValue::Text(v), &Type::JSONB | &Type::JSONPATH) => {
// `jsonb` and `jsonpath` share a wire format: a one-byte version
// header (1) then the text. A `jsonpath` param is bound as text,
// e.g. `JSONB_PATH_EXISTS(fields, ?)` when reading a custom
// profile field.
buf.extend_from_slice(&[1u8]);
buf.extend_from_slice(v.as_bytes());
Ok(IsNull::No)
}
(PgValue::Bytea(v), &Type::BYTEA) => {
bytea_to_sql(v, buf);
Ok(IsNull::No)
@@ -323,8 +339,10 @@ impl ToSql for PgValue {
}
fn accepts(ty: &Type) -> bool {
// Scalars, plus arrays of a supported scalar element type.
// Scalars, `json`/`jsonb`/`jsonpath` (a `Text` document binds to these),
// plus arrays of a supported scalar element type.
accepts_column_type(ty)
|| matches!(*ty, Type::JSON | Type::JSONB | Type::JSONPATH)
|| matches!(ty.kind(), Kind::Array(element) if accepts_column_type(element))
}
@@ -692,6 +710,20 @@ mod tests {
});
}
#[test]
fn to_sql_encodes_text_as_json_and_jsonb() {
// A JSON document passed as text binds to a `json` column verbatim, and
// to `jsonb` with the one-byte version header prepended.
assert_eq!(
encode(&PgValue::Text("[1, 2]".into()), &Type::JSON).0,
b"[1, 2]".to_vec()
);
assert_eq!(
encode(&PgValue::Text("[1, 2]".into()), &Type::JSONB).0,
b"\x01[1, 2]".to_vec()
);
}
#[test]
fn to_sql_encodes_arrays() {
// An `INT8[]` array encodes without error and produces a non-empty
+6 -6
View File
@@ -455,12 +455,10 @@ class ProfileWorkerStore(SQLBaseStore):
self._check_profile_size(txn, user_id, field_name, new_value)
if isinstance(self.database_engine, PostgresEngine):
from psycopg2.extras import Json
# Note that the || jsonb operator is not recursive, any duplicate
# keys will be taken from the second value.
sql = """
INSERT INTO profiles (user_id, full_user_id, fields) VALUES (?, ?, JSON_BUILD_OBJECT(?, ?::jsonb))
INSERT INTO profiles (user_id, full_user_id, fields) VALUES (?, ?, JSON_BUILD_OBJECT(?::text, ?::jsonb))
ON CONFLICT (user_id)
DO UPDATE SET full_user_id = EXCLUDED.full_user_id, fields = COALESCE(profiles.fields, '{}'::jsonb) || EXCLUDED.fields
"""
@@ -471,9 +469,11 @@ class ProfileWorkerStore(SQLBaseStore):
user_id.localpart,
user_id.to_string(),
field_name,
# Pass as a JSON object since we have passing bytes disabled
# at the database driver.
Json(json.loads(canonical_value)),
# The field value as a JSON document; the `?::jsonb` cast
# in the query turns it into jsonb. Passed as text rather
# than raw bytes, since binding bytes is disabled at the
# database driver.
canonical_value.decode("utf-8"),
),
)
else: