From 8f5ebef3a0baf5382a349aa138d0dafee96a01ca Mon Sep 17 00:00:00 2001
From: Suhas Karanth <sudo.suhas@gmail.com>
Date: Wed, 12 Oct 2022 13:17:00 +0530
Subject: [PATCH] fix: apply asset URN filter for fetching probes

---
 internal/store/postgres/asset_repository.go      |  1 +
 internal/store/postgres/asset_repository_test.go | 14 ++++++++++++++
 2 files changed, 15 insertions(+)

diff --git a/internal/store/postgres/asset_repository.go b/internal/store/postgres/asset_repository.go
index cb1a1702..41ad9f43 100644
--- a/internal/store/postgres/asset_repository.go
+++ b/internal/store/postgres/asset_repository.go
@@ -352,6 +352,7 @@ func (r *AssetRepository) GetProbes(ctx context.Context, assetURN string) ([]ass
 		"id", "asset_urn", "status", "status_reason", "metadata", "timestamp", "created_at",
 	).From("asset_probes").
 		OrderBy("created_at").
+		Where(sq.Eq{"asset_urn": assetURN}).
 		PlaceholderFormat(sq.Dollar).
 		ToSql()
 	if err != nil {
diff --git a/internal/store/postgres/asset_repository_test.go b/internal/store/postgres/asset_repository_test.go
index a9676c51..c736f12c 100644
--- a/internal/store/postgres/asset_repository_test.go
+++ b/internal/store/postgres/asset_repository_test.go
@@ -1266,12 +1266,23 @@ func (r *AssetRepositoryTestSuite) TestAddProbe() {
 			Service:   "optimus",
 			UpdatedBy: user.User{ID: defaultAssetUpdaterUserID},
 		}
+		otherAst := asset.Asset{
+			URN:       "urn-add-probe-3",
+			Type:      asset.TypeJob,
+			Service:   "airflow",
+			UpdatedBy: user.User{ID: defaultAssetUpdaterUserID},
+		}
 		probe := asset.Probe{
 			Status: "RUNNING",
 		}
+		otherProbe := asset.Probe{
+			Status: "RUNNING",
+		}
 
 		_, err := r.repository.Upsert(r.ctx, &ast)
 		r.Require().NoError(err)
+		_, err = r.repository.Upsert(r.ctx, &otherAst)
+		r.Require().NoError(err)
 
 		err = r.repository.AddProbe(r.ctx, ast.URN, &probe)
 		r.NoError(err)
@@ -1280,6 +1291,9 @@ func (r *AssetRepositoryTestSuite) TestAddProbe() {
 		r.False(probe.CreatedAt.IsZero())
 		r.Equal(probe.CreatedAt, probe.Timestamp)
 
+		err = r.repository.AddProbe(r.ctx, otherAst.URN, &otherProbe)
+		r.NoError(err)
+
 		// assert probe is persisted
 		probesFromDB, err := r.repository.GetProbes(r.ctx, ast.URN)
 		r.Require().NoError(err)