From 935e556cf6568d671431c9d1fafe2110e718e602 Mon Sep 17 00:00:00 2001 From: Daniel Hug <38571110+danielpatrickhug@users.noreply.github.com> Date: Thu, 5 Jan 2023 15:32:47 -0500 Subject: [PATCH] Add message tree state model (#414) * inital commit, created file for MessageTreeState table * added initial implementation of MessageTreeState SQl model and added it to__init__ * Ran Alembic revision for migration * reran precommit on * removed create_data and deleted flag from model * ran migration revision to remove unwant class variables date_created and deleted * removed unused imports * ran pre-commit * Updated States definiton to be Enums instead fo Constants * ran pre-commit formatting * Fixed Enum class * ran pre-commit --- ...161e384f83_added_messagetreestate_table.py | 46 +++++++++++++++++++ ..._removed_date_created_and_deleted_flag_.py | 42 +++++++++++++++++ backend/oasst_backend/models/__init__.py | 2 + .../models/message_tree_state.py | 44 ++++++++++++++++++ 4 files changed, 134 insertions(+) create mode 100644 backend/alembic/versions/2023_01_05_1144-d4161e384f83_added_messagetreestate_table.py create mode 100644 backend/alembic/versions/2023_01_05_1346-3b0adfadbef9_removed_date_created_and_deleted_flag_.py create mode 100644 backend/oasst_backend/models/message_tree_state.py diff --git a/backend/alembic/versions/2023_01_05_1144-d4161e384f83_added_messagetreestate_table.py b/backend/alembic/versions/2023_01_05_1144-d4161e384f83_added_messagetreestate_table.py new file mode 100644 index 00000000..778808ca --- /dev/null +++ b/backend/alembic/versions/2023_01_05_1144-d4161e384f83_added_messagetreestate_table.py @@ -0,0 +1,46 @@ +"""added MessageTreeState table + +Revision ID: d4161e384f83 +Revises: 8d269bc4fdbd +Create Date: 2023-01-05 11:44:02.630633 + +""" +import sqlalchemy as sa +import sqlmodel +from alembic import op +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "d4161e384f83" +down_revision = "8d269bc4fdbd" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "message_tree_state", + sa.Column("id", postgresql.UUID(as_uuid=True), server_default=sa.text("gen_random_uuid()"), nullable=False), + sa.Column("created_date", sa.DateTime(), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False), + sa.Column("deleted", sa.Boolean(), server_default=sa.text("false"), nullable=False), + sa.Column("message_tree_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column("state", sqlmodel.sql.sqltypes.AutoString(length=128), nullable=False), + sa.Column("goal_tree_size", sa.Integer(), nullable=False), + sa.Column("current_num_non_filtered_messages", sa.Integer(), nullable=False), + sa.Column("max_depth", sa.Integer(), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_message_tree_state_message_tree_id"), "message_tree_state", ["message_tree_id"], unique=False + ) + op.create_index("ix_message_tree_state_tree_id", "message_tree_state", ["message_tree_id"], unique=True) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index("ix_message_tree_state_tree_id", table_name="message_tree_state") + op.drop_index(op.f("ix_message_tree_state_message_tree_id"), table_name="message_tree_state") + op.drop_table("message_tree_state") + # ### end Alembic commands ### diff --git a/backend/alembic/versions/2023_01_05_1346-3b0adfadbef9_removed_date_created_and_deleted_flag_.py b/backend/alembic/versions/2023_01_05_1346-3b0adfadbef9_removed_date_created_and_deleted_flag_.py new file mode 100644 index 00000000..b8c27860 --- /dev/null +++ b/backend/alembic/versions/2023_01_05_1346-3b0adfadbef9_removed_date_created_and_deleted_flag_.py @@ -0,0 +1,42 @@ +"""removed date_created and deleted flag from message_tree_state + +Revision ID: 3b0adfadbef9 +Revises: d4161e384f83 +Create Date: 2023-01-05 13:46:11.338655 + +""" +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "3b0adfadbef9" +down_revision = "d4161e384f83" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("message_tree_state", "deleted") + op.drop_column("message_tree_state", "created_date") + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "message_tree_state", + sa.Column( + "created_date", + postgresql.TIMESTAMP(), + server_default=sa.text("CURRENT_TIMESTAMP"), + autoincrement=False, + nullable=False, + ), + ) + op.add_column( + "message_tree_state", + sa.Column("deleted", sa.BOOLEAN(), server_default=sa.text("false"), autoincrement=False, nullable=False), + ) + # ### end Alembic commands ### diff --git a/backend/oasst_backend/models/__init__.py b/backend/oasst_backend/models/__init__.py index 5818dbef..a856b155 100644 --- a/backend/oasst_backend/models/__init__.py +++ b/backend/oasst_backend/models/__init__.py @@ -2,6 +2,7 @@ from .api_client import ApiClient from .journal import Journal, JournalIntegration from .message import Message from .message_reaction import MessageReaction +from .message_tree_state import MessageTreeState from .task import Task from .text_labels import TextLabels from .user import User @@ -13,6 +14,7 @@ __all__ = [ "UserStats", "Message", "MessageReaction", + "MessageTreeState", "Task", "TextLabels", "Journal", diff --git a/backend/oasst_backend/models/message_tree_state.py b/backend/oasst_backend/models/message_tree_state.py new file mode 100644 index 00000000..386595e9 --- /dev/null +++ b/backend/oasst_backend/models/message_tree_state.py @@ -0,0 +1,44 @@ +from enum import Enum +from typing import Optional +from uuid import UUID, uuid4 + +import sqlalchemy as sa +import sqlalchemy.dialects.postgresql as pg +from sqlmodel import Field, Index, SQLModel + +# The types of States a message tree can have. + + +class States(Enum): + INITIAL = "initial" + BREEDING_PHASE = "breeding_phase" + RANKING_PHASE = "ranking_phase" + READY_FOR_SCORING = "ready_for_scoring" + CHILDREN_SCORED = "children_scored" + FINAL = "final" + + +VALID_STATES = ( + States.INITIAL, + States.BREEDING_PHASE, + States.RANKING_PHASE, + States.READY_FOR_SCORING, + States.CHILDREN_SCORED, + States.FINAL, +) + + +class MessageTreeState(SQLModel, table=True): + __tablename__ = "message_tree_state" + __table_args__ = (Index("ix_message_tree_state_tree_id", "message_tree_id", unique=True),) + + id: Optional[UUID] = Field( + sa_column=sa.Column( + pg.UUID(as_uuid=True), primary_key=True, default=uuid4, server_default=sa.text("gen_random_uuid()") + ), + ) + message_tree_id: UUID = Field(nullable=False, index=True) + state: str = Field(nullable=False, max_length=128) + goal_tree_size: int = Field(nullable=False) + current_num_non_filtered_messages: int = Field(nullable=False) + max_depth: int = Field(nullable=False)