Skip to content

Commit

Permalink
Allow postgresql connector to coerce type during table creation
Browse files Browse the repository at this point in the history
  • Loading branch information
@mdesmet @Praveen2112
mdesmet authored and Praveen2112 committed Sep 4, 2023
1 parent b7380f1 commit d904a6a
Show file tree
Hide file tree
Showing 3 changed files with 113 additions and 19 deletions.
41 changes: 24 additions & 17 deletions plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -661,6 +661,21 @@ private Optional<ColumnMapping> arrayToTrinoType(ConnectorSession session, Conne
throw new IllegalStateException("Unsupported array mapping type: " + arrayMapping);
}

@Override
public Optional<Type> getSupportedType(ConnectorSession session, Type type)
{
if (type instanceof TimeType timeType && timeType.getPrecision() > POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION) {
return Optional.of(createTimeType(POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION));
}
if (type instanceof TimestampType timestampType && timestampType.getPrecision() > POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION) {
return Optional.of(createTimestampType(POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION));
}
if (type instanceof TimestampWithTimeZoneType timestampWithTimeZoneType && timestampWithTimeZoneType.getPrecision() > POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION) {
return Optional.of(createTimestampWithTimeZoneType(POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION));
}
return Optional.empty();
}

@Override
public WriteMapping toWriteMapping(ConnectorSession session, Type type)
{
Expand Down Expand Up @@ -720,29 +735,21 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type)
}

if (type instanceof TimeType timeType) {
if (timeType.getPrecision() <= POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION) {
return WriteMapping.longMapping(format("time(%s)", timeType.getPrecision()), timeWriteFunction(timeType.getPrecision()));
}
return WriteMapping.longMapping(format("time(%s)", POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION), timeWriteFunction(POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION));
verify(timeType.getPrecision() <= POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION);
return WriteMapping.longMapping(format("time(%s)", timeType.getPrecision()), timeWriteFunction(timeType.getPrecision()));
}

if (type instanceof TimestampType timestampType) {
if (timestampType.getPrecision() <= POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION) {
verify(timestampType.getPrecision() <= TimestampType.MAX_SHORT_PRECISION);
return WriteMapping.longMapping(format("timestamp(%s)", timestampType.getPrecision()), PostgreSqlClient::shortTimestampWriteFunction);
}
verify(timestampType.getPrecision() > TimestampType.MAX_SHORT_PRECISION);
return WriteMapping.objectMapping(format("timestamp(%s)", POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION), longTimestampWriteFunction());
verify(timestampType.getPrecision() <= POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION);
return WriteMapping.longMapping(format("timestamp(%s)", timestampType.getPrecision()), PostgreSqlClient::shortTimestampWriteFunction);
}
if (type instanceof TimestampWithTimeZoneType timestampWithTimeZoneType) {
if (timestampWithTimeZoneType.getPrecision() <= POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION) {
String dataType = format("timestamptz(%d)", timestampWithTimeZoneType.getPrecision());
if (timestampWithTimeZoneType.getPrecision() <= TimestampWithTimeZoneType.MAX_SHORT_PRECISION) {
return WriteMapping.longMapping(dataType, shortTimestampWithTimeZoneWriteFunction());
}
return WriteMapping.objectMapping(dataType, longTimestampWithTimeZoneWriteFunction());
verify(timestampWithTimeZoneType.getPrecision() <= POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION);
String dataType = format("timestamptz(%d)", timestampWithTimeZoneType.getPrecision());
if (timestampWithTimeZoneType.isShort()) {
return WriteMapping.longMapping(dataType, shortTimestampWithTimeZoneWriteFunction());
}
return WriteMapping.objectMapping(format("timestamptz(%d)", POSTGRESQL_MAX_SUPPORTED_TIMESTAMP_PRECISION), longTimestampWithTimeZoneWriteFunction());
return WriteMapping.objectMapping(dataType, longTimestampWithTimeZoneWriteFunction());
}
if (type.equals(jsonType)) {
return WriteMapping.sliceMapping("jsonb", typedVarcharWriteFunction("json"));
Expand Down
89 changes: 89 additions & 0 deletions ...rino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConnectorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
import io.trino.testing.sql.TestTable;
import io.trino.testing.sql.TestView;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;

import java.sql.Connection;
Expand Down Expand Up @@ -73,6 +74,7 @@
import static java.util.stream.Collectors.joining;
import static java.util.stream.IntStream.range;
import static org.assertj.core.api.Assertions.assertThat;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;

public class TestPostgreSqlConnectorTest
Expand Down Expand Up @@ -142,6 +144,93 @@ protected TestTable createTableWithUnsupportedColumn()
"(one bigint, two decimal(50,0), three varchar(10))");
}

