diff options
| -rw-r--r-- | pgcli/packages/parseutils.py | 4 | ||||
| -rw-r--r-- | tests/test_parseutils.py | 5 |
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'] |