diff --git a/main/astrodatagui/db/StarsDB.py b/main/astrodatagui/db/StarsDB.py index 172a9f4..b0edabf 100644 --- a/main/astrodatagui/db/StarsDB.py +++ b/main/astrodatagui/db/StarsDB.py @@ -1,6 +1,8 @@ from enum import Enum import sqlite3 +supportedFitTypes = ["sine", "linear", "poly"] + class StarDBError(Enum): NO_ERROR = 0 DB_CONNECTION_FAILED = 1 @@ -87,6 +89,14 @@ class StarDB(): FOREIGN KEY (mainName) REFERENCES starnames(mainName) );""") + self.dbCursor.execute("""CREATE TABLE IF NOT EXISTS starFitType + ( + mainName TEXT, + fitType TEXT, + PRIMARY KEY (mainName), + FOREIGN KEY (mainName) REFERENCES starnames(mainName) + );""") + self.insertDefaultValues() self.dbCorrection() @@ -100,6 +110,9 @@ class StarDB(): self.dbCursor.execute("""UPDATE stars SET filename = replace(filename, '\\', '/') WHERE filename LIKE '%\\%';""") + self.dbCursor.execute("""UPDATE starFitType + SET fitType = 'sine' + WHERE fitType IS NULL OR fitType = '';""") def insertDefaultValues(self): @@ -140,6 +153,15 @@ class StarDB(): self.connection.commit() + def updateFitType(self, mainName, fitType: str): + if(fitType in supportedFitTypes): + self.dbCursor.execute(f"""UPDATE starFitType + SET fitType = '{fitType}' + WHERE mainName = '{mainName}';""") + self.connection.commit() + else: + raise Exception(f"Fit type {fitType} not supported, must be one of {supportedFitTypes}") + def getAllStars(self): res = self.dbCursor.execute("""SELECT DISTINCT mainName FROM stars ORDER BY mainName""") @@ -188,3 +210,12 @@ class StarDB(): sequence = {sequence} """) return res.fetchone()[0] + + def getFitType(self, mainName): + res = self.dbCursor.execute(f"""SELECT mainName FROM starFitType + WHERE mainName = \"{mainName}\";""") + ret = res.fetchone()[0] + if(ret in supportedFitTypes): + return ret + else: + return "sine" \ No newline at end of file