@Test(dataProvider = "testTimestampPrecisionOnCreateTable")
public void testTimestampPrecisionOnCreateTable(String inputType, String expectedType)
{
try (TestTable testTable = new TestTable(
getQueryRunner()::execute,
"test_coercion_show_create_table",
format("(a %s)", inputType))) {
assertEquals(getColumnType(testTable.getName(), "a"), expectedType);
}
}

@DataProvider(name = "testTimestampPrecisionOnCreateTable")
public static Object[][] timestampPrecisionOnCreateTableProvider()
{
return new Object[][]{
{"timestamp(0)", "timestamp(0)"},
{"timestamp(1)", "timestamp(1)"},
{"timestamp(2)", "timestamp(2)"},
{"timestamp(3)", "timestamp(3)"},
{"timestamp(4)", "timestamp(4)"},
{"timestamp(5)", "timestamp(5)"},
{"timestamp(6)", "timestamp(6)"},
{"timestamp(7)", "timestamp(6)"},
{"timestamp(8)", "timestamp(6)"},
{"timestamp(9)", "timestamp(6)"},
{"timestamp(10)", "timestamp(6)"},
{"timestamp(11)", "timestamp(6)"},
{"timestamp(12)", "timestamp(6)"}
};
}

@Test(dataProvider = "testTimestampPrecisionOnCreateTableAsSelect")
public void testTimestampPrecisionOnCreateTableAsSelect(String inputType, String tableType, String tableValue)
{
try (TestTable testTable = new TestTable(
getQueryRunner()::execute,
"test_coercion_show_create_table",
format("AS SELECT %s a", inputType))) {
assertEquals(getColumnType(testTable.getName(), "a"), tableType);
assertQuery(
format("SELECT * FROM %s", testTable.getName()),
format("VALUES (%s)", tableValue));
}
}

@Test(dataProvider = "testTimestampPrecisionOnCreateTableAsSelect")
public void testTimestampPrecisionOnCreateTableAsSelectWithNoData(String inputType, String tableType, String ignored)
{
try (TestTable testTable = new TestTable(
getQueryRunner()::execute,
"test_coercion_show_create_table",
format("AS SELECT %s a WITH NO DATA", inputType))) {
assertEquals(getColumnType(testTable.getName(), "a"), tableType);
}
}

