From 8ef6441737ab53cb74594a35cedcfcc1d3a6e191 Mon Sep 17 00:00:00 2001
From: Jorge Padilla <jorge.esteban.padilla@gmail.com>
Date: Wed, 13 Sep 2023 16:13:31 -0600
Subject: [PATCH] wip

---
 server/datastore/datastore_repository.go         |  9 ++++++---
 server/migrations/33_add_composite_pkey.down.sql |  3 ---
 server/migrations/33_add_composite_pkey.up.sql   |  3 ---
 server/pkg/sqlutil/tenant.go                     | 14 ++++++++++++++
 4 files changed, 20 insertions(+), 9 deletions(-)
 delete mode 100644 server/migrations/33_add_composite_pkey.down.sql
 delete mode 100644 server/migrations/33_add_composite_pkey.up.sql

diff --git a/server/datastore/datastore_repository.go b/server/datastore/datastore_repository.go
index a9ed1c8b0e..2e6f143154 100644
--- a/server/datastore/datastore_repository.go
+++ b/server/datastore/datastore_repository.go
@@ -91,7 +91,7 @@ func (r *Repository) Update(ctx context.Context, dataStore DataStore) (DataStore
 	}
 	defer tx.Rollback()
 
-	query, params := sqlutil.Tenant(ctx, deleteQuery, DataStoreSingleID)
+	query, params := sqlutil.TenantWithReplacedID(ctx, deleteQuery, DataStoreSingleID)
 	_, err = tx.ExecContext(ctx, query, params...)
 	if err != nil {
 		return DataStore{}, fmt.Errorf("datastore repository sql exec delete: %w", err)
@@ -103,6 +103,9 @@ func (r *Repository) Update(ctx context.Context, dataStore DataStore) (DataStore
 	}
 
 	tenantID := sqlutil.TenantID(ctx)
+	if tenantID != nil {
+		dataStore.ID = id.ID(*tenantID)
+	}
 
 	_, err = tx.ExecContext(ctx, insertQuery,
 		dataStore.ID,
@@ -132,7 +135,7 @@ func (r *Repository) Delete(ctx context.Context, id id.ID) error {
 	}
 	defer tx.Rollback()
 
-	query, params := sqlutil.Tenant(ctx, deleteQuery, id)
+	query, params := sqlutil.TenantWithReplacedID(ctx, deleteQuery, id)
 	_, err = tx.ExecContext(ctx, query, params...)
 	if err != nil {
 		return fmt.Errorf("datastore repository sql exec delete: %w", err)
@@ -167,7 +170,7 @@ func (r *Repository) Current(ctx context.Context) (DataStore, error) {
 }
 
 func (r *Repository) Get(ctx context.Context, id id.ID) (DataStore, error) {
-	query, params := sqlutil.Tenant(ctx, getQuery, id)
+	query, params := sqlutil.TenantWithReplacedID(ctx, getQuery, id)
 	row := r.db.QueryRowContext(ctx, query, params...)
 
 	dataStore, err := r.readRow(row)
diff --git a/server/migrations/33_add_composite_pkey.down.sql b/server/migrations/33_add_composite_pkey.down.sql
deleted file mode 100644
index a9d31a009e..0000000000
--- a/server/migrations/33_add_composite_pkey.down.sql
+++ /dev/null
@@ -1,3 +0,0 @@
-ALTER TABLE data_stores
-DROP CONSTRAINT data_stores_pkey,
-ADD PRIMARY KEY (id);
\ No newline at end of file
diff --git a/server/migrations/33_add_composite_pkey.up.sql b/server/migrations/33_add_composite_pkey.up.sql
deleted file mode 100644
index 50be1e0988..0000000000
--- a/server/migrations/33_add_composite_pkey.up.sql
+++ /dev/null
@@ -1,3 +0,0 @@
-ALTER TABLE data_stores
-DROP CONSTRAINT data_stores_pkey,
-ADD PRIMARY KEY (id, tenant_id);
\ No newline at end of file
diff --git a/server/pkg/sqlutil/tenant.go b/server/pkg/sqlutil/tenant.go
index 5ee968827b..eb9d69ad0e 100644
--- a/server/pkg/sqlutil/tenant.go
+++ b/server/pkg/sqlutil/tenant.go
@@ -35,6 +35,20 @@ func TenantWithPrefix(ctx context.Context, query string, prefix string, params .
 	return query + condition, append(params, *tenantID)
 }
 
+func TenantWithReplacedID(ctx context.Context, query string, params ...any) (string, []any) {
+	tenantID := TenantID(ctx)
+	if tenantID == nil {
+		return query, params
+	}
+
+	prefix := getQueryPrefix(query)
+	paramNumber := len(params) + 1
+	condition := fmt.Sprintf(" %s tenant_id = $%d", prefix, paramNumber)
+
+	var newParams []any
+	return query + condition, append(newParams, *tenantID, *tenantID)
+}
+
 func TenantID(ctx context.Context) *string {
 	tenantID := ctx.Value(middleware.TenantIDKey)