Refacto vendor associations queries

Signed-off-by: Sacha Al Himdani <sacha@getprobo.com>

apps/console/src/pages/organizations/assets/AssetView.tsx

@@ -208,7 +208,7 @@ function AssetViewContent({
           assetType: formData.assetType,
           dataTypesStored: formData.dataTypesStored,
           ownerId: formData.ownerId,
-          vendorIds: formData.selectedVendorIds.length > 0 ? formData.selectedVendorIds : undefined,
+          vendorIds: formData.selectedVendorIds,
         },
       },
       onCompleted: () => {

apps/console/src/pages/organizations/data/DatumView.tsx

@@ -190,7 +190,7 @@ function DatumViewContent({
           name: formData.name,
           dataClassification: formData.dataClassification,
           ownerId: formData.ownerId,
-          vendorIds: formData.selectedVendorIds.length > 0 ? formData.selectedVendorIds : undefined,
+          vendorIds: formData.selectedVendorIds,
         },
       },
       onCompleted: () => {

compose.yaml

@@ -1,6 +1,6 @@
 services:
   postgres:
-    image: "postgres:16"
+    image: "postgres:17.4"
     shm_size: "1g"
     command: >
       postgres -c "shared_buffers=4GB"

pkg/coredata/asset.go

@@ -365,96 +365,3 @@ WHERE
 
 	return nil
 }
-
-type AssetVendor struct {
-	AssetID   int64     `db:"asset_id"`
-	VendorID  int64     `db:"vendor_id"`
-	CreatedAt time.Time `db:"created_at"`
-}
-
-func (a *Asset) CreateWithVendors(
-	ctx context.Context,
-	conn pg.Conn,
-	scope Scoper,
-	organization *Organization,
-	vendorIDs []gid.GID,
-	now time.Time,
-) error {
-	if err := organization.LoadByID(ctx, conn, scope, a.OrganizationID); err != nil {
-		return fmt.Errorf("cannot load organization %q: %w", a.OrganizationID, err)
-	}
-
-	if err := a.Insert(ctx, conn, scope); err != nil {
-		return fmt.Errorf("cannot insert asset: %w", err)
-	}
-
-	if len(vendorIDs) > 0 {
-		for _, vendorID := range vendorIDs {
-			_, err := conn.Exec(ctx, `
-				INSERT INTO asset_vendors (tenant_id, asset_id, vendor_id, created_at)
-				VALUES ($1, $2, $3, $4)
-			`, scope.GetTenantID(), a.ID, vendorID, now)
-			if err != nil {
-				return fmt.Errorf("cannot insert asset vendor: %w", err)
-			}
-		}
-	}
-
-	return nil
-}
-
-func (a *Asset) UpdateWithVendors(
-	ctx context.Context,
-	conn pg.Conn,
-	scope Scoper,
-	vendorIDs []gid.GID,
-	now time.Time,
-) error {
-	existing := &Asset{}
-	if err := existing.LoadByID(ctx, conn, scope, a.ID); err != nil {
-		return fmt.Errorf("cannot load asset: %w", err)
-	}
-
-	a.OrganizationID = existing.OrganizationID
-	a.CreatedAt = existing.CreatedAt
-	a.UpdatedAt = now
-
-	if err := a.Update(ctx, conn, scope); err != nil {
-		return fmt.Errorf("cannot update asset: %w", err)
-	}
-
-	_, err := conn.Exec(ctx, `
-		DELETE FROM asset_vendors
-		WHERE tenant_id = $1 AND asset_id = $2
-	`, scope.GetTenantID(), a.ID)
-	if err != nil {
-		return fmt.Errorf("cannot delete asset vendors: %w", err)
-	}
-
-	if len(vendorIDs) > 0 {
-		for _, vendorID := range vendorIDs {
-			_, err := conn.Exec(ctx, `
-				INSERT INTO asset_vendors (tenant_id, asset_id, vendor_id, created_at)
-				VALUES ($1, $2, $3, $4)
-			`, scope.GetTenantID(), a.ID, vendorID, now)
-			if err != nil {
-				return fmt.Errorf("cannot insert asset vendor: %w", err)
-			}
-		}
-	}
-
-	return nil
-}
-
-// UpdateWithVendorsTx updates an asset and its vendor relationships in a single transaction
-func (a *Asset) UpdateWithVendorsTx(
-	ctx context.Context,
-	db *pg.Client,
-	scope Scoper,
-	vendorIDs []gid.GID,
-	now time.Time,
-) error {
-	return db.WithTx(ctx, func(conn pg.Conn) error {
-		return a.UpdateWithVendors(ctx, conn, scope, vendorIDs, now)
-	})
-}

pkg/coredata/asset_vendor.go

@@ -0,0 +1,114 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package coredata
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/getprobo/probo/pkg/gid"
+	"github.com/jackc/pgx/v5"
+	"go.gearno.de/kit/pg"
+)
+
+type (
+	AssetVendor struct {
+		AssetID   gid.GID      `db:"asset_id"`
+		VendorID  gid.GID      `db:"vendor_id"`
+		TenantID  gid.TenantID `db:"tenant_id"`
+		CreatedAt time.Time    `db:"created_at"`
+	}
+
+	AssetVendors []*AssetVendor
+)
+
+func (av AssetVendors) Merge(
+	ctx context.Context,
+	conn pg.Conn,
+	scope Scoper,
+	assetID gid.GID,
+	vendorIDs []gid.GID,
+) error {
+	q := `
+WITH vendor_ids AS (
+	SELECT
+		unnest(@vendor_ids::text[]) AS vendor_id,
+		@tenant_id AS tenant_id,
+		@asset_id AS asset_id,
+		@created_at::timestamptz AS created_at
+)
+MERGE INTO asset_vendors AS tgt
+USING vendor_ids AS src
+ON tgt.tenant_id = src.tenant_id
+	AND tgt.asset_id = src.asset_id
+	AND tgt.vendor_id = src.vendor_id
+WHEN NOT MATCHED
+	THEN INSERT (tenant_id, asset_id, vendor_id, created_at)
+		VALUES (src.tenant_id, src.asset_id, src.vendor_id, src.created_at)
+	WHEN NOT MATCHED BY SOURCE
+		AND tgt.tenant_id = @tenant_id AND tgt.asset_id = @asset_id
+		THEN DELETE
+	`
+
+	args := pgx.StrictNamedArgs{
+		"tenant_id":  scope.GetTenantID(),
+		"asset_id":   assetID,
+		"created_at": time.Now(),
+		"vendor_ids": vendorIDs,
+	}
+
+	_, err := conn.Exec(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot merge asset vendors: %w", err)
+	}
+
+	return nil
+}
+
+func (av AssetVendors) Insert(
+	ctx context.Context,
+	conn pg.Conn,
+	scope Scoper,
+	assetID gid.GID,
+	vendorIDs []gid.GID,
+) error {
+	q := `
+WITH vendor_ids AS (
+	SELECT unnest(@vendor_ids::text[]) AS vendor_id
+)
+INSERT INTO asset_vendors (tenant_id, asset_id, vendor_id, created_at)
+SELECT
+	@tenant_id AS tenant_id,
+	@asset_id AS asset_id,
+	vendor_id,
+	@created_at AS created_at
+FROM vendor_ids
+`
+
+	args := pgx.StrictNamedArgs{
+		"tenant_id":  scope.GetTenantID(),
+		"asset_id":   assetID,
+		"created_at": time.Now(),
+		"vendor_ids": vendorIDs,
+	}
+
+	_, err := conn.Exec(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot insert asset vendors: %w", err)
+	}
+
+	return nil
+}

