Skip to content

Commit

Permalink
Pass all GremlinConnectionProperties to SchemaHelperGremlinDataModel (#…
Browse files Browse the repository at this point in the history 
…250)

* pass optional connection propertie to sql SchemaHelper (#116)

* Update src/main/java/software/aws/neptune/common/gremlindatamodel/MetadataCache.java

apply reviewer comments

Co-authored-by: Lyndon Bauto <[email protected]>

* Checkstyle fixes

Import ordering, javadoc descriptions, indentation

* Add MetadataCache tests

Co-Authored-By: Lyndon Bauto <[email protected]>

---------

Co-authored-by: Artem Aliev <[email protected]>
Co-authored-by: Artem Aliev <[email protected]>
Co-authored-by: Lyndon Bauto <[email protected]>
Co-authored-by: Lyndon Bauto <[email protected]>
  • Loading branch information
@ryn5 @artem-aliev @lyndonbauto
5 people authored Feb 7, 2024
1 parent caa3be0 commit f65eca1
Show file tree
Hide file tree
Showing 4 changed files with 130 additions and 13 deletions.
20 changes: 16 additions & 4 deletions src/main/java/software/aws/neptune/common/gremlindatamodel/MetadataCache.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,21 @@ public static void updateCache(final String endpoint, final int port, final bool
}
}

/**
* Function to update the cache of the metadata.
* @param gremlinConnectionProperties GremlinConnectionProperties to use.
* @throws SQLException
*/
public static void updateGremlinMetadataCache(final GremlinConnectionProperties gremlinConnectionProperties)
throws SQLException {
final String endpoint = gremlinConnectionProperties.getContactPoint();
synchronized (LOCK) {
if (!GREMLIN_SCHEMAS.containsKey(endpoint)) {
GREMLIN_SCHEMAS.put(endpoint, SchemaHelperGremlinDataModel.getGremlinGraphSchema(gremlinConnectionProperties));
}
}
}

/**
* Function to update the cache of the metadata.
*
Expand All @@ -77,10 +92,7 @@ public static void updateCache(final String endpoint, final int port, final bool
public static void updateCacheIfNotUpdated(final GremlinConnectionProperties gremlinConnectionProperties)
throws SQLException {
if (!isMetadataCached(gremlinConnectionProperties.getContactPoint())) {
updateCache(gremlinConnectionProperties.getContactPoint(), gremlinConnectionProperties.getPort(),
(gremlinConnectionProperties.getAuthScheme() == AuthScheme.IAMSigV4),
gremlinConnectionProperties.getEnableSsl(),
MetadataCache.PathType.Gremlin, gremlinConnectionProperties.getScanType());
updateGremlinMetadataCache(gremlinConnectionProperties);
}
}

Expand Down
16 changes: 16 additions & 0 deletions src/main/java/software/aws/neptune/common/gremlindatamodel/SchemaHelperGremlinDataModel.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.aws.neptune.common.IAMHelper;
import software.aws.neptune.gremlin.GremlinConnectionProperties;
import software.aws.neptune.gremlin.GremlinQueryExecutor;
import software.aws.neptune.gremlin.adapter.converter.schema.SqlSchemaGrabber;
import software.aws.neptune.gremlin.adapter.converter.schema.calcite.GremlinSchema;
import software.aws.neptune.jdbc.utilities.SqlError;
Expand Down Expand Up @@ -89,4 +91,18 @@ public static GremlinSchema getGraphSchema(final String endpoint, final int port
traversal().withRemote(DriverRemoteConnection.using(getClient(adjustedEndpoint, port, useIAM, useSsl))),
scanType);
}

/**
* Function to get the schema of the graph through gremlin connection
*
* @param gremlinConnectionProperties Connection parameters.
* @return Graph Schema.
* @throws SQLException If graph schema cannot be obtained.s
*/
public static GremlinSchema getGremlinGraphSchema(final GremlinConnectionProperties gremlinConnectionProperties)
throws SQLException {
return SqlSchemaGrabber.getSchema(
traversal().withRemote(DriverRemoteConnection.using(GremlinQueryExecutor.getClient(gremlinConnectionProperties))),
gremlinConnectionProperties.getScanType());
}
}
4 changes: 2 additions & 2 deletions src/main/java/software/aws/neptune/gremlin/GremlinQueryExecutor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ public static Cluster.Builder createClusterBuilder(final GremlinConnectionProper
builder.serializer(properties.getSerializerObject());
} else if (properties.isSerializerEnum()) {
builder.serializer(properties.getSerializerEnum());
} else if (properties.isChannelizerString()) {
} else if (properties.isSerializerString()) {
builder.serializer(properties.getSerializerString());
}
}
Expand Down Expand Up @@ -215,7 +215,7 @@ public static void close() {
}
}

