diff --git a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/clients/ClientRegistry.java b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/clients/ClientRegistry.java index c779e0cc1..e7b993759 100644 --- a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/clients/ClientRegistry.java +++ b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/clients/ClientRegistry.java @@ -6,6 +6,7 @@ import org.hypertrace.core.datamodel.StructuredTrace; import org.hypertrace.entity.data.service.client.EdsCacheClient; import org.hypertrace.entity.data.service.rxclient.EntityDataClient; +import org.hypertrace.entity.query.service.v1.EntityQueryServiceGrpc.EntityQueryServiceBlockingStub; import org.hypertrace.trace.accessor.entities.TraceEntityAccessor; import org.hypertrace.trace.reader.attributes.TraceAttributeReader; import org.hypertrace.traceenricher.enrichment.enrichers.cache.EntityCache; @@ -26,6 +27,8 @@ public interface ClientRegistry { EntityDataClient getEntityDataClient(); + EntityQueryServiceBlockingStub getEntityQueryServiceClient(); + EntityCache getEntityCache(); CachingAttributeClient getCachingAttributeClient(); diff --git a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/clients/DefaultClientRegistry.java b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/clients/DefaultClientRegistry.java index 9d9507a50..483683036 100644 --- a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/clients/DefaultClientRegistry.java +++ b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/clients/DefaultClientRegistry.java @@ -9,9 +9,12 @@ import org.hypertrace.core.datamodel.Event; import org.hypertrace.core.datamodel.StructuredTrace; import org.hypertrace.core.grpcutils.client.GrpcChannelRegistry; +import org.hypertrace.core.grpcutils.client.RequestContextClientCallCredsProviderFactory; import org.hypertrace.entity.data.service.client.EdsCacheClient; import org.hypertrace.entity.data.service.client.EntityDataServiceClient; import org.hypertrace.entity.data.service.rxclient.EntityDataClient; +import org.hypertrace.entity.query.service.v1.EntityQueryServiceGrpc; +import org.hypertrace.entity.query.service.v1.EntityQueryServiceGrpc.EntityQueryServiceBlockingStub; import org.hypertrace.entity.service.client.config.EntityServiceClientConfig; import org.hypertrace.entity.type.service.rxclient.EntityTypeClient; import org.hypertrace.trace.accessor.entities.TraceEntityAccessor; @@ -113,6 +116,13 @@ public EntityDataClient getEntityDataClient() { return this.entityDataClient; } + @Override + public EntityQueryServiceBlockingStub getEntityQueryServiceClient() { + return EntityQueryServiceGrpc.newBlockingStub(entityServiceChannel) + .withCallCredentials( + RequestContextClientCallCredsProviderFactory.getClientCallCredsProvider().get()); + } + @Override public EntityCache getEntityCache() { return this.entityCache;