Define db schema using SQLAlchemy

Author: cc-a

pyproject.toml

@@ -36,7 +36,9 @@ dependencies = [
     "pypubsub",
     "tomlkit",
     "duckdb",
-    "fsspec"
+    "fsspec",
+    "sqlalchemy",
+    "duckdb-engine"
 ]
 dynamic = ["version"]
 

src/muse/new_input/readers.py

@@ -1,11 +1,48 @@
-import duckdb
 import numpy as np
 import xarray as xr
+from sqlalchemy import CheckConstraint, ForeignKey
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+
+
+class TableBase(DeclarativeBase):
+    pass
+
+
+class Regions(TableBase):
+    __tablename__ = "regions"
+
+    name: Mapped[str] = mapped_column(primary_key=True)
+
+
+class Commodities(TableBase):
+    __tablename__ = "commodities"
+
+    name: Mapped[str] = mapped_column(primary_key=True)
+    type: Mapped[str] = mapped_column(
+        CheckConstraint("type IN ('energy', 'service', 'material', 'environmental')")
+    )
+    unit: Mapped[str]
+
+
+class Demand(TableBase):
+    __tablename__ = "demand"
+
+    year: Mapped[int] = mapped_column(primary_key=True, autoincrement=False)
+    commodity: Mapped[Commodities] = mapped_column(
+        ForeignKey("commodities.name"), primary_key=True
+    )
+    region: Mapped[Regions] = mapped_column(
+        ForeignKey("regions.name"), primary_key=True
+    )
+    demand: Mapped[float]
 
 
 def read_inputs(data_dir):
-    data = {}
-    con = duckdb.connect(":memory:")
+    from sqlalchemy import create_engine
+
+    engine = create_engine("duckdb:///:memory:")
+    TableBase.metadata.create_all(engine)
+    con = engine.raw_connection().driver_connection
 
     with open(data_dir / "regions.csv") as f:
         regions = read_regions_csv(f, con)  # noqa: F841
@@ -16,32 +53,20 @@ def read_inputs(data_dir):
     with open(data_dir / "demand.csv") as f:
         demand = read_demand_csv(f, con)  # noqa: F841
 
+    data = {}
     data["global_commodities"] = calculate_global_commodities(commodities)
     return data
 
 
 def read_regions_csv(buffer_, con):
-    sql = """CREATE TABLE regions (
-      name VARCHAR PRIMARY KEY,
-    );
-    """
-    con.sql(sql)
     rel = con.read_csv(buffer_, header=True, delimiter=",")  # noqa: F841
-    con.sql("INSERT INTO regions SELECT name FROM rel;")
+    con.execute("INSERT INTO regions SELECT name FROM rel;")
     return con.sql("SELECT name from regions").fetchnumpy()
 
 
 def read_commodities_csv(buffer_, con):
-    sql = """CREATE TABLE commodities (
-      name VARCHAR PRIMARY KEY,
-      type VARCHAR CHECK (type IN ('energy', 'service', 'material', 'environmental')),
-      unit VARCHAR,
-    );
-    """
-    con.sql(sql)
     rel = con.read_csv(buffer_, header=True, delimiter=",")  # noqa: F841
     con.sql("INSERT INTO commodities SELECT name, type, unit FROM rel;")
-
     return con.sql("select name, type, unit from commodities").fetchnumpy()
 
 
@@ -63,14 +88,6 @@ def calculate_global_commodities(commodities):
 
 
 def read_demand_csv(buffer_, con):
-    sql = """CREATE TABLE demand (
-    year BIGINT,
-    commodity VARCHAR REFERENCES commodities(name),
-    region VARCHAR REFERENCES regions(name),
-    demand DOUBLE,
-    );
-    """
-    con.sql(sql)
     rel = con.read_csv(buffer_, header=True, delimiter=",")  # noqa: F841
     con.sql("INSERT INTO demand SELECT year, commodity_name, region, demand FROM rel;")
     return con.sql("SELECT * from demand").fetchnumpy()

tests/test_readers.py

@@ -868,7 +868,14 @@ def default_new_input(tmp_path):
 
 @fixture
 def con():
-    return duckdb.connect(":memory:")
+    from muse.new_input.readers import TableBase
+    from sqlalchemy import create_engine
+    from sqlalchemy.orm import Session
+
+    engine = create_engine("duckdb:///:memory:")
+    session = Session(engine)
+    TableBase.metadata.create_all(engine)
+    return session.connection().connection
 
 
 @fixture
@@ -899,7 +906,15 @@ def test_read_regions(populate_regions):
     assert populate_regions["name"] == np.array(["R1"])
 
 
-def test_read_new_global_commodities(populate_commodities):
+def test_read_regions_primary_key_constraint(default_new_input, con):
+    from muse.new_input.readers import read_regions_csv
+
+    csv = StringIO("name\nR1\nR1\n")
+    with raises(duckdb.ConstraintException, match=".*duplicate key.*"):
+        read_regions_csv(csv, con)
+
+
+def test_read_new_commodities(populate_commodities):
     data = populate_commodities
     assert list(data["name"]) == ["electricity", "gas", "heat", "wind", "CO2f"]
     assert list(data["type"]) == ["energy"] * 5
@@ -921,7 +936,15 @@ def test_calculate_global_commodities(populate_commodities):
     assert list(data.data_vars["unit"].values) == list(populate_commodities["unit"])
 
 
-def test_read_new_global_commodities_type_constraint(default_new_input, con):
+def test_read_new_commodities_primary_key_constraint(default_new_input, con):
+    from muse.new_input.readers import read_commodities_csv
+
+    csv = StringIO("name,type,unit\nfoo,energy,bar\nfoo,energy,bar\n")
+    with raises(duckdb.ConstraintException, match=".*duplicate key.*"):
+        read_commodities_csv(csv, con)
+
+
+def test_read_new_commodities_type_constraint(default_new_input, con):
     from muse.new_input.readers import read_commodities_csv
 
     csv = StringIO("name,type,unit\nfoo,invalid,bar\n")
@@ -957,6 +980,32 @@ def test_new_read_demand_csv_region_constraint(
         read_demand_csv(csv, con)
 
 
+def test_new_read_demand_csv_primary_key_constraint(
+    default_new_input, con, populate_commodities, populate_regions
+):
+    from muse.new_input.readers import read_demand_csv, read_regions_csv
+
+    # Add another region so we can test varying it as a primary key
+    csv = StringIO("name\nR2\n")
+    read_regions_csv(csv, con)
+
+    # all fine so long as one primary key column differs
+    csv = StringIO(
+        """year,commodity_name,region,demand
+2020,gas,R1,0
+2021,gas,R1,0
+2020,heat,R1,0
+2020,gas,R2,0
+"""
+    )
+    read_demand_csv(csv, con)
+
+    # no good if all primary key columns match a previous entry
+    csv = StringIO("year,commodity_name,region,demand\n2020,gas,R1,0")
+    with raises(duckdb.ConstraintException, match=".*duplicate key.*"):
+        read_demand_csv(csv, con)
+
+
 @mark.xfail
 def test_demand_dataset(default_new_input):
     import duckdb