pkg/coredata/datum.go

@@ -293,7 +293,7 @@ RETURNING
 		"name":                d.Name,
 		"owner_id":            d.OwnerID,
 		"data_classification": d.DataClassification,
-		"updated_at":          time.Now(),
+		"updated_at":          d.UpdatedAt,
 	}
 	maps.Copy(args, scope.SQLArguments())
 
@@ -336,89 +336,3 @@ WHERE
 
 	return nil
 }
-
-type DataVendor struct {
-	DatumID   gid.GID   `db:"datum_id"`
-	VendorID  gid.GID   `db:"vendor_id"`
-	CreatedAt time.Time `db:"created_at"`
-}
-
-func (d *Datum) CreateWithVendors(
-	ctx context.Context,
-	conn pg.Conn,
-	scope Scoper,
-	vendorIDs []gid.GID,
-	now time.Time,
-) error {
-	if err := d.Insert(ctx, conn, scope); err != nil {
-		return fmt.Errorf("cannot insert data: %w", err)
-	}
-
-	if len(vendorIDs) > 0 {
-		for _, vendorID := range vendorIDs {
-			_, err := conn.Exec(ctx, `
-				INSERT INTO data_vendors (tenant_id, datum_id, vendor_id, created_at)
-				VALUES ($1, $2, $3, $4)
-			`, scope.GetTenantID(), d.ID, vendorID, now)
-			if err != nil {
-				return fmt.Errorf("cannot insert data vendor: %w", err)
-			}
-		}
-	}
-
-	return nil
-}
-
-func (d *Datum) UpdateWithVendors(
-	ctx context.Context,
-	conn pg.Conn,
-	scope Scoper,
-	vendorIDs []gid.GID,
-	now time.Time,
-) error {
-	existing := &Datum{}
-	if err := existing.LoadByID(ctx, conn, scope, d.ID); err != nil {
-		return fmt.Errorf("cannot load data: %w", err)
-	}
-
-	d.CreatedAt = existing.CreatedAt
-	d.UpdatedAt = now
-
-	if err := d.Update(ctx, conn, scope); err != nil {
-		return fmt.Errorf("cannot update data: %w", err)
-	}
-
-	_, err := conn.Exec(ctx, `
-		DELETE FROM data_vendors
-		WHERE tenant_id = $1 AND datum_id = $2
-	`, scope.GetTenantID(), d.ID)
-	if err != nil {
-		return fmt.Errorf("cannot delete data vendors: %w", err)
-	}
-
-	if len(vendorIDs) > 0 {
-		for _, vendorID := range vendorIDs {
-			_, err := conn.Exec(ctx, `
-				INSERT INTO data_vendors (tenant_id, datum_id, vendor_id, created_at)
-				VALUES ($1, $2, $3, $4)
-			`, scope.GetTenantID(), d.ID, vendorID, now)
-			if err != nil {
-				return fmt.Errorf("cannot insert data vendor: %w", err)
-			}
-		}
-	}
-
-	return nil
-}
-
-func (d *Datum) UpdateWithVendorsTx(
-	ctx context.Context,
-	db *pg.Client,
-	scope Scoper,
-	vendorIDs []gid.GID,
-	now time.Time,
-) error {
-	return db.WithTx(ctx, func(conn pg.Conn) error {
-		return d.UpdateWithVendors(ctx, conn, scope, vendorIDs, now)
-	})
-}

