opensource/WrenAI
3
0
Fork 0
mirror of https://github.com/Canner/WrenAI.git synced 2026-08-30 18:00:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

Remove the legacy connector, Cache server and PG Wire Protocol (#639)

Browse Source 
* remove legacy connector and pg-wire related implementations

* modified the example and docker file

* fix the tests

* remove the sqlglot ci flow

* fix tests
This commit is contained in:
Jax Liu
2024-06-28 14:19:33 +08:00
committed by GitHub
parent b1037c264d
commit ddac47fcd7
369 changed files with 469 additions and 41682 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docker/Dockerfile
-6
View File
@@ -15,12 +15,6 @@ ENV ENV_WREN_VERSION=${WREN_VERSION}
ENV WREN_JAR=wren-server-${ENV_WREN_VERSION}-executable.jar
COPY ${WREN_JAR} ./
COPY wren-sqlglot-server/ ./wren-sqlglot-server/
RUN apt-get install -y python3 python3-pip
RUN pip3 install -r ./wren-sqlglot-server/requirements.txt
COPY entrypoint.sh ./
RUN chmod +x ./entrypoint.sh
docker/entrypoint.sh
-1
View File
@@ -1,7 +1,6 @@
#!/bin/bash
export ENV_MAX_HEAP_SIZE=$2
export ENV_MIN_HEAP_SIZE=$3
python3 wren-sqlglot-server/main.py &
# Required add-opens=java.nio=ALL-UNNAMED for Apache arrow in the Snowflake
java -Xmx${ENV_MAX_HEAP_SIZE:-"512m"} -Xms${ENV_MIN_HEAP_SIZE:-"64m"} -Dconfig=etc/config.properties \
examples/README.md
-1
View File
@@ -2,5 +2,4 @@
This is some example Docker Compose project for different data source types. Please go to the respective directory for more information.
- [BigQuery Example](bigquery-example/README.md)
- [DuckDB TPCH Example](duckdb-tpch-example/README.md)
examples/bigquery-example/.env
-1
View File
@@ -1 +0,0 @@
PLATFORM=linux/amd64
examples/bigquery-example/README.md
-23
View File
@@ -1,23 +0,0 @@
# Description
This is an example Docker Compose project for running the Wren Engine.
# How to use
1. Set up the platform in `.env` file. (`linux/amd64`, `linux/arm64`)
2. Configure settings in the `etc/config.properties` file.
3. Place your MDL in `etc/mdl` file after removing the sample MDL file `etc/mdl/sample.json`.
- The `mdl` directory should contain only one json file.
4. Set up the accounts if you needs or remove the sample accounts if you don't need.
- Sample accounts are provided in the `etc/accounts` directory.
5. Run the docker-compose in this directory.
```bash
docker compose --env-file .env up
```
6. Connect using psql or another PostgreSQL driver using port 7432.
- Sample usernames and passwords are `ina` and `wah`, or `azki` and `guess`.
- The default database name should match the catalog of the MDL file.
- The default schema name should match the schema of the MDL file.
```bash
psql 'host=localhost user=ina dbname=test_catalog port=7432 options=--search_path=test_schema'
```
examples/bigquery-example/docker-compose.yaml
-13
View File
@@ -1,13 +0,0 @@
version: '3.8'
services:
engine:
image: ghcr.io/canner/wren-engine:latest
platform: ${PLATFORM}
ports:
- 8080:8080
- 7432:7432
volumes:
- ./etc:/usr/src/app/etc
environment:
- SQLGLOT_PORT=8000
examples/bigquery-example/etc/accounts
-2
View File
@@ -1,2 +0,0 @@
ina wah
azki guess
examples/bigquery-example/etc/config.properties
-27
View File
@@ -1,27 +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.
# */
#
node.environment=production
wren.directory=/usr/src/app/etc/mdl
pg-wire-protocol.auth.file=/usr/src/app/etc/accounts
wren.experimental-enable-dynamic-fields=false
wren.datasource.type=bigquery
bigquery.project-id=
bigquery.credentials-key=
bigquery.location=
bigquery.bucket-name=
duckdb.storage.access-key=
duckdb.storage.secret-key=
pg-wire-protocol.enabled=true
examples/bigquery-example/etc/mdl/sample.json
-5
View File
@@ -1,5 +0,0 @@
{
"catalog": "test_catalog",
"schema": "test_schema",
"models": []
}
examples/duckdb-tpch-example/README.md
+48 -9
View File
@@ -11,16 +11,55 @@ You can learn how to set up the Wren Engine with DuckDB data source to analyze y
4. Set up your initial SQL and session SQL in `etc/duckdb-init.sql` and `etc/duckdb-session.sql` files.
5. Place your MDL in `etc/mdl` file after removing the sample MDL file `etc/mdl/sample.json`.
- The `mdl` directory should contain only one json file.
6. Set up the accounts if you needs or remove the sample accounts if you don't need.
- Sample accounts are provided in the `etc/accounts` directory.
7. Run the docker-compose in this directory.
```bash
docker compose --env-file .env up
```
8. Connect using psql or another PostgreSQL driver using port 7432.
- Sample usernames and passwords are `ina` and `wah`, or `azki` and `guess`.
- The default database name should match the catalog of the MDL file.
- The default schema name should match the schema of the MDL file.
```bash
psql 'host=localhost user=ina dbname=wren port=7432 options=--search_path=tpch'
```
8. Call the Wren Engine API to analyze your data with the request body as below.
- URL
```
GET http://localhost:8080/v1/mdl/preview
```
- Body: The manifest is the content of the MDL file.
```json
{
"manifest": {
"catalog": "wren",
"schema": "tpch",
"models": [
{
"name": "Orders",
"tableReference": {
"catalog": "memory",
"schema": "tpch",
"table": "orders"
},
"columns": [
{
"name": "orderkey",
"expression": "o_orderkey",
"type": "integer"
},
{
"name": "custkey",
"expression": "o_custkey",
"type": "integer"
},
{
"name": "orderstatus",
"expression": "o_orderstatus",
"type": "varchar"
},
{
"name": "totalprice",
"expression": "o_totalprice",
"type": "float"
}
],
"primaryKey": "orderkey"
}
]
},
"sql": "select * from Orders"
}
```
examples/duckdb-tpch-example/docker-compose.yaml
-3
View File
@@ -6,9 +6,6 @@ services:
platform: ${PLATFORM}
ports:
- 8080:8080
- 7432:7432
volumes:
- ./etc:/usr/src/app/etc
- ./data:/usr/src/app/data
environment:
- SQLGLOT_PORT=8000
examples/duckdb-tpch-example/etc/accounts
-2
View File
@@ -1,2 +0,0 @@
ina wah
azki guess
examples/duckdb-tpch-example/etc/config.properties
-2
View File
@@ -15,9 +15,7 @@
#
node.environment=production
wren.directory=/usr/src/app/etc/mdl
pg-wire-protocol.auth.file=/usr/src/app/etc/accounts
wren.experimental-enable-dynamic-fields=false
wren.datasource.type=duckdb
duckdb.connector.init-sql-path=/usr/src/app/etc/duckdb-init.sql
duckdb.connector.session-sql-path=/usr/src/app/etc/duckdb-session.sql
pg-wire-protocol.enabled=true
pom.xml
+1 -136
View File
@@ -29,11 +29,8 @@
<modules>
<module>trino-parser</module>
<module>wren-base</module>
<module>wren-cache</module>
<module>wren-main</module>
<module>wren-server</module>
<module>wren-shaded</module>
<module>wren-testing</module>
<module>wren-tests</module>
</modules>
@@ -92,15 +89,6 @@
<dependencyManagement>
<dependencies>
<dependency>
<groupId>com.google.cloud</groupId>
<artifactId>libraries-bom</artifactId>
<version>25.2.0</version>
<type>pom</type>
<scope>import</scope>
</dependency>
<dependency>
<groupId>org.jdbi</groupId>
<artifactId>jdbi3-bom</artifactId>
@@ -117,12 +105,6 @@
<scope>import</scope>
</dependency>
<dependency>
<groupId>com.carrotsearch</groupId>
<artifactId>hppc</artifactId>
<version>0.9.1</version>
</dependency>
<dependency>
<groupId>com.fasterxml.jackson.core</groupId>
<artifactId>jackson-annotations</artifactId>
@@ -147,30 +129,6 @@
<version>2.17.0</version>
</dependency>
<dependency>
<groupId>com.google.api</groupId>
<artifactId>gax</artifactId>
<version>2.16.0</version>
</dependency>
<dependency>
<groupId>com.google.auth</groupId>
<artifactId>google-auth-library-credentials</artifactId>
<version>1.23.0</version>
</dependency>
<dependency>
<groupId>com.google.auth</groupId>
<artifactId>google-auth-library-oauth2-http</artifactId>
<version>1.23.0</version>
</dependency>
<dependency>
<groupId>com.google.auto.value</groupId>
<artifactId>auto-value-annotations</artifactId>
<version>1.10.4</version>
</dependency>
<dependency>
<groupId>com.google.code.findbugs</groupId>
<artifactId>jsr305</artifactId>
@@ -195,30 +153,12 @@
<version>1.0.2</version>
</dependency>
<dependency>
<groupId>com.google.http-client</groupId>
<artifactId>google-http-client</artifactId>
<version>1.44.1</version>
</dependency>
<dependency>
<groupId>com.google.http-client</groupId>
<artifactId>google-http-client-gson</artifactId>
<version>1.44.1</version>
</dependency>
<dependency>
<groupId>com.google.protobuf</groupId>
<artifactId>protobuf-java</artifactId>
<version>3.19.4</version>
</dependency>
<dependency>
<groupId>com.h2database</groupId>
<artifactId>h2</artifactId>
<version>2.0.202</version>
</dependency>
<dependency>
<groupId>com.hubspot.jinjava</groupId>
<artifactId>jinjava</artifactId>
@@ -243,12 +183,6 @@
</exclusions>
</dependency>
<dependency>
<groupId>commons-codec</groupId>
<artifactId>commons-codec</artifactId>
<version>1.15</version>
</dependency>
<dependency>
<groupId>io.airlift</groupId>
<artifactId>bootstrap</artifactId>
@@ -339,36 +273,6 @@
<version>1.10</version>
</dependency>
<dependency>
<groupId>io.grpc</groupId>
<artifactId>grpc-api</artifactId>
<version>1.60.1</version>
</dependency>
<dependency>
<groupId>io.grpc</groupId>
<artifactId>grpc-context</artifactId>
<version>1.60.1</version>
</dependency>
<dependency>
<groupId>io.netty</groupId>
<artifactId>netty-all</artifactId>
<version>${netty.version}</version>
</dependency>
<dependency>
<groupId>io.opencensus</groupId>
<artifactId>opencensus-api</artifactId>
<version>0.31.1</version>
</dependency>
<dependency>
<groupId>io.opencensus</groupId>
<artifactId>opencensus-contrib-http-util</artifactId>
<version>0.31.1</version>
</dependency>
<dependency>
<groupId>io.wren</groupId>
<artifactId>trino-parser</artifactId>
@@ -381,11 +285,6 @@
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>io.wren</groupId>
<artifactId>wren-cache</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>io.wren</groupId>
<artifactId>wren-main</artifactId>
@@ -398,18 +297,6 @@
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>io.wren</groupId>
<artifactId>wren-shaded</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>io.wren</groupId>
<artifactId>wren-testing</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>io.wren</groupId>
<artifactId>wren-tests</artifactId>
@@ -422,36 +309,18 @@
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>it.unimi.dsi</groupId>
<artifactId>fastutil</artifactId>
<version>8.3.0</version>
</dependency>
<dependency>
<groupId>net.bytebuddy</groupId>
<artifactId>byte-buddy</artifactId>
<version>1.14.13</version>
</dependency>
<dependency>
<groupId>net.snowflake</groupId>
<artifactId>snowflake-jdbc</artifactId>
<version>3.15.0</version>
</dependency>
<dependency>
<groupId>org.antlr</groupId>
<artifactId>antlr4-runtime</artifactId>
<version>${dep.antlr.version}</version>
</dependency>
<dependency>
<groupId>org.apache.commons</groupId>
<artifactId>commons-lang3</artifactId>
<version>3.9</version>
</dependency>
<!-- TODO: move this to Airbase -->
<dependency>
<groupId>org.apache.commons</groupId>
@@ -498,7 +367,7 @@
<dependency>
<groupId>org.postgresql</groupId>
<artifactId>postgresql</artifactId>
<version>42.3.1</version>
<version>42.3.9</version>
</dependency>
<dependency>
@@ -710,10 +579,6 @@
<excludes>
<exclude>**/*jmhTest*.java</exclude>
<exclude>**/*jmhType*.java</exclude>
<exclude>TestTPCHWithSnowflake.java</exclude>
<exclude>TestDeploySnowflakeRuntime</exclude>
<exclude>TestWireProtocolTypeWithSnowflake</exclude>
<exclude>TestWrenWithSnowflake.java</exclude>
</excludes>
</configuration>
</plugin>
wren-base/pom.xml
-26
View File
@@ -31,12 +31,6 @@
</properties>
<dependencies>
<dependency>
<groupId>com.carrotsearch</groupId>
<artifactId>hppc</artifactId>
</dependency>
<dependency>
<groupId>com.fasterxml.jackson.core</groupId>
<artifactId>jackson-annotations</artifactId>
@@ -79,10 +73,6 @@
</exclusions>
</dependency>
<dependency>
<groupId>commons-codec</groupId>
<artifactId>commons-codec</artifactId>
</dependency>
<dependency>
<groupId>io.airlift</groupId>
<artifactId>configuration</artifactId>
@@ -103,11 +93,6 @@
<artifactId>units</artifactId>
</dependency>
<dependency>
<groupId>io.netty</groupId>
<artifactId>netty-all</artifactId>
</dependency>
<dependency>
<groupId>io.wren</groupId>
<artifactId>trino-parser</artifactId>
@@ -118,11 +103,6 @@
<artifactId>jakarta.validation-api</artifactId>
</dependency>
<dependency>
<groupId>joda-time</groupId>
<artifactId>joda-time</artifactId>
</dependency>
<dependency>
<groupId>org.antlr</groupId>
<artifactId>antlr4-runtime</artifactId>
@@ -156,12 +136,6 @@
<scope>test</scope>
</dependency>
<dependency>
<groupId>it.unimi.dsi</groupId>
<artifactId>fastutil</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.assertj</groupId>
<artifactId>assertj-core</artifactId>
wren-base/src/main/java/io/trino/execution/sql/SqlFormatterUtil.java
-60
View File
@@ -1,60 +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.trino.execution.sql;
import io.trino.sql.SqlFormatter;
import io.trino.sql.parser.ParsingException;
import io.trino.sql.parser.ParsingOptions;
import io.trino.sql.parser.SqlParser;
import io.trino.sql.tree.Statement;
import io.wren.base.WrenException;
import javax.annotation.Nullable;
import static io.trino.sql.parser.ParsingOptions.DecimalLiteralTreatment.REJECT;
import static io.wren.base.metadata.StandardErrorCode.GENERIC_INTERNAL_ERROR;
import static java.lang.String.format;
public final class SqlFormatterUtil
{
private SqlFormatterUtil() {}
public static String getFormattedSql(Statement statement, SqlParser sqlParser)
{
String sql = SqlFormatter.formatSql(statement);
// verify round-trip
Statement parsed;
try {
ParsingOptions parsingOptions = new ParsingOptions(REJECT /* formatted SQL should be unambiguous */);
parsed = sqlParser.createStatement(sql, parsingOptions);
}
catch (ParsingException e) {
throw formattingFailure(e, "Formatted query does not parse", statement, sql);
}
if (!statement.equals(parsed)) {
throw formattingFailure(null, "Query does not round-trip", statement, sql);
}
return sql;
}
private static WrenException formattingFailure(@Nullable Throwable cause, String message, Statement statement, String sql)
{
WrenException exception = new WrenException(GENERIC_INTERNAL_ERROR, message, cause);
exception.addSuppressed(new RuntimeException("Statement: " + statement));
exception.addSuppressed(new RuntimeException(format("Formatted: [%s]", sql)));
return exception;
}
}
wren-base/src/main/java/io/wren/base/Column.java
+40 -15
View File
@@ -16,14 +16,21 @@ package io.wren.base;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import io.wren.base.type.AnyType;
import io.wren.base.type.PGType;
import io.wren.base.type.PGTypes;
import java.util.Locale;
import java.util.Objects;
import static com.google.common.base.MoreObjects.toStringHelper;
public final class Column
{
public static Column column(String name, String type)
{
return new Column(name, type);
}
private final String name;
private final PGType<?> type;
private final String type;
@JsonCreator
public Column(
@@ -31,13 +38,7 @@ public final class Column
@JsonProperty("type") String type)
{
this.name = name;
this.type = PGTypes.nameToPgType(type).orElse(AnyType.ANY);
}
public Column(String name, PGType<?> type)
{
this.name = name;
this.type = type;
this.type = type.toUpperCase(Locale.ROOT);
}
@JsonProperty
@@ -46,14 +47,38 @@ public final class Column
return name;
}
public PGType<?> getType()
@JsonProperty
public String getType()
{
return type;
}
@JsonProperty("type")
public String getTypeName()
@Override
public boolean equals(Object o)
{
return type.typName();
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
Column column = (Column) o;
return Objects.equals(name, column.name) &&
Objects.equals(type, column.type);
}
@Override
public int hashCode()
{
return Objects.hash(name, type);
}
@Override
public String toString()
{
return toStringHelper(this)
.add("name", name)
.add("type", type)
.toString();
}
}
wren-base/src/main/java/io/wren/base/Parameter.java
+3 -5
View File
@@ -14,22 +14,20 @@
package io.wren.base;
import io.wren.base.type.PGType;
import java.util.Objects;
public class Parameter
{
private final PGType<?> type;
private final String type;
private final Object value;
public Parameter(PGType<?> type, Object value)
public Parameter(String type, Object value)
{
this.type = type;
this.value = value;
}
public PGType<?> getType()
public String getType()
{
return type;
}
wren-base/src/main/java/io/wren/base/client/Client.java
+2 -2
View File
@@ -14,8 +14,8 @@
package io.wren.base.client;
import io.wren.base.Column;
import io.wren.base.Parameter;
import io.wren.base.metadata.ColumnMetadata;
import java.sql.Connection;
import java.sql.SQLException;
@@ -34,7 +34,7 @@ public interface Client
void executeDDL(String sql);
List<ColumnMetadata> describe(String sql, List<Parameter> parameters);
List<Column> describe(String sql, List<Parameter> parameters);
List<String> listTables();
wren-base/src/main/java/io/wren/base/client/duckdb/CacheStorageConfig.java
-19
View File
@@ -1,19 +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.wren.base.client.duckdb;
public interface CacheStorageConfig
{
String generateDuckdbParquetStatement(String path, String tableName);
}
wren-base/src/main/java/io/wren/base/client/duckdb/DuckdbClient.java
+8 -30
View File
@@ -19,14 +19,13 @@ import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import io.airlift.log.Logger;
import io.airlift.units.DataSize;
import io.wren.base.Column;
import io.wren.base.Parameter;
import io.wren.base.WrenException;
import io.wren.base.client.AutoCloseableIterator;
import io.wren.base.client.Client;
import io.wren.base.client.jdbc.JdbcRecordIterator;
import io.wren.base.metadata.ColumnMetadata;
import io.wren.base.metadata.StandardErrorCode;
import io.wren.base.type.PGType;
import org.duckdb.DuckDBConnection;
import javax.annotation.Nullable;
@@ -42,7 +41,6 @@ import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import static io.wren.base.client.duckdb.DuckdbTypes.toPGType;
import static java.lang.String.format;
public final class DuckdbClient
@@ -50,18 +48,15 @@ public final class DuckdbClient
{
private static final Logger LOG = Logger.get(DuckdbClient.class);
private final DuckDBConfig duckDBConfig;
private final CacheStorageConfig cacheStorageConfig;
private final DuckDBSettingSQL duckDBSettingSQL;
private DuckDBConnection duckDBConnection;
private HikariDataSource connectionPool;
public DuckdbClient(
DuckDBConfig duckDBConfig,
@Nullable CacheStorageConfig cacheStorageConfig,
@Nullable DuckDBSettingSQL duckDBSettingSQL)
{
this.duckDBConfig = duckDBConfig;
this.cacheStorageConfig = cacheStorageConfig;
this.duckDBSettingSQL = duckDBSettingSQL;
init();
}
@@ -99,12 +94,11 @@ public final class DuckdbClient
public synchronized void initPool()
{
connectionPool = new HikariDataSource(getHikariConfig(duckDBConfig, cacheStorageConfig, duckDBConnection, duckDBSettingSQL));
connectionPool = new HikariDataSource(getHikariConfig(duckDBConfig, duckDBConnection, duckDBSettingSQL));
}
private static HikariConfig getHikariConfig(
DuckDBConfig duckDBConfig,
CacheStorageConfig cacheStorageConfig,
DuckDBConnection duckDBConnection,
DuckDBSettingSQL duckDBSettingSQL)
{
@@ -116,12 +110,12 @@ public final class DuckdbClient
config.setMinimumIdle(duckDBConfig.getMaxConcurrentTasks());
// remain some query slots for metadata queries
config.setMaximumPoolSize(duckDBConfig.getMaxConcurrentTasks() + duckDBConfig.getMaxConcurrentMetadataQueries());
String initSql = buildConnectionInitSql(duckDBSettingSQL, cacheStorageConfig, duckDBConfig);
String initSql = buildConnectionInitSql(duckDBSettingSQL, duckDBConfig);
config.setConnectionInitSql(initSql);
return config;
}
private static String buildConnectionInitSql(DuckDBSettingSQL duckDBSettingSQL, CacheStorageConfig cacheStorageConfig, DuckDBConfig duckDBConfig)
private static String buildConnectionInitSql(DuckDBSettingSQL duckDBSettingSQL, DuckDBConfig duckDBConfig)
{
List<String> sql = new ArrayList<>();
// Both of them should be true in default, however they're some issue in v0.10.3.
@@ -135,11 +129,6 @@ public final class DuckdbClient
}
else {
sql.add("SET search_path = 'main'");
if (cacheStorageConfig instanceof DuckdbS3StyleStorageConfig) {
DuckdbS3StyleStorageConfig duckdbS3StyleStorageConfig = (DuckdbS3StyleStorageConfig) cacheStorageConfig;
sql.add(format("SET s3_endpoint='%s'", duckdbS3StyleStorageConfig.getEndpoint()));
sql.add(format("SET s3_url_style='%s'", duckdbS3StyleStorageConfig.getUrlStyle()));
}
sql.add(format("SET home_directory='%s'", duckDBConfig.getHomeDirectory()));
}
return String.join(";", sql);
@@ -173,20 +162,16 @@ public final class DuckdbClient
* So we ignore the parameters here.
*/
@Override
public List<ColumnMetadata> describe(String sql, List<Parameter> ignored)
public List<Column> describe(String sql, List<Parameter> ignored)
{
try (Connection connection = createConnection()) {
PreparedStatement preparedStatement = connection.prepareStatement(sql);
ResultSetMetaData metaData = preparedStatement.getMetaData();
int columnCount = metaData.getColumnCount();
ImmutableList.Builder<ColumnMetadata> builder = ImmutableList.builder();
ImmutableList.Builder<Column> builder = ImmutableList.builder();
for (int i = 1; i <= columnCount; i++) {
PGType<?> type = toPGType(metaData, i);
builder.add(ColumnMetadata.builder()
.setName(metaData.getColumnName(i))
.setType(type)
.build());
builder.add(new Column(metaData.getColumnName(i), metaData.getColumnTypeName(i)));
}
return builder.build();
}
@@ -269,7 +254,6 @@ public final class DuckdbClient
public static class Builder
{
private DuckDBConfig duckDBConfig;
private CacheStorageConfig cacheStorageConfig;
private DuckDBSettingSQL duckDBSettingSQL;
public Builder setDuckDBConfig(DuckDBConfig duckDBConfig)
@@ -278,12 +262,6 @@ public final class DuckdbClient
return this;
}
public Builder setCacheStorageConfig(CacheStorageConfig cacheStorageConfig)
{
this.cacheStorageConfig = cacheStorageConfig;
return this;
}
public Builder setDuckDBSettingSQL(DuckDBSettingSQL duckDBSettingSQL)
{
this.duckDBSettingSQL = duckDBSettingSQL;
@@ -292,7 +270,7 @@ public final class DuckdbClient
public DuckdbClient build()
{
return new DuckdbClient(duckDBConfig, cacheStorageConfig, duckDBSettingSQL);
return new DuckdbClient(duckDBConfig, duckDBSettingSQL);
}
public Optional<DuckdbClient> buildSafely()
wren-base/src/main/java/io/wren/base/client/duckdb/DuckdbS3StyleStorageConfig.java
-121
View File
@@ -1,121 +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.wren.base.client.duckdb;
import io.airlift.configuration.Config;
import io.airlift.configuration.ConfigDescription;
import io.airlift.configuration.ConfigSecuritySensitive;
import java.util.Optional;
import static java.lang.String.format;
public class DuckdbS3StyleStorageConfig
implements CacheStorageConfig
{
public static final String DUCKDB_STORAGE_ENDPOINT = "duckdb.storage.endpoint";
public static final String DUCKDB_STORAGE_ACCESS_KEY = "duckdb.storage.access-key";
public static final String DUCKDB_STORAGE_SECRET_KEY = "duckdb.storage.secret-key";
public static final String DUCKDB_STORAGE_REGION = "duckdb.storage.region";
public static final String DUCKDB_STORAGE_URL_STYLE = "duckdb.storage.url-style";
// https://duckdb.org/docs/guides/import/s3_import.html
private String endpoint = "storage.googleapis.com";
private Optional<String> accessKey = Optional.empty();
private Optional<String> secretKey = Optional.empty();
private Optional<String> region = Optional.empty();
private String urlStyle = "path";
@Config(DUCKDB_STORAGE_ENDPOINT)
@ConfigDescription("The storage endpoint; default is storage.googleapis.com")
public DuckdbS3StyleStorageConfig setEndpoint(String endpoint)
{
this.endpoint = endpoint;
return this;
}
public String getEndpoint()
{
return endpoint;
}
@Config(DUCKDB_STORAGE_ACCESS_KEY)
@ConfigDescription("The storage access key")
@ConfigSecuritySensitive
public DuckdbS3StyleStorageConfig setAccessKey(String accessKey)
{
this.accessKey = Optional.of(accessKey);
return this;
}
public Optional<String> getAccessKey()
{
return accessKey;
}
@Config(DUCKDB_STORAGE_SECRET_KEY)
@ConfigDescription("The storage secret key")
@ConfigSecuritySensitive
public DuckdbS3StyleStorageConfig setSecretKey(String secretKey)
{
this.secretKey = Optional.of(secretKey);
return this;
}
public Optional<String> getSecretKey()
{
return secretKey;
}
@Config(DUCKDB_STORAGE_REGION)
@ConfigDescription("The storage region")
public DuckdbS3StyleStorageConfig setRegion(String region)
{
this.region = Optional.of(region);
return this;
}
public Optional<String> getRegion()
{
return region;
}
@Config(DUCKDB_STORAGE_URL_STYLE)
@ConfigDescription("The storage url style; default is path")
public DuckdbS3StyleStorageConfig setUrlStyle(String urlStyle)
{
this.urlStyle = urlStyle;
return this;
}
public String getUrlStyle()
{
return urlStyle;
}
@Override
public String generateDuckdbParquetStatement(String path, String tableName)
{
// ref: https://github.com/duckdb/duckdb/issues/1403
StringBuilder sb = new StringBuilder();
// TODO: check why can't we set s3 access key and secret key in Data source
accessKey.ifPresent(accessKey -> sb.append(format("SET s3_access_key_id='%s';\n", accessKey)));
secretKey.ifPresent(secretKey -> sb.append(format("SET s3_secret_access_key='%s';\n", secretKey)));
sb.append("BEGIN TRANSACTION;\n");
sb.append(format("CREATE TABLE \"%s\" AS SELECT * FROM read_parquet('s3://%s');", tableName, path));
sb.append("COMMIT;\n");
return sb.toString();
}
}
wren-base/src/main/java/io/wren/base/client/duckdb/DuckdbType.java
-43
View File
@@ -1,43 +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.wren.base.client.duckdb;
public class DuckdbType
{
private final int jdbcType;
private final String name;
public DuckdbType(int jdbcType, String name)
{
this.jdbcType = jdbcType;
this.name = name;
}
public int getJdbcType()
{
return jdbcType;
}
public String getName()
{
return name;
}
@Override
public String toString()
{
return name;
}
}
wren-base/src/main/java/io/wren/base/client/duckdb/DuckdbTypes.java
-213
View File
@@ -1,213 +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.wren.base.client.duckdb;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import io.wren.base.WrenException;
import io.wren.base.type.BigIntType;
import io.wren.base.type.BooleanType;
import io.wren.base.type.BpCharType;
import io.wren.base.type.ByteaType;
import io.wren.base.type.CharType;
import io.wren.base.type.DateType;
import io.wren.base.type.DoubleType;
import io.wren.base.type.InetType;
import io.wren.base.type.IntegerType;
import io.wren.base.type.IntervalType;
import io.wren.base.type.JsonType;
import io.wren.base.type.NumericType;
import io.wren.base.type.OidType;
import io.wren.base.type.PGType;
import io.wren.base.type.PGTypes;
import io.wren.base.type.RealType;
import io.wren.base.type.SmallIntType;
import io.wren.base.type.TimestampType;
import io.wren.base.type.TimestampWithTimeZoneType;
import io.wren.base.type.TinyIntType;
import io.wren.base.type.UuidType;
import io.wren.base.type.VarcharType;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Types;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import static io.wren.base.metadata.StandardErrorCode.NOT_SUPPORTED;
import static java.util.Locale.ENGLISH;
public final class DuckdbTypes
{
// other types LIST, ENUM, UTINYINT, USMALLINT, STRUCT, UUID, JSON, UINTEGER, UBIGINT, INTERVAL, MAP
public static final DuckdbType BOOLEAN = new DuckdbType(Types.BOOLEAN, "BOOLEAN");
public static final DuckdbType BIGINT = new DuckdbType(Types.BIGINT, "BIGINT");
public static final DuckdbType HUGEINT = new DuckdbType(Types.DECIMAL, "HUGEINT");
public static final DuckdbType BIT = new DuckdbType(Types.BIT, "BIT");
public static final DuckdbType BLOB = new DuckdbType(Types.BLOB, "BLOB");
public static final DuckdbType DATE = new DuckdbType(Types.DATE, "DATE");
public static final DuckdbType DOUBLE = new DuckdbType(Types.DOUBLE, "DOUBLE");
public static final DuckdbType REAL = new DuckdbType(Types.REAL, "REAL");
public static final DuckdbType FLOAT = new DuckdbType(Types.FLOAT, "FLOAT");
public static final DuckdbType DECIMAL = new DuckdbType(Types.DECIMAL, "DECIMAL");
public static final DuckdbType INTEGER = new DuckdbType(Types.INTEGER, "INTEGER");
public static final DuckdbType SMALLINT = new DuckdbType(Types.SMALLINT, "SMALLINT");
public static final DuckdbType TINYINT = new DuckdbType(Types.TINYINT, "TINYINT");
// TODO: check
public static final DuckdbType INTERVAL = new DuckdbType(Types.OTHER, "INTERVAL");
public static final DuckdbType TIME = new DuckdbType(Types.TIME, "TIME");
public static final DuckdbType TIMESTAMP = new DuckdbType(Types.TIMESTAMP, "TIMESTAMP");
public static final DuckdbType TIMESTAMP_WITH_TIMEZONE = new DuckdbType(Types.TIMESTAMP_WITH_TIMEZONE, "TIMESTAMP WITH TIMEZONE");
public static final DuckdbType VARCHAR = new DuckdbType(Types.VARCHAR, "VARCHAR");
public static final DuckdbType JSON = new DuckdbType(Types.STRUCT, "JSON");
public static final DuckdbType UUID = new DuckdbType(Types.JAVA_OBJECT, "UUID");
private static final Map<String, DuckdbType> duckdbTypes = ImmutableMap.<String, DuckdbType>builder()
.put(BOOLEAN.getName(), BOOLEAN)
.put(BIGINT.getName(), BIGINT)
.put(BIT.getName(), BIT)
.put(BLOB.getName(), BLOB)
.put(DATE.getName(), DATE)
.put(DOUBLE.getName(), DOUBLE)
.put(REAL.getName(), REAL)
.put(FLOAT.getName(), FLOAT)
.put(DECIMAL.getName(), DECIMAL)
.put(INTEGER.getName(), INTEGER)
.put(SMALLINT.getName(), SMALLINT)
.put(TINYINT.getName(), TINYINT)
.put(INTERVAL.getName(), INTERVAL)
.put(TIME.getName(), TIME)
.put(TIMESTAMP.getName(), TIMESTAMP)
.put(TIMESTAMP_WITH_TIMEZONE.getName(), TIMESTAMP_WITH_TIMEZONE)
.put(VARCHAR.getName(), VARCHAR)
.put(JSON.getName(), JSON)
.put(HUGEINT.getName(), HUGEINT)
.put(UUID.getName(), UUID)
.build();
private static final Map<Integer, PGType<?>> duckdbTypeToPgTypeMap = ImmutableMap.<Integer, PGType<?>>builder()
.put(BOOLEAN.getJdbcType(), BooleanType.BOOLEAN)
.put(BLOB.getJdbcType(), ByteaType.BYTEA)
.put(TINYINT.getJdbcType(), TinyIntType.TINYINT)
.put(SMALLINT.getJdbcType(), SmallIntType.SMALLINT)
.put(INTEGER.getJdbcType(), IntegerType.INTEGER)
.put(BIGINT.getJdbcType(), BigIntType.BIGINT)
.put(FLOAT.getJdbcType(), RealType.REAL)
.put(DOUBLE.getJdbcType(), DoubleType.DOUBLE)
.put(DECIMAL.getJdbcType(), NumericType.NUMERIC)
.put(VARCHAR.getJdbcType(), VarcharType.VARCHAR)
.put(DATE.getJdbcType(), DateType.DATE)
.put(TIMESTAMP.getJdbcType(), TimestampType.TIMESTAMP)
.put(TIMESTAMP_WITH_TIMEZONE.getJdbcType(), TimestampWithTimeZoneType.TIMESTAMP_WITH_TIMEZONE)
.put(JSON.getJdbcType(), JsonType.JSON)
.build();
// TODO: RECORD, HSTORE
private static final Map<PGType<?>, DuckdbType> pgTypeToDuckdbTypeMap = ImmutableMap.<PGType<?>, DuckdbType>builder()
.put(BooleanType.BOOLEAN, BOOLEAN)
.put(ByteaType.BYTEA, BLOB)
.put(TinyIntType.TINYINT, TINYINT)
.put(SmallIntType.SMALLINT, SMALLINT)
.put(IntegerType.INTEGER, INTEGER)
.put(BigIntType.BIGINT, BIGINT)
.put(RealType.REAL, FLOAT)
.put(DoubleType.DOUBLE, DOUBLE)
.put(NumericType.NUMERIC, DECIMAL)
.put(VarcharType.VARCHAR, VARCHAR)
.put(VarcharType.TextType.TEXT, VARCHAR)
.put(VarcharType.NameType.NAME, VARCHAR)
.put(CharType.CHAR, VARCHAR)
.put(DateType.DATE, DATE)
.put(TimestampType.TIMESTAMP, TIMESTAMP)
.put(TimestampWithTimeZoneType.TIMESTAMP_WITH_TIMEZONE, TIMESTAMP_WITH_TIMEZONE)
.put(IntervalType.INTERVAL, INTERVAL)
.put(JsonType.JSON, JSON)
.put(UuidType.UUID, UUID)
.put(OidType.OID_INSTANCE, BIGINT)
.put(InetType.INET, VARCHAR)
.put(BpCharType.BPCHAR, VARCHAR)
.build();
/**
* getColumnType only return LIST without inner type if the type is INT[]
* But getColumnTypeName will return `INT[]`
*/
public static PGType<?> toPGType(ResultSetMetaData metaData, int i)
throws SQLException
{
try {
return DuckdbTypes.toPGType(metaData.getColumnType(i));
}
catch (WrenException e) {
return DuckdbTypes.toPGType(metaData.getColumnTypeName(i));
}
}
public static PGType<?> toPGType(String typeName)
{
Optional<? extends PGType<?>> pgType = getDuckDBType(getEqualTypeName(typeName))
.map(DuckdbType::getJdbcType)
.map(duckdbTypeToPgTypeMap::get);
if (typeName.endsWith("[]")) {
pgType = pgType
.map(PGType::oid)
.map(PGTypes::getArrayType);
}
return pgType.orElseThrow(() -> new WrenException(NOT_SUPPORTED, "Unsupported Type: " + typeName));
}
private static String getEqualTypeName(String typeName)
{
if (typeName.toLowerCase(ENGLISH).startsWith("decimal")) {
return DECIMAL.toString();
}
if (typeName.toLowerCase(ENGLISH).startsWith("timestamp")) {
return TIMESTAMP.toString();
}
if (typeName.endsWith("[]")) {
return typeName.substring(0, typeName.length() - 2);
}
// TODO: MAP(INT, VARCHAR)
// TODO: STRUCT(i INT, j VARCHAR)
// TODO: UNION(num INT, text VARCHAR)
return typeName;
}
public static PGType<?> toPGType(int type)
{
return Optional.ofNullable(duckdbTypeToPgTypeMap.get(type))
.orElseThrow(() -> new WrenException(NOT_SUPPORTED, "DuckDB unsupported Type: " + type));
}
public static DuckdbType toDuckdbType(PGType<?> type)
{
return Optional.ofNullable(pgTypeToDuckdbTypeMap.get(type))
.orElseThrow(() -> new WrenException(NOT_SUPPORTED, "DuckDB unsupported Type: " + type));
}
public static List<String> getDuckDBTypeNames()
{
return ImmutableList.copyOf(duckdbTypes.keySet());
}
public static Optional<DuckdbType> getDuckDBType(String name)
{
return Optional.ofNullable(duckdbTypes.get(name));
}
private DuckdbTypes() {}
}
wren-base/src/main/java/io/wren/base/client/duckdb/DuckdbUtil.java
-65
View File
@@ -1,65 +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.wren.base.client.duckdb;
import io.airlift.units.DataSize;
import io.wren.base.WrenException;
import io.wren.base.metadata.StandardErrorCode;
import java.util.HashMap;
import java.util.Map;
import static java.util.Locale.ENGLISH;
public final class DuckdbUtil
{
private static final Map<String, DataSize.Unit> UNIT_MAP = new HashMap<>();
// https://github.com/duckdb/duckdb/blob/4c7cb20474baa3a8ca1d5d8ceb22beae0e4c0e4c/src/common/string_util.cpp#L178
static {
UNIT_MAP.put("BYTE", DataSize.Unit.BYTE);
UNIT_MAP.put("BYTES", DataSize.Unit.BYTE);
UNIT_MAP.put("KB", DataSize.Unit.KILOBYTE);
UNIT_MAP.put("KIB", DataSize.Unit.KILOBYTE);
UNIT_MAP.put("MB", DataSize.Unit.MEGABYTE);
UNIT_MAP.put("MIB", DataSize.Unit.MEGABYTE);
UNIT_MAP.put("GB", DataSize.Unit.GIGABYTE);
UNIT_MAP.put("GIB", DataSize.Unit.GIGABYTE);
UNIT_MAP.put("TB", DataSize.Unit.TERABYTE);
UNIT_MAP.put("TIB", DataSize.Unit.TERABYTE);
UNIT_MAP.put("PB", DataSize.Unit.PETABYTE);
UNIT_MAP.put("PIB", DataSize.Unit.PETABYTE);
}
private DuckdbUtil() {}
public static DataSize convertDuckDBUnits(String valueWithUnit)
{
try {
String valueWithUnitUpperCase = valueWithUnit.toUpperCase(ENGLISH);
for (Map.Entry<String, DataSize.Unit> entry : UNIT_MAP.entrySet()) {
String unit = entry.getKey();
if (valueWithUnitUpperCase.endsWith(unit)) {
double value = Double.parseDouble(valueWithUnitUpperCase.substring(0, valueWithUnitUpperCase.length() - unit.length()).trim());
return DataSize.of((long) value, entry.getValue());
}
}
}
catch (Exception e) {
throw new WrenException(StandardErrorCode.GENERIC_INTERNAL_ERROR, String.format("Failed to parse duckdb value %s", valueWithUnit), e);
}
throw new WrenException(StandardErrorCode.GENERIC_INTERNAL_ERROR, String.format("Failed to parse duckdb value %s", valueWithUnit));
}
}
wren-base/src/main/java/io/wren/base/client/jdbc/BaseJdbcRecordIterator.java
-7
View File
@@ -27,7 +27,6 @@ import java.util.List;
import java.util.NoSuchElementException;
import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
import static java.util.Collections.emptyList;
import static java.util.Objects.requireNonNull;
public abstract class BaseJdbcRecordIterator<T>
@@ -41,12 +40,6 @@ public abstract class BaseJdbcRecordIterator<T>
private boolean hasNext;
public BaseJdbcRecordIterator(Client client, String sql)
throws SQLException
{
this(client, sql, emptyList());
}
public BaseJdbcRecordIterator(Client client, String sql, List<Parameter> parameters)
throws SQLException
{
wren-base/src/main/java/io/wren/base/config/BigQueryConfig.java
-120
View File
@@ -1,120 +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.wren.base.config;
import io.airlift.configuration.Config;
import io.airlift.configuration.ConfigDescription;
import io.airlift.configuration.ConfigSecuritySensitive;
import io.airlift.configuration.validation.FileExists;
import jakarta.validation.constraints.NotNull;
import java.util.Optional;
public class BigQueryConfig
{
public static final String BIGQUERY_CRENDITALS_KEY = "bigquery.credentials-key";
public static final String BIGQUERY_CRENDITALS_FILE = "bigquery.credentials-file";
public static final String BIGQUERY_PROJECT_ID = "bigquery.project-id";
public static final String BIGQUERY_LOCATION = "bigquery.location";
public static final String BIGQUERY_BUCKET_NAME = "bigquery.bucket-name";
public static final String BIGQUERY_METADATA_SCHEMA_PREFIX = "bigquery.metadata.schema.prefix";
private Optional<String> credentialsKey = Optional.empty();
private Optional<String> credentialsFile = Optional.empty();
private Optional<String> projectId = Optional.empty();
private Optional<String> location = Optional.empty();
private Optional<String> bucketName = Optional.empty();
private String metadataSchemaPrefix = "";
public Optional<String> getCredentialsKey()
{
return credentialsKey;
}
@Config(BIGQUERY_CRENDITALS_KEY)
@ConfigDescription("The base64 encoded credentials key")
@ConfigSecuritySensitive
public BigQueryConfig setCredentialsKey(String credentialsKey)
{
this.credentialsKey = Optional.ofNullable(credentialsKey);
return this;
}
public Optional<@FileExists String> getCredentialsFile()
{
return credentialsFile;
}
@Config(BIGQUERY_CRENDITALS_FILE)
@ConfigDescription("The path to the JSON credentials file")
public BigQueryConfig setCredentialsFile(String credentialsFile)
{
this.credentialsFile = Optional.ofNullable(credentialsFile);
return this;
}
public Optional<String> getProjectId()
{
return projectId;
}
@Config(BIGQUERY_PROJECT_ID)
@ConfigDescription("The Google Cloud Project ID where the data reside")
public BigQueryConfig setProjectId(String projectId)
{
this.projectId = Optional.ofNullable(projectId);
return this;
}
public Optional<String> getLocation()
{
return location;
}
@Config(BIGQUERY_LOCATION)
@ConfigDescription("The Google Cloud Project ID where the data reside")
public BigQueryConfig setLocation(String location)
{
this.location = Optional.ofNullable(location);
return this;
}
public Optional<String> getBucketName()
{
return bucketName;
}
@Config(BIGQUERY_BUCKET_NAME)
@ConfigDescription("The Google Cloud bucket name used to temporarily store the metric cached results")
public BigQueryConfig setBucketName(String bucketName)
{
this.bucketName = Optional.ofNullable(bucketName);
return this;
}
@NotNull
public String getMetadataSchemaPrefix()
{
return metadataSchemaPrefix;
}
@Config(BIGQUERY_METADATA_SCHEMA_PREFIX)
@ConfigDescription("Wren needs to create two schemas in BigQuery: wren_temp, pg_catalog. This is a config to add a prefix to the names of these two schemas if it's set.")
public BigQueryConfig setMetadataSchemaPrefix(String metadataSchemaPrefix)
{
this.metadataSchemaPrefix = metadataSchemaPrefix;
return this;
}
}
wren-base/src/main/java/io/wren/base/config/ConfigManager.java
+4 -175
View File
@@ -20,10 +20,8 @@ import com.google.inject.Inject;
import io.airlift.log.Logger;
import io.airlift.units.DataSize;
import io.wren.base.WrenException;
import io.wren.base.client.duckdb.CacheStorageConfig;
import io.wren.base.client.duckdb.DuckDBConfig;
import io.wren.base.client.duckdb.DuckDBConnectorConfig;
import io.wren.base.client.duckdb.DuckdbS3StyleStorageConfig;
import java.io.File;
import java.io.FileWriter;
@@ -53,26 +51,7 @@ import static io.wren.base.client.duckdb.DuckDBConfig.DUCKDB_MEMORY_LIMIT;
import static io.wren.base.client.duckdb.DuckDBConfig.DUCKDB_TEMP_DIRECTORY;
import static io.wren.base.client.duckdb.DuckDBConnectorConfig.DUCKDB_CONNECTOR_INIT_SQL_PATH;
import static io.wren.base.client.duckdb.DuckDBConnectorConfig.DUCKDB_CONNECTOR_SESSION_SQL_PATH;
import static io.wren.base.client.duckdb.DuckdbS3StyleStorageConfig.DUCKDB_STORAGE_ACCESS_KEY;
import static io.wren.base.client.duckdb.DuckdbS3StyleStorageConfig.DUCKDB_STORAGE_ENDPOINT;
import static io.wren.base.client.duckdb.DuckdbS3StyleStorageConfig.DUCKDB_STORAGE_REGION;
import static io.wren.base.client.duckdb.DuckdbS3StyleStorageConfig.DUCKDB_STORAGE_SECRET_KEY;
import static io.wren.base.client.duckdb.DuckdbS3StyleStorageConfig.DUCKDB_STORAGE_URL_STYLE;
import static io.wren.base.client.duckdb.FileUtil.ARCHIVED;
import static io.wren.base.config.PostgresConfig.POSTGRES_JDBC_URL;
import static io.wren.base.config.PostgresConfig.POSTGRES_PASSWORD;
import static io.wren.base.config.PostgresConfig.POSTGRES_USER;
import static io.wren.base.config.PostgresWireProtocolConfig.PG_WIRE_PROTOCOL_AUTH_FILE;
import static io.wren.base.config.PostgresWireProtocolConfig.PG_WIRE_PROTOCOL_NETTY_THREAD_COUNT;
import static io.wren.base.config.PostgresWireProtocolConfig.PG_WIRE_PROTOCOL_PORT;
import static io.wren.base.config.PostgresWireProtocolConfig.PG_WIRE_PROTOCOL_SSL_ENABLED;
import static io.wren.base.config.SnowflakeConfig.SNOWFLAKE_DATABASE;
import static io.wren.base.config.SnowflakeConfig.SNOWFLAKE_JDBC_URL;
import static io.wren.base.config.SnowflakeConfig.SNOWFLAKE_PASSWORD;
import static io.wren.base.config.SnowflakeConfig.SNOWFLAKE_ROLE;
import static io.wren.base.config.SnowflakeConfig.SNOWFLAKE_SCHEMA;
import static io.wren.base.config.SnowflakeConfig.SNOWFLAKE_USER;
import static io.wren.base.config.SnowflakeConfig.SNOWFLAKE_WAREHOUSE;
import static io.wren.base.metadata.StandardErrorCode.NOT_FOUND;
import static java.lang.String.format;
import static java.util.stream.Collectors.toList;
@@ -82,13 +61,8 @@ public class ConfigManager
{
private static final Logger LOG = Logger.get(ConfigManager.class);
private Optional<WrenConfig> wrenConfig;
private Optional<PostgresConfig> postgresConfig;
private Optional<BigQueryConfig> bigQueryConfig;
private Optional<DuckDBConfig> duckDBConfig;
private Optional<PostgresWireProtocolConfig> postgresWireProtocolConfig;
private Optional<DuckdbS3StyleStorageConfig> duckdbS3StyleStorageConfig;
private Optional<DuckDBConnectorConfig> duckDBConnectorConfig;
private Optional<SnowflakeConfig> snowflakeConfig;
private final Map<String, String> configs = new HashMap<>();
// All configs set by user and config files. It's used to sync with config file.
@@ -100,32 +74,17 @@ public class ConfigManager
@Inject
public ConfigManager(
WrenConfig wrenConfig,
PostgresConfig postgresConfig,
BigQueryConfig bigQueryConfig,
DuckDBConfig duckDBConfig,
PostgresWireProtocolConfig postgresWireProtocolConfig,
DuckdbS3StyleStorageConfig duckdbS3StyleStorageConfig,
DuckDBConnectorConfig duckDBConnectorConfig,
SnowflakeConfig snowflakeConfig)
DuckDBConnectorConfig duckDBConnectorConfig)
{
this.wrenConfig = Optional.of(wrenConfig);
this.postgresConfig = Optional.of(postgresConfig);
this.bigQueryConfig = Optional.of(bigQueryConfig);
this.duckDBConfig = Optional.of(duckDBConfig);
this.postgresWireProtocolConfig = Optional.of(postgresWireProtocolConfig);
this.duckdbS3StyleStorageConfig = Optional.of(duckdbS3StyleStorageConfig);
this.duckDBConnectorConfig = Optional.of(duckDBConnectorConfig);
this.snowflakeConfig = Optional.of(snowflakeConfig);
initConfig(
wrenConfig,
postgresConfig,
bigQueryConfig,
duckDBConfig,
postgresWireProtocolConfig,
duckdbS3StyleStorageConfig,
duckDBConnectorConfig,
snowflakeConfig);
duckDBConnectorConfig);
try {
setConfigs.putAll(loadPropertiesFrom(configFile));
@@ -137,22 +96,12 @@ public class ConfigManager
private void initConfig(
WrenConfig wrenConfig,
PostgresConfig postgresConfig,
BigQueryConfig bigQueryConfig,
DuckDBConfig duckDBConfig,
PostgresWireProtocolConfig postgresWireProtocolConfig,
DuckdbS3StyleStorageConfig duckdbS3StyleStorageConfig,
DuckDBConnectorConfig duckDBConnectorConfig,
SnowflakeConfig snowflakeConfig)
DuckDBConnectorConfig duckDBConnectorConfig)
{
initConfig(WrenConfig.WREN_DIRECTORY, wrenConfig.getWrenMDLDirectory().getPath(), false, true);
initConfig(WrenConfig.WREN_DATASOURCE_TYPE, Optional.ofNullable(wrenConfig.getDataSourceType()).map(Enum::name).orElse(null), true, false);
initConfig(WrenConfig.WREN_ENABLE_DYNAMIC_FIELDS, Boolean.toString(wrenConfig.getEnableDynamicFields()), false, false);
initConfig(DUCKDB_STORAGE_ENDPOINT, duckdbS3StyleStorageConfig.getEndpoint(), false, true);
initConfig(DUCKDB_STORAGE_ACCESS_KEY, duckdbS3StyleStorageConfig.getAccessKey().orElse(null), true, false);
initConfig(DUCKDB_STORAGE_SECRET_KEY, duckdbS3StyleStorageConfig.getSecretKey().orElse(null), true, false);
initConfig(DUCKDB_STORAGE_REGION, duckdbS3StyleStorageConfig.getRegion().orElse(null), true, false);
initConfig(DUCKDB_STORAGE_URL_STYLE, duckdbS3StyleStorageConfig.getUrlStyle(), false, false);
initConfig(DUCKDB_MEMORY_LIMIT, duckDBConfig.getMemoryLimit().toString(), true, false);
initConfig(DUCKDB_HOME_DIRECTORY, duckDBConfig.getHomeDirectory(), true, false);
initConfig(DUCKDB_TEMP_DIRECTORY, duckDBConfig.getTempDirectory(), true, false);
@@ -161,28 +110,8 @@ public class ConfigManager
initConfig(DUCKDB_MAX_CONCURRENT_METADATA_QUERIES, Integer.toString(duckDBConfig.getMaxConcurrentMetadataQueries()), false, true);
initConfig(DUCKDB_MAX_CACHE_QUERY_TIMEOUT, Long.toString(duckDBConfig.getMaxCacheQueryTimeout()), false, true);
initConfig(DUCKDB_CACHE_TASK_RETRY_DELAY, Long.toString(duckDBConfig.getCacheTaskRetryDelay()), false, true);
initConfig(PG_WIRE_PROTOCOL_PORT, postgresWireProtocolConfig.getPort(), false, true);
initConfig(PG_WIRE_PROTOCOL_SSL_ENABLED, Boolean.toString(postgresWireProtocolConfig.isSslEnable()), false, true);
initConfig(PG_WIRE_PROTOCOL_NETTY_THREAD_COUNT, Integer.toString(postgresWireProtocolConfig.getNettyThreadCount()), false, true);
initConfig(PG_WIRE_PROTOCOL_AUTH_FILE, postgresWireProtocolConfig.getAuthFile().getPath(), false, true);
initConfig(BigQueryConfig.BIGQUERY_CRENDITALS_KEY, bigQueryConfig.getCredentialsKey().orElse(null), true, false);
initConfig(BigQueryConfig.BIGQUERY_CRENDITALS_FILE, bigQueryConfig.getCredentialsFile().orElse(null), true, false);
initConfig(BigQueryConfig.BIGQUERY_PROJECT_ID, bigQueryConfig.getProjectId().orElse(null), true, false);
initConfig(BigQueryConfig.BIGQUERY_LOCATION, bigQueryConfig.getLocation().orElse(null), true, false);
initConfig(BigQueryConfig.BIGQUERY_BUCKET_NAME, bigQueryConfig.getBucketName().orElse(null), true, false);
initConfig(BigQueryConfig.BIGQUERY_METADATA_SCHEMA_PREFIX, bigQueryConfig.getMetadataSchemaPrefix(), true, false);
initConfig(POSTGRES_JDBC_URL, postgresConfig.getJdbcUrl(), true, false);
initConfig(POSTGRES_USER, postgresConfig.getUser(), true, false);
initConfig(POSTGRES_PASSWORD, postgresConfig.getPassword(), true, false);
initConfig(DUCKDB_CONNECTOR_INIT_SQL_PATH, duckDBConnectorConfig.getInitSQLPath(), false, false);
initConfig(DUCKDB_CONNECTOR_SESSION_SQL_PATH, duckDBConnectorConfig.getSessionSQLPath(), false, false);
initConfig(SNOWFLAKE_JDBC_URL, snowflakeConfig.getJdbcUrl(), true, false);
initConfig(SNOWFLAKE_USER, snowflakeConfig.getUser(), true, false);
initConfig(SNOWFLAKE_PASSWORD, snowflakeConfig.getPassword(), true, false);
initConfig(SNOWFLAKE_DATABASE, snowflakeConfig.getDatabase().orElse(null), true, false);
initConfig(SNOWFLAKE_SCHEMA, snowflakeConfig.getSchema().orElse(null), true, false);
initConfig(SNOWFLAKE_WAREHOUSE, snowflakeConfig.getWarehouse().orElse(null), true, false);
initConfig(SNOWFLAKE_ROLE, snowflakeConfig.getRole().orElse(null), true, false);
}
private void initConfig(String key, String value, boolean requiredReload, boolean isStatic)
@@ -207,20 +136,6 @@ public class ConfigManager
return result;
});
}
if (config == BigQueryConfig.class) {
return (T) bigQueryConfig.orElseGet(() -> {
BigQueryConfig result = getBigQueryConfig();
bigQueryConfig = Optional.of(result);
return result;
});
}
if (config == PostgresConfig.class) {
return (T) postgresConfig.orElseGet(() -> {
PostgresConfig result = getPostgresConfig();
postgresConfig = Optional.of(result);
return result;
});
}
if (config == DuckDBConfig.class) {
return (T) duckDBConfig.orElseGet(() -> {
DuckDBConfig result = getDuckDBConfig();
@@ -228,21 +143,6 @@ public class ConfigManager
return result;
});
}
if (config == PostgresWireProtocolConfig.class) {
return (T) postgresWireProtocolConfig.orElseGet(() -> {
PostgresWireProtocolConfig result = getPostgresWireProtocolConfig();
postgresWireProtocolConfig = Optional.of(result);
return result;
});
}
if (config == CacheStorageConfig.class &&
wrenConfig.map(WrenConfig::getDataSourceType).stream().anyMatch(type -> type == WrenConfig.DataSourceType.BIGQUERY)) {
return (T) duckdbS3StyleStorageConfig.orElseGet(() -> {
DuckdbS3StyleStorageConfig result = getDuckdbS3StyleStorageConfig();
duckdbS3StyleStorageConfig = Optional.of(result);
return result;
});
}
if (config == DuckDBConnectorConfig.class) {
return (T) duckDBConnectorConfig.orElseGet(() -> {
DuckDBConnectorConfig result = getDuckDBConnectorConfig();
@@ -250,13 +150,6 @@ public class ConfigManager
return result;
});
}
if (config == SnowflakeConfig.class) {
return (T) snowflakeConfig.orElseGet(() -> {
SnowflakeConfig result = getSnowflakeConfig();
snowflakeConfig = Optional.of(result);
return result;
});
}
throw new RuntimeException("Unknown config class: " + config.getName());
}
@@ -270,27 +163,6 @@ public class ConfigManager
return result;
}
private BigQueryConfig getBigQueryConfig()
{
BigQueryConfig result = new BigQueryConfig();
result.setCredentialsKey(configs.get(BigQueryConfig.BIGQUERY_CRENDITALS_KEY));
result.setCredentialsFile(configs.get(BigQueryConfig.BIGQUERY_CRENDITALS_FILE));
result.setProjectId(configs.get(BigQueryConfig.BIGQUERY_PROJECT_ID));
result.setLocation(configs.get(BigQueryConfig.BIGQUERY_LOCATION));
result.setBucketName(configs.get(BigQueryConfig.BIGQUERY_BUCKET_NAME));
result.setMetadataSchemaPrefix(configs.get(BigQueryConfig.BIGQUERY_METADATA_SCHEMA_PREFIX));
return result;
}
private PostgresConfig getPostgresConfig()
{
PostgresConfig result = new PostgresConfig();
result.setJdbcUrl(configs.get(POSTGRES_JDBC_URL));
result.setUser(configs.get(POSTGRES_USER));
result.setPassword(configs.get(POSTGRES_PASSWORD));
return result;
}
private DuckDBConfig getDuckDBConfig()
{
DuckDBConfig result = new DuckDBConfig();
@@ -304,27 +176,6 @@ public class ConfigManager
return result;
}
private PostgresWireProtocolConfig getPostgresWireProtocolConfig()
{
PostgresWireProtocolConfig result = new PostgresWireProtocolConfig();
result.setPort(configs.get(PG_WIRE_PROTOCOL_PORT));
result.setSslEnable(Boolean.parseBoolean(configs.get(PG_WIRE_PROTOCOL_SSL_ENABLED)));
result.setNettyThreadCount(Integer.parseInt(configs.get(PG_WIRE_PROTOCOL_NETTY_THREAD_COUNT)));
result.setAuthFile(new File(configs.get(PG_WIRE_PROTOCOL_AUTH_FILE)));
return result;
}
private DuckdbS3StyleStorageConfig getDuckdbS3StyleStorageConfig()
{
DuckdbS3StyleStorageConfig result = new DuckdbS3StyleStorageConfig();
result.setEndpoint(configs.get(DUCKDB_STORAGE_ENDPOINT));
result.setAccessKey(configs.get(DUCKDB_STORAGE_ACCESS_KEY));
result.setSecretKey(configs.get(DUCKDB_STORAGE_SECRET_KEY));
result.setRegion(configs.get(DUCKDB_STORAGE_REGION));
result.setUrlStyle(configs.get(DUCKDB_STORAGE_URL_STYLE));
return result;
}
private DuckDBConnectorConfig getDuckDBConnectorConfig()
{
DuckDBConnectorConfig result = new DuckDBConnectorConfig();
@@ -333,19 +184,6 @@ public class ConfigManager
return result;
}
private SnowflakeConfig getSnowflakeConfig()
{
SnowflakeConfig config = new SnowflakeConfig();
config.setJdbcUrl(configs.get(SNOWFLAKE_JDBC_URL));
config.setUser(configs.get(SNOWFLAKE_USER));
config.setPassword(configs.get(SNOWFLAKE_PASSWORD));
config.setDatabase(configs.get(SNOWFLAKE_DATABASE));
config.setSchema(configs.get(SNOWFLAKE_SCHEMA));
config.setWarehouse(configs.get(SNOWFLAKE_WAREHOUSE));
config.setRole(configs.get(SNOWFLAKE_ROLE));
return config;
}
public synchronized boolean setConfigs(List<ConfigEntry> configEntries, boolean reset)
{
if (reset) {
@@ -386,11 +224,7 @@ public class ConfigManager
private void resetCache()
{
wrenConfig = Optional.empty();
postgresConfig = Optional.empty();
bigQueryConfig = Optional.empty();
duckDBConfig = Optional.empty();
postgresWireProtocolConfig = Optional.empty();
duckdbS3StyleStorageConfig = Optional.empty();
duckDBConnectorConfig = Optional.empty();
}
@@ -400,13 +234,8 @@ public class ConfigManager
setConfigs.clear();
initConfig(
new WrenConfig(),
new PostgresConfig(),
new BigQueryConfig(),
new DuckDBConfig(),
new PostgresWireProtocolConfig(),
new DuckdbS3StyleStorageConfig(),
new DuckDBConnectorConfig(),
new SnowflakeConfig());
new DuckDBConnectorConfig());
}
private void syncFile(Map<String, String> updated)
wren-base/src/main/java/io/wren/base/config/PostgresConfig.java
-63
View File
@@ -1,63 +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.wren.base.config;
import io.airlift.configuration.Config;
public class PostgresConfig
{
public static final String POSTGRES_JDBC_URL = "postgres.jdbc.url";
public static final String POSTGRES_USER = "postgres.user";
public static final String POSTGRES_PASSWORD = "postgres.password";
private String jdbcUrl;
private String user;
private String password;
public String getJdbcUrl()
{
return jdbcUrl;
}
@Config(POSTGRES_JDBC_URL)
public PostgresConfig setJdbcUrl(String jdbcUrl)
{
this.jdbcUrl = jdbcUrl;
return this;
}
public String getUser()
{
return user;
}
@Config(POSTGRES_USER)
public PostgresConfig setUser(String user)
{
this.user = user;
return this;
}
public String getPassword()
{
return password;
}
@Config(POSTGRES_PASSWORD)
public PostgresConfig setPassword(String password)
{
this.password = password;
return this;
}
}
wren-base/src/main/java/io/wren/base/config/PostgresWireProtocolConfig.java
-96
View File
@@ -1,96 +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.wren.base.config;
import io.airlift.configuration.Config;
import jakarta.validation.constraints.NotNull;
import java.io.File;
public class PostgresWireProtocolConfig
{
public static final String PG_WIRE_PROTOCOL_ENABLED = "pg-wire-protocol.enabled";
public static final String PG_WIRE_PROTOCOL_SSL_ENABLED = "pg-wire-protocol.ssl.enabled";
public static final String PG_WIRE_PROTOCOL_NETTY_THREAD_COUNT = "pg-wire-protocol.netty.thread.count";
public static final String PG_WIRE_PROTOCOL_AUTH_FILE = "pg-wire-protocol.auth.file";
public static final String PG_WIRE_PROTOCOL_PORT = "pg-wire-protocol.port";
private String port = "7432";
private boolean sslEnable;
private int nettyThreadCount;
private File authFile = new File("etc/accounts");
private boolean pgWireProtocolEnabled;
@NotNull
public String getPort()
{
return port;
}
@Config(PG_WIRE_PROTOCOL_PORT)
public PostgresWireProtocolConfig setPort(String port)
{
this.port = port;
return this;
}
@NotNull
public boolean isSslEnable()
{
return sslEnable;
}
@Config(PG_WIRE_PROTOCOL_SSL_ENABLED)
public PostgresWireProtocolConfig setSslEnable(boolean sslEnable)
{
this.sslEnable = sslEnable;
return this;
}
@NotNull
public int getNettyThreadCount()
{
return nettyThreadCount;
}
@Config(PG_WIRE_PROTOCOL_NETTY_THREAD_COUNT)
public PostgresWireProtocolConfig setNettyThreadCount(int nettyThreadCount)
{
this.nettyThreadCount = nettyThreadCount;
return this;
}
public File getAuthFile()
{
return authFile;
}
@Config(PG_WIRE_PROTOCOL_AUTH_FILE)
public PostgresWireProtocolConfig setAuthFile(File authFile)
{
this.authFile = authFile;
return this;
}
@Config(PG_WIRE_PROTOCOL_ENABLED)
public void setPgWireProtocolEnabled(boolean pgWireProtocolEnabled)
{
this.pgWireProtocolEnabled = pgWireProtocolEnabled;
}
public boolean isPgWireProtocolEnabled()
{
return pgWireProtocolEnabled;
}
}
wren-base/src/main/java/io/wren/base/config/SQLGlotConfig.java
-60
View File
@@ -1,60 +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.wren.base.config;
import io.airlift.configuration.Config;
import java.io.IOException;
import java.net.ServerSocket;
import java.util.Optional;
import static java.lang.System.getenv;
public class SQLGlotConfig
{
public static final String SQLGLOT_PORT = "sqlglot.port";
private int port = Optional.ofNullable(getenv("SQLGLOT_PORT"))
.map(Integer::parseInt)
.orElse(8000);
public int getPort()
{
return port;
}
@Config(SQLGLOT_PORT)
public void setPort(int port)
{
this.port = port;
}
public static SQLGlotConfig createConfigWithFreePort()
{
SQLGlotConfig config = new SQLGlotConfig();
config.setPort(findFreePort());
return config;
}
private static int findFreePort()
{
try (ServerSocket serverSocket = new ServerSocket(0)) {
return serverSocket.getLocalPort();
}
catch (IOException e) {
throw new AssertionError(e);
}
}
}
wren-base/src/main/java/io/wren/base/config/SnowflakeConfig.java
-122
View File
@@ -1,122 +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.wren.base.config;
import io.airlift.configuration.Config;
import java.util.Optional;
public class SnowflakeConfig
{
public static final String SNOWFLAKE_JDBC_URL = "snowflake.jdbc.url";
public static final String SNOWFLAKE_USER = "snowflake.user";
public static final String SNOWFLAKE_PASSWORD = "snowflake.password";
public static final String SNOWFLAKE_ROLE = "snowflake.role";
public static final String SNOWFLAKE_WAREHOUSE = "snowflake.warehouse";
public static final String SNOWFLAKE_DATABASE = "snowflake.database";
public static final String SNOWFLAKE_SCHEMA = "snowflake.schema";
private String jdbcUrl;
private String user;
private String password;
private String role;
private String warehouse;
private String database;
private String schema;
public String getJdbcUrl()
{
return jdbcUrl;
}
@Config(SNOWFLAKE_JDBC_URL)
public SnowflakeConfig setJdbcUrl(String jdbcUrl)
{
this.jdbcUrl = jdbcUrl;
return this;
}
public String getUser()
{
return user;
}
@Config(SNOWFLAKE_USER)
public SnowflakeConfig setUser(String user)
{
this.user = user;
return this;
}
public String getPassword()
{
return password;
}
@Config(SNOWFLAKE_PASSWORD)
public SnowflakeConfig setPassword(String password)
{
this.password = password;
return this;
}
public Optional<String> getRole()
{
return Optional.ofNullable(role);
}
@Config(SNOWFLAKE_ROLE)
public SnowflakeConfig setRole(String role)
{
this.role = role;
return this;
}
public Optional<String> getWarehouse()
{
return Optional.ofNullable(warehouse);
}
@Config(SNOWFLAKE_WAREHOUSE)
public SnowflakeConfig setWarehouse(String warehouse)
{
this.warehouse = warehouse;
return this;
}
public Optional<String> getDatabase()
{
return Optional.ofNullable(database);
}
@Config(SNOWFLAKE_DATABASE)
public SnowflakeConfig setDatabase(String database)
{
this.database = database;
return this;
}
public Optional<String> getSchema()
{
return Optional.ofNullable(schema);
}
@Config(SNOWFLAKE_SCHEMA)
public SnowflakeConfig setSchema(String schema)
{
this.schema = schema;
return this;
}
}
wren-base/src/main/java/io/wren/base/function/OperatorType.java
-54
View File
@@ -1,54 +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.wren.base.function;
public enum OperatorType
{
ADD("+", 2),
SUBTRACT("-", 2),
MULTIPLY("*", 2),
DIVIDE("/", 2),
MODULUS("%", 2),
NEGATION("-", 1),
EQUAL("=", 2),
COMPARISON("COMPARISON", 2),
LESS_THAN("<", 2),
LESS_THAN_OR_EQUAL("<=", 2),
CAST("CAST", 1),
SUBSCRIPT("[]", 2),
HASH_CODE("HASH CODE", 1),
SATURATED_FLOOR_CAST("SATURATED FLOOR CAST", 1),
IS_DISTINCT_FROM("IS DISTINCT FROM", 2),
XX_HASH_64("XX HASH 64", 1),
INDETERMINATE("INDETERMINATE", 1);
private final String operator;
private final int argumentCount;
OperatorType(String operator, int argumentCount)
{
this.operator = operator;
this.argumentCount = argumentCount;
}
public String getOperator()
{
return operator;
}
public int getArgumentCount()
{
return argumentCount;
}
}
wren-base/src/main/java/io/wren/base/metadata/BasicFunctionRegistry.java
-51
View File
@@ -1,51 +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.wren.base.metadata;
import com.google.common.collect.ImmutableList;
import io.wren.base.pgcatalog.function.FunctionRegistry;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
public class BasicFunctionRegistry
implements FunctionRegistry<Function>
{
private final List<Function> functions = ImmutableList.<Function>builder()
.add(BasicFunctions.DATE_TRUNC)
.build();
private final Map<FunctionKey, Function> simpleNameToFunction = new HashMap<>();
public BasicFunctionRegistry()
{
// TODO: handle function name overloading
// https://github.com/Canner/canner-metric-layer/issues/73
// use HashMap to handle multiple same key entries
functions.forEach(function -> simpleNameToFunction.put(FunctionKey.functionKey(function.getName(), function.getArguments().map(List::size).orElse(0)), function));
}
public List<Function> getFunctions()
{
return functions;
}
public Optional<Function> getFunction(String name, int numArgument)
{
return Optional.ofNullable(simpleNameToFunction.get(FunctionKey.functionKey(name, numArgument)));
}
}
wren-base/src/main/java/io/wren/base/metadata/BasicFunctions.java
-32
View File
@@ -1,32 +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.wren.base.metadata;
import com.google.common.collect.ImmutableList;
import io.wren.base.type.DateType;
import io.wren.base.type.VarcharType;
import static io.wren.base.metadata.Function.Argument.argument;
public class BasicFunctions
{
private BasicFunctions() {}
public static final Function DATE_TRUNC = Function.builder()
.setName("date_trunc")
.setArguments(ImmutableList.of(argument("field", VarcharType.VARCHAR), argument("source", DateType.DATE)))
.setReturnType(DateType.DATE)
.build();
}
wren-base/src/main/java/io/wren/base/metadata/ColumnMetadata.java
-102
View File
@@ -1,102 +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.wren.base.metadata;
import io.wren.base.type.PGType;
import java.util.Objects;
import static java.util.Objects.requireNonNull;
public class ColumnMetadata
{
private final String name;
private final PGType<?> type;
private ColumnMetadata(String name, PGType<?> type)
{
this.name = name;
this.type = type;
}
public String getName()
{
return name;
}
public PGType<?> getType()
{
return type;
}
@Override
public String toString()
{
StringBuilder sb = new StringBuilder("ColumnMetadata{");
sb.append("name='").append(name).append('\'');
sb.append(", type=").append(type);
sb.append('}');
return sb.toString();
}
@Override
public int hashCode()
{
return Objects.hash(name, type);
}
@Override
public boolean equals(Object obj)
{
if (this == obj) {
return true;
}
if (obj == null || getClass() != obj.getClass()) {
return false;
}
ColumnMetadata other = (ColumnMetadata) obj;
return Objects.equals(this.name, other.name) &&
Objects.equals(this.type, other.type);
}
public static Builder builder()
{
return new Builder();
}
public static class Builder
{
private String name;
private PGType<?> type;
private Builder() {}
public Builder setName(String name)
{
this.name = requireNonNull(name, "name is null");
return this;
}
public Builder setType(PGType<?> type)
{
this.type = requireNonNull(type, "type is null");
return this;
}
public ColumnMetadata build()
{
return new ColumnMetadata(name, type);
}
}
}
wren-base/src/main/java/io/wren/base/metadata/Function.java
-114
View File
@@ -1,114 +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.wren.base.metadata;
import io.wren.base.type.PGType;
import java.util.List;
import java.util.Optional;
import static java.util.Objects.requireNonNull;
public class Function
{
protected final String name;
protected final List<Argument> arguments;
protected final PGType returnType;
public Function(String name, List<Argument> arguments, PGType returnType)
{
this.name = name;
this.arguments = arguments;
this.returnType = returnType;
}
public String getName()
{
return name;
}
public Optional<List<Argument>> getArguments()
{
return Optional.ofNullable(arguments);
}
public Optional<PGType> getReturnType()
{
return Optional.ofNullable(returnType);
}
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 Builder builder()
{
return new Builder();
}
public static class Builder
{
private String name;
private List<Argument> arguments;
private PGType returnType;
public Builder setName(String name)
{
this.name = name;
return this;
}
public Builder setArguments(List<Argument> arguments)
{
this.arguments = arguments;
return this;
}
public Builder setReturnType(PGType returnType)
{
this.returnType = returnType;
return this;
}
public Function build()
{
requireNonNull(name, "name is null");
return new Function(name, arguments, returnType);
}
}
}
wren-base/src/main/java/io/wren/base/metadata/FunctionBundle.java
-42
View File
@@ -1,42 +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.wren.base.metadata;
import io.wren.base.pgcatalog.function.DataSourceFunctionRegistry;
import io.wren.base.pgcatalog.function.PgMetastoreFunctionRegistry;
import java.util.Optional;
public class FunctionBundle
{
private static final BasicFunctionRegistry basicFunctionRegistry;
private static final PgMetastoreFunctionRegistry pgMetastoreFunctionRegistry;
private static final DataSourceFunctionRegistry datSourceFunctionRegistry;
private FunctionBundle() {}
static {
basicFunctionRegistry = new BasicFunctionRegistry();
pgMetastoreFunctionRegistry = new PgMetastoreFunctionRegistry();
datSourceFunctionRegistry = new DataSourceFunctionRegistry();
}
public static Optional<Function> getFunction(String name, int numArgument)
{
return basicFunctionRegistry.getFunction(name, numArgument)
.or(() -> pgMetastoreFunctionRegistry.getFunction(name, numArgument))
.or(() -> datSourceFunctionRegistry.getFunction(name, numArgument));
}
}
wren-base/src/main/java/io/wren/base/metadata/FunctionKey.java
-66
View File
@@ -1,66 +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.wren.base.metadata;
import java.util.Objects;
/**
* TODO: analyze the type of argument expression
* https://github.com/Canner/canner-metric-layer/issues/92
* <p>
* We only support function overloading with different number of argument now. Because
* the work of analyze the type of argument is too huge to implement, FunctionKey only
* recognizes each function by its name and number of argument.
*/
public class FunctionKey
{
public static FunctionKey functionKey(String name, int numArgument)
{
return new FunctionKey(name, numArgument);
}
private final String name;
private final int numArgument;
private FunctionKey(String name, int numArgument)
{
this.name = name;
this.numArgument = numArgument;
}
public String getName()
{
return name;
}
@Override
public int hashCode()
{
return Objects.hash(name, numArgument);
}
@Override
public boolean equals(Object o)
{
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
FunctionKey that = (FunctionKey) o;
return numArgument == that.numArgument && Objects.equals(name, that.name);
}
}
wren-base/src/main/java/io/wren/base/metadata/TableMetadata.java
-88
View File
@@ -1,88 +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.wren.base.metadata;
import com.google.common.collect.ImmutableList;
import io.wren.base.type.PGType;
import java.util.List;
import static java.util.Objects.requireNonNull;
public class TableMetadata
{
private final SchemaTableName table;
private final List<ColumnMetadata> columns;
public static Builder builder(SchemaTableName tableName)
{
return new Builder(tableName);
}
public TableMetadata(SchemaTableName table, List<ColumnMetadata> columns)
{
this.table = requireNonNull(table, "table is null");
this.columns = List.copyOf(requireNonNull(columns, "columns is null"));
}
public SchemaTableName getTable()
{
return table;
}
public List<ColumnMetadata> getColumns()
{
return columns;
}
@Override
public String toString()
{
StringBuilder sb = new StringBuilder("ConnectorTableMetadata{");
sb.append("table=").append(table);
sb.append(", columns=").append(columns);
sb.append('}');
return sb.toString();
}
public static class Builder
{
private final SchemaTableName tableName;
private final ImmutableList.Builder<ColumnMetadata> columns = ImmutableList.builder();
private Builder(SchemaTableName tableName)
{
this.tableName = tableName;
}
public Builder column(String columnName, PGType<?> type)
{
columns.add(ColumnMetadata.builder()
.setName(columnName)
.setType(type).build());
return this;
}
public Builder column(ColumnMetadata columnMetadata)
{
columns.add(columnMetadata);
return this;
}
public TableMetadata build()
{
return new TableMetadata(tableName, columns.build());
}
}
}
wren-base/src/main/java/io/wren/base/pgcatalog/function/BigQueryFunctions.java
-80
View File
@@ -1,80 +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.wren.base.pgcatalog.function;
import io.wren.base.type.BigIntType;
import io.wren.base.type.PGArray;
import io.wren.base.type.RecordType;
import io.wren.base.type.VarcharType;
import java.util.List;
import static io.wren.base.metadata.Function.Argument.argument;
import static io.wren.base.pgcatalog.function.PgFunction.Language.SQL;
import static io.wren.base.type.AnyType.ANY;
import static io.wren.base.type.TimestampType.TIMESTAMP;
public final class BigQueryFunctions
{
private BigQueryFunctions() {}
// TODO Support more date/time format https://www.postgresql.org/docs/13/functions-formatting.html#FUNCTIONS-FORMATTING-DATETIME-TABLE
// TODO Support more timezone, now only support UTC
public static final PgFunction PG_TO_CHAR = PgFunction.builder()
.setName("to_char")
.setLanguage(SQL)
.setDefinition("WITH to_char AS (SELECT " +
"CONTAINS_SUBSTR(string_format, 'TZ') as contain_timezone, " +
"CAST(TIMESTAMP(value) AS STRING FORMAT REPLACE(REPLACE(string_format, 'MS', 'FF3'), 'TZ', '')) AS timestamp_with_format) " +
"SELECT CASE WHEN contain_timezone " +
"THEN CONCAT(timestamp_with_format, 'UTC') " +
"ELSE timestamp_with_format " +
"END " +
"FROM to_char")
.setSubquery(true)
.setArguments(List.of(argument("value", TIMESTAMP), argument("string_format", VarcharType.VARCHAR)))
.setReturnType(VarcharType.VARCHAR)
.build();
public static final PgFunction NOW = PgFunction.builder()
.setName("now")
.setLanguage(SQL)
.setDefinition("SELECT CURRENT_DATETIME")
.setReturnType(TIMESTAMP)
.build();
// TODO This is a mock function, need to be implemented
public static final PgFunction PG_EXPANDARRAY = PgFunction.builder()
.setName("_pg_expandarray")
.setLanguage(SQL)
.setDefinition("CASE WHEN (array_length(int_arr) > 0) THEN cast((int_arr[0], 1) as row(x int, n int)) ELSE NULL END")
.setArguments(List.of(argument("int_arr", PGArray.INT4_ARRAY)))
.setReturnType(new RecordType(List.of(BigIntType.BIGINT, BigIntType.BIGINT)))
.build();
// TODO If the input is a string only include number, it will be parsed as a number. So substring('123' from '1') would get the wrong answer '123', actual should be '1'
// https://github.com/Canner/wren/issues/329
public static final PgFunction SUBSTR = PgFunction.builder()
.setName("substr")
.setLanguage(SQL)
.setDefinition("SELECT " +
"CASE WHEN REGEXP_CONTAINS(SAFE_CAST(arg2 AS STRING), '^[0-9]*$') IS TRUE\n" +
"THEN SUBSTR(arg1, CAST(arg2 AS INT64))\n" +
"ELSE REGEXP_EXTRACT(arg1, CAST(arg2 AS STRING))\n" +
"END")
.setArguments(List.of(argument("arg1", VarcharType.VARCHAR), argument("arg2", ANY)))
.setReturnType(VarcharType.VARCHAR)
.build();
}
wren-base/src/main/java/io/wren/base/pgcatalog/function/DataSourceFunctionRegistry.java
-57
View File
@@ -1,57 +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.wren.base.pgcatalog.function;
import com.google.common.collect.ImmutableList;
import io.wren.base.metadata.FunctionKey;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import static io.wren.base.metadata.FunctionKey.functionKey;
public class DataSourceFunctionRegistry
implements FunctionRegistry<PgFunction>
{
private final List<PgFunction> functions = ImmutableList.<PgFunction>builder()
.add(BigQueryFunctions.PG_TO_CHAR)
.add(BigQueryFunctions.NOW)
.add(BigQueryFunctions.SUBSTR)
.build();
private final Map<FunctionKey, PgFunction> simpleNameToFunction = new HashMap<>();
public DataSourceFunctionRegistry()
{
// TODO: handle function name overloading
// https://github.com/Canner/canner-metric-layer/issues/73
// use HashMap to handle multiple same key entries
functions.forEach(function -> simpleNameToFunction.put(functionKey(function.getName(), function.getArguments().map(List::size).orElse(0)), function));
}
@Override
public List<PgFunction> getFunctions()
{
return functions;
}
@Override
public Optional<PgFunction> getFunction(String name, int numArgument)
{
return Optional.ofNullable(simpleNameToFunction.get(functionKey(name, numArgument)));
}
}
wren-base/src/main/java/io/wren/base/pgcatalog/function/DuckDBFunctions.java
-155
View File
@@ -1,155 +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.wren.base.pgcatalog.function;
import com.google.common.collect.ImmutableList;
import io.wren.base.type.BigIntType;
import io.wren.base.type.IntegerType;
import io.wren.base.type.PGArray;
import io.wren.base.type.RecordType;
import io.wren.base.type.VarcharType;
import java.util.List;
import static io.wren.base.metadata.Function.Argument.argument;
import static io.wren.base.type.BooleanType.BOOLEAN;
public class DuckDBFunctions
{
private DuckDBFunctions() {}
private static final String NULL = "NULL";
public static final PgFunction CURRENT_DATABASE = PgFunction.builder()
.setName("current_database")
.setLanguage(PgFunction.Language.SQL)
.setImplemented(true)
.build();
public static final PgFunction CURRENT_SCHEMAS = PgFunction.builder()
.setName("current_schemas")
.setLanguage(PgFunction.Language.SQL)
.setArguments(ImmutableList.of(argument("includeImplicit", BOOLEAN)))
.setImplemented(true)
.build();
public static final PgFunction PG_RELATION_SIZE__INT_VARCHAR___BIGINT = PgFunction.builder()
.setName("pg_relation_size")
.setLanguage(PgFunction.Language.SQL)
.setDefinition(NULL)
.setArguments(ImmutableList.of(argument("relOid", IntegerType.INTEGER), argument("text", VarcharType.VARCHAR)))
.setReturnType(BigIntType.BIGINT)
.build();
// It's an overloading of PG_RELATION_SIZE__INT_VARCHAR___BIGINT, no need to implement it.
public static final PgFunction PG_RELATION_SIZE__INT___BIGINT = PgFunction.builder()
.setName("pg_relation_size")
.setLanguage(PgFunction.Language.SQL)
.setArguments(ImmutableList.of(argument("relOid", IntegerType.INTEGER)))
.setReturnType(BigIntType.BIGINT)
.setImplemented(true)
.build();
public static final PgFunction ARRAY_IN = PgFunction.builder()
.setName("array_in")
.setLanguage(PgFunction.Language.SQL)
.setDefinition(NULL)
.setArguments(ImmutableList.of(argument("ignored", VarcharType.VARCHAR)))
.setReturnType(PGArray.VARCHAR_ARRAY)
.build();
public static final PgFunction ARRAY_OUT = PgFunction.builder()
.setName("array_out")
.setLanguage(PgFunction.Language.SQL)
.setDefinition(NULL)
.setArguments(ImmutableList.of(argument("ignored", PGArray.VARCHAR_ARRAY)))
.setReturnType(VarcharType.VARCHAR)
.build();
public static final PgFunction ARRAY_RECV = PgFunction.builder()
.setName("array_recv")
.setLanguage(PgFunction.Language.SQL)
.setDefinition(NULL)
.setArguments(ImmutableList.of(argument("ignored", VarcharType.VARCHAR)))
.setReturnType(PGArray.VARCHAR_ARRAY)
.build();
public static final PgFunction ARRAY_UPPER = PgFunction.builder()
.setName("array_upper")
.setLanguage(PgFunction.Language.SQL)
.setDefinition("CASE WHEN dim = 1 THEN array_length(input) ELSE NULL END")
.setArguments(ImmutableList.of(argument("input", PGArray.VARCHAR_ARRAY), argument("dim", BigIntType.BIGINT)))
.setReturnType(IntegerType.INTEGER)
.build();
public static final PgFunction PG_GET_EXPR = PgFunction.builder()
.setName("pg_get_expr")
.setLanguage(PgFunction.Language.SQL)
.setDefinition("''")
.setArguments(List.of(argument("pg_node", VarcharType.VARCHAR), argument("relation", IntegerType.INTEGER)))
.setReturnType(VarcharType.VARCHAR)
.setImplemented(true)
.build();
public static final PgFunction PG_GET_EXPR_PRETTY = PgFunction.builder()
.setName("pg_get_expr")
.setLanguage(PgFunction.Language.SQL)
.setDefinition("SELECT ''")
.setArguments(List.of(argument("pg_node", VarcharType.VARCHAR), argument("relation", IntegerType.INTEGER), argument("pretty", BOOLEAN)))
.setReturnType(VarcharType.VARCHAR)
.setImplemented(true)
.build();
public static final PgFunction FORMAT_TYPE = PgFunction.builder()
.setName("format_type")
.setLanguage(PgFunction.Language.SQL)
.setDefinition("(select format_pg_type(logical_type, t.type_name) from duckdb_types() t where t.type_name=lower(tname)) || case when typemod>0 then concat('(', typemod//1000, ',', typemod%1000, ')') else '' end")
.setArguments(List.of(argument("tname", VarcharType.VARCHAR), argument("typemod", IntegerType.INTEGER)))
.setReturnType(VarcharType.VARCHAR)
.build();
public static final PgFunction PG_GET_FUNCTION_RESULT = PgFunction.builder()
.setName("pg_get_function_result")
.setLanguage(PgFunction.Language.SQL)
.setDefinition("SELECT prorettype FROM pg_proc WHERE oid = func")
.setArguments(List.of(argument("func", IntegerType.INTEGER)))
.setReturnType(VarcharType.VARCHAR)
.build();
// TODO This is a mock function, need to be implemented
public static final PgFunction PG_EXPANDARRAY = PgFunction.builder()
.setName("_pg_expandarray")
.setLanguage(PgFunction.Language.SQL)
.setDefinition("CASE WHEN (array_length(int_arr) > 0) THEN cast((int_arr[0], 1) as row(x int, n int)) ELSE NULL END")
.setArguments(List.of(argument("int_arr", PGArray.INT4_ARRAY)))
.setReturnType(new RecordType(List.of(BigIntType.BIGINT, BigIntType.BIGINT)))
.build();
public static final PgFunction REGEXP_LIKE = PgFunction.builder()
.setName("regexp_like")
.setLanguage(PgFunction.Language.SQL)
.setDefinition("SELECT regexp_matches(arg1, arg2)")
.setArguments(List.of(argument("arg1", VarcharType.VARCHAR), argument("arg2", VarcharType.VARCHAR)))
.setReturnType(BOOLEAN)
.build();
public static final PgFunction GENERATE_ARRAY = PgFunction.builder()
.setName("generate_array")
.setLanguage(PgFunction.Language.SQL)
.setDefinition("generate_series(start, stop)")
.setArguments(List.of(argument("start", BigIntType.BIGINT), argument("stop", BigIntType.BIGINT)))
.setReturnType(PGArray.INT4_ARRAY)
.build();
}
wren-base/src/main/java/io/wren/base/pgcatalog/function/FunctionRegistry.java
-27
View File
@@ -1,27 +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.wren.base.pgcatalog.function;
import io.wren.base.metadata.Function;
import java.util.List;
import java.util.Optional;
public interface FunctionRegistry<F extends Function>
{
List<F> getFunctions();
Optional<F> getFunction(String name, int numArgument);
}
wren-base/src/main/java/io/wren/base/pgcatalog/function/PgFunction.java
-180
View File
@@ -1,180 +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.wren.base.pgcatalog.function;
import com.google.common.base.Joiner;
import io.wren.base.metadata.Function;
import io.wren.base.type.PGType;
import java.util.List;
import java.util.regex.Pattern;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static java.lang.String.format;
import static java.util.Objects.requireNonNull;
public class PgFunction
extends Function
{
public static final Pattern PG_FUNCTION_PATTERN = Pattern.compile("(?<functionName>[a-zA-Z]+(_[a-zA-Z0-9]+)*)(__(?<argsType>[a-zA-Z]+(_[a-zA-Z0-9]+)*))?(___(?<returnType>[a-zA-Z]+(_[a-zA-Z0-9]+)*))?");
public enum Language
{
SQL,
JS
}
public static Builder builder()
{
return new Builder();
}
private final Language language;
private final String definition;
private final boolean subquery;
// if the function is implemented in the database
private final boolean implemented;
public PgFunction(
String name,
Language language,
List<Argument> arguments,
PGType returnType,
String definition,
boolean subquery,
boolean implemented)
{
super(name, arguments, returnType);
this.language = language;
this.definition = definition;
this.subquery = subquery;
this.implemented = implemented;
}
/**
* 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 String getDefinition()
{
return definition;
}
public boolean isSubquery()
{
return subquery;
}
public boolean isImplemented()
{
return implemented;
}
@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 Builder
extends Function.Builder
{
private String name;
private Language language;
private String definition;
private List<Argument> arguments;
private PGType returnType;
private boolean subquery;
private boolean implemented;
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<Argument> arguments)
{
this.arguments = arguments;
return this;
}
public Builder setReturnType(PGType returnType)
{
this.returnType = returnType;
return this;
}
public Builder setSubquery(boolean subquery)
{
this.subquery = subquery;
return this;
}
public Builder setImplemented(boolean implemented)
{
this.implemented = implemented;
return this;
}
public PgFunction build()
{
requireNonNull(name, "name is null");
requireNonNull(language, "language is null");
return new PgFunction(name, language, arguments, returnType, definition, subquery, implemented);
}
}
}
wren-base/src/main/java/io/wren/base/pgcatalog/function/PgMetastoreFunctionRegistry.java
-84
View File
@@ -1,84 +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.wren.base.pgcatalog.function;
import com.google.common.collect.ImmutableList;
import io.wren.base.metadata.FunctionKey;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import static io.wren.base.metadata.FunctionKey.functionKey;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.ARRAY_IN;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.ARRAY_OUT;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.ARRAY_RECV;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.ARRAY_UPPER;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.CURRENT_DATABASE;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.CURRENT_SCHEMAS;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.FORMAT_TYPE;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.GENERATE_ARRAY;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.PG_EXPANDARRAY;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.PG_GET_EXPR;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.PG_GET_EXPR_PRETTY;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.PG_GET_FUNCTION_RESULT;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.PG_RELATION_SIZE__INT_VARCHAR___BIGINT;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.PG_RELATION_SIZE__INT___BIGINT;
import static io.wren.base.pgcatalog.function.DuckDBFunctions.REGEXP_LIKE;
public final class PgMetastoreFunctionRegistry
implements FunctionRegistry<PgFunction>
{
public final List<PgFunction> functions = ImmutableList.<PgFunction>builder()
.add(CURRENT_DATABASE)
.add(CURRENT_SCHEMAS)
.add(PG_RELATION_SIZE__INT___BIGINT)
.add(PG_RELATION_SIZE__INT_VARCHAR___BIGINT)
.add(ARRAY_IN)
.add(ARRAY_OUT)
.add(ARRAY_RECV)
.add(ARRAY_UPPER)
.add(FORMAT_TYPE)
.add(PG_GET_FUNCTION_RESULT)
.add(REGEXP_LIKE)
.add(GENERATE_ARRAY)
.add(PG_EXPANDARRAY)
.add(PG_GET_EXPR)
.add(PG_GET_EXPR_PRETTY)
.build();
private final Map<FunctionKey, PgFunction> simpleNameToFunction = new HashMap<>();
public PgMetastoreFunctionRegistry()
{
// TODO: handle function name overloading
// https://github.com/Canner/canner-metric-layer/issues/73
// use HashMap to handle multiple same key entries
functions.forEach(function -> simpleNameToFunction.put(functionKey(function.getName(), function.getArguments().map(List::size).orElse(0)), function));
}
@Override
public List<PgFunction> getFunctions()
{
return functions;
}
@Override
public Optional<PgFunction> getFunction(String name, int numArgument)
{
return Optional.ofNullable(simpleNameToFunction.get(functionKey(name, numArgument)));
}
}
wren-base/src/main/java/io/wren/base/sqlrewrite/analyzer/ExpressionTypeAnalyzer.java
-405
View File
@@ -1,405 +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.wren.base.sqlrewrite.analyzer;
import com.google.common.collect.ImmutableList;
import io.trino.sql.tree.ArithmeticBinaryExpression;
import io.trino.sql.tree.ArrayConstructor;
import io.trino.sql.tree.BetweenPredicate;
import io.trino.sql.tree.BinaryLiteral;
import io.trino.sql.tree.BooleanLiteral;
import io.trino.sql.tree.Cast;
import io.trino.sql.tree.CharLiteral;
import io.trino.sql.tree.ComparisonExpression;
import io.trino.sql.tree.CurrentCatalog;
import io.trino.sql.tree.CurrentPath;
import io.trino.sql.tree.CurrentSchema;
import io.trino.sql.tree.CurrentTime;
import io.trino.sql.tree.CurrentUser;
import io.trino.sql.tree.DateTimeDataType;
import io.trino.sql.tree.DecimalLiteral;
import io.trino.sql.tree.DefaultTraversalVisitor;
import io.trino.sql.tree.DereferenceExpression;
import io.trino.sql.tree.DoubleLiteral;
import io.trino.sql.tree.ExistsPredicate;
import io.trino.sql.tree.Expression;
import io.trino.sql.tree.FunctionCall;
import io.trino.sql.tree.GenericDataType;
import io.trino.sql.tree.GenericLiteral;
import io.trino.sql.tree.Identifier;
import io.trino.sql.tree.InPredicate;
import io.trino.sql.tree.IntervalDayTimeDataType;
import io.trino.sql.tree.IntervalLiteral;
import io.trino.sql.tree.IsNotNullPredicate;
import io.trino.sql.tree.IsNullPredicate;
import io.trino.sql.tree.LikePredicate;
import io.trino.sql.tree.LongLiteral;
import io.trino.sql.tree.NullLiteral;
import io.trino.sql.tree.QualifiedName;
import io.trino.sql.tree.QuantifiedComparisonExpression;
import io.trino.sql.tree.Row;
import io.trino.sql.tree.RowDataType;
import io.trino.sql.tree.StringLiteral;
import io.trino.sql.tree.SubscriptExpression;
import io.trino.sql.tree.TimeLiteral;
import io.trino.sql.tree.TimestampLiteral;
import io.wren.base.WrenMDL;
import io.wren.base.metadata.Function;
import io.wren.base.metadata.FunctionBundle;
import io.wren.base.type.BigIntType;
import io.wren.base.type.BooleanType;
import io.wren.base.type.ByteaType;
import io.wren.base.type.DateType;
import io.wren.base.type.DoubleType;
import io.wren.base.type.IntervalType;
import io.wren.base.type.NumericType;
import io.wren.base.type.PGType;
import io.wren.base.type.PGTypes;
import io.wren.base.type.RecordType;
import io.wren.base.type.TimestampType;
import io.wren.base.type.VarcharType;
import java.util.Optional;
import static io.trino.sql.tree.DereferenceExpression.getQualifiedName;
import static java.util.Objects.requireNonNull;
public class ExpressionTypeAnalyzer
extends DefaultTraversalVisitor<Void>
{
public static PGType<?> analyze(WrenMDL mdl, Scope scope, Expression expression)
{
ExpressionTypeAnalyzer analyzer = new ExpressionTypeAnalyzer(mdl, scope);
analyzer.process(expression);
return analyzer.result;
}
private final WrenMDL mdl;
private final Scope scope;
private PGType<?> result;
public ExpressionTypeAnalyzer(WrenMDL mdl, Scope scope)
{
this.mdl = requireNonNull(mdl, "mdl is null");
this.scope = requireNonNull(scope, "scope is null");
}
@Override
protected Void visitStringLiteral(StringLiteral node, Void context)
{
result = VarcharType.VARCHAR;
return null;
}
@Override
protected Void visitDoubleLiteral(DoubleLiteral node, Void context)
{
result = DoubleType.DOUBLE;
return null;
}
@Override
protected Void visitDecimalLiteral(DecimalLiteral node, Void context)
{
result = NumericType.NUMERIC;
return null;
}
@Override
protected Void visitGenericLiteral(GenericLiteral node, Void context)
{
PGTypes.nameToPgType(node.getType())
.ifPresent(pgType -> result = pgType);
return null;
}
@Override
protected Void visitTimeLiteral(TimeLiteral node, Void context)
{
// TODO: we don't support time type yet, so we treat it as timestamp type.
result = TimestampType.TIMESTAMP;
return null;
}
@Override
protected Void visitTimestampLiteral(TimestampLiteral node, Void context)
{
// TODO: timestamp literal may contain timezone, we need to handle it.
result = TimestampType.TIMESTAMP;
return null;
}
@Override
protected Void visitIntervalLiteral(IntervalLiteral node, Void context)
{
result = IntervalType.INTERVAL;
return null;
}
@Override
protected Void visitCharLiteral(CharLiteral node, Void context)
{
result = VarcharType.VARCHAR;
return null;
}
@Override
protected Void visitBinaryLiteral(BinaryLiteral node, Void context)
{
result = ByteaType.BYTEA;
return null;
}
@Override
protected Void visitBooleanLiteral(BooleanLiteral node, Void context)
{
result = BooleanType.BOOLEAN;
return null;
}
@Override
protected Void visitLongLiteral(LongLiteral node, Void context)
{
result = BigIntType.BIGINT;
return null;
}
@Override
protected Void visitNullLiteral(NullLiteral node, Void context)
{
return super.visitNullLiteral(node, context);
}
@Override
protected Void visitCast(Cast node, Void context)
{
process(node.getType());
// The type is the final output. We don't need to dig into the expression.
return null;
}
@Override
protected Void visitRowDataType(RowDataType node, Void context)
{
result = RecordType.EMPTY_RECORD;
return null;
}
@Override
protected Void visitDateTimeType(DateTimeDataType node, Void context)
{
// TODO: it may contain timezone, we need to handle it.
result = TimestampType.TIMESTAMP;
return null;
}
@Override
protected Void visitIntervalDataType(IntervalDayTimeDataType node, Void context)
{
result = IntervalType.INTERVAL;
return null;
}
@Override
protected Void visitGenericDataType(GenericDataType node, Void context)
{
PGTypes.nameToPgType(node.getName().getValue()).ifPresent(pgType -> result = pgType);
return null;
}
@Override
protected Void visitInPredicate(InPredicate node, Void context)
{
result = BooleanType.BOOLEAN;
return null;
}
@Override
protected Void visitLikePredicate(LikePredicate node, Void context)
{
result = BooleanType.BOOLEAN;
return null;
}
@Override
protected Void visitBetweenPredicate(BetweenPredicate node, Void context)
{
result = BooleanType.BOOLEAN;
return null;
}
@Override
protected Void visitIsNotNullPredicate(IsNotNullPredicate node, Void context)
{
result = BooleanType.BOOLEAN;
return null;
}
@Override
protected Void visitIsNullPredicate(IsNullPredicate node, Void context)
{
result = BooleanType.BOOLEAN;
return null;
}
@Override
protected Void visitExists(ExistsPredicate node, Void context)
{
result = BooleanType.BOOLEAN;
return null;
}
@Override
protected Void visitComparisonExpression(ComparisonExpression node, Void context)
{
result = BooleanType.BOOLEAN;
return null;
}
@Override
protected Void visitQuantifiedComparisonExpression(QuantifiedComparisonExpression node, Void context)
{
// TODO: we don't support quantified comparison yet.
return null;
}
@Override
protected Void visitFunctionCall(FunctionCall node, Void context)
{
FunctionBundle.getFunction(node.getName().getSuffix(), node.getArguments().size())
.flatMap(Function::getReturnType)
.ifPresent(type -> result = type);
// TODO: handle the remote name
if (node.getName().getSuffix().equalsIgnoreCase("now") ||
node.getName().getSuffix().equalsIgnoreCase("now___timestamp")) {
result = TimestampType.TIMESTAMP;
}
return null;
}
@Override
protected Void visitDereferenceExpression(DereferenceExpression node, Void context)
{
QualifiedName qualifiedName = getQualifiedName(node);
if (qualifiedName != null) {
Optional<Field> fieldOptional = scope.getRelationType().getFields().stream()
.filter(field -> field.canResolve(qualifiedName))
.findFirst();
fieldOptional.ifPresent(field -> result = getColumnType(field));
}
return null;
}
@Override
protected Void visitIdentifier(Identifier node, Void context)
{
QualifiedName qualifiedName = QualifiedName.of(ImmutableList.of(node.getValue()));
scope.getRelationType().getFields().stream()
.filter(field -> field.canResolve(qualifiedName))
.findFirst()
.ifPresent(field -> result = getColumnType(field));
return null;
}
private PGType<?> getColumnType(Field field)
{
String objectName = field.getTableName().getSchemaTableName().getTableName();
String columnName = field.getColumnName();
// TODO: support to analyze the column type of CTE.
// It could be a remote table or a custom CTE.
if (!mdl.isObjectExist(objectName)) {
return null;
}
return mdl.getColumnType(objectName, columnName).flatMap(PGTypes::nameToPgType).orElse(null);
}
@Override
protected Void visitRow(Row node, Void context)
{
result = RecordType.EMPTY_RECORD;
return null;
}
@Override
protected Void visitSubscriptExpression(SubscriptExpression node, Void context)
{
process(node.getBase(), context);
if (result != null) {
result = PGTypes.getArrayType(result.oid());
}
return null;
}
@Override
protected Void visitArithmeticBinary(ArithmeticBinaryExpression node, Void context)
{
// TODO: check the type coercion rule. For now, we just use the left type.
process(node.getLeft());
return null;
}
@Override
protected Void visitCurrentTime(CurrentTime node, Void context)
{
if (node.getFunction().equals(CurrentTime.Function.DATE)) {
result = DateType.DATE;
}
else {
result = TimestampType.TIMESTAMP;
}
return null;
}
@Override
protected Void visitCurrentUser(CurrentUser node, Void context)
{
result = VarcharType.VARCHAR;
return null;
}
@Override
protected Void visitCurrentSchema(CurrentSchema node, Void context)
{
result = VarcharType.VARCHAR;
return null;
}
@Override
protected Void visitCurrentCatalog(CurrentCatalog node, Void context)
{
result = VarcharType.VARCHAR;
return null;
}
@Override
protected Void visitCurrentPath(CurrentPath node, Void context)
{
result = VarcharType.VARCHAR;
return null;
}
@Override
protected Void visitArrayConstructor(ArrayConstructor node, Void context)
{
// ALl value should be same type in array, we only check first value type here.
process(node.getValues().get(0));
if (result != null) {
result = PGTypes.getArrayType(result.oid());
}
return null;
}
}
wren-base/src/main/java/io/wren/base/sqlrewrite/analyzer/StatementAnalyzer.java
+4 -32
View File
@@ -57,8 +57,6 @@ import io.wren.base.dto.Model;
import io.wren.base.dto.TimeUnit;
import io.wren.base.dto.View;
import javax.annotation.Nullable;
import java.util.List;
import java.util.Optional;
import java.util.Set;
@@ -84,13 +82,7 @@ public final class StatementAnalyzer
public static Scope analyze(Analysis analysis, Statement statement, SessionContext sessionContext, WrenMDL wrenMDL)
{
return analyze(analysis, statement, sessionContext, wrenMDL, null);
}
public static Scope analyze(Analysis analysis, Statement statement, SessionContext sessionContext, WrenMDL wrenMDL, TypeCoercion typeCoercion)
{
Scope queryScope = new Visitor(sessionContext, analysis, wrenMDL, typeCoercion).process(statement, Optional.empty());
Scope queryScope = new Visitor(sessionContext, analysis, wrenMDL).process(statement, Optional.empty());
// add models directly used in sql query
analysis.addModels(
wrenMDL.listModels().stream()
@@ -137,18 +129,15 @@ public final class StatementAnalyzer
private final SessionContext sessionContext;
private final Analysis analysis;
private final WrenMDL wrenMDL;
private final Optional<TypeCoercion> typeCoercionOptional;
public Visitor(
SessionContext sessionContext,
Analysis analysis,
WrenMDL wrenMDL,
@Nullable TypeCoercion typeCoercion)
WrenMDL wrenMDL)
{
this.sessionContext = requireNonNull(sessionContext, "sessionContext is null");
this.analysis = requireNonNull(analysis, "analysis is null");
this.wrenMDL = requireNonNull(wrenMDL, "wrenMDL is null");
this.typeCoercionOptional = Optional.ofNullable(typeCoercion);
}
public Scope process(Node node)
@@ -214,7 +203,7 @@ public final class StatementAnalyzer
{
Query query = withQuery.getQuery();
Analysis analyzed = new Analysis(query);
Optional<Scope> queryScope = Optional.ofNullable(analyze(analyzed, query, sessionContext, wrenMDL, typeCoercionOptional.orElse(null)));
Optional<Scope> queryScope = Optional.ofNullable(analyze(analyzed, query, sessionContext, wrenMDL));
List<Field> fields;
Optional<List<Identifier>> columnNames = withQuery.getColumnNames();
if (columnNames.isPresent()) {
@@ -395,11 +384,6 @@ public final class StatementAnalyzer
analysis.addRequiredSourceNode(scope.getRelationId().getSourceNode()
.orElseThrow(() -> new IllegalArgumentException("count(*) should have a followed source")));
}
typeCoercionOptional.ifPresent(typeCoercion -> {
Optional<Expression> coerced = typeCoercion.coerceExpression(singleColumn.getExpression(), scope);
coerced.ifPresent(expression -> analysis.addTypeCoercion(NodeRef.of(singleColumn.getExpression()), expression));
});
}
private Scope analyzeFrom(QuerySpecification node, Optional<Scope> scope)
@@ -413,8 +397,6 @@ public final class StatementAnalyzer
private void analyzeWhere(Expression node, Scope scope)
{
ExpressionAnalysis expressionAnalysis = analyzeExpression(node, scope);
typeCoercionOptional.flatMap(typeCoercion -> typeCoercion.coerceExpression(node, scope))
.ifPresent(expression -> analysis.addTypeCoercion(NodeRef.of(node), expression));
}
private void analyzeWindowSpecification(WindowSpecification windowSpecification, Scope scope)
@@ -522,16 +504,6 @@ public final class StatementAnalyzer
case JoinOn joinOn:
Expression expression = joinOn.getExpression();
analyzeExpression(expression, outputScope);
typeCoercionOptional.ifPresent(typeCoercion -> {
Optional<Expression> coerced = typeCoercion.coerceExpression(expression, outputScope);
if (coerced.isPresent()) {
JoinOn newJoinOn = new JoinOn(coerced.get());
Join newJoin = node.getLocation().isPresent() ?
new Join(node.getLocation().get(), node.getType(), node.getLeft(), node.getRight(), Optional.of(newJoinOn)) :
new Join(node.getType(), node.getLeft(), node.getRight(), Optional.of(newJoinOn));
analysis.addTypeCoercion(NodeRef.of(node), newJoin);
}
});
break;
case JoinUsing joinUsing:
joinUsing.getColumns().forEach(column -> analyzeExpression(column, outputScope));
@@ -574,7 +546,7 @@ public final class StatementAnalyzer
@Override
protected Scope visitTableSubquery(TableSubquery node, Optional<Scope> scope)
{
return Optional.ofNullable(analyze(analysis, node.getQuery(), sessionContext, wrenMDL, typeCoercionOptional.orElse(null)))
return Optional.ofNullable(analyze(analysis, node.getQuery(), sessionContext, wrenMDL))
.map(value -> createAndAssignScope(node, scope, value))
.orElseGet(() -> Scope.builder().parent(scope).build());
}
wren-base/src/main/java/io/wren/base/sqlrewrite/analyzer/TypeCoercion.java
-24
View File
@@ -1,24 +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.wren.base.sqlrewrite.analyzer;
import io.trino.sql.tree.Expression;
import java.util.Optional;
public interface TypeCoercion
{
Optional<Expression> coerceExpression(Expression expression, Scope scope);
}
wren-base/src/main/java/io/wren/base/type/AnyType.java
-86
View File
@@ -1,86 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import java.nio.charset.StandardCharsets;
public class AnyType
extends PGType<Integer>
{
// we represent the any data type as integer,
// because the postgresql typalign of any is integer.
public static final AnyType ANY = new AnyType();
private static final int TYPE_LEN = 4;
static final int OID = 2276;
private AnyType()
{
super(OID, TYPE_LEN, -1, "any");
}
@Override
public int typArray()
{
return 0;
}
@Override
public String type()
{
return Type.PSEUDO.code();
}
@Override
public String typeCategory()
{
return TypeCategory.PSEUDO.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull Integer value)
{
buffer.writeInt(TYPE_LEN);
buffer.writeInt(value);
return INT32_BYTE_SIZE + TYPE_LEN;
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull Integer value)
{
return Integer.toString(value).getBytes(StandardCharsets.UTF_8);
}
@Override
public Integer readBinaryValue(ByteBuf buffer, int valueLength)
{
return buffer.readInt();
}
@Override
public Integer decodeUTF8Text(byte[] bytes)
{
return Integer.parseInt(new String(bytes, StandardCharsets.UTF_8));
}
@Override
public Object getEmptyValue()
{
return 0;
}
}
wren-base/src/main/java/io/wren/base/type/BaseTimestampType.java
-61
View File
@@ -1,61 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import static com.google.common.base.Preconditions.checkArgument;
abstract class BaseTimestampType
extends PGType
{
protected static final int TYPE_LEN = 8;
protected static final int TYPE_MOD = -1;
BaseTimestampType(int oid, int typeLen, int typeMod, @Nonnull String typeName)
{
super(oid, typeLen, typeMod, typeName);
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull Object value)
{
buffer.writeInt(TYPE_LEN);
buffer.writeLong(PgDatetimeUtils.toPgTimestamp((long) value));
return INT32_BYTE_SIZE + TYPE_LEN;
}
@Override
public String typeCategory()
{
return TypeCategory.DATETIME.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public Object readBinaryValue(ByteBuf buffer, int valueLength)
{
checkArgument(valueLength == TYPE_LEN, "valueLength must be %s because timestamp is a 64 bit long. Actual length is %s", TYPE_LEN, valueLength);
long microSecondsSince2K = buffer.readLong();
return PgDatetimeUtils.toTrinoTimestamp(microSecondsSince2K);
}
}
wren-base/src/main/java/io/wren/base/type/BigIntType.java
-89
View File
@@ -1,89 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import java.nio.charset.StandardCharsets;
import static com.google.common.base.Preconditions.checkArgument;
public class BigIntType
extends PGType<Long>
{
public static final BigIntType BIGINT = new BigIntType();
static final int OID = 20;
private static final int TYPE_LEN = 8;
private static final int TYPE_MOD = -1;
private BigIntType()
{
super(OID, TYPE_LEN, TYPE_MOD, "int8");
}
@Override
public int typArray()
{
return PGArray.INT8_ARRAY.oid();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull Long value)
{
buffer.writeInt(TYPE_LEN);
buffer.writeLong(value);
return INT32_BYTE_SIZE + TYPE_LEN;
}
@Override
public String typeCategory()
{
return TypeCategory.NUMERIC.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull Long value)
{
return Long.toString(value).getBytes(StandardCharsets.UTF_8);
}
@Override
public Long readBinaryValue(ByteBuf buffer, int valueLength)
{
checkArgument(valueLength == TYPE_LEN, "length should be %s because long is int64. Actual length: %s", TYPE_LEN, valueLength);
return buffer.readLong();
}
@Override
public Long decodeUTF8Text(byte[] bytes)
{
return Long.parseLong(new String(bytes, StandardCharsets.UTF_8));
}
@Override
public Object getEmptyValue()
{
return 0;
}
}
wren-base/src/main/java/io/wren/base/type/BooleanType.java
-110
View File
@@ -1,110 +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.wren.base.type;
import com.google.common.collect.ImmutableSet;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import java.nio.ByteBuffer;
import java.util.Collection;
import static com.google.common.base.Preconditions.checkArgument;
public class BooleanType
extends PGType<Boolean>
{
public static final PGType BOOLEAN = new BooleanType();
static final int OID = 16;
private static final int TYPE_LEN = 1;
private static final int TYPE_MOD = -1;
private static final byte[] TEXT_TRUE = new byte[] {'t'};
private static final byte[] TEXT_FALSE = new byte[] {'f'};
private static final Collection<ByteBuffer> TRUTH_VALUES = ImmutableSet.of(
ByteBuffer.wrap(new byte[] {'1'}),
ByteBuffer.wrap(new byte[] {'t'}),
ByteBuffer.wrap(new byte[] {'T'}),
ByteBuffer.wrap(new byte[] {'t', 'r', 'u', 'e'}),
ByteBuffer.wrap(new byte[] {'T', 'R', 'U', 'E'}));
private BooleanType()
{
super(OID, TYPE_LEN, TYPE_MOD, "bool");
}
@Override
public int typArray()
{
return PGArray.BOOL_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.NUMERIC.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull Boolean value)
{
byte byteValue = (byte) (value ? 1 : 0);
buffer.writeInt(TYPE_LEN);
buffer.writeByte(byteValue);
return INT32_BYTE_SIZE + TYPE_LEN;
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull Boolean value)
{
return value ? TEXT_TRUE : TEXT_FALSE;
}
@Override
public Boolean readBinaryValue(ByteBuf buffer, int valueLength)
{
checkArgument(valueLength == TYPE_LEN, "length should be %s boolean is just a byte. Actual length: %s", TYPE_LEN, valueLength);
byte value = buffer.readByte();
switch (value) {
case 0:
return false;
case 1:
return true;
default:
throw new IllegalArgumentException("Unsupported binary bool: " + value);
}
}
@Override
public Boolean decodeUTF8Text(byte[] bytes)
{
return TRUTH_VALUES.contains(ByteBuffer.wrap(bytes));
}
@Override
public Object getEmptyValue()
{
return false;
}
}
wren-base/src/main/java/io/wren/base/type/BpCharType.java
-94
View File
@@ -1,94 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import java.nio.charset.StandardCharsets;
import static java.nio.charset.StandardCharsets.UTF_8;
public class BpCharType
extends PGType<String>
{
public static final BpCharType BPCHAR = new BpCharType();
static final int OID = 1042;
private BpCharType()
{
super(OID, -1, -1, "bpchar");
}
@Override
public int typArray()
{
return PGArray.BPCHAR_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.STRING.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull String value)
{
byte[] bytes = value.getBytes(UTF_8);
buffer.writeInt(bytes.length);
buffer.writeBytes(bytes);
return INT32_BYTE_SIZE + bytes.length;
}
@Override
public int writeAsText(ByteBuf buffer, @Nonnull String value)
{
return writeAsBinary(buffer, value);
}
@Override
public String readBinaryValue(ByteBuf buffer, int valueLength)
{
byte[] utf8 = new byte[valueLength];
buffer.readBytes(utf8);
return new String(utf8, StandardCharsets.UTF_8);
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull String value)
{
return value.getBytes(UTF_8);
}
@Override
public String decodeUTF8Text(byte[] bytes)
{
return new String(bytes, UTF_8);
}
@Override
public Object getEmptyValue()
{
return 0;
}
}
wren-base/src/main/java/io/wren/base/type/ByteaType.java
-104
View File
@@ -1,104 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import org.apache.commons.codec.binary.Hex;
import javax.annotation.Nonnull;
import static com.google.common.base.Preconditions.checkArgument;
import static java.nio.charset.StandardCharsets.UTF_8;
public class ByteaType
extends PGType<Object>
{
public static final int OID = 17;
private static final int TYPE_LEN = -1;
private static final int TYPE_MOD = -1;
public static final ByteaType BYTEA = new ByteaType();
private ByteaType()
{
super(OID, TYPE_LEN, TYPE_MOD, "bytea");
}
@Override
public int typArray()
{
return PGArray.BYTEA_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.USER_DEFINED_TYPES.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, Object value)
{
byte[] bytes = encodeHexString((byte[]) value).getBytes(UTF_8);
buffer.writeInt(bytes.length);
buffer.writeBytes(bytes);
return INT32_BYTE_SIZE + bytes.length;
}
@Override
public int writeAsText(ByteBuf buffer, @Nonnull Object value)
{
return writeAsBinary(buffer, value);
}
@Override
public Object readBinaryValue(ByteBuf buffer, int valueLength)
{
checkArgument(valueLength >= 1, "The length of bytea should be 1 at least.");
byte[] bytes = new byte[valueLength];
buffer.readBytes(bytes);
return bytes;
}
@Override
public byte[] encodeAsUTF8Text(Object value)
{
String strBuilder = "\\" + encodeHexString((byte[]) value);
return strBuilder.getBytes(UTF_8);
}
@Override
public Object decodeUTF8Text(byte[] bytes)
{
return bytes;
}
private String encodeHexString(byte[] decimalValue)
{
return "\\x" + Hex.encodeHexString(decimalValue);
}
@Override
public Object getEmptyValue()
{
return new byte[0];
}
}
wren-base/src/main/java/io/wren/base/type/CharType.java
-85
View File
@@ -1,85 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import static com.google.common.base.Preconditions.checkArgument;
import static java.nio.charset.StandardCharsets.UTF_8;
public class CharType
extends PGType<String>
{
public static final CharType CHAR = new CharType();
static final int OID = 18;
private CharType()
{
super(OID, 1, -1, "char");
}
@Override
public int typArray()
{
return PGArray.CHAR_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.STRING.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull String value)
{
buffer.writeInt(1);
buffer.writeBytes(value.getBytes(UTF_8));
return 5;
}
@Override
public String readBinaryValue(ByteBuf buffer, int valueLength)
{
checkArgument(valueLength == 1, "The length of char should be 1");
return new String(buffer.readBytes(valueLength).array(), UTF_8);
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull String value)
{
return value.getBytes(UTF_8);
}
@Override
public String decodeUTF8Text(byte[] bytes)
{
return new String(bytes, UTF_8);
}
@Override
public Object getEmptyValue()
{
return "";
}
}
wren-base/src/main/java/io/wren/base/type/DateType.java
-104
View File
@@ -1,104 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import java.time.LocalDate;
import java.time.format.DateTimeFormatter;
import java.time.format.DateTimeFormatterBuilder;
import java.time.format.ResolverStyle;
import java.util.Locale;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.time.format.DateTimeFormatter.ISO_LOCAL_DATE;
public class DateType
extends PGType<LocalDate>
{
public static final PGType DATE = new DateType();
private static final int OID = 1082;
private static final String NAME = "date";
private static final int TYPE_LEN = 4;
private static final int TYPE_MOD = -1;
private static final DateTimeFormatter ISO_FORMATTER = new DateTimeFormatterBuilder()
.parseCaseInsensitive()
.append(ISO_LOCAL_DATE)
.toFormatter(Locale.ENGLISH).withResolverStyle(ResolverStyle.STRICT);
private static final DateTimeFormatter ISO_FORMATTER_AD = new DateTimeFormatterBuilder()
.parseCaseInsensitive()
.appendPattern("yyyy-MM-dd")
.toFormatter(Locale.ENGLISH)
.withResolverStyle(ResolverStyle.STRICT);
private DateType()
{
super(OID, TYPE_LEN, TYPE_MOD, NAME);
}
@Override
public String typeCategory()
{
return TypeCategory.DATETIME.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int typArray()
{
return PGArray.DATE_ARRAY.oid();
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull LocalDate value)
{
return value.format(ISO_FORMATTER_AD).getBytes(UTF_8);
}
@Override
public LocalDate decodeUTF8Text(byte[] bytes)
{
String s = new String(bytes, UTF_8);
return LocalDate.parse(s, ISO_FORMATTER);
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull LocalDate value)
{
throw new UnsupportedOperationException();
}
@Override
public LocalDate readBinaryValue(ByteBuf buffer, int valueLength)
{
throw new UnsupportedOperationException();
}
@Override
public Object getEmptyValue()
{
return "1970-01-01";
}
}
wren-base/src/main/java/io/wren/base/type/DoubleType.java
-89
View File
@@ -1,89 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import java.nio.charset.StandardCharsets;
import static com.google.common.base.Preconditions.checkArgument;
public class DoubleType
extends PGType<Double>
{
public static final DoubleType DOUBLE = new DoubleType();
static final int OID = 701;
private static final int TYPE_LEN = 8;
private static final int TYPE_MOD = -1;
private DoubleType()
{
super(OID, TYPE_LEN, TYPE_MOD, "float8");
}
@Override
public int typArray()
{
return PGArray.FLOAT8_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.NUMERIC.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull Double value)
{
buffer.writeInt(TYPE_LEN);
buffer.writeDouble(value);
return INT32_BYTE_SIZE + TYPE_LEN;
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull Double value)
{
return Double.toString(value).getBytes(StandardCharsets.UTF_8);
}
@Override
public Double readBinaryValue(ByteBuf buffer, int valueLength)
{
checkArgument(valueLength == TYPE_LEN, "length should be %s because double is int64. Actual length: %s", TYPE_LEN, valueLength);
return buffer.readDouble();
}
@Override
public Double decodeUTF8Text(byte[] bytes)
{
return Double.parseDouble(new String(bytes, StandardCharsets.UTF_8));
}
@Override
public Object getEmptyValue()
{
return 0.0;
}
}
wren-base/src/main/java/io/wren/base/type/HstoreType.java
-131
View File
@@ -1,131 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import java.util.Map;
import static java.nio.charset.StandardCharsets.UTF_8;
/**
* A Hashmap type in PostgreSQL which keys and values are simple text value.
* We used it to represent the presto Map type.
* https://www.postgresql.org/docs/current/hstore.html
* <p>
* TODO: handle non-text keys and value.
* Because presto allows map with non-text key and value, we should handle this case.
*/
public class HstoreType
extends PGType<Map<Object, Object>>
{
public static final HstoreType HSTORE = new HstoreType();
HstoreType()
{
super(57640, -1, -1, "hstore");
}
@Override
public int typArray()
{
return PGArray.HSTORE_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.USER_DEFINED_TYPES.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
/**
* https://github.com/postgres/postgres/blob/master/contrib/hstore/hstore_io.c (hstore_send)
*/
@Override
public int writeAsBinary(ByteBuf buffer, Map<Object, Object> value)
{
int count = value.size();
final int lenIndex = buffer.writerIndex();
int bytesWritten = 4 + 4;
buffer.writeInt(0);
buffer.writeInt(count);
int valueLen = 0;
for (Map.Entry<Object, Object> entry : value.entrySet()) {
byte[] key = entry.getKey().toString().getBytes(UTF_8);
buffer.writeInt(key.length);
buffer.writeBytes(key);
valueLen += key.length + 4;
if (entry.getValue() == null) {
buffer.writeInt(-1);
valueLen += 4;
}
else {
byte[] val = entry.getValue().toString().getBytes(UTF_8);
buffer.writeInt(val.length);
buffer.writeBytes(val);
valueLen += val.length + 4;
}
}
int len = bytesWritten + valueLen;
buffer.setInt(lenIndex, len);
return INT32_BYTE_SIZE + len;
}
@Override
public Map<Object, Object> readBinaryValue(ByteBuf buffer, int valueLength)
{
throw new UnsupportedOperationException("Input of anonymous hstore type values is not implemented");
}
/**
* https://github.com/postgres/postgres/blob/master/contrib/hstore/hstore_io.c (hstore_out)
*/
@Override
public byte[] encodeAsUTF8Text(Map<Object, Object> value)
{
StringBuilder builder = new StringBuilder();
for (Map.Entry<Object, Object> entry : value.entrySet()) {
builder.append("\"").append(entry.getKey()).append("\"").append("=>");
if (entry.getValue() == null) {
builder.append("NULL");
}
else {
builder.append("\"").append(entry.getValue()).append("\"");
}
builder.append(",");
}
builder.setLength(builder.length() - 1);
return builder.toString().getBytes(UTF_8);
}
@Override
public Map<Object, Object> decodeUTF8Text(byte[] bytes)
{
throw new UnsupportedOperationException("Input of anonymous hstore type values is not implemented");
}
@Override
public Object getEmptyValue()
{
return Map.of();
}
}
wren-base/src/main/java/io/wren/base/type/InetType.java
-84
View File
@@ -1,84 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import static java.nio.charset.StandardCharsets.UTF_8;
public class InetType
extends PGType<String>
{
public static final InetType INET = new InetType();
static final int OID = 869;
private static final int TYPE_LEN = -1;
private static final int TYPE_MOD = -1;
private InetType()
{
super(OID, TYPE_LEN, TYPE_MOD, "inet");
}
@Override
public int typArray()
{
return PGArray.INET_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.NETWORK.code();
}
@Override
public String type()
{
return PGType.Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull String value)
{
throw new UnsupportedOperationException("InetType doesn't support binary format.");
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull String value)
{
return value.getBytes(UTF_8);
}
@Override
public String readBinaryValue(ByteBuf buffer, int valueLength)
{
throw new UnsupportedOperationException("InetType doesn't support binary format.");
}
@Override
public String decodeUTF8Text(byte[] bytes)
{
return new String(bytes, UTF_8);
}
@Override
public Object getEmptyValue()
{
return "";
}
}
wren-base/src/main/java/io/wren/base/type/IntegerType.java
-90
View File
@@ -1,90 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import java.nio.charset.StandardCharsets;
import static com.google.common.base.Preconditions.checkArgument;
public class IntegerType
extends PGType<Integer>
{
static final int OID = 23;
private static final int TYPE_LEN = 4;
private static final int TYPE_MOD = -1;
public static final IntegerType INTEGER = new IntegerType();
private IntegerType()
{
super(OID, TYPE_LEN, TYPE_MOD, "int4");
}
@Override
public int typArray()
{
return PGArray.INT4_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.NUMERIC.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull Integer value)
{
buffer.writeInt(TYPE_LEN);
buffer.writeInt(value);
return INT32_BYTE_SIZE + TYPE_LEN;
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull Integer value)
{
return Integer.toString(value).getBytes(StandardCharsets.UTF_8);
}
@Override
public Integer readBinaryValue(ByteBuf buffer, int valueLength)
{
checkArgument(valueLength == TYPE_LEN, "length should be %s because int is int32. Actual length: %s", TYPE_LEN, valueLength);
return buffer.readInt();
}
@Override
public Integer decodeUTF8Text(byte[] bytes)
{
return Integer.parseInt(new String(bytes, StandardCharsets.UTF_8));
}
@Override
public Object getEmptyValue()
{
return 0;
}
}
wren-base/src/main/java/io/wren/base/type/IntervalType.java
-180
View File
@@ -1,180 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import org.joda.time.Period;
import org.joda.time.format.PeriodFormatter;
import org.joda.time.format.PeriodFormatterBuilder;
import javax.annotation.Nonnull;
import java.nio.charset.StandardCharsets;
import static io.wren.base.Utils.checkArgument;
import static java.lang.Math.toIntExact;
public class IntervalType
extends PGType<Period>
{
private static final int OID = 1186;
private static final int TYPE_LEN = 16;
private static final int TYPE_MOD = -1;
public static final IntervalType INTERVAL = new IntervalType();
private static final PeriodFormatter DAY_FORMATTER = new PeriodFormatterBuilder()
.appendYears()
.appendSuffix(" year", " years")
.appendSeparator(" ")
.appendMonths()
.appendSuffix(" mon", " mons")
.appendSeparator(" ")
.appendWeeks()
.appendSuffix(" weeks")
.appendSeparator(" ")
.appendDays()
.appendSuffix(" day", " days")
.toFormatter();
private static final PeriodFormatter TIME_FORMATTER = new PeriodFormatterBuilder()
.printZeroAlways()
.minimumPrintedDigits(2)
.appendHours()
.appendSeparator(":")
.minimumPrintedDigits(2)
.printZeroAlways()
.appendMinutes()
.appendSeparator(":")
.minimumPrintedDigits(2)
.printZeroAlways()
.appendSecondsWithOptionalMillis()
.toFormatter();
private static final PeriodFormatter PG_INTERVAL_FORMATTER = new PeriodFormatterBuilder()
.appendYears()
.appendSuffix(" years ")
.appendMonths()
.appendSuffix(" mons ")
.appendDays()
.appendSuffix(" days ")
.appendHours()
.appendSuffix(" hours ")
.appendMinutes()
.appendSuffix(" mins ")
.appendSecondsWithOptionalMillis()
.appendSuffix(" secs")
.toFormatter();
private IntervalType()
{
super(OID, TYPE_LEN, TYPE_MOD, "interval");
}
@Override
public int typArray()
{
return PGArray.INTERVAL_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.TIMESPAN.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull Period period)
{
buffer.writeInt(TYPE_LEN);
// from PostgreSQL code:
// pq_sendint64(&buf, interval->time);
// pq_sendint32(&buf, interval->day);
// pq_sendint32(&buf, interval->month);
buffer.writeLong(
(period.getHours() * 60 * 60 * 1000_000L)
+ (period.getMinutes() * 60 * 1000_000L)
+ (period.getSeconds() * 1000_000L)
+ (period.getMillis() * 1000));
buffer.writeInt((period.getWeeks() * 7) + period.getDays());
buffer.writeInt((period.getYears() * 12) + period.getMonths());
return INT32_BYTE_SIZE + TYPE_LEN;
}
@Override
public Period readBinaryValue(ByteBuf buffer, int valueLength)
{
checkArgument(valueLength == TYPE_LEN, "length should be %s because interval is 16. Actual length: %s", TYPE_LEN, valueLength);
long micros = buffer.readLong();
int days = buffer.readInt();
int months = buffer.readInt();
long microsInAnHour = 60 * 60 * 1000_000L;
int hours = toIntExact(micros / microsInAnHour);
long microsWithoutHours = micros % microsInAnHour;
long microsInAMinute = 60 * 1000_000L;
int minutes = toIntExact(microsWithoutHours / microsInAMinute);
long microsWithoutMinutes = microsWithoutHours % microsInAMinute;
int seconds = toIntExact(microsWithoutMinutes / 1000_000);
int millis = toIntExact((microsWithoutMinutes % 1000_000) / 1000);
return new Period(
months / 12,
months % 12,
days / 7,
days % 7,
hours,
minutes,
seconds,
millis);
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull Period value)
{
StringBuilder sb = new StringBuilder();
sb.append(DAY_FORMATTER.print(value));
sb.append(" ");
// the negative sign need to be placed before the time, like -00:00:01
if (value.getHours() < 0 || value.getMinutes() < 0 || value.getSeconds() < 0 || value.getMillis() < 0) {
sb.append("-");
}
Period absValue = new Period(
Math.abs(value.getHours()),
Math.abs(value.getMinutes()),
Math.abs(value.getSeconds()),
Math.abs(value.getMillis()));
sb.append(TIME_FORMATTER.print(absValue));
return sb.toString().replace("00:00:00", "").trim().getBytes(StandardCharsets.UTF_8);
}
@Override
public Period decodeUTF8Text(byte[] bytes)
{
return PG_INTERVAL_FORMATTER.parsePeriod(new String(bytes, StandardCharsets.UTF_8));
}
@Override
public Object getEmptyValue()
{
return Period.ZERO;
}
}
wren-base/src/main/java/io/wren/base/type/JsonType.java
-89
View File
@@ -1,89 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import static java.nio.charset.StandardCharsets.UTF_8;
public class JsonType
extends PGType<String>
{
public static final JsonType JSON = new JsonType();
static final int OID = 114;
private static final int TYPE_LEN = -1;
private static final int TYPE_MOD = -1;
private JsonType()
{
super(OID, TYPE_LEN, TYPE_MOD, "json");
}
@Override
public int typArray()
{
return PGArray.JSON_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.USER_DEFINED_TYPES.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull String value)
{
byte[] bytes = encodeAsUTF8Text(value);
buffer.writeInt(bytes.length);
buffer.writeBytes(bytes);
return INT32_BYTE_SIZE + bytes.length;
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull String value)
{
return value.getBytes(UTF_8);
}
@Override
public String readBinaryValue(ByteBuf buffer, int valueLength)
{
byte[] bytes = new byte[valueLength];
buffer.readBytes(bytes);
return decodeUTF8Text(bytes);
}
@Override
public String decodeUTF8Text(byte[] bytes)
{
return new String(bytes, UTF_8);
}
@Override
public Object getEmptyValue()
{
return "";
}
}
wren-base/src/main/java/io/wren/base/type/NumericType.java
-204
View File
@@ -1,204 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.math.MathContext;
import java.nio.charset.StandardCharsets;
public class NumericType
extends PGType<Number>
{
static final int OID = 1700;
private static final int TYPE_LEN = -1;
private static final int TYPE_MOD = -1;
private static final short DEC_DIGITS = 4;
private static final short NUMERIC_POS = 0x0000;
private static final short NUMERIC_NEG = 0x4000;
private static final short NUMERIC_NAN = (short) 0xC000;
public static final NumericType NUMERIC = new NumericType();
private NumericType()
{
super(OID, TYPE_LEN, TYPE_MOD, "numeric");
}
@Override
public int typArray()
{
return PGArray.NUMERIC_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.NUMERIC.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull Number value)
{
return switch (value) {
case BigDecimal bigDecimal -> writeAsBinary(buffer, bigDecimal);
// TODO: customize the handling of BigInteger for performance.
case BigInteger bigInteger -> writeAsBinary(buffer, new BigDecimal(bigInteger));
default -> throw new IllegalArgumentException("Unsupported numeric type: " + value.getClass().getName());
};
}
private int writeAsBinary(ByteBuf buffer, @Nonnull BigDecimal value)
{
// Taken from https://github.com/cockroachdb/cockroach/blob/master/pkg/sql/pgwire/types.go#L336
// and https://github.com/postgres/postgres/blob/master/src/backend/utils/adt/numeric.c#L6760.
// The number is split into chunks of DEC_DIGITS short values while leading and trailing 0's are omitted.
// Examples:
// * 01234 -> [1234]
// * 1234567 -> [0123, 4567], scale 0
// * 1.234500 -> [0001, 2345], scale 4
// * 1234567.12 -> [0123, 4567, 1200], scale 1
// * 1234.0 -> [1234], scale 1
// * 0123.45 -> [0123, 4500], scale 2
char[] digits = value.unscaledValue().toString().toCharArray();
int start = 0;
int end = digits.length;
while (start < end && (digits[start] == '0' || digits[start] == '-')) {
start++;
}
int dWeight = end - start - value.scale() - 1;
while (start < end && digits[end - 1] == '0') {
end--;
}
int len = end - start;
short weight = 0; // Max DEC_DIGIT block index before decimal point
int offset = 0; // Offset inside the first block, e.g. 234.23 has and offset of 1
short nDigits = 0; // Number of DEC_DIGIT blocks
if (len != 0) {
if (dWeight >= 0) {
weight = (short) ((dWeight + 1 + DEC_DIGITS - 1) / DEC_DIGITS - 1);
}
else {
weight = (short) (-((-dWeight - 1) / DEC_DIGITS + 1));
}
offset = (weight + 1) * DEC_DIGITS - (dWeight + 1);
nDigits = (short) ((len + offset + DEC_DIGITS - 1) / DEC_DIGITS);
}
int typeLen = 2 * (4 + nDigits);
buffer.writeInt(typeLen);
buffer.writeShort(nDigits);
buffer.writeShort(weight);
switch (value.signum()) {
case -1:
buffer.writeShort(NUMERIC_NEG);
break;
case 0:
buffer.writeShort(NUMERIC_NAN);
break;
case 1:
buffer.writeShort(NUMERIC_POS);
break;
default:
buffer.writeShort(NUMERIC_POS);
}
buffer.writeShort(value.scale());
int digitIdx = -offset + start;
while (nDigits-- > 0) {
short ndigit = 0;
// Encode 4 digits into a 16 bit short value
for (int nextDigitIdx = digitIdx + DEC_DIGITS; digitIdx < nextDigitIdx; digitIdx++) {
ndigit *= 10;
if (digitIdx >= start && digitIdx < end) {
ndigit += digits[digitIdx] - '0';
}
}
buffer.writeShort(ndigit);
}
return INT32_BYTE_SIZE + typeLen;
}
@Override
public BigDecimal readBinaryValue(ByteBuf buffer, int valueLength)
{
// Number of DEC_DIGIT blocks
short nDigits = buffer.readShort();
// DEC_DIGIT blocks before decimal point
short weight = buffer.readShort();
short sign = buffer.readShort();
short scale = buffer.readShort();
if (sign == NUMERIC_NAN) {
throw new IllegalArgumentException("Infinite or NaN values are not supported");
}
boolean hasDp = scale > 0;
int sizeOfBytes = (nDigits * DEC_DIGITS) + (hasDp ? 1 : 0);
char[] decDigits = new char[sizeOfBytes];
int decDigitsIdx = 0;
for (int i = 0; i < nDigits; i++) {
int decDigit = buffer.readShort();
if (decDigit > 0) {
// Decode 4 digits from a 16 bit short
for (int j = 1000; j > 0 && decDigitsIdx < sizeOfBytes; j /= 10) {
int d1 = (decDigit / j);
decDigit -= d1 * j;
decDigits[decDigitsIdx++] = (char) (d1 + '0');
}
}
if (hasDp && i == weight) {
decDigits[decDigitsIdx++] = '.';
}
}
BigDecimal bd = new BigDecimal(decDigits)
.setScale(scale, MathContext.UNLIMITED.getRoundingMode());
return sign == NUMERIC_NEG ? bd.negate() : bd;
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull Number value)
{
return value.toString().getBytes(StandardCharsets.UTF_8);
}
@Override
public Number decodeUTF8Text(byte[] bytes)
{
return new BigDecimal(new String(bytes, StandardCharsets.UTF_8));
}
@Override
public Object getEmptyValue()
{
return BigDecimal.ZERO;
}
}
wren-base/src/main/java/io/wren/base/type/OidType.java
-88
View File
@@ -1,88 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import java.nio.charset.StandardCharsets;
import static com.google.common.base.Preconditions.checkArgument;
public class OidType
extends PGType<Long>
{
public static final OidType OID_INSTANCE = new OidType();
static final int OID = 26;
private static final int TYPE_LEN = 8;
private static final int TYPE_MOD = -1;
OidType()
{
super(OID, TYPE_LEN, TYPE_MOD, "oid");
}
@Override
public int typArray()
{
return PGArray.OID_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.NUMERIC.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, Long value)
{
buffer.writeInt(TYPE_LEN);
buffer.writeLong(value);
return INT32_BYTE_SIZE + TYPE_LEN;
}
@Override
public Long readBinaryValue(ByteBuf buffer, int valueLength)
{
checkArgument(valueLength == TYPE_LEN, "length should be %s because oid is int32. Actual length: %s", TYPE_LEN, valueLength);
return buffer.readLong();
}
@Override
public byte[] encodeAsUTF8Text(Long value)
{
return Long.toString(value).getBytes(StandardCharsets.UTF_8);
}
@Override
public Long decodeUTF8Text(byte[] bytes)
{
return Long.parseLong(new String(bytes, StandardCharsets.UTF_8));
}
@Override
public Object getEmptyValue()
{
return 0;
}
}
wren-base/src/main/java/io/wren/base/type/PGArray.java
-404
View File
@@ -1,404 +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.wren.base.type;
import com.google.common.collect.ImmutableList;
import com.google.common.primitives.Bytes;
import io.netty.buffer.ByteBuf;
import io.wren.base.type.parser.PgArrayParserWrapper;
import javax.annotation.Nonnull;
import java.util.ArrayList;
import java.util.List;
import static io.wren.base.type.BooleanType.BOOLEAN;
import static io.wren.base.type.DateType.DATE;
import static java.nio.charset.StandardCharsets.UTF_8;
public class PGArray
extends PGType<List<Object>>
{
public static final PGArray CHAR_ARRAY = new PGArray(1002, CharType.CHAR);
public static final PGArray BPCHAR_ARRAY = new PGArray(1014, BpCharType.BPCHAR);
public static final PGArray INT2_ARRAY = new PGArray(1005, SmallIntType.SMALLINT);
public static final PGArray INT4_ARRAY = new PGArray(1007, IntegerType.INTEGER);
public static final PGArray OID_ARRAY = new PGArray(1028, OidType.OID_INSTANCE);
public static final PGArray INT8_ARRAY = new PGArray(1016, BigIntType.BIGINT);
public static final PGArray FLOAT4_ARRAY = new PGArray(1021, RealType.REAL);
public static final PGArray FLOAT8_ARRAY = new PGArray(1022, DoubleType.DOUBLE);
public static final PGArray NUMERIC_ARRAY = new PGArray(1231, NumericType.NUMERIC);
public static final PGArray BOOL_ARRAY = new PGArray(1000, BOOLEAN);
public static final PGArray VARCHAR_ARRAY = new PGArray(1015, VarcharType.VARCHAR);
public static final PGArray TEXT_ARRAY = new PGArray(1009, VarcharType.TextType.TEXT);
public static final PGArray NAME_ARRAY = new PGArray(1003, VarcharType.NameType.NAME);
public static final PGArray JSON_ARRAY = new PGArray(199, JsonType.JSON);
public static final PGArray TIMESTAMP_WITH_TIMEZONE_ARRAY = new PGArray(1185, TimestampWithTimeZoneType.TIMESTAMP_WITH_TIMEZONE);
public static final PGArray TIMESTAMP_ARRAY = new PGArray(1115, TimestampType.TIMESTAMP);
public static final PGArray DATE_ARRAY = new PGArray(1182, DATE);
public static final PGArray UUID_ARRAY = new PGArray(2951, UuidType.UUID);
public static final PGArray BYTEA_ARRAY = new PGArray(1001, ByteaType.BYTEA);
public static final PGArray INET_ARRAY = new PGArray(1041, InetType.INET);
public static final PGArray EMPTY_RECORD_ARRAY = new PGArray(2287, RecordType.EMPTY_RECORD);
public static final PGArray HSTORE_ARRAY = new PGArray(57645, HstoreType.HSTORE);
public static final PGArray REGPROC_ARRAY = new PGArray(1008, RegprocType.REGPROC);
public static final PGArray INTERVAL_ARRAY = new PGArray(1187, IntervalType.INTERVAL);
// TODO:
// public static final PGArray TIMETZ_ARRAY = new PGArray(1270, TimeTZType.INSTANCE);
// public static final PGArray POINT_ARRAY = new PGArray(1017, PointType.INSTANCE);
// public static final PGArray ANY_ARRAY = new PGArray(
// 2277,
// AnyType.INSTANCE.typName() + "array",
// AnyType.INSTANCE)
// {
// @Override
// public String typeCategory()
// {
// return TypeCategory.PSEUDO.code();
// }
// };
private static final byte[] NULL_BYTES = new byte[] {'N', 'U', 'L', 'L'};
public static List<PGArray> allArray()
{
// TODO: support IntervalType array
return ImmutableList.of(
CHAR_ARRAY,
BPCHAR_ARRAY,
INT2_ARRAY,
INT4_ARRAY,
OID_ARRAY,
INT8_ARRAY,
FLOAT4_ARRAY,
FLOAT8_ARRAY,
NUMERIC_ARRAY,
BOOL_ARRAY,
VARCHAR_ARRAY,
TEXT_ARRAY,
JSON_ARRAY,
NAME_ARRAY,
TIMESTAMP_ARRAY,
TIMESTAMP_WITH_TIMEZONE_ARRAY,
DATE_ARRAY,
HSTORE_ARRAY,
INET_ARRAY,
EMPTY_RECORD_ARRAY,
UUID_ARRAY,
BYTEA_ARRAY);
}
private final PGType<?> innerType;
PGArray(int oid, String name, PGType<?> innerType)
{
super(oid, -1, -1, name);
this.innerType = innerType;
}
PGArray(int oid, PGType<?> innerType)
{
this(oid, "_" + innerType.typName(), innerType);
}
public PGType<?> getInnerType()
{
return innerType;
}
@Override
public int typArray()
{
return 0;
}
@Override
public String typeCategory()
{
return TypeCategory.ARRAY.code();
}
@Override
public String type()
{
return innerType.type();
}
@Override
public int typElem()
{
return innerType.oid();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull List<Object> value)
{
int dimensions = getDimensions(value);
List<Integer> dimensionsList = new ArrayList<>();
buildDimensions(value, dimensionsList, dimensions, 1);
int bytesWritten = 4 + 4 + 4;
final int lenIndex = buffer.writerIndex();
buffer.writeInt(0);
buffer.writeInt(dimensions);
buffer.writeInt(1); // flags bit 0: 0=no-nulls, 1=has-nulls
buffer.writeInt(typElem());
for (Integer dim : dimensionsList) {
buffer.writeInt(dim); // upper bound
buffer.writeInt(dim); // lower bound
bytesWritten += 8;
}
int len = bytesWritten + writeArrayAsBinary(buffer, value, dimensionsList, 1);
buffer.setInt(lenIndex, len);
return INT32_BYTE_SIZE + len; // add also the size of the length itself
}
private int getDimensions(@Nonnull Object value)
{
int dimensions = 0;
Object array = value;
do {
dimensions++;
List<?> arr = (List<?>) array;
if (arr.isEmpty()) {
break;
}
array = null;
for (Object o : arr) {
if (o == null) {
continue;
}
array = o;
}
}
while (array instanceof List);
return dimensions;
}
@Override
public List<Object> readBinaryValue(ByteBuf buffer, int valueLength)
{
int dimensions = buffer.readInt();
buffer.readInt(); // flags bit 0: 0=no-nulls, 1=has-nulls
buffer.readInt(); // element oid
if (dimensions == 0) {
return ImmutableList.of();
}
int[] dims = new int[dimensions];
for (int d = 0; d < dimensions; ++d) {
dims[d] = buffer.readInt();
buffer.readInt(); // lowerBound ignored
}
List<Object> values = new ArrayList<>(dims[0]);
readArrayAsBinary(buffer, values, dims, 0);
return values;
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull List<Object> array)
{
boolean isJson = JsonType.OID == innerType.oid();
List<Byte> encodedValues = new ArrayList<>();
encodedValues.add((byte) '{');
for (int i = 0; i < array.size(); i++) {
Object o = array.get(i);
if (o instanceof List) { // Nested Array -> recursive call
byte[] bytes = encodeAsUTF8Text((List) o);
for (byte b : bytes) {
encodedValues.add(b);
}
if (i == 0) {
encodedValues.add((byte) ',');
}
}
else {
if (i > 0) {
encodedValues.add((byte) ',');
}
byte[] bytes;
if (o == null) {
bytes = NULL_BYTES;
for (byte aByte : bytes) {
encodedValues.add(aByte);
}
}
else {
bytes = ((PGType) innerType).encodeAsUTF8Text(o);
if (needDoubleQuoteAround(innerType.oid(), bytes)) {
encodedValues.add((byte) '"');
}
if (isJson) {
for (byte aByte : bytes) {
// Escape double quotes with backslash for json
char ch = (char) aByte;
if (ch == '"' || ch == '\\') {
encodedValues.add((byte) '\\');
}
encodedValues.add(aByte);
}
}
else {
for (byte aByte : bytes) {
encodedValues.add(aByte);
}
}
if (needDoubleQuoteAround(innerType.oid(), bytes)) {
encodedValues.add((byte) '"');
}
}
}
}
encodedValues.add((byte) '}');
return Bytes.toArray(encodedValues);
}
@Override
public List<Object> decodeUTF8Text(byte[] bytes)
{
/*
* text representation:
*
* 1-dimension integer array:
* {10,NULL,NULL,20,30}
* {"10",NULL,NULL,"20","30"}
* 2-dimension integer array:
* {{"10","20"},{"30",NULL,"40}}
*
* 1-dimension json array:
* {"{"x": 10}","{"y": 20}"}
* 2-dimension json array:
* {{"{"x": 10}","{"y": 20}"},{"{"x": 30}","{"y": 40}"}}
*/
return (List<Object>) PgArrayParserWrapper.parse(bytes, innerType::decodeUTF8Text);
}
private int buildDimensions(List<Object> values, List<Integer> dimensionsList, int maxDimensions, int currentDimension)
{
if (values == null) {
return 1;
}
// While elements of array are also arrays
if (currentDimension < maxDimensions) {
int max = 0;
for (Object o : values) {
max = Math.max(max, buildDimensions((List<Object>) o, dimensionsList, maxDimensions, currentDimension + 1));
}
if (currentDimension == maxDimensions - 1) {
dimensionsList.add(max);
}
else {
Integer current = dimensionsList.get(0);
dimensionsList.set(0, Math.max(current, max));
}
}
// Add the dimensions of 1st dimension
if (currentDimension == 1) {
dimensionsList.add(0, values.size());
}
return values.size();
}
private int writeArrayAsBinary(ByteBuf buffer, List<Object> array, List<Integer> dimensionsList, int currentDimension)
{
int bytesWritten = 0;
if (array == null) {
for (int i = 0; i < dimensionsList.get(currentDimension - 1); i++) {
buffer.writeInt(-1);
bytesWritten += 4;
}
return bytesWritten;
}
// 2nd to last level
if (currentDimension == dimensionsList.size()) {
int i = 0;
for (Object o : array) {
if (o == null) {
buffer.writeInt(-1);
bytesWritten += 4;
}
else {
bytesWritten += ((PGType) innerType).writeAsBinary(buffer, o);
}
i++;
}
// Fill in with -1 for up to max dimensions
for (; i < dimensionsList.get(currentDimension - 1); i++) {
buffer.writeInt(-1);
bytesWritten += 4;
}
}
else {
for (Object o : array) {
bytesWritten += writeArrayAsBinary(buffer, (List<Object>) o, dimensionsList, currentDimension + 1);
}
}
return bytesWritten;
}
private void readArrayAsBinary(ByteBuf buffer,
final List<Object> array,
final int[] dims,
final int thisDimension)
{
if (thisDimension == dims.length - 1) {
for (int i = 0; i < dims[thisDimension]; ++i) {
int len = buffer.readInt();
if (len == -1) {
array.add(null);
}
else {
array.add(innerType.readBinaryValue(buffer, len));
}
}
}
else {
for (int i = 0; i < dims[thisDimension]; ++i) {
ArrayList<Object> list = new ArrayList<>(dims[thisDimension + 1]);
array.add(list);
readArrayAsBinary(buffer, list, dims, thisDimension + 1);
}
}
}
private static boolean needDoubleQuoteAround(int typeOid, byte[] element)
{
// The array output routine will put double quotes around element values
// if they are empty strings, contain curly braces, delimiter characters, double quotes, backslashes, or white space, or match the word NULL.
// https://www.postgresql.org/docs/13/arrays.html#ARRAYS-IO
for (byte b : element) {
if (b == '{' || b == '}' || b == '"' || b == '\\' || b == ',' || b == '\t' || b == '\n' || b == '\r' || b == '\f' || b == ' ') {
return true;
}
}
if (VarcharType.OID == typeOid) {
return element.length == 0 || new String(element, UTF_8).equalsIgnoreCase("NULL");
}
return false;
}
@Override
public Object getEmptyValue()
{
return ImmutableList.of();
}
}
wren-base/src/main/java/io/wren/base/type/PGType.java
-211
View File
@@ -1,211 +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.wren.base.type;
import io.airlift.log.Logger;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import java.nio.charset.StandardCharsets;
public abstract class PGType<T>
{
public enum Type
{
BASE("b"),
COMPOSITE("c"),
ENUM("e"),
DOMAIN("d"),
PSEUDO("p"),
RANGE("r");
private final String code;
Type(String code)
{
this.code = code;
}
public String code()
{
return code;
}
}
public enum TypeCategory
{
ARRAY("A"),
BOOLEAN("B"),
COMPOSITE("C"),
DATETIME("D"),
GEOMETRIC("G"),
NETWORK("I"),
NUMERIC("N"),
PSEUDO("P"),
RANGE("R"),
STRING("S"),
TIMESPAN("T"),
USER_DEFINED_TYPES("U"),
BIT_STRING("V"),
UNKNOWN("X");
private final String code;
TypeCategory(String code)
{
this.code = code;
}
public String code()
{
return code;
}
}
public static final int INT32_BYTE_SIZE = Integer.SIZE / 8;
private static final Logger LOGGER = Logger.get(PGType.class);
private final int oid;
private final int typeLen;
private final int typeMod;
private final String typName;
protected PGType(int oid, int typeLen, int typeMod, @Nonnull String typName)
{
this.oid = oid;
this.typeLen = typeLen;
this.typeMod = typeMod;
this.typName = typName;
}
public int oid()
{
return oid;
}
public short typeLen()
{
return (short) typeLen;
}
public abstract int typArray();
public int typeMod()
{
return typeMod;
}
public String typName()
{
return typName;
}
public String typInput()
{
if (typArray() == 0) {
return "array_in";
}
return "any_in";
}
public String typOutput()
{
if (typArray() == 0) {
return "array_out";
}
return "any_out";
}
public String typReceive()
{
if (typArray() == 0) {
return "array_recv";
}
return "any_recv";
}
public int typElem()
{
return 0;
}
public String typDelim()
{
return ",";
}
public abstract String typeCategory();
public abstract String type();
/**
* Write the value as text into the buffer.
* <p>
* Format:
* <pre>
* | int32 len (excluding len itself) | byte<b>N</b> value onto the buffer
* </pre>
*
* @return the number of bytes written. (4 (int32) + N)
*/
public int writeAsText(ByteBuf buffer, @Nonnull T value)
{
byte[] bytes = encodeAsUTF8Text(value);
buffer.writeInt(bytes.length);
buffer.writeBytes(bytes);
return INT32_BYTE_SIZE + bytes.length;
}
public T readTextValue(ByteBuf buffer, int valueLength)
{
byte[] bytes = new byte[valueLength];
buffer.readBytes(bytes);
try {
return decodeUTF8Text(bytes);
}
catch (Throwable t) {
LOGGER.warn("decodeUTF8Text failed. input=%s type=%s",
new String(bytes, StandardCharsets.UTF_8), typName);
throw t;
}
}
/**
* Write the value as binary into the buffer.
* <p>
* Format:
* <pre>
* | int32 len (excluding len itself) | byte<b>N</b> value onto the buffer
* </pre>
*
* @return the number of bytes written. (4 (int32) + N)
*/
public abstract int writeAsBinary(ByteBuf buffer, @Nonnull T value);
public abstract T readBinaryValue(ByteBuf buffer, int valueLength);
/**
* Return the UTF8 encoded text representation of the value
*/
public abstract byte[] encodeAsUTF8Text(@Nonnull T value);
/**
* Convert a UTF8 encoded text representation into the actual value
*/
public abstract T decodeUTF8Text(byte[] bytes);
public abstract Object getEmptyValue();
}
wren-base/src/main/java/io/wren/base/type/PGTypes.java
-132
View File
@@ -1,132 +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.wren.base.type;
import com.google.common.collect.ImmutableMap;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Locale;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import static io.wren.base.type.BooleanType.BOOLEAN;
import static io.wren.base.type.HstoreType.HSTORE;
import static io.wren.base.type.IntervalType.INTERVAL;
import static io.wren.base.type.RealType.REAL;
import static io.wren.base.type.TimestampType.TIMESTAMP;
import static java.lang.String.format;
public final class PGTypes
{
private PGTypes() {}
private static final Map<Integer, PGType<?>> TYPE_TABLE = new HashMap<>();
private static final Map<Integer, PGArray> INNER_TYPE_TO_ARRAY_TABLE;
private static final Set<PGType<?>> TYPES;
private static final Map<String, PGType<?>> TYPE_NAME_TABLE = new HashMap<>();
static {
TYPE_TABLE.put(BOOLEAN.oid(), BOOLEAN);
TYPE_TABLE.put(TinyIntType.TINYINT.oid(), TinyIntType.TINYINT);
TYPE_TABLE.put(SmallIntType.SMALLINT.oid(), SmallIntType.SMALLINT);
TYPE_TABLE.put(IntegerType.INTEGER.oid(), IntegerType.INTEGER);
TYPE_TABLE.put(BigIntType.BIGINT.oid(), BigIntType.BIGINT);
TYPE_TABLE.put(REAL.oid(), REAL);
TYPE_TABLE.put(DoubleType.DOUBLE.oid(), DoubleType.DOUBLE);
TYPE_TABLE.put(NumericType.NUMERIC.oid(), NumericType.NUMERIC);
TYPE_TABLE.put(VarcharType.VARCHAR.oid(), VarcharType.VARCHAR);
TYPE_TABLE.put(CharType.CHAR.oid(), CharType.CHAR);
TYPE_TABLE.put(JsonType.JSON.oid(), JsonType.JSON);
TYPE_TABLE.put(TIMESTAMP.oid(), TIMESTAMP);
TYPE_TABLE.put(TimestampWithTimeZoneType.TIMESTAMP_WITH_TIMEZONE.oid(), TimestampWithTimeZoneType.TIMESTAMP_WITH_TIMEZONE);
TYPE_TABLE.put(VarcharType.TextType.TEXT.oid(), VarcharType.TextType.TEXT);
TYPE_TABLE.put(VarcharType.NameType.NAME.oid(), VarcharType.NameType.NAME);
TYPE_TABLE.put(OidType.OID_INSTANCE.oid(), OidType.OID_INSTANCE);
TYPE_TABLE.put(DateType.DATE.oid(), DateType.DATE);
TYPE_TABLE.put(ByteaType.BYTEA.oid(), ByteaType.BYTEA);
TYPE_TABLE.put(BpCharType.BPCHAR.oid(), BpCharType.BPCHAR);
// we handle all unspecified type as text type.
TYPE_TABLE.put(0, VarcharType.TextType.TEXT);
TYPE_TABLE.put(InetType.INET.oid(), InetType.INET);
TYPE_TABLE.put(RecordType.EMPTY_RECORD.oid(), RecordType.EMPTY_RECORD);
// Just need a fake instance to do type mapping. We never use the field.
TYPE_TABLE.put(HSTORE.oid(), HSTORE);
TYPE_TABLE.put(UuidType.UUID.oid(), UuidType.UUID);
TYPE_TABLE.put(INTERVAL.oid(), INTERVAL);
ImmutableMap.Builder<Integer, PGArray> innerToPgTypeBuilder = ImmutableMap.builder();
// initial collection types
PGArray.allArray().forEach(array -> {
PGType<?> innerType = array.getInnerType();
TYPE_TABLE.put(array.oid(), array);
innerToPgTypeBuilder.put(innerType.oid(), array);
});
INNER_TYPE_TO_ARRAY_TABLE = innerToPgTypeBuilder.build();
TYPES = new HashSet<>(TYPE_TABLE.values());
// the following polymorphic types are added manually,
// because there are no corresponding data types in Cannerflow
TYPES.add(AnyType.ANY);
TYPES.forEach(type -> TYPE_NAME_TABLE.put(type.typName().toUpperCase(Locale.ROOT), type));
TYPE_NAME_TABLE.put("REAL", REAL);
TYPE_NAME_TABLE.put("DOUBLE", DoubleType.DOUBLE);
TYPE_NAME_TABLE.put("BOOLEAN", BOOLEAN);
TYPE_NAME_TABLE.put("INTEGER", IntegerType.INTEGER);
TYPE_NAME_TABLE.put("SMALLINT", SmallIntType.SMALLINT);
TYPE_NAME_TABLE.put("BIGINT", BigIntType.BIGINT);
TYPE_NAME_TABLE.put("STRING", VarcharType.VARCHAR);
TYPE_NAME_TABLE.put("DECIMAL", NumericType.NUMERIC);
}
public static Iterable<PGType<?>> pgTypes()
{
return TYPES;
}
public static PGType<?> oidToPgType(int oid)
{
PGType<?> pgType = TYPE_TABLE.get(oid);
if (pgType == null) {
throw new IllegalArgumentException(
format("No oid mapping from '%s' to pg_type", oid));
}
return pgType;
}
public static PGType<?> getArrayType(int innerOid)
{
PGType<?> arrayType = INNER_TYPE_TO_ARRAY_TABLE.get(innerOid);
if (arrayType == null) {
throw new IllegalArgumentException(
format("No array type mapping from '%s' to pg_type", innerOid));
}
return arrayType;
}
public static Optional<PGType<?>> nameToPgType(String name)
{
return Optional.ofNullable(TYPE_NAME_TABLE.get(name.toUpperCase(Locale.ROOT)));
}
public static PGType<?> toPgRecordArray(PGType<?> innerRecordType)
{
return new PGArray(2287, innerRecordType);
}
}
wren-base/src/main/java/io/wren/base/type/PgDatetimeUtils.java
-62
View File
@@ -1,62 +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.wren.base.type;
public final class PgDatetimeUtils
{
private PgDatetimeUtils() {}
public static final int SECS_PER_DAY = 86400;
// amount of seconds between 1970-01-01 and 2000-01-01
public static final int EPOCH_DIFF_IN_SEC = 946_684_800;
public static final long EPOCH_DIFF_IN_MS = EPOCH_DIFF_IN_SEC * 1000L;
public static final int EPOCH_DIFF_IN_DAY = EPOCH_DIFF_IN_SEC / SECS_PER_DAY;
/**
* Convert a trino date (unix timestamp in day) into a postgres date
* (int days since 2000-01-01)
*/
public static int toPgDate(int unixTsInDay)
{
return (unixTsInDay - EPOCH_DIFF_IN_DAY);
}
/**
* Convert a postgres date (days since 2000-01-01) into a trino
* date (unix timestamp in days).
*/
public static int toTrinoDate(int daysSince2k)
{
return daysSince2k + EPOCH_DIFF_IN_DAY;
}
/**
* Convert a presto timestamp (unix timestamp in ms) into a postgres timestamp
* (long microseconds since 2000-01-01)
*/
public static long toPgTimestamp(long unixTsInMs)
{
return (unixTsInMs - EPOCH_DIFF_IN_MS) * 1000;
}
/**
* Convert a postgres timestamp (seconds since 2000-01-01) into a presto
* timestamp (unix timestamp in ms).
*/
public static long toTrinoTimestamp(long microSecondsSince2k)
{
return (microSecondsSince2k / 1000) + EPOCH_DIFF_IN_MS;
}
}
wren-base/src/main/java/io/wren/base/type/PgOidVectorType.java
-102
View File
@@ -1,102 +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.wren.base.type;
import com.google.common.base.Joiner;
import io.netty.buffer.ByteBuf;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.StringTokenizer;
public class PgOidVectorType
extends PGType<List<Integer>>
{
public static final PgOidVectorType INSTANCE = new PgOidVectorType();
private static final int OID = 30;
PgOidVectorType()
{
super(OID, -1, -1, "oidvector");
}
@Override
public int typArray()
{
return 1013;
}
@Override
public int typElem()
{
return OidType.OID;
}
@Override
public String typeCategory()
{
return TypeCategory.ARRAY.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
@SuppressWarnings({"unchecked", "rawtypes"})
public int writeAsBinary(ByteBuf buffer, List<Integer> value)
{
return PGArray.INT4_ARRAY.writeAsBinary(buffer, (List) value);
}
@Override
@SuppressWarnings({"unchecked", "rawtypes"})
public List<Integer> readBinaryValue(ByteBuf buffer, int valueLength)
{
return (List<Integer>) (List) PGArray.INT4_ARRAY.readBinaryValue(buffer, valueLength);
}
@Override
public byte[] encodeAsUTF8Text(List<Integer> value)
{
return Joiner.on(" ").join(value).getBytes(StandardCharsets.UTF_8);
}
@Override
public List<Integer> decodeUTF8Text(byte[] bytes)
{
String string = new String(bytes, StandardCharsets.UTF_8);
return listFromOidVectorString(string);
}
public static List<Integer> listFromOidVectorString(String value)
{
StringTokenizer tokenizer = new StringTokenizer(value, " ");
ArrayList<Integer> oids = new ArrayList<>();
while (tokenizer.hasMoreTokens()) {
oids.add(Integer.parseInt(tokenizer.nextToken()));
}
return oids;
}
@Override
public Object getEmptyValue()
{
return List.of();
}
}
wren-base/src/main/java/io/wren/base/type/PgTypeUtils.java
-103
View File
@@ -1,103 +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.wren.base.type;
import com.google.common.collect.ImmutableMap;
import java.util.Map;
import java.util.Optional;
import static io.wren.base.type.BooleanType.BOOLEAN;
import static io.wren.base.type.ByteaType.BYTEA;
import static java.util.Locale.ENGLISH;
public final class PgTypeUtils
{
private PgTypeUtils() {}
private static final Map<String, PGType<?>> pgNameToTypeMap;
static {
ImmutableMap.Builder<String, PGType<?>> pgNameToTypeMapBuilder = ImmutableMap.<String, PGType<?>>builder();
// typName
Map<String, PGType<?>> typNameBuilder = ImmutableMap.<String, PGType<?>>builder()
.put("bool", BOOLEAN)
.put("int2", SmallIntType.SMALLINT)
.put("int4", IntegerType.INTEGER)
.put("int8", BigIntType.BIGINT)
.put("float4", RealType.REAL)
.put("float8", DoubleType.DOUBLE)
.put("numeric", NumericType.NUMERIC)
.put("varchar", VarcharType.VARCHAR)
.put("char", CharType.CHAR)
.put("date", DateType.DATE)
.put("timestamp", TimestampType.TIMESTAMP)
.put("json", JsonType.JSON)
.put("bytea", BYTEA)
.put("record", RecordType.EMPTY_RECORD)
.put("interval", IntervalType.INTERVAL)
.build();
// TODO
// .put("name", NAME)
// .put("text", TEXT)
// .put("inet", INET)
// .put("uuid", UUID)
// .put("timestamptz", TIMESTAMP_WITH_TIMEZONE)
// .put("timestamp with time zone", TIMESTAMP_WITH_TIMEZONE)
pgNameToTypeMapBuilder.putAll(typNameBuilder);
// alias name
Map<String, PGType<?>> aliasNameBuilder = ImmutableMap.<String, PGType<?>>builder()
.put("boolean", BOOLEAN)
.put("smallint", SmallIntType.SMALLINT)
.put("int", IntegerType.INTEGER)
.put("integer", IntegerType.INTEGER)
.put("bigint", BigIntType.BIGINT)
.put("real", RealType.REAL)
.put("decimal", NumericType.NUMERIC)
.put("double precision", DoubleType.DOUBLE)
.put("character", CharType.CHAR)
.put("character varying", VarcharType.VARCHAR)
.build();
pgNameToTypeMapBuilder.putAll(aliasNameBuilder);
// array name
typNameBuilder.forEach((name, type) -> {
// TODO Support interval array
if (type == IntervalType.INTERVAL) {
return;
}
PGType<?> arrayType = PGTypes.getArrayType(type.oid());
pgNameToTypeMapBuilder.put("_" + name, arrayType);
pgNameToTypeMapBuilder.put(name + "[]", arrayType);
pgNameToTypeMapBuilder.put(name + " array", arrayType);
});
aliasNameBuilder.forEach((name, type) -> {
PGType<?> arrayType = PGTypes.getArrayType(type.oid());
pgNameToTypeMapBuilder.put(name + "[]", arrayType);
pgNameToTypeMapBuilder.put(name + " array", arrayType);
});
pgNameToTypeMap = pgNameToTypeMapBuilder.build();
}
public static Optional<PGType<?>> pgNameToType(String name)
{
String lowerCaseName = name.toLowerCase(ENGLISH);
return pgNameToTypeMap.containsKey(lowerCaseName) ? Optional.of(pgNameToTypeMap.get(lowerCaseName)) : Optional.empty();
}
}
wren-base/src/main/java/io/wren/base/type/RealType.java
-89
View File
@@ -1,89 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import java.nio.charset.StandardCharsets;
import static com.google.common.base.Preconditions.checkArgument;
public class RealType
extends PGType<Float>
{
public static final RealType REAL = new RealType();
static final int OID = 700;
private static final int TYPE_LEN = 4;
private static final int TYPE_MOD = -1;
private RealType()
{
super(OID, TYPE_LEN, TYPE_MOD, "float4");
}
@Override
public int typArray()
{
return PGArray.FLOAT4_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.NUMERIC.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull Float value)
{
buffer.writeInt(TYPE_LEN);
buffer.writeFloat(value);
return INT32_BYTE_SIZE + TYPE_LEN;
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull Float value)
{
return Float.toString(value).getBytes(StandardCharsets.UTF_8);
}
@Override
public Float readBinaryValue(ByteBuf buffer, int valueLength)
{
checkArgument(valueLength == TYPE_LEN, "length should be %s because float is int32. Actual length: %s", TYPE_LEN, valueLength);
return buffer.readFloat();
}
@Override
public Float decodeUTF8Text(byte[] bytes)
{
return Float.parseFloat(new String(bytes, StandardCharsets.UTF_8));
}
@Override
public Object getEmptyValue()
{
return 0.0f;
}
}
wren-base/src/main/java/io/wren/base/type/RecordType.java
-142
View File
@@ -1,142 +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.wren.base.type;
import com.carrotsearch.hppc.ByteArrayList;
import io.netty.buffer.ByteBuf;
import java.util.List;
import java.util.Map;
import static com.google.common.collect.ImmutableList.toImmutableList;
public class RecordType
extends PGType<Map<String, Object>>
{
public static final RecordType EMPTY_RECORD = new RecordType(List.of());
private static final int OID = 2249;
private static final String NAME = "record";
private final List<PGType<?>> fieldTypes;
public RecordType(List<PGType<?>> fieldTypes)
{
super(OID, -1, -1, NAME);
this.fieldTypes = fieldTypes;
}
@Override
public int typArray()
{
return PGArray.EMPTY_RECORD_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.PSEUDO.code();
}
@Override
public String type()
{
return Type.PSEUDO.code();
}
@SuppressWarnings({"unchecked", "rawtypes"})
@Override
public int writeAsBinary(ByteBuf buffer, Map<String, Object> record)
{
final int startWriterIndex = buffer.writerIndex();
buffer.writeInt(0); // reserve space for the length of the record; updated later
buffer.writeInt(fieldTypes.size());
int bytesWritten = 4;
List<Map.Entry<String, Object>> entries = record.entrySet().stream().collect(toImmutableList());
for (int i = 0; i < fieldTypes.size(); i++) {
PGType fieldType = fieldTypes.get(i);
buffer.writeInt(fieldType.oid());
bytesWritten += 4;
Map.Entry<String, Object> entry = entries.get(i);
if (entry.getValue() == null) {
buffer.writeInt(-1); // -1 data length signals a NULL
bytesWritten += 4;
continue;
}
bytesWritten += fieldType.writeAsBinary(buffer, entry.getValue());
}
buffer.setInt(startWriterIndex, bytesWritten);
return 4 + bytesWritten;
}
@Override
public Map<String, Object> readBinaryValue(ByteBuf buffer, int valueLength)
{
throw new UnsupportedOperationException("Input of anonymous record type values is not implemented");
}
@SuppressWarnings({"unchecked", "rawtypes"})
@Override
public byte[] encodeAsUTF8Text(Map<String, Object> record)
{
ByteArrayList bytes = new ByteArrayList();
// See PostgreSQL src/backend/utils/adt/rowtypes.c record_out(PG_FUNCTION_ARGS)
bytes.add((byte) '(');
List<Map.Entry<String, Object>> rows = record.entrySet().stream().collect(toImmutableList());
for (int i = 0; i < record.size(); i++) {
PGType fieldType = fieldTypes.get(i);
Map.Entry<String, Object> row = rows.get(i);
if (i > 0) {
bytes.add((byte) ',');
}
if (row.getValue() == null) {
continue;
}
byte[] encodedValue = fieldType.encodeAsUTF8Text(row.getValue());
boolean needQuotes = encodedValue.length == 0;
for (byte b : encodedValue) {
char c = (char) b;
if (c == '"' || c == '\\' || c == '(' || c == ')' || c == ',' || Character.isWhitespace(c)) {
needQuotes = true;
break;
}
}
if (needQuotes) {
bytes.add((byte) '\"');
}
bytes.add(encodedValue);
if (needQuotes) {
bytes.add((byte) '\"');
}
}
bytes.add((byte) ')');
return bytes.toArray();
}
@Override
public Map<String, Object> decodeUTF8Text(byte[] bytes)
{
throw new UnsupportedOperationException("Input of record type values is not implemented");
}
@Override
public Object getEmptyValue()
{
return Map.of();
}
}
wren-base/src/main/java/io/wren/base/type/RegprocType.java
-96
View File
@@ -1,96 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import static java.nio.charset.StandardCharsets.UTF_8;
public class RegprocType
extends PGType<String>
{
public static final RegprocType REGPROC = new RegprocType();
private static final int OID = 24;
// TODO: It's 4 in PostgreSQL because pg use oid to present this type but we use name.
private static final int TYPE_LEN = -1;
private static final int TYPE_MOD = -1;
RegprocType()
{
super(OID, TYPE_LEN, TYPE_MOD, "regproc");
}
@Override
public int typArray()
{
return PGArray.REGPROC_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.NUMERIC.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull String value)
{
byte[] bytes = value.getBytes(UTF_8);
buffer.writeInt(bytes.length);
buffer.writeBytes(bytes);
return INT32_BYTE_SIZE + bytes.length;
}
@Override
public int writeAsText(ByteBuf buffer, @Nonnull String value)
{
return writeAsBinary(buffer, value);
}
@Override
public String readBinaryValue(ByteBuf buffer, int valueLength)
{
byte[] utf8 = new byte[valueLength];
buffer.readBytes(utf8);
return new String(utf8, UTF_8);
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull String value)
{
return value.getBytes(UTF_8);
}
@Override
public String decodeUTF8Text(byte[] bytes)
{
return new String(bytes, UTF_8);
}
@Override
public Object getEmptyValue()
{
return "";
}
}
wren-base/src/main/java/io/wren/base/type/SmallIntType.java
-89
View File
@@ -1,89 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import java.nio.charset.StandardCharsets;
import static io.wren.base.Utils.checkArgument;
public class SmallIntType
extends PGType<Short>
{
public static final SmallIntType SMALLINT = new SmallIntType();
private static final int OID = 21;
private static final int TYPE_LEN = 2;
private static final int TYPE_MOD = -1;
private SmallIntType()
{
super(OID, TYPE_LEN, TYPE_MOD, "int2");
}
@Override
public int typArray()
{
return PGArray.INT2_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.NUMERIC.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull Short value)
{
buffer.writeInt(TYPE_LEN);
buffer.writeShort(value);
return INT32_BYTE_SIZE + TYPE_LEN;
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull Short value)
{
return Short.toString(value).getBytes(StandardCharsets.UTF_8);
}
@Override
public Short readBinaryValue(ByteBuf buffer, int valueLength)
{
checkArgument(valueLength == TYPE_LEN, "length should be %s because short is int16. Actual length: %s", TYPE_LEN, valueLength);
return buffer.readShort();
}
@Override
public Short decodeUTF8Text(byte[] bytes)
{
return Short.parseShort(new String(bytes, StandardCharsets.UTF_8));
}
@Override
public Object getEmptyValue()
{
return 0;
}
}
wren-base/src/main/java/io/wren/base/type/StandardTypes.java
-52
View File
@@ -1,52 +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.wren.base.type;
public final class StandardTypes
{
public static final String BIGINT = "bigint";
public static final String INTEGER = "integer";
public static final String SMALLINT = "smallint";
public static final String TINYINT = "tinyint";
public static final String BOOLEAN = "boolean";
public static final String DATE = "date";
public static final String DECIMAL = "decimal";
public static final String REAL = "real";
public static final String DOUBLE = "double";
public static final String HYPER_LOG_LOG = "HyperLogLog";
public static final String QDIGEST = "qdigest";
public static final String TDIGEST = "tdigest";
public static final String P4_HYPER_LOG_LOG = "P4HyperLogLog";
public static final String INTERVAL_DAY_TO_SECOND = "interval day to second";
public static final String INTERVAL_YEAR_TO_MONTH = "interval year to month";
public static final String TIMESTAMP = "timestamp";
public static final String TIMESTAMP_WITH_TIME_ZONE = "timestamp with time zone";
public static final String TIME = "time";
public static final String TIME_WITH_TIME_ZONE = "time with time zone";
public static final String BYTEA = "bytea";
public static final String VARCHAR = "varchar";
public static final String CHAR = "char";
public static final String ROW = "row";
public static final String TEXT = "text";
public static final String NAME = "name";
public static final String ARRAY = "array";
public static final String MAP = "map";
public static final String JSON = "json";
public static final String IPADDRESS = "ipaddress";
public static final String GEOMETRY = "Geometry";
public static final String BING_TILE = "BingTile";
public static final String UUID = "uuid";
private StandardTypes() {}
}
wren-base/src/main/java/io/wren/base/type/TimestampType.java
-98
View File
@@ -1,98 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.time.format.DateTimeFormatterBuilder;
import java.time.format.ResolverStyle;
import java.util.Locale;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.time.format.DateTimeFormatter.ISO_LOCAL_DATE;
import static java.time.format.DateTimeFormatter.ISO_LOCAL_TIME;
import static java.util.Locale.ENGLISH;
public class TimestampType
extends BaseTimestampType
{
public static final PGType<?> TIMESTAMP = new TimestampType();
private static final int OID = 1114;
private static final String NAME = "timestamp";
// TODO support timestamp with precision dynamically
// BigQuery support precision with 6
private static final DateTimeFormatter PG_TIMESTAMP = new DateTimeFormatterBuilder()
.parseCaseInsensitive()
.appendPattern("yyyy-MM-dd HH:mm:ss.SSSSSS")
.toFormatter(ENGLISH)
.withResolverStyle(ResolverStyle.STRICT);
private static final DateTimeFormatter PARSER_WITH_OPTIONAL_ERA = new DateTimeFormatterBuilder()
.parseCaseInsensitive()
.append(ISO_LOCAL_DATE)
.optionalStart()
.appendLiteral(' ')
.append(ISO_LOCAL_TIME)
.optionalStart()
.appendPattern("[VV][x][xx][xxx]")
.optionalStart()
.appendLiteral(' ')
.appendPattern("G")
.toFormatter(Locale.ENGLISH).withResolverStyle(ResolverStyle.STRICT);
private TimestampType()
{
super(OID, TYPE_LEN, TYPE_MOD, NAME);
}
@Override
public int typArray()
{
return PGArray.TIMESTAMP_ARRAY.oid();
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull Object value)
{
LocalDateTime dt = (LocalDateTime) value;
return PG_TIMESTAMP.format(dt).getBytes(UTF_8);
}
@Override
public Object decodeUTF8Text(byte[] bytes)
{
String s = new String(bytes, UTF_8);
LocalDateTime dt = LocalDateTime.parse(s, PARSER_WITH_OPTIONAL_ERA);
return PG_TIMESTAMP.format(dt);
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull Object value)
{
throw new UnsupportedOperationException();
}
@Override
public Object getEmptyValue()
{
return "1970-01-01 00:00:00";
}
}
wren-base/src/main/java/io/wren/base/type/TimestampWithTimeZoneType.java
-106
View File
@@ -1,106 +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.wren.base.type;
import com.google.common.annotations.VisibleForTesting;
import io.netty.buffer.ByteBuf;
import org.joda.time.format.DateTimeFormat;
import javax.annotation.Nonnull;
import java.time.Instant;
import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.time.format.DateTimeFormatterBuilder;
import java.time.format.ResolverStyle;
import java.time.temporal.TemporalAccessor;
import java.util.Locale;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.time.format.DateTimeFormatter.ISO_LOCAL_DATE;
import static java.time.format.DateTimeFormatter.ISO_LOCAL_TIME;
import static java.time.temporal.ChronoField.MILLI_OF_SECOND;
public class TimestampWithTimeZoneType
extends BaseTimestampType
{
public static final PGType TIMESTAMP_WITH_TIMEZONE = new TimestampWithTimeZoneType();
public static final DateTimeFormatter PG_TIMESTAMP = new DateTimeFormatterBuilder()
.parseCaseInsensitive()
.append(ISO_LOCAL_DATE)
.optionalStart()
.appendLiteral(' ')
.append(ISO_LOCAL_TIME)
.optionalStart()
.appendLiteral(' ')
.optionalEnd()
.appendPattern("[VV][x][xx][xxx][z]")
.toFormatter(Locale.ENGLISH).withResolverStyle(ResolverStyle.STRICT);
public static final org.joda.time.format.DateTimeFormatter ISO_FORMATTER =
DateTimeFormat.forPattern("YYYY-MM-dd HH:mm:ss.SSSZ").withLocale(Locale.ENGLISH);
private static final int OID = 1184;
private static final String NAME = "timestamptz";
private TimestampWithTimeZoneType()
{
super(OID, TYPE_LEN, TYPE_MOD, NAME);
}
@Override
public int typArray()
{
return PGArray.TIMESTAMP_WITH_TIMEZONE_ARRAY.oid();
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull Object value)
{
// TODO: consider AD, BC and dynamic fraction precision.
if (value instanceof String) {
return ((String) value).getBytes(UTF_8);
}
return ISO_FORMATTER.print(Instant.from((TemporalAccessor) value).getLong(MILLI_OF_SECOND))
.getBytes(UTF_8);
}
@Override
public Object decodeUTF8Text(byte[] bytes)
{
String dtString = new String(bytes, UTF_8);
ZonedDateTime zonedDateTime = tryParse(dtString);
return ISO_FORMATTER.print(Instant.from(zonedDateTime).toEpochMilli());
}
@VisibleForTesting
ZonedDateTime tryParse(String timeString)
{
// Postgres TimestampTz format
return ZonedDateTime.parse(timeString, PG_TIMESTAMP);
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull Object value)
{
throw new UnsupportedOperationException("Not implemented yet");
}
@Override
public Object getEmptyValue()
{
return "1970-01-01 00:00:00.000000+01:00";
}
}
wren-base/src/main/java/io/wren/base/type/TinyIntType.java
-93
View File
@@ -1,93 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import io.wren.base.WrenException;
import javax.annotation.Nonnull;
import static io.wren.base.metadata.StandardErrorCode.GENERIC_INTERNAL_ERROR;
import static java.nio.charset.StandardCharsets.UTF_8;
/**
* This class is write-only class. We never read any data by it and won't show it to pg_type list.
* It is a wrap of SmallIntType. Actually, PostgreSQL doesn't have TinyInt.
* To handle Presto TinyInt, we use this to transform Presto TinyInt to Pg SmallInt.
*/
public class TinyIntType
extends PGType<Byte>
{
public static final TinyIntType TINYINT = new TinyIntType();
private static final int OID = 21;
private static final int TYPE_LEN = 2;
private static final int TYPE_MOD = -1;
private TinyIntType()
{
super(OID, TYPE_LEN, TYPE_MOD, "int2");
}
@Override
public int typArray()
{
return PGArray.INT2_ARRAY.oid();
}
@Override
public String typeCategory()
{
return TypeCategory.NUMERIC.code();
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull Byte value)
{
buffer.writeInt(TYPE_LEN);
buffer.writeShort(value);
return INT32_BYTE_SIZE + TYPE_LEN;
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull Byte value)
{
return Byte.toString(value).getBytes(UTF_8);
}
@Override
public Byte readBinaryValue(ByteBuf buffer, int valueLength)
{
throw new WrenException(GENERIC_INTERNAL_ERROR, new IllegalAccessException("PostgreSQL doesn't have TinyIntType. We never read TinyInt from client."));
}
@Override
public Byte decodeUTF8Text(byte[] bytes)
{
throw new WrenException(GENERIC_INTERNAL_ERROR, new IllegalAccessException("PostgreSQL doesn't have TinyIntType. We never read TinyInt from client."));
}
@Override
public Object getEmptyValue()
{
return 0;
}
}
wren-base/src/main/java/io/wren/base/type/UuidType.java
-89
View File
@@ -1,89 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import static java.nio.charset.StandardCharsets.UTF_8;
public class UuidType
extends PGType<String>
{
public static final UuidType UUID = new UuidType();
static final int OID = 2950;
private static final int TYPE_LEN = -1;
private static final int TYPE_MOD = -1;
private UuidType()
{
super(OID, TYPE_LEN, TYPE_MOD, "uuid");
}
@Override
public int typArray()
{
return PGArray.UUID_ARRAY.oid();
}
@Override
public String typeCategory()
{
return PGType.TypeCategory.USER_DEFINED_TYPES.code();
}
@Override
public String type()
{
return PGType.Type.BASE.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull String value)
{
byte[] bytes = encodeAsUTF8Text(value);
buffer.writeInt(bytes.length);
buffer.writeBytes(bytes);
return INT32_BYTE_SIZE + bytes.length;
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull String value)
{
return value.getBytes(UTF_8);
}
@Override
public String readBinaryValue(ByteBuf buffer, int valueLength)
{
byte[] bytes = new byte[valueLength];
buffer.readBytes(bytes);
return decodeUTF8Text(bytes);
}
@Override
public String decodeUTF8Text(byte[] bytes)
{
return new String(bytes, UTF_8);
}
@Override
public Object getEmptyValue()
{
return "";
}
}
wren-base/src/main/java/io/wren/base/type/VarcharType.java
-121
View File
@@ -1,121 +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.wren.base.type;
import io.netty.buffer.ByteBuf;
import javax.annotation.Nonnull;
import java.nio.charset.StandardCharsets;
public class VarcharType
extends PGType<String>
{
static final int OID = 1043;
private static final int ARRAY_OID = 1015;
private static final int TYPE_LEN = -1;
private static final int TYPE_MOD = -1;
public static final VarcharType VARCHAR = new VarcharType(ARRAY_OID);
private final int typArray;
private VarcharType(int typArray)
{
super(OID, TYPE_LEN, TYPE_MOD, "varchar");
this.typArray = typArray;
}
private VarcharType(int oid, int typArray, int maxLength, String aliasName)
{
super(oid, maxLength, TYPE_MOD, aliasName);
this.typArray = typArray;
}
@Override
public int typArray()
{
return typArray;
}
@Override
public String type()
{
return Type.BASE.code();
}
@Override
public String typeCategory()
{
return TypeCategory.STRING.code();
}
@Override
public int writeAsBinary(ByteBuf buffer, @Nonnull String value)
{
byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
buffer.writeInt(bytes.length);
buffer.writeBytes(bytes);
return INT32_BYTE_SIZE + bytes.length;
}
@Override
public int writeAsText(ByteBuf buffer, @Nonnull String value)
{
return writeAsBinary(buffer, value);
}
@Override
public byte[] encodeAsUTF8Text(@Nonnull String value)
{
return value.getBytes(StandardCharsets.UTF_8);
}
@Override
public String readBinaryValue(ByteBuf buffer, int valueLength)
{
byte[] utf8 = new byte[valueLength];
buffer.readBytes(utf8);
return new String(utf8, StandardCharsets.UTF_8);
}
@Override
public String decodeUTF8Text(byte[] bytes)
{
return new String(bytes, StandardCharsets.UTF_8);
}
@Override
public Object getEmptyValue()
{
return "";
}
public static class NameType
{
static final int OID = 19;
private static final int ARRAY_OID = -1;
private static final int TYPE_LEN = 64;
public static final VarcharType NAME = new VarcharType(OID, ARRAY_OID, TYPE_LEN, "name");
}
public static class TextType
{
static final int OID = 25;
static final int TEXT_ARRAY_OID = 1009;
public static final VarcharType TEXT = new VarcharType(OID, TEXT_ARRAY_OID, -1, "text");
}
}
wren-base/src/main/java/io/wren/base/type/parser/PgArrayASTVisitor.java
-85
View File
@@ -1,85 +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.wren.base.type.parser;
import com.carrotsearch.hppc.ByteArrayList;
import io.wren.base.type.parser.antlr.v4.PgArrayBaseVisitor;
import io.wren.base.type.parser.antlr.v4.PgArrayParser;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.function.Function;
class PgArrayASTVisitor
extends PgArrayBaseVisitor<Object>
{
private final Function<byte[], Object> convert;
PgArrayASTVisitor(Function<byte[], Object> convert)
{
this.convert = convert;
}
@Override
public Object visitArray(PgArrayParser.ArrayContext ctx)
{
ArrayList<Object> array = new ArrayList<>();
for (PgArrayParser.ItemContext value : ctx.item()) {
array.add(value.accept(this));
}
return array;
}
@Override
public Object visitUnquotedString(PgArrayParser.UnquotedStringContext ctx)
{
String text = ctx.getText();
return convert.apply(text.getBytes(StandardCharsets.UTF_8));
}
@Override
public Object visitNull(PgArrayParser.NullContext ctx)
{
return null;
}
@Override
public Object visitQuotedString(PgArrayParser.QuotedStringContext ctx)
{
String text = ctx.getText();
String withoutQuotes = text.substring(1, text.length() - 1);
return convert.apply(removeEscapes(withoutQuotes.getBytes(StandardCharsets.UTF_8)));
}
/**
* @param bytes {@code byte[]} that represent an array's item.
*/
private static byte[] removeEscapes(byte[] bytes)
{
ByteArrayList itemBytes = new ByteArrayList(bytes.length);
int end = bytes.length - 1;
for (int i = 0; i <= end; i++) {
char c = (char) bytes[i];
if (i < end) {
char next = (char) bytes[i + 1];
if (c == '\\' && (next == '\\' || next == '\"')) {
i++;
}
}
itemBytes.add(bytes[i]);
}
return itemBytes.toArray();
}
}
wren-base/src/main/java/io/wren/base/type/parser/PgArrayParserWrapper.java
-102
View File
@@ -1,102 +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.wren.base.type.parser;
import io.wren.base.WrenException;
import io.wren.base.metadata.StandardErrorCode;
import io.wren.base.type.parser.antlr.v4.PgArrayLexer;
import io.wren.base.type.parser.antlr.v4.PgArrayParser;
import org.antlr.v4.runtime.BaseErrorListener;
import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.ParserRuleContext;
import org.antlr.v4.runtime.RecognitionException;
import org.antlr.v4.runtime.Recognizer;
import org.antlr.v4.runtime.atn.PredictionMode;
import org.antlr.v4.runtime.misc.ParseCancellationException;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.function.Function;
public class PgArrayParserWrapper
{
private static final BaseErrorListener ERROR_LISTENER = new BaseErrorListener()
{
@Override
public void syntaxError(Recognizer<?, ?> recognizer,
Object offendingSymbol,
int line,
int charPositionInLine,
String message,
RecognitionException e)
{
throw new WrenException(StandardErrorCode.GENERIC_INTERNAL_ERROR, e);
}
};
private static final PgArrayParserWrapper INSTANCE = new PgArrayParserWrapper();
public static Object parse(byte[] bytes, Function<byte[], Object> convert)
{
return INSTANCE.invokeParser(
new ByteArrayInputStream(bytes),
PgArrayParser::array,
convert);
}
private Object invokeParser(InputStream inputStream,
Function<PgArrayParser, ParserRuleContext> parseFunction,
Function<byte[], Object> convert)
{
try {
PgArrayLexer lexer = new PgArrayLexer(CharStreams.fromStream(
inputStream,
StandardCharsets.UTF_8));
CommonTokenStream tokenStream = new CommonTokenStream(lexer);
PgArrayParser parser = new PgArrayParser(tokenStream);
lexer.removeErrorListeners();
lexer.addErrorListener(ERROR_LISTENER);
parser.removeErrorListeners();
parser.addErrorListener(ERROR_LISTENER);
ParserRuleContext tree;
try {
// first, try parsing with potentially faster SLL mode
parser.getInterpreter().setPredictionMode(PredictionMode.SLL);
tree = parseFunction.apply(parser);
}
catch (ParseCancellationException ex) {
// if we fail, parse with LL mode
tokenStream.seek(0); // rewind input stream
parser.reset();
parser.getInterpreter().setPredictionMode(PredictionMode.LL);
tree = parseFunction.apply(parser);
}
return tree.accept(new PgArrayASTVisitor(convert));
}
catch (StackOverflowError e) {
throw new WrenException(StandardErrorCode.GENERIC_INTERNAL_ERROR, "stack overflow while parsing: " + e.getLocalizedMessage());
}
catch (IOException e) {
return new IllegalArgumentException(e);
}
}
}
wren-base/src/main/java/io/wren/base/type/parser/PgDateTimeFormatParserWrapper.java
-176
View File
@@ -1,176 +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.wren.base.type.parser;
import io.wren.base.WrenException;
import io.wren.base.metadata.StandardErrorCode;
import io.wren.base.type.parser.antlr.v4.PgDateTimeFormatBaseVisitor;
import io.wren.base.type.parser.antlr.v4.PgDateTimeFormatLexer;
import io.wren.base.type.parser.antlr.v4.PgDateTimeFormatParser;
import org.antlr.v4.runtime.BaseErrorListener;
import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.Lexer;
import org.antlr.v4.runtime.RecognitionException;
import org.antlr.v4.runtime.Recognizer;
import org.antlr.v4.runtime.tree.ParseTree;
import java.util.Locale;
/**
* this parser adapts patterns from PostgreSQL to Joda.
* <p>
* for more info about patterns of PostgreSQL, refer to https://www.postgresql.org/docs/13/functions-formatting.html
* for more info about patterns of Joda, refer to https://help.gooddata.com/cloudconnect/manual/date-and-time-format.html
*/
public final class PgDateTimeFormatParserWrapper
{
private PgDateTimeFormatParserWrapper() {}
private static final BaseErrorListener ERROR_LISTENER = new BaseErrorListener()
{
@Override
public void syntaxError(Recognizer<?, ?> recognizer,
Object offendingSymbol,
int line,
int charPositionInLine,
String message,
RecognitionException e)
{
throw new WrenException(StandardErrorCode.GENERIC_INTERNAL_ERROR, e);
}
};
public static String parse(String origin)
{
Lexer lexer = new PgDateTimeFormatLexer(CharStreams.fromString(origin));
lexer.removeErrorListeners();
lexer.addErrorListener(ERROR_LISTENER);
PgDateTimeFormatParser parser = new PgDateTimeFormatParser(new CommonTokenStream(lexer));
parser.removeErrorListeners();
parser.addErrorListener(ERROR_LISTENER);
ParseTree tree = parser.format();
return new AstBuilder().visit(tree);
}
private static class AstBuilder
extends PgDateTimeFormatBaseVisitor<String>
{
@Override
public String visitFormat(PgDateTimeFormatParser.FormatContext ctx)
{
StringBuilder builder = new StringBuilder();
ctx.symbol().forEach(symbol -> builder.append(visit(symbol)));
return builder.toString();
}
@Override
public String visitSeparator(PgDateTimeFormatParser.SeparatorContext ctx)
{
return ctx.getText();
}
@Override
public String visitHourLiteral(PgDateTimeFormatParser.HourLiteralContext ctx)
{
if (ctx.getText().endsWith("24")) {
return "HH";
}
return "hh";
}
@Override
public String visitMinuteLiteral(PgDateTimeFormatParser.MinuteLiteralContext ctx)
{
return "mm";
}
@Override
public String visitSecondLiteral(PgDateTimeFormatParser.SecondLiteralContext ctx)
{
return "ss";
}
@Override
public String visitMilliSecondLiteral(PgDateTimeFormatParser.MilliSecondLiteralContext ctx)
{
return "SSS";
}
@Override
public String visitMeridiemMarkerLiteral(PgDateTimeFormatParser.MeridiemMarkerLiteralContext ctx)
{
return "a";
}
@Override
public String visitEraDesignatorLiteral(PgDateTimeFormatParser.EraDesignatorLiteralContext ctx)
{
return "G";
}
@Override
public String visitTimeZoneLiteral(PgDateTimeFormatParser.TimeZoneLiteralContext ctx)
{
return "z";
}
@Override
public String visitYearLiteral(PgDateTimeFormatParser.YearLiteralContext ctx)
{
return ctx.getText().toUpperCase(Locale.ROOT);
}
@Override
public String visitMonthLiteral(PgDateTimeFormatParser.MonthLiteralContext ctx)
{
String month = ctx.getText();
if (month.equals("Month")) {
return "MMMM";
}
if (month.equals("Mon")) {
return "MMM";
}
return month.toUpperCase(Locale.ROOT);
}
@Override
public String visitWeekLiteral(PgDateTimeFormatParser.WeekLiteralContext ctx)
{
return ctx.getText().toLowerCase(Locale.ROOT);
}
@Override
public String visitDayLiteral(PgDateTimeFormatParser.DayLiteralContext ctx)
{
String day = ctx.getText();
if (day.equals("Day")) {
return "EEEE";
}
if (day.equals("Dy")) {
return "EEE";
}
if (day.length() == 2) {
return "dd";
}
if (day.length() == 1) {
return "e";
}
return day.toUpperCase(Locale.ROOT);
}
}
}
wren-base/src/main/java/io/wren/base/type/parser/PgNumericFormatParserWrapper.java
-112
View File
@@ -1,112 +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.wren.base.type.parser;
import io.wren.base.WrenException;
import io.wren.base.metadata.StandardErrorCode;
import io.wren.base.type.parser.antlr.v4.PgNumericFormatBaseVisitor;
import io.wren.base.type.parser.antlr.v4.PgNumericFormatLexer;
import io.wren.base.type.parser.antlr.v4.PgNumericFormatParser;
import org.antlr.v4.runtime.BaseErrorListener;
import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.Lexer;
import org.antlr.v4.runtime.RecognitionException;
import org.antlr.v4.runtime.Recognizer;
import org.antlr.v4.runtime.tree.ParseTree;
public final class PgNumericFormatParserWrapper
{
private PgNumericFormatParserWrapper() {}
private static final BaseErrorListener ERROR_LISTENER = new BaseErrorListener()
{
@Override
public void syntaxError(Recognizer<?, ?> recognizer,
Object offendingSymbol,
int line,
int charPositionInLine,
String message,
RecognitionException e)
{
throw new WrenException(StandardErrorCode.GENERIC_INTERNAL_ERROR, e);
}
};
public static String parse(String origin)
{
Lexer lexer = new PgNumericFormatLexer(CharStreams.fromString(origin));
lexer.removeErrorListeners();
lexer.addErrorListener(ERROR_LISTENER);
PgNumericFormatParser parser = new PgNumericFormatParser(new CommonTokenStream(lexer));
parser.removeErrorListeners();
parser.addErrorListener(ERROR_LISTENER);
ParseTree tree = parser.format();
return new AstBuilder().visit(tree);
}
private static class AstBuilder
extends PgNumericFormatBaseVisitor<String>
{
@Override
public String visitFormat(PgNumericFormatParser.FormatContext ctx)
{
StringBuilder builder = new StringBuilder();
ctx.pattern().forEach(pattern -> builder.append(visit(pattern)));
return builder.toString();
}
@Override
public String visitDigitPattern(PgNumericFormatParser.DigitPatternContext ctx)
{
if (ctx.getText().equals("9")) {
return "#";
}
return "0";
}
@Override
public String visitDecimalPointPattern(PgNumericFormatParser.DecimalPointPatternContext ctx)
{
return ".";
}
@Override
public String visitGroupSeparatorPattern(PgNumericFormatParser.GroupSeparatorPatternContext ctx)
{
return ",";
}
@Override
public String visitCurrencySymbolPattern(PgNumericFormatParser.CurrencySymbolPatternContext ctx)
{
return "\u00A4";
}
@Override
public String visitExponentPattern(PgNumericFormatParser.ExponentPatternContext ctx)
{
return "E00";
}
@Override
public String visitNonReservedPattern(PgNumericFormatParser.NonReservedPatternContext ctx)
{
throw new WrenException(StandardErrorCode.GENERIC_INTERNAL_ERROR, String.format("we didn't support the pattern %s.", ctx.getText()));
}
}
}
wren-base/src/main/java/io/wren/base/wireprotocol/PgMetastore.java
-43
View File
@@ -1,43 +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.wren.base.wireprotocol;
import io.wren.base.ConnectorRecordIterator;
import io.wren.base.Parameter;
import io.wren.base.client.Client;
import io.wren.base.sql.SqlConverter;
import java.util.List;
public interface PgMetastore
{
void directDDL(String sql);
ConnectorRecordIterator directQuery(String sql, List<Parameter> parameters);
String handlePgType(String type);
String getPgCatalogName();
boolean isSchemaExist(String schemaName);
void dropTableIfExists(String name);
Client getClient();
SqlConverter getSqlConverter();
void close();
}
wren-base/src/test/java/io/wren/base/sqlrewrite/AbstractTestFramework.java
+1 -4
View File
@@ -15,13 +15,11 @@
package io.wren.base.sqlrewrite;
import com.google.common.collect.ImmutableList;
import io.trino.sql.parser.SqlParser;
import io.wren.base.SessionContext;
import io.wren.base.client.AutoCloseableIterator;
import io.wren.base.client.duckdb.DuckDBConfig;
import io.wren.base.client.duckdb.DuckDBSettingSQL;
import io.wren.base.client.duckdb.DuckdbClient;
import io.wren.base.client.duckdb.DuckdbS3StyleStorageConfig;
import io.wren.base.dto.Column;
import io.wren.base.dto.Manifest;
import io.wren.base.dto.Model;
@@ -38,7 +36,6 @@ import static io.wren.base.sqlrewrite.Utils.parseSql;
public abstract class AbstractTestFramework
{
private static final SqlParser SQL_PARSER = new SqlParser();
public static final SessionContext DEFAULT_SESSION_CONTEXT =
SessionContext.builder().setCatalog("wren").setSchema("test").build();
private DuckdbClient duckdbClient;
@@ -70,7 +67,7 @@ public abstract class AbstractTestFramework
@BeforeClass
public void init()
{
duckdbClient = new DuckdbClient(new DuckDBConfig(), new DuckdbS3StyleStorageConfig(), new DuckDBSettingSQL());
duckdbClient = new DuckdbClient(new DuckDBConfig(), new DuckDBSettingSQL());
prepareData();
}
wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestExpressionTypeAnalyzer.java
-161
View File
@@ -1,161 +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.wren.base.sqlrewrite.analyzer;
import io.wren.base.CatalogSchemaTableName;
import io.wren.base.WrenMDL;
import io.wren.base.WrenTypes;
import io.wren.base.dto.Column;
import io.wren.base.dto.Model;
import io.wren.base.sqlrewrite.AbstractTestFramework;
import io.wren.base.type.BigIntType;
import io.wren.base.type.BooleanType;
import io.wren.base.type.ByteaType;
import io.wren.base.type.DateType;
import io.wren.base.type.DoubleType;
import io.wren.base.type.IntegerType;
import io.wren.base.type.IntervalType;
import io.wren.base.type.PGArray;
import io.wren.base.type.RealType;
import io.wren.base.type.RecordType;
import io.wren.base.type.TimestampType;
import io.wren.base.type.VarcharType;
import org.testng.annotations.Test;
import java.util.List;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static io.wren.base.sqlrewrite.Utils.parseExpression;
import static io.wren.base.sqlrewrite.analyzer.ExpressionTypeAnalyzer.analyze;
import static org.assertj.core.api.Assertions.assertThat;
public class TestExpressionTypeAnalyzer
extends AbstractTestFramework
{
private static final Scope EMPTY_SCOPE = Scope.builder().build();
private static final WrenMDL EMPTY_MDL = WrenMDL.fromManifest(withDefaultCatalogSchema().build());
private final Model customer;
public TestExpressionTypeAnalyzer()
{
customer = Model.model("Customer",
"select * from main.customer",
List.of(
Column.column("custkey", WrenTypes.INTEGER, null, true),
Column.column("name", WrenTypes.VARCHAR, null, true),
Column.column("address", WrenTypes.VARCHAR, null, true),
Column.column("nationkey", WrenTypes.INTEGER, null, true),
Column.column("phone", WrenTypes.VARCHAR, null, true),
Column.column("acctbal", WrenTypes.INTEGER, null, true),
Column.column("mktsegment", WrenTypes.VARCHAR, null, true),
Column.column("comment", WrenTypes.VARCHAR, null, true)),
"custkey");
}
@Test
public void testLiteral()
{
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("1"))).isEqualTo(BigIntType.BIGINT);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("'abc'"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("INTERVAL '1 month'"))).isEqualTo(IntervalType.INTERVAL);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("NULL"))).isEqualTo(null);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("TIMESTAMP '2023-10-29 00:00:00.000000'"))).isEqualTo(TimestampType.TIMESTAMP);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("REAL '3.5'"))).isEqualTo(RealType.REAL);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("x'65683F'"))).isEqualTo(ByteaType.BYTEA);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("1.1"))).isEqualTo(DoubleType.DOUBLE);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("10.3e0"))).isEqualTo(DoubleType.DOUBLE);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("false"))).isEqualTo(BooleanType.BOOLEAN);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("cast(1.1 as DOUBLE)"))).isEqualTo(DoubleType.DOUBLE);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("ROW(1, 2e0)"))).isEqualTo(RecordType.EMPTY_RECORD);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("array[1,2,3]"))).isEqualTo(PGArray.INT8_ARRAY);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("array['a','b','c']"))).isEqualTo(PGArray.VARCHAR_ARRAY);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("current_user"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("current_schema"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("current_catalog"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("current_path"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("current_date"))).isEqualTo(DateType.DATE);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("current_time"))).isEqualTo(TimestampType.TIMESTAMP);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("current_timestamp"))).isEqualTo(TimestampType.TIMESTAMP);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("localtime"))).isEqualTo(TimestampType.TIMESTAMP);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("localtimestamp"))).isEqualTo(TimestampType.TIMESTAMP);
}
@Test
public void testPredicate()
{
assertPredicate("x > 1");
assertPredicate("x >= 1");
assertPredicate("x < 1");
assertPredicate("x <= 1");
assertPredicate("x = 1");
assertPredicate("x <> y");
assertPredicate("x != INTERVAL '1 month'");
assertPredicate("x IS NULL");
assertPredicate("x IS NOT NULL");
assertPredicate("x in (1, 2, 3)");
assertPredicate("x like 'abc'");
assertPredicate("x between 1 and 2");
assertPredicate("x > 1 and y < 2");
assertPredicate("x > 1 or y < 2");
assertPredicate("not x > 1");
assertPredicate("exists (select 1)");
}
private void assertPredicate(String expression)
{
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression(expression))).isEqualTo(BooleanType.BOOLEAN);
}
@Test
public void testFunction()
{
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("date_trunc('day', create_date)"))).isEqualTo(DateType.DATE);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("now()"))).isEqualTo(TimestampType.TIMESTAMP);
assertThat(analyze(EMPTY_MDL, EMPTY_SCOPE, parseExpression("now___timestamp()"))).isEqualTo(TimestampType.TIMESTAMP);
}
@Test
public void testColumns()
{
WrenMDL mdl = WrenMDL.fromManifest(withDefaultCatalogSchema().setModels(List.of(customer)).build());
List<Field> fields = customer.getColumns().stream()
.map(column -> Field.builder()
.tableName(new CatalogSchemaTableName(mdl.getCatalog(), mdl.getSchema(), customer.getName()))
.columnName(column.getName())
.name(column.getName())
.build())
.collect(toImmutableList());
Scope scope = Scope.builder().relationType(new RelationType(fields)).build();
assertThat(analyze(mdl, scope, parseExpression("custkey"))).isEqualTo(IntegerType.INTEGER);
assertThat(analyze(mdl, scope, parseExpression("name"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(mdl, scope, parseExpression("address"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(mdl, scope, parseExpression("nationkey"))).isEqualTo(IntegerType.INTEGER);
assertThat(analyze(mdl, scope, parseExpression("phone"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(mdl, scope, parseExpression("acctbal"))).isEqualTo(IntegerType.INTEGER);
assertThat(analyze(mdl, scope, parseExpression("mktsegment"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(mdl, scope, parseExpression("comment"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(mdl, scope, parseExpression("Customer.custkey"))).isEqualTo(IntegerType.INTEGER);
assertThat(analyze(mdl, scope, parseExpression("Customer.name"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(mdl, scope, parseExpression("Customer.address"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(mdl, scope, parseExpression("Customer.nationkey"))).isEqualTo(IntegerType.INTEGER);
assertThat(analyze(mdl, scope, parseExpression("Customer.phone"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(mdl, scope, parseExpression("Customer.acctbal"))).isEqualTo(IntegerType.INTEGER);
assertThat(analyze(mdl, scope, parseExpression("Customer.mktsegment"))).isEqualTo(VarcharType.VARCHAR);
assertThat(analyze(mdl, scope, parseExpression("Customer.comment"))).isEqualTo(VarcharType.VARCHAR);
}
}
wren-base/src/test/java/io/wren/base/type/TestTimestampWithTimeZoneTypeParsing.java
-75
View File
@@ -1,75 +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.
*/
/*
* 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.wren.base.type;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import static org.assertj.core.api.Assertions.assertThatNoException;
public class TestTimestampWithTimeZoneTypeParsing
{
@DataProvider
public Object[][] timestampWithTimeZone()
{
return new Object[][] {
{"2004-10-19 10:23:54 +02"},
{"1999-01-08 04:05:06 PST"},
{"1999-01-08 04:05:06 PST8PDT"},
{"1999-01-08 04:05:06 zulu"},
{"1999-01-08 04:05:06 z"},
{"1999-01-08 04:05:06 America/New_York"},
{"2004-10-19 10:23:54+02"},
{"1999-01-08 04:05:06PST"},
{"1999-01-08 04:05:06PST8PDT"},
{"1999-01-08 04:05:06zulu"},
{"1999-01-08 04:05:06z"},
// TODO: unsupported pg pattern
// {"1999-01-08 04:05:06 -8:00"},
// {"1999-01-08 04:05:06-8:00"},
// {"1999-01-08 04:05:06 -8:00:00"},
// {"1999-01-08 04:05:06 -8:00"},
// {"1999-01-08 04:05:06 -800"},
// {"1999-01-08 04:05:06-8:00:00"},
// {"1999-01-08 04:05:06-8:00"},
// {"1999-01-08 04:05:06-800"},
// {"1999-01-08 04:05:06-8"},
// {"1999-01-08 04:05:06 -8"},
// {"1999-01-08 04:05:06Americ/New_York"},
};
}
@Test(dataProvider = "timestampWithTimeZone")
public void testParsing(String timestampString)
{
assertThatNoException()
.isThrownBy(() -> TimestampWithTimeZoneType.PG_TIMESTAMP.parse(timestampString));
}
}
wren-cache/pom.xml
-85
View File
@@ -1,85 +0,0 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ 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.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
<groupId>io.wren</groupId>
<artifactId>wren-root</artifactId>
<version>0.5.3-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
<artifactId>wren-cache</artifactId>
<name>wren-cache</name>
<description>WrenEngine - Cache</description>
<properties>
<air.main.basedir>${project.parent.basedir}</air.main.basedir>
</properties>
<dependencies>
<dependency>
<groupId>com.fasterxml.jackson.core</groupId>
<artifactId>jackson-annotations</artifactId>
</dependency>
<dependency>
<groupId>com.google.guava</groupId>
<artifactId>guava</artifactId>
</dependency>
<dependency>
<groupId>com.google.inject</groupId>
<artifactId>guice</artifactId>
</dependency>
<dependency>
<groupId>io.airlift</groupId>
<artifactId>concurrent</artifactId>
</dependency>
<dependency>
<groupId>io.airlift</groupId>
<artifactId>configuration</artifactId>
</dependency>
<dependency>
<groupId>io.airlift</groupId>
<artifactId>log</artifactId>
</dependency>
<dependency>
<groupId>io.airlift</groupId>
<artifactId>units</artifactId>
</dependency>
<dependency>
<groupId>io.wren</groupId>
<artifactId>trino-parser</artifactId>
</dependency>
<dependency>
<groupId>io.wren</groupId>
<artifactId>wren-base</artifactId>
</dependency>
<dependency>
<groupId>org.duckdb</groupId>
<artifactId>duckdb_jdbc</artifactId>
</dependency>
</dependencies>
</project>
wren-cache/src/main/java/io/wren/cache/CacheInfoPair.java
Vendored
-68
View File
@@ -1,68 +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.wren.cache;
import io.wren.base.WrenException;
import io.wren.base.dto.CacheInfo;
import java.util.Optional;
import static io.wren.base.metadata.StandardErrorCode.GENERIC_USER_ERROR;
import static java.util.Objects.requireNonNull;
public class CacheInfoPair
{
private final CacheInfo cacheInfo;
private final Optional<String> tableName;
private final Optional<String> errorMessage;
private final long createTime;
protected CacheInfoPair(CacheInfo cacheInfo, String tableName, long createTime)
{
this(cacheInfo, Optional.of(tableName), Optional.empty(), createTime);
}
protected CacheInfoPair(CacheInfo cacheInfo, Optional<String> tableName, Optional<String> errorMessage, long createTime)
{
this.cacheInfo = requireNonNull(cacheInfo, "cacheInfo is null");
this.tableName = requireNonNull(tableName, "tableName is null");
this.errorMessage = requireNonNull(errorMessage, "errorMessage is null");
this.createTime = createTime;
}
public CacheInfo getCacheInfo()
{
return cacheInfo;
}
public String getRequiredTableName()
{
return tableName.orElseThrow(() -> new WrenException(GENERIC_USER_ERROR, "Mapping table name is refreshing or not exists"));
}
public Optional<String> getTableName()
{
return tableName;
}
public Optional<String> getErrorMessage()
{
return errorMessage;
}
public long getCreateTime()
{
return createTime;
}
}
wren-cache/src/main/java/io/wren/cache/CacheManager.java
Vendored
-80
View File
@@ -1,80 +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.wren.cache;
import com.google.common.annotations.VisibleForTesting;
import io.wren.base.AnalyzedMDL;
import io.wren.base.CatalogSchemaTableName;
import io.wren.base.ConnectorRecordIterator;
import io.wren.base.Parameter;
import io.wren.base.WrenException;
import io.wren.base.dto.CacheInfo;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import static io.wren.base.metadata.StandardErrorCode.GENERIC_INTERNAL_ERROR;
public interface CacheManager
{
default ConnectorRecordIterator query(String sql, List<Parameter> parameters)
{
throw new WrenException(GENERIC_INTERNAL_ERROR, "Enable Wren Protocol to use this feature");
}
default void removeCacheIfExist(String catalogName, String schemaName) {}
default void removeCacheIfExist(CatalogSchemaTableName catalogSchemaTableName) {}
default boolean cacheScheduledFutureExists(CatalogSchemaTableName catalogSchemaTableName)
{
return false;
}
@VisibleForTesting
default boolean retryScheduledFutureExists(CatalogSchemaTableName catalogSchemaTableName)
{
return false;
}
default CompletableFuture<List<TaskInfo>> createTask(AnalyzedMDL analyzedMDL)
{
return CompletableFuture.completedFuture(List.of());
}
default CompletableFuture<TaskInfo> createTask(AnalyzedMDL analyzedMDL, CacheInfo cacheInfo)
{
throw new WrenException(GENERIC_INTERNAL_ERROR, "Enable Wren Protocol to use this feature");
}
default CompletableFuture<List<TaskInfo>> listTaskInfo(String catalogName, String schemaName)
{
throw new WrenException(GENERIC_INTERNAL_ERROR, "Enable Wren Protocol to use this feature");
}
default CompletableFuture<Optional<TaskInfo>> getTaskInfo(CatalogSchemaTableName catalogSchemaTableName)
{
throw new WrenException(GENERIC_INTERNAL_ERROR, "Enable Wren Protocol to use this feature");
}
@VisibleForTesting
default void untilTaskDone(CatalogSchemaTableName name) {}
default List<Object> getDuckDBSettings()
{
throw new WrenException(GENERIC_INTERNAL_ERROR, "Enable Wren Protocol to use this feature");
}
}
wren-cache/src/main/java/io/wren/cache/CacheManagerImpl.java
Vendored
-434
View File
@@ -1,434 +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.wren.cache;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.inject.Inject;
import io.airlift.log.Logger;
import io.trino.sql.parser.SqlParser;
import io.trino.sql.tree.Statement;
import io.wren.base.AnalyzedMDL;
import io.wren.base.CatalogSchemaTableName;
import io.wren.base.ConnectorRecordIterator;
import io.wren.base.Parameter;
import io.wren.base.SessionContext;
import io.wren.base.WrenException;
import io.wren.base.WrenMDL;
import io.wren.base.client.duckdb.CacheStorageConfig;
import io.wren.base.client.duckdb.DuckDBConfig;
import io.wren.base.config.ConfigManager;
import io.wren.base.dto.CacheInfo;
import io.wren.base.sql.SqlConverter;
import io.wren.base.sqlrewrite.WrenPlanner;
import io.wren.base.wireprotocol.PgMetastore;
import io.wren.cache.dto.CachedTable;
import java.io.Closeable;
import java.time.Instant;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.function.Predicate;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static io.airlift.concurrent.Threads.daemonThreadsNamed;
import static io.airlift.concurrent.Threads.threadsNamed;
import static io.trino.execution.sql.SqlFormatterUtil.getFormattedSql;
import static io.wren.base.CatalogSchemaTableName.catalogSchemaTableName;
import static io.wren.base.metadata.StandardErrorCode.EXCEEDED_GLOBAL_MEMORY_LIMIT;
import static io.wren.base.metadata.StandardErrorCode.GENERIC_INTERNAL_ERROR;
import static io.wren.base.metadata.StandardErrorCode.GENERIC_USER_ERROR;
import static io.wren.base.sqlrewrite.Utils.parseSql;
import static io.wren.cache.EventLogger.Level.ERROR;
import static io.wren.cache.EventLogger.Level.INFO;
import static io.wren.cache.TaskInfo.TaskStatus.DONE;
import static io.wren.cache.TaskInfo.TaskStatus.QUEUED;
import static io.wren.cache.TaskInfo.TaskStatus.RUNNING;
import static java.lang.String.format;
import static java.lang.System.currentTimeMillis;
import static java.util.Objects.requireNonNull;
import static java.util.UUID.randomUUID;
import static java.util.concurrent.CompletableFuture.supplyAsync;
import static java.util.concurrent.Executors.newCachedThreadPool;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static java.util.concurrent.TimeUnit.SECONDS;
import static java.util.stream.Collectors.toList;
public class CacheManagerImpl
implements CacheManager, Closeable
{
private static final Logger LOG = Logger.get(CacheManager.class);
private final ExtraRewriter extraRewriter;
private final CacheService cacheService;
private final SqlParser sqlParser;
private final SqlConverter sqlConverter;
private final PgMetastore pgMetastore;
private final ConcurrentLinkedQueue<PathInfo> tempFileLocations = new ConcurrentLinkedQueue<>();
private final CachedTableMapping cachedTableMapping;
private final ConcurrentMap<CatalogSchemaTableName, ScheduledFuture<?>> cacheScheduledFutures = new ConcurrentHashMap<>();
private final ConcurrentMap<CatalogSchemaTableName, ScheduledFuture<?>> retryScheduledFutures = new ConcurrentHashMap<>();
private final ScheduledThreadPoolExecutor refreshExecutor = new ScheduledThreadPoolExecutor(5, daemonThreadsNamed("cache-refresh-%s"));
private final ScheduledThreadPoolExecutor retryExecutor = new ScheduledThreadPoolExecutor(5, daemonThreadsNamed("cache-retry-%s"));
private final ExecutorService executorService = newCachedThreadPool(threadsNamed("cache-manager-%s"));
private final ConcurrentHashMap<CatalogSchemaTableName, Task> tasks = new ConcurrentHashMap<>();
private final EventLogger eventLogger;
private final CacheTaskManager cacheTaskManager;
private final ConfigManager configManager;
@Inject
public CacheManagerImpl(
SqlConverter sqlConverter,
CacheService cacheService,
ExtraRewriter extraRewriter,
PgMetastore pgMetastore,
CachedTableMapping cachedTableMapping,
EventLogger eventLogger,
CacheTaskManager cacheTaskManager,
ConfigManager configManager)
{
this.sqlParser = new SqlParser();
this.sqlConverter = requireNonNull(sqlConverter, "sqlConverter is null");
this.cacheService = requireNonNull(cacheService, "cacheService is null");
this.extraRewriter = requireNonNull(extraRewriter, "extraRewriter is null");
this.pgMetastore = requireNonNull(pgMetastore, "pgMetastore is null");
this.cacheTaskManager = requireNonNull(cacheTaskManager, "cacheTaskManager is null");
this.cachedTableMapping = requireNonNull(cachedTableMapping, "cachedTableMapping is null");
this.eventLogger = requireNonNull(eventLogger, "eventLogger is null");
this.configManager = requireNonNull(configManager, "configManager is null");
refreshExecutor.setRemoveOnCancelPolicy(true);
}
private synchronized CompletableFuture<Void> refreshCache(AnalyzedMDL analyzedMDL, CacheInfo cacheInfo, TaskInfo taskInfo)
{
WrenMDL mdl = analyzedMDL.getWrenMDL();
CatalogSchemaTableName catalogSchemaTableName = catalogSchemaTableName(mdl.getCatalog(), mdl.getSchema(), cacheInfo.getName());
Optional<Task> taskOptional = Optional.ofNullable(tasks.get(catalogSchemaTableName));
if (taskOptional.isPresent() && taskOptional.get().getTaskInfo().inProgress()) {
throw new WrenException(GENERIC_USER_ERROR, format("cache is already running; catalogName: %s, schemaName: %s, tableName: %s", mdl.getCatalog(), mdl.getSchema(), cacheInfo.getName()));
}
removeCacheIfExist(catalogSchemaTableName);
return doCache(analyzedMDL, cacheInfo, taskInfo);
}
private CompletableFuture<Void> handleCache(AnalyzedMDL analyzedMDL, CacheInfo cacheInfo, TaskInfo taskInfo)
{
WrenMDL mdl = analyzedMDL.getWrenMDL();
CatalogSchemaTableName catalogSchemaTableName = new CatalogSchemaTableName(mdl.getCatalog(), mdl.getSchema(), cacheInfo.getName());
String duckdbTableName = format("%s_%s", cacheInfo.getName(), randomUUID().toString().replace("-", ""));
long createTime = currentTimeMillis();
return refreshCache(analyzedMDL, cacheInfo, taskInfo)
.thenRun(() -> {
if (cacheInfo.getRefreshTime().toMillis() > 0) {
cacheScheduledFutures.put(
catalogSchemaTableName,
refreshExecutor.scheduleWithFixedDelay(
() -> createTask(analyzedMDL, cacheInfo).join(),
cacheInfo.getRefreshTime().toMillis(),
cacheInfo.getRefreshTime().toMillis(),
MILLISECONDS));
}
})
.exceptionally(e -> {
String errMsg = format("Failed to do cache for cacheInfo %s; caused by %s", cacheInfo.getName(), e.getMessage());
// If the cache fails because DuckDB doesn't have sufficient memory, we'll attempt to retry it later.
if (e.getCause() instanceof WrenException && EXCEEDED_GLOBAL_MEMORY_LIMIT.toErrorCode().equals(((WrenException) e.getCause()).getErrorCode())) {
long delay = configManager.getConfig(DuckDBConfig.class).getCacheTaskRetryDelay();
retryScheduledFutures.put(
catalogSchemaTableName,
retryExecutor.schedule(
() -> createTask(analyzedMDL, cacheInfo).join(),
delay,
SECONDS));
errMsg += "; will retry after " + delay + " seconds";
}
pgMetastore.dropTableIfExists(duckdbTableName);
LOG.error(e, errMsg);
cachedTableMapping.putCachedTableMapping(catalogSchemaTableName, new CacheInfoPair(cacheInfo, Optional.empty(), Optional.of(errMsg), createTime));
return null;
});
}
@Override
public ConnectorRecordIterator query(String sql, List<Parameter> parameters)
{
return cacheTaskManager.addCacheQueryTask(() -> DuckdbRecordIterator.of(pgMetastore.getClient(), sql, parameters.stream().collect(toImmutableList())));
}
private CompletableFuture<Void> doCache(AnalyzedMDL analyzedMDL, CacheInfo cacheInfo, TaskInfo taskInfo)
{
WrenMDL mdl = analyzedMDL.getWrenMDL();
CatalogSchemaTableName catalogSchemaTableName = new CatalogSchemaTableName(mdl.getCatalog(), mdl.getSchema(), cacheInfo.getName());
String duckdbTableName = format("%s_%s", cacheInfo.getName(), randomUUID().toString().replace("-", ""));
long createTime = currentTimeMillis();
return cacheTaskManager.addCacheTask(() -> {
cacheTaskManager.checkCacheMemoryLimit();
taskInfo.setTaskStatus(RUNNING);
SessionContext sessionContext = SessionContext.builder()
.setCatalog(mdl.getCatalog())
.setSchema(mdl.getSchema())
.build();
String wrenRewritten = WrenPlanner.rewrite(
format("select * from %s", cacheInfo.getName()),
sessionContext,
analyzedMDL);
Statement parsedStatement = parseSql(wrenRewritten);
Statement rewrittenStatement = extraRewriter.rewrite(parsedStatement);
createCache(mdl, cacheInfo, sessionContext, rewrittenStatement, duckdbTableName);
cachedTableMapping.putCachedTableMapping(catalogSchemaTableName, new CacheInfoPair(cacheInfo, duckdbTableName, createTime));
});
}
private void createCache(
WrenMDL mdl,
CacheInfo cacheInfo,
SessionContext sessionContext,
Statement rewrittenStatement,
String duckdbTableName)
{
cacheService.createCache(
mdl.getCatalog(),
mdl.getSchema(),
cacheInfo.getName(),
sqlConverter.convert(getFormattedSql(rewrittenStatement, sqlParser), sessionContext))
.ifPresent(pathInfo -> {
try {
tempFileLocations.add(pathInfo);
refreshCacheInDuckDB(pathInfo.getPath() + "/" + pathInfo.getFilePattern(), duckdbTableName);
}
finally {
removeTempFile(pathInfo);
}
});
}
private void refreshCacheInDuckDB(String path, String tableName)
{
pgMetastore.directDDL(configManager.getConfig(CacheStorageConfig.class).generateDuckdbParquetStatement(path, tableName));
}
@Override
public void removeCacheIfExist(String catalogName, String schemaName)
{
requireNonNull(catalogName, "catalogName is null");
requireNonNull(schemaName, "schemaName is null");
cacheScheduledFutures.keySet().stream()
.filter(catalogSchemaTableName -> catalogSchemaTableName.getCatalogName().equals(catalogName)
&& catalogSchemaTableName.getSchemaTableName().getSchemaName().equals(schemaName))
.forEach(catalogSchemaTableName -> {
cacheScheduledFutures.get(catalogSchemaTableName).cancel(true);
cacheScheduledFutures.remove(catalogSchemaTableName);
});
retryScheduledFutures.keySet().stream()
.filter(catalogSchemaTableName -> catalogSchemaTableName.getCatalogName().equals(catalogName)
&& catalogSchemaTableName.getSchemaTableName().getSchemaName().equals(schemaName))
.forEach(catalogSchemaTableName -> {
retryScheduledFutures.get(catalogSchemaTableName).cancel(true);
retryScheduledFutures.remove(catalogSchemaTableName);
});
cachedTableMapping.entrySet().stream()
.filter(entry -> entry.getKey().getCatalogName().equals(catalogName)
&& entry.getKey().getSchemaTableName().getSchemaName().equals(schemaName))
.forEach(entry -> {
entry.getValue().getTableName().ifPresent(pgMetastore::dropTableIfExists);
cachedTableMapping.remove(entry.getKey());
});
tasks.keySet().stream()
.filter(catalogSchemaTableName -> catalogSchemaTableName.getCatalogName().equals(catalogName)
&& catalogSchemaTableName.getSchemaTableName().getSchemaName().equals(schemaName))
.forEach(tasks::remove);
}
@Override
public void removeCacheIfExist(CatalogSchemaTableName catalogSchemaTableName)
{
if (cacheScheduledFutures.containsKey(catalogSchemaTableName)) {
cacheScheduledFutures.get(catalogSchemaTableName).cancel(true);
cacheScheduledFutures.remove(catalogSchemaTableName);
}
if (retryScheduledFutures.containsKey(catalogSchemaTableName)) {
retryScheduledFutures.get(catalogSchemaTableName).cancel(true);
retryScheduledFutures.remove(catalogSchemaTableName);
}
Optional.ofNullable(cachedTableMapping.get(catalogSchemaTableName)).ifPresent(cacheInfoPair -> {
cacheInfoPair.getTableName().ifPresent(pgMetastore::dropTableIfExists);
cachedTableMapping.remove(catalogSchemaTableName);
});
Task task = tasks.remove(catalogSchemaTableName);
if (task != null) {
eventLogger.logEvent(INFO, "REMOVE_TASK", "Remove cache: " + catalogSchemaTableName);
}
}
@Override
public boolean cacheScheduledFutureExists(CatalogSchemaTableName catalogSchemaTableName)
{
return cacheScheduledFutures.containsKey(catalogSchemaTableName);
}
@Override
public boolean retryScheduledFutureExists(CatalogSchemaTableName catalogSchemaTableName)
{
return retryScheduledFutures.containsKey(catalogSchemaTableName);
}
@Override
public void close()
{
refreshExecutor.shutdownNow();
retryExecutor.shutdownNow();
executorService.shutdownNow();
cleanTempFiles();
}
private void cleanTempFiles()
{
try {
List<PathInfo> locations = ImmutableList.copyOf(tempFileLocations);
locations.forEach(this::removeTempFile);
}
catch (Exception e) {
LOG.error(e, "Failed to clean temp file");
}
}
private void removeTempFile(PathInfo pathInfo)
{
if (tempFileLocations.contains(pathInfo)) {
cacheService.deleteTarget(pathInfo);
tempFileLocations.remove(pathInfo);
}
}
@Override
public CompletableFuture<List<TaskInfo>> createTask(AnalyzedMDL analyzedMDL)
{
return supplyAsync(() ->
analyzedMDL.getWrenMDL().listCached().stream().map(cacheInfo -> createTask(analyzedMDL, cacheInfo).join()).collect(toList()));
}
@Override
public CompletableFuture<TaskInfo> createTask(AnalyzedMDL analyzedMDL, CacheInfo cacheInfo)
{
return supplyAsync(() -> {
WrenMDL mdl = analyzedMDL.getWrenMDL();
CatalogSchemaTableName catalogSchemaTableName = new CatalogSchemaTableName(mdl.getCatalog(), mdl.getSchema(), cacheInfo.getName());
TaskInfo taskInfo = new TaskInfo(mdl.getCatalog(), mdl.getSchema(), cacheInfo.getName(), QUEUED, Instant.now());
// To fix flaky test, we pass value to tasks instead of a reference;
Task task = new Task(TaskInfo.copyFrom(taskInfo), analyzedMDL, cacheInfo);
tasks.put(catalogSchemaTableName, task);
return taskInfo;
});
}
@Override
public CompletableFuture<List<TaskInfo>> listTaskInfo(String catalogName, String schemaName)
{
Predicate<TaskInfo> catalogNamePred = catalogName.isEmpty() ?
(t) -> true :
(t) -> catalogName.equals(t.getCatalogName());
Predicate<TaskInfo> schemaNamePred = schemaName.isEmpty() ?
(t) -> true :
(t) -> schemaName.equals(t.getSchemaName());
return supplyAsync(
() -> tasks.values().stream()
.map(Task::getTaskInfo)
.filter(catalogNamePred.and(schemaNamePred))
.collect(toList()),
executorService);
}
@Override
public CompletableFuture<Optional<TaskInfo>> getTaskInfo(CatalogSchemaTableName catalogSchemaTableName)
{
requireNonNull(catalogSchemaTableName);
return supplyAsync(
() -> Optional.ofNullable(tasks.get(catalogSchemaTableName)).map(Task::getTaskInfo),
executorService);
}
@VisibleForTesting
@Override
public void untilTaskDone(CatalogSchemaTableName name)
{
Optional.ofNullable(tasks.get(name)).ifPresent(Task::waitUntilDone);
}
public List<Object> getDuckDBSettings()
{
try (ConnectorRecordIterator iter = query("SELECT * FROM duckdb_settings()", List.of())) {
return ImmutableList.copyOf(iter);
}
catch (Exception e) {
LOG.error(e, "Failed to get duckdb settings");
throw new WrenException(GENERIC_INTERNAL_ERROR, e);
}
}
private class Task
{
private final TaskInfo taskInfo;
private final CompletableFuture<?> completableFuture;
public Task(TaskInfo taskInfo, AnalyzedMDL analyzedMDL, CacheInfo cacheInfo)
{
this.taskInfo = taskInfo;
this.completableFuture = handleCache(analyzedMDL, cacheInfo, taskInfo)
.thenRun(() -> {
CacheInfoPair cacheInfoPair = cachedTableMapping.getCacheInfoPair(
taskInfo.getCatalogName(),
taskInfo.getSchemaName(), taskInfo.getTableName());
taskInfo.setCachedTable(new CachedTable(
cacheInfoPair.getCacheInfo().getName(),
cacheInfoPair.getErrorMessage(),
cacheInfoPair.getCacheInfo().getRefreshTime(),
Instant.ofEpochMilli(cacheInfoPair.getCreateTime())));
taskInfo.setTaskStatus(DONE);
if (cacheInfoPair.getErrorMessage().isPresent()) {
eventLogger.logEvent(ERROR, "CREATE_TASK", taskInfo);
}
else {
eventLogger.logEvent(INFO, "CREATE_TASK", taskInfo);
}
});
}
public TaskInfo getTaskInfo()
{
return taskInfo;
}
public void waitUntilDone()
{
completableFuture.join();
}
}
}
wren-cache/src/main/java/io/wren/cache/CacheModule.java
Vendored
-32
View File
@@ -1,32 +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.wren.cache;
import com.google.inject.Binder;
import com.google.inject.Scopes;
import io.airlift.configuration.AbstractConfigurationAwareModule;
public class CacheModule
extends AbstractConfigurationAwareModule
{
@Override
protected void setup(Binder binder)
{
binder.bind(CacheManager.class).to(CacheManagerImpl.class).in(Scopes.SINGLETON);
binder.bind(CacheTaskManager.class).in(Scopes.SINGLETON);
binder.bind(EventLogger.class).to(Log4jEventLogger.class).in(Scopes.SINGLETON);
binder.bind(CachedTableMapping.class).to(DefaultCachedTableMapping.class).in(Scopes.SINGLETON);
}
}
wren-cache/src/main/java/io/wren/cache/CacheService.java
Vendored
-23
View File
@@ -1,23 +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.wren.cache;
import java.util.Optional;
public interface CacheService
{
Optional<PathInfo> createCache(String catalog, String schema, String name, String statement);
void deleteTarget(PathInfo pathInfo);
}
wren-cache/src/main/java/io/wren/cache/CacheTaskManager.java
Vendored
-117
View File
@@ -1,117 +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.wren.cache;
import com.google.common.collect.ImmutableList;
import com.google.inject.Inject;
import io.wren.base.ConnectorRecordIterator;
import io.wren.base.WrenException;
import io.wren.base.client.duckdb.DuckDBConfig;
import io.wren.base.wireprotocol.PgMetastore;
import java.io.Closeable;
import java.io.IOException;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeoutException;
import static io.airlift.concurrent.Threads.threadsNamed;
import static io.wren.base.client.duckdb.DuckdbUtil.convertDuckDBUnits;
import static io.wren.base.metadata.StandardErrorCode.EXCEEDED_GLOBAL_MEMORY_LIMIT;
import static io.wren.base.metadata.StandardErrorCode.EXCEEDED_TIME_LIMIT;
import static io.wren.base.metadata.StandardErrorCode.GENERIC_INTERNAL_ERROR;
import static java.util.Objects.requireNonNull;
import static java.util.concurrent.CompletableFuture.runAsync;
import static java.util.concurrent.Executors.newFixedThreadPool;
import static java.util.concurrent.TimeUnit.SECONDS;
public class CacheTaskManager
implements Closeable
{
private final PgMetastore pgMetastore;
private final ExecutorService taskExecutorService;
private final DuckDBConfig duckDBConfig;
private final double cacheMemoryLimit;
@Inject
public CacheTaskManager(DuckDBConfig duckDBConfig, PgMetastore pgMetastore)
{
this.duckDBConfig = requireNonNull(duckDBConfig, "duckDBConfig is null");
this.pgMetastore = requireNonNull(pgMetastore, "pgMetastore is null");
this.taskExecutorService = newFixedThreadPool(duckDBConfig.getMaxConcurrentTasks(), threadsNamed("duckdb-task-%s"));
this.cacheMemoryLimit = duckDBConfig.getMaxCacheTableSizeRatio() * duckDBConfig.getMemoryLimit().toBytes();
}
public CompletableFuture<Void> addCacheTask(Runnable runnable)
{
return runAsync(runnable, taskExecutorService);
}
public <T> T addCacheQueryTask(Callable<T> callable)
{
try {
return taskExecutorService.submit(callable).get(duckDBConfig.getMaxCacheQueryTimeout(), SECONDS);
}
catch (TimeoutException e) {
throw new WrenException(EXCEEDED_TIME_LIMIT, "Query time limit exceeded", e);
}
catch (InterruptedException | ExecutionException e) {
throw new WrenException(GENERIC_INTERNAL_ERROR, e);
}
}
// for canner use
public void addCacheQueryDDLTask(Runnable runnable)
{
try {
taskExecutorService.submit(runnable).get(duckDBConfig.getMaxCacheQueryTimeout(), SECONDS);
}
catch (TimeoutException e) {
throw new WrenException(EXCEEDED_TIME_LIMIT, "Query time limit exceeded", e);
}
catch (InterruptedException | ExecutionException e) {
throw new WrenException(GENERIC_INTERNAL_ERROR, e);
}
}
public long getMemoryUsageBytes()
{
try (ConnectorRecordIterator result = pgMetastore.directQuery("SELECT memory_usage FROM pragma_database_size()", ImmutableList.of())) {
Object[] row = result.next();
return convertDuckDBUnits(row[0].toString()).toBytes();
}
catch (Exception e) {
throw new WrenException(GENERIC_INTERNAL_ERROR, "Failed to get memory usage", e);
}
}
public void checkCacheMemoryLimit()
{
long usage = getMemoryUsageBytes();
if (usage >= cacheMemoryLimit) {
throw new WrenException(EXCEEDED_GLOBAL_MEMORY_LIMIT, "Cache memory limit exceeded. Usage: " + usage + " bytes, Limit: " + cacheMemoryLimit + " bytes");
}
}
@Override
public void close()
throws IOException
{
taskExecutorService.shutdownNow();
pgMetastore.close();
}
}
wren-cache/src/main/java/io/wren/cache/CachedTableMapping.java
Vendored
-38
View File
@@ -1,38 +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.wren.cache;
import io.wren.base.CatalogSchemaTableName;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
public interface CachedTableMapping
{
void putCachedTableMapping(CatalogSchemaTableName catalogSchemaTableName, CacheInfoPair cacheInfoPair);
CacheInfoPair get(CatalogSchemaTableName cachedTable);
void remove(CatalogSchemaTableName cachedTable);
CacheInfoPair getCacheInfoPair(String catalog, String schema, String table);
Optional<String> convertToCachedTable(CatalogSchemaTableName catalogSchemaTableName);
Set<Map.Entry<CatalogSchemaTableName, CacheInfoPair>> entrySet();
List<CacheInfoPair> getCacheInfoPairs(String catalogName, String schemaName);
}
wren-cache/src/main/java/io/wren/cache/DefaultCachedTableMapping.java
Vendored
-102
View File
@@ -1,102 +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.wren.cache;
import com.google.inject.Inject;
import io.wren.base.CatalogSchemaTableName;
import io.wren.base.wireprotocol.PgMetastore;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static java.util.Objects.requireNonNull;
public class DefaultCachedTableMapping
implements CachedTableMapping
{
private final PgMetastore pgMetastore;
private final ConcurrentMap<CatalogSchemaTableName, CacheInfoPair> cachedTableMapping = new ConcurrentHashMap<>();
@Inject
public DefaultCachedTableMapping(PgMetastore pgMetastore)
{
this.pgMetastore = requireNonNull(pgMetastore, "duckdbClient is null");
}
@Override
public void putCachedTableMapping(CatalogSchemaTableName catalogSchemaTableName, CacheInfoPair cacheInfoPair)
{
synchronized (cachedTableMapping) {
if (cachedTableMapping.containsKey(catalogSchemaTableName)) {
CacheInfoPair existedCacheInfoPair = cachedTableMapping.get(catalogSchemaTableName);
if (existedCacheInfoPair.getCreateTime() > cacheInfoPair.getCreateTime()) {
cacheInfoPair.getTableName().ifPresent(pgMetastore::dropTableIfExists);
return;
}
existedCacheInfoPair.getTableName().ifPresent(pgMetastore::dropTableIfExists);
}
cachedTableMapping.put(catalogSchemaTableName, cacheInfoPair);
}
}
@Override
public CacheInfoPair get(CatalogSchemaTableName cachedTable)
{
return cachedTableMapping.get(cachedTable);
}
@Override
public void remove(CatalogSchemaTableName cachedTable)
{
cachedTableMapping.remove(cachedTable);
}
@Override
public CacheInfoPair getCacheInfoPair(String catalog, String schema, String table)
{
return cachedTableMapping.get(new CatalogSchemaTableName(catalog, schema, table));
}
@Override
public Optional<String> convertToCachedTable(CatalogSchemaTableName catalogSchemaTableName)
{
return Optional.ofNullable(cachedTableMapping.get(catalogSchemaTableName))
.flatMap(CacheInfoPair::getTableName);
}
@Override
public Set<Map.Entry<CatalogSchemaTableName, CacheInfoPair>> entrySet()
{
return cachedTableMapping.entrySet();
}
@Override
public List<CacheInfoPair> getCacheInfoPairs(String catalogName, String schemaName)
{
requireNonNull(catalogName, "catalogName is null");
requireNonNull(schemaName, "schemaName is null");
return cachedTableMapping.entrySet()
.stream()
.filter(entry ->
entry.getKey().getCatalogName().equals(catalogName)
&& entry.getKey().getSchemaTableName().getSchemaName().equals(schemaName))
.map(Map.Entry::getValue)
.collect(toImmutableList());
}
}
wren-cache/src/main/java/io/wren/cache/EventLogger.java
Vendored
-29
View File
@@ -1,29 +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.wren.cache;
public interface EventLogger
{
enum Level
{
INFO,
WARN,
ERROR
}
void logEvent(Level level, String eventName, TaskInfo event);
void logEvent(Level level, String eventName, String description);
}
wren-cache/src/main/java/io/wren/cache/ExtraRewriter.java
Vendored
-24
View File
@@ -1,24 +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.wren.cache;
import io.trino.sql.tree.Statement;
public interface ExtraRewriter
{
default Statement rewrite(Statement statement)
{
return statement;
}
}

Some files were not shown because too many files have changed in this diff Show More