com.twosigma.beakerx.kernel.KernelFunctionality Java Examples

@Override
protected KernelFunctionality createKernel(String sessionId, KernelSocketsFactory kernelSocketsFactory, CloseKernelAction closeKernelAction) {
  autotranslationService = new NamespaceClientTest.AutotranslationServiceTestImpl();
  NamespaceClient nc = new NamespaceClient(autotranslationService, new DefaultBeakerXJsonSerializer(), new BeakerXCommRepositoryMock());
  MagicCommandConfigurationMock magicCommandConfiguration = new MagicCommandConfigurationMock();
  KotlinEvaluator evaluator = new KotlinEvaluator(sessionId,
          sessionId,
          cellExecutor(),
          getTestTempFolderFactory(),
          getEvaluatorParameters(),
          nc,
          magicCommandConfiguration.patterns(),
          new ClasspathScannerMock());
  return new Kotlin(sessionId,
          evaluator,
          new Configuration(
                  kernelSocketsFactory,
                  closeKernelAction,
                  getCacheFolderFactory(),
                  new CustomMagicCommandsEmptyImpl(),
                  new BeakerXCommRepositoryMock(),
                  BeakerXServerMock.create(),
                  magicCommandConfiguration,
                  new KernelTest.BeakerXJsonMock(),
                  new RuntimetoolsMock()));
}
 

private TryResult initSparkContext(Message parentMessage, Map sparkOptions) {
  KernelFunctionality kernel = KernelManager.get();
  if (singleSparkSession.isActive()) {
    return TryResult.createError(ONE_SPARK_SESSION_MSG_ERROR);
  } else {
    SparkVariable.putSparkUI(this);
    return createSparkContext(parentMessage, kernel, sparkOptions);
  }
}
 

public KernelSocketsZMQ(KernelFunctionality kernel, Config configuration, SocketCloseAction closeAction) {
  this.closeAction = closeAction;
  this.kernel = kernel;
  this.hmac = new HashedMessageAuthenticationCode(configuration.getKey());
  this.context = ZMQ.context(1);
  configureSockets(configuration);
}
 

public SparkFactoryImpl(KernelFunctionality kernel,
                        SparkEngineNoUIFactory sparkEngineNoUIFactory,
                        SparkEngineWithUIFactory sparkEngineWithUIFactory,
                        SparkUIFactory sparkUIFactory,
                        SparkUiDefaults sparkUiDefaults,
                        SparkSessionBuilderFactory sparkSessionBuilderFactory,
                        SparkListenerService sparkListenerService) {
  this.kernel = kernel;
  this.sparkEngineNoUIFactory = sparkEngineNoUIFactory;
  this.sparkEngineWithUIFactory = sparkEngineWithUIFactory;
  this.sparkUIFactory = sparkUIFactory;
  this.sparkUiDefaults = sparkUiDefaults;
  this.sparkSessionBuilderFactory = sparkSessionBuilderFactory;
  this.sparkListenerService = sparkListenerService;
}
 

private TryResult createSparkContext(Message parentMessage, KernelFunctionality kernel, Map sparkOptions) {
  try {
    TryResult result = sparkEngine.createSparkContext(kernel, this, parentMessage, sparkOptions);
    if (result.isError()) {
      return TryResult.createError(result.error());
    } else {
      singleSparkSession.activate();
      applicationStart();
      return TryResult.createResult("done");
    }
  } catch (Exception e) {
    return TryResult.createError(e.toString());
  }
}
 

@Override
public MagicCommandOutcomeItem addSparkexJar(KernelFunctionality kernel) {
  Optional<MagicCommandFunctionality> magic = CodeFactory.findMagicCommandFunctionality(kernel.getMagicCommandTypes(), ClasspathAddJarMagicCommand.CLASSPATH_ADD_JAR);
  String sparkexJar = getSparkexJar();
  MagicCommandOutcomeItem magicCommandOutcomeItem = ((ClasspathAddJarMagicCommand) magic.get()).addJar(sparkexJar);
  return magicCommandOutcomeItem;
}
 

