135 lines
4.5 KiB
Python
135 lines
4.5 KiB
Python
from typing import Any
|
|
|
|
from psycopg.types.json import Jsonb
|
|
|
|
from config.pgDb import pg_pool
|
|
|
|
|
|
def list_cube_report_states(user_id: str) -> dict[str, dict[str, Any]]:
|
|
with pg_pool.getConn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT dify_conversation_id, current_cube, limit_source,
|
|
response_message, updated_at
|
|
FROM ai_cube_report_states
|
|
WHERE created_by = %s
|
|
""",
|
|
(user_id,),
|
|
)
|
|
return {
|
|
str(row[0]): {
|
|
"query": row[1],
|
|
"limitSource": row[2],
|
|
"message": row[3],
|
|
"updatedAt": row[4].isoformat() if row[4] else None,
|
|
}
|
|
for row in cur.fetchall()
|
|
}
|
|
|
|
|
|
def get_cube_report_state(
|
|
conversation_id: str, user_id: str
|
|
) -> dict[str, Any] | None:
|
|
with pg_pool.getConn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT title, tenant_id, tenant_name, allow_global,
|
|
current_cube, limit_source, response_message,
|
|
created_at, updated_at
|
|
FROM ai_cube_report_states
|
|
WHERE dify_conversation_id = %s AND created_by = %s
|
|
""",
|
|
(conversation_id, user_id),
|
|
)
|
|
row = cur.fetchone()
|
|
if not row:
|
|
return None
|
|
return {
|
|
"conversationId": conversation_id,
|
|
"title": row[0],
|
|
"tenantId": str(row[1]) if row[1] is not None else None,
|
|
"tenantName": row[2],
|
|
"allowGlobal": row[3],
|
|
"query": row[4],
|
|
"limitSource": row[5],
|
|
"message": row[6],
|
|
"createdAt": row[7].isoformat() if row[7] else None,
|
|
"updatedAt": row[8].isoformat() if row[8] else None,
|
|
}
|
|
|
|
|
|
def upsert_cube_report_state(
|
|
*,
|
|
conversation_id: str,
|
|
user_id: str,
|
|
title: str,
|
|
tenant_id: str | None,
|
|
tenant_name: str,
|
|
allow_global: bool,
|
|
query: dict[str, Any],
|
|
limit_source: str,
|
|
response_message: str,
|
|
) -> None:
|
|
normalized_source = "user" if limit_source == "user" else "system"
|
|
with pg_pool.getConn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO ai_cube_report_states (
|
|
dify_conversation_id, created_by, title, tenant_id,
|
|
tenant_name, allow_global, current_cube, limit_source,
|
|
response_message, created_at, updated_at
|
|
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, NOW(), NOW())
|
|
ON CONFLICT (dify_conversation_id) DO UPDATE SET
|
|
title = EXCLUDED.title,
|
|
tenant_id = EXCLUDED.tenant_id,
|
|
tenant_name = EXCLUDED.tenant_name,
|
|
allow_global = EXCLUDED.allow_global,
|
|
current_cube = EXCLUDED.current_cube,
|
|
limit_source = EXCLUDED.limit_source,
|
|
response_message = EXCLUDED.response_message,
|
|
updated_at = NOW()
|
|
WHERE ai_cube_report_states.created_by = EXCLUDED.created_by
|
|
""",
|
|
(
|
|
conversation_id,
|
|
user_id,
|
|
title[:200] or "新数据对话",
|
|
tenant_id,
|
|
tenant_name[:200],
|
|
allow_global,
|
|
Jsonb(query),
|
|
normalized_source,
|
|
response_message,
|
|
),
|
|
)
|
|
|
|
|
|
def rename_cube_report_state(
|
|
conversation_id: str, user_id: str, title: str
|
|
) -> None:
|
|
with pg_pool.getConn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
UPDATE ai_cube_report_states
|
|
SET title = %s, updated_at = NOW()
|
|
WHERE dify_conversation_id = %s AND created_by = %s
|
|
""",
|
|
(title[:200], conversation_id, user_id),
|
|
)
|
|
|
|
|
|
def delete_cube_report_state(conversation_id: str, user_id: str) -> None:
|
|
with pg_pool.getConn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
DELETE FROM ai_cube_report_states
|
|
WHERE dify_conversation_id = %s AND created_by = %s
|
|
""",
|
|
(conversation_id, user_id),
|
|
)
|