diff --git a/docs/changelog.rst b/docs/changelog.rst index 0ef85ca..624808f 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -9,6 +9,7 @@ Unreleased ---------- +- ``table.transform()`` now preserves ``AUTOINCREMENT`` primary keys and their sequence high-water marks. Previously a transform removed ``AUTOINCREMENT`` and could reuse deleted row IDs. (:issue:`602`) - New ``table.checks``, ``table.column_checks`` and ``table.table_checks`` introspection properties expose column-level and table-level ``CHECK`` constraints. (:issue:`834`) - New ``sqlite_utils.ANY`` marker type for creating and introspecting SQLite ``ANY`` columns. The Python API and CLI can create, add and transform these columns, and ``table.transform()`` and ``table.extract()`` now preserve ``ANY`` columns and their values in ``STRICT`` tables. (:issue:`790`) - ``table.transform()`` now preserves ``CHECK`` constraints, including comments within their expressions. Renaming a column rewrites identifier references in checks without changing string literals or function names. Dropping a column drops a check owned by that column, and raises ``TransformError`` if a remaining check depends on it. (:issue:`762`) diff --git a/sqlite_utils/create_table_parser.py b/sqlite_utils/create_table_parser.py index 2d891ae..d426286 100644 --- a/sqlite_utils/create_table_parser.py +++ b/sqlite_utils/create_table_parser.py @@ -1,4 +1,4 @@ -"""Helpers for parsing CHECK constraints from SQLite CREATE TABLE SQL. +"""Helpers for parsing constraints from SQLite CREATE TABLE SQL. SQLite does not expose CHECK constraints through a pragma, so preserving them across a table rebuild requires reading ``sqlite_schema.sql``. This module is @@ -564,6 +564,34 @@ def parse_checks(create_sql: str) -> list[Check]: return checks +def parse_autoincrement(create_sql: str) -> str | None: + """Return the AUTOINCREMENT column from a valid CREATE TABLE statement.""" + body_info = _table_body(create_sql) + if body_info is None: + return None + body, _ = body_info + for item, _, _ in _split_spans(body, _lex(body)): + item_tokens = _meaningful(_lex(item)) + if not item_tokens: + continue + head = item_tokens[0] + if ( + head.kind == "word" and head.text.upper() in _TABLE_CONSTRAINT_KEYWORDS + ) or head.is_keyword("CONSTRAINT"): + continue + column = _unquote(head.text) + index = 1 + while index < len(item_tokens): + token = item_tokens[index] + if token.text == "(": + index = _matching_paren(item_tokens, index) + 1 + continue + if token.is_keyword("AUTOINCREMENT"): + return column + index += 1 + return None + + def parse_column_comments(create_sql: str) -> dict[str, ColumnComments]: """Return comments immediately before and after each column definition.""" body_info = _table_body(create_sql) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 9c0b402..48c5d5d 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -33,6 +33,7 @@ from .create_table_parser import ( ColumnComments, ParseError, check_references_identifier, + parse_autoincrement, parse_checks, parse_column_comments, rewrite_check_expression, @@ -1422,6 +1423,7 @@ class Database: strict: bool = False, _checks: Iterable[Check] | None = None, _column_comments: Mapping[str, ColumnComments] | None = None, + _autoincrement: str | None = None, ) -> str: """ Returns the SQL ``CREATE TABLE`` statement for creating the specified table. @@ -1525,10 +1527,22 @@ class Database: column_items.insert(0, (pk, int)) elif pk: pk = [resolve_casing(p, [c[0] for c in column_items]) for p in pk] + if _autoincrement is not None: + _autoincrement = resolve_casing( + _autoincrement, [c[0] for c in column_items] + ) + if _autoincrement != single_pk: + raise ValueError("AUTOINCREMENT requires a single-column primary key") for column_name, column_type in column_items: column_extras = [] if column_name == single_pk: column_extras.append("PRIMARY KEY") + if column_name == _autoincrement: + if COLUMN_TYPE_MAPPING[column_type] != "INTEGER": + raise ValueError( + "AUTOINCREMENT requires an INTEGER PRIMARY KEY column" + ) + column_extras.append("AUTOINCREMENT") if column_name in not_null: column_extras.append("NOT NULL") if column_name in defaults and defaults[column_name] is not None: @@ -2748,6 +2762,7 @@ class Table(Queryable): try: existing_checks = self.checks existing_column_comments = parse_column_comments(self.schema) + existing_autoincrement = parse_autoincrement(self.schema) except ParseError as ex: raise TransformError( f"Could not parse table schema for table {self.name!r}: {ex}" @@ -2870,6 +2885,11 @@ class Table(Queryable): new_column_pairs.append((new_name, type_)) copy_from_to[name] = new_name + if existing_autoincrement: + existing_autoincrement = resolve_casing( + existing_autoincrement, existing_columns + ) + if pk is DEFAULT: pks_renamed = tuple( rename.get(pk_name) or pk_name @@ -2880,6 +2900,28 @@ class Table(Queryable): else: pk = pks_renamed + create_table_autoincrement = None + if existing_autoincrement and existing_autoincrement not in drop: + renamed_autoincrement = ( + rename.get(existing_autoincrement) or existing_autoincrement + ) + single_pk = pk[0] if isinstance(pk, (list, tuple)) and len(pk) == 1 else pk + new_column_types = dict(new_column_pairs) + if ( + single_pk == renamed_autoincrement + and COLUMN_TYPE_MAPPING.get(new_column_types.get(renamed_autoincrement)) + == "INTEGER" + ): + create_table_autoincrement = renamed_autoincrement + + autoincrement_sequence = None + if create_table_autoincrement: + sequence_row = self.db.execute( + "SELECT seq FROM sqlite_sequence WHERE name = ?", [self.name] + ).fetchone() + if sequence_row is not None: + autoincrement_sequence = sequence_row[0] + # not_null may be a set or dict, need to convert to a set create_table_not_null = { rename.get(c.name) or c.name @@ -2931,6 +2973,7 @@ class Table(Queryable): strict=self.strict if strict is None else strict, _checks=create_table_checks, _column_comments=create_table_column_comments, + _autoincrement=create_table_autoincrement, ).strip() ) @@ -3053,6 +3096,23 @@ class Table(Queryable): "ON" if legacy_alter_table_was_on else "OFF" ) ) + if autoincrement_sequence is not None: + table_name_literal = self.db.quote(self.name) + sqls.extend( + ( + "UPDATE sqlite_sequence SET seq = MAX(seq, {sequence}) " + "WHERE name = {table_name};".format( + sequence=autoincrement_sequence, + table_name=table_name_literal, + ), + "INSERT INTO sqlite_sequence (name, seq) " + "SELECT {table_name}, {sequence} WHERE NOT EXISTS " + "(SELECT 1 FROM sqlite_sequence WHERE name = {table_name});".format( + sequence=autoincrement_sequence, + table_name=table_name_literal, + ), + ) + ) # Re-add existing indexes sqls.extend(index_create_sqls) return sqls diff --git a/tests/test_create_table_parser.py b/tests/test_create_table_parser.py index a7aa0c0..54bf221 100644 --- a/tests/test_create_table_parser.py +++ b/tests/test_create_table_parser.py @@ -8,6 +8,7 @@ from sqlite_utils.create_table_parser import ( Check, ColumnComments, ParseError, + parse_autoincrement, parse_checks, parse_column_comments, ) @@ -117,6 +118,36 @@ def test_virtual_table_has_no_checks(): ) +@pytest.mark.parametrize( + "sql,expected", + [ + ( + "CREATE TABLE t(id INTEGER PRIMARY KEY AUTOINCREMENT, value TEXT)", + "id", + ), + ( + 'CREATE TABLE t("quoted id" INTEGER PRIMARY KEY AUTOINCREMENT)', + "quoted id", + ), + ( + 'CREATE TABLE t("autoincrement" INTEGER PRIMARY KEY, value TEXT)', + None, + ), + ( + "CREATE TABLE t(id INTEGER PRIMARY KEY /* AUTOINCREMENT */, value TEXT)", + None, + ), + ( + "CREATE TABLE t(id INTEGER PRIMARY KEY, value TEXT CHECK(value != 'AUTOINCREMENT'))", + None, + ), + ], +) +def test_parse_autoincrement(sql, expected): + sqlite3.connect(":memory:").execute(sql) + assert parse_autoincrement(sql) == expected + + comment_or_space = st.sampled_from( [ " ", diff --git a/tests/test_transform.py b/tests/test_transform.py index 6a8a143..e6096cd 100644 --- a/tests/test_transform.py +++ b/tests/test_transform.py @@ -1053,6 +1053,24 @@ def test_transform_with_unique_constraint_implicit_index(fresh_db): ) +def test_transform_preserves_autoincrement_and_sequence(fresh_db): + fresh_db.execute( + "CREATE TABLE entries (id INTEGER PRIMARY KEY AUTOINCREMENT, value TEXT)" + ) + entries = fresh_db.table("entries") + entries.insert_all(({"value": "one"}, {"value": "two"})) + entries.delete(2) + + entries.transform(rename={"value": "label"}) + + assert "PRIMARY KEY AUTOINCREMENT" in entries.schema + entries.insert({"label": "three"}) + assert list(entries.rows) == [ + {"id": 1, "label": "one"}, + {"id": 3, "label": "three"}, + ] + + def test_transform_preserves_view(fresh_db): # https://github.com/simonw/sqlite-utils/issues/831 dogs = fresh_db.table("dogs")