Merge pull request #24987 from favoretti/favoretti/ste_select

Enhancement: add `select` argument to `azurerm_storage_table_entities`
hashicorp · Feb 26, 2024 · cef1f64 · cef1f64
2 parents cee315f + eb7e7f1
commit cef1f64
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 14 deletions.
diff --git a/internal/services/containerapps/container_app_resource_test.go b/internal/services/containerapps/container_app_resource_test.go
@@ -2207,29 +2207,29 @@ resource "azurerm_container_app" "test" {
 
 func (r ContainerAppResource) trafficBlockMoreThanOne() string {
 	return `
-    traffic_weight {
-      percentage      = 50
-    }
-    traffic_weight {
-      percentage      = 50
-    }
+traffic_weight {
+  percentage = 50
+}
+traffic_weight {
+  percentage = 50
+}
 `
 }
 
 func (r ContainerAppResource) trafficBlockLatestRevisionNotSet() string {
 	return `
-    traffic_weight {
-      percentage      = 100
-    }
+traffic_weight {
+  percentage = 100
+}
 `
 }
 
 func (r ContainerAppResource) trafficBlockRevisionSuffixSet() string {
 	return `
-    traffic_weight {
-      percentage      = 100
-	  latest_revision = true
-	  revision_suffix = "foo"
-    }
+traffic_weight {
+  percentage      = 100
+  latest_revision = true
+  revision_suffix = "foo"
+}
 `
 }
diff --git a/internal/services/storage/storage_table_entities_data_source.go b/internal/services/storage/storage_table_entities_data_source.go
@@ -26,6 +26,7 @@ type TableEntitiesDataSourceModel struct {
 	TableName          string                        `tfschema:"table_name"`
 	StorageAccountName string                        `tfschema:"storage_account_name"`
 	Filter             string                        `tfschema:"filter"`
+	Select             []string                      `tfschema:"select"`
 	Items              []TableEntitiyDataSourceModel `tfschema:"items"`
 }
 
@@ -54,6 +55,14 @@ func (k storageTableEntitiesDataSource) Arguments() map[string]*pluginsdk.Schema
 			Required:     true,
 			ValidateFunc: validation.StringIsNotEmpty,
 		},
+
+		"select": {
+			Type:     pluginsdk.TypeList,
+			Optional: true,
+			Elem: &pluginsdk.Schema{
+				Type: pluginsdk.TypeString,
+			},
+		},
 	}
 }
 
@@ -124,6 +133,11 @@ func (k storageTableEntitiesDataSource) Read() sdk.ResourceFunc {
 				MetaDataLevel: entities.MinimalMetaData,
 			}
 
+			if model.Select != nil {
+				model.Select = append(model.Select, "RowKey", "PartitionKey")
+				input.PropertyNamesToSelect = &model.Select
+			}
+
 			id := parse.NewStorageTableEntitiesId(model.StorageAccountName, storageClient.Environment.StorageEndpointSuffix, model.TableName, model.Filter)
 
 			result, err := client.Query(ctx, model.StorageAccountName, model.TableName, input)
@@ -134,6 +148,10 @@ func (k storageTableEntitiesDataSource) Read() sdk.ResourceFunc {
 			var flattenedEntities []TableEntitiyDataSourceModel
 			for _, entity := range result.Entities {
 				flattenedEntity := flattenEntityWithMetadata(entity)
+				if len(flattenedEntity.Properties) == 0 {
+					// if we use selector, we get empty objects back, skip them
+					continue
+				}
 				flattenedEntities = append(flattenedEntities, flattenedEntity)
 			}
 			model.Items = flattenedEntities

diff --git a/internal/services/storage/storage_table_entities_data_source_test.go b/internal/services/storage/storage_table_entities_data_source_test.go
@@ -26,6 +26,19 @@ func TestAccDataSourceStorageTableEntities_basic(t *testing.T) {
 	})
 }
 
+func TestAccDataSourceStorageTableEntities_withSelector(t *testing.T) {
+	data := acceptance.BuildTestData(t, "data.azurerm_storage_table_entities", "test")
+
+	data.DataSourceTest(t, []acceptance.TestStep{
+		{
+			Config: StorageTableEntitiesDataSource{}.basicWithDataSourceAndSelector(data),
+			Check: acceptance.ComposeTestCheckFunc(
+				check.That(data.ResourceName).Key("items.#").HasValue("1"),
+			),
+		},
+	})
+}
+
 func (d StorageTableEntitiesDataSource) basic(data acceptance.TestData) string {
 	return fmt.Sprintf(`
 provider "azurerm" {
@@ -76,6 +89,19 @@ resource "azurerm_storage_table_entity" "test2" {
     testkey = "testval2"
   }
 }
+
+resource "azurerm_storage_table_entity" "testselector" {
+  storage_account_name = azurerm_storage_account.test.name
+  table_name           = azurerm_storage_table.test.name
+
+  partition_key = "testselectorpartition"
+  row_key       = "testrow"
+
+  entity = {
+    testkey      = "testval"
+    testselector = "testselectorval"
+  }
+}
 `, data.RandomString, data.Locations.Primary, data.RandomString, data.RandomString)
 }
 
@@ -96,3 +122,23 @@ data "azurerm_storage_table_entities" "test" {
 }
 `, config)
 }
+
+func (d StorageTableEntitiesDataSource) basicWithDataSourceAndSelector(data acceptance.TestData) string {
+	config := d.basic(data)
+	return fmt.Sprintf(`
+%s
+
+data "azurerm_storage_table_entities" "test" {
+  table_name           = azurerm_storage_table_entity.test.table_name
+  storage_account_name = azurerm_storage_table_entity.test.storage_account_name
+  filter               = "PartitionKey eq 'testselectorpartition'"
+  select               = ["testselector"]
+
+  depends_on = [
+    azurerm_storage_table_entity.test,
+    azurerm_storage_table_entity.test2,
+    azurerm_storage_table_entity.testselector,
+  ]
+}
+`, config)
+}
diff --git a/website/docs/d/storage_table_entities.html.markdown b/website/docs/d/storage_table_entities.html.markdown
@@ -30,6 +30,8 @@ The following arguments are supported:
 
 * `filter` - The filter used to retrieve the entities.
 
+* `select` - (Optional) A list of properties to select from the returned Storage Table Entities.
+
 ## Attributes Reference
 
 * `id` - The ID of the storage table entity.