From e32777fcf7314fa40ba27985147c8473f58e352f Mon Sep 17 00:00:00 2001
From: Igor Shishkin <me@teran.dev>
Date: Sat, 13 Jul 2024 12:18:25 +0300
Subject: [PATCH] Add EnsureBlobKey method to metadata repository (#44)

Signed-off-by: Igor Shishkin <me@teran.dev>
---
 repositories/metadata/metadata.go             |  1 +
 repositories/metadata/mock/mock.go            |  5 +++++
 repositories/metadata/postgresql/blobs.go     | 19 +++++++++++++++++++
 .../metadata/postgresql/blobs_test.go         | 12 ++++++++++++
 4 files changed, 37 insertions(+)

diff --git a/repositories/metadata/metadata.go b/repositories/metadata/metadata.go
index 1b7642a..805c2d5 100644
--- a/repositories/metadata/metadata.go
+++ b/repositories/metadata/metadata.go
@@ -25,4 +25,5 @@ type Repository interface {
 
 	CreateBLOB(ctx context.Context, checksum string, size uint64, mimeType string) error
 	GetBlobKeyByObject(ctx context.Context, container, version, key string) (string, error)
+	EnsureBlobKey(ctx context.Context, key string, size uint64) error
 }
diff --git a/repositories/metadata/mock/mock.go b/repositories/metadata/mock/mock.go
index 3af7372..14cdc2b 100644
--- a/repositories/metadata/mock/mock.go
+++ b/repositories/metadata/mock/mock.go
@@ -81,3 +81,8 @@ func (m *Mock) GetBlobKeyByObject(_ context.Context, container, version, key str
 	args := m.Called(container, version, key)
 	return args.String(0), args.Error(1)
 }
+
+func (m *Mock) EnsureBlobKey(_ context.Context, key string, size uint64) error {
+	args := m.Called(key, size)
+	return args.Error(0)
+}
diff --git a/repositories/metadata/postgresql/blobs.go b/repositories/metadata/postgresql/blobs.go
index be99894..59d79bc 100644
--- a/repositories/metadata/postgresql/blobs.go
+++ b/repositories/metadata/postgresql/blobs.go
@@ -49,3 +49,22 @@ func (r *repository) GetBlobKeyByObject(ctx context.Context, container, version,
 
 	return checksum, nil
 }
+
+func (r *repository) EnsureBlobKey(ctx context.Context, key string, size uint64) error {
+	row := psql.
+		Select("id").
+		From("blobs").
+		Where(sq.Eq{
+			"checksum": key,
+			"size":     size,
+		}).
+		RunWith(r.db).
+		QueryRowContext(ctx)
+
+	var blobID uint
+	if err := row.Scan(&blobID); err != nil {
+		return mapSQLErrors(err)
+	}
+
+	return nil
+}
diff --git a/repositories/metadata/postgresql/blobs_test.go b/repositories/metadata/postgresql/blobs_test.go
index 18022b7..582c1ae 100644
--- a/repositories/metadata/postgresql/blobs_test.go
+++ b/repositories/metadata/postgresql/blobs_test.go
@@ -62,3 +62,15 @@ func (s *postgreSQLRepositoryTestSuite) TestGetBlobKeyByObjectErrors() {
 	s.Require().Error(err)
 	s.Require().Equal(metadata.ErrNotFound, err)
 }
+
+func (s *postgreSQLRepositoryTestSuite) TestEnsureBlobKey() {
+	err := s.repo.EnsureBlobKey(s.ctx, "deadbeef", 1234)
+	s.Require().Error(err)
+	s.Require().Equal(metadata.ErrNotFound, err)
+
+	err = s.repo.CreateBLOB(s.ctx, "deadbeef", 1234, "application/octet-stream")
+	s.Require().NoError(err)
+
+	err = s.repo.EnsureBlobKey(s.ctx, "deadbeef", 1234)
+	s.Require().NoError(err)
+}