cleanup and slight refactor

db.py

@@ -1,15 +1,8 @@
-from sqlalchemy import create_engine, Table, Column, MetaData, Integer, String, Boolean, DateTime, ForeignKey
-from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import sessionmaker, relationship, mapped_column
 import os
+from sqlalchemy import create_engine, Column, Integer, String, DateTime, Boolean
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
 
-# Set up SQLite database
-if os.getenv("SQLITE_FILE") is not None:
-    sqlite_file = os.getenv("SQLITE_FILE")
-else:
-    raise Exception("SQLite File not defined!")
-DATABASE_URI = f'sqlite:///{sqlite_file}'
-engine = create_engine(DATABASE_URI)
 Base = declarative_base()
 
 class DBGroupInstance(Base):
@@ -19,10 +12,14 @@ class DBGroupInstance(Base):
     instance_id = Column(String, nullable=False)
     world_name = Column(String, nullable=False)
     member_count = Column(Integer, nullable=False)
-
-# Create tables if they don't exist
-Base.metadata.create_all(engine)
+    delivered = Column(Boolean, nullable=False)
 
 def get_session():
+    sqlite_file = os.getenv("SQLITE_FILE")
+    if not sqlite_file:
+        raise Exception("SQLite File not defined!")
+    database_uri = f'sqlite:///{sqlite_file}'
+    engine = create_engine(database_uri)
+    Base.metadata.create_all(engine)
     Session = sessionmaker(bind=engine)
-    return Session()
+    return Session()

main.py

@@ -1,22 +1,35 @@
-from vrcapi import getGroupInstances, fakeGroupInstances
+import os
+import logging
+from vrcapi import get_group_instances
+from testing import fake_group_instances
 from datetime import datetime
 from db import DBGroupInstance, get_session
 
+log_level = os.getenv("LOG_LEVEL")
+logging.basicConfig(level=log_level)
 session = get_session()
 
+# Commit the group instance data in the DB
+def store_instance_data(instance):
+    try:
+        db_instance = DBGroupInstance(
+            timestamp=datetime.now(),
+            instance_id=instance.instance_id,
+            world_name=instance.world.name,
+            member_count=instance.member_count,
+            delivered=False # Discord delivery, always false until separate job marks it as True
+        )
+        session.add(db_instance)
+        session.commit()
+        logging.info("Instance log added")
+    except Exception as e:
+        logging.error(f"Exception while adding instance log: {e}")
+        session.rollback()
+
 if __name__ == "__main__":
-    group_instances = getGroupInstances()
+    group_instances = get_group_instances() # Call the VRChat API and pull group instance data
     if not group_instances:
-        print("No group instances were found! Nothing to do...")
+        logging.info("No group instances were found! Nothing to do...") # Since this will run in a cron job, the program ends here
     else:
         for instance in group_instances:
-            try:
-                db_instance = DBGroupInstance(timestamp=datetime.now(),
-                                              instance_id=instance.instance_id,
-                                              world_name=instance.world.name,
-                                              member_count=instance.member_count)
-                session.add(db_instance)
-                session.commit()
-                print("Instance log added")
-            except Exception as e:
-                print(f"Exception {e}")
+            store_instance_data(instance) # Send instance data to the DB

testing.py

@@ -0,0 +1,29 @@
+from vrchatapi.models.group_instance import GroupInstance
+from vrchatapi.models.world import World
+from random import randint
+from datetime import datetime
+
+def fake_group_instances():
+    group_instances = []
+    # For testing purposes
+
+    # Worlds API object that is valid enough to not cause errors in testing
+    fakeWorld = World(
+        "usr_123TEST-789-abcd-efghi-aaa69420",
+        "McLovin",
+        90,
+        50,
+        datetime.now(),
+        "A test world, it does not exist. No VRC API's were harmed",
+        0, False, 0, "world_AAAA-TEST-49494-69420", 
+        "http://null.example/hello.jpg", [], "4/20/2020",
+        "Fake world", None, 0, "vrchat", 0, None, 0, 0, "9/11/2021", "public", ["nothing", "testing"],
+        "http://testing.localhost/thisisnotreal.png", [], datetime.now(), 5, 90, [])
+
+    # Fake instance 
+    fakeInstance1 = GroupInstance("TESTTESTgrp_69420_aaaaaaaaa", 
+                                  "test", fakeWorld, randint(1, 50))
+
+    group_instances.append(fakeInstance1)
+
+    return group_instances

totp.py

@@ -1,7 +1,7 @@
 import pyotp
 import os
 
-def getTOTP():
+def get_totp():
     totp_key = os.getenv("TOTP_KEY")
     totp = pyotp.TOTP(totp_key)
     return totp.now()

vrcapi.py

@@ -1,41 +1,13 @@
+import os
+import logging
 import vrchatapi
-from vrchatapi.api import authentication_api
+from vrchatapi.api import authentication_api, groups_api
 from vrchatapi.exceptions import UnauthorizedException, ApiException
 from vrchatapi.models.two_factor_auth_code import TwoFactorAuthCode
 from vrchatapi.models.two_factor_email_code import TwoFactorEmailCode
-from vrchatapi.models.group_instance import GroupInstance
-from vrchatapi.models.world import World
-from totp import getTOTP
-import os
-from random import randint
-from datetime import datetime
+from totp import get_totp
 
