Skip to content

Commit

Permalink
Include active hibernation override info in deployment metadata (#1597)
Browse files Browse the repository at this point in the history 
* Include active hibernation override info in deployment metadata

* Add more tests for hibernation schedules
  • Loading branch information
@feluelle
feluelle committed Mar 18, 2024
1 parent e6f73d7 commit 2326208
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 95 deletions.
41 changes: 28 additions & 13 deletions cloud/deployment/inspect/inspect.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,17 +18,23 @@ import (
)

type deploymentMetadata struct {
DeploymentID *string `mapstructure:"deployment_id" yaml:"deployment_id" json:"deployment_id"`
WorkspaceID *string `mapstructure:"workspace_id" yaml:"workspace_id" json:"workspace_id"`
ClusterID *string `mapstructure:"cluster_id" yaml:"cluster_id" json:"cluster_id"`
ReleaseName *string `mapstructure:"release_name" yaml:"release_name" json:"release_name"`
AirflowVersion *string `mapstructure:"airflow_version" yaml:"airflow_version" json:"airflow_version"`
CurrentTag *string `mapstructure:"current_tag" yaml:"current_tag" json:"current_tag"`
Status *string `mapstructure:"status" yaml:"status" json:"status"`
CreatedAt *time.Time `mapstructure:"created_at" yaml:"created_at" json:"created_at"`
UpdatedAt *time.Time `mapstructure:"updated_at" yaml:"updated_at" json:"updated_at"`
DeploymentURL *string `mapstructure:"deployment_url" yaml:"deployment_url" json:"deployment_url"`
WebserverURL *string `mapstructure:"webserver_url" yaml:"webserver_url" json:"webserver_url"`
DeploymentID *string `mapstructure:"deployment_id" yaml:"deployment_id" json:"deployment_id"`
WorkspaceID *string `mapstructure:"workspace_id" yaml:"workspace_id" json:"workspace_id"`
ClusterID *string `mapstructure:"cluster_id" yaml:"cluster_id" json:"cluster_id"`
ReleaseName *string `mapstructure:"release_name" yaml:"release_name" json:"release_name"`
AirflowVersion *string `mapstructure:"airflow_version" yaml:"airflow_version" json:"airflow_version"`
CurrentTag *string `mapstructure:"current_tag" yaml:"current_tag" json:"current_tag"`
Status *string `mapstructure:"status" yaml:"status" json:"status"`
CreatedAt *time.Time `mapstructure:"created_at" yaml:"created_at" json:"created_at"`
UpdatedAt *time.Time `mapstructure:"updated_at" yaml:"updated_at" json:"updated_at"`
DeploymentURL *string `mapstructure:"deployment_url" yaml:"deployment_url" json:"deployment_url"`
WebserverURL *string `mapstructure:"webserver_url" yaml:"webserver_url" json:"webserver_url"`
HibernationOverride *HibernationOverride `mapstructure:"hibernation_override,omitempty" yaml:"hibernation_override,omitempty" json:"hibernation_override,omitempty"`
}

type HibernationOverride struct {
IsHibernating *bool `mapstructure:"is_hibernating,omitempty" yaml:"is_hibernating,omitempty" json:"is_hibernating,omitempty"`
OverrideUntil *time.Time `mapstructure:"override_until,omitempty" yaml:"override_until,omitempty" json:"override_until,omitempty"`
}

type deploymentConfig struct {
Expand Down Expand Up @@ -185,7 +191,7 @@ func getDeploymentInfo(coreDeployment astroplatformcore.Deployment) (map[string]
if deployment.IsDeploymentStandard(*coreDeployment.Type) {
clusterID = notApplicable
}
return map[string]interface{}{
metadata := map[string]interface{}{
"deployment_id": coreDeployment.Id,
"workspace_id": coreDeployment.WorkspaceId,
"cluster_id": clusterID,
Expand All @@ -197,7 +203,16 @@ func getDeploymentInfo(coreDeployment astroplatformcore.Deployment) (map[string]
"created_at": coreDeployment.CreatedAt,
"updated_at": coreDeployment.UpdatedAt,
"status": coreDeployment.Status,
}, nil
}
if coreDeployment.ScalingSpec != nil && coreDeployment.ScalingSpec.HibernationSpec != nil {
if override := coreDeployment.ScalingSpec.HibernationSpec.Override; override != nil && override.IsActive != nil && *override.IsActive {
metadata["hibernation_override"] = HibernationOverride{
IsHibernating: override.IsHibernating,
OverrideUntil: override.OverrideUntil,
}
}
}
return metadata, nil
}

func getDeploymentConfig(coreDeploymentPointer *astroplatformcore.Deployment, platformCoreClient astroplatformcore.CoreClient) (map[string]interface{}, error) {
Expand Down
143 changes: 61 additions & 82 deletions cloud/deployment/inspect/inspect_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,14 @@ var (
IsEnabled: true,
},
}
hibernationIsActive = true
hibernationIsHibernating = true
hibernationOverrideUntil = time.Now()
hibernationOverride = astroplatformcore.DeploymentHibernationOverride{
IsActive: &hibernationIsActive,
IsHibernating: &hibernationIsHibernating,
OverrideUntil: &hibernationOverrideUntil,
}
isDevelopmentMode = true
sourceDeployment = astroplatformcore.Deployment{
Id: deploymentID,
Expand Down Expand Up @@ -187,6 +195,7 @@ var (
IsDevelopmentMode: &isDevelopmentMode,
ScalingSpec: &astroplatformcore.DeploymentScalingSpec{
HibernationSpec: &astroplatformcore.DeploymentHibernationSpec{
Override: &hibernationOverride,
Schedules: &hibernationSchedules,
},
},
Expand Down Expand Up @@ -627,6 +636,7 @@ func TestFormatPrintableDeployment(t *testing.T) {
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["environment_variables"],
"hibernation_schedules": additional["hibernation_schedules"],
},
}
expectedDeployment := `deployment:
Expand Down Expand Up @@ -674,6 +684,9 @@ func TestFormatPrintableDeployment(t *testing.T) {
updated_at: 2022-11-17T13:25:55.275697-08:00
deployment_url: cloud.astronomer.io/test-ws-id/deployments/test-deployment-id/overview
webserver_url: some-url
hibernation_override:
is_hibernating: true
override_until: 2022-11-17T13:25:55.275697-08:00
alert_emails:
- email1
- email2
Expand Down Expand Up @@ -718,6 +731,7 @@ func TestFormatPrintableDeployment(t *testing.T) {
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["environment_variables"],
"hibernation_schedules": additional["hibernation_schedules"],
},
}
expectedDeployment := `deployment:
Expand Down Expand Up @@ -757,6 +771,11 @@ func TestFormatPrintableDeployment(t *testing.T) {
alert_emails:
- email1
- email2
hibernation_schedules:
- hibernate_at: 1 * * * *
wake_at: 2 * * * *
description: hibernation schedule 1
enabled: true
`
var orderedAndTaggedDeployment FormattedDeployment
actualPrintableDeployment, err := formatPrintableDeployment("", true, printableDeployment)
Expand Down Expand Up @@ -786,6 +805,8 @@ func TestFormatPrintableDeployment(t *testing.T) {
sourceDeployment2.CloudProvider = &provider
sourceDeployment2.ImageTag = "some-tag"
sourceDeployment2.Status = "UNHEALTHY"
overrideUntil := time.Date(2023, time.February, 1, 12, 0, 0, 0, time.UTC)
sourceDeployment2.ScalingSpec.HibernationSpec.Override.OverrideUntil = &overrideUntil

info, _ := getDeploymentInfo(sourceDeployment2)
config, err := getDeploymentConfig(&sourceDeployment2, mockPlatformCoreClient)
Expand All @@ -799,6 +820,7 @@ func TestFormatPrintableDeployment(t *testing.T) {
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["environment_variables"],
"hibernation_schedules": additional["hibernation_schedules"],
},
}
expectedDeployment := `deployment:
Expand Down Expand Up @@ -855,9 +877,17 @@ func TestFormatPrintableDeployment(t *testing.T) {
updated_at: 2023-02-01T12:00:00Z
deployment_url: cloud.astronomer.io/test-ws-id/deployments/test-deployment-id/overview
webserver_url: some-url
hibernation_override:
is_hibernating: true
override_until: 2023-02-01T12:00:00Z
alert_emails:
- email1
- email2
hibernation_schedules:
- hibernate_at: 1 * * * *
wake_at: 2 * * * *
description: hibernation schedule 1
enabled: true
`
var orderedAndTaggedDeployment FormattedDeployment
actualPrintableDeployment, err := formatPrintableDeployment("", false, printableDeployment)
Expand Down Expand Up @@ -886,6 +916,7 @@ func TestFormatPrintableDeployment(t *testing.T) {
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["environment_variables"],
"hibernation_schedules": additional["hibernation_schedules"],
},
}
expectedDeployment := `{
Expand Down Expand Up @@ -944,7 +975,11 @@ func TestFormatPrintableDeployment(t *testing.T) {
"created_at": "2022-11-17T12:26:45.362983-08:00",
"updated_at": "2022-11-17T12:26:45.362983-08:00",
"deployment_url": "cloud.astronomer.io/test-ws-id/deployments/test-deployment-id/overview",
"webserver_url": "some-url"
"webserver_url": "some-url",
"hibernation_override": {
"is_hibernating": true,
"override_until": "2022-11-17T12:26:45.362983-08:00"
}
},
"alert_emails": [
"email1",
Expand Down Expand Up @@ -988,6 +1023,7 @@ func TestFormatPrintableDeployment(t *testing.T) {
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["environment_variables"],
"hibernation_schedules": additional["hibernation_schedules"],
},
}

Expand Down Expand Up @@ -1023,6 +1059,14 @@ func TestFormatPrintableDeployment(t *testing.T) {
"alert_emails": [
"email1",
"email2"
],
"hibernation_schedules": [
{
"hibernate_at": "1 * * * *",
"wake_at": "2 * * * *",
"description": "hibernation schedule 1",
"enabled": true
}
]
}
}`
Expand Down Expand Up @@ -1082,136 +1126,71 @@ func TestGetSpecificField(t *testing.T) {
config, err := getDeploymentConfig(&sourceDeployment, mockPlatformCoreClient)
assert.NoError(t, err)
additional := getAdditionalNullableFields(&sourceDeployment, nodePools)
printableDeployment := map[string]interface{}{
"deployment": map[string]interface{}{
"metadata": info,
"configuration": config,
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["environment_variables"],
"hibernation_schedules": additional["hibernation_schedules"],
},
}
t.Run("returns a value if key is found in deployment.metadata", func(t *testing.T) {
requestedField := "metadata.workspace_id"
printableDeployment := map[string]interface{}{
"deployment": map[string]interface{}{
"metadata": info,
"configuration": config,
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["environment_variables"],
},
}
actual, err := getSpecificField(printableDeployment, requestedField)
assert.NoError(t, err)
assert.Equal(t, sourceDeployment.WorkspaceId, actual)
})
t.Run("returns a value if key is found in deployment.configuration", func(t *testing.T) {
requestedField := "configuration.scheduler_count"
printableDeployment := map[string]interface{}{
"deployment": map[string]interface{}{
"metadata": info,
"configuration": config,
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["environment_variables"],
},
}
actual, err := getSpecificField(printableDeployment, requestedField)
assert.NoError(t, err)
assert.Equal(t, sourceDeployment.SchedulerReplicas, actual)
})
t.Run("returns a value if key is alert_emails", func(t *testing.T) {
requestedField := "alert_emails"
printableDeployment := map[string]interface{}{
"deployment": map[string]interface{}{
"metadata": info,
"configuration": config,
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["environment_variables"],
},
}
actual, err := getSpecificField(printableDeployment, requestedField)
assert.NoError(t, err)
assert.Equal(t, sourceDeployment.ContactEmails, actual)
})
t.Run("returns a value if key is environment_variables", func(t *testing.T) {
requestedField := "environment_variables"
printableDeployment := map[string]interface{}{
"deployment": map[string]interface{}{
"metadata": info,
"configuration": config,
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["environment_variables"],
},
}
actual, err := getSpecificField(printableDeployment, requestedField)
assert.NoError(t, err)
assert.Equal(t, getVariablesMap(*sourceDeployment.EnvironmentVariables), actual)
})
t.Run("returns a value if key is worker_queues", func(t *testing.T) {
requestedField := "worker_queues"
printableDeployment := map[string]interface{}{
"deployment": map[string]interface{}{
"metadata": info,
"configuration": config,
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["environment_variables"],
},
}
actual, err := getSpecificField(printableDeployment, requestedField)
assert.NoError(t, err)
assert.Equal(t, getQMap(&sourceDeployment, nodePools), actual)
})
t.Run("returns a value if key is hibernation_schedules", func(t *testing.T) {
requestedField := "hibernation_schedules"
actual, err := getSpecificField(printableDeployment, requestedField)
assert.NoError(t, err)
assert.Equal(t, getHibernationSchedulesMap(*sourceDeployment.ScalingSpec.HibernationSpec.Schedules), actual)
})
t.Run("returns a value if key is metadata", func(t *testing.T) {
requestedField := "metadata"
printableDeployment := map[string]interface{}{
"deployment": map[string]interface{}{
"metadata": info,
"configuration": config,
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["environment_variables"],
},
}
actual, err := getSpecificField(printableDeployment, requestedField)
assert.NoError(t, err)
assert.Equal(t, info, actual)
})
t.Run("returns value regardless of upper or lower case key", func(t *testing.T) {
requestedField := "Configuration.Cluster_NAME"
printableDeployment := map[string]interface{}{
"deployment": map[string]interface{}{
"metadata": info,
"configuration": config,
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["environment_variables"],
},
}
actual, err := getSpecificField(printableDeployment, requestedField)
assert.NoError(t, err)
assert.Equal(t, *sourceDeployment.ClusterName, actual)
})
t.Run("returns error if no value is found", func(t *testing.T) {
printableDeployment := map[string]interface{}{
"deployment": map[string]interface{}{
"metadata": info,
"configuration": config,
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["astronomer_variables"],
},
}
requestedField := "does-not-exist"
actual, err := getSpecificField(printableDeployment, requestedField)
assert.ErrorContains(t, err, "requested key "+requestedField+" not found in deployment")
assert.Equal(t, nil, actual)
})
t.Run("returns error if incorrect field is requested", func(t *testing.T) {
printableDeployment := map[string]interface{}{
"deployment": map[string]interface{}{
"metadata": info,
"configuration": config,
"alert_emails": additional["alert_emails"],
"worker_queues": additional["worker_queues"],
"environment_variables": additional["astronomer_variables"],
},
}
requestedField := "configuration.does-not-exist"
actual, err := getSpecificField(printableDeployment, requestedField)
assert.ErrorIs(t, err, errKeyNotFound)
Expand Down

0 comments on commit 2326208

Please sign in to comment.