chore: Cleanup helpers poc

pkg/acceptance/helpers/database_client.go

@@ -0,0 +1,85 @@
+package helpers
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
+	"github.com/stretchr/testify/require"
+)
+
+type DatabaseClient struct {
+	context *TestClientContext
+}
+
+func NewDatabaseClient(context *TestClientContext) *DatabaseClient {
+	return &DatabaseClient{
+		context: context,
+	}
+}
+
+func (d *DatabaseClient) client() sdk.Databases {
+	return d.context.client.Databases
+}
+
+func (d *DatabaseClient) CreateDatabase(t *testing.T) (*sdk.Database, func()) {
+	t.Helper()
+	return d.CreateDatabaseWithOptions(t, sdk.RandomAccountObjectIdentifier(), &sdk.CreateDatabaseOptions{})
+}
+
+func (d *DatabaseClient) CreateDatabaseWithName(t *testing.T, name string) (*sdk.Database, func()) {
+	t.Helper()
+	return d.CreateDatabaseWithOptions(t, sdk.NewAccountObjectIdentifier(name), &sdk.CreateDatabaseOptions{})
+}
+
+func (d *DatabaseClient) CreateDatabaseWithOptions(t *testing.T, id sdk.AccountObjectIdentifier, opts *sdk.CreateDatabaseOptions) (*sdk.Database, func()) {
+	t.Helper()
+	ctx := context.Background()
+	err := d.client().Create(ctx, id, opts)
+	require.NoError(t, err)
+	database, err := d.client().ShowByID(ctx, id)
+	require.NoError(t, err)
+	return database, d.DropDatabaseFunc(t, id)
+}
+
+func (d *DatabaseClient) DropDatabaseFunc(t *testing.T, id sdk.AccountObjectIdentifier) func() {
+	t.Helper()
+	ctx := context.Background()
+
+	return func() {
+		err := d.client().Drop(ctx, id, &sdk.DropDatabaseOptions{IfExists: sdk.Bool(true)})
+		require.NoError(t, err)
+		err = d.context.client.Sessions.UseSchema(ctx, sdk.NewDatabaseObjectIdentifier(d.context.database, d.context.schema))
+		require.NoError(t, err)
+	}
+}
+
+func (d *DatabaseClient) CreateSecondaryDatabaseWithOptions(t *testing.T, id sdk.AccountObjectIdentifier, externalId sdk.ExternalObjectIdentifier, opts *sdk.CreateSecondaryDatabaseOptions) (*sdk.Database, func()) {
+	t.Helper()
+	ctx := context.Background()
+
+	// TODO [926148]: make this wait better with tests stabilization
+	// waiting because sometimes creating secondary db right after primary creation resulted in error
+	time.Sleep(1 * time.Second)
+
+	err := d.client().CreateSecondary(ctx, id, externalId, opts)
+	require.NoError(t, err)
+
+	// TODO [926148]: make this wait better with tests stabilization
+	// waiting because sometimes secondary database is not shown as SHOW REPLICATION DATABASES results right after creation
+	time.Sleep(1 * time.Second)
+
+	database, err := d.client().ShowByID(ctx, id)
+	require.NoError(t, err)
+	return database, func() {
+		err := d.client().Drop(ctx, id, nil)
+		require.NoError(t, err)
+
+		// TODO [926148]: make this wait better with tests stabilization
+		// waiting because sometimes dropping primary db right after dropping the secondary resulted in error
+		time.Sleep(1 * time.Second)
+		err = d.context.client.Sessions.UseSchema(ctx, sdk.NewDatabaseObjectIdentifier(d.context.database, d.context.schema))
+		require.NoError(t, err)
+	}
+}

pkg/acceptance/helpers/test_client.go

