Просмотр исходного кода

[beeswax] Fix parsing complex PartitionKey types

This fixes parsing complex types that contain a variety of sigils
in them.
Erick Tryzelaar 11 лет назад
Родитель
Сommit
5c799cc
2 измененных файлов с 34 добавлено и 22 удалено
  1. 15 7
      apps/beeswax/src/beeswax/server/hive_server2_lib.py
  2. 19 15
      apps/beeswax/src/beeswax/tests.py

+ 15 - 7
apps/beeswax/src/beeswax/server/hive_server2_lib.py

@@ -73,8 +73,8 @@ class HiveServerTable(Table):
     match = re.search('partitionKeys:\[([^\]]+)\]', describe)
     if match is not None:
       match = match.group(1)
-      return [PartitionKeyCompatible(partition)
-              for partition in re.findall('FieldSchema\((.+?)\)', match)]
+      return [PartitionKeyCompatible(*partition)
+          for partition in re.findall('FieldSchema\(name:(.+?), type:(.+?), comment:(.+?)\)', match)]
     else:
       return []
 
@@ -682,12 +682,20 @@ class ResultCompatible:
 
 class PartitionKeyCompatible:
 
-  def __init__(self, partition):
+  def __init__(self, name, type, comment):
     # Parses: ['name:datehour, type:int, comment:null']
-    name, type, comment = partition.split(', ', 2)
-    self.name = name.split(':', 1)[1]
-    self.type = type.split(':', 1)[1]
-    self.comment = comment.split(':', 1)[1]
+    self.name = name
+    self.type = type
+    self.comment = comment
+
+  def __eq__(self, other):
+    return isinstance(other, PartitionKeyCompatible) and \
+        self.name == other.name and \
+        self.type == other.type and \
+        self.comment == other.comment
+
+  def __repr__(self):
+    return 'PartitionKey(name:%s, type:%s, comment:%s)' % (self.name, self.type, self.comment)
 
 
 class PartitionValueCompatible:

+ 19 - 15
apps/beeswax/src/beeswax/tests.py

@@ -1789,19 +1789,6 @@ class MockHiveServerTable(HiveServerTable):
 
 class TestHiveServer2API():
 
-  def test_partition_keys(self):
-    table = MockHiveServerTable({'path_location': '/my/table'})
-
-    key = PartitionKeyCompatible('name:a_name, type:int, comment:null')
-    assert_equal('a_name', key.name)
-    assert_equal('int', key.type)
-    assert_equal('null', key.comment)
-
-    key = PartitionKeyCompatible('name:a_name, type:int, comment:this, has extra: sigils')
-    assert_equal('a_name', key.name)
-    assert_equal('int', key.type)
-    assert_equal('this, has extra: sigils', key.comment)
-
   def test_partition_values(self):
     table = MockHiveServerTable({'path_location': '/my/table'})
 
@@ -1856,7 +1843,9 @@ class TestHiveServer2API():
           '), '
           'partitionKeys:['
             'FieldSchema(name:dt, type:string, comment:null), '
-            'FieldSchema(name:country, type:string, comment:null)'
+            'FieldSchema(name:country, type:string, comment:null), '
+            'FieldSchema(name:decimal, type:decimal(9, 7), comment:this, has extra: sigils), '
+            'FieldSchema(name:complex, type:UNIONTYPE<int, double, array<string>, struct<a:int,b:string>>, comment:null), '
           '], '
           'parameters:{'
             'numPartitions=0, '
@@ -1897,7 +1886,16 @@ class TestHiveServer2API():
                     ['comment', 'null)'],
                     ['FieldSchema(name', 'country'],
                     ['type', 'string'],
-                    ['comment', 'null)]'],
+                    ['comment', 'null)'],
+                    ['FieldSchema(name', 'decimal'],
+                    ['type', 'decimal(9'],
+                    #['7)'],
+                    ['comment', 'this'],
+                    ['has extra', ' sigils)'],
+                    ['FieldSchema(name', 'complex'],
+                    ['type', 'UNIONTYPE<int'],
+                    ['struct<a:int,b', 'string>>'],
+                    ['comment', 'null)'],
                     ['parameters', '{numPartitions=0'],
                     ['numFiles', '1'],
                     ['transient_lastDdlTime', '1360732885'],
@@ -1907,6 +1905,12 @@ class TestHiveServer2API():
                     ['tableType', 'MANAGED_TABLE']
                   ],
                   table.properties)
+
+      assert_equal([PartitionKeyCompatible('dt', 'string', 'null'),
+                    PartitionKeyCompatible('country', 'string', 'null'),
+                    PartitionKeyCompatible('decimal', 'decimal(9, 7)', 'this, has extra: sigils'),
+                    PartitionKeyCompatible('complex', 'UNIONTYPE<int, double, array<string>, struct<a:int,b:string>>', 'null'),
+                   ], table.partition_keys)
     finally:
       setattr(table, 'extended_describe', prev_extended_describe)