diff --git a/timescale/db/backends/postgresql/schema.py b/timescale/db/backends/postgresql/schema.py index 368165c..5e56596 100644 --- a/timescale/db/backends/postgresql/schema.py +++ b/timescale/db/backends/postgresql/schema.py @@ -2,7 +2,9 @@ from django.db.backends.postgresql.schema import DatabaseSchemaEditor from timescale.db.models.fields import TimescaleDateTimeField +import re +reg = re.compile(r"[A-Z]") class TimescaleSchemaEditor(DatabaseSchemaEditor): sql_is_hypertable = '''SELECT * FROM timescaledb_information.hypertables @@ -70,6 +72,11 @@ def _assert_is_not_hypertable(self, model): self.execute(sql) + def quote_upper_value(self, name): + if reg.findall(name): + return '\'"%s"\'' % name + return self.quote_value(name) + def _drop_primary_key(self, model): """ Hypertables can't partition if the primary key is not @@ -78,12 +85,22 @@ def _drop_primary_key(self, model): """ db_table = model._meta.db_table table = self.quote_name(db_table) - pkey_length = self.connection.ops.max_name_length() - pkey = self.quote_name(f'{db_table[:pkey_length - 5]}_pkey') - - sql = self.sql_drop_primary_key.format(table=table, pkey=pkey) - - self.execute(sql) + ## cdchen-20250507: Get primary key from information_schema.table_constraints + # + # pkey_length = self.connection.ops.max_name_length() + # pkey = self.quote_name(f'{db_table[:pkey_length - 5]}_pkey') + # + # sql = self.sql_drop_primary_key.format(table=table, pkey=pkey) + # + # self.execute(sql) + with self.connection.cursor() as cursor: + sql = "SELECT constraint_name FROM information_schema.table_constraints WHERE constraint_type = 'PRIMARY' AND table_name = '{table}' ".format(table=table) + cursor.execute(sql) + row = cursor.fetchone() + if row: + pkey = row[0] + sql = self.sql_drop_primary_key.format(table=table, pkey=pkey) + cursor.execute(sql) def _create_hypertable(self, model, field, should_migrate=False): """ @@ -92,12 +109,15 @@ def _create_hypertable(self, model, field, should_migrate=False): # assert that the table is not already a hypertable self._assert_is_not_hypertable(model) - # drop primary key of the table - self._drop_primary_key(model) + # drop primary key of the table if not retain + if not field.retain_primary_key: + self._drop_primary_key(model) partition_column = self.quote_value(field.column) interval = self.quote_value(field.interval) - table = self.quote_value(model._meta.db_table) + # cdchen-2020427: Fix for upper case table name + # table = self.quote_value(model._meta.db_table) + table = self.quote_upper_value(model._meta.db_table) migrate = "true" if should_migrate else "false" if should_migrate and getattr(settings, "TIMESCALE_MIGRATE_HYPERTABLE_WITH_FRESH_TABLE", False): diff --git a/timescale/db/models/fields.py b/timescale/db/models/fields.py index 2b35ab2..0023bc9 100644 --- a/timescale/db/models/fields.py +++ b/timescale/db/models/fields.py @@ -2,12 +2,15 @@ class TimescaleDateTimeField(DateTimeField): - def __init__(self, *args, interval, **kwargs): + def __init__(self, *args, interval, retain_primary_key=False, **kwargs): self.interval = interval + self.retain_primary_key = retain_primary_key super().__init__(*args, **kwargs) def deconstruct(self): name, path, args, kwargs = super().deconstruct() - kwargs['interval'] = self.interval - + kwargs.update({ + 'interval': self.interval, + 'retain_primary_key': self.retain_primary_key, + }) return name, path, args, kwargs \ No newline at end of file