@@ -0,0 +1,29 @@
+package helpers
+
+import "github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
+
+type TestClient struct {
+	context *TestClientContext
+
+	Database *DatabaseClient
+}
+
+func NewTestClient(c *sdk.Client, database string, schema string, warehouse string) *TestClient {
+	context := &TestClientContext{
+		client:    c,
+		database:  database,
+		schema:    schema,
+		warehouse: warehouse,
+	}
+	return &TestClient{
+		context:  context,
+		Database: NewDatabaseClient(context),
+	}
+}
+
+type TestClientContext struct {
+	client    *sdk.Client
+	database  string
+	schema    string
+	warehouse string
+}

pkg/acceptance/testing.go

@@ -8,6 +8,7 @@ import (
 	"sync"
 	"testing"
 
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/helpers"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/testprofiles"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
@@ -72,12 +73,17 @@ func init() {
 		log.Panicf("Cannot instantiate new secondary client, err: %v", err)
 	}
 	atc.secondaryClient = secondaryClient
+
+	atc.testClient = helpers.NewTestClient(client, TestDatabaseName, TestSchemaName, TestWarehouseName)
+	atc.secondaryTestClient = helpers.NewTestClient(secondaryClient, TestDatabaseName, TestSchemaName, TestWarehouseName)
 }
 
 type acceptanceTestContext struct {
-	config          *gosnowflake.Config
-	client          *sdk.Client
-	secondaryClient *sdk.Client
+	config              *gosnowflake.Config
+	client              *sdk.Client
+	secondaryClient     *sdk.Client
+	testClient          *helpers.TestClient
+	secondaryTestClient *helpers.TestClient
 }
 
 var TestAccProtoV6ProviderFactories = map[string]func() (tfprotov6.ProviderServer, error){
@@ -164,3 +170,11 @@ func DefaultConfig(t *testing.T) *gosnowflake.Config {
 	t.Helper()
 	return atc.config
 }
+
+func TestClient() *helpers.TestClient {
+	return atc.testClient
+}
+
+func SecondaryTestClient() *helpers.TestClient {
+	return atc.secondaryTestClient
+}

pkg/resources/database_acceptance_test.go

@@ -109,7 +109,7 @@ func TestAcc_Database(t *testing.T) {
 }
 
 func TestAcc_DatabaseRemovedOutsideOfTerraform(t *testing.T) {
-	id := generateUnsafeExecuteTestDatabaseName(t)
+	name := generateUnsafeExecuteTestDatabaseName(t)
 
 	resource.Test(t, resource.TestCase{
 		ProtoV6ProviderFactories: acc.TestAccProtoV6ProviderFactories,
@@ -122,30 +122,30 @@ func TestAcc_DatabaseRemovedOutsideOfTerraform(t *testing.T) {
 			{
 				ConfigDirectory: config.TestNameDirectory(),
 				ConfigVariables: map[string]config.Variable{
-					"db": config.StringVariable(id),
+					"db": config.StringVariable(name),
 				},
 				ConfigPlanChecks: resource.ConfigPlanChecks{
 					PreApply: []plancheck.PlanCheck{plancheck.ExpectNonEmptyPlan()},
 				},
 				Check: resource.ComposeTestCheckFunc(
-					resource.TestCheckResourceAttr("snowflake_database.db", "name", id),
+					resource.TestCheckResourceAttr("snowflake_database.db", "name", name),
 					resource.TestCheckResourceAttr("snowflake_database.db", "comment", "test comment"),
-					testAccCheckDatabaseExistence(t, id, true),
+					testAccCheckDatabaseExistence(t, name, true),
 				),
 			},
 			{
-				PreConfig:       func() { dropDatabaseOutsideTerraform(t, id) },
+				PreConfig:       func() { acc.TestClient().Database.DropDatabaseFunc(t, sdk.NewAccountObjectIdentifier(name))() },
 				ConfigDirectory: config.TestNameDirectory(),
 				ConfigVariables: map[string]config.Variable{
-					"db": config.StringVariable(id),
+					"db": config.StringVariable(name),
 				},
 				ConfigPlanChecks: resource.ConfigPlanChecks{
 					PreApply: []plancheck.PlanCheck{plancheck.ExpectNonEmptyPlan()},
 				},
 				Check: resource.ComposeTestCheckFunc(
-					resource.TestCheckResourceAttr("snowflake_database.db", "name", id),
+					resource.TestCheckResourceAttr("snowflake_database.db", "name", name),
 					resource.TestCheckResourceAttr("snowflake_database.db", "comment", "test comment"),
-					testAccCheckDatabaseExistence(t, id, true),
+					testAccCheckDatabaseExistence(t, name, true),
 				),
 			},
 		},
@@ -376,16 +376,6 @@ resource "snowflake_database" "db" {
 	return fmt.Sprintf(s, prefix, secondaryAccountName)
 }
 
-func dropDatabaseOutsideTerraform(t *testing.T, id string) {
-	t.Helper()
-
-	client := acc.Client(t)
-	ctx := context.Background()
-
-	err := client.Databases.Drop(ctx, sdk.NewAccountObjectIdentifier(id), &sdk.DropDatabaseOptions{})
-	require.NoError(t, err)
-}
-
 func getSecondaryAccount(t *testing.T) string {
 	t.Helper()
 
@@ -477,25 +467,3 @@ func checkAccountAndDatabaseDataRetentionTime(id sdk.AccountObjectIdentifier, ex
 		return nil
 	}
 }
-
-func createDatabaseOutsideTerraform(t *testing.T, name string) func() {
-	t.Helper()
-	client := acc.Client(t)
-	ctx := context.Background()
-
-	if err := client.Databases.Create(ctx, sdk.NewAccountObjectIdentifier(name), new(sdk.CreateDatabaseOptions)); err != nil {
-		if err != nil {
-			t.Fatal(err)
-		}
-	}
-
-	return func() {
-		opts := new(sdk.DropDatabaseOptions)
-		opts.IfExists = sdk.Bool(true)
-		if err := client.Databases.Drop(ctx, sdk.NewAccountObjectIdentifier(name), opts); err != nil {
-			if err != nil {
-				t.Fatal(err)
-			}
-		}
-	}
-}

pkg/resources/grant_privileges_to_account_role_acceptance_test.go

@@ -1377,7 +1377,7 @@ func TestAcc_GrantPrivilegesToAccountRole_RemoveGrantedObjectOutsideTerraform(t
 		Steps: []resource.TestStep{
 			{
 				PreConfig: func() {
-					databaseCleanup = createDatabaseOutsideTerraform(t, databaseName)
+					_, databaseCleanup = acc.TestClient().Database.CreateDatabaseWithName(t, databaseName)
 					t.Cleanup(createAccountRoleOutsideTerraform(t, name))
 				},
 				ConfigDirectory: acc.ConfigurationDirectory("TestAcc_GrantPrivilegesToAccountRole/OnAccountObject"),
@@ -1420,7 +1420,8 @@ func TestAcc_GrantPrivilegesToAccountRole_RemoveAccountRoleOutsideTerraform(t *t
 		Steps: []resource.TestStep{
 			{
 				PreConfig: func() {
-					t.Cleanup(createDatabaseOutsideTerraform(t, databaseName))
+					_, dbCleanup := acc.TestClient().Database.CreateDatabaseWithName(t, databaseName)
+					t.Cleanup(dbCleanup)
 					roleCleanup = createAccountRoleOutsideTerraform(t, name)
 					t.Cleanup(roleCleanup)
 				},

pkg/resources/grant_privileges_to_database_role_acceptance_test.go

@@ -1053,7 +1053,7 @@ func TestAcc_GrantPrivilegesToDatabaseRole_RemoveGrantedObjectOutsideTerraform(t
 		Steps: []resource.TestStep{
 			{
 				PreConfig: func() {
-					databaseCleanup = createDatabaseOutsideTerraform(t, databaseName)
+					_, databaseCleanup = acc.TestClient().Database.CreateDatabaseWithName(t, databaseName)
 					t.Cleanup(databaseCleanup)
 					// no need to clean the role, because database will be dropped
 					createDatabaseRoleOutsideTerraform(t, databaseName, name)
@@ -1096,7 +1096,8 @@ func TestAcc_GrantPrivilegesToDatabaseRole_RemoveDatabaseRoleOutsideTerraform(t
 		Steps: []resource.TestStep{
 			{
 				PreConfig: func() {
-					t.Cleanup(createDatabaseOutsideTerraform(t, databaseName))
+					_, dbCleanup := acc.TestClient().Database.CreateDatabaseWithName(t, databaseName)
+					t.Cleanup(dbCleanup)
 					databaseRoleCleanup = createDatabaseRoleOutsideTerraform(t, databaseName, name)
 				},
 				ConfigDirectory: acc.ConfigurationDirectory("TestAcc_GrantPrivilegesToDatabaseRole/OnDatabase"),

pkg/resources/schema_acceptance_test.go

@@ -357,7 +357,8 @@ func TestAcc_Schema_RemoveSchemaOutsideOfTerraform(t *testing.T) {
 		"database_name": config.StringVariable(databaseName),
 	}
 
-	cleanupDatabase := createDatabaseOutsideTerraform(t, databaseName)
+	_, cleanupDatabase := acc.TestClient().Database.CreateDatabaseWithName(t, databaseName)
+	t.Cleanup(cleanupDatabase)
 
 	resource.Test(t, resource.TestCase{
 		ProtoV6ProviderFactories: acc.TestAccProtoV6ProviderFactories,

pkg/sdk/testint/client_unsafe_extensions_integration_test.go

@@ -30,11 +30,11 @@ func TestInt_Client_UnsafeQuery(t *testing.T) {
 	})
 
 	t.Run("test more results", func(t *testing.T) {
-		db1, db1Cleanup := createDatabase(t, client)
+		db1, db1Cleanup := testClientHelper().Database.CreateDatabase(t)
 		t.Cleanup(db1Cleanup)
-		db2, db2Cleanup := createDatabase(t, client)
+		db2, db2Cleanup := testClientHelper().Database.CreateDatabase(t)
 		t.Cleanup(db2Cleanup)
-		db3, db3Cleanup := createDatabase(t, client)
+		db3, db3Cleanup := testClientHelper().Database.CreateDatabase(t)
 		t.Cleanup(db3Cleanup)
 
 		sql := "SHOW DATABASES"

pkg/sdk/testint/context_functions_integration_test.go

@@ -66,7 +66,7 @@ func TestInt_CurrentSessionDetails(t *testing.T) {
 func TestInt_CurrentDatabase(t *testing.T) {
 	client := testClient(t)
 	ctx := testContext(t)
-	databaseTest, databaseCleanup := createDatabase(t, client)
+	databaseTest, databaseCleanup := testClientHelper().Database.CreateDatabase(t)
 	t.Cleanup(databaseCleanup)
 	err := client.Sessions.UseDatabase(ctx, databaseTest.ID())
 	require.NoError(t, err)
@@ -80,7 +80,7 @@ func TestInt_CurrentSchema(t *testing.T) {
 	ctx := testContext(t)
 
 	// new database and schema created on purpose
-	databaseTest, databaseCleanup := createDatabase(t, client)
+	databaseTest, databaseCleanup := testClientHelper().Database.CreateDatabase(t)
 	t.Cleanup(databaseCleanup)
 	schemaTest, schemaCleanup := createSchema(t, client, databaseTest)
 	t.Cleanup(schemaCleanup)

pkg/sdk/testint/database_role_integration_test.go

@@ -154,7 +154,7 @@ func TestInt_DatabaseRoles(t *testing.T) {
 	})
 
 	t.Run("alter database_role: rename to other database", func(t *testing.T) {
-		secondDatabase, secondDatabaseCleanup := createDatabase(t, client)
+		secondDatabase, secondDatabaseCleanup := testClientHelper().Database.CreateDatabase(t)
 		t.Cleanup(secondDatabaseCleanup)
 
 		name := random.String()