2929import java .sql .ResultSetMetaData ;
3030import java .sql .SQLException ;
3131import java .sql .Statement ;
32+ import java .util .Arrays ;
3233import java .util .Collections ;
3334import java .util .HashMap ;
3435import java .util .HashSet ;
3738import java .util .Set ;
3839
3940import org .apache .arrow .memory .RootAllocator ;
41+ import org .apache .arrow .vector .types .Types ;
42+ import org .apache .arrow .vector .types .pojo .ArrowType ;
43+ import org .apache .arrow .vector .types .pojo .Field ;
44+ import org .apache .arrow .vector .types .pojo .FieldType ;
4045import org .apache .arrow .vector .types .pojo .Schema ;
41- import org .apache .arrow .vector .util .ObjectMapperFactory ;
4246import org .junit .After ;
4347import org .junit .Before ;
4448import org .junit .Test ;
4549
46- import com .fasterxml .jackson .databind .ObjectWriter ;
47-
4850public class JdbcToArrowCommentMetadataTest {
4951
5052 private static final String COMMENT = "comment" ; //use this metadata key for interoperability with Spark StructType
51- private final ObjectWriter schemaSerializer = ObjectMapperFactory .newObjectMapper ().writerWithDefaultPrettyPrinter ();
5253 private Connection conn = null ;
5354
5455 /**
@@ -73,26 +74,85 @@ public void tearDown() throws SQLException {
7374 }
7475 }
7576
77+ private static Field field (String name , boolean nullable , ArrowType type , Map <String , String > metadata ) {
78+ return new Field (name , new FieldType (nullable , type , null , metadata ), Collections .emptyList ());
79+ }
80+
81+ private static Map <String , String > metadata (String ... entries ) {
82+ if (entries .length % 2 != 0 ) {
83+ throw new IllegalArgumentException ("Map must have equal number of keys and values" );
84+ }
85+
86+ final Map <String , String > result = new HashMap <>();
87+ for (int i = 0 ; i < entries .length ; i += 2 ) {
88+ result .put (entries [i ], entries [i + 1 ]);
89+ }
90+ return result ;
91+ }
92+
7693 @ Test
7794 public void schemaComment () throws Exception {
7895 boolean includeMetadata = false ;
79- String schemaJson = schemaSerializer .writeValueAsString (getSchemaWithCommentFromQuery (includeMetadata ));
80- String expectedSchema = getExpectedSchema ("/h2/expectedSchemaWithComments.json" );
81- assertThat (schemaJson ).isEqualTo (expectedSchema );
96+ Schema schema = getSchemaWithCommentFromQuery (includeMetadata );
97+ Schema expectedSchema = new Schema (Arrays .asList (
98+ field ("ID" , false , Types .MinorType .BIGINT .getType (),
99+ metadata ("comment" , "Record identifier" )),
100+ field ("NAME" , true , Types .MinorType .VARCHAR .getType (),
101+ metadata ("comment" , "Name of record" )),
102+ field ("COLUMN1" , true , Types .MinorType .BIT .getType (),
103+ metadata ()),
104+ field ("COLUMNN" , true , Types .MinorType .INT .getType (),
105+ metadata ("comment" , "Informative description of columnN" ))
106+ ), metadata ("comment" , "This is super special table with valuable data" ));
107+ assertThat (schema ).isEqualTo (expectedSchema );
82108 }
83109
84110 @ Test
85111 public void schemaCommentWithDatabaseMetadata () throws Exception {
86112 boolean includeMetadata = true ;
87- String schemaJson = schemaSerializer .writeValueAsString (getSchemaWithCommentFromQuery (includeMetadata ));
88- String expectedSchema = getExpectedSchema ("/h2/expectedSchemaWithCommentsAndJdbcMeta.json" );
113+ Schema schema = getSchemaWithCommentFromQuery (includeMetadata );
114+ Schema expectedSchema = new Schema (Arrays .asList (
115+ field ("ID" , false , Types .MinorType .BIGINT .getType (),
116+ metadata (
117+ "SQL_CATALOG_NAME" , "JDBCTOARROWTEST?CHARACTERENCODING=UTF-8" ,
118+ "SQL_SCHEMA_NAME" , "PUBLIC" ,
119+ "SQL_TABLE_NAME" , "TABLE1" ,
120+ "SQL_COLUMN_NAME" , "ID" ,
121+ "SQL_TYPE" , "BIGINT" ,
122+ "comment" , "Record identifier"
123+ )),
124+ field ("NAME" , true , Types .MinorType .VARCHAR .getType (),
125+ metadata (
126+ "SQL_CATALOG_NAME" , "JDBCTOARROWTEST?CHARACTERENCODING=UTF-8" ,
127+ "SQL_SCHEMA_NAME" , "PUBLIC" ,
128+ "SQL_TABLE_NAME" , "TABLE1" ,
129+ "SQL_COLUMN_NAME" , "NAME" ,
130+ "SQL_TYPE" , "VARCHAR" ,
131+ "comment" , "Name of record" )),
132+ field ("COLUMN1" , true , Types .MinorType .BIT .getType (),
133+ metadata (
134+ "SQL_CATALOG_NAME" , "JDBCTOARROWTEST?CHARACTERENCODING=UTF-8" ,
135+ "SQL_SCHEMA_NAME" , "PUBLIC" ,
136+ "SQL_TABLE_NAME" , "TABLE1" ,
137+ "SQL_COLUMN_NAME" , "COLUMN1" ,
138+ "SQL_TYPE" , "BOOLEAN" )),
139+ field ("COLUMNN" , true , Types .MinorType .INT .getType (),
140+ metadata (
141+ "SQL_CATALOG_NAME" , "JDBCTOARROWTEST?CHARACTERENCODING=UTF-8" ,
142+ "SQL_SCHEMA_NAME" , "PUBLIC" ,
143+ "SQL_TABLE_NAME" , "TABLE1" ,
144+ "SQL_COLUMN_NAME" , "COLUMNN" ,
145+ "SQL_TYPE" , "INTEGER" ,
146+ "comment" , "Informative description of columnN" ))
147+ ), metadata ("comment" , "This is super special table with valuable data" ));
148+ assertThat (schema ).isEqualTo (expectedSchema );
89149 /* corresponding Apache Spark DDL after conversion:
90150 ID BIGINT NOT NULL COMMENT 'Record identifier',
91151 NAME STRING COMMENT 'Name of record',
92152 COLUMN1 BOOLEAN,
93153 COLUMNN INT COMMENT 'Informative description of columnN'
94154 */
95- assertThat (schemaJson ).isEqualTo (expectedSchema );
155+ assertThat (schema ).isEqualTo (expectedSchema );
96156 }
97157
98158 private Schema getSchemaWithCommentFromQuery (boolean includeMetadata ) throws SQLException {
0 commit comments