io.sentry.spring.jakarta.graphql.SentryBatchLoaderRegistry Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of sentry-spring-jakarta Show documentation
SDK for sentry.io
There is a newer version: 8.0.0-rc.3
Show newest version
package io.sentry.spring.jakarta.graphql;

import static io.sentry.graphql.SentryGraphqlInstrumentation.SENTRY_SCOPES_CONTEXT_KEY;

import graphql.GraphQLContext;
import io.sentry.Breadcrumb;
import io.sentry.IScopes;
import io.sentry.NoOpScopes;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import org.dataloader.BatchLoaderEnvironment;
import org.dataloader.DataLoaderOptions;
import org.dataloader.DataLoaderRegistry;
import org.jetbrains.annotations.ApiStatus;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.graphql.execution.BatchLoaderRegistry;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

@ApiStatus.Internal
public final class SentryBatchLoaderRegistry implements BatchLoaderRegistry {

  private final @NotNull BatchLoaderRegistry delegate;

  SentryBatchLoaderRegistry(final @NotNull BatchLoaderRegistry delegate) {
    this.delegate = delegate;
  }

  @Override
  public  RegistrationSpec forTypePair(Class keyType, Class valueType) {
    return new SentryRegistrationSpec(
        delegate.forTypePair(keyType, valueType), keyType, valueType);
  }

  @Override
  public  RegistrationSpec forName(String name) {
    return new SentryRegistrationSpec(delegate.forName(name), name);
  }

  @Override
  public void registerDataLoaders(DataLoaderRegistry registry, GraphQLContext context) {
    delegate.registerDataLoaders(registry, context);
  }

  public static final class SentryRegistrationSpec
      implements BatchLoaderRegistry.RegistrationSpec {

    private final @NotNull RegistrationSpec delegate;
    private final @Nullable String name;
    private final @Nullable Class keyType;
    private final @Nullable Class valueType;

    public SentryRegistrationSpec(
        final @NotNull RegistrationSpec delegate, Class keyType, Class valueType) {
      this.delegate = delegate;
      this.keyType = keyType;
      this.valueType = valueType;
      this.name = null;
    }

    public SentryRegistrationSpec(final @NotNull RegistrationSpec delegate, String name) {
      this.delegate = delegate;
      this.name = name;
      this.keyType = null;
      this.valueType = null;
    }

    @Override
    public BatchLoaderRegistry.RegistrationSpec withName(String name) {
      return delegate.withName(name);
    }

    @Override
    public BatchLoaderRegistry.RegistrationSpec withOptions(
        Consumer optionsConsumer) {
      return delegate.withOptions(optionsConsumer);
    }

    @Override
    public BatchLoaderRegistry.RegistrationSpec withOptions(DataLoaderOptions options) {
      return delegate.withOptions(options);
    }

    @Override
    public void registerBatchLoader(BiFunction, BatchLoaderEnvironment, Flux> loader) {
      delegate.registerBatchLoader(
          (keys, batchLoaderEnvironment) -> {
            scopesFromContext(batchLoaderEnvironment)
                .addBreadcrumb(Breadcrumb.graphqlDataLoader(keys, keyType, valueType, name));
            return loader.apply(keys, batchLoaderEnvironment);
          });
    }

    @Override
    public void registerMappedBatchLoader(
        BiFunction, BatchLoaderEnvironment, Mono>> loader) {
      delegate.registerMappedBatchLoader(
          (keys, batchLoaderEnvironment) -> {
            scopesFromContext(batchLoaderEnvironment)
                .addBreadcrumb(Breadcrumb.graphqlDataLoader(keys, keyType, valueType, name));
            return loader.apply(keys, batchLoaderEnvironment);
          });
    }

    private @NotNull IScopes scopesFromContext(final @NotNull BatchLoaderEnvironment environment) {
      Object context = environment.getContext();
      if (context instanceof GraphQLContext) {
        GraphQLContext graphqlContext = (GraphQLContext) context;
        return graphqlContext.getOrDefault(SENTRY_SCOPES_CONTEXT_KEY, NoOpScopes.getInstance());
      }

      return NoOpScopes.getInstance();
    }
  }
}