Skip to content

Commit 7fc7018

Browse files
ayushtknayushtkn
authored
Fix JDBC writeEntities issuing a wasteful full-row lookup per entity (apache#5134)
1 parent 0d164ef commit 7fc7018

2 files changed

Lines changed: 162 additions & 20 deletions

File tree

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

Lines changed: 41 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import java.util.List;
3232
import java.util.Map;
3333
import java.util.Optional;
34+
import java.util.concurrent.atomic.AtomicBoolean;
3435
import java.util.concurrent.atomic.AtomicReference;
3536
import java.util.function.Function;
3637
import java.util.function.Predicate;
@@ -93,6 +94,21 @@ public class JdbcBasePersistenceImpl implements BasePersistence, IntegrationPers
9394
// The max number of components a location can have before the optimized sibling check is not used
9495
private static final int MAX_LOCATION_COMPONENTS = 40;
9596

97+
// Converter for SELECT 1 existence probes: only the presence of a row matters, so every row maps
98+
// to 1 and toMap is never used (existence queries are read-only).
99+
private static final Converter<Integer> ROW_EXISTS_CONVERTER =
100+
new Converter<>() {
101+
@Override
102+
public Integer fromResultSet(ResultSet rs) {
103+
return 1;
104+
}
105+
106+
@Override
107+
public Map<String, Object> toMap(DatabaseType databaseType) {
108+
throw new UnsupportedOperationException();
109+
}
110+
};
111+
96112
public JdbcBasePersistenceImpl(
97113
PolarisDiagnostics diagnostics,
98114
DatasourceOperations databaseOperations,
@@ -143,15 +159,7 @@ public void writeEntities(
143159
PolarisBaseEntity entity = entities.get(i);
144160
PolarisBaseEntity originalEntity =
145161
originalEntities != null ? originalEntities.get(i) : null;
146-
// first, check if the entity has already been created, in which case we will simply
147-
// return it.
148-
PolarisBaseEntity entityFound =
149-
lookupEntity(
150-
callCtx, entity.getCatalogId(), entity.getId(), entity.getTypeCode());
151-
if (entityFound != null && originalEntity == null) {
152-
// probably the client retried, simply return it
153-
// TODO: Check correctness of returning entityFound vs entity here. It may have
154-
// already been updated after the creation.
162+
if (originalEntity == null && entityExists(connection, entity.getId())) {
155163
continue;
156164
}
157165
persistEntity(
@@ -167,6 +175,25 @@ public void writeEntities(
167175
}
168176
}
169177

178+
/**
179+
* Returns whether an entity row already exists, using a {@code SELECT 1 ... LIMIT 1} on the
180+
* supplied transaction connection. Only existence is needed by {@link #writeEntities} on the
181+
* create path, so this avoids fetching the full entity row (including the large JSON property
182+
* blobs).
183+
*/
184+
private boolean entityExists(@NonNull Connection connection, long entityId) throws SQLException {
185+
AtomicBoolean exists = new AtomicBoolean(false);
186+
datasourceOperations.executeSelectOverStream(
187+
connection,
188+
QueryGenerator.generateExistsQuery(
189+
ModelEntity.getAllColumnNames(schemaVersion),
190+
ModelEntity.TABLE_NAME,
191+
entityKeyParams(entityId)),
192+
ROW_EXISTS_CONVERTER,
193+
stream -> exists.set(stream.findAny().isPresent()));
194+
return exists.get();
195+
}
196+
170197
private void persistEntity(
171198
@NonNull PolarisCallContext callCtx,
172199
@NonNull PolarisBaseEntity entity,
@@ -422,6 +449,10 @@ public PolarisBaseEntity lookupEntity(
422449
ModelEntity.getAllColumnNames(schemaVersion), ModelEntity.TABLE_NAME, params));
423450
}
424451

452+
private Map<String, Object> entityKeyParams(long entityId) {
453+
return Map.of("id", entityId, "realm_id", realmId);
454+
}
455+
425456
@Override
426457
public PolarisBaseEntity lookupEntityByName(
427458
@NonNull PolarisCallContext callCtx,
@@ -740,17 +771,7 @@ public boolean hasChildren(
740771
datasourceOperations.executeSelect(
741772
QueryGenerator.generateExistsQuery(
742773
ModelEntity.getAllColumnNames(schemaVersion), ModelEntity.TABLE_NAME, params),
743-
new Converter<Integer>() {
744-
@Override
745-
public Integer fromResultSet(ResultSet rs) {
746-
return 1;
747-
}
748-
749-
@Override
750-
public Map<String, Object> toMap(DatabaseType databaseType) {
751-
throw new UnsupportedOperationException();
752-
}
753-
});
774+
ROW_EXISTS_CONVERTER);
754775
return results != null && !results.isEmpty();
755776
} catch (SQLException e) {
756777
throw new RuntimeException(

‎persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImplTest.java‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,11 +25,13 @@
2525
import static org.mockito.ArgumentMatchers.any;
2626
import static org.mockito.Mockito.doCallRealMethod;
2727
import static org.mockito.Mockito.doThrow;
28+
import static org.mockito.Mockito.never;
2829
import static org.mockito.Mockito.verify;
2930
import static org.mockito.Mockito.when;
3031

3132
import java.io.IOException;
3233
import java.io.InputStream;
34+
import java.sql.Connection;
3335
import java.sql.SQLException;
3436
import java.util.List;
3537
import java.util.Optional;
@@ -43,6 +45,7 @@
4345
import org.apache.polaris.core.entity.PolarisEntityType;
4446
import org.apache.polaris.core.entity.PolarisGrantRecord;
4547
import org.h2.jdbcx.JdbcConnectionPool;
48+
import org.junit.jupiter.api.Test;
4649
import org.junit.jupiter.params.ParameterizedTest;
4750
import org.junit.jupiter.params.provider.ValueSource;
4851
import org.mockito.ArgumentCaptor;
@@ -289,6 +292,124 @@ private static PolarisBaseEntity newTestEntity(
289292
.build();
290293
}
291294

295+
/**
296+
* On the create path (no original entity) {@code writeEntities} must detect an already-created
297+
* entity with a cheap {@code SELECT 1 ... LIMIT 1} existence check that does not fetch the large
298+
* JSON property blobs, and a retried create must stay idempotent.
299+
*/
300+
@Test
301+
void writeEntitiesCreatePathUsesExistsCheckNotFullRowFetch() throws SQLException, IOException {
302+
int schemaVersion = 5;
303+
JdbcConnectionPool dataSource =
304+
JdbcConnectionPool.create(
305+
"jdbc:h2:mem:write_entities_create_v"
306+
+ schemaVersion
307+
+ "_"
308+
+ System.nanoTime()
309+
+ ";DB_CLOSE_DELAY=-1",
310+
"sa",
311+
"");
312+
DatasourceOperations real = new DatasourceOperations(dataSource, new TestJdbcConfiguration());
313+
try (InputStream script = DatabaseType.H2.openInitScriptResource(schemaVersion)) {
314+
real.executeScript(script);
315+
}
316+
DatasourceOperations spy = Mockito.spy(real);
317+
318+
JdbcBasePersistenceImpl impl =
319+
new JdbcBasePersistenceImpl(
320+
new PolarisDefaultDiagServiceImpl(),
321+
spy,
322+
RANDOM_SECRETS,
323+
REALM_CONTEXT.getRealmIdentifier(),
324+
schemaVersion);
325+
PolarisCallContext callCtx = new PolarisCallContext(REALM_CONTEXT, impl);
326+
327+
PolarisBaseEntity entity = principalEntity(1L, "create-me", 1);
328+
impl.writeEntities(callCtx, List.of(entity), null);
329+
330+
// The create-path existence check must run as SELECT 1 ... LIMIT 1 (no JSON property columns),
331+
// on the transaction connection.
332+
ArgumentCaptor<QueryGenerator.PreparedQuery> captor =
333+
ArgumentCaptor.forClass(QueryGenerator.PreparedQuery.class);
334+
verify(spy).executeSelectOverStream(any(Connection.class), captor.capture(), any(), any());
335+
String sql = captor.getValue().sql();
336+
assertThat(sql).contains("SELECT 1").contains("LIMIT 1");
337+
assertThat(sql).doesNotContain("properties");
338+
339+
// The entity was actually inserted.
340+
assertThat(
341+
impl.lookupEntity(callCtx, entity.getCatalogId(), entity.getId(), entity.getTypeCode()))
342+
.isNotNull();
343+
344+
// A retried create is idempotent (the existence check short-circuits the insert).
345+
assertThatCode(() -> impl.writeEntities(callCtx, List.of(entity), null))
346+
.doesNotThrowAnyException();
347+
}
348+
349+
/**
350+
* On the update path (original entity supplied) {@code writeEntities} already knows the entity
351+
* exists, so it must not issue any pre-write existence/lookup SELECT before the CAS update.
352+
*/
353+
@Test
354+
void writeEntitiesUpdatePathSkipsExistenceLookup() throws SQLException, IOException {
355+
int schemaVersion = 5;
356+
JdbcConnectionPool dataSource =
357+
JdbcConnectionPool.create(
358+
"jdbc:h2:mem:write_entities_update_v"
359+
+ schemaVersion
360+
+ "_"
361+
+ System.nanoTime()
362+
+ ";DB_CLOSE_DELAY=-1",
363+
"sa",
364+
"");
365+
DatasourceOperations real = new DatasourceOperations(dataSource, new TestJdbcConfiguration());
366+
try (InputStream script = DatabaseType.H2.openInitScriptResource(schemaVersion)) {
367+
real.executeScript(script);
368+
}
369+
DatasourceOperations spy = Mockito.spy(real);
370+
371+
JdbcBasePersistenceImpl impl =
372+
new JdbcBasePersistenceImpl(
373+
new PolarisDefaultDiagServiceImpl(),
374+
spy,
375+
RANDOM_SECRETS,
376+
REALM_CONTEXT.getRealmIdentifier(),
377+
schemaVersion);
378+
PolarisCallContext callCtx = new PolarisCallContext(REALM_CONTEXT, impl);
379+
380+
PolarisBaseEntity original = principalEntity(1L, "update-me", 1);
381+
impl.writeEntities(callCtx, List.of(original), null);
382+
383+
// Only observe the update below.
384+
Mockito.clearInvocations(spy);
385+
386+
PolarisBaseEntity updated = principalEntity(1L, "update-me", 2);
387+
impl.writeEntities(callCtx, List.of(updated), List.of(original));
388+
389+
// No existence/lookup SELECT is issued before the CAS update.
390+
verify(spy, never()).executeSelectOverStream(any(Connection.class), any(), any(), any());
391+
392+
// The CAS update was applied (original v1 -> v2).
393+
PolarisBaseEntity reloaded =
394+
impl.lookupEntity(callCtx, updated.getCatalogId(), updated.getId(), updated.getTypeCode());
395+
assertThat(reloaded).isNotNull();
396+
assertThat(reloaded.getEntityVersion()).isEqualTo(2);
397+
}
398+
399+
private static PolarisBaseEntity principalEntity(long id, String name, int entityVersion) {
400+
return new PolarisBaseEntity.Builder()
401+
.id(id)
402+
.catalogId(0L)
403+
.parentId(0L)
404+
.typeCode(PolarisEntityType.PRINCIPAL.getCode())
405+
.subTypeCode(PolarisEntitySubType.NULL_SUBTYPE.getCode())
406+
.name(name)
407+
.entityVersion(entityVersion)
408+
.grantRecordsVersion(1)
409+
.createTimestamp(System.currentTimeMillis())
410+
.build();
411+
}
412+
292413
private static final class TestJdbcConfiguration implements RelationalJdbcConfiguration {
293414
@Override
294415
public Optional<Integer> maxRetries() {

0 commit comments

Comments
 (0)