ynab-portfolio-monitor/providers/questrade/providerImpl.go

package questrade

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"strconv"
	"time"
)

type persistentData struct {
	QuestradeRefreshToken string `json:"questradeRefreshToken"` // Questrade API OAuth2 refresh token
}

type Provider struct {
	questradeAccountIDs []int           // Slice of Questrade account numbers this provider monitors
	ynabAccountIDs      []string        // Slice of YNAB account ID's this provider updates - index position maps with questradeAccountIDs
	data                *persistentData // Data stored on disk and loaded when program starts
	client              *client         // HTTP client for interacting with Questrade API
	lastRefresh         time.Time
}

func (p *Provider) Name() string {
	return "Questrade"
}

// Configures the provider for usage via environment variables and persistentData
// If an error is returned, the provider will not be used
func (p *Provider) Configure() error {
	var err error

	p.questradeAccountIDs = make([]int, 0)
	p.ynabAccountIDs = make([]string, 0)
	// Load environment variables in continous series with suffix starting at 0
	// Multiple accounts can be configured, (eg _1, _2)
	// As soon as the series is interrupted, we assume we're done
	for i := 0; true; i++ {
		questradeAccountIDString := os.Getenv(fmt.Sprintf("questrade_account_%d", i))
		ynabAccountID := os.Getenv(fmt.Sprintf("questrade_ynab_account_%d", i))
		if questradeAccountIDString == "" || ynabAccountID == "" {
			if i == 0 {
				return fmt.Errorf("this account provider is not configured")
			}
			break
		}
		questradeAccountID, err := strconv.Atoi(questradeAccountIDString)
		if err != nil {
			return fmt.Errorf("failed to convert environment variable questrade_account_%d with value of '%s' to integer: %v", i, questradeAccountIDString, err)
		}
		p.questradeAccountIDs = append(p.questradeAccountIDs, questradeAccountID)
		p.ynabAccountIDs = append(p.ynabAccountIDs, ynabAccountID)
	}

	// Load persistent data from disk - the OAuth2.0 refresh tokens are one time use
	p.data, err = loadPersistentData()
	if err != nil {
		return fmt.Errorf("failed to load questrade configuration: %v", err)
	}

	// Create new HTTP client and login to API - will error if login fails
	err = p.refresh()
	if err != nil {
		return fmt.Errorf("failed to refresh http client: %v", err)
	}

	return nil
}

// Returns slices of account balances and mapped YNAB account IDs, along with an error
func (p *Provider) GetBalances() ([]int, []string, error) {
	// Refresh credentials if past half way until expiration
	if p.lastRefresh.Add(time.Second * time.Duration(p.client.Credentials.ExpiresIn) / 2).Before(time.Now()) {
		err := p.refresh()
		if err != nil {
			return make([]int, 0), make([]string, 0), fmt.Errorf("failed to refresh http client: %v", err)
		}
	}

	// Gather account balances from Questrade API
	balances := make([]int, 0)
	for _, questradeAccountID := range p.questradeAccountIDs {
		balance, err := p.client.GetQuestradeAccountBalance(questradeAccountID)
		if err != nil {
			return balances, p.ynabAccountIDs, fmt.Errorf("failed to get questrade account balance for account ID '%d': %v", questradeAccountID, err)
		}
		balances = append(balances, balance)
	}

	return balances, p.ynabAccountIDs, nil
}

func (p *Provider) refresh() error {
	var err error
	// Create new HTTP client and login to API - will error if login fails
	p.client, err = newClient(p.data.QuestradeRefreshToken)
	if err != nil {
		return fmt.Errorf("failed to create new questrade client: %v", err)
	}
	p.lastRefresh = time.Now()

	// After logging in, we get a new refresh token - save it for next login
	p.data.QuestradeRefreshToken = p.client.Credentials.RefreshToken
	err = savePersistentData(p.data)
	if err != nil {
		return fmt.Errorf("failed to save persistent data: %v", err)
	}
	return nil

}

// Load persistent data from disk, if it fails it initializes using environment variables
func loadPersistentData() (*persistentData, error) {
	data := &persistentData{}

	f, err := os.Open("data/questrade-data.json")
	if errors.Is(err, os.ErrNotExist) {
		// handle the case where the file doesn't exist
		data.QuestradeRefreshToken = os.Getenv("questrade_refresh_token")
		return data, nil
	}
	defer f.Close()
	b, err := io.ReadAll(f)
	if err != nil {
		return nil, fmt.Errorf("failed to read file data/questrade-data.json: %v", err)
	}

	err = json.Unmarshal(b, data)
	if err != nil {
		return nil, fmt.Errorf("failed to unmarshal data/questrade-data.json to PersistentData struct: %v", err)
	}
	return data, nil
}

// Save persistent data to disk, this should be done any time the data changes to ensure it can be loaded on next run
func savePersistentData(data *persistentData) error {
	b, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("failed to marshal persistentData to bytes: %v", err)
	}
	err = os.WriteFile("data/questrade-data.json", b, 0644)
	if err != nil {
		return fmt.Errorf("failed to write file data/questrade-data.json: %v", err)
	}
	return nil
}