|
18 | 18 | * #L% |
19 | 19 | */ |
20 | 20 | package com.amazonaws.athena.connectors.snowflake; |
21 | | -import com.amazonaws.athena.connector.lambda.domain.Split; |
22 | | -import org.junit.Assert; |
23 | | -import org.junit.Test; |
| 21 | + |
| 22 | +import com.amazonaws.athena.connector.lambda.domain.predicate.Constraints; |
| 23 | +import com.amazonaws.athena.connectors.jdbc.manager.TypeAndValue; |
| 24 | +import org.apache.arrow.vector.types.DateUnit; |
| 25 | +import org.apache.arrow.vector.types.TimeUnit; |
| 26 | +import org.apache.arrow.vector.types.pojo.ArrowType; |
| 27 | +import org.apache.arrow.vector.types.pojo.Field; |
| 28 | +import org.apache.arrow.vector.types.pojo.FieldType; |
| 29 | +import org.apache.arrow.vector.types.pojo.Schema; |
| 30 | +import org.junit.jupiter.api.BeforeEach; |
| 31 | +import org.junit.jupiter.api.Test; |
24 | 32 | import org.mockito.Mock; |
25 | | -import org.mockito.Mockito; |
| 33 | +import org.mockito.MockitoAnnotations; |
26 | 34 |
|
| 35 | +import java.math.BigDecimal; |
| 36 | +import java.sql.Connection; |
| 37 | +import java.sql.SQLException; |
27 | 38 | import java.util.ArrayList; |
28 | | -import java.util.Collections; |
29 | 39 | import java.util.List; |
30 | 40 |
|
31 | | -import static com.amazonaws.athena.connectors.snowflake.SnowflakeConstants.SNOWFLAKE_QUOTE_CHARACTER; |
| 41 | +import static org.junit.jupiter.api.Assertions.assertEquals; |
| 42 | +import static org.junit.jupiter.api.Assertions.assertThrows; |
| 43 | +import static org.junit.jupiter.api.Assertions.assertTrue; |
| 44 | +import static org.mockito.Mockito.mock; |
| 45 | +import static org.mockito.Mockito.when; |
| 46 | + |
| 47 | +class SnowflakeQueryStringBuilderTest { |
| 48 | + |
| 49 | + private SnowflakeQueryStringBuilder queryBuilder; |
32 | 50 |
|
33 | | -public class SnowflakeQueryStringBuilderTest |
34 | | -{ |
35 | 51 | @Mock |
36 | | - Split split; |
37 | | - |
38 | | - //@Test |
39 | | -// public void testQueryBuilderNew() |
40 | | -// { |
41 | | -// Split split = Mockito.mock(Split.class); |
42 | | -// SnowflakeQueryStringBuilder builder = new SnowflakeQueryStringBuilder(SNOWFLAKE_QUOTE_CHARACTER, new SnowflakeFederationExpressionParser(SNOWFLAKE_QUOTE_CHARACTER)); |
43 | | -// Mockito.when(split.getProperties()).thenReturn(Collections.singletonMap("partition", "p0")); |
44 | | -// Mockito.when(split.getProperty(Mockito.eq("partition"))).thenReturn("p1-p2-p3-p4-p5-p6-p7"); |
45 | | -// builder.getFromClauseWithSplit("default", "", "table", split); |
46 | | -// builder.appendLimitOffset(split); |
47 | | -// } |
48 | | - |
49 | | - //@Test |
50 | | -// public void testGetPartitionWhereClauses() |
51 | | -// { |
52 | | -// SnowflakeQueryStringBuilder builder = new SnowflakeQueryStringBuilder(SNOWFLAKE_QUOTE_CHARACTER, new SnowflakeFederationExpressionParser(SNOWFLAKE_QUOTE_CHARACTER)); |
53 | | -// List<String> fromClauseWithSplit = builder.getPartitionWhereClauses(split); |
54 | | -// List<String> expected = new ArrayList<>(); |
55 | | -// Assert.assertEquals(expected, fromClauseWithSplit); |
56 | | -// } |
57 | | -} |
| 52 | + private Connection mockConnection; |
| 53 | + |
| 54 | + @BeforeEach |
| 55 | + void setUp() { |
| 56 | + MockitoAnnotations.openMocks(this); |
| 57 | + queryBuilder = new SnowflakeQueryStringBuilder("\"", null); |
| 58 | + } |
| 59 | + |
| 60 | + @Test |
| 61 | + void testGetFromClauseWithSplit() { |
| 62 | + String result = queryBuilder.getFromClauseWithSplit(null, "public", "users", null); |
| 63 | + assertEquals(" FROM \"public\".\"users\" ", result); |
| 64 | + } |
| 65 | + |
| 66 | + @Test |
| 67 | + void testGetPartitionWhereClauses() { |
| 68 | + List<String> result = queryBuilder.getPartitionWhereClauses(null); |
| 69 | + assertTrue(result.isEmpty()); |
| 70 | + } |
| 71 | + |
| 72 | + @Test |
| 73 | + void testBuildSqlString_NoConstraints() throws SQLException { |
| 74 | + Schema tableSchema = new Schema(List.of(new Field("id", new FieldType(true, new ArrowType.Int(32, true), null), null))); |
| 75 | + |
| 76 | + Constraints constraints = mock(Constraints.class); |
| 77 | + when(constraints.getLimit()).thenReturn(0L); |
| 78 | + |
| 79 | + String sql = queryBuilder.buildSqlString(mockConnection, null, "public", "users", tableSchema, constraints, null); |
| 80 | + assertTrue(sql.contains("SELECT \"id\" FROM \"public\".\"users\" ")); |
| 81 | + } |
| 82 | + |
| 83 | + @Test |
| 84 | + void testBuildSqlString_WithConstraints() throws SQLException { |
| 85 | + Schema tableSchema = new Schema(List.of(new Field("id", new FieldType(true, new ArrowType.Int(32, true), null), null))); |
58 | 86 |
|
| 87 | + Constraints constraints = mock(Constraints.class); |
| 88 | + when(constraints.getLimit()).thenReturn(10L); |
| 89 | + |
| 90 | + String sql = queryBuilder.buildSqlString(mockConnection, null, "public", "users", tableSchema, constraints, null); |
| 91 | + assertTrue(sql.contains("LIMIT 10")); |
| 92 | + } |
| 93 | + |
| 94 | + @Test |
| 95 | + void testQuote() { |
| 96 | + String result = queryBuilder.quote("users"); |
| 97 | + assertEquals("\"users\"", result); |
| 98 | + } |
| 99 | + |
| 100 | + @Test |
| 101 | + void testSingleQuote() { |
| 102 | + String result = queryBuilder.singleQuote("O'Reilly"); |
| 103 | + assertEquals("'O''Reilly'", result); |
| 104 | + } |
| 105 | + |
| 106 | + @Test |
| 107 | + void testToPredicate_SingleValue() { |
| 108 | + List<TypeAndValue> accumulator = new ArrayList<>(); |
| 109 | + String predicate = queryBuilder.toPredicate("age", "=", 30, new ArrowType.Int(32, true)); |
| 110 | + assertEquals("age = 30", predicate); |
| 111 | + } |
| 112 | + |
| 113 | + @Test |
| 114 | + void testGetObjectForWhereClause_Int() { |
| 115 | + Object result = queryBuilder.getObjectForWhereClause("age", 42, new ArrowType.Int(32, true)); |
| 116 | + assertEquals(42L, result); |
| 117 | + } |
| 118 | + |
| 119 | + @Test |
| 120 | + void testGetObjectForWhereClause_Decimal() { |
| 121 | + Object result = queryBuilder.getObjectForWhereClause("price", new BigDecimal("99.99"), new ArrowType.Decimal(10, 2)); |
| 122 | + assertEquals(new BigDecimal("99.99"), result); |
| 123 | + } |
| 124 | + |
| 125 | + @Test |
| 126 | + void testGetObjectForWhereClause_Date() { |
| 127 | + Object result = queryBuilder.getObjectForWhereClause("date", "2023-03-15T00:00", new ArrowType.Date(DateUnit.DAY)); |
| 128 | + assertEquals("2023-03-15 00:00:00", result); |
| 129 | + } |
| 130 | + |
| 131 | + @Test |
| 132 | + void testGetObjectForWhereClause_Timestamp() { |
| 133 | + Object result = queryBuilder.getObjectForWhereClause("created_at", 1711929600000L, new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")); |
| 134 | + assertEquals("2024-04-01 00:00:00", result); |
| 135 | + } |
| 136 | + |
| 137 | + @Test |
| 138 | + void testToPredicateWithUnsupportedType() { |
| 139 | + assertThrows(UnsupportedOperationException.class, () -> |
| 140 | + queryBuilder.getObjectForWhereClause("unsupported", "value", new ArrowType.Struct()) |
| 141 | + ); |
| 142 | + } |
| 143 | +} |
0 commit comments