Files
relspecgo/pkg/readers/pgsql/reader.go
T
warkanum ee5c009234 feat(main): add silent mode to suppress output messages
* implement --silent flag to control output verbosity
* update error handling to restore stderr after silent mode
* enhance progress reporting in PostgreSQL reader
2026-09-09 21:30:15 +02:00

415 lines
13 KiB
Go

package pgsql
import (
"context"
"fmt"
"strings"
"github.com/jackc/pgx/v5"
"git.warky.dev/wdevs/relspecgo/pkg/models"
"git.warky.dev/wdevs/relspecgo/pkg/pgsql"
"git.warky.dev/wdevs/relspecgo/pkg/readers"
)
// Reader implements the readers.Reader interface for PostgreSQL databases
type Reader struct {
options *readers.ReaderOptions
conn *pgx.Conn
ctx context.Context
}
// NewReader creates a new PostgreSQL reader
func NewReader(options *readers.ReaderOptions) *Reader {
return &Reader{
options: options,
ctx: context.Background(),
}
}
// ReadDatabase reads the entire database schema from PostgreSQL
func (r *Reader) ReadDatabase() (*models.Database, error) {
// Validate connection string
if r.options.ConnectionString == "" {
return nil, fmt.Errorf("connection string is required")
}
// Connect to the database. This can take noticeable time across a slow network,
// so report it before the driver starts the connection attempt.
r.progress("Connecting to PostgreSQL...")
if err := r.connect(); err != nil {
return nil, fmt.Errorf("failed to connect: %w", err)
}
defer r.close()
r.progress("Connected. Reading database metadata...")
// Get database name from connection
var dbName string
err := r.conn.QueryRow(r.ctx, "SELECT current_database()").Scan(&dbName)
if err != nil {
return nil, fmt.Errorf("failed to get database name: %w", err)
}
// Initialize database model
db := models.InitDatabase(dbName)
db.DatabaseType = models.PostgresqlDatabaseType
db.SourceFormat = "pgsql"
// Get PostgreSQL version
var version string
err = r.conn.QueryRow(r.ctx, "SELECT version()").Scan(&version)
if err == nil {
db.DatabaseVersion = version
}
// Query all schemas
r.progress("Discovering schemas...")
schemas, err := r.querySchemas()
if err != nil {
return nil, fmt.Errorf("failed to query schemas: %w", err)
}
// Process each schema
for schemaIndex, schema := range schemas {
r.progress(fmt.Sprintf("Reading schema %q (%d/%d): tables...", schema.Name, schemaIndex+1, len(schemas)))
// Query tables for this schema
tables, err := r.queryTables(schema.Name)
if err != nil {
return nil, fmt.Errorf("failed to query tables for schema %s: %w", schema.Name, err)
}
schema.Tables = tables
r.progress(fmt.Sprintf("Reading schema %q: found %d table(s).", schema.Name, len(tables)))
r.progress(fmt.Sprintf("Reading schema %q: views...", schema.Name))
// Query views for this schema
views, err := r.queryViews(schema.Name)
if err != nil {
return nil, fmt.Errorf("failed to query views for schema %s: %w", schema.Name, err)
}
schema.Views = views
r.progress(fmt.Sprintf("Reading schema %q: found %d view(s).", schema.Name, len(views)))
r.progress(fmt.Sprintf("Reading schema %q: sequences...", schema.Name))
// Query sequences for this schema
sequences, err := r.querySequences(schema.Name)
if err != nil {
return nil, fmt.Errorf("failed to query sequences for schema %s: %w", schema.Name, err)
}
schema.Sequences = sequences
r.progress(fmt.Sprintf("Reading schema %q: found %d sequence(s).", schema.Name, len(sequences)))
r.progress(fmt.Sprintf("Reading schema %q: extensions...", schema.Name))
// Query extensions installed into this schema
extensions, err := r.queryExtensions(schema.Name)
if err != nil {
return nil, fmt.Errorf("failed to query extensions for schema %s: %w", schema.Name, err)
}
if len(extensions) > 0 {
if schema.Metadata == nil {
schema.Metadata = make(map[string]any)
}
schema.Metadata["extensions"] = extensions
}
r.progress(fmt.Sprintf("Reading schema %q: found %d extension(s).", schema.Name, len(extensions)))
r.progress(fmt.Sprintf("Reading schema %q: columns...", schema.Name))
// Query columns for tables and views
columnsMap, err := r.queryColumns(schema.Name)
if err != nil {
return nil, fmt.Errorf("failed to query columns for schema %s: %w", schema.Name, err)
}
r.progress(fmt.Sprintf("Reading schema %q: found %d column(s).", schema.Name, countColumns(columnsMap)))
// Populate table columns
for _, table := range schema.Tables {
tableKey := schema.Name + "." + table.Name
if cols, exists := columnsMap[tableKey]; exists {
table.Columns = cols
}
}
// Populate view columns
for _, view := range schema.Views {
viewKey := schema.Name + "." + view.Name
if cols, exists := columnsMap[viewKey]; exists {
view.Columns = cols
}
}
r.progress(fmt.Sprintf("Reading schema %q: primary keys...", schema.Name))
// Query primary keys
primaryKeys, err := r.queryPrimaryKeys(schema.Name)
if err != nil {
return nil, fmt.Errorf("failed to query primary keys for schema %s: %w", schema.Name, err)
}
r.progress(fmt.Sprintf("Reading schema %q: found %d primary key(s).", schema.Name, len(primaryKeys)))
// Apply primary keys to tables
for _, table := range schema.Tables {
tableKey := schema.Name + "." + table.Name
if pk, exists := primaryKeys[tableKey]; exists {
table.Constraints[pk.Name] = pk
// Mark columns as primary key and not null
for _, colName := range pk.Columns {
if col, colExists := table.Columns[colName]; colExists {
col.IsPrimaryKey = true
col.NotNull = true
}
}
}
}
r.progress(fmt.Sprintf("Reading schema %q: foreign keys...", schema.Name))
// Query foreign keys
foreignKeys, err := r.queryForeignKeys(schema.Name)
if err != nil {
return nil, fmt.Errorf("failed to query foreign keys for schema %s: %w", schema.Name, err)
}
r.progress(fmt.Sprintf("Reading schema %q: found %d foreign key(s).", schema.Name, countConstraints(foreignKeys)))
// Apply foreign keys to tables
for _, table := range schema.Tables {
tableKey := schema.Name + "." + table.Name
if fks, exists := foreignKeys[tableKey]; exists {
for _, fk := range fks {
table.Constraints[fk.Name] = fk
// Derive relationship from foreign key
r.deriveRelationship(table, fk)
}
}
}
r.progress(fmt.Sprintf("Reading schema %q: unique constraints...", schema.Name))
// Query unique constraints
uniqueConstraints, err := r.queryUniqueConstraints(schema.Name)
if err != nil {
return nil, fmt.Errorf("failed to query unique constraints for schema %s: %w", schema.Name, err)
}
r.progress(fmt.Sprintf("Reading schema %q: found %d unique constraint(s).", schema.Name, countConstraints(uniqueConstraints)))
// Apply unique constraints to tables
for _, table := range schema.Tables {
tableKey := schema.Name + "." + table.Name
if ucs, exists := uniqueConstraints[tableKey]; exists {
for _, uc := range ucs {
table.Constraints[uc.Name] = uc
}
}
}
r.progress(fmt.Sprintf("Reading schema %q: check constraints...", schema.Name))
// Query check constraints
checkConstraints, err := r.queryCheckConstraints(schema.Name)
if err != nil {
return nil, fmt.Errorf("failed to query check constraints for schema %s: %w", schema.Name, err)
}
r.progress(fmt.Sprintf("Reading schema %q: found %d check constraint(s).", schema.Name, countConstraints(checkConstraints)))
// Apply check constraints to tables
for _, table := range schema.Tables {
tableKey := schema.Name + "." + table.Name
if ccs, exists := checkConstraints[tableKey]; exists {
for _, cc := range ccs {
table.Constraints[cc.Name] = cc
}
}
}
r.progress(fmt.Sprintf("Reading schema %q: indexes...", schema.Name))
// Query indexes
indexes, err := r.queryIndexes(schema.Name)
if err != nil {
return nil, fmt.Errorf("failed to query indexes for schema %s: %w", schema.Name, err)
}
r.progress(fmt.Sprintf("Reading schema %q: found %d index(es).", schema.Name, countIndexes(indexes)))
// Apply indexes to tables
for _, table := range schema.Tables {
tableKey := schema.Name + "." + table.Name
if idxs, exists := indexes[tableKey]; exists {
for _, idx := range idxs {
table.Indexes[idx.Name] = idx
}
}
}
// Set RefDatabase for schema
schema.RefDatabase = db
// Set RefSchema for tables and views
for _, table := range schema.Tables {
table.RefSchema = schema
}
for _, view := range schema.Views {
view.RefSchema = schema
}
for _, seq := range schema.Sequences {
seq.RefSchema = schema
}
// Add schema to database
db.Schemas = append(db.Schemas, schema)
}
r.progress("PostgreSQL schema read complete.")
return db, nil
}
func (r *Reader) progress(message string) {
if r.options.Progress != nil {
r.options.Progress(message)
}
}
func countColumns(columns map[string]map[string]*models.Column) int {
total := 0
for _, tableColumns := range columns {
total += len(tableColumns)
}
return total
}
func countConstraints(constraints map[string][]*models.Constraint) int {
total := 0
for _, tableConstraints := range constraints {
total += len(tableConstraints)
}
return total
}
func countIndexes(indexes map[string][]*models.Index) int {
total := 0
for _, tableIndexes := range indexes {
total += len(tableIndexes)
}
return total
}
// ReadSchema reads a single schema (returns the first schema from the database)
func (r *Reader) ReadSchema() (*models.Schema, error) {
db, err := r.ReadDatabase()
if err != nil {
return nil, err
}
if len(db.Schemas) == 0 {
return nil, fmt.Errorf("no schemas found in database")
}
return db.Schemas[0], nil
}
// ReadTable reads a single table (returns the first table from the first schema)
func (r *Reader) ReadTable() (*models.Table, error) {
schema, err := r.ReadSchema()
if err != nil {
return nil, err
}
if len(schema.Tables) == 0 {
return nil, fmt.Errorf("no tables found in schema")
}
return schema.Tables[0], nil
}
// connect establishes a connection to the PostgreSQL database
func (r *Reader) connect() error {
conn, err := pgsql.Connect(r.ctx, r.options.ConnectionString, "reader-pgsql")
if err != nil {
return err
}
r.conn = conn
return nil
}
// close closes the database connection
func (r *Reader) close() {
if r.conn != nil {
r.conn.Close(r.ctx)
}
}
// mapDataType maps a PostgreSQL data type to its canonical RelSpec name.
// For known built-in types, dimensions are stripped from the type string (they are
// stored separately in column.Length/Precision/Scale). For custom types (e.g.
// vector(1536), postgis geometries), the full formatted type is preserved.
func (r *Reader) mapDataType(pgType, udtName, formattedType string, hasNextval bool) string {
normalizedPGType := strings.ToLower(strings.TrimSpace(pgType))
// Detect serial types from nextval defaults before anything else.
if hasNextval {
switch normalizedPGType {
case "integer", "int", "int4":
return "serial"
case "bigint", "int8":
return "bigserial"
case "smallint", "int2":
return "smallserial"
}
}
// Use the database-formatted type when available. For known built-in types, strip
// embedded dimensions (they are stored in column.Length/Precision/Scale separately).
// For unknown/custom types, keep the full formatted string (e.g. vector(1536)).
if strings.TrimSpace(formattedType) != "" {
lower := strings.ToLower(strings.TrimSpace(formattedType))
isArray := strings.HasSuffix(lower, "[]")
base := strings.TrimSuffix(lower, "[]")
if idx := strings.Index(base, "("); idx >= 0 {
base = strings.TrimSpace(base[:idx])
}
canonical := pgsql.NormalizePGType(base)
if pgsql.IsKnownPGBaseType(canonical) {
if isArray {
return canonical + "[]"
}
return canonical
}
return formattedType
}
// information_schema reports arrays generically as "ARRAY" with udt_name like "_text".
// Only reached when the catalog-formatted type is unavailable, which is the one case
// where the element modifier (e.g. geometry(Point,4326)[]) cannot be recovered.
if strings.EqualFold(pgType, "ARRAY") && strings.HasPrefix(udtName, "_") && len(udtName) > 1 {
return udtName[1:] + "[]"
}
// Fall back to normalizing the information_schema type name directly.
canonical := pgsql.NormalizePGType(normalizedPGType)
if pgsql.IsKnownPGBaseType(canonical) {
return canonical
}
// Return UDT name for custom types.
if udtName != "" {
if strings.HasPrefix(udtName, "_") && len(udtName) > 1 {
return udtName[1:] + "[]"
}
return udtName
}
return pgType
}
// deriveRelationship creates a relationship from a foreign key constraint
func (r *Reader) deriveRelationship(table *models.Table, fk *models.Constraint) {
relationshipName := fmt.Sprintf("%s_to_%s", table.Name, fk.ReferencedTable)
relationship := models.InitRelationship(relationshipName, models.OneToMany)
relationship.FromTable = table.Name
relationship.FromSchema = table.Schema
relationship.FromColumns = append([]string(nil), fk.Columns...)
relationship.ToTable = fk.ReferencedTable
relationship.ToSchema = fk.ReferencedSchema
relationship.ToColumns = append([]string(nil), fk.ReferencedColumns...)
relationship.ForeignKey = fk.Name
// Store constraint actions in properties
if fk.OnDelete != "" {
relationship.Properties["on_delete"] = fk.OnDelete
}
if fk.OnUpdate != "" {
relationship.Properties["on_update"] = fk.OnUpdate
}
table.Relationships[relationshipName] = relationship
}