Improve handling of quoted indentifiers (fixes issue78).

andialbrecht · andialbrecht · commit 39ddba76709c · 2012-11-03T13:05:25.000+01:00
diff --git a/CHANGES b/CHANGES
@@ -1,3 +1,15 @@
+Development Version
+-------------------
+
+Bug Fixes
+ * Improve handling of quoted identifiers (issue78).
+
+Other
+ * Deprecated sqlparse.SQLParseError. Please use
+   sqlparse.exceptions.SQLParseError instead.
+
+
+
 Release 0.1.4 (Apr 20, 2012)
 ----------------------------
 
@@ -14,10 +26,6 @@ Bug Fixes
  * Pretty format comments in identifier lists (issue59).
  * Several minor bug fixes and improvements.
 
-Other
- * Deprecated sqlparse.SQLParseError. Please use
-   sqlparse.exceptions.SQLParseError instead.
-
 
 Release 0.1.3 (Jul 29, 2011)
 ----------------------------
diff --git a/sqlparse/sql.py b/sqlparse/sql.py
@@ -169,6 +169,14 @@ def _pprint_tree(self, max_depth=None, depth=0):
             if (token.is_group() and (max_depth is None or depth < max_depth)):
                 token._pprint_tree(max_depth, depth + 1)
 
+    def _remove_quotes(self, val):
+        """Helper that removes surrounding quotes from strings."""
+        if not val:
+            return val
+        if val[0] in ('"', '\'') and val[-1] == val[0]:
+            val = val[1:-1]
+        return val
+
     def flatten(self):
         """Generator yielding ungrouped tokens.
 
@@ -362,11 +370,13 @@ def get_alias(self):
         else:
             next_ = self.token_next_by_instance(0, Identifier)
             if next_ is None:
-                return None
+                next_ = self.token_next_by_type(0, T.String.Symbol)
+                if next_ is None:
+                    return None
             alias = next_
         if isinstance(alias, Identifier):
             return alias.get_name()
-        return unicode(alias)
+        return self._remove_quotes(unicode(alias))
 
     def get_name(self):
         """Returns the name of this identifier.
@@ -385,13 +395,16 @@ def get_real_name(self):
         # a.b
         dot = self.token_next_match(0, T.Punctuation, '.')
         if dot is None:
-            return self.token_next_by_type(0, T.Name).value
+            next_ = self.token_next_by_type(0, T.Name)
+            if next_ is not None:
+                return self._remove_quotes(next_.value)
+            return None
 
         next_ = self.token_next_by_type(self.token_index(dot),
-                                        (T.Name, T.Wildcard))
+                                        (T.Name, T.Wildcard, T.String.Symbol))
         if next_ is None:  # invalid identifier, e.g. "a."
             return None
-        return next_.value
+        return self._remove_quotes(next_.value)
 
 
 class Statement(TokenList):
@@ -437,7 +450,7 @@ def get_parent_name(self):
         prev_ = self.token_prev(self.token_index(dot))
         if prev_ is None:  # something must be verry wrong here..
             return None
-        return prev_.value
+        return self._remove_quotes(prev_.value)
 
     def is_wildcard(self):
         """Return ``True`` if this identifier contains a wildcard."""
diff --git a/tests/test_parse.py b/tests/test_parse.py
@@ -106,3 +106,10 @@ def test_keyword_like_identifier(self):  # see issue47
         t = sqlparse.parse('foo.key')[0].tokens
         self.assertEqual(len(t), 1)
         self.assert_(isinstance(t[0], sqlparse.sql.Identifier))
+
+
+def test_quoted_identifier():
+    t = sqlparse.parse('select x.y as "z" from foo')[0].tokens
+    assert isinstance(t[2], sqlparse.sql.Identifier)
+    assert t[2].get_name() == 'z'
+    assert t[2].get_real_name() == 'y'
diff --git a/tests/test_regressions.py b/tests/test_regressions.py
@@ -106,3 +106,31 @@ def test_issue40(self):
              '  (SELECT id,\n'
              '          name\n'
              '   FROM bar) as foo'))
+
+
+def test_issue78():
+    # the bug author provided this nice examples, let's use them!
+    def _get_identifier(sql):
+        p = sqlparse.parse(sql)[0]
+        return p.tokens[2]
+    results = (('get_name', 'z'),
+               ('get_real_name', 'y'),
+               ('get_parent_name', 'x'),
+               ('get_alias', 'z'),
+               ('get_typecast', 'text'))
+    variants = (
+        'select x.y::text as z from foo',
+        'select x.y::text as "z" from foo',
+        'select x."y"::text as z from foo',
+        'select x."y"::text as "z" from foo',
+        'select "x".y::text as z from foo',
+        'select "x".y::text as "z" from foo',
+        'select "x"."y"::text as z from foo',
+        'select "x"."y"::text as "z" from foo',
+    )
+    for variant in variants:
+        i = _get_identifier(variant)
+        assert isinstance(i, sql.Identifier)
+        for func_name, result in results:
+            func = getattr(i, func_name)
+            assert func() == result
