add select statement parser

mysqltokenparser/constant.py

@@ -117,6 +117,13 @@
 DML_TYPE_DOSTATEMENT = 'do'
 DML_TYPE_HANDLERSTATEMENT = 'handler'
 
+# select type
+SELECT_TYPE_SIMPLESELECT = 'simpleselect'
+SELECT_TYPE_PARENTHESISSELECT = 'parenthesisselect'
+SELECT_TYPE_UNIONSELECT = 'unionselect'
+SELECT_TYPE_UNIONPARENTHESISSELECT = 'unionparenthesisselect'
+
+
 # table option
 TABLE_OPTION_ENGINE = 'engine'
 TABLE_OPTION_AUTOINCREMENT = 'autoincrement'

mysqltokenparser/mysqltokenparser.py

@@ -8,18 +8,23 @@
 from MySqlLexer import MySqlLexer
 from MySqlParser import MySqlParser
 from MySqlParserListener import MySqlParserListener
-from sqltypemixins.altertable import AlterTableMixin
-from sqltypemixins.createtable import CreateTableMixin
-from constant import SQL_TYPE_DDL
+from sqltypemixins import (
+    AlterTableMixin, CreateTableMixin, SelectMixin
+)
+from constant import *
 
 
-class MyListener(AlterTableMixin, CreateTableMixin, MySqlParserListener):
+class MyListener(
+    AlterTableMixin, CreateTableMixin, SelectMixin, MySqlParserListener):
     def __init__(self, ret):
         self.ret = ret
 
     def enterDdlStatement(self, ctx):
         self.ret['type'] = SQL_TYPE_DDL
 
+    def enterDmlStatement(self, ctx):
+        self.ret['type'] = SQL_TYPE_DML
+
     @staticmethod
     def _get_last_name(ctx):
         while hasattr(ctx, 'children'):

mysqltokenparser/sqltypemixins/__init__.py

@@ -1,2 +1,5 @@
 # coding: utf-8
 
+from altertable import *
+from createtable import *
+from select import *

mysqltokenparser/sqltypemixins/select.py

@@ -0,0 +1,76 @@
+# coding: utf-8
+from mysqltokenparser.utils import iterchild
+from mysqltokenparser.MySqlParser import MySqlParser
+from mysqltokenparser.constant import *
+
+
+class SelectMixin:
+    """
+    select type:
+        simpleSelect parenthesisSelect unionSelect unionParenthesisSelect
+    """
+    def enterSimpleSelect(self, ctx):
+        data = {}
+        self.ret['data'] = {
+            'type': DML_TYPE_SELECTSTATEMENT,
+            'data': {
+                'type': SELECT_TYPE_SIMPLESELECT,
+                'data': data
+            }
+        }
+
+        children = ctx.children
+        for child in children:
+            if isinstance(child, MySqlParser.QuerySpecificationContext):
+                data.update(self._enterQuerySpecification(child))
+
+    @iterchild
+    def _enterQuerySpecification(self, child, ret):
+        if isinstance(child, MySqlParser.SelectElementsContext):
+            ret.update(self._enterSelectElements(child))
+
+        if isinstance(child, MySqlParser.FromClauseContext):
+            ret['from_clause'] = self._enterFromClause(child)
+
+    @iterchild
+    def _enterFromClause(self, child, ret):
+        if isinstance(child, MySqlParser.TableSourcesContext):
+            ret.update(self._enterTableSources(child))
+
+        if isinstance(child, MySqlParser.LogicalExpressionContext):
+            ret.update(self._enterLogicalExpression(child))
+
+    @iterchild
+    def _enterLogicalExpression(self, child, ret):
+        pass
+
+    @iterchild
+    def _enterTableSources(self, child, ret):
+        select_table = ret.setdefault('from_tables', [])
+        if isinstance(child, MySqlParser.TableSourceBaseContext):
+            select_table.append(self._enterTableSourceBase(child))
+
+    @iterchild
+    def _enterTableSourceBase(self, child, ret):
+        if isinstance(child, MySqlParser.AtomTableItemContext):
+            ret.update(self._enterAtomTableItem(child))
+
+    @iterchild
+    def _enterAtomTableItem(self, child, ret):
+        if isinstance(child, MySqlParser.TableNameContext):
+            ret[TABLE_NAME] = self._get_last_name(child)
+        if isinstance(child, MySqlParser.UidContext):
+            ret['as'] = self._get_last_name(child)
+
+    @iterchild
+    def _enterSelectElements(self, child, ret):
+        select_element = ret.setdefault('select_element', [])
+        if isinstance(child, MySqlParser.SelectColumnElementContext):
+            select_element.append(self._enterSelectColumnElement(child))
+
+    @iterchild
+    def _enterSelectColumnElement(self, child, ret):
+        if isinstance(child, MySqlParser.FullColumnNameContext):
+            ret[COLUMN_NAME] = self._get_last_name(child)
+        if isinstance(child, MySqlParser.UidContext):
+            ret['as'] = self._get_last_name(child)

tests/test_select.py

@@ -0,0 +1,18 @@
+#!/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"""
+    select id as qq, name as ww, age as ee, class as rr
+    from student as s, teacher as t
+    where name='cs' and age=26;
+    """
+
+    tokens = mtp.mysql_token_parser(sql)
+    print(tokens)
+    assert isinstance(tokens, dict)