From a58e2b0d129e4bc6612543d620e26cd8b2324bbb Mon Sep 17 00:00:00 2001 From: Mark McDonald Date: Thu, 28 Sep 2023 17:19:39 -0400 Subject: [PATCH 1/4] fix: update aws_us_gov_east_1 region map and handling of accounts in orgs spanning multiple region groups" --- pkg/snowflake/current_account.go | 7 +++++-- pkg/snowflake/current_account_test.go | 14 +++++++++++--- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/pkg/snowflake/current_account.go b/pkg/snowflake/current_account.go index c3540b5a71..d004a7e296 100644 --- a/pkg/snowflake/current_account.go +++ b/pkg/snowflake/current_account.go @@ -13,7 +13,7 @@ var regionMapping = map[string]string{ "aws_us_west_2": "", // left black as this is the default "aws_us_east_2": "us-east-2.aws", "aws_us_east_1": "us-east-1", - "aws_us_east_1_gov": "us-east-1-gov.aws", + "aws_us_gov_east_1": "us-east-1-gov.aws", "aws_ca_central_1": "ca-central-1.aws", "aws_sa_east_1": "sa-east-1.aws", "aws_eu_west_1": "eu-west-1", @@ -48,12 +48,15 @@ var regionMapping = map[string]string{ "azure_australiaeast": "australia-east.azure", } +// SelectCurrentAccount returns the query that will return the current account, region_group, and region +// the CURRENT_REGION() function returns the format region_group.region (e.g. PUBLIC.AWS_US_WEST_2) only when organizations have accounts in multiple region groups. Otherwise, this function returns the snowflake region without the region_group. func SelectCurrentAccount() string { - return `SELECT CURRENT_ACCOUNT() AS "account", CURRENT_REGION() AS "region";` + return `SELECT CURRENT_ACCOUNT() as "account", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN LEFT(CURRENT_REGION(), POSITION('.' IN CURRENT_REGION()) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN RIGHT(CURRENT_REGION(), LENGTH(CURRENT_REGION()) - POSITION('.' IN CURRENT_REGION())) ELSE CURRENT_REGION() END AS "region";` } type CurrentAccount struct { Account string `db:"account"` + RegionGroup string `db:"region_group"` Region string `db:"region"` } diff --git a/pkg/snowflake/current_account_test.go b/pkg/snowflake/current_account_test.go index 0273157f21..7000c02fc5 100644 --- a/pkg/snowflake/current_account_test.go +++ b/pkg/snowflake/current_account_test.go @@ -1,6 +1,7 @@ package snowflake_test import ( + "regexp" "testing" sqlmock "github.com/DATA-DOG/go-sqlmock" @@ -11,12 +12,13 @@ import ( func TestCurrentAccountSelect(t *testing.T) { r := require.New(t) - r.Equal(`SELECT CURRENT_ACCOUNT() AS "account", CURRENT_REGION() AS "region";`, snowflake.SelectCurrentAccount()) + r.Equal(`SELECT CURRENT_ACCOUNT() as "account", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN LEFT(CURRENT_REGION(), POSITION('.' IN CURRENT_REGION()) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN RIGHT(CURRENT_REGION(), LENGTH(CURRENT_REGION()) - POSITION('.' IN CURRENT_REGION())) ELSE CURRENT_REGION() END AS "region";`, snowflake.SelectCurrentAccount()) } func TestCurrentAccountRead(t *testing.T) { type testCaseEntry struct { account string + region_group string region string url string } @@ -24,26 +26,31 @@ func TestCurrentAccountRead(t *testing.T) { testCases := map[string]testCaseEntry{ "aws oregon": { "ab1234", + "PUBLIC", "AWS_US_WEST_2", "https://ab1234.snowflakecomputing.com", }, "aws n virginia": { "cd5678", + "PUBLIC", "AWS_US_EAST_1", "https://cd5678.us-east-1.snowflakecomputing.com", }, "aws canada central": { "ef9012", + "PUBLIC", "AWS_CA_CENTRAL_1", "https://ef9012.ca-central-1.aws.snowflakecomputing.com", }, "gcp canada central": { "gh3456", + "PUBLIC", "gcp_us_central1", "https://gh3456.us-central1.gcp.snowflakecomputing.com", }, "azure washington": { "ij7890", + "PUBLIC", "azure_westus2", "https://ij7890.west-us-2.azure.snowflakecomputing.com", }, @@ -58,12 +65,13 @@ func TestCurrentAccountRead(t *testing.T) { defer mockDB.Close() sqlxDB := sqlx.NewDb(mockDB, "sqlmock") - rows := sqlmock.NewRows([]string{"account", "region"}).AddRow(tc.account, tc.region) - mock.ExpectQuery(`SELECT CURRENT_ACCOUNT\(\) AS "account", CURRENT_REGION\(\) AS "region";`).WillReturnRows(rows) + rows := sqlmock.NewRows([]string{"account", "region_group", "region"}).AddRow(tc.account, tc.region_group, tc.region) + mock.ExpectQuery(regexp.QuoteMeta(`SELECT CURRENT_ACCOUNT() as "account", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN LEFT(CURRENT_REGION(), POSITION('.' IN CURRENT_REGION()) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN RIGHT(CURRENT_REGION(), LENGTH(CURRENT_REGION()) - POSITION('.' IN CURRENT_REGION())) ELSE CURRENT_REGION() END AS "region";`)).WillReturnRows(rows) acc, err := snowflake.ReadCurrentAccount(sqlxDB.DB) r.NoError(err) r.Equal(tc.account, acc.Account) + r.Equal(tc.region_group, acc.RegionGroup) r.Equal(tc.region, acc.Region) url, err := acc.AccountURL() r.NoError(err) From 50f2495a6c334c48264c2749c18f1f9e7bf2dc47 Mon Sep 17 00:00:00 2001 From: Mark McDonald Date: Thu, 28 Sep 2023 17:34:44 -0400 Subject: [PATCH 2/4] style: formatting --- pkg/snowflake/current_account.go | 4 ++-- pkg/snowflake/current_account_test.go | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pkg/snowflake/current_account.go b/pkg/snowflake/current_account.go index d004a7e296..dec5028df8 100644 --- a/pkg/snowflake/current_account.go +++ b/pkg/snowflake/current_account.go @@ -55,9 +55,9 @@ func SelectCurrentAccount() string { } type CurrentAccount struct { - Account string `db:"account"` + Account string `db:"account"` RegionGroup string `db:"region_group"` - Region string `db:"region"` + Region string `db:"region"` } func ScanCurrentAccount(row *sqlx.Row) (*CurrentAccount, error) { diff --git a/pkg/snowflake/current_account_test.go b/pkg/snowflake/current_account_test.go index 7000c02fc5..3579f5758f 100644 --- a/pkg/snowflake/current_account_test.go +++ b/pkg/snowflake/current_account_test.go @@ -17,10 +17,10 @@ func TestCurrentAccountSelect(t *testing.T) { func TestCurrentAccountRead(t *testing.T) { type testCaseEntry struct { - account string + account string region_group string - region string - url string + region string + url string } testCases := map[string]testCaseEntry{ From aef30b263ea33af06877c65331ffba058f3163cb Mon Sep 17 00:00:00 2001 From: Mark McDonald Date: Thu, 28 Sep 2023 18:33:24 -0400 Subject: [PATCH 3/4] style: formatting --- pkg/snowflake/current_account.go | 2 +- pkg/snowflake/current_account_test.go | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pkg/snowflake/current_account.go b/pkg/snowflake/current_account.go index dec5028df8..6a9d2eadc1 100644 --- a/pkg/snowflake/current_account.go +++ b/pkg/snowflake/current_account.go @@ -51,7 +51,7 @@ var regionMapping = map[string]string{ // SelectCurrentAccount returns the query that will return the current account, region_group, and region // the CURRENT_REGION() function returns the format region_group.region (e.g. PUBLIC.AWS_US_WEST_2) only when organizations have accounts in multiple region groups. Otherwise, this function returns the snowflake region without the region_group. func SelectCurrentAccount() string { - return `SELECT CURRENT_ACCOUNT() as "account", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN LEFT(CURRENT_REGION(), POSITION('.' IN CURRENT_REGION()) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN RIGHT(CURRENT_REGION(), LENGTH(CURRENT_REGION()) - POSITION('.' IN CURRENT_REGION())) ELSE CURRENT_REGION() END AS "region";` + return `SELECT CURRENT_ACCOUNT() AS "account", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN LEFT(CURRENT_REGION(), POSITION('.' IN CURRENT_REGION()) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN RIGHT(CURRENT_REGION(), LENGTH(CURRENT_REGION()) - POSITION('.' IN CURRENT_REGION())) ELSE CURRENT_REGION() END AS "region";` } type CurrentAccount struct { diff --git a/pkg/snowflake/current_account_test.go b/pkg/snowflake/current_account_test.go index 3579f5758f..7c11df38e1 100644 --- a/pkg/snowflake/current_account_test.go +++ b/pkg/snowflake/current_account_test.go @@ -12,7 +12,7 @@ import ( func TestCurrentAccountSelect(t *testing.T) { r := require.New(t) - r.Equal(`SELECT CURRENT_ACCOUNT() as "account", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN LEFT(CURRENT_REGION(), POSITION('.' IN CURRENT_REGION()) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN RIGHT(CURRENT_REGION(), LENGTH(CURRENT_REGION()) - POSITION('.' IN CURRENT_REGION())) ELSE CURRENT_REGION() END AS "region";`, snowflake.SelectCurrentAccount()) + r.Equal(`SELECT CURRENT_ACCOUNT() AS "account", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN LEFT(CURRENT_REGION(), POSITION('.' IN CURRENT_REGION()) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN RIGHT(CURRENT_REGION(), LENGTH(CURRENT_REGION()) - POSITION('.' IN CURRENT_REGION())) ELSE CURRENT_REGION() END AS "region";`, snowflake.SelectCurrentAccount()) } func TestCurrentAccountRead(t *testing.T) { @@ -66,7 +66,7 @@ func TestCurrentAccountRead(t *testing.T) { sqlxDB := sqlx.NewDb(mockDB, "sqlmock") rows := sqlmock.NewRows([]string{"account", "region_group", "region"}).AddRow(tc.account, tc.region_group, tc.region) - mock.ExpectQuery(regexp.QuoteMeta(`SELECT CURRENT_ACCOUNT() as "account", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN LEFT(CURRENT_REGION(), POSITION('.' IN CURRENT_REGION()) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN RIGHT(CURRENT_REGION(), LENGTH(CURRENT_REGION()) - POSITION('.' IN CURRENT_REGION())) ELSE CURRENT_REGION() END AS "region";`)).WillReturnRows(rows) + mock.ExpectQuery(regexp.QuoteMeta(`SELECT CURRENT_ACCOUNT() AS "account", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN LEFT(CURRENT_REGION(), POSITION('.' IN CURRENT_REGION()) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN RIGHT(CURRENT_REGION(), LENGTH(CURRENT_REGION()) - POSITION('.' IN CURRENT_REGION())) ELSE CURRENT_REGION() END AS "region";`)).WillReturnRows(rows) acc, err := snowflake.ReadCurrentAccount(sqlxDB.DB) r.NoError(err) From b79a0289844069bedae8bca829f7251de96bc776 Mon Sep 17 00:00:00 2001 From: Mark McDonald Date: Thu, 28 Sep 2023 22:01:00 -0400 Subject: [PATCH 4/4] refactor: improved query --- pkg/snowflake/current_account.go | 2 +- pkg/snowflake/current_account_test.go | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pkg/snowflake/current_account.go b/pkg/snowflake/current_account.go index 6a9d2eadc1..a4044e136a 100644 --- a/pkg/snowflake/current_account.go +++ b/pkg/snowflake/current_account.go @@ -51,7 +51,7 @@ var regionMapping = map[string]string{ // SelectCurrentAccount returns the query that will return the current account, region_group, and region // the CURRENT_REGION() function returns the format region_group.region (e.g. PUBLIC.AWS_US_WEST_2) only when organizations have accounts in multiple region groups. Otherwise, this function returns the snowflake region without the region_group. func SelectCurrentAccount() string { - return `SELECT CURRENT_ACCOUNT() AS "account", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN LEFT(CURRENT_REGION(), POSITION('.' IN CURRENT_REGION()) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN RIGHT(CURRENT_REGION(), LENGTH(CURRENT_REGION()) - POSITION('.' IN CURRENT_REGION())) ELSE CURRENT_REGION() END AS "region";` + return `SELECT current_account AS "account", CASE WHEN CONTAINS(current_region, '.') THEN LEFT(current_region, POSITION('.' IN current_region) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(current_region, '.') THEN RIGHT(current_region, LENGTH(current_region) - POSITION('.' IN current_region)) ELSE current_region END AS "region" FROM (SELECT CURRENT_ACCOUNT() AS current_account, CURRENT_REGION() AS current_region);` } type CurrentAccount struct { diff --git a/pkg/snowflake/current_account_test.go b/pkg/snowflake/current_account_test.go index 7c11df38e1..e528bf95d9 100644 --- a/pkg/snowflake/current_account_test.go +++ b/pkg/snowflake/current_account_test.go @@ -12,7 +12,7 @@ import ( func TestCurrentAccountSelect(t *testing.T) { r := require.New(t) - r.Equal(`SELECT CURRENT_ACCOUNT() AS "account", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN LEFT(CURRENT_REGION(), POSITION('.' IN CURRENT_REGION()) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN RIGHT(CURRENT_REGION(), LENGTH(CURRENT_REGION()) - POSITION('.' IN CURRENT_REGION())) ELSE CURRENT_REGION() END AS "region";`, snowflake.SelectCurrentAccount()) + r.Equal(`SELECT current_account AS "account", CASE WHEN CONTAINS(current_region, '.') THEN LEFT(current_region, POSITION('.' IN current_region) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(current_region, '.') THEN RIGHT(current_region, LENGTH(current_region) - POSITION('.' IN current_region)) ELSE current_region END AS "region" FROM (SELECT CURRENT_ACCOUNT() AS current_account, CURRENT_REGION() AS current_region);`, snowflake.SelectCurrentAccount()) } func TestCurrentAccountRead(t *testing.T) { @@ -66,7 +66,7 @@ func TestCurrentAccountRead(t *testing.T) { sqlxDB := sqlx.NewDb(mockDB, "sqlmock") rows := sqlmock.NewRows([]string{"account", "region_group", "region"}).AddRow(tc.account, tc.region_group, tc.region) - mock.ExpectQuery(regexp.QuoteMeta(`SELECT CURRENT_ACCOUNT() AS "account", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN LEFT(CURRENT_REGION(), POSITION('.' IN CURRENT_REGION()) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(CURRENT_REGION(), '.') THEN RIGHT(CURRENT_REGION(), LENGTH(CURRENT_REGION()) - POSITION('.' IN CURRENT_REGION())) ELSE CURRENT_REGION() END AS "region";`)).WillReturnRows(rows) + mock.ExpectQuery(regexp.QuoteMeta(`SELECT current_account AS "account", CASE WHEN CONTAINS(current_region, '.') THEN LEFT(current_region, POSITION('.' IN current_region) - 1) ELSE 'PUBLIC' END AS "region_group", CASE WHEN CONTAINS(current_region, '.') THEN RIGHT(current_region, LENGTH(current_region) - POSITION('.' IN current_region)) ELSE current_region END AS "region" FROM (SELECT CURRENT_ACCOUNT() AS current_account, CURRENT_REGION() AS current_region);`)).WillReturnRows(rows) acc, err := snowflake.ReadCurrentAccount(sqlxDB.DB) r.NoError(err)