pkg/coredata/datum_vendor.go

@@ -0,0 +1,114 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package coredata
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/getprobo/probo/pkg/gid"
+	"github.com/jackc/pgx/v5"
+	"go.gearno.de/kit/pg"
+)
+
+type (
+	DatumVendor struct {
+		DatumID   gid.GID      `db:"datum_id"`
+		VendorID  gid.GID      `db:"vendor_id"`
+		TenantID  gid.TenantID `db:"tenant_id"`
+		CreatedAt time.Time    `db:"created_at"`
+	}
+
+	DatumVendors []*DatumVendor
+)
+
+func (dv DatumVendors) Merge(
+	ctx context.Context,
+	conn pg.Conn,
+	scope Scoper,
+	datumID gid.GID,
+	vendorIDs []gid.GID,
+) error {
+	q := `
+WITH vendor_ids AS (
+	SELECT
+		unnest(@vendor_ids::text[]) AS vendor_id,
+		@tenant_id AS tenant_id,
+		@datum_id AS datum_id,
+		@created_at::timestamptz AS created_at
+)
+MERGE INTO data_vendors AS tgt
+USING vendor_ids AS src
+ON tgt.tenant_id = src.tenant_id
+	AND tgt.datum_id = src.datum_id
+	AND tgt.vendor_id = src.vendor_id
+WHEN NOT MATCHED THEN
+	INSERT (tenant_id, datum_id, vendor_id, created_at)
+	VALUES (src.tenant_id, src.datum_id, src.vendor_id, src.created_at)
+WHEN NOT MATCHED BY SOURCE
+	AND tgt.tenant_id = @tenant_id AND tgt.datum_id = @datum_id
+	THEN DELETE
+	`
+
+	args := pgx.StrictNamedArgs{
+		"tenant_id":  scope.GetTenantID(),
+		"datum_id":   datumID,
+		"created_at": time.Now(),
+		"vendor_ids": vendorIDs,
+	}
+
+	_, err := conn.Exec(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot merge data vendors: %w", err)
+	}
+
+	return nil
+}
+
+func (dv DatumVendors) Insert(
+	ctx context.Context,
+	conn pg.Conn,
+	scope Scoper,
+	datumID gid.GID,
+	vendorIDs []gid.GID,
+) error {
+	q := `
+WITH vendor_ids AS (
+	SELECT unnest(@vendor_ids::text[]) AS vendor_id
+)
+INSERT INTO data_vendors (tenant_id, datum_id, vendor_id, created_at)
+SELECT
+	@tenant_id::text AS tenant_id,
+	@datum_id::text AS datum_id,
+	vendor_id,
+	@created_at::timestamptz AS created_at
+FROM vendor_ids
+`
+
+	args := pgx.StrictNamedArgs{
+		"tenant_id":  scope.GetTenantID(),
+		"datum_id":   datumID,
+		"created_at": time.Now(),
+		"vendor_ids": vendorIDs,
+	}
+
+	_, err := conn.Exec(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot insert data vendors: %w", err)
+	}
+
+	return nil
+}

