From 85129d8b7af77a7e23608bb89917653e4a72c836 Mon Sep 17 00:00:00 2001
From: StNicolay <103897650+StNicolay@users.noreply.github.com>
Date: Fri, 14 Oct 2022 17:12:24 +0300
Subject: [PATCH] Changes to db: fixed error in __init__, looking for more
 precise errors

---
 src/database/add.py     | 6 +++---
 src/database/get.py     | 8 ++++----
 src/database/prepare.py | 5 ++++-
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/database/add.py b/src/database/add.py
index 54d67d9..e64ec44 100644
--- a/src/database/add.py
+++ b/src/database/add.py
@@ -1,5 +1,5 @@
 import sqlmodel
-import mariadb
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.future import Engine
 
 from . import models
@@ -21,7 +21,7 @@ def add_account(
         with sqlmodel.Session(engine) as session:
             session.add(account)
             session.commit()
-    except Exception:
+    except IntegrityError:
         return False
     else:
         return True
@@ -34,7 +34,7 @@ def add_master_pass(engine: Engine, user_id: int, salt: bytes, passwd: bytes) ->
         with sqlmodel.Session(engine) as session:
             session.add(master_pass)
             session.commit()
-    except Exception:
+    except IntegrityError:
         return False
     else:
         return True
diff --git a/src/database/get.py b/src/database/get.py
index 5772d71..000c4e2 100644
--- a/src/database/get.py
+++ b/src/database/get.py
@@ -22,7 +22,7 @@ def get_accounts(engine: Engine, user_id: int) -> list[str]:
     statement = sqlmodel.select(models.Account).where(models.Account.user_id == user_id)
     with sqlmodel.Session(engine) as session:
         result = session.exec(statement)
-    return [account.name for account in result]
+        return [account.name for account in result]
 
 
 def get_account_info(
@@ -34,6 +34,6 @@ def get_account_info(
     )
     with sqlmodel.Session(engine) as session:
         result = session.exec(statement).first()
-    if result is None:
-        return
-    return (result.salt, result.enc_login, result.enc_pass)
+        if result is None:
+            return
+        return (result.salt, result.enc_login, result.enc_pass)
diff --git a/src/database/prepare.py b/src/database/prepare.py
index b171aca..29630fa 100644
--- a/src/database/prepare.py
+++ b/src/database/prepare.py
@@ -13,4 +13,7 @@ def get_engine(host: str, user: str, passwd: str, db: str) -> Engine:
 
 
 def prepare(engine: Engine) -> None:
-    sqlmodel.SQLModel.metadata.create_all(engine, [models.Account, models.MasterPass])
+    sqlmodel.SQLModel.metadata.create_all(
+        engine,
+        # [models.Account, models.MasterPass]
+    )