Skip to content

Commit d977fa6

Browse files
linhdmnlinhdmn
committed
feat: enhance postgres connection
1 parent 86a230a commit d977fa6

9 files changed

+590
-104
lines changed

‎config.json

+39-3
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,11 @@
77
"port": 3306,
88
"name": "db1",
99
"user": "user1",
10-
"password": "password1"
10+
"password": "password1",
11+
"max_open_conns": 20,
12+
"max_idle_conns": 5,
13+
"conn_max_lifetime_seconds": 300,
14+
"conn_max_idle_time_seconds": 60
1115
},
1216
{
1317
"id": "mysql2",
@@ -16,7 +20,11 @@
1620
"port": 3306,
1721
"name": "db2",
1822
"user": "user2",
19-
"password": "password2"
23+
"password": "password2",
24+
"max_open_conns": 20,
25+
"max_idle_conns": 5,
26+
"conn_max_lifetime_seconds": 300,
27+
"conn_max_idle_time_seconds": 60
2028
},
2129
{
2230
"id": "postgres1",
@@ -25,7 +33,35 @@
2533
"port": 5432,
2634
"name": "db1",
2735
"user": "user1",
28-
"password": "password1"
36+
"password": "password1",
37+
"ssl_mode": "disable",
38+
"application_name": "db-mcp-server",
39+
"connect_timeout": 10,
40+
"max_open_conns": 20,
41+
"max_idle_conns": 5,
42+
"conn_max_lifetime_seconds": 300,
43+
"conn_max_idle_time_seconds": 60
44+
},
45+
{
46+
"id": "postgres17",
47+
"type": "postgres",
48+
"host": "postgres17",
49+
"port": 5432,
50+
"name": "db2",
51+
"user": "user2",
52+
"password": "password2",
53+
"ssl_mode": "prefer",
54+
"application_name": "db-mcp-server",
55+
"connect_timeout": 15,
56+
"target_session_attrs": "any",
57+
"options": {
58+
"application_name": "db-mcp-server",
59+
"client_encoding": "UTF8"
60+
},
61+
"max_open_conns": 25,
62+
"max_idle_conns": 5,
63+
"conn_max_lifetime_seconds": 300,
64+
"conn_max_idle_time_seconds": 60
2965
}
3066
]
3167
}

‎docker-compose.test.yml

+26-2
Original file line numberDiff line numberDiff line change
@@ -40,19 +40,43 @@ services:
4040
postgres1:
4141
image: postgres:15
4242
container_name: mcp_postgres1
43+
environment:
44+
POSTGRES_USER: user1
45+
POSTGRES_PASSWORD: password1
46+
POSTGRES_DB: db1
47+
ports:
48+
- "15432:5432"
49+
healthcheck:
50+
test: ["CMD-SHELL", "pg_isready -U user1 -d db1"]
51+
interval: 5s
52+
timeout: 5s
53+
retries: 5
54+
volumes:
55+
- postgres1_data:/var/lib/postgresql/data
56+
57+
postgres17:
58+
image: postgres:17beta1
59+
container_name: mcp_postgres17
4360
environment:
4461
POSTGRES_USER: user2
4562
POSTGRES_PASSWORD: password2
4663
POSTGRES_DB: db2
4764
ports:
48-
- "15432:5432"
65+
- "15433:5432"
4966
healthcheck:
5067
test: ["CMD-SHELL", "pg_isready -U user2 -d db2"]
5168
interval: 5s
5269
timeout: 5s
5370
retries: 5
71+
volumes:
72+
- postgres17_data:/var/lib/postgresql/data
5473

5574

5675
networks:
5776
default:
58-
name: mcp_test_network
77+
name: mcp_test_network
78+
79+
volumes:
80+
mysql1_data:
81+
postgres1_data:
82+
postgres17_data:

‎docker-compose.yml

+1-1
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ services:
5656
retries: 10
5757

5858
postgres1:
59-
image: postgres:15
59+
image: postgres:17
6060
environment:
6161
POSTGRES_USER: user1
6262
POSTGRES_PASSWORD: password1

‎examples/postgres_connection.go

