From 226a85de500295dc0590ee1e4e949b040f231d6a Mon Sep 17 00:00:00 2001 From: Justin Sheu Date: Thu, 8 May 2025 21:59:25 -0700 Subject: [PATCH 1/2] more robust root user creation for existing projects + return api key and org id --- judgment/cli.py | 8 +-- judgment/command_utils/self_host.py | 10 +++- judgment/self_host/supabase/supabase.py | 77 +++++++++++++++++++++---- 3 files changed, 75 insertions(+), 20 deletions(-) diff --git a/judgment/cli.py b/judgment/cli.py index 6ff8dcf..da7a8c2 100644 --- a/judgment/cli.py +++ b/judgment/cli.py @@ -10,7 +10,7 @@ from typing_extensions import Annotated -app = typer.Typer(help="Judgment CLI tool for managing self-hosted instances.", add_completion=False) +app = typer.Typer(help="Judgment CLI tool for managing self-hosted instances.", add_completion=False, pretty_exceptions_enable=False) self_host_app = typer.Typer(help="Commands for self-hosting Judgment", add_completion=False) app.add_typer(self_host_app, name="self-host") @@ -167,11 +167,7 @@ def self_host( typer.echo() # Run the deployment - try: - deploy(creds, supabase_compute_size, root_judgment_email, root_judgment_password, domain_name, invitation_email_service) - except Exception as e: - typer.echo(f"Error during deployment: {str(e)}", err=True) - raise typer.Exit(1) + deploy(creds, supabase_compute_size, root_judgment_email, root_judgment_password, domain_name, invitation_email_service) @self_host_app.command(name="https-listener") def https_listener(): diff --git a/judgment/command_utils/self_host.py b/judgment/command_utils/self_host.py index 38a85e7..f33301b 100644 --- a/judgment/command_utils/self_host.py +++ b/judgment/command_utils/self_host.py @@ -10,14 +10,18 @@ def deploy(creds: dict, supabase_compute_size: str, root_judgment_email: str, ro """Deploy a self-hosted instance of Judgment.""" supabase_token = creds["supabase_token"] org_id = creds["org_id"] - project_name = "Judgment Database" + project_name = "Judgment Database 2" db_password = creds["db_password"] # Create Supabase project and get secrets supabase_client = SupabaseClient(supabase_token, org_id, db_password) supabase_secrets, project_exists = supabase_client.create_project_and_get_secrets(project_name, supabase_compute_size) - if not project_exists: - supabase_client.create_root_user(supabase_secrets["supabase_url"], supabase_secrets["supabase_service_role_key"], root_judgment_email, root_judgment_password) + root_user_id = supabase_client.create_root_user(supabase_secrets["supabase_url"], supabase_secrets["supabase_service_role_key"], root_judgment_email, root_judgment_password) + root_user_api_key, root_user_org_id = supabase_client.get_user_api_key_and_org(supabase_secrets["supabase_url"], supabase_secrets["supabase_service_role_key"], root_user_id) + print(f"Root user Judgment API key: {root_user_api_key}") + print(f"Root user organization ID: {root_user_org_id}") + print(f"Make sure to save these credentials in a secure location. You will need them to run the Judgeval SDK against your self-hosted Judgment instance.") + if not typer.confirm("Would you like to proceed with AWS infrastructure deployment?"): print("Exiting... You can run the same command you just ran to deploy the AWS infrastructure with the Supabase project that was just created. Just enter 'y' when prompted to use the existing project.") raise typer.Exit(0) diff --git a/judgment/self_host/supabase/supabase.py b/judgment/self_host/supabase/supabase.py index afcae02..8c3f2a5 100644 --- a/judgment/self_host/supabase/supabase.py +++ b/judgment/self_host/supabase/supabase.py @@ -101,24 +101,79 @@ def load_schema(self, db_url: str): cursor.execute(schema_sql) def create_root_user(self, supabase_url: str, supabase_service_role_key: str, email: str, password: str): - """Create a root user using Supabase Auth API.""" - print("Creating root user...") + """Create a root user using Supabase Auth API or return existing root user if credentials match. + + Args: + supabase_url: The Supabase project URL + supabase_service_role_key: The service role key for authentication + email: The email of the root user + password: The password of the root user + + Returns: + The user ID of the root user (either existing or newly created) + """ + print("Checking for existing root user with provided credentials...") supabase: Client = create_client(supabase_url, supabase_service_role_key) - # Create the user try: - response = supabase.auth.admin.create_user({ + # Try to sign in first + sign_in_response = supabase.auth.sign_in_with_password({ "email": email, - "password": password, - "email_confirm": True, # Auto-confirm the email - "user_metadata": { - "role": "root" - } + "password": password }) + user = sign_in_response.user + + # Verify this is a root user + if user.user_metadata.get("role") != "root": + raise Exception("The provided credentials belong to a non-root user. Please run the command again with root user credentials or new credentials to create a root user with.") + + print("Found existing root user with matching credentials!") + return user.id except Exception as e: - raise Exception(f"Failed to create root user: {e}") + if "non-root user" in str(e): + raise e + # If sign in fails, try to create the user + print("Creating new root user...") + try: + response = supabase.auth.admin.create_user({ + "email": email, + "password": password, + "email_confirm": True, # Auto-confirm the email + "user_metadata": { + "role": "root" + } + }) + print("Root user created successfully!") + return response.user.id + except Exception as e: + raise Exception(f"Failed to create root user: {e}") + + def get_user_api_key_and_org(self, supabase_url: str, supabase_service_role_key: str, user_id: str) -> tuple[str, str]: + """Get the user's API key and organization ID after creation. + + Args: + supabase_url: The Supabase project URL + supabase_service_role_key: The service role key for authentication + user_id: The ID of the user to get information for + + Returns: + A tuple containing (api_key, organization_id) + """ + supabase: Client = create_client(supabase_url, supabase_service_role_key) + + # Get the API key from user_data + api_key_result = supabase.table("user_data").select("judgment_api_key").eq("id", user_id).execute() + if not api_key_result.data: + raise Exception(f"No user data found for user {user_id}") + api_key = api_key_result.data[0]["judgment_api_key"] + + # Get the organization ID from user_organizations + org_result = supabase.table("user_organizations").select("organization_id").eq("user_id", user_id).execute() + if not org_result.data: + raise Exception(f"No organization found for user {user_id}") + organization_id = org_result.data[0]["organization_id"] - print("Root user created successfully!") + return api_key, organization_id def create_project_and_get_secrets(self, project_name: str = "Judgment Database", supabase_compute_size: str = "small") -> Dict: # Check for existing project From 5159f84b1dc264fa679d57f76234fc41949e7894 Mon Sep 17 00:00:00 2001 From: Justin Sheu Date: Thu, 8 May 2025 22:19:05 -0700 Subject: [PATCH 2/2] change exception handling + changes supabase proj name --- judgment/command_utils/self_host.py | 10 +++------- judgment/self_host/supabase/supabase.py | 16 ++++++++++++---- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/judgment/command_utils/self_host.py b/judgment/command_utils/self_host.py index f33301b..38a85e7 100644 --- a/judgment/command_utils/self_host.py +++ b/judgment/command_utils/self_host.py @@ -10,18 +10,14 @@ def deploy(creds: dict, supabase_compute_size: str, root_judgment_email: str, ro """Deploy a self-hosted instance of Judgment.""" supabase_token = creds["supabase_token"] org_id = creds["org_id"] - project_name = "Judgment Database 2" + project_name = "Judgment Database" db_password = creds["db_password"] # Create Supabase project and get secrets supabase_client = SupabaseClient(supabase_token, org_id, db_password) supabase_secrets, project_exists = supabase_client.create_project_and_get_secrets(project_name, supabase_compute_size) - root_user_id = supabase_client.create_root_user(supabase_secrets["supabase_url"], supabase_secrets["supabase_service_role_key"], root_judgment_email, root_judgment_password) - root_user_api_key, root_user_org_id = supabase_client.get_user_api_key_and_org(supabase_secrets["supabase_url"], supabase_secrets["supabase_service_role_key"], root_user_id) + if not project_exists: + supabase_client.create_root_user(supabase_secrets["supabase_url"], supabase_secrets["supabase_service_role_key"], root_judgment_email, root_judgment_password) - print(f"Root user Judgment API key: {root_user_api_key}") - print(f"Root user organization ID: {root_user_org_id}") - print(f"Make sure to save these credentials in a secure location. You will need them to run the Judgeval SDK against your self-hosted Judgment instance.") - if not typer.confirm("Would you like to proceed with AWS infrastructure deployment?"): print("Exiting... You can run the same command you just ran to deploy the AWS infrastructure with the Supabase project that was just created. Just enter 'y' when prompted to use the existing project.") raise typer.Exit(0) diff --git a/judgment/self_host/supabase/supabase.py b/judgment/self_host/supabase/supabase.py index f3b1bf9..ae48a12 100644 --- a/judgment/self_host/supabase/supabase.py +++ b/judgment/self_host/supabase/supabase.py @@ -11,6 +11,10 @@ load_dotenv() +class NonRootUserError(Exception): + """Exception raised when a non-root user attempts to use root-only functionality.""" + pass + class SupabaseClient: def __init__(self, supabase_token: str, org_id: str, db_password: str): self.supabase_token = supabase_token @@ -121,6 +125,10 @@ def create_root_user(self, supabase_url: str, supabase_service_role_key: str, em Returns: The user ID of the root user (either existing or newly created) + + Raises: + NonRootUserError: If the provided credentials belong to a non-root user + Exception: For other errors during user creation or authentication """ print("Checking for existing root user with provided credentials...") supabase: Client = create_client(supabase_url, supabase_service_role_key) @@ -135,13 +143,13 @@ def create_root_user(self, supabase_url: str, supabase_service_role_key: str, em # Verify this is a root user if user.user_metadata.get("role") != "root": - raise Exception("The provided credentials belong to a non-root user. Please run the command again with root user credentials or new credentials to create a root user with.") + raise NonRootUserError("The provided credentials belong to a non-root user. Please run the command again with root user credentials or new credentials to create a root user with.") print("Found existing root user with matching credentials!") return user.id - except Exception as e: - if "non-root user" in str(e): - raise e + except NonRootUserError: + raise + except Exception: # If sign in fails, try to create the user print("Creating new root user...") try: