diff --git a/demo.py b/demo.py index d074658..228c38c 100644 --- a/demo.py +++ b/demo.py @@ -29,8 +29,8 @@ class UserModel(BaseDBModel): class Meta: """Override the default options for the UserModel.""" - create_pk: bool = False # Disable auto-increment ID - primary_key: str = "slug" # Use 'slug' as the primary key + create_pk: bool = True # Disable auto-increment ID + # primary_key: str = "slug" # Use 'slug' as the PK # noqa: ERA001 table_name: str = "users" # Explicitly define the table name @@ -41,7 +41,7 @@ def main() -> None: level=logging.DEBUG, format="%(levelname)-8s%(message)s" ) - db = SqliterDB(memory=True, auto_commit=True, debug=True) + db = SqliterDB("poop.db", memory=False, auto_commit=True, debug=True) with db: db.create_table(UserModel) # Create the users table user1 = UserModel( diff --git a/sqliter/query/query.py b/sqliter/query/query.py index 704c30d..ff6b2b3 100644 --- a/sqliter/query/query.py +++ b/sqliter/query/query.py @@ -146,6 +146,15 @@ def fields(self, fields: Optional[list[str]] = None) -> QueryBuilder: """ if fields: self._fields = fields + # Ensure the primary key is included and placed first if create_pk + # is True + if self.model_class.should_create_pk(): + primary_key = self.model_class.get_primary_key() + if primary_key in self._fields: + self._fields.remove( + primary_key + ) # Remove if it's already there + self._fields.insert(0, primary_key) # Add it at the beginning self._validate_fields() return self diff --git a/sqliter/sqliter.py b/sqliter/sqliter.py index 7e2a392..bf274da 100644 --- a/sqliter/sqliter.py +++ b/sqliter/sqliter.py @@ -492,6 +492,14 @@ def select( Returns: A QueryBuilder instance for further query construction. """ + # Ensure the primary key is included and placed first + if model_class.should_create_pk(): + primary_key = model_class.get_primary_key() + if fields: + if primary_key in fields: + fields.remove(primary_key) + fields.insert(0, primary_key) + query_builder = QueryBuilder(self, model_class, fields) # If exclude is provided, apply the exclude method