diff --git a/web/pgadmin/utils/driver/psycopg3/encoding.py b/web/pgadmin/utils/driver/psycopg3/encoding.py index da317e008..b994f9ecb 100644 --- a/web/pgadmin/utils/driver/psycopg3/encoding.py +++ b/web/pgadmin/utils/driver/psycopg3/encoding.py @@ -9,12 +9,67 @@ # Get Postgres and Python encoding +import codecs +import encodings.raw_unicode_escape + import psycopg from flask import current_app +# Python codec used for SQL_ASCII databases. It behaves exactly like +# raw_unicode_escape, except that a backslash sequence which is not a valid +# \uXXXX or \UXXXXXXXX escape (for example the '\U' in a column named +# 'C:\User') is kept as the literal bytes rather than raising +# UnicodeDecodeError. psycopg decodes column names and query text with the +# connection codec, so a strict decode there breaks View/Edit Data. +SQL_ASCII_CODEC = 'pgadmin_sql_ascii' +_SQL_ASCII_ERRORS = 'pgadmin_sql_ascii_literal' + + +def _keep_invalid_escape(exc): + if not isinstance(exc, UnicodeDecodeError): + raise exc + # Every byte maps to the same code point in latin-1, which is also what + # raw_unicode_escape produces for the bytes it does not unescape. + return exc.object[exc.start:exc.end].decode('latin-1'), exc.end + + +def _sql_ascii_errors(errors): + # psycopg decodes server error messages with errors='replace', so keep + # invalid escapes literally there too, rather than showing U+FFFD in + # place of part of the object name. + if errors in ('strict', 'replace'): + return _SQL_ASCII_ERRORS + return errors + + +def _sql_ascii_decode(data, errors='strict'): + return codecs.raw_unicode_escape_decode(data, _sql_ascii_errors(errors)) + + +class _SQLAsciiIncrementalDecoder( + encodings.raw_unicode_escape.IncrementalDecoder): + def _buffer_decode(self, data, errors, final): + return super()._buffer_decode(data, _sql_ascii_errors(errors), final) + + +def _sql_ascii_search(name): + if name != SQL_ASCII_CODEC: + return None + return codecs.CodecInfo( + name=SQL_ASCII_CODEC, + encode=codecs.raw_unicode_escape_encode, + decode=_sql_ascii_decode, + incrementalencoder=encodings.raw_unicode_escape.IncrementalEncoder, + incrementaldecoder=_SQLAsciiIncrementalDecoder, + ) + + +codecs.register_error(_SQL_ASCII_ERRORS, _keep_invalid_escape) +codecs.register(_sql_ascii_search) + encode_dict = { - 'SQL_ASCII': ['SQL_ASCII', 'raw-unicode-escape'], - 'SQLASCII': ['SQLASCII', 'raw-unicode-escape'], + 'SQL_ASCII': ['SQL_ASCII', SQL_ASCII_CODEC], + 'SQLASCII': ['SQLASCII', SQL_ASCII_CODEC], # EUC_TW Not availble in Python, # so psycopg3 do not support it, we are on our own 'EUC_TW': ['BIG5', 'big5'], @@ -35,7 +90,7 @@ def get_encoding(key): # Reference: https://www.postgresql.org/docs/11/multibyte.html # if key == 'ascii': - key = 'raw_unicode_escape' + key = SQL_ASCII_CODEC try: postgres_encoding = psycopg._encodings.py2pgenc(key).decode() except Exception as e: diff --git a/web/pgadmin/utils/driver/psycopg3/tests/test_encoding.py b/web/pgadmin/utils/driver/psycopg3/tests/test_encoding.py index 47554408c..f2467c4fc 100644 --- a/web/pgadmin/utils/driver/psycopg3/tests/test_encoding.py +++ b/web/pgadmin/utils/driver/psycopg3/tests/test_encoding.py @@ -44,8 +44,8 @@ class TestEncoding(BaseTestGenerator): dict(key='EUC_TW', expected=['utf-8', 'utf-8'])), ('UNICODE falls back to the pgAdmin-only override', dict(key='UNICODE', expected=['utf-8', 'utf-8'])), - ('the ascii special-case resolves via raw_unicode_escape', - dict(key='ascii', expected=['SQLASCII', 'raw-unicode-escape'])), + ('the ascii special-case resolves via the SQL_ASCII codec', + dict(key='ascii', expected=['SQLASCII', 'pgadmin_sql_ascii'])), ] # Only set by scenarios where more than one result is correct. diff --git a/web/pgadmin/utils/driver/psycopg3/tests/test_sql_ascii_codec.py b/web/pgadmin/utils/driver/psycopg3/tests/test_sql_ascii_codec.py new file mode 100644 index 000000000..dc2355625 --- /dev/null +++ b/web/pgadmin/utils/driver/psycopg3/tests/test_sql_ascii_codec.py @@ -0,0 +1,82 @@ +########################################################################## +# +# pgAdmin 4 - PostgreSQL Tools +# +# Copyright (C) 2013 - 2026, The pgAdmin Development Team +# This software is released under the PostgreSQL Licence +# +########################################################################## + +""" +Unit tests for the codec used for SQL_ASCII databases (#5940). + +Identifiers or data containing a backslash followed by 'u' or 'U' (for +example a table or column called 'C:\\User\\te') used to make +raw_unicode_escape raise "truncated \\UXXXXXXXX escape" whilst psycopg +decoded the column names, breaking View/Edit Data. +""" + +import codecs + +from pgadmin.utils.driver.psycopg3.encoding import SQL_ASCII_CODEC +from pgadmin.utils.route import BaseTestGenerator + + +class TestSQLAsciiCodec(BaseTestGenerator): + """Decode/encode behaviour of the SQL_ASCII codec.""" + + scenarios = [ + ('Invalid \\U escape is kept literally', + dict(data=b'C:\\User\\te', expected='C:\\User\\te')), + ('Invalid \\u escape is kept literally', + dict(data=b'C:\\users', expected='C:\\users')), + ('Truncated escape at the end is kept literally', + dict(data=b'abc\\u12', expected='abc\\u12')), + ('Valid \\u escape is still unescaped', + dict(data=b'caf\\u00e9', expected='caf\u00e9', + round_trip=False)), + ('Valid escape after an invalid one is still unescaped', + dict(data=b'\\User \\u00e9', expected='\\User \u00e9', + round_trip=False)), + ('High bytes decode as latin-1, as with raw_unicode_escape', + dict(data=b'\xe9\\Ux', expected='\xe9\\Ux')), + ('Other backslash sequences are left alone', + dict(data=b'C:\\temp\\new', expected='C:\\temp\\new')), + ] + + # raw_unicode_escape writes code points below 256 as single bytes, so + # an unescaped \u00e9 does not encode back to the original escape. + round_trip = True + + def setUp(self): + # No DB connection needed for this test. + pass + + def runTest(self): + decoded = self.data.decode(SQL_ASCII_CODEC) + self.assertEqual(decoded, self.expected) + self.assertEqual(str(self.data, SQL_ASCII_CODEC), self.expected) + + # Encoding must match raw_unicode_escape, so queries built from + # the decoded names are sent to the server unchanged. + self.assertEqual(decoded.encode(SQL_ASCII_CODEC), + decoded.encode('raw_unicode_escape')) + if self.round_trip: + self.assertEqual(decoded.encode(SQL_ASCII_CODEC), self.data) + + # psycopg decodes server error messages with errors='replace', + # which must keep invalid escapes literally as well. + self.assertEqual(self.data.decode(SQL_ASCII_CODEC, 'replace'), + self.expected) + + # Decoding a byte at a time, as a stream reader would, gives the + # same result as decoding the whole value at once. + decoder = codecs.getincrementaldecoder(SQL_ASCII_CODEC)() + self.assertEqual( + ''.join(decoder.decode(self.data[i:i + 1]) + for i in range(len(self.data))) + + decoder.decode(b'', final=True), + self.expected) + self.assertEqual( + codecs.getincrementalencoder(SQL_ASCII_CODEC)().encode(decoded), + decoded.encode(SQL_ASCII_CODEC))