96 lines
2.3 KiB
Go
96 lines
2.3 KiB
Go
package remotefileupload
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/fiskerinc/cloud-services/pkg/logger"
|
|
az "github.com/Azure/azure-storage-blob-go/azblob"
|
|
pqAZ "github.com/xitongsys/parquet-go-source/azblob"
|
|
"github.com/xitongsys/parquet-go/source"
|
|
"github.com/xitongsys/parquet-go/writer"
|
|
)
|
|
|
|
var (
|
|
parquetThreadCount int64 = 4
|
|
)
|
|
|
|
var (
|
|
errOnCloseWriter = "Unable to close writer"
|
|
)
|
|
|
|
// Required struct to intake compressed parquet files which lists fields as optional
|
|
//
|
|
// hence the pointers to int,string
|
|
type ParquetCANMessage struct {
|
|
TimestampUSec *int64 `json:"epoch_usec" parquet:"name=epoch_usec, type=INT64"`
|
|
ID *int32 `json:"id" parquet:"name=id, type=INT32"`
|
|
Data *string `json:"data" parquet:"name=data, type=BYTE_ARRAY"`
|
|
}
|
|
|
|
// NewAzureParquetBlobWriter creates a new instance of ParquetBlobWriter that can be used to write Parquet files to Azure Blob Storage.
|
|
//
|
|
// Parameters:
|
|
// - blobUrl: The URL of the Azure Blob Storage container where the Parquet files will be stored.
|
|
//
|
|
// Returns:
|
|
// - ParquetBlobWriter: An instance of ParquetBlobWriter.
|
|
// - error: An error if there was a problem creating the writer.
|
|
func NewAzureParquetBlobWriter(blobUrl, azureAccount, azureAccountKey string) (ParquetBlobWriter, error) {
|
|
creds, err := az.NewSharedKeyCredential(azureAccount, azureAccountKey)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
fr, err := pqAZ.NewAzBlobFileWriter(
|
|
context.Background(),
|
|
blobUrl,
|
|
creds,
|
|
pqAZ.WriterOptions{},
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
pr, err := writer.NewParquetWriter(fr, new(ParquetCANMessage), parquetThreadCount)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &AzureParquetBlobWriter{blob: fr, fileWriter: pr}, nil
|
|
}
|
|
|
|
type ParquetBlobWriter interface {
|
|
Write(payload interface{}) error
|
|
Size() int64
|
|
Close()
|
|
}
|
|
|
|
type AzureParquetBlobWriter struct {
|
|
blob source.ParquetFile
|
|
fileWriter *writer.ParquetWriter
|
|
}
|
|
|
|
func (w *AzureParquetBlobWriter) Write(payload interface{}) error {
|
|
err := w.fileWriter.Write(payload)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (w *AzureParquetBlobWriter) Size() int64 {
|
|
return w.fileWriter.Size
|
|
}
|
|
|
|
func (w *AzureParquetBlobWriter) Close() {
|
|
err := w.fileWriter.WriteStop()
|
|
if err != nil {
|
|
logger.Debug().Msgf("%v: %s", err, errOnCloseWriter)
|
|
}
|
|
|
|
err = w.blob.Close()
|
|
if err != nil {
|
|
logger.Debug().Msgf("%v:%s", err, errOnCloseWriter)
|
|
}
|
|
}
|