content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import unittest
  
  from pyparsing import ParseException
  
  from databuilder.models.table_metadata import ColumnMetadata
  from databuilder.models.type_metadata import (
      ArrayTypeMetadata, MapTypeMetadata, ScalarTypeMetadata, StructTypeMetadata,
  )
  from databuilder.utils.trino_complex_type_parser import parse_trino_type
  
  
  class TestTrinoComplexTypeParser(unittest.TestCase):
      def setUp(self) -> None:
          self.column_key = 'trino://gold.test_schema/test_table/col1'
  
      def test_transform_no_complex_type(self) -> None:
          column = ColumnMetadata('col1', None, 'int', 0)
          column.set_column_key(self.column_key)
  
          scalar_type = ScalarTypeMetadata(name='col1',
                                           parent=column,
                                           type_str='int')
  
          actual = parse_trino_type(column.type, column.name, column)
          self.assertEqual(actual, scalar_type)
  
      def test_transform_array_type(self) -> None:
          column = ColumnMetadata('col1', None, 'array(array(int))', 0)
          column.set_column_key(self.column_key)
  
          array_type = ArrayTypeMetadata(name='col1',
                                         parent=column,
                                         type_str='array(array(int))')
          inner_array = ArrayTypeMetadata(name='_inner_',
                                          parent=array_type,
                                          type_str='array(int)')
  
          array_type.array_inner_type = inner_array
  
          actual = parse_trino_type(column.type, column.name, column)
          self.assertEqual(actual, array_type)
  
      def test_transform_array_map_nested_type(self) -> None:
          column = ColumnMetadata('col1', None, 'array(map(string,int))', 0)
          column.set_column_key(self.column_key)
  
          array_type = ArrayTypeMetadata(name='col1',
                                         parent=column,
                                         type_str='array(map(string,int))')
          inner_map = MapTypeMetadata(name='_inner_',
                                      parent=array_type,
                                      type_str='map(string,int)')
          inner_map_key = ScalarTypeMetadata(name='_map_key',
                                             parent=inner_map,
                                             type_str='string')
          inner_scalar = ScalarTypeMetadata(name='_map_value',
                                            parent=inner_map,
                                            type_str='int')
  
          array_type.array_inner_type = inner_map
          inner_map.map_key_type = inner_map_key
          inner_map.map_value_type = inner_scalar
  
          actual = parse_trino_type(column.type, column.name, column)
          self.assertEqual(actual, array_type)
  
      def test_transform_array_struct_nested_type(self) -> None:
          column = ColumnMetadata('col1', None, 'array(row(nest1 int,nest2 int))', 0)
          column.set_column_key(self.column_key)
  
          array_type = ArrayTypeMetadata(name='col1',
                                         parent=column,
                                         type_str='array(row(nest1 int,nest2 int))')
          inner_struct = StructTypeMetadata(name='_inner_',
                                            parent=array_type,
                                            type_str='row(nest1 int,nest2 int)')
          inner_scalar_nest1 = ScalarTypeMetadata(name='nest1',
                                                  parent=inner_struct,
                                                  type_str='int')
          inner_scalar_nest2 = ScalarTypeMetadata(name='nest2',
                                                  parent=inner_struct,
                                                  type_str='int')
  
          array_type.array_inner_type = inner_struct
          inner_struct.struct_items = {'nest1': inner_scalar_nest1, 'nest2': inner_scalar_nest2}
          inner_scalar_nest1.sort_order = 0
          inner_scalar_nest2.sort_order = 1
  
          actual = parse_trino_type(column.type, column.name, column)
          self.assertEqual(actual, array_type)
  
      def test_transform_map_type(self) -> None:
          column = ColumnMetadata('col1', None, 'map(string,map(string,int))', 0)
          column.set_column_key(self.column_key)
  
          map_type = MapTypeMetadata(name='col1',
                                     parent=column,
                                     type_str='map(string,map(string,int))')
          map_key = ScalarTypeMetadata(name='_map_key',
                                       parent=map_type,
                                       type_str='string')
          map_value = MapTypeMetadata(name='_map_value',
                                      parent=map_type,
                                      type_str='map(string,int)')
          inner_map_key = ScalarTypeMetadata(name='_map_key',
                                             parent=map_value,
                                             type_str='string')
          inner_scalar = ScalarTypeMetadata(name='_map_value',
                                            parent=map_value,
                                            type_str='int')
  
          map_type.map_key_type = map_key
          map_type.map_value_type = map_value
          map_value.map_key_type = inner_map_key
          map_value.map_value_type = inner_scalar
  
          actual = parse_trino_type(column.type, column.name, column)
          self.assertEqual(actual, map_type)
  
      def test_transform_map_struct_nested_type(self) -> None:
          column = ColumnMetadata('col1', None, 'map(string,row(nest1 int,nest2 int))', 0)
          column.set_column_key(self.column_key)
  
          map_type = MapTypeMetadata(name='col1',
                                     parent=column,
                                     type_str='map(string,row(nest1 int,nest2 int))')
          map_key = ScalarTypeMetadata(name='_map_key',
                                       parent=map_type,
                                       type_str='string')
          inner_struct = StructTypeMetadata(name='_map_value',
                                            parent=map_type,
                                            type_str='row(nest1 int,nest2 int)')
          inner_scalar_nest1 = ScalarTypeMetadata(name='nest1',
                                                  parent=inner_struct,
                                                  type_str='int')
          inner_scalar_nest2 = ScalarTypeMetadata(name='nest2',
                                                  parent=inner_struct,
                                                  type_str='int')
  
          map_type.map_key_type = map_key
          map_type.map_value_type = inner_struct
          inner_struct.struct_items = {'nest1': inner_scalar_nest1, 'nest2': inner_scalar_nest2}
          inner_scalar_nest1.sort_order = 0
          inner_scalar_nest2.sort_order = 1
  
          actual = parse_trino_type(column.type, column.name, column)
          self.assertEqual(actual, map_type)
  
      def test_transform_struct_type(self) -> None:
          column = ColumnMetadata('col1', None, 'row(nest1 int,nest2 int)', 0)
          column.set_column_key(self.column_key)
  
          struct_type = StructTypeMetadata(name='col1',
                                           parent=column,
                                           type_str='row(nest1 int,nest2 int)')
          inner_scalar_nest1 = ScalarTypeMetadata(name='nest1',
                                                  parent=struct_type,
                                                  type_str='int')
          inner_scalar_nest2 = ScalarTypeMetadata(name='nest2',
                                                  parent=struct_type,
                                                  type_str='int')
  
          struct_type.struct_items = {'nest1': inner_scalar_nest1, 'nest2': inner_scalar_nest2}
          inner_scalar_nest1.sort_order = 0
          inner_scalar_nest2.sort_order = 1
  
          actual = parse_trino_type(column.type, column.name, column)
          self.assertEqual(actual, struct_type)
  
      def test_transform_struct_map_array_nested_type(self) -> None:
          column = ColumnMetadata('col1', None, 'row(nest1 map(string,array(int)),nest2 array(string))', 0)
          column.set_column_key(self.column_key)
  
          struct_type = StructTypeMetadata(name='col1',
                                           parent=column,
                                           type_str='row(nest1 map(string,array(int)),nest2 array(string))')
          inner_map = MapTypeMetadata(name='nest1',
                                      parent=struct_type,
                                      type_str='map(string,array(int))')
          inner_map_key = ScalarTypeMetadata(name='_map_key',
                                             parent=inner_map,
                                             type_str='string')
          inner_map_array = ArrayTypeMetadata(name='_map_value',
                                              parent=inner_map,
                                              type_str='array(int)')
          inner_struct_array = ArrayTypeMetadata(name='nest2',
                                                 parent=struct_type,
                                                 type_str='array(string)')
  
          struct_type.struct_items = {'nest1': inner_map, 'nest2': inner_struct_array}
          inner_map.map_key_type = inner_map_key
          inner_map.map_value_type = inner_map_array
          inner_map.sort_order = 0
          inner_struct_array.sort_order = 1
  
          actual = parse_trino_type(column.type, column.name, column)
          self.assertEqual(actual, struct_type)
  
      def test_transform_struct_nested_type_with_quoted_names(self) -> None:
          column = ColumnMetadata('col1', None, 'row("nest1" varchar,"nest2" row("nest3" varchar,'
                                                '"nest4" timestamp(3),"nest5" timestamp(3)))', 0)
          column.set_column_key(self.column_key)
  
          struct_type = StructTypeMetadata(name='col1',
                                           parent=column,
                                           type_str='row(nest1 varchar,nest2 row(nest3 varchar,'
                                                    'nest4 timestamp(3),nest5 timestamp(3)))')
          inner_scalar = ScalarTypeMetadata(name='nest1',
                                            parent=struct_type,
                                            type_str='varchar')
          inner_struct = StructTypeMetadata(name='nest2',
                                            parent=struct_type,
                                            type_str='row(nest3 varchar,nest4 timestamp(3),nest5 timestamp(3))')
          inner_scalar_1 = ScalarTypeMetadata(name='nest3',
                                              parent=inner_struct,
                                              type_str='varchar')
          inner_scalar_2 = ScalarTypeMetadata(name='nest4',
                                              parent=inner_struct,
                                              type_str='timestamp(3)')
          inner_scalar_3 = ScalarTypeMetadata(name='nest5',
                                              parent=inner_struct,
                                              type_str='timestamp(3)')
  
          struct_type.struct_items = {'nest1': inner_scalar, 'nest2': inner_struct}
          inner_struct.struct_items = {'nest3': inner_scalar_1, 'nest4': inner_scalar_2, 'nest5': inner_scalar_3}
          inner_scalar.sort_order = 0
          inner_struct.sort_order = 1
          inner_scalar_1.sort_order = 0
          inner_scalar_2.sort_order = 1
          inner_scalar_3.sort_order = 2
  
          actual = parse_trino_type(column.type, column.name, column)
          self.assertEqual(actual, struct_type)
  
      def test_transform_non_alpha_only_types(self) -> None:
          column = ColumnMetadata('col1', None, 'row(nest1 decimal(10,2),nest2 double precision,'
                                                'nest3 varchar(32),nest4 map(varchar(32),decimal(10,2)),'
                                                'nest5 interval_day_time)', 0)
          column.set_column_key(self.column_key)
  
          struct_type = StructTypeMetadata(name='col1',
                                           parent=column,
                                           type_str='row(nest1 decimal(10,2),nest2 double precision,'
                                                    'nest3 varchar(32),nest4 map(varchar(32),decimal(10,2)),'
                                                    'nest5 interval_day_time)')
          inner_scalar_nest1 = ScalarTypeMetadata(name='nest1',
                                                  parent=struct_type,
                                                  type_str='decimal(10,2)')
          inner_scalar_nest2 = ScalarTypeMetadata(name='nest2',
                                                  parent=struct_type,
                                                  type_str='double precision')
          inner_scalar_nest3 = ScalarTypeMetadata(name='nest3',
                                                  parent=struct_type,
                                                  type_str='varchar(32)')
          inner_map_nest4 = MapTypeMetadata(name='nest4',
                                            parent=struct_type,
                                            type_str='map(varchar(32),decimal(10,2))')
          inner_map_nest4_key = ScalarTypeMetadata(name='_map_key',
                                                   parent=inner_map_nest4,
                                                   type_str='varchar(32)')
          inner_map_nest4_value = ScalarTypeMetadata(name='_map_value',
                                                     parent=inner_map_nest4,
                                                     type_str='decimal(10,2)')
          inner_scalar_nest5 = ScalarTypeMetadata(name='nest5',
                                                  parent=struct_type,
                                                  type_str='interval_day_time')
  
          struct_type.struct_items = {'nest1': inner_scalar_nest1, 'nest2': inner_scalar_nest2,
                                      'nest3': inner_scalar_nest3, 'nest4': inner_map_nest4,
                                      'nest5': inner_scalar_nest5}
          inner_map_nest4.map_key_type = inner_map_nest4_key
          inner_map_nest4.map_value_type = inner_map_nest4_value
          inner_scalar_nest1.sort_order = 0
          inner_scalar_nest2.sort_order = 1
          inner_scalar_nest3.sort_order = 2
          inner_map_nest4.sort_order = 3
          inner_scalar_nest5.sort_order = 4
  
          actual = parse_trino_type(column.type, column.name, column)
          self.assertEqual(actual, struct_type)
  
      def test_transform_invalid_array_inner_type(self) -> None:
          column = ColumnMetadata('col1', None, 'array(array(int*))', 0)
          column.set_column_key(self.column_key)
  
          with self.assertRaises(ParseException):
              parse_trino_type(column.type, column.name, column)
  
      def test_transform_invalid_struct_inner_type(self) -> None:
          column = ColumnMetadata('col1', None, 'row(nest1 varchar(256)å,'
                                                'nest2 (derived from deserializer))', 0)
          column.set_column_key(self.column_key)
  
          with self.assertRaises(ParseException):
              parse_trino_type(column.type, column.name, column)
  
  
  if __name__ == '__main__':
      unittest.main()
  