Browse Source

Formatting + tweak table name parser for safety

Kirk Trombley 4 years ago
parent
commit
c1be77d72b
2 changed files with 8 additions and 4 deletions
  1. 1 0
      lib/rollbot/decorators/on_event.py
  2. 7 4
      lib/rollbot/injection/data.py

+ 1 - 0
lib/rollbot/decorators/on_event.py

@@ -4,6 +4,7 @@ decorated_startup: list[StartupShutdownType] = []
 decorated_shutdown: list[StartupShutdownType] = []
 # TODO - injection? might be useful to just inject with a special event Message
 
+
 def on_startup(fn: StartupShutdownType) -> StartupShutdownType:
     decorated_startup.append(fn)
     return fn

+ 7 - 4
lib/rollbot/injection/data.py

@@ -33,7 +33,9 @@ class DataStore(Generic[DataType]):
             raise ValueError
         self.datatype = datatype
         self.connection = connection
-        self.table_name = "".join(("_" + c.lower()) if "A" <= c <= "Z" else c for c in datatype.__name__).strip("_")
+        self.table_name = "".join(
+            ("_" + c.lower()) if "A" <= c <= "Z" else c for c in datatype.__name__ if c.isalnum()
+        ).strip("_")
 
     async def setup(self):
         await self.connection.execute(
@@ -45,7 +47,9 @@ class DataStore(Generic[DataType]):
         await self.connection.commit()
 
     async def load(self, key: str) -> Optional[DataType]:
-        async with self.connection.execute(f"SELECT body FROM {self.table_name} WHERE key = ?", (key,)) as cursor:
+        async with self.connection.execute(
+            f"SELECT body FROM {self.table_name} WHERE key = ?", (key,)
+        ) as cursor:
             found = await cursor.fetchone()
         if found is None:
             return found
@@ -69,7 +73,7 @@ class DataStore(Generic[DataType]):
         await self.connection.execute(
             f"INSERT INTO {self.table_name} VALUES (:key, :body) \
                 ON CONFLICT(key) DO UPDATE SET body=:body",
-            { "key": key, "body": blob }
+            {"key": key, "body": blob},
         )
         await self.connection.commit()
 
@@ -100,4 +104,3 @@ class Data(InjectorWithCleanup[DataStore[DataType]]):
 
     async def cleanup(self, store: DataStore[DataType]):
         await store.connection.close()
-