diff --git a/ee/api/chalicelib/core/traces.py b/ee/api/chalicelib/core/traces.py index fd0ae6c2b..d77b0f580 100644 --- a/ee/api/chalicelib/core/traces.py +++ b/ee/api/chalicelib/core/traces.py @@ -9,7 +9,8 @@ from pydantic import BaseModel, Field from starlette.background import BackgroundTask import app as main_app -from chalicelib.utils import pg_client +import schemas_ee +from chalicelib.utils import pg_client, helper from chalicelib.utils.TimeUTC import TimeUTC from schemas import CurrentContext @@ -151,6 +152,29 @@ async def process_traces_queue(): await write_traces_batch(traces) +def get_all(tenant_id, data: schemas_ee.TrailSearchPayloadSchema): + with pg_client.PostgresClient() as cur: + cur.execute( + cur.mogrify( + """SELECT COUNT(*) AS count, + COALESCE(JSONB_AGG(full_traces) + FILTER (WHERE rn > %(p_start)s AND rn <= %(p_end)s), '[]'::JSONB) AS sessions + FROM (SELECT *, ROW_NUMBER() OVER (ORDER BY created_at) AS rn + FROM traces + WHERE tenant_id=%(tenant_id)s + AND created_at>=%(startDate)s + AND created_at<=%(endDate)s + ORDER BY created_at) AS full_traces;""", + {"tenant_id": tenant_id, + "startDate": data.startDate, + "endDate": data.endDate, + "p_start": (data.page - 1) * data.limit, + "p_end": data.page * data.limit}) + ) + rows = cur.fetchall() + return helper.list_to_camel_case(rows) + + cron_jobs = [ {"func": process_traces_queue, "trigger": "interval", "seconds": config("traces_period", cast=int, default=60), "misfire_grace_time": 20} diff --git a/ee/api/routers/ee.py b/ee/api/routers/ee.py index 1a9589eaa..f63d0dd3a 100644 --- a/ee/api/routers/ee.py +++ b/ee/api/routers/ee.py @@ -1,6 +1,7 @@ -from chalicelib.core import roles +from chalicelib.core import roles, traces from chalicelib.core import unlock from chalicelib.utils import assist_helper +from chalicelib.utils.TimeUTC import TimeUTC unlock.check() @@ -58,3 +59,11 @@ def delete_role(roleId: int, context: schemas.CurrentContext = Depends(OR_contex @app.get('/assist/credentials', tags=["assist"]) def get_assist_credentials(): return {"data": assist_helper.get_full_config()} + + +@app.post('/trails', tags=["traces", "trails"]) +def get_trails(data: schemas_ee.TrailSearchPayloadSchema = Body(...), + context: schemas.CurrentContext = Depends(OR_context)): + return { + 'data': traces.get_all(tenant_id=context.tenant_id, data=data) + } diff --git a/ee/api/schemas_ee.py b/ee/api/schemas_ee.py index 59a58f94b..06ae8f2ba 100644 --- a/ee/api/schemas_ee.py +++ b/ee/api/schemas_ee.py @@ -3,6 +3,7 @@ from typing import Optional, List from pydantic import BaseModel, Field import schemas +from chalicelib.utils.TimeUTC import TimeUTC class RolePayloadSchema(BaseModel): @@ -22,3 +23,12 @@ class CreateMemberSchema(schemas.CreateMemberSchema): class EditMemberSchema(schemas.EditMemberSchema): roleId: int = Field(...) + + +class TrailSearchPayloadSchema(schemas._PaginatedSchema): + startDate: int = Field(default=TimeUTC.now(-7)) + endDate: int = Field(default=TimeUTC.now(1)) + user_id: Optional[int] = Field(default=None) + + class Config: + alias_generator = schemas.attribute_to_camel_case diff --git a/ee/scripts/helm/db/init_dbs/postgresql/1.6.1/1.6.1.sql b/ee/scripts/helm/db/init_dbs/postgresql/1.6.1/1.6.1.sql index e94ccc4e1..00d871cac 100644 --- a/ee/scripts/helm/db/init_dbs/postgresql/1.6.1/1.6.1.sql +++ b/ee/scripts/helm/db/init_dbs/postgresql/1.6.1/1.6.1.sql @@ -9,4 +9,6 @@ $$ LANGUAGE sql IMMUTABLE; ALTER TABLE IF EXISTS dashboards ADD COLUMN IF NOT EXISTS description text NOT NULL DEFAULT ''; + +CREATE INDEX IF NOT EXISTS traces_created_at_idx ON traces (created_at); COMMIT; \ No newline at end of file diff --git a/ee/scripts/helm/db/init_dbs/postgresql/init_schema.sql b/ee/scripts/helm/db/init_dbs/postgresql/init_schema.sql index 7d6bdece7..d78a99c27 100644 --- a/ee/scripts/helm/db/init_dbs/postgresql/init_schema.sql +++ b/ee/scripts/helm/db/init_dbs/postgresql/init_schema.sql @@ -785,6 +785,7 @@ $$ ); CREATE INDEX IF NOT EXISTS traces_user_id_idx ON traces (user_id); CREATE INDEX IF NOT EXISTS traces_tenant_id_idx ON traces (tenant_id); + CREATE INDEX IF NOT EXISTS traces_created_at_idx ON traces (created_at); CREATE TYPE metric_type AS ENUM ('timeseries','table', 'predefined'); CREATE TYPE metric_view_type AS ENUM ('lineChart','progress','table','pieChart','areaChart','barChart','stackedBarChart','stackedBarLineChart','overview','map');