feat: split function

Author: dungdm93

sqlalchemy_trino/datatype.py

@@ -1,4 +1,5 @@
 import re
+from typing import *
 
 from sqlalchemy import util
 from sqlalchemy.sql import sqltypes
@@ -53,24 +54,74 @@
 }
 
 
-def parse_sqltype(type_str: str, column: str) -> TypeEngine:
-    type_str = type_str.lower()
-    m = re.match(r'^([\w\s]+)(?:\(([\d,\s]*)\))?', type_str)
-    if m is None:
+class MAP(TypeEngine):
+    pass
+
+
+class ROW(TypeEngine):
+    pass
+
+
+def split(string: str, delimiter: str = ',',
+          quote: str = '"', escaped_quote: str = r'\"',
+          open_bracket: str = '(', close_bracket: str = ')') -> Iterator[str]:
+    """
+    A split function that is aware of quotes and brackets/parentheses.
+
+    :param string: string to split
+    :param delimiter: string defining where to split, usually a comma or space
+    :param quote: string, either a single or a double quote
+    :param escaped_quote: string representing an escaped quote
+    :param open_bracket: string, either [, {, < or (
+    :param close_bracket: string, either ], }, > or )
+    """
+    parens = 0
+    quotes = False
+    i = 0
+    for j, character in enumerate(string):
+        complete = parens == 0 and not quotes
+        if complete and character == delimiter:
+            yield string[i:j]
+            i = j + len(delimiter)
+        elif character == open_bracket:
+            parens += 1
+        elif character == close_bracket:
+            parens -= 1
+        elif character == quote:
+            if quotes and string[j - len(escaped_quote) + 1: j + 1] != escaped_quote:
+                quotes = False
+            elif not quotes:
+                quotes = True
+    yield string[i:]
+
+
+def parse_sqltype(type_str: str) -> TypeEngine:
+    type_str = type_str.strip().lower()
+    match = re.match(r'^(?P<type>\w+)\s*(?:\((?P<options>.*)\))?', type_str)
+    if not match:
         util.warn(f"Could not parse type name '{type_str}'")
         return sqltypes.NULLTYPE
-    type_name, type_opts = m.groups()  # type: str, str
-    type_name = type_name.strip()
+    type_name = match.group("type")
+    type_opts = match.group("options")
+
+    if type_name == "array":
+        item_type = parse_sqltype(type_opts)
+        return sqltypes.ARRAY(item_type)
+    elif type_name == "map":
+        key_type_str, value_type_str = split(type_opts)
+        key_type = parse_sqltype(key_type_str)
+        value_type = parse_sqltype(value_type_str)
+        return MAP(key_type, value_type)
+    elif type_name == "row":
+        attr_types = split(type_opts)
+        return ROW()  # TODO
+
     if type_name not in _type_map:
-        util.warn(f"Did not recognize type '{type_name}' of column '{column}'")
+        util.warn(f"Did not recognize type '{type_name}'")
         return sqltypes.NULLTYPE
     type_class = _type_map[type_name]
     type_args = [int(o.strip()) for o in type_opts.split(',')] if type_opts else []
     if type_name in ('time', 'timestamp'):
         type_kwargs = dict(timezone=type_str.endswith("with time zone"))
-        # TODO: handle time/timestamp(p) precision
-        return type_class(**type_kwargs)
-    if type_name in ('array', 'map', 'row'):
-        # TODO
-        return sqltypes.NULLTYPE
+        return type_class(**type_kwargs)  # TODO: handle time/timestamp(p) precision
     return type_class(*type_args)

sqlalchemy_trino/dialect.py

@@ -97,7 +97,7 @@ def get_columns(self, connection: Connection,
             for row in rows:
                 columns.append(dict(
                     name=row.Column,
-                    type=datatype.parse_sqltype(row.Type, row.Column),
+                    type=datatype.parse_sqltype(row.Type),
                     nullable=getattr(row, 'Null', True),
                     default=None,
                 ))