Skip to content

Commit b15b964

Browse files
singhpk234singhpk234
committed
[WIP] Support multiple DS
1 parent 66f75a9 commit b15b964

File tree

12 files changed

+142
-42
lines changed

12 files changed

+142
-42
lines changed

‎extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/DatasourceOperations.java

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,8 @@
3838

3939
public class DatasourceOperations {
4040

41-
private static final String CONSTRAINT_VIOLATION_SQL_CODE = "23505";
41+
private static final String CONSTRAINT_VIOLATION_SQL_CODE_POSTGRES = "23505";
42+
private static final String TABLE_DOES_NOT_EXIST_SQL_CODE_POSTGRES = "42P01";
4243

4344
private final DataSource datasource;
4445

@@ -182,9 +183,14 @@ public interface TransactionCallback {
182183
}
183184

184185
public boolean isConstraintViolation(SQLException e) {
185-
return CONSTRAINT_VIOLATION_SQL_CODE.equals(e.getSQLState());
186+
return CONSTRAINT_VIOLATION_SQL_CODE_POSTGRES.equals(e.getSQLState());
186187
}
187188

189+
public boolean isTableNotExists(SQLException e) {
190+
return TABLE_DOES_NOT_EXIST_SQL_CODE_POSTGRES.equals(e.getSQLState());
191+
}
192+
193+
188194
private Connection borrowConnection() throws SQLException {
189195
return datasource.getConnection();
190196
}

‎extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcDatasource.java renamed to ‎extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/DatasourceSupplier.java

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,10 @@
1616
* specific language governing permissions and limitations
1717
* under the License.
1818
*/
19-
2019
package org.apache.polaris.extension.persistence.relational.jdbc;
2120

2221
import javax.sql.DataSource;
2322

24-
public interface JdbcDatasource {
25-
DataSource fromRealmId(String realmId);
23+
public interface DatasourceSupplier {
24+
DataSource fromRealmId(String realmId);
2625
}

‎extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -304,6 +304,12 @@ private PolarisBaseEntity getPolarisBaseEntity(String query) {
304304
return results.getFirst();
305305
}
306306
} catch (SQLException e) {
307+
// This look-up is used for checking if the realm is bootstrap or not.
308+
// If we have 1 DB per realm it might happen that the realm is not boostrap
309+
// at all.
310+
if (datasourceOperations.isTableNotExists(e)) {
311+
return null;
312+
}
307313
throw new RuntimeException(
308314
String.format("Failed to retrieve polaris entity due to %s", e.getMessage()), e);
309315
}

‎extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcMetaStoreManagerFactory.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
import io.smallrye.common.annotation.Identifier;
2222
import jakarta.annotation.Nullable;
2323
import jakarta.enterprise.context.ApplicationScoped;
24-
import jakarta.enterprise.inject.Instance;
2524
import jakarta.inject.Inject;
2625
import java.sql.Connection;
2726
import java.sql.SQLException;
@@ -72,7 +71,7 @@ public class JdbcMetaStoreManagerFactory implements MetaStoreManagerFactory {
7271
protected final PolarisDiagnostics diagServices = new PolarisDefaultDiagServiceImpl();
7372

7473
@Inject PolarisStorageIntegrationProvider storageIntegrationProvider;
75-
@Inject JdbcDatasource jdbcDatasource;
74+
@Inject DatasourceSupplier jdbcDatasource;
7675

7776
protected JdbcMetaStoreManagerFactory() {}
7877

@@ -92,7 +91,8 @@ protected PolarisMetaStoreManager createNewMetaStoreManager() {
9291

9392
private void initializeForRealm(
9493
RealmContext realmContext, RootCredentialsSet rootCredentialsSet, boolean isBootstrap) {
95-
DatasourceOperations databaseOperations = getDatasourceOperations(isBootstrap, realmContext.getRealmIdentifier());
94+
DatasourceOperations databaseOperations =
95+
getDatasourceOperations(isBootstrap, realmContext.getRealmIdentifier());
9696
sessionSupplierMap.put(
9797
realmContext.getRealmIdentifier(),
9898
() ->

‎quarkus/admin/src/main/java/org/apache/polaris/admintool/config/QuarkusDatasourceSupplier.java

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.polaris.admintool.config;
20+
21+
import io.quarkus.arc.InstanceHandle;
22+
import java.util.List;
23+
import javax.sql.DataSource;
24+
import org.apache.polaris.extension.persistence.relational.jdbc.DatasourceSupplier;
25+
import org.postgresql.shaded.com.ongres.scram.common.util.Preconditions;
26+
27+
public class QuarkusDatasourceSupplier implements DatasourceSupplier {
28+
private final List<InstanceHandle<DataSource>> dataSources;
29+
private final String realmIsolationType;
30+
31+
public static final String DEFAULT_DATA_SOURCE_NAME = "<default>";
32+
public static final String REALM_ISOLATION_TYPE_DB = "db";
33+
34+
public QuarkusDatasourceSupplier(String realmIsolationType, List<InstanceHandle<DataSource>> dataSources) {
35+
this.realmIsolationType = realmIsolationType;
36+
this.dataSources = dataSources;
37+
}
38+
39+
@Override
40+
public DataSource fromRealmId(String realmId) {
41+
Preconditions.checkNotNull(REALM_ISOLATION_TYPE_DB.equals(realmIsolationType) && dataSources.size() == 1, "Misconfigured datasource's as per realm requirements");
42+
for (InstanceHandle<DataSource> handle : dataSources) {
43+
String name = handle.getBean().getName();
44+
name = name == null ? DEFAULT_DATA_SOURCE_NAME : unquoteDataSourceName(name);
45+
// if realm isolation is DB then there should be only one DS configured.
46+
if (name.equals(realmId) || (REALM_ISOLATION_TYPE_DB.equals(realmIsolationType) && name.equals(DEFAULT_DATA_SOURCE_NAME))) {
47+
return handle.get();
48+
}
49+
}
50+
throw new IllegalStateException("No datasource configured with name: " + realmId);
51+
}
52+
53+
public static String unquoteDataSourceName(String dataSourceName) {
54+
if (dataSourceName.startsWith("\"") && dataSourceName.endsWith("\"")) {
55+
dataSourceName = dataSourceName.substring(1, dataSourceName.length() - 1);
56+
}
57+
return dataSourceName;
58+
}
59+
}

‎quarkus/admin/src/main/java/org/apache/polaris/admintool/config/QuarkusProducers.java

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,19 +28,17 @@
2828
import jakarta.enterprise.inject.Produces;
2929
import java.time.Clock;
3030
import java.util.List;
31-
31+
import javax.sql.DataSource;
3232
import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
3333
import org.apache.polaris.core.PolarisDiagnostics;
3434
import org.apache.polaris.core.config.PolarisConfigurationStore;
3535
import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
3636
import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo;
3737
import org.apache.polaris.core.storage.PolarisStorageIntegration;
3838
import org.apache.polaris.core.storage.PolarisStorageIntegrationProvider;
39-
import org.apache.polaris.extension.persistence.relational.jdbc.JdbcDatasource;
39+
import org.apache.polaris.extension.persistence.relational.jdbc.DatasourceSupplier;
4040
import org.eclipse.microprofile.config.inject.ConfigProperty;
4141

42-
import javax.sql.DataSource;
43-
4442
public class QuarkusProducers {
4543

4644
@Produces
@@ -85,7 +83,7 @@ public PolarisConfigurationStore configurationStore() {
8583
}
8684

8785
@Produces
88-
public JdbcDatasource jdbcDatasource(@All List<InstanceHandle<DataSource>> dataSources) {
89-
return new QuarkusJdbcDataSource(dataSources);
86+
public DatasourceSupplier jdbcDatasource(@ConfigProperty(name = "polaris.persistence.realm.isolation.type") String realmIsolationType, @All List<InstanceHandle<DataSource>> dataSources) {
87+
return new QuarkusDatasourceSupplier(realmIsolationType, dataSources);
9088
}
9189
}

‎quarkus/admin/src/main/resources/application.properties

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,10 @@ quarkus.index-dependency.guava.artifact-id=guava
5050
quarkus.index-dependency.protobuf.group-id=com.google.protobuf
5151
quarkus.index-dependency.protobuf.artifact-id=protobuf-java
5252

53+
#quarkus.datasource.db-kind=pgsql
54+
#quarkus.datasource.jdbc.url=polaris
55+
#quarkus.datasource.username=polaris
56+
#quarkus.datasource.password=polaris
5357
quarkus.datasource.\"realm1\".db-kind=pgsql
5458
quarkus.datasource.\"realm1\".jdbc.url=polaris
5559
quarkus.datasource.\"realm1\".username=polaris

‎quarkus/admin/src/test/java/org/apache/polaris/admintool/relational/jdbc/RelationalJdbcAdminProfile.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ public List<TestResourceEntry> testResources() {
3636
return List.of(
3737
new TestResourceEntry(
3838
PostgresRelationalJdbcLifeCycleManagement.class,
39-
Map.of(INIT_SCRIPT, "org/apache/polaris/admintool/jdbc/init.sql")));
39+
Map.of(
40+
INIT_SCRIPT, "org/apache/polaris/admintool/jdbc/init.sql", "databases", "realm1,realm2,realm3")));
4041
}
4142
}

‎quarkus/service/build.gradle.kts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ dependencies {
108108
testImplementation("software.amazon.awssdk:kms")
109109
testImplementation("software.amazon.awssdk:dynamodb")
110110

111-
runtimeOnly(project(":polaris-relational-jdbc"))
111+
implementation(project(":polaris-relational-jdbc"))
112112
runtimeOnly("io.quarkus:quarkus-jdbc-postgresql") {
113113
exclude(group = "org.antlr", module = "antlr4-runtime")
114114
exclude(group = "org.scala-lang", module = "scala-library")

‎quarkus/admin/src/main/java/org/apache/polaris/admintool/config/QuarkusJdbcDataSource.java renamed to ‎quarkus/service/src/main/java/org/apache/polaris/service/quarkus/config/QuarkusDatasourceSupplier.java

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,30 +17,35 @@
1717
* under the License.
1818
*/
1919

20-
package org.apache.polaris.admintool.config;
20+
package org.apache.polaris.service.quarkus.config;
2121

22+
import com.google.api.client.util.Preconditions;
2223
import io.quarkus.arc.InstanceHandle;
23-
import org.apache.polaris.extension.persistence.relational.jdbc.JdbcDatasource;
24+
import org.apache.polaris.extension.persistence.relational.jdbc.DatasourceSupplier;
2425

2526
import javax.sql.DataSource;
26-
import java.util.ArrayList;
2727
import java.util.List;
2828

29-
public class QuarkusJdbcDataSource implements JdbcDatasource {
30-
private List<InstanceHandle<DataSource>> dataSources;
29+
public class QuarkusDatasourceSupplier implements DatasourceSupplier {
30+
private final List<InstanceHandle<DataSource>> dataSources;
31+
private final String realmIsolationType;
3132

3233
public static final String DEFAULT_DATA_SOURCE_NAME = "<default>";
34+
public static final String REALM_ISOLATION_TYPE_DB = "db";
3335

34-
public QuarkusJdbcDataSource(List<InstanceHandle<DataSource>> dataSources) {
36+
public QuarkusDatasourceSupplier(String realmIsolationType, List<InstanceHandle<DataSource>> dataSources) {
37+
this.realmIsolationType = realmIsolationType;
3538
this.dataSources = dataSources;
3639
}
3740

3841
@Override
3942
public DataSource fromRealmId(String realmId) {
43+
Preconditions.checkNotNull(REALM_ISOLATION_TYPE_DB.equals(realmIsolationType) && dataSources.size() == 1, "Misconfigured datasource's as per realm requirements");
4044
for (InstanceHandle<DataSource> handle : dataSources) {
4145
String name = handle.getBean().getName();
4246
name = name == null ? DEFAULT_DATA_SOURCE_NAME : unquoteDataSourceName(name);
43-
if (name.equals(realmId)) {
47+
// if realm isolation is DB then there should be only one DS configured.
48+
if (name.equals(realmId) || (REALM_ISOLATION_TYPE_DB.equals(realmIsolationType) && name.equals(DEFAULT_DATA_SOURCE_NAME))) {
4449
return handle.get();
4550
}
4651
}

‎quarkus/service/src/main/java/org/apache/polaris/service/quarkus/config/QuarkusProducers.java

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
*/
1919
package org.apache.polaris.service.quarkus.config;
2020

21+
import io.quarkus.arc.All;
22+
import io.quarkus.arc.InstanceHandle;
2123
import io.quarkus.runtime.StartupEvent;
2224
import io.smallrye.common.annotation.Identifier;
2325
import io.smallrye.context.SmallRyeManagedExecutor;
@@ -32,6 +34,8 @@
3234
import jakarta.ws.rs.container.ContainerRequestContext;
3335
import jakarta.ws.rs.core.Context;
3436
import java.time.Clock;
37+
import java.util.List;
38+
3539
import org.apache.polaris.core.PolarisCallContext;
3640
import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
3741
import org.apache.polaris.core.PolarisDiagnostics;
@@ -49,6 +53,7 @@
4953
import org.apache.polaris.core.secrets.UserSecretsManager;
5054
import org.apache.polaris.core.secrets.UserSecretsManagerFactory;
5155
import org.apache.polaris.core.storage.cache.StorageCredentialCache;
56+
import org.apache.polaris.extension.persistence.relational.jdbc.DatasourceSupplier;
5257
import org.apache.polaris.service.auth.ActiveRolesProvider;
5358
import org.apache.polaris.service.auth.Authenticator;
5459
import org.apache.polaris.service.auth.TokenBrokerFactory;
@@ -72,6 +77,8 @@
7277
import org.eclipse.microprofile.context.ManagedExecutor;
7378
import org.eclipse.microprofile.context.ThreadContext;
7479

80+
import javax.sql.DataSource;
81+
7582
public class QuarkusProducers {
7683

7784
@Produces
@@ -261,4 +268,9 @@ public ActiveRolesProvider activeRolesProvider(
261268
public void closeTaskExecutor(@Disposes @Identifier("task-executor") ManagedExecutor executor) {
262269
executor.close();
263270
}
271+
272+
@Produces
273+
public DatasourceSupplier jdbcDatasource(@ConfigProperty(name = "polaris.persistence.realm.isolation.type") String realmIsolationType, @All List<InstanceHandle<DataSource>> dataSources) {
274+
return new QuarkusDatasourceSupplier(realmIsolationType, dataSources);
275+
}
264276
}

‎quarkus/test-commons/src/main/java/org/apache/polaris/test/commons/PostgresRelationalJdbcLifeCycleManagement.java

Lines changed: 29 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -21,21 +21,28 @@
2121

2222
import io.quarkus.test.common.DevServicesContext;
2323
import io.quarkus.test.common.QuarkusTestResourceLifecycleManager;
24+
25+
import java.util.Arrays;
26+
import java.util.HashMap;
2427
import java.util.Map;
28+
import java.util.List;
2529
import org.testcontainers.containers.PostgreSQLContainer;
2630
import org.testcontainers.utility.DockerImageName;
2731

2832
public class PostgresRelationalJdbcLifeCycleManagement
2933
implements QuarkusTestResourceLifecycleManager, DevServicesContext.ContextAware {
3034
public static final String INIT_SCRIPT = "init-script";
35+
public static final String DATABASES = "databases";
3136

3237
private PostgreSQLContainer<?> postgres;
3338
private String initScript;
39+
private List<String> databases;
3440
private DevServicesContext context;
3541

3642
@Override
3743
public void init(Map<String, String> initArgs) {
3844
initScript = initArgs.get(INIT_SCRIPT);
45+
databases = Arrays.stream(initArgs.getOrDefault(DATABASES, "").split(",")).toList();
3946
}
4047

4148
@Override
@@ -53,25 +60,28 @@ public Map<String, String> start() {
5360

5461
context.containerNetworkId().ifPresent(postgres::withNetworkMode);
5562
postgres.start();
56-
// Use Map.ofEntries to create the map with more than 10 entries
57-
return Map.ofEntries(
58-
Map.entry("polaris.persistence.type", "relational-jdbc"),
59-
Map.entry("quarkus.datasource.realm1.db-kind", "pgsql"),
60-
Map.entry("quarkus.datasource.realm1.active", "true"),
61-
Map.entry("quarkus.datasource.realm1.jdbc.url", postgres.getJdbcUrl()),
62-
Map.entry("quarkus.datasource.realm1.username", postgres.getUsername()),
63-
Map.entry("quarkus.datasource.realm1.password", postgres.getPassword()),
64-
Map.entry("quarkus.datasource.realm2.db-kind", "pgsql"),
65-
Map.entry("quarkus.datasource.realm2.active", "true"),
66-
Map.entry("quarkus.datasource.realm2.jdbc.url", postgres.getJdbcUrl().replace("realm1", "realm2")),
67-
Map.entry("quarkus.datasource.realm2.username", postgres.getUsername()),
68-
Map.entry("quarkus.datasource.realm2.password", postgres.getPassword()),
69-
Map.entry("quarkus.datasource.realm3.db-kind", "pgsql"),
70-
Map.entry("quarkus.datasource.realm3.active", "true"),
71-
Map.entry("quarkus.datasource.realm3.jdbc.url", postgres.getJdbcUrl().replace("realm1", "realm3")),
72-
Map.entry("quarkus.datasource.realm3.username", postgres.getUsername()),
73-
Map.entry("quarkus.datasource.realm3.password", postgres.getPassword())
74-
);
63+
Map<String, String> props = new HashMap<>();
64+
props.put("polaris.persistence.type", "relational-jdbc");
65+
props.put("polaris.persistence.realm.isolation.type", "table");
66+
67+
if (!databases.isEmpty()) {
68+
for (String database : databases) {
69+
props.put("polaris.persistence.realm.isolation.type", "db");
70+
props.put(String.format("quarkus.datasource.%s.db-kind", database), "pgsql");
71+
props.put(String.format("quarkus.datasource.%s.active", database), "true");
72+
props.put(String.format("quarkus.datasource.%s.jdbc.url", database), postgres.getJdbcUrl().replace("realm1", database));
73+
props.put(String.format("quarkus.datasource.%s.username", database), postgres.getUsername());
74+
props.put(String.format("quarkus.datasource.%s.password", database), postgres.getPassword());
75+
}
76+
} else {
77+
props.put("quarkus.datasource.db-kind", "pgsql");
78+
props.put("quarkus.datasource.active", "true");
79+
props.put("quarkus.datasource.jdbc.url", postgres.getJdbcUrl());
80+
props.put("quarkus.datasource.jdbc.username", postgres.getUsername());
81+
props.put("quarkus.datasource.jdbc.password", postgres.getPassword());
82+
}
83+
84+
return props;
7585
}
7686

7787
@Override

0 commit comments

Comments
 (0)