26
26
from collections import defaultdict
27
27
from collections .abc import Generator , Iterable , Iterator , Mapping , Sequence
28
28
from sqlite3 import Connection
29
- from typing import TYPE_CHECKING , Any , AnyStr , Callable , Generic
29
+ from typing import TYPE_CHECKING , Any , AnyStr , Callable , Generic , NamedTuple
30
30
31
31
from typing_extensions import TypeVar # default value support
32
32
from unidecode import unidecode
@@ -287,7 +287,7 @@ class Model(ABC, Generic[D]):
287
287
terms.
288
288
"""
289
289
290
- _indices : Sequence [types . Index ] = ()
290
+ _indices : Sequence [Index ] = ()
291
291
"""A sequence of `Index` objects that describe the indices to be
292
292
created for this table.
293
293
"""
@@ -1036,10 +1036,9 @@ def __init__(self, path, timeout: float = 5.0):
1036
1036
1037
1037
# Set up database schema.
1038
1038
for model_cls in self ._models :
1039
- self ._make_table (
1040
- model_cls ._table , model_cls ._fields , model_cls ._indices
1041
- )
1039
+ self ._make_table (model_cls ._table , model_cls ._fields )
1042
1040
self ._make_attribute_table (model_cls ._flex_table )
1041
+ self ._migrate_indices (model_cls ._table , model_cls ._indices )
1043
1042
1044
1043
# Primitive access control: connections and transactions.
1045
1044
@@ -1164,28 +1163,20 @@ def _make_table(
1164
1163
self ,
1165
1164
table : str ,
1166
1165
fields : Mapping [str , types .Type ],
1167
- indices : Sequence [types .Index ],
1168
1166
):
1169
1167
"""Set up the schema of the database. `fields` is a mapping
1170
1168
from field names to `Type`s. Columns are added if necessary.
1171
1169
"""
1172
- # Get current schema and existing indexes
1170
+ # Get current schema.
1173
1171
with self .transaction () as tx :
1174
1172
rows = tx .query ("PRAGMA table_info(%s)" % table )
1175
- current_fields = {row [1 ] for row in rows }
1176
- index_rows = tx .query (f"PRAGMA index_list({ table } )" )
1177
- current_indices = {row [1 ] for row in index_rows }
1173
+ current_fields = {row [1 ] for row in rows }
1178
1174
1179
- # Skip table creation if the current schema matches the
1180
- # requested schema (and no indexes are missing).
1181
1175
field_names = set (fields .keys ())
1182
- index_names = {index .name for index in indices }
1183
- if current_fields .issuperset (
1184
- field_names
1185
- ) and current_indices .issuperset (index_names ):
1176
+ if current_fields .issuperset (field_names ):
1177
+ # Table exists and has all the required columns.
1186
1178
return
1187
1179
1188
- # Table schema handling
1189
1180
if not current_fields :
1190
1181
# No table exists.
1191
1182
columns = []
@@ -1208,17 +1199,6 @@ def _make_table(
1208
1199
with self .transaction () as tx :
1209
1200
tx .script (setup_sql )
1210
1201
1211
- # Index handling
1212
- with self .transaction () as tx :
1213
- for index in indices :
1214
- if index .name in current_indices :
1215
- continue
1216
-
1217
- columns_str = ", " .join (index .columns )
1218
- tx .script (
1219
- f"CREATE INDEX { index .name } ON { table } ({ columns_str } )"
1220
- )
1221
-
1222
1202
def _make_attribute_table (self , flex_table : str ):
1223
1203
"""Create a table and associated index for flexible attributes
1224
1204
for the given entity (if they don't exist).
@@ -1237,6 +1217,33 @@ def _make_attribute_table(self, flex_table: str):
1237
1217
""" .format (flex_table )
1238
1218
)
1239
1219
1220
+ def _migrate_indices (
1221
+ self ,
1222
+ table : str ,
1223
+ indices : Sequence [Index ],
1224
+ ):
1225
+ """Create or replace indices for the given table.
1226
+
1227
+ If the indices already exists and are up to date (i.e., the
1228
+ index name and columns match), nothing is done. Otherwise, the
1229
+ indices are created or replaced.
1230
+ """
1231
+ with self .transaction () as tx :
1232
+ index_rows = tx .query (f"PRAGMA index_list({ table } )" )
1233
+ current_indices = {Index .from_db (tx , row [1 ]) for row in index_rows }
1234
+
1235
+ _indices = set (indices )
1236
+
1237
+ if current_indices .issuperset (_indices ):
1238
+ return
1239
+
1240
+ # May also include missing indices.
1241
+ changed_indices = _indices - current_indices
1242
+
1243
+ with self .transaction () as tx :
1244
+ for index in changed_indices :
1245
+ index .recreate (tx , table )
1246
+
1240
1247
# Querying.
1241
1248
1242
1249
def _fetch (
@@ -1306,3 +1313,42 @@ def _get(
1306
1313
exist.
1307
1314
"""
1308
1315
return self ._fetch (model_cls , MatchQuery ("id" , id )).get ()
1316
+
1317
+
1318
+ class Index (NamedTuple ):
1319
+ """A helper class to represent the index
1320
+ information in the database schema.
1321
+ """
1322
+
1323
+ name : str
1324
+ columns : Sequence [str ]
1325
+
1326
+ def recreate (self , tx : Transaction , table : str ) -> None :
1327
+ """Recreate the index in the database.
1328
+
1329
+ This is useful when the index has been changed and needs to be
1330
+ updated.
1331
+ """
1332
+ tx .script (f"DROP INDEX IF EXISTS { self .name } " )
1333
+ self .create (tx , table )
1334
+
1335
+ def create (self , tx : Transaction , table : str ) -> None :
1336
+ """Create the index in the database."""
1337
+ return tx .script (
1338
+ f"CREATE INDEX { self .name } ON { table } ({ ', ' .join (self .columns )} )"
1339
+ )
1340
+
1341
+ @classmethod
1342
+ def from_db (cls , tx : Transaction , name : str ) -> Index :
1343
+ """Create an Index object from the database if it exists.
1344
+
1345
+ The name has to exists in the database! Otherwise, an
1346
+ Error will be raised.
1347
+ """
1348
+ rows = tx .query (f"PRAGMA index_info({ name } )" )
1349
+ columns = [row [2 ] for row in rows ]
1350
+ return cls (name , columns )
1351
+
1352
+ def __hash__ (self ) -> int :
1353
+ """Unique hash for the index based on its name and columns."""
1354
+ return hash ((self .name , tuple (self .columns )))
0 commit comments