probo/cmd/migrate-third-party-snapshots-to-documents/main.go

// Copyright (c) 2026 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.

// Command migrate-thirdParty-snapshots-to-documents creates documents and document
// versions from existing thirdParty snapshots. For each organization that has thirdParty
// snapshots, it generates a thirdParty list document using the same ProseMirror
// builder as the publish flow.
package main

import (
	"context"
	"flag"
	"fmt"
	"net"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
	"go.gearno.de/kit/pg"
	"go.probo.inc/probo/pkg/coredata"
	"go.probo.inc/probo/pkg/docgen"
	"go.probo.inc/probo/pkg/gid"
	"go.probo.inc/probo/pkg/probo"
)

func main() {
	if err := run(); err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}
}

func run() error {
	var (
		pgDSN  string
		dryRun bool
	)

	flag.StringVar(
		&pgDSN,
		"pg-dsn",
		os.Getenv("DATABASE_URL"),
		"PostgreSQL connection URL (default: DATABASE_URL env)",
	)
	flag.BoolVar(&dryRun, "dry-run", false, "show what would be done without writing")
	flag.Parse()

	if pgDSN == "" {
		return fmt.Errorf("set -pg-dsn or DATABASE_URL")
	}

	ctx := context.Background()

	pgClient, err := newPgClientFromDSN(pgDSN)
	if err != nil {
		return fmt.Errorf("cannot create pg client: %w", err)
	}

	return migrate(ctx, pgClient, dryRun)
}

type orgWithThirdPartySnapshots struct {
	organizationID   gid.GID
	tenantID         gid.TenantID
	organizationName string
}

type thirdPartySnapshot struct {
	snapshotID  string
	publishedAt time.Time
}

func migrate(ctx context.Context, pgClient *pg.Client, dryRun bool) error {
	var orgs []orgWithThirdPartySnapshots
	err := pgClient.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error {
		var err error
		orgs, err = loadOrgsWithThirdPartySnapshots(ctx, conn)
		return err
	})
	if err != nil {
		return err
	}

	if len(orgs) == 0 {
		fmt.Println("no organizations with thirdParty snapshots to migrate")
		return nil
	}

	var stats struct {
		documents, versions, failed int
	}

	for _, org := range orgs {
		if dryRun {
			var count int
			err := pgClient.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error {
				snapshots, err := loadThirdPartySnapshots(ctx, conn, org.organizationID)
				count = len(snapshots)
				return err
			})
			if err != nil {
				return err
			}
			fmt.Printf("would migrate org %s (%s) — %d thirdParty snapshot(s)\n",
				org.organizationID, org.organizationName, count)
			continue
		}

		err := pgClient.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error {
			return migrateOrg(ctx, tx, org)
		})
		if err != nil {
			fmt.Fprintf(os.Stderr, "FAIL org %s (%s): %v\n",
				org.organizationID, org.organizationName, err)
			stats.failed++
			continue
		}

		stats.documents++
	}

	if dryRun {
		fmt.Printf("\n%d organization(s) would be migrated\n", len(orgs))
		return nil
	}

	fmt.Printf("\nmigrated %d organization(s), %d failed\n",
		stats.documents, stats.failed)

	return nil
}

