Glue partitions refactor

warehouse/datalake/schema-repository/glue.go

@@ -115,11 +115,17 @@ func (gl *GlueSchemaRepository) CreateSchema() (err error) {
 }
 
 func (gl *GlueSchemaRepository) getPartitionKeys() []*glue.Column {
-	timeWindowFormat, _ := gl.Warehouse.Destination.Config["timeWindowFormat"].(string)
-	if timeWindowFormat != "" {
-		// Assumes a well-formed partitioning format
-		columnName := strings.Split(timeWindowFormat, "=")[0]
-		return []*glue.Column{&glue.Column{Name: aws.String(columnName), Type: aws.String("date")}}
+	var partitionColumns []*glue.Column
+	for _, partitionPair := range gl.Warehouse.DestPartitionKeys {
+		// Assumes a well-formed partitioning pair
+		columnName := partitionPair["key"]
+		if columnName != "" {
+			partitionColumns = append(partitionColumns, &glue.Column{Name: aws.String(columnName), Type: aws.String("date")})
+		}
+	}
+
+	if len(partitionColumns) > 0 {
+		return partitionColumns
 	}
 	return nil
 }

warehouse/upload.go

@@ -404,9 +404,8 @@ func (job *UploadJobT) run() (err error) {
 			job.matchRowsInStagingAndLoadFiles()
 			job.recordLoadFileGenerationTimeStat(startLoadFileID, endLoadFileID)
 			if job.warehouse.Type == "S3_DATALAKE" {
-				timeWindowFormatI, timeWindowFormatAvailable := job.warehouse.Destination.Config["timeWindowFormat"]
-				if timeWindowFormatAvailable {
-					timeWindowFormat, _ := timeWindowFormatI.(string)
+				timeWindowFormat := warehouseutils.GetTimeWindowFormat(job.warehouse.DestPartitionKeys)
+				if timeWindowFormat != "" {
 					for tableName := range job.upload.UploadSchema {
 						loadFiles := job.GetLoadFilesMetadata(warehouseutils.GetLoadFilesOptionsT{
 							Table:   tableName,
@@ -415,7 +414,7 @@ func (job *UploadJobT) run() (err error) {
 						})
 						// This is best done every 100 files, since it's a batch request for updates in Glue
 						partitionBatchSize := 99
-						for i := 0; i < len(loadFiles) && timeWindowFormat != ""; i += partitionBatchSize {
+						for i := 0; i < len(loadFiles); i += partitionBatchSize {
 							end := i + partitionBatchSize
 
 							if end > len(loadFiles) {
@@ -1672,7 +1671,7 @@ func (job *UploadJobT) createLoadFiles(generateAll bool) (startLoadFileID int64,
 			}
 
 			if job.warehouse.Type == "S3_DATALAKE" {
-				timeWindowFormat, _ := job.warehouse.Destination.Config["timeWindowFormat"].(string)
+				timeWindowFormat := warehouseutils.GetTimeWindowFormat(job.warehouse.DestPartitionKeys)
 				if timeWindowFormat != "" {
 					payload.LoadFilePrefix = stagingFile.TimeWindow.Format(timeWindowFormat)
 				}

warehouse/utils/utils.go

@@ -133,11 +133,12 @@ func loadConfig() {
 }
 
 type WarehouseT struct {
-	Source      backendconfig.SourceT
-	Destination backendconfig.DestinationT
-	Namespace   string
-	Type        string
-	Identifier  string
+	Source            backendconfig.SourceT
+	Destination       backendconfig.DestinationT
+	Namespace         string
+	Type              string
+	Identifier        string
+	DestPartitionKeys []map[string]string
 }
 
 type DestinationT struct {
@@ -844,3 +845,29 @@ func GetSSLKeyDirPath(destinationID string) (whSSLRootDir string) {
 	sslDirPath := fmt.Sprintf("%s/dest-ssls/%s", directoryName, destinationID)
 	return sslDirPath
 }
+
+func GetPartitionKeysFromConfig(config map[string]interface{}) (partitionKeys []map[string]string) {
+	partitionKeysI, partitionKeysIAvailable := config["partitionKeys"]
+	if partitionKeysIAvailable {
+		partitionKeysString, _ := json.Marshal(partitionKeysI)
+		err := json.Unmarshal(partitionKeysString, &partitionKeys)
+		if err != nil {
+			pkgLogger.Errorf("Error extracting partition keys from destination config %v", err)
+		}
+	}
+	return
+}
+
+func GetTimeWindowFormat(partitionKeys []map[string]string) (timeWindowFormat string) {
+	// supports only single partition
+	// TODO: extend to allow multiple partition keys on table
+	if len(partitionKeys) > 0 {
+		partitionPair := partitionKeys[0]
+		key := partitionPair["key"]
+		val := partitionPair["value"]
+		if key != "" && val != "" {
+			timeWindowFormat = fmt.Sprintf("%v=%v", key, val)
+		}
+	}
+	return
+}

warehouse/warehouse.go

@@ -263,11 +263,12 @@ func (wh *HandleT) backendConfigSubscriber() {
 				}
 				namespace := wh.getNamespace(destination.Config, source, destination, wh.destType)
 				warehouse := warehouseutils.WarehouseT{
-					Source:      source,
-					Destination: destination,
-					Namespace:   namespace,
-					Type:        wh.destType,
-					Identifier:  warehouseutils.GetWarehouseIdentifier(wh.destType, source.ID, destination.ID),
+					Source:            source,
+					Destination:       destination,
+					Namespace:         namespace,
+					Type:              wh.destType,
+					Identifier:        warehouseutils.GetWarehouseIdentifier(wh.destType, source.ID, destination.ID),
+					DestPartitionKeys: warehouseutils.GetPartitionKeysFromConfig(destination.Config),
 				}
 				wh.warehouses = append(wh.warehouses, warehouse)
 
@@ -1105,11 +1106,12 @@ func minimalConfigSubscriber() {
 						connectionsMap[destination.ID] = map[string]warehouseutils.WarehouseT{}
 					}
 					connectionsMap[destination.ID][source.ID] = warehouseutils.WarehouseT{
-						Destination: destination,
-						Namespace:   namespace,
-						Type:        wh.destType,
-						Source:      source,
-						Identifier:  warehouseutils.GetWarehouseIdentifier(wh.destType, source.ID, destination.ID),
+						Destination:       destination,
+						Namespace:         namespace,
+						Type:              wh.destType,
+						Source:            source,
+						Identifier:        warehouseutils.GetWarehouseIdentifier(wh.destType, source.ID, destination.ID),
+						DestPartitionKeys: warehouseutils.GetPartitionKeysFromConfig(destination.Config),
 					}
 					connectionsMapLock.Unlock()
 				}