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