-def fakeGroupInstances():
-    group_instances = []
-    # For testing purposes
-
-    # Worlds API object that is valid enough to not cause errors in testing
-    fakeWorld = World(
-        "usr_123TEST-789-abcd-efghi-aaa69420",
-        "McLovin",
-        90,
-        50,
-        datetime.now(),
-        "A test world, it does not exist. No VRC API's were harmed",
-        0, False, 0, "world_AAAA-TEST-49494-69420", 
-        "http://null.example/hello.jpg", [], "4/20/2020",
-        "Fake world", None, 0, "vrchat", 0, None, 0, 0, "9/11/2021", "public", ["nothing", "testing"],
-        "http://testing.localhost/thisisnotreal.png", [], datetime.now(), 5, 90, [])
-
-    # Fake instance 
-    fakeInstance1 = GroupInstance("TESTTESTgrp_69420_aaaaaaaaa", 
-                                  "test", fakeWorld, randint(1, 50))
-
-    group_instances.append(fakeInstance1)
-
-    return group_instances
-
-def getGroupInstances():
+def get_group_instances():
     group_instances = []
 
     configuration = vrchatapi.Configuration( 
@@ -47,33 +19,15 @@ def getGroupInstances():
     group_id = os.getenv("GROUP_ID")
 
     with vrchatapi.ApiClient(configuration) as api_client:
-        # Set our User-Agent as per VRChat Usage Policy 
+        # Set our User-Agent as per VRChat Usage Policy and call the Auth API
         api_client.user_agent = "GroupInstanceLogger/0.1alpha me@williamtpeebles.com"
-
-        # Instantiate instances of API classes
         auth_api = authentication_api.AuthenticationApi(api_client)
 
         try:
             # Calling getCurrentUser on Authentication API logs you in if the user isn't already logged in.
             current_user = auth_api.get_current_user()
         except UnauthorizedException as e:
-            if e.status == 200:
-                if "Email 2 Factor Authentication" in e.reason:
-                    # Calling email verify2fa if the account has 2FA disabled, give warning that automatic runs cannot be done
-                    print("IMPORTANT: Your VRChat Account is only enabled for **email** based 2-factor authentication! "
-                          "You will not be able to automate runs of GroupInstanceLogger\n"
-                          "Set up your account to use TOTP-based 2FA so the application can log you in automatically\n\n")
-                    auth_api.verify2_fa_email_code(two_factor_email_code=TwoFactorEmailCode(input("Email 2FA Code: ")))
-                elif "2 Factor Authentication" in e.reason:
-                    # Check if TOTP code is defined in env variables
-                    if os.getenv("TOTP_KEY") is not None:
-                        print("Attempting to authenticate with defined TOTP key")
-                        auth_api.verify2_fa(two_factor_auth_code=TwoFactorAuthCode(getTOTP()))
-                    else:
-                        # TOTP prompt if not configured
-                        # Calling verify2fa if the account has 2FA enabled
-                        auth_api.verify2_fa(two_factor_auth_code=TwoFactorAuthCode(input("Enter 2FA Code: ")))
-                    
+                handle_auth_exception(auth_api, e)
                 current_user = auth_api.get_current_user()
 
                 # Call groups APIs
@@ -86,9 +40,35 @@ def getGroupInstances():
                     return group_instances
                 except ApiException as e:
                     print("Exception when calling GroupsApi->get_group: %s\n" % e)
-            else:
-                print("Exception when calling API: %s\n", e)
-        except vrchatapi.ApiException as e:
-            print("Exception when calling API: %s\n", e)
+        logging.info(f"Logged in as: {current_user.display_name}")
 
-        print("Logged in as:", current_user.display_name)
+        try:
+             api_response = groups_api.GroupsApi(api_client),get_group_instances(group_id)
+             group_instances.extend(api_response)
+        except ApiException as e:
+             logging.error(f"Exception when calling GroupsApi->get_group_instances: {e}")
+        except Exception as e:
+             logging.error(f"Unexpected Error: {e}")
+    return group_instances
+
+        
+
+def handle_auth_exception(auth_api, e):
+    if e.status == 200:
+        if "Email 2 Factor Authentication" in e.reason:
+                # Calling email verify2fa if the account has 2FA disabled, give warning that automatic runs cannot be done
+                logging.warning("IMPORTANT: Your VRChat Account is only enabled for **email** based 2-factor authentication! "
+                                "You will not be able to automate runs of GroupInstanceLogger\n"
+                                "Set up your account to use TOTP-based 2FA so the application can log you in automatically\n\n")
+                auth_api.verify2_fa_email_code(two_factor_email_code=TwoFactorEmailCode(input("Email 2FA Code: ")))
+        elif "2 Factor Authentication" in e.reason:
+                # Check if TOTP code is defined in env variables
+                totp_key = os.getenv("TOTP_KEY")
+                if totp_key:
+                    logging.info("Attempting to authenticate with defined TOTP key")
+                    auth_api.verify2_fa(two_factor_auth_code=TwoFactorAuthCode(get_totp()))
+                else:
+                    logging.info("No TOTP key defined, prompting manual input of 2FA code")
+                    auth_api.verify2_fa(two_factor_auth_code=TwoFactorAuthCode(input("Enter 2FA Code: ")))
+    else:
+         logging.error(f"UnauthorizedException: {e}")