From 492ab455db1d164cce2e49f95e75227fb952e000 Mon Sep 17 00:00:00 2001 From: Kirk Hansen Date: Fri, 30 Aug 2024 10:32:01 -0500 Subject: [PATCH] Trino Dialect: update ARRAY type handling Trino supports type casting arrays, and the syntax accepts ARRAY and ARRAY(DTYPE) --- src/sqlfluff/dialects/dialect_trino.py | 41 +++++++++-- test/fixtures/dialects/trino/array.sql | 6 ++ test/fixtures/dialects/trino/array.yml | 99 +++++++++++++++++++++++++- 3 files changed, 140 insertions(+), 6 deletions(-) diff --git a/src/sqlfluff/dialects/dialect_trino.py b/src/sqlfluff/dialects/dialect_trino.py index 0a565d42100..9f5a9e1be1c 100644 --- a/src/sqlfluff/dialects/dialect_trino.py +++ b/src/sqlfluff/dialects/dialect_trino.py @@ -61,6 +61,16 @@ trino_dialect.add( RightArrowOperator=StringParser("->", SymbolSegment, type="binary_operator"), + StartAngleBracketSegment=StringParser( + "<", SymbolSegment, type="start_angle_bracket" + ), + EndAngleBracketSegment=StringParser(">", SymbolSegment, type="end_angle_bracket"), +) + +trino_dialect.bracket_sets("angle_bracket_pairs").update( + [ + ("angle", "StartAngleBracketSegment", "EndAngleBracketSegment", False), + ] ) trino_dialect.patch_lexer_matchers( @@ -258,7 +268,7 @@ class DatatypeSegment(BaseSegment): Sequence(OneOf("WITH", "WITHOUT"), "TIME", "ZONE", optional=True), ), # Structural - "ARRAY", + Ref("ArrayTypeSegment"), "MAP", Ref("RowTypeSegment"), # Others @@ -466,13 +476,34 @@ class ListaggOverflowClauseSegment(BaseSegment): class ArrayTypeSegment(ansi.ArrayTypeSegment): - """Prefix for array literals. + """Prefix for array literals optionally specifying the type.""" + + type = "array_type" + match_grammar = Sequence( + "ARRAY", + Ref("ArrayTypeSchemaSegment", optional=True), + ) + - Trino supports "ARRAY" +class ArrayTypeSchemaSegment(ansi.ArrayTypeSegment): + """Data type segment of the array. + + Trino supports ARRAY(DATA_TYPE) and ARRAY """ - type = "array_type" - match_grammar = Ref.keyword("ARRAY") + type = "array_type_schema" + match_grammar = OneOf( + Bracketed( + Ref("DatatypeSegment"), + bracket_pairs_set="angle_bracket_pairs", + bracket_type="angle", + ), + Bracketed( + Ref("DatatypeSegment"), + bracket_pairs_set="bracket_pairs", + bracket_type="round", + ), + ) class GroupByClauseSegment(BaseSegment): diff --git a/test/fixtures/dialects/trino/array.sql b/test/fixtures/dialects/trino/array.sql index 067fd727feb..ea28b6b2ee8 100644 --- a/test/fixtures/dialects/trino/array.sql +++ b/test/fixtures/dialects/trino/array.sql @@ -9,3 +9,9 @@ SELECT FILTER(ARRAY[5, -6, NULL, 7], x -> x > 0); SELECT ANY_MATCH(ARRAY[5, -6, NULL, 7], x -> x > 0); SELECT ELEMENT_AT(ARRAY['apple', 'banana', 'orange'], 2); + +SELECT CAST(ARRAY[1,2] AS ARRAY); + +SELECT CAST(ARRAY[1,2] AS ARRAY(DOUBLE)); + +SELECT CAST(JSON_PARSE('[1, 2, 3, 4, 5]') AS ARRAY); diff --git a/test/fixtures/dialects/trino/array.yml b/test/fixtures/dialects/trino/array.yml index f5e93b77b41..8cd0021e8f3 100644 --- a/test/fixtures/dialects/trino/array.yml +++ b/test/fixtures/dialects/trino/array.yml @@ -3,7 +3,7 @@ # computed by SQLFluff when running the tests. Please run # `python test/generate_parse_fixture_yml.py` to generate them after adding or # altering SQL files. -_hash: 08a22fad6dd8609d2246337125972253f4c412c21f1df040fd2349a0e9e4a85c +_hash: d0bf6b8b24c65e7cdb438e4729eb7ce5c42cb21c0eacd859a76bfac0b021f319 file: - statement: select_statement: @@ -200,3 +200,100 @@ file: numeric_literal: '2' - end_bracket: ) - statement_terminator: ; +- statement: + select_statement: + select_clause: + keyword: SELECT + select_clause_element: + function: + function_name: + function_name_identifier: CAST + function_contents: + bracketed: + start_bracket: ( + expression: + typed_array_literal: + array_type: + keyword: ARRAY + array_literal: + - start_square_bracket: '[' + - numeric_literal: '1' + - comma: ',' + - numeric_literal: '2' + - end_square_bracket: ']' + keyword: AS + data_type: + array_type: + keyword: ARRAY + array_type_schema: + start_angle_bracket: < + data_type: + keyword: DOUBLE + end_angle_bracket: '>' + end_bracket: ) +- statement_terminator: ; +- statement: + select_statement: + select_clause: + keyword: SELECT + select_clause_element: + function: + function_name: + function_name_identifier: CAST + function_contents: + bracketed: + start_bracket: ( + expression: + typed_array_literal: + array_type: + keyword: ARRAY + array_literal: + - start_square_bracket: '[' + - numeric_literal: '1' + - comma: ',' + - numeric_literal: '2' + - end_square_bracket: ']' + keyword: AS + data_type: + array_type: + keyword: ARRAY + array_type_schema: + bracketed: + start_bracket: ( + data_type: + keyword: DOUBLE + end_bracket: ) + end_bracket: ) +- statement_terminator: ; +- statement: + select_statement: + select_clause: + keyword: SELECT + select_clause_element: + function: + function_name: + function_name_identifier: CAST + function_contents: + bracketed: + start_bracket: ( + expression: + function: + function_name: + function_name_identifier: JSON_PARSE + function_contents: + bracketed: + start_bracket: ( + expression: + quoted_literal: "'[1, 2, 3, 4, 5]'" + end_bracket: ) + keyword: AS + data_type: + array_type: + keyword: ARRAY + array_type_schema: + start_angle_bracket: < + data_type: + keyword: DOUBLE + end_angle_bracket: '>' + end_bracket: ) +- statement_terminator: ;