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