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:
rkuo-danswer 2024-11-13 11:35:08 -08:00 committed by GitHub
parent a50a3944b3
commit dcbea883ae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 103 additions and 8 deletions

View File

@ -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")

View File

@ -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,

View File

@ -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"

View File

@ -255,7 +255,7 @@ def index_doc_batch_prepare(
)
@log_function_time()
@log_function_time(debug_only=True)
def index_doc_batch(
*,
chunker: Chunker,

View File

@ -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,
)

View File

@ -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:

View File

@ -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")

View File

@ -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!"

View File

@ -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>

View File

@ -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 {