sqlconvert/process_tokens.py

   1
   2 from sqlparse.sql import Comment
   3 from sqlobject.converters import sqlrepr
   4 from sqlparse import parse
   5 from sqlparse.compat import PY3
   6 from sqlparse import tokens as T
   7
   8
   9 def find_error(token_list):
  10     """Find an error"""
  11     for token in token_list.flatten():
  12         if token.ttype is T.Error:
  13             return True
  14     return False
  15
  16
  17 def is_comment_or_space(token):
  18     return isinstance(token, Comment) or \
  19         token.ttype in (T.Comment, T.Comment.Single, T.Comment.Multiline,
  20                         T.Newline, T.Whitespace)
  21
  22
  23 def is_newline_statement(statement):
  24     for token in statement.tokens[:]:
  25         if token.ttype is not T.Newline:
  26             return False
  27     return True
  28
  29
  30 def escape_strings(token_list, dbname):
  31     """Escape strings"""
  32     for token in token_list.flatten():
  33         if token.ttype is T.String.Single:
  34             value = token.value[1:-1]  # unquote by removing apostrophes
  35             value = sqlrepr(value, dbname)
  36             token.normalized = token.value = value
  37
  38
  39 if PY3:
  40     xrange = range
  41
  42
  43 class StatementGrouper(object):
  44     """Collect lines and reparse until the last statement is complete"""
  45
  46     def __init__(self, encoding=None):
  47         self.lines = []
  48         self.statements = []
  49         self.encoding = encoding
  50
  51     def process_line(self, line):
  52         self.lines.append(line)
  53         self.process_lines()
  54
  55     def process_lines(self):
  56         statements = parse(''.join(self.lines), encoding=self.encoding)
  57         last_stmt = statements[-1]
  58         for i in xrange(len(last_stmt.tokens) - 1, 0, -1):
  59             token = last_stmt.tokens[i]
  60             if is_comment_or_space(token):
  61                 continue
  62             if token.ttype is T.Punctuation and token.value == ';':
  63                 break  # The last statement is complete
  64             # The last statement is still incomplete - wait for the next line
  65             return
  66         self.lines = []
  67         self.statements = statements
  68
  69     def get_statements(self):
  70         for stmt in self.statements:
  71             yield stmt
  72         self.statements = []
  73         return
  74
  75     def close(self):
  76         if not self.lines:
  77             return
  78         tokens = parse(''.join(self.lines), encoding=self.encoding)
  79         for token in tokens:
  80             if not is_comment_or_space(token):
  81                 raise ValueError("Incomplete SQL statement: %s" %
  82                                  tokens)
  83         self.lines = []
  84         self.statements = []
  85         return tokens