Skip to content

Commit a1cab63

Browse files
jordanstephensjordanstephens
authored
add additional heroku context for compose generation (#1373)
* add heroku pg info to compose generation context * add dyno size info to the compose generate context * add release tasks to heroku compose generation context
1 parent 64b7ee5 commit a1cab63

File tree

2 files changed

+165
-4
lines changed

2 files changed

+165
-4
lines changed

‎src/pkg/migrate/heroku.go‎

Lines changed: 99 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,12 @@ import (
1717
)
1818

1919
type HerokuApplicationInfo struct {
20-
Addons []HerokuAddon `json:"addons"`
21-
Dynos []HerokuDyno `json:"dynos"`
22-
ConfigVars HerokuConfigVars `json:"config_vars"`
20+
Addons []HerokuAddon `json:"addons"`
21+
Dynos []HerokuDyno `json:"dynos"`
22+
ConfigVars HerokuConfigVars `json:"config_vars"`
23+
PGInfo []PGInfo `json:"pg_info"`
24+
DynoSizes map[string]HerokuDynoSize `json:"dyno_sizes"`
25+
ReleaseTasks []HerokuReleaseTask `json:"release_tasks"`
2326
}
2427

2528
func collectHerokuApplicationInfo(ctx context.Context, client HerokuClientInterface, appName string) (HerokuApplicationInfo, error) {
@@ -37,15 +40,44 @@ func collectHerokuApplicationInfo(ctx context.Context, client HerokuClientInterf
3740
return HerokuApplicationInfo{}, fmt.Errorf("failed to list Heroku addons: %w", err)
3841
}
3942
applicationInfo.Addons = addons
40-
4143
term.Debugf("Addons for the selected application: %+v\n", addons)
4244

45+
for _, addon := range addons {
46+
if addon.AddonService.Name == "heroku-postgresql" {
47+
pgInfo, err := client.GetPGInfo(ctx, addon.ID)
48+
if err != nil {
49+
return HerokuApplicationInfo{}, fmt.Errorf("failed to get Postgres info for addon %s: %w", addon.Name, err)
50+
}
51+
applicationInfo.PGInfo = append(applicationInfo.PGInfo, pgInfo)
52+
}
53+
}
54+
55+
term.Debugf("Postgres info for the selected application: %+v\n", applicationInfo.PGInfo)
56+
dynoSizes := make(map[string]HerokuDynoSize)
57+
58+
// for each dyno, get the dyno size,
59+
for _, dyno := range dynos {
60+
dynoSize, err := client.GetDynoSize(ctx, dyno.Size)
61+
if err != nil {
62+
return HerokuApplicationInfo{}, fmt.Errorf("failed to get dyno size for dyno %s: %w", dyno.Name, err)
63+
}
64+
dynoSizes[dyno.Name] = dynoSize
65+
}
66+
67+
applicationInfo.DynoSizes = dynoSizes
68+
4369
configVars, err := client.ListConfigVars(ctx, appName)
4470
if err != nil {
4571
return HerokuApplicationInfo{}, fmt.Errorf("failed to list Heroku config vars: %w", err)
4672
}
4773
applicationInfo.ConfigVars = configVars
4874

75+
releaseTasks, err := client.GetReleaseTasks(ctx, appName)
76+
if err != nil {
77+
return HerokuApplicationInfo{}, fmt.Errorf("failed to get Heroku release tasks: %w", err)
78+
}
79+
applicationInfo.ReleaseTasks = releaseTasks
80+
4981
return applicationInfo, nil
5082
}
5183

@@ -75,7 +107,10 @@ type HerokuClientInterface interface {
75107
ListApps(ctx context.Context) ([]HerokuApplication, error)
76108
ListDynos(ctx context.Context, appName string) ([]HerokuDyno, error)
77109
ListAddons(ctx context.Context, appName string) ([]HerokuAddon, error)
110+
GetDynoSize(ctx context.Context, dynoSizeName string) (HerokuDynoSize, error)
111+
GetPGInfo(ctx context.Context, addonID string) (PGInfo, error)
78112
ListConfigVars(ctx context.Context, appName string) (HerokuConfigVars, error)
113+
GetReleaseTasks(ctx context.Context, appName string) ([]HerokuReleaseTask, error)
79114
}
80115

81116
// HerokuClient represents the Heroku API client
@@ -158,6 +193,37 @@ func (h *HerokuClient) ListConfigVars(ctx context.Context, appName string) (Hero
158193
return herokuGet[HerokuConfigVars](ctx, h, url)
159194
}
160195

196+
type HerokuFormation struct {
197+
Command string `json:"command"`
198+
Type string `json:"type"`
199+
Size string `json:"size"`
200+
}
201+
202+
type HerokuReleaseTask = HerokuFormation
203+
204+
func (h *HerokuClient) GetFormation(ctx context.Context, appName string) ([]HerokuFormation, error) {
205+
endpoint := fmt.Sprintf("/apps/%s/formation", appName)
206+
url := h.BaseURL + endpoint
207+
return herokuGet[[]HerokuFormation](ctx, h, url)
208+
}
209+
210+
func (h *HerokuClient) GetReleaseTasks(ctx context.Context, appName string) ([]HerokuReleaseTask, error) {
211+
formationList, err := h.GetFormation(ctx, appName)
212+
if err != nil {
213+
return nil, err
214+
}
215+
216+
releaseTasks := []HerokuReleaseTask{}
217+
218+
for _, formation := range formationList {
219+
if formation.Type == "release" {
220+
releaseTasks = append(releaseTasks, formation)
221+
}
222+
}
223+
224+
return releaseTasks, nil
225+
}
226+
161227
type HerokuDyno struct {
162228
Name string `json:"name"`
163229
Command string `json:"command"`
@@ -172,6 +238,35 @@ func (h *HerokuClient) ListDynos(ctx context.Context, appName string) ([]HerokuD
172238
return herokuGet[[]HerokuDyno](ctx, h, url)
173239
}
174240

241+
type HerokuDynoSize struct {
242+
Architecture string `json:"architecture"`
243+
Compute int `json:"compute"`
244+
PreciseDynoUnits float64 `json:"precise_dyno_units"`
245+
Memory float64 `json:"memory"`
246+
Name string `json:"name"`
247+
}
248+
249+
func (h *HerokuClient) GetDynoSize(ctx context.Context, dynoSizeName string) (HerokuDynoSize, error) {
250+
endpoint := "/dyno-sizes/" + dynoSizeName
251+
url := h.BaseURL + endpoint
252+
return herokuGet[HerokuDynoSize](ctx, h, url)
253+
}
254+
255+
type PGInfo struct {
256+
DatabaseName string `json:"database_name"`
257+
NumBytes int64 `json:"num_bytes"`
258+
Info []struct {
259+
Name string `json:"name"`
260+
Values []string `json:"values"`
261+
} `json:"info"`
262+
}
263+
264+
func (h *HerokuClient) GetPGInfo(ctx context.Context, addonID string) (PGInfo, error) {
265+
endpoint := "/client/v11/databases/" + addonID
266+
url := "https://postgres-api.heroku.com" + endpoint
267+
return herokuGet[PGInfo](ctx, h, url)
268+
}
269+
175270
func herokuGet[T any](ctx context.Context, h *HerokuClient, url string) (T, error) {
176271
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
177272
if err != nil {

‎src/pkg/migrate/migrate_test.go‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,15 @@ func (m *MockHerokuClient) ListDynos(ctx context.Context, appName string) ([]Her
7676
return dynos, args.Error(1)
7777
}
7878

79+
func (m *MockHerokuClient) GetDynoSize(ctx context.Context, dynoSizeName string) (HerokuDynoSize, error) {
80+
args := m.Called(ctx, dynoSizeName)
81+
dynoSize, ok := args.Get(0).(HerokuDynoSize)
82+
if !ok {
83+
return HerokuDynoSize{}, errors.New("failed to cast to HerokuDynoSize")
84+
}
85+
return dynoSize, args.Error(1)
86+
}
87+
7988
func (m *MockHerokuClient) ListAddons(ctx context.Context, appName string) ([]HerokuAddon, error) {
8089
args := m.Called(ctx, appName)
8190
addons, ok := args.Get(0).([]HerokuAddon)
@@ -85,6 +94,15 @@ func (m *MockHerokuClient) ListAddons(ctx context.Context, appName string) ([]He
8594
return addons, args.Error(1)
8695
}
8796

97+
func (m *MockHerokuClient) GetPGInfo(ctx context.Context, addonID string) (PGInfo, error) {
98+
args := m.Called(ctx, addonID)
99+
pgInfo, ok := args.Get(0).(PGInfo)
100+
if !ok {
101+
return PGInfo{}, errors.New("failed to cast to *PGInfo")
102+
}
103+
return pgInfo, args.Error(1)
104+
}
105+
88106
func (m *MockHerokuClient) ListConfigVars(ctx context.Context, appName string) (HerokuConfigVars, error) {
89107
args := m.Called(ctx, appName)
90108
configVars, ok := args.Get(0).(HerokuConfigVars)
@@ -94,6 +112,15 @@ func (m *MockHerokuClient) ListConfigVars(ctx context.Context, appName string) (
94112
return configVars, args.Error(1)
95113
}
96114

115+
func (m *MockHerokuClient) GetReleaseTasks(ctx context.Context, appName string) ([]HerokuReleaseTask, error) {
116+
args := m.Called(ctx, appName)
117+
releaseTasks, ok := args.Get(0).([]HerokuReleaseTask)
118+
if !ok {
119+
return nil, errors.New("failed to cast to []HerokuReleaseTask")
120+
}
121+
return releaseTasks, args.Error(1)
122+
}
123+
97124
func (m *MockHerokuClient) SetToken(token string) {
98125
m.Called(token)
99126
}
@@ -106,8 +133,11 @@ func TestInteractiveSetup(t *testing.T) {
106133
herokuToken string
107134
herokuApps []HerokuApplication
108135
herokuDynos []HerokuDyno
136+
herokuDynoSize HerokuDynoSize
109137
herokuAddons []HerokuAddon
138+
herokuPGInfo PGInfo
110139
herokuConfigVars HerokuConfigVars
140+
herokuReleaseTasks []HerokuReleaseTask
111141
composeResponse *defangv1.GenerateComposeResponse
112142
expectedComposeFileContents string
113143
composeError error
@@ -128,6 +158,7 @@ func TestInteractiveSetup(t *testing.T) {
128158
herokuDynos: []HerokuDyno{
129159
{Name: "web.1", Command: "npm start", Size: "Standard-1X", Type: "web"},
130160
},
161+
herokuDynoSize: HerokuDynoSize{Architecture: "amd64", Name: "Eco", Memory: 0.5, Compute: 1, PreciseDynoUnits: 0.28},
131162
herokuAddons: []HerokuAddon{
132163
{
133164
Name: "postgresql-addon-123",
@@ -145,6 +176,23 @@ func TestInteractiveSetup(t *testing.T) {
145176
State: "provisioned",
146177
},
147178
},
179+
herokuReleaseTasks: []HerokuReleaseTask{
180+
{
181+
Command: "rails db:migrate",
182+
Type: "release",
183+
Size: "Eco",
184+
},
185+
},
186+
herokuPGInfo: PGInfo{
187+
DatabaseName: "mydb",
188+
NumBytes: 12345,
189+
Info: []struct {
190+
Name string `json:"name"`
191+
Values []string `json:"values"`
192+
}{
193+
{Name: "PG Version", Values: []string{"17.4"}},
194+
},
195+
},
148196
herokuConfigVars: HerokuConfigVars{
149197
"NODE_ENV": "production",
150198
"PORT": "3000",
@@ -171,6 +219,7 @@ func TestInteractiveSetup(t *testing.T) {
171219
{Name: "web.1", Command: "node server.js", Size: "Standard-2X", Type: "web"},
172220
{Name: "web.2", Command: "node server.js", Size: "Standard-2X", Type: "web"},
173221
},
222+
herokuDynoSize: HerokuDynoSize{Architecture: "amd64", Name: "Eco", Memory: 0.5, Compute: 1, PreciseDynoUnits: 0.28},
174223
herokuAddons: []HerokuAddon{
175224
{
176225
Name: "redis-addon-456",
@@ -209,6 +258,7 @@ func TestInteractiveSetup(t *testing.T) {
209258
{Name: "failing-app", ID: "app-789"},
210259
},
211260
herokuDynos: []HerokuDyno{{Name: "web.1", Command: "python app.py", Type: "web", Size: "Standard-1X"}},
261+
herokuDynoSize: HerokuDynoSize{Architecture: "amd64", Name: "Eco", Memory: 0.5, Compute: 1, PreciseDynoUnits: 0.28},
212262
herokuAddons: []HerokuAddon{},
213263
herokuConfigVars: HerokuConfigVars{},
214264
composeError: errors.New("fabric service unavailable"),
@@ -226,6 +276,7 @@ func TestInteractiveSetup(t *testing.T) {
226276
{Name: "yaml-invalid-app", ID: "app-invalid"},
227277
},
228278
herokuDynos: []HerokuDyno{{Name: "web.1", Command: "python app.py", Type: "web", Size: "Standard-1X"}},
279+
herokuDynoSize: HerokuDynoSize{Architecture: "amd64", Name: "Eco", Memory: 0.5, Compute: 1, PreciseDynoUnits: 0.28},
229280
herokuAddons: []HerokuAddon{},
230281
herokuConfigVars: HerokuConfigVars{},
231282
composeResponse: &defangv1.GenerateComposeResponse{
@@ -280,8 +331,23 @@ func TestInteractiveSetup(t *testing.T) {
280331
mockHerokuClient.On("SetToken", tt.herokuToken).Once()
281332
mockHerokuClient.On("ListApps", mock.Anything).Return(tt.herokuApps, nil)
282333
mockHerokuClient.On("ListDynos", mock.Anything, mock.Anything).Return(tt.herokuDynos, nil)
334+
335+
// GetDynoSize is called once for each dyno
336+
for range tt.herokuDynos {
337+
mockHerokuClient.On("GetDynoSize", mock.Anything, mock.Anything).Return(tt.herokuDynoSize, nil).Once()
338+
}
339+
283340
mockHerokuClient.On("ListAddons", mock.Anything, mock.Anything).Return(tt.herokuAddons, nil)
341+
342+
// GetPGInfo is only called for Postgres addons
343+
for _, addon := range tt.herokuAddons {
344+
if addon.AddonService.Name == "heroku-postgresql" {
345+
mockHerokuClient.On("GetPGInfo", mock.Anything, addon.ID).Return(tt.herokuPGInfo, nil).Once()
346+
}
347+
}
348+
284349
mockHerokuClient.On("ListConfigVars", mock.Anything, mock.Anything).Return(tt.herokuConfigVars, nil)
350+
mockHerokuClient.On("GetReleaseTasks", mock.Anything, mock.Anything).Return(tt.herokuReleaseTasks, nil)
285351

286352
// Execute the function under test
287353
ctx := context.Background()

0 commit comments

Comments
 (0)