plug-sdk/pkg/plug/server.go

package plug

import (
	"context"
	"fmt"
	"log"
	"net"

	"git.dragse.it/anthrove/otter-space-sdk/v2/pkg/database"
	"git.dragse.it/anthrove/otter-space-sdk/v2/pkg/models"

	pb "git.dragse.it/anthrove/plug-sdk/v2/pkg/grpc"
	"google.golang.org/grpc"
)

type TaskExecution func(ctx context.Context, database database.OtterSpace, sourceUsername string, anthroveUser models.User, deepScrape bool, apiKey string, cancelFunction func()) error
type SendMessageExecution func(ctx context.Context, userSourceID string, message string) error

type Plug struct {
	address               string
	port                  string
	ctx                   context.Context
	database              database.OtterSpace
	taskExecutionFunction TaskExecution
	sendMessageExecution  SendMessageExecution
	source                models.Source
}

func NewPlug(ctx context.Context, address string, port string, source models.Source) Plug {
	return Plug{
		ctx:     ctx,
		address: address,
		port:    port,
		source:  source,
	}
}

func (p *Plug) Listen() error {
	var err error

	plugSource, err := p.database.GetSourceByDomain(p.ctx, models.AnthroveSourceDomain(p.source.Domain))
	if err != nil {
		return err
	}

	if plugSource == nil {
		log.Printf("no source found for domain %s, initilazing source!", p.source.Domain)
		err := p.database.CreateSource(p.ctx, &p.source)
		if err != nil {
			return err
		}
	}

	lis, err := net.Listen("tcp", fmt.Sprintf("%s:%s", p.address, p.port))
	if err != nil {
		return err
	}

	grpcServer := grpc.NewServer()

	pb.RegisterPlugConnectorServer(grpcServer, NewGrpcServer(p.database, p.taskExecutionFunction, p.sendMessageExecution))

	err = grpcServer.Serve(lis)
	if err != nil {
		return err
	}

	return nil
}

func (p *Plug) WithOtterSpace(graph database.OtterSpace) {
	p.database = graph
}

func (p *Plug) TaskExecutionFunction(function TaskExecution) {
	p.taskExecutionFunction = function
}

func (p *Plug) SendMessageExecution(function SendMessageExecution) {
	p.sendMessageExecution = function
}