From 669c8c3cca76a01384eae7010db510a03f3894d5 Mon Sep 17 00:00:00 2001 From: Jyong <76649700+JohnJyong@users.noreply.github.com> Date: Thu, 28 Mar 2024 17:02:52 +0800 Subject: [PATCH] some optimization for admin api key, create tenant and reset-encrypt-key-pair command (#3013) Co-authored-by: jyong --- api/commands.py | 21 +++++++++++---------- api/libs/login.py | 2 +- api/services/account_service.py | 11 +++++------ 3 files changed, 17 insertions(+), 17 deletions(-) diff --git a/api/commands.py b/api/commands.py index b82d4d5d5d97a..9f1dc9528190f 100644 --- a/api/commands.py +++ b/api/commands.py @@ -109,19 +109,20 @@ def reset_encrypt_key_pair(): click.echo(click.style('Sorry, only support SELF_HOSTED mode.', fg='red')) return - tenant = db.session.query(Tenant).first() - if not tenant: - click.echo(click.style('Sorry, no workspace found. Please enter /install to initialize.', fg='red')) - return + tenants = db.session.query(Tenant).all() + for tenant in tenants: + if not tenant: + click.echo(click.style('Sorry, no workspace found. Please enter /install to initialize.', fg='red')) + return - tenant.encrypt_public_key = generate_key_pair(tenant.id) + tenant.encrypt_public_key = generate_key_pair(tenant.id) - db.session.query(Provider).filter(Provider.provider_type == 'custom').delete() - db.session.query(ProviderModel).delete() - db.session.commit() + db.session.query(Provider).filter(Provider.provider_type == 'custom', Provider.tenant_id == tenant.id).delete() + db.session.query(ProviderModel).filter(ProviderModel.tenant_id == tenant.id).delete() + db.session.commit() - click.echo(click.style('Congratulations! ' - 'the asymmetric key pair of workspace {} has been reset.'.format(tenant.id), fg='green')) + click.echo(click.style('Congratulations! ' + 'the asymmetric key pair of workspace {} has been reset.'.format(tenant.id), fg='green')) @click.command('vdb-migrate', help='migrate vector db.') diff --git a/api/libs/login.py b/api/libs/login.py index 5c03cfe957544..14085fe603530 100644 --- a/api/libs/login.py +++ b/api/libs/login.py @@ -53,7 +53,7 @@ def post(): def decorated_view(*args, **kwargs): auth_header = request.headers.get('Authorization') admin_api_key_enable = os.getenv('ADMIN_API_KEY_ENABLE', default='False') - if admin_api_key_enable: + if admin_api_key_enable.lower() == 'true': if auth_header: if ' ' not in auth_header: raise Unauthorized('Invalid Authorization header format. Expected \'Bearer \' format.') diff --git a/api/services/account_service.py b/api/services/account_service.py index 103af7f79c0ba..7fc61e40e3cb4 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -435,11 +435,13 @@ def register(cls, email, name, password: str = None, open_id: str = None, provid if open_id is not None or provider is not None: AccountService.link_account_integrate(provider, open_id, account) + if current_app.config['EDITION'] != 'SELF_HOSTED': + tenant = TenantService.create_tenant(f"{account.name}'s Workspace") - tenant = TenantService.create_tenant(f"{account.name}'s Workspace") + TenantService.create_tenant_member(tenant, account, role='owner') + account.current_tenant = tenant - TenantService.create_tenant_member(tenant, account, role='owner') - account.current_tenant = tenant + tenant_was_created.send(tenant) db.session.commit() except Exception as e: @@ -447,8 +449,6 @@ def register(cls, email, name, password: str = None, open_id: str = None, provid logging.error(f'Register failed: {e}') raise AccountRegisterError(f'Registration failed: {e}') from e - tenant_was_created.send(tenant) - return account @classmethod @@ -461,7 +461,6 @@ def invite_new_member(cls, tenant: Tenant, email: str, language: str, role: str name = email.split('@')[0] account = cls.register(email=email, name=name, language=language, status=AccountStatus.PENDING) - # Create new tenant member for invited tenant TenantService.create_tenant_member(tenant, account, role) TenantService.switch_tenant(account, tenant.id)