feat: Add support for Functions and Stored Procedures (Fixes #2)

Author: Meru143

internal/db/queries.go

@@ -388,3 +388,31 @@ func getMaterializedViews(ctx context.Context, db *DB, schemaName string) ([]typ
 	}
 	return matViews, rows.Err()
 }
+
+func getFunctions(ctx context.Context, db *DB, schemaName string) ([]types.Function, error) {
+	// Query to extract functions, handling overloading correctly
+	query := `
+		SELECT 
+			p.proname as name,
+			pg_get_function_identity_arguments(p.oid) as arguments,
+			pg_get_functiondef(p.oid) as definition
+		FROM pg_proc p
+		JOIN pg_namespace n ON p.pronamespace = n.oid
+		WHERE n.nspname = $1
+	`
+	rows, err := db.Query(ctx, query, schemaName)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var functions []types.Function
+	for rows.Next() {
+		var f types.Function
+		if err := rows.Scan(&f.Name, &f.Arguments, &f.Definition); err != nil {
+			return nil, err
+		}
+		functions = append(functions, f)
+	}
+	return functions, rows.Err()
+}

internal/db/schema.go

@@ -34,6 +34,7 @@ func (i *Introspector) Introspect(ctx context.Context) (*types.Schema, error) {
 		Types:             make([]types.Type, 0),
 		Views:             make([]types.View, 0),
 		MaterializedViews: make([]types.MaterializedView, 0),
+		Functions:         make([]types.Function, 0),
 	}
 
 	// Get tables
@@ -127,9 +128,19 @@ func (i *Introspector) Introspect(ctx context.Context) (*types.Schema, error) {
 	}
 	schema.MaterializedViews = matViews
 
+	// Get functions
 	if i.verbose {
-		log.Printf("Introspection complete: %d tables, %d sequences, %d types, %d views, %d materialized views",
-			len(schema.Tables), len(schema.Sequences), len(schema.Types), len(schema.Views), len(schema.MaterializedViews))
+		log.Println("Introspecting functions...")
+	}
+	functions, err := getFunctions(ctx, i.db, i.schemaName)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get functions: %w", err)
+	}
+	schema.Functions = functions
+
+	if i.verbose {
+		log.Printf("Introspection complete: %d tables, %d sequences, %d types, %d views, %d materialized views, %d functions",
+			len(schema.Tables), len(schema.Sequences), len(schema.Types), len(schema.Views), len(schema.MaterializedViews), len(schema.Functions))
 	}
 
 	return schema, nil

internal/diff/engine.go

@@ -1,6 +1,7 @@
 package diff
 
 import (
+	"fmt"
 	"strconv"
 	"strings"
 
@@ -40,6 +41,9 @@ func (e *DiffEngine) Compare() types.DiffList {
 	// Materialized View comparison
 	differences = append(differences, e.compareMaterializedViews()...)
 
+	// Function comparison
+	differences = append(differences, e.compareFunctions()...)
+
 	return differences
 }
 
@@ -725,3 +729,55 @@ func (e *DiffEngine) compareMaterializedViews() types.DiffList {
 
 	return differences
 }
+
+func (e *DiffEngine) compareFunctions() types.DiffList {
+	var differences types.DiffList
+
+	sourceFuncs := make(map[string]*types.Function)
+	targetFuncs := make(map[string]*types.Function)
+
+	for i := range e.source.Functions {
+		f := &e.source.Functions[i]
+		key := fmt.Sprintf("%s(%s)", f.Name, f.Arguments)
+		sourceFuncs[key] = f
+	}
+	for i := range e.target.Functions {
+		f := &e.target.Functions[i]
+		key := fmt.Sprintf("%s(%s)", f.Name, f.Arguments)
+		targetFuncs[key] = f
+	}
+
+	// New and altered functions
+	for key, sourceFunc := range sourceFuncs {
+		if targetFunc, exists := targetFuncs[key]; !exists {
+			differences = append(differences, types.Diff{
+				Type:     types.DiffAdd,
+				Object:   types.ObjectFunction,
+				Name:     key,                   // key includes func_name(args)
+				NewValue: sourceFunc.Definition, // complete CREATE OR REPLACE snippet
+			})
+		} else if sourceFunc.Definition != targetFunc.Definition {
+			// Definition changed
+			differences = append(differences, types.Diff{
+				Type:     types.DiffAlter,
+				Object:   types.ObjectFunction,
+				Name:     key,
+				OldValue: targetFunc.Definition,
+				NewValue: sourceFunc.Definition,
+			})
+		}
+	}
+
+	// Dropped functions
+	for key := range targetFuncs {
+		if _, exists := sourceFuncs[key]; !exists {
+			differences = append(differences, types.Diff{
+				Type:   types.DiffDrop,
+				Object: types.ObjectFunction,
+				Name:   key,
+			})
+		}
+	}
+
+	return differences
+}

internal/diff/generator.go

@@ -473,6 +473,25 @@ func (g *SQLGenerator) generateMaterializedViewDiff(diff *types.Diff) string {
 	}
 }
 
+func (g *SQLGenerator) generateFunctionDiff(diff *types.Diff) string {
+	schema := schemaPrefix(g.schemaName)
+	definition := strings.TrimRight(diff.NewValue, " \t\r\n;")
+	// Notice: diff.Name contains func_name(args) to handle overloads specifically for drop
+	switch diff.Type {
+	case types.DiffAdd:
+		return fmt.Sprintf("-- Create function: %s\n%s;",
+			diff.Name, definition)
+	case types.DiffAlter:
+		return fmt.Sprintf("-- Alter function: %s\n%s;",
+			diff.Name, definition)
+	case types.DiffDrop:
+		return fmt.Sprintf("-- Drop function: %s\nDROP FUNCTION IF EXISTS %s%s CASCADE;",
+			diff.Name, schema, diff.Name)
+	default:
+		return ""
+	}
+}
+
 // GenerateSQL is a legacy function
 func GenerateSQL(diffs types.DiffList, schemaName string) string {
 	gen := NewSQLGenerator(diffs, schemaName)

pkg/types/types.go

@@ -10,6 +10,7 @@ type Schema struct {
 	Types             []Type             `json:"types,omitempty"`
 	Views             []View             `json:"views,omitempty"`
 	MaterializedViews []MaterializedView `json:"materialized_views,omitempty"`
+	Functions         []Function         `json:"functions,omitempty"`
 }
 
 type Table struct {
@@ -81,6 +82,13 @@ type MaterializedView struct {
 	Definition string `json:"definition"`
 }
 
+// Function represents a PostgreSQL stored procedure or function
+type Function struct {
+	Name       string `json:"name"`
+	Arguments  string `json:"arguments"`
+	Definition string `json:"definition"`
+}
+
 type DiffType string
 
 const (
@@ -102,6 +110,7 @@ const (
 	ObjectType             DiffObject = "TYPE"
 	ObjectView             DiffObject = "VIEW"
 	ObjectMaterializedView DiffObject = "MATERIALIZED_VIEW"
+	ObjectFunction         DiffObject = "FUNCTION"
 )
 
 type Diff struct {