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'))
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.')

View File

@ -53,7 +53,7 @@ def login_required(func):
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 <api-key>\' format.')

View File

@ -435,11 +435,13 @@ class RegisterService:
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 @@ class RegisterService:
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 @@ class RegisterService:
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)