Updates to hypertable creation, compression, retention

Author: jmitchel3

README.md

@@ -10,14 +10,101 @@ Looking for Django? [Check out django-timescaledb](https://github.com/jamessewel
 pip install timescaledb
 ```
 
-## Usage 
+## Quickstart
 
+The timescaledb python package provides helpers for creating hypertables, configuring compression, retention policies, and more.
 
-### Create a TimescaleDB Model
+## Two ways to create a TimescaleDB Model
 
-The `TimescaleModel` class inherits from `SQLModel` but incldues a `time` field for TimescaleDB hypertables and an `id` field for primary keys.
+- Automatically via `TimescaleModel`
+- Manually via `create_hypertable` on any table with a `time` column
+
+Let's take a look at the manual way first.
+
+
+### Manually Create a Hypertable
 
-`models.py`
+```python
+from sqlmodel import create_engine, Field, SQLModel
+import timescaledb
+
+TIMESCALE_DATABASE_URL = "postgresql://user:password@localhost:5432/timescaledb"
+engine = create_engine(TIMESCALE_DATABASE_URL)
+
+class Sensor(SQLModel, table=True):
+    id: int = Field(default=None, primary_key=True)
+    time: datetime = Field(default=None, primary_key=True)
+    sensor_id: int = Field(index=True)
+    value: float
+
+    __tablename__ = "my_time_series_table"
+
+
+hypertable_options = {
+    "time_column": "time",
+    "compress_orderby": "time DESC",
+    "compress_segmentby": "sensor_id",
+    "chunk_time_interval": "7 days",
+    "drop_after": "1 year",
+    "migrate_data": True,
+    "if_not_exists": True,
+}
+
+# Create the table and the hypertable
+with Session(engine) as session:
+    # Create the table in the database
+    SQLModel.metadata.create_all(engine)
+    # Create the hypertable
+    table_name="my_time_series_table",
+    timescaledb.create_hypertable(session, commit=True, table_name=table_name, hypertable_options=hypertable_options)
+    # Add compression policy
+    timescaledb.add_compression_policy(session, commit=True, table_name=table_name, interval="7 days")
+    # Add retention policy
+    timescaledb.add_retention_policy(session, commit=True, table_name=table_name, drop_after=hypertable_options.get('drop_after'))
+```
+
+
+### Automatically via `TimescaleModel`
+
+
+```python
+from sqlmodel import Field
+
+import timescaledb
+from timescaledb import create_engine, TimescaleModel
+
+TIMESCALE_DATABASE_URL = "postgresql://user:password@localhost:5432/timescaledb"
+engine = create_engine(TIMESCALE_DATABASE_URL, timezone="UTC")
+
+class SensorDos(TimescaleModel, table=True):
+    sensor_id: int = Field(index=True)
+    value: float
+
+    __enable_compression__ = True
+    __compress_orderby__ = "time DESC"
+    __compress_segmentby__ = "sensor_id"  
+    __chunk_time_interval__ = "7 days"
+    __drop_after__ = "1 year"
+    __migrate_data__ = True
+    __if_not_exists__ = True
+
+
+# Create the table and the hypertable
+with Session(engine) as session:
+    # Create the table in the database
+    SQLModel.metadata.create_all(engine)
+    # Creates all hypertable, add compression policies, and add retention policy
+    timescaledb.metadata.create_all(engine)
+```
+
+
+
+
+## Sample Usage 
+
+Below is a sample of using `timescaledb` in a FastAPI app much like the example in [./sample_project](./sample_project).
+
+`src/models.py`
 ```python
 from sqlmodel import Field, SQLModel
 
@@ -27,6 +114,10 @@ from timescaledb import TimescaleModel
 class Metric(TimescaleModel, table=True):
     temp: float
 
+    __enable_compression__ = True
+    __chunk_time_interval__ = "2 weeks"
+    __drop_after__ = "1 year"
+
 
 class MetricCreate(Metric):
     # not a table but a Pydantic model
@@ -45,7 +136,7 @@ class MetricRead(Metric):
 
 The `timescaledb.create_engine` is a wrapper around `sqlmodel.create_engine` (which is a wrapper around `sqlalchemy.create_engine`) that ensures a timezone is set for your database. 
 
-`database.py`
+`src/database.py`
 ```python
 import timescaledb
 from sqlmodel import Session, SQLModel
@@ -78,7 +169,7 @@ def init_db():
 
 Put it all together in a FastAPI app.
 
-`main.py`
+`src/main.py`
 ```python
 from fastapi import FastAPI
 
@@ -114,14 +205,3 @@ def list_metrics(session: Session = Depends(get_session)):
     return metrics
 ```
 
-### Review the Sample Project
-
-
-In [./sample_project](./sample_project) you can review a complete example that includes:
-- A TimescaleDB model
-- A FastAPI app
-- Sample queries (`time_bucket_gapfill_query`, `time_bucket_query`)
-
-
-
-

src/timescaledb/__init__.py

@@ -12,6 +12,7 @@
 )
 from .models import TimescaleModel
 from .queries import time_bucket_gapfill_query, time_bucket_query
+from .retention import add_retention_policy
 
 __all__ = [
     "metadata",

src/timescaledb/cleaners.py

@@ -0,0 +1,32 @@
+from datetime import timedelta
+
+
+def clean_interval(interval: str | int | timedelta) -> tuple[str | int, bool]:
+    """
+    Given an interval, return the interval type and a valid internal value
+
+    Args:
+        interval: The interval to extract the type and value from
+
+    Returns:
+        If valid, a tuple containing the interval type and a valid internal value
+        If invalid, the original interval and the string "INVALID"
+    """
+    if isinstance(interval, timedelta):
+        # Convert to microseconds
+        cleaned_interval = int(interval.total_seconds())
+        return cleaned_interval, "INTEGER"
+    elif isinstance(interval, int):
+        # Microseconds
+        cleaned_interval = interval
+        return cleaned_interval, "INTEGER"
+    elif isinstance(interval, str):
+        # Such as INTERVAL 1 day
+        # or INTERVAL '2 weeks'
+        # pop the term "INTERVAL"
+        cleaned_interval = interval.replace("INTERVAL", "").strip()
+        # remove any extra quotes
+        cleaned_interval = cleaned_interval.replace("'", "").replace('"', "")
+        return cleaned_interval, "INTERVAL"
+    else:
+        return interval, "INVALID"

src/timescaledb/compression/add.py

@@ -12,7 +12,7 @@
 
 
 def add_compression_policy(
-    session: Session, model: Type[SQLModel], commit: bool = True
+    session: Session, model: Type[SQLModel], commit: bool = True, auto_enable=False
 ) -> None:
     """
     Enable compression for a hypertable

src/timescaledb/hypertables/chunks/__init__.py

@@ -4,62 +4,88 @@
 import sqlalchemy
 from sqlmodel import Session, SQLModel
 
-from timescaledb.hypertables import sql_statements as sql
-from timescaledb.hypertables import validators
+from timescaledb.hypertables.extractors import extract_model_hypertable_params
+from timescaledb.hypertables.schemas import HypertableCreateSchema
 
-HYPERTABLE_INTERVAL_TYPE_SQL = {
-    "INTERVAL": sql.CREATE_HYPERTABLE_SQL_VIA_INTERVAL,
-    "TIMESTAMP": sql.CREATE_HYPERTABLE_SQL_VIA_TIMESTAMP,
-}
+# from timescaledb.hypertables.schemas import HypertableParams
 
 
 def create_hypertable(
     session: Session,
-    model: Type[SQLModel],
-    if_not_exists: bool = True,
-    migrate_data: bool = True,
     commit: bool = True,
+    model: Type[SQLModel] = None,
+    table_name: str = None,
+    hypertable_options: dict = {
+        "if_not_exists": True,
+        "migrate_data": True,
+    },
+    overwrite_model_params: bool = False,
 ) -> None:
-    """
-    Create a hypertable from a SQLModel class
-    """
-    time_column = getattr(model, "__time_column__", None)
-    validators.validate_time_column(model, time_column)
-    interval = getattr(model, "__chunk_time_interval__", None)
-    validators.validate_chunk_time_interval(model, time_column, interval)
-    sql_template = None
-    cleaned_interval = None
-    if isinstance(interval, timedelta):
-        # Convert to microseconds
-        cleaned_interval = int(interval.total_seconds())
-        sql_template = HYPERTABLE_INTERVAL_TYPE_SQL["TIMESTAMP"]
-    elif isinstance(interval, int):
-        # Microseconds
-        cleaned_interval = interval
-        sql_template = HYPERTABLE_INTERVAL_TYPE_SQL["TIMESTAMP"]
-    elif isinstance(interval, str):
-        # Such as INTERVAL 1 day
-        # or INTERVAL '2 weeks'
-        # pop the term "INTERVAL"
-        cleaned_interval = interval.replace("INTERVAL", "").strip()
-        # remove any extra quotes
-        cleaned_interval = cleaned_interval.replace("'", "").replace('"', "")
-        sql_template = HYPERTABLE_INTERVAL_TYPE_SQL["INTERVAL"]
-    else:
-        raise ValueError("Invalid interval type")
-    if sql_template is None or cleaned_interval is None:
-        raise ValueError("Invalid interval type")
+    """Create a TimescaleDB hypertable from a SQLModel class.
+
+    This function converts a regular table into a TimescaleDB hypertable. Hypertable parameters
+    can be specified either through model configuration or directly via hypertable_options.
+
+    Args:
+        session (Session): SQLAlchemy session instance
+        commit (bool, optional): Whether to commit the transaction after creating the hypertable.
+            Defaults to True.
+        model (Type[SQLModel]): The SQLModel class to convert into a hypertable. Must have
+            hypertable parameters defined either in the model or via hypertable_options.
+        table_name (str, optional): Override the table name. If None, uses the model's table name.
+        hypertable_options (dict, optional): Additional hypertable configuration options. Defaults to:
+            {
+                "if_not_exists": True,  # Skip if hypertable already exists
+                "migrate_data": True,   # Migrate existing data to the hypertable
+            }
+        overwrite_model_params (bool, optional): If True, hypertable_options will override any
+            conflicting parameters defined in the model. If False, model parameters take precedence.
+            Defaults to False.
+
+    Raises:
+        ValueError: If model parameter is None.
 
-    table_name = getattr(model, "__tablename__", None)
+    Example:
+        ```python
+        from sqlmodel import SQLModel
+
+        class Metrics(SQLModel, table=True):
+            __hypertable_params__ = {
+                "time_column": "timestamp",
+                "chunk_time_interval": timedelta(days=7)
+            }
+            # ... model fields ...
+
+        # Create hypertable using model parameters
+        create_hypertable(session, model=Metrics)
+
+        # Create hypertable with custom options
+        create_hypertable(
+            session,
+            model=Metrics,
+            hypertable_options={
+                "chunk_time_interval": timedelta(days=1),
+                "if_not_exists": True
+            },
+            overwrite_model_params=True
+        )
+        ```
+    """
+    if model is None and table_name is None:
+        raise ValueError("model or table_name is required")
     params = {
+        "model": None,
         "table_name": table_name,
-        "time_column": time_column,
-        "chunk_time_interval": cleaned_interval,
-        "if_not_exists": "true" if if_not_exists else "false",
-        "migrate_data": "true" if migrate_data else "false",
+        **hypertable_options,
     }
-    query = sqlalchemy.text(sql_template).bindparams(**params)
-    compiled_query = str(query.compile(compile_kwargs={"literal_binds": True}))
-    session.execute(sqlalchemy.text(compiled_query))
+    if model is not None:
+        model_params = extract_model_hypertable_params(model)
+        params = {**model_params}
+        if overwrite_model_params:
+            params.update(**hypertable_options)
+
+    schema = HypertableCreateSchema(**params)
+    query = schema.to_sql_query()
+    session.execute(sqlalchemy.text(query))
     if commit:
         session.commit()

src/timescaledb/hypertables/chunks/show.py

@@ -0,0 +1,27 @@
+from typing import Type
+
+from sqlmodel import SQLModel
+
+
+def extract_model_hypertable_params(
+    model: Type[SQLModel],
+) -> dict:
+    """
+    Format the SQL query based on the model's retention policy
+
+    Returns:
+        str: The formatted SQL query ready for execution
+    """
+    time_column = getattr(model, "__time_column__", None)
+    chunk_time_interval = getattr(model, "__chunk_time_interval__", None)
+    table_name = getattr(model, "__tablename__", None)
+    if_not_exists = getattr(model, "__if_not_exists__", False)
+    migrate_data = getattr(model, "__migrate_data__", False)
+    return {
+        "model": model,
+        "table_name": table_name,
+        "time_column": time_column,
+        "chunk_time_interval": chunk_time_interval,
+        "if_not_exists": if_not_exists,
+        "migrate_data": migrate_data,
+    }