joelgilbert/NL2SQL
0
1"""2Pre-deployment health check script.3 4This script validates your environment configuration and tests5all critical components before deployment.6 7Usage:8 python scripts/health_check.py9"""10 11import sys12import os13 14# Add parent directory to path15sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))16 17import logging18from colorama import init, Fore, Style19 20# Initialize colorama for colored output21init(autoreset=True)22 23logging.basicConfig(level=logging.WARNING)24 25 26def print_header(text):27 """Print section header."""28 print(f"\n{Fore.CYAN}{'='*60}")29 print(f"{Fore.CYAN}{text:^60}")30 print(f"{Fore.CYAN}{'='*60}{Style.RESET_ALL}\n")31 32 33def print_success(text):34 """Print success message."""35 print(f"{Fore.GREEN}✅ {text}{Style.RESET_ALL}")36 37 38def print_error(text):39 """Print error message."""40 print(f"{Fore.RED}❌ {text}{Style.RESET_ALL}")41 42 43def print_warning(text):44 """Print warning message."""45 print(f"{Fore.YELLOW}⚠️ {text}{Style.RESET_ALL}")46 47 48def check_environment_variables():49 """Check if all required environment variables are set."""50 print_header("Environment Variables Check")51 52 required_vars = [53 "NEON_READONLY_CONNECTION_STRING",54 "NEON_DBA_CONNECTION_STRING",55 "GROQ_API_KEY",56 "CLOUDFLARE_ACCOUNT_ID",57 "CLOUDFLARE_AUTH_TOKEN",58 "UPSTASH_VECTOR_URL",59 "UPSTASH_VECTOR_TOKEN",60 "DBA_PASSWORD"61 ]62 63 all_set = True64 65 for var in required_vars:66 value = os.getenv(var)67 if value:68 # Show first/last few chars only for security69 masked = f"{value[:8]}...{value[-4:]}" if len(value) > 15 else "***"70 print_success(f"{var}: {masked}")71 else:72 print_error(f"{var}: NOT SET")73 all_set = False74 75 return all_set76 77 78def check_database_connection():79 """Test database connections."""80 print_header("Database Connection Check")81 82 try:83 from database.connection import db84 85 # Test readonly connection86 if db.test_connection():87 print_success("Read-only database connection successful")88 else:89 print_error("Read-only database connection failed")90 return False91 92 # Test DBA connection93 if db.test_dba_connection():94 print_success("DBA database connection successful")95 else:96 print_warning("DBA database connection failed (check credentials)")97 98 return True99 100 except Exception as e:101 print_error(f"Database connection error: {e}")102 return False103 104 105def check_api_keys():106 """Test API connections."""107 print_header("API Keys Check")108 109 # Test Groq110 try:111 from agents.gatekeeper import gatekeeper112 result = gatekeeper.classify_intent("Hello")113 if result and "intent" in result:114 print_success("Groq API connection successful")115 else:116 print_error("Groq API test failed")117 except Exception as e:118 print_error(f"Groq API error: {e}")119 120 # Test Cloudflare121 try:122 from agents.sql_generator import sql_generator123 # This is a basic connectivity test124 print_success("Cloudflare Workers AI configured")125 except Exception as e:126 print_error(f"Cloudflare Workers AI error: {e}")127 128 # Test Upstash129 try:130 from vector_store.upstash_client import upstash_client131 print_success("Upstash Vector configured")132 except Exception as e:133 print_error(f"Upstash Vector error: {e}")134 135 return True136 137 138def check_dependencies():139 """Check if all required packages are installed."""140 print_header("Dependencies Check")141 142 required_packages = [143 "streamlit",144 "psycopg2",145 "groq",146 "requests",147 "sqlparse",148 "pydantic",149 "pytest",150 "pandas",151 "tabulate"152 ]153 154 all_installed = True155 156 for package in required_packages:157 try:158 __import__(package.replace("-", "_"))159 print_success(f"{package} installed")160 except ImportError:161 print_error(f"{package} NOT installed")162 all_installed = False163 164 return all_installed165 166 167def check_file_structure():168 """Verify project file structure."""169 print_header("File Structure Check")170 171 required_files = [172 "app.py",173 "requirements.txt",174 ".env",175 ".gitignore",176 "config/settings.py",177 "config/prompts.py",178 "database/connection.py",179 "database/schema_manager.py",180 "database/query_executor.py",181 "agents/gatekeeper.py",182 "agents/sql_generator.py",183 "agents/explainer.py",184 "security/validator.py",185 "security/auth.py",186 "security/audit_logger.py",187 ".streamlit/config.toml"188 ]189 190 all_exist = True191 base_dir = os.path.dirname(os.path.dirname(__file__))192 193 for file_path in required_files:194 full_path = os.path.join(base_dir, file_path)195 if os.path.exists(full_path):196 print_success(f"{file_path}")197 else:198 print_error(f"{file_path} NOT FOUND")199 all_exist = False200 201 return all_exist202 203 204def check_vector_store():205 """Check if vector store is initialized."""206 print_header("Vector Store Check")207 208 try:209 from vector_store.upstash_client import upstash_client210 211 # Try to search (this will fail gracefully if not initialized)212 results = upstash_client.search_similar_schemas("test query", top_k=1)213 214 if results:215 print_success(f"Vector store initialized ({len(results)} schemas found)")216 return True217 else:218 print_warning("Vector store appears empty - run: python scripts/init_vector_store.py")219 return False220 221 except Exception as e:222 print_error(f"Vector store error: {e}")223 return False224 225 226def run_basic_tests():227 """Run basic unit tests."""228 print_header("Basic Tests")229 230 try:231 import pytest232 233 # Run tests234 result = pytest.main([235 "tests/test_setup.py",236 "-v",237 "--tb=short",238 "-q"239 ])240 241 if result == 0:242 print_success("All basic tests passed")243 return True244 else:245 print_error("Some tests failed")246 return False247 248 except Exception as e:249 print_error(f"Test execution error: {e}")250 return False251 252 253def main():254 """Run all health checks."""255 print(f"\n{Fore.MAGENTA}{'='*60}")256 print(f"{Fore.MAGENTA}{'NL2SQL System Health Check':^60}")257 print(f"{Fore.MAGENTA}{'='*60}{Style.RESET_ALL}")258 259 checks = [260 ("Environment Variables", check_environment_variables),261 ("File Structure", check_file_structure),262 ("Dependencies", check_dependencies),263 ("Database Connection", check_database_connection),264 ("API Keys", check_api_keys),265 ("Vector Store", check_vector_store),266 ("Basic Tests", run_basic_tests),267 ]268 269 results = {}270 271 for check_name, check_func in checks:272 try:273 results[check_name] = check_func()274 except Exception as e:275 print_error(f"Check '{check_name}' failed with error: {e}")276 results[check_name] = False277 278 # Summary279 print_header("Summary")280 281 passed = sum(1 for v in results.values() if v)282 total = len(results)283 284 for check_name, result in results.items():285 if result:286 print_success(f"{check_name}: PASSED")287 else:288 print_error(f"{check_name}: FAILED")289 290 print(f"\n{Fore.CYAN}Overall: {passed}/{total} checks passed{Style.RESET_ALL}")291 292 if passed == total:293 print(f"\n{Fore.GREEN}{'='*60}")294 print(f"{Fore.GREEN}{'✅ ALL CHECKS PASSED!':^60}")295 print(f"{Fore.GREEN}{'Your system is ready for deployment':^60}")296 print(f"{Fore.GREEN}{'='*60}{Style.RESET_ALL}\n")297 return 0298 else:299 print(f"\n{Fore.RED}{'='*60}")300 print(f"{Fore.RED}{'❌ SOME CHECKS FAILED':^60}")301 print(f"{Fore.RED}{'Please fix the issues before deployment':^60}")302 print(f"{Fore.RED}{'='*60}{Style.RESET_ALL}\n")303 return 1304 305 306if __name__ == "__main__":307 try:308 sys.exit(main())309 except KeyboardInterrupt:310 print(f"\n{Fore.YELLOW}Health check cancelled by user{Style.RESET_ALL}")311 sys.exit(1)312