From 6ecf48e1635c3d4118698da718dab02818e9cb41 Mon Sep 17 00:00:00 2001
From: Andrew Nester <andrew.nester@databricks.com>
Date: Fri, 27 Sep 2024 13:21:15 +0200
Subject: [PATCH] Add an error if state files grow bigger than the export limit

---
 bundle/deploy/state_push.go           | 11 +++++++++++
 bundle/deploy/terraform/state_push.go |  9 +++++++++
 2 files changed, 20 insertions(+)

diff --git a/bundle/deploy/state_push.go b/bundle/deploy/state_push.go
index 176a907c8..873739a76 100644
--- a/bundle/deploy/state_push.go
+++ b/bundle/deploy/state_push.go
@@ -10,6 +10,8 @@ import (
 	"github.com/databricks/cli/libs/log"
 )
 
+const MaxStateFileSize = 10 * 1024 * 1024 // 10MB
+
 type statePush struct {
 	filerFactory FilerFactory
 }
@@ -35,6 +37,15 @@ func (s *statePush) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostic
 	}
 	defer local.Close()
 
+	state, err := local.Stat()
+	if err != nil {
+		return diag.FromErr(err)
+	}
+
+	if state.Size() > MaxStateFileSize {
+		return diag.Errorf("Deployment state file size exceeds the maximum allowed size of %d bytes. Please reduce the number of resources in your bundle and/or split your bundle into multiple.", MaxStateFileSize)
+	}
+
 	log.Infof(ctx, "Writing local deployment state file to remote state directory")
 	err = f.Write(ctx, DeploymentStateFileName, local, filer.CreateParentDirectories, filer.OverwriteIfExists)
 	if err != nil {
diff --git a/bundle/deploy/terraform/state_push.go b/bundle/deploy/terraform/state_push.go
index 6cdde1371..5dfcde265 100644
--- a/bundle/deploy/terraform/state_push.go
+++ b/bundle/deploy/terraform/state_push.go
@@ -47,6 +47,15 @@ func (l *statePush) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostic
 	}
 	defer local.Close()
 
+	state, err := local.Stat()
+	if err != nil {
+		return diag.FromErr(err)
+	}
+
+	if state.Size() > deploy.MaxStateFileSize {
+		return diag.Errorf("Terraform state file size exceeds the maximum allowed size of %d bytes. Please reduce the number of resources in your bundle and/or split your bundle into multiple.", deploy.MaxStateFileSize)
+	}
+
 	// Upload state file from local cache directory to filer.
 	cmdio.LogString(ctx, "Updating deployment state...")
 	log.Infof(ctx, "Writing local state file to remote state directory")