patx/pickledb

closes #110

Commit ac5441c · patx · 2026-05-09T20:15:28-04:00

Changeset

ac5441c6fea2e161232678d650993808e1a918b1

Parents

0cd28935317412785ab863b45d01b60a985b7736

View source at this commit

Comments

No comments yet.

Log in to comment

Diff

diff --git a/pickledb.py b/pickledb.py
index 306808b..f1914a2 100644
--- a/pickledb.py
+++ b/pickledb.py
@@ -6,6 +6,7 @@ Licensed - BSD 3 Clause (see LICENSE)
 
 import asyncio
 import os
+import re
 from typing import Any
 import uuid
 
@@ -20,6 +21,7 @@ except ImportError:
 
 
 MISSING = object()
+_SQLITE_TABLE_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
 
 
 def in_async() -> bool:
@@ -161,6 +163,15 @@ if sqlite_enable:
             sqlite_path: str = "pickledb.sqlite3",
             table_name: str = "kv",
         ) -> None:
+            if (
+                not isinstance(table_name, str)
+                or not _SQLITE_TABLE_NAME_RE.fullmatch(table_name)
+            ):
+                raise ValueError(
+                    "Invalid table name. Use only ASCII letters, numbers, and "
+                    "underscores, and start with a letter or underscore."
+                )
+
             self.sqlite_path = sqlite_path
             self.table_name = table_name
 
@@ -387,4 +398,3 @@ else:
                 "PickleDBSQLite requires `aiosqlite`. "
                 "Install it via `pip install \"pickledb[sqlite]\"`."
             )
-
diff --git a/test_pickledb.py b/test_pickledb.py
index a3c6346..9899a3d 100644
--- a/test_pickledb.py
+++ b/test_pickledb.py
@@ -366,6 +366,26 @@ def test_sqlite_sync_set_get_and_all(tmp_path):
     kv.close()
 
 
+def test_sqlite_custom_table_name(tmp_path):
+    sqlite_path = tmp_path / "custom.sqlite3"
+    kv = PickleDBSQLite(str(sqlite_path), table_name="kv_custom_1")
+
+    assert kv.set("explicit", {"ok": True}) == "explicit"
+    assert kv.get("explicit") == {"ok": True}
+
+    kv.close()
+
+
+def test_sqlite_rejects_invalid_table_name(tmp_path):
+    sqlite_path = tmp_path / "invalid.sqlite3"
+
+    with pytest.raises(ValueError, match="Invalid table name"):
+        PickleDBSQLite(
+            str(sqlite_path),
+            table_name="kv (key TEXT); DROP TABLE kv; --",
+        )
+
+
 @pytest.mark.asyncio
 async def test_sqlite_async_set_get_and_purge(tmp_path):
     sqlite_path = tmp_path / "kv_async.sqlite3"
@@ -387,4 +407,3 @@ async def test_sqlite_async_set_get_and_purge(tmp_path):
     assert await kv.all() == []
 
     await kv.close()
-