fix(jobs): use on_conflict_do_update to avoid unnecessary select query

Author: M03ED

app/db/base.py

@@ -28,16 +28,6 @@
 
 SessionLocal = async_sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
-# Determine dialect once at startup based on connection URL
-if SQLALCHEMY_DATABASE_URL.startswith("sqlite"):
-    DATABASE_DIALECT = "sqlite"
-elif SQLALCHEMY_DATABASE_URL.startswith("postgresql"):
-    DATABASE_DIALECT = "postgresql"
-elif SQLALCHEMY_DATABASE_URL.startswith("mysql"):
-    DATABASE_DIALECT = "mysql"
-else:
-    raise ValueError("Unsupported database URL")
-
 
 class Base(DeclarativeBase, MappedAsDataclass, AsyncAttrs):
     pass

app/db/crud/bulk.py

@@ -5,7 +5,6 @@
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from app.db.base import DATABASE_DIALECT
 from app.db.models import (
     Admin,
     Group,
@@ -266,7 +265,7 @@ async def update_users_expire(db: AsyncSession, bulk_model: BulkUser) -> tuple[l
         await db.execute(select(func.count(User.id)).where(and_(final_filter, User.expire.isnot(None))))
     ).scalar_one_or_none() or 0
     # Get database-specific datetime addition expression
-    new_expire = get_datetime_add_expression(User.expire, bulk_model.amount)
+    new_expire = get_datetime_add_expression(db, User.expire, bulk_model.amount)
     current_time = dt.now(tz.utc)
 
     # First, get the users that will have status changes BEFORE updating
@@ -377,8 +376,10 @@ async def update_users_proxy_settings(
     if not users_to_update:
         return [], count_effctive_users
 
+    dialect = db.bind.dialect.name
+
     # Prepare the update statement
-    if DATABASE_DIALECT == "postgresql":
+    if dialect == "postgresql":
         proxy_settings_expr = cast(User.proxy_settings, JSONB)
         if bulk_model.flow is not None:
             proxy_settings_expr = func.jsonb_set(

app/db/crud/general.py

@@ -1,7 +1,6 @@
 from sqlalchemy import String, func, or_, select, text
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from app.db.base import DATABASE_DIALECT
 from app.db.models import JWT, System
 from app.models.stats import Period
 
@@ -19,39 +18,43 @@
 }
 
 
-def _build_trunc_expression(period: Period, column):
-    """Builds the appropriate truncation SQL expression based on DATABASE_DIALECT and period."""
-    if DATABASE_DIALECT == "postgresql":
+def _build_trunc_expression(db: AsyncSession, period: Period, column):
+    dialect = db.bind.dialect.name
+
+    """Builds the appropriate truncation SQL expression based on dialect and period."""
+    if dialect == "postgresql":
         return func.date_trunc(period.value, column)
-    elif DATABASE_DIALECT == "mysql":
+    elif dialect == "mysql":
         return func.date_format(column, MYSQL_FORMATS[period.value])
-    elif DATABASE_DIALECT == "sqlite":
+    elif dialect == "sqlite":
         return func.strftime(SQLITE_FORMATS[period.value], column)
 
-    raise ValueError(f"Unsupported dialect: {DATABASE_DIALECT}")
+    raise ValueError(f"Unsupported dialect: {dialect}")
 
 
-def get_datetime_add_expression(datetime_column, seconds: int):
+def get_datetime_add_expression(db: AsyncSession, datetime_column, seconds: int):
     """
     Get database-specific datetime addition expression
     """
-    if DATABASE_DIALECT == "mysql":
+    dialect = db.bind.dialect.name
+    if dialect == "mysql":
         return func.date_add(datetime_column, text("INTERVAL :seconds SECOND").bindparams(seconds=seconds))
-    elif DATABASE_DIALECT == "postgresql":
+    elif dialect == "postgresql":
         return datetime_column + func.make_interval(0, 0, 0, 0, 0, 0, seconds)
-    elif DATABASE_DIALECT == "sqlite":
+    elif dialect == "sqlite":
         return func.datetime(func.strftime("%s", datetime_column) + seconds, "unixepoch")
 
-    raise ValueError(f"Unsupported dialect: {DATABASE_DIALECT}")
+    raise ValueError(f"Unsupported dialect: {dialect}")
 
 
-def json_extract(column, path: str):
+def json_extract(db: AsyncSession, column, path: str):
     """
     Args:
         column: The JSON column in your model
         path: JSON path (e.g., '$.theme')
     """
-    match DATABASE_DIALECT:
+    dialect = db.bind.dialect.name
+    match dialect:
         case "postgresql":
             keys = path.replace("$.", "").split(".")
             expr = column
@@ -64,14 +67,14 @@ def json_extract(column, path: str):
             return func.json_extract(column, path).cast(String)
 
 
-def build_json_proxy_settings_search_condition(column, value: str):
+def build_json_proxy_settings_search_condition(db: AsyncSession, column, value: str):
     """
     Builds a condition to search JSON column for UUIDs or passwords.
     Supports PostgresSQL, MySQL, SQLite.
     """
     return or_(
         *[
-            json_extract(column, field) == value
+            json_extract(db, column, field) == value
             for field in ("$.vmess.id", "$.vless.id", "$.trojan.password", "$.shadowsocks.password")
         ],
     )

app/db/crud/node.py

@@ -99,7 +99,7 @@ async def get_nodes_usage(
     Returns:
         NodeUsageStatsList: A NodeUsageStatsList contain list of NodeUsageResponse objects containing usage data.
     """
-    trunc_expr = _build_trunc_expression(period, NodeUsage.created_at)
+    trunc_expr = _build_trunc_expression(db, period, NodeUsage.created_at)
 
     conditions = [NodeUsage.created_at >= start, NodeUsage.created_at <= end]
 
@@ -147,7 +147,7 @@ async def get_nodes_usage(
 async def get_node_stats(
     db: AsyncSession, node_id: int, start: datetime, end: datetime, period: Period
 ) -> NodeStatsList:
-    trunc_expr = _build_trunc_expression(period, NodeStat.created_at)
+    trunc_expr = _build_trunc_expression(db, period, NodeStat.created_at)
     conditions = [NodeStat.created_at >= start, NodeStat.created_at <= end, NodeStat.node_id == node_id]
 
     stmt = (

app/db/crud/user.py

@@ -157,7 +157,7 @@ async def get_users(
     if group_ids:
         filters.append(User.groups.any(Group.id.in_(group_ids)))
     if proxy_id:
-        filters.append(build_json_proxy_settings_search_condition(User.proxy_settings, proxy_id))
+        filters.append(build_json_proxy_settings_search_condition(db, User.proxy_settings, proxy_id))
 
     if filters:
         stmt = stmt.where(and_(*filters))
@@ -337,7 +337,7 @@ async def get_user_usages(
     """
 
     # Build the appropriate truncation expression
-    trunc_expr = _build_trunc_expression(period, NodeUserUsage.created_at)
+    trunc_expr = _build_trunc_expression(db, period, NodeUserUsage.created_at)
 
     conditions = [
         NodeUserUsage.created_at >= start,
@@ -862,7 +862,7 @@ async def get_all_users_usages(
     admin_users = {user.id for user in await get_users(db=db, admins=admin)}
 
     # Build the appropriate truncation expression
-    trunc_expr = _build_trunc_expression(period, NodeUserUsage.created_at)
+    trunc_expr = _build_trunc_expression(db, period, NodeUserUsage.created_at)
 
     conditions = [
         NodeUserUsage.created_at >= start,

app/jobs/cleanup_subscription_updates.py

@@ -2,7 +2,6 @@
 
 from app import scheduler
 from app.db import GetDB
-from app.db.base import DATABASE_DIALECT
 from app.db.models import UserSubscriptionUpdate
 from app.utils.logger import get_logger
 from config import USER_SUBSCRIPTION_CLIENTS_LIMIT, JOB_CLEANUP_SUBSCRIPTION_UPDATES_INTERVAL
@@ -26,8 +25,10 @@ async def cleanup_user_subscription_updates():
             logger.info("No users with excess subscription updates")
             return
 
+        dialect = db.bind.dialect.name
+
         # Second query: Use different approaches based on database type
-        if DATABASE_DIALECT == "mysql":
+        if dialect == "mysql":
             # MySQL/MariaDB: Use correlated subquery without LIMIT
             total_deleted = 0
             for user_id in user_ids: