feat(elasticsearch): make service as an index (#136)

Author: mabdh

cli/migrate.go

@@ -5,19 +5,13 @@ import (
 	"fmt"
 	"os/signal"
 	"syscall"
-	"time"
 
 	"github.com/MakeNowJust/heredoc"
-	"github.com/odpf/compass/core/asset"
 	"github.com/odpf/compass/internal/store/postgres"
 	"github.com/odpf/salt/log"
 	"github.com/spf13/cobra"
 )
 
-const (
-	esMigrationTimeout = 5 * time.Second
-)
-
 func cmdMigrate() *cobra.Command {
 	return &cobra.Command{
 		Use:   "migrate",
@@ -55,11 +49,6 @@ func runMigrations(ctx context.Context, config Config) error {
 	}
 	logger.Info("Migration Postgres done.")
 
-	logger.Info("Migrating ES...")
-	if err := migrateElasticsearch(logger, config); err != nil {
-		return err
-	}
-	logger.Info("Migration ES done.")
 	return nil
 }
 
@@ -79,23 +68,3 @@ func migratePostgres(logger log.Logger, config Config) (err error) {
 
 	return nil
 }
-
-func migrateElasticsearch(logger log.Logger, config Config) error {
-	logger.Info("Initiating ES client...")
-	esClient, err := initElasticsearch(logger, config.Elasticsearch)
-	if err != nil {
-		return err
-	}
-
-	for _, supportedType := range asset.AllSupportedTypes {
-		logger.Info("Migrating type\n", "type", supportedType)
-		ctx, cancel := context.WithTimeout(context.Background(), esMigrationTimeout)
-		defer cancel()
-		err := esClient.Migrate(ctx, supportedType)
-		if err != nil {
-			return fmt.Errorf("error creating/replacing type %q: %w", supportedType, err)
-		}
-		logger.Info("created/updated type\n", "type", supportedType)
-	}
-	return nil
-}

internal/store/elasticsearch/discovery_repository.go

@@ -6,25 +6,15 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
-	"strconv"
 	"strings"
 
 	"github.com/odpf/compass/core/asset"
 )
 
-var indexTypeMap = map[asset.Type]string{
-	asset.TypeTable:     "table",
-	asset.TypeTopic:     "topic",
-	asset.TypeJob:       "job",
-	asset.TypeDashboard: "dashboard",
-}
-
 // DiscoveryRepository implements discovery.Repository
 // with elasticsearch as the backing store.
 type DiscoveryRepository struct {
-	cli              *Client
-	typeWhiteList    []string
-	typeWhiteListSet map[string]bool
+	cli *Client
 }
 
 func NewDiscoveryRepository(cli *Client) *DiscoveryRepository {
@@ -41,6 +31,17 @@ func (repo *DiscoveryRepository) Upsert(ctx context.Context, ast asset.Asset) er
 		return asset.ErrUnknownType
 	}
 
+	idxExists, err := repo.cli.indexExists(ctx, ast.Service)
+	if err != nil {
+		return elasticSearchError(err)
+	}
+
+	if !idxExists {
+		if err := repo.cli.CreateIdx(ctx, ast.Service); err != nil {
+			return err
+		}
+	}
+
 	body, err := repo.createUpsertBody(ast)
 	if err != nil {
 		return fmt.Errorf("error serialising payload: %w", err)
@@ -82,9 +83,9 @@ func (repo *DiscoveryRepository) Delete(ctx context.Context, assetID string) err
 }
 
 func (repo *DiscoveryRepository) GetTypes(ctx context.Context) (map[asset.Type]int, error) {
-	resp, err := repo.cli.client.Cat.Indices(
-		repo.cli.client.Cat.Indices.WithFormat("json"),
-		repo.cli.client.Cat.Indices.WithContext(ctx),
+	resp, err := repo.cli.client.Search(
+		repo.cli.client.Search.WithContext(ctx),
+		repo.cli.client.Search.WithBody(strings.NewReader(`{"size": 0,"aggs":{"aggregation_name":{"terms":{"field":"type.keyword"}}}}`)),
 	)
 	if err != nil {
 		return nil, fmt.Errorf("error from es client %w", err)
@@ -93,23 +94,20 @@ func (repo *DiscoveryRepository) GetTypes(ctx context.Context) (map[asset.Type]i
 	if resp.IsError() {
 		return nil, fmt.Errorf("error from es server: %s", errorReasonFromResponse(resp))
 	}
-	var indices []esIndex
-	err = json.NewDecoder(resp.Body).Decode(&indices)
+
+	var response searchResponse
+	err = json.NewDecoder(resp.Body).Decode(&response)
 	if err != nil {
-		return nil, fmt.Errorf("error decoding es response %w", err)
+		return nil, fmt.Errorf("error decoding aggregate search response %w", err)
 	}
 
 	results := map[asset.Type]int{}
-	for _, index := range indices {
-		count, err := strconv.Atoi(index.DocsCount)
-		if err != nil {
-			return results, fmt.Errorf("error converting docs count to a number: %w", err)
-		}
-		typName := asset.Type(index.Index)
+	for _, bucket := range response.Aggregations.AggregationName.Buckets {
+		typName := asset.Type(bucket.Key)
 		if !typName.IsValid() {
 			continue
 		}
-		results[typName] = count
+		results[typName] = bucket.DocCount
 	}
 
 	return results, nil
@@ -132,7 +130,7 @@ func (repo *DiscoveryRepository) createUpsertBody(ast asset.Asset) (io.Reader, e
 func (repo *DiscoveryRepository) writeInsertAction(w io.Writer, ast asset.Asset) error {
 	action := map[string]interface{}{
 		"index": map[string]interface{}{
-			"_index": indexTypeMap[ast.Type],
+			"_index": ast.Service,
 			"_id":    ast.ID,
 		},
 	}

internal/store/elasticsearch/discovery_repository_test.go

@@ -14,7 +14,10 @@ import (
 )
 
 func TestDiscoveryRepositoryUpsert(t *testing.T) {
-	ctx := context.Background()
+	var (
+		ctx             = context.Background()
+		bigqueryService = "bigquery-test"
+	)
 
 	t.Run("should return error if id empty", func(t *testing.T) {
 		cli, err := esTestServer.NewClient()
@@ -28,8 +31,9 @@ func TestDiscoveryRepositoryUpsert(t *testing.T) {
 
 		repo := store.NewDiscoveryRepository(esClient)
 		err = repo.Upsert(ctx, asset.Asset{
-			ID:   "",
-			Type: asset.TypeTable,
+			ID:      "",
+			Type:    asset.TypeTable,
+			Service: bigqueryService,
 		})
 		assert.ErrorIs(t, err, asset.ErrEmptyID)
 	})
@@ -46,18 +50,19 @@ func TestDiscoveryRepositoryUpsert(t *testing.T) {
 
 		repo := store.NewDiscoveryRepository(esClient)
 		err = repo.Upsert(ctx, asset.Asset{
-			ID:   "sample-id",
-			Type: asset.Type("unknown-type"),
+			ID:      "sample-id",
+			Type:    asset.Type("unknown-type"),
+			Service: bigqueryService,
 		})
 		assert.ErrorIs(t, err, asset.ErrUnknownType)
 	})
 
-	t.Run("should insert asset to the correct index by its type", func(t *testing.T) {
+	t.Run("should insert asset to the correct index by its service", func(t *testing.T) {
 		ast := asset.Asset{
 			ID:          "sample-id",
 			URN:         "sample-urn",
 			Type:        asset.TypeTable,
-			Service:     "bigquery",
+			Service:     bigqueryService,
 			Name:        "sample-name",
 			Description: "sample-description",
 			Data: map[string]interface{}{
@@ -85,7 +90,7 @@ func TestDiscoveryRepositoryUpsert(t *testing.T) {
 		err = repo.Upsert(ctx, ast)
 		assert.NoError(t, err)
 
-		res, err := cli.API.Get("table", ast.ID)
+		res, err := cli.API.Get(bigqueryService, ast.ID)
 		require.NoError(t, err)
 		require.False(t, res.IsError())
 
@@ -112,7 +117,7 @@ func TestDiscoveryRepositoryUpsert(t *testing.T) {
 			ID:          "existing-id",
 			URN:         "existing-urn",
 			Type:        asset.TypeTable,
-			Service:     "bigquery",
+			Service:     bigqueryService,
 			Name:        "existing-name",
 			Description: "existing-description",
 		}
@@ -137,7 +142,7 @@ func TestDiscoveryRepositoryUpsert(t *testing.T) {
 		err = repo.Upsert(ctx, newAsset)
 		assert.NoError(t, err)
 
-		res, err := cli.API.Get("table", existingAsset.ID)
+		res, err := cli.API.Get(bigqueryService, existingAsset.ID)
 		require.NoError(t, err)
 		require.False(t, res.IsError())
 
@@ -155,7 +160,10 @@ func TestDiscoveryRepositoryUpsert(t *testing.T) {
 }
 
 func TestDiscoveryRepositoryDelete(t *testing.T) {
-	ctx := context.Background()
+	var (
+		ctx             = context.Background()
+		bigqueryService = "bigquery-test"
+	)
 
 	t.Run("should return error if id empty", func(t *testing.T) {
 		cli, err := esTestServer.NewClient()
@@ -174,9 +182,10 @@ func TestDiscoveryRepositoryDelete(t *testing.T) {
 
 	t.Run("should not return error on success", func(t *testing.T) {
 		ast := asset.Asset{
-			ID:   "delete-id",
-			Type: asset.TypeTable,
-			URN:  "some-urn",
+			ID:      "delete-id",
+			Type:    asset.TypeTable,
+			Service: bigqueryService,
+			URN:     "some-urn",
 		}
 
 		cli, err := esTestServer.NewClient()
@@ -199,7 +208,10 @@ func TestDiscoveryRepositoryDelete(t *testing.T) {
 }
 
 func TestDiscoveryRepositoryGetTypes(t *testing.T) {
-	ctx := context.Background()
+	var (
+		ctx             = context.Background()
+		bigqueryService = "bigquery-test"
+	)
 
 	t.Run("should return empty map if no type is available", func(t *testing.T) {
 		cli, err := esTestServer.NewClient()
@@ -218,8 +230,7 @@ func TestDiscoveryRepositoryGetTypes(t *testing.T) {
 		assert.Equal(t, map[asset.Type]int{}, counts)
 	})
 
-	t.Run("should return map with 0 count if type has not been populated yet", func(t *testing.T) {
-		typ := asset.TypeTable
+	t.Run("should return empty map if type has not been populated yet", func(t *testing.T) {
 		cli, err := esTestServer.NewClient()
 		require.NoError(t, err)
 		esClient, err := store.NewClient(
@@ -229,25 +240,27 @@ func TestDiscoveryRepositoryGetTypes(t *testing.T) {
 		)
 		require.NoError(t, err)
 
-		err = esClient.Migrate(ctx, typ)
+		err = esClient.CreateIdx(ctx, bigqueryService)
 		require.NoError(t, err)
 
 		repo := store.NewDiscoveryRepository(esClient)
 		counts, err := repo.GetTypes(ctx)
 		require.NoError(t, err)
 
-		expected := map[asset.Type]int{
-			asset.TypeTable: 0,
-		}
+		expected := map[asset.Type]int{}
 		assert.Equal(t, expected, counts)
 	})
 
 	t.Run("should return maps of asset count with valid type as its key", func(t *testing.T) {
-		typ := asset.TypeDashboard
+		var (
+			typ            = asset.TypeDashboard
+			tableauService = "tableau-test"
+		)
+
 		assets := []asset.Asset{
-			{ID: "id-asset-1", URN: "asset-1", Name: "asset-1", Type: typ},
-			{ID: "id-asset-2", URN: "asset-2", Name: "asset-2", Type: typ},
-			{ID: "id-asset-3", URN: "asset-3", Name: "asset-3", Type: typ},
+			{ID: "id-asset-1", URN: "asset-1", Name: "asset-1", Type: typ, Service: tableauService},
+			{ID: "id-asset-2", URN: "asset-2", Name: "asset-2", Type: typ, Service: tableauService},
+			{ID: "id-asset-3", URN: "asset-3", Name: "asset-3", Type: typ, Service: tableauService},
 		}
 
 		cli, err := esTestServer.NewClient()
@@ -259,13 +272,6 @@ func TestDiscoveryRepositoryGetTypes(t *testing.T) {
 		)
 		require.NoError(t, err)
 
-		err = esClient.Migrate(ctx, asset.TypeDashboard)
-		require.NoError(t, err)
-
-		invalidType := "invalid-type"
-		err = esClient.Migrate(ctx, asset.Type(invalidType))
-		require.NoError(t, err)
-
 		repo := store.NewDiscoveryRepository(esClient)
 		_, err = repo.GetTypes(ctx)
 		require.NoError(t, err)