2025-04-02 10:16:07 +08:00

202 lines
5.0 KiB
Go

package db
import (
"context"
"errors"
"fmt"
"management/internal/pkg/config"
"github.com/jackc/pgconn"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
)
// ****************** conn ******************
func newDsn(conf config.DB) string {
return fmt.Sprintf("postgresql://%s:%s@%s:%d/%s?sslmode=disable",
conf.Username,
conf.Password,
conf.Host,
conf.Port,
conf.DBName)
}
// ****************** errors ******************
const (
foreignKeyViolation = "23503"
uniqueViolation = "23505"
)
// var ErrUniqueViolation = &pgconn.PgError{
// Code: UniqueViolation,
// }
func ErrorCode(err error) string {
var pgErr *pgconn.PgError
if errors.As(err, &pgErr) {
return pgErr.Code
}
return ""
}
// ****************** store ******************
var Engine Store
// Store defines all functions to execute db queries and transactions
type Store interface {
Querier
ExecTx(context.Context, func(*Queries) error) error
Pool() *pgxpool.Pool
SysDepartmentRebuildPath(ctx context.Context) error
SysRoleRebuildPath(ctx context.Context) error
CategoryRebuildPath(ctx context.Context) error
ListCustomerCondition(ctx context.Context, arg *ListCustomerConditionParam) ([]*CustomerView, int64, error)
ListProjectCondition(ctx context.Context, arg *ListProjectConditionParam) ([]*ProjectView, int64, error)
ListBudgetCondition(ctx context.Context, arg *ListBudgetConditionParam) ([]*BudgetView, int64, error)
ListIncomeCondition(ctx context.Context, arg *ListIncomeConditionParam) ([]*IncomeView, int64, error)
ListExpenseCondition(ctx context.Context, arg *ListExpenseConditionParam) ([]*ExpenseView, int64, error)
}
// SQLStore provides all functions to execute SQL queries and transactions
type SQLStore struct {
connPool *pgxpool.Pool
*Queries
}
func NewIStore(ctx context.Context, conf config.DB) (Store, error) {
pool, err := pgxpool.New(ctx, newDsn(conf))
if err != nil {
return nil, err
}
err = pool.Ping(ctx)
if err != nil {
return nil, err
}
return &SQLStore{
connPool: pool,
Queries: New(pool),
}, nil
}
// NewStore creates a new store
func NewStore(ctx context.Context) error {
pool, err := pgxpool.New(ctx, newDsn(config.File.DB))
if err != nil {
return err
}
err = pool.Ping(ctx)
if err != nil {
return err
}
Engine = &SQLStore{
connPool: pool,
Queries: New(pool),
}
return nil
}
func (store *SQLStore) Pool() *pgxpool.Pool {
return store.connPool
}
func IsUniqueViolation(err error) bool {
return ErrorCode(err) == uniqueViolation
}
func IsForeignKeyViolation(err error) bool {
return ErrorCode(err) == foreignKeyViolation
}
func IsNoRows(err error) bool {
return errors.Is(err, pgx.ErrNoRows)
}
func (store *SQLStore) ExecTx(ctx context.Context, fn func(*Queries) error) error {
tx, err := store.connPool.Begin(ctx)
if err != nil {
return err
}
q := New(tx)
err = fn(q)
if err != nil {
if rbErr := tx.Rollback(ctx); rbErr != nil {
return fmt.Errorf("tx err: %v, rb err: %v", err, rbErr)
}
return err
}
return tx.Commit(ctx)
}
func (store *SQLStore) SysDepartmentRebuildPath(ctx context.Context) error {
query := `
UPDATE sys_department AS tm
SET parent_path = (SELECT ',' || string_agg(cast(t.parent_id AS VARCHAR), ',') || ','
FROM (WITH RECURSIVE temp (id, parent_id) AS (SELECT id, tm.parent_id
FROM sys_department
WHERE id = tm.id
UNION ALL
SELECT sys_department.id, sys_department.parent_id
FROM sys_department,
temp
WHERE sys_department.id = temp.parent_id)
SELECT id, parent_id
FROM temp
ORDER BY id) AS t)
WHERE tm.status = 0 ;`
_, err := store.db.Exec(ctx, query)
return err
}
func (store *SQLStore) SysRoleRebuildPath(ctx context.Context) error {
query := `
UPDATE sys_role AS tm
SET parent_path = (SELECT ',' || string_agg(cast(t.parent_id AS VARCHAR), ',') || ','
FROM (WITH RECURSIVE temp (id, parent_id) AS (SELECT id, tm.parent_id
FROM sys_role
WHERE id = tm.id
UNION ALL
SELECT sys_role.id, sys_role.parent_id
FROM sys_role,
temp
WHERE sys_role.id = temp.parent_id)
SELECT id, parent_id
FROM temp
ORDER BY id) AS t)
WHERE tm.status = 0;`
_, err := store.db.Exec(ctx, query)
return err
}
func (store *SQLStore) CategoryRebuildPath(ctx context.Context) error {
query := `
UPDATE categories AS tm
SET parent_path = (SELECT ',' || string_agg(cast(t.parent_id AS VARCHAR), ',') || ','
FROM (WITH RECURSIVE temp (id, parent_id) AS (SELECT id, tm.parent_id
FROM categories
WHERE id = tm.id
UNION ALL
SELECT categories.id, categories.parent_id
FROM categories,
temp
WHERE categories.id = temp.parent_id)
SELECT id, parent_id
FROM temp
ORDER BY id) AS t)
WHERE tm.status = 0;`
_, err := store.db.Exec(ctx, query)
return err
}