diff --git a/src/ewoksutils/sqlite3_utils.py b/src/ewoksutils/sqlite3_utils.py index ecf0d5d..cef0aaf 100644 --- a/src/ewoksutils/sqlite3_utils.py +++ b/src/ewoksutils/sqlite3_utils.py @@ -16,17 +16,30 @@ from . import uri_utils from .datetime_utils import fromisoformat +_SQL_TYPES = {"INTEGER", "REAL", "TEXT", "BLOB"} + def ensure_table_query(table: str, field_sql_types: Dict[str, str]) -> str: + table = _validate_identifier(table) + s = f"CREATE TABLE IF NOT EXISTS {table}" if not field_sql_types: return s - lst = [f"{k} {v}" for k, v in field_sql_types.items()] + + lst = [] + for field, sql_type in field_sql_types.items(): + field = _validate_identifier(field) + sql_type = sql_type.upper() + if sql_type not in _SQL_TYPES: + raise ValueError(f"{sql_type!r} is not a supported SQL type") + lst.append(f"{field} {sql_type}") + columns = ", ".join(lst) return f"{s} ({columns})" -def insert_query(table: str, nfields: int): +def insert_query(table: str, nfields: int) -> str: + table = _validate_identifier(table) values = ("?," * nfields)[:-1] return f"INSERT INTO {table} VALUES({values})" diff --git a/src/ewoksutils/tests/test_sqlite3_utils.py b/src/ewoksutils/tests/test_sqlite3_utils.py index 5d748a2..a8a8d78 100644 --- a/src/ewoksutils/tests/test_sqlite3_utils.py +++ b/src/ewoksutils/tests/test_sqlite3_utils.py @@ -185,3 +185,28 @@ def test_sqlite3_select_rejects_invalid_filter_key(): with sqlite3_utils.connect(":memory:") as conn: with pytest.raises(ValueError): list(sqlite3_utils.select(conn, "test", **{"bad name; --": "value"})) + + +def test_ensure_table_query_rejects_invalid_table(): + with pytest.raises(ValueError): + sqlite3_utils.ensure_table_query("foo; DROP TABLE bar; --", {"id": "INTEGER"}) + + +def test_ensure_table_query_rejects_invalid_column(): + with pytest.raises(ValueError): + sqlite3_utils.ensure_table_query("foo", {"id; DROP TABLE bar; --": "INTEGER"}) + + +def test_ensure_table_query_rejects_invalid_sql_type(): + with pytest.raises(ValueError): + sqlite3_utils.ensure_table_query("foo", {"id": "INTEGER); DROP TABLE bar; --"}) + + +def test_ensure_table_query_accepts_supported_types(): + assert ( + sqlite3_utils.ensure_table_query( + "foo", + {"id": "INTEGER", "name": "TEXT"}, + ) + == "CREATE TABLE IF NOT EXISTS foo (id INTEGER, name TEXT)" + )