adding tags to rut

Author: matusdrobuliak66

packages/models-library/src/models_library/rabbitmq_messages.py

@@ -201,7 +201,7 @@ class RabbitResourceTrackingProjectSyncMessage(RabbitMessageBase):
 
     project_id: ProjectID
     project_name: str | None = None
-    project_tags_names: list[str] | None = None
+    project_tags: list[tuple[int, str]] | None = None
     created_at: datetime.datetime = Field(
         default_factory=lambda: arrow.utcnow().datetime,
         description="message creation datetime",

packages/postgres-database/src/simcore_postgres_database/migration/versions/da145a5fe27b_introducing_rut_project_metadata.py

@@ -0,0 +1,59 @@
+"""introducing rut project metadata
+
+Revision ID: da145a5fe27b
+Revises: 8bfe65a5e294
+Create Date: 2024-11-12 16:03:41.595377+00:00
+
+"""
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "da145a5fe27b"
+down_revision = "8bfe65a5e294"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "resource_tracker_project_metadata",
+        sa.Column("project_id", sa.String(), nullable=False),
+        sa.Column("project_name", sa.String(), nullable=False),
+        sa.Column(
+            "project_tags",
+            postgresql.JSONB(astext_type=sa.Text()),
+            server_default=sa.text("'{}'::jsonb"),
+            nullable=False,
+        ),
+        sa.Column(
+            "modified",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=False,
+        ),
+        sa.PrimaryKeyConstraint("project_id"),
+    )
+    op.create_foreign_key(
+        "fk_resource_tracker_service_runs_project_metadata_project_id",
+        "resource_tracker_service_runs",
+        "resource_tracker_project_metadata",
+        ["project_id"],
+        ["project_id"],
+        onupdate="CASCADE",
+        ondelete="RESTRICT",
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(
+        "fk_resource_tracker_service_runs_project_metadata_project_id",
+        "resource_tracker_service_runs",
+        type_="foreignkey",
+    )
+    op.drop_table("resource_tracker_project_metadata")
+    # ### end Alembic commands ###

packages/postgres-database/src/simcore_postgres_database/models/resource_tracker_project_metadata.py

@@ -0,0 +1,33 @@
+""" resource_tracker_service_runs table
+"""
+
+import sqlalchemy as sa
+from sqlalchemy.dialects.postgresql import JSONB
+
+from ._common import column_modified_datetime
+from .base import metadata
+
+resource_tracker_project_metadata = sa.Table(
+    "resource_tracker_project_metadata",
+    metadata,
+    sa.Column(
+        "project_id",  # UUID
+        sa.String,
+        nullable=False,
+        primary_key=True,
+    ),
+    sa.Column(
+        "project_name",
+        sa.String,
+        nullable=False,
+        doc="we want to store the project name for tracking/billing purposes and be sure it stays there even when the project is deleted (that's also reason why we do not introduce foreign key)",
+    ),
+    sa.Column(
+        "project_tags",
+        JSONB,
+        nullable=False,
+        server_default=sa.text("'{}'::jsonb"),
+        doc="we want to store the project tags for tracking/billing purposes and be sure it stays there even when the project is deleted (that's also reason why we do not introduce foreign key)",
+    ),
+    column_modified_datetime(timezone=True),
+)

packages/postgres-database/src/simcore_postgres_database/models/resource_tracker_service_runs.py

@@ -98,6 +98,12 @@ class ResourceTrackerServiceRunStatus(str, enum.Enum):
     sa.Column(
         "project_id",  # UUID
         sa.String,
+        sa.ForeignKey(
+            "resource_tracker_project_metadata.project_id",
+            name="fk_resource_tracker_service_runs_project_metadata_project_id",
+            onupdate="CASCADE",
+            ondelete="RESTRICT",
+        ),
         nullable=False,
         doc="We want to store the project id for tracking/billing purposes and be sure it stays there even when the project is deleted (that's also reason why we do not introduce foreign key)",
     ),
@@ -236,30 +242,3 @@ class ResourceTrackerServiceRunStatus(str, enum.Enum):
         == ResourceTrackerServiceRunStatus.RUNNING
     ),
 )
