From 3ea333fd527070de11171d7a1dcc25b42607d093 Mon Sep 17 00:00:00 2001
From: c9s <yoanlin93@gmail.com>
Date: Mon, 20 Nov 2023 16:14:09 +0800
Subject: [PATCH] bbgo: add DisableStartupBalanceQuery option

---
 pkg/bbgo/config.go   |  9 +++++++--
 pkg/bbgo/session.go  | 23 +++++++++++++++--------
 pkg/types/account.go | 20 +++++++++++++++++++-
 3 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/pkg/bbgo/config.go b/pkg/bbgo/config.go
index 1fac9e314..f892f2138 100644
--- a/pkg/bbgo/config.go
+++ b/pkg/bbgo/config.go
@@ -329,8 +329,13 @@ type ServiceConfig struct {
 type EnvironmentConfig struct {
 	DisableDefaultKLineSubscription bool `json:"disableDefaultKLineSubscription"`
 	DisableHistoryKLinePreload      bool `json:"disableHistoryKLinePreload"`
-	DisableSessionTradeBuffer       bool `json:"disableSessionTradeBuffer"`
-	MaxSessionTradeBufferSize       int  `json:"maxSessionTradeBufferSize"`
+
+	// DisableStartUpBalanceQuery disables the balance query in the startup process
+	// which initializes the session.Account with the QueryAccount method.
+	DisableStartupBalanceQuery bool `json:"disableStartupBalanceQuery"`
+
+	DisableSessionTradeBuffer bool `json:"disableSessionTradeBuffer"`
+	MaxSessionTradeBufferSize int  `json:"maxSessionTradeBufferSize"`
 }
 
 type Config struct {
diff --git a/pkg/bbgo/session.go b/pkg/bbgo/session.go
index 75012fb03..ba10b1e1e 100644
--- a/pkg/bbgo/session.go
+++ b/pkg/bbgo/session.go
@@ -256,15 +256,22 @@ func (session *ExchangeSession) Init(ctx context.Context, environ *Environment)
 			}
 		}
 
-		logger.Infof("querying account balances...")
-		account, err := session.Exchange.QueryAccount(ctx)
-		if err != nil {
-			return err
-		}
+		disableStartupBalanceQuery := environ.environmentConfig != nil && environ.environmentConfig.DisableStartupBalanceQuery
+		if disableStartupBalanceQuery {
+			session.accountMutex.Lock()
+			session.Account = types.NewAccount()
+			session.accountMutex.Unlock()
+		} else {
+			logger.Infof("querying account balances...")
+			account, err := session.Exchange.QueryAccount(ctx)
+			if err != nil {
+				return err
+			}
 
-		session.accountMutex.Lock()
-		session.Account = account
-		session.accountMutex.Unlock()
+			session.accountMutex.Lock()
+			session.Account = account
+			session.accountMutex.Unlock()
+		}
 
 		logger.Infof("account %s balances:\n%s", session.Name, account.Balances().String())
 
diff --git a/pkg/types/account.go b/pkg/types/account.go
index cd590f530..5cb4ca92d 100644
--- a/pkg/types/account.go
+++ b/pkg/types/account.go
@@ -103,8 +103,26 @@ type IsolatedMarginAccountInfo struct {
 
 func NewAccount() *Account {
 	return &Account{
-		balances: make(BalanceMap),
+		AccountType:        "spot",
+		FuturesInfo:        nil,
+		MarginInfo:         nil,
+		IsolatedMarginInfo: nil,
+		MarginLevel:        fixedpoint.Zero,
+		MarginTolerance:    fixedpoint.Zero,
+		BorrowEnabled:      false,
+		TransferEnabled:    false,
+		MarginRatio:        fixedpoint.Zero,
+		LiquidationPrice:   fixedpoint.Zero,
+		LiquidationRate:    fixedpoint.Zero,
+		MakerFeeRate:       fixedpoint.Zero,
+		TakerFeeRate:       fixedpoint.Zero,
+		TotalAccountValue:  fixedpoint.Zero,
+		CanDeposit:         false,
+		CanTrade:           false,
+		CanWithdraw:        false,
+		balances:           make(BalanceMap),
 	}
+
 }
 
 // Balances lock the balances and returned the copied balances