Cloudpods/backend/pkg/cloudcommon/db/quotas/models.go

// Copyright 2019 Yunion
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package quotas

import (
	"context"
	"database/sql"
	"reflect"
	"sort"

	"yunion.io/x/log"
	"yunion.io/x/pkg/errors"
	"yunion.io/x/pkg/util/rbacscope"
	"yunion.io/x/pkg/util/reflectutils"
	"yunion.io/x/sqlchemy"

	identityapi "yunion.io/x/onecloud/pkg/apis/identity"
	"yunion.io/x/onecloud/pkg/cloudcommon/db"
	"yunion.io/x/onecloud/pkg/mcclient"
	"yunion.io/x/onecloud/pkg/util/ctx"
)

type SQuotaBaseManager struct {
	db.SResourceBaseManager

	pendingStore IQuotaStore
	usageStore   IQuotaStore

	nonNegative bool

	scope rbacscope.TRbacScope
}

func NewQuotaBaseManager(model interface{}, scope rbacscope.TRbacScope, tableName string, pendingStore IQuotaStore, usageStore IQuotaStore, keyword, keywordPlural string) SQuotaBaseManager {
	pendingStore.SetVirtualObject(pendingStore)
	usageStore.SetVirtualObject(usageStore)
	return SQuotaBaseManager{
		SResourceBaseManager: db.NewResourceBaseManager(model, tableName, keyword, keywordPlural),
		pendingStore:         pendingStore,
		usageStore:           usageStore,
		nonNegative:          false,
		scope:                scope,
	}
}

func NewQuotaUsageManager(model interface{}, scope rbacscope.TRbacScope, tableName string, keyword, keywordPlural string) SQuotaBaseManager {
	return SQuotaBaseManager{
		SResourceBaseManager: db.NewResourceBaseManager(model, tableName, keyword, keywordPlural),
		nonNegative:          true,
		scope:                scope,
	}
}

type SQuotaBase struct {
	db.SResourceBase
}

func (manager *SQuotaBaseManager) GetIQuotaManager() IQuotaManager {
	return manager.GetIResourceModelManager().(IQuotaManager)
}

func (manager *SQuotaBaseManager) FetchIdNames(ctx context.Context, idMap map[string]map[string]string) (map[string]map[string]string, error) {
	return idMap, nil
}

func (manager *SQuotaBaseManager) getQuotaByKeys(ctx context.Context, keys IQuotaKeys, quota IQuota) error {
	q := manager.Query()

	fields := keys.Fields()
	values := keys.Values()
	for i := range fields {
		if len(values[i]) == 0 {
			q = q.IsNullOrEmpty(fields[i])
		} else {
			q = q.Equals(fields[i], values[i])
		}
	}

	err := q.First(quota)
	if manager.nonNegative {
		quota.ResetNegative()
	}
	if err != nil {
		return errors.Wrap(err, "q.Query")
	}
	return nil
}

func filterParentByKey(q *sqlchemy.SQuery, fieldName string, value string) *sqlchemy.SQuery {
	if len(value) > 0 {
		q = q.Filter(sqlchemy.OR(
			sqlchemy.IsNullOrEmpty(q.Field(fieldName)),
			sqlchemy.Equals(q.Field(fieldName), value),
		))
	} else {
		q = q.Filter(sqlchemy.IsNullOrEmpty(q.Field(fieldName)))
	}
	return q
}

func filterChildrenByKey(q *sqlchemy.SQuery, fieldName string, value string) *sqlchemy.SQuery {
	if len(value) > 0 {
		q = q.Equals(fieldName, value)
	}
	return q
}

func (manager *SQuotaBaseManager) getQuotasInternal(ctx context.Context, keys IQuotaKeys, isParent bool) ([]IQuota, error) {
	q := manager.Query()

	fields := keys.Fields()
	values := keys.Values()
	for i := range fields {
		if isParent {
			q = filterParentByKey(q, fields[i], values[i])
		} else {
			q = filterChildrenByKey(q, fields[i], values[i])
		}
	}

	rows, err := q.Rows()
	if err != nil {
		if errors.Cause(err) == sql.ErrNoRows {
			return nil, nil
		} else {
			return nil, errors.Wrap(err, "q.Rows")
		}
	}
	defer rows.Close()
	results := make([]IQuota, 0)
	for rows.Next() {
		r := manager.newQuota()
		err := q.Row2Struct(rows, r)
		if err != nil {
			return nil, errors.Wrap(err, "q.Row2Struct")
		}
		if manager.nonNegative {
			r.ResetNegative()
		}
		results = append(results, r)
	}
	sort.Sort(TQuotaList(results))
	return results, nil
}

func (manager *SQuotaBaseManager) setQuotaInternal(ctx context.Context, userCred mcclient.TokenCredential, quota IQuota) error {
	err := manager.TableSpec().InsertOrUpdate(ctx, quota)
	if err != nil {
		return errors.Wrap(err, "InsertOrUpdate")
	}
	if manager.nonNegative {
		quota.ResetNegative()
	}
	return nil
}

