Merge pull request #94 from staadecker/save_scenario

Add script to save new scenarios to the database

Author: pesap

docs/Database.md

@@ -41,8 +41,18 @@ The command to enter PostgreSQL while SSH'd into the server is `psql wecc`.
 
 ## Making changes to the database
 
-Whether it's adding data to the database or changing its schema, it's important
-to proceed carefully when making changes to the database. Always make sure to
+### Common operations
+
+Some database changes are common operations that are used repeatedly.
+For these operations, it's helpful to make a command line tool to make the process easy.
+
+For example, adding a scenario to the database is a common operation. Therefore, we have built
+the `switch db save_scenario` command to easily do this. (Run `switch db save_scenario --help` for details).
+
+### Custom operations
+
+When running custom operations on the database it's very important
+to proceed carefully. Always make sure to
 **test and keep track of your changes**.
 
 Here are the steps to make a change.
@@ -61,9 +71,7 @@ to the convention `YYYY-MM-DD_<script_name>`.
 5. Open a pull request to add your script to the repository (see [`docs/Contribute.md`](Contribute.md))
 so we can keep track of the changes that have been made.
 
-### Bigger changes
-
-Sometimes, it isn't feasible to have your entire change as a single SQL script.
+Note that sometimes, it isn't feasible to have your entire change as a single SQL script.
 One way to make bigger changes is to use a Python script. Use the same process
 as for the SQL scripts. That is save your Python scripts in the `/database` folder.
 

switch_model/tools/templates/config.yaml

@@ -21,6 +21,8 @@ get_inputs:
   #  the following parameters are optional and will override the defaults for that scenario
   #  this should only be used for preliminary exploration of scenarios or for testing
   #  if a scenario is part of research, it should be included in the database as a row in the scenarios table.
+  #
+  # -1 is equivalent to null
   #  study_timeframe_id:
   #  time_sample_id:
   #  demand_scenario_id:

switch_model/wecc/__main__.py

@@ -5,17 +5,25 @@
 from __future__ import print_function
 
 import argparse
+import importlib
 import sys
-import switch_model.wecc.sampling.cli as sample
+
+
+def get_module_runner(module):
+    def runner():
+        importlib.import_module(module).main()
+
+    return runner
 
 
 cmds = {
-    "sample": sample.main,
+    "sample": get_module_runner("switch_model.wecc.sampling.cli"),
+    "save_scenario": get_module_runner("switch_model.wecc.save_scenario"),
 }
 
 
 def main(args=None):
-    parser = argparse.ArgumentParser()
+    parser = argparse.ArgumentParser(add_help=False)
     parser.add_argument(
         "subcommand", choices=cmds.keys(), help="The possible switch subcommands"
     )

switch_model/wecc/get_inputs/scenario.py

@@ -0,0 +1,53 @@
+class ScenarioParams:
+    def __init__(self, scenario_id=None):
+        # This list of attributes defines the columns that we query from the database.
+        # Therefore attribute names matter!
+        self.scenario_id = scenario_id
+        self.name = None
+        self.description = None
+        self.study_timeframe_id = None
+        self.time_sample_id = None
+        self.demand_scenario_id = None
+        self.fuel_simple_price_scenario_id = None
+        self.generation_plant_scenario_id = None
+        self.generation_plant_cost_scenario_id = None
+        self.generation_plant_existing_and_planned_scenario_id = None
+        self.hydro_simple_scenario_id = None
+        self.carbon_cap_scenario_id = None
+        self.supply_curves_scenario_id = None
+        self.regional_fuel_market_scenario_id = None
+        self.rps_scenario_id = None
+        self.enable_dr = None
+        self.enable_ev = None
+        self.transmission_base_capital_cost_scenario_id = None
+        self.ca_policies_scenario_id = None
+        self.enable_planning_reserves = None
+        self.generation_plant_technologies_scenario_id = None
+        self.variable_o_m_cost_scenario_id = None
+        self.wind_to_solar_ratio = None
+
+
+def load_scenario_from_config(config, db_cursor) -> ScenarioParams:
+    config = config["get_inputs"]  # Only keep the config that matters
+
+    # Create the ScenarioParams object
+    params = ScenarioParams(scenario_id=config["scenario_id"])
+    param_names = list(params.__dict__.keys())
+    print(param_names)
+
+    # Read from the database all the parameters.
+    db_cursor.execute(
+        f"""SELECT {",".join(param_names)}
+            FROM scenario
+            WHERE scenario_id = {params.scenario_id};"""
+    )
+    db_values = list(db_cursor.fetchone())
+
+    # Allow overriding from config
+    for i, param_name in enumerate(param_names):
+        value = config[param_name] if param_name in config else db_values[i]
+        if value == -1:
+            value = None
+        setattr(params, param_name, value)
+
+    return params

