Usage caps (#412)

micro · Apr 1, 2022 · af07943 · af07943
1 parent 2b72a66
commit af07943
Show file tree

Hide file tree

Showing 22 changed files with 488 additions and 57 deletions.
diff --git a/app/handler/google.go b/app/handler/google.go
@@ -802,6 +802,10 @@ func (e *GoogleApp) DeleteData(ctx context.Context, request *adminpb.DeleteDataR
 	return nil
 }
 
+func (e *GoogleApp) Usage(ctx context.Context, request *adminpb.UsageRequest, response *adminpb.UsageResponse) error {
+	return nil
+}
+
 func (e *GoogleApp) List(ctx context.Context, req *pb.ListRequest, rsp *pb.ListResponse) error {
 	log.Info("Received App.List request")
 

diff --git a/cache/handler/cache.go b/cache/handler/cache.go
@@ -44,7 +44,6 @@ func (c *Cache) Get(ctx context.Context, req *pb.GetRequest, rsp *pb.GetResponse
 		rsp.Ttl = 0
 	}
 
-
 	return nil
 }
 
@@ -163,3 +162,7 @@ func (c *Cache) DeleteData(ctx context.Context, request *adminpb.DeleteDataReque
 	log.Infof("Deleted %d keys for %s", len(keys), request.TenantId)
 	return nil
 }
+
+func (c *Cache) Usage(ctx context.Context, request *adminpb.UsageRequest, response *adminpb.UsageResponse) error {
+	return nil
+}
diff --git a/comments/handler/comments.go b/comments/handler/comments.go
@@ -321,3 +321,7 @@ func (h *Comments) DeleteData(ctx context.Context, request *adminpb.DeleteDataRe
 	logger.Infof("Deleted %d keys for %s", len(keys), request.TenantId)
 	return nil
 }
+
+func (h *Comments) Usage(ctx context.Context, request *adminpb.UsageRequest, response *adminpb.UsageResponse) error {
+	return nil
+}
diff --git a/contact/handler/contact.go b/contact/handler/contact.go
@@ -168,3 +168,7 @@ func (c *contact) DeleteData(ctx context.Context, request *adminpb.DeleteDataReq
 	logger.Infof("Deleted %d keys for %s", len(keys), request.TenantId)
 	return nil
 }
+
+func (c *contact) Usage(ctx context.Context, request *adminpb.UsageRequest, response *adminpb.UsageResponse) error {
+	return nil
+}
diff --git a/db/handler/db.go b/db/handler/db.go
@@ -31,6 +31,7 @@ const renameTableStmt = `ALTER TABLE "%v" RENAME TO "%v"`
 
 var re = regexp.MustCompile("^[a-zA-Z0-9_]*$")
 var c = cache.New(5*time.Minute, 10*time.Minute)
+var usageCache = cache.New(30*time.Second, 10*time.Minute)
 
 type Record struct {
 	ID   string
@@ -91,7 +92,6 @@ func (e *Db) tableName(ctx context.Context, t string) (string, error) {
 	return tableName, nil
 }
 
-// Call is a single request handler called via client.Call or the generated client code
 func (e *Db) Create(ctx context.Context, req *db.CreateRequest, rsp *db.CreateResponse) error {
 	if len(req.Record.AsMap()) == 0 {
 		return errors.BadRequest("db.create", "missing record")
@@ -484,3 +484,57 @@ func (e *Db) DeleteData(ctx context.Context, request *adminpb.DeleteDataRequest,
 	logger.Infof("Deleted %d tables for %s", dropCount, request.TenantId)
 	return nil
 }
+
+func (e *Db) Usage(ctx context.Context, request *adminpb.UsageRequest, response *adminpb.UsageResponse) error {
+	method := "admin.Usage"
+	_, err := pauth.VerifyMicroAdmin(ctx, method)
+	if err != nil {
+		return err
+	}
+
+	if len(request.TenantId) < 10 { // deliberate length check so we don't grab all the things
+		return errors.BadRequest(method, "Missing tenant ID")
+	}
+
+	split := strings.Split(request.TenantId, "/")
+	tctx := tenant.NewContext(split[1], split[0], split[1])
+
+	tenantId := request.TenantId
+	tenantId = strings.Replace(strings.Replace(tenantId, "/", "_", -1), "-", "_", -1)
+
+	// Saving load on DB - do we have a cached response?
+	if ret, ok := usageCache.Get(tenantId); ok {
+		response.Usage = ret.(map[string]*adminpb.Usage)
+
+		return nil
+	}
+
+	db, err := e.GetDBConn(tctx)
+	if err != nil {
+		return err
+	}
+
+	var tables []string
+	if err := db.Table("information_schema.tables").Select("table_name").Where("table_schema = ?", "public").Find(&tables).Error; err != nil {
+		return err
+	}
+	var rowCount int64
+	for _, v := range tables {
+		if !strings.HasPrefix(v, tenantId) {
+			continue
+		}
+		var a int64
+		err = db.Table(v).Model(Record{}).Count(&a).Error
+		if err != nil {
+			return err
+		}
+		rowCount += a
+	}
+	response.Usage = map[string]*adminpb.Usage{
+		"Db.Create": &adminpb.Usage{Usage: rowCount, Units: "rows"},
+		// all other methods don't add rows so are not usage capped
+	}
+	usageCache.Set(tenantId, response.Usage, 0)
+	return nil
+
+}
diff --git a/file/handler/files.go b/file/handler/files.go
@@ -232,3 +232,7 @@ func (e *File) DeleteData(ctx context.Context, request *adminpb.DeleteDataReques
 
 	return nil
 }
+
+func (e *File) Usage(ctx context.Context, request *adminpb.UsageRequest, response *adminpb.UsageResponse) error {
+	return nil
+}
diff --git a/function/handler/google.go b/function/handler/google.go
@@ -1016,6 +1016,10 @@ func (e *GoogleFunction) DeleteData(ctx context.Context, request *adminpb.Delete
 	return nil
 }
 
+func (e *GoogleFunction) Usage(ctx context.Context, request *adminpb.UsageRequest, response *adminpb.UsageResponse) error {
+	return nil
+}
+
 var (
 	logsFuncMap = map[string]func(e *GoogleFunction, ctx context.Context, req *function.LogsRequest, rsp *function.LogsResponse) error{
 		"build": buildLogs, // TODO add runtime logs

diff --git a/github/handler/github.go b/github/handler/github.go
@@ -353,3 +353,7 @@ func (e *Github) DeleteData(ctx context.Context, request *adminpb.DeleteDataRequ
 
 	return nil
 }
+
+func (e *Github) Usage(ctx context.Context, request *adminpb.UsageRequest, response *adminpb.UsageResponse) error {
+	return nil
+}
diff --git a/image/handler/image.go b/image/handler/image.go
@@ -351,3 +351,7 @@ func (e *Image) DeleteData(ctx context.Context, request *adminpb.DeleteDataReque
 
 	return nil
 }
+
+func (e *Image) Usage(ctx context.Context, request *adminpb.UsageRequest, response *adminpb.UsageResponse) error {
+	return nil
+}
diff --git a/lists/handler/lists.go b/lists/handler/lists.go
@@ -321,3 +321,7 @@ func (h *Lists) DeleteData(ctx context.Context, request *adminpb.DeleteDataReque
 	logger.Infof("Deleted %d keys for %s", len(keys), request.TenantId)
 	return nil
 }
+
+func (h *Lists) Usage(ctx context.Context, request *adminpb.UsageRequest, response *adminpb.UsageResponse) error {
+	return nil
+}
diff --git a/location/handler/handler.go b/location/handler/handler.go
@@ -91,3 +91,7 @@ func (l *Location) DeleteData(ctx context.Context, request *adminpb.DeleteDataRe
 	logger.Infof("Deleted index for %s", request.TenantId)
 	return nil
 }
+
+func (l *Location) Usage(ctx context.Context, request *adminpb.UsageRequest, response *adminpb.UsageResponse) error {
+	return nil
+}
diff --git a/notes/handler/notes.go b/notes/handler/notes.go
@@ -321,3 +321,7 @@ func (h *Notes) DeleteData(ctx context.Context, request *adminpb.DeleteDataReque
 	logger.Infof("Deleted %d keys for %s", len(keys), request.TenantId)
 	return nil
 }
+
+func (h *Notes) Usage(ctx context.Context, request *adminpb.UsageRequest, response *adminpb.UsageResponse) error {
+	return nil
+}
diff --git a/otp/handler/otp.go b/otp/handler/otp.go
@@ -152,3 +152,7 @@ func (e *Otp) DeleteData(ctx context.Context, request *adminpb.DeleteDataRequest
 	logger.Infof("Deleted %d keys for %s", len(keys), request.TenantId)
 	return nil
 }
+
+func (e *Otp) Usage(ctx context.Context, request *adminpb.UsageRequest, response *adminpb.UsageResponse) error {
+	return nil
+}