protected TryResult initSparkContextInShell(KernelFunctionality kernel, Message parent) {
  String addSc = String.format(("import com.twosigma.beakerx.widget.SparkVariable\n" +
                  "val %s = SparkVariable.getSparkSession()\n" +
                  "val %s = %s.sparkContext\n" +
                  "import org.apache.spark.SparkContext._\n" +
                  "import %s.implicits._\n" +
                  "import %s.sql\n" +
                  "import org.apache.spark.sql.functions._\n"),
          SPARK_SESSION_NAME, SPARK_CONTEXT_NAME, SPARK_SESSION_NAME, SPARK_SESSION_NAME, SPARK_SESSION_NAME);

  SimpleEvaluationObject seo = createSimpleEvaluationObject(addSc, kernel, new Message(new Header(JupyterMessages.COMM_MSG, parent.getHeader().getSession())), 1);
  return kernel.executeCode(addSc, seo);
}
 

@Override
public void sendRepliesWithStatus(KernelFunctionality kernel, Message message, int executionCount) {
  if (getStatus().equals(MagicCommandOutcomeItem.Status.OK)) {
    publishMessage(kernel, message, executionCount);
    kernel.send(MessageCreator.buildReplyWithOkStatus(message, executionCount));
  } else {
    publishMessage(kernel, message, executionCount);
    kernel.send(MessageCreator.buildReplyWithErrorStatus(message, executionCount));
  }
}
 

public SparkUI(Comm comm, SparkEngineWithUI sparkEngine, SparkUiDefaults sparkUiDefaults, SingleSparkSession singleSparkSession, KernelFunctionality kernel) {
  super(comm, new ArrayList<>());
  this.sparkUiDefaults = sparkUiDefaults;
  this.sparkEngine = sparkEngine;
  this.singleSparkSession = singleSparkSession;
  this.kernel = kernel;
  getComm().addMsgCallbackList((Handler<Message>) this::handleMessage);
  SparkVariable.putSparkUI(this);
  this.configureRESTMapping();
  openComm();
}
 

@Override
protected KernelFunctionality createKernel(String sessionId, KernelSocketsFactory kernelSocketsFactory, CloseKernelAction closeKernelAction) {
  return new Groovy(sessionId,
          groovyEvaluator(),
          new Configuration(
                  kernelSocketsFactory,
                  closeKernelAction,
                  getCacheFolderFactory(),
                  new CustomMagicCommandsEmptyImpl(),
                  new BeakerXCommRepositoryMock(),
                  BeakerXServerMock.create(),
                  new MagicCommandConfigurationMock(),
                  new KernelTest.BeakerXJsonMock(),
                  new RuntimetoolsMock()));
}
 

@Override
public TryResult configure(KernelFunctionality kernel, Message parentMessage) {
  final SparkConf sparkConf = newSparkConf();
  this.sparkSessionBuilder = this.sparkSessionBuilderFactory.newInstance(sparkConf);
  TryResult sparkSessionTry = createSparkSession();
  if (sparkSessionTry.isError()) {
    return sparkSessionTry;
  }
  SparkVariable.putSparkSession(this.sparkSessionBuilder.getOrCreate());
  TryResult tryResultSparkContext = initSparkContextInShell(kernel, parentMessage);
  if (!tryResultSparkContext.isError()) {
    kernel.registerCancelHook(SparkVariable::cancelAllJobs);
  }
  return tryResultSparkContext;
}
 

@Override
public void sendRepliesWithStatus(KernelFunctionality kernel, Message message, int executionCount) {
  if (getStatus().equals(MagicCommandOutcomeItem.Status.OK)) {
    sendHTML(message);
    kernel.send(MessageCreator.buildReplyWithOkStatus(message, executionCount));
  } else {
    sendHTML(message);
    kernel.send(MessageCreator.buildReplyWithErrorStatus(message, executionCount));
  }
}
 