func migrateOrg(ctx context.Context, tx pg.Tx, org orgWithThirdPartySnapshots) error {
	snapshots, err := loadThirdPartySnapshots(ctx, tx, org.organizationID)
	if err != nil {
		return err
	}

	if len(snapshots) == 0 {
		return nil
	}

	documentID := gid.New(org.tenantID, coredata.DocumentEntityType)
	now := time.Now()

	_, err = tx.Exec(
		ctx,
		`
INSERT INTO documents (
    id, tenant_id, organization_id, write_mode,
    current_published_major, current_published_minor,
    trust_center_visibility, status, created_at, updated_at
) VALUES (
    @id, @tenant_id, @organization_id,
    'GENERATED'::document_write_mode,
    @current_published_major, 0,
    'NONE'::trust_center_visibility,
    'ACTIVE'::document_status,
    @created_at, @updated_at
)`,
		pgx.NamedArgs{
			"id":                      documentID,
			"tenant_id":               org.tenantID,
			"organization_id":         org.organizationID,
			"current_published_major": len(snapshots),
			"created_at":              now,
			"updated_at":              now,
		},
	)
	if err != nil {
		return fmt.Errorf("cannot insert document: %w", err)
	}

	_, err = tx.Exec(
		ctx,
		`INSERT INTO generated_documents (organization_id, tenant_id, third_parties_document_id, created_at, updated_at)
VALUES (@organization_id, @tenant_id, @third_parties_document_id, @created_at, @updated_at)
ON CONFLICT (organization_id) DO UPDATE SET third_parties_document_id = @third_parties_document_id, updated_at = @updated_at`,
		pgx.NamedArgs{
			"organization_id":           org.organizationID,
			"tenant_id":                 org.tenantID,
			"third_parties_document_id": documentID,
			"created_at":                now,
			"updated_at":                now,
		},
	)
	if err != nil {
		return fmt.Errorf("cannot link document: %w", err)
	}

	for major, snap := range snapshots {
		content, err := buildSnapshotContent(ctx, tx, snap.snapshotID, org.organizationName, snap.publishedAt)
		if err != nil {
			return fmt.Errorf("cannot build content for snapshot %s: %w", snap.snapshotID, err)
		}

		versionID := gid.New(org.tenantID, coredata.DocumentVersionEntityType)

		_, err = tx.Exec(
			ctx,
			`
INSERT INTO document_versions (
    id, tenant_id, organization_id, document_id,
    title, major, minor, classification, document_type,
    content, changelog, status, orientation,
    pdf_attempt_count,
    published_at, created_at, updated_at
) VALUES (
    @id, @tenant_id, @organization_id, @document_id,
    @title, @major, 0,
    'CONFIDENTIAL'::document_classification,
    'REGISTER'::document_type,
    @content, '',
    'PUBLISHED'::document_version_status,
    'PORTRAIT'::document_version_orientation,
    0,
    @published_at, @published_at, @published_at
)`,
			pgx.NamedArgs{
				"id":              versionID,
				"tenant_id":       org.tenantID,
				"organization_id": org.organizationID,
				"document_id":     documentID,
				"title":           "ThirdParties",
				"major":           major + 1,
				"content":         content,
				"published_at":    snap.publishedAt,
			},
		)
		if err != nil {
			return fmt.Errorf("cannot insert version for snapshot %s: %w", snap.snapshotID, err)
		}
	}

	fmt.Printf("OK   org %s (%s) — %d version(s)\n",
		org.organizationID, org.organizationName, len(snapshots))

	return nil
}

func loadOrgsWithThirdPartySnapshots(ctx context.Context, conn pg.Querier) ([]orgWithThirdPartySnapshots, error) {
	rows, err := conn.Query(
		ctx,
		`
SELECT DISTINCT
    o.id,
    o.tenant_id,
    o.name,
    o.created_at
FROM organizations o
WHERE NOT EXISTS (
        SELECT 1 FROM generated_documents gd
        WHERE gd.organization_id = o.id AND gd.third_parties_document_id IS NOT NULL
    )
    AND EXISTS (
        SELECT 1 FROM snapshots s
        WHERE s.organization_id = o.id AND s.type = 'VENDORS'
    )
ORDER BY o.created_at;
`,
	)
	if err != nil {
		return nil, fmt.Errorf("cannot query organizations with thirdParty snapshots: %w", err)
	}
	defer rows.Close()

	var result []orgWithThirdPartySnapshots
	for rows.Next() {
		var o orgWithThirdPartySnapshots
		var createdAt time.Time
		if err := rows.Scan(&o.organizationID, &o.tenantID, &o.organizationName, &createdAt); err != nil {
			return nil, fmt.Errorf("cannot scan organization: %w", err)
		}
		result = append(result, o)
	}

	return result, rows.Err()
}

func loadThirdPartySnapshots(ctx context.Context, conn pg.Querier, organizationID gid.GID) ([]thirdPartySnapshot, error) {
	rows, err := conn.Query(
		ctx,
		`
SELECT DISTINCT
    s.id,
    s.created_at
FROM snapshots s
WHERE s.organization_id = @organization_id
    AND s.type = 'VENDORS'
ORDER BY s.created_at ASC;
`,
		pgx.NamedArgs{"organization_id": organizationID},
	)
	if err != nil {
		return nil, fmt.Errorf("cannot query thirdParty snapshots for org %s: %w", organizationID, err)
	}
	defer rows.Close()

	var result []thirdPartySnapshot
	for rows.Next() {
		var s thirdPartySnapshot
		if err := rows.Scan(&s.snapshotID, &s.publishedAt); err != nil {
			return nil, fmt.Errorf("cannot scan snapshot: %w", err)
		}
		result = append(result, s)
	}

	return result, rows.Err()
}