@DataProvider(name = "testTimestampPrecisionOnCreateTableAsSelect")
public static Object[][] timestampPrecisionOnCreateTableAsSelectProvider()
{
return new Object[][] {
{"TIMESTAMP '1970-01-01 00:00:00'", "timestamp(0)", "TIMESTAMP '1970-01-01 00:00:00'"},
{"TIMESTAMP '1970-01-01 00:00:00.9'", "timestamp(1)", "TIMESTAMP '1970-01-01 00:00:00.9'"},
{"TIMESTAMP '1970-01-01 00:00:00.56'", "timestamp(2)", "TIMESTAMP '1970-01-01 00:00:00.56'"},
{"TIMESTAMP '1970-01-01 00:00:00.123'", "timestamp(3)", "TIMESTAMP '1970-01-01 00:00:00.123'"},
{"TIMESTAMP '1970-01-01 00:00:00.4896'", "timestamp(4)", "TIMESTAMP '1970-01-01 00:00:00.4896'"},
{"TIMESTAMP '1970-01-01 00:00:00.89356'", "timestamp(5)", "TIMESTAMP '1970-01-01 00:00:00.89356'"},
{"TIMESTAMP '1970-01-01 00:00:00.123000'", "timestamp(6)", "TIMESTAMP '1970-01-01 00:00:00.123000'"},
{"TIMESTAMP '1970-01-01 00:00:00.999'", "timestamp(3)", "TIMESTAMP '1970-01-01 00:00:00.999'"},
{"TIMESTAMP '1970-01-01 00:00:00.123456'", "timestamp(6)", "TIMESTAMP '1970-01-01 00:00:00.123456'"},
{"TIMESTAMP '2020-09-27 12:34:56.1'", "timestamp(1)", "TIMESTAMP '2020-09-27 12:34:56.1'"},
{"TIMESTAMP '2020-09-27 12:34:56.9'", "timestamp(1)", "TIMESTAMP '2020-09-27 12:34:56.9'"},
{"TIMESTAMP '2020-09-27 12:34:56.123'", "timestamp(3)", "TIMESTAMP '2020-09-27 12:34:56.123'"},
{"TIMESTAMP '2020-09-27 12:34:56.123000'", "timestamp(6)", "TIMESTAMP '2020-09-27 12:34:56.123000'"},
{"TIMESTAMP '2020-09-27 12:34:56.999'", "timestamp(3)", "TIMESTAMP '2020-09-27 12:34:56.999'"},
{"TIMESTAMP '2020-09-27 12:34:56.123456'", "timestamp(6)", "TIMESTAMP '2020-09-27 12:34:56.123456'"},
{"TIMESTAMP '1970-01-01 00:00:00.1234561'", "timestamp(6)", "TIMESTAMP '1970-01-01 00:00:00.123456'"},
{"TIMESTAMP '1970-01-01 00:00:00.123456499'", "timestamp(6)", "TIMESTAMP '1970-01-01 00:00:00.123456'"},
{"TIMESTAMP '1970-01-01 00:00:00.123456499999'", "timestamp(6)", "TIMESTAMP '1970-01-01 00:00:00.123456'"},
{"TIMESTAMP '1970-01-01 00:00:00.1234565'", "timestamp(6)", "TIMESTAMP '1970-01-01 00:00:00.123457'"},
{"TIMESTAMP '1970-01-01 00:00:00.111222333444'", "timestamp(6)", "TIMESTAMP '1970-01-01 00:00:00.111222'"},
{"TIMESTAMP '1970-01-01 00:00:00.9999995'", "timestamp(6)", "TIMESTAMP '1970-01-01 00:00:01.000000'"},
{"TIMESTAMP '1970-01-01 23:59:59.9999995'", "timestamp(6)", "TIMESTAMP '1970-01-02 00:00:00.000000'"},
{"TIMESTAMP '1969-12-31 23:59:59.9999995'", "timestamp(6)", "TIMESTAMP '1970-01-01 00:00:00.000000'"},
{"TIMESTAMP '1969-12-31 23:59:59.999999499999'", "timestamp(6)", "TIMESTAMP '1969-12-31 23:59:59.999999'"},
{"TIMESTAMP '1969-12-31 23:59:59.9999994'", "timestamp(6)", "TIMESTAMP '1969-12-31 23:59:59.999999'"}};
}

@Override
protected void verifyAddNotNullColumnToNonEmptyTableFailurePermissible(Throwable e)
{
Expand Down
2 changes: 0 additions & 2 deletions .../trino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlTypeMapping.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1421,9 +1421,7 @@ public void testTimestampCoercion()
.addRoundTrip("TIMESTAMP '1969-12-31 23:59:59.999999499999'", "TIMESTAMP '1969-12-31 23:59:59.999999'")
.addRoundTrip("TIMESTAMP '1969-12-31 23:59:59.9999994'", "TIMESTAMP '1969-12-31 23:59:59.999999'")

// CTAS with Trino, where the coercion is done by the connector
.execute(getQueryRunner(), trinoCreateAsSelect("test_timestamp_coercion"))
// INSERT with Trino, where the coercion is done by the engine
.execute(getQueryRunner(), trinoCreateAndInsert("test_timestamp_coercion"));
}

Expand Down

0 comments on commit d904a6a

Please sign in to comment.