Merge pull request #31 from muroon/specify-catalog-query-execution

Specify catalog query execution

Author: muroon

conn.go

@@ -178,10 +178,22 @@ func (c *conn) dropCTASTable(ctx context.Context, table string) func() error {
 
 // startQuery starts an Athena query and returns its ID.
 func (c *conn) startQuery(ctx context.Context, query string) (string, error) {
+	// resolve catalog from context, fallback to connection-level catalog
+	catalog := c.catalog
+	if cat, ok := getCatalog(ctx); ok {
+		catalog = cat
+	}
+
+	var catalogPtr *string
+	if catalog != "" {
+		catalogPtr = aws.String(catalog)
+	}
+
 	resp, err := c.athena.StartQueryExecution(ctx, &athena.StartQueryExecutionInput{
 		QueryString: aws.String(query),
 		QueryExecutionContext: &types.QueryExecutionContext{
 			Database: aws.String(c.db),
+			Catalog:  catalogPtr,
 		},
 		ResultConfiguration: &types.ResultConfiguration{
 			OutputLocation: aws.String(c.OutputLocation),

rows_gzip.go

@@ -162,8 +162,13 @@ func (r *rowsGzipDL) downloadCompressedData(ctx context.Context, cfg aws.Config,
 }
 
 func (r *rowsGzipDL) getTableAsync(ctx context.Context, errCh chan error) {
+	catalog := r.catalog
+	if catalog == "" {
+		catalog = CATALOG_AWS_DATA_CATALOG
+	}
+
 	data, err := r.athena.GetTableMetadata(ctx, &athena.GetTableMetadataInput{
-		CatalogName:  aws.String(r.catalog),
+		CatalogName:  aws.String(catalog),
 		DatabaseName: aws.String(r.db),
 		TableName:    aws.String(r.ctasTable),
 	})