rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refacto vendor associations queries

Browse Source 
Signed-off-by: Sacha Al Himdani <sacha@getprobo.com>
This commit is contained in:
Sacha Al Himdani
2025-06-10 14:50:29 -07:00
committed by Bryan Frimin
parent e547a4e96e
commit c7f6dba576
10 changed files with 325 additions and 255 deletions
Download Patch File Download Diff File Expand all files Collapse all files

93
pkg/coredata/asset.go
View File

@@ -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)
})
}

114
pkg/coredata/asset_vendor.go Normal file
View File

@@ -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
}

88
pkg/coredata/data.go → pkg/coredata/datum.go
View File

@@ -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)
})
}

114
pkg/coredata/datum_vendor.go Normal file
View File

@@ -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
}

97
pkg/probo/asset_service.go
View File

@@ -135,48 +135,48 @@ func (s AssetService) Update(
req UpdateAssetRequest,
) (*coredata.Asset, error) {
now := time.Now()
asset := &coredata.Asset{ID: req.ID}
assetVendors := &coredata.AssetVendors{}
existing := &coredata.Asset{}
if err := s.svc.pg.WithConn(ctx, func(conn pg.Conn) error {
return existing.LoadByID(ctx, conn, s.svc.scope, req.ID)
}); err != nil {
return nil, fmt.Errorf("cannot load asset: %w", err)
}
err := s.svc.pg.WithTx(ctx, func(conn pg.Conn) error {
if err := asset.LoadByID(ctx, conn, s.svc.scope, req.ID); err != nil {
return fmt.Errorf("cannot load asset: %w", err)
}
asset := &coredata.Asset{
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,
}
asset.UpdatedAt = now
if req.Name != nil {
asset.Name = *req.Name
}
if req.Amount != nil {
asset.Amount = *req.Amount
}
if req.OwnerID != nil {
asset.OwnerID = *req.OwnerID
}
if req.Criticity != nil {
asset.Criticity = *req.Criticity
}
if req.AssetType != nil {
asset.AssetType = *req.AssetType
}
if req.DataTypesStored != nil {
asset.DataTypesStored = *req.DataTypesStored
}
// Update fields from request
if req.Name != nil {
asset.Name = *req.Name
}
if req.Amount != nil {
asset.Amount = *req.Amount
}
if req.OwnerID != nil {
asset.OwnerID = *req.OwnerID
}
if req.Criticity != nil {
asset.Criticity = *req.Criticity
}
if req.AssetType != nil {
asset.AssetType = *req.AssetType
}
if req.DataTypesStored != nil {
asset.DataTypesStored = *req.DataTypesStored
}
if err := asset.Update(ctx, conn, s.svc.scope); err != nil {
return fmt.Errorf("cannot update asset: %w", err)
}
if err := asset.UpdateWithVendorsTx(ctx, s.svc.pg, s.svc.scope, req.VendorIDs, now); err != nil {
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(
ctx,
func(conn pg.Conn) error {
return asset.CreateWithVendors(ctx, conn, s.svc.scope, organization, req.VendorIDs, now)
},
)
err := s.svc.pg.WithTx(ctx, func(conn pg.Conn) error {
if err := asset.Insert(ctx, conn, s.svc.scope); err != nil {
return fmt.Errorf("cannot insert asset: %w", err)
}
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

68
pkg/probo/datum_service.go
View File

@@ -143,36 +143,39 @@ func (s DatumService) Update(
req UpdateDatumRequest,
) (*coredata.Datum, error) {
now := time.Now()
datum := &coredata.Datum{}
datumVendors := &coredata.DatumVendors{}
existing := &coredata.Datum{}
if err := s.svc.pg.WithConn(ctx, func(conn pg.Conn) error {
return existing.LoadByID(ctx, conn, s.svc.scope, req.ID)
}); err != nil {
return nil, fmt.Errorf("cannot load data: %w", err)
}
err := s.svc.pg.WithTx(ctx, func(conn pg.Conn) error {
if err := datum.LoadByID(ctx, conn, s.svc.scope, req.ID); err != nil {
return 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,
}
if req.Name != nil {
datum.Name = *req.Name
}
if req.DataClassification != nil {
datum.DataClassification = *req.DataClassification
}
if req.OwnerID != nil {
datum.OwnerID = *req.OwnerID
}
datum.UpdatedAt = now
// Update fields from request
if req.Name != nil {
datum.Name = *req.Name
}
if req.DataClassification != nil {
datum.DataClassification = *req.DataClassification
}
if req.OwnerID != nil {
datum.OwnerID = *req.OwnerID
}
if err := datum.Update(ctx, conn, s.svc.scope); err != nil {
return fmt.Errorf("cannot update data: %w", err)
}
if err := datum.UpdateWithVendorsTx(ctx, s.svc.pg, s.svc.scope, req.VendorIDs, now); err != nil {
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
},
)

0
pkg/server/api/console/v1/types/data.go → pkg/server/api/console/v1/types/datum.go
View File