diff --git a/pfbudget/core/manager.py b/pfbudget/core/manager.py index b5145d4..9877e09 100644 --- a/pfbudget/core/manager.py +++ b/pfbudget/core/manager.py @@ -28,9 +28,9 @@ from pfbudget.db.model import ( TransactionCategory, ) from pfbudget.extract.credentials import Credentials -from pfbudget.extract.extract import Extract -from pfbudget.extract.psd2 import PSD2Client +from pfbudget.extract.nordigen import NordigenClient from pfbudget.extract.parsers import parse_data +from pfbudget.extract.psd2 import PSD2Extractor dotenv.load_dotenv() @@ -85,16 +85,18 @@ class Manager: else: banks = session.get(Bank, Bank.name, params[3]) session.expunge_all() - client.start = params[0] - client.end = params[1] - transactions = client.extract(banks) + + extractor = PSD2Extractor(client) + transactions = [] + for bank in banks: + transactions.extend(extractor.extract(bank, params[0], params[1])) # dry-run if not params[2]: with self.db.session() as session: session.add(sorted(transactions)) else: - print(transactions) + print(sorted(transactions)) case Operation.Categorize: with self.db.session() as session: @@ -419,10 +421,10 @@ class Manager: self._db = url @staticmethod - def nordigen_client() -> Extract: + def nordigen_client() -> NordigenClient: credentials = Credentials( os.environ.get("SECRET_ID"), os.environ.get("SECRET_KEY"), os.environ.get("TOKEN"), ) - return PSD2Client(credentials) + return NordigenClient(credentials) diff --git a/pfbudget/extract/exceptions.py b/pfbudget/extract/exceptions.py index cb21212..0129509 100644 --- a/pfbudget/extract/exceptions.py +++ b/pfbudget/extract/exceptions.py @@ -6,5 +6,13 @@ class BankError(ExtractError): pass -class CredentialsError(ExtractError): +class PSD2ClientError(ExtractError): + pass + + +class CredentialsError(PSD2ClientError): + pass + + +class DownloadError(PSD2ClientError): pass diff --git a/pfbudget/extract/nordigen.py b/pfbudget/extract/nordigen.py new file mode 100644 index 0000000..7edb378 --- /dev/null +++ b/pfbudget/extract/nordigen.py @@ -0,0 +1,92 @@ +import json +import nordigen +import requests +import time +import uuid + +from .credentials import Credentials +from .exceptions import CredentialsError, DownloadError + + +class NordigenClient: + redirect_url = "https://murta.dev" + + def __init__(self, credentials: Credentials): + super().__init__() + + if not credentials.valid(): + raise CredentialsError + + self._client = nordigen.NordigenClient( + secret_key=credentials.key, secret_id=credentials.id, timeout=5 + ) + + if credentials.token: + self._client.token = credentials.token + + def download(self, requisition_id): + try: + requisition = self._client.requisition.get_requisition_by_id(requisition_id) + print(requisition) + except requests.HTTPError as e: + raise DownloadError(e) + + transactions = {} + for acc in requisition["accounts"]: + account = self._client.account_api(acc) + + retries = 0 + while retries < 3: + try: + downloaded = account.get_transactions() + break + except requests.ReadTimeout: + retries += 1 + print(f"Request #{retries} timed-out, retrying in 1s") + time.sleep(1) + + if not downloaded: + print(f"Couldn't download transactions for {account}") + continue + + transactions.update(downloaded) + + return transactions + + def dump(self, bank, downloaded): + with open("json/" + bank.name + ".json", "w") as f: + json.dump(downloaded, f) + + def generate_token(self): + self.token = self._client.generate_token() + print(f"New access token: {self.token}") + return self.token + + def requisition(self, id: str, country: str = "PT"): + requisition = self._client.initialize_session( + redirect_uri=self.redirect_url, + institution_id=id, + reference_id=str(uuid.uuid4()), + ) + return requisition.link, requisition.requisition_id + + def country_banks(self, country: str): + return self._client.institution.get_institutions(country) + + # def __token(self): + # if token := os.environ.get("TOKEN"): + # return token + # else: + # token = self._client.generate_token() + # print(f"New access token: {token}") + # return token["access"] + + @property + def token(self): + return self._token + + @token.setter + def token(self, value): + if self._token: + print("Replacing existing token with {value}") + self._token = value diff --git a/pfbudget/extract/psd2.py b/pfbudget/extract/psd2.py index dde9ead..6bffafc 100644 --- a/pfbudget/extract/psd2.py +++ b/pfbudget/extract/psd2.py @@ -1,142 +1,39 @@ import datetime as dt -import json -import nordigen -import requests -import time -import uuid from typing import Sequence -import pfbudget.db.model as t +from pfbudget.db.model import Bank, BankTransaction from pfbudget.utils.converters import convert -from .credentials import Credentials -from .exceptions import BankError, CredentialsError, ExtractError +from .exceptions import BankError, DownloadError, ExtractError from .extract import Extract +from .nordigen import NordigenClient -class PSD2Client(Extract): - redirect_url = "https://murta.dev" +class PSD2Extractor(Extract): + def __init__(self, client: NordigenClient): + self.__client = client - def __init__(self, credentials: Credentials): - super().__init__() + def extract( + self, bank: Bank, start=dt.date.min, end=dt.date.max + ) -> Sequence[BankTransaction]: + if not bank.nordigen: + raise BankError("Bank doesn't have Nordigen info") - if not credentials.valid(): - raise CredentialsError + try: + print(f"Downloading from {bank}...") + downloaded = self.__client.download(bank.nordigen.requisition_id) + except DownloadError as e: + print(f"There was an issue downloading from {bank.name}\n{e}") + raise ExtractError(e) - self._client = nordigen.NordigenClient( - secret_key=credentials.key, secret_id=credentials.id, timeout=5 - ) + self.__client.dump(bank, downloaded) - if credentials.token: - self._client.token = credentials.token + return [ + t + for t in self.convert(bank, downloaded, start, end) + if start <= t.date <= end + ] - self._start = dt.date.min - self._end = dt.date.max - - def extract(self, banks: Sequence[t.Bank]) -> list[t.BankTransaction]: - transactions = [] - if not banks or any(not b.nordigen for b in banks): - raise BankError - - for bank in banks: - downloaded = None - try: - print(f"Downloading from {bank}...") - downloaded = self.download(bank.nordigen.requisition_id) - except requests.HTTPError as e: - print(f"There was an issue downloading from {bank.name} -> {e}") - raise ExtractError(e) - - if downloaded: - self.dump(bank, downloaded) - - converted = [ - convert(t, bank) for t in downloaded["transactions"]["booked"] - ] - - transactions.extend( - [t for t in converted if self._start <= t.date <= self._end] - ) - - return sorted(transactions) - - def download(self, requisition_id): - requisition = self._client.requisition.get_requisition_by_id(requisition_id) - print(requisition) - - transactions = {} - for acc in requisition["accounts"]: - account = self._client.account_api(acc) - - retries = 0 - while retries < 3: - try: - downloaded = account.get_transactions() - break - except requests.ReadTimeout: - retries += 1 - print(f"Request #{retries} timed-out, retrying in 1s") - time.sleep(1) - - if not downloaded: - print(f"Couldn't download transactions for {account}") - continue - - transactions.update(downloaded) - - return transactions - - def dump(self, bank, downloaded): - with open("json/" + bank.name + ".json", "w") as f: - json.dump(downloaded, f) - - def generate_token(self): - self.token = self._client.generate_token() - print(f"New access token: {self.token}") - return self.token - - def requisition(self, id: str, country: str = "PT"): - requisition = self._client.initialize_session( - redirect_uri=self.redirect_url, - institution_id=id, - reference_id=str(uuid.uuid4()), - ) - return requisition.link, requisition.requisition_id - - def country_banks(self, country: str): - return self._client.institution.get_institutions(country) - - @property - def start(self): - return self._start - - @start.setter - def start(self, value): - self._start = value - - @property - def end(self): - return self._end - - @end.setter - def end(self, value): - self._end = value - - # def __token(self): - # if token := os.environ.get("TOKEN"): - # return token - # else: - # token = self._client.generate_token() - # print(f"New access token: {token}") - # return token["access"] - - @property - def token(self): - return self._token - - @token.setter - def token(self, value): - if self._token: - print("Replacing existing token with {value}") - self._token = value + def convert(self, bank, downloaded, start, end): + return [convert(t, bank) for t in downloaded["transactions"]["booked"]] diff --git a/tests/test_psd2.py b/tests/test_psd2.py index 80006a0..eb5c2ba 100644 --- a/tests/test_psd2.py +++ b/tests/test_psd2.py @@ -8,7 +8,8 @@ import mocks.nordigen as mock from pfbudget.db.model import Bank, BankTransaction, Nordigen from pfbudget.extract.credentials import Credentials from pfbudget.extract.exceptions import BankError, CredentialsError -from pfbudget.extract.psd2 import PSD2Client +from pfbudget.extract.nordigen import NordigenClient +from pfbudget.extract.psd2 import PSD2Extractor class MockGet: @@ -55,41 +56,37 @@ def mock_requests(monkeypatch): @pytest.fixture -def client() -> PSD2Client: +def extractor() -> NordigenClient: credentials = Credentials("ID", "KEY", "TOKEN") - return PSD2Client(credentials) + return PSD2Extractor(NordigenClient(credentials)) @pytest.fixture -def banks() -> list[Bank]: +def bank() -> list[Bank]: bank = Bank("Bank#1", "", "") bank.nordigen = Nordigen("", "", mock.id, False) - return [bank] + return bank class TestExtractPSD2: def test_empty_credentials(self): cred = Credentials("", "") with pytest.raises(CredentialsError): - PSD2Client(cred) + NordigenClient(cred) - def test_empty_banks(self, client): + def test_no_psd2_bank(self, extractor): with pytest.raises(BankError): - client.extract([]) + extractor.extract(Bank("", "", "")) - def test_no_psd2_bank(self, client): - with pytest.raises(BankError): - client.extract([Bank("", "", "")]) - - def test_timeout(self, monkeypatch, client, banks): + def test_timeout(self, monkeypatch, extractor, bank): monkeypatch.setattr( "requests.get", MockGet(mock_exception=requests.ReadTimeout) ) with pytest.raises(requests.Timeout): - client.extract(banks) + extractor.extract(bank) - def test_extract(self, client, banks): - assert client.extract(banks) == [ + def test_extract(self, extractor, bank): + assert extractor.extract(bank) == [ BankTransaction( dt.date(2023, 1, 14), "string", Decimal("328.18"), "Bank#1" ),