private MagicCommandType enableSparkSupportMagicCommand(KernelFunctionality kernel) {
  SparkexJarServiceImpl sparkexJarService = new SparkexJarServiceImpl();
  EnableSparkSupportActionOptions supportActionOptions = new EnableSparkSupportActionOptionsImpl(kernel);
  SparkInitCommandFactoryImpl sparkInitCommandFactory = new SparkInitCommandFactoryImpl(
          kernel,
          sparkexJarService,
          new EnableSparkSupportOptions(supportActionOptions),
          supportActionOptions);
  return new MagicCommandType(
          EnableSparkSupportMagicCommand.ENABLE_SPARK_SUPPORT,
          "<>",
          new EnableSparkSupportMagicCommand(kernel, sparkInitCommandFactory));
}
 

@Override
public void sendMagicCommandOutcome(KernelFunctionality kernel, Message message, int executionCount) {
  if (getMIMEContainer().isPresent()) {
    boolean hasError = getStatus().equals(MagicCommandOutcomeItem.Status.ERROR);
    kernel.publish(Collections.singletonList(MessageCreator.buildOutputMessage(message, (String) getMIMEContainer().get().getData(), hasError)));
  }
}
 

public EnableSparkSupportMagicSparkConfiguration(KernelFunctionality kernel) {
  this.kernel = kernel;
}
 

private void publishMessage(KernelFunctionality kernel, Message message, int executionCount) {
    for (Message msg : responses) {
        msg.setParentHeader(message.getHeader());
    }
    kernel.publish(responses);
}
 

private MagicCommandType timeItLine(KernelFunctionality kernel) {
  return new MagicCommandType(TimeItLineModeMagicCommand.TIMEIT_LINE, "", new TimeItLineModeMagicCommand(kernel));
}
 

private static MagicCommandType addClasspathReset(KernelFunctionality kernel, FileService fileService) {
  return new MagicCommandType(ClasspathResetMagicCommand.CLASSPATH_RESET, "", new ClasspathResetMagicCommand(kernel, fileService));
}
 

@Override
public void createMapping(Javalin app, KernelFunctionality kernel) {
}
 

private MagicCommandType timeItCell(KernelFunctionality kernel) {
  return new MagicCommandType(TimeItCellModeMagicCommand.TIMEIT_CELL, "", new TimeItCellModeMagicCommand(kernel));
}
 

public ClasspathAddJarMagicCommand(KernelFunctionality kernel) {
  super(kernel);
}
 

public SymjaMMAKernelInfoHandler(KernelFunctionality kernel) {
  super(kernel);
}
 

public ScalaKernelInfoHandler(KernelFunctionality kernel) {
  super(kernel);
}
 

public TimeCellModeMagicCommand(KernelFunctionality kernel) {
  super(kernel);
}
 

public ClasspathAddRepoMagicCommand(KernelFunctionality kernel) {
  this.kernel = kernel;
}
 

private MagicCommandType addImport(KernelFunctionality kernel) {
  return new MagicCommandType(AddImportMagicCommand.IMPORT, "<classpath>", new AddImportMagicCommand(kernel));
}
 

public MavenJarResolver.ResolverParams mavenResolverParam(KernelFunctionality kernel) {
  return new MavenJarResolver.ResolverParams(
          new File(TEST_MVN_CACHE).getAbsolutePath(),
          kernel.getTempFolder().toString() + MavenJarResolver.MVN_DIR,
          true);
}
 

public ClasspathMagicCommand(KernelFunctionality kernel) {
  this.kernel = kernel;
}
 

public TimeItLineModeMagicCommand(KernelFunctionality kernel) {
  super(kernel);
}
 

public UnImportMagicCommand(KernelFunctionality kernel) {
  this.kernel = kernel;
}