check uri exists

Signed-off-by: Yicheng-Lu-llll <[email protected]>

dataproxy/service.go

@@ -132,6 +132,18 @@ func (s Service) CreateDownloadLink(ctx context.Context, req *service.CreateDown
 		return nil, errors.NewFlyteAdminErrorf(codes.Internal, "no deckUrl found for request [%+v]", req)
 	}
 
+	// Check if the native url exists
+	metadata, err := s.dataStore.Head(ctx, storage.DataReference(nativeURL))
+	if err != nil {
+		return nil, errors.NewFlyteAdminErrorf(codes.Internal, "Failed to check the existence of the URL [%s]. Error: %v", nativeURL, err)
+	}
+	if !metadata.Exists() {
+		return nil, errors.NewFlyteAdminErrorf(
+			codes.NotFound,
+			"URL [%s] does not exist yet. Please try again later. If you are using the real-time deck, this could be because the 'persist' function has not been called yet.",
+			nativeURL)
+	}
+
 	signedURLResp, err := s.dataStore.CreateSignedURL(ctx, storage.DataReference(nativeURL), storage.SignedURLProperties{
 		Scope:     stow.ClientMethodGet,
 		ExpiresIn: req.ExpiresIn.AsDuration(),

dataproxy/service_test.go

@@ -87,6 +87,9 @@ func TestCreateUploadLocation(t *testing.T) {
 
 func TestCreateDownloadLink(t *testing.T) {
 	dataStore := commonMocks.GetMockStorageClient()
+	dataStore.ComposedProtobufStore.(*commonMocks.TestDataStore).HeadCb = func(ctx context.Context, reference storage.DataReference) (storage.Metadata, error) {
+		return existsMetadata{}, nil
+	}
 	nodeExecutionManager := &mocks.MockNodeExecutionManager{}
 	nodeExecutionManager.SetGetNodeExecutionFunc(func(ctx context.Context, request admin.NodeExecutionGetRequest) (*admin.NodeExecution, error) {
 		return &admin.NodeExecution{
@@ -127,6 +130,19 @@ func TestCreateDownloadLink(t *testing.T) {
 		})
 		assert.NoError(t, err)
 	})
+
+	t.Run("nonexistent URI", func(t *testing.T) {
+		dataStore.ComposedProtobufStore.(*commonMocks.TestDataStore).HeadCb = func(ctx context.Context, reference storage.DataReference) (storage.Metadata, error) {
+			return nonexistentMetadata{}, nil
+		}
+		_, err = s.CreateDownloadLink(context.Background(), &service.CreateDownloadLinkRequest{
+			ArtifactType: service.ArtifactType_ARTIFACT_TYPE_DECK,
+			Source: &service.CreateDownloadLinkRequest_NodeExecutionId{
+				NodeExecutionId: &core.NodeExecutionIdentifier{},
+			},
+		})
+		assert.Error(t, err)
+	})
 }
 
 func TestCreateDownloadLocation(t *testing.T) {
@@ -350,3 +366,23 @@ func TestService_Error(t *testing.T) {
 		assert.Error(t, err, "no task executions")
 	})
 }
+
+type existsMetadata struct{}
+
+func (e existsMetadata) Exists() bool {
+	return true
+}
+
+func (e existsMetadata) Size() int64 {
+	return int64(1)
+}
+
+type nonexistentMetadata struct{}
+
+func (e nonexistentMetadata) Exists() bool {
+	return false
+}
+
+func (e nonexistentMetadata) Size() int64 {
+	return int64(0)
+}