Type checking fix: Restrict the keys / values for ArtifactLocation and centralize the common ones

This commit is contained in:
Michael Hohn
2024-07-08 12:07:46 -07:00
committed by =Michael Hohn
parent b3cf7a4f65
commit 3566f5169e
7 changed files with 45 additions and 29 deletions

View File

@@ -5,10 +5,21 @@ import (
"mrvacommander/pkg/common"
)
// XX: static types: split by type?
// Restrict the keys / values for ArtifactLocation and centralize the common ones
// here
const (
AF_VAL_BUCKET_RESULTS = "results"
AF_VAL_BUCKET_PACKS = "packs"
AF_KEY_BUCKET = "bucket"
AF_KEY_KEY = "key"
)
type ArtifactLocation struct {
// Data is a map of key-value pairs that describe the location of the artifact.
// For example, a simple key-value pair could be "path" -> "/path/to/artifact.tgz".
// Alternatively, a more complex example could be "bucket" -> "example", "key" -> "UNIQUE_ARTIFACT_IDENTIFIER".
// XX: static types
Data map[string]string `json:"data"`
}

View File

@@ -25,7 +25,7 @@ func (store *InMemoryArtifactStore) GetQueryPack(location ArtifactLocation) ([]b
store.mu.Lock()
defer store.mu.Unlock()
key := location.Data["key"]
key := location.Data[AF_KEY_KEY]
data, exists := store.packs[key]
if !exists {
return nil, fmt.Errorf("query pack not found: %s", key)
@@ -41,10 +41,11 @@ func (store *InMemoryArtifactStore) SaveQueryPack(sessionId int, data []byte) (A
key := deriveKeyFromSessionId(sessionId)
store.packs[key] = data
// XX: static types
location := ArtifactLocation{
Data: map[string]string{
"bucket": "packs",
"key": key,
AF_KEY_BUCKET: AF_VAL_BUCKET_PACKS,
AF_KEY_KEY: key,
},
}
return location, nil
@@ -55,7 +56,7 @@ func (store *InMemoryArtifactStore) GetResult(location ArtifactLocation) ([]byte
store.mu.Lock()
defer store.mu.Unlock()
key := location.Data["key"]
key := location.Data[AF_KEY_KEY]
data, exists := store.results[key]
if !exists {
return nil, fmt.Errorf("result not found: %s", key)
@@ -68,7 +69,7 @@ func (store *InMemoryArtifactStore) GetResultSize(location ArtifactLocation) (in
store.mu.Lock()
defer store.mu.Unlock()
key := location.Data["key"]
key := location.Data[AF_KEY_KEY]
data, exists := store.results[key]
if !exists {
return 0, fmt.Errorf("result not found: %s", key)
@@ -84,10 +85,11 @@ func (store *InMemoryArtifactStore) SaveResult(jobSpec common.JobSpec, data []by
key := deriveKeyFromJobSpec(jobSpec)
store.results[key] = data
// XX: static types
location := ArtifactLocation{
Data: map[string]string{
"bucket": "results",
"key": key,
AF_KEY_BUCKET: AF_VAL_BUCKET_RESULTS,
AF_KEY_KEY: key,
},
}
return location, nil

View File

@@ -13,11 +13,6 @@ import (
"github.com/minio/minio-go/v7/pkg/credentials"
)
const (
RESULTS_BUCKET_NAME = "results"
PACKS_BUCKET_NAME = "packs"
)
type MinIOArtifactStore struct {
client *minio.Client
}
@@ -34,12 +29,12 @@ func NewMinIOArtifactStore(endpoint, id, secret string) (*MinIOArtifactStore, er
slog.Info("Connected to MinIO artifact store server")
// Create "results" bucket
if err := common.CreateMinIOBucketIfNotExists(minioClient, RESULTS_BUCKET_NAME); err != nil {
if err := common.CreateMinIOBucketIfNotExists(minioClient, AF_VAL_BUCKET_RESULTS); err != nil {
return nil, fmt.Errorf("could not create results bucket: %v", err)
}
// Create "packs" bucket
if err := common.CreateMinIOBucketIfNotExists(minioClient, PACKS_BUCKET_NAME); err != nil {
if err := common.CreateMinIOBucketIfNotExists(minioClient, AF_VAL_BUCKET_PACKS); err != nil {
return nil, fmt.Errorf("could not create packs bucket: %v", err)
}
@@ -53,7 +48,7 @@ func (store *MinIOArtifactStore) GetQueryPack(location ArtifactLocation) ([]byte
}
func (store *MinIOArtifactStore) SaveQueryPack(jobId int, data []byte) (ArtifactLocation, error) {
return store.saveArtifact(PACKS_BUCKET_NAME, deriveKeyFromSessionId(jobId), data, "application/gzip")
return store.saveArtifact(AF_VAL_BUCKET_PACKS, deriveKeyFromSessionId(jobId), data, "application/gzip")
}
func (store *MinIOArtifactStore) GetResult(location ArtifactLocation) ([]byte, error) {
@@ -61,8 +56,8 @@ func (store *MinIOArtifactStore) GetResult(location ArtifactLocation) ([]byte, e
}
func (store *MinIOArtifactStore) GetResultSize(location ArtifactLocation) (int, error) {
bucket := location.Data["bucket"]
key := location.Data["key"]
bucket := location.Data[AF_KEY_BUCKET]
key := location.Data[AF_KEY_KEY]
objectInfo, err := store.client.StatObject(context.Background(), bucket, key, minio.StatObjectOptions{})
if err != nil {
@@ -77,12 +72,12 @@ func (store *MinIOArtifactStore) GetResultSize(location ArtifactLocation) (int,
}
func (store *MinIOArtifactStore) SaveResult(jobSpec common.JobSpec, data []byte) (ArtifactLocation, error) {
return store.saveArtifact(RESULTS_BUCKET_NAME, deriveKeyFromJobSpec(jobSpec), data, "application/zip")
return store.saveArtifact(AF_VAL_BUCKET_RESULTS, deriveKeyFromJobSpec(jobSpec), data, "application/zip")
}
func (store *MinIOArtifactStore) getArtifact(location ArtifactLocation) ([]byte, error) {
bucket := location.Data["bucket"]
key := location.Data["key"]
bucket := location.Data[AF_KEY_BUCKET]
key := location.Data[AF_KEY_KEY]
object, err := store.client.GetObject(context.Background(), bucket, key, minio.GetObjectOptions{})
if err != nil {
@@ -98,18 +93,21 @@ func (store *MinIOArtifactStore) getArtifact(location ArtifactLocation) ([]byte,
return data, nil
}
func (store *MinIOArtifactStore) saveArtifact(bucket, key string, data []byte, contentType string) (ArtifactLocation, error) {
_, err := store.client.PutObject(context.Background(), bucket, key, bytes.NewReader(data), int64(len(data)), minio.PutObjectOptions{
ContentType: contentType,
})
func (store *MinIOArtifactStore) saveArtifact(bucket, key string, data []byte,
contentType string) (ArtifactLocation, error) {
_, err := store.client.PutObject(context.Background(), bucket, key,
bytes.NewReader(data), int64(len(data)), minio.PutObjectOptions{
ContentType: contentType,
})
if err != nil {
return ArtifactLocation{}, err
}
// XX: static types
location := ArtifactLocation{
Data: map[string]string{
"bucket": bucket,
"key": key,
AF_KEY_BUCKET: bucket,
AF_KEY_KEY: key,
},
}

View File

@@ -26,8 +26,10 @@ const (
)
type SarifRun struct {
// XX: static types
VersionControlProvenance []interface{} `json:"versionControlProvenance,omitempty"`
Results []interface{} `json:"results"`
// XX: never set, only read
Results []interface{} `json:"results"`
}
type Sarif struct {

View File

@@ -8,6 +8,7 @@ type CodeQLDatabaseLocation struct {
// `data` is a map of key-value pairs that describe the location of the database.
// For example, a simple key-value pair could be "path" -> "/path/to/database.zip".
// A more complex implementation could be "bucket" -> "example", "key" -> "unique_identifier".
// XX: static types
data map[string]string
}

View File

@@ -5,6 +5,7 @@ import (
"fmt"
"io"
"log/slog"
"mrvacommander/pkg/artifactstore"
"mrvacommander/pkg/common"
"github.com/minio/minio-go/v7"
@@ -58,8 +59,8 @@ func (store *MinIOCodeQLDatabaseStore) FindAvailableDBs(analysisReposRequested [
}
func (store *MinIOCodeQLDatabaseStore) GetDatabase(location CodeQLDatabaseLocation) ([]byte, error) {
bucket := location.data["bucket"]
key := location.data["key"]
bucket := location.data[artifactstore.AF_KEY_BUCKET]
key := location.data[artifactstore.AF_KEY_KEY]
object, err := store.client.GetObject(context.Background(), bucket, key, minio.GetObjectOptions{})
if err != nil {

View File

@@ -34,6 +34,7 @@ func NewRabbitMQQueue(
const (
tryCount = 5
retryDelaySec = 3
// XX: static typing?
jobsQueueName = "tasks"
resultsQueueName = "results"
)