type thirdPartyInfo struct {
	id       string
	name     string
	category string

	legalName                     *string
	description                   *string
	headquarterAddress            *string
	websiteURL                    *string
	privacyPolicyURL              *string
	serviceLevelAgreementURL      *string
	dataProcessingAgreementURL    *string
	businessAssociateAgreementURL *string
	subprocessorsListURL          *string
	statusPageURL                 *string
	termsOfServiceURL             *string
	securityPageURL               *string
	trustPageURL                  *string
	certifications                []string
	countries                     []string
	businessOwnerName             string
	securityOwnerName             string
}

func buildSnapshotContent(
	ctx context.Context,
	tx pg.Tx,
	snapshotID string,
	orgName string,
	publishedAt time.Time,
) (string, error) {
	thirdPartyRows, err := tx.Query(
		ctx,
		`
SELECT
    v.id,
    v.name,
    v.category,
    v.legal_name,
    v.description,
    v.headquarter_address,
    v.website_url,
    v.privacy_policy_url,
    v.service_level_agreement_url,
    v.data_processing_agreement_url,
    v.business_associate_agreement_url,
    v.subprocessors_list_url,
    v.status_page_url,
    v.terms_of_service_url,
    v.security_page_url,
    v.trust_page_url,
    v.certifications,
    v.countries,
    COALESCE(bo.full_name, 'Not assigned'),
    COALESCE(so.full_name, 'Not assigned')
FROM third_parties v
LEFT JOIN iam_membership_profiles bo ON bo.id = v.business_owner_profile_id
LEFT JOIN iam_membership_profiles so ON so.id = v.security_owner_profile_id
WHERE v.snapshot_id = @snapshot_id
ORDER BY v.name ASC;
`,
		pgx.NamedArgs{"snapshot_id": snapshotID},
	)
	if err != nil {
		return "", fmt.Errorf("cannot load snapshot thirdParties: %w", err)
	}
	defer thirdPartyRows.Close()

	var thirdParties []thirdPartyInfo
	for thirdPartyRows.Next() {
		var v thirdPartyInfo
		if err := thirdPartyRows.Scan(
			&v.id, &v.name, &v.category,
			&v.legalName, &v.description, &v.headquarterAddress,
			&v.websiteURL, &v.privacyPolicyURL, &v.serviceLevelAgreementURL,
			&v.dataProcessingAgreementURL, &v.businessAssociateAgreementURL,
			&v.subprocessorsListURL, &v.statusPageURL, &v.termsOfServiceURL,
			&v.securityPageURL, &v.trustPageURL,
			&v.certifications, &v.countries,
			&v.businessOwnerName, &v.securityOwnerName,
		); err != nil {
			return "", fmt.Errorf("cannot scan thirdParty: %w", err)
		}
		thirdParties = append(thirdParties, v)
	}
	if err := thirdPartyRows.Err(); err != nil {
		return "", err
	}

	thirdPartyIDs := make([]string, len(thirdParties))
	for i, v := range thirdParties {
		thirdPartyIDs[i] = v.id
	}

	servicesByThirdParty, err := loadSnapshotServices(ctx, tx, snapshotID, thirdPartyIDs)
	if err != nil {
		return "", err
	}

	contactsByThirdParty, err := loadSnapshotContacts(ctx, tx, snapshotID, thirdPartyIDs)
	if err != nil {
		return "", err
	}

	assessmentsByThirdParty, err := loadSnapshotRiskAssessments(ctx, tx, snapshotID, thirdPartyIDs)
	if err != nil {
		return "", err
	}

	reportsByThirdParty, err := loadSnapshotComplianceReports(ctx, tx, snapshotID, thirdPartyIDs)
	if err != nil {
		return "", err
	}

	baaByThirdParty, err := loadSnapshotBAAs(ctx, tx, snapshotID, thirdPartyIDs)
	if err != nil {
		return "", err
	}

	dpaByThirdParty, err := loadSnapshotDPAs(ctx, tx, snapshotID, thirdPartyIDs)
	if err != nil {
		return "", err
	}

	rows := make([]docgen.ThirdPartyListRow, 0, len(thirdParties))
	for _, v := range thirdParties {
		row := docgen.ThirdPartyListRow{
			Name:                          v.name,
			LegalName:                     deref(v.legalName),
			Description:                   deref(v.description),
			Category:                      formatCategory(v.category),
			HeadquarterAddress:            deref(v.headquarterAddress),
			WebsiteURL:                    deref(v.websiteURL),
			PrivacyPolicyURL:              deref(v.privacyPolicyURL),
			ServiceLevelAgreementURL:      deref(v.serviceLevelAgreementURL),
			DataProcessingAgreementURL:    deref(v.dataProcessingAgreementURL),
			BusinessAssociateAgreementURL: deref(v.businessAssociateAgreementURL),
			SubprocessorsListURL:          deref(v.subprocessorsListURL),
			StatusPageURL:                 deref(v.statusPageURL),
			TermsOfServiceURL:             deref(v.termsOfServiceURL),
			SecurityPageURL:               deref(v.securityPageURL),
			TrustPageURL:                  deref(v.trustPageURL),
			Certifications:                joinOrDefault(v.certifications),
			Countries:                     joinOrDefault(v.countries),
			BusinessOwner:                 v.businessOwnerName,
			SecurityOwner:                 v.securityOwnerName,
			Services:                      servicesByThirdParty[v.id],
			Contacts:                      contactsByThirdParty[v.id],
			RiskAssessments:               assessmentsByThirdParty[v.id],
			ComplianceReports:             reportsByThirdParty[v.id],
			BusinessAssociateAgreement:    baaByThirdParty[v.id],
			DataPrivacyAgreement:          dpaByThirdParty[v.id],
		}
		rows = append(rows, row)
	}

	docData := docgen.ThirdPartyListData{
		Title:             "ThirdParties",
		OrganizationName:  orgName,
		CreatedAt:         publishedAt,
		TotalThirdParties: len(rows),
		Rows:              rows,
	}

	return probo.BuildThirdPartyListDocument(docData)
}

