Skip to content

fix(ebean-dao): honor isTestMode in batchUpsert and create paths #624

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jphui merged 1 commit into from
May 16, 2026
Merged

fix(ebean-dao): honor isTestMode in batchUpsert and create paths #624

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -256,7 +256,7 @@ public <ASPECT_UNION extends RecordTemplate> int create(

// Use comprehensive helper to prepare SqlUpdate with all common logic
SqlUpdate sqlUpdate = prepareMultiColumnInsert(urn, aspectValues, aspectCreateLambdas,
auditStamp, ingestionTrackingContext, onDuplicateKeyClause);
auditStamp, ingestionTrackingContext, onDuplicateKeyClause, isTestMode);

return sqlUpdate.execute();
}
Expand Down Expand Up @@ -294,7 +294,7 @@ public <ASPECT_UNION extends RecordTemplate> int batchUpsert(

// Use comprehensive helper to prepare SqlUpdate with all common logic
SqlUpdate sqlUpdate = prepareMultiColumnInsert(urn, aspectValues, aspectUpdateLambdas,
auditStamp, ingestionTrackingContext, onDuplicateKeyClause);
auditStamp, ingestionTrackingContext, onDuplicateKeyClause, isTestMode);

return sqlUpdate.execute();
}
Expand Down Expand Up @@ -889,7 +889,8 @@ private SqlUpdate prepareMultiColumnInsert(
@Nonnull List<? extends BaseLocalDAO.AspectUpdateLambda<? extends RecordTemplate>> aspectLambdas,
@Nonnull AuditStamp auditStamp,
@Nullable IngestionTrackingContext ingestionTrackingContext,
@Nonnull String onDuplicateKeyClause) {
@Nonnull String onDuplicateKeyClause,
boolean isTestMode) {

// Validate that aspectValues and aspectLambdas have the same size
if (aspectValues.size() != aspectLambdas.size()) {
Expand Down Expand Up @@ -942,7 +943,8 @@ private SqlUpdate prepareMultiColumnInsert(

// Build complete SQL statement with ON DUPLICATE KEY clause
String insertStatement = insertIntoSql.toString() + insertSqlValues.toString() + onDuplicateKeyClause;
insertStatement = String.format(insertStatement, getTableName(urn));
insertStatement = String.format(insertStatement,
isTestMode ? getTestTableName(urn) : getTableName(urn));

// Create SqlUpdate and set all parameters
SqlUpdate sqlUpdate = _server.createSqlUpdate(insertStatement);
Expand Down
32 changes: 30 additions & 2 deletions dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalAccessTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -701,12 +701,40 @@ public void testBatchUpsertWithIngestionTrackingContext() {
assertEquals("{\"value\":\"tracked_value\"}", results.get(0).getMetadata());

// Verify IngestionTrackingContext fields are persisted and readable
assertEquals("test-emitter", results.get(0).getEmitter(),
assertEquals("test-emitter", results.get(0).getEmitter(),
"Emitter from IngestionTrackingContext should be persisted");
assertEquals(Long.valueOf(emitTime), results.get(0).getEmitTime(),
assertEquals(Long.valueOf(emitTime), results.get(0).getEmitTime(),
"EmitTime from IngestionTrackingContext should be persisted");
}

/**
* Tests that batchUpsert() with isTestMode=true writes to the test table rather than the
* production table. Regression coverage for the previously-dropped flag in
* {@code prepareMultiColumnInsert}, which used to hardcode {@code getTableName(urn)}.
*/
@Test
public void testBatchUpsertWithTestModeWritesToTestTable() {
// Arrange
FooUrn fooUrn = makeFooUrn(310);
AspectFoo foo = new AspectFoo().setValue("test_mode_value");
List<BaseLocalDAO.AspectUpdateContext<RecordTemplate>> updateContexts =
Collections.singletonList(new BaseLocalDAO.AspectUpdateContext<>(null, foo,
new BaseLocalDAO.AspectUpdateLambda<>(foo)));
AuditStamp auditStamp = makeAuditStamp("actor", _now);

// Act
int result = _ebeanLocalAccessFoo.batchUpsert(fooUrn, updateContexts, auditStamp, null, true);

// Assert - row should land in the test table, readable via batchGetUnion with isTestMode=true
assertEquals(result, 1);
AspectKey<FooUrn, AspectFoo> aspectKey = new AspectKey<>(AspectFoo.class, fooUrn, 0L);
List<EbeanMetadataAspect> testTableResults = _ebeanLocalAccessFoo.batchGetUnion(
Collections.singletonList(aspectKey), 1, 0, false, true);
assertEquals(1, testTableResults.size());
assertEquals("{\"value\":\"test_mode_value\"}", testTableResults.get(0).getMetadata());
assertEquals(fooUrn.toString(), testTableResults.get(0).getKey().getUrn());
}

// ==================== readDeletionInfoBatch tests ====================

/**
Expand Down
13 changes: 13 additions & 0 deletions ...src/test/resources/ebean-local-access-create-all-with-non-dollar-virtual-column-names.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
DROP TABLE IF EXISTS metadata_entity_foo;
DROP TABLE IF EXISTS metadata_entity_foo_test;
DROP TABLE IF EXISTS metadata_entity_bar;
DROP TABLE IF EXISTS metadata_entity_burger;
DROP TABLE IF EXISTS metadata_aspect;
Expand All @@ -16,6 +17,16 @@ CREATE TABLE IF NOT EXISTS metadata_entity_foo (
CONSTRAINT pk_metadata_entity_foo PRIMARY KEY (urn)
);

-- initialize foo test entity table (mirror of metadata_entity_foo, used for isTestMode writes)
CREATE TABLE IF NOT EXISTS metadata_entity_foo_test (
urn VARCHAR(100) NOT NULL,
lastmodifiedon TIMESTAMP NOT NULL,
lastmodifiedby VARCHAR(255) NOT NULL,
createdfor VARCHAR(255),
deleted_ts datetime(6) DEFAULT NULL,
CONSTRAINT pk_metadata_entity_foo_test PRIMARY KEY (urn)
);

-- initialize bar entity table
CREATE TABLE IF NOT EXISTS metadata_entity_bar (
urn VARCHAR(100) NOT NULL,
Expand Down Expand Up @@ -68,6 +79,7 @@ CREATE TABLE IF NOT EXISTS metadata_relationship_belongsto (
);

ALTER TABLE metadata_entity_foo ADD a_urn JSON;
ALTER TABLE metadata_entity_foo_test ADD a_urn JSON;
ALTER TABLE metadata_entity_bar ADD a_urn JSON;

ALTER TABLE metadata_entity_foo ADD COLUMN i_urn0fooId VARCHAR(255)
Expand All @@ -78,6 +90,7 @@ ALTER TABLE metadata_entity_foo ADD a_status JSON;

-- add foo aspect to foo entity
ALTER TABLE metadata_entity_foo ADD a_aspectfoo JSON;
ALTER TABLE metadata_entity_foo_test ADD a_aspectfoo JSON;

-- add foo aspect to bar entity
ALTER TABLE metadata_entity_bar ADD a_aspectfoo JSON;
Expand Down
13 changes: 13 additions & 0 deletions dao-impl/ebean-dao/src/test/resources/ebean-local-access-create-all.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
DROP TABLE IF EXISTS metadata_entity_foo;
DROP TABLE IF EXISTS metadata_entity_foo_test;
DROP TABLE IF EXISTS metadata_entity_bar;
DROP TABLE IF EXISTS metadata_entity_burger;
DROP TABLE IF EXISTS metadata_aspect;
Expand All @@ -16,6 +17,16 @@ CREATE TABLE IF NOT EXISTS metadata_entity_foo (
CONSTRAINT pk_metadata_entity_foo PRIMARY KEY (urn)
);

-- initialize foo test entity table (mirror of metadata_entity_foo, used for isTestMode writes)
CREATE TABLE IF NOT EXISTS metadata_entity_foo_test (
urn VARCHAR(100) NOT NULL,
lastmodifiedon TIMESTAMP NOT NULL,
lastmodifiedby VARCHAR(255) NOT NULL,
createdfor VARCHAR(255),
deleted_ts datetime(6) DEFAULT NULL,
CONSTRAINT pk_metadata_entity_foo_test PRIMARY KEY (urn)
);

-- initialize bar entity table
CREATE TABLE IF NOT EXISTS metadata_entity_bar (
urn VARCHAR(100) NOT NULL,
Expand Down Expand Up @@ -68,6 +79,7 @@ CREATE TABLE IF NOT EXISTS metadata_relationship_belongsto (
);

ALTER TABLE metadata_entity_foo ADD a_urn JSON;
ALTER TABLE metadata_entity_foo_test ADD a_urn JSON;
ALTER TABLE metadata_entity_bar ADD a_urn JSON;

ALTER TABLE metadata_entity_foo ADD COLUMN i_urn$fooId VARCHAR(255)
Expand All @@ -78,6 +90,7 @@ ALTER TABLE metadata_entity_foo ADD a_status JSON;

-- add foo aspect to foo entity
ALTER TABLE metadata_entity_foo ADD a_aspectfoo JSON;
ALTER TABLE metadata_entity_foo_test ADD a_aspectfoo JSON;

-- add foo aspect to bar entity
ALTER TABLE metadata_entity_bar ADD a_aspectfoo JSON;
Expand Down
Loading