Merge pull request #513 from m-lab/sandbox-dataset

choose dataset based on datatype and batch

Author: gfr10598

bq/insert.go

@@ -60,6 +60,9 @@ func NewInserter(dataset string, dt etl.DataType, partition time.Time) (etl.Inse
 		suffix = "$" + partition.Format("20060102")
 	}
 
+	// TODO - remove dataset parameter.
+	dataset = etl.DataTypeToDataset(dt)
+
 	return NewBQInserter(
 		etl.InserterParams{Dataset: dataset, Table: table, Suffix: suffix,
 			Timeout: 15 * time.Minute, BufferSize: dt.BQBufferSize(), RetryDelay: 30 * time.Second},

cmd/etl_worker/etl_worker.go

@@ -209,11 +209,7 @@ func subworker(rawFileName string, executionCount, retryCount int) (status int,
 	dateFormat := "20060102"
 	date, err := time.Parse(dateFormat, data.PackedDate)
 
-	dataset, ok := os.LookupEnv("BIGQUERY_DATASET")
-	if !ok {
-		// TODO - make this fatal.
-		dataset = "mlab_sandbox"
-	}
+	dataset := etl.DataTypeToDataset(dataType)
 	ins, err := bq.NewInserter(dataset, dataType, date)
 	if err != nil {
 		metrics.TaskCount.WithLabelValues(data.TableBase(), string(dataType), "NewInserterError").Inc()

etl/globals.go

@@ -13,6 +13,17 @@ import (
 	"github.com/m-lab/etl/metrics"
 )
 
+// IsBatch indicates this process is a batch processing service.
+var IsBatch bool
+
+// OmitDeltas indicates we should NOT process all snapshots.
+var OmitDeltas bool
+
+func init() {
+	IsBatch, _ = strconv.ParseBool(os.Getenv("NDT_BATCH"))
+	OmitDeltas, _ = strconv.ParseBool(os.Getenv("NDT_OMIT_DELTAS"))
+}
+
 // YYYYMMDD is a regexp string for identifying dense dates.
 const YYYYMMDD = `\d{4}[01]\d[0123]\d`
 
@@ -109,8 +120,7 @@ func (fn *DataPath) TableBase() string {
 // IsBatchService return true if this is a NDT batch service.
 // TODO - update this to BATCH_SERVICE, so it makes sense for other pipelines.
 func IsBatchService() bool {
-	isBatch, _ := strconv.ParseBool(os.Getenv("NDT_BATCH"))
-	return isBatch
+	return IsBatch
 }
 
 // GetMetroName extracts metro name like "acc" from file name like
@@ -167,8 +177,7 @@ type DataType string
 func (dt DataType) BQBufferSize() int {
 	// Special case for NDT when omitting deltas.
 	if dt == NDT {
-		omitDeltas, _ := strconv.ParseBool(os.Getenv("NDT_OMIT_DELTAS"))
-		if omitDeltas {
+		if OmitDeltas {
 			return dataTypeToBQBufferSize[NDT_OMIT_DELTAS]
 		}
 	}
@@ -218,8 +227,21 @@ var (
 	// queue_pusher.go
 )
 
-// AddPanicMetric captures panics, increments the
-// panic metric, and then repanics.
+// DataTypeToDataset returns the appropriate dataset to use.
+// This is a bit of a hack, but works for our current needs.
+func DataTypeToDataset(dt DataType) string {
+	if dt == SS {
+		return "private"
+	}
+
+	if IsBatchService() {
+		return "batch"
+	}
+
+	return "base_tables"
+}
+
+// CountPanics updates the PanicCount metric, then repanics.
 // It must be wrapped in a defer.
 // Examples:
 //  For function that returns an error:

etl/globals_test.go

@@ -198,3 +198,17 @@ func TestCountPanics(t *testing.T) {
 
 	rePanic()
 }
+
+func TestSidestreamDataset(t *testing.T) {
+	if etl.DataTypeToDataset(etl.SS) != "private" {
+		t.Error("For SS, should be private:", etl.DataTypeToDataset(etl.SS))
+	}
+	etl.IsBatch = true
+	if etl.DataTypeToDataset(etl.NDT) != "batch" {
+		t.Error("For IsBatchService, should be batch:", etl.DataTypeToDataset(etl.NDT))
+	}
+	etl.IsBatch = false
+	if etl.DataTypeToDataset(etl.NDT) != "base_tables" {
+		t.Error("For IsBatchService, should be batch:", etl.DataTypeToDataset(etl.NDT))
+	}
+}