plug-sdk/pkg/plug/server.go

package plug

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

	"git.anthrove.art/Anthrove/otter-space-sdk/v3/pkg/database"
	"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, userSource models.UserSource, deepScrape bool, apiKey string, cancelFunction func()) error
type SendMessageExecution func(ctx context.Context, userSource models.UserSource, message string) error
type GetMessageExecution func(ctx context.Context, userSource models.UserSource) ([]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

	source, err = database.GetSourceByDomain(p.ctx, p.source.Domain)
	if err != nil {
		if err.Error() == "Database error: NoDataFound" {
			log.Printf("Initalizing source!")
			source, err = database.CreateSource(p.ctx, p.source)
			if err != nil {
				panic(err)
			}
		} else {
			panic(err)
		}
	}

	p.source = source

	log.Print("Source exists")

	// Start the Server

	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
}