diff --git a/ee/api/chalicelib/core/assist_records.py b/ee/api/chalicelib/core/assist_records.py index c5a7dd0dd..7683f5645 100644 --- a/ee/api/chalicelib/core/assist_records.py +++ b/ee/api/chalicelib/core/assist_records.py @@ -5,23 +5,17 @@ from chalicelib.utils import s3, pg_client from chalicelib.utils.TimeUTC import TimeUTC -def presign_records(project_id, data: schemas_ee.AssistRecordUploadPayloadSchema, context: schemas_ee.CurrentContext): - results = [] - params = {"user_id": context.user_id, "project_id": project_id} +def presign_records(project_id, data: schemas_ee.AssistRecordPayloadSchema, context: schemas_ee.CurrentContext): + params = {"user_id": context.user_id, "project_id": project_id, **data.dict()} - for i, r in enumerate(data.records): - key = f"{TimeUTC.now() + i}-{r.name}" - results.append(s3.get_presigned_url_for_upload(bucket=config('ASSIST_RECORDS_BUCKET'), - expires_in=1800, - key=s3.generate_file_key(project_id=project_id, key=key))) - params[f"name_{i}"] = r.name - params[f"duration_{i}"] = r.duration - params[f"session_id_{i}"] = r.session_id - params[f"key_{i}"] = key + key = f"{TimeUTC.now()}-{data.name}" + presigned_url = s3.get_presigned_url_for_upload(bucket=config('ASSIST_RECORDS_BUCKET'), expires_in=1800, + key=s3.generate_file_key(project_id=project_id, key=key)) + params["key"] = key with pg_client.PostgresClient() as cur: - values = [f"(%(project_id)s, %(user_id)s, %(name_{i})s, %(key_{i})s, %(duration_{i})s, %(session_id_{i})s)" - for i in range(len(data.records))] query = cur.mogrify(f"""INSERT INTO assist_records(project_id, user_id, name, file_key, duration, session_id) - VALUES {",".join(values)}""", params) + VALUES (%(project_id)s, %(user_id)s, %(name)s, %(key)s, + %(duration)s, %(session_id)s);""", + params) cur.execute(query) - return results + return presigned_url diff --git a/ee/api/routers/ee.py b/ee/api/routers/ee.py index 915c016df..e04795ae0 100644 --- a/ee/api/routers/ee.py +++ b/ee/api/routers/ee.py @@ -1,6 +1,6 @@ from typing import Union -from chalicelib.core import roles, traces, projects, sourcemaps, assist_records +from chalicelib.core import roles, traces, projects, sourcemaps, assist_records, sessions from chalicelib.core import unlock from chalicelib.utils import assist_helper @@ -76,6 +76,8 @@ def get_available_trail_actions(context: schemas_ee.CurrentContext = Depends(OR_ @app.put('/{projectId}/assist/save/', tags=["assist"]) @app.put('/{projectId}/assist/save', tags=["assist"]) -def sign_record_for_upload(projectId: int, data: schemas_ee.AssistRecordUploadPayloadSchema = Body(...), +def sign_record_for_upload(projectId: int, data: schemas_ee.AssistRecordPayloadSchema = Body(...), context: schemas_ee.CurrentContext = Depends(OR_context)): - return {"data": assist_records.presign_records(project_id=projectId, data=data, context=context)} + if not sessions.session_exists(project_id=projectId, session_id=data.session_id): + return {"errors": ["Session not found"]} + return {"data": {"URL": assist_records.presign_records(project_id=projectId, data=data, context=context)}} diff --git a/ee/api/schemas_ee.py b/ee/api/schemas_ee.py index 620074e61..193fdc2a5 100644 --- a/ee/api/schemas_ee.py +++ b/ee/api/schemas_ee.py @@ -83,14 +83,10 @@ class SessionModel(BaseModel): metadata: dict = Field(default={}) -class AssistRecord(BaseModel): +class AssistRecordPayloadSchema(BaseModel): name: str = Field(...) duration: int = Field(...) session_id: int = Field(...) class Config: alias_generator = schemas.attribute_to_camel_case - - -class AssistRecordUploadPayloadSchema(BaseModel): - records: List[AssistRecord] = Field(default=[]) diff --git a/ee/scripts/schema/db/init_dbs/postgresql/1.9.5/1.9.5.sql b/ee/scripts/schema/db/init_dbs/postgresql/1.9.5/1.9.5.sql index fde3e381f..3d9a9e4d0 100644 --- a/ee/scripts/schema/db/init_dbs/postgresql/1.9.5/1.9.5.sql +++ b/ee/scripts/schema/db/init_dbs/postgresql/1.9.5/1.9.5.sql @@ -10,7 +10,7 @@ CREATE TABLE IF NOT EXISTS assist_records record_id integer generated BY DEFAULT AS IDENTITY PRIMARY KEY, project_id integer NOT NULL REFERENCES projects (project_id) ON DELETE CASCADE, user_id integer NOT NULL REFERENCES users (user_id) ON DELETE SET NULL, - session_id integer NOT NULL REFERENCES sessions (session_id) ON DELETE SET NULL, + session_id bigint NOT NULL REFERENCES sessions (session_id) ON DELETE SET NULL, created_at timestamp without time zone NOT NULL default (now() at time zone 'utc'), deleted_at timestamp without time zone NULL DEFAULT NULL, name text NOT NULL, diff --git a/ee/scripts/schema/db/init_dbs/postgresql/init_schema.sql b/ee/scripts/schema/db/init_dbs/postgresql/init_schema.sql index 0a1b238b6..c473e8fe5 100644 --- a/ee/scripts/schema/db/init_dbs/postgresql/init_schema.sql +++ b/ee/scripts/schema/db/init_dbs/postgresql/init_schema.sql @@ -1251,7 +1251,7 @@ $$ record_id integer generated BY DEFAULT AS IDENTITY PRIMARY KEY, project_id integer NOT NULL REFERENCES projects (project_id) ON DELETE CASCADE, user_id integer NOT NULL REFERENCES users (user_id) ON DELETE SET NULL, - session_id integer NOT NULL REFERENCES sessions (session_id) ON DELETE SET NULL, + session_id bigint NOT NULL REFERENCES sessions (session_id) ON DELETE SET NULL, created_at timestamp without time zone NOT NULL default (now() at time zone 'utc'), deleted_at timestamp without time zone NULL DEFAULT NULL, name text NOT NULL,