Add environment tag management (#196)

docs/data-sources/service.md

@@ -21,6 +21,7 @@ Service data source
 
 ### Optional
 
+- `environment_tag` (String) Environment tag for this service.
 - `vpc_id` (Number) VPC ID this service is linked to.
 
 ### Read-Only

docs/resources/service.md

@@ -38,6 +38,7 @@ resource "timescale_service" "read_replica" {
 
 - `connection_pooler_enabled` (Boolean) Set connection pooler status for this service.
 - `enable_ha_replica` (Boolean) Enable HA Replica
+- `environment_tag` (String) Set environment tag for this service.
 - `memory_gb` (Number) Memory GB
 - `milli_cpu` (Number) Milli CPU
 - `name` (String) Service Name is the configurable name assigned to this resource. If none is provided, a default will be generated by the provider.

internal/client/client.go

@@ -28,6 +28,8 @@ var (
 	ToggleServiceMutation string
 	//go:embed queries/toggle_connection_pooler.graphql
 	ToggleConnectionPoolerMutation string
+	//go:embed queries/set_env_tag.graphql
+	SetEnvironmentTagMutation string
 	//go:embed queries/get_all_services.graphql
 	GetAllServicesQuery string
 	//go:embed queries/get_service.graphql

internal/client/queries/create_service.graphql

@@ -1,6 +1,6 @@
 mutation CreateService($projectId: ID!, $name: String!, $type: Type!, $resourceConfig:
     ResourceConfig, $regionCode: String!, $vpcId: ID, $forkConfig: ForkConfig, 
-    $enableConnectionPooler: Boolean) {
+    $enableConnectionPooler: Boolean, $environmentTag:ServiceEnvironment) {
     createService(data:{
         projectId:$projectId,
         name:$name,
@@ -9,7 +9,8 @@ mutation CreateService($projectId: ID!, $name: String!, $type: Type!, $resourceC
         regionCode:$regionCode,
         forkConfig:$forkConfig,
         enableConnectionPooler: $enableConnectionPooler,
-        vpcId: $vpcId
+        vpcId: $vpcId,
+        environmentTag: $environmentTag
     }){
         initialPassword
         service {
@@ -42,6 +43,9 @@ mutation CreateService($projectId: ID!, $name: String!, $type: Type!, $resourceC
                 }
             }
             regionCode
+            metadata {
+                environment
+            }
         }
     }
 }

internal/client/queries/get_all_services.graphql

@@ -40,5 +40,8 @@ query GetAllServices($projectId: ID!) {
             serviceId
             isStandby
         }
+        metadata {
+            environment
+        }
     }
 }

internal/client/queries/get_service.graphql

@@ -43,5 +43,8 @@ query GetService($projectId: ID!, $serviceId: ID!) {
             serviceId
             isStandby
         }
+        metadata {
+            environment
+        }
     }
 }

internal/client/queries/set_env_tag.graphql

@@ -0,0 +1,7 @@
+mutation SetEnvironmentTag($projectId: ID!, $serviceId: ID!, $environment: ServiceEnvironment!) {
+    setServiceEnvironmentTag (data:{
+        serviceId: $serviceId,
+        projectId: $projectId,
+        environment: $environment
+    })
+}

internal/client/service.go

@@ -23,6 +23,7 @@ type Service struct {
 	ReplicaStatus string         `json:"replicaStatus"`
 	VPCEndpoint   *VPCEndpoint   `json:"vpcEndpoint"`
 	ForkSpec      *ForkSpec      `json:"forkedFromId"`
+	Metadata      *Metadata      `json:"metadata"`
 }
 
 type ServiceSpec struct {
@@ -49,6 +50,10 @@ type ResourceSpec struct {
 	} `json:"spec"`
 }
 
