<commit_before>package bql

import (
	"errors"
	"fmt"
	"pfi/sensorbee/sensorbee/bql/execution"
	"pfi/sensorbee/sensorbee/bql/parser"
	"pfi/sensorbee/sensorbee/bql/udf"
	"pfi/sensorbee/sensorbee/core"
	"pfi/sensorbee/sensorbee/data"
	"sync"
	"sync/atomic"
)

type TopologyBuilder struct {
	topology       core.Topology
	Reg            udf.FunctionManager
	UDSFCreators   udf.UDSFCreatorRegistry
	UDSCreators    udf.UDSCreatorRegistry
	SourceCreators SourceCreatorRegistry
	SinkCreators   SinkCreatorRegistry
}

// TODO: Provide AtomicTopologyBuilder which support building multiple nodes
// in an atomic manner (kind of transactionally)

// NewTopologyBuilder creates a new TopologyBuilder which dynamically creates
// nodes from BQL statements. The target Topology can be shared by
// multiple TopologyBuilders.
//
// TopologyBuilder doesn't support atomic topology building. For example,
// when a user wants to add three statement and the second statement fails,
// only the node created from the first statement is registered to the topology
// and it starts to generate tuples. Others won't be registered.
func NewTopologyBuilder(t core.Topology) (*TopologyBuilder, error) {
	udsfs, err := udf.CopyGlobalUDSFCreatorRegistry()
	if err != nil {
		return nil, err
	}

	udss, err := udf.CopyGlobalUDSCreatorRegistry()
	if err != nil {
		return nil, err
	}

	srcs, err := CopyGlobalSourceCreatorRegistry()
	if err != nil {
		return nil, err
	}
	// node_statuses builtin source can only be registered here because it
	// requires a topology.
	if err := srcs.Register("node_statuses", createNodeStatusSourceCreator(t)); err != nil {
		return nil, err
	}

	sinks, err := CopyGlobalSinkCreatorRegistry()
	if err != nil {
		return nil, err
	}

	tb := &TopologyBuilder{
		topology:       t,
		Reg:            udf.CopyGlobalUDFRegistry(t.Context()),
		UDSFCreators:   udsfs,
		UDSCreators:    udss,
		SourceCreators: srcs,
		SinkCreators:   sinks,
	}
	return tb, nil
}

// TODO: if IDs are shared by distributed processes, they should have a process
// ID of each process in it or they should be generated by a central server
// to be globally unique. Currently, this id is only used temporarily and
// doesn't have to be strictly unique nor recoverable (durable).
var (
	topologyBuilderTemporaryID int64
)

func topologyBuilderNextTemporaryID() int64 {
	return atomic.AddInt64(&topologyBuilderTemporaryID, 1)
}

func (tb *TopologyBuilder) Topology() core.Topology {
	return tb.topology
}

