Examples with MutableFunctionRegistry io.confluent.ksql.function.MutableFunctionRegistry used on opensource projects

Example 1 with MutableFunctionRegistry

use of io.confluent.ksql.function.MutableFunctionRegistry in project ksql by confluentinc.

the class FlatMapNodeTest method setUp.

@Before
public void setUp() {
    when(source.getSchema()).thenReturn(SOURCE_SCHEMA);
    when(source.getNodeOutputType()).thenReturn(DataSourceType.KSTREAM);
    when(analysis.getTableFunctions()).thenReturn(ImmutableList.of(A_TABLE_FUNCTION));
    when(analysis.getSelectItems()).thenReturn(ImmutableList.of(new AllColumns(Optional.empty()), new SingleColumn(new UnqualifiedColumnReferenceExp(COL0), Optional.empty()), new SingleColumn(A_TABLE_FUNCTION, Optional.empty())));
    final MutableFunctionRegistry functionRegistry = new InternalFunctionRegistry();
    new UdtfLoader(functionRegistry, Optional.empty(), SqlTypeParser.create(TypeRegistry.EMPTY), true).loadUdtfFromClass(Explode.class, "load path");
    flatMapNode = new FlatMapNode(PLAN_ID, source, functionRegistry, analysis);
}

Example 2 with MutableFunctionRegistry

use of io.confluent.ksql.function.MutableFunctionRegistry in project ksql by confluentinc.

the class KsqlContext method create.

/**
 * Create a KSQL context object with the given properties. A KSQL context has it's own metastore
 * valid during the life of the object.
 */
public static KsqlContext create(final KsqlConfig ksqlConfig, final ProcessingLogContext processingLogContext, final MetricCollectors metricCollectors) {
    Objects.requireNonNull(ksqlConfig, "ksqlConfig cannot be null.");
    final ServiceContext serviceContext = ServiceContextFactory.create(ksqlConfig, DisabledKsqlClient::instance);
    final MutableFunctionRegistry functionRegistry = new InternalFunctionRegistry();
    UserFunctionLoader.newInstance(ksqlConfig, functionRegistry, ".", metricCollectors.getMetrics()).load();
    final ServiceInfo serviceInfo = ServiceInfo.create(ksqlConfig);
    final KsqlEngine engine = new KsqlEngine(serviceContext, processingLogContext, functionRegistry, serviceInfo, new SequentialQueryIdGenerator(), ksqlConfig, Collections.emptyList(), metricCollectors);
    return new KsqlContext(serviceContext, ksqlConfig, engine, Injectors.DEFAULT);
}

Example 3 with MutableFunctionRegistry

use of io.confluent.ksql.function.MutableFunctionRegistry in project ksql by confluentinc.

the class TestExecutor method getKsqlEngine.

static KsqlEngine getKsqlEngine(final ServiceContext serviceContext, final Optional<String> extensionDir) {
    final FunctionRegistry functionRegistry;
    if (extensionDir.isPresent()) {
        final MutableFunctionRegistry mutable = new InternalFunctionRegistry();
        UdfLoaderUtil.load(mutable, extensionDir.get());
        functionRegistry = mutable;
    } else {
        functionRegistry = TestFunctionRegistry.INSTANCE.get();
    }
    final MutableMetaStore metaStore = new MetaStoreImpl(functionRegistry);
    final MetricCollectors metricCollectors = new MetricCollectors();
    return new KsqlEngine(serviceContext, ProcessingLogContext.create(), "test_instance_", metaStore, (engine) -> new KsqlEngineMetrics("", engine, Collections.emptyMap(), Optional.empty(), metricCollectors), new SequentialQueryIdGenerator(), KsqlConfig.empty(), Collections.emptyList(), metricCollectors);
}

Example 4 with MutableFunctionRegistry

use of io.confluent.ksql.function.MutableFunctionRegistry in project ksql by confluentinc.

the class KsqlResourceTest method setUp.

