KeyError in cursor.execute with pgbouncer, works with psycopg2

[j-bennet]

An error happens when trying to run a command on pgbouncer database (virtual database).

Traceback (most recent call last):
  File "/Users/irinatruong/.pyenv/versions/pgcli/lib/python3.9/site-packages/psycopg/errors.py", line 310, in _class_for_state
    return lookup(sqlstate)
  File "/Users/irinatruong/.pyenv/versions/pgcli/lib/python3.9/site-packages/psycopg/errors.py", line 293, in lookup
    return _sqlcodes[sqlstate.upper()]
KeyError: ''

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/irinatruong/src/pgcli/t2.py", line 8, in <module>
    cur.execute("show help;")
  File "/Users/irinatruong/.pyenv/versions/pgcli/lib/python3.9/site-packages/psycopg/cursor.py", line 551, in execute
    self._conn.wait(
  File "/Users/irinatruong/.pyenv/versions/pgcli/lib/python3.9/site-packages/psycopg/connection.py", line 776, in wait
    return waiting.wait(gen, self.pgconn.socket, timeout=timeout)
  File "/Users/irinatruong/.pyenv/versions/pgcli/lib/python3.9/site-packages/psycopg/waiting.py", line 62, in wait_selector
    s = gen.send(ready)
  File "/Users/irinatruong/.pyenv/versions/pgcli/lib/python3.9/site-packages/psycopg/cursor.py", line 200, in _execute_gen
    self._check_results(results)
  File "/Users/irinatruong/.pyenv/versions/pgcli/lib/python3.9/site-packages/psycopg/cursor.py", line 405, in _check_results
    self._raise_for_result(res)
  File "/Users/irinatruong/.pyenv/versions/pgcli/lib/python3.9/site-packages/psycopg/cursor.py", line 412, in _raise_for_result
    raise e.error_from_result(result, encoding=self._encoding)
  File "/Users/irinatruong/.pyenv/versions/pgcli/lib/python3.9/site-packages/psycopg/errors.py", line 300, in error_from_result
    cls = _class_for_state(state.decode("ascii"))
  File "/Users/irinatruong/.pyenv/versions/pgcli/lib/python3.9/site-packages/psycopg/errors.py", line 312, in _class_for_state
    return get_base_exception(sqlstate)
  File "/Users/irinatruong/.pyenv/versions/pgcli/lib/python3.9/site-packages/psycopg/errors.py", line 318, in get_base_exception
    or _base_exc_map.get(sqlstate[0])
IndexError: string index out of range

Code:

# t2.py
import psycopg


if __name__ == "__main__":
    conn = psycopg.connect("user=postgres port=6432 password=postgres dbname=pgbouncer")
    conn.autocommit = True
    cur = conn.cursor()
    cur.execute("show help;")
    for n in conn.notifies():
        print(n)
    print(cur.statusmessage)
    cur.close()
    conn.close()

Same code, but with psycopg2 works:

# t1.py

import psycopg2


if __name__ == "__main__":
    conn = psycopg2.connect("user=postgres port=6432 password=postgres dbname=pgbouncer")
    conn.autocommit = True
    cur = conn.cursor()
    cur.execute("show help;")
    for n in conn.notices:
        print(n)
    print(cur.statusmessage)
    cur.close()
    conn.close()

Output:

> python t1.py
NOTICE:  Console usage
DETAIL:
	SHOW HELP|CONFIG|DATABASES|POOLS|CLIENTS|SERVERS|USERS|VERSION
	SHOW FDS|SOCKETS|ACTIVE_SOCKETS|LISTS|MEM
	SHOW DNS_HOSTS|DNS_ZONES
	SHOW STATS|STATS_TOTALS|STATS_AVERAGES|TOTALS
	SET key = arg
	RELOAD
	PAUSE [<db>]
	RESUME [<db>]
	DISABLE <db>
	ENABLE <db>
	RECONNECT [<db>]
	KILL <db>
	SUSPEND
	SHUTDOWN


SHOW

Steps to reproduce

1. Install pgbouncer (brew install pgbouncer or get one here).
2. Put this in pgbouncer.ini:

# pgbouncer.ini

 [databases]
 test = host=localhost port=5432 dbname=test

 [pgbouncer]
 listen_port = 6432
 listen_addr = localhost
 auth_type = md5
 auth_file = userlist.txt
 logfile = pgbouncer.log
 pidfile = pgbouncer.pid
 admin_users = postgres

3. Put this in userlist.txt:

"postgres" "postgres"

4. Start pgbouncer.
5. Run the code snippet t2.py.

[dvarrazzo]

Hello,

could you please test with the patch:

diff --git a/psycopg/psycopg/errors.py b/psycopg/psycopg/errors.py
index d23ce750..21fbd7e3 100644
--- a/psycopg/psycopg/errors.py
+++ b/psycopg/psycopg/errors.py
@@ -315,7 +315,7 @@ def _class_for_state(sqlstate: str) -> Type[Error]:
 def get_base_exception(sqlstate: str) -> Type[Error]:
     return (
         _base_exc_map.get(sqlstate[:2])
-        or _base_exc_map.get(sqlstate[0])
+        or _base_exc_map.get(sqlstate[:1])
         or DatabaseError
     )

and let me know if it works?

Thank you very much!

[dvarrazzo]

I run a test, and there seem to be some other problem:

>>> import psycopg
>>> conn = psycopg.connect("port=6432 host=localhost dbname=pgbouncer user=postgres password=postgres sslmode=disable", autocommit=True)
>>> conn.execute("show help")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/piro/dev/psycopg3/psycopg/psycopg/connection.py", line 861, in execute
    raise ex.with_traceback(None)
psycopg.DatabaseError: message contents do not agree with length in message type "N"

However, sniffing the communication with wireshark, I don't see an error in that message:

0000   00 00 00 00 00 00 00 00 00 00 00 00 08 00 45 00   ..............E.
0010   01 a0 b1 fa 40 00 40 06 89 5b 7f 00 00 01 7f 00   ....@.@..[......
0020   00 01 19 20 aa ee 21 07 ab 45 24 0e cd f4 80 18   ... ..!..E$.....
0030   02 00 ff 94 00 00 01 01 08 0a 6b 49 75 23 6b 49   ..........kIu#kI
0040   75 22 4e 00 00 01 6b 53 4e 4f 54 49 43 45 00 43   u"N...kSNOTICE.C
0050   30 30 30 30 30 00 4d 43 6f 6e 73 6f 6c 65 20 75   00000.MConsole u
0060   73 61 67 65 00 44 0a 09 53 48 4f 57 20 48 45 4c   sage.D..SHOW HEL
0070   50 7c 43 4f 4e 46 49 47 7c 44 41 54 41 42 41 53   P|CONFIG|DATABAS
0080   45 53 7c 50 4f 4f 4c 53 7c 43 4c 49 45 4e 54 53   ES|POOLS|CLIENTS
0090   7c 53 45 52 56 45 52 53 7c 55 53 45 52 53 7c 56   |SERVERS|USERS|V
00a0   45 52 53 49 4f 4e 0a 09 53 48 4f 57 20 46 44 53   ERSION..SHOW FDS
00b0   7c 53 4f 43 4b 45 54 53 7c 41 43 54 49 56 45 5f   |SOCKETS|ACTIVE_
00c0   53 4f 43 4b 45 54 53 7c 4c 49 53 54 53 7c 4d 45   SOCKETS|LISTS|ME
00d0   4d 0a 09 53 48 4f 57 20 44 4e 53 5f 48 4f 53 54   M..SHOW DNS_HOST
00e0   53 7c 44 4e 53 5f 5a 4f 4e 45 53 0a 09 53 48 4f   S|DNS_ZONES..SHO
00f0   57 20 53 54 41 54 53 7c 53 54 41 54 53 5f 54 4f   W STATS|STATS_TO
0100   54 41 4c 53 7c 53 54 41 54 53 5f 41 56 45 52 41   TALS|STATS_AVERA
0110   47 45 53 7c 54 4f 54 41 4c 53 0a 09 53 45 54 20   GES|TOTALS..SET 
0120   6b 65 79 20 3d 20 61 72 67 0a 09 52 45 4c 4f 41   key = arg..RELOA
0130   44 0a 09 50 41 55 53 45 20 5b 3c 64 62 3e 5d 0a   D..PAUSE [<db>].
0140   09 52 45 53 55 4d 45 20 5b 3c 64 62 3e 5d 0a 09   .RESUME [<db>]..
0150   44 49 53 41 42 4c 45 20 3c 64 62 3e 0a 09 45 4e   DISABLE <db>..EN
0160   41 42 4c 45 20 3c 64 62 3e 0a 09 52 45 43 4f 4e   ABLE <db>..RECON
0170   4e 45 43 54 20 5b 3c 64 62 3e 5d 0a 09 4b 49 4c   NECT [<db>]..KIL
0180   4c 20 3c 64 62 3e 0a 09 53 55 53 50 45 4e 44 0a   L <db>..SUSPEND.
0190   09 53 48 55 54 44 4f 57 4e 0a 00 09 57 41 49 54   .SHUTDOWN...WAIT
01a0   5f 43 4c 4f 53 45 20 5b 3c 64 62 3e 5d 00         _CLOSE [<db>].

the message data length is 364 bytes. It starts at the N (0x4e) in position 0042, and the following 4 bytes report a length of 0x16b = 363, which matches the message length without the message type byte.

The error message is raised by the libpq. This is the same libpq used by psql, which can manage this message ok. So I wonder what is psycopg doing with it that makes it unhappy...

[dvarrazzo]

In psycopg 2, this works using the basic query protocol (using PQexec), but fails using the extended protocol (using PQsendQuery).

>>> import psycopg2
>>> import psycopg2.extras
>>> psycopg2.extensions.set_wait_callback(psycopg2.extras.wait_select)  # go async
>>> cnn = psycopg2.connect("port=6432 host=localhost dbname=pgbouncer user=postgres password=postgres sslmode=disable")
>>> cnn.autocommit=True
>>> cur = cnn.cursor()
>>> cur.execute("show help")
# Traceback (most recent call last)
# ...
# DatabaseError: message contents do not agree with length 359 in message type "N"
# 367 != 364

The last message comes from the libpq instrumented with extra diagnostics:

--- a/src/interfaces/libpq/fe-protocol3.c
+++ b/src/interfaces/libpq/fe-protocol3.c
@@ -468,8 +468,9 @@ pqParseInput3(PGconn *conn)
                {
                        /* Trouble --- report it */
                        appendPQExpBuffer(&conn->errorMessage,
-                                                         libpq_gettext("message contents do not agree with length in message type \"%c\"\n"),
-                                                         id);
+                                                         libpq_gettext("message contents do not agree with length %d in message type \"%c\"\n"
+                                                         "%d != %d"),
+                                                         msgLength, id, conn->inCursor, conn->inStart + 5 + msgLength);
                        /* build an error result holding the error message */
                        pqSaveErrorResult(conn);
                        conn->asyncStatus = PGASYNC_READY;

wireshark shows that the messages are the same:

Fe > Be: Q (15 bytes)
Fe < Be: N (364 bytes)
Fe < Be: C + Z + I (16 bytes)

[dvarrazzo]

Using a notice hander to look better at the notice received by the client:

import psycopg

def notice_handler(n):
    for m in psycopg.errors.DiagnosticField:
        v = getattr(n, m.name.lower())
        if v is not None:
            print(f"{m.name.lower()}: {v!r}")

conn = psycopg.connect("port=6432 host=localhost dbname=pgbouncer user=postgres password=postgres sslmode=disable", autocommit=True)
conn.add_notice_handler(notice_handler)
conn.pgconn.exec_(b"show help")
severity: 'NOTICE'
sqlstate: ''
message_primary: 'Console usage'
message_detail: '\n\tSHOW HELP|CONFIG|DATABASES|POOLS|CLIENTS|SERVERS|USERS|VERSION\n\tSHOW FDS|SOCKETS|ACTIVE_SOCKETS|LISTS|MEM\n\tSHOW DNS_HOSTS|DNS_ZONES\n\tSHOW STATS|STATS_TOTALS|STATS_AVERAGES|TOTALS\n\tSET key = arg\n\tRELOAD\n\tPAUSE [<db>]\n\tRESUME [<db>]\n\tDISABLE <db>\n\tENABLE <db>\n\tRECONNECT [<db>]\n\tKILL <db>\n\tSUSPEND\n\tSHUTDOWN\n'

there is a mismatch with what seen on wireshark: sqlstate should have been 00000 (the C field in the N message), but is received as the empty string instead.

Simulating a similar message, in a normal database:

conn2 = psycopg.connect("host=localhost")
conn2.add_notice_handler(notice_handler)
conn2.pgconn.exec_(br"""do $$
begin
    raise notice 'Console usage' using detail=e'\n\tSHOW HELP|CONFIG|DATABASES|POOLS|CLIENTS|SERVERS|USERS|VERSION\n\tSHOW FDS|SOCKETS|ACTIVE_SOCKETS|LISTS|MEM\n\tSHOW DNS_HOSTS|DNS_ZONES\n\tSHOW STATS|STATS_TOTALS|STATS_AVERAGES|TOTALS\n\tSET key = arg\n\tRELOAD\n\tPAUSE [<db>]\n\tRESUME [<db>]\n\tDISABLE <db>\n\tENABLE <db>\n\tRECONNECT [<db>]\n\tKILL <db>\n\tSUSPEND\n\tSHUTDOWN\n';
end
$$ language plpgsql""")
severity: 'NOTICE'
severity_nonlocalized: 'NOTICE'
sqlstate: '00000'
message_primary: 'Console usage'
message_detail: '\n\tSHOW HELP|CONFIG|DATABASES|POOLS|CLIENTS|SERVERS|USERS|VERSION\n\tSHOW FDS|SOCKETS|ACTIVE_SOCKETS|LISTS|MEM\n\tSHOW DNS_HOSTS|DNS_ZONES\n\tSHOW STATS|STATS_TOTALS|STATS_AVERAGES|TOTALS\n\tSET key = arg\n\tRELOAD\n\tPAUSE [<db>]\n\tRESUME [<db>]\n\tDISABLE <db>\n\tENABLE <db>\n\tRECONNECT [<db>]\n\tKILL <db>\n\tSUSPEND\n\tSHUTDOWN\n'
context: 'PL/pgSQL function inline_code_block line 3 at RAISE'
source_file: 'pl_exec.c'
source_line: '3841'
source_function: 'exec_stmt_raise'

Sniffing this message:

0000   4e 00 00 01 b8 53 4e 4f 54 49 43 45 00 56 4e 4f   N....SNOTICE.VNO
0010   54 49 43 45 00 43 30 30 30 30 30 00 4d 43 6f 6e   TICE.C00000.MCon
0020   73 6f 6c 65 20 75 73 61 67 65 00 44 0a 09 53 48   sole usage.D..SH
0030   4f 57 20 48 45 4c 50 7c 43 4f 4e 46 49 47 7c 44   OW HELP|CONFIG|D
0040   41 54 41 42 41 53 45 53 7c 50 4f 4f 4c 53 7c 43   ATABASES|POOLS|C
0050   4c 49 45 4e 54 53 7c 53 45 52 56 45 52 53 7c 55   LIENTS|SERVERS|U
0060   53 45 52 53 7c 56 45 52 53 49 4f 4e 0a 09 53 48   SERS|VERSION..SH
0070   4f 57 20 46 44 53 7c 53 4f 43 4b 45 54 53 7c 41   OW FDS|SOCKETS|A
0080   43 54 49 56 45 5f 53 4f 43 4b 45 54 53 7c 4c 49   CTIVE_SOCKETS|LI
0090   53 54 53 7c 4d 45 4d 0a 09 53 48 4f 57 20 44 4e   STS|MEM..SHOW DN
00a0   53 5f 48 4f 53 54 53 7c 44 4e 53 5f 5a 4f 4e 45   S_HOSTS|DNS_ZONE
00b0   53 0a 09 53 48 4f 57 20 53 54 41 54 53 7c 53 54   S..SHOW STATS|ST
00c0   41 54 53 5f 54 4f 54 41 4c 53 7c 53 54 41 54 53   ATS_TOTALS|STATS
00d0   5f 41 56 45 52 41 47 45 53 7c 54 4f 54 41 4c 53   _AVERAGES|TOTALS
00e0   0a 09 53 45 54 20 6b 65 79 20 3d 20 61 72 67 0a   ..SET key = arg.
00f0   09 52 45 4c 4f 41 44 0a 09 50 41 55 53 45 20 5b   .RELOAD..PAUSE [
0100   3c 64 62 3e 5d 0a 09 52 45 53 55 4d 45 20 5b 3c   <db>]..RESUME [<
0110   64 62 3e 5d 0a 09 44 49 53 41 42 4c 45 20 3c 64   db>]..DISABLE <d
0120   62 3e 0a 09 45 4e 41 42 4c 45 20 3c 64 62 3e 0a   b>..ENABLE <db>.
0130   09 52 45 43 4f 4e 4e 45 43 54 20 5b 3c 64 62 3e   .RECONNECT [<db>
0140   5d 0a 09 4b 49 4c 4c 20 3c 64 62 3e 0a 09 53 55   ]..KILL <db>..SU
0150   53 50 45 4e 44 0a 09 53 48 55 54 44 4f 57 4e 0a   SPEND..SHUTDOWN.
0160   00 57 50 4c 2f 70 67 53 51 4c 20 66 75 6e 63 74   .WPL/pgSQL funct
0170   69 6f 6e 20 69 6e 6c 69 6e 65 5f 63 6f 64 65 5f   ion inline_code_
0180   62 6c 6f 63 6b 20 6c 69 6e 65 20 33 20 61 74 20   block line 3 at 
0190   52 41 49 53 45 00 46 70 6c 5f 65 78 65 63 2e 63   RAISE.Fpl_exec.c
01a0   00 4c 33 38 34 31 00 52 65 78 65 63 5f 73 74 6d   .L3841.Rexec_stm
01b0   74 5f 72 61 69 73 65 00 00                        t_raise..

looks similar to the previous, as it's 441 bytes of TCP payload, and the notice is reported as the same N in position 0x42 followed by 0x1b8 = 440 which is the message length minus the first char.

The content of the message shouldn't matter, however, it can be seen that the sqlstate (C) = 00000 this time makes it all the way to the client.

[dvarrazzo]

Low level repro using libpq functions:

# test_pgbouncer.py
import select
import psycopg

def notice_handler(n):
    for m in psycopg.errors.DiagnosticField:
        v = getattr(n, m.name.lower())
        if v is not None:
            print(f"notice: {m.name.lower()}: {v!r}")

conn = psycopg.connect(
    "port=6432 host=localhost dbname=pgbouncer user=postgres password=postgres sslmode=disable",
    autocommit=True,
)
conn.add_notice_handler(notice_handler)

print("Running 'SHOW HELP' using PQexec")
result = conn.pgconn.exec_(b"SHOW HELP")
print(f"result with PQexec: {result}")

print("\nRunning 'SHOW HELP' using PQsendQuery")
conn.pgconn.send_query(b"SHOW HELP")
while True:
    rv = conn.pgconn.flush()
    print(f"PQflush returned: {rv}")
    if not rv:
        break

rv = conn.pgconn.is_busy()
print(f"PQisBusy returned: {rv}")
if rv:
    select.select([conn], [], [], 5)
    while True:
        conn.pgconn.consume_input()
        rv = conn.pgconn.is_busy()
        print(f"PQisBusy returned: {rv}")
        if not rv:
            break
        select.select([conn], [], [], 5)

result = conn.pgconn.get_result()
print(f"result with PQsendQuery: {result}")
print(f"{result.error_message.decode()=}")

The output looks like:

$ python test_pgbouncer.py 
Running 'SHOW HELP' using PQexec
notice severity: 'NOTICE'
notice sqlstate: ''
notice message_primary: 'Console usage'
notice message_detail: '\n\tSHOW HELP|CONFIG|DATABASES|POOLS|CLIENTS|SERVERS|USERS|VERSION\n\tSHOW FDS|SOCKETS|ACTIVE_SOCKETS|LISTS|MEM\n\tSHOW DNS_HOSTS|DNS_ZONES\n\tSHOW STATS|STATS_TOTALS|STATS_AVERAGES|TOTALS\n\tSET key = arg\n\tRELOAD\n\tPAUSE [<db>]\n\tRESUME [<db>]\n\tDISABLE <db>\n\tENABLE <db>\n\tRECONNECT [<db>]\n\tKILL <db>\n\tSUSPEND\n\tSHUTDOWN\n'
result with PQexec: <psycopg_c.pq.PGresult [COMMAND_OK] at 0x7fec49fd3db0>

Running 'SHOW HELP' using PQsendQuery
PQflush returned: 0
PQisBusy returned: 1
PQisBusy returned: 1
notice severity: 'NOTICE'
notice sqlstate: ''
notice message_primary: 'Console usage'
notice message_detail: '\n\tSHOW HELP|CONFIG|DATABASES|POOLS|CLIENTS|SERVERS|USERS|VERSION\n\tSHOW FDS|SOCKETS|ACTIVE_SOCKETS|LISTS|MEM\n\tSHOW DNS_HOSTS|DNS_ZONES\n\tSHOW STATS|STATS_TOTALS|STATS_AVERAGES|TOTALS\n\tSET key = arg\n\tRELOAD\n\tPAUSE [<db>]\n\tRESUME [<db>]\n\tDISABLE <db>\n\tENABLE <db>\n\tRECONNECT [<db>]\n\tKILL <db>\n\tSUSPEND\n\tSHUTDOWN\n'
PQisBusy returned: 0
result with PQsendQuery: <psycopg_c.pq.PGresult [FATAL_ERROR] at 0x7fec49fd3e10>
result.error_message.decode()='message contents do not agree with length in message type "N"\n'

[dvarrazzo]

I have submitted the issue to pgbouncer as pgbouncer/pgbouncer#718

As far as I can see, the notice generated by pgbouncer makes the libpq angry, as can be seen by psycopg2 in async mode too and which can be reproduced by the libpq only.

I have tested pgbouncer further and it seems working as expected:

>>> conn.execute("show config").fetchone()
('admin_users', 'postgres', '', 'yes')

The only workaround I have in mind in psycopg 3 is to special-case SHOW HELP and run it via connection.pgconn.exec_(); in psycopg 2 the green mode should be suspended temporarily as done for COPY.

I would like to help more, but I have seen enough evidence to think that the issue is either in the message sent by pgbouncer or in the way such message is processed by the libpq, and not in psycopg. Closing the ticket here, but let me know if we can help further.

[dvarrazzo]

(Actually, another workaround for pgcli is to just ignore the error if the query was "show help"...)

[j-bennet]

@dvarrazzo Thank you, great investigation. Is there a high-level way with psycopg to tell the connection to use the basic query protocol, or can it be used only via the protected connection.pgconn.exec_()?

[dvarrazzo]

No, there is no high-level psycopg code path using pgconn.exec_(), because the PQexec() libpq call is blocking and I have tried to avoid using any blocking function throughout the entire codebase. It is the reason why in pgcli+psycopg2 you have to use set_wait_callback(): doing so you get to use an async path and so Ctrl-C becomes responsive; if not, in order to stop PQexec("SELECT pg_sleep(60)") it takes a cannon. So far, I hadn't seen any case where this would have created a problem.

If the issue is only limited to SHOW HELP in pgbouncer, I would resist to add any other query execution path. After all, you do receive the notice with the bulk of the information... the only problem is that then psycopg will raise a DatabaseError (I am releasing psycopg 3.0.14, where I've fixed the KeyError/IndexError in the original report), but you can manage it easily if you special-case the query.

I am open to consider using PQexec if other cases arise and they are not acknowledged as bug by the projects generating them. Just this one is too minor to warrant extending the interface.

[dvarrazzo]

Psycopg 3.0.14 released, which fixes the IndexError, leaving the DatabaseError.