Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 0 additions & 59 deletions src/mcp_services/postgres/postgres_state_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,65 +77,6 @@ def _test_connection(self):
conn = psycopg2.connect(**self.conn_params, database="postgres")
conn.close()

def _setup_database(self):
"""Setup all required databases by downloading and restoring from backup."""
databases = ['employees', 'chinook', 'dvdrental', 'sports', 'lego']

for db_name in databases:
if not self._database_exists(db_name):
logger.info(f"Setting up {db_name} database...")

# Path to backup file
backup_dir = Path(__file__).parent.parent.parent.parent / "postgres_state"
backup_file = backup_dir / f"{db_name}.backup"

# Download backup if not exists
if not backup_file.exists():
backup_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Downloading {db_name} backup...")
try:
import urllib.request
urllib.request.urlretrieve(
f'https://storage.mcpmark.ai/postgres/{db_name}.backup',
str(backup_file)
)
logger.info(f"{db_name} backup downloaded")
except Exception as e:
logger.warning(f"Failed to download {db_name} backup: {e}")
continue

# Create database
try:
self._create_empty_database(db_name)
logger.info(f"Created {db_name} database")
except Exception as e:
logger.warning(f"Failed to create {db_name} database: {e}")
continue

# Restore from backup
env = os.environ.copy()
env['PGPASSWORD'] = self.password

try:
result = subprocess.run([
'pg_restore',
'-h', str(self.host),
'-p', str(self.port),
'-U', self.username,
'-d', db_name,
'-v',
str(backup_file)
], env=env, capture_output=True, text=True)

if result.returncode != 0:
logger.warning(f"pg_restore had errors for {db_name}: {result.stderr}")
else:
logger.info(f"{db_name} database restored successfully")
except Exception as e:
logger.warning(f"Failed to restore {db_name} database: {e}")
else:
logger.debug(f"{db_name} database already exists")

def _setup_database(self):
"""Setup all required databases by downloading and restoring from backup."""
databases = ['employees', 'chinook', 'dvdrental', 'sports', 'lego']
Expand Down