Handle naive datetime in database with custom database field
This commit is contained in:
@@ -9,17 +9,18 @@ from peewee import (
|
||||
AutoField,
|
||||
CharField,
|
||||
DatabaseProxy,
|
||||
Model,
|
||||
DateTimeField,
|
||||
DoesNotExist,
|
||||
fn,
|
||||
ForeignKeyField,
|
||||
IntegerField,
|
||||
Model,
|
||||
SqliteDatabase,
|
||||
PeeweeException as DatabaseError,
|
||||
)
|
||||
from playhouse.pool import PooledMySQLDatabase
|
||||
|
||||
|
||||
# Re-export PeeweeException as DatabseException, DoesNotExist and
|
||||
# EncodingError for convenience
|
||||
__all__ = [
|
||||
@@ -64,6 +65,18 @@ MIGRATION_COLUMN_MAPS = {
|
||||
}
|
||||
|
||||
|
||||
class DateTimeFieldUTC(DateTimeField):
|
||||
def db_value(self, value):
|
||||
if value:
|
||||
value = datetime_naive_utc(value)
|
||||
return super().db_value(value)
|
||||
|
||||
def python_value(self, value):
|
||||
value = super().python_value(value)
|
||||
if value:
|
||||
return datetime_aware_utc(value)
|
||||
|
||||
|
||||
class BaseModel(Model):
|
||||
"""Base model with database binding."""
|
||||
|
||||
@@ -78,19 +91,11 @@ class User(BaseModel):
|
||||
username = CharField(max_length=64, unique=True)
|
||||
password_hash = CharField(max_length=128)
|
||||
email = CharField(max_length=255)
|
||||
created_at = DateTimeField(default=now_utc)
|
||||
created_at = DateTimeFieldUTC(default=now_utc)
|
||||
|
||||
class Meta:
|
||||
table_name = "users"
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.created_at = datetime_aware_utc(self.created_at)
|
||||
|
||||
def save(self, *args, **kwargs):
|
||||
self.created_at = datetime_naive_utc(self.created_at)
|
||||
return super().save(*args, **kwargs)
|
||||
|
||||
|
||||
class Hostname(BaseModel):
|
||||
"""Hostname model for DNS records."""
|
||||
@@ -102,9 +107,9 @@ class Hostname(BaseModel):
|
||||
dns_ttl = IntegerField()
|
||||
expiry_ttl = IntegerField()
|
||||
last_ipv4 = CharField(max_length=15, null=True)
|
||||
last_ipv4_update = DateTimeField(null=True)
|
||||
last_ipv4_update = DateTimeFieldUTC(null=True)
|
||||
last_ipv6 = CharField(max_length=45, null=True)
|
||||
last_ipv6_update = DateTimeField(null=True)
|
||||
last_ipv6_update = DateTimeFieldUTC(null=True)
|
||||
|
||||
class Meta:
|
||||
table_name = "hostnames"
|
||||
@@ -112,21 +117,6 @@ class Hostname(BaseModel):
|
||||
(('hostname', 'zone'), True),
|
||||
)
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.last_ipv4_update = datetime_aware_utc(self.last_ipv4_update)
|
||||
self.last_ipv6_update = datetime_aware_utc(self.last_ipv6_update)
|
||||
|
||||
def save(self, *args, **kwargs):
|
||||
"""Validate and encode hostname/zone before saving."""
|
||||
if self.hostname:
|
||||
self.hostname = encode_dnsname(self.hostname)
|
||||
if self.zone:
|
||||
self.zone = encode_dnsname(self.zone)
|
||||
self.last_ipv4_update = datetime_naive_utc(self.last_ipv4_update)
|
||||
self.last_ipv6_update = datetime_naive_utc(self.last_ipv6_update)
|
||||
return super().save(*args, **kwargs)
|
||||
|
||||
|
||||
class Version(BaseModel):
|
||||
"""Database schema version for migrations."""
|
||||
|
||||
Reference in New Issue
Block a user