diff --git a/core/cml-main/src/main/java/io/cml/connector/bigquery/BigQueryConnector.java b/core/cml-main/src/main/java/io/cml/connector/bigquery/BigQueryConnector.java index 4bd3ceb03..389cf4029 100644 --- a/core/cml-main/src/main/java/io/cml/connector/bigquery/BigQueryConnector.java +++ b/core/cml-main/src/main/java/io/cml/connector/bigquery/BigQueryConnector.java @@ -20,6 +20,7 @@ import com.google.cloud.bigquery.BigQueryException; import com.google.cloud.bigquery.Dataset; import com.google.cloud.bigquery.DatasetInfo; import com.google.cloud.bigquery.QueryJobConfiguration; +import com.google.cloud.bigquery.Routine; import com.google.cloud.bigquery.Table; import com.google.cloud.bigquery.TableResult; import com.google.common.collect.Streams; @@ -100,6 +101,20 @@ public class BigQueryConnector }).collect(toImmutableList()); } + @Override + public List listFunctionNames(String schemaName) + { + Optional dataset = getDataset(schemaName); + if (dataset.isEmpty()) { + throw new CmlException(NOT_FOUND, format("Dataset %s is not found", schemaName)); + } + Page routines = bigQuery.listRoutines(dataset.get().getDatasetId(), BigQuery.RoutineListOption.pageSize(100)); + if (routines == null) { + throw new CmlException(NOT_FOUND, format("Dataset %s doesn't contain any routines.", dataset.get().getDatasetId())); + } + return Streams.stream(routines.iterateAll()).map(routine -> routine.getRoutineId().getRoutine()).collect(toImmutableList()); + } + @Override public boolean directDDL(String sql) { diff --git a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/PgCatalogTableManager.java b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/PgCatalogManager.java similarity index 62% rename from core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/PgCatalogTableManager.java rename to core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/PgCatalogManager.java index 42cfc040b..d0fa49e60 100644 --- a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/PgCatalogTableManager.java +++ b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/PgCatalogManager.java @@ -16,7 +16,9 @@ package io.cml.pgcatalog; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.inject.Inject; -import io.cml.pgcatalog.builder.PgCatalogBuilder; +import io.cml.pgcatalog.builder.PgCatalogTableBuilder; +import io.cml.pgcatalog.builder.PgFunctionBuilder; +import io.cml.pgcatalog.function.PgFunction; import io.cml.pgcatalog.table.CharacterSets; import io.cml.pgcatalog.table.KeyColumnUsage; import io.cml.pgcatalog.table.PgAmTable; @@ -47,22 +49,30 @@ import java.util.Map; import static com.google.common.collect.ImmutableList.toImmutableList; import static io.cml.pgcatalog.PgCatalogUtils.CML_TEMP_NAME; import static io.cml.pgcatalog.PgCatalogUtils.PG_CATALOG_NAME; +import static io.cml.pgcatalog.function.PgFunctions.CURRENT_DATABASE; +import static io.cml.pgcatalog.function.PgFunctions.PG_RELATION_SIZE__INT_VARCHAR___BIGINT; +import static io.cml.pgcatalog.function.PgFunctions.PG_RELATION_SIZE__INT___BIGINT; import static java.util.Objects.requireNonNull; -public class PgCatalogTableManager +public class PgCatalogManager { private final Map tables; + private final List functions; + private final Connector connector; - private final PgCatalogBuilder pgCatalogBuilder; + private final PgCatalogTableBuilder pgCatalogTableBuilder; + private final PgFunctionBuilder pgFunctionBuilder; private final List highPriorityTableName = ImmutableList.of(PgTypeTable.NAME); @Inject - public PgCatalogTableManager(Connector connector, PgCatalogBuilder pgCatalogBuilder) + public PgCatalogManager(Connector connector, PgCatalogTableBuilder pgCatalogTableBuilder, PgFunctionBuilder pgFunctionBuilder) { this.tables = initTables(); + this.functions = initFunctions(); this.connector = requireNonNull(connector, "connector is null"); - this.pgCatalogBuilder = requireNonNull(pgCatalogBuilder, "pgCatalogBuilder is null"); + this.pgCatalogTableBuilder = requireNonNull(pgCatalogTableBuilder, "pgCatalogBuilder is null"); + this.pgFunctionBuilder = requireNonNull(pgFunctionBuilder, "pgFunctionBuilder is null"); } private Map initTables() @@ -91,30 +101,47 @@ public class PgCatalogTableManager .build(); } - public Connector getConnector() + private List initFunctions() { - return this.connector; + return ImmutableList.builder() + .add(CURRENT_DATABASE) + .add(PG_RELATION_SIZE__INT___BIGINT) + .add(PG_RELATION_SIZE__INT_VARCHAR___BIGINT) + .build(); } public void initPgCatalog() { createCatalogIfNotExist(PG_CATALOG_NAME); if (!isPgCatalogValid()) { - createCatalogIfNotExist(CML_TEMP_NAME); + initPgTables(); + initPgFunctions(); + } + } - // Some table has dependency with the high priority table. - // Create them first. - for (String tableName : highPriorityTableName) { - createPgCatalogTable(tables.get(tableName)); - } + public void initPgTables() + { + createCatalogIfNotExist(CML_TEMP_NAME); - List lowPriorityTable = tables.values().stream() - .filter(pgCatalogTable -> !highPriorityTableName.contains(pgCatalogTable.getName())) - .collect(toImmutableList()); + // Some table has dependency with the high priority table. + // Create them first. + for (String tableName : highPriorityTableName) { + createPgCatalogTable(tables.get(tableName)); + } - for (PgCatalogTable pgCatalogTable : lowPriorityTable) { - createPgCatalogTable(pgCatalogTable); - } + List lowPriorityTable = tables.values().stream() + .filter(pgCatalogTable -> !highPriorityTableName.contains(pgCatalogTable.getName())) + .collect(toImmutableList()); + + for (PgCatalogTable pgCatalogTable : lowPriorityTable) { + createPgCatalogTable(pgCatalogTable); + } + } + + public void initPgFunctions() + { + for (PgFunction pgFunction : functions) { + pgFunctionBuilder.createPgFunction(pgFunction); } } @@ -128,11 +155,25 @@ public class PgCatalogTableManager private boolean isPgCatalogValid() { List remoteTables = connector.listTables(PG_CATALOG_NAME); - return remoteTables.size() == tables.values().size(); + if (remoteTables.size() != tables.values().size()) { + return false; + } + + List remoteFunctions = connector.listFunctionNames(PG_CATALOG_NAME); + if (functions.size() != remoteFunctions.size()) { + return false; + } + + return true; } private void createPgCatalogTable(PgCatalogTable pgCatalogTable) { - pgCatalogBuilder.createPgTable(pgCatalogTable); + pgCatalogTableBuilder.createPgTable(pgCatalogTable); + } + + public List getFunctions() + { + return functions; } } diff --git a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQueryPgCatalogBuilder.java b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQueryPgCatalogTableBuilder.java similarity index 79% rename from core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQueryPgCatalogBuilder.java rename to core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQueryPgCatalogTableBuilder.java index 30de7418a..1116a9137 100644 --- a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQueryPgCatalogBuilder.java +++ b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQueryPgCatalogTableBuilder.java @@ -19,7 +19,6 @@ import io.cml.pgcatalog.table.PgCatalogTable; import io.cml.spi.connector.Connector; import io.cml.spi.metadata.ColumnMetadata; import io.cml.spi.type.PGType; -import io.cml.type.PGArray; import javax.inject.Inject; @@ -29,38 +28,21 @@ import java.util.stream.IntStream; import static io.cml.pgcatalog.PgCatalogUtils.CML_TEMP_NAME; import static io.cml.pgcatalog.PgCatalogUtils.PG_CATALOG_NAME; -import static io.cml.pgcatalog.builder.BigQuerySqls.createOrReplaceAllColumn; -import static io.cml.pgcatalog.builder.BigQuerySqls.createOrReplaceAllTable; -import static io.cml.pgcatalog.builder.PgCatalogBuilderUtils.generatePgTypeRecords; -import static io.cml.type.BigIntType.BIGINT; -import static io.cml.type.BooleanType.BOOLEAN; -import static io.cml.type.BpCharType.BPCHAR; -import static io.cml.type.ByteaType.BYTEA; +import static io.cml.pgcatalog.builder.BigQueryUtils.createOrReplaceAllColumn; +import static io.cml.pgcatalog.builder.BigQueryUtils.createOrReplaceAllTable; +import static io.cml.pgcatalog.builder.BigQueryUtils.getOidToBqType; +import static io.cml.pgcatalog.builder.PgCatalogTableBuilderUtils.generatePgTypeRecords; import static io.cml.type.CharType.CHAR; -import static io.cml.type.DateType.DATE; -import static io.cml.type.DoubleType.DOUBLE; -import static io.cml.type.InetType.INET; -import static io.cml.type.IntegerType.INTEGER; -import static io.cml.type.JsonType.JSON; -import static io.cml.type.NumericType.NUMERIC; -import static io.cml.type.OidType.OID_INSTANCE; -import static io.cml.type.RealType.REAL; import static io.cml.type.RegprocType.REGPROC; -import static io.cml.type.SmallIntType.SMALLINT; -import static io.cml.type.TimestampType.TIMESTAMP; -import static io.cml.type.TimestampWithTimeZoneType.TIMESTAMP_WITH_TIMEZONE; -import static io.cml.type.UuidType.UUID; -import static io.cml.type.VarcharType.NameType.NAME; -import static io.cml.type.VarcharType.TextType.TEXT; import static io.cml.type.VarcharType.VARCHAR; import static java.lang.String.format; import static java.util.stream.Collectors.joining; -public final class BigQueryPgCatalogBuilder - extends PgCatalogBuilder +public final class BigQueryPgCatalogTableBuilder + extends PgCatalogTableBuilder { @Inject - public BigQueryPgCatalogBuilder(Connector connector) + public BigQueryPgCatalogTableBuilder(Connector connector) { super(connector); } @@ -83,39 +65,7 @@ public final class BigQueryPgCatalogBuilder @Override protected Map initOidToTypeMap() { - ImmutableMap.Builder builder = ImmutableMap.builder() - .put(BOOLEAN.oid(), "BOOL") - .put(SMALLINT.oid(), "SMALLINT") - .put(INTEGER.oid(), "INTEGER") - .put(BIGINT.oid(), "BIGINT") - .put(REAL.oid(), "FLOAT64") // BigQuery only has FLOAT64 for floating point type - .put(DOUBLE.oid(), "FLOAT64") - .put(NUMERIC.oid(), "NUMERIC") - .put(VARCHAR.oid(), "STRING") - .put(CHAR.oid(), "STRING") - .put(JSON.oid(), "JSON") - .put(TIMESTAMP.oid(), "TIMESTAMP") - .put(TIMESTAMP_WITH_TIMEZONE.oid(), "TIMESTAMP") - .put(TEXT.oid(), "STRING") - .put(NAME.oid(), "STRING") - .put(OID_INSTANCE.oid(), "INTEGER") - .put(DATE.oid(), "DATE") - .put(BYTEA.oid(), "BYTES") - .put(BPCHAR.oid(), "STRING") - .put(INET.oid(), "INET") - .put(UUID.oid(), "STRING") - .put(REGPROC.oid(), "STRING"); - // TODO: support record type, hstore - // .put(EMPTY_RECORD.oid(), "STRUCT") - // .put(HSTORE.oid(), "STRUCT") - - Map simpleTypeMap = builder.build(); - - for (PGArray pgArray : PGArray.allArray()) { - String innerType = simpleTypeMap.get(pgArray.getInnerType().oid()); - builder.put(pgArray.oid(), format("ARRAY<%s>", innerType)); - } - return builder.build(); + return getOidToBqType(); } @Override diff --git a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQueryPgFunctionBuilder.java b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQueryPgFunctionBuilder.java new file mode 100644 index 000000000..a26ce0136 --- /dev/null +++ b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQueryPgFunctionBuilder.java @@ -0,0 +1,87 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.cml.pgcatalog.builder; + +import io.cml.pgcatalog.function.PgFunction; +import io.cml.spi.CmlException; +import io.cml.spi.connector.Connector; + +import javax.inject.Inject; + +import java.util.Locale; + +import static io.cml.pgcatalog.PgCatalogUtils.PG_CATALOG_NAME; +import static io.cml.pgcatalog.builder.BigQueryUtils.getOidToBqType; +import static io.cml.spi.metadata.StandardErrorCode.GENERIC_INTERNAL_ERROR; +import static io.cml.type.VarcharType.VARCHAR; +import static java.lang.String.format; + +public class BigQueryPgFunctionBuilder + extends PgFunctionBuilder +{ + @Inject + public BigQueryPgFunctionBuilder(Connector connector) + { + super(connector); + } + + @Override + protected String generateCreateFunction(PgFunction pgFunction) + { + switch (pgFunction.getLanguage()) { + case SQL: + return generateCreateSqlFunction(pgFunction); + case JS: + return generateCreateJsFunction(pgFunction); + } + throw new CmlException(GENERIC_INTERNAL_ERROR, "Unsupported language: " + pgFunction.getLanguage()); + } + + private String generateCreateSqlFunction(PgFunction pgFunction) + { + StringBuilder parameterBuilder = new StringBuilder(); + if (pgFunction.getArguments().isPresent()) { + for (PgFunction.Argument argument : pgFunction.getArguments().get()) { + parameterBuilder + .append(argument.getName()).append(" ") + .append(getOidToBqType().get(argument.getType().oid())).append(","); + } + parameterBuilder.setLength(parameterBuilder.length() - 1); + } + + return format("CREATE OR REPLACE FUNCTION %s.%s(%s) AS ((%s))", PG_CATALOG_NAME, pgFunction.getRemoteName(), parameterBuilder, pgFunction.getDefinition()); + } + + private String generateCreateJsFunction(PgFunction pgFunction) + { + StringBuilder parameterBuilder = new StringBuilder(); + if (pgFunction.getArguments().isPresent()) { + for (PgFunction.Argument argument : pgFunction.getArguments().get()) { + parameterBuilder + .append(argument.getName()).append(" ") + .append(getOidToBqType().get(argument.getType().oid())).append(","); + } + parameterBuilder.setLength(parameterBuilder.length() - 1); + } + + return format("CREATE OR REPLACE FUNCTION %s.%s(%s) RETURNS %s LANGUAGE %s AS r\"\"\"%s\"\"\"", + PG_CATALOG_NAME, + pgFunction.getRemoteName(), + parameterBuilder, + getOidToBqType().get(pgFunction.getReturnType().orElse(VARCHAR).oid()), + pgFunction.getLanguage().name().toLowerCase(Locale.ROOT), + pgFunction.getDefinition()); + } +} diff --git a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQuerySqls.java b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQuerySqls.java deleted file mode 100644 index a54eea7dc..000000000 --- a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQuerySqls.java +++ /dev/null @@ -1,56 +0,0 @@ -/* - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package io.cml.pgcatalog.builder; - -import io.cml.spi.connector.Connector; - -import java.util.List; - -import static io.cml.pgcatalog.PgCatalogUtils.CML_TEMP_NAME; -import static java.lang.String.format; - -public final class BigQuerySqls -{ - private BigQuerySqls() {} - - public static String createOrReplaceAllTable(Connector connector) - { - List schemas = connector.listSchemas(); - StringBuilder builder = new StringBuilder(); - builder.append(format("CREATE OR REPLACE VIEW `%s.all_tables` AS ", CML_TEMP_NAME)); - for (String schema : schemas) { - builder.append(format("SELECT * FROM `%s`.INFORMATION_SCHEMA.TABLES UNION ALL ", schema)); - } - builder.setLength(builder.length() - "UNION ALL ".length()); - builder.append(";"); - return builder.toString(); - } - - public static String createOrReplaceAllColumn(Connector connector) - { - // TODO: we should check if pg_type has created or not. - List schemas = connector.listSchemas(); - StringBuilder builder = new StringBuilder(); - builder.append(format("CREATE OR REPLACE VIEW `%s.all_columns` AS ", CML_TEMP_NAME)); - for (String schema : schemas) { - builder.append(format("SELECT col.column_name, col.ordinal_position, ptype.oid as typoid, ptype.typlen " + - "FROM `%s`.INFORMATION_SCHEMA.COLUMNS col, `pg_catalog.pg_type` ptype " + - "WHERE col.data_type = ptype.remotetype UNION ALL ", schema)); - } - builder.setLength(builder.length() - "UNION ALL ".length()); - builder.append(";"); - return builder.toString(); - } -} diff --git a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQueryUtils.java b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQueryUtils.java new file mode 100644 index 000000000..48e5a39dc --- /dev/null +++ b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/BigQueryUtils.java @@ -0,0 +1,123 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.cml.pgcatalog.builder; + +import com.google.common.collect.ImmutableMap; +import io.cml.spi.connector.Connector; +import io.cml.type.PGArray; + +import java.util.List; +import java.util.Map; + +import static io.cml.pgcatalog.PgCatalogUtils.CML_TEMP_NAME; +import static io.cml.type.BigIntType.BIGINT; +import static io.cml.type.BooleanType.BOOLEAN; +import static io.cml.type.BpCharType.BPCHAR; +import static io.cml.type.ByteaType.BYTEA; +import static io.cml.type.CharType.CHAR; +import static io.cml.type.DateType.DATE; +import static io.cml.type.DoubleType.DOUBLE; +import static io.cml.type.InetType.INET; +import static io.cml.type.IntegerType.INTEGER; +import static io.cml.type.JsonType.JSON; +import static io.cml.type.NumericType.NUMERIC; +import static io.cml.type.OidType.OID_INSTANCE; +import static io.cml.type.RealType.REAL; +import static io.cml.type.RegprocType.REGPROC; +import static io.cml.type.SmallIntType.SMALLINT; +import static io.cml.type.TimestampType.TIMESTAMP; +import static io.cml.type.TimestampWithTimeZoneType.TIMESTAMP_WITH_TIMEZONE; +import static io.cml.type.UuidType.UUID; +import static io.cml.type.VarcharType.NameType.NAME; +import static io.cml.type.VarcharType.TextType.TEXT; +import static io.cml.type.VarcharType.VARCHAR; +import static java.lang.String.format; + +public final class BigQueryUtils +{ + private static final Map oidToBqType; + + static { + ImmutableMap.Builder builder = ImmutableMap.builder() + .put(BOOLEAN.oid(), "BOOL") + .put(SMALLINT.oid(), "SMALLINT") + .put(INTEGER.oid(), "INTEGER") + .put(BIGINT.oid(), "BIGINT") + .put(REAL.oid(), "FLOAT64") // BigQuery only has FLOAT64 for floating point type + .put(DOUBLE.oid(), "FLOAT64") + .put(NUMERIC.oid(), "NUMERIC") + .put(VARCHAR.oid(), "STRING") + .put(CHAR.oid(), "STRING") + .put(JSON.oid(), "JSON") + .put(TIMESTAMP.oid(), "TIMESTAMP") + .put(TIMESTAMP_WITH_TIMEZONE.oid(), "TIMESTAMP") + .put(TEXT.oid(), "STRING") + .put(NAME.oid(), "STRING") + .put(OID_INSTANCE.oid(), "INTEGER") + .put(DATE.oid(), "DATE") + .put(BYTEA.oid(), "BYTES") + .put(BPCHAR.oid(), "STRING") + .put(INET.oid(), "INET") + .put(UUID.oid(), "STRING") + .put(REGPROC.oid(), "STRING"); + // TODO: support record type, hstore + // .put(EMPTY_RECORD.oid(), "STRUCT") + // .put(HSTORE.oid(), "STRUCT") + + Map simpleTypeMap = builder.build(); + + for (PGArray pgArray : PGArray.allArray()) { + String innerType = simpleTypeMap.get(pgArray.getInnerType().oid()); + builder.put(pgArray.oid(), format("ARRAY<%s>", innerType)); + } + oidToBqType = builder.build(); + } + + private BigQueryUtils() {} + + public static String createOrReplaceAllTable(Connector connector) + { + List schemas = connector.listSchemas(); + StringBuilder builder = new StringBuilder(); + builder.append(format("CREATE OR REPLACE VIEW `%s.all_tables` AS ", CML_TEMP_NAME)); + for (String schema : schemas) { + builder.append(format("SELECT * FROM `%s`.INFORMATION_SCHEMA.TABLES UNION ALL ", schema)); + } + builder.setLength(builder.length() - "UNION ALL ".length()); + builder.append(";"); + return builder.toString(); + } + + public static String createOrReplaceAllColumn(Connector connector) + { + // TODO: we should check if pg_type has created or not. + List schemas = connector.listSchemas(); + StringBuilder builder = new StringBuilder(); + builder.append(format("CREATE OR REPLACE VIEW `%s.all_columns` AS ", CML_TEMP_NAME)); + for (String schema : schemas) { + builder.append(format("SELECT col.column_name, col.ordinal_position, ptype.oid as typoid, ptype.typlen " + + "FROM `%s`.INFORMATION_SCHEMA.COLUMNS col, `pg_catalog.pg_type` ptype " + + "WHERE col.data_type = ptype.remotetype UNION ALL ", schema)); + } + builder.setLength(builder.length() - "UNION ALL ".length()); + builder.append(";"); + return builder.toString(); + } + + public static Map getOidToBqType() + { + return oidToBqType; + } +} diff --git a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgCatalogBuilder.java b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgCatalogTableBuilder.java similarity index 97% rename from core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgCatalogBuilder.java rename to core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgCatalogTableBuilder.java index a88b80951..d35d46bf2 100644 --- a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgCatalogBuilder.java +++ b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgCatalogTableBuilder.java @@ -46,15 +46,15 @@ import static io.cml.spi.metadata.StandardErrorCode.GENERIC_INTERNAL_ERROR; import static java.lang.String.format; import static java.util.Objects.requireNonNull; -public abstract class PgCatalogBuilder +public abstract class PgCatalogTableBuilder { - private static final Logger LOG = Logger.get(PgCatalogBuilder.class); + private static final Logger LOG = Logger.get(PgCatalogTableBuilder.class); private final Connector connector; private final Map replaceMap; private final Map oidToTypeMap; private final StrSubstitutor strSubstitutor; - public PgCatalogBuilder(Connector connector) + public PgCatalogTableBuilder(Connector connector) { this.connector = requireNonNull(connector, "connector is null"); this.replaceMap = initReplaceMap(); diff --git a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgCatalogBuilderUtils.java b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgCatalogTableBuilderUtils.java similarity index 97% rename from core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgCatalogBuilderUtils.java rename to core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgCatalogTableBuilderUtils.java index 75320015f..3f8387211 100644 --- a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgCatalogBuilderUtils.java +++ b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgCatalogTableBuilderUtils.java @@ -26,9 +26,9 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static io.cml.pgcatalog.table.PgCatalogTableUtils.DEFAULT_AUTH; import static io.cml.pgcatalog.table.PgCatalogTableUtils.PG_CATALOG; -public final class PgCatalogBuilderUtils +public final class PgCatalogTableBuilderUtils { - private PgCatalogBuilderUtils() {} + private PgCatalogTableBuilderUtils() {} public static List generatePgTypeRecords(PgCatalogTable pgCatalogTable, Map oidToTypeMap) { diff --git a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgFunctionBuilder.java b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgFunctionBuilder.java new file mode 100644 index 000000000..1372a1b30 --- /dev/null +++ b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/builder/PgFunctionBuilder.java @@ -0,0 +1,39 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.cml.pgcatalog.builder; + +import io.airlift.log.Logger; +import io.cml.pgcatalog.function.PgFunction; +import io.cml.spi.connector.Connector; + +public abstract class PgFunctionBuilder +{ + private static final Logger LOG = Logger.get(PgFunctionBuilder.class); + private final Connector connector; + + public PgFunctionBuilder(Connector connector) + { + this.connector = connector; + } + + public void createPgFunction(PgFunction pgFunction) + { + String sql = generateCreateFunction(pgFunction); + connector.directDDL(sql); + LOG.info("pg_catalog.%s has created or updated", pgFunction.getName()); + } + + protected abstract String generateCreateFunction(PgFunction pgFunction); +} diff --git a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/function/PgFunction.java b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/function/PgFunction.java new file mode 100644 index 000000000..2076ccad9 --- /dev/null +++ b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/function/PgFunction.java @@ -0,0 +1,190 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.cml.pgcatalog.function; + +import com.google.common.base.Joiner; +import io.cml.spi.type.PGType; + +import java.util.List; +import java.util.Optional; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +public class PgFunction +{ + public enum Language + { + SQL, + JS + } + + public static PgFunction.Builder builder() + { + return new Builder(); + } + + private final String name; + private final Language language; + + private final List arguments; + + private final PGType returnType; + + private final String definition; + + public PgFunction(String name, Language language, List arguments, PGType returnType, String definition) + { + this.name = name; + this.language = language; + this.arguments = arguments; + this.returnType = returnType; + this.definition = definition; + } + + public String getName() + { + return name; + } + + /** + * Some data warehouse(BigQuery) doesn't support function overloading. We should name the function with its argument's type and return type. + * For example: + * pg_relation_size(relOid int)bigint -> pg_relation_size__int___bigint(relOid int) + * pg_relation_size(relOid int, text varchar)bigint -> pg_relation_size__int_varchar___bigint(relOid int, text varchar) + * + * @return the name used by the remote database. + */ + public String getRemoteName() + { + String argString = getArguments().isPresent() ? "__" + Joiner.on("_").join(arguments.stream().map(Argument::getType).map(PGType::typName).collect(toImmutableList())) : ""; + String returnString = getReturnType().isPresent() ? "___" + returnType.typName() : ""; + return getName() + argString + returnString; + } + + public Language getLanguage() + { + return language; + } + + public Optional> getArguments() + { + return Optional.ofNullable(arguments); + } + + public Optional getReturnType() + { + return Optional.ofNullable(returnType); + } + + public String getDefinition() + { + return definition; + } + + @Override + public String toString() + { + StringBuilder parameterBuilder = new StringBuilder(); + if (getArguments().isPresent()) { + for (Argument argument : getArguments().get()) { + parameterBuilder + .append(argument.getName()).append(" ") + .append(argument.getType()).append(","); + } + parameterBuilder.setLength(parameterBuilder.length() - 1); + } + + return format("%s(%s)%s", getName(), parameterBuilder, getReturnType().isPresent() ? returnType.typName() : "void"); + } + + public static class Argument + { + public static Argument argument(String name, PGType type) + { + return new Argument(name, type); + } + + private final String name; + private final PGType type; + + public Argument(String name, PGType type) + { + this.name = name; + this.type = type; + } + + public String getName() + { + return name; + } + + public PGType getType() + { + return type; + } + } + + public static class Builder + { + private String name; + private Language language; + + private String definition; + + private List arguments; + + private PGType returnType; + + public Builder setName(String name) + { + this.name = name; + return this; + } + + public Builder setLanguage(Language language) + { + this.language = language; + return this; + } + + public Builder setDefinition(String definition) + { + this.definition = definition; + return this; + } + + public Builder setArguments(List arguments) + { + this.arguments = arguments; + return this; + } + + public Builder setReturnType(PGType returnType) + { + this.returnType = returnType; + return this; + } + + public PgFunction build() + { + requireNonNull(name, "name is null"); + requireNonNull(language, "language is null"); + requireNonNull(definition, "definition is null"); + return new PgFunction(name, language, arguments, returnType, definition); + } + } +} diff --git a/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/function/PgFunctions.java b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/function/PgFunctions.java new file mode 100644 index 000000000..7cacec1ff --- /dev/null +++ b/core/cml-postgres-wire-protocol/src/main/java/io/cml/pgcatalog/function/PgFunctions.java @@ -0,0 +1,53 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.cml.pgcatalog.function; + +import com.google.common.collect.ImmutableList; + +import static io.cml.pgcatalog.function.PgFunction.Argument.argument; +import static io.cml.pgcatalog.function.PgFunction.Language.SQL; +import static io.cml.pgcatalog.function.PgFunction.builder; +import static io.cml.type.BigIntType.BIGINT; +import static io.cml.type.IntegerType.INTEGER; +import static io.cml.type.VarcharType.VARCHAR; + +public final class PgFunctions +{ + private static final String EMPTY_STATEMENT = "SELECT null LIMIT 0"; + + private PgFunctions() {} + + public static final PgFunction CURRENT_DATABASE = builder() + .setName("current_database") + .setLanguage(SQL) + .setDefinition("SELECT DISTINCT catalog_name FROM INFORMATION_SCHEMA.SCHEMATA") + .build(); + + public static final PgFunction PG_RELATION_SIZE__INT___BIGINT = builder() + .setName("pg_relation_size") + .setLanguage(SQL) + .setDefinition(EMPTY_STATEMENT) + .setArguments(ImmutableList.of(argument("relOid", INTEGER))) + .setReturnType(BIGINT) + .build(); + + public static final PgFunction PG_RELATION_SIZE__INT_VARCHAR___BIGINT = builder() + .setName("pg_relation_size") + .setLanguage(SQL) + .setDefinition(EMPTY_STATEMENT) + .setArguments(ImmutableList.of(argument("relOid", INTEGER), argument("text", VARCHAR))) + .setReturnType(BIGINT) + .build(); +} diff --git a/core/cml-server/src/main/java/io/cml/server/CmlServer.java b/core/cml-server/src/main/java/io/cml/server/CmlServer.java index bcd0b37fe..488774fa7 100644 --- a/core/cml-server/src/main/java/io/cml/server/CmlServer.java +++ b/core/cml-server/src/main/java/io/cml/server/CmlServer.java @@ -17,7 +17,7 @@ package io.cml.server; import com.google.common.collect.ImmutableList; import com.google.inject.Injector; import com.google.inject.Module; -import io.cml.pgcatalog.PgCatalogTableManager; +import io.cml.pgcatalog.PgCatalogManager; import io.cml.server.module.BigQueryConnectorModule; import io.cml.server.module.PostgresWireProtocolModule; import io.cml.wireprotocol.ssl.EmptyTlsDataProvider; @@ -33,8 +33,8 @@ public class CmlServer @Override protected void configure(Injector injector) { - PgCatalogTableManager pgCatalogTableManager = injector.getInstance(PgCatalogTableManager.class); - pgCatalogTableManager.initPgCatalog(); + PgCatalogManager pgCatalogManager = injector.getInstance(PgCatalogManager.class); + pgCatalogManager.initPgCatalog(); } @Override diff --git a/core/cml-server/src/main/java/io/cml/server/module/BigQueryConnectorModule.java b/core/cml-server/src/main/java/io/cml/server/module/BigQueryConnectorModule.java index 187d7e0cb..f5bf45ff6 100644 --- a/core/cml-server/src/main/java/io/cml/server/module/BigQueryConnectorModule.java +++ b/core/cml-server/src/main/java/io/cml/server/module/BigQueryConnectorModule.java @@ -28,8 +28,10 @@ import io.airlift.configuration.AbstractConfigurationAwareModule; import io.cml.connector.bigquery.BigQueryConfig; import io.cml.connector.bigquery.BigQueryConnector; import io.cml.connector.bigquery.BigQueryCredentialsSupplier; -import io.cml.pgcatalog.builder.BigQueryPgCatalogBuilder; -import io.cml.pgcatalog.builder.PgCatalogBuilder; +import io.cml.pgcatalog.builder.BigQueryPgCatalogTableBuilder; +import io.cml.pgcatalog.builder.BigQueryPgFunctionBuilder; +import io.cml.pgcatalog.builder.PgCatalogTableBuilder; +import io.cml.pgcatalog.builder.PgFunctionBuilder; import io.cml.pgcatalog.regtype.BigQueryPgMetadata; import io.cml.pgcatalog.regtype.PgMetadata; import io.cml.spi.connector.Connector; @@ -45,7 +47,8 @@ public class BigQueryConnectorModule protected void setup(Binder binder) { binder.bind(Connector.class).to(BigQueryConnector.class); - binder.bind(PgCatalogBuilder.class).to(BigQueryPgCatalogBuilder.class).in(Scopes.SINGLETON); + binder.bind(PgCatalogTableBuilder.class).to(BigQueryPgCatalogTableBuilder.class).in(Scopes.SINGLETON); + binder.bind(PgFunctionBuilder.class).to(BigQueryPgFunctionBuilder.class).in(Scopes.SINGLETON); binder.bind(PgMetadata.class).to(BigQueryPgMetadata.class).in(Scopes.SINGLETON); configBinder(binder).bindConfig(BigQueryConfig.class); } diff --git a/core/cml-server/src/main/java/io/cml/server/module/PostgresWireProtocolModule.java b/core/cml-server/src/main/java/io/cml/server/module/PostgresWireProtocolModule.java index 3148702f3..c2b20fd45 100644 --- a/core/cml-server/src/main/java/io/cml/server/module/PostgresWireProtocolModule.java +++ b/core/cml-server/src/main/java/io/cml/server/module/PostgresWireProtocolModule.java @@ -19,7 +19,7 @@ import com.google.inject.Scopes; import io.airlift.configuration.AbstractConfigurationAwareModule; import io.cml.PostgresNettyProvider; import io.cml.PostgresWireProtocolConfig; -import io.cml.pgcatalog.PgCatalogTableManager; +import io.cml.pgcatalog.PgCatalogManager; import io.cml.pgcatalog.regtype.RegObjectFactory; import io.cml.wireprotocol.PostgresNetty; import io.cml.wireprotocol.ssl.SslContextProvider; @@ -45,7 +45,7 @@ public class PostgresWireProtocolModule binder.bind(SqlParser.class).in(Scopes.SINGLETON); binder.bind(TlsDataProvider.class).toInstance(tlsDataProvider); binder.bind(SslContextProvider.class).in(Scopes.SINGLETON); - binder.bind(PgCatalogTableManager.class).in(Scopes.SINGLETON); + binder.bind(PgCatalogManager.class).in(Scopes.SINGLETON); binder.bind(RegObjectFactory.class).in((Scopes.SINGLETON)); binder.bind(PostgresNetty.class).toProvider(PostgresNettyProvider.class).in(Scopes.SINGLETON); } diff --git a/core/cml-spi/src/main/java/io/cml/spi/connector/Connector.java b/core/cml-spi/src/main/java/io/cml/spi/connector/Connector.java index 069f68a52..edd107dd9 100644 --- a/core/cml-spi/src/main/java/io/cml/spi/connector/Connector.java +++ b/core/cml-spi/src/main/java/io/cml/spi/connector/Connector.java @@ -28,6 +28,8 @@ public interface Connector List listTables(String schemaName); + List listFunctionNames(String schemaName); + boolean directDDL(String sql); Iterable directQuery(String sql);