Fix the table parsing for update and add a test.

2 files changed, 7 insertions, 2 deletions

diff --git a/pgcli/packages/parseutils.py b/pgcli/packages/parseutils.py
index a3de836e..7dc96116 100644
--- a/pgcli/packages/parseutils.py
+++ b/pgcli/packages/parseutils.py
@@ -78,8 +78,8 @@ def extract_from_part(parsed):
                 raise StopIteration
             else:
                 yield item
-        elif item.ttype is Keyword and item.value.upper() in ('FROM', 'INTO',
-                'UPDATE', 'TABLE', ):
+        elif ((item.ttype is Keyword or item.ttype is Keyword.DML) and
+                item.value.upper() in ('FROM', 'INTO', 'UPDATE', 'TABLE', )):
             tbl_prefix_seen = True
         # 'SELECT a, FROM abc' will detect FROM as part of the column list.
         # So this check here is necessary.
diff --git a/tests/test_parseutils.py b/tests/test_parseutils.py
index 63198d57..0d299786 100644
--- a/tests/test_parseutils.py
+++ b/tests/test_parseutils.py
@@ -27,3 +27,8 @@ def test_select_with_hanging_comma_multiple_tables():
 def test_simple_insert_single_table():
     tables = extract_tables('insert into abc (id, name) values (1, "def")')
     assert tables == ['abc']
+
+def test_simple_update_table():
+    import pdb; pdb.set_trace()
+    tables = extract_tables('update abc set id = 1')
+    assert tables == ['abc']