func loadSnapshotServices(ctx context.Context, tx pg.Tx, snapshotID string, thirdPartyIDs []string) (map[string][]docgen.ThirdPartyListService, error) {
	rows, err := tx.Query(ctx,
		`SELECT vs.third_party_id, vs.name, COALESCE(vs.description, 'Not specified')
		FROM third_party_services vs
		WHERE vs.snapshot_id = @snapshot_id AND vs.third_party_id = ANY(@third_party_ids)
		ORDER BY vs.third_party_id, vs.name ASC`,
		pgx.NamedArgs{"snapshot_id": snapshotID, "third_party_ids": thirdPartyIDs})
	if err != nil {
		return nil, fmt.Errorf("cannot load snapshot services: %w", err)
	}
	defer rows.Close()

	result := make(map[string][]docgen.ThirdPartyListService)
	for rows.Next() {
		var thirdPartyID, name, desc string
		if err := rows.Scan(&thirdPartyID, &name, &desc); err != nil {
			return nil, fmt.Errorf("cannot scan service: %w", err)
		}
		result[thirdPartyID] = append(result[thirdPartyID], docgen.ThirdPartyListService{Name: name, Description: desc})
	}
	return result, rows.Err()
}

func loadSnapshotContacts(ctx context.Context, tx pg.Tx, snapshotID string, thirdPartyIDs []string) (map[string][]docgen.ThirdPartyListContact, error) {
	rows, err := tx.Query(ctx,
		`SELECT vc.third_party_id,
		    COALESCE(vc.full_name, 'Not specified'),
		    COALESCE(vc.email, 'Not specified'),
		    COALESCE(vc.phone, 'Not specified'),
		    COALESCE(vc.role, 'Not specified')
		FROM third_party_contacts vc
		WHERE vc.snapshot_id = @snapshot_id AND vc.third_party_id = ANY(@third_party_ids)
		ORDER BY vc.third_party_id, vc.full_name ASC`,
		pgx.NamedArgs{"snapshot_id": snapshotID, "third_party_ids": thirdPartyIDs})
	if err != nil {
		return nil, fmt.Errorf("cannot load snapshot contacts: %w", err)
	}
	defer rows.Close()

	result := make(map[string][]docgen.ThirdPartyListContact)
	for rows.Next() {
		var thirdPartyID, name, email, phone, role string
		if err := rows.Scan(&thirdPartyID, &name, &email, &phone, &role); err != nil {
			return nil, fmt.Errorf("cannot scan contact: %w", err)
		}
		result[thirdPartyID] = append(result[thirdPartyID], docgen.ThirdPartyListContact{
			FullName: name, Email: email, Phone: phone, Role: role,
		})
	}
	return result, rows.Err()
}

