Fix encoding error when database encoding is ASCII

pgcli/pgexecute.py

@@ -1,18 +1,20 @@
+from collections import namedtuple
+import contextlib
 import ipaddress
 import logging
-import traceback
-from collections import namedtuple
 import re
+import traceback
+
 import pgspecial as special
 import psycopg
-import psycopg.sql
 from psycopg.conninfo import make_conninfo
+import psycopg.sql
 import sqlparse
 
 sqlparse.engine.grouping.MAX_GROUPING_DEPTH = None
 sqlparse.engine.grouping.MAX_GROUPING_TOKENS = None
 
-from .packages.parseutils.meta import FunctionMetadata, ForeignKey
+from .packages.parseutils.meta import ForeignKey, FunctionMetadata
 
 _logger = logging.getLogger(__name__)
 
@@ -287,16 +289,6 @@ def short_host(self):
         short_host, _, _ = host.partition(".")
         return short_host
 
-    def _select_one(self, cur, sql):
-        """
-        Helper method to run a select and retrieve a single field value
-        :param cur: cursor
-        :param sql: string
-        :return: string
-        """
-        cur.execute(sql)
-        return cur.fetchone()
-
     def failed_transaction(self):
         return self.conn.info.transaction_status == psycopg.pq.TransactionStatus.INERROR
 
@@ -491,14 +483,16 @@ def search_path(self):
         try:
             with self.conn.cursor() as cur:
                 _logger.debug("Search path query. sql: %r", self.search_path_query)
-                cur.execute(self.search_path_query)
-                return [x[0] for x in cur.fetchall()]
+                with override_client_encoding(self.conn, "utf8"):
+                    cur.execute(self.search_path_query)
+                    return [x[0] for x in cur.fetchall()]
         except psycopg.ProgrammingError:
             fallback = "SELECT * FROM current_schemas(true)"
             with self.conn.cursor() as cur:
                 _logger.debug("Search path query. sql: %r", fallback)
-                cur.execute(fallback)
-                return cur.fetchone()[0]
+                with override_client_encoding(self.conn, "utf8"):
+                    cur.execute(fallback)
+                    return cur.fetchone()[0]
 
     def view_definition(self, spec):
         """Returns the SQL defining views described by `spec`"""
@@ -510,7 +504,8 @@ def view_definition(self, spec):
             sql = self.view_definition_query
             _logger.debug("View Definition Query. sql: %r\nspec: %r", sql, spec)
             try:
-                cur.execute(sql, (spec,))
+                with override_client_encoding(self.conn, "utf8"):
+                    cur.execute(sql, (spec,))
             except psycopg.ProgrammingError:
                 raise RuntimeError(f"View {spec} does not exist.")
             result = ViewDef(*cur.fetchone())
@@ -535,9 +530,10 @@ def function_definition(self, spec):
             sql = self.function_definition_query
             _logger.debug("Function Definition Query. sql: %r\nspec: %r", sql, spec)
             try:
-                cur.execute(sql, (spec,))
-                result = cur.fetchone()
-                return result[0]
+                with override_client_encoding(self.conn, "utf8"):
+                    cur.execute(sql, (spec,))
+                    result = cur.fetchone()
+                    return result[0]
             except psycopg.ProgrammingError:
                 raise RuntimeError(f"Function {spec} does not exist.")
 
@@ -546,8 +542,9 @@ def schemata(self):
 
         with self.conn.cursor() as cur:
             _logger.debug("Schemata Query. sql: %r", self.schemata_query)
