724 bi collection improvments

api/api_controller.go

@@ -53,32 +53,32 @@ const (
 )
 
 type Dependencies struct {
-	ctx          context.Context
-	Cataloger    catalog.Cataloger
-	Auth         auth.Service
-	BlockAdapter block.Adapter
-	Stats        stats.Collector
-	Retention    retention.Service
-	Dedup        *dedup.Cleaner
-	Meta         auth.MetadataManager
-	Migrator     db.Migrator
-	Collector    stats.Collector
-	logger       logging.Logger
+	ctx             context.Context
+	Cataloger       catalog.Cataloger
+	Auth            auth.Service
+	BlockAdapter    block.Adapter
+	Stats           stats.Collector
+	Retention       retention.Service
+	Dedup           *dedup.Cleaner
+	MetadataManager auth.MetadataManager
+	Migrator        db.Migrator
+	Collector       stats.Collector
+	logger          logging.Logger
 }
 
 func (d *Dependencies) WithContext(ctx context.Context) *Dependencies {
 	return &Dependencies{
-		ctx:          ctx,
-		Cataloger:    d.Cataloger,
-		Auth:         d.Auth,
-		BlockAdapter: d.BlockAdapter.WithContext(ctx),
-		Stats:        d.Stats,
-		Retention:    d.Retention,
-		Dedup:        d.Dedup,
-		Meta:         d.Meta,
-		Migrator:     d.Migrator,
-		Collector:    d.Collector,
-		logger:       d.logger.WithContext(ctx),
+		ctx:             ctx,
+		Cataloger:       d.Cataloger,
+		Auth:            d.Auth,
+		BlockAdapter:    d.BlockAdapter.WithContext(ctx),
+		Stats:           d.Stats,
+		Retention:       d.Retention,
+		Dedup:           d.Dedup,
+		MetadataManager: d.MetadataManager,
+		Migrator:        d.Migrator,
+		Collector:       d.Collector,
+		logger:          d.logger.WithContext(ctx),
 	}
 }
 
