diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/Document.java b/document-store/src/main/java/org/hypertrace/core/documentstore/Document.java index 944913aa..fe9c84b8 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/Document.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/Document.java @@ -3,4 +3,8 @@ public interface Document { String toJson(); + + default DocumentType getDocumentType() { + return DocumentType.NESTED; + } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/DocumentType.java b/document-store/src/main/java/org/hypertrace/core/documentstore/DocumentType.java new file mode 100644 index 00000000..948b3898 --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/DocumentType.java @@ -0,0 +1,8 @@ +package org.hypertrace.core.documentstore; + +public enum DocumentType { + // FLAT documents contain individual columns for each attribute + FLAT, + // NESTED documents contains attributes as a nested JSON document + NESTED +} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/JSONDocument.java b/document-store/src/main/java/org/hypertrace/core/documentstore/JSONDocument.java index c04d806d..85d23d70 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/JSONDocument.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/JSONDocument.java @@ -11,19 +11,40 @@ public class JSONDocument implements Document { private static ObjectMapper mapper = new ObjectMapper(); private JsonNode node; + private DocumentType documentType = DocumentType.NESTED; public JSONDocument(String json) throws IOException { node = mapper.readTree(json); } + public JSONDocument(String json, DocumentType documentType) throws IOException { + node = mapper.readTree(json); + this.documentType = documentType; + } + public JSONDocument(Object object) throws IOException { node = mapper.readTree(mapper.writeValueAsString(object)); } + public JSONDocument(Object object, DocumentType documentType) throws IOException { + node = mapper.readTree(mapper.writeValueAsString(object)); + this.documentType = documentType; + } + public JSONDocument(JsonNode node) { this.node = node; } + public JSONDocument(JsonNode node, DocumentType documentType) { + this.node = node; + this.documentType = documentType; + } + + @Override + public DocumentType getDocumentType() { + return this.documentType; + } + @Override public String toJson() { try { @@ -54,6 +75,6 @@ public boolean equals(Object obj) { } JSONDocument other = (JSONDocument) obj; - return Objects.equals(node, other.node); + return Objects.equals(node, other.node) && documentType == other.documentType; } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java index 5582889a..8c6444be 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java @@ -61,6 +61,7 @@ import org.hypertrace.core.documentstore.Collection; import org.hypertrace.core.documentstore.CreateResult; import org.hypertrace.core.documentstore.Document; +import org.hypertrace.core.documentstore.DocumentType; import org.hypertrace.core.documentstore.Filter; import org.hypertrace.core.documentstore.JSONDocument; import org.hypertrace.core.documentstore.Key; @@ -1264,6 +1265,10 @@ public PostgresResultIteratorWithBasicTypes(ResultSet resultSet) { super(resultSet); } + public PostgresResultIteratorWithBasicTypes(ResultSet resultSet, DocumentType documentType) { + super(resultSet, documentType); + } + @Override public Document next() { try { @@ -1299,7 +1304,7 @@ protected Document prepareDocument() throws SQLException, IOException { jsonNode.remove(DOCUMENT_ID); } - return new JSONDocument(MAPPER.writeValueAsString(jsonNode)); + return new JSONDocument(MAPPER.writeValueAsString(jsonNode), documentType); } private void addColumnToJsonNode( @@ -1386,17 +1391,29 @@ static class PostgresResultIterator implements CloseableIterator { protected final ObjectMapper MAPPER = new ObjectMapper(); protected ResultSet resultSet; - private final boolean removeDocumentId; protected boolean cursorMovedForward = false; protected boolean hasNext = false; + private final boolean removeDocumentId; + protected DocumentType documentType; + public PostgresResultIterator(ResultSet resultSet) { this(resultSet, true); } PostgresResultIterator(ResultSet resultSet, boolean removeDocumentId) { + this(resultSet, removeDocumentId, DocumentType.NESTED); + } + + public PostgresResultIterator(ResultSet resultSet, DocumentType documentType) { + this(resultSet, true, documentType); + } + + PostgresResultIterator( + ResultSet resultSet, boolean removeDocumentId, DocumentType documentType) { this.resultSet = resultSet; this.removeDocumentId = removeDocumentId; + this.documentType = documentType; } @Override @@ -1447,7 +1464,7 @@ protected Document prepareDocument() throws SQLException, IOException { jsonNode.put(CREATED_AT, String.valueOf(createdAt)); jsonNode.put(UPDATED_AT, String.valueOf(updatedAt)); - return new JSONDocument(MAPPER.writeValueAsString(jsonNode)); + return new JSONDocument(MAPPER.writeValueAsString(jsonNode), documentType); } protected void closeResultSet() { @@ -1508,7 +1525,7 @@ protected Document prepareDocument() throws SQLException, IOException { } } } - return new JSONDocument(MAPPER.writeValueAsString(jsonNode)); + return new JSONDocument(MAPPER.writeValueAsString(jsonNode), documentType); } private String getColumnValue( diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java index 389a52b8..33ea637c 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryExecutor.java @@ -10,6 +10,7 @@ import lombok.extern.slf4j.Slf4j; import org.hypertrace.core.documentstore.CloseableIterator; import org.hypertrace.core.documentstore.Document; +import org.hypertrace.core.documentstore.DocumentType; import org.hypertrace.core.documentstore.postgres.PostgresCollection.PostgresResultIterator; import org.hypertrace.core.documentstore.postgres.PostgresCollection.PostgresResultIteratorWithMetaData; import org.hypertrace.core.documentstore.postgres.query.v1.transformer.PostgresQueryTransformer; @@ -18,6 +19,7 @@ @Slf4j @AllArgsConstructor public class PostgresQueryExecutor { + private final PostgresTableIdentifier tableIdentifier; public CloseableIterator execute(final Connection connection, final Query query) { @@ -37,7 +39,8 @@ public CloseableIterator execute( final ResultSet resultSet = preparedStatement.executeQuery(); if ((tableIdentifier.getTableName().equals(flatStructureCollectionName))) { - return new PostgresCollection.PostgresResultIteratorWithBasicTypes(resultSet); + return new PostgresCollection.PostgresResultIteratorWithBasicTypes( + resultSet, DocumentType.FLAT); } return query.getSelections().size() > 0 ? new PostgresResultIteratorWithMetaData(resultSet) diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/JSONDocumentTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/JSONDocumentTest.java index 5cd3be9a..e4ec510d 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/JSONDocumentTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/JSONDocumentTest.java @@ -1,11 +1,15 @@ package org.hypertrace.core.documentstore; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; import java.util.Map; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; public class JSONDocumentTest { + private static final ObjectMapper mapper = new ObjectMapper(); + @Test public void testJSONDocument() throws Exception { Map data = Map.of("key1", "value1", "key2", "value2"); @@ -14,4 +18,92 @@ public void testJSONDocument() throws Exception { Assertions.assertEquals(document1, document2); Assertions.assertEquals(document1.toJson(), document2.toJson()); } + + @Test + public void testJSONDocumentWithDefaultDocumentType() throws Exception { + Map data = Map.of("key1", "value1", "key2", "value2"); + JSONDocument document = new JSONDocument(data); + Assertions.assertEquals(DocumentType.NESTED, document.getDocumentType()); + } + + @Test + public void testJSONDocumentWithExplicitDocumentType() throws Exception { + Map data = Map.of("key1", "value1", "key2", "value2"); + + JSONDocument nestedDocument = new JSONDocument(data, DocumentType.NESTED); + Assertions.assertEquals(DocumentType.NESTED, nestedDocument.getDocumentType()); + + JSONDocument flatDocument = new JSONDocument(data, DocumentType.FLAT); + Assertions.assertEquals(DocumentType.FLAT, flatDocument.getDocumentType()); + } + + @Test + public void testJSONDocumentConstructorsWithDocumentType() throws Exception { + String json = "{\"key1\":\"value1\",\"key2\":\"value2\"}"; + JsonNode node = mapper.readTree(json); + Map data = Map.of("key1", "value1", "key2", "value2"); + + // Test string constructor with DocumentType + JSONDocument stringDoc = new JSONDocument(json, DocumentType.FLAT); + Assertions.assertEquals(DocumentType.FLAT, stringDoc.getDocumentType()); + + // Test object constructor with DocumentType + JSONDocument objectDoc = new JSONDocument(data, DocumentType.FLAT); + Assertions.assertEquals(DocumentType.FLAT, objectDoc.getDocumentType()); + + // Test JsonNode constructor with DocumentType + JSONDocument nodeDoc = new JSONDocument(node, DocumentType.FLAT); + Assertions.assertEquals(DocumentType.FLAT, nodeDoc.getDocumentType()); + } + + @Test + public void testEqualsWithSameContentDifferentDocumentType() throws Exception { + Map data = Map.of("key1", "value1", "key2", "value2"); + + JSONDocument nestedDoc = new JSONDocument(data, DocumentType.NESTED); + JSONDocument flatDoc = new JSONDocument(data, DocumentType.FLAT); + + // Current implementation only compares JsonNode, not DocumentType + // This test documents the current behavior - documents with same content but different types + // are equal + Assertions.assertNotEquals(nestedDoc, flatDoc); + Assertions.assertEquals(nestedDoc.toJson(), flatDoc.toJson()); + } + + @Test + public void testEqualsWithSameContentSameDocumentType() throws Exception { + Map data = Map.of("key1", "value1", "key2", "value2"); + + JSONDocument doc1 = new JSONDocument(data, DocumentType.FLAT); + JSONDocument doc2 = new JSONDocument(data, DocumentType.FLAT); + + Assertions.assertEquals(doc1, doc2); + Assertions.assertEquals(doc1.getDocumentType(), doc2.getDocumentType()); + } + + @Test + public void testErrorDocument() throws Exception { + String errorMessage = "Test error message"; + JSONDocument errorDoc = JSONDocument.errorDocument(errorMessage); + + // Verify default document type + Assertions.assertEquals(DocumentType.NESTED, errorDoc.getDocumentType()); + + // Verify error message is in the JSON + String expectedJson = "{\"errorMessage\":\"" + errorMessage + "\"}"; + Assertions.assertEquals(expectedJson, errorDoc.toJson()); + + // Test error document equality + JSONDocument anotherErrorDoc = JSONDocument.errorDocument(errorMessage); + Assertions.assertEquals(errorDoc, anotherErrorDoc); + } + + @Test + public void testToStringMethod() throws Exception { + Map data = Map.of("key1", "value1"); + JSONDocument document = new JSONDocument(data, DocumentType.FLAT); + + // toString should return the same as toJson + Assertions.assertEquals(document.toJson(), document.toString()); + } } diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresCollectionTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresCollectionTest.java index e94f6009..fb5a71b0 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresCollectionTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresCollectionTest.java @@ -3,6 +3,7 @@ import static java.sql.Types.INTEGER; import static java.sql.Types.VARCHAR; import static java.util.Collections.emptyList; +import static org.bson.assertions.Assertions.assertNotNull; import static org.hypertrace.core.documentstore.expression.operators.LogicalOperator.AND; import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.EQ; import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.LT; @@ -25,6 +26,7 @@ import static org.mockito.internal.verification.VerificationModeFactory.times; import java.io.IOException; +import java.sql.Array; import java.sql.Connection; import java.sql.PreparedStatement; import java.sql.ResultSet; @@ -36,6 +38,7 @@ import java.util.UUID; import org.hypertrace.core.documentstore.CloseableIterator; import org.hypertrace.core.documentstore.Document; +import org.hypertrace.core.documentstore.DocumentType; import org.hypertrace.core.documentstore.Filter; import org.hypertrace.core.documentstore.JSONDocument; import org.hypertrace.core.documentstore.Key; @@ -186,6 +189,7 @@ void testUpdateAtomicWithFilter() throws IOException, SQLException { assertTrue(oldDocument.isPresent()); assertEquals(document, oldDocument.get()); + assertEquals(DocumentType.NESTED, document.getDocumentType()); verify(mockClient, times(1)).getPooledConnection(); verify(mockConnection, times(1)).prepareStatement(selectQuery); @@ -937,6 +941,127 @@ private void mockResultSetMetadata(final String id) throws SQLException { when(mockResultSet.getString(5)).thenReturn(id); } + @Test + void testPostgresResultIteratorWithBasicTypesUsage() throws SQLException, IOException { + // Test that PostgresResultIteratorWithBasicTypes can handle various column types + when(mockResultSet.next()).thenReturn(true, false); + when(mockResultSet.getMetaData()).thenReturn(mockResultSetMetaData); + when(mockResultSetMetaData.getColumnCount()).thenReturn(6); + + // Setup different column types to test the addColumnToJsonNode method + when(mockResultSetMetaData.getColumnName(1)).thenReturn("boolean_field"); + when(mockResultSetMetaData.getColumnTypeName(1)).thenReturn("bool"); + when(mockResultSet.getBoolean(1)).thenReturn(true); + when(mockResultSet.wasNull()).thenReturn(false); + + when(mockResultSetMetaData.getColumnName(2)).thenReturn("integer_field"); + when(mockResultSetMetaData.getColumnTypeName(2)).thenReturn("int4"); + when(mockResultSet.getInt(2)).thenReturn(42); + + when(mockResultSetMetaData.getColumnName(3)).thenReturn("bigint_field"); + when(mockResultSetMetaData.getColumnTypeName(3)).thenReturn("int8"); + when(mockResultSet.getLong(3)).thenReturn(123456789L); + + when(mockResultSetMetaData.getColumnName(4)).thenReturn("double_field"); + when(mockResultSetMetaData.getColumnTypeName(4)).thenReturn("float8"); + when(mockResultSet.getDouble(4)).thenReturn(3.14159); + + when(mockResultSetMetaData.getColumnName(5)).thenReturn("text_field"); + when(mockResultSetMetaData.getColumnTypeName(5)).thenReturn("text"); + when(mockResultSet.getString(5)).thenReturn("sample text"); + + when(mockResultSetMetaData.getColumnName(6)).thenReturn("jsonb_field"); + when(mockResultSetMetaData.getColumnTypeName(6)).thenReturn("jsonb"); + when(mockResultSet.getString(6)).thenReturn("{\"nested\":\"value\"}"); + + // Create and test the iterator directly + PostgresCollection.PostgresResultIteratorWithBasicTypes iterator = + new PostgresCollection.PostgresResultIteratorWithBasicTypes( + mockResultSet, DocumentType.FLAT); + + assertTrue(iterator.hasNext()); + Document result = iterator.next(); + + assertNotNull(result); + assertEquals(DocumentType.FLAT, result.getDocumentType()); + + String json = result.toJson(); + assertTrue(json.contains("\"boolean_field\":true")); + assertTrue(json.contains("\"integer_field\":42")); + assertTrue(json.contains("\"bigint_field\":123456789")); + assertTrue(json.contains("\"double_field\":3.14159")); + assertTrue(json.contains("\"text_field\":\"sample text\"")); + assertTrue(json.contains("\"jsonb_field\":{\"nested\":\"value\"}")); + + assertFalse(iterator.hasNext()); + iterator.close(); + } + + @Test + void testPostgresResultIteratorWithBasicTypesNullHandling() throws SQLException, IOException { + // Test null value handling in PostgresResultIteratorWithBasicTypes + when(mockResultSet.next()).thenReturn(true, false); + when(mockResultSet.getMetaData()).thenReturn(mockResultSetMetaData); + when(mockResultSetMetaData.getColumnCount()).thenReturn(3); + + when(mockResultSetMetaData.getColumnName(1)).thenReturn("nullable_int"); + when(mockResultSetMetaData.getColumnTypeName(1)).thenReturn("int4"); + when(mockResultSet.getInt(1)).thenReturn(0); + when(mockResultSet.wasNull()).thenReturn(true); + + when(mockResultSetMetaData.getColumnName(2)).thenReturn("nullable_text"); + when(mockResultSetMetaData.getColumnTypeName(2)).thenReturn("text"); + when(mockResultSet.getString(2)).thenReturn(null); + + when(mockResultSetMetaData.getColumnName(3)).thenReturn("valid_text"); + when(mockResultSetMetaData.getColumnTypeName(3)).thenReturn("text"); + when(mockResultSet.getString(3)).thenReturn("not null"); + + PostgresCollection.PostgresResultIteratorWithBasicTypes iterator = + new PostgresCollection.PostgresResultIteratorWithBasicTypes(mockResultSet); + + assertTrue(iterator.hasNext()); + Document result = iterator.next(); + + assertNotNull(result); + String json = result.toJson(); + + // Null values should not appear in the JSON + assertFalse(json.contains("nullable_int")); + assertFalse(json.contains("nullable_text")); + assertTrue(json.contains("\"valid_text\":\"not null\"")); + + iterator.close(); + } + + @Test + void testPostgresResultIteratorWithBasicTypesArrayColumn() throws SQLException, IOException { + // Test array column handling + when(mockResultSet.next()).thenReturn(true, false); + when(mockResultSet.getMetaData()).thenReturn(mockResultSetMetaData); + when(mockResultSetMetaData.getColumnCount()).thenReturn(1); + + when(mockResultSetMetaData.getColumnName(1)).thenReturn("array_field"); + when(mockResultSetMetaData.getColumnTypeName(1)).thenReturn("_text"); + + Array mockArray = mock(Array.class); + String[] arrayData = {"item1", "item2", "item3"}; + when(mockResultSet.getArray(1)).thenReturn(mockArray); + when(mockArray.getArray()).thenReturn(arrayData); + + PostgresCollection.PostgresResultIteratorWithBasicTypes iterator = + new PostgresCollection.PostgresResultIteratorWithBasicTypes(mockResultSet); + + assertTrue(iterator.hasNext()); + Document result = iterator.next(); + + assertNotNull(result); + String json = result.toJson(); + assertTrue(json.contains("\"array_field\":[\"item1\",\"item2\",\"item3\"]")); + + iterator.close(); + } + private void mockResultSetMetadata() throws SQLException { when(mockResultSetMetaData.getColumnName(1)).thenReturn("quantity"); when(mockResultSetMetaData.getColumnType(1)).thenReturn(INTEGER);