Merge pull request #2141 from a7i/athena-catalog

feat: athena configuration support for Catalog

pkg/cloud/aws/athenaconfiguration.go

@@ -12,6 +12,7 @@ type AthenaConfiguration struct {
 	Bucket     string     `json:"bucket"`
 	Region     string     `json:"region"`
 	Database   string     `json:"database"`
+	Catalog    string     `json:"catalog""`
 	Table      string     `json:"table"`
 	Workgroup  string     `json:"workgroup"`
 	Account    string     `json:"account"`
@@ -85,6 +86,10 @@ func (ac *AthenaConfiguration) Equals(config config.Config) bool {
 		return false
 	}
 
+	if ac.Catalog != thatConfig.Catalog {
+		return false
+	}
+
 	if ac.Table != thatConfig.Table {
 		return false
 	}
@@ -105,6 +110,7 @@ func (ac *AthenaConfiguration) Sanitize() config.Config {
 		Bucket:     ac.Bucket,
 		Region:     ac.Region,
 		Database:   ac.Database,
+		Catalog:    ac.Catalog,
 		Table:      ac.Table,
 		Workgroup:  ac.Workgroup,
 		Account:    ac.Account,
@@ -143,6 +149,12 @@ func (ac *AthenaConfiguration) UnmarshalJSON(b []byte) error {
 	}
 	ac.Database = database
 
+	catalog, err := config.GetInterfaceValue[string](fmap, "catalog")
+	if err != nil {
+		return fmt.Errorf("AthenaConfiguration: UnmarshalJSON: %s", err.Error())
+	}
+	ac.Catalog = catalog
+
 	table, err := config.GetInterfaceValue[string](fmap, "table")
 	if err != nil {
 		return fmt.Errorf("AthenaConfiguration: UnmarshalJSON: %s", err.Error())
@@ -204,6 +216,7 @@ func ConvertAwsAthenaInfoToConfig(aai AwsAthenaInfo) config.KeyedConfig {
 		config = &AthenaConfiguration{
 			Bucket:     aai.AthenaBucketName,
 			Region:     aai.AthenaRegion,
+			Catalog:    aai.AthenaCatalog,
 			Database:   aai.AthenaDatabase,
 			Table:      aai.AthenaTable,
 			Workgroup:  aai.AthenaWorkgroup,

pkg/cloud/aws/athenaconfiguration_test.go

@@ -19,6 +19,7 @@ func TestAthenaConfiguration_Validate(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -30,6 +31,19 @@ func TestAthenaConfiguration_Validate(t *testing.T) {
 			expected: nil,
 		},
 		"valid config service account": {
+			config: AthenaConfiguration{
+				Bucket:     "bucket",
+				Region:     "region",
+				Database:   "database",
+				Catalog:    "catalog",
+				Table:      "table",
+				Workgroup:  "workgroup",
+				Account:    "account",
+				Authorizer: &ServiceAccount{},
+			},
+			expected: nil,
+		},
+		"valid missing catalog": {
 			config: AthenaConfiguration{
 				Bucket:     "bucket",
 				Region:     "region",
@@ -46,6 +60,7 @@ func TestAthenaConfiguration_Validate(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -60,6 +75,7 @@ func TestAthenaConfiguration_Validate(t *testing.T) {
 				Bucket:     "bucket",
 				Region:     "region",
 				Database:   "database",
+				Catalog:    "catalog",
 				Table:      "table",
 				Workgroup:  "workgroup",
 				Account:    "account",
@@ -72,6 +88,7 @@ func TestAthenaConfiguration_Validate(t *testing.T) {
 				Bucket:     "",
 				Region:     "region",
 				Database:   "database",
+				Catalog:    "catalog",
 				Table:      "table",
 				Workgroup:  "workgroup",
 				Account:    "account",
@@ -84,6 +101,7 @@ func TestAthenaConfiguration_Validate(t *testing.T) {
 				Bucket:     "bucket",
 				Region:     "",
 				Database:   "database",
+				Catalog:    "catalog",
 				Table:      "table",
 				Workgroup:  "workgroup",
 				Account:    "account",
@@ -96,6 +114,7 @@ func TestAthenaConfiguration_Validate(t *testing.T) {
 				Bucket:     "bucket",
 				Region:     "region",
 				Database:   "",
+				Catalog:    "catalog",
 				Table:      "table",
 				Workgroup:  "workgroup",
 				Account:    "account",
@@ -109,6 +128,7 @@ func TestAthenaConfiguration_Validate(t *testing.T) {
 				Region:     "region",
 				Database:   "database",
 				Table:      "",
+				Catalog:    "catalog",
 				Workgroup:  "workgroup",
 				Account:    "account",
 				Authorizer: &ServiceAccount{},
@@ -120,6 +140,7 @@ func TestAthenaConfiguration_Validate(t *testing.T) {
 				Bucket:     "bucket",
 				Region:     "region",
 				Database:   "database",
+				Catalog:    "catalog",
 				Table:      "table",
 				Workgroup:  "",
 				Account:    "account",
@@ -132,6 +153,7 @@ func TestAthenaConfiguration_Validate(t *testing.T) {
 				Bucket:     "bucket",
 				Region:     "region",
 				Database:   "database",
+				Catalog:    "catalog",
 				Table:      "table",
 				Workgroup:  "workgroup",
 				Account:    "",
@@ -170,6 +192,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -182,6 +205,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -197,6 +221,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -209,6 +234,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:     "bucket",
 				Region:     "region",
 				Database:   "database",
+				Catalog:    "catalog",
 				Table:      "table",
 				Workgroup:  "workgroup",
 				Account:    "account",
@@ -221,6 +247,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:     "bucket",
 				Region:     "region",
 				Database:   "database",
+				Catalog:    "catalog",
 				Table:      "table",
 				Workgroup:  "workgroup",
 				Account:    "account",
@@ -230,6 +257,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:     "bucket",
 				Region:     "region",
 				Database:   "database",
+				Catalog:    "catalog",
 				Table:      "table",
 				Workgroup:  "workgroup",
 				Account:    "account",
@@ -242,6 +270,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:     "bucket",
 				Region:     "region",
 				Database:   "database",
+				Catalog:    "catalog",
 				Table:      "table",
 				Workgroup:  "workgroup",
 				Account:    "account",
@@ -251,6 +280,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:     "bucket",
 				Region:     "region",
 				Database:   "database",
+				Catalog:    "catalog",
 				Table:      "table",
 				Workgroup:  "workgroup",
 				Account:    "account",
@@ -263,6 +293,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -275,6 +306,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:     "bucket",
 				Region:     "region",
 				Database:   "database",
+				Catalog:    "catalog",
 				Table:      "table",
 				Workgroup:  "workgroup",
 				Account:    "account",
@@ -287,6 +319,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -299,6 +332,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:    "bucket2",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -314,6 +348,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -326,6 +361,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region2",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -341,6 +377,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -353,6 +390,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database2",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -368,6 +406,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -380,6 +419,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table2",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -390,11 +430,41 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 			},
 			expected: false,
 		},
+		"different catalog": {
+			left: AthenaConfiguration{
+				Bucket:    "bucket",
+				Region:    "region",
+				Database:  "database",
+				Catalog:   "catalog",
+				Table:     "table",
+				Workgroup: "workgroup",
+				Account:   "account",
+				Authorizer: &AccessKey{
+					ID:     "id",
+					Secret: "secret",
+				},
+			},
+			right: &AthenaConfiguration{
+				Bucket:    "bucket",
+				Region:    "region",
+				Database:  "database",
+				Catalog:   "catalog2",
+				Table:     "table",
+				Workgroup: "workgroup",
+				Account:   "account",
+				Authorizer: &AccessKey{
+					ID:     "id",
+					Secret: "secret",
+				},
+			},
+			expected: false,
+		},
 		"different workgroup": {
 			left: AthenaConfiguration{
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -407,6 +477,7 @@ func TestAthenaConfiguration_Equals(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup2",
 				Account:   "account",
@@ -487,6 +558,7 @@ func TestAthenaConfiguration_JSON(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -502,6 +574,7 @@ func TestAthenaConfiguration_JSON(t *testing.T) {
 				Bucket:     "bucket",
 				Region:     "region",
 				Database:   "database",
+				Catalog:    "catalog",
 				Table:      "table",
 				Workgroup:  "workgroup",
 				Account:    "account",
@@ -513,6 +586,7 @@ func TestAthenaConfiguration_JSON(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -530,6 +604,7 @@ func TestAthenaConfiguration_JSON(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -544,6 +619,7 @@ func TestAthenaConfiguration_JSON(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",
@@ -558,6 +634,7 @@ func TestAthenaConfiguration_JSON(t *testing.T) {
 				Bucket:    "bucket",
 				Region:    "region",
 				Database:  "database",
+				Catalog:   "catalog",
 				Table:     "table",
 				Workgroup: "workgroup",
 				Account:   "account",

pkg/cloud/aws/athenaquerier.go

@@ -94,6 +94,9 @@ func (aq *AthenaQuerier) queryAthenaPaginated(ctx context.Context, query string,
 		Database: aws.String(aq.Database),
 	}
 
+	if aq.Catalog != "" {
+		queryExecutionCtx.Catalog = aws.String(aq.Catalog)
+	}
 	resultConfiguration := &types.ResultConfiguration{
 		OutputLocation: aws.String(aq.Bucket),
 	}

pkg/cloud/aws/provider.go

@@ -401,6 +401,7 @@ type AwsAthenaInfo struct {
 	AthenaBucketName string `json:"athenaBucketName"`
 	AthenaRegion     string `json:"athenaRegion"`
 	AthenaDatabase   string `json:"athenaDatabase"`
+	AthenaCatalog    string `json:"athenaCatalog"`
 	AthenaTable      string `json:"athenaTable"`
 	AthenaWorkgroup  string `json:"athenaWorkgroup"`
 	ServiceKeyName   string `json:"serviceKeyName"`
@@ -414,6 +415,7 @@ func (aai *AwsAthenaInfo) IsEmpty() bool {
 	return aai.AthenaBucketName == "" &&
 		aai.AthenaRegion == "" &&
 		aai.AthenaDatabase == "" &&
+		aai.AthenaCatalog == "" &&
 		aai.AthenaTable == "" &&
 		aai.AthenaWorkgroup == "" &&
 		aai.ServiceKeyName == "" &&
@@ -514,6 +516,7 @@ func (aws *AWS) GetAWSAthenaInfo() (*AwsAthenaInfo, error) {
 		AthenaBucketName: config.AthenaBucketName,
 		AthenaRegion:     config.AthenaRegion,
 		AthenaDatabase:   config.AthenaDatabase,
+		AthenaCatalog:    config.AthenaCatalog,
 		AthenaTable:      config.AthenaTable,
 		AthenaWorkgroup:  config.AthenaWorkgroup,
 		ServiceKeyName:   aak.AccessKeyID,
@@ -569,6 +572,7 @@ func (aws *AWS) UpdateConfig(r io.Reader, updateType string) (*models.CustomPric
 			c.AthenaBucketName = aai.AthenaBucketName
 			c.AthenaRegion = aai.AthenaRegion
 			c.AthenaDatabase = aai.AthenaDatabase
+			c.AthenaCatalog = aai.AthenaCatalog
 			c.AthenaTable = aai.AthenaTable
 			c.AthenaWorkgroup = aai.AthenaWorkgroup
 			c.ServiceKeyName = aai.ServiceKeyName
@@ -1902,6 +1906,10 @@ func (aws *AWS) QueryAthenaPaginated(ctx context.Context, query string, fn func(
 		Database: awsSDK.String(awsAthenaInfo.AthenaDatabase),
 	}
 
+	if awsAthenaInfo.AthenaCatalog != "" {
+		queryExecutionCtx.Catalog = awsSDK.String(awsAthenaInfo.AthenaCatalog)
+	}
+
 	resultConfiguration := &athenaTypes.ResultConfiguration{
 		OutputLocation: awsSDK.String(awsAthenaInfo.AthenaBucketName),
 	}

pkg/cloud/gcp/provider.go

@@ -277,6 +277,7 @@ func (gcp *GCP) UpdateConfig(r io.Reader, updateType string) (*models.CustomPric
 			c.AthenaBucketName = a.AthenaBucketName
 			c.AthenaRegion = a.AthenaRegion
 			c.AthenaDatabase = a.AthenaDatabase
+			c.AthenaCatalog = a.AthenaCatalog
 			c.AthenaTable = a.AthenaTable
 			c.AthenaWorkgroup = a.AthenaWorkgroup
 			c.ServiceKeyName = a.ServiceKeyName

pkg/cloud/models/models.go

@@ -160,6 +160,7 @@ type CustomPricing struct {
 	AthenaBucketName             string `json:"athenaBucketName"`
 	AthenaRegion                 string `json:"athenaRegion"`
 	AthenaDatabase               string `json:"athenaDatabase"`
+	AthenaCatalog                string `json:"athenaCatalog"`
 	AthenaTable                  string `json:"athenaTable"`
 	AthenaWorkgroup              string `json:"athenaWorkgroup"`
 	MasterPayerARN               string `json:"masterPayerARN"`