From c77f9e5674bd92ff4cbe98f3fcc227dbd935f349 Mon Sep 17 00:00:00 2001
From: Benjamin Renard <brenard@easter-eggs.com>
Date: Sun, 7 Nov 2021 21:57:00 +0100
Subject: [PATCH] PgSQL: fix quoting table name with schema prefix

---
 mylib/pgsql.py      | 23 ++++++++++++++++-------
 tests/test_pgsql.py |  5 +++++
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/mylib/pgsql.py b/mylib/pgsql.py
index b8e559a..7c0f847 100644
--- a/mylib/pgsql.py
+++ b/mylib/pgsql.py
@@ -254,10 +254,19 @@ class PgDB:
             sql += " WHERE " + sql_where
         return (sql, params)
 
+    @staticmethod
+    def _quote_table_name(table):
+        """ Quote table name """
+        return '"{0}"'.format(
+            '"."'.join(
+                table.split('.')
+            )
+        )
+
     def insert(self, table, values, just_try=False):
         """ Run INSERT SQL query """
-        sql = 'INSERT INTO "{0}" ("{1}") VALUES ({2})'.format(
-            table,
+        sql = 'INSERT INTO {0} ("{1}") VALUES ({2})'.format(
+            self._quote_table_name(table),
             '", "'.join(values.keys()),
             ", ".join([
                 '%({0})s'.format(key)
@@ -277,8 +286,8 @@ class PgDB:
 
     def update(self, table, values, where_clauses, where_op=None, just_try=False):
         """ Run UPDATE SQL query """
-        sql = 'UPDATE "{0}" SET {1}'.format(
-            table,
+        sql = 'UPDATE {0} SET {1}'.format(
+            self._quote_table_name(table),
             ", ".join([
                 '"{0}" = %({0})s'.format(key)
                 for key in values
@@ -304,7 +313,7 @@ class PgDB:
 
     def delete(self, table, where_clauses, where_op='AND', just_try=False):
         """ Run DELETE SQL query """
-        sql = 'DELETE FROM "{0}"'.format(table)
+        sql = 'DELETE FROM {0}'.format(self._quote_table_name(table))
         params = dict()
 
         try:
@@ -326,7 +335,7 @@ class PgDB:
     def truncate(self, table, just_try=False):
         """ Run TRUNCATE SQL query """
 
-        sql = 'TRUNCATE "{0}"'.format(table)
+        sql = 'TRUNCATE {0}'.format(self._quote_table_name(table))
 
         if just_try:
             log.debug("Just-try mode: execute TRUNCATE query: %s", sql)
@@ -348,7 +357,7 @@ class PgDB:
         else:
             sql += '"{0}"'.format('", "'.join(fields))
 
-        sql += ' FROM "{0}"'.format(table)
+        sql += ' FROM {0}'.format(self._quote_table_name(table))
         params = dict()
 
         try:
diff --git a/tests/test_pgsql.py b/tests/test_pgsql.py
index 254d94d..c2d719f 100644
--- a/tests/test_pgsql.py
+++ b/tests/test_pgsql.py
@@ -251,6 +251,11 @@ def test_add_where_clauses_with_duplicated_field():
     )
 
 
+def test_quote_table_name():
+    assert PgDB._quote_table_name("mytable") == '"mytable"'
+    assert PgDB._quote_table_name("myschema.mytable") == '"myschema"."mytable"'
+
+
 def test_insert(mocker, test_pgdb):
     values = dict(test1=1, test2=2)
     mocker.patch(