ceph/src/arrow/java/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/h2/JdbcAliasToArrowTest.java

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one or more
   3  * contributor license agreements.  See the NOTICE file distributed with
   4  * this work for additional information regarding copyright ownership.
   5  * The ASF licenses this file to You under the Apache License, Version 2.0
   6  * (the "License"); you may not use this file except in compliance with
   7  * the License.  You may obtain a copy of the License at
   8  *
   9  *    http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 package org.apache.arrow.adapter.jdbc.h2;
  19
  20 import static org.apache.arrow.adapter.jdbc.AbstractJdbcToArrowTest.sqlToArrow;
  21 import static org.junit.Assert.assertEquals;
  22
  23 import java.sql.Connection;
  24 import java.sql.DriverManager;
  25 import java.sql.PreparedStatement;
  26 import java.sql.ResultSet;
  27 import java.sql.ResultSetMetaData;
  28 import java.sql.SQLException;
  29 import java.sql.Statement;
  30 import java.util.List;
  31
  32 import org.apache.arrow.memory.RootAllocator;
  33 import org.apache.arrow.vector.VectorSchemaRoot;
  34 import org.apache.arrow.vector.types.pojo.Field;
  35 import org.apache.arrow.vector.types.pojo.Schema;
  36 import org.junit.After;
  37 import org.junit.Before;
  38 import org.junit.Test;
  39
  40 public class JdbcAliasToArrowTest {
  41   private Connection conn = null;
  42
  43   private static final String CREATE_STATEMENT =
  44       "CREATE TABLE example_table (id INTEGER);";
  45   private static final String INSERT_STATEMENT =
  46       "INSERT INTO example_table (id) VALUES (?);";
  47   private static final String QUERY = "SELECT id as a, id as b FROM example_table;";
  48   private static final String DROP_STATEMENT = "DROP TABLE example_table;";
  49   private static final String ORIGINAL_COLUMN_NAME = "ID";
  50   private static final String COLUMN_A = "A";
  51   private static final String COLUMN_B = "B";
  52
  53   @Before
  54   public void setUp() throws Exception {
  55     String url = "jdbc:h2:mem:JdbcAliasToArrowTest";
  56     String driver = "org.h2.Driver";
  57     Class.forName(driver);
  58     conn = DriverManager.getConnection(url);
  59     try (Statement stmt = conn.createStatement()) {
  60       stmt.executeUpdate(CREATE_STATEMENT);
  61     }
  62   }
  63
  64   /**
  65    * Test h2 database query with alias for column name and column label.
  66    * To vetify reading field alias from an H2 database works as expected.
  67    * If this test fails, something is either wrong with the setup,
  68    * or the H2 SQL behavior changed.
  69    */
  70   @Test
  71   public void testReadH2Alias() throws Exception {
  72     // insert rows
  73     int rowCount = 4;
  74     insertRows(rowCount);
  75
  76     try (ResultSet resultSet = conn.createStatement().executeQuery(QUERY)) {
  77       ResultSetMetaData rsmd = resultSet.getMetaData();
  78       assertEquals(2, rsmd.getColumnCount());
  79
  80       // check column name and column label
  81       assertEquals(ORIGINAL_COLUMN_NAME, rsmd.getColumnName(1));
  82       assertEquals(COLUMN_A, rsmd.getColumnLabel(1));
  83       assertEquals(ORIGINAL_COLUMN_NAME, rsmd.getColumnName(2));
  84       assertEquals(COLUMN_B, rsmd.getColumnLabel(2));
  85
  86       int rowNum = 0;
  87
  88       while (resultSet.next()) {
  89         assertEquals(rowNum, resultSet.getInt(COLUMN_A));
  90         assertEquals(rowNum, resultSet.getInt(COLUMN_B));
  91         ++rowNum;
  92       }
  93
  94       assertEquals(rowCount, rowNum);
  95     }
  96   }
  97
  98   /**
  99    * Test jdbc query results with alias to arrow works expected.
 100    * Arrow result schema name should be field alias name.
 101    */
 102   @Test
 103   public void testJdbcAliasToArrow() throws Exception {
 104     int rowCount = 4;
 105     insertRows(rowCount);
 106
 107     try (ResultSet resultSet = conn.createStatement().executeQuery(QUERY)) {
 108       final VectorSchemaRoot vector =
 109           sqlToArrow(resultSet, new RootAllocator(Integer.MAX_VALUE));
 110
 111       assertEquals(rowCount, vector.getRowCount());
 112       Schema vectorSchema = vector.getSchema();
 113       List<Field> vectorFields = vectorSchema.getFields();
 114       assertEquals(vectorFields.get(0).getName(), COLUMN_A);
 115       assertEquals(vectorFields.get(1).getName(), COLUMN_B);
 116     }
 117   }
 118
 119   @After
 120   public void tearDown() throws SQLException {
 121     try (Statement stmt = conn.createStatement()) {
 122       stmt.executeUpdate(DROP_STATEMENT);
 123     } finally {
 124       if (conn != null) {
 125         conn.close();
 126         conn = null;
 127       }
 128     }
 129   }
 130
 131   private void insertRows(int numRows) throws SQLException {
 132     // Insert [numRows] Rows
 133     try (PreparedStatement stmt = conn.prepareStatement(INSERT_STATEMENT)) {
 134       for (int i = 0; i < numRows; ++i) {
 135         stmt.setInt(1, i);
 136         stmt.executeUpdate();
 137       }
 138     }
 139   }
 140 }