support update statement

README.rst

@@ -71,7 +71,7 @@ Quickstart
 Features
 --------
 
-* Current version only support DDL(CREATE table, ALTER table), DML(simple SELECT, INSERT, DELETE).
+* Current version only support DDL(CREATE table, ALTER table), DML(simple SELECT, INSERT, DELETE, UPDATE).
 * TODO：SUPPORT MORE SQL STATEMENTS
 
 Credits

mysqltokenparser/mysqltokenparser.py

@@ -9,14 +9,14 @@
 from MySqlParser import MySqlParser
 from MySqlParserListener import MySqlParserListener
 from sqltypemixins import (
-    AlterTableMixin, CreateTableMixin, SelectMixin, InsertMixin, DeleteMixin
+    AlterTableMixin, CreateTableMixin, SelectMixin, InsertMixin, DeleteMixin, UpdateMixin
 )
 from constant import *
 
 
 class MyListener(
-    AlterTableMixin, CreateTableMixin, SelectMixin, InsertMixin, DeleteMixin,
-    MySqlParserListener
+    AlterTableMixin, CreateTableMixin, SelectMixin, InsertMixin,
+    DeleteMixin, UpdateMixin, MySqlParserListener
 ):
     def __init__(self, ret):
         self.ret = ret

mysqltokenparser/sqltypemixins/__init__.py

@@ -5,3 +5,4 @@
 from select import *
 from insert import *
 from delete import *
+from update import *

mysqltokenparser/sqltypemixins/select.py

@@ -64,6 +64,26 @@ def _enterPredicateExpression(self, child, ret):
         if isinstance(child, MySqlParser.ExpressionAtomPredicateContext):
             where_expression.extend(self._enterExpressionAtomPredicate(child).get('where_expression'))
 
+    @iterchild
+    def _enterExpressionAtomPredicateForFullColumnNameExpressionAtomContext(self, child, ret):
+        full_columns = ret.setdefault('full_columns', [])
+        if isinstance(child, MySqlParser.FullColumnNameExpressionAtomContext):
+            full_columns.append(''.join(self._enterFullColumnName(child).get('full_column')))
+        if isinstance(child, MySqlParser.ConstantExpressionAtomContext):
+            full_columns.append(self._enterConstantExpressionAtom(child).get('constant_expression'))
+
+    @iterchild
+    def _enterFullColumnName(self, child, ret):
+        if isinstance(child, MySqlParser.FullColumnNameContext):
+            ret.update(self._enterFullColumnNameDetail(child))
+
+    @iterchild
+    def _enterFullColumnNameDetail(self, child, ret):
+        full_column = ret.setdefault('full_column', [])
+        if isinstance(child, MySqlParser.UidContext) or\
+            isinstance(child, MySqlParser.DottedIdContext):
+            full_column.append(self._get_last_name(child))
+
     @iterchild
     def _enterExpressionAtomPredicate(self, child, ret):
         where_expression = ret.setdefault('where_expression', [])
@@ -90,7 +110,11 @@ def _enterNestedExpressionAtom(self, child, ret):
     def _enterBinaryComparasionPredicate(self, child, ret):
         where_expression = ret.setdefault('where_expression', [])
         if isinstance(child, MySqlParser.ExpressionAtomPredicateContext):
-            where_expression.append(self._get_last_name(child))
+            where_expression.extend(
+                self._enterExpressionAtomPredicateForFullColumnNameExpressionAtomContext(
+                    child
+                ).get('full_columns', [])
+            )
 
         if isinstance(child, MySqlParser.ComparisonOperatorContext):
             where_expression.append(''.join(self._enterComparisonOperator(child).get('comparison_oper')))

mysqltokenparser/sqltypemixins/update.py

@@ -0,0 +1,50 @@
+# coding: utf-8
+import antlr4
+from mysqltokenparser.utils import iterchild
+from mysqltokenparser.MySqlParser import MySqlParser
+from mysqltokenparser.constant import *
+
+
+class UpdateMixin(object):
+    """
+    delete type:
+        singleDeleteStatement multipleDeleteStatement
+    """
+    def enterUpdateStatement(self, ctx):
+        data = {}
+        self.ret['data'] = {
+            'type': DML_TYPE_UPDATESTATEMENT,
+            'data': data
+        }
+
+        children = ctx.children
+        for child in children:
+            if isinstance(child, MySqlParser.MultipleUpdateStatementContext):
+                data.update(self._enterMultipleUpdateStatement(child))
+
+            if isinstance(child, MySqlParser.SingleUpdateStatementContext):
+                data.update(self._enterSingleUpdateStatement(child))
+
+    @iterchild
+    def _enterMultipleUpdateStatement(self, child, ret):
+        pass
+
+    @iterchild
+    def _enterSingleUpdateStatement(self, child, ret):
+        if isinstance(child, MySqlParser.TableNameContext):
+            ret[TABLE_NAME] = self._get_last_name(child)
+
+        if isinstance(child, MySqlParser.PredicateExpressionContext):
+            ret.update(self._enterPredicateExpression(child))
+
+        if isinstance(child, MySqlParser.UpdatedElementContext):
+            update_elements = ret.setdefault('update_elements', [])
+            update_elements.append(self._enterUpdatedElement(child).get('update_element'))
+
+    @iterchild
+    def _enterUpdatedElement(self, child, ret):
+        update_element = ret.setdefault('update_element', [])
+        if isinstance(child, MySqlParser.FullColumnNameContext):
+            update_element.append(self._get_last_name(child))
+        if isinstance(child, MySqlParser.PredicateExpressionContext):
+            update_element.extend(self._enterPredicateExpression(child).get('where_expression'))

tests/test_update.py

@@ -0,0 +1,16 @@
+#!/usr/bin/env python
+# coding:utf-8
+
+import pytest
+from helper import mysqltokenparser as mtp
+from helper import constant as _c
+
+
+def test_simpleselect():
+    sql = u"""
+    UPDATE table_name SET field1='new-value1', field2='new-value2' WHERE id=23 and name="css";
+    """
+
+    tokens = mtp.mysql_token_parser(sql)
+    print(tokens)
+    assert isinstance(tokens, dict)