func loadSnapshotRiskAssessments(ctx context.Context, tx pg.Tx, snapshotID string, thirdPartyIDs []string) (map[string][]docgen.ThirdPartyListRiskAssessment, error) {
	rows, err := tx.Query(ctx,
		`SELECT vra.third_party_id, vra.created_at, vra.expires_at, vra.data_sensitivity, vra.business_impact, COALESCE(vra.notes, 'Not specified')
		FROM third_party_risk_assessments vra
		WHERE vra.snapshot_id = @snapshot_id AND vra.third_party_id = ANY(@third_party_ids)
		ORDER BY vra.third_party_id, vra.created_at DESC`,
		pgx.NamedArgs{"snapshot_id": snapshotID, "third_party_ids": thirdPartyIDs})
	if err != nil {
		return nil, fmt.Errorf("cannot load snapshot risk assessments: %w", err)
	}
	defer rows.Close()

	result := make(map[string][]docgen.ThirdPartyListRiskAssessment)
	for rows.Next() {
		var thirdPartyID, sensitivity, impact, notes string
		var assessedAt, expiresAt time.Time
		if err := rows.Scan(&thirdPartyID, &assessedAt, &expiresAt, &sensitivity, &impact, &notes); err != nil {
			return nil, fmt.Errorf("cannot scan risk assessment: %w", err)
		}
		result[thirdPartyID] = append(result[thirdPartyID], docgen.ThirdPartyListRiskAssessment{
			AssessedAt:      assessedAt.Format("2006-01-02"),
			ExpiresAt:       expiresAt.Format("2006-01-02"),
			DataSensitivity: sensitivity,
			BusinessImpact:  impact,
			Notes:           notes,
		})
	}
	return result, rows.Err()
}

func loadSnapshotComplianceReports(ctx context.Context, tx pg.Tx, snapshotID string, thirdPartyIDs []string) (map[string][]docgen.ThirdPartyListComplianceReport, error) {
	rows, err := tx.Query(ctx,
		`SELECT vcr.third_party_id, vcr.report_name, vcr.report_date, vcr.valid_until
		FROM third_party_compliance_reports vcr
		WHERE vcr.snapshot_id = @snapshot_id AND vcr.third_party_id = ANY(@third_party_ids)
		ORDER BY vcr.third_party_id, vcr.report_date DESC`,
		pgx.NamedArgs{"snapshot_id": snapshotID, "third_party_ids": thirdPartyIDs})
	if err != nil {
		return nil, fmt.Errorf("cannot load snapshot compliance reports: %w", err)
	}
	defer rows.Close()

	result := make(map[string][]docgen.ThirdPartyListComplianceReport)
	for rows.Next() {
		var thirdPartyID, name string
		var reportDate time.Time
		var validUntil *time.Time
		if err := rows.Scan(&thirdPartyID, &name, &reportDate, &validUntil); err != nil {
			return nil, fmt.Errorf("cannot scan compliance report: %w", err)
		}
		vu := "Not specified"
		if validUntil != nil {
			vu = validUntil.Format("2006-01-02")
		}
		result[thirdPartyID] = append(result[thirdPartyID], docgen.ThirdPartyListComplianceReport{
			ReportName: name, ReportDate: reportDate.Format("2006-01-02"), ValidUntil: vu,
		})
	}
	return result, rows.Err()
}

func loadSnapshotBAAs(ctx context.Context, tx pg.Tx, snapshotID string, thirdPartyIDs []string) (map[string]*docgen.ThirdPartyListAgreement, error) {
	rows, err := tx.Query(ctx,
		`SELECT vbaa.third_party_id, vbaa.valid_from, vbaa.valid_until
		FROM third_party_business_associate_agreements vbaa
		WHERE vbaa.snapshot_id = @snapshot_id AND vbaa.third_party_id = ANY(@third_party_ids)`,
		pgx.NamedArgs{"snapshot_id": snapshotID, "third_party_ids": thirdPartyIDs})
	if err != nil {
		return nil, fmt.Errorf("cannot load snapshot BAAs: %w", err)
	}
	defer rows.Close()

	result := make(map[string]*docgen.ThirdPartyListAgreement)
	for rows.Next() {
		var thirdPartyID string
		var validFrom, validUntil *time.Time
		if err := rows.Scan(&thirdPartyID, &validFrom, &validUntil); err != nil {
			return nil, fmt.Errorf("cannot scan BAA: %w", err)
		}
		result[thirdPartyID] = &docgen.ThirdPartyListAgreement{
			ValidFrom: fmtTime(validFrom), ValidUntil: fmtTime(validUntil),
		}
	}
	return result, rows.Err()
}

