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(