+180
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
package main
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"log"
7+
"os"
8+
"time"
9+
10+
"github.com/FreePeak/db-mcp-server/pkg/db"
11+
)
12+
13+
func main() {
14+
// Example 1: Direct PostgreSQL 17 connection
15+
connectDirectly()
16+
17+
// Example 2: Using the DB Manager with configuration file
18+
connectWithManager()
19+
}
20+
21+
func connectDirectly() {
22+
fmt.Println("=== Example 1: Direct PostgreSQL 17 Connection ===")
23+
24+
// Create configuration for PostgreSQL 17
25+
config := db.Config{
26+
Type: "postgres",
27+
Host: getEnv("POSTGRES_HOST", "localhost"),
28+
Port: 5432,
29+
User: getEnv("POSTGRES_USER", "postgres"),
30+
Password: getEnv("POSTGRES_PASSWORD", "postgres"),
31+
Name: getEnv("POSTGRES_DB", "postgres"),
32+
33+
// PostgreSQL 17 specific options
34+
SSLMode: db.SSLPrefer,
35+
ApplicationName: "db-mcp-example",
36+
ConnectTimeout: 10,
37+
TargetSessionAttrs: "any", // Works with PostgreSQL 10+
38+
39+
// Additional options
40+
Options: map[string]string{
41+
"client_encoding": "UTF8",
42+
},
43+
44+
// Connection pool settings
45+
MaxOpenConns: 10,
46+
MaxIdleConns: 5,
47+
ConnMaxLifetime: 5 * time.Minute,
48+
ConnMaxIdleTime: 5 * time.Minute,
49+
}
50+
51+
// Create database connection
52+
database, err := db.NewDatabase(config)
53+
if err != nil {
54+
log.Fatalf("Failed to create database instance: %v", err)
55+
}
56+
57+
// Connect to the database
58+
fmt.Println("Connecting to PostgreSQL...")
59+
if err := database.Connect(); err != nil {
60+
log.Fatalf("Failed to connect to database: %v", err)
61+
}
62+
defer database.Close()
63+
64+
fmt.Println("Successfully connected to PostgreSQL")
65+
fmt.Println("Connection string (masked): ", database.ConnectionString())
66+
67+
// Query PostgreSQL version to verify compatibility
68+
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
69+
defer cancel()
70+
71+
var version string
72+
err = database.QueryRow(ctx, "SELECT version()").Scan(&version)
73+
if err != nil {
74+
log.Fatalf("Failed to query PostgreSQL version: %v", err)
75+
}
76+
77+
fmt.Printf("Connected to: %s\n", version)
78+
79+
// Run a sample query with PostgreSQL-style placeholders
80+
rows, err := database.Query(ctx, "SELECT datname FROM pg_database WHERE datistemplate = $1", false)
81+
if err != nil {
82+
log.Fatalf("Query failed: %v", err)
83+
}
84+
defer rows.Close()
85+
86+
fmt.Println("\nAvailable databases:")
87+
for rows.Next() {
88+
var dbName string
89+
if err := rows.Scan(&dbName); err != nil {
90+
log.Printf("Failed to scan row: %v", err)
91+
continue
92+
}
93+
fmt.Printf("- %s\n", dbName)
94+
}
95+
96+
if err = rows.Err(); err != nil {
97+
log.Printf("Error during row iteration: %v", err)
98+
}
99+
100+
fmt.Println()
101+
}
102+
103+
func connectWithManager() {
104+
fmt.Println("=== Example 2: Using DB Manager with Configuration ===")
105+
106+
// Create a database manager
107+
manager := db.NewDBManager()
108+
109+
// Create sample configuration with PostgreSQL 17 settings
110+
config := []byte(`{
111+
"connections": [
112+
{
113+
"id": "postgres17",
114+
"type": "postgres",
115+
"host": "localhost",
116+
"port": 5432,
117+
"name": "postgres",
118+
"user": "postgres",
119+
"password": "postgres",
120+
"ssl_mode": "prefer",
121+
"application_name": "db-mcp-example",
122+
"connect_timeout": 10,
123+
"target_session_attrs": "any",
124+
"options": {
125+
"client_encoding": "UTF8"
126+
},
127+
"max_open_conns": 10,
128+
"max_idle_conns": 5,
129+
"conn_max_lifetime_seconds": 300,
130+
"conn_max_idle_time_seconds": 60
131+
}
132+
]
133+
}`)
134+
135+
// Update with environment variables
136+
// In a real application, you would load this from a file
137+
// and use proper environment variable substitution
138+
139+
// Load configuration
140+
if err := manager.LoadConfig(config); err != nil {
141+
log.Fatalf("Failed to load database config: %v", err)
142+
}
143+
144+
// Connect to databases
145+
fmt.Println("Connecting to all configured databases...")
146+
if err := manager.Connect(); err != nil {
147+
log.Fatalf("Failed to connect to databases: %v", err)
148+
}
149+
defer manager.CloseAll()
150+
151+
// Get a specific database connection
152+
database, err := manager.GetDatabase("postgres17")
153+
if err != nil {
154+
log.Fatalf("Failed to get database: %v", err)
155+
}
156+
157+
fmt.Println("Successfully connected to PostgreSQL via manager")
158+
fmt.Println("Connection string (masked): ", database.ConnectionString())
159+
160+
// Query PostgreSQL version to verify compatibility
161+
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
162+
defer cancel()
163+
164+
var version string
165+
err = database.QueryRow(ctx, "SELECT version()").Scan(&version)
166+
if err != nil {
167+
log.Fatalf("Failed to query PostgreSQL version: %v", err)
168+
}
169+
170+
fmt.Printf("Connected to: %s\n", version)
171+
fmt.Println()
172+
}
173+
174+
// Helper function to get environment variable with fallback
175+
func getEnv(key, fallback string) string {
176+
if value, exists := os.LookupEnv(key); exists {
177+
return value
178+
}
179+
return fallback
180+
}

0 commit comments

Comments
 (0)