Skip to content

Commit d9020db

Browse files
Mariamalmesferpratyakshsharmaagrawalreetika
authored and
Mariam-Almesfer
committed
Add support for geometry and geography in postgres connector data type in postgres connector
Co-authored-by: pratyakshsharma <[email protected]>[email protected]> Co-authored-by: agrawalreetika <[email protected]>
1 parent 21a37b8 commit d9020db

4 files changed

Lines changed: 195 additions & 2 deletions

File tree

‎presto-docs/src/main/sphinx/connector/postgresql.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,8 @@ The connector maps PostgreSQL types to the corresponding PrestoDB types:
141141
- ``JSON``
142142
* - ``JSONB``
143143
- ``JSON``
144+
* - ``GEOMETRY``
145+
- ``VARCHAR``
144146

145147
No other types are supported.
146148

‎presto-postgresql/pom.xml‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,30 @@
9999
<scope>provided</scope>
100100
</dependency>
101101

102+
<dependency>
103+
<groupId>com.facebook.presto</groupId>
104+
<artifactId>presto-geospatial-toolkit</artifactId>
105+
</dependency>
106+
107+
<dependency>
108+
<groupId>org.openjdk.jol</groupId>
109+
<artifactId>jol-core</artifactId>
110+
<scope>provided</scope>
111+
</dependency>
112+
113+
<dependency>
114+
<groupId>com.esri.geometry</groupId>
115+
<artifactId>esri-geometry-api</artifactId>
116+
</dependency>
117+
102118
<!-- for testing -->
119+
120+
<dependency>
121+
<groupId>com.h2database</groupId>
122+
<artifactId>h2</artifactId>
123+
<scope>test</scope>
124+
</dependency>
125+
103126
<dependency>
104127
<groupId>com.facebook.presto</groupId>
105128
<artifactId>presto-testng-services</artifactId>

‎presto-postgresql/src/main/java/com/facebook/presto/plugin/postgresql/PostgreSqlClient.java‎

Lines changed: 74 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
*/
1414
package com.facebook.presto.plugin.postgresql;
1515

16+
import com.esri.core.geometry.ogc.OGCGeometry;
1617
import com.facebook.airlift.json.JsonObjectMapperProvider;
1718
import com.facebook.presto.common.type.StandardTypes;
1819
import com.facebook.presto.common.type.Type;
@@ -21,9 +22,12 @@
2122
import com.facebook.presto.plugin.jdbc.BaseJdbcClient;
2223
import com.facebook.presto.plugin.jdbc.BaseJdbcConfig;
2324
import com.facebook.presto.plugin.jdbc.DriverConnectionFactory;
25+
import com.facebook.presto.plugin.jdbc.JdbcColumnHandle;
2426
import com.facebook.presto.plugin.jdbc.JdbcConnectorId;
2527
import com.facebook.presto.plugin.jdbc.JdbcIdentity;
28+
import com.facebook.presto.plugin.jdbc.JdbcSplit;
2629
import com.facebook.presto.plugin.jdbc.JdbcTypeHandle;
30+
import com.facebook.presto.plugin.jdbc.QueryBuilder;
2731
import com.facebook.presto.plugin.jdbc.ReadMapping;
2832
import com.facebook.presto.spi.ConnectorSession;
2933
import com.facebook.presto.spi.ConnectorTableMetadata;
@@ -48,20 +52,35 @@
4852
import java.sql.PreparedStatement;
4953
import java.sql.ResultSet;
5054
import java.sql.SQLException;
55+
import java.util.List;
56+
import java.util.Map;
5157
import java.util.Optional;
5258
import java.util.UUID;
5359