+type Metadata struct {
+	Environment string `json:"environment"`
+}
+
 type CreateServiceRequest struct {
 	Name     string
 	MilliCPU string
@@ -62,6 +67,7 @@ type CreateServiceRequest struct {
 	ForkConfig   *ForkConfig
 
 	EnableConnectionPooler bool
+	EnvironmentTag         string
 }
 
 type ForkConfig struct {
@@ -133,7 +139,9 @@ func (c *Client) CreateService(ctx context.Context, request CreateServiceRequest
 	if request.ForkConfig != nil {
 		variables["forkConfig"] = request.ForkConfig
 	}
-
+	if request.EnvironmentTag != "" {
+		variables["environmentTag"] = request.EnvironmentTag
+	}
 	req := map[string]interface{}{
 		"operationName": "CreateService",
 		"query":         CreateServiceMutation,
@@ -352,3 +360,27 @@ func (c *Client) ToggleConnectionPooler(ctx context.Context, serviceID string, e
 	}
 	return nil
 }
+
+func (c *Client) SetEnvironmentTag(ctx context.Context, serviceID, environment string) error {
+	tflog.Trace(ctx, "Client.SetEnvironmentTag")
+	req := map[string]interface{}{
+		"operationName": "SetEnvironmentTag",
+		"query":         SetEnvironmentTagMutation,
+		"variables": map[string]any{
+			"projectId":   c.projectID,
+			"serviceId":   serviceID,
+			"environment": environment,
+		},
+	}
+	var resp Response[any]
+	if err := c.do(ctx, req, &resp); err != nil {
+		return err
+	}
+	if len(resp.Errors) > 0 {
+		return resp.Errors[0]
+	}
+	if resp.Data == nil {
+		return errors.New("no response found")
+	}
+	return nil
+}

internal/provider/provider_test.go

@@ -165,13 +165,19 @@ type ServiceConfig struct {
 	VpcID             int64
 	ReadReplicaSource string
 	Pooler            bool
+	Environment       string
 }
 
 func (c *ServiceConfig) WithName(name string) *ServiceConfig {
 	c.Name = name
 	return c
 }
 
+func (c *ServiceConfig) WithEnvironment(name string) *ServiceConfig {
+	c.Environment = name
+	return c
+}
+
 func (c *ServiceConfig) WithSpec(milliCPU, memoryGB int64) *ServiceConfig {
 	c.MilliCPU = milliCPU
 	c.MemoryGB = memoryGB
@@ -218,6 +224,9 @@ func (c *ServiceConfig) String(t *testing.T) string {
 	if c.Pooler {
 		write("connection_pooler_enabled = %t \n", c.Pooler)
 	}
+	if c.Environment != "" {
+		write("environment_tag = %q \n", c.Environment)
+	}
 	if c.RegionCode != "" {
 		write("region_code = %q \n", c.RegionCode)
 	}

internal/provider/service_data_source.go

@@ -5,10 +5,12 @@ import (
 	"fmt"
 	"strconv"
 
+	"github.com/hashicorp/terraform-plugin-framework-validators/stringvalidator"
 	"github.com/hashicorp/terraform-plugin-framework/datasource"
 	"github.com/hashicorp/terraform-plugin-framework/datasource/schema"
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/path"
+	"github.com/hashicorp/terraform-plugin-framework/schema/validator"
 	"github.com/hashicorp/terraform-plugin-framework/types"
 	"github.com/hashicorp/terraform-plugin-log/tflog"
 
@@ -37,6 +39,8 @@ type ServiceDataSourceModel struct {
 	Resources  []ResourceModel `tfsdk:"resources"`
 	Created    types.String    `tfsdk:"created"`
 	VpcID      types.Int64     `tfsdk:"vpc_id"`
+
+	EnvironmentTag types.String `tfsdk:"environment_tag"`
 }
 
 type SpecModel struct {
@@ -153,6 +157,13 @@ func (d *ServiceDataSource) Schema(_ context.Context, _ datasource.SchemaRequest
 				Optional:            true,
 				Computed:            true,
 			},
+			"environment_tag": schema.StringAttribute{
+				MarkdownDescription: "Environment tag for this service.",
+				Description:         "Environment tag for this service.",
+				Optional:            true,
+				Computed:            true,
+				Validators:          []validator.String{stringvalidator.OneOf("DEV", "PROD")},
+			},
 		},
 	}
 }
@@ -234,5 +245,8 @@ func serviceToDataModel(diag diag.Diagnostics, s *tsClient.Service) ServiceDataS
 			},
 		})
 	}
+	if s.Metadata != nil {
+		serviceModel.EnvironmentTag = types.StringValue(s.Metadata.Environment)
+	}
 	return serviceModel
 }

internal/provider/service_resource.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/hashicorp/terraform-plugin-framework-timeouts/resource/timeouts"
 	"github.com/hashicorp/terraform-plugin-framework-validators/int64validator"
+	"github.com/hashicorp/terraform-plugin-framework-validators/stringvalidator"
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/path"
 	"github.com/hashicorp/terraform-plugin-framework/resource"
