Add env vars for Workspace, Region and Branch

internal/integration-tests/table_client_test.go

@@ -5,7 +5,6 @@ package integrationtests
 import (
 	"context"
 	"fmt"
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -15,106 +14,50 @@ import (
 )
 
 func Test_tableClient(t *testing.T) {
-	apiKey, found := os.LookupEnv("XATA_API_KEY")
-	if !found {
-		t.Skipf("%s not found in env vars", "XATA_API_KEY")
+	cfg, err := setupDatabase()
+	if err != nil {
+		t.Fatalf("unable to setup database: %v", err)
 	}
 
-	ctx := context.Background()
-	t.Run("should create/delete, get schema and columns of a table and add/delete column", func(t *testing.T) {
-		httpCli := retryablehttp.NewClient().StandardClient()
-
-		workspaceCli, err := xata.NewWorkspacesClient(
-			xata.WithAPIKey(apiKey),
-			xata.WithHTTPClient(httpCli),
-		)
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		testID := testIdentifier()
-
-		ws, err := workspaceCli.Create(
-			context.Background(),
-			&xata.WorkspaceMeta{Name: "ws_name_" + testID},
-		)
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		t.Cleanup(func() {
-			err := workspaceCli.Delete(ctx, ws.Id)
-			if err != nil {
-				t.Fatal(err)
-			}
-		})
-
-		databaseCli, err := xata.NewDatabasesClient(
-			xata.WithAPIKey(apiKey),
-			xata.WithHTTPClient(httpCli),
-		)
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		listRegionsResponse, err := databaseCli.GetRegionsWithWorkspaceID(ctx, ws.Id)
+	t.Cleanup(func() {
+		err = cleanup(cfg)
 		if err != nil {
-			t.Fatal(err)
-		}
-
-		regionID := listRegionsResponse.Regions[0].Id
-
-		db, err := databaseCli.Create(ctx, xata.CreateDatabaseRequest{
-			DatabaseName: "db_name_" + testID,
-			WorkspaceID:  xata.String(ws.Id),
-			Region:       &regionID,
-			UI:           &xata.UI{Color: xata.String("RED")},
-			BranchMetaData: &xata.BranchMetadata{
-				Repository: xata.String("github.com/my/repository"),
-				Branch:     xata.String("feature-branch"),
-				Stage:      xata.String("testing"),
-				Labels:     &[]string{"development"},
-			},
-		})
-		if err != nil {
-			t.Fatal(err)
+			t.Fatalf("unable to cleanup test setup: %v", err)
 		}
+	})
 
-		t.Cleanup(func() {
-			_, err = databaseCli.Delete(ctx, xata.DeleteDatabaseRequest{
-				WorkspaceID:  xata.String(ws.Id),
-				DatabaseName: db.DatabaseName,
-			})
-			if err != nil {
-				t.Fatal(err)
-			}
-		})
+	tableCli, err := xata.NewTableClient(
+		xata.WithAPIKey(cfg.apiKey),
+		xata.WithBaseURL(fmt.Sprintf(
+			"https://%s.%s.xata.sh",
+			cfg.wsID,
+			cfg.region,
+		)),
+		xata.WithHTTPClient(retryablehttp.NewClient().StandardClient()),
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
 
-		tableCli, err := xata.NewTableClient(
-			xata.WithAPIKey(apiKey),
-			xata.WithBaseURL(fmt.Sprintf(
-				"https://%s.%s.xata.sh",
-				ws.Id,
-				regionID,
-			)),
-			xata.WithHTTPClient(retryablehttp.NewClient().StandardClient()),
-		)
-		if err != nil {
-			t.Fatal(err)
-		}
+	ctx := context.Background()
+	t.Run("should create/delete, get schema and columns of a table and add/delete column", func(t *testing.T) {
+		testID := testIdentifier()
+		dbName := cfg.databaseName
+		tableName := "table_" + testID
 
 		createTableResponse, err := tableCli.Create(ctx, xata.TableRequest{
-			DatabaseName: xata.String(db.DatabaseName),
-			TableName:    "table-integration-test_" + testID,
+			DatabaseName: xata.String(dbName),
+			TableName:    tableName,
 		})
 		if err != nil {
 			t.Fatal(err)
 		}
+		assert.Equal(t, createTableResponse.TableName, tableName)
 
 		_, err = tableCli.AddColumn(ctx, xata.AddColumnRequest{
 			TableRequest: xata.TableRequest{
-				TableName:    createTableResponse.TableName,
-				DatabaseName: xata.String(db.DatabaseName),
+				TableName:    tableName,
+				DatabaseName: xata.String(dbName),
 			},
 			Column: &xata.Column{
 				Name:         stringColumn,
@@ -129,23 +72,23 @@ func Test_tableClient(t *testing.T) {
 		}
 
 		schema, err := tableCli.GetSchema(ctx, xata.TableRequest{
-			TableName:    createTableResponse.TableName,
-			DatabaseName: xata.String(db.DatabaseName),
+			TableName:    tableName,
+			DatabaseName: xata.String(dbName),
 		})
 		assert.NoError(t, err)
 		assert.Equal(t, stringColumn, schema.Columns[0].Name)
 
 		columns, err := tableCli.GetColumns(ctx, xata.TableRequest{
-			TableName:    createTableResponse.TableName,
-			DatabaseName: xata.String(db.DatabaseName),
+			TableName:    tableName,
+			DatabaseName: xata.String(dbName),
 		})
 		assert.NoError(t, err)
 		assert.Equal(t, stringColumn, columns.Columns[0].Name)
 
 		_, err = tableCli.DeleteColumn(ctx, xata.DeleteColumnRequest{
 			TableRequest: xata.TableRequest{
-				TableName:    createTableResponse.TableName,
-				DatabaseName: xata.String(db.DatabaseName),
+				TableName:    tableName,
+				DatabaseName: xata.String(dbName),
 			},
 			ColumnName: stringColumn,
 		})

xata/utils.go

@@ -161,11 +161,11 @@ func loadConfig(fieName string) (config, error) {
 	return cfg, nil
 }
 
-// getBranchName retrieves the branch name.
-// If not found, falls back to defaultBranchName
-func getBranchName() string {
-	if branchName, found := os.LookupEnv(branchNameEnvVar); found {
-		return branchName
+// Get value from env var with fallback to godotenv
+// return default value if not found
+func getEnvVar(name string, defaultValue string) string {
+	if val, found := os.LookupEnv(name); found {
+		return val
 	}
 
 	var myEnv map[string]string
@@ -177,11 +177,22 @@ func getBranchName() string {
 		}
 	}
 
-	if branchName, found := myEnv[branchNameEnvVar]; found {
-		return branchName
+	if val, found := myEnv[name]; found {
+		return val
 	}
+	return defaultValue
+}
 
-	return defaultBranchName
+// getBranchName retrieves the branch name.
+// If not found, falls back to defaultBranchName
+func getBranchName() string {
+	return getEnvVar(branchNameEnvVar, defaultBranchName)
+}
+
+// Get the region if the corresponding env var `XATA_REGION` is set
+// otherwise return the default region: us-east-1
+func getRegion() string {
+	return getEnvVar("XATA_REGION", defaultRegion)
 }
 
 // loadDatabaseConfig will return config with defaults if the error is not nil.
@@ -191,6 +202,21 @@ func loadDatabaseConfig() (databaseConfig, error) {
 		branchName:      defaultBranchName,
 		domainWorkspace: defaultDataPlaneDomain,
 	}
+
+	// Setup with env var
+	// XATA_WORKSPACE_ID to set the workspace Id
+	wsID := getEnvVar("XATA_WORKSPACE_ID", "")
+	if wsID != "" {
+		region := getRegion()
+		branch := getBranchName()
+		db := databaseConfig{
+			workspaceID: wsID,
+			region:      region,
+			branchName:  branch,
+		}
+		return db, nil
+	}
+
 	cfg, err := loadConfig(configFileName)
 	if err != nil {
 		return defaultDBConfig, err

xata/utils_test.go

@@ -26,6 +26,50 @@ func TestClientOptions_getAPIKey(t *testing.T) {
 	})
 }
 
+func Test_getBranchName(t *testing.T) {
+	// default state
+	t.Run("should be default branch name", func(t *testing.T) {
+		gotBranchName := getBranchName()
+		assert.Equal(t, gotBranchName, defaultBranchName)
+	})
+
+	setBranchName := "feature-042"
+	err := os.Setenv(branchNameEnvVar, setBranchName)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// from env var
+	t.Run("should be branch name from env var", func(t *testing.T) {
+		gotBranchName := getBranchName()
+		assert.Equal(t, gotBranchName, setBranchName)
+	})
+
+	t.Cleanup(func() { os.Unsetenv(branchNameEnvVar) })
+}
+
+func Test_getRegion(t *testing.T) {
+	// default state
+	t.Run("should be default region", func(t *testing.T) {
+		gotRegion := getRegion()
+		assert.Equal(t, gotRegion, defaultRegion)
+	})
+
+	setRegion := "eu-west-3"
+	err := os.Setenv("XATA_REGION", setRegion)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// from env var
+	t.Run("should be region from the env var", func(t *testing.T) {
+		gotRegion := getRegion()
+		assert.Equal(t, gotRegion, setRegion)
+	})
+
+	t.Cleanup(func() { os.Unsetenv("XATA_REGION") })
+}
+
 func Test_parseDatabaseURL(t *testing.T) {
 	tests := []struct {
 		name    string
@@ -95,6 +139,7 @@ func Test_parseDatabaseURL(t *testing.T) {
 }
 
 func Test_loadConfig(t *testing.T) {
+	// from .xatarc
 	t.Run("should read database URL", func(t *testing.T) {
 		// Create a temporary JSON file for testing
 		tempFile, err := os.CreateTemp("", "config_test.json")
@@ -128,3 +173,64 @@ func Test_loadConfig(t *testing.T) {
 		}
 	})
 }
+
+func Test_loadDatabaseConfig_with_envvars(t *testing.T) {
+	setWsId := "workspace-0lac00"
+	err := os.Setenv("XATA_WORKSPACE_ID", setWsId)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// test workspace id from env var
+	t.Run("load config from WORKSPACE_ID env var", func(t *testing.T) {
+		dbCfg, err := loadDatabaseConfig()
+		if err != nil {
+			t.Fatalf("Error loading config: %v", err)
+		}
+
+		if dbCfg.workspaceID != setWsId {
+			t.Fatalf("Expected Workspace ID: %s, got: %s", setWsId, dbCfg.workspaceID)
+		}
+		if dbCfg.branchName != defaultBranchName {
+			t.Fatalf("Expected branch name: %s, got: %s", defaultBranchName, dbCfg.branchName)
+		}
+		if dbCfg.region != defaultRegion {
+			t.Fatalf("Expected region: %s, got: %s", defaultRegion, dbCfg.region)
+		}
+	})
+
+	setBranch := "branch123"
+	err2 := os.Setenv(branchNameEnvVar, setBranch)
+	if err2 != nil {
+		t.Fatal(err2)
+	}
+	setRegion := "ap-southeast-16"
+	err3 := os.Setenv("XATA_REGION", setRegion)
+	if err3 != nil {
+		t.Fatal(err3)
+	}
+
+	// with branch and region env vars
+	t.Run("load config from XATA_WORKSPACE_ID, XATA_REGION and XATA_BRANCH env vars", func(t *testing.T) {
+		dbCfg, err := loadDatabaseConfig()
+		if err != nil {
+			t.Fatalf("Error loading config: %v", err)
+		}
+
+		if dbCfg.workspaceID != setWsId {
+			t.Fatalf("Expected Workspace ID: %s, got: %s", setWsId, dbCfg.workspaceID)
+		}
+		if dbCfg.branchName != setBranch {
+			t.Fatalf("Expected branch name: %s, got: %s", setBranch, dbCfg.branchName)
+		}
+		if dbCfg.region != setRegion {
+			t.Fatalf("Expected region: %s, got: %s", setRegion, dbCfg.region)
+		}
+	})
+
+	t.Cleanup(func() {
+		os.Unsetenv("XATA_WORKSPACE_ID")
+		os.Unsetenv("XATA_BRANCH")
+		os.Unsetenv("XATA_REGION")
+	})
+}