@Before
public void setUp() throws IOException, RestClientException {
    commandStatus = new QueuedCommandStatus(0, new CommandStatusFuture(new CommandId(TOPIC, "whateva", CREATE)));
    commandStatus1 = new QueuedCommandStatus(1, new CommandStatusFuture(new CommandId(TABLE, "something", DROP)));
    final QueuedCommandStatus commandStatus2 = new QueuedCommandStatus(2, new CommandStatusFuture(new CommandId(STREAM, "something", EXECUTE)));
    kafkaTopicClient = new FakeKafkaTopicClient();
    kafkaConsumerGroupClient = new FakeKafkaConsumerGroupClient();
    serviceContext = TestServiceContext.create(kafkaTopicClient, kafkaConsumerGroupClient);
    schemaRegistryClient = serviceContext.getSchemaRegistryClient();
    registerValueSchema(schemaRegistryClient);
    ksqlRestConfig = new KsqlRestConfig(getDefaultKsqlConfig());
    ksqlConfig = new KsqlConfig(ksqlRestConfig.getKsqlConfigProperties());
    final KsqlExecutionContext.ExecuteResult result = mock(KsqlExecutionContext.ExecuteResult.class);
    when(sandbox.execute(any(), any(ConfiguredKsqlPlan.class))).thenReturn(result);
    when(result.getQuery()).thenReturn(Optional.empty());
    MutableFunctionRegistry fnRegistry = new InternalFunctionRegistry();
    final Metrics metrics = new Metrics();
    UserFunctionLoader.newInstance(ksqlConfig, fnRegistry, ".", metrics).load();
    metaStore = new MetaStoreImpl(fnRegistry);
    final MetricCollectors metricCollectors = new MetricCollectors(metrics);
    realEngine = KsqlEngineTestUtil.createKsqlEngine(serviceContext, metaStore, (engine) -> new KsqlEngineMetrics("", engine, Collections.emptyMap(), Optional.empty(), metricCollectors), new SequentialQueryIdGenerator(), ksqlConfig, metricCollectors);
    securityContext = new KsqlSecurityContext(Optional.empty(), serviceContext);
    when(commandRunner.getCommandQueue()).thenReturn(commandStore);
    when(commandRunnerWarning.get()).thenReturn("");
    when(commandStore.createTransactionalProducer()).thenReturn(transactionalProducer);
    ksqlEngine = realEngine;
    when(sandbox.getMetaStore()).thenAnswer(inv -> metaStore.copy());
    addTestTopicAndSources();
    when(commandStore.enqueueCommand(any(), any(), any(Producer.class))).thenReturn(commandStatus).thenReturn(commandStatus1).thenReturn(commandStatus2);
    streamName = KsqlIdentifierTestUtil.uniqueIdentifierName();
    when(schemaInjectorFactory.apply(any())).thenReturn(sandboxSchemaInjector);
    when(schemaInjectorFactory.apply(serviceContext)).thenReturn(schemaInjector);
    when(topicInjectorFactory.apply(any())).thenReturn(sandboxTopicInjector);
    when(topicInjectorFactory.apply(ksqlEngine)).thenReturn(topicInjector);
    when(sandboxSchemaInjector.inject(any())).thenAnswer(inv -> inv.getArgument(0));
    when(schemaInjector.inject(any())).thenAnswer(inv -> inv.getArgument(0));
    when(sandboxTopicInjector.inject(any())).thenAnswer(inv -> inv.getArgument(0));
    when(topicInjector.inject(any())).thenAnswer(inv -> inv.getArgument(0));
    when(errorsHandler.generateResponse(any(), any())).thenAnswer(new Answer<EndpointResponse>() {

        @Override
        public EndpointResponse answer(final InvocationOnMock invocation) throws Throwable {
            final Object[] args = invocation.getArguments();
            return (EndpointResponse) args[1];
        }
    });
    setUpKsqlResource();
}

Example 5 with MutableFunctionRegistry

use of io.confluent.ksql.function.MutableFunctionRegistry in project ksql by confluentinc.

the class KsqlRestApplication method buildApplication.

