Development of an internal social media platform with personalised dashboards for students
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

statement_splitter.py 3.6KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105
  1. # -*- coding: utf-8 -*-
  2. #
  3. # Copyright (C) 2016 Andi Albrecht, albrecht.andi@gmail.com
  4. #
  5. # This module is part of python-sqlparse and is released under
  6. # the BSD License: https://opensource.org/licenses/BSD-3-Clause
  7. from sqlparse import sql, tokens as T
  8. class StatementSplitter(object):
  9. """Filter that split stream at individual statements"""
  10. def __init__(self):
  11. self._reset()
  12. def _reset(self):
  13. """Set the filter attributes to its default values"""
  14. self._in_declare = False
  15. self._is_create = False
  16. self._begin_depth = 0
  17. self.consume_ws = False
  18. self.tokens = []
  19. self.level = 0
  20. def _change_splitlevel(self, ttype, value):
  21. """Get the new split level (increase, decrease or remain equal)"""
  22. # ANSI
  23. # if normal token return
  24. # wouldn't parenthesis increase/decrease a level?
  25. # no, inside a paranthesis can't start new statement
  26. if ttype not in T.Keyword:
  27. return 0
  28. # Everything after here is ttype = T.Keyword
  29. # Also to note, once entered an If statement you are done and basically
  30. # returning
  31. unified = value.upper()
  32. # three keywords begin with CREATE, but only one of them is DDL
  33. # DDL Create though can contain more words such as "or replace"
  34. if ttype is T.Keyword.DDL and unified.startswith('CREATE'):
  35. self._is_create = True
  36. return 0
  37. # can have nested declare inside of being...
  38. if unified == 'DECLARE' and self._is_create and self._begin_depth == 0:
  39. self._in_declare = True
  40. return 1
  41. if unified == 'BEGIN':
  42. self._begin_depth += 1
  43. if self._is_create:
  44. # FIXME(andi): This makes no sense.
  45. return 1
  46. return 0
  47. # Should this respect a preceeding BEGIN?
  48. # In CASE ... WHEN ... END this results in a split level -1.
  49. # Would having multiple CASE WHEN END and a Assigment Operator
  50. # cause the statement to cut off prematurely?
  51. if unified == 'END':
  52. self._begin_depth = max(0, self._begin_depth - 1)
  53. return -1
  54. if (unified in ('IF', 'FOR', 'WHILE') and
  55. self._is_create and self._begin_depth > 0):
  56. return 1
  57. if unified in ('END IF', 'END FOR', 'END WHILE'):
  58. return -1
  59. # Default
  60. return 0
  61. def process(self, stream):
  62. """Process the stream"""
  63. EOS_TTYPE = T.Whitespace, T.Comment.Single
  64. # Run over all stream tokens
  65. for ttype, value in stream:
  66. # Yield token if we finished a statement and there's no whitespaces
  67. # It will count newline token as a non whitespace. In this context
  68. # whitespace ignores newlines.
  69. # why don't multi line comments also count?
  70. if self.consume_ws and ttype not in EOS_TTYPE:
  71. yield sql.Statement(self.tokens)
  72. # Reset filter and prepare to process next statement
  73. self._reset()
  74. # Change current split level (increase, decrease or remain equal)
  75. self.level += self._change_splitlevel(ttype, value)
  76. # Append the token to the current statement
  77. self.tokens.append(sql.Token(ttype, value))
  78. # Check if we get the end of a statement
  79. if self.level <= 0 and ttype is T.Punctuation and value == ';':
  80. self.consume_ws = True
  81. # Yield pending statement (if any)
  82. if self.tokens:
  83. yield sql.Statement(self.tokens)