Skip to content

Commit

Permalink
feature: add support for schema check when running migrations
Browse files Browse the repository at this point in the history
  • Loading branch information
eyadmba committed Jan 22, 2022
1 parent 7039236 commit 43ac06b
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 6 deletions.
29 changes: 26 additions & 3 deletions timescale/db/backends/postgis/schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,8 @@


class TimescaleSchemaEditor(PostGISSchemaEditor):
sql_is_hypertable = 'SELECT * FROM timescaledb_information.hypertables WHERE hypertable_name = {table}'
sql_is_hypertable = '''SELECT * FROM timescaledb_information.hypertables
WHERE hypertable_name = {table}{extra_condition}'''

sql_assert_is_hypertable = (
'DO $do$ BEGIN '
Expand Down Expand Up @@ -38,14 +39,20 @@ class TimescaleSchemaEditor(PostGISSchemaEditor):

sql_set_chunk_time_interval = 'SELECT set_chunk_time_interval({table}, interval {interval})'

sql_hypertable_is_in_schema = '''hypertable_schema = {schema_name}'''

def _assert_is_hypertable(self, model):
"""
Assert if the table is a hyper table
"""
table = self.quote_value(model._meta.db_table)
error_message = self.quote_value("assert failed - " + table + " should be a hyper table")

sql = self.sql_assert_is_hypertable.format(table=table, error_message=error_message)
extra_condition = self._get_extra_condition()

sql = self.sql_assert_is_hypertable.format(table=table, error_message=error_message,
extra_condition=extra_condition)

self.execute(sql)

def _assert_is_not_hypertable(self, model):
Expand All @@ -55,7 +62,11 @@ def _assert_is_not_hypertable(self, model):
table = self.quote_value(model._meta.db_table)
error_message = self.quote_value("assert failed - " + table + " should not be a hyper table")

sql = self.sql_assert_is_not_hypertable.format(table=table, error_message=error_message)
extra_condition = self._get_extra_condition()

sql = self.sql_assert_is_not_hypertable.format(table=table, error_message=error_message,
extra_condition=extra_condition)

self.execute(sql)

def _drop_primary_key(self, model):
Expand Down Expand Up @@ -139,3 +150,15 @@ def alter_field(self, model, old_field, new_field, strict=False):
and old_field.interval != new_field.interval:
# change chunk-size
self._set_chunk_time_interval(model, new_field)

def _get_extra_condition(self):
extra_condition = ''

try:
if self.connection.schema_name:
schema_name = self.quote_value(self.connection.schema_name)
extra_condition = ' AND ' + self.sql_hypertable_is_in_schema.format(schema_name=schema_name)
except:
pass

return extra_condition
29 changes: 26 additions & 3 deletions timescale/db/backends/postgresql/schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@


class TimescaleSchemaEditor(DatabaseSchemaEditor):
sql_is_hypertable = 'SELECT * FROM timescaledb_information.hypertables WHERE hypertable_name = {table}'
sql_is_hypertable = '''SELECT * FROM timescaledb_information.hypertables
WHERE hypertable_name = {table}{extra_condition}'''

sql_assert_is_hypertable = (
'DO $do$ BEGIN '
Expand Down Expand Up @@ -39,14 +40,20 @@ class TimescaleSchemaEditor(DatabaseSchemaEditor):

sql_set_chunk_time_interval = 'SELECT set_chunk_time_interval({table}, interval {interval})'

sql_hypertable_is_in_schema = '''hypertable_schema = {schema_name}'''

def _assert_is_hypertable(self, model):
"""
Assert if the table is a hyper table
"""
table = self.quote_value(model._meta.db_table)
error_message = self.quote_value("assert failed - " + table + " should be a hyper table")

sql = self.sql_assert_is_hypertable.format(table=table, error_message=error_message)
extra_condition = self._get_extra_condition()

sql = self.sql_assert_is_hypertable.format(table=table, error_message=error_message,
extra_condition=extra_condition)

self.execute(sql)

def _assert_is_not_hypertable(self, model):
Expand All @@ -56,7 +63,11 @@ def _assert_is_not_hypertable(self, model):
table = self.quote_value(model._meta.db_table)
error_message = self.quote_value("assert failed - " + table + " should not be a hyper table")

sql = self.sql_assert_is_not_hypertable.format(table=table, error_message=error_message)
extra_condition = self._get_extra_condition()

sql = self.sql_assert_is_not_hypertable.format(table=table, error_message=error_message,
extra_condition=extra_condition)

self.execute(sql)

def _drop_primary_key(self, model):
Expand Down Expand Up @@ -140,3 +151,15 @@ def alter_field(self, model, old_field, new_field, strict=False):
and old_field.interval != new_field.interval:
# change chunk-size
self._set_chunk_time_interval(model, new_field)

def _get_extra_condition(self):
extra_condition = ''

try:
if self.connection.schema_name:
schema_name = self.quote_value(self.connection.schema_name)
extra_condition = ' AND ' + self.sql_hypertable_is_in_schema.format(schema_name=schema_name)
except:
pass

return extra_condition

0 comments on commit 43ac06b

Please sign in to comment.