func (manager *SQuotaBaseManager) addQuotaInternal(ctx context.Context, userCred mcclient.TokenCredential, diff IQuota) error {
	keys := diff.GetKeys()
	quota := manager.newQuota()
	quota.SetKeys(keys)
	err := manager.getQuotaByKeys(ctx, keys, quota)
	if err != nil {
		if errors.Cause(err) == sql.ErrNoRows {
			// insert one
		} else {
			return errors.Wrap(err, "manager.getQuotaByKeys")
		}
	}
	quota.Add(diff)
	return manager.setQuotaInternal(ctx, userCred, quota)
}

func (manager *SQuotaBaseManager) subQuotaInternal(ctx context.Context, userCred mcclient.TokenCredential, diff IQuota) error {
	keys := diff.GetKeys()
	quota := manager.newQuota()
	quota.SetKeys(keys)
	err := manager.getQuotaByKeys(ctx, keys, quota)
	if err != nil {
		if errors.Cause(err) == sql.ErrNoRows {
			// insert one
		} else {
			return errors.Wrap(err, "manager.getQuotaByKeys")
		}
	}
	quota.Sub(diff)
	return manager.setQuotaInternal(ctx, userCred, quota)
}

func (manager *SQuotaBaseManager) deleteQuotaByKeys(ctx context.Context, userCred mcclient.TokenCredential, keys IQuotaKeys) error {
	quota := manager.newQuota()
	quota.SetKeys(keys)
	_, err := db.Update(quota.(db.IModel), func() error {
		return quota.(db.IModel).MarkDelete()
	})
	if err != nil {
		if errors.Cause(err) != sql.ErrNoRows {
			return errors.Wrap(err, "Delete")
		}
	}
	return nil
}

func (manager *SQuotaBaseManager) deleteAllQuotas(ctx context.Context, userCred mcclient.TokenCredential, keys IQuotaKeys) error {
	quotas, err := manager.getQuotasInternal(ctx, keys, false)
	if err != nil {
		if errors.Cause(err) != sql.ErrNoRows {
			return errors.Wrap(err, "manager.getQuotasInternal")
		} else {
			return nil
		}
	}
	for i := range quotas {
		err := manager.deleteQuotaByKeys(ctx, userCred, quotas[i].GetKeys())
		if err != nil {
			if errors.Cause(err) != sql.ErrNoRows {
				return errors.Wrapf(err, "manager.deleteQuotaByKeys %s", QuotaKeyString(quotas[i].GetKeys()))
			}
		}
	}
	return nil
}

func (manager *SQuotaBaseManager) InitializeData() error {
	q := manager.Query()
	quotaCnt, err := q.CountWithError()
	if err != nil {
		return errors.Wrap(err, "SQuotaManager.CountWithError")
	}
	if quotaCnt > 0 {
		// initlaized, quit
		return nil
	}

	log.Debugf("%s", q.String())

	metaQuota := newDBQuotaStore()

	tenants := make([]db.STenant, 0)
	err = db.TenantCacheManager.Query().All(&tenants)
	if err != nil && err != sql.ErrNoRows {
		return errors.Wrap(err, "Query")
	}

	for i := range tenants {
		obj := tenants[i]
		var scope rbacscope.TRbacScope
		var ownerId mcclient.IIdentityProvider
		if obj.DomainId == identityapi.KeystoneDomainRoot {
			// domain
			scope = rbacscope.ScopeDomain
			ownerId = &db.SOwnerId{
				DomainId: tenants[i].Id,
				Domain:   tenants[i].Name,
			}
		} else {
			// project
			scope = rbacscope.ScopeProject
			ownerId = &db.SOwnerId{
				DomainId:  tenants[i].DomainId,
				Domain:    tenants[i].Domain,
				ProjectId: tenants[i].Id,
				Project:   tenants[i].Name,
			}
		}

		quota := manager.newQuota()
		var baseKeys IQuotaKeys
		if manager.scope == rbacscope.ScopeDomain {
			baseKeys = OwnerIdDomainQuotaKeys(ownerId)
		} else {
			baseKeys = OwnerIdProjectQuotaKeys(scope, ownerId)
		}
		if !reflectutils.FillEmbededStructValue(reflect.Indirect(reflect.ValueOf(quota)), reflect.ValueOf(baseKeys)) {
			log.Fatalf("invalid quota??? fail to find SBaseQuotaKey")
		}
		err := metaQuota.GetQuota(ctx.CtxWithTime(), scope, ownerId, quota)
		if err != nil && err != sql.ErrNoRows {
			log.Errorf("metaQuota.GetQuota error %s for %s", err, ownerId)
			continue
		}
		if quota.IsEmpty() {
			quota.FetchSystemQuota()
		}
		err = manager.TableSpec().Insert(ctx.CtxWithTime(), quota)
		if err != nil {
			log.Errorf("%s insert error %s", manager.KeywordPlural(), err)
			continue
		} else {
			log.Infof("Insert %s", quota)
		}
	}

	return nil
}