|  | 
| 2 | 2 | # Distributed under the terms of the AGPLv3 license, see LICENSE. | 
| 3 | 3 | import pytest | 
| 4 | 4 | 
 | 
| 5 |  | -from cratedb_toolkit.model import DatabaseAddress | 
| 6 |  | -from cratedb_toolkit.retention.model import JobSettings, RetentionPolicy, RetentionStrategy | 
| 7 |  | -from cratedb_toolkit.retention.setup.schema import setup_schema | 
| 8 |  | -from cratedb_toolkit.retention.store import RetentionPolicyStore | 
| 9 | 5 | from cratedb_toolkit.testing.testcontainers.azurite import ExtendedAzuriteContainer | 
| 10 | 6 | from cratedb_toolkit.testing.testcontainers.cratedb import CrateDBContainer | 
| 11 | 7 | from cratedb_toolkit.testing.testcontainers.minio import ExtendedMinioContainer | 
|  | 8 | +from cratedb_toolkit.util import DatabaseAdapter | 
| 12 | 9 | from cratedb_toolkit.util.common import setup_logging | 
| 13 |  | -from cratedb_toolkit.util.database import DatabaseAdapter, run_sql | 
| 14 | 10 | 
 | 
| 15 | 11 | # Use different schemas both for storing the retention policy table, and | 
| 16 | 12 | # the test data, so that they do not accidentally touch the default `doc` | 
| @@ -107,186 +103,4 @@ def azurite(): | 
| 107 | 103 |         yield azurite | 
| 108 | 104 | 
 | 
| 109 | 105 | 
 | 
