122 lines
3.3 KiB
Python
122 lines
3.3 KiB
Python
"""
|
|
Main script to run all seeders in sequence.
|
|
Checks dependencies between seeders and runs them in the correct order.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import logging
|
|
import argparse
|
|
from dotenv import load_dotenv
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Import seeders
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
from scripts.seeders.admin_seeder import create_admin_user
|
|
from scripts.seeders.client_seeder import create_demo_client_and_user
|
|
from scripts.seeders.mcp_server_seeder import create_mcp_servers
|
|
from scripts.seeders.tool_seeder import create_tools
|
|
|
|
|
|
def setup_environment():
|
|
"""Configure the environment for seeders"""
|
|
load_dotenv()
|
|
|
|
# Check if essential environment variables are defined
|
|
required_vars = ["POSTGRES_CONNECTION_STRING"]
|
|
missing_vars = [var for var in required_vars if not os.getenv(var)]
|
|
|
|
if missing_vars:
|
|
logger.error(
|
|
f"Required environment variables not defined: {', '.join(missing_vars)}"
|
|
)
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def run_seeders(seeders):
|
|
"""
|
|
Run the specified seeders
|
|
|
|
Args:
|
|
seeders (list): List of seeders to run
|
|
|
|
Returns:
|
|
bool: True if all seeders were executed successfully, False otherwise
|
|
"""
|
|
all_seeders = {
|
|
"admin": create_admin_user,
|
|
"client": create_demo_client_and_user,
|
|
"mcp_servers": create_mcp_servers,
|
|
"tools": create_tools,
|
|
}
|
|
|
|
# Define the correct execution order (dependencies)
|
|
seeder_order = ["admin", "client", "mcp_servers", "tools"]
|
|
|
|
# If no seeder is specified, run all
|
|
if not seeders:
|
|
seeders = seeder_order
|
|
else:
|
|
# Check if all specified seeders exist
|
|
invalid_seeders = [s for s in seeders if s not in all_seeders]
|
|
if invalid_seeders:
|
|
logger.error(f"Invalid seeders: {', '.join(invalid_seeders)}")
|
|
logger.info(f"Available seeders: {', '.join(all_seeders.keys())}")
|
|
return False
|
|
|
|
# Ensure seeders are executed in the correct order
|
|
seeders = [s for s in seeder_order if s in seeders]
|
|
|
|
# Run seeders
|
|
success = True
|
|
for seeder_name in seeders:
|
|
logger.info(f"Running seeder: {seeder_name}")
|
|
|
|
try:
|
|
seeder_func = all_seeders[seeder_name]
|
|
if not seeder_func():
|
|
logger.error(f"Failed to run seeder: {seeder_name}")
|
|
success = False
|
|
except Exception as e:
|
|
logger.error(f"Error running seeder {seeder_name}: {str(e)}")
|
|
success = False
|
|
|
|
return success
|
|
|
|
|
|
def main():
|
|
"""Main function"""
|
|
parser = argparse.ArgumentParser(description="Run seeders to populate the database")
|
|
parser.add_argument(
|
|
"--seeders",
|
|
nargs="+",
|
|
help="Seeders to run (admin, client, mcp_servers, tools)",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# Configure environment
|
|
if not setup_environment():
|
|
sys.exit(1)
|
|
|
|
# Run seeders
|
|
success = run_seeders(args.seeders)
|
|
|
|
# Output
|
|
if success:
|
|
logger.info("All seeders were executed successfully")
|
|
sys.exit(0)
|
|
else:
|
|
logger.error("There were errors running the seeders")
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|