// AddStmt add a node created from a statement to the topology. It returns
// a created node. It returns a nil node when the statement is CREATE STATE.
func (tb *TopologyBuilder) AddStmt(stmt interface{}) (core.Node, error) {
	// TODO: Enable StopOnDisconnect properly

	// check the type of statement
	switch stmt := stmt.(type) {
	case parser.CreateSourceStmt:
		// load params into map for faster access
		paramsMap := tb.mkParamsMap(stmt.Params)

		// check if we know whis type of source
		creator, err := tb.SourceCreators.Lookup(string(stmt.Type))
		if err != nil {
			return nil, err
		}

		// if so, try to create such a source
		source, err := creator.CreateSource(tb.topology.Context(), &IOParams{
			TypeName: string(stmt.Type),
			Name:     string(stmt.Name),
		}, paramsMap)
		if err != nil {
			return nil, err
		}
		return tb.topology.AddSource(string(stmt.Name), source, &core.SourceConfig{
			PausedOnStartup: stmt.Paused == parser.Yes,
		})

	case parser.CreateStreamAsSelectStmt:
		return tb.createStreamAsSelectStmt(&stmt)

	case parser.CreateStreamAsSelectUnionStmt:
		// idea: create an intermediate box for each SELECT substatement,
		// then connect them with a simple forwarder box
		names := make([]string, 0, len(stmt.Selects))
		removeTmpNodes := func() {
			for _, name := range names {
				tb.topology.Remove(name)
			}
		}
		for _, selStmt := range stmt.Selects {
			// create a stream with a generated name and recurse
			tmpName := fmt.Sprintf("sensorbee_tmp_%v", topologyBuilderNextTemporaryID())
			tmpStmt := parser.CreateStreamAsSelectStmt{
				parser.StreamIdentifier(tmpName),
				selStmt,
			}
			box, err := tb.AddStmt(tmpStmt)
			if err != nil {
				removeTmpNodes()
				return nil, err
			}
			names = append(names, tmpName)
			box.(core.BoxNode).StopOnDisconnect(core.Inbound | core.Outbound)
			box.(core.BoxNode).RemoveOnStop()
		}
		// simple forwarder box
		forwardBox := core.BoxFunc(func(ctx *core.Context, t *core.Tuple, w core.Writer) error {
			return w.Write(ctx, t)
		})
		node, err := tb.topology.AddBox(string(stmt.Name), forwardBox, nil)
		if err != nil {
			removeTmpNodes()
			return nil, err
		}
		// TODO if we add
		//   node.StopOnDisconnect(core.Inbound)
		//   node.RemoveOnStop()
		// here, then all the tests will fail. this should
		// be investigated...?
		// connect inputs
		for _, name := range names {
			if err := node.Input(name, nil); err != nil {
				removeTmpNodes()
				return nil, err
			}
		}
		return node, nil

	case parser.CreateSinkStmt:
		// load params into map for faster access
		paramsMap := tb.mkParamsMap(stmt.Params)

		// check if we know whis type of sink
		creator, err := tb.SinkCreators.Lookup(string(stmt.Type))
		if err != nil {
			return nil, err
		}

		// if so, try to create such a sink
		sink, err := creator.CreateSink(tb.topology.Context(), &IOParams{
			TypeName: string(stmt.Type),
			Name:     string(stmt.Name),
		}, paramsMap)<commit_msg>move the StopOnDisconnect()/RemoveOnStop() calls to the bottom<commit_after>package bql

import (
	"errors"
	"fmt"
	"pfi/sensorbee/sensorbee/bql/execution"
	"pfi/sensorbee/sensorbee/bql/parser"
	"pfi/sensorbee/sensorbee/bql/udf"
	"pfi/sensorbee/sensorbee/core"
	"pfi/sensorbee/sensorbee/data"
	"sync"
	"sync/atomic"
)

type TopologyBuilder struct {
	topology       core.Topology
	Reg            udf.FunctionManager
	UDSFCreators   udf.UDSFCreatorRegistry
	UDSCreators    udf.UDSCreatorRegistry
	SourceCreators SourceCreatorRegistry
	SinkCreators   SinkCreatorRegistry
}

// TODO: Provide AtomicTopologyBuilder which support building multiple nodes
// in an atomic manner (kind of transactionally)

// NewTopologyBuilder creates a new TopologyBuilder which dynamically creates
// nodes from BQL statements. The target Topology can be shared by
// multiple TopologyBuilders.
//
// TopologyBuilder doesn't support atomic topology building. For example,
// when a user wants to add three statement and the second statement fails,
// only the node created from the first statement is registered to the topology
// and it starts to generate tuples. Others won't be registered.
func NewTopologyBuilder(t core.Topology) (*TopologyBuilder, error) {
	udsfs, err := udf.CopyGlobalUDSFCreatorRegistry()
	if err != nil {
		return nil, err
	}

	udss, err := udf.CopyGlobalUDSCreatorRegistry()
	if err != nil {
		return nil, err
	}

	srcs, err := CopyGlobalSourceCreatorRegistry()
	if err != nil {
		return nil, err
	}
	// node_statuses builtin source can only be registered here because it
	// requires a topology.
	if err := srcs.Register("node_statuses", createNodeStatusSourceCreator(t)); err != nil {
		return nil, err
	}

	sinks, err := CopyGlobalSinkCreatorRegistry()
	if err != nil {
		return nil, err
	}

	tb := &TopologyBuilder{
		topology:       t,
		Reg:            udf.CopyGlobalUDFRegistry(t.Context()),
		UDSFCreators:   udsfs,
		UDSCreators:    udss,
		SourceCreators: srcs,
		SinkCreators:   sinks,
	}
	return tb, nil
}

// TODO: if IDs are shared by distributed processes, they should have a process
// ID of each process in it or they should be generated by a central server
// to be globally unique. Currently, this id is only used temporarily and
// doesn't have to be strictly unique nor recoverable (durable).
var (
	topologyBuilderTemporaryID int64
)

func topologyBuilderNextTemporaryID() int64 {
	return atomic.AddInt64(&topologyBuilderTemporaryID, 1)
}

func (tb *TopologyBuilder) Topology() core.Topology {
	return tb.topology
}

// AddStmt add a node created from a statement to the topology. It returns
// a created node. It returns a nil node when the statement is CREATE STATE.
func (tb *TopologyBuilder) AddStmt(stmt interface{}) (core.Node, error) {
	// TODO: Enable StopOnDisconnect properly

	// check the type of statement
	switch stmt := stmt.(type) {
	case parser.CreateSourceStmt:
		// load params into map for faster access
		paramsMap := tb.mkParamsMap(stmt.Params)

		// check if we know whis type of source
		creator, err := tb.SourceCreators.Lookup(string(stmt.Type))
		if err != nil {
			return nil, err
		}

		// if so, try to create such a source
		source, err := creator.CreateSource(tb.topology.Context(), &IOParams{
			TypeName: string(stmt.Type),
			Name:     string(stmt.Name),
		}, paramsMap)
		if err != nil {
			return nil, err
		}
		return tb.topology.AddSource(string(stmt.Name), source, &core.SourceConfig{
			PausedOnStartup: stmt.Paused == parser.Yes,
		})

	case parser.CreateStreamAsSelectStmt:
		return tb.createStreamAsSelectStmt(&stmt)

	case parser.CreateStreamAsSelectUnionStmt:
		// idea: create an intermediate box for each SELECT substatement,
		// then connect them with a simple forwarder box
		names := make([]string, 0, len(stmt.Selects))
		nodes := make([]core.BoxNode, 0, len(stmt.Selects))
		removeTmpNodes := func() {
			for _, name := range names {
				tb.topology.Remove(name)
			}
		}
		for _, selStmt := range stmt.Selects {
			// create a stream with a generated name and recurse
			tmpName := fmt.Sprintf("sensorbee_tmp_%v", topologyBuilderNextTemporaryID())
			tmpStmt := parser.CreateStreamAsSelectStmt{
				parser.StreamIdentifier(tmpName),
				selStmt,
			}
			box, err := tb.AddStmt(tmpStmt)
			if err != nil {
				removeTmpNodes()
				return nil, err
			}
			names = append(names, tmpName)
			nodes = append(nodes, box.(core.BoxNode))
		}
		// simple forwarder box
		forwardBox := core.BoxFunc(func(ctx *core.Context, t *core.Tuple, w core.Writer) error {
			return w.Write(ctx, t)
		})
		node, err := tb.topology.AddBox(string(stmt.Name), forwardBox, nil)
		if err != nil {
			removeTmpNodes()
			return nil, err
		}
		// connect inputs
		for _, name := range names {
			if err := node.Input(name, nil); err != nil {
				removeTmpNodes()
				return nil, err
			}
		}
		for _, node := range nodes {
			node.StopOnDisconnect(core.Inbound | core.Outbound)
			node.RemoveOnStop()
		}
		node.StopOnDisconnect(core.Inbound)
		node.RemoveOnStop()
		return node, nil

	case parser.CreateSinkStmt:
		// load params into map for faster access
		paramsMap := tb.mkParamsMap(stmt.Params)

		// check if we know whis type of sink
		creator, err := tb.SinkCreators.Lookup(string(stmt.Type))
		if err != nil {
			return nil, err
		}

		// if so, try to create such a sink
		sink, err := creator.CreateSink(tb.topology.Context(), &IOParams{
			TypeName: string(stmt.Type),
			Name:     string(stmt.Name),
		}, paramsMap)
{'fraction_non_alphanumeric': 0.10407695831424645, 'fraction_numerical': 0.0015574896930829133, 'mean_word_length': 3.6490630323679727, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4856512', 'n_tokens_mistral': 3774, 'n_tokens_neox': 3485, 'n_words': 1067}