| 110 |  | -@pytest.fixture() | 
| 111 |  | -def database(cratedb, settings): | 
| 112 |  | -    """ | 
| 113 |  | -    Provide a client database adapter, which is connected to the test database instance. | 
| 114 |  | -    """ | 
| 115 |  | -    yield DatabaseAdapter(dburi=settings.database.dburi) | 
| 116 |  | - | 
| 117 |  | - | 
| 118 |  | -@pytest.fixture() | 
| 119 |  | -def store(database, settings): | 
| 120 |  | -    """ | 
| 121 |  | -    Provide a client database adapter, which is connected to the test database instance. | 
| 122 |  | -    The retention policy database table schema has been established. | 
| 123 |  | -    """ | 
| 124 |  | -    setup_schema(settings=settings) | 
| 125 |  | -    rps = RetentionPolicyStore(settings=settings) | 
| 126 |  | -    yield rps | 
| 127 |  | - | 
| 128 |  | - | 
| 129 |  | -@pytest.fixture() | 
| 130 |  | -def settings(cratedb): | 
| 131 |  | -    """ | 
| 132 |  | -    Provide configuration and runtime settings object, parameterized for the test suite. | 
| 133 |  | -    """ | 
| 134 |  | -    database_url = cratedb.get_connection_url() | 
| 135 |  | -    job_settings = JobSettings(database=DatabaseAddress.from_string(database_url)) | 
| 136 |  | -    job_settings.policy_table.schema = TESTDRIVE_EXT_SCHEMA | 
| 137 |  | -    return job_settings | 
| 138 |  | - | 
| 139 |  | - | 
| 140 |  | -@pytest.fixture(scope="function") | 
| 141 |  | -def policies(cratedb, settings, store): | 
| 142 |  | -    """ | 
| 143 |  | -    Populate the retention policy table. | 
| 144 |  | -    """ | 
| 145 |  | -    database_url = cratedb.get_connection_url() | 
| 146 |  | -    rules = [ | 
| 147 |  | -        # Retention policy rule for the DELETE strategy. | 
| 148 |  | -        RetentionPolicy( | 
| 149 |  | -            strategy=RetentionStrategy.DELETE, | 
| 150 |  | -            table_schema=TESTDRIVE_DATA_SCHEMA, | 
| 151 |  | -            table_name="raw_metrics", | 
| 152 |  | -            partition_column="ts_day", | 
| 153 |  | -            retention_period=1, | 
| 154 |  | -        ), | 
| 155 |  | -        # Retention policy rule for the DELETE strategy, using tags. | 
| 156 |  | -        RetentionPolicy( | 
| 157 |  | -            strategy=RetentionStrategy.DELETE, | 
| 158 |  | -            tags={"foo", "bar"}, | 
| 159 |  | -            table_schema=TESTDRIVE_DATA_SCHEMA, | 
| 160 |  | -            table_name="sensor_readings", | 
| 161 |  | -            partition_column="time_month", | 
| 162 |  | -            retention_period=1, | 
| 163 |  | -        ), | 
| 164 |  | -    ] | 
| 165 |  | -    for rule in rules: | 
| 166 |  | -        store.create(rule, ignore="DuplicateKeyException") | 
| 167 |  | - | 
| 168 |  | -    # Synchronize data. | 
| 169 |  | -    run_sql(database_url, f"REFRESH TABLE {settings.policy_table.fullname};") | 
| 170 |  | - | 
| 171 |  | - | 
| 172 |  | -@pytest.fixture(scope="function") | 
| 173 |  | -def raw_metrics(cratedb, settings, store): | 
| 174 |  | -    """ | 
| 175 |  | -    Populate the `raw_metrics` table. | 
| 176 |  | -    """ | 
| 177 |  | - | 
| 178 |  | -    tablename_full = f'"{TESTDRIVE_DATA_SCHEMA}"."raw_metrics"' | 
| 179 |  | - | 
| 180 |  | -    database_url = cratedb.get_connection_url() | 
| 181 |  | -    ddl = f""" | 
| 182 |  | -        CREATE TABLE {tablename_full} ( | 
| 183 |  | -           "variable" TEXT, | 
| 184 |  | -           "timestamp" TIMESTAMP WITH TIME ZONE, | 
| 185 |  | -           "ts_day" TIMESTAMP GENERATED ALWAYS AS date_trunc('day', "timestamp"), | 
| 186 |  | -           "value" REAL, | 
| 187 |  | -           "quality" INTEGER, | 
| 188 |  | -           PRIMARY KEY ("variable", "timestamp", "ts_day") | 
| 189 |  | -        ) | 
| 190 |  | -        PARTITIONED BY ("ts_day") | 
| 191 |  | -        WITH ("routing.allocation.require.storage" = 'hot') | 
| 192 |  | -        ; | 
| 193 |  | -    """ | 
| 194 |  | - | 
| 195 |  | -    dml = f""" | 
| 196 |  | -        INSERT INTO {tablename_full} | 
| 197 |  | -            (variable, timestamp, value, quality) | 
| 198 |  | -        SELECT | 
| 199 |  | -            'temperature' AS variable, | 
| 200 |  | -            generate_series AS timestamp, | 
| 201 |  | -            RANDOM()*100 AS value, | 
| 202 |  | -            0 AS quality | 
| 203 |  | -        FROM generate_series('2023-06-01', '2023-06-30', '5 days'::INTERVAL); | 
| 204 |  | -    """ | 
| 205 |  | - | 
| 206 |  | -    run_sql(database_url, ddl) | 
| 207 |  | -    run_sql(database_url, dml) | 
| 208 |  | -    run_sql(database_url, f"REFRESH TABLE {tablename_full};") | 
| 209 |  | - | 
| 210 |  | -    return tablename_full | 
| 211 |  | - | 
| 212 |  | - | 
| 213 |  | -@pytest.fixture(scope="function") | 
| 214 |  | -def sensor_readings(cratedb, settings, store): | 
| 215 |  | -    """ | 
| 216 |  | -    Populate the `sensor_readings` table. | 
| 217 |  | -    """ | 
| 218 |  | - | 
| 219 |  | -    tablename_full = f'"{TESTDRIVE_DATA_SCHEMA}"."sensor_readings"' | 
| 220 |  | - | 
| 221 |  | -    database_url = cratedb.get_connection_url() | 
| 222 |  | -    ddl = f""" | 
| 223 |  | -        CREATE TABLE {tablename_full} ( | 
| 224 |  | -           time TIMESTAMP WITH TIME ZONE NOT NULL, | 
| 225 |  | -           time_month TIMESTAMP WITH TIME ZONE GENERATED ALWAYS AS DATE_TRUNC('month', "time"), | 
| 226 |  | -           sensor_id TEXT NOT NULL, | 
| 227 |  | -           battery_level DOUBLE PRECISION, | 
| 228 |  | -           battery_status TEXT, | 
| 229 |  | -           battery_temperature DOUBLE PRECISION | 
| 230 |  | -        ) | 
| 231 |  | -        PARTITIONED BY (time_month); | 
| 232 |  | -    """ | 
| 233 |  | - | 
| 234 |  | -    dml = f""" | 
| 235 |  | -        INSERT INTO {tablename_full} | 
| 236 |  | -            (time, sensor_id, battery_level, battery_status, battery_temperature) | 
| 237 |  | -        SELECT | 
| 238 |  | -            generate_series AS time, | 
| 239 |  | -            'batt01' AS sensor_id, | 
| 240 |  | -            RANDOM()*100 AS battery_level, | 
| 241 |  | -            'FULL' AS battery_status, | 
| 242 |  | -            RANDOM()*100 AS battery_temperature | 
| 243 |  | -        FROM generate_series( | 
| 244 |  | -            '2023-05-01'::TIMESTAMPTZ, | 
| 245 |  | -            '2023-06-30'::TIMESTAMPTZ, | 
| 246 |  | -            '7 days'::INTERVAL | 
| 247 |  | -        ); | 
| 248 |  | -    """ | 
| 249 |  | - | 
| 250 |  | -    run_sql(database_url, ddl) | 
| 251 |  | -    run_sql(database_url, dml) | 
| 252 |  | -    run_sql(database_url, f"REFRESH TABLE {tablename_full};") | 
| 253 |  | - | 
| 254 |  | -    return tablename_full | 
| 255 |  | - | 
| 256 |  | - | 
| 257 |  | -@pytest.fixture(scope="function") | 
| 258 |  | -def raw_metrics_reallocate_policy(store): | 
| 259 |  | -    """ | 
| 260 |  | -    Populate the retention policy table. | 
| 261 |  | -    """ | 
| 262 |  | -    # Retention policy rule for the REALLOCATE strategy. | 
| 263 |  | -    rule = RetentionPolicy( | 
| 264 |  | -        strategy=RetentionStrategy.REALLOCATE, | 
| 265 |  | -        table_schema=TESTDRIVE_DATA_SCHEMA, | 
| 266 |  | -        table_name="raw_metrics", | 
| 267 |  | -        partition_column="ts_day", | 
| 268 |  | -        retention_period=60, | 
| 269 |  | -        reallocation_attribute_name="storage", | 
| 270 |  | -        reallocation_attribute_value="warm", | 
| 271 |  | -    ) | 
| 272 |  | -    store.create(rule, ignore="DuplicateKeyException") | 
| 273 |  | - | 
| 274 |  | - | 
| 275 |  | -@pytest.fixture(scope="function") | 
| 276 |  | -def sensor_readings_snapshot_policy(store): | 
| 277 |  | -    """ | 
| 278 |  | -    Populate the retention policy table. | 
| 279 |  | -    """ | 
| 280 |  | -    # Retention policy rule for the SNAPSHOT strategy. | 
| 281 |  | -    rule = RetentionPolicy( | 
| 282 |  | -        strategy=RetentionStrategy.SNAPSHOT, | 
| 283 |  | -        table_schema=TESTDRIVE_DATA_SCHEMA, | 
| 284 |  | -        table_name="sensor_readings", | 
| 285 |  | -        partition_column="time_month", | 
| 286 |  | -        retention_period=365, | 
| 287 |  | -        target_repository_name="export_cold", | 
| 288 |  | -    ) | 
| 289 |  | -    store.create(rule, ignore="DuplicateKeyException") | 
| 290 |  | - | 
| 291 |  | - | 
| 292 | 106 | setup_logging() | 
0 commit comments