test_connections.py•6.88 kB
#!/usr/bin/env python3
"""
Database Connection Test Script
This script tests the database connections for all supported database types.
It helps verify that the Universal SQL MCP Server can connect to your databases.
"""
import os
import sys
import logging
from typing import Dict, Any
# Add the current directory to the path so we can import our modules
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from database import DatabaseConfig, DatabaseManager
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
def test_database_connection(db_type: str, config_overrides: Dict[str, str] = None) -> bool:
"""Test connection to a specific database type"""
# Save original environment variables
original_env = {}
for key in ['DB_TYPE', 'DB_HOST', 'DB_PORT', 'DB_USER', 'DB_PASSWORD', 'DB_NAME', 'DB_DRIVER']:
original_env[key] = os.getenv(key)
try:
# Set test configuration
test_configs = {
'mysql': {
'DB_TYPE': 'mysql',
'DB_HOST': 'localhost',
'DB_PORT': '3306',
'DB_USER': 'testuser',
'DB_PASSWORD': 'testpass',
'DB_NAME': 'testdb'
},
'postgresql': {
'DB_TYPE': 'postgresql',
'DB_HOST': 'localhost',
'DB_PORT': '5432',
'DB_USER': 'postgres',
'DB_PASSWORD': 'testpass',
'DB_NAME': 'testdb'
},
'sqlite': {
'DB_TYPE': 'sqlite',
'DB_NAME': './test_database.db'
},
'sqlserver': {
'DB_TYPE': 'sqlserver',
'DB_HOST': 'localhost',
'DB_PORT': '1433',
'DB_USER': 'sa',
'DB_PASSWORD': 'TestPass123',
'DB_NAME': 'testdb',
'DB_DRIVER': 'ODBC Driver 17 for SQL Server'
}
}
config = test_configs.get(db_type, {})
if config_overrides:
config.update(config_overrides)
# Apply configuration
for key, value in config.items():
os.environ[key] = value
# Test connection
logger.info(f"Testing {db_type.upper()} connection...")
db_manager = DatabaseManager()
if db_manager.test_connection():
logger.info(f"✅ {db_type.upper()} connection successful!")
# Test basic query
try:
if db_type == 'sqlite':
# Create a simple table for SQLite test
with db_manager.get_connection() as conn:
cursor = conn.cursor()
cursor.execute("CREATE TABLE IF NOT EXISTS test_table (id INTEGER PRIMARY KEY, name TEXT)")
cursor.execute("INSERT OR REPLACE INTO test_table (id, name) VALUES (1, 'test')")
conn.commit()
cursor.close()
# Test schema retrieval
tables = db_manager.get_table_descriptions()
logger.info(f"✅ Schema retrieval successful! Found {len(tables)} tables")
# Test SELECT query
if db_type == 'sqlite':
results = db_manager.execute_query("SELECT * FROM test_table LIMIT 1")
else:
results = db_manager.execute_query("SELECT 1 as test_column")
logger.info(f"✅ Query execution successful! Returned {len(results)} rows")
return True
except Exception as e:
logger.warning(f"⚠️ Basic operations failed for {db_type.upper()}: {e}")
return False
else:
logger.error(f"❌ {db_type.upper()} connection failed!")
return False
except Exception as e:
logger.error(f"❌ {db_type.upper()} test failed: {e}")
return False
finally:
# Restore original environment variables
for key, value in original_env.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = value
def main():
"""Main test function"""
print("🔍 Universal SQL MCP Server - Database Connection Test")
print("=" * 60)
# Test configurations from command line or environment
test_mysql = os.getenv('TEST_MYSQL', 'false').lower() == 'true'
test_postgresql = os.getenv('TEST_POSTGRESQL', 'false').lower() == 'true'
test_sqlite = os.getenv('TEST_SQLITE', 'true').lower() == 'true' # Default to true as it needs no setup
test_sqlserver = os.getenv('TEST_SQLSERVER', 'false').lower() == 'true'
results = {}
if test_sqlite:
print("\n📁 Testing SQLite...")
results['sqlite'] = test_database_connection('sqlite')
if test_mysql:
print("\n🐬 Testing MySQL...")
results['mysql'] = test_database_connection('mysql')
if test_postgresql:
print("\n🐘 Testing PostgreSQL...")
results['postgresql'] = test_database_connection('postgresql')
if test_sqlserver:
print("\n🏢 Testing SQL Server...")
results['sqlserver'] = test_database_connection('sqlserver')
# Summary
print("\n" + "=" * 60)
print("📊 Test Results Summary:")
print("-" * 30)
for db_type, success in results.items():
status = "✅ PASS" if success else "❌ FAIL"
print(f"{db_type.upper():12} : {status}")
successful_tests = sum(1 for success in results.values() if success)
total_tests = len(results)
print(f"\nTotal: {successful_tests}/{total_tests} tests passed")
if successful_tests == total_tests and total_tests > 0:
print("🎉 All enabled database tests passed!")
return 0
elif successful_tests > 0:
print("⚠️ Some database tests failed. Check the logs above.")
return 1
else:
print("❌ All database tests failed or no tests were enabled.")
return 1
if __name__ == "__main__":
print("Usage:")
print(" # Test only SQLite (default)")
print(" python test_connections.py")
print("")
print(" # Test specific databases")
print(" TEST_MYSQL=true python test_connections.py")
print(" TEST_POSTGRESQL=true python test_connections.py")
print(" TEST_SQLSERVER=true python test_connections.py")
print("")
print(" # Test multiple databases")
print(" TEST_SQLITE=true TEST_MYSQL=true python test_connections.py")
print("")
sys.exit(main())