1from sqlalchemy.exc import SQLAlchemyError
 2from sqlalchemy.orm import Session
 3
 4from .base import Base, SessionLocal, engine  # noqa
 5
 6from .crud import (
 7    create_admin,
 8    create_notification_reminder,  # noqa
 9    create_user,
10    delete_notification_reminder,
11    get_admin,
12    get_admins,
13    get_jwt_secret_key,
14    get_notification_reminder,
15    assure_node_inbounds,
16    get_system_usage,
17    get_tls_certificate,
18    get_user,
19    get_user_by_id,
20    get_users,
21    get_users_count,
22    remove_admin,
23    remove_user,
24    revoke_user_sub,
25    set_owner,
26    update_admin,
27    update_user,
28    update_user_status,
29    update_user_sub,
30    start_user_expire,
31)
32
33from .models import JWT, System, User  # noqa
34
35
36class GetDB:  # Context Manager
37    def __init__(self):
38        self.db = SessionLocal()
39
40    def __enter__(self):
41        return self.db
42
43    def __exit__(self, _, exc_value, traceback):
44        if isinstance(exc_value, SQLAlchemyError):
45            self.db.rollback()  # rollback on exception
46
47        self.db.close()
48
49
50__all__ = [
51    "get_user",
52    "get_user_by_id",
53    "get_users",
54    "get_users_count",
55    "create_user",
56    "remove_user",
57    "update_user",
58    "update_user_status",
59    "start_user_expire",
60    "update_user_sub",
61    "revoke_user_sub",
62    "set_owner",
63    "get_system_usage",
64    "get_jwt_secret_key",
65    "get_tls_certificate",
66    "get_admin",
67    "create_admin",
68    "update_admin",
69    "remove_admin",
70    "get_admins",
71    "create_notification_reminder",
72    "get_notification_reminder",
73    "delete_notification_reminder",
74    "GetDB",
75    "User",
76    "System",
77    "JWT",
78    "Base",
79    "Session",
80]