highperfocused/danswer
1
0
Fork 0
mirror of https://github.com/danswer-ai/danswer.git synced 2025-03-19 14:22:10 +01:00
Code Issues Packages Projects Releases Wiki Activity

update drive

Browse Source
This commit is contained in:
pablodanswer 2024-10-10 17:52:21 -07:00
parent 63e45d3323
commit b23c5a08a1
2 changed files with 47 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

59
backend/danswer/connectors/google_drive/connector_auth.py
View File

@ -11,11 +11,8 @@ from google_auth_oauthlib.flow import InstalledAppFlow # type: ignore
from sqlalchemy.orm import Session
from danswer.configs.app_configs import ENTERPRISE_EDITION_ENABLED
from danswer.configs.app_configs import MULTI_TENANT
from danswer.configs.app_configs import WEB_DOMAIN
from danswer.configs.constants import DocumentSource
from danswer.configs.constants import KV_CLOUD_GOOGLE_DRIVE_CRED_KEY
from danswer.configs.constants import KV_CLOUD_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY
from danswer.configs.constants import KV_CRED_KEY
from danswer.configs.constants import KV_GOOGLE_DRIVE_CRED_KEY
from danswer.configs.constants import KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY
@ -145,13 +142,7 @@ def verify_csrf(credential_id: int, state: str) -> None:
def get_auth_url(credential_id: int) -> str:
creds_str = str(
get_kv_store().load(
KV_GOOGLE_DRIVE_CRED_KEY
if not MULTI_TENANT
else KV_CLOUD_GOOGLE_DRIVE_CRED_KEY
)
)
creds_str = str(get_kv_store().load(KV_GOOGLE_DRIVE_CRED_KEY))
credential_json = json.loads(creds_str)
flow = InstalledAppFlow.from_client_config(
credential_json,
@ -211,58 +202,28 @@ def build_service_account_creds(
def get_google_app_cred() -> GoogleAppCredentials:
if MULTI_TENANT:
creds_str = str(get_kv_store().load(KV_CLOUD_GOOGLE_DRIVE_CRED_KEY))
else:
creds_str = str(get_kv_store().load(KV_GOOGLE_DRIVE_CRED_KEY))
creds_str = str(get_kv_store().load(KV_GOOGLE_DRIVE_CRED_KEY))
return GoogleAppCredentials(**json.loads(creds_str))
def upsert_google_app_cred(
app_credentials: GoogleAppCredentials, cloud_enabled: bool
) -> None:
if cloud_enabled:
get_kv_store().store(
KV_CLOUD_GOOGLE_DRIVE_CRED_KEY, app_credentials.json(), encrypt=True
)
else:
get_kv_store().store(
KV_GOOGLE_DRIVE_CRED_KEY, app_credentials.json(), encrypt=True
)
def upsert_google_app_cred(app_credentials: GoogleAppCredentials) -> None:
get_kv_store().store(KV_GOOGLE_DRIVE_CRED_KEY, app_credentials.json(), encrypt=True)
def delete_google_app_cred() -> None:
if MULTI_TENANT:
get_kv_store().delete(KV_CLOUD_GOOGLE_DRIVE_CRED_KEY)
else:
get_kv_store().delete(KV_GOOGLE_DRIVE_CRED_KEY)
get_kv_store().delete(KV_GOOGLE_DRIVE_CRED_KEY)
def get_service_account_key() -> GoogleServiceAccountKey:
if MULTI_TENANT:
creds_str = str(get_kv_store().load(KV_CLOUD_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY))
else:
creds_str = str(get_kv_store().load(KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY))
creds_str = str(get_kv_store().load(KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY))
return GoogleServiceAccountKey(**json.loads(creds_str))
def upsert_service_account_key(service_account_key: GoogleServiceAccountKey) -> None:
if MULTI_TENANT:
get_kv_store().store(
KV_CLOUD_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY,
service_account_key.json(),
encrypt=True,
)
else:
get_kv_store().store(
KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY,
service_account_key.json(),
encrypt=True,
)
get_kv_store().store(
KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY, service_account_key.json(), encrypt=True
)
def delete_service_account_key() -> None:
if MULTI_TENANT:
get_kv_store().delete(KV_CLOUD_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY)
else:
get_kv_store().delete(KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY)
get_kv_store().delete(KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY)

39
backend/danswer/server/documents/connector.py
View File

@ -125,6 +125,12 @@ def check_google_app_gmail_credentials_exist(
def upsert_google_app_gmail_credentials(
app_credentials: GoogleAppCredentials, _: User = Depends(current_admin_user)
) -> StatusResponse:
if MULTI_TENANT:
raise HTTPException(
status_code=400,
detail="Modifying Google App Credentials is not supported in multi-tenant mode",
)
try:
upsert_google_app_gmail_cred(app_credentials)
except ValueError as e:
@ -163,10 +169,15 @@ def check_google_app_credentials_exist(
def upsert_google_app_credentials(
app_credentials: GoogleAppCredentials,
_: User = Depends(current_admin_user),
cloud_enabled: bool = Query(MULTI_TENANT, alias="cloud_enabled"),
) -> StatusResponse:
if MULTI_TENANT:
raise HTTPException(
status_code=400,
detail="Modifying Google App Credentials is not supported in multi-tenant mode",
)
try:
upsert_google_app_cred(app_credentials, cloud_enabled)
upsert_google_app_cred(app_credentials)
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
@ -205,6 +216,12 @@ def check_google_service_gmail_account_key_exist(
def upsert_google_service_gmail_account_key(
service_account_key: GoogleServiceAccountKey, _: User = Depends(current_admin_user)
) -> StatusResponse:
if MULTI_TENANT:
raise HTTPException(
status_code=400,
detail="Modifying Google App Credentials is not supported in multi-tenant mode",
)
try:
upsert_gmail_service_account_key(service_account_key)
except ValueError as e:
@ -245,6 +262,12 @@ def check_google_service_account_key_exist(
def upsert_google_service_account_key(
service_account_key: GoogleServiceAccountKey, _: User = Depends(current_admin_user)
) -> StatusResponse:
if MULTI_TENANT:
raise HTTPException(
status_code=400,
detail="Modifying Google App Credentials is not supported in multi-tenant mode",
)
try:
upsert_service_account_key(service_account_key)
except ValueError as e:
@ -278,6 +301,12 @@ def upsert_service_account_credential(
"""Special API which allows the creation of a credential for a service account.
Combines the input with the saved service account key to create an entry in the
`Credential` table."""
if MULTI_TENANT:
raise HTTPException(
status_code=400,
detail="Modifying Google App Credentials is not supported in multi-tenant mode",
)
try:
credential_base = build_service_account_creds(
DocumentSource.GOOGLE_DRIVE,
@ -304,6 +333,12 @@ def upsert_gmail_service_account_credential(
"""Special API which allows the creation of a credential for a service account.
Combines the input with the saved service account key to create an entry in the
`Credential` table."""
if MULTI_TENANT:
raise HTTPException(
status_code=400,
detail="Modifying Google App Credentials is not supported in multi-tenant mode",
)
try:
credential_base = build_service_account_creds(
DocumentSource.GMAIL,