plug-sdk/pkg/plug/server.go

package plug

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

	"git.anthrove.art/Anthrove/otter-space-sdk/v3/pkg/database"
	otterError "git.anthrove.art/Anthrove/otter-space-sdk/v3/pkg/error"
	"git.anthrove.art/Anthrove/otter-space-sdk/v3/pkg/models"

	"github.com/golang/protobuf/ptypes/timestamp"

	pb "git.anthrove.art/Anthrove/plug-sdk/v3/pkg/grpc"
	"google.golang.org/grpc"
)

type Message struct {
	Title     string
	Body      string
	CreatedAt *timestamp.Timestamp
}

type TaskExecution func(ctx context.Context, source models.Source, userSourceUsername string, anthroveUser models.User, deepScrape bool, apiKey string, cancelFunction func()) error
type SendMessageExecution func(ctx context.Context, source models.Source, userSourceID string, userSourceUsername string, message string) error
type GetMessageExecution func(ctx context.Context, source models.Source, userSourceID string, userSourceUsername string) ([]Message, error)

type Plug struct {
	address               string
	port                  string
	ctx                   context.Context
	taskExecutionFunction TaskExecution
	sendMessageExecution  SendMessageExecution
	getMessageExecution   GetMessageExecution
	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
	var source models.Source

	log.Print("Check if source exists")
	source, err = database.GetSourceByDomain(p.ctx, p.source.Domain)
	if err != nil {
		if err.Error() == otterError.NoDataFound {
			log.Printf("Initalizing source!")
			if source, err = database.CreateSource(p.ctx, p.source); err != nil {
				panic(err)
			}
		} else {
			panic(err)
		}
	}

	p.source = source

	log.Print("Source exists")

	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.source, p.taskExecutionFunction, p.sendMessageExecution, p.getMessageExecution))

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

	return nil
}

func (p *Plug) GetSource() models.Source {
	return p.source
}

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

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

func (p *Plug) GetMessageExecution(function GetMessageExecution) {
	p.getMessageExecution = function
}