-            cur.execute(self.schemata_query)
-            return [x[0] for x in cur.fetchall()]
+            with override_client_encoding(self.conn, "utf8"):
+                cur.execute(self.schemata_query)
+                return [x[0] for x in cur.fetchall()]
 
     def _relations(self, kinds=("r", "p", "f", "v", "m")):
         """Get table or view name metadata
@@ -564,8 +561,9 @@ def _relations(self, kinds=("r", "p", "f", "v", "m")):
         with self.conn.cursor() as cur:
             # sql = cur.mogrify(self.tables_query, kinds)
             # _logger.debug("Tables Query. sql: %r", sql)
-            cur.execute(self.tables_query, [kinds])
-            yield from cur
+            with override_client_encoding(self.conn, "utf8"):
+                cur.execute(self.tables_query, [kinds])
+                yield from cur
 
     def tables(self):
         """Yields (schema_name, table_name) tuples"""
@@ -633,8 +631,9 @@ def _columns(self, kinds=("r", "p", "f", "v", "m")):
         with self.conn.cursor() as cur:
             # sql = cur.mogrify(columns_query, kinds)
             # _logger.debug("Columns Query. sql: %r", sql)
-            cur.execute(columns_query, [kinds])
-            yield from cur
+            with override_client_encoding(self.conn, "utf8"):
+                cur.execute(columns_query, [kinds])
+                yield from cur
 
     def table_columns(self):
         yield from self._columns(kinds=["r", "p", "f"])
@@ -645,8 +644,9 @@ def view_columns(self):
     def databases(self):
         with self.conn.cursor() as cur:
             _logger.debug("Databases Query. sql: %r", self.databases_query)
-            cur.execute(self.databases_query)
-            return [x[0] for x in cur.fetchall()]
+            with override_client_encoding(self.conn, "utf8"):
+                cur.execute(self.databases_query)
+                return [x[0] for x in cur.fetchall()]
 
     def full_databases(self):
         with self.conn.cursor() as cur:
@@ -705,9 +705,10 @@ def foreignkeys(self):
                 WHERE fk.contype = 'f';
                 """
             _logger.debug("Functions Query. sql: %r", query)
-            cur.execute(query)
-            for row in cur:
-                yield ForeignKey(*row)
+            with override_client_encoding(self.conn, "utf8"):
+                cur.execute(query)
+                for row in cur:
+                    yield ForeignKey(*row)
 
     def functions(self):
         """Yields FunctionMetadata named tuples"""
@@ -795,9 +796,10 @@ def functions(self):
 
         with self.conn.cursor() as cur:
             _logger.debug("Functions Query. sql: %r", query)
-            cur.execute(query)
-            for row in cur:
-                yield FunctionMetadata(*row)
+            with override_client_encoding(self.conn, "utf8"):
+                cur.execute(query)
+                for row in cur:
+                    yield FunctionMetadata(*row)
 
     def datatypes(self):
         """Yields tuples of (schema_name, type_name)"""
@@ -840,8 +842,9 @@ def datatypes(self):
                     ORDER BY 1, 2;
                 """
             _logger.debug("Datatypes Query. sql: %r", query)
-            cur.execute(query)
-            yield from cur
+            with override_client_encoding(self.conn, "utf8"):
+                cur.execute(query)
+                yield from cur
 
     def casing(self):
         """Yields the most common casing for names used in db functions"""
@@ -888,20 +891,43 @@ def casing(self):
             AND Row_Number = 1;
             """
             _logger.debug("Casing Query. sql: %r", query)
-            cur.execute(query)
-            for row in cur:
-                yield row[0]
+            with override_client_encoding(self.conn, "utf8"):
+                cur.execute(query)
+                for row in cur:
+                    yield row[0]
 
     def explain_prefix(self):
         return "EXPLAIN (ANALYZE, COSTS, VERBOSE, BUFFERS, FORMAT JSON) "
 
     def get_timezone(self) -> str:
         query = psycopg.sql.SQL("show time zone")
         with self.conn.cursor() as cur:
-            cur.execute(query)
-            return cur.fetchone()[0]
+            with override_client_encoding(self.conn, "utf8"):
+                cur.execute(query)
+                return cur.fetchone()[0]
 
     def set_timezone(self, timezone: str):
         query = psycopg.sql.SQL("set time zone {}").format(psycopg.sql.Identifier(timezone))
         with self.conn.cursor() as cur:
             cur.execute(query)
+
+
+@contextlib.contextmanager
+def override_client_encoding(cur, encoding):
+    initial_encoding = cur.connection.info.encoding
+    if encoding == initial_encoding:
+        yield
+
+    else:
+        try:
+            cur.execute(f"set client_encoding to {encoding}")
+            yield
+        finally:
+            # What is stored in `psycopg.ConnectionInfo.encoding` does
+            # not seem to always be an encoding name that is supported
+            # by PostgreSQL (I did not search why).
+            initial_encoding = {
+                "ascii": "sql_ascii",
+                "utf-8": "utf8",
+            }.get(initial_encoding, initial_encoding)
+            cur.execute(f"set client_encoding to {initial_encoding}")