[go] Add blob support for MakeLiteralForType (flyteorg#131)

The MakeLiteralForType client method currently doesn't support the blob type. This adds support for that. Signed-off-by: Sam Lai <[email protected]>
eapolinario · Mar 30, 2021 · 0ebee65 · 0ebee65
1 parent 243bee9
commit 0ebee65
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 0 deletions.
diff --git a/clients/go/coreutils/extract_literal.go b/clients/go/coreutils/extract_literal.go
@@ -51,7 +51,12 @@ func ExtractFromLiteral(literal *core.Literal) (interface{}, error) {
 			case *core.Primitive_Duration:
 				scalarPrimitiveDuration := scalarPrimitive.Duration
 				return scalarPrimitiveDuration, nil
+			default:
+				return nil, fmt.Errorf("unsupported literal scalar primitive type %T", scalarValue)
 			}
+		case *core.Scalar_Blob:
+			return scalarValue.Blob.Uri, nil
+		default:
 			return nil, fmt.Errorf("unsupported literal scalar type %T", scalarValue)
 		}
 	case *core.Literal_Collection:

diff --git a/clients/go/coreutils/literals.go b/clients/go/coreutils/literals.go
@@ -488,6 +488,11 @@ func MakeLiteralForType(t *core.LiteralType, v interface{}) (*core.Literal, erro
 			return nil, err
 		}
 		return lv, nil
+	case *core.LiteralType_Blob:
+		newT := t.Type.(*core.LiteralType_Blob)
+		isDir := newT.Blob.Dimensionality == core.BlobType_MULTIPART
+		lv := MakeLiteralForBlob(storage.DataReference(fmt.Sprintf("%v", v)), isDir, newT.Blob.Format)
+		return lv, nil
 	default:
 		return nil, fmt.Errorf("unsupported type %s", t.String())
 	}

diff --git a/clients/go/coreutils/literals_test.go b/clients/go/coreutils/literals_test.go
@@ -516,6 +516,62 @@ func TestMakeLiteralForType(t *testing.T) {
 		assert.Equal(t, expectedVal, actualVal)
 	})
 
+	t.Run("Blob", func(t *testing.T) {
+		var literalType = &core.LiteralType{Type: &core.LiteralType_Blob{Blob: &core.BlobType{
+			Dimensionality: core.BlobType_SINGLE,
+		}}}
+		expectedLV := &core.Literal{Value: &core.Literal_Scalar{Scalar: &core.Scalar{
+			Value: &core.Scalar_Blob{
+				Blob: &core.Blob{
+					Uri: "s3://blah/blah/blah",
+					Metadata: &core.BlobMetadata{
+						Type: &core.BlobType{
+							Dimensionality: core.BlobType_SINGLE,
+						},
+					},
+				},
+			},
+		}}}
+		lv, err := MakeLiteralForType(literalType, "s3://blah/blah/blah")
+		assert.NoError(t, err)
+
+		assert.Equal(t, expectedLV, lv)
+
+		expectedVal, err := ExtractFromLiteral(expectedLV)
+		assert.NoError(t, err)
+		actualVal, err := ExtractFromLiteral(lv)
+		assert.NoError(t, err)
+		assert.Equal(t, expectedVal, actualVal)
+	})
+
+	t.Run("MultipartBlob", func(t *testing.T) {
+		var literalType = &core.LiteralType{Type: &core.LiteralType_Blob{Blob: &core.BlobType{
+			Dimensionality: core.BlobType_MULTIPART,
+		}}}
+		expectedLV := &core.Literal{Value: &core.Literal_Scalar{Scalar: &core.Scalar{
+			Value: &core.Scalar_Blob{
+				Blob: &core.Blob{
+					Uri: "s3://blah/blah/blah",
+					Metadata: &core.BlobMetadata{
+						Type: &core.BlobType{
+							Dimensionality: core.BlobType_MULTIPART,
+						},
+					},
+				},
+			},
+		}}}
+		lv, err := MakeLiteralForType(literalType, "s3://blah/blah/blah")
+		assert.NoError(t, err)
+
+		assert.Equal(t, expectedLV, lv)
+
+		expectedVal, err := ExtractFromLiteral(expectedLV)
+		assert.NoError(t, err)
+		actualVal, err := ExtractFromLiteral(lv)
+		assert.NoError(t, err)
+		assert.Equal(t, expectedVal, actualVal)
+	})
+
 	t.Run("UnsupportedType", func(t *testing.T) {
 		var literalType = &core.LiteralType{Type: &core.LiteralType_Schema{
 			Schema: &core.SchemaType{}}}