mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-04-07 19:38:19 +02:00
add creator id to cc pair (#3121)
* add creator id to cc pair * fix alembic head * show email instead of UUID * safer check on email * make foreign key relationships optional * always allow creator to edit (per hagen) * use primary join * no index_doc_batch spam * try this again --------- Co-authored-by: Richard Kuo <rkuo@rkuo.com>
This commit is contained in:
parent
a50a3944b3
commit
dcbea883ae
@ -0,0 +1,30 @@
|
||||
"""add creator to cc pair
|
||||
|
||||
Revision ID: 9cf5c00f72fe
|
||||
Revises: c0fd6e4da83a
|
||||
Create Date: 2024-11-12 15:16:42.682902
|
||||
|
||||
"""
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "9cf5c00f72fe"
|
||||
down_revision = "26b931506ecb"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column(
|
||||
"connector_credential_pair",
|
||||
sa.Column(
|
||||
"creator_id",
|
||||
sa.UUID(as_uuid=True),
|
||||
nullable=True,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("connector_credential_pair", "creator_id")
|
@ -76,6 +76,7 @@ def _add_user_filters(
|
||||
.where(~UG__CCpair.user_group_id.in_(user_groups))
|
||||
.correlate(ConnectorCredentialPair)
|
||||
)
|
||||
where_clause |= ConnectorCredentialPair.creator_id == user.id
|
||||
else:
|
||||
where_clause |= ConnectorCredentialPair.access_type == AccessType.PUBLIC
|
||||
where_clause |= ConnectorCredentialPair.access_type == AccessType.SYNC
|
||||
@ -388,6 +389,7 @@ def add_credential_to_connector(
|
||||
)
|
||||
|
||||
association = ConnectorCredentialPair(
|
||||
creator_id=user.id if user else None,
|
||||
connector_id=connector_id,
|
||||
credential_id=credential_id,
|
||||
name=cc_pair_name,
|
||||
|
@ -173,6 +173,11 @@ class User(SQLAlchemyBaseUserTableUUID, Base):
|
||||
)
|
||||
# Whether the user has logged in via web. False if user has only used Danswer through Slack bot
|
||||
has_web_login: Mapped[bool] = mapped_column(Boolean, default=True)
|
||||
cc_pairs: Mapped[list["ConnectorCredentialPair"]] = relationship(
|
||||
"ConnectorCredentialPair",
|
||||
back_populates="creator",
|
||||
primaryjoin="User.id == foreign(ConnectorCredentialPair.creator_id)",
|
||||
)
|
||||
|
||||
|
||||
class InputPrompt(Base):
|
||||
@ -455,6 +460,14 @@ class ConnectorCredentialPair(Base):
|
||||
"IndexAttempt", back_populates="connector_credential_pair"
|
||||
)
|
||||
|
||||
# the user id of the user that created this cc pair
|
||||
creator_id: Mapped[UUID | None] = mapped_column(nullable=True)
|
||||
creator: Mapped["User"] = relationship(
|
||||
"User",
|
||||
back_populates="cc_pairs",
|
||||
primaryjoin="foreign(ConnectorCredentialPair.creator_id) == remote(User.id)",
|
||||
)
|
||||
|
||||
|
||||
class Document(Base):
|
||||
__tablename__ = "document"
|
||||
|
@ -255,7 +255,7 @@ def index_doc_batch_prepare(
|
||||
)
|
||||
|
||||
|
||||
@log_function_time()
|
||||
@log_function_time(debug_only=True)
|
||||
def index_doc_batch(
|
||||
*,
|
||||
chunker: Chunker,
|
||||
|
@ -237,6 +237,8 @@ class CCPairFullInfo(BaseModel):
|
||||
is_editable_for_current_user: bool
|
||||
deletion_failure_message: str | None
|
||||
indexing: bool
|
||||
creator: UUID | None
|
||||
creator_email: str | None
|
||||
|
||||
@classmethod
|
||||
def from_models(
|
||||
@ -282,6 +284,10 @@ class CCPairFullInfo(BaseModel):
|
||||
is_editable_for_current_user=is_editable_for_current_user,
|
||||
deletion_failure_message=cc_pair_model.deletion_failure_message,
|
||||
indexing=indexing,
|
||||
creator=cc_pair_model.creator_id,
|
||||
creator_email=cc_pair_model.creator.email
|
||||
if cc_pair_model.creator
|
||||
else None,
|
||||
)
|
||||
|
||||
|
||||
|
@ -8,6 +8,7 @@ import requests
|
||||
from danswer.connectors.models import InputType
|
||||
from danswer.db.enums import AccessType
|
||||
from danswer.db.enums import ConnectorCredentialPairStatus
|
||||
from danswer.server.documents.models import CCPairFullInfo
|
||||
from danswer.server.documents.models import ConnectorCredentialPairIdentifier
|
||||
from danswer.server.documents.models import ConnectorIndexingStatus
|
||||
from danswer.server.documents.models import DocumentSource
|
||||
@ -146,7 +147,22 @@ class CCPairManager:
|
||||
result.raise_for_status()
|
||||
|
||||
@staticmethod
|
||||
def get_one(
|
||||
def get_single(
|
||||
cc_pair_id: int,
|
||||
user_performing_action: DATestUser | None = None,
|
||||
) -> CCPairFullInfo | None:
|
||||
response = requests.get(
|
||||
f"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair_id}",
|
||||
headers=user_performing_action.headers
|
||||
if user_performing_action
|
||||
else GENERAL_HEADERS,
|
||||
)
|
||||
response.raise_for_status()
|
||||
cc_pair_json = response.json()
|
||||
return CCPairFullInfo(**cc_pair_json)
|
||||
|
||||
@staticmethod
|
||||
def get_indexing_status_by_id(
|
||||
cc_pair_id: int,
|
||||
user_performing_action: DATestUser | None = None,
|
||||
) -> ConnectorIndexingStatus | None:
|
||||
@ -165,7 +181,7 @@ class CCPairManager:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_all(
|
||||
def get_indexing_statuses(
|
||||
user_performing_action: DATestUser | None = None,
|
||||
) -> list[ConnectorIndexingStatus]:
|
||||
response = requests.get(
|
||||
@ -183,7 +199,7 @@ class CCPairManager:
|
||||
verify_deleted: bool = False,
|
||||
user_performing_action: DATestUser | None = None,
|
||||
) -> None:
|
||||
all_cc_pairs = CCPairManager.get_all(user_performing_action)
|
||||
all_cc_pairs = CCPairManager.get_indexing_statuses(user_performing_action)
|
||||
for retrieved_cc_pair in all_cc_pairs:
|
||||
if retrieved_cc_pair.cc_pair_id == cc_pair.id:
|
||||
if verify_deleted:
|
||||
@ -233,7 +249,9 @@ class CCPairManager:
|
||||
"""after: Wait for an indexing success time after this time"""
|
||||
start = time.monotonic()
|
||||
while True:
|
||||
fetched_cc_pairs = CCPairManager.get_all(user_performing_action)
|
||||
fetched_cc_pairs = CCPairManager.get_indexing_statuses(
|
||||
user_performing_action
|
||||
)
|
||||
for fetched_cc_pair in fetched_cc_pairs:
|
||||
if fetched_cc_pair.cc_pair_id != cc_pair.id:
|
||||
continue
|
||||
@ -467,7 +485,7 @@ class CCPairManager:
|
||||
cc_pair_id is good to do."""
|
||||
start = time.monotonic()
|
||||
while True:
|
||||
cc_pairs = CCPairManager.get_all(user_performing_action)
|
||||
cc_pairs = CCPairManager.get_indexing_statuses(user_performing_action)
|
||||
if cc_pair_id:
|
||||
found = False
|
||||
for cc_pair in cc_pairs:
|
||||
|
@ -29,6 +29,25 @@ from tests.integration.common_utils.test_models import DATestUserGroup
|
||||
from tests.integration.common_utils.vespa import vespa_fixture
|
||||
|
||||
|
||||
def test_connector_creation(reset: None) -> None:
|
||||
# Creating an admin user (first user created is automatically an admin)
|
||||
admin_user: DATestUser = UserManager.create(name="admin_user")
|
||||
|
||||
# create connectors
|
||||
cc_pair_1 = CCPairManager.create_from_scratch(
|
||||
source=DocumentSource.INGESTION_API,
|
||||
user_performing_action=admin_user,
|
||||
)
|
||||
|
||||
cc_pair_info = CCPairManager.get_single(
|
||||
cc_pair_1.id, user_performing_action=admin_user
|
||||
)
|
||||
assert cc_pair_info
|
||||
assert cc_pair_info.creator
|
||||
assert str(cc_pair_info.creator) == admin_user.id
|
||||
assert cc_pair_info.creator_email == admin_user.email
|
||||
|
||||
|
||||
def test_connector_deletion(reset: None, vespa_client: vespa_fixture) -> None:
|
||||
# Creating an admin user (first user created is automatically an admin)
|
||||
admin_user: DATestUser = UserManager.create(name="admin_user")
|
||||
|
@ -139,7 +139,7 @@ def test_web_pruning(reset: None, vespa_client: vespa_fixture) -> None:
|
||||
cc_pair_1, now, timeout=60, user_performing_action=admin_user
|
||||
)
|
||||
|
||||
selected_cc_pair = CCPairManager.get_one(
|
||||
selected_cc_pair = CCPairManager.get_indexing_status_by_id(
|
||||
cc_pair_1.id, user_performing_action=admin_user
|
||||
)
|
||||
assert selected_cc_pair is not None, "cc_pair not found after indexing!"
|
||||
@ -156,7 +156,7 @@ def test_web_pruning(reset: None, vespa_client: vespa_fixture) -> None:
|
||||
cc_pair_1, now, timeout=60, user_performing_action=admin_user
|
||||
)
|
||||
|
||||
selected_cc_pair = CCPairManager.get_one(
|
||||
selected_cc_pair = CCPairManager.get_indexing_status_by_id(
|
||||
cc_pair_1.id, user_performing_action=admin_user
|
||||
)
|
||||
assert selected_cc_pair is not None, "cc_pair not found after pruning!"
|
||||
|
@ -208,6 +208,10 @@ function Main({ ccPairId }: { ccPairId: number }) {
|
||||
disabled={ccPair.status === ConnectorCredentialPairStatus.PAUSED}
|
||||
isDeleting={isDeleting}
|
||||
/>
|
||||
<div className="text-sm mt-1">
|
||||
Creator:{" "}
|
||||
<b className="text-emphasis">{ccPair.creator_email ?? "Unknown"}</b>
|
||||
</div>
|
||||
<div className="text-sm mt-1">
|
||||
Total Documents Indexed:{" "}
|
||||
<b className="text-emphasis">{ccPair.num_docs_indexed}</b>
|
||||
|
@ -6,6 +6,7 @@ import {
|
||||
ValidStatuses,
|
||||
AccessType,
|
||||
} from "@/lib/types";
|
||||
import { UUID } from "crypto";
|
||||
|
||||
export enum ConnectorCredentialPairStatus {
|
||||
ACTIVE = "ACTIVE",
|
||||
@ -27,6 +28,8 @@ export interface CCPairFullInfo {
|
||||
is_editable_for_current_user: boolean;
|
||||
deletion_failure_message: string | null;
|
||||
indexing: boolean;
|
||||
creator: UUID | null;
|
||||
creator_email: string | null;
|
||||
}
|
||||
|
||||
export interface PaginatedIndexAttempts {
|
||||
|
Loading…
x
Reference in New Issue
Block a user