Source code for squabble.rules.require_primary_key

import pglast
from pglast.enums import ConstrType

import squabble.rule
from squabble.message import Message
from squabble.rules import BaseRule

[docs]class RequirePrimaryKey(BaseRule): """ Require that all new tables specify a ``PRIMARY KEY`` constraint. Configuration: :: { "RequirePrimaryKey": {} } """
[docs] class MissingPrimaryKey(Message): """ When creating a new table, it's usually a good idea to define a primary key, as it can guarantee a unique, fast lookup into the table. If no single column will uniquely identify a row, creating a composite primary key is also possible. .. code-block:: sql CREATE TABLE users (email VARCHAR(128) PRIMARY KEY); -- Also valid CREATE TABLE users ( email VARCHAR(128), -- ... PRIMARY KEY(email) ); """ CODE = 1002 TEMPLATE = 'table "{tbl}" does not name a primary key'
def enable(self, ctx, _config): ctx.register('CreateStmt', self._create_table()) @squabble.rule.node_visitor def _create_table(self, ctx, node): seen_pk = False def _check_constraint(_ctx, constraint): nonlocal seen_pk if constraint.contype == ConstrType.CONSTR_PRIMARY: seen_pk = True def _check_column(_ctx, col): if col.constraints == pglast.Missing: return for c in col.constraints: _check_constraint(_ctx, c) ctx.register('ColumnDef', _check_column) ctx.register('Constraint', _check_constraint) ctx.register_exit(lambda c: self._check_pk(c, seen_pk, node)) def _check_pk(self, ctx, seen_pk, table_node): """ Make sure we've seen a primary key constraint by the time the `CREATE TABLE` statement is finished. """ if seen_pk: return # Use the table's name as the reported error's location node = table_node.relation, node=node)