60+
import static com.esri.core.geometry.ogc.OGCGeometry.fromBinary;
61+
import static com.facebook.presto.common.type.StandardTypes.GEOMETRY;
62+
import static com.facebook.presto.common.type.StandardTypes.SPHERICAL_GEOGRAPHY;
5463
import static com.facebook.presto.common.type.VarbinaryType.VARBINARY;
64+
import static com.facebook.presto.common.type.VarcharType.VARCHAR;
65+
import static com.facebook.presto.geospatial.GeometryUtils.wktFromJtsGeometry;
66+
import static com.facebook.presto.geospatial.serde.EsriGeometrySerde.serialize;
67+
import static com.facebook.presto.geospatial.serde.JtsGeometrySerde.deserialize;
5568
import static com.facebook.presto.plugin.jdbc.JdbcErrorCode.JDBC_ERROR;
69+
import static com.facebook.presto.plugin.jdbc.QueryBuilder.quote;
70+
import static com.facebook.presto.plugin.jdbc.ReadMapping.sliceReadMapping;
5671
import static com.facebook.presto.spi.StandardErrorCode.ALREADY_EXISTS;
5772
import static com.facebook.presto.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
5873
import static com.fasterxml.jackson.core.JsonFactory.Feature.CANONICALIZE_FIELD_NAMES;
5974
import static com.fasterxml.jackson.databind.SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS;
75+
import static com.google.common.collect.ImmutableMap.toImmutableMap;
6076
import static io.airlift.slice.Slices.utf8Slice;
77+
import static io.airlift.slice.Slices.wrappedBuffer;
6178
import static io.airlift.slice.Slices.wrappedLongArray;
6279
import static java.lang.Long.reverseBytes;
6380
import static java.lang.String.format;
6481
import static java.nio.charset.StandardCharsets.UTF_8;
82+
import static java.util.Objects.requireNonNull;
83+
import static java.util.function.Function.identity;
6584

6685
public class PostgreSqlClient
6786
extends BaseJdbcClient
@@ -114,19 +133,72 @@ protected String toSqlType(Type type)
114133
return super.toSqlType(type);
115134
}
116135

136+
@Override
137+
public PreparedStatement buildSql(ConnectorSession session, Connection connection, JdbcSplit split, List<JdbcColumnHandle> columnHandles) throws SQLException
138+
{
139+
Map<String, String> columnExpressions = columnHandles.stream()
140+
.filter(handle -> handle.getJdbcTypeHandle().getJdbcTypeName().equalsIgnoreCase(GEOMETRY))
141+
.map(JdbcColumnHandle::getColumnName)
142+
.collect(toImmutableMap(
143+
identity(),
144+
columnName -> "ST_AsBinary(" + quote(identifierQuote, columnName) + ")"));
145+
146+
return new QueryBuilder(identifierQuote).buildSql(
147+
this,
148+
session,
149+
connection,
150+
split.getCatalogName(),
151+
split.getSchemaName(),
152+
split.getTableName(),
153+
columnHandles,
154+
columnExpressions,
155+
split.getTupleDomain(),
156+
split.getAdditionalPredicate());
157+
}
158+
117159
@Override
118160
public Optional<ReadMapping> toPrestoType(ConnectorSession session, JdbcTypeHandle typeHandle)
119161
{
162+
String typeName = typeHandle.getJdbcTypeName();
163+
120164
if (typeHandle.getJdbcTypeName().equals("jsonb") || typeHandle.getJdbcTypeName().equals("json")) {
121165
return Optional.of(jsonColumnMapping());
122166
}
123167

124168
else if (typeHandle.getJdbcTypeName().equals("uuid")) {
125169
return Optional.of(uuidColumnMapping());
126170
}
171+
else if (typeName.equalsIgnoreCase(GEOMETRY) || typeName.equalsIgnoreCase(SPHERICAL_GEOGRAPHY)) {
172+
return Optional.of(geometryReadMapping());
173+
}
127174
return super.toPrestoType(session, typeHandle);
128175
}
129176

