Skip to content

Commit 7df4703

Browse files
jphuijphui
authored
fix(ebean-dao): honor isTestMode in batchUpsert and create paths (#624)
1 parent b025046 commit 7df4703

4 files changed

Lines changed: 62 additions & 6 deletions

File tree

‎dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -256,7 +256,7 @@ public <ASPECT_UNION extends RecordTemplate> int create(
256256

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

261261
return sqlUpdate.execute();
262262
}
@@ -294,7 +294,7 @@ public <ASPECT_UNION extends RecordTemplate> int batchUpsert(
294294

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

299299
return sqlUpdate.execute();
300300
}
@@ -889,7 +889,8 @@ private SqlUpdate prepareMultiColumnInsert(
889889
@Nonnull List<? extends BaseLocalDAO.AspectUpdateLambda<? extends RecordTemplate>> aspectLambdas,
890890
@Nonnull AuditStamp auditStamp,
891891
@Nullable IngestionTrackingContext ingestionTrackingContext,
892-
@Nonnull String onDuplicateKeyClause) {
892+
@Nonnull String onDuplicateKeyClause,
893+
boolean isTestMode) {
893894

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

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

947949
// Create SqlUpdate and set all parameters
948950
SqlUpdate sqlUpdate = _server.createSqlUpdate(insertStatement);

‎dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalAccessTest.java‎

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -701,12 +701,40 @@ public void testBatchUpsertWithIngestionTrackingContext() {
701701
assertEquals("{\"value\":\"tracked_value\"}", results.get(0).getMetadata());
702702

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

710+
/**
711+
* Tests that batchUpsert() with isTestMode=true writes to the test table rather than the
712+
* production table. Regression coverage for the previously-dropped flag in
713+
* {@code prepareMultiColumnInsert}, which used to hardcode {@code getTableName(urn)}.
714+
*/
715+
@Test
716+
public void testBatchUpsertWithTestModeWritesToTestTable() {
717+
// Arrange
718+
FooUrn fooUrn = makeFooUrn(310);
719+
AspectFoo foo = new AspectFoo().setValue("test_mode_value");
720+
List<BaseLocalDAO.AspectUpdateContext<RecordTemplate>> updateContexts =
721+
Collections.singletonList(new BaseLocalDAO.AspectUpdateContext<>(null, foo,
722+
new BaseLocalDAO.AspectUpdateLambda<>(foo)));
723+
AuditStamp auditStamp = makeAuditStamp("actor", _now);
724+
725+
// Act
726+
int result = _ebeanLocalAccessFoo.batchUpsert(fooUrn, updateContexts, auditStamp, null, true);
727+
728+
// Assert - row should land in the test table, readable via batchGetUnion with isTestMode=true
729+
assertEquals(result, 1);
730+
AspectKey<FooUrn, AspectFoo> aspectKey = new AspectKey<>(AspectFoo.class, fooUrn, 0L);
731+
List<EbeanMetadataAspect> testTableResults = _ebeanLocalAccessFoo.batchGetUnion(
732+
Collections.singletonList(aspectKey), 1, 0, false, true);
733+
assertEquals(1, testTableResults.size());
734+
assertEquals("{\"value\":\"test_mode_value\"}", testTableResults.get(0).getMetadata());
735+
assertEquals(fooUrn.toString(), testTableResults.get(0).getKey().getUrn());
736+
}
737+
710738
// ==================== readDeletionInfoBatch tests ====================
711739

712740
/**

‎dao-impl/ebean-dao/src/test/resources/ebean-local-access-create-all-with-non-dollar-virtual-column-names.sql‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
DROP TABLE IF EXISTS metadata_entity_foo;
2+
DROP TABLE IF EXISTS metadata_entity_foo_test;
23
DROP TABLE IF EXISTS metadata_entity_bar;
34
DROP TABLE IF EXISTS metadata_entity_burger;
45
DROP TABLE IF EXISTS metadata_aspect;
@@ -16,6 +17,16 @@ CREATE TABLE IF NOT EXISTS metadata_entity_foo (
1617
CONSTRAINT pk_metadata_entity_foo PRIMARY KEY (urn)
1718
);
1819

20+
-- initialize foo test entity table (mirror of metadata_entity_foo, used for isTestMode writes)
21+
CREATE TABLE IF NOT EXISTS metadata_entity_foo_test (
22+
urn VARCHAR(100) NOT NULL,
23+
lastmodifiedon TIMESTAMP NOT NULL,
24+
lastmodifiedby VARCHAR(255) NOT NULL,
25+
createdfor VARCHAR(255),
26+
deleted_ts datetime(6) DEFAULT NULL,
27+
CONSTRAINT pk_metadata_entity_foo_test PRIMARY KEY (urn)
28+
);
29+
1930
-- initialize bar entity table
2031
CREATE TABLE IF NOT EXISTS metadata_entity_bar (
2132
urn VARCHAR(100) NOT NULL,
@@ -68,6 +79,7 @@ CREATE TABLE IF NOT EXISTS metadata_relationship_belongsto (
6879
);
6980

7081
ALTER TABLE metadata_entity_foo ADD a_urn JSON;
82+
ALTER TABLE metadata_entity_foo_test ADD a_urn JSON;
7183
ALTER TABLE metadata_entity_bar ADD a_urn JSON;
7284

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

7991
-- add foo aspect to foo entity
8092
ALTER TABLE metadata_entity_foo ADD a_aspectfoo JSON;
93+
ALTER TABLE metadata_entity_foo_test ADD a_aspectfoo JSON;
8194

8295
-- add foo aspect to bar entity
8396
ALTER TABLE metadata_entity_bar ADD a_aspectfoo JSON;

‎dao-impl/ebean-dao/src/test/resources/ebean-local-access-create-all.sql‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
DROP TABLE IF EXISTS metadata_entity_foo;
2+
DROP TABLE IF EXISTS metadata_entity_foo_test;
23
DROP TABLE IF EXISTS metadata_entity_bar;
34
DROP TABLE IF EXISTS metadata_entity_burger;
45
DROP TABLE IF EXISTS metadata_aspect;
@@ -16,6 +17,16 @@ CREATE TABLE IF NOT EXISTS metadata_entity_foo (
1617
CONSTRAINT pk_metadata_entity_foo PRIMARY KEY (urn)
1718
);
1819

20+
-- initialize foo test entity table (mirror of metadata_entity_foo, used for isTestMode writes)
21+
CREATE TABLE IF NOT EXISTS metadata_entity_foo_test (
22+
urn VARCHAR(100) NOT NULL,
23+
lastmodifiedon TIMESTAMP NOT NULL,
24+
lastmodifiedby VARCHAR(255) NOT NULL,
25+
createdfor VARCHAR(255),
26+
deleted_ts datetime(6) DEFAULT NULL,
27+
CONSTRAINT pk_metadata_entity_foo_test PRIMARY KEY (urn)
28+
);
29+
1930
-- initialize bar entity table
2031
CREATE TABLE IF NOT EXISTS metadata_entity_bar (
2132
urn VARCHAR(100) NOT NULL,
@@ -68,6 +79,7 @@ CREATE TABLE IF NOT EXISTS metadata_relationship_belongsto (
6879
);
6980

7081
ALTER TABLE metadata_entity_foo ADD a_urn JSON;
82+
ALTER TABLE metadata_entity_foo_test ADD a_urn JSON;
7183
ALTER TABLE metadata_entity_bar ADD a_urn JSON;
7284

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

7991
-- add foo aspect to foo entity
8092
ALTER TABLE metadata_entity_foo ADD a_aspectfoo JSON;
93+
ALTER TABLE metadata_entity_foo_test ADD a_aspectfoo JSON;
8194

8295
-- add foo aspect to bar entity
8396
ALTER TABLE metadata_entity_bar ADD a_aspectfoo JSON;

0 commit comments

Comments
 (0)