LAX JSON functions not found

[joshbode]

The LAX family of functions (LAX_INT64, LAX_STRING, LAX_BOOL, LAX_FLOAT64) are not recognised, showing a message like: INVALID_ARGUMENT: Function not found: LAX_STRING; Did you mean string?.

Additionally (happy to report as a separate issue) diagnostics are being generated for only the first error in the query, i.e. I can see the "error" for LAX_STRING but not for the subsequent unknown functions until I comment out the LAX_STRING line.
Diagnostics for separate queries (separated by ;) in the same file are still being generated, though.

Example query:

WITH
  x AS (
    SELECT
      PARSE_JSON(
        '{"id": 555, "flag": true, "value": 123.4, "message": "Hi kitagry"}'
      ) data
  )
SELECT
  LAX_STRING(JSON_QUERY(x.data, '$.message')) message,
  LAX_FLOAT64(JSON_QUERY(x.data, '$.value')) value,
  LAX_BOOL(JSON_QUERY(x.data, '$.flag')) flag,
  LAX_INT64(JSON_QUERY(x.data, '$.id')) id
FROM
  x;

Diagnostics in Neovim:

I'm running with the latest main.