ETL Token Limit (#294)

* chore: update submodules

* perf(alembic): add PAT action log

* fix(alembic): add timezone to created_at

* chore: update submodules

* Added token limit to whitelist

* fix: org token_limit defaults

* chore: update submodules

* chore: update submodules

---------

Co-authored-by: Lina <lina.lumburovska@kern.ai>

Author: andhreljaKern

alembic/versions/455cb5890ac1_adds_etl_token_limit_tables.py

@@ -0,0 +1,100 @@
+"""adds etl token limit tables
+
+Revision ID: 455cb5890ac1
+Revises: ab45734be591
+Create Date: 2025-03-10 18:01:29.434112
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "455cb5890ac1"
+down_revision = "ab45734be591"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "personal_access_token_activity_log_etl",
+        sa.Column("id", postgresql.UUID(as_uuid=True), nullable=False),
+        sa.Column("created_at", sa.DateTime(timezone=True), nullable=True),
+        sa.Column("action", sa.String(), nullable=True),
+        sa.Column("quantity", sa.Integer(), nullable=True),
+        sa.Column("endpoint", sa.String(), nullable=True),
+        sa.Column("organization_id", postgresql.UUID(as_uuid=True), nullable=True),
+        sa.Column("token_scope_id", postgresql.UUID(as_uuid=True), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["organization_id"], ["organization.id"], ondelete="CASCADE"
+        ),
+        sa.ForeignKeyConstraint(
+            ["token_scope_id"],
+            ["cognition.personal_access_token_scope_etl.id"],
+            ondelete="SET NULL",
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        schema="cognition",
+    )
+    op.create_index(
+        "idx_personal_access_token_activity_log_etl_created_at",
+        "personal_access_token_activity_log_etl",
+        ["created_at"],
+        unique=False,
+        schema="cognition",
+        postgresql_using="brin",
+    )
+    op.create_index(
+        op.f("ix_cognition_personal_access_token_activity_log_etl_action"),
+        "personal_access_token_activity_log_etl",
+        ["action"],
+        unique=False,
+        schema="cognition",
+    )
+    op.create_index(
+        op.f("ix_cognition_personal_access_token_activity_log_etl_organization_id"),
+        "personal_access_token_activity_log_etl",
+        ["organization_id"],
+        unique=False,
+        schema="cognition",
+    )
+    op.create_index(
+        op.f("ix_cognition_personal_access_token_activity_log_etl_token_scope_id"),
+        "personal_access_token_activity_log_etl",
+        ["token_scope_id"],
+        unique=False,
+        schema="cognition",
+    )
+    op.add_column("organization", sa.Column("token_limit", sa.JSON(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("organization", "token_limit")
+    op.drop_index(
+        op.f("ix_cognition_personal_access_token_activity_log_etl_token_scope_id"),
+        table_name="personal_access_token_activity_log_etl",
+        schema="cognition",
+    )
+    op.drop_index(
+        op.f("ix_cognition_personal_access_token_activity_log_etl_organization_id"),
+        table_name="personal_access_token_activity_log_etl",
+        schema="cognition",
+    )
+    op.drop_index(
+        op.f("ix_cognition_personal_access_token_activity_log_etl_action"),
+        table_name="personal_access_token_activity_log_etl",
+        schema="cognition",
+    )
+    op.drop_index(
+        "idx_personal_access_token_activity_log_etl_created_at",
+        table_name="personal_access_token_activity_log_etl",
+        schema="cognition",
+        postgresql_using="brin",
+    )
+    op.drop_table("personal_access_token_activity_log_etl", schema="cognition")
+    # ### end Alembic commands ###

fast_api/routes/organization.py

@@ -60,6 +60,7 @@
     "log_admin_requests",
     "conversation_lifespan_days",
     "file_lifespan_days",
+    "token_limit",
 }