-
-
-resource_tracker_project_metadata = sa.Table(
-    "resource_tracker_project_metadata",
-    metadata,
-    sa.Column(
-        "project_id",  # UUID
-        sa.String,
-        nullable=False,
-        doc="We want to store the project id for tracking/billing purposes and be sure it stays there even when the project is deleted (that's also reason why we do not introduce foreign key)",
-        primary_key=True,
-    ),
-    # sa.Column(
-    #     "project_name",
-    #     sa.String,
-    #     nullable=False,
-    #     doc="we want to store the project name for tracking/billing purposes and be sure it stays there even when the project is deleted (that's also reason why we do not introduce foreign key)",
-    # ),
-    sa.Column(
-        "project_tags_names",
-        JSONB,
-        nullable=False,
-        server_default=sa.text("'{}'::jsonb"),
-        doc="we want to store the project name for tracking/billing purposes and be sure it stays there even when the project is deleted (that's also reason why we do not introduce foreign key)",
-    ),
-    column_modified_datetime(timezone=True),
-)

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/modules/db/repositories/resource_tracker.py

@@ -9,6 +9,7 @@
 )
 from models_library.api_schemas_storage import S3BucketName
 from models_library.products import ProductName
+from models_library.projects import ProjectID
 from models_library.resource_tracker import (
     CreditClassification,
     CreditTransactionId,
@@ -43,10 +44,12 @@
 from simcore_postgres_database.models.resource_tracker_pricing_units import (
     resource_tracker_pricing_units,
 )
+from simcore_postgres_database.models.resource_tracker_project_metadata import (
+    resource_tracker_project_metadata,
+)
 from simcore_postgres_database.models.resource_tracker_service_runs import (
     resource_tracker_service_runs,
 )
-from simcore_service_resource_usage_tracker.services.utils import ProjectID
 from sqlalchemy.dialects.postgresql import ARRAY, INTEGER
 
 from .....exceptions.errors import (
@@ -204,31 +207,27 @@ async def insert_rut_project_metadata(
         self,
         project_id: ProjectID,
         project_name: str,
-        project_tags_names: list[str],
+        project_tags_db: dict[str, dict[str, str]],
     ) -> None:
         async with self.db_engine.begin() as conn:
-            insert_stmt = (
-                resource_tracker_service_runs.insert()
-                .values(
-                    project_id=f"{project_id}",
-                    project_name=project_name,
-                    project_tags_names=project_tags_names,
-                    modified=sa.func.now(),
-                )
-                .returning(resource_tracker_credit_transactions.c.transaction_id)
+            insert_stmt = resource_tracker_project_metadata.insert().values(
+                project_id=f"{project_id}",
+                project_name=project_name,
+                project_tags=project_tags_db,
+                modified=sa.func.now(),
             )
             await conn.execute(insert_stmt)
 
     async def update_rut_project_metadata(
         self,
         project_id: ProjectID,
         project_name: str | None = None,
-        project_tags_names: list[str] | None = None,
+        project_tags: list[str] | None = None,
     ) -> None:
 
         _update_data = {
             "project_name": project_name,
-            "project_tags_names": project_tags_names,
+            "project_tags": project_tags,
         }
         _update_data_clean = {k: v for k, v in _update_data.items() if v is not None}
 

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/process_message_running_service.py

@@ -95,6 +95,15 @@ async def _process_start_event(
         )
         pricing_unit_cost = pricing_unit_cost_db.cost_per_unit
 
+    project_tags_db: dict[str, dict[str, str]] = {}
+    for tag in msg.project_tags:
+        project_tags_db[f"{tag[0]}"] = {"name": tag[1]}
+    await resource_tracker_repo.insert_rut_project_metadata(
+        project_id=msg.project_id,
+        project_name=msg.project_name,
+        project_tags_db=jsonable_encoder(project_tags_db),
+    )
+
     create_service_run = ServiceRunCreate(
         product_name=msg.product_name,
         service_run_id=msg.service_run_id,
@@ -126,9 +135,6 @@ async def _process_start_event(
         last_heartbeat_at=msg.created_at,
     )
     service_run_id = await resource_tracker_repo.create_service_run(create_service_run)
-    await resource_tracker_repo.insert_rut_project_metadata(
-        project_id=msg.project_id, project_name=msg.project_name, project_tags_names=[]
-    )
 
     if msg.wallet_id and msg.wallet_name:
         transaction_create = CreditTransactionCreate(

services/resource-usage-tracker/tests/unit/with_dbs/conftest.py

@@ -14,6 +14,7 @@
 from asgi_lifespan import LifespanManager
 from faker import Faker
 from fastapi import FastAPI
+from models_library.projects import ProjectID
 from models_library.rabbitmq_messages import (
     RabbitResourceTrackingHeartbeatMessage,
     RabbitResourceTrackingMessageType,
@@ -26,6 +27,9 @@
 from simcore_postgres_database.models.resource_tracker_credit_transactions import (
     resource_tracker_credit_transactions,
 )
+from simcore_postgres_database.models.resource_tracker_project_metadata import (
+    resource_tracker_project_metadata,
+)
 from simcore_postgres_database.models.resource_tracker_service_runs import (
     resource_tracker_service_runs,
 )
@@ -51,6 +55,7 @@ def mock_env(monkeypatch: pytest.MonkeyPatch) -> EnvVarsDict:
         "SC_BOOT_MODE": "production",
         "POSTGRES_CLIENT_NAME": "postgres_test_client",
         "RESOURCE_USAGE_TRACKER_MISSED_HEARTBEAT_CHECK_ENABLED": "0",
+        "RESOURCE_USAGE_TRACKER_TRACING": "null",
     }
     setenvs_from_dict(monkeypatch, env_vars)
     return env_vars
@@ -181,6 +186,25 @@ async def assert_service_runs_db_row(
     raise ValueError
 
 
+async def assert_project_metadata_db_row(postgres_db, project_id: ProjectID) -> None:
+    async for attempt in AsyncRetrying(
+        wait=wait_fixed(0.2),
+        stop=stop_after_delay(10),
+        retry=retry_if_exception_type(AssertionError),
+        reraise=True,
+    ):
+        with attempt, postgres_db.connect() as con:
+            result = con.execute(
+                sa.select(resource_tracker_project_metadata).where(
+                    resource_tracker_project_metadata.c.project_id == f"{project_id}"
+                )
+            )
+            row = result.first()
+            assert row
+            return
+    raise ValueError
+
+
 async def assert_credit_transactions_db_row(
     postgres_db, service_run_id: str, modified_at: datetime | None = None
 ) -> CreditTransactionDB:
@@ -239,6 +263,7 @@ def _creator(**kwargs: dict[str, Any]) -> RabbitResourceTrackingStartedMessage:
             "user_email": faker.email(),
             "project_id": faker.uuid4(),
             "project_name": faker.pystr(),
+            "project_tags": [(faker.pyint(), faker.pystr())],
             "node_id": faker.uuid4(),
             "node_name": faker.pystr(),
             "parent_project_id": faker.uuid4(),

services/resource-usage-tracker/tests/unit/with_dbs/test_process_rabbitmq_message.py

@@ -17,7 +17,7 @@
     _process_stop_event,
 )
 
-from .conftest import assert_service_runs_db_row
+from .conftest import assert_project_metadata_db_row, assert_service_runs_db_row
 
 pytest_simcore_core_services_selection = ["postgres", "rabbit"]
 pytest_simcore_ops_services_selection = [
@@ -48,6 +48,7 @@ async def test_process_event_functions(
     )
     await _process_start_event(resource_tracker_repo, msg, publisher)
     output = await assert_service_runs_db_row(postgres_db, msg.service_run_id)
+    await assert_project_metadata_db_row(postgres_db, msg.project_id)
     assert output.stopped_at is None
     assert output.service_run_status == "RUNNING"
     first_occurence_of_last_heartbeat_at = output.last_heartbeat_at