From 6aac2ecccedc83b6cff92dab4d9757bec2b813b4 Mon Sep 17 00:00:00 2001
From: Brian Flad <bflad417@gmail.com>
Date: Tue, 14 Dec 2021 11:13:03 -0500
Subject: [PATCH] data-source/external: Make Read function context aware

Reference: https://github.com/hashicorp/terraform-provider-external/issues/14
Reference: https://github.com/hashicorp/terraform-provider-external/issues/22

This will enable the data source to exit immediately when receiving an interupt signal from Terraform. This also prepares the data source for future enhancements, such as including Terraform Plugin SDK context with logging messages and including additional information in return diagnostics.
---
 internal/provider/data_source.go | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/internal/provider/data_source.go b/internal/provider/data_source.go
index a8cc1e08..5d8cef9c 100644
--- a/internal/provider/data_source.go
+++ b/internal/provider/data_source.go
@@ -2,11 +2,12 @@ package provider
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
-	"fmt"
 	"log"
 	"os/exec"
 
+	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 )
 
@@ -27,7 +28,7 @@ func dataSource() *schema.Resource {
 			"or external programs beyond standard shell utilities, so it is not recommended to use this data source " +
 			"within configurations that are applied within Terraform Enterprise.",
 
-		Read: dataSourceRead,
+		ReadContext: dataSourceRead,
 
 		Schema: map[string]*schema.Schema{
 			"program": {
@@ -72,7 +73,7 @@ func dataSource() *schema.Resource {
 	}
 }
 
-func dataSourceRead(d *schema.ResourceData, meta interface{}) error {
+func dataSourceRead(ctx context.Context, d *schema.ResourceData, meta interface{}) diag.Diagnostics {
 
 	programI := d.Get("program").([]interface{})
 	workingDir := d.Get("working_dir").(string)
@@ -81,7 +82,7 @@ func dataSourceRead(d *schema.ResourceData, meta interface{}) error {
 	// This would be a ValidateFunc if helper/schema allowed these
 	// to be applied to lists.
 	if err := validateProgramAttr(programI); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 
 	program := make([]string, len(programI))
@@ -89,7 +90,7 @@ func dataSourceRead(d *schema.ResourceData, meta interface{}) error {
 		program[i] = vI.(string)
 	}
 
-	cmd := exec.Command(program[0], program[1:]...)
+	cmd := exec.CommandContext(ctx, program[0], program[1:]...)
 
 	cmd.Dir = workingDir
 
@@ -97,7 +98,7 @@ func dataSourceRead(d *schema.ResourceData, meta interface{}) error {
 	if err != nil {
 		// Should never happen, since we know query will always be a map
 		// from string to string, as guaranteed by d.Get and our schema.
-		return err
+		return diag.FromErr(err)
 	}
 
 	cmd.Stdin = bytes.NewReader(queryJson)
@@ -107,18 +108,18 @@ func dataSourceRead(d *schema.ResourceData, meta interface{}) error {
 	if err != nil {
 		if exitErr, ok := err.(*exec.ExitError); ok {
 			if exitErr.Stderr != nil && len(exitErr.Stderr) > 0 {
-				return fmt.Errorf("failed to execute %q: %s", program[0], string(exitErr.Stderr))
+				return diag.Errorf("failed to execute %q: %s", program[0], string(exitErr.Stderr))
 			}
-			return fmt.Errorf("command %q failed with no error message", program[0])
+			return diag.Errorf("command %q failed with no error message", program[0])
 		} else {
-			return fmt.Errorf("failed to execute %q: %s", program[0], err)
+			return diag.Errorf("failed to execute %q: %s", program[0], err)
 		}
 	}
 
 	result := map[string]string{}
 	err = json.Unmarshal(resultJson, &result)
 	if err != nil {
-		return fmt.Errorf("command %q produced invalid JSON: %s", program[0], err)
+		return diag.Errorf("command %q produced invalid JSON: %s", program[0], err)
 	}
 
 	d.Set("result", result)