Skip to content

fix: support cost attribution env based query labeling #95

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Mryashbhardwaj wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: support cost attribution env based query labeling #95

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions mc2mc/internal/config/config.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ type ConfigEnv struct {
LoadMethod string `env:"LOAD_METHOD" envDefault:"APPEND"`
QueryFilePath string `env:"QUERY_FILE_PATH" envDefault:"/data/in/query.sql"`
DestinationTableID string `env:"DESTINATION_TABLE_ID"`
CostAttributionTeam string `env:"COST_ATTRIBUTION_TEAM"`
DStart string `env:"DSTART"`
DEnd string `env:"DEND"`
ExecutionProject string `env:"EXECUTION_PROJECT"`
Expand Down
27 changes: 21 additions & 6 deletions mc2mc/internal/query/builder.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,11 @@ type Builder struct {

query string

method Method
destinationTableID string
orderedColumns []string
overridedValues map[string]string
method Method
destinationTableID string
costAttributionTeam string
orderedColumns []string
overridedValues map[string]string

enableAutoPartition bool
enablePartitionValue bool
Expand Down Expand Up @@ -54,6 +55,10 @@ func (b *Builder) SetOptions(options ...Option) *Builder {
return b
}

func getCostAttributionComment(teamName string) string {
return fmt.Sprintf("--cost_attribution_team=%s\n", teamName)
}

// Build constructs the final query with the given options
func (b *Builder) Build() (string, error) {
if b.query == "" {
Expand All @@ -64,9 +69,12 @@ func (b *Builder) Build() (string, error) {
// split query components
hrs, vars, queries := SplitQueryComponents(b.query)
if len(queries) <= 1 {
if b.costAttributionTeam != "" {
return fmt.Sprintf("%s\n%s", b.query, getCostAttributionComment(b.costAttributionTeam)), nil
}
return b.query, nil
}
query := b.constructMergeQuery(hrs, vars, queries)
query := b.constructMergeQuery(hrs, vars, queries, b.costAttributionTeam)
return query, nil
}

Expand Down Expand Up @@ -143,7 +151,11 @@ func (b *Builder) Build() (string, error) {
if varsAndUDFs != "" {
varsAndUDFs += "\n"
}

query = fmt.Sprintf("%s%s%s%s", hr, drops, varsAndUDFs, query)
if b.costAttributionTeam != "" {
query = fmt.Sprintf("%s\n%s\n", query, getCostAttributionComment(b.costAttributionTeam))
}
return query, nil
}

Expand Down Expand Up @@ -188,7 +200,7 @@ func (b *Builder) constructOverridedValues(query string) (string, error) {
}

// constructMergeQueries constructs merge queries with headers and variables
func (b *Builder) constructMergeQuery(hrs, vars, queries []string) string {
func (b *Builder) constructMergeQuery(hrs, vars, queries []string, costAttributionTeam string) string {
Copy link
Member

@deryrahman deryrahman Dec 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no need to pass costAttributionTeam argument here, we can use var in builder object itself

builder := strings.Builder{}
for i, q := range queries {
q = strings.TrimSpace(q)
Expand All @@ -204,6 +216,9 @@ func (b *Builder) constructMergeQuery(hrs, vars, queries []string) string {
builder.WriteString(fmt.Sprintf("%s\n", variables))
}
builder.WriteString(fmt.Sprintf("%s\n;", q))
if costAttributionTeam != "" {
builder.WriteString(fmt.Sprintf("\n%s\n", getCostAttributionComment(costAttributionTeam)))
}
if i < len(queries)-1 {
builder.WriteString(fmt.Sprintf("\n%s\n", BREAK_MARKER))
}
Expand Down
6 changes: 6 additions & 0 deletions mc2mc/internal/query/option.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,3 +46,9 @@ func WithOverridedValue(field, value string) Option {
b.overridedValues[field] = value
}
}

func WithCostAttributionLabel(teamName string) Option {
return func(b *Builder) {
b.costAttributionTeam = teamName
}
}
3 changes: 3 additions & 0 deletions mc2mc/mc2mc.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,7 @@ func mc2mc(envs []string) error {
query.WithOverridedValue("_partitiondate", fmt.Sprintf("DATE(timestamp('%s'))", dstart)),
query.WithAutoPartition(cfg.DevEnableAutoPartition == "true"),
query.WithPartitionValue(cfg.DevEnablePartitionValue == "true"),
query.WithCostAttributionLabel(cfg.CostAttributionTeam),
query.WithColumnOrder(),
).Build()
if err != nil {
Expand All @@ -97,6 +98,7 @@ func mc2mc(envs []string) error {
query.WithDestination(cfg.DestinationTableID),
query.WithAutoPartition(cfg.DevEnableAutoPartition == "true"),
query.WithPartitionValue(cfg.DevEnablePartitionValue == "true"),
query.WithCostAttributionLabel(cfg.CostAttributionTeam),
query.WithColumnOrder(),
)

Expand Down Expand Up @@ -151,6 +153,7 @@ func mc2mc(envs []string) error {
l,
client.NewODPSClient(l, cfg.GenOdps()),
query.WithQuery(string(raw)),
query.WithCostAttributionLabel(cfg.CostAttributionTeam),
query.WithMethod(query.MERGE),
).Build()
if err != nil {
Expand Down
Loading