@SuppressWarnings({ "checkstyle:JavaNCSS", "checkstyle:MethodLength", "checkstyle:ParameterNumber" })
static KsqlRestApplication buildApplication(final String metricsPrefix, final KsqlRestConfig restConfig, final Function<Supplier<Boolean>, VersionCheckerAgent> versionCheckerFactory, final int maxStatementRetries, final ServiceContext serviceContext, final Supplier<SchemaRegistryClient> schemaRegistryClientFactory, final ConnectClientFactory connectClientFactory, final Vertx vertx, final KsqlClient sharedClient, final DefaultServiceContextFactory defaultServiceContextFactory, final UserServiceContextFactory userServiceContextFactory, final MetricCollectors metricCollectors) {
    final String ksqlInstallDir = restConfig.getString(KsqlRestConfig.INSTALL_DIR_CONFIG);
    final KsqlConfig ksqlConfig = new KsqlConfig(restConfig.getKsqlConfigProperties());
    final ProcessingLogConfig processingLogConfig = new ProcessingLogConfig(restConfig.getOriginals());
    final ProcessingLogContext processingLogContext = ProcessingLogContext.create(processingLogConfig);
    final MutableFunctionRegistry functionRegistry = new InternalFunctionRegistry();
    if (restConfig.getBoolean(KsqlRestConfig.KSQL_SERVER_ENABLE_UNCAUGHT_EXCEPTION_HANDLER)) {
        Thread.setDefaultUncaughtExceptionHandler(new KsqlUncaughtExceptionHandler(LogManager::shutdown));
    }
    final SpecificQueryIdGenerator specificQueryIdGenerator = new SpecificQueryIdGenerator();
    final String stateDir = ksqlConfig.getKsqlStreamConfigProps().getOrDefault(StreamsConfig.STATE_DIR_CONFIG, StreamsConfig.configDef().defaultValues().get(StreamsConfig.STATE_DIR_CONFIG)).toString();
    final ServiceInfo serviceInfo = ServiceInfo.create(ksqlConfig, metricsPrefix);
    final Map<String, String> metricsTags = ImmutableMap.<String, String>builder().putAll(serviceInfo.customMetricsTags()).put(KsqlConstants.KSQL_SERVICE_ID_METRICS_TAG, serviceInfo.serviceId()).build();
    StorageUtilizationMetricsReporter.configureShared(new File(stateDir), metricCollectors.getMetrics(), ksqlConfig.getStringAsMap(KsqlConfig.KSQL_CUSTOM_METRICS_TAGS));
    final ScheduledExecutorService executorService = Executors.newScheduledThreadPool(1, new ThreadFactoryBuilder().setNameFormat("ksql-csu-metrics-reporter-%d").build());
    final ScheduledExecutorService leakedResourcesReporter = Executors.newScheduledThreadPool(1, new ThreadFactoryBuilder().setNameFormat("ksql-leaked-resources-metrics-reporter-%d").build());
    final KsqlEngine ksqlEngine = new KsqlEngine(serviceContext, processingLogContext, functionRegistry, serviceInfo, specificQueryIdGenerator, new KsqlConfig(restConfig.getKsqlConfigProperties()), Collections.emptyList(), metricCollectors);
    final PersistentQuerySaturationMetrics saturation = new PersistentQuerySaturationMetrics(ksqlEngine, new JmxDataPointsReporter(metricCollectors.getMetrics(), "ksqldb_utilization", Duration.ofMinutes(1)), Duration.ofMinutes(5), Duration.ofSeconds(30), ksqlConfig.getStringAsMap(KsqlConfig.KSQL_CUSTOM_METRICS_TAGS));
    executorService.scheduleAtFixedRate(saturation, 0, Duration.ofMinutes(1).toMillis(), TimeUnit.MILLISECONDS);
    final int transientQueryCleanupServicePeriod = ksqlConfig.getInt(KsqlConfig.KSQL_TRANSIENT_QUERY_CLEANUP_SERVICE_PERIOD_SECONDS);
    final LeakedResourcesMetrics leaked = new LeakedResourcesMetrics(ksqlEngine, new JmxDataPointsReporter(metricCollectors.getMetrics(), ReservedInternalTopics.KSQL_INTERNAL_TOPIC_PREFIX + ksqlConfig.getString(KsqlConfig.KSQL_SERVICE_ID_CONFIG) + ".leaked_resources_metrics", Duration.ofSeconds(transientQueryCleanupServicePeriod)), ksqlConfig.getStringAsMap(KsqlConfig.KSQL_CUSTOM_METRICS_TAGS));
    leakedResourcesReporter.scheduleAtFixedRate(leaked, 0, transientQueryCleanupServicePeriod, TimeUnit.SECONDS);
    UserFunctionLoader.newInstance(ksqlConfig, functionRegistry, ksqlInstallDir, metricCollectors.getMetrics()).load();
    final String commandTopicName = ReservedInternalTopics.commandTopic(ksqlConfig);
    final Admin internalAdmin = createCommandTopicAdminClient(restConfig, ksqlConfig);
    final KafkaTopicClient internalTopicClient = new KafkaTopicClientImpl(() -> internalAdmin);
    final CommandStore commandStore = CommandStore.Factory.create(ksqlConfig, commandTopicName, Duration.ofMillis(restConfig.getLong(DISTRIBUTED_COMMAND_RESPONSE_TIMEOUT_MS_CONFIG)), ksqlConfig.addConfluentMetricsContextConfigsKafka(restConfig.getCommandConsumerProperties()), ksqlConfig.addConfluentMetricsContextConfigsKafka(restConfig.getCommandProducerProperties()), internalTopicClient);
    final InteractiveStatementExecutor statementExecutor = new InteractiveStatementExecutor(serviceContext, ksqlEngine, specificQueryIdGenerator);
    final StatusResource statusResource = new StatusResource(statementExecutor);
    final VersionCheckerAgent versionChecker = versionCheckerFactory.apply(ksqlEngine::hasActiveQueries);
    final ServerState serverState = new ServerState();
    final KsqlSecurityExtension securityExtension = loadSecurityExtension(ksqlConfig);
    final KsqlSecurityContextProvider ksqlSecurityContextProvider = new DefaultKsqlSecurityContextProvider(securityExtension, defaultServiceContextFactory, userServiceContextFactory, ksqlConfig, schemaRegistryClientFactory, connectClientFactory, sharedClient);
    final Optional<AuthenticationPlugin> securityHandlerPlugin = loadAuthenticationPlugin(restConfig);
    final Optional<KsqlAuthorizationValidator> authorizationValidator = KsqlAuthorizationValidatorFactory.create(ksqlConfig, serviceContext, securityExtension.getAuthorizationProvider());
    final Errors errorHandler = new Errors(restConfig.getConfiguredInstance(KsqlRestConfig.KSQL_SERVER_ERROR_MESSAGES, ErrorMessages.class));
    final ConnectServerErrors connectErrorHandler = loadConnectErrorHandler(ksqlConfig);
    final Optional<LagReportingAgent> lagReportingAgent = initializeLagReportingAgent(restConfig, ksqlEngine, serviceContext);
    final Optional<HeartbeatAgent> heartbeatAgent = initializeHeartbeatAgent(restConfig, ksqlEngine, serviceContext, lagReportingAgent);
    final RoutingFilterFactory routingFilterFactory = initializeRoutingFilterFactory(ksqlConfig, heartbeatAgent, lagReportingAgent);
    final RateLimiter pullQueryRateLimiter = new RateLimiter(ksqlConfig.getInt(KsqlConfig.KSQL_QUERY_PULL_MAX_QPS_CONFIG), "pull", metricCollectors.getMetrics(), metricsTags);
    final ConcurrencyLimiter pullQueryConcurrencyLimiter = new ConcurrencyLimiter(ksqlConfig.getInt(KsqlConfig.KSQL_QUERY_PULL_MAX_CONCURRENT_REQUESTS_CONFIG), "pull", metricCollectors.getMetrics(), metricsTags);
    final SlidingWindowRateLimiter pullBandRateLimiter = new SlidingWindowRateLimiter(ksqlConfig.getInt(KsqlConfig.KSQL_QUERY_PULL_MAX_HOURLY_BANDWIDTH_MEGABYTES_CONFIG), NUM_MILLISECONDS_IN_HOUR, "pull", metricCollectors.getMetrics(), metricsTags);
    final SlidingWindowRateLimiter scalablePushBandRateLimiter = new SlidingWindowRateLimiter(ksqlConfig.getInt(KsqlConfig.KSQL_QUERY_PUSH_V2_MAX_HOURLY_BANDWIDTH_MEGABYTES_CONFIG), NUM_MILLISECONDS_IN_HOUR, "push", metricCollectors.getMetrics(), metricsTags);
    final DenyListPropertyValidator denyListPropertyValidator = new DenyListPropertyValidator(ksqlConfig.getList(KsqlConfig.KSQL_PROPERTIES_OVERRIDES_DENYLIST));
    final Optional<PullQueryExecutorMetrics> pullQueryMetrics = ksqlConfig.getBoolean(KsqlConfig.KSQL_QUERY_PULL_METRICS_ENABLED) ? Optional.of(new PullQueryExecutorMetrics(ksqlEngine.getServiceId(), ksqlConfig.getStringAsMap(KsqlConfig.KSQL_CUSTOM_METRICS_TAGS), Time.SYSTEM, metricCollectors.getMetrics())) : Optional.empty();
    final Optional<ScalablePushQueryMetrics> scalablePushQueryMetrics = ksqlConfig.getBoolean(KsqlConfig.KSQL_QUERY_PUSH_V2_ENABLED) ? Optional.of(new ScalablePushQueryMetrics(ksqlEngine.getServiceId(), ksqlConfig.getStringAsMap(KsqlConfig.KSQL_CUSTOM_METRICS_TAGS), Time.SYSTEM, metricCollectors.getMetrics())) : Optional.empty();
    final HARouting pullQueryRouting = new HARouting(routingFilterFactory, pullQueryMetrics, ksqlConfig);
    final PushRouting pushQueryRouting = new PushRouting();
    final Optional<LocalCommands> localCommands = createLocalCommands(restConfig, ksqlEngine);
    final QueryExecutor queryExecutor = new QueryExecutor(ksqlEngine, restConfig, ksqlConfig, pullQueryMetrics, scalablePushQueryMetrics, pullQueryRateLimiter, pullQueryConcurrencyLimiter, pullBandRateLimiter, scalablePushBandRateLimiter, pullQueryRouting, pushQueryRouting, localCommands);
    final StreamedQueryResource streamedQueryResource = new StreamedQueryResource(ksqlEngine, restConfig, commandStore, Duration.ofMillis(restConfig.getLong(KsqlRestConfig.STREAMED_QUERY_DISCONNECT_CHECK_MS_CONFIG)), Duration.ofMillis(restConfig.getLong(DISTRIBUTED_COMMAND_RESPONSE_TIMEOUT_MS_CONFIG)), versionChecker::updateLastRequestTime, authorizationValidator, errorHandler, denyListPropertyValidator, queryExecutor);
    final List<String> managedTopics = new LinkedList<>();
    managedTopics.add(commandTopicName);
    if (processingLogConfig.getBoolean(ProcessingLogConfig.TOPIC_AUTO_CREATE)) {
        managedTopics.add(ProcessingLogServerUtils.getTopicName(processingLogConfig, ksqlConfig));
    }
    final CommandRunner commandRunner = new CommandRunner(statementExecutor, commandStore, maxStatementRetries, new ClusterTerminator(ksqlEngine, serviceContext, managedTopics), serverState, ksqlConfig.getString(KsqlConfig.KSQL_SERVICE_ID_CONFIG), Duration.ofMillis(restConfig.getLong(KsqlRestConfig.KSQL_COMMAND_RUNNER_BLOCKED_THRESHHOLD_ERROR_MS)), metricsPrefix, InternalTopicSerdes.deserializer(Command.class), errorHandler, internalTopicClient, commandTopicName, metricCollectors.getMetrics());
    final KsqlResource ksqlResource = new KsqlResource(ksqlEngine, commandRunner, Duration.ofMillis(restConfig.getLong(DISTRIBUTED_COMMAND_RESPONSE_TIMEOUT_MS_CONFIG)), versionChecker::updateLastRequestTime, authorizationValidator, errorHandler, connectErrorHandler, denyListPropertyValidator);
    final List<KsqlServerPrecondition> preconditions = restConfig.getConfiguredInstances(KsqlRestConfig.KSQL_SERVER_PRECONDITIONS, KsqlServerPrecondition.class);
    final List<KsqlConfigurable> configurables = ImmutableList.of(ksqlResource, streamedQueryResource, statementExecutor);
    final Consumer<KsqlConfig> rocksDBConfigSetterHandler = RocksDBConfigSetterHandler::maybeConfigureRocksDBConfigSetter;
    return new KsqlRestApplication(serviceContext, ksqlEngine, ksqlConfig, restConfig, commandRunner, commandStore, statusResource, streamedQueryResource, ksqlResource, versionChecker, ksqlSecurityContextProvider, securityExtension, securityHandlerPlugin, serverState, processingLogContext, preconditions, configurables, rocksDBConfigSetterHandler, heartbeatAgent, lagReportingAgent, vertx, denyListPropertyValidator, pullQueryMetrics, scalablePushQueryMetrics, localCommands, queryExecutor, metricCollectors, internalTopicClient, internalAdmin);
}