func loadSnapshotDPAs(ctx context.Context, tx pg.Tx, snapshotID string, thirdPartyIDs []string) (map[string]*docgen.ThirdPartyListAgreement, error) {
	rows, err := tx.Query(ctx,
		`SELECT vdpa.third_party_id, vdpa.valid_from, vdpa.valid_until
		FROM third_party_data_privacy_agreements vdpa
		WHERE vdpa.snapshot_id = @snapshot_id AND vdpa.third_party_id = ANY(@third_party_ids)`,
		pgx.NamedArgs{"snapshot_id": snapshotID, "third_party_ids": thirdPartyIDs})
	if err != nil {
		return nil, fmt.Errorf("cannot load snapshot DPAs: %w", err)
	}
	defer rows.Close()

	result := make(map[string]*docgen.ThirdPartyListAgreement)
	for rows.Next() {
		var thirdPartyID string
		var validFrom, validUntil *time.Time
		if err := rows.Scan(&thirdPartyID, &validFrom, &validUntil); err != nil {
			return nil, fmt.Errorf("cannot scan DPA: %w", err)
		}
		result[thirdPartyID] = &docgen.ThirdPartyListAgreement{
			ValidFrom: fmtTime(validFrom), ValidUntil: fmtTime(validUntil),
		}
	}
	return result, rows.Err()
}

func deref(s *string) string {
	if s == nil || *s == "" {
		return "Not specified"
	}
	return *s
}

func joinOrDefault(items []string) string {
	if len(items) == 0 {
		return "Not specified"
	}
	return strings.Join(items, ", ")
}

func fmtTime(t *time.Time) string {
	if t == nil {
		return "Not specified"
	}
	return t.Format("2006-01-02")
}

func formatCategory(c string) string {
	switch c {
	case "ANALYTICS":
		return "Analytics"
	case "CLOUD_MONITORING":
		return "Cloud Monitoring"
	case "CLOUD_PROVIDER":
		return "Cloud Provider"
	case "COLLABORATION":
		return "Collaboration"
	case "CUSTOMER_SUPPORT":
		return "Customer Support"
	case "DATA_STORAGE_AND_PROCESSING":
		return "Data Storage and Processing"
	case "DOCUMENT_MANAGEMENT":
		return "Document Management"
	case "EMPLOYEE_MANAGEMENT":
		return "Employee Management"
	case "ENGINEERING":
		return "Engineering"
	case "FINANCE":
		return "Finance"
	case "IDENTITY_PROVIDER":
		return "Identity Provider"
	case "IT":
		return "IT"
	case "MARKETING":
		return "Marketing"
	case "OFFICE_OPERATIONS":
		return "Office Operations"
	case "OTHER":
		return "Other"
	case "PASSWORD_MANAGEMENT":
		return "Password Management"
	case "PRODUCT_AND_DESIGN":
		return "Product and Design"
	case "PROFESSIONAL_SERVICES":
		return "Professional Services"
	case "RECRUITING":
		return "Recruiting"
	case "SALES":
		return "Sales"
	case "SECURITY":
		return "Security"
	case "VERSION_CONTROL":
		return "Version Control"
	default:
		return c
	}
}

func newPgClientFromDSN(dsn string) (*pg.Client, error) {
	u, err := url.Parse(dsn)
	if err != nil {
		return nil, fmt.Errorf("cannot parse DSN")
	}

	var opts []pg.Option

	if u.Host != "" {
		host := u.Host
		if u.Port() == "" {
			host = net.JoinHostPort(u.Hostname(), "5432")
		}
		opts = append(opts, pg.WithAddr(host))
	}

	if u.User != nil {
		opts = append(opts, pg.WithUser(u.User.Username()))
		if password, ok := u.User.Password(); ok {
			opts = append(opts, pg.WithPassword(password))
		}
	}

	if len(u.Path) > 1 {
		opts = append(opts, pg.WithDatabase(u.Path[1:]))
	}

	return pg.NewClient(opts...)
}