switch_model/wecc/save_scenario.py

@@ -0,0 +1,81 @@
+import argparse
+import sys
+
+import yaml
+
+from switch_model.utilities import query_yes_no
+from switch_model.wecc.get_inputs.scenario import load_scenario_from_config
+from switch_model.wecc.utilities import connect, load_config
+
+
+def main():
+    # Start CLI
+    parser = argparse.ArgumentParser(
+        description="Creates a new scenario in the database by using the values in"
+        " config.yaml. Therefore the new scenario will have the same values"
+        " as the base scenario but you can override specific columns by "
+        " specifying them in config.yaml."
+    )
+    parser.add_argument(
+        "scenario_id", type=int, help="The id of the new scenario to add to db."
+    )
+    parser.add_argument(
+        "--name", required=True, help="The name of the new scenario to add in db."
+    )
+    parser.add_argument(
+        "--description",
+        required=True,
+        help="The new scenario description to add in db.",
+    )
+    parser.add_argument(
+        "--db-env-var", default="DB_URL", help="The connection environment variable."
+    )
+
+    # Optional arguments
+    parser.add_argument(
+        "--config_file",
+        default="config.yaml",
+        type=str,
+        help="Configuration file to use.",
+    )
+
+    args = parser.parse_args()
+
+    # Start db connection
+    db_conn = connect(connection_env_var=args.db_env_var)
+    db_cursor = db_conn.cursor()
+
+    # # Exit if you are not sure if you want to overwrite
+    # if args.overwrite:
+
+    config = load_config()
+    scenario_params = load_scenario_from_config(config, db_cursor)
+
+    # Override the given parameters
+    scenario_params.name = args.name
+    scenario_params.description = args.description
+    scenario_params.scenario_id = args.scenario_id
+
+    ordered_params = list(
+        filter(lambda v: v[1] is not None, scenario_params.__dict__.items())
+    )
+    columns = ",".join(v[0] for v in ordered_params)
+    values = ",".join(
+        f"'{v[1]}'" if type(v[1]) == str else str(v[1]) for v in ordered_params
+    )
+
+    query = f"""INSERT INTO scenario({columns}) VALUES  ({values});"""
+
+    print(f"\n{query}\n")
+
+    if not query_yes_no(
+        f"Are you sure you want to run the above query.?", default="no"
+    ):
+        sys.exit()
+
+    db_cursor.execute(query)
+    db_conn.commit()
+    db_cursor.close()
+    db_conn.close()
+
+    print(f"Ran query.")

switch_model/wecc/utilities.py

@@ -25,7 +25,7 @@ def load_dotenv():
         pass
 
 
-def connect(schema="switch"):
+def connect(schema="switch", connection_env_var="DB_URL"):
     """Connects to the Postgres DB
 
     This function uses the environment variables to get the URL to connect to the DB. Both
@@ -34,16 +34,17 @@ def connect(schema="switch"):
     Parameters
     ----------
     schema: str Schema of the DB to look for tables. Default is switch
+    connection_env_var: The environment variable to use as the connection string
 
     Returns
     -------
     conn: Database connection object from psycopg2
     """
     load_dotenv()
-    db_url = os.getenv("DB_URL")
+    db_url = os.getenv(connection_env_var)
     if db_url is None:
         raise Exception(
-            "Please set the environment variable 'DB_URL' to the database URL."
+            f"Please set the environment variable '{connection_env_var}' to the database URL."
             "The format is normally: postgresql://<user>:<password>@<host>:5432/<database>"
         )