@@ -95,20 +95,20 @@ type Controller struct {
 }
 
 func NewController(cataloger catalog.Cataloger, auth auth.Service, blockAdapter block.Adapter, stats stats.Collector, retention retention.Service,
-	dedupCleaner *dedup.Cleaner, meta auth.MetadataManager, migrator db.Migrator, collector stats.Collector, logger logging.Logger) *Controller {
+	dedupCleaner *dedup.Cleaner, metadataManager auth.MetadataManager, migrator db.Migrator, collector stats.Collector, logger logging.Logger) *Controller {
 	c := &Controller{
 		deps: &Dependencies{
-			ctx:          context.Background(),
-			Cataloger:    cataloger,
-			Auth:         auth,
-			BlockAdapter: blockAdapter,
-			Stats:        stats,
-			Retention:    retention,
-			Dedup:        dedupCleaner,
-			Meta:         meta,
-			Migrator:     migrator,
-			Collector:    collector,
-			logger:       logger,
+			ctx:             context.Background(),
+			Cataloger:       cataloger,
+			Auth:            auth,
+			BlockAdapter:    blockAdapter,
+			Stats:           stats,
+			Retention:       retention,
+			Dedup:           dedupCleaner,
+			MetadataManager: metadataManager,
+			Migrator:        migrator,
+			Collector:       collector,
+			logger:          logger,
 		},
 	}
 	return c
@@ -235,7 +235,7 @@ func (c *Controller) SetupLakeFSHandler() setupop.SetupLakeFSHandler {
 		}
 
 		// check if previous setup completed
-		if ts, _ := c.deps.Meta.SetupTimestamp(); !ts.IsZero() {
+		if ts, _ := c.deps.MetadataManager.SetupTimestamp(); !ts.IsZero() {
 			return setupop.NewSetupLakeFSConflict().
 				WithPayload(&models.Error{
 					Message: "lakeFS already initialized",
@@ -269,7 +269,7 @@ func (c *Controller) SetupLakeFSHandler() setupop.SetupLakeFSHandler {
 		}
 
 		// update setup completed timestamp
-		if err := c.deps.Meta.UpdateSetupTimestamp(time.Now()); err != nil {
+		if err := c.deps.MetadataManager.UpdateSetupTimestamp(time.Now()); err != nil {
 			c.deps.logger.WithError(err).Error("Failed the update setup timestamp")
 		}
 

api/handler.go

@@ -41,24 +41,24 @@ var (
 )
 
 type Handler struct {
-	meta         auth.MetadataManager
-	cataloger    catalog.Cataloger
-	blockStore   block.Adapter
-	authService  auth.Service
-	stats        stats.Collector
-	retention    retention.Service
-	migrator     db.Migrator
-	apiServer    *restapi.Server
-	handler      *http.ServeMux
-	dedupCleaner *dedup.Cleaner
-	logger       logging.Logger
+	metadataManager auth.MetadataManager
+	cataloger       catalog.Cataloger
+	blockStore      block.Adapter
+	authService     auth.Service
+	stats           stats.Collector
+	retention       retention.Service
+	migrator        db.Migrator
+	apiServer       *restapi.Server
+	handler         *http.ServeMux
+	dedupCleaner    *dedup.Cleaner
+	logger          logging.Logger
 }
 
 func NewHandler(
 	cataloger catalog.Cataloger,
 	blockStore block.Adapter,
 	authService auth.Service,
-	meta auth.MetadataManager,
+	metadataManager auth.MetadataManager,
 	stats stats.Collector,
 	retention retention.Service,
 	migrator db.Migrator,
@@ -67,15 +67,15 @@ func NewHandler(
 ) http.Handler {
 	logger.Info("initialized OpenAPI server")
 	s := &Handler{
-		cataloger:    cataloger,
-		blockStore:   blockStore,
-		authService:  authService,
-		meta:         meta,
-		stats:        stats,
-		retention:    retention,
-		migrator:     migrator,
-		dedupCleaner: dedupCleaner,
-		logger:       logger,
+		cataloger:       cataloger,
+		blockStore:      blockStore,
+		authService:     authService,
+		metadataManager: metadataManager,
+		stats:           stats,
+		retention:       retention,
+		migrator:        migrator,
+		dedupCleaner:    dedupCleaner,
+		logger:          logger,
 	}
 	s.buildAPI()
 	return s.handler
@@ -167,7 +167,7 @@ func (s *Handler) buildAPI() {
 	api.BasicAuthAuth = s.BasicAuth()
 	api.JwtTokenAuth = s.JwtTokenAuth()
 	// bind our handlers to the server
-	NewController(s.cataloger, s.authService, s.blockStore, s.stats, s.retention, s.dedupCleaner, s.meta, s.migrator, s.stats, s.logger).Configure(api)
+	NewController(s.cataloger, s.authService, s.blockStore, s.stats, s.retention, s.dedupCleaner, s.metadataManager, s.migrator, s.stats, s.logger).Configure(api)
 
 	// setup host/port
 	s.apiServer = restapi.NewServer(api)

api/handler_test.go

@@ -11,6 +11,7 @@ import (
 
 	dbparams "github.com/treeverse/lakefs/db/params"
 	"github.com/treeverse/lakefs/dedup"
+	"github.com/treeverse/lakefs/stats"
 
 	"github.com/go-openapi/runtime"
 	httptransport "github.com/go-openapi/runtime/client"
@@ -71,7 +72,7 @@ func createDefaultAdminUser(authService auth.Service, t *testing.T) *authmodel.C
 
 type mockCollector struct{}
 
-func (m *mockCollector) CollectMetadata(_ map[string]string) {}
+func (m *mockCollector) CollectMetadata(_ *stats.Metadata) {}
 
 func (m *mockCollector) CollectEvent(_, _ string) {}
 

block/s3/inventory.go

@@ -12,7 +12,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/s3"
 	"github.com/aws/aws-sdk-go/service/s3/s3iface"
 	"github.com/treeverse/lakefs/block"
-	inventorys3 "github.com/treeverse/lakefs/inventory/s3"
+	inventorys3 "github.com/treeverse/lakefs/cloud/aws/s3_inventory"
 	"github.com/treeverse/lakefs/logging"
 )
 

block/s3/inventory_iterator.go

@@ -7,8 +7,8 @@ import (
 	"time"
 
 	"github.com/treeverse/lakefs/block"
+	inventorys3 "github.com/treeverse/lakefs/cloud/aws/s3_inventory"
 	"github.com/treeverse/lakefs/cmdutils"
-	inventorys3 "github.com/treeverse/lakefs/inventory/s3"
 )
 
 var ErrInventoryNotSorted = errors.New("got unsorted s3 inventory")

block/s3/inventory_test.go

@@ -15,7 +15,7 @@ import (
 	"github.com/go-openapi/swag"
 	"github.com/treeverse/lakefs/block"
 	"github.com/treeverse/lakefs/block/s3"
-	inventorys3 "github.com/treeverse/lakefs/inventory/s3"
+	inventorys3 "github.com/treeverse/lakefs/cloud/aws/s3_inventory"
 	"github.com/treeverse/lakefs/logging"
 )
 

cloud/aws/metadata.go

@@ -0,0 +1,34 @@
+package aws
+
+import (
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/s3"
+	"github.com/aws/aws-sdk-go/service/sts"
+	"github.com/treeverse/lakefs/logging"
+)
+
+type MetadataProvider struct {
+	logger    logging.Logger
+	awsConfig *aws.Config
+}
+
+func NewMetadataProvider(logger logging.Logger, awsConfig *aws.Config) *MetadataProvider {
+	return &MetadataProvider{logger: logger, awsConfig: awsConfig}
+}
+
+func (m *MetadataProvider) GetMetadata() map[string]string {
+	sess, err := session.NewSession(m.awsConfig)
+	if err != nil {
+		m.logger.Errorf("%v: failed to create AWS session for BI", err)
+		return nil
+	}
+	sess.ClientConfig(s3.ServiceName)
+	stsClient := sts.New(sess)
+	identity, err := stsClient.GetCallerIdentity(&sts.GetCallerIdentityInput{})
+	if err != nil {
+		m.logger.Errorf("%v: failed to get AWS account ID for BI", err)
+		return nil
+	}
+	return map[string]string{"aws_account_id": *identity.Account}
+}

inventory/s3/orc_reader.go → cloud/aws/s3_inventory/orc_reader.go

@@ -1,4 +1,4 @@
-package s3
+package s3_inventory
 
 import (
 	"context"

inventory/s3/orc_utils.go → cloud/aws/s3_inventory/orc_utils.go

@@ -1,4 +1,4 @@
-package s3
+package s3_inventory
 
 import (
 	"context"

inventory/s3/parquet_reader.go → cloud/aws/s3_inventory/parquet_reader.go

@@ -1,4 +1,4 @@
-package s3
+package s3_inventory
 
 import "github.com/xitongsys/parquet-go/reader"
 

inventory/s3/reader.go → cloud/aws/s3_inventory/reader.go

@@ -1,4 +1,4 @@
-package s3
+package s3_inventory
 
 import (
 	"context"

inventory/s3/reader_test.go → cloud/aws/s3_inventory/reader_test.go

@@ -1,4 +1,4 @@
-package s3
+package s3_inventory
 
 import (
 	"context"

cloud/gcp/metadata.go

@@ -0,0 +1,23 @@
+package gcp
+
+import (
+	"cloud.google.com/go/compute/metadata"
+	"github.com/treeverse/lakefs/logging"
+)
+
+type MetadataProvider struct {
+	logger logging.Logger
+}
+
+func NewMetadataProvider(logger logging.Logger) *MetadataProvider {
+	return &MetadataProvider{logger: logger}
+}
+
+func (m *MetadataProvider) GetMetadata() map[string]string {
+	projectID, err := metadata.NumericProjectID()
+	if err != nil {
+		m.logger.Errorf("%v: failed to get Google numeric project ID from instance metadata", err)
+		return nil
+	}
+	return map[string]string{"google_numeric_project_id": projectID}
+}

cloud/metadata.go

@@ -0,0 +1,25 @@
+package cloud
+
+import (
+	"github.com/treeverse/lakefs/block/gs"
+	s3a "github.com/treeverse/lakefs/block/s3"
+	"github.com/treeverse/lakefs/cloud/aws"
+	"github.com/treeverse/lakefs/cloud/gcp"
+	"github.com/treeverse/lakefs/config"
+	"github.com/treeverse/lakefs/logging"
+)
+
+type MetadataProvider interface {
+	GetMetadata() map[string]string
+}
+
+func BuildMetadataProvider(logger logging.Logger, c *config.Config) MetadataProvider {
+	switch c.GetBlockstoreType() {
+	case gs.BlockstoreType:
+		return gcp.NewMetadataProvider(logger)
+	case s3a.BlockstoreType:
+		return aws.NewMetadataProvider(logger, c.GetAwsConfig())
+	default:
+		return nil
+	}
+}

cmd/lakefs/cmd/init.go

@@ -10,8 +10,10 @@ import (
 	"github.com/treeverse/lakefs/auth"
 	"github.com/treeverse/lakefs/auth/crypt"
 	"github.com/treeverse/lakefs/auth/model"
+	"github.com/treeverse/lakefs/cloud"
 	"github.com/treeverse/lakefs/config"
 	"github.com/treeverse/lakefs/db"
+	"github.com/treeverse/lakefs/logging"
 	"github.com/treeverse/lakefs/stats"
 )
 
@@ -38,14 +40,9 @@ var initCmd = &cobra.Command{
 			dbPool,
 			crypt.NewSecretStore(cfg.GetAuthEncryptionSecret()),
 			cfg.GetAuthCacheConfig())
-
-		metaManager := auth.NewDBMetadataManager(config.Version, dbPool)
-		metadata, err := metaManager.Write()
-		if err != nil {
-			fmt.Printf("failed to write initial setup metadata: %s\n", err)
-			os.Exit(1)
-		}
-
+		authMetadataManager := auth.NewDBMetadataManager(config.Version, dbPool)
+		cloudMetadataProvider := cloud.BuildMetadataProvider(logging.Default(), cfg)
+		metadata := stats.NewMetadata(logging.Default(), cfg, authMetadataManager, cloudMetadataProvider)
 		credentials, err := auth.SetupAdminUser(authService, &model.User{
 			CreatedAt: time.Now(),
 			Username:  userName,
@@ -56,8 +53,7 @@ var initCmd = &cobra.Command{
 		}
 
 		ctx, cancelFn := context.WithCancel(context.Background())
-		processID, bufferedCollectorArgs := cfg.GetStatsBufferedCollectorArgs()
-		stats := stats.NewBufferedCollector(metadata[auth.InstallationIDKeyName], processID, bufferedCollectorArgs...)
+		stats := stats.NewBufferedCollector(metadata.InstallationID, cfg)
 		go stats.Run(ctx)
 		stats.CollectMetadata(metadata)
 		stats.CollectEvent("global", "init")