<commit_before>// Package jobmonitor contains types and functions for aggregation job monitoring.
package jobmonitor

import (
	"context"
	"fmt"
	"time"

	"cloud.google.com/go/firestore"
)

// Paths should be used when writing to Firestore.
const (
	ProdPath = "jobs"
	TestPath = "jobs-test"
)

// PipelineJob represent a Beam pipeline job on an aggregator for a certain level of a aggregation job.
type PipelineJob struct {
	Created time.Time `firestore:"created,omitempty"`
	Message string    `firestore:"message,omitempty"`
	Result  string    `firestore:"result,omitempty"`
	Status  string    `firestore:"status,omitempty"`
	Updated time.Time `firestore:"updated,omitempty"`
}

// AggregatorJobs contains the pipeline jobs for different hierarchical levels.
type AggregatorJobs struct {
	// The sub jobs are keyed by the 0-based level.
	LevelJobs map[int]*PipelineJob
}

// AggregationJob represent an aggregation job.
type AggregationJob struct {
	// The aggregator is represented by its origin string.
	Aggregators map[string]*AggregatorJobs
	// Overall status of a job.
	Created time.Time `firestore:"created,omitempty"`
}

// WriteJobs writes a list of jobs to Firestore. The input jobs are keyed by the query IDs.
func WriteJobs(ctx context.Context, client *firestore.Client, path string, jobs map[string]*AggregationJob) error {
	for queryID, job := range jobs {
		_, err := client.Collection(path).Doc(queryID).Set(ctx, map[string]interface{}{
			"created": job.Created,
		})
		if err != nil {
			return err
		}
		for origin, aggjobs := range job.Aggregators {
			for level, subjob := range aggjobs.LevelJobs {
				_, err := client.Collection(path).Doc(queryID).Collection("aggregators").Doc(origin).Collection("levels").Doc(fmt.Sprintf("level-%d", level)).Set(ctx, subjob)
				if err != nil {
					return err
				}
			}
		}
	}
	return nil
}<commit_msg>Create aggregator docs in Firestore explicitly<commit_after>// Package jobmonitor contains types and functions for aggregation job monitoring.
package jobmonitor

import (
	"context"
	"fmt"
	"time"

	"cloud.google.com/go/firestore"
)

// Paths should be used when writing to Firestore.
const (
	ProdPath = "jobs"
	TestPath = "jobs-test"
)

// PipelineJob represent a Beam pipeline job on an aggregator for a certain level of a aggregation job.
type PipelineJob struct {
	Created time.Time `firestore:"created,omitempty"`
	Message string    `firestore:"message,omitempty"`
	Result  string    `firestore:"result,omitempty"`
	Status  string    `firestore:"status,omitempty"`
	Updated time.Time `firestore:"updated,omitempty"`
}

// AggregatorJobs contains the pipeline jobs for different hierarchical levels.
type AggregatorJobs struct {
	// The sub jobs are keyed by the 0-based level.
	LevelJobs map[int]*PipelineJob
}

// AggregationJob represent an aggregation job.
type AggregationJob struct {
	// The aggregator is represented by its origin string.
	Aggregators map[string]*AggregatorJobs
	// Overall status of a job.
	Created time.Time `firestore:"created,omitempty"`
}

// WriteJobs writes a list of jobs to Firestore. The input jobs are keyed by the query IDs.
func WriteJobs(ctx context.Context, client *firestore.Client, path string, jobs map[string]*AggregationJob) error {
	for queryID, job := range jobs {
		_, err := client.Collection(path).Doc(queryID).Set(ctx, map[string]interface{}{
			"created": job.Created,
		})
		if err != nil {
			return err
		}
		for origin, aggjobs := range job.Aggregators {
			_, err := client.Collection(path).Doc(queryID).Collection("aggregators").Doc(origin).Create(ctx, map[string]interface{}{})
			if err != nil {
				return err
			}
			for level, subjob := range aggjobs.LevelJobs {
				_, err := client.Collection(path).Doc(queryID).Collection("aggregators").Doc(origin).Collection("levels").Doc(fmt.Sprintf("level-%d", level)).Set(ctx, subjob)
				if err != nil {
					return err
				}
			}
		}
	}
	return nil
}
{'fraction_non_alphanumeric': 0.11238532110091744, 'fraction_numerical': 0.0005096839959225281, 'mean_word_length': 4.29689608636977, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7620839', 'n_tokens_mistral': 1263, 'n_tokens_neox': 1166, 'n_words': 369}