Skip to content

Commit 1fe5899

Browse files
committed
Add AlchemyPowerStorage
1 parent ddf4d48 commit 1fe5899

File tree

3 files changed

+403
-0
lines changed

3 files changed

+403
-0
lines changed
Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
from __future__ import annotations
2+
3+
import json
4+
from typing import ClassVar
5+
6+
from sqlalchemy import Float, Integer, String, UniqueConstraint
7+
from sqlalchemy.orm import Mapped, mapped_column
8+
from typing_extensions import override
9+
10+
from pysatl_experiment.persistence.db_store.base import ModelBase, SessionType
11+
from pysatl_experiment.persistence.db_store.model import AbstractDbStore
12+
from pysatl_experiment.persistence.model.power.power import IPowerStorage, PowerModel, PowerQuery
13+
14+
15+
class AlchemyPower(ModelBase):
16+
__tablename__ = "power"
17+
18+
id: Mapped[int] = mapped_column(Integer, primary_key=True) # type: ignore
19+
experiment_id: Mapped[int] = mapped_column(Integer, nullable=False) # type: ignore
20+
criterion_code: Mapped[str] = mapped_column(String, nullable=False, index=True) # type: ignore
21+
criterion_parameters: Mapped[str] = mapped_column(String, nullable=False, index=True) # type: ignore
22+
sample_size: Mapped[int] = mapped_column(Integer, nullable=False, index=True) # type: ignore
23+
alternative_code: Mapped[str] = mapped_column(String, nullable=False, index=True) # type: ignore
24+
alternative_parameters: Mapped[str] = mapped_column(String, nullable=False, index=True) # type: ignore
25+
monte_carlo_count: Mapped[int] = mapped_column(Integer, nullable=False, index=True) # type: ignore
26+
significance_level: Mapped[float] = mapped_column(Float, nullable=False, index=True) # type: ignore
27+
results_criteria: Mapped[str] = mapped_column(String, nullable=False) # type: ignore
28+
29+
__table_args__ = (
30+
UniqueConstraint(
31+
"criterion_code",
32+
"criterion_parameters",
33+
"sample_size",
34+
"alternative_code",
35+
"alternative_parameters",
36+
"monte_carlo_count",
37+
"significance_level",
38+
name="uq_power_unique",
39+
),
40+
)
41+
42+
43+
class AlchemyPowerStorage(AbstractDbStore, IPowerStorage):
44+
session: ClassVar[SessionType]
45+
46+
def __init__(self, db_url: str):
47+
super().__init__(db_url=db_url)
48+
self._initialized: bool = False
49+
50+
@override
51+
def init(self) -> None:
52+
super().init()
53+
self._initialized = True
54+
55+
def _get_session(self) -> SessionType:
56+
if not getattr(self, "_initialized", False):
57+
raise RuntimeError("Storage not initialized. Call init() first.")
58+
return AlchemyPowerStorage.session
59+
60+
@override
61+
def get_data(self, query: PowerQuery) -> PowerModel | None:
62+
params_json = json.dumps(query.criterion_parameters)
63+
alt_params_json = json.dumps(query.alternative_parameters)
64+
row: AlchemyPower | None = (
65+
self._get_session()
66+
.query(AlchemyPower)
67+
.filter(
68+
AlchemyPower.criterion_code == query.criterion_code,
69+
AlchemyPower.criterion_parameters == params_json,
70+
AlchemyPower.sample_size == int(query.sample_size),
71+
AlchemyPower.alternative_code == query.alternative_code,
72+
AlchemyPower.alternative_parameters == alt_params_json,
73+
AlchemyPower.monte_carlo_count == int(query.monte_carlo_count),
74+
AlchemyPower.significance_level == float(query.significance_level),
75+
)
76+
.one_or_none()
77+
)
78+
if row is None:
79+
return None
80+
return PowerModel(
81+
experiment_id=int(row.experiment_id),
82+
criterion_code=query.criterion_code,
83+
criterion_parameters=query.criterion_parameters,
84+
sample_size=query.sample_size,
85+
alternative_code=query.alternative_code,
86+
alternative_parameters=query.alternative_parameters,
87+
monte_carlo_count=query.monte_carlo_count,
88+
significance_level=query.significance_level,
89+
results_criteria=json.loads(row.results_criteria),
90+
)
91+
92+
@override
93+
def insert_data(self, data: PowerModel) -> None:
94+
params_json = json.dumps(data.criterion_parameters)
95+
alt_params_json = json.dumps(data.alternative_parameters)
96+
existing: AlchemyPower | None = (
97+
self._get_session()
98+
.query(AlchemyPower)
99+
.filter(
100+
AlchemyPower.criterion_code == data.criterion_code,
101+
AlchemyPower.criterion_parameters == params_json,
102+
AlchemyPower.sample_size == int(data.sample_size),
103+
AlchemyPower.alternative_code == data.alternative_code,
104+
AlchemyPower.alternative_parameters == alt_params_json,
105+
AlchemyPower.monte_carlo_count == int(data.monte_carlo_count),
106+
AlchemyPower.significance_level == float(data.significance_level),
107+
)
108+
.one_or_none()
109+
)
110+
if existing is None:
111+
entity = AlchemyPower(
112+
experiment_id=int(data.experiment_id),
113+
criterion_code=data.criterion_code,
114+
criterion_parameters=params_json,
115+
sample_size=int(data.sample_size),
116+
alternative_code=data.alternative_code,
117+
alternative_parameters=alt_params_json,
118+
monte_carlo_count=int(data.monte_carlo_count),
119+
significance_level=float(data.significance_level),
120+
results_criteria=json.dumps(data.results_criteria),
121+
)
122+
self._get_session().add(entity)
123+
else:
124+
existing.experiment_id = int(data.experiment_id)
125+
existing.results_criteria = json.dumps(data.results_criteria)
126+
self._get_session().commit()
127+
128+
@override
129+
def delete_data(self, query: PowerQuery) -> None:
130+
params_json = json.dumps(query.criterion_parameters)
131+
alt_params_json = json.dumps(query.alternative_parameters)
132+
(
133+
self._get_session()
134+
.query(AlchemyPower)
135+
.filter(
136+
AlchemyPower.criterion_code == query.criterion_code,
137+
AlchemyPower.criterion_parameters == params_json,
138+
AlchemyPower.sample_size == int(query.sample_size),
139+
AlchemyPower.alternative_code == query.alternative_code,
140+
AlchemyPower.alternative_parameters == alt_params_json,
141+
AlchemyPower.monte_carlo_count == int(query.monte_carlo_count),
142+
AlchemyPower.significance_level == float(query.significance_level),
143+
)
144+
.delete()
145+
)
146+
self._get_session().commit()
Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
from __future__ import annotations
2+
3+
import pytest
4+
5+
from pysatl_experiment.persistence.model.power.power import PowerModel, PowerQuery
6+
from pysatl_experiment.persistence.power.alchemy.alchemy import AlchemyPowerStorage
7+
8+
9+
@pytest.fixture()
10+
def db_url() -> str:
11+
# Use in-memory SQLite with StaticPool as configured by init_db
12+
return "sqlite://"
13+
14+
15+
@pytest.fixture()
16+
def storage(db_url: str) -> AlchemyPowerStorage:
17+
store = AlchemyPowerStorage(db_url)
18+
store.init()
19+
return store
20+
21+
22+
def test_guard_requires_init(db_url: str) -> None:
23+
store = AlchemyPowerStorage(db_url)
24+
with pytest.raises(RuntimeError):
25+
_ = store.get_data(
26+
PowerQuery(
27+
criterion_code="crit_A",
28+
criterion_parameters=[0.1, 0.2],
29+
sample_size=10,
30+
alternative_code="alt_A",
31+
alternative_parameters=[1.0],
32+
monte_carlo_count=100,
33+
significance_level=0.05,
34+
)
35+
)
36+
37+
38+
def test_get_data_empty_returns_none(storage: AlchemyPowerStorage) -> None:
39+
query = PowerQuery(
40+
criterion_code="crit_A",
41+
criterion_parameters=[0.1, 0.2],
42+
sample_size=10,
43+
alternative_code="alt_A",
44+
alternative_parameters=[1.0],
45+
monte_carlo_count=100,
46+
significance_level=0.05,
47+
)
48+
assert storage.get_data(query) is None
49+
50+
51+
def test_insert_and_get(storage: AlchemyPowerStorage) -> None:
52+
model = PowerModel(
53+
experiment_id=123,
54+
criterion_code="crit_A",
55+
criterion_parameters=[0.1, 0.2],
56+
sample_size=10,
57+
alternative_code="alt_A",
58+
alternative_parameters=[1.0],
59+
monte_carlo_count=100,
60+
significance_level=0.05,
61+
results_criteria=[True, False, True],
62+
)
63+
storage.insert_data(model)
64+
65+
got = storage.get_data(
66+
PowerQuery(
67+
criterion_code="crit_A",
68+
criterion_parameters=[0.1, 0.2],
69+
sample_size=10,
70+
alternative_code="alt_A",
71+
alternative_parameters=[1.0],
72+
monte_carlo_count=100,
73+
significance_level=0.05,
74+
)
75+
)
76+
77+
assert got is not None
78+
assert got.experiment_id == model.experiment_id
79+
assert got.criterion_code == model.criterion_code
80+
assert got.criterion_parameters == model.criterion_parameters
81+
assert got.sample_size == model.sample_size
82+
assert got.alternative_code == model.alternative_code
83+
assert got.alternative_parameters == model.alternative_parameters
84+
assert got.monte_carlo_count == model.monte_carlo_count
85+
assert got.significance_level == model.significance_level
86+
assert got.results_criteria == model.results_criteria
87+
88+
89+
def test_delete_data(storage: AlchemyPowerStorage) -> None:
90+
model = PowerModel(
91+
experiment_id=7,
92+
criterion_code="crit_B",
93+
criterion_parameters=[0.3],
94+
sample_size=5,
95+
alternative_code="alt_B",
96+
alternative_parameters=[2.0, 3.0],
97+
monte_carlo_count=50,
98+
significance_level=0.1,
99+
results_criteria=[False, False],
100+
)
101+
storage.insert_data(model)
102+
103+
storage.delete_data(
104+
PowerQuery(
105+
criterion_code="crit_B",
106+
criterion_parameters=[0.3],
107+
sample_size=5,
108+
alternative_code="alt_B",
109+
alternative_parameters=[2.0, 3.0],
110+
monte_carlo_count=50,
111+
significance_level=0.1,
112+
)
113+
)
114+
115+
assert (
116+
storage.get_data(
117+
PowerQuery(
118+
criterion_code="crit_B",
119+
criterion_parameters=[0.3],
120+
sample_size=5,
121+
alternative_code="alt_B",
122+
alternative_parameters=[2.0, 3.0],
123+
monte_carlo_count=50,
124+
significance_level=0.1,
125+
)
126+
)
127+
is None
128+
)

0 commit comments

Comments
 (0)