forked from trinodb/trino
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- implemented cached adb connection factory using hikariCP
- Loading branch information
1 parent
e81f023
commit 75f6f18
Showing
17 changed files
with
636 additions
and
38 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
67 changes: 67 additions & 0 deletions
67
...rc/main/java/io/trino/plugin/adb/connector/connection/AdbCachedJdbcConnectionFactory.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,67 @@ | ||
| /* | ||
| * 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.plugin.adb.connector.connection; | ||
|
|
||
| import io.airlift.log.Logger; | ||
| import io.trino.plugin.adb.connector.connection.pool.CachedDataSourceEntry; | ||
| import io.trino.plugin.adb.connector.connection.pool.JdbcDataSourcePool; | ||
| import io.trino.plugin.jdbc.ConnectionFactory; | ||
| import io.trino.spi.StandardErrorCode; | ||
| import io.trino.spi.TrinoException; | ||
| import io.trino.spi.connector.ConnectorSession; | ||
|
|
||
| import java.sql.Connection; | ||
| import java.sql.SQLException; | ||
| import java.util.Optional; | ||
|
|
||
| import static java.lang.String.format; | ||
|
|
||
| public class AdbCachedJdbcConnectionFactory | ||
| implements ConnectionFactory | ||
| { | ||
| private static final Logger log = Logger.get(AdbCachedJdbcConnectionFactory.class); | ||
| private static final int RETRY_GETTING_CONNECTION_DELAY_MS = 100; | ||
| private final AdbConnectionConfig connectionConfig; | ||
| private final JdbcDataSourcePool connectionPool; | ||
|
|
||
| public AdbCachedJdbcConnectionFactory(AdbConnectionConfig connectionConfig, JdbcDataSourcePool connectionPool) | ||
| { | ||
| this.connectionConfig = connectionConfig; | ||
| this.connectionPool = connectionPool; | ||
| } | ||
|
|
||
| @Override | ||
| public Connection openConnection(ConnectorSession session) | ||
| throws SQLException | ||
| { | ||
| long startTimeMs = System.currentTimeMillis(); | ||
| while ((startTimeMs + connectionConfig.getOpenConnectionTimeout().toMillis() > System.currentTimeMillis())) { | ||
| try { | ||
| CachedDataSourceEntry cachedDataSourceEntry = connectionPool.get(session); | ||
| Optional<Connection> connectionOptional = cachedDataSourceEntry.tryOpenConnection(); | ||
| if (connectionOptional.isPresent()) { | ||
| return connectionOptional.get(); | ||
| } | ||
| Thread.sleep(RETRY_GETTING_CONNECTION_DELAY_MS); | ||
| } | ||
| catch (Exception e) { | ||
| String errMsg = "Failed to open connection" + e.getMessage(); | ||
| log.error(errMsg, e); | ||
| throw new TrinoException(StandardErrorCode.GENERIC_INTERNAL_ERROR, errMsg, e); | ||
| } | ||
| } | ||
| throw new SQLException( | ||
| format("Failed to open connection. Timeout %s exceeded", connectionConfig.getOpenConnectionTimeout())); | ||
| } | ||
| } |
52 changes: 52 additions & 0 deletions
52
...trino-adb/src/main/java/io/trino/plugin/adb/connector/connection/AdbConnectionConfig.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,52 @@ | ||
| /* | ||
| * 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.plugin.adb.connector.connection; | ||
|
|
||
| import io.airlift.configuration.Config; | ||
| import io.airlift.configuration.ConfigDescription; | ||
| import io.airlift.units.Duration; | ||
|
|
||
| import java.util.concurrent.TimeUnit; | ||
|
|
||
| public class AdbConnectionConfig | ||
| { | ||
| private boolean isConnectionPoolEnabled = true; | ||
| private Duration openConnectionTimeout = new Duration(30, TimeUnit.SECONDS); | ||
|
|
||
| public boolean getConnectionPoolEnabled() | ||
| { | ||
| return isConnectionPoolEnabled; | ||
| } | ||
|
|
||
| @Config("adb.jdbc.connection.pool.enabled") | ||
| @ConfigDescription("Using connection pool. Default is true") | ||
| public AdbConnectionConfig setConnectionPoolEnabled(boolean isConnectionPoolEnabled) | ||
| { | ||
| this.isConnectionPoolEnabled = isConnectionPoolEnabled; | ||
| return this; | ||
| } | ||
|
|
||
| public Duration getOpenConnectionTimeout() | ||
| { | ||
| return this.openConnectionTimeout; | ||
| } | ||
|
|
||
| @Config("adb.jdbc.connection.open-timeout") | ||
| @ConfigDescription("Max time of trying to open connection to adb. Default is 30 seconds") | ||
| public AdbConnectionConfig setOpenConnectionTimeout(Duration openConnectionTimeout) | ||
| { | ||
| this.openConnectionTimeout = openConnectionTimeout; | ||
| return this; | ||
| } | ||
| } |
75 changes: 75 additions & 0 deletions
75
...db/src/main/java/io/trino/plugin/adb/connector/connection/AdbConnectionFactoryModule.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| /* | ||
| * 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.plugin.adb.connector.connection; | ||
|
|
||
| import com.google.inject.Binder; | ||
| import com.google.inject.Provides; | ||
| import com.google.inject.Singleton; | ||
| import io.airlift.configuration.AbstractConfigurationAwareModule; | ||
| import io.airlift.configuration.ConfigBinder; | ||
| import io.opentelemetry.api.OpenTelemetry; | ||
| import io.trino.plugin.adb.connector.connection.pool.AdbJdbcDataSourcePool; | ||
| import io.trino.plugin.adb.connector.connection.pool.AdbJdbcConnectionPoolConfig; | ||
| import io.trino.plugin.jdbc.BaseJdbcConfig; | ||
| import io.trino.plugin.jdbc.ConnectionFactory; | ||
| import io.trino.plugin.jdbc.DriverConnectionFactory; | ||
| import io.trino.plugin.jdbc.ForBaseJdbc; | ||
| import io.trino.plugin.jdbc.IdentityCacheMapping; | ||
| import io.trino.plugin.jdbc.credential.CredentialProvider; | ||
| import io.trino.spi.catalog.CatalogName; | ||
| import org.postgresql.Driver; | ||
| import org.postgresql.PGProperty; | ||
|
|
||
| import java.util.Properties; | ||
|
|
||
| public class AdbConnectionFactoryModule | ||
| extends AbstractConfigurationAwareModule | ||
| { | ||
| @Override | ||
| protected void setup(Binder binder) | ||
| { | ||
| ConfigBinder.configBinder(binder).bindConfig(AdbConnectionConfig.class); | ||
| ConfigBinder.configBinder(binder).bindConfig(AdbJdbcConnectionPoolConfig.class); | ||
| } | ||
|
|
||
| @Provides | ||
| @Singleton | ||
| @ForBaseJdbc | ||
| public static ConnectionFactory getConnectionFactory(BaseJdbcConfig config, | ||
| OpenTelemetry openTelemetry, | ||
| CredentialProvider credentialProvider, | ||
| CatalogName catalogName, | ||
| IdentityCacheMapping identityCacheMapping, | ||
| AdbJdbcConnectionPoolConfig connectionPoolConfig, | ||
| AdbConnectionConfig connectionConfig) | ||
| { | ||
| Properties connectionProperties = new Properties(); | ||
| connectionProperties.put(PGProperty.REWRITE_BATCHED_INSERTS.getName(), "true"); | ||
| ConnectionFactory delegate = | ||
| DriverConnectionFactory.builder(new Driver(), config.getConnectionUrl(), credentialProvider) | ||
| .setConnectionProperties(connectionProperties) | ||
| .setOpenTelemetry(openTelemetry) | ||
| .build(); | ||
| if (connectionConfig.getConnectionPoolEnabled()) { | ||
| return new AdbCachedJdbcConnectionFactory(connectionConfig, | ||
| new AdbJdbcDataSourcePool(connectionPoolConfig, | ||
| delegate, | ||
| catalogName, | ||
| identityCacheMapping)); | ||
| } | ||
| else { | ||
| return delegate; | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.