Refactor gateway codes + server initialization code (#104)

* Refactor gateway codes + server initilalization code * Change environment variable validation logic to applicable for prod only for now. --------- Co-authored-by: James Kwon <[email protected]>
Comfy-Org · Dec 31, 2024 · 87546ee · 87546ee
1 parent 5bf6223
commit 87546ee
Show file tree

Hide file tree

Showing 12 changed files with 262 additions and 176 deletions.
diff --git a/config/config.go b/config/config.go
@@ -11,4 +11,5 @@ type Config struct {
 	IDTokenAudience               string
 	AlgoliaAppID                  string
 	AlgoliaAPIKey                 string
+	CloudStorageBucketName        string
 }
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -21,3 +21,4 @@ services:
       ALGOLIA_APP_ID: 
       ALGOLIA_API_KEY: 
       SECRET_SCANNER_URL: ""
+      CLOUD_STORAGE_BUCKET_NAME: "staging-comfy-registry"
diff --git a/gateways/algolia/algolia-noop.go b/gateways/algolia/algolia-noop.go
@@ -9,39 +9,34 @@ import (
 
 var _ AlgoliaService = (*algolianoop)(nil)
 
+// No-op implementation for AlgoliaService, logging calls instead of performing operations
 type algolianoop struct{}
 
-// DeleteNode implements AlgoliaService.
 func (a *algolianoop) DeleteNode(ctx context.Context, node *ent.Node) error {
 	log.Ctx(ctx).Info().Msgf("algolia noop: delete node: %s", node.ID)
 	return nil
 }
 
-// IndexNodes implements AlgoliaService.
 func (a *algolianoop) IndexNodes(ctx context.Context, nodes ...*ent.Node) error {
 	log.Ctx(ctx).Info().Msgf("algolia noop: index nodes: %d number of nodes", len(nodes))
 	return nil
 }
 
-// SearchNodes implements AlgoliaService.
 func (a *algolianoop) SearchNodes(ctx context.Context, query string, opts ...interface{}) ([]*ent.Node, error) {
 	log.Ctx(ctx).Info().Msgf("algolia noop: search nodes: %s", query)
 	return nil, nil
 }
 
-// DeleteNodeVersion implements AlgoliaService.
 func (a *algolianoop) DeleteNodeVersions(ctx context.Context, nodes ...*ent.NodeVersion) error {
-	log.Ctx(ctx).Info().Msgf("algolia noop: delete node version:  %d number of node versions", len(nodes))
+	log.Ctx(ctx).Info().Msgf("algolia noop: delete node version: %d number of node versions", len(nodes))
 	return nil
 }
 
-// IndexNodeVersions implements AlgoliaService.
 func (a *algolianoop) IndexNodeVersions(ctx context.Context, nodes ...*ent.NodeVersion) error {
 	log.Ctx(ctx).Info().Msgf("algolia noop: index node versions: %d number of node versions", len(nodes))
 	return nil
 }
 
-// SearchNodeVersions implements AlgoliaService.
 func (a *algolianoop) SearchNodeVersions(ctx context.Context, query string, opts ...interface{}) ([]*ent.NodeVersion, error) {
 	log.Ctx(ctx).Info().Msgf("algolia noop: search node versions: %s", query)
 	return nil, nil

diff --git a/gateways/algolia/algolia.go b/gateways/algolia/algolia.go
@@ -3,10 +3,10 @@ package algolia
 import (
 	"context"
 	"fmt"
-	"os"
-	"registry-backend/ent"
-
 	"github.com/algolia/algoliasearch-client-go/v3/algolia/search"
+	"github.com/rs/zerolog/log"
+	"registry-backend/config" // assuming a config package exists to hold config values
+	"registry-backend/ent"
 )
 
 // AlgoliaService defines the interface for interacting with Algolia search.
@@ -27,35 +27,21 @@ type algolia struct {
 	client *search.Client
 }
 
-// New creates a new Algolia service with the provided app ID and API key.
-func New(appid, apikey string) (AlgoliaService, error) {
-	return &algolia{
-		client: search.NewClient(appid, apikey),
-	}, nil
-}
-
-// NewFromEnv creates a new Algolia service using environment variables for configuration.
-func NewFromEnv() (AlgoliaService, error) {
-	appid, ok := os.LookupEnv("ALGOLIA_APP_ID")
-	if !ok {
-		return nil, fmt.Errorf("required env variable ALGOLIA_APP_ID is not set")
-	}
-	apikey, ok := os.LookupEnv("ALGOLIA_API_KEY")
-	if !ok {
-		return nil, fmt.Errorf("required env variable ALGOLIA_API_KEY is not set")
-	}
-	return New(appid, apikey)
-}
-
-// NewFromEnvOrNoop creates a new Algolia service using environment variables or noop implementation if no environment found
-func NewFromEnvOrNoop() (AlgoliaService, error) {
-	id := os.Getenv("ALGOLIA_APP_ID")
-	key := os.Getenv("ALGOLIA_API_KEY")
-	if id == "" && key == "" {
+// NewAlgoliaService creates a new Algolia service using the provided config or returns a noop implementation if the config is missing.
+func NewAlgoliaService(cfg *config.Config) (AlgoliaService, error) {
+	if cfg == nil || cfg.AlgoliaAppID == "" || cfg.AlgoliaAPIKey == "" {
+		// Return a noop implementation if config is nil or missing keys
+		log.Info().Msg("No Algolia configuration found, using noop implementation")
 		return &algolianoop{}, nil
 	}
 
-	return NewFromEnv()
+	// Fetch the Algolia app ID and API key from the provided config
+	appID := cfg.AlgoliaAppID
+	apiKey := cfg.AlgoliaAPIKey
+
+	// Initialize the Algolia client
+	client := search.NewClient(appID, apiKey)
+	return &algolia{client: client}, nil
 }
 
 // IndexNodes indexes the provided nodes in Algolia.
@@ -67,7 +53,7 @@ func (a *algolia) IndexNodes(ctx context.Context, nodes ...*ent.Node) error {
 		o := map[string]interface{}{
 			"objectID":       n.ID,
 			"name":           n.Name,
-			"publiser_id":    n.PublisherID,
+			"publisher_id":   n.PublisherID,
 			"description":    n.Description,
 			"id":             n.ID,
 			"create_time":    n.CreateTime,
@@ -139,13 +125,13 @@ func (a *algolia) IndexNodeVersions(ctx context.Context, nodes ...*ent.NodeVersi
 
 	res, err := index.SaveObjects(objects)
 	if err != nil {
-		return fmt.Errorf("failed to index nodes: %w", err)
+		return fmt.Errorf("failed to index node versions: %w", err)
 	}
 
 	return res.Wait()
 }
 
-// DeleteNodeVersion implements AlgoliaService.
+// DeleteNodeVersions implements AlgoliaService.
 func (a *algolia) DeleteNodeVersions(ctx context.Context, nodes ...*ent.NodeVersion) error {
 	index := a.client.InitIndex("node_versions_index")
 	ids := []string{}
@@ -154,7 +140,7 @@ func (a *algolia) DeleteNodeVersions(ctx context.Context, nodes ...*ent.NodeVers
 	}
 	res, err := index.DeleteObjects(ids)
 	if err != nil {
-		return fmt.Errorf("failed to delete node: %w", err)
+		return fmt.Errorf("failed to delete node versions: %w", err)
 	}
 	return res.Wait()
 }
@@ -164,7 +150,7 @@ func (a *algolia) SearchNodeVersions(ctx context.Context, query string, opts ...
 	index := a.client.InitIndex("node_versions_index")
 	res, err := index.Search(query, opts...)
 	if err != nil {
-		return nil, fmt.Errorf("failed to search nodes: %w", err)
+		return nil, fmt.Errorf("failed to search node versions: %w", err)
 	}
 
 	var nodes []*ent.NodeVersion

diff --git a/gateways/algolia/algolia_test.go b/gateways/algolia/algolia_test.go
@@ -3,6 +3,7 @@ package algolia
 import (
 	"context"
 	"os"
+	"registry-backend/config"
 	"registry-backend/ent"
 	"registry-backend/ent/schema"
 	"testing"
@@ -23,7 +24,10 @@ func TestIndex(t *testing.T) {
 		t.Skip("Required env variables `ALGOLIA_API_KEY` is not set")
 	}
 
-	algolia, err := NewFromEnv()
+	algolia, err := NewAlgoliaService(&config.Config{
+		AlgoliaAppID:  os.Getenv("ALGOLIA_APP_ID"),
+		AlgoliaAPIKey: os.Getenv("ALGOLIA_API_KEY"),
+	})
 	require.NoError(t, err)
 
 	t.Run("node", func(t *testing.T) {
@@ -74,9 +78,7 @@ func TestIndex(t *testing.T) {
 }
 
 func TestNoop(t *testing.T) {
-	t.Setenv("ALGOLIA_APP_ID", "")
-	t.Setenv("ALGOLIA_API_KEY", "")
-	a, err := NewFromEnvOrNoop()
+	a, err := NewAlgoliaService(&config.Config{})
 	require.NoError(t, err)
 	require.NoError(t, a.IndexNodes(context.Background(), &ent.Node{}))
 	require.NoError(t, a.DeleteNode(context.Background(), &ent.Node{}))

diff --git a/gateways/slack/slack.go b/gateways/slack/slack.go
@@ -1,45 +1,61 @@
-package gateway
+package slack
 
 import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"github.com/rs/zerolog/log"
 	"net/http"
 	"registry-backend/config"
 )
 
+// SlackService defines the interface for interacting with Slack notifications.
 type SlackService interface {
 	SendRegistryMessageToSlack(msg string) error
 }
 
-type DripSlackService struct {
+// Ensure slackService struct implements SlackService interface
+var _ SlackService = (*slackService)(nil)
+
+// slackService struct holds the configuration and webhook URL.
+type slackService struct {
 	registrySlackChannelWebhook string
 	config                      *config.Config
 }
 
-func NewSlackService(config *config.Config) *DripSlackService {
-	return &DripSlackService{
-		config:                      config,
-		registrySlackChannelWebhook: config.SlackRegistryChannelWebhook,
+// NewSlackService creates a new Slack service using the provided config or returns a noop implementation if the config is missing.
+func NewSlackService(cfg *config.Config) SlackService {
+	if cfg == nil || cfg.SlackRegistryChannelWebhook == "" {
+		// Return a noop implementation if config is nil or missing keys
+		log.Info().Msg("No Slack configuration found, using noop implementation")
+		return &slackNoop{}
 	}
 
+	return &slackService{
+		config:                      cfg,
+		registrySlackChannelWebhook: cfg.SlackRegistryChannelWebhook,
+	}
 }
 
 type slackRequestBody struct {
 	Text string `json:"text"`
 }
 
-func (s *DripSlackService) SendRegistryMessageToSlack(msg string) error {
-	if s.config.DripEnv == "prod" {
-		return sendSlackNotification(msg, s.registrySlackChannelWebhook)
+// SendRegistryMessageToSlack sends a message to the registry Slack channel.
+func (s *slackService) SendRegistryMessageToSlack(msg string) error {
+	// Skip sending messages in non-production environments
+	if s.config.DripEnv != "prod" {
+		log.Info().Msg("Skipping Slack operations in non-prod environment")
+		return nil
 	}
-	return nil
+
+	return sendSlackNotification(msg, s.registrySlackChannelWebhook)
 }
 
+// sendSlackNotification sends the message to the provided Slack webhook URL.
 func sendSlackNotification(msg string, slackWebhookURL string) error {
 	if slackWebhookURL == "" {
-		println("No Slack webhook URL provided, skipping sending message to Slack")
-		return nil
+		return fmt.Errorf("no Slack webhook URL provided, skipping sending message to Slack")
 	}
 
 	body, err := json.Marshal(slackRequestBody{Text: msg})
@@ -69,3 +85,14 @@ func sendSlackNotification(msg string, slackWebhookURL string) error {
 
 	return nil
 }
+
+// slackNoop is a noop implementation of the SlackService interface.
+// It does nothing and is used when no valid config or the environment is non-production.
+type slackNoop struct{}
+
+// Implement all SlackService methods for noop behavior.
+
+func (s *slackNoop) SendRegistryMessageToSlack(msg string) error {
+	// No-op, just return nil to avoid any side-effects
+	return nil
+}