177+
protected static ReadMapping geometryReadMapping()
178+
{
179+
return sliceReadMapping(VARCHAR,
180+
(resultSet, columnIndex) -> getAsText(stGeomFromBinary(wrappedBuffer(resultSet.getBytes(columnIndex)))));
181+
}
182+
183+
protected static Slice getAsText(Slice input)
184+
{
185+
return utf8Slice(wktFromJtsGeometry(deserialize(input)));
186+
}
187+
188+
private static Slice stGeomFromBinary(Slice input)
189+
{
190+
requireNonNull(input, "input is null");
191+
OGCGeometry geometry;
192+
try {
193+
geometry = fromBinary(input.toByteBuffer().slice());
194+
}
195+
catch (IllegalArgumentException | IndexOutOfBoundsException e) {
196+
throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Invalid Well-Known Binary (WKB)", e);
197+
}
198+
geometry.setSpatialReference(null);
199+
return serialize(geometry);
200+
}
201+
130202
@Override
131203
public void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
132204
{
@@ -159,7 +231,7 @@ protected void renameTable(JdbcIdentity identity, String catalogName, SchemaTabl
159231

160232
private ReadMapping jsonColumnMapping()
161233
{
162-
return ReadMapping.sliceReadMapping(
234+
return sliceReadMapping(
163235
jsonType,
164236
(resultSet, columnIndex) -> jsonParse(utf8Slice(resultSet.getString(columnIndex))));
165237
}
@@ -190,7 +262,7 @@ public static JsonParser createJsonParser(JsonFactory factory, Slice json)
190262

191263
private ReadMapping uuidColumnMapping()
192264
{
193-
return ReadMapping.sliceReadMapping(
265+
return sliceReadMapping(
194266
uuidType,
195267
(resultSet, columnIndex) -> uuidSlice((UUID) resultSet.getObject(columnIndex)));
196268
}

‎presto-postgresql/src/test/java/com/facebook/presto/plugin/postgresql/TestPostgreSqlClient.java‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package com.facebook.presto.plugin.postgresql;
15+
16+
import com.esri.core.geometry.Point;
17+
import com.esri.core.geometry.ogc.OGCGeometry;
18+
import com.esri.core.geometry.ogc.OGCPoint;
19+
import com.facebook.presto.plugin.jdbc.SliceReadFunction;
20+
import com.facebook.presto.spi.PrestoException;
21+
import io.airlift.slice.Slice;
22+
import io.airlift.slice.Slices;
23+
import org.h2.tools.SimpleResultSet;
24+
import org.testng.annotations.Test;
25+
26+
import java.nio.ByteBuffer;
27+
import java.sql.SQLException;
28+
29+
import static com.facebook.presto.geospatial.GeoFunctions.stGeomFromBinary;
30+
import static com.facebook.presto.plugin.postgresql.PostgreSqlClient.geometryReadMapping;
31+
import static com.facebook.presto.plugin.postgresql.PostgreSqlClient.getAsText;
32+
import static org.testng.Assert.assertEquals;
33+
import static org.testng.Assert.fail;
34+
35+
public class TestPostgreSqlClient
36+
{
37+
@Test
38+
public void testValidGeometryReadMapping()
39+
{
40+
OGCGeometry geometry = new OGCPoint(new Point(1.0, 2.0), null);
41+
ByteBuffer buffer = geometry.asBinary();
42+
43+
Slice value = getAsText(stGeomFromBinary(Slices.wrappedBuffer(buffer)));
44+
45+
assertEquals(value.toStringUtf8(), "POINT (1 2)");
46+
}
47+
48+
@Test
49+
public void testInvalidGeometryReadMapping()
50+
{
51+
byte[] invalidWkb = new byte[]{0x00, 0x01, 0x02};
52+
try {
53+
getAsText(stGeomFromBinary(Slices.wrappedBuffer(invalidWkb)));
54+
fail("stGeomFromBinary should throw");
55+
}
56+
catch (PrestoException e) {
57+
assertEquals(e.getMessage(), "Invalid WKB");
58+
}
59+
}
60+
61+
@Test
62+
public void testReadMapping() throws SQLException
63+
{
64+
SliceReadFunction fn = (SliceReadFunction) geometryReadMapping().getReadFunction();
65+
OGCGeometry geometry = new OGCPoint(new Point(1.0, 2.0), null);
66+
ByteBuffer buffer = geometry.asBinary();
67+
Slice value = fn.readSlice(new MockResultSet(buffer.array()), 1);
68+
assertEquals(value.toStringUtf8(), "POINT (1 2)");
69+
70+
byte[] invalid = new byte[] {0x01, 0x02, 0x03};
71+
try {
72+
fn.readSlice(new MockResultSet(invalid), 1);
73+
fail("stGeomFromBinary should throw");
74+
}
75+
catch (PrestoException e) {
76+
assertEquals(e.getMessage(), "Invalid Well-Known Binary (WKB)");
77+
}
78+
}
79+
80+
private static class MockResultSet
81+
extends SimpleResultSet
82+
{
83+
private final byte[] bytes;
84+
85+
public MockResultSet(byte[] bytes)
86+
{
87+
this.bytes = bytes;
88+
}
89+
90+
@Override
91+
public byte[] getBytes(int columnIndex)
92+
{
93+
return bytes;
94+
}
95+
}
96+
}

0 commit comments

Comments
 (0)