codekingpro/portable-devtools

1##########################################################################2#3# pgAdmin 4 - PostgreSQL Tools4#5# Copyright (C) 2013 - 2024, The pgAdmin Development Team6# This software is released under the PostgreSQL Licence7#8##########################################################################9 10#  Get Postgres and Python encoding11 12import psycopg13from flask import current_app14 15encode_dict = {16    'SQL_ASCII': ['SQL_ASCII', 'raw-unicode-escape'],17    'SQLASCII': ['SQLASCII', 'raw-unicode-escape'],18    # EUC_TW Not availble in Python,19    # so psycopg3 do not support it, we are on our own20    'EUC_TW': ['BIG5', 'big5'],21    'EUCTW': ['BIG5', 'big5'],22    # psycopg3 do not support unicode23    'UNICODE': ['utf-8', 'utf-8']24}25 26 27def get_encoding(key):28    """29    :param key: Database Encoding30    :return:31    [Postgres_encoding, Python_encoding] -32    Postgres encoding, Python encoding, type cast encoding33    """34    #35    # Reference: https://www.postgresql.org/docs/11/multibyte.html36    #37    if key == 'ascii':38        key = 'raw_unicode_escape'39    try:40        postgres_encoding = psycopg._encodings.py2pgenc(key).decode()41    except Exception as e:42        # For unsupported encodings by psycopg like UNICODE43        current_app.logger.error(e)44        postgres_encoding = 'utf-8'45 46    python_encoding = psycopg._encodings._py_codecs.get(postgres_encoding,47                                                        'utf-8')48 49    _dict = encode_dict.get(postgres_encoding.upper(),50                            [postgres_encoding,51                             python_encoding])52    return _dict53 54 55def configure_driver_encodings(encodings):56    # Replace the python encoding for original name and renamed encodings57    # psycopg removes the underscore in conn.encoding58    # Setting the encodings dict value will only help for select statements59    # because for parameterized DML, param values are converted based on60    # python encoding of pyscopg's internal encodings dict.61 62    for key, val in encode_dict.items():63        _, python_encoding = val64        psycopg._encodings._py_codecs[key] = python_encoding65 66    encodings.update((k.encode(), v67                      ) for k, v in psycopg._encodings._py_codecs.items())68    psycopg._encodings.pg_codecs = {69        v: k.encode() for k, v in psycopg._encodings._py_codecs.items()}70