some optimization for admin api key, create tenant and reset-encrypt-key-pair command (#3013)

Co-authored-by: jyong <jyong@dify.ai>
This commit is contained in:
Jyong 2024-03-28 17:02:52 +08:00 committed by GitHub
parent b0b0cc045f
commit 669c8c3cca
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 17 additions and 17 deletions

View File

@ -109,19 +109,20 @@ def reset_encrypt_key_pair():
click.echo(click.style('Sorry, only support SELF_HOSTED mode.', fg='red')) click.echo(click.style('Sorry, only support SELF_HOSTED mode.', fg='red'))
return return
tenant = db.session.query(Tenant).first() tenants = db.session.query(Tenant).all()
if not tenant: for tenant in tenants:
click.echo(click.style('Sorry, no workspace found. Please enter /install to initialize.', fg='red')) if not tenant:
return 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(Provider).filter(Provider.provider_type == 'custom', Provider.tenant_id == tenant.id).delete()
db.session.query(ProviderModel).delete() db.session.query(ProviderModel).filter(ProviderModel.tenant_id == tenant.id).delete()
db.session.commit() db.session.commit()
click.echo(click.style('Congratulations! ' click.echo(click.style('Congratulations! '
'the asymmetric key pair of workspace {} has been reset.'.format(tenant.id), fg='green')) 'the asymmetric key pair of workspace {} has been reset.'.format(tenant.id), fg='green'))
@click.command('vdb-migrate', help='migrate vector db.') @click.command('vdb-migrate', help='migrate vector db.')

View File

@ -53,7 +53,7 @@ def login_required(func):
def decorated_view(*args, **kwargs): def decorated_view(*args, **kwargs):
auth_header = request.headers.get('Authorization') auth_header = request.headers.get('Authorization')
admin_api_key_enable = os.getenv('ADMIN_API_KEY_ENABLE', default='False') 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 auth_header:
if ' ' not in auth_header: if ' ' not in auth_header:
raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.') raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')

View File

@ -435,11 +435,13 @@ class RegisterService:
if open_id is not None or provider is not None: if open_id is not None or provider is not None:
AccountService.link_account_integrate(provider, open_id, account) 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') tenant_was_created.send(tenant)
account.current_tenant = tenant
db.session.commit() db.session.commit()
except Exception as e: except Exception as e:
@ -447,8 +449,6 @@ class RegisterService:
logging.error(f'Register failed: {e}') logging.error(f'Register failed: {e}')
raise AccountRegisterError(f'Registration failed: {e}') from e raise AccountRegisterError(f'Registration failed: {e}') from e
tenant_was_created.send(tenant)
return account return account
@classmethod @classmethod
@ -461,7 +461,6 @@ class RegisterService:
name = email.split('@')[0] name = email.split('@')[0]
account = cls.register(email=email, name=name, language=language, status=AccountStatus.PENDING) account = cls.register(email=email, name=name, language=language, status=AccountStatus.PENDING)
# Create new tenant member for invited tenant # Create new tenant member for invited tenant
TenantService.create_tenant_member(tenant, account, role) TenantService.create_tenant_member(tenant, account, role)
TenantService.switch_tenant(account, tenant.id) TenantService.switch_tenant(account, tenant.id)