pkg/probo/asset_service.go

@@ -135,28 +135,15 @@ func (s AssetService) Update(
 	req UpdateAssetRequest,
 ) (*coredata.Asset, error) {
 	now := time.Now()
+	asset := &coredata.Asset{ID: req.ID}
+	assetVendors := &coredata.AssetVendors{}
 
-	existing := &coredata.Asset{}
+	err := s.svc.pg.WithTx(ctx, func(conn pg.Conn) error {
-	if err := s.svc.pg.WithConn(ctx, func(conn pg.Conn) error {
+		if err := asset.LoadByID(ctx, conn, s.svc.scope, req.ID); err != nil {
-		return existing.LoadByID(ctx, conn, s.svc.scope, req.ID)
+			return fmt.Errorf("cannot load asset: %w", err)
-	}); err != nil {
-		return nil, fmt.Errorf("cannot load asset: %w", err)
 		}
 
-	asset := &coredata.Asset{
+		asset.UpdatedAt = now
-		ID:              req.ID,
-		OrganizationID:  existing.OrganizationID,
-		Name:            existing.Name,
-		Amount:          existing.Amount,
-		OwnerID:         existing.OwnerID,
-		Criticity:       existing.Criticity,
-		AssetType:       existing.AssetType,
-		DataTypesStored: existing.DataTypesStored,
-		CreatedAt:       existing.CreatedAt,
-		UpdatedAt:       now,
-	}
-
-	// Update fields from request
 		if req.Name != nil {
 			asset.Name = *req.Name
 		}
@@ -176,7 +163,20 @@ func (s AssetService) Update(
 			asset.DataTypesStored = *req.DataTypesStored
 		}
 
-	if err := asset.UpdateWithVendorsTx(ctx, s.svc.pg, s.svc.scope, req.VendorIDs, now); err != nil {
+		if err := asset.Update(ctx, conn, s.svc.scope); err != nil {
+			return fmt.Errorf("cannot update asset: %w", err)
+		}
+
+		if req.VendorIDs != nil {
+			if err := assetVendors.Merge(ctx, conn, s.svc.scope, asset.ID, req.VendorIDs); err != nil {
+				return fmt.Errorf("cannot update asset vendors: %w", err)
+			}
+		}
+
+		return nil
+	})
+
+	if err != nil {
 		return nil, err
 	}
 
@@ -189,8 +189,8 @@ func (s AssetService) Create(
 ) (*coredata.Asset, error) {
 	now := time.Now()
 	assetID := gid.New(s.svc.scope.GetTenantID(), coredata.AssetEntityType)
+	assetVendors := &coredata.AssetVendors{}
 
-	organization := &coredata.Organization{}
 	asset := &coredata.Asset{
 		ID:              assetID,
 		OrganizationID:  req.OrganizationID,
@@ -204,12 +204,19 @@ func (s AssetService) Create(
 		UpdatedAt:       now,
 	}
 
-	err := s.svc.pg.WithTx(
+	err := s.svc.pg.WithTx(ctx, func(conn pg.Conn) error {
-		ctx,
+		if err := asset.Insert(ctx, conn, s.svc.scope); err != nil {
-		func(conn pg.Conn) error {
+			return fmt.Errorf("cannot insert asset: %w", err)
-			return asset.CreateWithVendors(ctx, conn, s.svc.scope, organization, req.VendorIDs, now)
+		}
-		},
+
-	)
+		if len(req.VendorIDs) > 0 {
+			if err := assetVendors.Insert(ctx, conn, s.svc.scope, asset.ID, req.VendorIDs); err != nil {
+				return fmt.Errorf("cannot create asset vendors: %w", err)
+			}
+		}
+
+		return nil
+	})
 
 	if err != nil {
 		return nil, err

pkg/probo/datum_service.go

@@ -143,25 +143,14 @@ func (s DatumService) Update(
 	req UpdateDatumRequest,
 ) (*coredata.Datum, error) {
 	now := time.Now()
+	datum := &coredata.Datum{}
+	datumVendors := &coredata.DatumVendors{}
 
-	existing := &coredata.Datum{}
+	err := s.svc.pg.WithTx(ctx, func(conn pg.Conn) error {
-	if err := s.svc.pg.WithConn(ctx, func(conn pg.Conn) error {
+		if err := datum.LoadByID(ctx, conn, s.svc.scope, req.ID); err != nil {
-		return existing.LoadByID(ctx, conn, s.svc.scope, req.ID)
+			return fmt.Errorf("cannot load data: %w", err)
-	}); err != nil {
-		return nil, fmt.Errorf("cannot load data: %w", err)
 		}
 
-	datum := &coredata.Datum{
-		ID:                 req.ID,
-		OrganizationID:     existing.OrganizationID,
-		Name:               existing.Name,
-		DataClassification: existing.DataClassification,
-		OwnerID:            existing.OwnerID,
-		CreatedAt:          existing.CreatedAt,
-		UpdatedAt:          now,
-	}
-
-	// Update fields from request
 		if req.Name != nil {
 			datum.Name = *req.Name
 		}
@@ -171,8 +160,22 @@ func (s DatumService) Update(
 		if req.OwnerID != nil {
 			datum.OwnerID = *req.OwnerID
 		}
+		datum.UpdatedAt = now
 
-	if err := datum.UpdateWithVendorsTx(ctx, s.svc.pg, s.svc.scope, req.VendorIDs, now); err != nil {
+		if err := datum.Update(ctx, conn, s.svc.scope); err != nil {
+			return fmt.Errorf("cannot update data: %w", err)
+		}
+
+		if req.VendorIDs != nil {
+			if err := datumVendors.Merge(ctx, conn, s.svc.scope, datum.ID, req.VendorIDs); err != nil {
+				return fmt.Errorf("cannot update data vendors: %w", err)
+			}
+		}
+
+		return nil
+	})
+
+	if err != nil {
 		return nil, err
 	}
 
@@ -185,6 +188,7 @@ func (s DatumService) Create(
 ) (*coredata.Datum, error) {
 	now := time.Now()
 	datumID := gid.New(s.svc.scope.GetTenantID(), coredata.DatumEntityType)
+	datumVendors := &coredata.DatumVendors{}
 
 	datum := &coredata.Datum{
 		ID:                 datumID,
@@ -199,7 +203,17 @@ func (s DatumService) Create(
 	err := s.svc.pg.WithTx(
 		ctx,
 		func(conn pg.Conn) error {
-			return datum.CreateWithVendors(ctx, conn, s.svc.scope, req.VendorIDs, now)
+			if err := datum.Insert(ctx, conn, s.svc.scope); err != nil {
+				return fmt.Errorf("cannot insert datum: %w", err)
+			}
+
+			if len(req.VendorIDs) > 0 {
+				if err := datumVendors.Insert(ctx, conn, s.svc.scope, datum.ID, req.VendorIDs); err != nil {
+					return fmt.Errorf("cannot create data vendors: %w", err)
+				}
+			}
+
+			return nil
 		},
 	)