import json
import re
from contextlib import closing
from contextlib import contextmanager
from datetime import datetime
from numbers import Integral
from numbers import Real
from typing import Any
from typing import Dict
from typing import Generator
from typing import Iterator
from typing import Optional
from typing import Union
try:
import sqlite3
except ImportError as exc:
raise ImportError(
"Python's sqlite3 module is required but is not available. "
"This usually means that the Python interpreter was built without SQLite "
"support. Install the SQLite development libraries required by your "
"operating system, then reinstall or rebuild Python and recreate the "
"virtual environment."
) from exc
from . import uri_utils
from .datetime_utils import fromisoformat
_SQL_TYPES = {"INTEGER", "REAL", "TEXT", "BLOB"}
[docs]
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 = []
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})"
[docs]
def insert_query(table: str, nfields: int) -> str:
table = _validate_identifier(table)
values = ("?," * nfields)[:-1]
return f"INSERT INTO {table} VALUES({values})"
[docs]
def python_to_sql_type(value: Any) -> str:
if isinstance(value, (Integral, bool)):
return "INTEGER"
elif isinstance(value, Real):
return "REAL"
elif isinstance(value, (str, datetime)):
return "TEXT"
else:
return "BLOB"
[docs]
def python_to_sql_types(field_types: Optional[Dict]) -> dict:
if not field_types:
return dict()
return {k: python_to_sql_type(v) for k, v in field_types.items()}
[docs]
def serialize(value: Any, sql_type: Optional[str] = None):
if value is not None and sql_type is not None:
vsql_type = python_to_sql_type(value)
if sql_type != vsql_type:
raise TypeError(f"value {value} does not have SQL type {sql_type}")
if isinstance(value, (Integral, Real, bool, str)):
return value
elif isinstance(value, datetime):
return value.isoformat()
else:
return json.dumps(value).encode()
[docs]
def deserialize(sql_value, field_type: Optional[str] = None):
if isinstance(sql_value, bytes):
sql_value = sql_value.decode()
if sql_value == "null" or sql_value is None:
return None
elif isinstance(field_type, bool):
return bool(sql_value)
elif isinstance(field_type, (Integral, Real, str)):
return sql_value
elif isinstance(field_type, datetime):
return fromisoformat(sql_value)
else:
return json.loads(sql_value)
[docs]
def select(
conn,
table: str,
field_types: Optional[Dict] = None,
sql_types: Optional[Dict] = None,
starttime: Optional[Union[str, datetime]] = None,
endtime: Optional[Union[str, datetime]] = None,
**is_equal_filter,
) -> Iterator[dict]:
table = _validate_identifier(table)
conditions = []
params = []
if is_equal_filter:
if sql_types is None:
sql_types = python_to_sql_types(field_types)
for k, v in is_equal_filter.items():
conditions.append(f"{_validate_identifier(k)} = ?")
params.append(serialize(v, sql_types.get(k)))
if starttime:
if isinstance(starttime, str):
starttime = fromisoformat(starttime)
conditions.append("time >= ?")
params.append(starttime.isoformat())
if endtime:
if isinstance(endtime, str):
endtime = fromisoformat(endtime)
conditions.append("time <= ?")
params.append(endtime.isoformat())
if conditions:
search_condition = " AND ".join(conditions)
# table/columns are validated by _validate_identifier(); values are parameterized
query = f"SELECT * FROM {table} WHERE {search_condition}" # noqa: S608
else:
query = f"SELECT * FROM {table}" # noqa: S608
with closing(conn.cursor()) as cursor:
try:
cursor.execute(query, params)
except sqlite3.OperationalError as e:
if "no such table" in str(e):
return
raise # Re-raise other errors
rows = cursor.fetchall()
conn.commit()
if cursor.description is None:
return
fields = [col[0] for col in cursor.description]
if field_types is None:
field_types = {}
for values in rows:
yield {
k: deserialize(v, field_types.get(k)) for k, v in zip(fields, values)
}
[docs]
@contextmanager
def connect(database: str, **kwargs) -> Generator[sqlite3.Connection, None, None]:
if database != ":memory:":
_ensure_directory_exists(database)
with closing(sqlite3.connect(database, **kwargs)) as conn:
yield conn
def _ensure_directory_exists(uri: str) -> None:
parsed = uri_utils.parse_uri(uri)
if parsed.scheme == "file":
path = uri_utils.path_from_uri(uri)
path.parent.mkdir(parents=True, exist_ok=True)
# Table/column names cannot be parameterized in DBAPI queries, so restrict
# them to simple identifiers before interpolating them into SQL strings.
_IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
def _validate_identifier(name: str) -> str:
if not _IDENTIFIER_RE.match(name):
raise ValueError(f"{name!r} is not a valid SQL identifier")
return name