@@ -77,7 +78,8 @@ type serviceResourceModel struct {
 	ReadReplicaSource types.String   `tfsdk:"read_replica_source"`
 	VpcID             types.Int64    `tfsdk:"vpc_id"`
 
-	ConnectionPoolerEnabled types.Bool `tfsdk:"connection_pooler_enabled"`
+	ConnectionPoolerEnabled types.Bool   `tfsdk:"connection_pooler_enabled"`
+	EnvironmentTag          types.String `tfsdk:"environment_tag"`
 }
 
 func (r *ServiceResource) Metadata(ctx context.Context, req resource.MetadataRequest, resp *resource.MetadataResponse) {
@@ -194,6 +196,16 @@ The change has been taken into account but must still be propagated. You can run
 					boolplanmodifier.UseStateForUnknown(),
 				},
 			},
+			"environment_tag": schema.StringAttribute{
+				MarkdownDescription: "Set environment tag for this service.",
+				Description:         "Set environment tag for this service.",
+				Optional:            true,
+				Computed:            true,
+				PlanModifiers: []planmodifier.String{
+					stringplanmodifier.UseStateForUnknown(),
+				},
+				Validators: []validator.String{stringvalidator.OneOf("DEV", "PROD")},
+			},
 			"username": schema.StringAttribute{
 				Description:         "The Postgres user for this service",
 				MarkdownDescription: "The Postgres user for this service",
@@ -277,6 +289,7 @@ func (r *ServiceResource) Create(ctx context.Context, req resource.CreateRequest
 		RegionCode:             plan.RegionCode.ValueString(),
 		ReplicaCount:           strconv.FormatInt(replicaCount, 10),
 		EnableConnectionPooler: plan.ConnectionPoolerEnabled.ValueBool(),
+		EnvironmentTag:         plan.EnvironmentTag.ValueString(),
 	}
 	if !plan.VpcID.IsNull() {
 		request.VpcID = plan.VpcID.ValueInt64()
@@ -469,6 +482,12 @@ func (r *ServiceResource) Update(ctx context.Context, req resource.UpdateRequest
 			return
 		}
 	}
+	if plan.EnvironmentTag != state.EnvironmentTag {
+		if err := r.client.SetEnvironmentTag(ctx, serviceID, plan.EnvironmentTag.ValueString()); err != nil {
+			resp.Diagnostics.AddError("Failed to set environment tag", err.Error())
+			return
+		}
+	}
 	if !plan.PoolerHostname.IsUnknown() {
 		resp.Diagnostics.AddError(ErrUpdateService, "Do not support pooler hostname change")
 		return
@@ -614,6 +633,9 @@ func serviceToResource(diag diag.Diagnostics, s *tsClient.Service, state service
 		model.Hostname = types.StringValue(s.VPCEndpoint.Host)
 		model.Port = types.Int64Value(s.VPCEndpoint.Port)
 	}
+	if s.Metadata != nil {
+		model.EnvironmentTag = types.StringValue(s.Metadata.Environment)
+	}
 
 	return model
 }

internal/provider/service_resource_test.go

@@ -37,6 +37,7 @@ func TestServiceResource_Default_Success(t *testing.T) {
 					resource.TestCheckResourceAttr("timescale_service.resource", "region_code", "us-east-1"),
 					resource.TestCheckResourceAttr("timescale_service.resource", "enable_ha_replica", "false"),
 					resource.TestCheckResourceAttr("timescale_service.resource", "connection_pooler_enabled", "false"),
+					resource.TestCheckResourceAttr("timescale_service.resource", "environment_tag", "DEV"),
 					resource.TestCheckNoResourceAttr("timescale_service.resource", "vpc_id"),
 				),
 			},
@@ -55,6 +56,13 @@ func TestServiceResource_Default_Success(t *testing.T) {
 					resource.TestCheckResourceAttr("timescale_service.resource", "name", "service resource test update"),
 				),
 			},
+			// Update tag
+			{
+				Config: getServiceConfig(t, config.WithEnvironment("PROD")),
+				Check: resource.ComposeAggregateTestCheckFunc(
+					resource.TestCheckResourceAttr("timescale_service.resource", "environment_tag", "PROD"),
+				),
+			},
 			// Enable pooler
 			{
 				Config: getServiceConfig(t, config.WithPooler(true)),