Add reserved words support for select

Author: passren

pydynamodb/__init__.py

@@ -6,7 +6,7 @@
 if TYPE_CHECKING:
     from .connection import Connection
 
-__version__: str = "0.7.4"
+__version__: str = "0.7.5"
 
 # Globals https://www.python.org/dev/peps/pep-0249/#globals
 apilevel: str = "2.0"

pydynamodb/sql/common.py

@@ -351,6 +351,12 @@ def get_query_type(sql: str) -> QueryType:
     raise LookupError("Not supported query type")
 
 
+def escape_keyword(word: str) -> str:
+    if word.upper() in RESERVED_WORDS:
+        return f"\"{word}\""
+    else:
+        return word
+
 # DynamoDB reserved words
 # https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/ReservedWords.html
 RESERVED_WORDS = [

pydynamodb/sql/dml_select.py

@@ -36,7 +36,7 @@
 from abc import ABCMeta
 from collections import OrderedDict
 from .dml_sql import DmlBase, DmlFunction
-from .common import KeyWords, Tokens
+from .common import KeyWords, Tokens, escape_keyword
 from pyparsing import Opt, Forward, Group, ZeroOrMore, delimited_list, Regex
 from typing import Any, Dict, List, Optional
 
@@ -212,7 +212,7 @@ def _construct_columns(self, columns: List[Any]) -> str:
                 return "*"
 
             column_ = []
-            column_.append(column["column_name"])
+            column_.append(escape_keyword(column["column_name"]))
 
             for rcolumn in column["column_ops"]:
                 column_.append(rcolumn["arithmetic_operators"])

pydynamodb/sql/dml_sql.py

@@ -36,7 +36,7 @@ class DmlBase(Base):
         ],
     )
 
-    ATTR_NAME = Opt('"') + Word(alphanums + "_-") + Opt('"')
+    ATTR_NAME = Opt('"') + Word(alphanums + "_-")("attr_name").set_name("attr_name") + Opt('"')
     ATTR_ARRAY_NAME = ATTR_NAME + "[" + Word(nums) + "]"
 
     _COLUMN_NAME = (

tests/test_cursor_dml.py

@@ -379,7 +379,7 @@ def test_reserved_words(self, cursor):
         sql_reserved_words_1 = (
             """
         INSERT INTO %s VALUE {
-            'key_partition': ?, 'key_sort': ?, 'username': ?, 'password': ?
+            'key_partition': ?, 'key_sort': ?, 'username': ?, 'password': ?, 'default': ?, 'comment': ?
         }
         """
             % USER_TABLE
@@ -389,14 +389,27 @@ def test_reserved_words(self, cursor):
             0,
             "admin",
             "admin",
+            1,
+            "",
         ]
         cursor.execute(sql_reserved_words_1, params_1_)
 
+        cursor.execute(
+            """
+        SELECT username, password, "default", "comment" FROM %s
+        WHERE key_partition=? AND key_sort=?
+        """
+            % USER_TABLE,
+            ["test_user_row_1", 0],
+        )
+        assert cursor.fetchone() == ("admin", "admin", 1, "")
+
         sql_reserved_words_2 = (
             """
         UPDATE %s
             SET username=?
             SET password=?
+            SET "default"=?
             WHERE key_partition=? AND key_sort=?
             RETURNING ALL OLD *
         """
@@ -405,11 +418,22 @@ def test_reserved_words(self, cursor):
         params_2_ = [
             "admin1",
             "admin1",
+            0,
             "test_user_row_1",
             0,
         ]
         cursor.execute(sql_reserved_words_2, params_2_)
 
+        cursor.execute(
+            """
+        SELECT username, password, "default" FROM %s
+        WHERE key_partition=? AND key_sort=?
+        """
+            % USER_TABLE,
+            ["test_user_row_1", 0],
+        )
+        assert cursor.fetchone() == ("admin1", "admin1", 0)
+
         sql_reserved_words_3 = (
             """
         DELETE FROM %s

tests/test_dml_select.py

@@ -187,7 +187,7 @@ def test_parse_completed_case(self):
         """
         ret = SQLParser(sql).transform()
         assert ret == {
-            "Statement": 'SELECT IssueId,Total,Content.DateWatched[0] FROM "Issues"."CreateDateIndex" '
+            "Statement": 'SELECT IssueId,"Total",Content.DateWatched[0] FROM "Issues"."CreateDateIndex" '
             + "WHERE IssueId IN [100,300,234] "
             + "AND Title = 'some title' "
             + "AND Content[0] >= 100 "

tests/test_sqlalchemy_dynamodb.py

@@ -7,6 +7,7 @@
 Base = declarative_base()
 
 TESTCASE02_TABLE = "pydynamodb_test_case02"
+USER_TABLE = "user"
 
 
 # Declarative Mapping
@@ -19,6 +20,16 @@ class _TestCase02(Base):
     col_num = Column(Numeric)
     col_nested = Column()
 
+class _User(Base):
+    __tablename__ = USER_TABLE
+
+    key_partition = Column(String, primary_key=True)
+    key_sort = Column(Integer, primary_key=True)
+    username = Column(String)
+    password = Column(String)
+    default = Column(Numeric)
+    comment = Column(String)
+
 
 class TestSQLAlchemyDynamoDB:
     def test_ping(self, engine):
@@ -370,3 +381,60 @@ def test_data_limit(self, engine):
         )
         rows = conn.execute(table.select().limit(1)).fetchall()
         assert len(rows) == 1
+
+    def test_reserved_word_table_insert(self, engine):
+        engine, conn = engine
+
+        with Session(engine) as session:
+            for i in range(0, 5):
+                user = _User()
+                user.key_partition = "test_user_row_2"
+                user.key_sort = i
+                user.username = "user" + str(i)
+                user.password = "pwd" + str(i)
+                user.default = 1
+                user.comment = "user account"
+                session.add(user)
+            session.commit()
+
+        rows = conn.execute(
+            text(
+                """
+            SELECT * FROM %s WHERE key_partition = :pk
+            """
+                % USER_TABLE
+            ),
+            {"pk": "test_user_row_2"},
+        ).fetchall()
+        assert len(rows) == 5
+
+    def test_reserved_word_table_update(self, engine):
+        engine, conn = engine
+
+        with Session(engine) as session:
+            user = session.scalars(
+                select(_User).where(
+                    _User.key_partition == "test_user_row_2",
+                    _User.key_sort == 1,
+                )
+            ).one()
+            user.username = "user_updated"
+            user.default = 0
+            user.comment = "user account updated"
+            session.commit()
+
+        rows = conn.execute(
+            text(
+                """
+            SELECT username, "default", "comment" FROM %s
+            WHERE key_partition = :pk
+            AND key_sort = :sk
+            """
+                % USER_TABLE
+            ),
+            {"pk": "test_user_row_2", "sk": 1},
+        ).fetchall()
+        assert len(rows) == 1
+        assert rows[0][0] == "user_updated"
+        assert rows[0][1] == 0
+        assert rows[0][2] == "user account updated"