protected static Client getClient(final GremlinConnectionProperties gremlinConnectionProperties)
public static Client getClient(final GremlinConnectionProperties gremlinConnectionProperties)
throws SQLException {
synchronized (CLUSTER_LOCK) {
cluster = getCluster(gremlinConnectionProperties);
Expand Down
103 changes: 96 additions & 7 deletions src/test/java/software/aws/neptune/common/gremlindatamodel/MetadataCacheTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,16 +16,26 @@

package software.aws.neptune.common.gremlindatamodel;

import org.apache.calcite.avatica.Meta;
import org.apache.calcite.util.Pair;
import org.apache.tinkerpop.gremlin.driver.Cluster;
import org.apache.tinkerpop.gremlin.driver.ser.GraphSONMessageSerializerV3d0;
import org.apache.tinkerpop.gremlin.tinkergraph.structure.TinkerGraph;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
import software.aws.neptune.gremlin.GremlinConnectionProperties;
import software.aws.neptune.gremlin.GremlinQueryExecutor;
import software.aws.neptune.gremlin.adapter.converter.schema.SqlSchemaGrabber;
import software.aws.neptune.gremlin.adapter.converter.schema.calcite.GremlinSchema;
import software.aws.neptune.gremlin.adapter.converter.schema.gremlin.GremlinEdgeTable;
import software.aws.neptune.gremlin.adapter.converter.schema.gremlin.GremlinProperty;
import software.aws.neptune.gremlin.adapter.converter.schema.gremlin.GremlinVertexTable;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.sql.SQLException;
import java.util.ArrayList;
Expand All @@ -36,29 +46,42 @@

public class MetadataCacheTest {
private static final String ENDPOINT = "mockEndpoint";
private static final String SERIALIZER = "GRAPHSON_V3D0";
private final GremlinVertexTable testTableVertex;
private final GremlinVertexTable testTableVertexBeta;
private final GremlinEdgeTable testTableEdge;
private final GremlinEdgeTable testTableEdgeBeta;
private final GremlinSchema testFullSchema;

@Test
void testMetadataCacheFiltering() throws SQLException {
MetadataCacheTest() {
// Set up four tables, two vertex table with label "vertex" and "vertexBeta", and two edge tables with labels "edge", "edgeBeta".
final GremlinVertexTable testTableVertex = new GremlinVertexTable("vertex",
testTableVertex = new GremlinVertexTable("vertex",
new ArrayList<>(Collections.singletonList(new GremlinProperty("testVertex", "string"))),
new ArrayList<>(Collections.singletonList("edgeIn")),
new ArrayList<>(Collections.singletonList("edgeOut")));
final GremlinVertexTable testTableVertexBeta = new GremlinVertexTable("vertexBeta",
testTableVertexBeta = new GremlinVertexTable("vertexBeta",
new ArrayList<>(Collections.singletonList(new GremlinProperty("testVertex", "string"))),
new ArrayList<>(Collections.singletonList("edgeIn")),
new ArrayList<>(Collections.singletonList("edgeOut")));
final GremlinEdgeTable testTableEdge = new GremlinEdgeTable("edge",
testTableEdge = new GremlinEdgeTable("edge",
new ArrayList<>(Collections.singletonList(new GremlinProperty("testEdge", "string"))),
new ArrayList<>(Collections.singletonList(new Pair<>("vertexIn", "vertexOut"))));
final GremlinEdgeTable testTableEdgeBeta = new GremlinEdgeTable("edgeBeta",
testTableEdgeBeta = new GremlinEdgeTable("edgeBeta",
new ArrayList<>(Collections.singletonList(new GremlinProperty("testEdge", "string"))),
new ArrayList<>(Collections.singletonList(new Pair<>("vertexIn", "vertexOut"))));

// Create a full schema to call the method on.
final GremlinSchema testFullSchema = new GremlinSchema(new ArrayList<>(Arrays.asList(testTableVertex, testTableVertexBeta)),
testFullSchema = new GremlinSchema(new ArrayList<>(Arrays.asList(testTableVertex, testTableVertexBeta)),
new ArrayList<>(Arrays.asList(testTableEdge, testTableEdgeBeta)));
}

@BeforeEach
void refreshCache() {
MetadataCache.getGremlinSchemas().clear();
}

@Test
void testMockMetadataCache() throws SQLException {
try (MockedStatic<MetadataCache> mockMetadataCache = Mockito.mockStatic(MetadataCache.class, invocation -> {
final Method method = invocation.getMethod();
if ("getGremlinSchemas".equals(method.getName())) {
Expand All @@ -71,6 +94,22 @@ void testMetadataCacheFiltering() throws SQLException {
schemaMap.put(ENDPOINT, testFullSchema);
mockMetadataCache.when(MetadataCache::getGremlinSchemas).thenReturn(schemaMap);
Assertions.assertEquals(schemaMap, MetadataCache.getGremlinSchemas());
}
}

@Test
void testMetadataCacheFiltering() throws SQLException {
try (MockedStatic<MetadataCache> mockMetadataCache = Mockito.mockStatic(MetadataCache.class, invocation -> {
final Method method = invocation.getMethod();
if ("getGremlinSchemas".equals(method.getName())) {
return invocation.getMock();
} else {
return invocation.callRealMethod();
}
})) {
final Map<String, GremlinSchema> schemaMap = new HashMap<>();
schemaMap.put(ENDPOINT, testFullSchema);
mockMetadataCache.when(MetadataCache::getGremlinSchemas).thenReturn(schemaMap);

// Assert that filtering based label only gets the specified table.
final GremlinSchema generatedVertexSchema = MetadataCache.getFilteredCacheNodeColumnInfos("vertex", ENDPOINT);
Expand Down Expand Up @@ -98,4 +137,54 @@ void testMetadataCacheFiltering() throws SQLException {
Assertions.assertEquals(1, generatedEdgeBetaSchema.getEdges().size());
}
}

@Test
void testUpdateGremlinCacheConnectionless() throws SQLException {
GremlinConnectionProperties properties = new GremlinConnectionProperties();
properties.setProperty(GremlinConnectionProperties.CONTACT_POINT_KEY, ENDPOINT);

Assertions.assertFalse(MetadataCache.isMetadataCached(ENDPOINT));
if (!MetadataCache.getGremlinSchemas().containsKey(ENDPOINT)) {
MetadataCache.getGremlinSchemas().put(ENDPOINT, SqlSchemaGrabber.getSchema(TinkerGraph.open().traversal(), properties.getScanType()));
}
Assertions.assertTrue(MetadataCache.isMetadataCached(ENDPOINT));
}

@Test
void testUpdateGremlinCacheAlreadyCached() throws SQLException {
GremlinConnectionProperties properties = new GremlinConnectionProperties();
properties.setProperty(GremlinConnectionProperties.CONTACT_POINT_KEY, ENDPOINT);

MetadataCache.getGremlinSchemas().put(ENDPOINT, testFullSchema);
Assertions.assertTrue(MetadataCache.isMetadataCached(ENDPOINT));

MetadataCache.updateGremlinMetadataCache(properties);
Assertions.assertEquals(testFullSchema, MetadataCache.getGremlinSchema(ENDPOINT));
}

@Test
void testUpdateCacheIfNotUpdatedAlreadyCached() throws SQLException {
GremlinConnectionProperties properties = new GremlinConnectionProperties();
properties.setProperty(GremlinConnectionProperties.CONTACT_POINT_KEY, ENDPOINT);

MetadataCache.getGremlinSchemas().put(ENDPOINT, testFullSchema);
Assertions.assertTrue(MetadataCache.isMetadataCached(ENDPOINT));

MetadataCache.updateCacheIfNotUpdated(properties);
Assertions.assertEquals(testFullSchema, MetadataCache.getGremlinSchema(ENDPOINT));
}

@Test
void testReceivingSerializerProperty() throws SQLException, NoSuchMethodException, InvocationTargetException, IllegalAccessException {
GremlinConnectionProperties properties = new GremlinConnectionProperties();
properties.setProperty(GremlinConnectionProperties.SERIALIZER_KEY, SERIALIZER);

Method getClusterMethod = GremlinQueryExecutor.class.getDeclaredMethod("getCluster", GremlinConnectionProperties.class);
getClusterMethod.setAccessible(true);
Cluster cluster = (Cluster) getClusterMethod.invoke(null, properties);

Method getSerializerMethod = Cluster.class.getDeclaredMethod("getSerializer");
getSerializerMethod.setAccessible(true);
Assertions.assertEquals(GraphSONMessageSerializerV3d0.class, getSerializerMethod.invoke(cluster).getClass());
}
}

0 comments on commit f65eca1

Please sign in to comment.