From 716457363336ca05fb3512e00a293579968ed1ff Mon Sep 17 00:00:00 2001 From: Zohar Malamant Date: Fri, 28 May 2021 12:57:31 +0200 Subject: [PATCH 01/43] Move source to src/ directory --- src/ert/dark_storage/__init__.py | 0 src/ert/dark_storage/__main__.py | 105 +++ src/ert/dark_storage/_alembic/alembic.ini | 87 ++ src/ert/dark_storage/_alembic/alembic/README | 1 + src/ert/dark_storage/_alembic/alembic/env.py | 86 ++ .../_alembic/alembic/script.py.mako | 24 + .../versions/021d7514e351_add_recordinfo.py | 127 +++ .../1a4a83dfb895_add_metadata_column.py | 34 + ...76d1e2_inline_parameters_into_ensembles.py | 60 ++ .../versions/55307b70da13_add_record_class.py | 44 + .../versions/56177419dc39_add_priors.py | 77 ++ .../65569f4f3756_add_name_to_record.py | 28 + .../73541c7b7aa6_add_size_to_ensemble.py | 28 + .../7812741cc469_recordtype_as_enum.py | 51 ++ .../79137bad215d_add_matrix_labels.py | 38 + .../versions/896a28a0652f_add_file_block.py | 54 ++ ...29af73_record_as_inputs_and_outputs_of_.py | 87 ++ ...344bcf5_add_realization_index_to_record.py | 28 + .../9a218c7df2ac_separate_into_file_matrix.py | 79 ++ .../versions/a3f7ae51a72b_initial_schema.py | 95 ++ .../versions/a88924af33b0_add_experiment.py | 49 + ...05904381_add_backend_agnostic_uuid_type.py | 245 +++++ .../bea15008ea91_add_obs_trans_and_updates.py | 62 ++ .../c1b95721c62a_add_azure_blob_storage.py | 32 + .../c96ec073f9ee_change_prior_enum.py | 54 ++ .../d72e51c6bdf2_remove_record_class.py | 61 ++ .../versions/f46d5afa7639_add_observations.py | 68 ++ ...7aef9b17f8_add_record_unique_constraint.py | 30 + src/ert/dark_storage/app.py | 73 ++ src/ert/dark_storage/compute/__init__.py | 1 + src/ert/dark_storage/compute/misfits.py | 35 + src/ert/dark_storage/database.py | 66 ++ .../dark_storage/database_schema/__init__.py | 9 + .../dark_storage/database_schema/ensemble.py | 51 ++ .../database_schema/experiment.py | 36 + .../database_schema/metadatafield.py | 12 + .../database_schema/observation.py | 60 ++ src/ert/dark_storage/database_schema/prior.py | 44 + .../dark_storage/database_schema/record.py | 110 +++ .../database_schema/record_info.py | 46 + .../dark_storage/database_schema/update.py | 39 + src/ert/dark_storage/endpoints/__init__.py | 19 + .../endpoints/compute/__init__.py | 0 .../dark_storage/endpoints/compute/misfits.py | 90 ++ src/ert/dark_storage/endpoints/ensembles.py | 81 ++ src/ert/dark_storage/endpoints/experiments.py | 117 +++ .../dark_storage/endpoints/observations.py | 148 +++ src/ert/dark_storage/endpoints/priors.py | 57 ++ src/ert/dark_storage/endpoints/records.py | 856 ++++++++++++++++++ src/ert/dark_storage/endpoints/responses.py | 44 + src/ert/dark_storage/endpoints/updates.py | 71 ++ src/ert/dark_storage/ext/__init__.py | 3 + .../dark_storage/ext/graphene_sqlalchemy.py | 132 +++ src/ert/dark_storage/ext/sqlalchemy_arrays.py | 44 + src/ert/dark_storage/ext/uuid.py | 55 ++ src/ert/dark_storage/graphql/__init__.py | 61 ++ src/ert/dark_storage/graphql/ensembles.py | 106 +++ src/ert/dark_storage/graphql/experiments.py | 51 ++ src/ert/dark_storage/graphql/parameters.py | 26 + src/ert/dark_storage/graphql/responses.py | 19 + .../dark_storage/graphql/unique_responses.py | 13 + src/ert/dark_storage/graphql/updates.py | 15 + src/ert/dark_storage/json_schema/__init__.py | 11 + src/ert/dark_storage/json_schema/ensemble.py | 35 + .../dark_storage/json_schema/experiment.py | 22 + .../dark_storage/json_schema/observation.py | 42 + src/ert/dark_storage/json_schema/prior.py | 151 +++ src/ert/dark_storage/json_schema/record.py | 16 + src/ert/dark_storage/json_schema/update.py | 25 + src/ert/dark_storage/py.typed | 1 + src/ert/dark_storage/security.py | 25 + src/ert/dark_storage/testing/__init__.py | 4 + src/ert/dark_storage/testing/pytest11.py | 13 + src/ert/dark_storage/testing/testclient.py | 326 +++++++ 74 files changed, 4995 insertions(+) create mode 100644 src/ert/dark_storage/__init__.py create mode 100644 src/ert/dark_storage/__main__.py create mode 100644 src/ert/dark_storage/_alembic/alembic.ini create mode 100644 src/ert/dark_storage/_alembic/alembic/README create mode 100644 src/ert/dark_storage/_alembic/alembic/env.py create mode 100644 src/ert/dark_storage/_alembic/alembic/script.py.mako create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/021d7514e351_add_recordinfo.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/1a4a83dfb895_add_metadata_column.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/4effed76d1e2_inline_parameters_into_ensembles.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/55307b70da13_add_record_class.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/56177419dc39_add_priors.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/65569f4f3756_add_name_to_record.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/73541c7b7aa6_add_size_to_ensemble.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/7812741cc469_recordtype_as_enum.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/79137bad215d_add_matrix_labels.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/896a28a0652f_add_file_block.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/9064db29af73_record_as_inputs_and_outputs_of_.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/98601344bcf5_add_realization_index_to_record.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/9a218c7df2ac_separate_into_file_matrix.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/a3f7ae51a72b_initial_schema.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/a88924af33b0_add_experiment.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/ba7905904381_add_backend_agnostic_uuid_type.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/bea15008ea91_add_obs_trans_and_updates.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/c1b95721c62a_add_azure_blob_storage.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/c96ec073f9ee_change_prior_enum.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/d72e51c6bdf2_remove_record_class.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/f46d5afa7639_add_observations.py create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/fc7aef9b17f8_add_record_unique_constraint.py create mode 100644 src/ert/dark_storage/app.py create mode 100644 src/ert/dark_storage/compute/__init__.py create mode 100644 src/ert/dark_storage/compute/misfits.py create mode 100644 src/ert/dark_storage/database.py create mode 100644 src/ert/dark_storage/database_schema/__init__.py create mode 100644 src/ert/dark_storage/database_schema/ensemble.py create mode 100644 src/ert/dark_storage/database_schema/experiment.py create mode 100644 src/ert/dark_storage/database_schema/metadatafield.py create mode 100644 src/ert/dark_storage/database_schema/observation.py create mode 100644 src/ert/dark_storage/database_schema/prior.py create mode 100644 src/ert/dark_storage/database_schema/record.py create mode 100644 src/ert/dark_storage/database_schema/record_info.py create mode 100644 src/ert/dark_storage/database_schema/update.py create mode 100644 src/ert/dark_storage/endpoints/__init__.py create mode 100644 src/ert/dark_storage/endpoints/compute/__init__.py create mode 100644 src/ert/dark_storage/endpoints/compute/misfits.py create mode 100644 src/ert/dark_storage/endpoints/ensembles.py create mode 100644 src/ert/dark_storage/endpoints/experiments.py create mode 100644 src/ert/dark_storage/endpoints/observations.py create mode 100644 src/ert/dark_storage/endpoints/priors.py create mode 100644 src/ert/dark_storage/endpoints/records.py create mode 100644 src/ert/dark_storage/endpoints/responses.py create mode 100644 src/ert/dark_storage/endpoints/updates.py create mode 100644 src/ert/dark_storage/ext/__init__.py create mode 100644 src/ert/dark_storage/ext/graphene_sqlalchemy.py create mode 100644 src/ert/dark_storage/ext/sqlalchemy_arrays.py create mode 100644 src/ert/dark_storage/ext/uuid.py create mode 100644 src/ert/dark_storage/graphql/__init__.py create mode 100644 src/ert/dark_storage/graphql/ensembles.py create mode 100644 src/ert/dark_storage/graphql/experiments.py create mode 100644 src/ert/dark_storage/graphql/parameters.py create mode 100644 src/ert/dark_storage/graphql/responses.py create mode 100644 src/ert/dark_storage/graphql/unique_responses.py create mode 100644 src/ert/dark_storage/graphql/updates.py create mode 100644 src/ert/dark_storage/json_schema/__init__.py create mode 100644 src/ert/dark_storage/json_schema/ensemble.py create mode 100644 src/ert/dark_storage/json_schema/experiment.py create mode 100644 src/ert/dark_storage/json_schema/observation.py create mode 100644 src/ert/dark_storage/json_schema/prior.py create mode 100644 src/ert/dark_storage/json_schema/record.py create mode 100644 src/ert/dark_storage/json_schema/update.py create mode 100644 src/ert/dark_storage/py.typed create mode 100644 src/ert/dark_storage/security.py create mode 100644 src/ert/dark_storage/testing/__init__.py create mode 100644 src/ert/dark_storage/testing/pytest11.py create mode 100644 src/ert/dark_storage/testing/testclient.py diff --git a/src/ert/dark_storage/__init__.py b/src/ert/dark_storage/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/src/ert/dark_storage/__main__.py b/src/ert/dark_storage/__main__.py new file mode 100644 index 00000000000..2c1327d8795 --- /dev/null +++ b/src/ert/dark_storage/__main__.py @@ -0,0 +1,105 @@ +""" +Start a debug uvicorn server +""" +import os +import sys +import uvicorn +from typing import List, Optional +from tempfile import mkdtemp +from shutil import rmtree + + +def run_server() -> None: + database_dir: Optional[str] = None + if "ERT_STORAGE_DATABASE_URL" not in os.environ: + print( + "Environment variable 'ERT_STORAGE_DATABASE_URL' not set.\n" + "Defaulting to development SQLite temporary database.\n" + "Configure:\n" + "1. File-based SQLite (development):\n" + "\tERT_STORAGE_DATABASE_URL=sqlite:///ert.db\n" + "2. PostgreSQL (production):\n" + "\tERT_STORAGE_DATABASE_URL=postgresql:///:@:/\n", + file=sys.stderr, + ) + database_dir = mkdtemp(prefix="ert-storage_") + os.environ["ERT_STORAGE_DATABASE_URL"] = f"sqlite:///{database_dir}/ert.db" + + if "ERT_STORAGE_AZURE_CONNECTION_STRING" not in os.environ: + print( + "Environment variable 'ERT_STORAGE_AZURE_CONNECTION_STRING' not set.\n" + "Not using Azure Blob Storage. Blob data will be stored in the RDBMS.\n" + ) + + try: + uvicorn.run( + "ert_storage.app:app", reload=True, reload_dirs=[os.path.dirname(__file__)] + ) + finally: + if database_dir is not None: + rmtree(database_dir) + + +def run_alembic(args: List[str]) -> None: + """ + Forward arguments to alembic + """ + from alembic.config import main + + dbkey = "ERT_STORAGE_DATABASE_URL" + dburl = os.getenv(dbkey) + if dburl is None: + sys.exit( + f"Environment variable '{dbkey}' not set.\n" + "It needs to point to a PostgreSQL server for alembic to work." + ) + if not dburl.startswith("postgresql"): + sys.exit( + f"Environment variable '{dbkey}' does not point to a postgresql database.\n" + "Only PostgreSQL is supported for alembic migrations at the moment.\n" + f"Its value is: {dburl}" + ) + + argv = [ + "-c", + os.path.join(os.path.dirname(__file__), "_alembic", "alembic.ini"), + *args, + ] + + try: + main(argv=argv, prog="ert-storage alembic") + except FileNotFoundError as exc: + if os.path.basename(exc.filename) == "script.py.mako": + sys.exit( + f"\nAlembic could not find 'script.py.mako' in location:\n" + f"\n{exc.filename}\n\n" + "This is most likely because you've installed ert-storage without --edit mode\n" + "Reinstall ert-storage with: pip install -e " + ) + else: + raise + sys.exit(0) + + +def print_usage() -> None: + sys.exit( + "Usage: ert-storage [alembic...]\n\n" + "If alembic is given as the first argument, forward the rest of the\n" + "arguments to alembic. Otherwise start ERT Storage in development mode." + ) + + +def main(args: Optional[List[str]] = None) -> None: + if args is None: + args = sys.argv[1:] + + if len(args) > 0: + if args[0] == "alembic": + run_alembic(args[1:]) + else: + print_usage() + run_server() + + +if __name__ == "__main__": + main() diff --git a/src/ert/dark_storage/_alembic/alembic.ini b/src/ert/dark_storage/_alembic/alembic.ini new file mode 100644 index 00000000000..63c8f499024 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic.ini @@ -0,0 +1,87 @@ +# A generic, single database configuration. + +[alembic] +# path to migration scripts +script_location = ert_storage:_alembic/alembic + +# template used to generate migration files +# file_template = %%(rev)s_%%(slug)s + +# sys.path path, will be prepended to sys.path if present. +# defaults to the current working directory. +prepend_sys_path = . + +# timezone to use when rendering the date +# within the migration file as well as the filename. +# string value is passed to dateutil.tz.gettz() +# leave blank for localtime +# timezone = + +# max length of characters to apply to the +# "slug" field +# truncate_slug_length = 40 + +# set to 'true' to run the environment during +# the 'revision' command, regardless of autogenerate +# revision_environment = false + +# set to 'true' to allow .pyc and .pyo files without +# a source .py file to be detected as revisions in the +# versions/ directory +# sourceless = false + +# version location specification; this defaults +# to alembic/versions. When using multiple version +# directories, initial revisions must be specified with --version-path +# version_locations = %(here)s/bar %(here)s/bat alembic/versions + +# the output encoding used when revision files +# are written from script.py.mako +# output_encoding = utf-8 + + +[post_write_hooks] +# post_write_hooks defines scripts or Python functions that are run +# on newly generated revision scripts. See the documentation for further +# detail and examples + +# format using "black" - use the console_scripts runner, against the "black" entrypoint +# hooks=black +# black.type=console_scripts +# black.entrypoint=black +# black.options=-l 79 + +# Logging configuration +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARN +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARN +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/src/ert/dark_storage/_alembic/alembic/README b/src/ert/dark_storage/_alembic/alembic/README new file mode 100644 index 00000000000..98e4f9c44ef --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/README @@ -0,0 +1 @@ +Generic single-database configuration. \ No newline at end of file diff --git a/src/ert/dark_storage/_alembic/alembic/env.py b/src/ert/dark_storage/_alembic/alembic/env.py new file mode 100644 index 00000000000..ffb52010677 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/env.py @@ -0,0 +1,86 @@ +import os +from logging.config import fileConfig + +from sqlalchemy import engine_from_config +from sqlalchemy import pool + +from alembic import context + +from ert_storage.database import ENV_RDBMS +from ert_storage.database_schema import Base + +# this is the Alembic Config object, which provides +# access to the values within the .ini file in use. +config = context.config + +# Interpret the config file for Python logging. +# This line sets up loggers basically. +fileConfig(config.config_file_name) + +# add your model's MetaData object here +# for 'autogenerate' support +# from myapp import mymodel +# target_metadata = mymodel.Base.metadata +target_metadata = Base.metadata + +# other values from the config, defined by the needs of env.py, +# can be acquired: +# my_important_option = config.get_main_option("my_important_option") +# ... etc. + + +def run_migrations_offline(): + """Run migrations in 'offline' mode. + + This configures the context with just a URL + and not an Engine, though an Engine is acceptable + here as well. By skipping the Engine creation + we don't even need a DBAPI to be available. + + Calls to context.execute() here emit the given string to the + script output. + + """ + # Alembic uses sprintf somewhere. Escape the '%'s so that alembic doesn't + # think they're part of the format string. + url = os.environ[ENV_RDBMS].replace("%", "%%") + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online(): + """Run migrations in 'online' mode. + + In this scenario we need to create an Engine + and associate a connection with the context. + + """ + # Alembic uses sprintf somewhere. Escape the '%'s so that alembic doesn't + # think they're part of the format string. + url = os.environ[ENV_RDBMS].replace("%", "%%") + + config.set_section_option(config.config_ini_section, "sqlalchemy.url", str(url)) + connectable = engine_from_config( + config.get_section(config.config_ini_section), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + + with connectable.connect() as connection: + context.configure(connection=connection, target_metadata=target_metadata) + + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/src/ert/dark_storage/_alembic/alembic/script.py.mako b/src/ert/dark_storage/_alembic/alembic/script.py.mako new file mode 100644 index 00000000000..2c0156303a8 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/script.py.mako @@ -0,0 +1,24 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision = ${repr(up_revision)} +down_revision = ${repr(down_revision)} +branch_labels = ${repr(branch_labels)} +depends_on = ${repr(depends_on)} + + +def upgrade(): + ${upgrades if upgrades else "pass"} + + +def downgrade(): + ${downgrades if downgrades else "pass"} diff --git a/src/ert/dark_storage/_alembic/alembic/versions/021d7514e351_add_recordinfo.py b/src/ert/dark_storage/_alembic/alembic/versions/021d7514e351_add_recordinfo.py new file mode 100644 index 00000000000..fa34a3aaae6 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/021d7514e351_add_recordinfo.py @@ -0,0 +1,127 @@ +"""Add RecordInfo + +Revision ID: 021d7514e351 +Revises: 7812741cc469 +Create Date: 2021-05-18 10:09:39.833270 + +""" +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "021d7514e351" +down_revision = "7812741cc469" +branch_labels = None +depends_on = None + + +def upgrade(): + op.execute("ALTER TYPE recordtype RENAME VALUE 'float_vector' TO 'f64_matrix'") + op.drop_column("record", "record_class") + op.execute("DROP TYPE recordclass") + recordtype_postgres = postgresql.ENUM( + "f64_matrix", "file", name="recordtype", create_type=False + ) + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "record_info", + sa.Column("pk", sa.Integer(), nullable=False), + sa.Column( + "time_created", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column( + "time_updated", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column("ensemble_pk", sa.Integer(), nullable=False), + sa.Column("name", sa.String(), nullable=False), + sa.Column( + "record_type", + recordtype_postgres, + nullable=False, + ), + sa.Column( + "record_class", + sa.Enum("parameter", "response", "other", name="recordclass"), + nullable=False, + ), + sa.Column("prior_pk", sa.Integer(), nullable=True), + sa.ForeignKeyConstraint( + ["ensemble_pk"], + ["ensemble.pk"], + ), + sa.ForeignKeyConstraint( + ["prior_pk"], + ["prior.pk"], + ), + sa.PrimaryKeyConstraint("pk"), + sa.UniqueConstraint("name", "ensemble_pk"), + ) + op.add_column( + "ensemble", sa.Column("parameter_names", sa.ARRAY(sa.String()), nullable=False) + ) + op.add_column( + "ensemble", sa.Column("response_names", sa.ARRAY(sa.String()), nullable=False) + ) + op.drop_column("ensemble", "inputs") + op.add_column("record", sa.Column("record_info_pk", sa.Integer(), nullable=True)) + op.drop_constraint("record_prior_pk_fkey", "record", type_="foreignkey") + op.drop_constraint("record_ensemble_id_fkey", "record", type_="foreignkey") + op.create_foreign_key(None, "record", "record_info", ["record_info_pk"], ["pk"]) + op.drop_column("record", "ensemble_pk") + op.drop_column("record", "prior_pk") + op.drop_column("record", "record_type") + op.drop_column("record", "name") + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "record", sa.Column("name", sa.VARCHAR(), autoincrement=False, nullable=False) + ) + op.add_column( + "record", + sa.Column( + "record_class", + postgresql.ENUM("parameter", "response", "other", name="recordclass"), + autoincrement=False, + nullable=True, + ), + ) + op.add_column( + "record", + sa.Column("record_type", sa.INTEGER(), autoincrement=False, nullable=False), + ) + op.add_column( + "record", + sa.Column("prior_pk", sa.INTEGER(), autoincrement=False, nullable=True), + ) + op.add_column( + "record", + sa.Column("ensemble_pk", sa.INTEGER(), autoincrement=False, nullable=True), + ) + op.drop_constraint(None, "record", type_="foreignkey") + op.create_foreign_key( + "record_ensemble_id_fkey", "record", "ensemble", ["ensemble_pk"], ["pk"] + ) + op.create_foreign_key( + "record_prior_pk_fkey", "record", "prior", ["prior_pk"], ["pk"] + ) + op.drop_column("record", "record_info_pk") + op.add_column( + "ensemble", + sa.Column( + "inputs", postgresql.ARRAY(sa.VARCHAR()), autoincrement=False, nullable=True + ), + ) + op.drop_column("ensemble", "response_names") + op.drop_column("ensemble", "parameter_names") + op.drop_table("record_info") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/1a4a83dfb895_add_metadata_column.py b/src/ert/dark_storage/_alembic/alembic/versions/1a4a83dfb895_add_metadata_column.py new file mode 100644 index 00000000000..7f6933887a8 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/1a4a83dfb895_add_metadata_column.py @@ -0,0 +1,34 @@ +"""Add metadata column + +Revision ID: 1a4a83dfb895 +Revises: bea15008ea91 +Create Date: 2021-04-13 10:41:35.187857 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "1a4a83dfb895" +down_revision = "bea15008ea91" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("ensemble", sa.Column("metadata", sa.JSON(), nullable=True)) + op.add_column("experiment", sa.Column("metadata", sa.JSON(), nullable=True)) + op.add_column("observation", sa.Column("metadata", sa.JSON(), nullable=True)) + op.add_column("record", sa.Column("metadata", sa.JSON(), nullable=True)) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("record", "metadata") + op.drop_column("observation", "metadata") + op.drop_column("experiment", "metadata") + op.drop_column("ensemble", "metadata") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/4effed76d1e2_inline_parameters_into_ensembles.py b/src/ert/dark_storage/_alembic/alembic/versions/4effed76d1e2_inline_parameters_into_ensembles.py new file mode 100644 index 00000000000..42305afdd26 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/4effed76d1e2_inline_parameters_into_ensembles.py @@ -0,0 +1,60 @@ +"""Inline parameters into ensembles + +Revision ID: 4effed76d1e2 +Revises: 98601344bcf5 +Create Date: 2021-02-24 16:48:36.601870 + +""" +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "4effed76d1e2" +down_revision = "98601344bcf5" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("parameter") + op.add_column( + "ensemble", sa.Column("parameters", sa.ARRAY(sa.FLOAT()), nullable=False) + ) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("ensemble", "parameters") + op.create_table( + "parameter", + sa.Column("id", sa.INTEGER(), autoincrement=True, nullable=False), + sa.Column( + "time_created", + postgresql.TIMESTAMP(), + server_default=sa.text("now()"), + autoincrement=False, + nullable=True, + ), + sa.Column( + "time_updated", + postgresql.TIMESTAMP(), + server_default=sa.text("now()"), + autoincrement=False, + nullable=True, + ), + sa.Column( + "values", + postgresql.ARRAY(postgresql.DOUBLE_PRECISION(precision=53)), + autoincrement=False, + nullable=False, + ), + sa.Column("ensemble_id", sa.INTEGER(), autoincrement=False, nullable=False), + sa.ForeignKeyConstraint( + ["ensemble_id"], ["ensemble.id"], name="parameter_ensemble_id_fkey" + ), + sa.PrimaryKeyConstraint("id", name="parameter_pkey"), + ) + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/55307b70da13_add_record_class.py b/src/ert/dark_storage/_alembic/alembic/versions/55307b70da13_add_record_class.py new file mode 100644 index 00000000000..2d74c30d36f --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/55307b70da13_add_record_class.py @@ -0,0 +1,44 @@ +"""add record class + +Revision ID: 55307b70da13 +Revises: f46d5afa7639 +Create Date: 2021-04-06 13:15:02.926798 + +""" +from alembic import op +import sqlalchemy as sa +from enum import Enum + +# revision identifiers, used by Alembic. +revision = "55307b70da13" +down_revision = "f46d5afa7639" +branch_labels = None +depends_on = None + + +class RecordClass(Enum): + parameter = 1 + response = 2 + other = 3 + + +def upgrade(): + enum_type = sa.dialects.postgresql.ENUM(RecordClass, name="recordclass") + enum_type.create(op.get_bind()) + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "record", + sa.Column( + "record_class", + sa.Enum(RecordClass), + nullable=True, + ), + ) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("record", "record_class") + op.execute("DROP TYPE recordclass") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/56177419dc39_add_priors.py b/src/ert/dark_storage/_alembic/alembic/versions/56177419dc39_add_priors.py new file mode 100644 index 00000000000..c43d83ff74a --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/56177419dc39_add_priors.py @@ -0,0 +1,77 @@ +"""Add priors + +Revision ID: 56177419dc39 +Revises: 79137bad215d +Create Date: 2021-04-23 17:56:48.270141 + +""" +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "56177419dc39" +down_revision = "79137bad215d" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "prior", + sa.Column("metadata", sa.JSON(), nullable=True), + sa.Column("pk", sa.Integer(), nullable=False), + sa.Column("id", postgresql.UUID(), nullable=True), + sa.Column( + "time_created", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column( + "time_updated", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column("name", sa.String(), nullable=False), + sa.Column( + "function", + sa.Enum( + "const", + "trig", + "normal", + "lognormal", + "truncnormal", + "stdnormal", + "uniform", + "duniform", + "loguniform", + "erf", + "derf", + name="priorfunction", + ), + nullable=False, + ), + sa.Column("argument_names", sa.ARRAY(sa.String), nullable=False), + sa.Column("argument_values", sa.ARRAY(sa.Float), nullable=False), + sa.Column("experiment_pk", sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ["experiment_pk"], + ["experiment.pk"], + ), + sa.PrimaryKeyConstraint("pk"), + sa.UniqueConstraint("id"), + ) + op.add_column("record", sa.Column("prior_pk", sa.Integer(), nullable=True)) + op.create_foreign_key(None, "record", "prior", ["prior_pk"], ["pk"]) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_constraint(None, "record", type_="foreignkey") + op.drop_column("record", "prior_pk") + op.drop_table("prior") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/65569f4f3756_add_name_to_record.py b/src/ert/dark_storage/_alembic/alembic/versions/65569f4f3756_add_name_to_record.py new file mode 100644 index 00000000000..29629bbc272 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/65569f4f3756_add_name_to_record.py @@ -0,0 +1,28 @@ +"""Add name to record + +Revision ID: 65569f4f3756 +Revises: a3f7ae51a72b +Create Date: 2021-02-24 14:52:14.989676 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "65569f4f3756" +down_revision = "a3f7ae51a72b" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("record", sa.Column("name", sa.String(), nullable=False)) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("record", "name") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/73541c7b7aa6_add_size_to_ensemble.py b/src/ert/dark_storage/_alembic/alembic/versions/73541c7b7aa6_add_size_to_ensemble.py new file mode 100644 index 00000000000..7ec6727a581 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/73541c7b7aa6_add_size_to_ensemble.py @@ -0,0 +1,28 @@ +"""Add size to ensemble + +Revision ID: 73541c7b7aa6 +Revises: 79137bad215d +Create Date: 2021-04-27 14:18:53.641564 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "73541c7b7aa6" +down_revision = "c96ec073f9ee" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("ensemble", sa.Column("size", sa.Integer(), nullable=False)) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("ensemble", "size") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/7812741cc469_recordtype_as_enum.py b/src/ert/dark_storage/_alembic/alembic/versions/7812741cc469_recordtype_as_enum.py new file mode 100644 index 00000000000..5d6d8fecf49 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/7812741cc469_recordtype_as_enum.py @@ -0,0 +1,51 @@ +"""RecordType as Enum + +Revision ID: 7812741cc469 +Revises: 73541c7b7aa6 +Create Date: 2021-05-19 12:46:11.497572 + +""" +from alembic import op +import sqlalchemy as sa +from enum import Enum + +# revision identifiers, used by Alembic. +revision = "7812741cc469" +down_revision = "73541c7b7aa6" +branch_labels = None +depends_on = None + + +class RecordType(Enum): + float_vector = 1 + file = 2 + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + enum_type = sa.dialects.postgresql.ENUM(RecordType, name="recordtype") + enum_type.create(op.get_bind()) + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column( + table_name="record", + column_name="record_type", + type_=sa.Enum(RecordType), + postgresql_using=""" +case record_type + when 1 then 'float_vector' + when 2 then 'file' +end :: recordtype; +""", + ) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column( + "record", + "record_type", + type_=sa.Integer, + ) + op.execute("DROP TYPE recordtype") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/79137bad215d_add_matrix_labels.py b/src/ert/dark_storage/_alembic/alembic/versions/79137bad215d_add_matrix_labels.py new file mode 100644 index 00000000000..e590eb88954 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/79137bad215d_add_matrix_labels.py @@ -0,0 +1,38 @@ +"""add matrix labels + +Revision ID: 79137bad215d +Revises: ba7905904381 +Create Date: 2021-04-20 16:03:53.539775 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "79137bad215d" +down_revision = "ba7905904381" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("f64_matrix", sa.Column("labels", sa.PickleType(), nullable=True)) + op.drop_constraint("uq_observation_name", "observation", type_="unique") + op.create_unique_constraint( + "uq_observation_name", "observation", ["name", "experiment_pk"] + ) + op.create_unique_constraint("uq_update_result_pk", "update", ["ensemble_result_pk"]) + op.drop_constraint("uq_update_result_id", "update", type_="unique") + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_unique_constraint("uq_update_result_id", "update", ["ensemble_result_pk"]) + op.drop_constraint("uq_update_result_pk", "update", type_="unique") + op.drop_constraint("uq_observation_name", "observation", type_="unique") + op.create_unique_constraint("uq_observation_name", "observation", ["name"]) + op.drop_column("f64_matrix", "labels") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/896a28a0652f_add_file_block.py b/src/ert/dark_storage/_alembic/alembic/versions/896a28a0652f_add_file_block.py new file mode 100644 index 00000000000..5c14cc7be08 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/896a28a0652f_add_file_block.py @@ -0,0 +1,54 @@ +"""Add file block + +Revision ID: 896a28a0652f +Revises: c1b95721c62a +Create Date: 2021-03-24 11:56:27.012796 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "896a28a0652f" +down_revision = "c1b95721c62a" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "file_block", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column( + "time_created", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column( + "time_updated", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column("block_id", sa.String(), nullable=False), + sa.Column("block_index", sa.Integer(), nullable=False), + sa.Column("record_name", sa.String(), nullable=False), + sa.Column("realization_index", sa.Integer(), nullable=True), + sa.Column("ensemble_id", sa.Integer(), nullable=True), + sa.Column("content", sa.LargeBinary(), nullable=True), + sa.ForeignKeyConstraint( + ["ensemble_id"], + ["ensemble.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("file_block") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/9064db29af73_record_as_inputs_and_outputs_of_.py b/src/ert/dark_storage/_alembic/alembic/versions/9064db29af73_record_as_inputs_and_outputs_of_.py new file mode 100644 index 00000000000..f22e654dffc --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/9064db29af73_record_as_inputs_and_outputs_of_.py @@ -0,0 +1,87 @@ +"""Record as inputs and outputs of ensembles + +Revision ID: 9064db29af73 +Revises: 9a218c7df2ac +Create Date: 2021-03-04 10:20:17.101013 + +""" +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql +from enum import Enum + + +class RecordClass(str, Enum): + normal = "normal" + response = "response" + parameter = "parameter" + + +# revision identifiers, used by Alembic. +revision = "9064db29af73" +down_revision = "9a218c7df2ac" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("ensemble", "parameters") + op.drop_column("ensemble", "num_realizations") + op.add_column("record", sa.Column("consumer_id", sa.Integer(), nullable=True)) + op.add_column("record", sa.Column("producer_id", sa.Integer(), nullable=True)) + enum_type = postgresql.ENUM(RecordClass, name="recordclass") + enum_type.create(op.get_bind()) + op.add_column( + "record", sa.Column("record_class", sa.Enum(RecordClass), nullable=False) + ) + op.drop_constraint( + "record_ensemble_id_realization_index_name_key", "record", type_="unique" + ) + op.drop_constraint("record_ensemble_id_fkey", "record", type_="foreignkey") + op.create_foreign_key(None, "record", "ensemble", ["consumer_id"], ["id"]) + op.create_foreign_key(None, "record", "ensemble", ["producer_id"], ["id"]) + op.drop_column("record", "is_response") + op.drop_column("record", "ensemble_id") + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "record", + sa.Column("ensemble_id", sa.INTEGER(), autoincrement=False, nullable=False), + ) + op.add_column( + "record", + sa.Column("is_response", sa.BOOLEAN(), autoincrement=False, nullable=False), + ) + op.drop_constraint(None, "record", type_="foreignkey") + op.drop_constraint(None, "record", type_="foreignkey") + op.create_foreign_key( + "record_ensemble_id_fkey", "record", "ensemble", ["ensemble_id"], ["id"] + ) + op.create_unique_constraint( + "record_ensemble_id_realization_index_name_key", + "record", + ["ensemble_id", "realization_index", "name"], + ) + op.drop_column("record", "record_class") + op.drop_column("record", "producer_id") + op.drop_column("record", "consumer_id") + op.add_column( + "ensemble", + sa.Column( + "num_realizations", sa.INTEGER(), autoincrement=False, nullable=False + ), + ) + op.add_column( + "ensemble", + sa.Column( + "parameters", + postgresql.ARRAY(postgresql.DOUBLE_PRECISION(precision=53)), + autoincrement=False, + nullable=False, + ), + ) + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/98601344bcf5_add_realization_index_to_record.py b/src/ert/dark_storage/_alembic/alembic/versions/98601344bcf5_add_realization_index_to_record.py new file mode 100644 index 00000000000..80f36468414 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/98601344bcf5_add_realization_index_to_record.py @@ -0,0 +1,28 @@ +"""Add realization_index to record + +Revision ID: 98601344bcf5 +Revises: 65569f4f3756 +Create Date: 2021-02-24 14:53:07.632942 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "98601344bcf5" +down_revision = "65569f4f3756" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("record", sa.Column("realization_index", sa.Integer(), nullable=True)) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("record", "realization_index") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/9a218c7df2ac_separate_into_file_matrix.py b/src/ert/dark_storage/_alembic/alembic/versions/9a218c7df2ac_separate_into_file_matrix.py new file mode 100644 index 00000000000..242132c60b2 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/9a218c7df2ac_separate_into_file_matrix.py @@ -0,0 +1,79 @@ +"""Separate into file, matrix + +Revision ID: 9a218c7df2ac +Revises: fc7aef9b17f8 +Create Date: 2021-02-26 15:17:07.083450 + +""" +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "9a218c7df2ac" +down_revision = "fc7aef9b17f8" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "f64_matrix", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column( + "time_created", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column( + "time_updated", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column("content", sa.ARRAY(sa.FLOAT()), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + op.create_table( + "file", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column( + "time_created", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column( + "time_updated", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column("filename", sa.String(), nullable=False), + sa.Column("mimetype", sa.String(), nullable=False), + sa.Column("content", sa.LargeBinary(), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + op.add_column("record", sa.Column("f64_matrix_id", sa.Integer(), nullable=True)) + op.add_column("record", sa.Column("file_id", sa.Integer(), nullable=True)) + op.create_foreign_key(None, "record", "file", ["file_id"], ["id"]) + op.create_foreign_key(None, "record", "f64_matrix", ["f64_matrix_id"], ["id"]) + op.drop_column("record", "data") + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "record", + sa.Column("data", postgresql.BYTEA(), autoincrement=False, nullable=False), + ) + op.drop_constraint(None, "record", type_="foreignkey") + op.drop_constraint(None, "record", type_="foreignkey") + op.drop_column("record", "file_id") + op.drop_column("record", "f64_matrix_id") + op.drop_table("file") + op.drop_table("f64_matrix") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/a3f7ae51a72b_initial_schema.py b/src/ert/dark_storage/_alembic/alembic/versions/a3f7ae51a72b_initial_schema.py new file mode 100644 index 00000000000..3608051d6e3 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/a3f7ae51a72b_initial_schema.py @@ -0,0 +1,95 @@ +"""Initial schema + +Revision ID: a3f7ae51a72b +Revises: +Create Date: 2021-02-24 14:49:27.508365 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "a3f7ae51a72b" +down_revision = None +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "ensemble", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column( + "time_created", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column( + "time_updated", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column("num_realizations", sa.Integer(), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + op.create_table( + "parameter", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column( + "time_created", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column( + "time_updated", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column("values", sa.ARRAY(sa.FLOAT()), nullable=False), + sa.Column("ensemble_id", sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ["ensemble_id"], + ["ensemble.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_table( + "record", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column( + "time_created", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column( + "time_updated", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column("ensemble_id", sa.Integer(), nullable=False), + sa.Column("record_type", sa.Integer(), nullable=False), + sa.Column("data", sa.PickleType(), nullable=False), + sa.Column("is_response", sa.Boolean(), nullable=False), + sa.ForeignKeyConstraint( + ["ensemble_id"], + ["ensemble.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("record") + op.drop_table("parameter") + op.drop_table("ensemble") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/a88924af33b0_add_experiment.py b/src/ert/dark_storage/_alembic/alembic/versions/a88924af33b0_add_experiment.py new file mode 100644 index 00000000000..1eb3c140e74 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/a88924af33b0_add_experiment.py @@ -0,0 +1,49 @@ +"""Add experiment + +Revision ID: a88924af33b0 +Revises: 896a28a0652f +Create Date: 2021-03-24 12:51:30.807505 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "a88924af33b0" +down_revision = "896a28a0652f" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "experiment", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column( + "time_created", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column( + "time_updated", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column("name", sa.String(), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + op.add_column("ensemble", sa.Column("experiment_id", sa.Integer(), nullable=False)) + op.create_foreign_key(None, "ensemble", "experiment", ["experiment_id"], ["id"]) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_constraint(None, "ensemble", type_="foreignkey") + op.drop_column("ensemble", "experiment_id") + op.drop_table("experiment") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/ba7905904381_add_backend_agnostic_uuid_type.py b/src/ert/dark_storage/_alembic/alembic/versions/ba7905904381_add_backend_agnostic_uuid_type.py new file mode 100644 index 00000000000..c471ec4a9b4 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/ba7905904381_add_backend_agnostic_uuid_type.py @@ -0,0 +1,245 @@ +"""Add backend agnostic UUID type + +Revision ID: ba7905904381 +Revises: 1a4a83dfb895 +Create Date: 2021-04-14 16:54:52.038776 + +""" +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects.postgresql import UUID + + +# revision identifiers, used by Alembic. +revision = "ba7905904381" +down_revision = "1a4a83dfb895" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + + # experiment + op.alter_column("experiment", "id", new_column_name="pk") + op.add_column("experiment", sa.Column("id", UUID(), unique=True, nullable=False)) + + # ensemble + op.alter_column("ensemble", "id", new_column_name="pk") + op.add_column("ensemble", sa.Column("id", UUID(), unique=True, nullable=False)) + op.alter_column( + "ensemble", + "experiment_id", + new_column_name="experiment_pk", + ) + + # f64_matrix + op.alter_column("f64_matrix", "id", new_column_name="pk") + op.add_column("f64_matrix", sa.Column("id", UUID(), unique=True, nullable=False)) + + # file + op.alter_column("file", "id", new_column_name="pk") + op.add_column("file", sa.Column("id", UUID(), unique=True, nullable=False)) + + # file_block + op.alter_column("file_block", "id", new_column_name="pk") + op.add_column("file_block", sa.Column("id", UUID(), unique=True, nullable=False)) + op.alter_column( + "file_block", + "ensemble_id", + new_column_name="ensemble_pk", + ) + + # observation + op.alter_column("observation", "id", new_column_name="pk") + op.add_column("observation", sa.Column("id", UUID(), unique=True, nullable=False)) + op.alter_column( + "observation", + "experiment_id", + new_column_name="experiment_pk", + ) + + # observation_record_assiociation + op.alter_column( + "observation_record_association", + "observation_id", + new_column_name="observation_pk", + ) + op.alter_column( + "observation_record_association", "record_id", new_column_name="record_pk" + ) + + # observation_transformation + op.alter_column( + "observation_transformation", + "id", + new_column_name="pk", + ) + op.add_column( + "observation_transformation", + sa.Column("id", UUID(), unique=True, nullable=False), + ) + op.alter_column( + "observation_transformation", + "observation_id", + new_column_name="observation_pk", + ) + op.alter_column( + "observation_transformation", + "update_id", + new_column_name="update_pk", + ) + + # record + op.alter_column( + "record", + "id", + new_column_name="pk", + ) + op.add_column("record", sa.Column("id", UUID(), unique=True, nullable=False)) + op.alter_column( + "record", + "f64_matrix_id", + new_column_name="f64_matrix_pk", + ) + op.alter_column( + "record", + "file_id", + new_column_name="file_pk", + ) + op.alter_column( + "record", + "ensemble_id", + new_column_name="ensemble_pk", + ) + + # update + op.alter_column( + "update", + "id", + new_column_name="pk", + ) + op.add_column("update", sa.Column("id", UUID(), unique=True, nullable=False)) + op.alter_column( + "update", + "ensemble_reference_id", + new_column_name="ensemble_reference_pk", + ) + op.alter_column( + "update", + "ensemble_result_id", + new_column_name="ensemble_result_pk", + ) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + + # update + op.drop_column("update", "id") + op.alter_column( + "update", + "pk", + new_column_name="id", + ) + op.alter_column( + "update", + "ensemble_reference_pk", + new_column_name="ensemble_reference_id", + ) + op.alter_column( + "update", + "ensemble_result_pk", + new_column_name="ensemble_result_id", + ) + + # record + op.drop_column("record", "id") + op.alter_column( + "record", + "pk", + new_column_name="id", + ) + op.alter_column( + "record", + "f64_matrix_pk", + new_column_name="f64_matrix_id", + ) + op.alter_column( + "record", + "file_pk", + new_column_name="file_id", + ) + op.alter_column( + "record", + "ensemble_pk", + new_column_name="ensemble_id", + ) + + # observation_transformation + op.drop_column("observation_transformation", "id") + op.alter_column( + "observation_transformation", + "pk", + new_column_name="id", + ) + op.alter_column( + "observation_transformation", + "observation_pk", + new_column_name="observation_id", + ) + op.alter_column( + "observation_transformation", + "update_pk", + new_column_name="update_id", + ) + + # observation_record_association + op.alter_column( + "observation_record_association", + "observation_pk", + new_column_name="observation_id", + ) + op.alter_column( + "observation_record_association", "record_pk", new_column_name="record_id" + ) + + # observation + op.drop_column("observation", "id") + op.alter_column("observation", "pk", new_column_name="id") + op.alter_column( + "observation", + "experiment_pk", + new_column_name="experiment_id", + ) + + # file_block + op.drop_column("file_block", "id") + op.alter_column("file_block", "pk", new_column_name="id") + op.alter_column( + "file_block", + "ensemble_pk", + new_column_name="ensemble_id", + ) + + # file + op.drop_column("file", "id") + op.alter_column("file", "pk", new_column_name="id") + + # f64_matrix + op.drop_column("f64_matrix", "id") + op.alter_column("f64_matrix", "pk", new_column_name="id") + + # ensemble + op.drop_column("ensemble", "id") + op.alter_column("ensemble", "pk", new_column_name="id") + op.alter_column( + "ensemble", + "experiment_pk", + new_column_name="experiment_id", + ) + # experiment + op.drop_column("experiment", "id") + op.alter_column("experiment", "pk", new_column_name="id") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/bea15008ea91_add_obs_trans_and_updates.py b/src/ert/dark_storage/_alembic/alembic/versions/bea15008ea91_add_obs_trans_and_updates.py new file mode 100644 index 00000000000..967a81271a1 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/bea15008ea91_add_obs_trans_and_updates.py @@ -0,0 +1,62 @@ +"""Add obs trans and updates + +Revision ID: bea15008ea91 +Revises: 55307b70da13 +Create Date: 2021-04-08 15:57:25.088170 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "bea15008ea91" +down_revision = "55307b70da13" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "update", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("algorithm", sa.String(), nullable=False), + sa.Column("ensemble_reference_id", sa.Integer(), nullable=True), + sa.Column("ensemble_result_id", sa.Integer(), nullable=True), + sa.ForeignKeyConstraint( + ["ensemble_reference_id"], + ["ensemble.id"], + ), + sa.ForeignKeyConstraint( + ["ensemble_result_id"], + ["ensemble.id"], + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("ensemble_result_id", name="uq_update_result_id"), + ) + op.create_table( + "observation_transformation", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("active_list", sa.PickleType(), nullable=False), + sa.Column("scale_list", sa.PickleType(), nullable=False), + sa.Column("observation_id", sa.Integer(), nullable=False), + sa.Column("update_id", sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ["observation_id"], + ["observation.id"], + ), + sa.ForeignKeyConstraint( + ["update_id"], + ["update.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("observation_transformation") + op.drop_table("update") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/c1b95721c62a_add_azure_blob_storage.py b/src/ert/dark_storage/_alembic/alembic/versions/c1b95721c62a_add_azure_blob_storage.py new file mode 100644 index 00000000000..e7206c291f1 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/c1b95721c62a_add_azure_blob_storage.py @@ -0,0 +1,32 @@ +"""Add azure blob storage + +Revision ID: c1b95721c62a +Revises: d72e51c6bdf2 +Create Date: 2021-03-10 14:30:59.473787 + +""" +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "c1b95721c62a" +down_revision = "d72e51c6bdf2" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("file", sa.Column("az_blob", sa.String(), nullable=True)) + op.add_column("file", sa.Column("az_container", sa.String(), nullable=True)) + op.alter_column("file", "content", existing_type=postgresql.BYTEA(), nullable=True) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column("file", "content", existing_type=postgresql.BYTEA(), nullable=False) + op.drop_column("file", "az_container") + op.drop_column("file", "az_blob") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/c96ec073f9ee_change_prior_enum.py b/src/ert/dark_storage/_alembic/alembic/versions/c96ec073f9ee_change_prior_enum.py new file mode 100644 index 00000000000..822696768a5 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/c96ec073f9ee_change_prior_enum.py @@ -0,0 +1,54 @@ +"""Change Prior Enum + +Revision ID: c96ec073f9ee +Revises: 56177419dc39 +Create Date: 2021-04-28 09:05:08.830307 + +""" +from alembic import op +import sqlalchemy as sa +from enum import Enum + + +# revision identifiers, used by Alembic. +revision = "c96ec073f9ee" +down_revision = "56177419dc39" +branch_labels = None +depends_on = None + + +class PriorFunction(Enum): + const = 1 + trig = 2 + normal = 3 + lognormal = 4 + ert_truncnormal = 5 + stdnormal = 6 + uniform = 7 + ert_duniform = 8 + loguniform = 9 + ert_erf = 10 + ert_derf = 11 + + +def upgrade(): + op.drop_column("prior", "function") + op.execute("DROP TYPE priorfunction") + + enum_type = sa.dialects.postgresql.ENUM(PriorFunction, name="priorfunction") + enum_type.create(op.get_bind()) + + op.add_column( + "prior", + sa.Column( + "function", + sa.Enum( + PriorFunction, + ), + nullable=False, + ), + ) + + +def downgrade(): + raise NotImplementedError("Downgrade not implemented") diff --git a/src/ert/dark_storage/_alembic/alembic/versions/d72e51c6bdf2_remove_record_class.py b/src/ert/dark_storage/_alembic/alembic/versions/d72e51c6bdf2_remove_record_class.py new file mode 100644 index 00000000000..33a8ffd6b35 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/d72e51c6bdf2_remove_record_class.py @@ -0,0 +1,61 @@ +"""Remove record class + +Revision ID: d72e51c6bdf2 +Revises: 9064db29af73 +Create Date: 2021-03-10 10:34:45.613538 + +""" +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "d72e51c6bdf2" +down_revision = "9064db29af73" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("ensemble", sa.Column("inputs", sa.ARRAY(sa.String()), nullable=True)) + op.add_column("record", sa.Column("ensemble_id", sa.Integer(), nullable=True)) + op.drop_constraint("record_consumer_id_fkey", "record", type_="foreignkey") + op.drop_constraint("record_producer_id_fkey", "record", type_="foreignkey") + op.create_foreign_key(None, "record", "ensemble", ["ensemble_id"], ["id"]) + op.drop_column("record", "consumer_id") + op.drop_column("record", "producer_id") + op.drop_column("record", "record_class") + op.execute("DROP TYPE recordclass") + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "record", + sa.Column( + "record_class", + postgresql.ENUM("normal", "response", "parameter", name="recordclass"), + autoincrement=False, + nullable=False, + ), + ) + op.add_column( + "record", + sa.Column("producer_id", sa.INTEGER(), autoincrement=False, nullable=True), + ) + op.add_column( + "record", + sa.Column("consumer_id", sa.INTEGER(), autoincrement=False, nullable=True), + ) + op.drop_constraint(None, "record", type_="foreignkey") + op.create_foreign_key( + "record_producer_id_fkey", "record", "ensemble", ["producer_id"], ["id"] + ) + op.create_foreign_key( + "record_consumer_id_fkey", "record", "ensemble", ["consumer_id"], ["id"] + ) + op.drop_column("record", "ensemble_id") + op.drop_column("ensemble", "inputs") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/f46d5afa7639_add_observations.py b/src/ert/dark_storage/_alembic/alembic/versions/f46d5afa7639_add_observations.py new file mode 100644 index 00000000000..fc4522f54a4 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/f46d5afa7639_add_observations.py @@ -0,0 +1,68 @@ +"""add observations + +Revision ID: f46d5afa7639 +Revises: a88924af33b0 +Create Date: 2021-04-06 09:58:45.107064 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "f46d5afa7639" +down_revision = "a88924af33b0" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "observation", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column( + "time_created", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column( + "time_updated", + sa.DateTime(), + server_default=sa.text("now()"), + nullable=True, + ), + sa.Column("name", sa.String(), nullable=False), + sa.Column("x_axis", sa.ARRAY(sa.String()), nullable=False), + sa.Column("values", sa.ARRAY(sa.FLOAT()), nullable=False), + sa.Column("errors", sa.ARRAY(sa.FLOAT()), nullable=False), + sa.Column("experiment_id", sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ["experiment_id"], + ["experiment.id"], + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("name", name="uq_observation_name"), + ) + op.create_table( + "observation_record_association", + sa.Column("observation_id", sa.Integer(), nullable=True), + sa.Column("record_id", sa.Integer(), nullable=True), + sa.ForeignKeyConstraint( + ["observation_id"], + ["observation.id"], + ), + sa.ForeignKeyConstraint( + ["record_id"], + ["record.id"], + ), + ) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("observation_record_association") + op.drop_table("observation") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/_alembic/alembic/versions/fc7aef9b17f8_add_record_unique_constraint.py b/src/ert/dark_storage/_alembic/alembic/versions/fc7aef9b17f8_add_record_unique_constraint.py new file mode 100644 index 00000000000..3a7f1ca6706 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/fc7aef9b17f8_add_record_unique_constraint.py @@ -0,0 +1,30 @@ +"""Add record unique constraint + +Revision ID: fc7aef9b17f8 +Revises: 4effed76d1e2 +Create Date: 2021-02-24 17:05:51.046173 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "fc7aef9b17f8" +down_revision = "4effed76d1e2" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_unique_constraint( + None, "record", ["ensemble_id", "realization_index", "name"] + ) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_constraint(None, "record", type_="unique") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/app.py b/src/ert/dark_storage/app.py new file mode 100644 index 00000000000..aa0b9a2912f --- /dev/null +++ b/src/ert/dark_storage/app.py @@ -0,0 +1,73 @@ +import json +from typing import Any +from fastapi import FastAPI, Request, status +from fastapi.responses import Response, RedirectResponse +from starlette.graphql import GraphQLApp + +from ert_storage.endpoints import router as endpoints_router +from ert_storage.graphql import router as graphql_router + +from sqlalchemy.orm.exc import NoResultFound + + +class JSONResponse(Response): + """A replacement for Starlette's JSONResponse that permits NaNs.""" + + media_type = "application/json" + + def render(self, content: Any) -> bytes: + return json.dumps( + content, + ensure_ascii=False, + allow_nan=True, + indent=None, + separators=(",", ":"), + ).encode("utf-8") + + +app = FastAPI( + title="ERT Storage API", + version="0.1.2", + debug=True, + default_response_class=JSONResponse, +) + + +@app.on_event("startup") +async def initialize_database() -> None: + from ert_storage.database import engine, IS_SQLITE, HAS_AZURE_BLOB_STORAGE + from ert_storage.database_schema import Base + + if IS_SQLITE: + # Our SQLite backend doesn't support migrations, so create the database on the fly. + Base.metadata.create_all(bind=engine) + if HAS_AZURE_BLOB_STORAGE: + from ert_storage.database import create_container_if_not_exist + + await create_container_if_not_exist() + + +@app.exception_handler(NoResultFound) +async def sqlalchemy_exception_handler( + request: Request, exc: NoResultFound +) -> JSONResponse: + """Automatically catch and convert an SQLAlchemy NoResultFound exception (when + using `.one()`, for example) to an HTTP 404 message + """ + return JSONResponse( + {"detail": "Item not found"}, status_code=status.HTTP_404_NOT_FOUND + ) + + +@app.get("/") +async def root() -> RedirectResponse: + return RedirectResponse("/docs") + + +@app.get("/healthcheck") +async def healthcheck() -> str: + return "ALL OK!" + + +app.include_router(endpoints_router) +app.include_router(graphql_router) diff --git a/src/ert/dark_storage/compute/__init__.py b/src/ert/dark_storage/compute/__init__.py new file mode 100644 index 00000000000..80c0c5f8cf7 --- /dev/null +++ b/src/ert/dark_storage/compute/__init__.py @@ -0,0 +1 @@ +from .misfits import calculate_misfits_from_pandas diff --git a/src/ert/dark_storage/compute/misfits.py b/src/ert/dark_storage/compute/misfits.py new file mode 100644 index 00000000000..bd198da4d4d --- /dev/null +++ b/src/ert/dark_storage/compute/misfits.py @@ -0,0 +1,35 @@ +import numpy as np +import pandas as pd +from uuid import UUID +from typing import Any, Mapping, List, Optional + + +def _calculate_misfit( + obs_value: np.ndarray, response_value: np.ndarray, obs_std: np.ndarray +) -> List[float]: + difference = response_value - obs_value + misfit = (difference / obs_std) ** 2 + return (misfit * np.sign(difference)).tolist() + + +def calculate_misfits_from_pandas( + reponses_dict: Mapping[int, pd.DataFrame], + observation: pd.DataFrame, + summary_misfits: bool = False, +) -> pd.DataFrame: + """ + Compute misfits from reponses_dict (real_id, values in dataframe) + and observation + """ + misfits_dict = {} + for realization_index in reponses_dict: + misfits_dict[realization_index] = _calculate_misfit( + observation["values"], + reponses_dict[realization_index].loc[:, observation.index].values.flatten(), + observation["errors"], + ) + + df = pd.DataFrame(data=misfits_dict, index=observation.index) + if summary_misfits: + df = pd.DataFrame([df.abs().sum(axis=0)], columns=df.columns, index=[0]) + return df.T diff --git a/src/ert/dark_storage/database.py b/src/ert/dark_storage/database.py new file mode 100644 index 00000000000..9ff6c7f8056 --- /dev/null +++ b/src/ert/dark_storage/database.py @@ -0,0 +1,66 @@ +import os +import sys +from typing import Any, Callable, Type +from fastapi import Depends +from sqlalchemy import create_engine +from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy.orm import sessionmaker +from ert_storage.security import security + + +ENV_RDBMS = "ERT_STORAGE_DATABASE_URL" +ENV_BLOB = "ERT_STORAGE_AZURE_CONNECTION_STRING" +ENV_BLOB_CONTAINER = "ERT_STORAGE_AZURE_BLOB_CONTAINER" + +if ENV_RDBMS not in os.environ: + sys.exit(f"Environment variable '{ENV_RDBMS}' not set") + + +URI_RDBMS = os.environ[ENV_RDBMS] +IS_SQLITE = URI_RDBMS.startswith("sqlite") +IS_POSTGRES = URI_RDBMS.startswith("postgres") +HAS_AZURE_BLOB_STORAGE = ENV_BLOB in os.environ +BLOB_CONTAINER = os.getenv(ENV_BLOB_CONTAINER, "ert") + + +if IS_SQLITE: + engine = create_engine(URI_RDBMS, connect_args={"check_same_thread": False}) +else: + engine = create_engine(URI_RDBMS, pool_size=50, max_overflow=100) +Session = sessionmaker(autocommit=False, autoflush=False, bind=engine) + + +Base = declarative_base() + + +async def get_db(*, _: None = Depends(security)) -> Any: + db = Session() + + # Make PostgreSQL return float8 columns with highest precision. If we don't + # do this, we may lose up to 3 of the least significant digits. + if IS_POSTGRES: + db.execute("SET extra_float_digits=3") + try: + yield db + db.commit() + db.close() + except: + db.rollback() + db.close() + raise + + +if HAS_AZURE_BLOB_STORAGE: + import asyncio + from azure.core.exceptions import ResourceNotFoundError + from azure.storage.blob.aio import ContainerClient + + azure_blob_container = ContainerClient.from_connection_string( + os.environ[ENV_BLOB], BLOB_CONTAINER + ) + + async def create_container_if_not_exist() -> None: + try: + await azure_blob_container.get_container_properties() + except ResourceNotFoundError: + await azure_blob_container.create_container() diff --git a/src/ert/dark_storage/database_schema/__init__.py b/src/ert/dark_storage/database_schema/__init__.py new file mode 100644 index 00000000000..213082c8a94 --- /dev/null +++ b/src/ert/dark_storage/database_schema/__init__.py @@ -0,0 +1,9 @@ +from .metadatafield import MetadataField +from .record_info import RecordInfo, RecordType, RecordClass +from .record import Record, F64Matrix, File, FileBlock +from .ensemble import Ensemble +from .experiment import Experiment +from .observation import Observation, ObservationTransformation +from .update import Update +from .prior import Prior, PriorFunction +from ert_storage.database import Base diff --git a/src/ert/dark_storage/database_schema/ensemble.py b/src/ert/dark_storage/database_schema/ensemble.py new file mode 100644 index 00000000000..e8c28786f01 --- /dev/null +++ b/src/ert/dark_storage/database_schema/ensemble.py @@ -0,0 +1,51 @@ +from typing import List +from uuid import uuid4, UUID as PyUUID +import sqlalchemy as sa +from sqlalchemy.orm import relationship +from sqlalchemy.sql import func +from ert_storage.database import Base +from .metadatafield import MetadataField +from ert_storage.ext.uuid import UUID as UUID +from ert_storage.ext.sqlalchemy_arrays import StringArray + + +class Ensemble(Base, MetadataField): + __tablename__ = "ensemble" + + pk = sa.Column(sa.Integer, primary_key=True) + id = sa.Column(UUID, unique=True, default=uuid4, nullable=False) + size = sa.Column(sa.Integer, nullable=False) + time_created = sa.Column(sa.DateTime, server_default=func.now()) + time_updated = sa.Column( + sa.DateTime, server_default=func.now(), onupdate=func.now() + ) + parameter_names = sa.Column(StringArray, nullable=False) + response_names = sa.Column(StringArray, nullable=False) + record_infos = relationship( + "RecordInfo", + foreign_keys="[RecordInfo.ensemble_pk]", + cascade="all, delete-orphan", + lazy="dynamic", + ) + experiment_pk = sa.Column( + sa.Integer, sa.ForeignKey("experiment.pk"), nullable=False + ) + experiment = relationship("Experiment", back_populates="ensembles") + children = relationship( + "Update", + foreign_keys="[Update.ensemble_reference_pk]", + ) + parent = relationship( + "Update", + uselist=False, + foreign_keys="[Update.ensemble_result_pk]", + cascade="all, delete-orphan", + ) + + @property + def parent_ensemble_id(self) -> PyUUID: + return self.parent.ensemble_reference.id + + @property + def child_ensemble_ids(self) -> List[PyUUID]: + return [x.ensemble_result.id for x in self.children] diff --git a/src/ert/dark_storage/database_schema/experiment.py b/src/ert/dark_storage/database_schema/experiment.py new file mode 100644 index 00000000000..3f8e871d318 --- /dev/null +++ b/src/ert/dark_storage/database_schema/experiment.py @@ -0,0 +1,36 @@ +import sqlalchemy as sa +from uuid import uuid4 +from sqlalchemy.sql import func +from sqlalchemy.orm import relationship +from ert_storage.database import Base +from ert_storage.ext.uuid import UUID +from .metadatafield import MetadataField + + +class Experiment(Base, MetadataField): + __tablename__ = "experiment" + + pk = sa.Column(sa.Integer, primary_key=True) + id = sa.Column(UUID, unique=True, default=uuid4, nullable=False) + time_created = sa.Column(sa.DateTime, server_default=func.now()) + time_updated = sa.Column( + sa.DateTime, server_default=func.now(), onupdate=func.now() + ) + name = sa.Column(sa.String) + ensembles = relationship( + "Ensemble", + foreign_keys="[Ensemble.experiment_pk]", + cascade="all, delete-orphan", + ) + observations = relationship( + "Observation", + foreign_keys="[Observation.experiment_pk]", + cascade="all, delete-orphan", + back_populates="experiment", + ) + priors = relationship( + "Prior", + foreign_keys="[Prior.experiment_pk]", + cascade="all, delete-orphan", + back_populates="experiment", + ) diff --git a/src/ert/dark_storage/database_schema/metadatafield.py b/src/ert/dark_storage/database_schema/metadatafield.py new file mode 100644 index 00000000000..aff5786e673 --- /dev/null +++ b/src/ert/dark_storage/database_schema/metadatafield.py @@ -0,0 +1,12 @@ +from typing import Any, Mapping +import sqlalchemy as sa + + +class MetadataField: + _metadata = sa.Column("metadata", sa.JSON, nullable=True) + + @property + def metadata_dict(self) -> Mapping[str, Any]: + if self._metadata is None: + return dict() + return self._metadata diff --git a/src/ert/dark_storage/database_schema/observation.py b/src/ert/dark_storage/database_schema/observation.py new file mode 100644 index 00000000000..04ade5dfa72 --- /dev/null +++ b/src/ert/dark_storage/database_schema/observation.py @@ -0,0 +1,60 @@ +import sqlalchemy as sa +from ert_storage.database import Base +from ert_storage.ext.sqlalchemy_arrays import StringArray, FloatArray +from ert_storage.ext.uuid import UUID +from sqlalchemy.sql import func +from sqlalchemy.orm import relationship +from uuid import uuid4 +from .metadatafield import MetadataField + +observation_record_association = sa.Table( + "observation_record_association", + Base.metadata, + sa.Column("observation_pk", sa.Integer, sa.ForeignKey("observation.pk")), + sa.Column("record_pk", sa.Integer, sa.ForeignKey("record.pk")), +) + + +class Observation(Base, MetadataField): + __tablename__ = "observation" + __table_args__ = ( + sa.UniqueConstraint("name", "experiment_pk", name="uq_observation_name"), + ) + + pk = sa.Column(sa.Integer, primary_key=True) + id = sa.Column(UUID, unique=True, default=uuid4, nullable=False) + time_created = sa.Column(sa.DateTime, server_default=func.now()) + time_updated = sa.Column( + sa.DateTime, server_default=func.now(), onupdate=func.now() + ) + name = sa.Column(sa.String, nullable=False) + x_axis = sa.Column(StringArray, nullable=False) + values = sa.Column(FloatArray, nullable=False) + errors = sa.Column(FloatArray, nullable=False) + + records = relationship( + "Record", + secondary=observation_record_association, + back_populates="observations", + ) + experiment_pk = sa.Column( + sa.Integer, sa.ForeignKey("experiment.pk"), nullable=False + ) + experiment = relationship("Experiment") + + +class ObservationTransformation(Base): + __tablename__ = "observation_transformation" + + pk = sa.Column(sa.Integer, primary_key=True) + id = sa.Column(UUID, unique=True, default=uuid4, nullable=False) + active_list = sa.Column(sa.PickleType, nullable=False) + scale_list = sa.Column(sa.PickleType, nullable=False) + + observation_pk = sa.Column( + sa.Integer, sa.ForeignKey("observation.pk"), nullable=False + ) + observation = relationship("Observation") + + update_pk = sa.Column(sa.Integer, sa.ForeignKey("update.pk"), nullable=False) + update = relationship("Update", back_populates="observation_transformations") diff --git a/src/ert/dark_storage/database_schema/prior.py b/src/ert/dark_storage/database_schema/prior.py new file mode 100644 index 00000000000..781c45eacde --- /dev/null +++ b/src/ert/dark_storage/database_schema/prior.py @@ -0,0 +1,44 @@ +from enum import Enum +import sqlalchemy as sa +from sqlalchemy.sql import func +from sqlalchemy.orm import relationship +from uuid import uuid4 +from .metadatafield import MetadataField +from ert_storage.database import Base +from ert_storage.ext.uuid import UUID +from ert_storage.ext.sqlalchemy_arrays import StringArray, FloatArray + + +class PriorFunction(Enum): + const = 1 + trig = 2 + normal = 3 + lognormal = 4 + ert_truncnormal = 5 + stdnormal = 6 + uniform = 7 + ert_duniform = 8 + loguniform = 9 + ert_erf = 10 + ert_derf = 11 + + +class Prior(Base, MetadataField): + __tablename__ = "prior" + + pk = sa.Column(sa.Integer, primary_key=True) + id = sa.Column(UUID, unique=True, default=uuid4) + time_created = sa.Column(sa.DateTime, server_default=func.now()) + time_updated = sa.Column( + sa.DateTime, server_default=func.now(), onupdate=func.now() + ) + + name = sa.Column(sa.String, nullable=False) + function = sa.Column(sa.Enum(PriorFunction), nullable=False) + argument_names = sa.Column(StringArray, nullable=False) + argument_values = sa.Column(FloatArray, nullable=False) + + experiment_pk = sa.Column( + sa.Integer, sa.ForeignKey("experiment.pk"), nullable=False + ) + experiment = relationship("Experiment") diff --git a/src/ert/dark_storage/database_schema/record.py b/src/ert/dark_storage/database_schema/record.py new file mode 100644 index 00000000000..0a581b838fc --- /dev/null +++ b/src/ert/dark_storage/database_schema/record.py @@ -0,0 +1,110 @@ +from typing import Any +from uuid import uuid4 + +import sqlalchemy as sa +from sqlalchemy.sql import func +from sqlalchemy.orm import relationship +from sqlalchemy.ext.hybrid import hybrid_property + +from ert_storage.ext.sqlalchemy_arrays import FloatArray +from ert_storage.ext.uuid import UUID +from ert_storage.database import Base + +from .metadatafield import MetadataField +from .observation import observation_record_association +from .record_info import RecordType + + +class Record(Base, MetadataField): + __tablename__ = "record" + + pk = sa.Column(sa.Integer, primary_key=True) + id = sa.Column(UUID, unique=True, default=uuid4, nullable=False) + time_created = sa.Column(sa.DateTime, server_default=func.now()) + time_updated = sa.Column( + sa.DateTime, server_default=func.now(), onupdate=func.now() + ) + + realization_index = sa.Column(sa.Integer, nullable=True) + + record_info_pk = sa.Column( + sa.Integer, sa.ForeignKey("record_info.pk"), nullable=True + ) + record_info = relationship("RecordInfo", back_populates="records") + + file_pk = sa.Column(sa.Integer, sa.ForeignKey("file.pk")) + f64_matrix_pk = sa.Column(sa.Integer, sa.ForeignKey("f64_matrix.pk")) + + file = relationship("File", cascade="all") + f64_matrix = relationship("F64Matrix", cascade="all") + + observations = relationship( + "Observation", + secondary=observation_record_association, + back_populates="records", + ) + + @property + def data(self) -> Any: + info = self.record_info + if info.record_type == RecordType.file: + return self.file.content + elif info.record_type == RecordType.f64_matrix: + return self.f64_matrix.content + else: + raise NotImplementedError( + f"The record type {self.record_type} is not yet implemented" + ) + + @property + def name(self) -> str: + return self.record_info.name + + +class File(Base): + __tablename__ = "file" + + pk = sa.Column(sa.Integer, primary_key=True) + id = sa.Column(UUID, unique=True, default=uuid4, nullable=False) + time_created = sa.Column(sa.DateTime, server_default=func.now()) + time_updated = sa.Column( + sa.DateTime, server_default=func.now(), onupdate=func.now() + ) + + filename = sa.Column(sa.String, nullable=False) + mimetype = sa.Column(sa.String, nullable=False) + + content = sa.Column(sa.LargeBinary) + az_container = sa.Column(sa.String) + az_blob = sa.Column(sa.String) + + +class F64Matrix(Base): + __tablename__ = "f64_matrix" + + pk = sa.Column(sa.Integer, primary_key=True) + id = sa.Column(UUID, unique=True, default=uuid4, nullable=False) + time_created = sa.Column(sa.DateTime, server_default=func.now()) + time_updated = sa.Column( + sa.DateTime, server_default=func.now(), onupdate=func.now() + ) + content = sa.Column(FloatArray, nullable=False) + labels = sa.Column(sa.PickleType) + + +class FileBlock(Base): + __tablename__ = "file_block" + + pk = sa.Column(sa.Integer, primary_key=True) + id = sa.Column(UUID, unique=True, default=uuid4, nullable=False) + time_created = sa.Column(sa.DateTime, server_default=func.now()) + time_updated = sa.Column( + sa.DateTime, server_default=func.now(), onupdate=func.now() + ) + block_id = sa.Column(sa.String, nullable=False) + block_index = sa.Column(sa.Integer, nullable=False) + record_name = sa.Column(sa.String, nullable=False) + realization_index = sa.Column(sa.Integer, nullable=True) + ensemble_pk = sa.Column(sa.Integer, sa.ForeignKey("ensemble.pk"), nullable=True) + ensemble = relationship("Ensemble") + content = sa.Column(sa.LargeBinary, nullable=True) diff --git a/src/ert/dark_storage/database_schema/record_info.py b/src/ert/dark_storage/database_schema/record_info.py new file mode 100644 index 00000000000..fdccc5d428d --- /dev/null +++ b/src/ert/dark_storage/database_schema/record_info.py @@ -0,0 +1,46 @@ +from enum import Enum + +import sqlalchemy as sa +from sqlalchemy.orm import relationship +from sqlalchemy.sql import func + +from ert_storage.database import Base + + +class RecordType(Enum): + f64_matrix = 1 + file = 2 + + +class RecordClass(Enum): + parameter = 1 + response = 2 + other = 3 + + +class RecordInfo(Base): + __tablename__ = "record_info" + __table_args__ = (sa.UniqueConstraint("name", "ensemble_pk"),) + + pk = sa.Column(sa.Integer, primary_key=True) + time_created = sa.Column(sa.DateTime, server_default=func.now()) + time_updated = sa.Column( + sa.DateTime, server_default=func.now(), onupdate=func.now() + ) + + ensemble_pk = sa.Column(sa.Integer, sa.ForeignKey("ensemble.pk"), nullable=False) + ensemble = relationship("Ensemble") + + records = relationship( + "Record", + foreign_keys="[Record.record_info_pk]", + cascade="all, delete-orphan", + ) + + name = sa.Column(sa.String, nullable=False) + record_type = sa.Column(sa.Enum(RecordType), nullable=False) + record_class = sa.Column(sa.Enum(RecordClass), nullable=False) + + # Parameter-specific data + prior_pk = sa.Column(sa.Integer, sa.ForeignKey("prior.pk"), nullable=True) + prior = relationship("Prior") diff --git a/src/ert/dark_storage/database_schema/update.py b/src/ert/dark_storage/database_schema/update.py new file mode 100644 index 00000000000..01b27a33583 --- /dev/null +++ b/src/ert/dark_storage/database_schema/update.py @@ -0,0 +1,39 @@ +from ert_storage.database import Base +import sqlalchemy as sa +from sqlalchemy.orm import relationship +from uuid import uuid4 +from ert_storage.ext.uuid import UUID + + +class Update(Base): + __tablename__ = "update" + __table_args__ = ( + sa.UniqueConstraint("ensemble_result_pk", name="uq_update_result_pk"), + ) + + pk = sa.Column(sa.Integer, primary_key=True) + id = sa.Column(UUID, unique=True, default=uuid4, nullable=False) + algorithm = sa.Column(sa.String, nullable=False) + ensemble_reference_pk = sa.Column( + sa.Integer, sa.ForeignKey("ensemble.pk"), nullable=True + ) + ensemble_result_pk = sa.Column( + sa.Integer, sa.ForeignKey("ensemble.pk"), nullable=True + ) + + ensemble_reference = relationship( + "Ensemble", + foreign_keys=[ensemble_reference_pk], + back_populates="children", + ) + ensemble_result = relationship( + "Ensemble", + foreign_keys=[ensemble_result_pk], + uselist=False, + back_populates="parent", + ) + observation_transformations = relationship( + "ObservationTransformation", + foreign_keys="[ObservationTransformation.update_pk]", + cascade="all, delete-orphan", + ) diff --git a/src/ert/dark_storage/endpoints/__init__.py b/src/ert/dark_storage/endpoints/__init__.py new file mode 100644 index 00000000000..64d4ef11e08 --- /dev/null +++ b/src/ert/dark_storage/endpoints/__init__.py @@ -0,0 +1,19 @@ +from fastapi import APIRouter +from .ensembles import router as ensembles_router +from .records import router as records_router +from .experiments import router as experiments_router +from .observations import router as observations_router +from .updates import router as updates_router +from .priors import router as priors_router +from .compute.misfits import router as misfits_router +from .responses import router as response_router + +router = APIRouter() +router.include_router(experiments_router) +router.include_router(ensembles_router) +router.include_router(records_router) +router.include_router(observations_router) +router.include_router(updates_router) +router.include_router(priors_router) +router.include_router(misfits_router) +router.include_router(response_router) diff --git a/src/ert/dark_storage/endpoints/compute/__init__.py b/src/ert/dark_storage/endpoints/compute/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/src/ert/dark_storage/endpoints/compute/misfits.py b/src/ert/dark_storage/endpoints/compute/misfits.py new file mode 100644 index 00000000000..463cb14c963 --- /dev/null +++ b/src/ert/dark_storage/endpoints/compute/misfits.py @@ -0,0 +1,90 @@ +from ert_storage.database_schema.record import F64Matrix +import numpy as np +import pandas as pd +from uuid import UUID +from typing import Any, Optional, List +import sqlalchemy as sa +from fastapi.responses import Response +from fastapi import APIRouter, Depends, HTTPException, status +from ert_storage.database import Session, get_db +from ert_storage import database_schema as ds, json_schema as js +from ert_storage.compute import calculate_misfits_from_pandas, misfits + +router = APIRouter(tags=["misfits"]) + + +@router.get( + "/compute/misfits", + responses={ + status.HTTP_200_OK: { + "content": {"application/x-dataframe": {}}, + "description": "Return misfits as csv, where columns are realizations.", + } + }, +) +async def get_response_misfits( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + response_name: str, + realization_index: Optional[int] = None, + summary_misfits: bool = False, +) -> Response: + """ + Compute univariate misfits for response(s) + """ + + response_query = ( + db.query(ds.Record) + .filter(ds.Record.observations != None) + .join(ds.RecordInfo) + .filter_by( + name=response_name, + record_type=ds.RecordType.f64_matrix, + ) + .join(ds.Ensemble) + .filter_by(id=ensemble_id) + ) + if realization_index is not None: + responses = [ + response_query.filter( + ds.Record.realization_index == realization_index + ).one() + ] + else: + responses = response_query.all() + + observation_df = None + response_dict = {} + for response in responses: + data_df = pd.DataFrame(response.f64_matrix.content) + labels = response.f64_matrix.labels + if labels is not None: + data_df.columns = labels[0] + data_df.index = labels[1] + response_dict[response.realization_index] = data_df + if observation_df is None: + # currently we expect only a single observation object, while + # later in the future this might change + obs = response.observations[0] + observation_df = pd.DataFrame( + data={"values": obs.values, "errors": obs.errors}, index=obs.x_axis + ) + + try: + result_df = calculate_misfits_from_pandas( + response_dict, observation_df, summary_misfits + ) + except Exception as misfits_exc: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={ + "error": f"Unable to compute misfits: {misfits_exc}", + "name": response_name, + "ensemble_id": str(ensemble_id), + }, + ) + return Response( + content=result_df.to_csv().encode(), + media_type="text/csv", + ) diff --git a/src/ert/dark_storage/endpoints/ensembles.py b/src/ert/dark_storage/endpoints/ensembles.py new file mode 100644 index 00000000000..e20e2d9244f --- /dev/null +++ b/src/ert/dark_storage/endpoints/ensembles.py @@ -0,0 +1,81 @@ +from uuid import UUID + +from fastapi import APIRouter, Depends, Body +from sqlalchemy.orm.attributes import flag_modified +from ert_storage.database import Session, get_db +from ert_storage import database_schema as ds, json_schema as js +from typing import Any, Mapping + +router = APIRouter(tags=["ensemble"]) + + +@router.post("/experiments/{experiment_id}/ensembles", response_model=js.EnsembleOut) +def post_ensemble( + *, db: Session = Depends(get_db), ens_in: js.EnsembleIn, experiment_id: UUID +) -> ds.Ensemble: + + experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() + ens = ds.Ensemble( + parameter_names=ens_in.parameter_names, + response_names=ens_in.response_names, + experiment=experiment, + size=ens_in.size, + _metadata=ens_in.metadata, + ) + db.add(ens) + + if ens_in.update_id: + update_obj = db.query(ds.Update).filter_by(id=ens_in.update_id).one() + update_obj.ensemble_result = ens + db.commit() + + return ens + + +@router.get("/ensembles/{ensemble_id}", response_model=js.EnsembleOut) +def get_ensemble(*, db: Session = Depends(get_db), ensemble_id: UUID) -> ds.Ensemble: + return db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + + +@router.put("/ensembles/{ensemble_id}/metadata") +async def replace_ensemble_metadata( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + body: Any = Body(...), +) -> None: + """ + Assign new metadata json + """ + ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + ensemble._metadata = body + db.commit() + + +@router.patch("/ensembles/{ensemble_id}/metadata") +async def patch_ensemble_metadata( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + body: Any = Body(...), +) -> None: + """ + Update metadata json + """ + ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + ensemble._metadata.update(body) + flag_modified(ensemble, "_metadata") + db.commit() + + +@router.get("/ensembles/{ensemble_id}/metadata", response_model=Mapping[str, Any]) +async def get_ensemble_metadata( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, +) -> Mapping[str, Any]: + """ + Get metadata json + """ + ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + return ensemble.metadata_dict diff --git a/src/ert/dark_storage/endpoints/experiments.py b/src/ert/dark_storage/endpoints/experiments.py new file mode 100644 index 00000000000..3286bf280df --- /dev/null +++ b/src/ert/dark_storage/endpoints/experiments.py @@ -0,0 +1,117 @@ +from uuid import UUID +from fastapi import APIRouter, Depends, Body +from sqlalchemy.orm.attributes import flag_modified +from ert_storage.database import Session, get_db +from ert_storage import database_schema as ds, json_schema as js +from typing import Any, Mapping, Optional, List + + +router = APIRouter(tags=["experiment"]) + + +@router.get("/experiments", response_model=List[js.ExperimentOut]) +def get_experiments( + *, + db: Session = Depends(get_db), +) -> List[js.ExperimentOut]: + return [ + js.ExperimentOut( + id=exp.id, + name=exp.name, + ensembles=[ens.id for ens in exp.ensembles], + priors={pri.name: pri.id for pri in exp.priors}, + metadata=exp.metadata_dict, + ) + for exp in db.query(ds.Experiment).all() + ] + + +@router.post("/experiments", response_model=js.ExperimentOut) +def post_experiments( + *, + db: Session = Depends(get_db), + ens_in: js.ExperimentIn, +) -> js.ExperimentOut: + experiment = ds.Experiment(name=ens_in.name) + + if ens_in.priors: + db.add_all( + ds.Prior( + function=ds.PriorFunction.__members__[prior.function], + experiment=experiment, + name=name, + argument_names=[x[0] for x in prior if isinstance(x[1], (float, int))], + argument_values=[x[1] for x in prior if isinstance(x[1], (float, int))], + ) + for name, prior in ens_in.priors.items() + ) + + db.add(experiment) + db.commit() + return js.ExperimentOut( + id=experiment.id, + name=experiment.name, + ensembles=[ens.id for ens in experiment.ensembles], + priors={pri.name: pri.id for pri in experiment.priors}, + metadata=experiment.metadata_dict, + ) + + +@router.get( + "/experiments/{experiment_id}/ensembles", response_model=List[js.EnsembleOut] +) +def get_experiment_ensembles( + *, db: Session = Depends(get_db), experiment_id: UUID +) -> List[ds.Ensemble]: + return db.query(ds.Ensemble).join(ds.Experiment).filter_by(id=experiment_id).all() + + +@router.put("/experiments/{experiment_id}/metadata") +async def replace_experiment_metadata( + *, + db: Session = Depends(get_db), + experiment_id: UUID, + body: Any = Body(...), +) -> None: + """ + Assign new metadata json + """ + experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() + experiment._metadata = body + db.commit() + + +@router.patch("/experiments/{experiment_id}/metadata") +async def patch_experiment_metadata( + *, + db: Session = Depends(get_db), + experiment_id: UUID, + body: Any = Body(...), +) -> None: + """ + Update metadata json + """ + experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() + experiment._metadata.update(body) + flag_modified(experiment, "_metadata") + db.commit() + + +@router.get("/experiments/{experiment_id}/metadata", response_model=Mapping[str, Any]) +async def get_experiment_metadata( + *, + db: Session = Depends(get_db), + experiment_id: UUID, +) -> Mapping[str, Any]: + """ + Get metadata json + """ + experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() + return experiment.metadata_dict + + +@router.delete("/experiments/{experiment_id}") +def delete_experiment(*, db: Session = Depends(get_db), experiment_id: UUID) -> None: + experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() + db.delete(experiment) + db.commit() diff --git a/src/ert/dark_storage/endpoints/observations.py b/src/ert/dark_storage/endpoints/observations.py new file mode 100644 index 00000000000..55f78253cbb --- /dev/null +++ b/src/ert/dark_storage/endpoints/observations.py @@ -0,0 +1,148 @@ +from uuid import UUID + +from fastapi import APIRouter, Depends, Body +from typing import List, Any, Mapping +from sqlalchemy.orm.attributes import flag_modified +from ert_storage.database import Session, get_db +from ert_storage import database_schema as ds, json_schema as js + + +router = APIRouter(tags=["ensemble"]) + + +@router.post( + "/experiments/{experiment_id}/observations", response_model=js.ObservationOut +) +def post_observation( + *, db: Session = Depends(get_db), obs_in: js.ObservationIn, experiment_id: UUID +) -> js.ObservationOut: + experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() + records = ( + [db.query(ds.Record).filter_by(id=rec_id).one() for rec_id in obs_in.records] + if obs_in.records is not None + else [] + ) + obs = ds.Observation( + name=obs_in.name, + x_axis=obs_in.x_axis, + errors=obs_in.errors, + values=obs_in.values, + experiment=experiment, + records=records, + ) + + db.add(obs) + db.commit() + + return js.ObservationOut( + id=obs.id, + name=obs.name, + x_axis=obs.x_axis, + errors=obs.errors, + values=obs.values, + records=[rec.id for rec in obs.records], + metadata=obs.metadata_dict, + ) + + +@router.get( + "/experiments/{experiment_id}/observations", response_model=List[js.ObservationOut] +) +def get_observations( + *, db: Session = Depends(get_db), experiment_id: UUID +) -> List[js.ObservationOut]: + experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() + return [ + js.ObservationOut( + id=obs.id, + name=obs.name, + x_axis=obs.x_axis, + errors=obs.errors, + values=obs.values, + records=[rec.id for rec in obs.records], + metadata=obs.metadata_dict, + ) + for obs in experiment.observations + ] + + +@router.get( + "/ensembles/{ensemble_id}/observations", response_model=List[js.ObservationOut] +) +def get_observations_with_transformation( + *, db: Session = Depends(get_db), ensemble_id: UUID +) -> List[js.ObservationOut]: + ens = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + experiment = ens.experiment + update = ens.parent + transformations = ( + {trans.observation.name: trans for trans in update.observation_transformations} + if update is not None + else {} + ) + + return [ + js.ObservationOut( + id=obs.id, + name=obs.name, + x_axis=obs.x_axis, + errors=obs.errors, + values=obs.values, + records=[rec.id for rec in obs.records], + metadata=obs.metadata_dict, + transformation=js.ObservationTransformationOut( + id=transformations[obs.name].id, + name=obs.name, + observation_id=obs.id, + scale=transformations[obs.name].scale_list, + active=transformations[obs.name].active_list, + ) + if obs.name in transformations + else None, + ) + for obs in experiment.observations + ] + + +@router.put("/observations/{obs_id}/metadata") +async def replace_observation_metadata( + *, + db: Session = Depends(get_db), + obs_id: UUID, + body: Any = Body(...), +) -> None: + """ + Assign new metadata json + """ + obs = db.query(ds.Observation).filter_by(id=obs_id).one() + obs._metadata = body + db.commit() + + +@router.patch("/observations/{obs_id}/metadata") +async def patch_observation_metadata( + *, + db: Session = Depends(get_db), + obs_id: UUID, + body: Any = Body(...), +) -> None: + """ + Update metadata json + """ + obs = db.query(ds.Observation).filter_by(id=obs_id).one() + obs._metadata.update(body) + flag_modified(obs, "_metadata") + db.commit() + + +@router.get("/observations/{obs_id}/metadata", response_model=Mapping[str, Any]) +async def get_observation_metadata( + *, + db: Session = Depends(get_db), + obs_id: UUID, +) -> Mapping[str, Any]: + """ + Get metadata json + """ + obs = db.query(ds.Observation).filter_by(id=obs_id).one() + return obs.metadata_dict diff --git a/src/ert/dark_storage/endpoints/priors.py b/src/ert/dark_storage/endpoints/priors.py new file mode 100644 index 00000000000..4cc6a023fc0 --- /dev/null +++ b/src/ert/dark_storage/endpoints/priors.py @@ -0,0 +1,57 @@ +from uuid import UUID +from typing import Mapping, Type +from fastapi import APIRouter, Depends +from ert_storage.database import Session, get_db +from ert_storage import database_schema as ds, json_schema as js +from ert_storage.json_schema.prior import ( + PriorConst, + PriorTrig, + PriorNormal, + PriorLogNormal, + PriorErtTruncNormal, + PriorStdNormal, + PriorUniform, + PriorErtDUniform, + PriorLogUniform, + PriorErtErf, + PriorErtDErf, +) + +router = APIRouter(tags=["prior"]) + + +@router.get("/experiments/{experiment_id}/priors") +def get_priors( + *, db: Session = Depends(get_db), experiment_id: UUID +) -> Mapping[str, dict]: + experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() + return experiment_priors_to_dict(experiment) + + +PRIOR_FUNCTION_TO_PYDANTIC: Mapping[ds.PriorFunction, Type[js.Prior]] = { + ds.PriorFunction.const: PriorConst, + ds.PriorFunction.trig: PriorTrig, + ds.PriorFunction.normal: PriorNormal, + ds.PriorFunction.lognormal: PriorLogNormal, + ds.PriorFunction.ert_truncnormal: PriorErtTruncNormal, + ds.PriorFunction.stdnormal: PriorStdNormal, + ds.PriorFunction.uniform: PriorUniform, + ds.PriorFunction.ert_duniform: PriorErtDUniform, + ds.PriorFunction.loguniform: PriorLogUniform, + ds.PriorFunction.ert_erf: PriorErtErf, + ds.PriorFunction.ert_derf: PriorErtDErf, +} + + +def prior_to_dict(prior: ds.Prior) -> dict: + return ( + PRIOR_FUNCTION_TO_PYDANTIC[prior.function] + .parse_obj( + {key: val for key, val in zip(prior.argument_names, prior.argument_values)} + ) + .dict() + ) + + +def experiment_priors_to_dict(experiment: ds.Experiment) -> Mapping[str, dict]: + return {p.name: prior_to_dict(p) for p in experiment.priors} diff --git a/src/ert/dark_storage/endpoints/records.py b/src/ert/dark_storage/endpoints/records.py new file mode 100644 index 00000000000..089befe0820 --- /dev/null +++ b/src/ert/dark_storage/endpoints/records.py @@ -0,0 +1,856 @@ +from uuid import uuid4, UUID +import io +import numpy as np +import pandas as pd +from enum import Enum +from typing import Any, Mapping, Optional, List, AsyncGenerator +import sqlalchemy as sa +from fastapi import ( + APIRouter, + Body, + Depends, + File, + HTTPException, + Header, + Request, + UploadFile, + status, +) +from fastapi.responses import Response, StreamingResponse +from pydantic import BaseModel +from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm.exc import NoResultFound +from sqlalchemy.orm.attributes import flag_modified +from ert_storage.database import Session, get_db, HAS_AZURE_BLOB_STORAGE, BLOB_CONTAINER +from ert_storage import database_schema as ds, json_schema as js + +from fastapi.logger import logger + +if HAS_AZURE_BLOB_STORAGE: + from ert_storage.database import azure_blob_container + + +router = APIRouter(tags=["record"]) + + +class ListRecords(BaseModel): + ensemble: Mapping[str, str] + forward_model: Mapping[str, str] + + +@router.post("/ensembles/{ensemble_id}/records/{name}/file") +async def post_ensemble_record_file( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + name: str, + realization_index: Optional[int] = None, + file: UploadFile = File(...), +) -> None: + """ + Assign an arbitrary file to the given `name` record. + """ + ensemble = _get_and_assert_ensemble(db, ensemble_id, name, realization_index) + + file_obj = ds.File( + filename=file.filename, + mimetype=file.content_type, + ) + if HAS_AZURE_BLOB_STORAGE: + key = f"{name}@{realization_index}@{uuid4()}" + blob = azure_blob_container.get_blob_client(key) + await blob.upload_blob(file.file) + + file_obj.az_container = azure_blob_container.container_name + file_obj.az_blob = key + else: + file_obj.content = await file.read() + + db.add(file_obj) + _create_record( + db, + ensemble, + name, + ds.RecordType.file, + realization_index=realization_index, + file=file_obj, + ) + + +@router.put("/ensembles/{ensemble_id}/records/{name}/blob") +async def add_block( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + name: str, + realization_index: Optional[int] = None, + request: Request, + block_index: int, +) -> None: + """ + Stage blocks to an existing azure blob record. + """ + + ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + block_id = str(uuid4()) + + file_block_obj = ds.FileBlock( + ensemble=ensemble, + block_id=block_id, + block_index=block_index, + record_name=name, + realization_index=realization_index, + ) + + record_obj = ( + db.query(ds.Record) + .filter_by(realization_index=realization_index) + .join(ds.RecordInfo) + .filter_by(ensemble_pk=ensemble.pk, name=name) + .one() + ) + if HAS_AZURE_BLOB_STORAGE: + key = record_obj.file.az_blob + blob = azure_blob_container.get_blob_client(key) + await blob.stage_block(block_id, await request.body()) + else: + file_block_obj.content = await request.body() + + db.add(file_block_obj) + db.commit() + + +@router.post("/ensembles/{ensemble_id}/records/{name}/blob") +async def create_blob( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + name: str, + realization_index: Optional[int] = None, +) -> None: + """ + Create a record which points to a blob on Azure Blob Storage. + """ + + ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + file_obj = ds.File( + filename="test", + mimetype="mime/type", + ) + if HAS_AZURE_BLOB_STORAGE: + key = f"{name}@{realization_index}@{uuid4()}" + blob = azure_blob_container.get_blob_client(key) + file_obj.az_container = (azure_blob_container.container_name,) + file_obj.az_blob = (key,) + else: + pass + + db.add(file_obj) + + _create_record( + db, + ensemble, + name, + ds.RecordType.file, + realization_index=realization_index, + file=file_obj, + ) + + +@router.patch("/ensembles/{ensemble_id}/records/{name}/blob") +async def finalize_blob( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + name: str, + realization_index: Optional[int] = None, +) -> None: + """ + Commit all staged blocks to a blob record + """ + + ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + + record_obj = ( + db.query(ds.Record) + .filter_by(realization_index=realization_index) + .join(ds.RecordInfo) + .filter_by(ensemble_pk=ensemble.pk, name=name) + .one() + ) + + submitted_blocks = list( + db.query(ds.FileBlock) + .filter_by( + record_name=name, + ensemble_pk=ensemble.pk, + realization_index=realization_index, + ) + .all() + ) + + if HAS_AZURE_BLOB_STORAGE: + key = record_obj.file.az_blob + blob = azure_blob_container.get_blob_client(key) + block_ids = [ + block.block_id + for block in sorted(submitted_blocks, key=lambda x: x.block_index) + ] + await blob.commit_block_list(block_ids) + else: + data = b"".join([block.content for block in submitted_blocks]) + record_obj.file.content = data + + +@router.post( + "/ensembles/{ensemble_id}/records/{name}/matrix", response_model=js.RecordOut +) +async def post_ensemble_record_matrix( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + name: str, + realization_index: Optional[int] = None, + content_type: str = Header("application/json"), + prior_id: Optional[UUID] = None, + request: Request, +) -> js.RecordOut: + """ + Assign an n-dimensional float matrix, encoded in JSON, to the given `name` record. + """ + if content_type == "application/x-dataframe": + logger.warning( + "Content-Type with 'application/x-dataframe' is deprecated. Use 'text/csv' instead." + ) + content_type = "text/csv" + + ensemble = _get_and_assert_ensemble(db, ensemble_id, name, realization_index) + is_parameter = name in ensemble.parameter_names + is_response = name in ensemble.response_names + if prior_id is not None and not is_parameter: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={ + "error": "Priors can only be specified for parameter records", + "name": name, + "ensemble_id": str(ensemble_id), + "realization_index": realization_index, + "prior_id": str(prior_id), + }, + ) + + labels = None + prior = ( + ( + db.query(ds.Prior) + .filter_by(id=prior_id, experiment_pk=ensemble.experiment_pk) + .one() + ) + if prior_id + else None + ) + + try: + if content_type == "application/json": + content = np.array(await request.json(), dtype=np.float64) + elif content_type == "application/x-numpy": + from numpy.lib.format import read_array + + stream = io.BytesIO(await request.body()) + content = read_array(stream) + elif content_type == "text/csv": + stream = io.BytesIO(await request.body()) + df = pd.read_csv(stream, index_col=0, float_precision="round_trip") + content = df.values + labels = [ + [str(v) for v in df.columns.values], + [str(v) for v in df.index.values], + ] + else: + raise ValueError() + except ValueError: + if realization_index is None: + message = f"Ensemble-wide record '{name}' for ensemble '{ensemble_id}' needs to be a matrix" + else: + message = f"Forward-model record '{name}' for ensemble '{ensemble_id}', realization {realization_index} needs to be a matrix" + + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={ + "error": message, + "name": name, + "ensemble_id": str(ensemble_id), + "realization_index": realization_index, + }, + ) + + # Require that the dimensionality of an ensemble-wide parameter matrix is at least 2 + if realization_index is None and is_parameter: + if content.ndim <= 1: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={ + "error": "Ensemble-wide parameter record '{name}' for ensemble '{ensemble_id}'" + "must have dimensionality of at least 2", + "name": name, + "ensemble_id": str(ensemble_id), + "realization_index": realization_index, + }, + ) + + matrix_obj = ds.F64Matrix(content=content.tolist(), labels=labels) + db.add(matrix_obj) + + record_class = ds.RecordClass.other + if is_parameter: + record_class = ds.RecordClass.parameter + if is_response: + record_class = ds.RecordClass.response + + return _create_record( + db, + ensemble, + name, + ds.RecordType.f64_matrix, + record_class, + prior, + f64_matrix=matrix_obj, + realization_index=realization_index, + ) + + +@router.put("/ensembles/{ensemble_id}/records/{name}/metadata") +async def replace_record_metadata( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + name: str, + realization_index: Optional[int] = None, + body: Any = Body(...), +) -> None: + """ + Assign new metadata json + """ + if realization_index is None: + record_obj = _get_ensemble_record(db, ensemble_id, name) + else: + record_obj = _get_forward_model_record(db, ensemble_id, name, realization_index) + record_obj._metadata = body + db.commit() + + +@router.patch("/ensembles/{ensemble_id}/records/{name}/metadata") +async def patch_record_metadata( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + name: str, + realization_index: Optional[int] = None, + body: Any = Body(...), +) -> None: + """ + Update metadata json + """ + if realization_index is None: + record_obj = _get_ensemble_record(db, ensemble_id, name) + else: + record_obj = _get_forward_model_record(db, ensemble_id, name, realization_index) + + record_obj._metadata.update(body) + flag_modified(record_obj, "_metadata") + db.commit() + + +@router.get( + "/ensembles/{ensemble_id}/records/{name}/metadata", response_model=Mapping[str, Any] +) +async def get_record_metadata( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + name: str, + realization_index: Optional[int] = None, +) -> Mapping[str, Any]: + """ + Get metadata json + """ + if realization_index is None: + bundle = _get_ensemble_record(db, ensemble_id, name) + else: + bundle = _get_forward_model_record(db, ensemble_id, name, realization_index) + + return bundle.metadata_dict + + +@router.post("/ensembles/{ensemble_id}/records/{name}/observations") +async def post_record_observations( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + name: str, + realization_index: Optional[int] = None, + observation_ids: List[UUID] = Body(...), +) -> None: + + if realization_index is None: + record_obj = _get_ensemble_record(db, ensemble_id, name) + else: + record_obj = _get_forward_model_record(db, ensemble_id, name, realization_index) + + observations = ( + db.query(ds.Observation).filter(ds.Observation.id.in_(observation_ids)).all() + ) + if observations: + record_obj.observations = observations + flag_modified(record_obj, "observations") + db.commit() + else: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={ + "error": f"Observations {observation_ids} not found!", + "name": name, + "ensemble_id": str(ensemble_id), + }, + ) + + +@router.get("/ensembles/{ensemble_id}/records/{name}/observations") +async def get_record_observations( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + name: str, + realization_index: Optional[int] = None, +) -> List[js.ObservationOut]: + if realization_index is None: + bundle = _get_ensemble_record(db, ensemble_id, name) + else: + bundle = _get_forward_model_record(db, ensemble_id, name, realization_index) + if bundle.observations: + return [ + js.ObservationOut( + id=obs.id, + name=obs.name, + x_axis=obs.x_axis, + errors=obs.errors, + values=obs.values, + records=[rec.id for rec in obs.records], + metadata=obs.metadata_dict, + ) + for obs in bundle.observations + ] + return [] + + +GET_CSV_NO_LABELS_DESCRIPTION = """\ +If the matrix data contained no label information, the returned CSV uses +integer ranges as the column and row labels. + +If the record is a parameter matrix, the rows are the different realizations. +""" + + +GET_CSV_ALL_LABELS_DESCRIPTION = """\ +Returned CSV contains strings for column and row labels. The data itself is numeric. + +If the record is a parameter matrix, the rows are the different realizations. +""" + + +GET_NUMPY_DESCRIPTION = """\ +Data encoded with `numpy.lib.format.write_array` and `numpy.lib.format.read_array`. + +To parse data using Python, assuming ERT Storage is running on `http://localhost:8000` : + +```python + import io + import requests + from numpy.lib.format import read_array + + resp = requests.get( + "http://localhost:8000/ensembles/{ENSEMBLE_ID}/records/{RECORD_NAME}", + headers={"Accept": "application/x-numpy"} + ) + stream = io.BytesIO(resp.content) + nparray = read_array(stream) + + # Print the numpy array + print(nparray) +``` +""" + + +@router.get( + "/ensembles/{ensemble_id}/records/{name}", + responses={ + 200: { + "description": "Successful fetch of record", + "content": { + "application/json": { + "examples": { + "no-labels": { + "value": [[11.5, 12.5, 13.5], [21.5, 22.5, 23.5]], + } + } + }, + "text/csv": { + "examples": { + "no-labels": { + "value": ",\t0,\t1,\t2\n" + "0,\t11.5,\t12.5,\t13.5\n" + "1,\t21.5,\t22.5,\t23.5\n", + "description": GET_CSV_NO_LABELS_DESCRIPTION, + }, + "all-labels": { + "value": ",\t2010-01-01,\t2015-01-01,\t2020-01-01\n" + "real_0,\t11.5,\t\t12.5,\t\t13.5\n" + "real_1,\t21.5,\t\t22.5,\t\t23.5\n", + "description": GET_CSV_ALL_LABELS_DESCRIPTION, + }, + } + }, + "application/x-numpy": { + "examples": { + "success": { + "summary": "Fetch data encoded in NPY array format", + "description": GET_NUMPY_DESCRIPTION, + } + } + }, + }, + } + }, +) +async def get_ensemble_record( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + realization_index: Optional[int] = None, + name: str, + accept: str = Header("application/json"), +) -> Any: + """ + Get record with a given `name`. If `realization_index` is not set, look for + the ensemble-wide record. If it is set, look first for one created by a + forward-model for the given realization index and then the ensemble-wide + record. + + Records support multiple data formats. In particular: + - Matrix: + Will return n-dimensional float matrix, where n is arbitrary. + - File: + Will return the file that was uploaded. + """ + if accept == "application/x-dataframe": + logger.warning( + "Accept with 'application/x-dataframe' is deprecated. Use 'text/csv' instead." + ) + accept = "text/csv" + + if realization_index is None: + bundle = _get_ensemble_record(db, ensemble_id, name) + else: + try: + bundle = _get_forward_model_record(db, ensemble_id, name, realization_index) + realization_index = None + except HTTPException as exc: + bundle = _get_ensemble_record(db, ensemble_id, name) + + # Only parameter records can be "up-casted" to ensemble-wide + is_matrix = bundle.record_info.record_type == ds.RecordType.f64_matrix + is_parameter = name in bundle.record_info.ensemble.parameter_names + if not is_matrix or not is_parameter: + raise exc + return await _get_record_data(bundle, accept, realization_index) + + +@router.get("/ensembles/{ensemble_id}/parameters", response_model=List[str]) +async def get_ensemble_parameters( + *, db: Session = Depends(get_db), ensemble_id: UUID +) -> List[str]: + ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + return ensemble.parameter_names + + +@router.get( + "/ensembles/{ensemble_id}/records", response_model=Mapping[str, js.RecordOut] +) +async def get_ensemble_records( + *, db: Session = Depends(get_db), ensemble_id: UUID +) -> Mapping[str, ds.Record]: + return { + rec.name: rec + for rec in ( + db.query(ds.Record) + .join(ds.RecordInfo) + .join(ds.Ensemble) + .filter_by(id=ensemble_id) + .all() + ) + } + + +@router.get("/records/{record_id}", response_model=js.RecordOut) +async def get_record(*, db: Session = Depends(get_db), record_id: UUID) -> ds.Record: + return db.query(ds.Record).filter_by(id=record_id).one() + + +@router.get("/records/{record_id}/data") +async def get_record_data( + *, + db: Session = Depends(get_db), + record_id: UUID, + accept: Optional[str] = Header(default="application/json"), +) -> Any: + if accept == "application/x-dataframe": + logger.warning( + "Accept with 'application/x-dataframe' is deprecated. Use 'text/csv' instead." + ) + accept = "text/csv" + + record = db.query(ds.Record).filter_by(id=record_id).one() + return await _get_record_data(record, accept) + + +@router.get( + "/ensembles/{ensemble_id}/responses", response_model=Mapping[str, js.RecordOut] +) +def get_ensemble_responses( + *, db: Session = Depends(get_db), ensemble_id: UUID +) -> Mapping[str, ds.Record]: + return { + rec.name: rec + for rec in ( + db.query(ds.Record) + .join(ds.RecordInfo) + .filter_by(record_class=ds.RecordClass.response) + .join(ds.Ensemble) + .filter_by(id=ensemble_id) + .all() + ) + } + + +def _get_ensemble_record(db: Session, ensemble_id: UUID, name: str) -> ds.Record: + try: + ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + return ( + db.query(ds.Record) + .filter_by(realization_index=None) + .join(ds.RecordInfo) + .filter_by( + ensemble_pk=ensemble.pk, + name=name, + ) + .one() + ) + except NoResultFound: + raise HTTPException( + status_code=404, + detail={ + "error": f"Ensemble-wide record '{name}' for ensemble '{ensemble_id}' not found!", + "name": name, + "ensemble_id": str(ensemble_id), + }, + ) + + +def _get_forward_model_record( + db: Session, ensemble_id: UUID, name: str, realization_index: int +) -> ds.Record: + try: + ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + return ( + db.query(ds.Record) + .filter_by(realization_index=realization_index) + .join(ds.RecordInfo) + .filter_by( + ensemble_pk=ensemble.pk, + name=name, + ) + .one() + ) + except NoResultFound: + raise HTTPException( + status_code=404, + detail={ + "error": f"Forward-model record '{name}' for ensemble '{ensemble_id}', realization {realization_index} not found!", + "name": name, + "ensemble_id": str(ensemble_id), + }, + ) + + +def _get_and_assert_ensemble( + db: Session, ensemble_id: UUID, name: str, realization_index: Optional[int] +) -> ds.Ensemble: + """ + Get ensemble and verify that no record with the name `name` exists + """ + ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + + q = ( + db.query(ds.Record) + .join(ds.RecordInfo) + .filter_by(ensemble_pk=ensemble.pk, name=name) + ) + if realization_index is not None: + if realization_index not in range(ensemble.size) and ensemble.size != -1: + raise HTTPException( + status_code=status.HTTP_417_EXPECTATION_FAILED, + detail={ + "error": f"Ensemble '{name}' ('{ensemble_id}') does have a 'size' " + f"of {ensemble.size}. The posted record is targeting " + f"'realization_index' {realization_index} which is out " + f"of bounds.", + "name": name, + "ensemble_id": str(ensemble_id), + }, + ) + + q = q.filter( + (ds.Record.realization_index == None) + | (ds.Record.realization_index == realization_index) + ) + + if q.count() > 0: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "error": f"Ensemble-wide record '{name}' for ensemble '{ensemble_id}' already exists", + "name": name, + "ensemble_id": str(ensemble_id), + }, + ) + + return ensemble + + +async def _get_record_data( + record: ds.Record, accept: Optional[str], realization_index: Optional[int] = None +) -> Response: + type_ = record.record_info.record_type + if type_ == ds.RecordType.f64_matrix: + if realization_index is None: + content = record.f64_matrix.content + else: + content = record.f64_matrix.content[realization_index] + + if accept == "application/x-numpy": + from numpy.lib.format import write_array + + stream = io.BytesIO() + write_array(stream, np.array(content)) + + return Response( + content=stream.getvalue(), + media_type="application/x-numpy", + ) + if accept == "text/csv": + data = pd.DataFrame(content) + labels = record.f64_matrix.labels + if labels is not None and realization_index is None: + data.columns = labels[0] + data.index = labels[1] + elif labels is not None and realization_index is not None: + # The output is such that rows are realizations. Because + # `content` is a 1d list in this case, it treats each element as + # its own row. We transpose the data so that all of the data + # falls on the same row. + data = data.T + data.columns = labels[0] + data.index = [realization_index] + + return Response( + content=data.to_csv().encode(), + media_type="text/csv", + ) + else: + return content + if type_ == ds.RecordType.file: + f = record.file + if f.content is not None: + return Response( + content=f.content, + media_type=f.mimetype, + headers={"Content-Disposition": f'attachment; filename="{f.filename}"'}, + ) + elif f.az_container is not None and f.az_blob is not None: + blob = azure_blob_container.get_blob_client(f.az_blob) + download = await blob.download_blob() + + async def chunk_generator() -> AsyncGenerator[bytes, None]: + async for chunk in download.chunks(): + yield chunk + + return StreamingResponse( + chunk_generator(), + media_type=f.mimetype, + headers={"Content-Disposition": f'attachment; filename="{f.filename}"'}, + ) + raise NotImplementedError( + f"Getting record data for type {type_} and Accept header {accept} not implemented" + ) + + +def _create_record( + db: Session, + ensemble: ds.Ensemble, + name: str, + record_type: ds.RecordType, + record_class: ds.RecordClass = ds.RecordClass.other, + prior: Optional[ds.Prior] = None, + **kwargs: Any, +) -> ds.Record: + record_info = ds.RecordInfo( + ensemble=ensemble, + name=name, + record_class=record_class, + record_type=record_type, + prior=prior, + ) + record = ds.Record(record_info=record_info, **kwargs) + + nested = db.begin_nested() + try: + db.add(record) + db.commit() + except IntegrityError: + # Assuming this is a UNIQUE constraint failure due to an existing + # record_info with the same name and ensemble. Try to fetch the + # record_info + nested.rollback() + old_record_info = ( + db.query(ds.RecordInfo).filter_by(ensemble=ensemble, name=name).one() + ) + + # Check that the parameters match + if record_info.record_class != old_record_info.record_class: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "error": "Record class of new record does not match previous record class", + "new_record_class": str(record_info.record_class), + "old_record_class": str(old_record_info.record_class), + "name": name, + "ensemble_id": str(ensemble.id), + }, + ) + if record_info.record_type != old_record_info.record_type: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "error": "Record type of new record does not match previous record type", + "new_record_type": str(record_info.record_type), + "old_record_type": str(old_record_info.record_type), + "name": name, + "ensemble_id": str(ensemble.id), + }, + ) + + record = ds.Record(record_info=old_record_info, **kwargs) + db.add(record) + db.commit() + + return record diff --git a/src/ert/dark_storage/endpoints/responses.py b/src/ert/dark_storage/endpoints/responses.py new file mode 100644 index 00000000000..8374d4d116f --- /dev/null +++ b/src/ert/dark_storage/endpoints/responses.py @@ -0,0 +1,44 @@ +from uuid import uuid4, UUID +import pandas as pd +from fastapi import ( + APIRouter, + Depends, +) +from fastapi.responses import Response +from pandas.core.frame import DataFrame +from ert_storage.database import Session, get_db, HAS_AZURE_BLOB_STORAGE +from ert_storage import database_schema as ds + +router = APIRouter(tags=["response"]) + + +@router.get("/ensembles/{ensemble_id}/responses/{response_name}/data") +async def get_ensemble_response_dataframe( + *, db: Session = Depends(get_db), ensemble_id: UUID, response_name: str +) -> Response: + ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + records = ( + db.query(ds.Record) + .filter(ds.Record.realization_index != None) + .join(ds.RecordInfo) + .filter_by( + ensemble_pk=ensemble.pk, + name=response_name, + record_class=ds.RecordClass.response, + ) + ).all() + df_list = [] + for record in records: + data_df = pd.DataFrame(record.f64_matrix.content) + labels = record.f64_matrix.labels + if labels is not None: + # if the realization is more than 1D array + # the next line will produce ValueError exception + data_df.index = [record.realization_index] + data_df.columns = labels[0] + df_list.append(data_df) + + return Response( + content=pd.concat(df_list, axis=0).to_csv().encode(), + media_type="text/csv", + ) diff --git a/src/ert/dark_storage/endpoints/updates.py b/src/ert/dark_storage/endpoints/updates.py new file mode 100644 index 00000000000..931d8dfc67f --- /dev/null +++ b/src/ert/dark_storage/endpoints/updates.py @@ -0,0 +1,71 @@ +from fastapi import APIRouter, Depends +from ert_storage.database import Session, get_db +from ert_storage import database_schema as ds, json_schema as js +from uuid import UUID + +router = APIRouter(tags=["ensemble"]) + + +@router.post("/updates", response_model=js.UpdateOut) +def create_update( + *, + db: Session = Depends(get_db), + update: js.UpdateIn, +) -> js.UpdateOut: + + ensemble = db.query(ds.Ensemble).filter_by(id=update.ensemble_reference_id).one() + update_obj = ds.Update( + algorithm=update.algorithm, + ensemble_reference_pk=ensemble.pk, + ) + db.add(update_obj) + + if update.observation_transformations: + transformations = {t.name: t for t in update.observation_transformations} + + observation_ids = [t.observation_id for t in transformations.values()] + observations = ( + db.query(ds.Observation) + .filter(ds.Observation.id.in_(observation_ids)) + .all() + ) + + observation_transformations = [ + ds.ObservationTransformation( + active_list=transformations[observation.name].active, + scale_list=transformations[observation.name].scale, + observation=observation, + update=update_obj, + ) + for observation in observations + ] + + db.add_all(observation_transformations) + + db.commit() + return js.UpdateOut( + id=update_obj.id, + experiment_id=ensemble.experiment.id, + algorithm=update_obj.algorithm, + ensemble_reference_id=ensemble.id, + ) + + +@router.get("/updates/{update_id}", response_model=js.UpdateOut) +def get_update( + *, + db: Session = Depends(get_db), + update_id: UUID, +) -> js.UpdateOut: + update_obj = db.query(ds.Update).filter_by(id=update_id).one() + return js.UpdateOut( + id=update_obj.id, + experiment_id=update_obj.ensemble_reference.experiment.id, + algorithm=update_obj.algorithm, + ensemble_reference_id=update_obj.ensemble_reference.id + if update_obj.ensemble_reference is not None + else None, + ensemble_result_id=update_obj.ensemble_result.id + if update_obj.ensemble_result is not None + else None, + ) diff --git a/src/ert/dark_storage/ext/__init__.py b/src/ert/dark_storage/ext/__init__.py new file mode 100644 index 00000000000..9634eb377a5 --- /dev/null +++ b/src/ert/dark_storage/ext/__init__.py @@ -0,0 +1,3 @@ +""" +Extensions to external library APIs +""" diff --git a/src/ert/dark_storage/ext/graphene_sqlalchemy.py b/src/ert/dark_storage/ext/graphene_sqlalchemy.py new file mode 100644 index 00000000000..485d9b087fb --- /dev/null +++ b/src/ert/dark_storage/ext/graphene_sqlalchemy.py @@ -0,0 +1,132 @@ +""" +This module adds the SQLAlchemyMutation class, a graphene.Mutation whose +output mirrors an SQLAlchemyObjectType. This allows us to create mutation that behave +""" + + +from collections import OrderedDict +from typing import Any, Type, Iterable, Callable, Optional, TYPE_CHECKING, Dict + +from graphene_sqlalchemy import SQLAlchemyObjectType as _SQLAlchemyObjectType +from graphene.types.mutation import MutationOptions +from graphene.types.utils import yank_fields_from_attrs +from graphene.types.interface import Interface +from graphene.utils.get_unbound_function import get_unbound_function +from graphene.utils.props import props +from graphene.types.objecttype import ObjectTypeOptions +from graphene import ObjectType, Field, Interface +from graphql import ResolveInfo + + +__all__ = ["SQLAlchemyObjectType", "SQLAlchemyMutation"] + + +if TYPE_CHECKING: + from graphene_sqlalchemy.types.argument import Argument + import graphene.types.field + + +class SQLAlchemyObjectType(_SQLAlchemyObjectType): + class Meta: + abstract = True + + def resolve_id(self, info: ResolveInfo) -> str: + return str(self.id) + + +class SQLAlchemyMutation(SQLAlchemyObjectType): + """ + Object Type Definition (mutation field), appropriated from graphene.Mutation + + SQLAlchemyMutation is a convenience type that helps us build a Field which + takes Arguments and returns a mutation Output SQLAlchemyObjectType. + + Meta class options (optional): + resolver (Callable resolver method): Or ``mutate`` method on Mutation class. Perform data + change and return output. + arguments (Dict[str, graphene.Argument]): Or ``Arguments`` inner class with attributes on + Mutation class. Arguments to use for the mutation Field. + name (str): Name of the GraphQL type (must be unique in schema). Defaults to class + name. + description (str): Description of the GraphQL type in the schema. Defaults to class + docstring. + interfaces (Iterable[graphene.Interface]): GraphQL interfaces to extend with the payload + object. All fields from interface will be included in this object's schema. + fields (Dict[str, graphene.Field]): Dictionary of field name to Field. Not recommended to + use (prefer class attributes or ``Meta.output``). + + """ + + class Meta: + abstract = True + + @classmethod + def __init_subclass_with_meta__( + cls, + interfaces: Iterable[Type[Interface]] = (), + resolver: Callable = None, + arguments: Dict[str, "Argument"] = None, + _meta: Optional[ObjectTypeOptions] = None, + **options: Any + ) -> None: + if not _meta: + _meta = MutationOptions(cls) + + fields = {} + + for interface in interfaces: + assert issubclass(interface, Interface), ( + 'All interfaces of {} must be a subclass of Interface. Received "{}".' + ).format(cls.__name__, interface) + fields.update(interface._meta.fields) + + fields = OrderedDict() + for base in reversed(cls.__mro__): + fields.update(yank_fields_from_attrs(base.__dict__, _as=Field)) + + if not arguments: + input_class = getattr(cls, "Arguments", None) + + if input_class: + arguments = props(input_class) + else: + arguments = {} + + if not resolver: + mutate = getattr(cls, "mutate", None) + assert mutate, "All mutations must define a mutate method in it" + resolver = get_unbound_function(mutate) + + if _meta.fields: + _meta.fields.update(fields) + else: + _meta.fields = fields + + _meta.interfaces = interfaces + _meta.resolver = resolver + _meta.arguments = arguments + + super(SQLAlchemyMutation, cls).__init_subclass_with_meta__( + _meta=_meta, **options + ) + + @classmethod + def Field( + cls, + name: Optional[str] = None, + description: Optional[str] = None, + deprecation_reason: Optional[str] = None, + required: bool = False, + **kwargs: Any + ) -> "graphene.types.field.Field": + """Mount instance of mutation Field.""" + return Field( + cls, + args=cls._meta.arguments, + resolver=cls._meta.resolver, + name=name, + description=description or cls._meta.description, + deprecation_reason=deprecation_reason, + required=required, + **kwargs + ) diff --git a/src/ert/dark_storage/ext/sqlalchemy_arrays.py b/src/ert/dark_storage/ext/sqlalchemy_arrays.py new file mode 100644 index 00000000000..edd3b3f4f43 --- /dev/null +++ b/src/ert/dark_storage/ext/sqlalchemy_arrays.py @@ -0,0 +1,44 @@ +""" +This module adds thte FloatArray and StringArray column types. In Postgresql, +both are native `sqlalchemy.ARRAY`s, while on SQLite, they are `PickleType`s. + +In order to have graphene_sqlalchemy dump the arrays as arrays and not strings, +we need to subclass `PickleType`, and then use +`convert_sqlalchemy_type.register` much in the same way that graphene_sqlalchemy +does it internally for its other types. +""" +from typing import Optional, Type, Union +import sqlalchemy as sa + +from ert_storage.database import IS_POSTGRES + +import graphene +from graphene_sqlalchemy.converter import convert_sqlalchemy_type +from graphene_sqlalchemy.registry import Registry + + +__all__ = ["FloatArray", "StringArray"] + + +SQLAlchemyColumn = Union[sa.types.TypeEngine, Type[sa.types.TypeEngine]] +FloatArray: SQLAlchemyColumn +StringArray: SQLAlchemyColumn + +if IS_POSTGRES: + FloatArray = sa.ARRAY(sa.FLOAT) + StringArray = sa.ARRAY(sa.String) +else: + FloatArray = type("FloatArray", (sa.PickleType,), dict(sa.PickleType.__dict__)) + StringArray = type("StringArray", (sa.PickleType,), dict(sa.PickleType.__dict__)) + + @convert_sqlalchemy_type.register(StringArray) + def convert_column_to_string_array( + type: SQLAlchemyColumn, column: sa.Column, registry: Optional[Registry] = None + ) -> graphene.types.structures.Structure: + return graphene.List(graphene.String) + + @convert_sqlalchemy_type.register(FloatArray) + def convert_column_to_float_array( + type: SQLAlchemyColumn, column: sa.Column, registry: Optional[Registry] = None + ) -> graphene.types.structures.Structure: + return graphene.List(graphene.Float) diff --git a/src/ert/dark_storage/ext/uuid.py b/src/ert/dark_storage/ext/uuid.py new file mode 100644 index 00000000000..89089511c8e --- /dev/null +++ b/src/ert/dark_storage/ext/uuid.py @@ -0,0 +1,55 @@ +from uuid import UUID as SystemUUID +from typing import Any, Optional + +import sqlalchemy as sa +from sqlalchemy.dialects.postgresql import UUID as PostgresUUID +from sqlalchemy.engine import Dialect +from sqlalchemy.types import TypeDecorator, CHAR + +import graphene +from graphene import UUID as GrapheneUUID +from graphene_sqlalchemy.converter import convert_sqlalchemy_type +from graphene_sqlalchemy.registry import Registry + + +class UUID(TypeDecorator): + """Platform-independent UUID type. + + Uses PostgreSQL's UUID type, otherwise uses + CHAR(32), storing as stringified hex values. + + """ + + impl = CHAR + + def load_dialect_impl(self, dialect: Dialect) -> Any: + if dialect.name == "postgresql": + return dialect.type_descriptor(PostgresUUID()) + else: + return dialect.type_descriptor(CHAR(32)) + + def process_bind_param(self, value: Any, dialect: Dialect) -> Any: + if value is None: + return value + elif dialect.name == "postgresql": + return str(value) + else: + if not isinstance(value, SystemUUID): + return "%.32x" % SystemUUID(value).int + else: + return "%.32x" % value.int + + def process_result_value(self, value: Any, dialect: Dialect) -> Any: + if value is None: + return value + else: + if not isinstance(value, SystemUUID): + value = SystemUUID(value) + return value + + +@convert_sqlalchemy_type.register(UUID) +def convert_column_to_uuid( + type: UUID, column: sa.Column, registry: Optional[Registry] = None +) -> graphene.types.structures.Structure: + return GrapheneUUID diff --git a/src/ert/dark_storage/graphql/__init__.py b/src/ert/dark_storage/graphql/__init__.py new file mode 100644 index 00000000000..5e4a119fb56 --- /dev/null +++ b/src/ert/dark_storage/graphql/__init__.py @@ -0,0 +1,61 @@ +from typing import Any, Optional +from starlette.graphql import GraphQLApp +from fastapi import APIRouter +from ert_storage.database import sessionmaker, Session +import graphene as gr +from graphql.execution.base import ExecutionResult, ResolveInfo + +from ert_storage.graphql.ensembles import Ensemble, CreateEnsemble +from ert_storage.graphql.experiments import Experiment, CreateExperiment +from ert_storage import database_schema as ds + + +class Mutations(gr.ObjectType): + create_experiment = CreateExperiment.Field() + create_ensemble = CreateEnsemble.Field(experiment_id=gr.ID(required=True)) + + +class Query(gr.ObjectType): + experiments = gr.List(Experiment) + experiment = gr.Field(Experiment, id=gr.ID(required=True)) + ensemble = gr.Field(Ensemble, id=gr.ID(required=True)) + + @staticmethod + def resolve_experiments(root: None, info: ResolveInfo) -> ds.Experiment: + return Experiment.get_query(info).all() + + @staticmethod + def resolve_experiment(root: None, info: ResolveInfo, id: str) -> ds.Experiment: + return Experiment.get_query(info).filter_by(id=id).one() + + @staticmethod + def resolve_ensemble(root: None, info: ResolveInfo, id: str) -> ds.Ensemble: + return Ensemble.get_query(info).filter_by(id=id).one() + + +class Schema(gr.Schema): + """ + Extended graphene Schema class, where `execute` creates a database session + and passes it on further. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.override_session: Optional[sessionmaker] = None + + def execute(self, *args: Any, **kwargs: Any) -> ExecutionResult: + kwargs.setdefault("context_value", {}) + if self.override_session is not None: + session_obj = self.override_session() + else: + session_obj = Session() + with session_obj as session: + kwargs["context_value"]["session"] = session + + return super().execute(*args, **kwargs) + + +schema = Schema(query=Query, mutation=Mutations) +graphql_app = GraphQLApp(schema=schema) +router = APIRouter(tags=["graphql"]) +router.add_route("/gql", graphql_app) diff --git a/src/ert/dark_storage/graphql/ensembles.py b/src/ert/dark_storage/graphql/ensembles.py new file mode 100644 index 00000000000..a2f912d2fcb --- /dev/null +++ b/src/ert/dark_storage/graphql/ensembles.py @@ -0,0 +1,106 @@ +from typing import List, Iterable, Optional, TYPE_CHECKING +import graphene as gr +from graphene_sqlalchemy.utils import get_session + +from ert_storage.ext.graphene_sqlalchemy import SQLAlchemyObjectType, SQLAlchemyMutation +from ert_storage import database_schema as ds + + +if TYPE_CHECKING: + from graphql.execution.base import ResolveInfo + from ert_storage.graphql.experiments import Experiment + + +class Ensemble(SQLAlchemyObjectType): + class Meta: + model = ds.Ensemble + + child_ensembles = gr.List(lambda: Ensemble) + parent_ensemble = gr.Field(lambda: Ensemble) + responses = gr.Field( + gr.List("ert_storage.graphql.responses.Response"), + names=gr.Argument(gr.List(gr.String), required=False, default_value=None), + ) + unique_responses = gr.List("ert_storage.graphql.unique_responses.UniqueResponse") + + parameters = gr.List("ert_storage.graphql.parameters.Parameter") + + def resolve_child_ensembles( + root: ds.Ensemble, info: "ResolveInfo" + ) -> List[ds.Ensemble]: + return [x.ensemble_result for x in root.children] + + def resolve_parent_ensemble( + root: ds.Ensemble, info: "ResolveInfo" + ) -> Optional[ds.Ensemble]: + update = root.parent + if update is not None: + return update.ensemble_reference + return None + + def resolve_unique_responses( + root: ds.Ensemble, info: "ResolveInfo" + ) -> Iterable[ds.RecordInfo]: + session = info.context["session"] # type: ignore + return root.record_infos.filter_by(record_class=ds.RecordClass.response) + + def resolve_responses( + root: ds.Ensemble, info: "ResolveInfo", names: Optional[Iterable[str]] = None + ) -> Iterable[ds.Record]: + session = info.context["session"] # type: ignore + q = ( + session.query(ds.Record) + .join(ds.RecordInfo) + .filter_by(ensemble=root, record_class=ds.RecordClass.response) + ) + if names is not None: + q = q.filter(ds.RecordInfo.name.in_(names)) + return q.all() + + def resolve_parameters( + root: ds.Ensemble, info: "ResolveInfo" + ) -> Iterable[ds.Record]: + session = info.context["session"] # type: ignore + return ( + session.query(ds.Record) + .join(ds.RecordInfo) + .filter_by(ensemble=root, record_class=ds.RecordClass.parameter) + .all() + ) + + +class CreateEnsemble(SQLAlchemyMutation): + class Meta: + model = ds.Ensemble + + class Arguments: + parameter_names = gr.List(gr.String) + size = gr.Int() + + @staticmethod + def mutate( + root: Optional["Experiment"], + info: "ResolveInfo", + parameter_names: List[str], + size: int, + experiment_id: Optional[str] = None, + ) -> ds.Ensemble: + db = get_session(info.context) + + if experiment_id is not None: + experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() + elif hasattr(root, "id"): + experiment = root + else: + raise ValueError("ID is required") + + ensemble = ds.Ensemble( + parameter_names=parameter_names, + response_names=[], + experiment=experiment, + size=size, + ) + + db.add(ensemble) + db.commit() + return ensemble diff --git a/src/ert/dark_storage/graphql/experiments.py b/src/ert/dark_storage/graphql/experiments.py new file mode 100644 index 00000000000..0c44ebbaca0 --- /dev/null +++ b/src/ert/dark_storage/graphql/experiments.py @@ -0,0 +1,51 @@ +from typing import List, Mapping, TYPE_CHECKING +import graphene as gr +from graphene_sqlalchemy.utils import get_session + +from ert_storage.ext.graphene_sqlalchemy import SQLAlchemyObjectType, SQLAlchemyMutation + +from ert_storage.graphql.ensembles import Ensemble, CreateEnsemble +from ert_storage.graphql.updates import Update +from ert_storage import database_schema as ds, json_schema as js +from ert_storage.endpoints.priors import experiment_priors_to_dict + + +if TYPE_CHECKING: + from graphql.execution.base import ResolveInfo + + +class Experiment(SQLAlchemyObjectType): + class Meta: + model = ds.Experiment + + ensembles = gr.List(Ensemble) + priors = gr.JSONString() + + def resolve_ensembles( + root: ds.Experiment, info: "ResolveInfo" + ) -> List[ds.Ensemble]: + return root.ensembles + + def resolve_priors(root: ds.Experiment, info: "ResolveInfo") -> Mapping[str, dict]: + return experiment_priors_to_dict(root) + + +class CreateExperiment(SQLAlchemyMutation): + class Arguments: + name = gr.String() + + class Meta: + model = ds.Experiment + + create_ensemble = CreateEnsemble.Field() + + @staticmethod + def mutate(root: None, info: "ResolveInfo", name: str) -> ds.Experiment: + db = get_session(info.context) + + experiment = ds.Experiment(name=name) + + db.add(experiment) + db.commit() + + return experiment diff --git a/src/ert/dark_storage/graphql/parameters.py b/src/ert/dark_storage/graphql/parameters.py new file mode 100644 index 00000000000..79e0046a063 --- /dev/null +++ b/src/ert/dark_storage/graphql/parameters.py @@ -0,0 +1,26 @@ +from typing import Any, List, Optional, TYPE_CHECKING +import graphene as gr +from sqlalchemy.orm import Session + +from ert_storage.ext.graphene_sqlalchemy import SQLAlchemyObjectType +from ert_storage import database_schema as ds +from ert_storage.endpoints.priors import prior_to_dict + + +if TYPE_CHECKING: + from graphql.execution.base import ResolveInfo + + +class Parameter(SQLAlchemyObjectType): + class Meta: + model = ds.Record + + name = gr.String() + prior = gr.JSONString() + + def resolve_name(root: ds.Record, info: "ResolveInfo") -> str: + return root.name + + def resolve_prior(root: ds.Record, info: "ResolveInfo") -> Optional[dict]: + prior = root.record_info.prior + return prior_to_dict(prior) if prior is not None else None diff --git a/src/ert/dark_storage/graphql/responses.py b/src/ert/dark_storage/graphql/responses.py new file mode 100644 index 00000000000..8154009ad51 --- /dev/null +++ b/src/ert/dark_storage/graphql/responses.py @@ -0,0 +1,19 @@ +from typing import List, Optional, TYPE_CHECKING +import graphene as gr + +from ert_storage.ext.graphene_sqlalchemy import SQLAlchemyObjectType +from ert_storage import database_schema as ds + + +if TYPE_CHECKING: + from graphql.execution.base import ResolveInfo + + +class Response(SQLAlchemyObjectType): + class Meta: + model = ds.Record + + name = gr.String() + + def resolve_name(root: ds.Record, info: "ResolveInfo") -> str: + return root.name diff --git a/src/ert/dark_storage/graphql/unique_responses.py b/src/ert/dark_storage/graphql/unique_responses.py new file mode 100644 index 00000000000..4d8acac50af --- /dev/null +++ b/src/ert/dark_storage/graphql/unique_responses.py @@ -0,0 +1,13 @@ +from typing import TYPE_CHECKING + +from ert_storage.ext.graphene_sqlalchemy import SQLAlchemyObjectType +from ert_storage import database_schema as ds + + +if TYPE_CHECKING: + from graphql.execution.base import ResolveInfo + + +class UniqueResponse(SQLAlchemyObjectType): + class Meta: + model = ds.RecordInfo diff --git a/src/ert/dark_storage/graphql/updates.py b/src/ert/dark_storage/graphql/updates.py new file mode 100644 index 00000000000..84008735d6f --- /dev/null +++ b/src/ert/dark_storage/graphql/updates.py @@ -0,0 +1,15 @@ +from typing import List, TYPE_CHECKING +import graphene as gr +from graphene_sqlalchemy.utils import get_session + +from ert_storage.ext.graphene_sqlalchemy import SQLAlchemyObjectType, SQLAlchemyMutation +from ert_storage import database_schema as ds + + +if TYPE_CHECKING: + from graphql.execution.base import ResolveInfo + + +class Update(SQLAlchemyObjectType): + class Meta: + model = ds.Update diff --git a/src/ert/dark_storage/json_schema/__init__.py b/src/ert/dark_storage/json_schema/__init__.py new file mode 100644 index 00000000000..74893127df8 --- /dev/null +++ b/src/ert/dark_storage/json_schema/__init__.py @@ -0,0 +1,11 @@ +from .ensemble import EnsembleIn, EnsembleOut +from .record import RecordOut +from .experiment import ExperimentIn, ExperimentOut +from .observation import ( + ObservationIn, + ObservationOut, + ObservationTransformationIn, + ObservationTransformationOut, +) +from .update import UpdateIn, UpdateOut +from .prior import Prior diff --git a/src/ert/dark_storage/json_schema/ensemble.py b/src/ert/dark_storage/json_schema/ensemble.py new file mode 100644 index 00000000000..d1f4fba3907 --- /dev/null +++ b/src/ert/dark_storage/json_schema/ensemble.py @@ -0,0 +1,35 @@ +from uuid import UUID +from typing import List, Optional, Any, Mapping +from pydantic import BaseModel, Field, root_validator + + +class _Ensemble(BaseModel): + size: int + parameter_names: List[str] + response_names: List[str] + + +class EnsembleIn(_Ensemble): + update_id: Optional[UUID] = None + metadata: Optional[Any] + + @root_validator + def _check_names_no_overlap(cls, values: Mapping[str, Any]) -> Mapping[str, Any]: + """ + Verify that `parameter_names` and `response_names` don't overlap. Ie, no + record can be both a parameter and a response. + """ + if not set(values["parameter_names"]).isdisjoint(set(values["response_names"])): + raise ValueError("parameters and responses cannot have a name in common") + return values + + +class EnsembleOut(_Ensemble): + id: UUID + children: List[UUID] = Field(alias="child_ensemble_ids") + parent: Optional[UUID] = Field(alias="parent_ensemble_id") + experiment_id: Optional[UUID] = None + metadata: Mapping[str, Any] = Field(alias="metadata_dict") + + class Config: + orm_mode = True diff --git a/src/ert/dark_storage/json_schema/experiment.py b/src/ert/dark_storage/json_schema/experiment.py new file mode 100644 index 00000000000..78154b8c438 --- /dev/null +++ b/src/ert/dark_storage/json_schema/experiment.py @@ -0,0 +1,22 @@ +from uuid import UUID +from typing import List, Mapping, Any +from pydantic import BaseModel +from .prior import Prior + + +class _Experiment(BaseModel): + name: str + + +class ExperimentIn(_Experiment): + priors: Mapping[str, Prior] = {} + + +class ExperimentOut(_Experiment): + id: UUID + ensembles: List[UUID] + priors: Mapping[str, UUID] + metadata: Mapping[str, Any] = {} + + class Config: + orm_mode = True diff --git a/src/ert/dark_storage/json_schema/observation.py b/src/ert/dark_storage/json_schema/observation.py new file mode 100644 index 00000000000..5ea986f468f --- /dev/null +++ b/src/ert/dark_storage/json_schema/observation.py @@ -0,0 +1,42 @@ +from uuid import UUID +from typing import List, Optional, Any, Mapping +from pydantic import BaseModel + + +class _ObservationTransformation(BaseModel): + name: str + active: List[bool] + scale: List[float] + observation_id: UUID + + +class ObservationTransformationIn(_ObservationTransformation): + pass + + +class ObservationTransformationOut(_ObservationTransformation): + id: UUID + + class Config: + orm_mode = True + + +class _Observation(BaseModel): + name: str + errors: List[float] + values: List[float] + x_axis: List[Any] + records: Optional[List[UUID]] = None + + +class ObservationIn(_Observation): + pass + + +class ObservationOut(_Observation): + id: UUID + transformation: Optional[ObservationTransformationOut] = None + metadata: Mapping[str, Any] = {} + + class Config: + orm_mode = True diff --git a/src/ert/dark_storage/json_schema/prior.py b/src/ert/dark_storage/json_schema/prior.py new file mode 100644 index 00000000000..10150946294 --- /dev/null +++ b/src/ert/dark_storage/json_schema/prior.py @@ -0,0 +1,151 @@ +import sys +from pydantic import BaseModel +from typing import Union + +if sys.version_info < (3, 8): + from typing_extensions import Literal +else: + from typing import Literal + + +class PriorConst(BaseModel): + """ + Constant parameter prior + """ + + function: Literal["const"] = "const" + value: float + + +class PriorTrig(BaseModel): + """ + Triangular distribution parameter prior + """ + + function: Literal["trig"] = "trig" + min: float + max: float + mode: float + + +class PriorNormal(BaseModel): + """ + Normal distribution parameter prior + """ + + function: Literal["normal"] = "normal" + mean: float + std: float + + +class PriorLogNormal(BaseModel): + """ + Log-normal distribution parameter prior + """ + + function: Literal["lognormal"] = "lognormal" + mean: float + std: float + + +class PriorErtTruncNormal(BaseModel): + """ + ERT Truncated normal distribution parameter prior + + ERT differs from the usual distribution by that it simply clamps on `min` + and `max`, which gives a bias towards the extremes. + + """ + + function: Literal["ert_truncnormal"] = "ert_truncnormal" + mean: float + std: float + min: float + max: float + + +class PriorStdNormal(BaseModel): + """ + Standard normal distribution parameter prior + + Normal distribution with mean of 0 and standard deviation of 1 + """ + + function: Literal["stdnormal"] = "stdnormal" + + +class PriorUniform(BaseModel): + """ + Uniform distribution parameter prior + """ + + function: Literal["uniform"] = "uniform" + min: float + max: float + + +class PriorErtDUniform(BaseModel): + """ + ERT Discrete uniform distribution parameter prior + + This discrete uniform distribution differs from the standard by using the + `bins` parameter. Normally, `a`, and `b` are integers, and the sample space + are the integers between. ERT allows `a` and `b` to be arbitrary floats, + where the sample space is binned. + + """ + + function: Literal["ert_duniform"] = "ert_duniform" + bins: int + min: float + max: float + + +class PriorLogUniform(BaseModel): + """ + Logarithmic uniform distribution parameter prior + """ + + function: Literal["loguniform"] = "loguniform" + min: float + max: float + + +class PriorErtErf(BaseModel): + """ + ERT Error function distribution parameter prior + """ + + function: Literal["ert_erf"] = "ert_erf" + min: float + max: float + skewness: float + width: float + + +class PriorErtDErf(BaseModel): + """ + ERT Discrete error function distribution parameter prior + """ + + function: Literal["ert_derf"] = "ert_derf" + bins: int + min: float + max: float + skewness: float + width: float + + +Prior = Union[ + PriorConst, + PriorTrig, + PriorNormal, + PriorLogNormal, + PriorErtTruncNormal, + PriorStdNormal, + PriorUniform, + PriorErtDUniform, + PriorLogUniform, + PriorErtErf, + PriorErtDErf, +] diff --git a/src/ert/dark_storage/json_schema/record.py b/src/ert/dark_storage/json_schema/record.py new file mode 100644 index 00000000000..646608ddb90 --- /dev/null +++ b/src/ert/dark_storage/json_schema/record.py @@ -0,0 +1,16 @@ +from uuid import UUID +from typing import Any, Mapping +from pydantic import BaseModel, Field + + +class _Record(BaseModel): + pass + + +class RecordOut(_Record): + id: UUID + name: str + metadata: Mapping[str, Any] = Field(alias="metadata_dict") + + class Config: + orm_mode = True diff --git a/src/ert/dark_storage/json_schema/update.py b/src/ert/dark_storage/json_schema/update.py new file mode 100644 index 00000000000..b4c6b2fef53 --- /dev/null +++ b/src/ert/dark_storage/json_schema/update.py @@ -0,0 +1,25 @@ +from uuid import UUID + +from pydantic import BaseModel +from typing import List, Union, Optional +from .observation import ( + ObservationTransformationIn, +) + + +class _Update(BaseModel): + algorithm: str + ensemble_result_id: Union[UUID, None] + ensemble_reference_id: Union[UUID, None] + + +class UpdateIn(_Update): + observation_transformations: Optional[List[ObservationTransformationIn]] = None + + +class UpdateOut(_Update): + id: UUID + experiment_id: UUID + + class Config: + orm_mode = True diff --git a/src/ert/dark_storage/py.typed b/src/ert/dark_storage/py.typed new file mode 100644 index 00000000000..a2c47e10281 --- /dev/null +++ b/src/ert/dark_storage/py.typed @@ -0,0 +1 @@ +This file marks that ert_storage is PEP-561 and type-compliant. Used by eg. mypy to enable type checking. diff --git a/src/ert/dark_storage/security.py b/src/ert/dark_storage/security.py new file mode 100644 index 00000000000..91a841f7fef --- /dev/null +++ b/src/ert/dark_storage/security.py @@ -0,0 +1,25 @@ +import os +from typing import Optional +from fastapi import Security, HTTPException, status +from fastapi.security import APIKeyHeader + + +DEFAULT_TOKEN = "hunter2" +_security_header = APIKeyHeader(name="Token", auto_error=False) + + +async def security(*, token: Optional[str] = Security(_security_header)) -> None: + if os.getenv("ERT_STORAGE_NO_TOKEN"): + return + if not token: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, detail="Not authenticated" + ) + real_token = os.getenv("ERT_STORAGE_TOKEN", DEFAULT_TOKEN) + if token == real_token: + # Success + return + + # HTTP 403 is when the user has authorized themselves, but aren't allowed to + # access this resource + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Invalid token") diff --git a/src/ert/dark_storage/testing/__init__.py b/src/ert/dark_storage/testing/__init__.py new file mode 100644 index 00000000000..927e5eb7475 --- /dev/null +++ b/src/ert/dark_storage/testing/__init__.py @@ -0,0 +1,4 @@ +""" +Utilities for running ERT Storage integration tests in external software. +""" +from ert_storage.testing.testclient import testclient_factory, ClientError diff --git a/src/ert/dark_storage/testing/pytest11.py b/src/ert/dark_storage/testing/pytest11.py new file mode 100644 index 00000000000..10d8de0f88f --- /dev/null +++ b/src/ert/dark_storage/testing/pytest11.py @@ -0,0 +1,13 @@ +import pytest +from typing import Generator, TYPE_CHECKING +from ert_storage.testing import testclient_factory + + +if TYPE_CHECKING: + from ert_storage.testing.testclient import _TestClient + + +@pytest.fixture +def ert_storage_client() -> Generator["_TestClient", None, None]: + with testclient_factory() as testclient: + yield testclient diff --git a/src/ert/dark_storage/testing/testclient.py b/src/ert/dark_storage/testing/testclient.py new file mode 100644 index 00000000000..cf4dff13385 --- /dev/null +++ b/src/ert/dark_storage/testing/testclient.py @@ -0,0 +1,326 @@ +import os +import requests +from typing import ( + Any, + AsyncGenerator, + Generator, + Mapping, + MutableMapping, + Optional, + TYPE_CHECKING, + Tuple, + Union, +) +from pprint import pformat +from contextlib import contextmanager +from fastapi import Depends +from sqlalchemy.orm.session import Session +from sqlalchemy.engine.base import Transaction +from starlette.testclient import ( + TestClient as StarletteTestClient, + ASGI2App, + ASGI3App, + Cookies, + Params, + DataType, + TimeOut, + FileType, +) +from sqlalchemy.orm import sessionmaker +from graphene import Schema as GrapheneSchema +from graphene.test import Client as GrapheneClient + +from ert_storage.security import security + + +if TYPE_CHECKING: + from promise import Promise + from rx import Observable + from graphql.execution import ExecutionResult + + GraphQLResult = Union[ExecutionResult, Observable, Promise[ExecutionResult]] + + +class ClientError(RuntimeError): + pass + + +class _TestClient: + __test__ = False # Pytest should ignore this class + + def __init__( + self, + app: Union[ASGI2App, ASGI3App], + session: sessionmaker, + gql_schema: GrapheneSchema, + base_url: str = "http://testserver", + raise_server_exceptions: bool = True, + root_path: str = "", + ) -> None: + self.raise_on_client_error = True + self.http_client = StarletteTestClient( + app, base_url, raise_server_exceptions, root_path + ) + self.gql_client = GrapheneClient(gql_schema) + self.session = session + + def get( + self, + url: str, + params: Params = None, + headers: MutableMapping[str, str] = None, + cookies: Cookies = None, + files: FileType = None, + timeout: TimeOut = None, + allow_redirects: bool = None, + stream: bool = None, + check_status_code: Optional[int] = 200, + ) -> requests.Response: + resp = self.http_client.get( + url, + params=params, + headers=headers, + cookies=cookies, + files=files, + timeout=timeout, + allow_redirects=allow_redirects, + stream=stream, + ) + self._check(check_status_code, resp) + return resp + + def post( + self, + url: str, + params: Params = None, + data: DataType = None, + headers: MutableMapping[str, str] = None, + cookies: Cookies = None, + files: FileType = None, + timeout: TimeOut = None, + allow_redirects: bool = None, + stream: bool = None, + json: Any = None, + check_status_code: Optional[int] = 200, + ) -> requests.Response: + resp = self.http_client.post( + url, + params=params, + data=data, + headers=headers, + cookies=cookies, + files=files, + timeout=timeout, + allow_redirects=allow_redirects, + stream=stream, + json=json, + ) + self._check(check_status_code, resp) + return resp + + def put( + self, + url: str, + params: Params = None, + data: DataType = None, + headers: MutableMapping[str, str] = None, + cookies: Cookies = None, + files: FileType = None, + timeout: TimeOut = None, + allow_redirects: bool = None, + stream: bool = None, + json: Any = None, + check_status_code: Optional[int] = 200, + ) -> requests.Response: + resp = self.http_client.put( + url, + params=params, + data=data, + headers=headers, + cookies=cookies, + files=files, + timeout=timeout, + allow_redirects=allow_redirects, + stream=stream, + json=json, + ) + self._check(check_status_code, resp) + return resp + + def patch( + self, + url: str, + params: Params = None, + data: DataType = None, + headers: MutableMapping[str, str] = None, + cookies: Cookies = None, + files: FileType = None, + timeout: TimeOut = None, + allow_redirects: bool = None, + stream: bool = None, + json: Any = None, + check_status_code: Optional[int] = 200, + ) -> requests.Response: + resp = self.http_client.patch( + url, + params=params, + data=data, + headers=headers, + cookies=cookies, + files=files, + timeout=timeout, + allow_redirects=allow_redirects, + stream=stream, + json=json, + ) + self._check(check_status_code, resp) + return resp + + def delete( + self, + url: str, + params: Params = None, + headers: MutableMapping[str, str] = None, + cookies: Cookies = None, + timeout: TimeOut = None, + allow_redirects: bool = None, + check_status_code: Optional[int] = 200, + ) -> requests.Response: + resp = self.http_client.delete( + url, + params=params, + headers=headers, + cookies=cookies, + timeout=timeout, + allow_redirects=allow_redirects, + ) + self._check(check_status_code, resp) + return resp + + def gql_execute( + self, + request_string: str, + variable_values: Optional[Mapping[str, Any]] = None, + check: bool = True, + ) -> dict: + doc = self.gql_client.execute(request_string, variable_values=variable_values) + if self.raise_on_client_error and check and "errors" in doc: + raise ClientError(f"GraphQL query returned an error:\n{pformat(doc)}") + + return doc + + def _check( + self, check_status_code: Optional[int], response: requests.Response + ) -> None: + if ( + not self.raise_on_client_error + or check_status_code is None + or response.status_code == check_status_code + ): + return + + try: + doc = response.json() + except: + doc = response.content + raise ClientError( + f"Status code was {response.status_code}, expected {check_status_code}:\n{doc}" + ) + + +@contextmanager +def testclient_factory() -> Generator[_TestClient, None, None]: + env_key = "ERT_STORAGE_DATABASE_URL" + env_unset = False + if env_key not in os.environ: + os.environ[env_key] = "sqlite:///:memory:" + env_unset = True + print("Using in-memory SQLite database for tests") + + if os.getenv("ERT_STORAGE_NO_ROLLBACK", ""): + print( + "Environment variable 'ERT_STORAGE_NO_ROLLBACK' is set.\n" + "Will keep data in database." + ) + rollback = False + else: + rollback = True + + from ert_storage.app import app + from ert_storage.graphql import schema + + session, transaction, connection = _begin_transaction() + schema.override_session = session + + yield _TestClient(app, session=session, gql_schema=schema) + + schema.override_session = None + _end_transaction(transaction, connection, rollback) + + if env_unset: + del os.environ[env_key] + + +_TransactionInfo = Tuple[sessionmaker, Transaction, Any] + + +def _override_get_db(session: sessionmaker) -> None: + from ert_storage.app import app + from ert_storage.database import ( + get_db, + IS_POSTGRES, + ) + + async def override_get_db( + *, _: None = Depends(security) + ) -> AsyncGenerator[Session, None]: + db = session() + + # Make PostgreSQL return float8 columns with highest precision. If we don't + # do this, we may lose up to 3 of the least significant digits. + if IS_POSTGRES: + db.execute("SET extra_float_digits=3") + try: + yield db + db.commit() + db.close() + except: + db.rollback() + db.close() + raise + + app.dependency_overrides[get_db] = override_get_db + + +def _begin_transaction() -> _TransactionInfo: + from ert_storage.database import ( + engine, + IS_SQLITE, + HAS_AZURE_BLOB_STORAGE, + ) + from ert_storage.database_schema import Base + + if IS_SQLITE: + Base.metadata.create_all(bind=engine) + if HAS_AZURE_BLOB_STORAGE: + import asyncio + from ert_storage.database import create_container_if_not_exist + + loop = asyncio.get_event_loop() + loop.run_until_complete(create_container_if_not_exist()) + + connection = engine.connect() + transaction = connection.begin() + session = sessionmaker(autocommit=False, autoflush=False, bind=connection) + + _override_get_db(session) + + return session, transaction, connection + + +def _end_transaction(transaction: Transaction, connection: Any, rollback: bool) -> None: + if rollback: + transaction.rollback() + else: + transaction.commit() + connection.close() From 6e3cf6d5203801cd8584918fba32d43b954616dc Mon Sep 17 00:00:00 2001 From: Zohar Malamant Date: Fri, 28 May 2021 14:21:00 +0200 Subject: [PATCH 02/43] Replace HTTPException with our own exceptions.py --- src/ert/dark_storage/app.py | 47 +++++-- .../dark_storage/endpoints/compute/misfits.py | 12 +- src/ert/dark_storage/endpoints/records.py | 117 +++++------------- src/ert/dark_storage/exceptions.py | 29 +++++ 4 files changed, 101 insertions(+), 104 deletions(-) create mode 100644 src/ert/dark_storage/exceptions.py diff --git a/src/ert/dark_storage/app.py b/src/ert/dark_storage/app.py index aa0b9a2912f..b646c26e197 100644 --- a/src/ert/dark_storage/app.py +++ b/src/ert/dark_storage/app.py @@ -1,4 +1,5 @@ import json +from enum import Enum from typing import Any from fastapi import FastAPI, Request, status from fastapi.responses import Response, RedirectResponse @@ -6,23 +7,38 @@ from ert_storage.endpoints import router as endpoints_router from ert_storage.graphql import router as graphql_router +from ert_storage.exceptions import ErtStorageError from sqlalchemy.orm.exc import NoResultFound +class JSONEncoder(json.JSONEncoder): + """ + Custom JSON encoder with support for Python 3.4 enums + """ + + def default(self, obj: Any) -> Any: + if isinstance(obj, Enum): + return obj.name + return super().default(obj) + + class JSONResponse(Response): """A replacement for Starlette's JSONResponse that permits NaNs.""" media_type = "application/json" def render(self, content: Any) -> bytes: - return json.dumps( - content, - ensure_ascii=False, - allow_nan=True, - indent=None, - separators=(",", ":"), - ).encode("utf-8") + return ( + JSONEncoder( + ensure_ascii=False, + allow_nan=True, + indent=None, + separators=(",", ":"), + ) + .encode(content) + .encode("utf-8") + ) app = FastAPI( @@ -59,6 +75,23 @@ async def sqlalchemy_exception_handler( ) +@app.exception_handler(ErtStorageError) +async def ert_storage_error_handler( + request: Request, exc: ErtStorageError +) -> JSONResponse: + return JSONResponse( + { + "detail": { + **request.query_params, + **request.path_params, + **exc.args[1], + "error": exc.args[0], + } + }, + status_code=exc.__status_code__, + ) + + @app.get("/") async def root() -> RedirectResponse: return RedirectResponse("/docs") diff --git a/src/ert/dark_storage/endpoints/compute/misfits.py b/src/ert/dark_storage/endpoints/compute/misfits.py index 463cb14c963..e9f82513afa 100644 --- a/src/ert/dark_storage/endpoints/compute/misfits.py +++ b/src/ert/dark_storage/endpoints/compute/misfits.py @@ -5,9 +5,10 @@ from typing import Any, Optional, List import sqlalchemy as sa from fastapi.responses import Response -from fastapi import APIRouter, Depends, HTTPException, status +from fastapi import APIRouter, Depends, status from ert_storage.database import Session, get_db from ert_storage import database_schema as ds, json_schema as js +from ert_storage import exceptions as exc from ert_storage.compute import calculate_misfits_from_pandas, misfits router = APIRouter(tags=["misfits"]) @@ -76,14 +77,7 @@ async def get_response_misfits( response_dict, observation_df, summary_misfits ) except Exception as misfits_exc: - raise HTTPException( - status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, - detail={ - "error": f"Unable to compute misfits: {misfits_exc}", - "name": response_name, - "ensemble_id": str(ensemble_id), - }, - ) + raise exc.UnprocessableError(f"Unable to compute misfits: {misfits_exc}") return Response( content=result_df.to_csv().encode(), media_type="text/csv", diff --git a/src/ert/dark_storage/endpoints/records.py b/src/ert/dark_storage/endpoints/records.py index 089befe0820..ae3a7819b5c 100644 --- a/src/ert/dark_storage/endpoints/records.py +++ b/src/ert/dark_storage/endpoints/records.py @@ -10,7 +10,6 @@ Body, Depends, File, - HTTPException, Header, Request, UploadFile, @@ -23,6 +22,7 @@ from sqlalchemy.orm.attributes import flag_modified from ert_storage.database import Session, get_db, HAS_AZURE_BLOB_STORAGE, BLOB_CONTAINER from ert_storage import database_schema as ds, json_schema as js +from ert_storage import exceptions as exc from fastapi.logger import logger @@ -228,15 +228,8 @@ async def post_ensemble_record_matrix( is_parameter = name in ensemble.parameter_names is_response = name in ensemble.response_names if prior_id is not None and not is_parameter: - raise HTTPException( - status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, - detail={ - "error": "Priors can only be specified for parameter records", - "name": name, - "ensemble_id": str(ensemble_id), - "realization_index": realization_index, - "prior_id": str(prior_id), - }, + raise exc.UnprocessableError( + "Priors can only be specified for parameter records" ) labels = None @@ -273,29 +266,14 @@ async def post_ensemble_record_matrix( message = f"Ensemble-wide record '{name}' for ensemble '{ensemble_id}' needs to be a matrix" else: message = f"Forward-model record '{name}' for ensemble '{ensemble_id}', realization {realization_index} needs to be a matrix" - - raise HTTPException( - status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, - detail={ - "error": message, - "name": name, - "ensemble_id": str(ensemble_id), - "realization_index": realization_index, - }, - ) + raise exc.UnprocessableError(message) # Require that the dimensionality of an ensemble-wide parameter matrix is at least 2 if realization_index is None and is_parameter: if content.ndim <= 1: - raise HTTPException( - status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, - detail={ - "error": "Ensemble-wide parameter record '{name}' for ensemble '{ensemble_id}'" - "must have dimensionality of at least 2", - "name": name, - "ensemble_id": str(ensemble_id), - "realization_index": realization_index, - }, + raise exc.UnprocessableError( + f"Ensemble-wide parameter record '{name}' for ensemble '{ensemble_id}'" + "must have dimensionality of at least 2" ) matrix_obj = ds.F64Matrix(content=content.tolist(), labels=labels) @@ -405,14 +383,7 @@ async def post_record_observations( flag_modified(record_obj, "observations") db.commit() else: - raise HTTPException( - status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, - detail={ - "error": f"Observations {observation_ids} not found!", - "name": name, - "ensemble_id": str(ensemble_id), - }, - ) + raise exc.UnprocessableError(f"Observations {observation_ids} not found!") @router.get("/ensembles/{ensemble_id}/records/{name}/observations") @@ -554,14 +525,14 @@ async def get_ensemble_record( try: bundle = _get_forward_model_record(db, ensemble_id, name, realization_index) realization_index = None - except HTTPException as exc: + except exc.NotFoundError: bundle = _get_ensemble_record(db, ensemble_id, name) # Only parameter records can be "up-casted" to ensemble-wide is_matrix = bundle.record_info.record_type == ds.RecordType.f64_matrix is_parameter = name in bundle.record_info.ensemble.parameter_names if not is_matrix or not is_parameter: - raise exc + raise return await _get_record_data(bundle, accept, realization_index) @@ -646,13 +617,8 @@ def _get_ensemble_record(db: Session, ensemble_id: UUID, name: str) -> ds.Record .one() ) except NoResultFound: - raise HTTPException( - status_code=404, - detail={ - "error": f"Ensemble-wide record '{name}' for ensemble '{ensemble_id}' not found!", - "name": name, - "ensemble_id": str(ensemble_id), - }, + raise exc.NotFoundError( + f"Ensemble-wide record '{name}' for ensemble '{ensemble_id}' not found!" ) @@ -672,13 +638,8 @@ def _get_forward_model_record( .one() ) except NoResultFound: - raise HTTPException( - status_code=404, - detail={ - "error": f"Forward-model record '{name}' for ensemble '{ensemble_id}', realization {realization_index} not found!", - "name": name, - "ensemble_id": str(ensemble_id), - }, + raise exc.NotFoundError( + f"Forward-model record '{name}' for ensemble '{ensemble_id}', realization {realization_index} not found!" ) @@ -697,16 +658,11 @@ def _get_and_assert_ensemble( ) if realization_index is not None: if realization_index not in range(ensemble.size) and ensemble.size != -1: - raise HTTPException( - status_code=status.HTTP_417_EXPECTATION_FAILED, - detail={ - "error": f"Ensemble '{name}' ('{ensemble_id}') does have a 'size' " - f"of {ensemble.size}. The posted record is targeting " - f"'realization_index' {realization_index} which is out " - f"of bounds.", - "name": name, - "ensemble_id": str(ensemble_id), - }, + raise exc.ExpectationError( + f"Ensemble '{name}' ('{ensemble_id}') does have a 'size' " + f"of {ensemble.size}. The posted record is targeting " + f"'realization_index' {realization_index} which is out " + f"of bounds." ) q = q.filter( @@ -715,13 +671,8 @@ def _get_and_assert_ensemble( ) if q.count() > 0: - raise HTTPException( - status_code=status.HTTP_409_CONFLICT, - detail={ - "error": f"Ensemble-wide record '{name}' for ensemble '{ensemble_id}' already exists", - "name": name, - "ensemble_id": str(ensemble_id), - }, + raise exc.ConflictError( + f"Ensemble-wide record '{name}' for ensemble '{ensemble_id}' already exists" ) return ensemble @@ -827,26 +778,16 @@ def _create_record( # Check that the parameters match if record_info.record_class != old_record_info.record_class: - raise HTTPException( - status_code=status.HTTP_409_CONFLICT, - detail={ - "error": "Record class of new record does not match previous record class", - "new_record_class": str(record_info.record_class), - "old_record_class": str(old_record_info.record_class), - "name": name, - "ensemble_id": str(ensemble.id), - }, + raise exc.ConflictError( + "Record class of new record does not match previous record class", + new_record_class=record_info.record_class, + old_record_class=old_record_info.record_class, ) if record_info.record_type != old_record_info.record_type: - raise HTTPException( - status_code=status.HTTP_409_CONFLICT, - detail={ - "error": "Record type of new record does not match previous record type", - "new_record_type": str(record_info.record_type), - "old_record_type": str(old_record_info.record_type), - "name": name, - "ensemble_id": str(ensemble.id), - }, + raise exc.ConflictError( + "Record type of new record does not match previous record type", + new_record_type=record_info.record_type, + old_record_type=old_record_info.record_type, ) record = ds.Record(record_info=old_record_info, **kwargs) diff --git a/src/ert/dark_storage/exceptions.py b/src/ert/dark_storage/exceptions.py new file mode 100644 index 00000000000..71be62c4ad2 --- /dev/null +++ b/src/ert/dark_storage/exceptions.py @@ -0,0 +1,29 @@ +from typing import Any +from fastapi import status + + +class ErtStorageError(RuntimeError): + """ + Base error class for all the rest of errors + """ + + __status_code__ = status.HTTP_200_OK + + def __init__(self, message: str, **kwargs: Any): + super().__init__(message, kwargs) + + +class NotFoundError(ErtStorageError): + __status_code__ = status.HTTP_404_NOT_FOUND + + +class ConflictError(ErtStorageError): + __status_code__ = status.HTTP_409_CONFLICT + + +class ExpectationError(ErtStorageError): + __status_code__ = status.HTTP_417_EXPECTATION_FAILED + + +class UnprocessableError(ErtStorageError): + __status_code__ = status.HTTP_422_UNPROCESSABLE_ENTITY From 67466632da596eb1400e303c7e241762650e9995 Mon Sep 17 00:00:00 2001 From: Zohar Malamant Date: Fri, 28 May 2021 15:15:50 +0200 Subject: [PATCH 03/43] Use FastAPI Depends in records.py --- .../dark_storage/database_schema/record.py | 10 +- src/ert/dark_storage/endpoints/records.py | 399 ++++++++---------- 2 files changed, 175 insertions(+), 234 deletions(-) diff --git a/src/ert/dark_storage/database_schema/record.py b/src/ert/dark_storage/database_schema/record.py index 0a581b838fc..1400036bc01 100644 --- a/src/ert/dark_storage/database_schema/record.py +++ b/src/ert/dark_storage/database_schema/record.py @@ -12,7 +12,7 @@ from .metadatafield import MetadataField from .observation import observation_record_association -from .record_info import RecordType +from .record_info import RecordType, RecordClass class Record(Base, MetadataField): @@ -60,6 +60,14 @@ def data(self) -> Any: def name(self) -> str: return self.record_info.name + @property + def record_type(self) -> RecordType: + return self.record_info.record_type + + @property + def record_class(self) -> RecordClass: + return self.record_info.record_class + class File(Base): __tablename__ = "file" diff --git a/src/ert/dark_storage/endpoints/records.py b/src/ert/dark_storage/endpoints/records.py index ae3a7819b5c..f165026ebee 100644 --- a/src/ert/dark_storage/endpoints/records.py +++ b/src/ert/dark_storage/endpoints/records.py @@ -38,26 +38,140 @@ class ListRecords(BaseModel): forward_model: Mapping[str, str] -@router.post("/ensembles/{ensemble_id}/records/{name}/file") -async def post_ensemble_record_file( +def get_record_by_name( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + name: str, + realization_index: Optional[int] = None, +) -> ds.Record: + try: + return ( + db.query(ds.Record) + .filter_by(realization_index=realization_index) + .join(ds.RecordInfo) + .filter_by(name=name) + .join(ds.Ensemble) + .filter_by(id=ensemble_id) + .one() + ) + except NoResultFound as e: + pass + + if realization_index is not None: + return ( + db.query(ds.Record) + .filter_by(realization_index=None) + .join(ds.RecordInfo) + .filter_by( + name=name, + record_class=ds.RecordClass.parameter, + record_type=ds.RecordType.f64_matrix, + ) + .join(ds.Ensemble) + .filter_by(id=ensemble_id) + .one() + ) + raise exc.NotFoundError(f"Record not found") + + +def new_record( *, db: Session = Depends(get_db), ensemble_id: UUID, name: str, realization_index: Optional[int] = None, +) -> ds.Record: + """ + Get ensemble and verify that no record with the name `name` exists + """ + ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() + + q = ( + db.query(ds.Record) + .join(ds.RecordInfo) + .filter_by(ensemble_pk=ensemble.pk, name=name) + ) + if realization_index is not None: + if realization_index not in range(ensemble.size) and ensemble.size != -1: + raise exc.ExpectationError( + f"Ensemble '{name}' ('{ensemble_id}') does have a 'size' " + f"of {ensemble.size}. The posted record is targeting " + f"'realization_index' {realization_index} which is out " + f"of bounds." + ) + + q = q.filter( + (ds.Record.realization_index == None) + | (ds.Record.realization_index == realization_index) + ) + + if q.count() > 0: + raise exc.ConflictError( + f"Ensemble-wide record '{name}' for ensemble '{ensemble_id}' already exists", + ) + + return ds.Record( + record_info=ds.RecordInfo( + ensemble=ensemble, + name=name, + ), + realization_index=realization_index, + ) + + +def new_record_file( + *, + db: Session = Depends(get_db), + record: ds.Record = Depends(new_record), +) -> ds.Record: + record.record_info.record_class = ds.RecordClass.other + record.record_info.record_type = ds.RecordType.file + return record + + +def new_record_matrix( + *, + db: Session = Depends(get_db), + record: ds.Record = Depends(new_record), + prior_id: Optional[UUID] = None, +) -> ds.Record: + ensemble = record.record_info.ensemble + if record.name in ensemble.parameter_names: + record_class = ds.RecordClass.parameter + elif record.name in ensemble.response_names: + record_class = ds.RecordClass.response + else: + record_class = ds.RecordClass.other + + if prior_id is not None: + if record_class is not ds.RecordClass.parameter: + raise exc.UnprocessableError( + "Priors can only be specified for parameter records" + ) + record.record_info.prior = db.query(ds.Prior).filter_by(id=prior_id).one() + + record.record_info.record_class = record_class + record.record_info.record_type = ds.RecordType.f64_matrix + return record + + +@router.post("/ensembles/{ensemble_id}/records/{name}/file") +async def post_ensemble_record_file( + *, + db: Session = Depends(get_db), + record: ds.Record = Depends(new_record_file), file: UploadFile = File(...), ) -> None: """ Assign an arbitrary file to the given `name` record. """ - ensemble = _get_and_assert_ensemble(db, ensemble_id, name, realization_index) - file_obj = ds.File( filename=file.filename, mimetype=file.content_type, ) if HAS_AZURE_BLOB_STORAGE: - key = f"{name}@{realization_index}@{uuid4()}" + key = f"{record.name}@{record.realization_index}@{uuid4()}" blob = azure_blob_container.get_blob_client(key) await blob.upload_blob(file.file) @@ -66,15 +180,8 @@ async def post_ensemble_record_file( else: file_obj.content = await file.read() - db.add(file_obj) - _create_record( - db, - ensemble, - name, - ds.RecordType.file, - realization_index=realization_index, - file=file_obj, - ) + record.file = file_obj + _create_record(db, record) @router.put("/ensembles/{ensemble_id}/records/{name}/blob") @@ -122,39 +229,25 @@ async def add_block( @router.post("/ensembles/{ensemble_id}/records/{name}/blob") async def create_blob( - *, - db: Session = Depends(get_db), - ensemble_id: UUID, - name: str, - realization_index: Optional[int] = None, + *, db: Session = Depends(get_db), record: ds.Record = Depends(new_record_file) ) -> None: """ Create a record which points to a blob on Azure Blob Storage. """ - - ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() file_obj = ds.File( filename="test", mimetype="mime/type", ) if HAS_AZURE_BLOB_STORAGE: - key = f"{name}@{realization_index}@{uuid4()}" + key = f"{record.name}@{record.realization_index}@{uuid4()}" blob = azure_blob_container.get_blob_client(key) file_obj.az_container = (azure_blob_container.container_name,) file_obj.az_blob = (key,) else: pass - db.add(file_obj) - - _create_record( - db, - ensemble, - name, - ds.RecordType.file, - realization_index=realization_index, - file=file_obj, - ) + record.file = file_obj + _create_record(db, record) @router.patch("/ensembles/{ensemble_id}/records/{name}/blob") @@ -208,11 +301,8 @@ async def finalize_blob( async def post_ensemble_record_matrix( *, db: Session = Depends(get_db), - ensemble_id: UUID, - name: str, - realization_index: Optional[int] = None, + record: ds.Record = Depends(new_record_matrix), content_type: str = Header("application/json"), - prior_id: Optional[UUID] = None, request: Request, ) -> js.RecordOut: """ @@ -224,24 +314,7 @@ async def post_ensemble_record_matrix( ) content_type = "text/csv" - ensemble = _get_and_assert_ensemble(db, ensemble_id, name, realization_index) - is_parameter = name in ensemble.parameter_names - is_response = name in ensemble.response_names - if prior_id is not None and not is_parameter: - raise exc.UnprocessableError( - "Priors can only be specified for parameter records" - ) - labels = None - prior = ( - ( - db.query(ds.Prior) - .filter_by(id=prior_id, experiment_pk=ensemble.experiment_pk) - .one() - ) - if prior_id - else None - ) try: if content_type == "application/json": @@ -262,58 +335,41 @@ async def post_ensemble_record_matrix( else: raise ValueError() except ValueError: - if realization_index is None: - message = f"Ensemble-wide record '{name}' for ensemble '{ensemble_id}' needs to be a matrix" + if record.realization_index is None: + message = f"Ensemble-wide record '{record.name}' for needs to be a matrix" else: - message = f"Forward-model record '{name}' for ensemble '{ensemble_id}', realization {realization_index} needs to be a matrix" + message = f"Forward-model record '{record.name}' for realization {record.realization_index} needs to be a matrix" + raise exc.UnprocessableError(message) # Require that the dimensionality of an ensemble-wide parameter matrix is at least 2 - if realization_index is None and is_parameter: + if ( + record.realization_index is None + and record.record_class is ds.RecordClass.parameter + ): if content.ndim <= 1: raise exc.UnprocessableError( - f"Ensemble-wide parameter record '{name}' for ensemble '{ensemble_id}'" + f"Ensemble-wide parameter record '{record.name}' for ensemble '{record.record_info.ensemble.id}'" "must have dimensionality of at least 2" ) matrix_obj = ds.F64Matrix(content=content.tolist(), labels=labels) - db.add(matrix_obj) - record_class = ds.RecordClass.other - if is_parameter: - record_class = ds.RecordClass.parameter - if is_response: - record_class = ds.RecordClass.response - - return _create_record( - db, - ensemble, - name, - ds.RecordType.f64_matrix, - record_class, - prior, - f64_matrix=matrix_obj, - realization_index=realization_index, - ) + record.f64_matrix = matrix_obj + return _create_record(db, record) @router.put("/ensembles/{ensemble_id}/records/{name}/metadata") async def replace_record_metadata( *, db: Session = Depends(get_db), - ensemble_id: UUID, - name: str, - realization_index: Optional[int] = None, + record: ds.Record = Depends(get_record_by_name), body: Any = Body(...), ) -> None: """ Assign new metadata json """ - if realization_index is None: - record_obj = _get_ensemble_record(db, ensemble_id, name) - else: - record_obj = _get_forward_model_record(db, ensemble_id, name, realization_index) - record_obj._metadata = body + record._metadata = body db.commit() @@ -321,21 +377,14 @@ async def replace_record_metadata( async def patch_record_metadata( *, db: Session = Depends(get_db), - ensemble_id: UUID, - name: str, - realization_index: Optional[int] = None, + record: ds.Record = Depends(get_record_by_name), body: Any = Body(...), ) -> None: """ Update metadata json """ - if realization_index is None: - record_obj = _get_ensemble_record(db, ensemble_id, name) - else: - record_obj = _get_forward_model_record(db, ensemble_id, name, realization_index) - - record_obj._metadata.update(body) - flag_modified(record_obj, "_metadata") + record._metadata.update(body) + flag_modified(record, "_metadata") db.commit() @@ -344,43 +393,26 @@ async def patch_record_metadata( ) async def get_record_metadata( *, - db: Session = Depends(get_db), - ensemble_id: UUID, - name: str, - realization_index: Optional[int] = None, + record: ds.Record = Depends(get_record_by_name), ) -> Mapping[str, Any]: """ Get metadata json """ - if realization_index is None: - bundle = _get_ensemble_record(db, ensemble_id, name) - else: - bundle = _get_forward_model_record(db, ensemble_id, name, realization_index) - - return bundle.metadata_dict + return record.metadata_dict @router.post("/ensembles/{ensemble_id}/records/{name}/observations") async def post_record_observations( *, db: Session = Depends(get_db), - ensemble_id: UUID, - name: str, - realization_index: Optional[int] = None, + record: ds.Record = Depends(get_record_by_name), observation_ids: List[UUID] = Body(...), ) -> None: - - if realization_index is None: - record_obj = _get_ensemble_record(db, ensemble_id, name) - else: - record_obj = _get_forward_model_record(db, ensemble_id, name, realization_index) - observations = ( db.query(ds.Observation).filter(ds.Observation.id.in_(observation_ids)).all() ) if observations: - record_obj.observations = observations - flag_modified(record_obj, "observations") + record.observations = observations db.commit() else: raise exc.UnprocessableError(f"Observations {observation_ids} not found!") @@ -390,15 +422,9 @@ async def post_record_observations( async def get_record_observations( *, db: Session = Depends(get_db), - ensemble_id: UUID, - name: str, - realization_index: Optional[int] = None, + record: ds.Record = Depends(get_record_by_name), ) -> List[js.ObservationOut]: - if realization_index is None: - bundle = _get_ensemble_record(db, ensemble_id, name) - else: - bundle = _get_forward_model_record(db, ensemble_id, name, realization_index) - if bundle.observations: + if record.observations: return [ js.ObservationOut( id=obs.id, @@ -409,7 +435,7 @@ async def get_record_observations( records=[rec.id for rec in obs.records], metadata=obs.metadata_dict, ) - for obs in bundle.observations + for obs in record.observations ] return [] @@ -496,10 +522,9 @@ async def get_record_observations( async def get_ensemble_record( *, db: Session = Depends(get_db), - ensemble_id: UUID, - realization_index: Optional[int] = None, - name: str, + record: ds.Record = Depends(get_record_by_name), accept: str = Header("application/json"), + realization_index: Optional[int] = None, ) -> Any: """ Get record with a given `name`. If `realization_index` is not set, look for @@ -519,21 +544,10 @@ async def get_ensemble_record( ) accept = "text/csv" - if realization_index is None: - bundle = _get_ensemble_record(db, ensemble_id, name) - else: - try: - bundle = _get_forward_model_record(db, ensemble_id, name, realization_index) - realization_index = None - except exc.NotFoundError: - bundle = _get_ensemble_record(db, ensemble_id, name) - - # Only parameter records can be "up-casted" to ensemble-wide - is_matrix = bundle.record_info.record_type == ds.RecordType.f64_matrix - is_parameter = name in bundle.record_info.ensemble.parameter_names - if not is_matrix or not is_parameter: - raise - return await _get_record_data(bundle, accept, realization_index) + new_realization_index = ( + realization_index if record.realization_index is None else None + ) + return await _get_record_data(record, accept, new_realization_index) @router.get("/ensembles/{ensemble_id}/parameters", response_model=List[str]) @@ -603,81 +617,6 @@ def get_ensemble_responses( } -def _get_ensemble_record(db: Session, ensemble_id: UUID, name: str) -> ds.Record: - try: - ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() - return ( - db.query(ds.Record) - .filter_by(realization_index=None) - .join(ds.RecordInfo) - .filter_by( - ensemble_pk=ensemble.pk, - name=name, - ) - .one() - ) - except NoResultFound: - raise exc.NotFoundError( - f"Ensemble-wide record '{name}' for ensemble '{ensemble_id}' not found!" - ) - - -def _get_forward_model_record( - db: Session, ensemble_id: UUID, name: str, realization_index: int -) -> ds.Record: - try: - ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() - return ( - db.query(ds.Record) - .filter_by(realization_index=realization_index) - .join(ds.RecordInfo) - .filter_by( - ensemble_pk=ensemble.pk, - name=name, - ) - .one() - ) - except NoResultFound: - raise exc.NotFoundError( - f"Forward-model record '{name}' for ensemble '{ensemble_id}', realization {realization_index} not found!" - ) - - -def _get_and_assert_ensemble( - db: Session, ensemble_id: UUID, name: str, realization_index: Optional[int] -) -> ds.Ensemble: - """ - Get ensemble and verify that no record with the name `name` exists - """ - ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() - - q = ( - db.query(ds.Record) - .join(ds.RecordInfo) - .filter_by(ensemble_pk=ensemble.pk, name=name) - ) - if realization_index is not None: - if realization_index not in range(ensemble.size) and ensemble.size != -1: - raise exc.ExpectationError( - f"Ensemble '{name}' ('{ensemble_id}') does have a 'size' " - f"of {ensemble.size}. The posted record is targeting " - f"'realization_index' {realization_index} which is out " - f"of bounds." - ) - - q = q.filter( - (ds.Record.realization_index == None) - | (ds.Record.realization_index == realization_index) - ) - - if q.count() > 0: - raise exc.ConflictError( - f"Ensemble-wide record '{name}' for ensemble '{ensemble_id}' already exists" - ) - - return ensemble - - async def _get_record_data( record: ds.Record, accept: Optional[str], realization_index: Optional[int] = None ) -> Response: @@ -747,22 +686,8 @@ async def chunk_generator() -> AsyncGenerator[bytes, None]: def _create_record( db: Session, - ensemble: ds.Ensemble, - name: str, - record_type: ds.RecordType, - record_class: ds.RecordClass = ds.RecordClass.other, - prior: Optional[ds.Prior] = None, - **kwargs: Any, + record: ds.Record, ) -> ds.Record: - record_info = ds.RecordInfo( - ensemble=ensemble, - name=name, - record_class=record_class, - record_type=record_type, - prior=prior, - ) - record = ds.Record(record_info=record_info, **kwargs) - nested = db.begin_nested() try: db.add(record) @@ -772,8 +697,11 @@ def _create_record( # record_info with the same name and ensemble. Try to fetch the # record_info nested.rollback() + record_info = record.record_info old_record_info = ( - db.query(ds.RecordInfo).filter_by(ensemble=ensemble, name=name).one() + db.query(ds.RecordInfo) + .filter_by(ensemble=record_info.ensemble, name=record_info.name) + .one() ) # Check that the parameters match @@ -790,7 +718,12 @@ def _create_record( old_record_type=old_record_info.record_type, ) - record = ds.Record(record_info=old_record_info, **kwargs) + record = ds.Record( + record_info=old_record_info, + f64_matrix=record.f64_matrix, + file=record.file, + realization_index=record.realization_index, + ) db.add(record) db.commit() From fc6ad4e1ab340f782701c7f8b18af24328ab3892 Mon Sep 17 00:00:00 2001 From: Zohar Malamant Date: Mon, 31 May 2021 15:53:02 +0200 Subject: [PATCH 04/43] Move record blob logic into dedicated BlobHandler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Øystein Olai Heggen --- .../dark_storage/database_schema/record.py | 4 + .../dark_storage/endpoints/_records_blob.py | 185 ++++++++++++++++++ src/ert/dark_storage/endpoints/records.py | 147 +++----------- 3 files changed, 220 insertions(+), 116 deletions(-) create mode 100644 src/ert/dark_storage/endpoints/_records_blob.py diff --git a/src/ert/dark_storage/database_schema/record.py b/src/ert/dark_storage/database_schema/record.py index 1400036bc01..6eaa0e68384 100644 --- a/src/ert/dark_storage/database_schema/record.py +++ b/src/ert/dark_storage/database_schema/record.py @@ -56,6 +56,10 @@ def data(self) -> Any: f"The record type {self.record_type} is not yet implemented" ) + @property + def ensemble_pk(self) -> int: + return self.record_info.ensemble_pk + @property def name(self) -> str: return self.record_info.name diff --git a/src/ert/dark_storage/endpoints/_records_blob.py b/src/ert/dark_storage/endpoints/_records_blob.py new file mode 100644 index 00000000000..3d4aec90182 --- /dev/null +++ b/src/ert/dark_storage/endpoints/_records_blob.py @@ -0,0 +1,185 @@ +import io +from typing import Optional, List, Type, AsyncGenerator +from uuid import uuid4, UUID + +import numpy as np +import pandas as pd +from fastapi import ( + Request, + UploadFile, + Depends, +) +from fastapi.logger import logger +from fastapi.responses import Response, StreamingResponse + +from ert_storage import database_schema as ds +from ert_storage.database import Session, get_db, HAS_AZURE_BLOB_STORAGE + +if HAS_AZURE_BLOB_STORAGE: + from ert_storage.database import azure_blob_container + + +class BlobHandler: + def __init__( + self, + db: Session, + name: Optional[str], + ensemble_id: Optional[UUID], + realization_index: Optional[int], + ): + self._db = db + self._name = name + self._ensemble_id = ensemble_id + self._realization_index = realization_index + + async def upload_file( + self, + file: UploadFile, + ) -> ds.File: + return ds.File( + filename=file.filename, + mimetype=file.content_type, + content=await file.read(), + ) + + async def stage_blob( + self, + record: ds.Record, + request: Request, + block_index: int, + ) -> ds.FileBlock: + ensemble = self._db.query(ds.Ensemble).filter_by(id=self._ensemble_id).one() + block_id = str(uuid4()) + + return ds.FileBlock( + ensemble=ensemble, + block_id=block_id, + block_index=block_index, + record_name=self._name, + realization_index=self._realization_index, + content=await request.body(), + ) + + def create_blob(self) -> ds.File: + return ds.File( + filename="test", + mimetype="mime/type", + ) + + async def finalize_blob( + self, submitted_blocks: List[ds.FileBlock], record: ds.Record + ) -> None: + record.file.content = b"".join([block.content for block in submitted_blocks]) + + async def get_content(self, record: ds.Record) -> Response: + assert record.record_type == ds.RecordType.file + return Response( + content=record.file.content, + media_type=record.file.mimetype, + headers={ + "Content-Disposition": f'attachment; filename="{record.file.filename}"' + }, + ) + + +class AzureBlobHandler(BlobHandler): + async def upload_file( + self, + file: UploadFile, + ) -> ds.File: + key = f"{self._name}@{self._realization_index}@{uuid4()}" + blob = azure_blob_container.get_blob_client(key) + await blob.upload_blob(file.file) + + return ds.File( + filename=file.filename, + mimetype=file.content_type, + az_container=azure_blob_container.container_name, + az_blob=key, + ) + + async def stage_blob( + self, + record: ds.Record, + request: Request, + block_index: int, + ) -> ds.FileBlock: + block_id = str(uuid4()) + blob = azure_blob_container.get_blob_client(record.file.az_blob) + await blob.stage_block(block_id, await request.body()) + + return ds.FileBlock( + ensemble_pk=record.ensemble_pk, + block_id=block_id, + block_index=block_index, + record_name=self._name, + realization_index=self._realization_index, + ) + + def create_blob(self) -> ds.File: + key = f"{self._name}@{self._realization_index}@{uuid4()}" + blob = azure_blob_container.get_blob_client(key) + + return ds.File( + filename="test", + mimetype="mime/type", + az_container=azure_blob_container.container_name, + az_blob=key, + ) + + async def finalize_blob( + self, submitted_blocks: List[ds.FileBlock], record: ds.Record + ) -> None: + blob = azure_blob_container.get_blob_client(record.file.az_blob) + block_ids = [ + block.block_id + for block in sorted(submitted_blocks, key=lambda x: x.block_index) + ] + await blob.commit_block_list(block_ids) + + async def get_content(self, record: ds.Record) -> Response: + blob = azure_blob_container.get_blob_client(record.file.az_blob) + download = await blob.download_blob() + + async def chunk_generator() -> AsyncGenerator[bytes, None]: + async for chunk in download.chunks(): + yield chunk + + return StreamingResponse( + chunk_generator(), + media_type=record.file.mimetype, + headers={ + "Content-Disposition": f'attachment; filename="{record.file.filename}"' + }, + ) + + +def get_blob_handler( + *, + db: Session = Depends(get_db), + name: str, + ensemble_id: UUID, + realization_index: Optional[int] = None, +) -> BlobHandler: + blob_handler: Type[BlobHandler] + if HAS_AZURE_BLOB_STORAGE: + blob_handler = AzureBlobHandler + else: + blob_handler = BlobHandler + return blob_handler( + db=db, name=name, ensemble_id=ensemble_id, realization_index=realization_index + ) + + +def get_blob_handler_from_record(db: Session, record: ds.Record) -> BlobHandler: + blob_handler: Type[BlobHandler] + if HAS_AZURE_BLOB_STORAGE: + blob_handler = AzureBlobHandler + else: + blob_handler = BlobHandler + return blob_handler( + db=db, + name=record.name, + ensemble_id=record.record_info.ensemble.id, + realization_index=record.realization_index, + ) diff --git a/src/ert/dark_storage/endpoints/records.py b/src/ert/dark_storage/endpoints/records.py index f165026ebee..3528ff9da4f 100644 --- a/src/ert/dark_storage/endpoints/records.py +++ b/src/ert/dark_storage/endpoints/records.py @@ -20,15 +20,17 @@ from sqlalchemy.exc import IntegrityError from sqlalchemy.orm.exc import NoResultFound from sqlalchemy.orm.attributes import flag_modified -from ert_storage.database import Session, get_db, HAS_AZURE_BLOB_STORAGE, BLOB_CONTAINER +from ert_storage.database import Session, get_db from ert_storage import database_schema as ds, json_schema as js from ert_storage import exceptions as exc +from ert_storage.endpoints._records_blob import ( + get_blob_handler, + get_blob_handler_from_record, + BlobHandler, +) from fastapi.logger import logger -if HAS_AZURE_BLOB_STORAGE: - from ert_storage.database import azure_blob_container - router = APIRouter(tags=["record"]) @@ -160,27 +162,14 @@ def new_record_matrix( async def post_ensemble_record_file( *, db: Session = Depends(get_db), + bh: BlobHandler = Depends(get_blob_handler), record: ds.Record = Depends(new_record_file), file: UploadFile = File(...), ) -> None: """ Assign an arbitrary file to the given `name` record. """ - file_obj = ds.File( - filename=file.filename, - mimetype=file.content_type, - ) - if HAS_AZURE_BLOB_STORAGE: - key = f"{record.name}@{record.realization_index}@{uuid4()}" - blob = azure_blob_container.get_blob_client(key) - await blob.upload_blob(file.file) - - file_obj.az_container = azure_blob_container.container_name - file_obj.az_blob = key - else: - file_obj.content = await file.read() - - record.file = file_obj + record.file = await bh.upload_file(file) _create_record(db, record) @@ -188,65 +177,28 @@ async def post_ensemble_record_file( async def add_block( *, db: Session = Depends(get_db), - ensemble_id: UUID, - name: str, - realization_index: Optional[int] = None, + bh: BlobHandler = Depends(get_blob_handler), + record: ds.Record = Depends(get_record_by_name), request: Request, block_index: int, ) -> None: """ Stage blocks to an existing azure blob record. """ - - ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() - block_id = str(uuid4()) - - file_block_obj = ds.FileBlock( - ensemble=ensemble, - block_id=block_id, - block_index=block_index, - record_name=name, - realization_index=realization_index, - ) - - record_obj = ( - db.query(ds.Record) - .filter_by(realization_index=realization_index) - .join(ds.RecordInfo) - .filter_by(ensemble_pk=ensemble.pk, name=name) - .one() - ) - if HAS_AZURE_BLOB_STORAGE: - key = record_obj.file.az_blob - blob = azure_blob_container.get_blob_client(key) - await blob.stage_block(block_id, await request.body()) - else: - file_block_obj.content = await request.body() - - db.add(file_block_obj) - db.commit() + db.add(await bh.stage_blob(record, request, block_index)) @router.post("/ensembles/{ensemble_id}/records/{name}/blob") async def create_blob( - *, db: Session = Depends(get_db), record: ds.Record = Depends(new_record_file) + *, + db: Session = Depends(get_db), + bh: BlobHandler = Depends(get_blob_handler), + record: ds.Record = Depends(new_record_file), ) -> None: """ Create a record which points to a blob on Azure Blob Storage. """ - file_obj = ds.File( - filename="test", - mimetype="mime/type", - ) - if HAS_AZURE_BLOB_STORAGE: - key = f"{record.name}@{record.realization_index}@{uuid4()}" - blob = azure_blob_container.get_blob_client(key) - file_obj.az_container = (azure_blob_container.container_name,) - file_obj.az_blob = (key,) - else: - pass - - record.file = file_obj + record.file = bh.create_blob() _create_record(db, record) @@ -254,45 +206,22 @@ async def create_blob( async def finalize_blob( *, db: Session = Depends(get_db), - ensemble_id: UUID, - name: str, - realization_index: Optional[int] = None, + bh: BlobHandler = Depends(get_blob_handler), + record: ds.Record = Depends(get_record_by_name), ) -> None: """ Commit all staged blocks to a blob record """ - - ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() - - record_obj = ( - db.query(ds.Record) - .filter_by(realization_index=realization_index) - .join(ds.RecordInfo) - .filter_by(ensemble_pk=ensemble.pk, name=name) - .one() - ) - submitted_blocks = list( db.query(ds.FileBlock) .filter_by( - record_name=name, - ensemble_pk=ensemble.pk, - realization_index=realization_index, + record_name=record.name, + ensemble_pk=record.ensemble_pk, + realization_index=record.realization_index, ) .all() ) - - if HAS_AZURE_BLOB_STORAGE: - key = record_obj.file.az_blob - blob = azure_blob_container.get_blob_client(key) - block_ids = [ - block.block_id - for block in sorted(submitted_blocks, key=lambda x: x.block_index) - ] - await blob.commit_block_list(block_ids) - else: - data = b"".join([block.content for block in submitted_blocks]) - record_obj.file.content = data + await bh.finalize_blob(submitted_blocks, record) @router.post( @@ -522,6 +451,7 @@ async def get_record_observations( async def get_ensemble_record( *, db: Session = Depends(get_db), + bh: BlobHandler = Depends(get_blob_handler), record: ds.Record = Depends(get_record_by_name), accept: str = Header("application/json"), realization_index: Optional[int] = None, @@ -547,7 +477,7 @@ async def get_ensemble_record( new_realization_index = ( realization_index if record.realization_index is None else None ) - return await _get_record_data(record, accept, new_realization_index) + return await _get_record_data(bh, record, accept, new_realization_index) @router.get("/ensembles/{ensemble_id}/parameters", response_model=List[str]) @@ -595,7 +525,8 @@ async def get_record_data( accept = "text/csv" record = db.query(ds.Record).filter_by(id=record_id).one() - return await _get_record_data(record, accept) + bh = get_blob_handler_from_record(db, record) + return await _get_record_data(bh, record, accept) @router.get( @@ -618,7 +549,10 @@ def get_ensemble_responses( async def _get_record_data( - record: ds.Record, accept: Optional[str], realization_index: Optional[int] = None + bh: BlobHandler, + record: ds.Record, + accept: Optional[str], + realization_index: Optional[int] = None, ) -> Response: type_ = record.record_info.record_type if type_ == ds.RecordType.f64_matrix: @@ -659,26 +593,7 @@ async def _get_record_data( else: return content if type_ == ds.RecordType.file: - f = record.file - if f.content is not None: - return Response( - content=f.content, - media_type=f.mimetype, - headers={"Content-Disposition": f'attachment; filename="{f.filename}"'}, - ) - elif f.az_container is not None and f.az_blob is not None: - blob = azure_blob_container.get_blob_client(f.az_blob) - download = await blob.download_blob() - - async def chunk_generator() -> AsyncGenerator[bytes, None]: - async for chunk in download.chunks(): - yield chunk - - return StreamingResponse( - chunk_generator(), - media_type=f.mimetype, - headers={"Content-Disposition": f'attachment; filename="{f.filename}"'}, - ) + return await bh.get_content(record) raise NotImplementedError( f"Getting record data for type {type_} and Accept header {accept} not implemented" ) From 3d8423d60eeb94b5f6abcab88b913f957b89d74d Mon Sep 17 00:00:00 2001 From: Zohar Malamant Date: Tue, 1 Jun 2021 09:55:26 +0200 Subject: [PATCH 05/43] Fix out-of-order chunked blob upload The code used to combine the chunks in upload order rather than block_index order. This commit fixes this issue. --- src/ert/dark_storage/endpoints/records.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/ert/dark_storage/endpoints/records.py b/src/ert/dark_storage/endpoints/records.py index 3528ff9da4f..0db3b94ff4f 100644 --- a/src/ert/dark_storage/endpoints/records.py +++ b/src/ert/dark_storage/endpoints/records.py @@ -219,6 +219,7 @@ async def finalize_blob( ensemble_pk=record.ensemble_pk, realization_index=record.realization_index, ) + .order_by(ds.FileBlock.block_index) .all() ) await bh.finalize_blob(submitted_blocks, record) From 865468ae9ce79555941d609ad20beff97f0ec23c Mon Sep 17 00:00:00 2001 From: Julius Parulek Date: Thu, 20 May 2021 14:06:51 +0200 Subject: [PATCH 06/43] Refactor json outputs with helper functions - Add observationOut helper function - Refactor experiment/prioirs - Add helper function to Update endpoint - Make use of ensemble_ids - Change relative import to absolute - Fix priors into new_record_matrix --- .../database_schema/experiment.py | 5 ++ src/ert/dark_storage/endpoints/experiments.py | 31 ++++----- .../dark_storage/endpoints/observations.py | 69 ++++++++----------- src/ert/dark_storage/endpoints/records.py | 6 +- src/ert/dark_storage/endpoints/updates.py | 11 ++- .../dark_storage/json_schema/experiment.py | 6 +- 6 files changed, 58 insertions(+), 70 deletions(-) diff --git a/src/ert/dark_storage/database_schema/experiment.py b/src/ert/dark_storage/database_schema/experiment.py index 3f8e871d318..9d335eb8321 100644 --- a/src/ert/dark_storage/database_schema/experiment.py +++ b/src/ert/dark_storage/database_schema/experiment.py @@ -1,5 +1,6 @@ import sqlalchemy as sa from uuid import uuid4 +from typing import List from sqlalchemy.sql import func from sqlalchemy.orm import relationship from ert_storage.database import Base @@ -34,3 +35,7 @@ class Experiment(Base, MetadataField): cascade="all, delete-orphan", back_populates="experiment", ) + + @property + def ensemble_ids(self) -> List[UUID]: + return [ens.id for ens in self.ensembles] diff --git a/src/ert/dark_storage/endpoints/experiments.py b/src/ert/dark_storage/endpoints/experiments.py index 3286bf280df..15e07c948aa 100644 --- a/src/ert/dark_storage/endpoints/experiments.py +++ b/src/ert/dark_storage/endpoints/experiments.py @@ -3,6 +3,7 @@ from sqlalchemy.orm.attributes import flag_modified from ert_storage.database import Session, get_db from ert_storage import database_schema as ds, json_schema as js +from ert_storage.endpoints.priors import experiment_priors_to_dict from typing import Any, Mapping, Optional, List @@ -14,16 +15,8 @@ def get_experiments( *, db: Session = Depends(get_db), ) -> List[js.ExperimentOut]: - return [ - js.ExperimentOut( - id=exp.id, - name=exp.name, - ensembles=[ens.id for ens in exp.ensembles], - priors={pri.name: pri.id for pri in exp.priors}, - metadata=exp.metadata_dict, - ) - for exp in db.query(ds.Experiment).all() - ] + experiments = db.query(ds.Experiment).all() + return [_experiment_from_db(exp) for exp in experiments] @router.post("/experiments", response_model=js.ExperimentOut) @@ -48,13 +41,7 @@ def post_experiments( db.add(experiment) db.commit() - return js.ExperimentOut( - id=experiment.id, - name=experiment.name, - ensembles=[ens.id for ens in experiment.ensembles], - priors={pri.name: pri.id for pri in experiment.priors}, - metadata=experiment.metadata_dict, - ) + return _experiment_from_db(experiment) @router.get( @@ -115,3 +102,13 @@ def delete_experiment(*, db: Session = Depends(get_db), experiment_id: UUID) -> experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() db.delete(experiment) db.commit() + + +def _experiment_from_db(exp: ds.Experiment) -> js.ExperimentOut: + return js.ExperimentOut( + id=exp.id, + name=exp.name, + ensemble_ids=exp.ensemble_ids, + priors=experiment_priors_to_dict(exp), + metadata=exp.metadata_dict, + ) diff --git a/src/ert/dark_storage/endpoints/observations.py b/src/ert/dark_storage/endpoints/observations.py index 55f78253cbb..5fedd34e91e 100644 --- a/src/ert/dark_storage/endpoints/observations.py +++ b/src/ert/dark_storage/endpoints/observations.py @@ -1,7 +1,7 @@ from uuid import UUID from fastapi import APIRouter, Depends, Body -from typing import List, Any, Mapping +from typing import List, Any, Mapping, Optional from sqlalchemy.orm.attributes import flag_modified from ert_storage.database import Session, get_db from ert_storage import database_schema as ds, json_schema as js @@ -34,15 +34,7 @@ def post_observation( db.add(obs) db.commit() - return js.ObservationOut( - id=obs.id, - name=obs.name, - x_axis=obs.x_axis, - errors=obs.errors, - values=obs.values, - records=[rec.id for rec in obs.records], - metadata=obs.metadata_dict, - ) + return _observation_from_db(obs) @router.get( @@ -52,18 +44,7 @@ def get_observations( *, db: Session = Depends(get_db), experiment_id: UUID ) -> List[js.ObservationOut]: experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() - return [ - js.ObservationOut( - id=obs.id, - name=obs.name, - x_axis=obs.x_axis, - errors=obs.errors, - values=obs.values, - records=[rec.id for rec in obs.records], - metadata=obs.metadata_dict, - ) - for obs in experiment.observations - ] + return [_observation_from_db(obs) for obs in experiment.observations] @router.get( @@ -82,25 +63,7 @@ def get_observations_with_transformation( ) return [ - js.ObservationOut( - id=obs.id, - name=obs.name, - x_axis=obs.x_axis, - errors=obs.errors, - values=obs.values, - records=[rec.id for rec in obs.records], - metadata=obs.metadata_dict, - transformation=js.ObservationTransformationOut( - id=transformations[obs.name].id, - name=obs.name, - observation_id=obs.id, - scale=transformations[obs.name].scale_list, - active=transformations[obs.name].active_list, - ) - if obs.name in transformations - else None, - ) - for obs in experiment.observations + _observation_from_db(obs, transformations) for obs in experiment.observations ] @@ -146,3 +109,27 @@ async def get_observation_metadata( """ obs = db.query(ds.Observation).filter_by(id=obs_id).one() return obs.metadata_dict + + +def _observation_from_db( + obs: ds.Observation, transformations: Optional[Mapping[str, Any]] = None +) -> js.ObservationOut: + transformation = None + if transformations is not None and obs.name in transformations: + transformation = js.ObservationTransformationOut( + id=transformations[obs.name].id, + name=obs.name, + observation_id=obs.id, + scale=transformations[obs.name].scale_list, + active=transformations[obs.name].active_list, + ) + return js.ObservationOut( + id=obs.id, + name=obs.name, + x_axis=obs.x_axis, + errors=obs.errors, + values=obs.values, + records=[rec.id for rec in obs.records], + metadata=obs.metadata_dict, + transformation=transformation, + ) diff --git a/src/ert/dark_storage/endpoints/records.py b/src/ert/dark_storage/endpoints/records.py index 0db3b94ff4f..ab9fae151e3 100644 --- a/src/ert/dark_storage/endpoints/records.py +++ b/src/ert/dark_storage/endpoints/records.py @@ -136,7 +136,7 @@ def new_record_matrix( *, db: Session = Depends(get_db), record: ds.Record = Depends(new_record), - prior_id: Optional[UUID] = None, + prior: Optional[str] = None, ) -> ds.Record: ensemble = record.record_info.ensemble if record.name in ensemble.parameter_names: @@ -146,12 +146,12 @@ def new_record_matrix( else: record_class = ds.RecordClass.other - if prior_id is not None: + if prior is not None: if record_class is not ds.RecordClass.parameter: raise exc.UnprocessableError( "Priors can only be specified for parameter records" ) - record.record_info.prior = db.query(ds.Prior).filter_by(id=prior_id).one() + record.record_info.prior = db.query(ds.Prior).filter_by(name=prior).one() record.record_info.record_class = record_class record.record_info.record_type = ds.RecordType.f64_matrix diff --git a/src/ert/dark_storage/endpoints/updates.py b/src/ert/dark_storage/endpoints/updates.py index 931d8dfc67f..ea54c836ed1 100644 --- a/src/ert/dark_storage/endpoints/updates.py +++ b/src/ert/dark_storage/endpoints/updates.py @@ -43,12 +43,7 @@ def create_update( db.add_all(observation_transformations) db.commit() - return js.UpdateOut( - id=update_obj.id, - experiment_id=ensemble.experiment.id, - algorithm=update_obj.algorithm, - ensemble_reference_id=ensemble.id, - ) + return _update_from_db(update_obj) @router.get("/updates/{update_id}", response_model=js.UpdateOut) @@ -58,6 +53,10 @@ def get_update( update_id: UUID, ) -> js.UpdateOut: update_obj = db.query(ds.Update).filter_by(id=update_id).one() + return _update_from_db(update_obj) + + +def _update_from_db(update_obj: ds.Update) -> js.UpdateOut: return js.UpdateOut( id=update_obj.id, experiment_id=update_obj.ensemble_reference.experiment.id, diff --git a/src/ert/dark_storage/json_schema/experiment.py b/src/ert/dark_storage/json_schema/experiment.py index 78154b8c438..7e6e5dfc59b 100644 --- a/src/ert/dark_storage/json_schema/experiment.py +++ b/src/ert/dark_storage/json_schema/experiment.py @@ -14,9 +14,9 @@ class ExperimentIn(_Experiment): class ExperimentOut(_Experiment): id: UUID - ensembles: List[UUID] - priors: Mapping[str, UUID] - metadata: Mapping[str, Any] = {} + ensemble_ids: List[UUID] + priors: Mapping[str, dict] + metadata: Mapping[str, Any] class Config: orm_mode = True From cdc50c9d41566635cf0460a99724b86c26c08f44 Mon Sep 17 00:00:00 2001 From: Julius Parulek Date: Wed, 2 Jun 2021 16:35:55 +0200 Subject: [PATCH 07/43] Remove prior endpoint - Migrate prior endpoint tests into experiment - Add get single experiment endpoint --- src/ert/dark_storage/endpoints/__init__.py | 2 - src/ert/dark_storage/endpoints/experiments.py | 53 ++++++++++++++++- src/ert/dark_storage/endpoints/priors.py | 57 ------------------- src/ert/dark_storage/graphql/experiments.py | 2 +- src/ert/dark_storage/graphql/parameters.py | 2 +- 5 files changed, 53 insertions(+), 63 deletions(-) delete mode 100644 src/ert/dark_storage/endpoints/priors.py diff --git a/src/ert/dark_storage/endpoints/__init__.py b/src/ert/dark_storage/endpoints/__init__.py index 64d4ef11e08..b63a5711241 100644 --- a/src/ert/dark_storage/endpoints/__init__.py +++ b/src/ert/dark_storage/endpoints/__init__.py @@ -4,7 +4,6 @@ from .experiments import router as experiments_router from .observations import router as observations_router from .updates import router as updates_router -from .priors import router as priors_router from .compute.misfits import router as misfits_router from .responses import router as response_router @@ -14,6 +13,5 @@ router.include_router(records_router) router.include_router(observations_router) router.include_router(updates_router) -router.include_router(priors_router) router.include_router(misfits_router) router.include_router(response_router) diff --git a/src/ert/dark_storage/endpoints/experiments.py b/src/ert/dark_storage/endpoints/experiments.py index 15e07c948aa..c6f51644b0d 100644 --- a/src/ert/dark_storage/endpoints/experiments.py +++ b/src/ert/dark_storage/endpoints/experiments.py @@ -3,8 +3,20 @@ from sqlalchemy.orm.attributes import flag_modified from ert_storage.database import Session, get_db from ert_storage import database_schema as ds, json_schema as js -from ert_storage.endpoints.priors import experiment_priors_to_dict -from typing import Any, Mapping, Optional, List +from ert_storage.json_schema.prior import ( + PriorConst, + PriorTrig, + PriorNormal, + PriorLogNormal, + PriorErtTruncNormal, + PriorStdNormal, + PriorUniform, + PriorErtDUniform, + PriorLogUniform, + PriorErtErf, + PriorErtDErf, +) +from typing import Any, Mapping, List, Type router = APIRouter(tags=["experiment"]) @@ -19,6 +31,14 @@ def get_experiments( return [_experiment_from_db(exp) for exp in experiments] +@router.get("/experiments/{experiment_id}", response_model=js.ExperimentOut) +def get_experiment_by_id( + *, db: Session = Depends(get_db), experiment_id: UUID +) -> js.ExperimentOut: + experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() + return _experiment_from_db(experiment) + + @router.post("/experiments", response_model=js.ExperimentOut) def post_experiments( *, @@ -104,6 +124,35 @@ def delete_experiment(*, db: Session = Depends(get_db), experiment_id: UUID) -> db.commit() +PRIOR_FUNCTION_TO_PYDANTIC: Mapping[ds.PriorFunction, Type[js.Prior]] = { + ds.PriorFunction.const: PriorConst, + ds.PriorFunction.trig: PriorTrig, + ds.PriorFunction.normal: PriorNormal, + ds.PriorFunction.lognormal: PriorLogNormal, + ds.PriorFunction.ert_truncnormal: PriorErtTruncNormal, + ds.PriorFunction.stdnormal: PriorStdNormal, + ds.PriorFunction.uniform: PriorUniform, + ds.PriorFunction.ert_duniform: PriorErtDUniform, + ds.PriorFunction.loguniform: PriorLogUniform, + ds.PriorFunction.ert_erf: PriorErtErf, + ds.PriorFunction.ert_derf: PriorErtDErf, +} + + +def prior_to_dict(prior: ds.Prior) -> dict: + return ( + PRIOR_FUNCTION_TO_PYDANTIC[prior.function] + .parse_obj( + {key: val for key, val in zip(prior.argument_names, prior.argument_values)} + ) + .dict() + ) + + +def experiment_priors_to_dict(experiment: ds.Experiment) -> Mapping[str, dict]: + return {p.name: prior_to_dict(p) for p in experiment.priors} + + def _experiment_from_db(exp: ds.Experiment) -> js.ExperimentOut: return js.ExperimentOut( id=exp.id, diff --git a/src/ert/dark_storage/endpoints/priors.py b/src/ert/dark_storage/endpoints/priors.py deleted file mode 100644 index 4cc6a023fc0..00000000000 --- a/src/ert/dark_storage/endpoints/priors.py +++ /dev/null @@ -1,57 +0,0 @@ -from uuid import UUID -from typing import Mapping, Type -from fastapi import APIRouter, Depends -from ert_storage.database import Session, get_db -from ert_storage import database_schema as ds, json_schema as js -from ert_storage.json_schema.prior import ( - PriorConst, - PriorTrig, - PriorNormal, - PriorLogNormal, - PriorErtTruncNormal, - PriorStdNormal, - PriorUniform, - PriorErtDUniform, - PriorLogUniform, - PriorErtErf, - PriorErtDErf, -) - -router = APIRouter(tags=["prior"]) - - -@router.get("/experiments/{experiment_id}/priors") -def get_priors( - *, db: Session = Depends(get_db), experiment_id: UUID -) -> Mapping[str, dict]: - experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() - return experiment_priors_to_dict(experiment) - - -PRIOR_FUNCTION_TO_PYDANTIC: Mapping[ds.PriorFunction, Type[js.Prior]] = { - ds.PriorFunction.const: PriorConst, - ds.PriorFunction.trig: PriorTrig, - ds.PriorFunction.normal: PriorNormal, - ds.PriorFunction.lognormal: PriorLogNormal, - ds.PriorFunction.ert_truncnormal: PriorErtTruncNormal, - ds.PriorFunction.stdnormal: PriorStdNormal, - ds.PriorFunction.uniform: PriorUniform, - ds.PriorFunction.ert_duniform: PriorErtDUniform, - ds.PriorFunction.loguniform: PriorLogUniform, - ds.PriorFunction.ert_erf: PriorErtErf, - ds.PriorFunction.ert_derf: PriorErtDErf, -} - - -def prior_to_dict(prior: ds.Prior) -> dict: - return ( - PRIOR_FUNCTION_TO_PYDANTIC[prior.function] - .parse_obj( - {key: val for key, val in zip(prior.argument_names, prior.argument_values)} - ) - .dict() - ) - - -def experiment_priors_to_dict(experiment: ds.Experiment) -> Mapping[str, dict]: - return {p.name: prior_to_dict(p) for p in experiment.priors} diff --git a/src/ert/dark_storage/graphql/experiments.py b/src/ert/dark_storage/graphql/experiments.py index 0c44ebbaca0..e8654a777c1 100644 --- a/src/ert/dark_storage/graphql/experiments.py +++ b/src/ert/dark_storage/graphql/experiments.py @@ -7,7 +7,7 @@ from ert_storage.graphql.ensembles import Ensemble, CreateEnsemble from ert_storage.graphql.updates import Update from ert_storage import database_schema as ds, json_schema as js -from ert_storage.endpoints.priors import experiment_priors_to_dict +from ert_storage.endpoints.experiments import experiment_priors_to_dict if TYPE_CHECKING: diff --git a/src/ert/dark_storage/graphql/parameters.py b/src/ert/dark_storage/graphql/parameters.py index 79e0046a063..45e47fd0f8b 100644 --- a/src/ert/dark_storage/graphql/parameters.py +++ b/src/ert/dark_storage/graphql/parameters.py @@ -4,7 +4,7 @@ from ert_storage.ext.graphene_sqlalchemy import SQLAlchemyObjectType from ert_storage import database_schema as ds -from ert_storage.endpoints.priors import prior_to_dict +from ert_storage.endpoints.experiments import prior_to_dict if TYPE_CHECKING: From 761886aad4fb9ec1d26486c1b12516a06f73e8e0 Mon Sep 17 00:00:00 2001 From: Zohar Malamant Date: Mon, 7 Jun 2021 10:15:56 +0200 Subject: [PATCH 08/43] Specify UUID cache_ok SQLAlchemy complains about cache_ok not existing in UUID and spamming me with warnings. We tell it that UUID is cacheable and that it's all fine. --- src/ert/dark_storage/ext/uuid.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/ert/dark_storage/ext/uuid.py b/src/ert/dark_storage/ext/uuid.py index 89089511c8e..e9527a4cd4b 100644 --- a/src/ert/dark_storage/ext/uuid.py +++ b/src/ert/dark_storage/ext/uuid.py @@ -21,6 +21,7 @@ class UUID(TypeDecorator): """ impl = CHAR + cache_ok = True def load_dialect_impl(self, dialect: Dialect) -> Any: if dialect.name == "postgresql": From 4a76923ea6e3e03d220ff4f1d68a0ca3e58ce6ec Mon Sep 17 00:00:00 2001 From: Zohar Malamant Date: Thu, 3 Jun 2021 12:14:37 +0200 Subject: [PATCH 09/43] Change `metadata` to non-nullable `userdata` The `Metadata` field that points to the `_metadata` column in SQL used to be auto-generated by the Graphene-SQLAlchemy package. Now there's only the `userdata` field. Further, this commit removes the `metadata_dict` property in favour of having the column default to `{}`. --- ...057555237953_make_metadata_not_nullable.py | 96 +++++++++++++++++++ .../dark_storage/database_schema/__init__.py | 1 - .../database_schema/_userdata_field.py | 6 ++ .../dark_storage/database_schema/ensemble.py | 4 +- .../database_schema/experiment.py | 4 +- .../database_schema/metadatafield.py | 12 --- .../database_schema/observation.py | 4 +- src/ert/dark_storage/database_schema/prior.py | 4 +- .../dark_storage/database_schema/record.py | 4 +- src/ert/dark_storage/endpoints/ensembles.py | 28 +++--- src/ert/dark_storage/endpoints/experiments.py | 28 +++--- .../dark_storage/endpoints/observations.py | 28 +++--- src/ert/dark_storage/endpoints/records.py | 28 +++--- .../dark_storage/ext/graphene_sqlalchemy.py | 6 +- src/ert/dark_storage/json_schema/ensemble.py | 4 +- .../dark_storage/json_schema/experiment.py | 2 +- src/ert/dark_storage/json_schema/record.py | 2 +- 17 files changed, 175 insertions(+), 86 deletions(-) create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/057555237953_make_metadata_not_nullable.py create mode 100644 src/ert/dark_storage/database_schema/_userdata_field.py delete mode 100644 src/ert/dark_storage/database_schema/metadatafield.py diff --git a/src/ert/dark_storage/_alembic/alembic/versions/057555237953_make_metadata_not_nullable.py b/src/ert/dark_storage/_alembic/alembic/versions/057555237953_make_metadata_not_nullable.py new file mode 100644 index 00000000000..1c0239e0610 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/057555237953_make_metadata_not_nullable.py @@ -0,0 +1,96 @@ +"""Make metadata not nullable and rename to userdata + +Revision ID: 057555237953 +Revises: 021d7514e351 +Create Date: 2021-06-03 14:04:02.998740 + +""" +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "057555237953" +down_revision = "021d7514e351" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column( + "ensemble", + "metadata", + new_column_name="userdata", + existing_type=postgresql.JSON(astext_type=sa.Text()), + nullable=False, + ) + op.alter_column( + "experiment", + "metadata", + new_column_name="userdata", + existing_type=postgresql.JSON(astext_type=sa.Text()), + nullable=False, + ) + op.alter_column( + "observation", + "metadata", + new_column_name="userdata", + existing_type=postgresql.JSON(astext_type=sa.Text()), + nullable=False, + ) + op.alter_column( + "prior", + "metadata", + new_column_name="userdata", + existing_type=postgresql.JSON(astext_type=sa.Text()), + nullable=False, + ) + op.alter_column( + "record", + "metadata", + new_column_name="userdata", + existing_type=postgresql.JSON(astext_type=sa.Text()), + nullable=False, + ) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column( + "record", + "userdata", + new_column_name="metadata", + existing_type=postgresql.JSON(astext_type=sa.Text()), + nullable=True, + ) + op.alter_column( + "prior", + "userdata", + new_column_name="metadata", + existing_type=postgresql.JSON(astext_type=sa.Text()), + nullable=True, + ) + op.alter_column( + "observation", + "userdata", + new_column_name="metadata", + existing_type=postgresql.JSON(astext_type=sa.Text()), + nullable=True, + ) + op.alter_column( + "experiment", + "userdata", + new_column_name="metadata", + existing_type=postgresql.JSON(astext_type=sa.Text()), + nullable=True, + ) + op.alter_column( + "ensemble", + "userdata", + new_column_name="metadata", + existing_type=postgresql.JSON(astext_type=sa.Text()), + nullable=True, + ) + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/database_schema/__init__.py b/src/ert/dark_storage/database_schema/__init__.py index 213082c8a94..64f1967d69e 100644 --- a/src/ert/dark_storage/database_schema/__init__.py +++ b/src/ert/dark_storage/database_schema/__init__.py @@ -1,4 +1,3 @@ -from .metadatafield import MetadataField from .record_info import RecordInfo, RecordType, RecordClass from .record import Record, F64Matrix, File, FileBlock from .ensemble import Ensemble diff --git a/src/ert/dark_storage/database_schema/_userdata_field.py b/src/ert/dark_storage/database_schema/_userdata_field.py new file mode 100644 index 00000000000..f58155c15c3 --- /dev/null +++ b/src/ert/dark_storage/database_schema/_userdata_field.py @@ -0,0 +1,6 @@ +from typing import Any, Mapping +import sqlalchemy as sa + + +class UserdataField: + userdata = sa.Column(sa.JSON, nullable=False, default=dict) diff --git a/src/ert/dark_storage/database_schema/ensemble.py b/src/ert/dark_storage/database_schema/ensemble.py index e8c28786f01..b7c0c83e6a1 100644 --- a/src/ert/dark_storage/database_schema/ensemble.py +++ b/src/ert/dark_storage/database_schema/ensemble.py @@ -4,12 +4,12 @@ from sqlalchemy.orm import relationship from sqlalchemy.sql import func from ert_storage.database import Base -from .metadatafield import MetadataField from ert_storage.ext.uuid import UUID as UUID from ert_storage.ext.sqlalchemy_arrays import StringArray +from ._userdata_field import UserdataField -class Ensemble(Base, MetadataField): +class Ensemble(Base, UserdataField): __tablename__ = "ensemble" pk = sa.Column(sa.Integer, primary_key=True) diff --git a/src/ert/dark_storage/database_schema/experiment.py b/src/ert/dark_storage/database_schema/experiment.py index 9d335eb8321..eefa1d6831b 100644 --- a/src/ert/dark_storage/database_schema/experiment.py +++ b/src/ert/dark_storage/database_schema/experiment.py @@ -5,10 +5,10 @@ from sqlalchemy.orm import relationship from ert_storage.database import Base from ert_storage.ext.uuid import UUID -from .metadatafield import MetadataField +from ._userdata_field import UserdataField -class Experiment(Base, MetadataField): +class Experiment(Base, UserdataField): __tablename__ = "experiment" pk = sa.Column(sa.Integer, primary_key=True) diff --git a/src/ert/dark_storage/database_schema/metadatafield.py b/src/ert/dark_storage/database_schema/metadatafield.py deleted file mode 100644 index aff5786e673..00000000000 --- a/src/ert/dark_storage/database_schema/metadatafield.py +++ /dev/null @@ -1,12 +0,0 @@ -from typing import Any, Mapping -import sqlalchemy as sa - - -class MetadataField: - _metadata = sa.Column("metadata", sa.JSON, nullable=True) - - @property - def metadata_dict(self) -> Mapping[str, Any]: - if self._metadata is None: - return dict() - return self._metadata diff --git a/src/ert/dark_storage/database_schema/observation.py b/src/ert/dark_storage/database_schema/observation.py index 04ade5dfa72..a9e01cae67e 100644 --- a/src/ert/dark_storage/database_schema/observation.py +++ b/src/ert/dark_storage/database_schema/observation.py @@ -5,7 +5,7 @@ from sqlalchemy.sql import func from sqlalchemy.orm import relationship from uuid import uuid4 -from .metadatafield import MetadataField +from ._userdata_field import UserdataField observation_record_association = sa.Table( "observation_record_association", @@ -15,7 +15,7 @@ ) -class Observation(Base, MetadataField): +class Observation(Base, UserdataField): __tablename__ = "observation" __table_args__ = ( sa.UniqueConstraint("name", "experiment_pk", name="uq_observation_name"), diff --git a/src/ert/dark_storage/database_schema/prior.py b/src/ert/dark_storage/database_schema/prior.py index 781c45eacde..2eb55f4cede 100644 --- a/src/ert/dark_storage/database_schema/prior.py +++ b/src/ert/dark_storage/database_schema/prior.py @@ -3,10 +3,10 @@ from sqlalchemy.sql import func from sqlalchemy.orm import relationship from uuid import uuid4 -from .metadatafield import MetadataField from ert_storage.database import Base from ert_storage.ext.uuid import UUID from ert_storage.ext.sqlalchemy_arrays import StringArray, FloatArray +from ._userdata_field import UserdataField class PriorFunction(Enum): @@ -23,7 +23,7 @@ class PriorFunction(Enum): ert_derf = 11 -class Prior(Base, MetadataField): +class Prior(Base, UserdataField): __tablename__ = "prior" pk = sa.Column(sa.Integer, primary_key=True) diff --git a/src/ert/dark_storage/database_schema/record.py b/src/ert/dark_storage/database_schema/record.py index 6eaa0e68384..385d822d44b 100644 --- a/src/ert/dark_storage/database_schema/record.py +++ b/src/ert/dark_storage/database_schema/record.py @@ -10,12 +10,12 @@ from ert_storage.ext.uuid import UUID from ert_storage.database import Base -from .metadatafield import MetadataField +from ._userdata_field import UserdataField from .observation import observation_record_association from .record_info import RecordType, RecordClass -class Record(Base, MetadataField): +class Record(Base, UserdataField): __tablename__ = "record" pk = sa.Column(sa.Integer, primary_key=True) diff --git a/src/ert/dark_storage/endpoints/ensembles.py b/src/ert/dark_storage/endpoints/ensembles.py index e20e2d9244f..cba3e70cd6b 100644 --- a/src/ert/dark_storage/endpoints/ensembles.py +++ b/src/ert/dark_storage/endpoints/ensembles.py @@ -20,7 +20,7 @@ def post_ensemble( response_names=ens_in.response_names, experiment=experiment, size=ens_in.size, - _metadata=ens_in.metadata, + userdata=ens_in.userdata, ) db.add(ens) @@ -37,45 +37,45 @@ def get_ensemble(*, db: Session = Depends(get_db), ensemble_id: UUID) -> ds.Ense return db.query(ds.Ensemble).filter_by(id=ensemble_id).one() -@router.put("/ensembles/{ensemble_id}/metadata") -async def replace_ensemble_metadata( +@router.put("/ensembles/{ensemble_id}/userdata") +async def replace_ensemble_userdata( *, db: Session = Depends(get_db), ensemble_id: UUID, body: Any = Body(...), ) -> None: """ - Assign new metadata json + Assign new userdata json """ ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() - ensemble._metadata = body + ensemble.userdata = body db.commit() -@router.patch("/ensembles/{ensemble_id}/metadata") -async def patch_ensemble_metadata( +@router.patch("/ensembles/{ensemble_id}/userdata") +async def patch_ensemble_userdata( *, db: Session = Depends(get_db), ensemble_id: UUID, body: Any = Body(...), ) -> None: """ - Update metadata json + Update userdata json """ ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() - ensemble._metadata.update(body) - flag_modified(ensemble, "_metadata") + ensemble.userdata.update(body) + flag_modified(ensemble, "userdata") db.commit() -@router.get("/ensembles/{ensemble_id}/metadata", response_model=Mapping[str, Any]) -async def get_ensemble_metadata( +@router.get("/ensembles/{ensemble_id}/userdata", response_model=Mapping[str, Any]) +async def get_ensemble_userdata( *, db: Session = Depends(get_db), ensemble_id: UUID, ) -> Mapping[str, Any]: """ - Get metadata json + Get userdata json """ ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() - return ensemble.metadata_dict + return ensemble.userdata diff --git a/src/ert/dark_storage/endpoints/experiments.py b/src/ert/dark_storage/endpoints/experiments.py index c6f51644b0d..2218f3db007 100644 --- a/src/ert/dark_storage/endpoints/experiments.py +++ b/src/ert/dark_storage/endpoints/experiments.py @@ -73,48 +73,48 @@ def get_experiment_ensembles( return db.query(ds.Ensemble).join(ds.Experiment).filter_by(id=experiment_id).all() -@router.put("/experiments/{experiment_id}/metadata") -async def replace_experiment_metadata( +@router.put("/experiments/{experiment_id}/userdata") +async def replace_experiment_userdata( *, db: Session = Depends(get_db), experiment_id: UUID, body: Any = Body(...), ) -> None: """ - Assign new metadata json + Assign new userdata json """ experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() - experiment._metadata = body + experiment.userdata = body db.commit() -@router.patch("/experiments/{experiment_id}/metadata") -async def patch_experiment_metadata( +@router.patch("/experiments/{experiment_id}/userdata") +async def patch_experiment_userdata( *, db: Session = Depends(get_db), experiment_id: UUID, body: Any = Body(...), ) -> None: """ - Update metadata json + Update userdata json """ experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() - experiment._metadata.update(body) - flag_modified(experiment, "_metadata") + experiment.userdata.update(body) + flag_modified(experiment, "userdata") db.commit() -@router.get("/experiments/{experiment_id}/metadata", response_model=Mapping[str, Any]) -async def get_experiment_metadata( +@router.get("/experiments/{experiment_id}/userdata", response_model=Mapping[str, Any]) +async def get_experiment_userdata( *, db: Session = Depends(get_db), experiment_id: UUID, ) -> Mapping[str, Any]: """ - Get metadata json + Get userdata json """ experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() - return experiment.metadata_dict + return experiment.userdata @router.delete("/experiments/{experiment_id}") @@ -159,5 +159,5 @@ def _experiment_from_db(exp: ds.Experiment) -> js.ExperimentOut: name=exp.name, ensemble_ids=exp.ensemble_ids, priors=experiment_priors_to_dict(exp), - metadata=exp.metadata_dict, + userdata=exp.userdata, ) diff --git a/src/ert/dark_storage/endpoints/observations.py b/src/ert/dark_storage/endpoints/observations.py index 5fedd34e91e..37cf5232727 100644 --- a/src/ert/dark_storage/endpoints/observations.py +++ b/src/ert/dark_storage/endpoints/observations.py @@ -67,48 +67,48 @@ def get_observations_with_transformation( ] -@router.put("/observations/{obs_id}/metadata") -async def replace_observation_metadata( +@router.put("/observations/{obs_id}/userdata") +async def replace_observation_userdata( *, db: Session = Depends(get_db), obs_id: UUID, body: Any = Body(...), ) -> None: """ - Assign new metadata json + Assign new userdata json """ obs = db.query(ds.Observation).filter_by(id=obs_id).one() - obs._metadata = body + obs.userdata = body db.commit() -@router.patch("/observations/{obs_id}/metadata") -async def patch_observation_metadata( +@router.patch("/observations/{obs_id}/userdata") +async def patch_observation_userdata( *, db: Session = Depends(get_db), obs_id: UUID, body: Any = Body(...), ) -> None: """ - Update metadata json + Update userdata json """ obs = db.query(ds.Observation).filter_by(id=obs_id).one() - obs._metadata.update(body) - flag_modified(obs, "_metadata") + obs.userdata.update(body) + flag_modified(obs, "userdata") db.commit() -@router.get("/observations/{obs_id}/metadata", response_model=Mapping[str, Any]) -async def get_observation_metadata( +@router.get("/observations/{obs_id}/userdata", response_model=Mapping[str, Any]) +async def get_observation_userdata( *, db: Session = Depends(get_db), obs_id: UUID, ) -> Mapping[str, Any]: """ - Get metadata json + Get userdata json """ obs = db.query(ds.Observation).filter_by(id=obs_id).one() - return obs.metadata_dict + return obs.userdata def _observation_from_db( @@ -130,6 +130,6 @@ def _observation_from_db( errors=obs.errors, values=obs.values, records=[rec.id for rec in obs.records], - metadata=obs.metadata_dict, + userdata=obs.userdata, transformation=transformation, ) diff --git a/src/ert/dark_storage/endpoints/records.py b/src/ert/dark_storage/endpoints/records.py index ab9fae151e3..5ce7f55b286 100644 --- a/src/ert/dark_storage/endpoints/records.py +++ b/src/ert/dark_storage/endpoints/records.py @@ -289,46 +289,46 @@ async def post_ensemble_record_matrix( return _create_record(db, record) -@router.put("/ensembles/{ensemble_id}/records/{name}/metadata") -async def replace_record_metadata( +@router.put("/ensembles/{ensemble_id}/records/{name}/userdata") +async def replace_record_userdata( *, db: Session = Depends(get_db), record: ds.Record = Depends(get_record_by_name), body: Any = Body(...), ) -> None: """ - Assign new metadata json + Assign new userdata json """ - record._metadata = body + record.userdata = body db.commit() -@router.patch("/ensembles/{ensemble_id}/records/{name}/metadata") -async def patch_record_metadata( +@router.patch("/ensembles/{ensemble_id}/records/{name}/userdata") +async def patch_record_userdata( *, db: Session = Depends(get_db), record: ds.Record = Depends(get_record_by_name), body: Any = Body(...), ) -> None: """ - Update metadata json + Update userdata json """ - record._metadata.update(body) - flag_modified(record, "_metadata") + record.userdata.update(body) + flag_modified(record, "userdata") db.commit() @router.get( - "/ensembles/{ensemble_id}/records/{name}/metadata", response_model=Mapping[str, Any] + "/ensembles/{ensemble_id}/records/{name}/userdata", response_model=Mapping[str, Any] ) -async def get_record_metadata( +async def get_record_userdata( *, record: ds.Record = Depends(get_record_by_name), ) -> Mapping[str, Any]: """ - Get metadata json + Get userdata json """ - return record.metadata_dict + return record.userdata @router.post("/ensembles/{ensemble_id}/records/{name}/observations") @@ -363,7 +363,7 @@ async def get_record_observations( errors=obs.errors, values=obs.values, records=[rec.id for rec in obs.records], - metadata=obs.metadata_dict, + userdata=obs.userdata, ) for obs in record.observations ] diff --git a/src/ert/dark_storage/ext/graphene_sqlalchemy.py b/src/ert/dark_storage/ext/graphene_sqlalchemy.py index 485d9b087fb..6dcaa671e8c 100644 --- a/src/ert/dark_storage/ext/graphene_sqlalchemy.py +++ b/src/ert/dark_storage/ext/graphene_sqlalchemy.py @@ -67,7 +67,7 @@ def __init_subclass_with_meta__( resolver: Callable = None, arguments: Dict[str, "Argument"] = None, _meta: Optional[ObjectTypeOptions] = None, - **options: Any + **options: Any, ) -> None: if not _meta: _meta = MutationOptions(cls) @@ -117,7 +117,7 @@ def Field( description: Optional[str] = None, deprecation_reason: Optional[str] = None, required: bool = False, - **kwargs: Any + **kwargs: Any, ) -> "graphene.types.field.Field": """Mount instance of mutation Field.""" return Field( @@ -128,5 +128,5 @@ def Field( description=description or cls._meta.description, deprecation_reason=deprecation_reason, required=required, - **kwargs + **kwargs, ) diff --git a/src/ert/dark_storage/json_schema/ensemble.py b/src/ert/dark_storage/json_schema/ensemble.py index d1f4fba3907..b12681c3dc5 100644 --- a/src/ert/dark_storage/json_schema/ensemble.py +++ b/src/ert/dark_storage/json_schema/ensemble.py @@ -11,7 +11,7 @@ class _Ensemble(BaseModel): class EnsembleIn(_Ensemble): update_id: Optional[UUID] = None - metadata: Optional[Any] + userdata: Mapping[str, Any] = {} @root_validator def _check_names_no_overlap(cls, values: Mapping[str, Any]) -> Mapping[str, Any]: @@ -29,7 +29,7 @@ class EnsembleOut(_Ensemble): children: List[UUID] = Field(alias="child_ensemble_ids") parent: Optional[UUID] = Field(alias="parent_ensemble_id") experiment_id: Optional[UUID] = None - metadata: Mapping[str, Any] = Field(alias="metadata_dict") + userdata: Mapping[str, Any] class Config: orm_mode = True diff --git a/src/ert/dark_storage/json_schema/experiment.py b/src/ert/dark_storage/json_schema/experiment.py index 7e6e5dfc59b..9cbe2c77f4b 100644 --- a/src/ert/dark_storage/json_schema/experiment.py +++ b/src/ert/dark_storage/json_schema/experiment.py @@ -16,7 +16,7 @@ class ExperimentOut(_Experiment): id: UUID ensemble_ids: List[UUID] priors: Mapping[str, dict] - metadata: Mapping[str, Any] + userdata: Mapping[str, Any] class Config: orm_mode = True diff --git a/src/ert/dark_storage/json_schema/record.py b/src/ert/dark_storage/json_schema/record.py index 646608ddb90..92977e1c36d 100644 --- a/src/ert/dark_storage/json_schema/record.py +++ b/src/ert/dark_storage/json_schema/record.py @@ -10,7 +10,7 @@ class _Record(BaseModel): class RecordOut(_Record): id: UUID name: str - metadata: Mapping[str, Any] = Field(alias="metadata_dict") + userdata: Mapping[str, Any] class Config: orm_mode = True From 2b060e4cde4a5670893ebff8d230f2adaf3d8580 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=85smund=20Birkeland?= Date: Wed, 16 Jun 2021 13:53:45 +0200 Subject: [PATCH 10/43] Add types-requests to test requirements Modify testclient post, get, .. to comply with requests types --- src/ert/dark_storage/testing/testclient.py | 91 +++------------------- 1 file changed, 10 insertions(+), 81 deletions(-) diff --git a/src/ert/dark_storage/testing/testclient.py b/src/ert/dark_storage/testing/testclient.py index cf4dff13385..a8efcd9e58e 100644 --- a/src/ert/dark_storage/testing/testclient.py +++ b/src/ert/dark_storage/testing/testclient.py @@ -67,24 +67,12 @@ def __init__( def get( self, url: str, - params: Params = None, - headers: MutableMapping[str, str] = None, - cookies: Cookies = None, - files: FileType = None, - timeout: TimeOut = None, - allow_redirects: bool = None, - stream: bool = None, check_status_code: Optional[int] = 200, + **kwargs: Any, ) -> requests.Response: resp = self.http_client.get( url, - params=params, - headers=headers, - cookies=cookies, - files=files, - timeout=timeout, - allow_redirects=allow_redirects, - stream=stream, + **kwargs, ) self._check(check_status_code, resp) return resp @@ -92,57 +80,22 @@ def get( def post( self, url: str, - params: Params = None, - data: DataType = None, - headers: MutableMapping[str, str] = None, - cookies: Cookies = None, - files: FileType = None, - timeout: TimeOut = None, - allow_redirects: bool = None, - stream: bool = None, - json: Any = None, check_status_code: Optional[int] = 200, + **kwargs: Any, ) -> requests.Response: - resp = self.http_client.post( - url, - params=params, - data=data, - headers=headers, - cookies=cookies, - files=files, - timeout=timeout, - allow_redirects=allow_redirects, - stream=stream, - json=json, - ) + resp = self.http_client.post(url, **kwargs) self._check(check_status_code, resp) return resp def put( self, url: str, - params: Params = None, - data: DataType = None, - headers: MutableMapping[str, str] = None, - cookies: Cookies = None, - files: FileType = None, - timeout: TimeOut = None, - allow_redirects: bool = None, - stream: bool = None, - json: Any = None, check_status_code: Optional[int] = 200, + **kwargs: Any, ) -> requests.Response: resp = self.http_client.put( url, - params=params, - data=data, - headers=headers, - cookies=cookies, - files=files, - timeout=timeout, - allow_redirects=allow_redirects, - stream=stream, - json=json, + **kwargs, ) self._check(check_status_code, resp) return resp @@ -150,28 +103,12 @@ def put( def patch( self, url: str, - params: Params = None, - data: DataType = None, - headers: MutableMapping[str, str] = None, - cookies: Cookies = None, - files: FileType = None, - timeout: TimeOut = None, - allow_redirects: bool = None, - stream: bool = None, - json: Any = None, check_status_code: Optional[int] = 200, + **kwargs: Any, ) -> requests.Response: resp = self.http_client.patch( url, - params=params, - data=data, - headers=headers, - cookies=cookies, - files=files, - timeout=timeout, - allow_redirects=allow_redirects, - stream=stream, - json=json, + **kwargs, ) self._check(check_status_code, resp) return resp @@ -179,20 +116,12 @@ def patch( def delete( self, url: str, - params: Params = None, - headers: MutableMapping[str, str] = None, - cookies: Cookies = None, - timeout: TimeOut = None, - allow_redirects: bool = None, check_status_code: Optional[int] = 200, + **kwargs: Any, ) -> requests.Response: resp = self.http_client.delete( url, - params=params, - headers=headers, - cookies=cookies, - timeout=timeout, - allow_redirects=allow_redirects, + **kwargs, ) self._check(check_status_code, resp) return resp From 290cd79319a06c1115c832a85b3ac74383a728e8 Mon Sep 17 00:00:00 2001 From: xjules Date: Wed, 16 Jun 2021 13:43:26 +0200 Subject: [PATCH 11/43] Replace metadata with userdata in observation --- src/ert/dark_storage/json_schema/observation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ert/dark_storage/json_schema/observation.py b/src/ert/dark_storage/json_schema/observation.py index 5ea986f468f..43083928bac 100644 --- a/src/ert/dark_storage/json_schema/observation.py +++ b/src/ert/dark_storage/json_schema/observation.py @@ -36,7 +36,7 @@ class ObservationIn(_Observation): class ObservationOut(_Observation): id: UUID transformation: Optional[ObservationTransformationOut] = None - metadata: Mapping[str, Any] = {} + userdata: Mapping[str, Any] = {} class Config: orm_mode = True From 12940abdcc34dae99951e665ad180d4112230f51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=85smund=20Birkeland?= Date: Mon, 5 Jul 2021 13:03:35 +0200 Subject: [PATCH 12/43] Allow fetching records by rel_idx for ens-wide matrix records --- src/ert/dark_storage/endpoints/records.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/ert/dark_storage/endpoints/records.py b/src/ert/dark_storage/endpoints/records.py index 5ce7f55b286..cd8ed3acf15 100644 --- a/src/ert/dark_storage/endpoints/records.py +++ b/src/ert/dark_storage/endpoints/records.py @@ -67,7 +67,6 @@ def get_record_by_name( .join(ds.RecordInfo) .filter_by( name=name, - record_class=ds.RecordClass.parameter, record_type=ds.RecordType.f64_matrix, ) .join(ds.Ensemble) From 8fa8fc64ac3cf52fa27d7bc4374cbf08ba03e54a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=85smund=20Birkeland?= Date: Mon, 9 Aug 2021 12:09:51 +0200 Subject: [PATCH 13/43] Add application/x-parquet support for uploading matrix records --- .../dark_storage/endpoints/compute/misfits.py | 2 +- src/ert/dark_storage/endpoints/records.py | 53 +++++++++++++------ 2 files changed, 39 insertions(+), 16 deletions(-) diff --git a/src/ert/dark_storage/endpoints/compute/misfits.py b/src/ert/dark_storage/endpoints/compute/misfits.py index e9f82513afa..38678c97fc1 100644 --- a/src/ert/dark_storage/endpoints/compute/misfits.py +++ b/src/ert/dark_storage/endpoints/compute/misfits.py @@ -18,7 +18,7 @@ "/compute/misfits", responses={ status.HTTP_200_OK: { - "content": {"application/x-dataframe": {}}, + "content": {"text/csv": {}}, "description": "Return misfits as csv, where columns are realizations.", } }, diff --git a/src/ert/dark_storage/endpoints/records.py b/src/ert/dark_storage/endpoints/records.py index cd8ed3acf15..1d95ad0f777 100644 --- a/src/ert/dark_storage/endpoints/records.py +++ b/src/ert/dark_storage/endpoints/records.py @@ -261,6 +261,14 @@ async def post_ensemble_record_matrix( [str(v) for v in df.columns.values], [str(v) for v in df.index.values], ] + elif content_type == "application/x-parquet": + stream = io.BytesIO(await request.body()) + df = pd.read_parquet(stream) + content = df.values + labels = [ + [v for v in df.columns.values], + [v for v in df.index.values], + ] else: raise ValueError() except ValueError: @@ -548,6 +556,25 @@ def get_ensemble_responses( } +def _get_dataframe( + content: Any, record: ds.Record, realization_index: Optional[int] +) -> pd.DataFrame: + data = pd.DataFrame(content) + labels = record.f64_matrix.labels + if labels is not None and realization_index is None: + data.columns = labels[0] + data.index = labels[1] + elif labels is not None and realization_index is not None: + # The output is such that rows are realizations. Because + # `content` is a 1d list in this case, it treats each element as + # its own row. We transpose the data so that all of the data + # falls on the same row. + data = data.T + data.columns = labels[0] + data.index = [realization_index] + return data + + async def _get_record_data( bh: BlobHandler, record: ds.Record, @@ -569,26 +596,22 @@ async def _get_record_data( return Response( content=stream.getvalue(), - media_type="application/x-numpy", + media_type=accept, ) if accept == "text/csv": - data = pd.DataFrame(content) - labels = record.f64_matrix.labels - if labels is not None and realization_index is None: - data.columns = labels[0] - data.index = labels[1] - elif labels is not None and realization_index is not None: - # The output is such that rows are realizations. Because - # `content` is a 1d list in this case, it treats each element as - # its own row. We transpose the data so that all of the data - # falls on the same row. - data = data.T - data.columns = labels[0] - data.index = [realization_index] + data = _get_dataframe(content, record, realization_index) return Response( content=data.to_csv().encode(), - media_type="text/csv", + media_type=accept, + ) + if accept == "application/x-parquet": + data = _get_dataframe(content, record, realization_index) + stream = io.BytesIO() + data.to_parquet(stream) + return Response( + content=stream.getvalue(), + media_type=accept, ) else: return content From b5f3b448b0977b44d605acc95ffad17e69876dd0 Mon Sep 17 00:00:00 2001 From: xjules Date: Mon, 16 Aug 2021 14:26:33 +0200 Subject: [PATCH 14/43] Add back_populates to Ensemble record_infos to allow many-to-one semantics --- src/ert/dark_storage/database_schema/ensemble.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/ert/dark_storage/database_schema/ensemble.py b/src/ert/dark_storage/database_schema/ensemble.py index b7c0c83e6a1..2f6a5bf2757 100644 --- a/src/ert/dark_storage/database_schema/ensemble.py +++ b/src/ert/dark_storage/database_schema/ensemble.py @@ -26,6 +26,7 @@ class Ensemble(Base, UserdataField): foreign_keys="[RecordInfo.ensemble_pk]", cascade="all, delete-orphan", lazy="dynamic", + back_populates="ensemble", ) experiment_pk = sa.Column( sa.Integer, sa.ForeignKey("experiment.pk"), nullable=False From 2fd9882cdc23d3dce5eafd1299e3bed02f62fe3e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=98ystein=20Olai=20Heggen?= Date: Thu, 26 Aug 2021 09:12:28 +0200 Subject: [PATCH 15/43] Replace sys.exit with Exception when ENV_RDBMS missing --- src/ert/dark_storage/database.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/ert/dark_storage/database.py b/src/ert/dark_storage/database.py index 9ff6c7f8056..ec2d412ff1f 100644 --- a/src/ert/dark_storage/database.py +++ b/src/ert/dark_storage/database.py @@ -1,6 +1,5 @@ import os -import sys -from typing import Any, Callable, Type +from typing import Any from fastapi import Depends from sqlalchemy import create_engine from sqlalchemy.ext.declarative import declarative_base @@ -12,11 +11,14 @@ ENV_BLOB = "ERT_STORAGE_AZURE_CONNECTION_STRING" ENV_BLOB_CONTAINER = "ERT_STORAGE_AZURE_BLOB_CONTAINER" -if ENV_RDBMS not in os.environ: - sys.exit(f"Environment variable '{ENV_RDBMS}' not set") +def get_env_rdbms() -> str: + if ENV_RDBMS not in os.environ: + raise EnvironmentError(f"Environment variable '{ENV_RDBMS}' not set") + return os.environ[ENV_RDBMS] -URI_RDBMS = os.environ[ENV_RDBMS] + +URI_RDBMS = get_env_rdbms() IS_SQLITE = URI_RDBMS.startswith("sqlite") IS_POSTGRES = URI_RDBMS.startswith("postgres") HAS_AZURE_BLOB_STORAGE = ENV_BLOB in os.environ From aaccf35bb76b14b40fe520df8cc2b27d7f21ff1b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=98ystein=20Olai=20Heggen?= Date: Wed, 15 Sep 2021 09:58:29 +0200 Subject: [PATCH 16/43] lazy load database configuration --- src/ert/dark_storage/__main__.py | 4 +- src/ert/dark_storage/_alembic/alembic/env.py | 6 +- src/ert/dark_storage/app.py | 9 +- src/ert/dark_storage/database.py | 96 +++++++++++++------ .../dark_storage/database_schema/record.py | 1 - .../dark_storage/endpoints/_records_blob.py | 12 +-- src/ert/dark_storage/endpoints/responses.py | 5 +- src/ert/dark_storage/ext/sqlalchemy_arrays.py | 11 ++- src/ert/dark_storage/testing/testclient.py | 24 ++--- 9 files changed, 97 insertions(+), 71 deletions(-) diff --git a/src/ert/dark_storage/__main__.py b/src/ert/dark_storage/__main__.py index 2c1327d8795..2a9efac7753 100644 --- a/src/ert/dark_storage/__main__.py +++ b/src/ert/dark_storage/__main__.py @@ -44,7 +44,7 @@ def run_alembic(args: List[str]) -> None: """ Forward arguments to alembic """ - from alembic.config import main + from alembic.config import main as alembic_main dbkey = "ERT_STORAGE_DATABASE_URL" dburl = os.getenv(dbkey) @@ -67,7 +67,7 @@ def run_alembic(args: List[str]) -> None: ] try: - main(argv=argv, prog="ert-storage alembic") + alembic_main(argv=argv, prog="ert-storage alembic") except FileNotFoundError as exc: if os.path.basename(exc.filename) == "script.py.mako": sys.exit( diff --git a/src/ert/dark_storage/_alembic/alembic/env.py b/src/ert/dark_storage/_alembic/alembic/env.py index ffb52010677..d6b1e95ce35 100644 --- a/src/ert/dark_storage/_alembic/alembic/env.py +++ b/src/ert/dark_storage/_alembic/alembic/env.py @@ -6,7 +6,7 @@ from alembic import context -from ert_storage.database import ENV_RDBMS +from ert_storage.database import database_config from ert_storage.database_schema import Base # this is the Alembic Config object, which provides @@ -43,7 +43,7 @@ def run_migrations_offline(): """ # Alembic uses sprintf somewhere. Escape the '%'s so that alembic doesn't # think they're part of the format string. - url = os.environ[ENV_RDBMS].replace("%", "%%") + url = os.environ[database_config.ENV_RDBMS].replace("%", "%%") context.configure( url=url, target_metadata=target_metadata, @@ -64,7 +64,7 @@ def run_migrations_online(): """ # Alembic uses sprintf somewhere. Escape the '%'s so that alembic doesn't # think they're part of the format string. - url = os.environ[ENV_RDBMS].replace("%", "%%") + url = os.environ[database_config.ENV_RDBMS].replace("%", "%%") config.set_section_option(config.config_ini_section, "sqlalchemy.url", str(url)) connectable = engine_from_config( diff --git a/src/ert/dark_storage/app.py b/src/ert/dark_storage/app.py index b646c26e197..e9a65a59119 100644 --- a/src/ert/dark_storage/app.py +++ b/src/ert/dark_storage/app.py @@ -3,7 +3,6 @@ from typing import Any from fastapi import FastAPI, Request, status from fastapi.responses import Response, RedirectResponse -from starlette.graphql import GraphQLApp from ert_storage.endpoints import router as endpoints_router from ert_storage.graphql import router as graphql_router @@ -51,13 +50,13 @@ def render(self, content: Any) -> bytes: @app.on_event("startup") async def initialize_database() -> None: - from ert_storage.database import engine, IS_SQLITE, HAS_AZURE_BLOB_STORAGE + from ert_storage.database import database_config from ert_storage.database_schema import Base - if IS_SQLITE: + if database_config.IS_SQLITE: # Our SQLite backend doesn't support migrations, so create the database on the fly. - Base.metadata.create_all(bind=engine) - if HAS_AZURE_BLOB_STORAGE: + Base.metadata.create_all(bind=database_config.engine) + if database_config.HAS_AZURE_BLOB_STORAGE: from ert_storage.database import create_container_if_not_exist await create_container_if_not_exist() diff --git a/src/ert/dark_storage/database.py b/src/ert/dark_storage/database.py index ec2d412ff1f..785e24a1554 100644 --- a/src/ert/dark_storage/database.py +++ b/src/ert/dark_storage/database.py @@ -2,45 +2,81 @@ from typing import Any from fastapi import Depends from sqlalchemy import create_engine +from sqlalchemy.engine import Engine from sqlalchemy.ext.declarative import declarative_base -from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm import sessionmaker, Session from ert_storage.security import security -ENV_RDBMS = "ERT_STORAGE_DATABASE_URL" -ENV_BLOB = "ERT_STORAGE_AZURE_CONNECTION_STRING" -ENV_BLOB_CONTAINER = "ERT_STORAGE_AZURE_BLOB_CONTAINER" - - -def get_env_rdbms() -> str: - if ENV_RDBMS not in os.environ: - raise EnvironmentError(f"Environment variable '{ENV_RDBMS}' not set") - return os.environ[ENV_RDBMS] - - -URI_RDBMS = get_env_rdbms() -IS_SQLITE = URI_RDBMS.startswith("sqlite") -IS_POSTGRES = URI_RDBMS.startswith("postgres") -HAS_AZURE_BLOB_STORAGE = ENV_BLOB in os.environ -BLOB_CONTAINER = os.getenv(ENV_BLOB_CONTAINER, "ert") - - -if IS_SQLITE: - engine = create_engine(URI_RDBMS, connect_args={"check_same_thread": False}) -else: - engine = create_engine(URI_RDBMS, pool_size=50, max_overflow=100) -Session = sessionmaker(autocommit=False, autoflush=False, bind=engine) - - +class DatabaseConfig: + def __init__(self) -> None: + self.ENV_RDBMS = "ERT_STORAGE_DATABASE_URL" + self.ENV_BLOB = "ERT_STORAGE_AZURE_CONNECTION_STRING" + self.ENV_BLOB_CONTAINER = "ERT_STORAGE_AZURE_BLOB_CONTAINER" + self._session = None + self._engine = None + + def get_env_rdbms(self) -> str: + if not self.ENV_RDBMS_AVAILABLE: + raise EnvironmentError(f"Environment variable '{self.ENV_RDBMS}' not set") + return os.environ[self.ENV_RDBMS] + + @property + def ENV_RDBMS_AVAILABLE(self) -> bool: + return self.ENV_RDBMS in os.environ + + @property + def URI_RDBMS(self) -> str: + return self.get_env_rdbms() + + @property + def IS_SQLITE(self) -> bool: + return self.URI_RDBMS.startswith("sqlite") + + @property + def IS_POSTGRES(self) -> bool: + return self.URI_RDBMS.startswith("postgres") + + @property + def HAS_AZURE_BLOB_STORAGE(self) -> bool: + return self.ENV_BLOB in os.environ + + @property + def BLOB_CONTAINER(self) -> str: + return os.getenv(self.ENV_BLOB_CONTAINER, "ert") + + @property + def engine(self) -> Engine: + if self._engine is None: + if self.IS_SQLITE: + self._engine = create_engine( + self.URI_RDBMS, connect_args={"check_same_thread": False} + ) + else: + self._engine = create_engine( + self.URI_RDBMS, pool_size=50, max_overflow=100 + ) + return self._engine + + @property + def Session(self) -> sessionmaker: + if self._session is None: + self._session = sessionmaker( + autocommit=False, autoflush=False, bind=self.engine + ) + return self._session + + +database_config = DatabaseConfig() Base = declarative_base() async def get_db(*, _: None = Depends(security)) -> Any: - db = Session() + db = database_config.Session() # Make PostgreSQL return float8 columns with highest precision. If we don't # do this, we may lose up to 3 of the least significant digits. - if IS_POSTGRES: + if database_config.IS_POSTGRES: db.execute("SET extra_float_digits=3") try: yield db @@ -52,13 +88,13 @@ async def get_db(*, _: None = Depends(security)) -> Any: raise -if HAS_AZURE_BLOB_STORAGE: +if database_config.HAS_AZURE_BLOB_STORAGE: import asyncio from azure.core.exceptions import ResourceNotFoundError from azure.storage.blob.aio import ContainerClient azure_blob_container = ContainerClient.from_connection_string( - os.environ[ENV_BLOB], BLOB_CONTAINER + os.environ[database_config.ENV_BLOB], database_config.BLOB_CONTAINER ) async def create_container_if_not_exist() -> None: diff --git a/src/ert/dark_storage/database_schema/record.py b/src/ert/dark_storage/database_schema/record.py index 385d822d44b..dcaa5200e7a 100644 --- a/src/ert/dark_storage/database_schema/record.py +++ b/src/ert/dark_storage/database_schema/record.py @@ -4,7 +4,6 @@ import sqlalchemy as sa from sqlalchemy.sql import func from sqlalchemy.orm import relationship -from sqlalchemy.ext.hybrid import hybrid_property from ert_storage.ext.sqlalchemy_arrays import FloatArray from ert_storage.ext.uuid import UUID diff --git a/src/ert/dark_storage/endpoints/_records_blob.py b/src/ert/dark_storage/endpoints/_records_blob.py index 3d4aec90182..0e8bddbf6c6 100644 --- a/src/ert/dark_storage/endpoints/_records_blob.py +++ b/src/ert/dark_storage/endpoints/_records_blob.py @@ -1,21 +1,17 @@ -import io from typing import Optional, List, Type, AsyncGenerator from uuid import uuid4, UUID -import numpy as np -import pandas as pd from fastapi import ( Request, UploadFile, Depends, ) -from fastapi.logger import logger from fastapi.responses import Response, StreamingResponse from ert_storage import database_schema as ds -from ert_storage.database import Session, get_db, HAS_AZURE_BLOB_STORAGE +from ert_storage.database import Session, get_db, database_config -if HAS_AZURE_BLOB_STORAGE: +if database_config.HAS_AZURE_BLOB_STORAGE: from ert_storage.database import azure_blob_container @@ -162,7 +158,7 @@ def get_blob_handler( realization_index: Optional[int] = None, ) -> BlobHandler: blob_handler: Type[BlobHandler] - if HAS_AZURE_BLOB_STORAGE: + if database_config.HAS_AZURE_BLOB_STORAGE: blob_handler = AzureBlobHandler else: blob_handler = BlobHandler @@ -173,7 +169,7 @@ def get_blob_handler( def get_blob_handler_from_record(db: Session, record: ds.Record) -> BlobHandler: blob_handler: Type[BlobHandler] - if HAS_AZURE_BLOB_STORAGE: + if database_config.HAS_AZURE_BLOB_STORAGE: blob_handler = AzureBlobHandler else: blob_handler = BlobHandler diff --git a/src/ert/dark_storage/endpoints/responses.py b/src/ert/dark_storage/endpoints/responses.py index 8374d4d116f..abfdcace3d2 100644 --- a/src/ert/dark_storage/endpoints/responses.py +++ b/src/ert/dark_storage/endpoints/responses.py @@ -1,12 +1,11 @@ -from uuid import uuid4, UUID +from uuid import UUID import pandas as pd from fastapi import ( APIRouter, Depends, ) from fastapi.responses import Response -from pandas.core.frame import DataFrame -from ert_storage.database import Session, get_db, HAS_AZURE_BLOB_STORAGE +from ert_storage.database import Session, get_db from ert_storage import database_schema as ds router = APIRouter(tags=["response"]) diff --git a/src/ert/dark_storage/ext/sqlalchemy_arrays.py b/src/ert/dark_storage/ext/sqlalchemy_arrays.py index edd3b3f4f43..be61a93a7dd 100644 --- a/src/ert/dark_storage/ext/sqlalchemy_arrays.py +++ b/src/ert/dark_storage/ext/sqlalchemy_arrays.py @@ -10,7 +10,7 @@ from typing import Optional, Type, Union import sqlalchemy as sa -from ert_storage.database import IS_POSTGRES +from ert_storage.database import database_config import graphene from graphene_sqlalchemy.converter import convert_sqlalchemy_type @@ -24,7 +24,14 @@ FloatArray: SQLAlchemyColumn StringArray: SQLAlchemyColumn -if IS_POSTGRES: +# IS_POSTGRES will fail with exception if ENV_DBMS is not available. +# This module is loaded during compability-check from ERT ( dark-storage api / +# ert-storage api ) and we need it to complete the loading without exception. It +# make no sense setting an environment variable in ERT when testing to specify which +# database is to be used when dark-storage is backed by no database at all. To +# achieve this we are checking if ENV_DBMS is available. The API itself should not +# change depending on which database is implemented so this should not be a problem. +if database_config.ENV_RDBMS_AVAILABLE and database_config.IS_POSTGRES: FloatArray = sa.ARRAY(sa.FLOAT) StringArray = sa.ARRAY(sa.String) else: diff --git a/src/ert/dark_storage/testing/testclient.py b/src/ert/dark_storage/testing/testclient.py index a8efcd9e58e..136caca3e65 100644 --- a/src/ert/dark_storage/testing/testclient.py +++ b/src/ert/dark_storage/testing/testclient.py @@ -5,7 +5,6 @@ AsyncGenerator, Generator, Mapping, - MutableMapping, Optional, TYPE_CHECKING, Tuple, @@ -20,11 +19,6 @@ TestClient as StarletteTestClient, ASGI2App, ASGI3App, - Cookies, - Params, - DataType, - TimeOut, - FileType, ) from sqlalchemy.orm import sessionmaker from graphene import Schema as GrapheneSchema @@ -197,7 +191,7 @@ def _override_get_db(session: sessionmaker) -> None: from ert_storage.app import app from ert_storage.database import ( get_db, - IS_POSTGRES, + database_config, ) async def override_get_db( @@ -207,7 +201,7 @@ async def override_get_db( # Make PostgreSQL return float8 columns with highest precision. If we don't # do this, we may lose up to 3 of the least significant digits. - if IS_POSTGRES: + if database_config.IS_POSTGRES: db.execute("SET extra_float_digits=3") try: yield db @@ -222,23 +216,19 @@ async def override_get_db( def _begin_transaction() -> _TransactionInfo: - from ert_storage.database import ( - engine, - IS_SQLITE, - HAS_AZURE_BLOB_STORAGE, - ) + from ert_storage.database import database_config from ert_storage.database_schema import Base - if IS_SQLITE: - Base.metadata.create_all(bind=engine) - if HAS_AZURE_BLOB_STORAGE: + if database_config.IS_SQLITE: + Base.metadata.create_all(bind=database_config.engine) + if database_config.HAS_AZURE_BLOB_STORAGE: import asyncio from ert_storage.database import create_container_if_not_exist loop = asyncio.get_event_loop() loop.run_until_complete(create_container_if_not_exist()) - connection = engine.connect() + connection = database_config.engine.connect() transaction = connection.begin() session = sessionmaker(autocommit=False, autoflush=False, bind=connection) From 5c95295bcd39d38b35f912c0bd372515c528489d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=98ystein=20Olai=20Heggen?= Date: Wed, 22 Sep 2021 10:46:22 +0200 Subject: [PATCH 17/43] Revert "lazy load database configuration" This reverts commit aaccf35bb76b14b40fe520df8cc2b27d7f21ff1b. --- src/ert/dark_storage/__main__.py | 4 +- src/ert/dark_storage/_alembic/alembic/env.py | 6 +- src/ert/dark_storage/app.py | 9 +- src/ert/dark_storage/database.py | 96 ++++++------------- .../dark_storage/database_schema/record.py | 1 + .../dark_storage/endpoints/_records_blob.py | 12 ++- src/ert/dark_storage/endpoints/responses.py | 5 +- src/ert/dark_storage/ext/sqlalchemy_arrays.py | 11 +-- src/ert/dark_storage/testing/testclient.py | 24 +++-- 9 files changed, 71 insertions(+), 97 deletions(-) diff --git a/src/ert/dark_storage/__main__.py b/src/ert/dark_storage/__main__.py index 2a9efac7753..2c1327d8795 100644 --- a/src/ert/dark_storage/__main__.py +++ b/src/ert/dark_storage/__main__.py @@ -44,7 +44,7 @@ def run_alembic(args: List[str]) -> None: """ Forward arguments to alembic """ - from alembic.config import main as alembic_main + from alembic.config import main dbkey = "ERT_STORAGE_DATABASE_URL" dburl = os.getenv(dbkey) @@ -67,7 +67,7 @@ def run_alembic(args: List[str]) -> None: ] try: - alembic_main(argv=argv, prog="ert-storage alembic") + main(argv=argv, prog="ert-storage alembic") except FileNotFoundError as exc: if os.path.basename(exc.filename) == "script.py.mako": sys.exit( diff --git a/src/ert/dark_storage/_alembic/alembic/env.py b/src/ert/dark_storage/_alembic/alembic/env.py index d6b1e95ce35..ffb52010677 100644 --- a/src/ert/dark_storage/_alembic/alembic/env.py +++ b/src/ert/dark_storage/_alembic/alembic/env.py @@ -6,7 +6,7 @@ from alembic import context -from ert_storage.database import database_config +from ert_storage.database import ENV_RDBMS from ert_storage.database_schema import Base # this is the Alembic Config object, which provides @@ -43,7 +43,7 @@ def run_migrations_offline(): """ # Alembic uses sprintf somewhere. Escape the '%'s so that alembic doesn't # think they're part of the format string. - url = os.environ[database_config.ENV_RDBMS].replace("%", "%%") + url = os.environ[ENV_RDBMS].replace("%", "%%") context.configure( url=url, target_metadata=target_metadata, @@ -64,7 +64,7 @@ def run_migrations_online(): """ # Alembic uses sprintf somewhere. Escape the '%'s so that alembic doesn't # think they're part of the format string. - url = os.environ[database_config.ENV_RDBMS].replace("%", "%%") + url = os.environ[ENV_RDBMS].replace("%", "%%") config.set_section_option(config.config_ini_section, "sqlalchemy.url", str(url)) connectable = engine_from_config( diff --git a/src/ert/dark_storage/app.py b/src/ert/dark_storage/app.py index e9a65a59119..b646c26e197 100644 --- a/src/ert/dark_storage/app.py +++ b/src/ert/dark_storage/app.py @@ -3,6 +3,7 @@ from typing import Any from fastapi import FastAPI, Request, status from fastapi.responses import Response, RedirectResponse +from starlette.graphql import GraphQLApp from ert_storage.endpoints import router as endpoints_router from ert_storage.graphql import router as graphql_router @@ -50,13 +51,13 @@ def render(self, content: Any) -> bytes: @app.on_event("startup") async def initialize_database() -> None: - from ert_storage.database import database_config + from ert_storage.database import engine, IS_SQLITE, HAS_AZURE_BLOB_STORAGE from ert_storage.database_schema import Base - if database_config.IS_SQLITE: + if IS_SQLITE: # Our SQLite backend doesn't support migrations, so create the database on the fly. - Base.metadata.create_all(bind=database_config.engine) - if database_config.HAS_AZURE_BLOB_STORAGE: + Base.metadata.create_all(bind=engine) + if HAS_AZURE_BLOB_STORAGE: from ert_storage.database import create_container_if_not_exist await create_container_if_not_exist() diff --git a/src/ert/dark_storage/database.py b/src/ert/dark_storage/database.py index 785e24a1554..ec2d412ff1f 100644 --- a/src/ert/dark_storage/database.py +++ b/src/ert/dark_storage/database.py @@ -2,81 +2,45 @@ from typing import Any from fastapi import Depends from sqlalchemy import create_engine -from sqlalchemy.engine import Engine from sqlalchemy.ext.declarative import declarative_base -from sqlalchemy.orm import sessionmaker, Session +from sqlalchemy.orm import sessionmaker from ert_storage.security import security -class DatabaseConfig: - def __init__(self) -> None: - self.ENV_RDBMS = "ERT_STORAGE_DATABASE_URL" - self.ENV_BLOB = "ERT_STORAGE_AZURE_CONNECTION_STRING" - self.ENV_BLOB_CONTAINER = "ERT_STORAGE_AZURE_BLOB_CONTAINER" - self._session = None - self._engine = None - - def get_env_rdbms(self) -> str: - if not self.ENV_RDBMS_AVAILABLE: - raise EnvironmentError(f"Environment variable '{self.ENV_RDBMS}' not set") - return os.environ[self.ENV_RDBMS] - - @property - def ENV_RDBMS_AVAILABLE(self) -> bool: - return self.ENV_RDBMS in os.environ - - @property - def URI_RDBMS(self) -> str: - return self.get_env_rdbms() - - @property - def IS_SQLITE(self) -> bool: - return self.URI_RDBMS.startswith("sqlite") - - @property - def IS_POSTGRES(self) -> bool: - return self.URI_RDBMS.startswith("postgres") - - @property - def HAS_AZURE_BLOB_STORAGE(self) -> bool: - return self.ENV_BLOB in os.environ - - @property - def BLOB_CONTAINER(self) -> str: - return os.getenv(self.ENV_BLOB_CONTAINER, "ert") - - @property - def engine(self) -> Engine: - if self._engine is None: - if self.IS_SQLITE: - self._engine = create_engine( - self.URI_RDBMS, connect_args={"check_same_thread": False} - ) - else: - self._engine = create_engine( - self.URI_RDBMS, pool_size=50, max_overflow=100 - ) - return self._engine - - @property - def Session(self) -> sessionmaker: - if self._session is None: - self._session = sessionmaker( - autocommit=False, autoflush=False, bind=self.engine - ) - return self._session - - -database_config = DatabaseConfig() +ENV_RDBMS = "ERT_STORAGE_DATABASE_URL" +ENV_BLOB = "ERT_STORAGE_AZURE_CONNECTION_STRING" +ENV_BLOB_CONTAINER = "ERT_STORAGE_AZURE_BLOB_CONTAINER" + + +def get_env_rdbms() -> str: + if ENV_RDBMS not in os.environ: + raise EnvironmentError(f"Environment variable '{ENV_RDBMS}' not set") + return os.environ[ENV_RDBMS] + + +URI_RDBMS = get_env_rdbms() +IS_SQLITE = URI_RDBMS.startswith("sqlite") +IS_POSTGRES = URI_RDBMS.startswith("postgres") +HAS_AZURE_BLOB_STORAGE = ENV_BLOB in os.environ +BLOB_CONTAINER = os.getenv(ENV_BLOB_CONTAINER, "ert") + + +if IS_SQLITE: + engine = create_engine(URI_RDBMS, connect_args={"check_same_thread": False}) +else: + engine = create_engine(URI_RDBMS, pool_size=50, max_overflow=100) +Session = sessionmaker(autocommit=False, autoflush=False, bind=engine) + + Base = declarative_base() async def get_db(*, _: None = Depends(security)) -> Any: - db = database_config.Session() + db = Session() # Make PostgreSQL return float8 columns with highest precision. If we don't # do this, we may lose up to 3 of the least significant digits. - if database_config.IS_POSTGRES: + if IS_POSTGRES: db.execute("SET extra_float_digits=3") try: yield db @@ -88,13 +52,13 @@ async def get_db(*, _: None = Depends(security)) -> Any: raise -if database_config.HAS_AZURE_BLOB_STORAGE: +if HAS_AZURE_BLOB_STORAGE: import asyncio from azure.core.exceptions import ResourceNotFoundError from azure.storage.blob.aio import ContainerClient azure_blob_container = ContainerClient.from_connection_string( - os.environ[database_config.ENV_BLOB], database_config.BLOB_CONTAINER + os.environ[ENV_BLOB], BLOB_CONTAINER ) async def create_container_if_not_exist() -> None: diff --git a/src/ert/dark_storage/database_schema/record.py b/src/ert/dark_storage/database_schema/record.py index dcaa5200e7a..385d822d44b 100644 --- a/src/ert/dark_storage/database_schema/record.py +++ b/src/ert/dark_storage/database_schema/record.py @@ -4,6 +4,7 @@ import sqlalchemy as sa from sqlalchemy.sql import func from sqlalchemy.orm import relationship +from sqlalchemy.ext.hybrid import hybrid_property from ert_storage.ext.sqlalchemy_arrays import FloatArray from ert_storage.ext.uuid import UUID diff --git a/src/ert/dark_storage/endpoints/_records_blob.py b/src/ert/dark_storage/endpoints/_records_blob.py index 0e8bddbf6c6..3d4aec90182 100644 --- a/src/ert/dark_storage/endpoints/_records_blob.py +++ b/src/ert/dark_storage/endpoints/_records_blob.py @@ -1,17 +1,21 @@ +import io from typing import Optional, List, Type, AsyncGenerator from uuid import uuid4, UUID +import numpy as np +import pandas as pd from fastapi import ( Request, UploadFile, Depends, ) +from fastapi.logger import logger from fastapi.responses import Response, StreamingResponse from ert_storage import database_schema as ds -from ert_storage.database import Session, get_db, database_config +from ert_storage.database import Session, get_db, HAS_AZURE_BLOB_STORAGE -if database_config.HAS_AZURE_BLOB_STORAGE: +if HAS_AZURE_BLOB_STORAGE: from ert_storage.database import azure_blob_container @@ -158,7 +162,7 @@ def get_blob_handler( realization_index: Optional[int] = None, ) -> BlobHandler: blob_handler: Type[BlobHandler] - if database_config.HAS_AZURE_BLOB_STORAGE: + if HAS_AZURE_BLOB_STORAGE: blob_handler = AzureBlobHandler else: blob_handler = BlobHandler @@ -169,7 +173,7 @@ def get_blob_handler( def get_blob_handler_from_record(db: Session, record: ds.Record) -> BlobHandler: blob_handler: Type[BlobHandler] - if database_config.HAS_AZURE_BLOB_STORAGE: + if HAS_AZURE_BLOB_STORAGE: blob_handler = AzureBlobHandler else: blob_handler = BlobHandler diff --git a/src/ert/dark_storage/endpoints/responses.py b/src/ert/dark_storage/endpoints/responses.py index abfdcace3d2..8374d4d116f 100644 --- a/src/ert/dark_storage/endpoints/responses.py +++ b/src/ert/dark_storage/endpoints/responses.py @@ -1,11 +1,12 @@ -from uuid import UUID +from uuid import uuid4, UUID import pandas as pd from fastapi import ( APIRouter, Depends, ) from fastapi.responses import Response -from ert_storage.database import Session, get_db +from pandas.core.frame import DataFrame +from ert_storage.database import Session, get_db, HAS_AZURE_BLOB_STORAGE from ert_storage import database_schema as ds router = APIRouter(tags=["response"]) diff --git a/src/ert/dark_storage/ext/sqlalchemy_arrays.py b/src/ert/dark_storage/ext/sqlalchemy_arrays.py index be61a93a7dd..edd3b3f4f43 100644 --- a/src/ert/dark_storage/ext/sqlalchemy_arrays.py +++ b/src/ert/dark_storage/ext/sqlalchemy_arrays.py @@ -10,7 +10,7 @@ from typing import Optional, Type, Union import sqlalchemy as sa -from ert_storage.database import database_config +from ert_storage.database import IS_POSTGRES import graphene from graphene_sqlalchemy.converter import convert_sqlalchemy_type @@ -24,14 +24,7 @@ FloatArray: SQLAlchemyColumn StringArray: SQLAlchemyColumn -# IS_POSTGRES will fail with exception if ENV_DBMS is not available. -# This module is loaded during compability-check from ERT ( dark-storage api / -# ert-storage api ) and we need it to complete the loading without exception. It -# make no sense setting an environment variable in ERT when testing to specify which -# database is to be used when dark-storage is backed by no database at all. To -# achieve this we are checking if ENV_DBMS is available. The API itself should not -# change depending on which database is implemented so this should not be a problem. -if database_config.ENV_RDBMS_AVAILABLE and database_config.IS_POSTGRES: +if IS_POSTGRES: FloatArray = sa.ARRAY(sa.FLOAT) StringArray = sa.ARRAY(sa.String) else: diff --git a/src/ert/dark_storage/testing/testclient.py b/src/ert/dark_storage/testing/testclient.py index 136caca3e65..a8efcd9e58e 100644 --- a/src/ert/dark_storage/testing/testclient.py +++ b/src/ert/dark_storage/testing/testclient.py @@ -5,6 +5,7 @@ AsyncGenerator, Generator, Mapping, + MutableMapping, Optional, TYPE_CHECKING, Tuple, @@ -19,6 +20,11 @@ TestClient as StarletteTestClient, ASGI2App, ASGI3App, + Cookies, + Params, + DataType, + TimeOut, + FileType, ) from sqlalchemy.orm import sessionmaker from graphene import Schema as GrapheneSchema @@ -191,7 +197,7 @@ def _override_get_db(session: sessionmaker) -> None: from ert_storage.app import app from ert_storage.database import ( get_db, - database_config, + IS_POSTGRES, ) async def override_get_db( @@ -201,7 +207,7 @@ async def override_get_db( # Make PostgreSQL return float8 columns with highest precision. If we don't # do this, we may lose up to 3 of the least significant digits. - if database_config.IS_POSTGRES: + if IS_POSTGRES: db.execute("SET extra_float_digits=3") try: yield db @@ -216,19 +222,23 @@ async def override_get_db( def _begin_transaction() -> _TransactionInfo: - from ert_storage.database import database_config + from ert_storage.database import ( + engine, + IS_SQLITE, + HAS_AZURE_BLOB_STORAGE, + ) from ert_storage.database_schema import Base - if database_config.IS_SQLITE: - Base.metadata.create_all(bind=database_config.engine) - if database_config.HAS_AZURE_BLOB_STORAGE: + if IS_SQLITE: + Base.metadata.create_all(bind=engine) + if HAS_AZURE_BLOB_STORAGE: import asyncio from ert_storage.database import create_container_if_not_exist loop = asyncio.get_event_loop() loop.run_until_complete(create_container_if_not_exist()) - connection = database_config.engine.connect() + connection = engine.connect() transaction = connection.begin() session = sessionmaker(autocommit=False, autoflush=False, bind=connection) From 375d8a0909c60d43d140c4234d28ace27311dd15 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=98ystein=20Olai=20Heggen?= Date: Wed, 22 Sep 2021 12:23:36 +0200 Subject: [PATCH 18/43] Distinguish between alembic main and module main --- src/ert/dark_storage/__main__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ert/dark_storage/__main__.py b/src/ert/dark_storage/__main__.py index 2c1327d8795..2a9efac7753 100644 --- a/src/ert/dark_storage/__main__.py +++ b/src/ert/dark_storage/__main__.py @@ -44,7 +44,7 @@ def run_alembic(args: List[str]) -> None: """ Forward arguments to alembic """ - from alembic.config import main + from alembic.config import main as alembic_main dbkey = "ERT_STORAGE_DATABASE_URL" dburl = os.getenv(dbkey) @@ -67,7 +67,7 @@ def run_alembic(args: List[str]) -> None: ] try: - main(argv=argv, prog="ert-storage alembic") + alembic_main(argv=argv, prog="ert-storage alembic") except FileNotFoundError as exc: if os.path.basename(exc.filename) == "script.py.mako": sys.exit( From edb2b7e37163fc31bff4f7e056b22beb996f092d Mon Sep 17 00:00:00 2001 From: DanSava Date: Tue, 12 Oct 2021 17:31:35 +0300 Subject: [PATCH 19/43] Allow custom active realization index when adding records --- ...826_add_active_realizations_to_ensemble.py | 31 +++++++++++++++++++ .../dark_storage/database_schema/ensemble.py | 3 +- src/ert/dark_storage/endpoints/ensembles.py | 18 +++++++++++ src/ert/dark_storage/endpoints/records.py | 23 +++++++------- src/ert/dark_storage/ext/sqlalchemy_arrays.py | 13 ++++++-- src/ert/dark_storage/graphql/ensembles.py | 6 ++++ src/ert/dark_storage/json_schema/ensemble.py | 1 + 7 files changed, 80 insertions(+), 15 deletions(-) create mode 100644 src/ert/dark_storage/_alembic/alembic/versions/abccdeea2826_add_active_realizations_to_ensemble.py diff --git a/src/ert/dark_storage/_alembic/alembic/versions/abccdeea2826_add_active_realizations_to_ensemble.py b/src/ert/dark_storage/_alembic/alembic/versions/abccdeea2826_add_active_realizations_to_ensemble.py new file mode 100644 index 00000000000..133ca083111 --- /dev/null +++ b/src/ert/dark_storage/_alembic/alembic/versions/abccdeea2826_add_active_realizations_to_ensemble.py @@ -0,0 +1,31 @@ +"""add active realizations to ensemble + +Revision ID: abccdeea2826 +Revises: 057555237953 +Create Date: 2021-10-14 11:30:46.804752 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "abccdeea2826" +down_revision = "057555237953" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "ensemble", + sa.Column("active_realizations", sa.ARRAY(sa.Integer()), nullable=False), + ) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("ensemble", "active_realizations") + # ### end Alembic commands ### diff --git a/src/ert/dark_storage/database_schema/ensemble.py b/src/ert/dark_storage/database_schema/ensemble.py index 2f6a5bf2757..8975c041ea7 100644 --- a/src/ert/dark_storage/database_schema/ensemble.py +++ b/src/ert/dark_storage/database_schema/ensemble.py @@ -5,7 +5,7 @@ from sqlalchemy.sql import func from ert_storage.database import Base from ert_storage.ext.uuid import UUID as UUID -from ert_storage.ext.sqlalchemy_arrays import StringArray +from ert_storage.ext.sqlalchemy_arrays import StringArray, IntArray from ._userdata_field import UserdataField @@ -15,6 +15,7 @@ class Ensemble(Base, UserdataField): pk = sa.Column(sa.Integer, primary_key=True) id = sa.Column(UUID, unique=True, default=uuid4, nullable=False) size = sa.Column(sa.Integer, nullable=False) + active_realizations = sa.Column(IntArray, nullable=True, default=[]) time_created = sa.Column(sa.DateTime, server_default=func.now()) time_updated = sa.Column( sa.DateTime, server_default=func.now(), onupdate=func.now() diff --git a/src/ert/dark_storage/endpoints/ensembles.py b/src/ert/dark_storage/endpoints/ensembles.py index cba3e70cd6b..d59eb483f70 100644 --- a/src/ert/dark_storage/endpoints/ensembles.py +++ b/src/ert/dark_storage/endpoints/ensembles.py @@ -5,6 +5,7 @@ from ert_storage.database import Session, get_db from ert_storage import database_schema as ds, json_schema as js from typing import Any, Mapping +from ert_storage import exceptions as exc router = APIRouter(tags=["ensemble"]) @@ -15,12 +16,29 @@ def post_ensemble( ) -> ds.Ensemble: experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() + active_reals = ( + ens_in.active_realizations + if ens_in.active_realizations + else list(range(ens_in.size)) + ) + + if ens_in.size > 0: + if max(active_reals) > ens_in.size - 1: + raise exc.ExpectationError( + f"Ensemble active realization index {max(active_reals)} out of realization range [0,{ ens_in.size - 1}]" + ) + if len(set(active_reals)) != len(active_reals): + raise exc.ExpectationError( + f"Non unique active realization index list not allowed {active_reals}" + ) + ens = ds.Ensemble( parameter_names=ens_in.parameter_names, response_names=ens_in.response_names, experiment=experiment, size=ens_in.size, userdata=ens_in.userdata, + active_realizations=active_reals, ) db.add(ens) diff --git a/src/ert/dark_storage/endpoints/records.py b/src/ert/dark_storage/endpoints/records.py index 1d95ad0f777..73fd44155b0 100644 --- a/src/ert/dark_storage/endpoints/records.py +++ b/src/ert/dark_storage/endpoints/records.py @@ -93,19 +93,18 @@ def new_record( .join(ds.RecordInfo) .filter_by(ensemble_pk=ensemble.pk, name=name) ) - if realization_index is not None: - if realization_index not in range(ensemble.size) and ensemble.size != -1: - raise exc.ExpectationError( - f"Ensemble '{name}' ('{ensemble_id}') does have a 'size' " - f"of {ensemble.size}. The posted record is targeting " - f"'realization_index' {realization_index} which is out " - f"of bounds." - ) - - q = q.filter( - (ds.Record.realization_index == None) - | (ds.Record.realization_index == realization_index) + if ( + ensemble.size != -1 + and realization_index is not None + and realization_index not in ensemble.active_realizations + ): + raise exc.ExpectationError( + f"Realization index {realization_index} outside of allowed realization indices {ensemble.active_realizations}" ) + q = q.filter( + (ds.Record.realization_index == None) + | (ds.Record.realization_index == realization_index) + ) if q.count() > 0: raise exc.ConflictError( diff --git a/src/ert/dark_storage/ext/sqlalchemy_arrays.py b/src/ert/dark_storage/ext/sqlalchemy_arrays.py index edd3b3f4f43..1d254c25f42 100644 --- a/src/ert/dark_storage/ext/sqlalchemy_arrays.py +++ b/src/ert/dark_storage/ext/sqlalchemy_arrays.py @@ -1,5 +1,5 @@ """ -This module adds thte FloatArray and StringArray column types. In Postgresql, +This module adds the FloatArray, StringArray and IntArray column types. In Postgresql, both are native `sqlalchemy.ARRAY`s, while on SQLite, they are `PickleType`s. In order to have graphene_sqlalchemy dump the arrays as arrays and not strings, @@ -17,19 +17,22 @@ from graphene_sqlalchemy.registry import Registry -__all__ = ["FloatArray", "StringArray"] +__all__ = ["FloatArray", "StringArray", "IntArray"] SQLAlchemyColumn = Union[sa.types.TypeEngine, Type[sa.types.TypeEngine]] FloatArray: SQLAlchemyColumn StringArray: SQLAlchemyColumn +IntArray: SQLAlchemyColumn if IS_POSTGRES: FloatArray = sa.ARRAY(sa.FLOAT) StringArray = sa.ARRAY(sa.String) + IntArray = sa.ARRAY(sa.Integer) else: FloatArray = type("FloatArray", (sa.PickleType,), dict(sa.PickleType.__dict__)) StringArray = type("StringArray", (sa.PickleType,), dict(sa.PickleType.__dict__)) + IntArray = type("IntArray", (sa.PickleType,), dict(sa.PickleType.__dict__)) @convert_sqlalchemy_type.register(StringArray) def convert_column_to_string_array( @@ -42,3 +45,9 @@ def convert_column_to_float_array( type: SQLAlchemyColumn, column: sa.Column, registry: Optional[Registry] = None ) -> graphene.types.structures.Structure: return graphene.List(graphene.Float) + + @convert_sqlalchemy_type.register(IntArray) + def convert_column_to_int_array( + type: SQLAlchemyColumn, column: sa.Column, registry: Optional[Registry] = None + ) -> graphene.types.structures.Structure: + return graphene.List(graphene.Int) diff --git a/src/ert/dark_storage/graphql/ensembles.py b/src/ert/dark_storage/graphql/ensembles.py index a2f912d2fcb..fe123bd72a8 100644 --- a/src/ert/dark_storage/graphql/ensembles.py +++ b/src/ert/dark_storage/graphql/ensembles.py @@ -76,6 +76,7 @@ class Meta: class Arguments: parameter_names = gr.List(gr.String) size = gr.Int() + active_realizations = gr.List(gr.Int) @staticmethod def mutate( @@ -83,6 +84,7 @@ def mutate( info: "ResolveInfo", parameter_names: List[str], size: int, + active_realizations: Optional[List[int]] = None, experiment_id: Optional[str] = None, ) -> ds.Ensemble: db = get_session(info.context) @@ -94,11 +96,15 @@ def mutate( else: raise ValueError("ID is required") + if active_realizations is None: + active_realizations = list(range(size)) + ensemble = ds.Ensemble( parameter_names=parameter_names, response_names=[], experiment=experiment, size=size, + active_realizations=active_realizations, ) db.add(ensemble) diff --git a/src/ert/dark_storage/json_schema/ensemble.py b/src/ert/dark_storage/json_schema/ensemble.py index b12681c3dc5..6d567766047 100644 --- a/src/ert/dark_storage/json_schema/ensemble.py +++ b/src/ert/dark_storage/json_schema/ensemble.py @@ -7,6 +7,7 @@ class _Ensemble(BaseModel): size: int parameter_names: List[str] response_names: List[str] + active_realizations: List[int] = [] class EnsembleIn(_Ensemble): From 94f95abaf7564415870250b299310d6c5b0ededc Mon Sep 17 00:00:00 2001 From: Terry Hannant Date: Wed, 10 Nov 2021 15:56:27 +0100 Subject: [PATCH 20/43] Initial implementation of client Session --- src/ert/dark_storage/client/__init__.py | 3 + src/ert/dark_storage/client/session.py | 91 +++++++++++++++++++++++++ 2 files changed, 94 insertions(+) create mode 100644 src/ert/dark_storage/client/__init__.py create mode 100644 src/ert/dark_storage/client/session.py diff --git a/src/ert/dark_storage/client/__init__.py b/src/ert/dark_storage/client/__init__.py new file mode 100644 index 00000000000..414557f5a1d --- /dev/null +++ b/src/ert/dark_storage/client/__init__.py @@ -0,0 +1,3 @@ +from .session import Session + +__all__ = ["Session"] diff --git a/src/ert/dark_storage/client/session.py b/src/ert/dark_storage/client/session.py new file mode 100644 index 00000000000..af3fea4a5ef --- /dev/null +++ b/src/ert/dark_storage/client/session.py @@ -0,0 +1,91 @@ +import os +import requests +import json +from pathlib import Path +from urllib.parse import urljoin +from typing import Any, Dict, Type +from types import TracebackType + + +class Session(requests.Session): + """Wrapper class for requests.Session to configure base url and auth token. + + The base url and auth token are read from either: + The file storage_server.json, in the current working directory or + The environment variable ERT_STORAGE_CONNECTION_STRING + + In both cases the configuration is represeted by JSON: + + { + "urls": [list of base urls] + "authtoken": auth_token + } + + The list of base urls are tested via the /healthcheck api, + with the first url returning a 200 status code is used. + + If both file and environment variable exist the environment variable is used. + """ + + def __init__(self) -> None: + super().__init__() + self._base_url: str = "" + self._headers: Dict[str, str] = {} + self._connection_info: Dict[str, Any] = {} + + connection_string = None + + connection_config = Path.cwd() / "storage_server.json" + if connection_config.exists(): + with open(connection_config) as f: + connection_string = f.read() + + # Allow env var to overide config file + if "ERT_STORAGE_CONNECTION_STRING" in os.environ: + connection_string = os.environ["ERT_STORAGE_CONNECTION_STRING"] + + if connection_string is None: + raise RuntimeError("No Storage Connection configuration found") + + try: + self._connection_info = json.loads(connection_string) + except json.JSONDecodeError: + raise RuntimeError("Invalid Storage Connection configuration") + + if {"urls", "authtoken"} <= self._connection_info.keys(): + self._base_url = self._resolve_url() + self._headers = {"Token": self._connection_info["authtoken"]} + else: + raise RuntimeError("Invalid Storage Connection configuration") + + def __enter__(self) -> "Session": + return self + + def __exit__( # type: ignore[override] + self, + exc_value: BaseException, + exc_type: Type[BaseException], + traceback: TracebackType, + ) -> bool: + pass + + def _resolve_url(self) -> str: + """Resolve which of the candidate base urls to use.""" + for url in self._connection_info["urls"]: + try: + # Original code has auth token passed but is it actually used? + resp = requests.get(f"{url}/healthcheck") + if resp.status_code == 200: + return url + except requests.ConnectionError: + pass + # Needs better exception message + raise RuntimeError("None of the Storage URLs provided worked") + + def request( # type: ignore[override] + self, method: str, url: str, *args: Any, **kwargs: Any + ) -> requests.Response: + """Perform HTTP request with preconfigured base url and auth token.""" + kwargs.setdefault("headers", {}) + kwargs["headers"].update(self._headers) + return super().request(method, urljoin(self._base_url, url), *args, **kwargs) From 9e869924407a15d209f7bcddf35ea5a8c43bb676 Mon Sep 17 00:00:00 2001 From: Terry Hannant Date: Fri, 26 Nov 2021 08:52:30 +0100 Subject: [PATCH 21/43] Add temp debuging for tests --- src/ert/dark_storage/client/session.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/ert/dark_storage/client/session.py b/src/ert/dark_storage/client/session.py index af3fea4a5ef..31c180ee032 100644 --- a/src/ert/dark_storage/client/session.py +++ b/src/ert/dark_storage/client/session.py @@ -73,9 +73,13 @@ def _resolve_url(self) -> str: """Resolve which of the candidate base urls to use.""" for url in self._connection_info["urls"]: try: + print(f"Testing {url}") # Original code has auth token passed but is it actually used? resp = requests.get(f"{url}/healthcheck") + print(f"Response code {resp.status}") if resp.status_code == 200: + print(f"200 status code for {url}") + print(f"Response {resp.text}") return url except requests.ConnectionError: pass From 0bae0736d01b0f395a3a61b813b07defc2cb3e6b Mon Sep 17 00:00:00 2001 From: Terry Hannant Date: Fri, 26 Nov 2021 08:58:09 +0100 Subject: [PATCH 22/43] Fix temp debuging for tests --- src/ert/dark_storage/client/session.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ert/dark_storage/client/session.py b/src/ert/dark_storage/client/session.py index 31c180ee032..ebb640177b8 100644 --- a/src/ert/dark_storage/client/session.py +++ b/src/ert/dark_storage/client/session.py @@ -76,7 +76,7 @@ def _resolve_url(self) -> str: print(f"Testing {url}") # Original code has auth token passed but is it actually used? resp = requests.get(f"{url}/healthcheck") - print(f"Response code {resp.status}") + print(f"Response code {resp.status_code}") if resp.status_code == 200: print(f"200 status code for {url}") print(f"Response {resp.text}") From 21c64d23b753d0ece22a0fc85ba902edd3a15a20 Mon Sep 17 00:00:00 2001 From: Terry Hannant Date: Fri, 26 Nov 2021 09:53:13 +0100 Subject: [PATCH 23/43] Add additional test debugging --- src/ert/dark_storage/client/session.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/ert/dark_storage/client/session.py b/src/ert/dark_storage/client/session.py index ebb640177b8..196e7d863ce 100644 --- a/src/ert/dark_storage/client/session.py +++ b/src/ert/dark_storage/client/session.py @@ -29,6 +29,7 @@ class Session(requests.Session): def __init__(self) -> None: super().__init__() + print("Session init") self._base_url: str = "" self._headers: Dict[str, str] = {} self._connection_info: Dict[str, Any] = {} @@ -47,11 +48,12 @@ def __init__(self) -> None: if connection_string is None: raise RuntimeError("No Storage Connection configuration found") + print(f"Session using {connection_string}") try: self._connection_info = json.loads(connection_string) except json.JSONDecodeError: raise RuntimeError("Invalid Storage Connection configuration") - + if {"urls", "authtoken"} <= self._connection_info.keys(): self._base_url = self._resolve_url() self._headers = {"Token": self._connection_info["authtoken"]} @@ -59,6 +61,7 @@ def __init__(self) -> None: raise RuntimeError("Invalid Storage Connection configuration") def __enter__(self) -> "Session": + print("Enter Session") return self def __exit__( # type: ignore[override] From 2995218a1d30ab363a89782130fca46cc5c6b4a8 Mon Sep 17 00:00:00 2001 From: Terry Hannant Date: Fri, 26 Nov 2021 10:01:58 +0100 Subject: [PATCH 24/43] Fix formatting --- src/ert/dark_storage/client/session.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ert/dark_storage/client/session.py b/src/ert/dark_storage/client/session.py index 196e7d863ce..15ebfd737cf 100644 --- a/src/ert/dark_storage/client/session.py +++ b/src/ert/dark_storage/client/session.py @@ -53,7 +53,7 @@ def __init__(self) -> None: self._connection_info = json.loads(connection_string) except json.JSONDecodeError: raise RuntimeError("Invalid Storage Connection configuration") - + if {"urls", "authtoken"} <= self._connection_info.keys(): self._base_url = self._resolve_url() self._headers = {"Token": self._connection_info["authtoken"]} From f1542545ff7a5b0a63ea21e3a7f3f8db2c350842 Mon Sep 17 00:00:00 2001 From: Terry Hannant Date: Fri, 26 Nov 2021 12:15:17 +0100 Subject: [PATCH 25/43] Revert "Fix formatting" This reverts commit 2995218a1d30ab363a89782130fca46cc5c6b4a8. --- src/ert/dark_storage/client/session.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ert/dark_storage/client/session.py b/src/ert/dark_storage/client/session.py index 15ebfd737cf..196e7d863ce 100644 --- a/src/ert/dark_storage/client/session.py +++ b/src/ert/dark_storage/client/session.py @@ -53,7 +53,7 @@ def __init__(self) -> None: self._connection_info = json.loads(connection_string) except json.JSONDecodeError: raise RuntimeError("Invalid Storage Connection configuration") - + if {"urls", "authtoken"} <= self._connection_info.keys(): self._base_url = self._resolve_url() self._headers = {"Token": self._connection_info["authtoken"]} From 4ec9b34e78a73aa75b8e229ae31e4c7cbc8e32f6 Mon Sep 17 00:00:00 2001 From: Terry Hannant Date: Fri, 26 Nov 2021 12:15:17 +0100 Subject: [PATCH 26/43] Revert "Add additional test debugging" This reverts commit 21c64d23b753d0ece22a0fc85ba902edd3a15a20. --- src/ert/dark_storage/client/session.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/ert/dark_storage/client/session.py b/src/ert/dark_storage/client/session.py index 196e7d863ce..ebb640177b8 100644 --- a/src/ert/dark_storage/client/session.py +++ b/src/ert/dark_storage/client/session.py @@ -29,7 +29,6 @@ class Session(requests.Session): def __init__(self) -> None: super().__init__() - print("Session init") self._base_url: str = "" self._headers: Dict[str, str] = {} self._connection_info: Dict[str, Any] = {} @@ -48,12 +47,11 @@ def __init__(self) -> None: if connection_string is None: raise RuntimeError("No Storage Connection configuration found") - print(f"Session using {connection_string}") try: self._connection_info = json.loads(connection_string) except json.JSONDecodeError: raise RuntimeError("Invalid Storage Connection configuration") - + if {"urls", "authtoken"} <= self._connection_info.keys(): self._base_url = self._resolve_url() self._headers = {"Token": self._connection_info["authtoken"]} @@ -61,7 +59,6 @@ def __init__(self) -> None: raise RuntimeError("Invalid Storage Connection configuration") def __enter__(self) -> "Session": - print("Enter Session") return self def __exit__( # type: ignore[override] From d1c56a251040d31d5cf9b95c21441a49d0e4add3 Mon Sep 17 00:00:00 2001 From: Terry Hannant Date: Fri, 26 Nov 2021 12:15:27 +0100 Subject: [PATCH 27/43] Revert "Fix temp debuging for tests" This reverts commit 0bae0736d01b0f395a3a61b813b07defc2cb3e6b. --- src/ert/dark_storage/client/session.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ert/dark_storage/client/session.py b/src/ert/dark_storage/client/session.py index ebb640177b8..31c180ee032 100644 --- a/src/ert/dark_storage/client/session.py +++ b/src/ert/dark_storage/client/session.py @@ -76,7 +76,7 @@ def _resolve_url(self) -> str: print(f"Testing {url}") # Original code has auth token passed but is it actually used? resp = requests.get(f"{url}/healthcheck") - print(f"Response code {resp.status_code}") + print(f"Response code {resp.status}") if resp.status_code == 200: print(f"200 status code for {url}") print(f"Response {resp.text}") From 29c560222491f59eaf1d62ad39d144838a4cea9f Mon Sep 17 00:00:00 2001 From: Terry Hannant Date: Fri, 26 Nov 2021 12:15:27 +0100 Subject: [PATCH 28/43] Revert "Add temp debuging for tests" This reverts commit 9e869924407a15d209f7bcddf35ea5a8c43bb676. --- src/ert/dark_storage/client/session.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/ert/dark_storage/client/session.py b/src/ert/dark_storage/client/session.py index 31c180ee032..af3fea4a5ef 100644 --- a/src/ert/dark_storage/client/session.py +++ b/src/ert/dark_storage/client/session.py @@ -73,13 +73,9 @@ def _resolve_url(self) -> str: """Resolve which of the candidate base urls to use.""" for url in self._connection_info["urls"]: try: - print(f"Testing {url}") # Original code has auth token passed but is it actually used? resp = requests.get(f"{url}/healthcheck") - print(f"Response code {resp.status}") if resp.status_code == 200: - print(f"200 status code for {url}") - print(f"Response {resp.text}") return url except requests.ConnectionError: pass From 685645c1faaea1de9544449e82d1cbcb4f86177b Mon Sep 17 00:00:00 2001 From: DanSava Date: Mon, 29 Nov 2021 17:21:03 +0200 Subject: [PATCH 29/43] Support retrieving all records as ensemble wide records --- src/ert/dark_storage/endpoints/records.py | 234 ++++++++++++++++------ 1 file changed, 175 insertions(+), 59 deletions(-) diff --git a/src/ert/dark_storage/endpoints/records.py b/src/ert/dark_storage/endpoints/records.py index 73fd44155b0..f77d1186d79 100644 --- a/src/ert/dark_storage/endpoints/records.py +++ b/src/ert/dark_storage/endpoints/records.py @@ -1,3 +1,4 @@ +import json from uuid import uuid4, UUID import io import numpy as np @@ -76,6 +77,50 @@ def get_record_by_name( raise exc.NotFoundError(f"Record not found") +def get_records_by_name( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + name: str, + realization_index: Optional[int] = None, +) -> List[ds.Record]: + records = ( + db.query(ds.Record) + .filter_by(realization_index=realization_index) + .join(ds.RecordInfo) + .filter_by(name=name) + .join(ds.Ensemble) + .filter_by(id=ensemble_id) + ).all() + + if not records: + records = ( + db.query(ds.Record) + .join(ds.RecordInfo) + .filter_by( + name=name, + record_type=ds.RecordType.f64_matrix, + ) + .join(ds.Ensemble) + .filter_by(id=ensemble_id) + ).all() + + if not records: + records = ( + db.query(ds.Record) + .filter_by(realization_index=None) + .join(ds.RecordInfo) + .filter_by(name=name) + .join(ds.Ensemble) + .filter_by(id=ensemble_id) + ).all() + + if not records: + raise exc.NotFoundError(f"Record not found") + + return records + + def new_record( *, db: Session = Depends(get_db), @@ -459,15 +504,19 @@ async def get_ensemble_record( *, db: Session = Depends(get_db), bh: BlobHandler = Depends(get_blob_handler), - record: ds.Record = Depends(get_record_by_name), + records: List[ds.Record] = Depends(get_records_by_name), accept: str = Header("application/json"), realization_index: Optional[int] = None, + label: Optional[str] = None, ) -> Any: """ Get record with a given `name`. If `realization_index` is not set, look for the ensemble-wide record. If it is set, look first for one created by a forward-model for the given realization index and then the ensemble-wide record. + If label is provided it is assumed the record data is of the form {"a": 1, "b": 2} + and will return only the data for the provided label (i.e. label = "a" -> return: [[1]]) + Records support multiple data formats. In particular: - Matrix: @@ -481,10 +530,57 @@ async def get_ensemble_record( ) accept = "text/csv" - new_realization_index = ( - realization_index if record.realization_index is None else None + _type = records[0].record_info.record_type + if _type == ds.RecordType.file: + return await bh.get_content(records[0]) + + df_list = [] + for record in records: + data_df = _get_record_dataframe(record, realization_index, label) + df_list.append(data_df) + + # Combine data for each realization into one dataframe + data_frame = pd.concat(df_list, axis=0) + # Sort data by realization number + data_frame.sort_index(axis=0, inplace=True) + + return await _get_record_resonse(data_frame, accept) + + +@router.get("/ensembles/{ensemble_id}/records/{name}/labels", response_model=List[str]) +async def get_record_labels( + *, + db: Session = Depends(get_db), + ensemble_id: UUID, + name: str, +) -> List[str]: + """ + Get the list of record data labels. If the record is not a group record the list of labels will + contain only the name of the record + + Example + - Group record: + data - {"a": 4, "b":42, "c": 2} + return: ["a", "b", "c"] + - Ensemble record: + data - [4, 42, 2, 32] + return: [] + """ + + record = ( + db.query(ds.Record) + .join(ds.RecordInfo) + .filter_by(name=name) + .join(ds.Ensemble) + .filter_by(id=ensemble_id) + .first() ) - return await _get_record_data(bh, record, accept, new_realization_index) + if record is None: + raise exc.NotFoundError(f"Record not found") + + if record.f64_matrix and record.f64_matrix.labels: + return record.f64_matrix.labels[0] + return [] @router.get("/ensembles/{ensemble_id}/parameters", response_model=List[str]) @@ -532,8 +628,12 @@ async def get_record_data( accept = "text/csv" record = db.query(ds.Record).filter_by(id=record_id).one() - bh = get_blob_handler_from_record(db, record) - return await _get_record_data(bh, record, accept) + if record.record_info.record_type == ds.RecordType.file: + bh = get_blob_handler_from_record(db, record) + return await bh.get_content(record) + + dataframe = _get_record_dataframe(record, None, None) + return await _get_record_resonse(dataframe, accept) @router.get( @@ -555,70 +655,86 @@ def get_ensemble_responses( } -def _get_dataframe( - content: Any, record: ds.Record, realization_index: Optional[int] +def _get_record_dataframe( + record: ds.Record, + realization_index: Optional[int], + label: Optional[str], ) -> pd.DataFrame: - data = pd.DataFrame(content) + type_ = record.record_info.record_type + if type_ != ds.RecordType.f64_matrix: + raise exc.ExpectationError("Non matrix record not supported") + labels = record.f64_matrix.labels - if labels is not None and realization_index is None: - data.columns = labels[0] - data.index = labels[1] - elif labels is not None and realization_index is not None: - # The output is such that rows are realizations. Because - # `content` is a 1d list in this case, it treats each element as - # its own row. We transpose the data so that all of the data - # falls on the same row. - data = data.T + content_is_labeled = labels is not None + label_specified = label is not None + + if content_is_labeled and label_specified and label not in labels[0]: + raise exc.UnprocessableError(f"Record label '{label}' not found!") + + if realization_index is None or record.realization_index is not None: + matrix_content = record.f64_matrix.content + elif record.realization_index is None: + matrix_content = record.f64_matrix.content[realization_index] + if not isinstance(matrix_content[0], List): + matrix_content = [matrix_content] + + if content_is_labeled and label_specified: + lbl_idx = labels[0].index(label) + data = pd.DataFrame([[c[lbl_idx]] for c in matrix_content]) + data.columns = [label] + elif content_is_labeled: + data = pd.DataFrame(matrix_content) data.columns = labels[0] - data.index = [realization_index] + else: + data = pd.DataFrame(matrix_content) + + # Set data index for labled content + if content_is_labeled: + if record.realization_index is not None: + data.index = [record.realization_index] + elif realization_index is not None: + data.index = [realization_index] + else: + data.index = labels[1] + return data -async def _get_record_data( - bh: BlobHandler, - record: ds.Record, +async def _get_record_resonse( + dataframe: pd.DataFrame, accept: Optional[str], - realization_index: Optional[int] = None, ) -> Response: - type_ = record.record_info.record_type - if type_ == ds.RecordType.f64_matrix: - if realization_index is None: - content = record.f64_matrix.content - else: - content = record.f64_matrix.content[realization_index] + if accept == "application/x-numpy": + from numpy.lib.format import write_array - if accept == "application/x-numpy": - from numpy.lib.format import write_array + stream = io.BytesIO() + write_array(stream, np.array(dataframe.values.tolist())) - stream = io.BytesIO() - write_array(stream, np.array(content)) - - return Response( - content=stream.getvalue(), - media_type=accept, - ) - if accept == "text/csv": - data = _get_dataframe(content, record, realization_index) - - return Response( - content=data.to_csv().encode(), - media_type=accept, - ) - if accept == "application/x-parquet": - data = _get_dataframe(content, record, realization_index) - stream = io.BytesIO() - data.to_parquet(stream) - return Response( - content=stream.getvalue(), - media_type=accept, - ) + return Response( + content=stream.getvalue(), + media_type=accept, + ) + if accept == "text/csv": + return Response( + content=dataframe.to_csv().encode(), + media_type=accept, + ) + if accept == "application/x-parquet": + stream = io.BytesIO() + dataframe.to_parquet(stream) + return Response( + content=stream.getvalue(), + media_type=accept, + ) + else: + if dataframe.values.shape[0] == 1: + content = dataframe.values[0].tolist() else: - return content - if type_ == ds.RecordType.file: - return await bh.get_content(record) - raise NotImplementedError( - f"Getting record data for type {type_} and Accept header {accept} not implemented" - ) + content = dataframe.values.tolist() + return Response( + content=json.dumps(content), + media_type="application/json", + ) def _create_record( From 7b69b61590cfb9bb7429dca9f39a8021533a7dfa Mon Sep 17 00:00:00 2001 From: Zohar Malamant Date: Wed, 22 Dec 2021 10:06:50 +0100 Subject: [PATCH 30/43] Add client.Client and client.AsyncClient These replace client.Session with versions that: 1. Cache a global connection information 2. Accept a `ConnInfo` object so that we don't have to create files or set environment variables. 3. Look for `storage_server.json` in any parent directory --- src/ert/dark_storage/client/__init__.py | 6 +- src/ert/dark_storage/client/_session.py | 62 ++++++++++++++ src/ert/dark_storage/client/async_client.py | 21 +++++ src/ert/dark_storage/client/client.py | 21 +++++ src/ert/dark_storage/client/session.py | 91 --------------------- 5 files changed, 108 insertions(+), 93 deletions(-) create mode 100644 src/ert/dark_storage/client/_session.py create mode 100644 src/ert/dark_storage/client/async_client.py create mode 100644 src/ert/dark_storage/client/client.py delete mode 100644 src/ert/dark_storage/client/session.py diff --git a/src/ert/dark_storage/client/__init__.py b/src/ert/dark_storage/client/__init__.py index 414557f5a1d..839ae06609b 100644 --- a/src/ert/dark_storage/client/__init__.py +++ b/src/ert/dark_storage/client/__init__.py @@ -1,3 +1,5 @@ -from .session import Session +from ._session import ConnInfo +from .async_client import AsyncClient +from .client import Client -__all__ = ["Session"] +__all__ = ["AsyncClient", "Client", "ConnInfo"] diff --git a/src/ert/dark_storage/client/_session.py b/src/ert/dark_storage/client/_session.py new file mode 100644 index 00000000000..90cdb97c3c3 --- /dev/null +++ b/src/ert/dark_storage/client/_session.py @@ -0,0 +1,62 @@ +import os +import json +from typing import Optional +from pathlib import Path +from pydantic import BaseModel, ValidationError + + +class ConnInfo(BaseModel): + base_url: str + auth_token: Optional[str] = None + + +ENV_VAR = "ERT_STORAGE_CONNECTION_STRING" + +# Avoid searching for the connection information on every request. We assume +# that a single client process will only ever want to connect to a single ERT +# Storage server during its lifetime, so we don't provide an API for managing +# this cache. +_CACHED_CONN_INFO: Optional[ConnInfo] = None + + +def find_conn_info() -> ConnInfo: + """ + The base url and auth token are read from either: + The file `storage_server.json`, starting from the current working directory + or the environment variable `ERT_STORAGE_CONNECTION_STRING` + + In both cases the configuration is represented by JSON representation of the + `ConnInfo` pydantic model. + + In the event that nothing is found, a RuntimeError is raised. + """ + global _CACHED_CONN_INFO + if _CACHED_CONN_INFO is not None: + return _CACHED_CONN_INFO + + conn_str = os.environ.get(ENV_VAR) + + # This could be an empty string rather than None, as by the shell + # invocation: env ERT_STORAGE_CONNECTION_STRING= python + if not conn_str: + # Look for `storage_server.json` from cwd up to root. + root = Path("/") + path = Path.cwd() + while path != root: + try: + conn_str = (path / "storage_server.json").read_text() + break + except FileNotFoundError: + path = path.parent + + if not conn_str: + raise RuntimeError("No Storage connection configuration found") + + try: + conn_info = ConnInfo.parse_obj(json.loads(conn_str)) + _CACHED_CONN_INFO = conn_info + return conn_info + except json.JSONDecodeError: + raise RuntimeError("Invalid storage conneciton configuration") + except ValidationError: + raise RuntimeError("Invalid storage conneciton configuration") diff --git a/src/ert/dark_storage/client/async_client.py b/src/ert/dark_storage/client/async_client.py new file mode 100644 index 00000000000..54df45730f5 --- /dev/null +++ b/src/ert/dark_storage/client/async_client.py @@ -0,0 +1,21 @@ +import httpx +from typing import Optional + +from ._session import ConnInfo, find_conn_info + + +class AsyncClient(httpx.AsyncClient): + """ + Wrapper class for httpx.AsyncClient that provides a user-friendly way to + interact with ERT Storage's API + """ + + def __init__(self, conn_info: Optional[ConnInfo] = None) -> None: + if conn_info is None: + conn_info = find_conn_info() + + headers = {} + if conn_info.auth_token is not None: + headers = {"Token": conn_info.auth_token} + + super().__init__(base_url=conn_info.base_url, headers=headers) diff --git a/src/ert/dark_storage/client/client.py b/src/ert/dark_storage/client/client.py new file mode 100644 index 00000000000..54a2a565de5 --- /dev/null +++ b/src/ert/dark_storage/client/client.py @@ -0,0 +1,21 @@ +import httpx +from typing import Optional + +from ._session import ConnInfo, find_conn_info + + +class Client(httpx.Client): + """ + Wrapper class for httpx.Client that provides a user-friendly way to + interact with ERT Storage's API + """ + + def __init__(self, conn_info: Optional[ConnInfo] = None) -> None: + if conn_info is None: + conn_info = find_conn_info() + + headers = {} + if conn_info.auth_token is not None: + headers = {"Token": conn_info.auth_token} + + super().__init__(base_url=conn_info.base_url, headers=headers) diff --git a/src/ert/dark_storage/client/session.py b/src/ert/dark_storage/client/session.py deleted file mode 100644 index af3fea4a5ef..00000000000 --- a/src/ert/dark_storage/client/session.py +++ /dev/null @@ -1,91 +0,0 @@ -import os -import requests -import json -from pathlib import Path -from urllib.parse import urljoin -from typing import Any, Dict, Type -from types import TracebackType - - -class Session(requests.Session): - """Wrapper class for requests.Session to configure base url and auth token. - - The base url and auth token are read from either: - The file storage_server.json, in the current working directory or - The environment variable ERT_STORAGE_CONNECTION_STRING - - In both cases the configuration is represeted by JSON: - - { - "urls": [list of base urls] - "authtoken": auth_token - } - - The list of base urls are tested via the /healthcheck api, - with the first url returning a 200 status code is used. - - If both file and environment variable exist the environment variable is used. - """ - - def __init__(self) -> None: - super().__init__() - self._base_url: str = "" - self._headers: Dict[str, str] = {} - self._connection_info: Dict[str, Any] = {} - - connection_string = None - - connection_config = Path.cwd() / "storage_server.json" - if connection_config.exists(): - with open(connection_config) as f: - connection_string = f.read() - - # Allow env var to overide config file - if "ERT_STORAGE_CONNECTION_STRING" in os.environ: - connection_string = os.environ["ERT_STORAGE_CONNECTION_STRING"] - - if connection_string is None: - raise RuntimeError("No Storage Connection configuration found") - - try: - self._connection_info = json.loads(connection_string) - except json.JSONDecodeError: - raise RuntimeError("Invalid Storage Connection configuration") - - if {"urls", "authtoken"} <= self._connection_info.keys(): - self._base_url = self._resolve_url() - self._headers = {"Token": self._connection_info["authtoken"]} - else: - raise RuntimeError("Invalid Storage Connection configuration") - - def __enter__(self) -> "Session": - return self - - def __exit__( # type: ignore[override] - self, - exc_value: BaseException, - exc_type: Type[BaseException], - traceback: TracebackType, - ) -> bool: - pass - - def _resolve_url(self) -> str: - """Resolve which of the candidate base urls to use.""" - for url in self._connection_info["urls"]: - try: - # Original code has auth token passed but is it actually used? - resp = requests.get(f"{url}/healthcheck") - if resp.status_code == 200: - return url - except requests.ConnectionError: - pass - # Needs better exception message - raise RuntimeError("None of the Storage URLs provided worked") - - def request( # type: ignore[override] - self, method: str, url: str, *args: Any, **kwargs: Any - ) -> requests.Response: - """Perform HTTP request with preconfigured base url and auth token.""" - kwargs.setdefault("headers", {}) - kwargs["headers"].update(self._headers) - return super().request(method, urljoin(self._base_url, url), *args, **kwargs) From f862461d181c4c62b8e5a92eaecc17b8290fe343 Mon Sep 17 00:00:00 2001 From: DanSava Date: Fri, 4 Feb 2022 17:12:21 +0200 Subject: [PATCH 31/43] Return parameters with labels if they exists --- src/ert/dark_storage/endpoints/records.py | 27 +++++++++++++++++++---- 1 file changed, 23 insertions(+), 4 deletions(-) diff --git a/src/ert/dark_storage/endpoints/records.py b/src/ert/dark_storage/endpoints/records.py index f77d1186d79..cbeadb467b6 100644 --- a/src/ert/dark_storage/endpoints/records.py +++ b/src/ert/dark_storage/endpoints/records.py @@ -4,7 +4,7 @@ import numpy as np import pandas as pd from enum import Enum -from typing import Any, Mapping, Optional, List, AsyncGenerator +from typing import Any, Mapping, Dict, Optional, List, AsyncGenerator import sqlalchemy as sa from fastapi import ( APIRouter, @@ -583,12 +583,31 @@ async def get_record_labels( return [] -@router.get("/ensembles/{ensemble_id}/parameters", response_model=List[str]) +@router.get("/ensembles/{ensemble_id}/parameters", response_model=List[Dict[str, Any]]) async def get_ensemble_parameters( *, db: Session = Depends(get_db), ensemble_id: UUID -) -> List[str]: +) -> List[Dict[str, Any]]: ensemble = db.query(ds.Ensemble).filter_by(id=ensemble_id).one() - return ensemble.parameter_names + parameters = [] + for name in ensemble.parameter_names: + param = {"name": name, "labels": []} + record = ( + db.query(ds.Record) + .join(ds.RecordInfo) + .filter_by(name=name) + .join(ds.Ensemble) + .filter_by(id=ensemble_id) + .first() + ) + if record is None: + parameters.append(param) + continue + + if record.f64_matrix and record.f64_matrix.labels: + param["labels"] = record.f64_matrix.labels[0] + + parameters.append(param) + return parameters @router.get( From cf648e5135b4b419c12c7ed62e6b505b6cbacc2f Mon Sep 17 00:00:00 2001 From: Frode Aarstad Date: Wed, 16 Feb 2022 08:35:41 +0100 Subject: [PATCH 32/43] Add observations to the record --- src/ert/dark_storage/json_schema/record.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/ert/dark_storage/json_schema/record.py b/src/ert/dark_storage/json_schema/record.py index 92977e1c36d..56eeeba16a3 100644 --- a/src/ert/dark_storage/json_schema/record.py +++ b/src/ert/dark_storage/json_schema/record.py @@ -1,5 +1,5 @@ from uuid import UUID -from typing import Any, Mapping +from typing import Any, Mapping, List from pydantic import BaseModel, Field @@ -11,6 +11,7 @@ class RecordOut(_Record): id: UUID name: str userdata: Mapping[str, Any] + observations: List[Any] class Config: orm_mode = True From 614352fb103a0ca4bf839156e6d738f682561911 Mon Sep 17 00:00:00 2001 From: Frode Aarstad Date: Tue, 1 Mar 2022 11:45:11 +0100 Subject: [PATCH 33/43] Add has_observations flag to RecordOut --- src/ert/dark_storage/json_schema/record.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ert/dark_storage/json_schema/record.py b/src/ert/dark_storage/json_schema/record.py index 56eeeba16a3..2cb1d808710 100644 --- a/src/ert/dark_storage/json_schema/record.py +++ b/src/ert/dark_storage/json_schema/record.py @@ -1,5 +1,5 @@ from uuid import UUID -from typing import Any, Mapping, List +from typing import Any, Mapping, Optional from pydantic import BaseModel, Field @@ -11,7 +11,7 @@ class RecordOut(_Record): id: UUID name: str userdata: Mapping[str, Any] - observations: List[Any] + has_observations: Optional[bool] class Config: orm_mode = True From ea3eb71f6df729d04207486a194dff696391974f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=98ystein=20Olai=20Heggen?= Date: Tue, 15 Mar 2022 12:56:52 +0100 Subject: [PATCH 34/43] Add has_observations flag to Record --- src/ert/dark_storage/database_schema/record.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/ert/dark_storage/database_schema/record.py b/src/ert/dark_storage/database_schema/record.py index 385d822d44b..646cd9fc37d 100644 --- a/src/ert/dark_storage/database_schema/record.py +++ b/src/ert/dark_storage/database_schema/record.py @@ -72,6 +72,10 @@ def record_type(self) -> RecordType: def record_class(self) -> RecordClass: return self.record_info.record_class + @property + def has_observations(self) -> bool: + return len(self.observations) > 0 + class File(Base): __tablename__ = "file" From b4fa57a7246d2f551759bfa05531dc90efa8e253 Mon Sep 17 00:00:00 2001 From: DanSava Date: Thu, 12 May 2022 16:49:22 +0300 Subject: [PATCH 35/43] Add refresh facade endpoint --- src/ert/dark_storage/endpoints/updates.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/ert/dark_storage/endpoints/updates.py b/src/ert/dark_storage/endpoints/updates.py index ea54c836ed1..55f71637794 100644 --- a/src/ert/dark_storage/endpoints/updates.py +++ b/src/ert/dark_storage/endpoints/updates.py @@ -56,6 +56,14 @@ def get_update( return _update_from_db(update_obj) +@router.post("/updates/facade") +def refresh_facade( + *, + db: Session = Depends(get_db), +) -> None: + raise NotImplementedError + + def _update_from_db(update_obj: ds.Update) -> js.UpdateOut: return js.UpdateOut( id=update_obj.id, From 5cf4e2ddf34a8aeac7292e30495ff102c235ea85 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=98ystein=20Olai=20Heggen?= Date: Thu, 16 Jun 2022 09:28:07 +0200 Subject: [PATCH 36/43] Add property experiment_id on ensemble model --- src/ert/dark_storage/database_schema/ensemble.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/ert/dark_storage/database_schema/ensemble.py b/src/ert/dark_storage/database_schema/ensemble.py index 8975c041ea7..0815eaccd16 100644 --- a/src/ert/dark_storage/database_schema/ensemble.py +++ b/src/ert/dark_storage/database_schema/ensemble.py @@ -51,3 +51,7 @@ def parent_ensemble_id(self) -> PyUUID: @property def child_ensemble_ids(self) -> List[PyUUID]: return [x.ensemble_result.id for x in self.children] + + @property + def experiment_id(self) -> PyUUID: + return self.experiment.id From 145e21c97dc8954bbc5f681d3a8172e679151786 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=98ystein=20Olai=20Heggen?= Date: Fri, 10 Jun 2022 14:04:16 +0200 Subject: [PATCH 37/43] Remove all graphql related code --- src/ert/dark_storage/app.py | 3 - .../dark_storage/ext/graphene_sqlalchemy.py | 132 ------------------ src/ert/dark_storage/ext/sqlalchemy_arrays.py | 25 +--- src/ert/dark_storage/ext/uuid.py | 12 -- src/ert/dark_storage/graphql/__init__.py | 61 -------- src/ert/dark_storage/graphql/ensembles.py | 112 --------------- src/ert/dark_storage/graphql/experiments.py | 51 ------- src/ert/dark_storage/graphql/parameters.py | 26 ---- src/ert/dark_storage/graphql/responses.py | 19 --- .../dark_storage/graphql/unique_responses.py | 13 -- src/ert/dark_storage/graphql/updates.py | 15 -- src/ert/dark_storage/testing/testclient.py | 38 +---- 12 files changed, 3 insertions(+), 504 deletions(-) delete mode 100644 src/ert/dark_storage/ext/graphene_sqlalchemy.py delete mode 100644 src/ert/dark_storage/graphql/__init__.py delete mode 100644 src/ert/dark_storage/graphql/ensembles.py delete mode 100644 src/ert/dark_storage/graphql/experiments.py delete mode 100644 src/ert/dark_storage/graphql/parameters.py delete mode 100644 src/ert/dark_storage/graphql/responses.py delete mode 100644 src/ert/dark_storage/graphql/unique_responses.py delete mode 100644 src/ert/dark_storage/graphql/updates.py diff --git a/src/ert/dark_storage/app.py b/src/ert/dark_storage/app.py index b646c26e197..a7bbc39dbc8 100644 --- a/src/ert/dark_storage/app.py +++ b/src/ert/dark_storage/app.py @@ -3,10 +3,8 @@ from typing import Any from fastapi import FastAPI, Request, status from fastapi.responses import Response, RedirectResponse -from starlette.graphql import GraphQLApp from ert_storage.endpoints import router as endpoints_router -from ert_storage.graphql import router as graphql_router from ert_storage.exceptions import ErtStorageError from sqlalchemy.orm.exc import NoResultFound @@ -103,4 +101,3 @@ async def healthcheck() -> str: app.include_router(endpoints_router) -app.include_router(graphql_router) diff --git a/src/ert/dark_storage/ext/graphene_sqlalchemy.py b/src/ert/dark_storage/ext/graphene_sqlalchemy.py deleted file mode 100644 index 6dcaa671e8c..00000000000 --- a/src/ert/dark_storage/ext/graphene_sqlalchemy.py +++ /dev/null @@ -1,132 +0,0 @@ -""" -This module adds the SQLAlchemyMutation class, a graphene.Mutation whose -output mirrors an SQLAlchemyObjectType. This allows us to create mutation that behave -""" - - -from collections import OrderedDict -from typing import Any, Type, Iterable, Callable, Optional, TYPE_CHECKING, Dict - -from graphene_sqlalchemy import SQLAlchemyObjectType as _SQLAlchemyObjectType -from graphene.types.mutation import MutationOptions -from graphene.types.utils import yank_fields_from_attrs -from graphene.types.interface import Interface -from graphene.utils.get_unbound_function import get_unbound_function -from graphene.utils.props import props -from graphene.types.objecttype import ObjectTypeOptions -from graphene import ObjectType, Field, Interface -from graphql import ResolveInfo - - -__all__ = ["SQLAlchemyObjectType", "SQLAlchemyMutation"] - - -if TYPE_CHECKING: - from graphene_sqlalchemy.types.argument import Argument - import graphene.types.field - - -class SQLAlchemyObjectType(_SQLAlchemyObjectType): - class Meta: - abstract = True - - def resolve_id(self, info: ResolveInfo) -> str: - return str(self.id) - - -class SQLAlchemyMutation(SQLAlchemyObjectType): - """ - Object Type Definition (mutation field), appropriated from graphene.Mutation - - SQLAlchemyMutation is a convenience type that helps us build a Field which - takes Arguments and returns a mutation Output SQLAlchemyObjectType. - - Meta class options (optional): - resolver (Callable resolver method): Or ``mutate`` method on Mutation class. Perform data - change and return output. - arguments (Dict[str, graphene.Argument]): Or ``Arguments`` inner class with attributes on - Mutation class. Arguments to use for the mutation Field. - name (str): Name of the GraphQL type (must be unique in schema). Defaults to class - name. - description (str): Description of the GraphQL type in the schema. Defaults to class - docstring. - interfaces (Iterable[graphene.Interface]): GraphQL interfaces to extend with the payload - object. All fields from interface will be included in this object's schema. - fields (Dict[str, graphene.Field]): Dictionary of field name to Field. Not recommended to - use (prefer class attributes or ``Meta.output``). - - """ - - class Meta: - abstract = True - - @classmethod - def __init_subclass_with_meta__( - cls, - interfaces: Iterable[Type[Interface]] = (), - resolver: Callable = None, - arguments: Dict[str, "Argument"] = None, - _meta: Optional[ObjectTypeOptions] = None, - **options: Any, - ) -> None: - if not _meta: - _meta = MutationOptions(cls) - - fields = {} - - for interface in interfaces: - assert issubclass(interface, Interface), ( - 'All interfaces of {} must be a subclass of Interface. Received "{}".' - ).format(cls.__name__, interface) - fields.update(interface._meta.fields) - - fields = OrderedDict() - for base in reversed(cls.__mro__): - fields.update(yank_fields_from_attrs(base.__dict__, _as=Field)) - - if not arguments: - input_class = getattr(cls, "Arguments", None) - - if input_class: - arguments = props(input_class) - else: - arguments = {} - - if not resolver: - mutate = getattr(cls, "mutate", None) - assert mutate, "All mutations must define a mutate method in it" - resolver = get_unbound_function(mutate) - - if _meta.fields: - _meta.fields.update(fields) - else: - _meta.fields = fields - - _meta.interfaces = interfaces - _meta.resolver = resolver - _meta.arguments = arguments - - super(SQLAlchemyMutation, cls).__init_subclass_with_meta__( - _meta=_meta, **options - ) - - @classmethod - def Field( - cls, - name: Optional[str] = None, - description: Optional[str] = None, - deprecation_reason: Optional[str] = None, - required: bool = False, - **kwargs: Any, - ) -> "graphene.types.field.Field": - """Mount instance of mutation Field.""" - return Field( - cls, - args=cls._meta.arguments, - resolver=cls._meta.resolver, - name=name, - description=description or cls._meta.description, - deprecation_reason=deprecation_reason, - required=required, - **kwargs, - ) diff --git a/src/ert/dark_storage/ext/sqlalchemy_arrays.py b/src/ert/dark_storage/ext/sqlalchemy_arrays.py index 1d254c25f42..c47eb23149e 100644 --- a/src/ert/dark_storage/ext/sqlalchemy_arrays.py +++ b/src/ert/dark_storage/ext/sqlalchemy_arrays.py @@ -7,16 +7,11 @@ `convert_sqlalchemy_type.register` much in the same way that graphene_sqlalchemy does it internally for its other types. """ -from typing import Optional, Type, Union +from typing import Type, Union import sqlalchemy as sa from ert_storage.database import IS_POSTGRES -import graphene -from graphene_sqlalchemy.converter import convert_sqlalchemy_type -from graphene_sqlalchemy.registry import Registry - - __all__ = ["FloatArray", "StringArray", "IntArray"] @@ -33,21 +28,3 @@ FloatArray = type("FloatArray", (sa.PickleType,), dict(sa.PickleType.__dict__)) StringArray = type("StringArray", (sa.PickleType,), dict(sa.PickleType.__dict__)) IntArray = type("IntArray", (sa.PickleType,), dict(sa.PickleType.__dict__)) - - @convert_sqlalchemy_type.register(StringArray) - def convert_column_to_string_array( - type: SQLAlchemyColumn, column: sa.Column, registry: Optional[Registry] = None - ) -> graphene.types.structures.Structure: - return graphene.List(graphene.String) - - @convert_sqlalchemy_type.register(FloatArray) - def convert_column_to_float_array( - type: SQLAlchemyColumn, column: sa.Column, registry: Optional[Registry] = None - ) -> graphene.types.structures.Structure: - return graphene.List(graphene.Float) - - @convert_sqlalchemy_type.register(IntArray) - def convert_column_to_int_array( - type: SQLAlchemyColumn, column: sa.Column, registry: Optional[Registry] = None - ) -> graphene.types.structures.Structure: - return graphene.List(graphene.Int) diff --git a/src/ert/dark_storage/ext/uuid.py b/src/ert/dark_storage/ext/uuid.py index e9527a4cd4b..c223158abb8 100644 --- a/src/ert/dark_storage/ext/uuid.py +++ b/src/ert/dark_storage/ext/uuid.py @@ -6,11 +6,6 @@ from sqlalchemy.engine import Dialect from sqlalchemy.types import TypeDecorator, CHAR -import graphene -from graphene import UUID as GrapheneUUID -from graphene_sqlalchemy.converter import convert_sqlalchemy_type -from graphene_sqlalchemy.registry import Registry - class UUID(TypeDecorator): """Platform-independent UUID type. @@ -47,10 +42,3 @@ def process_result_value(self, value: Any, dialect: Dialect) -> Any: if not isinstance(value, SystemUUID): value = SystemUUID(value) return value - - -@convert_sqlalchemy_type.register(UUID) -def convert_column_to_uuid( - type: UUID, column: sa.Column, registry: Optional[Registry] = None -) -> graphene.types.structures.Structure: - return GrapheneUUID diff --git a/src/ert/dark_storage/graphql/__init__.py b/src/ert/dark_storage/graphql/__init__.py deleted file mode 100644 index 5e4a119fb56..00000000000 --- a/src/ert/dark_storage/graphql/__init__.py +++ /dev/null @@ -1,61 +0,0 @@ -from typing import Any, Optional -from starlette.graphql import GraphQLApp -from fastapi import APIRouter -from ert_storage.database import sessionmaker, Session -import graphene as gr -from graphql.execution.base import ExecutionResult, ResolveInfo - -from ert_storage.graphql.ensembles import Ensemble, CreateEnsemble -from ert_storage.graphql.experiments import Experiment, CreateExperiment -from ert_storage import database_schema as ds - - -class Mutations(gr.ObjectType): - create_experiment = CreateExperiment.Field() - create_ensemble = CreateEnsemble.Field(experiment_id=gr.ID(required=True)) - - -class Query(gr.ObjectType): - experiments = gr.List(Experiment) - experiment = gr.Field(Experiment, id=gr.ID(required=True)) - ensemble = gr.Field(Ensemble, id=gr.ID(required=True)) - - @staticmethod - def resolve_experiments(root: None, info: ResolveInfo) -> ds.Experiment: - return Experiment.get_query(info).all() - - @staticmethod - def resolve_experiment(root: None, info: ResolveInfo, id: str) -> ds.Experiment: - return Experiment.get_query(info).filter_by(id=id).one() - - @staticmethod - def resolve_ensemble(root: None, info: ResolveInfo, id: str) -> ds.Ensemble: - return Ensemble.get_query(info).filter_by(id=id).one() - - -class Schema(gr.Schema): - """ - Extended graphene Schema class, where `execute` creates a database session - and passes it on further. - """ - - def __init__(self, *args: Any, **kwargs: Any) -> None: - super().__init__(*args, **kwargs) - self.override_session: Optional[sessionmaker] = None - - def execute(self, *args: Any, **kwargs: Any) -> ExecutionResult: - kwargs.setdefault("context_value", {}) - if self.override_session is not None: - session_obj = self.override_session() - else: - session_obj = Session() - with session_obj as session: - kwargs["context_value"]["session"] = session - - return super().execute(*args, **kwargs) - - -schema = Schema(query=Query, mutation=Mutations) -graphql_app = GraphQLApp(schema=schema) -router = APIRouter(tags=["graphql"]) -router.add_route("/gql", graphql_app) diff --git a/src/ert/dark_storage/graphql/ensembles.py b/src/ert/dark_storage/graphql/ensembles.py deleted file mode 100644 index fe123bd72a8..00000000000 --- a/src/ert/dark_storage/graphql/ensembles.py +++ /dev/null @@ -1,112 +0,0 @@ -from typing import List, Iterable, Optional, TYPE_CHECKING -import graphene as gr -from graphene_sqlalchemy.utils import get_session - -from ert_storage.ext.graphene_sqlalchemy import SQLAlchemyObjectType, SQLAlchemyMutation -from ert_storage import database_schema as ds - - -if TYPE_CHECKING: - from graphql.execution.base import ResolveInfo - from ert_storage.graphql.experiments import Experiment - - -class Ensemble(SQLAlchemyObjectType): - class Meta: - model = ds.Ensemble - - child_ensembles = gr.List(lambda: Ensemble) - parent_ensemble = gr.Field(lambda: Ensemble) - responses = gr.Field( - gr.List("ert_storage.graphql.responses.Response"), - names=gr.Argument(gr.List(gr.String), required=False, default_value=None), - ) - unique_responses = gr.List("ert_storage.graphql.unique_responses.UniqueResponse") - - parameters = gr.List("ert_storage.graphql.parameters.Parameter") - - def resolve_child_ensembles( - root: ds.Ensemble, info: "ResolveInfo" - ) -> List[ds.Ensemble]: - return [x.ensemble_result for x in root.children] - - def resolve_parent_ensemble( - root: ds.Ensemble, info: "ResolveInfo" - ) -> Optional[ds.Ensemble]: - update = root.parent - if update is not None: - return update.ensemble_reference - return None - - def resolve_unique_responses( - root: ds.Ensemble, info: "ResolveInfo" - ) -> Iterable[ds.RecordInfo]: - session = info.context["session"] # type: ignore - return root.record_infos.filter_by(record_class=ds.RecordClass.response) - - def resolve_responses( - root: ds.Ensemble, info: "ResolveInfo", names: Optional[Iterable[str]] = None - ) -> Iterable[ds.Record]: - session = info.context["session"] # type: ignore - q = ( - session.query(ds.Record) - .join(ds.RecordInfo) - .filter_by(ensemble=root, record_class=ds.RecordClass.response) - ) - if names is not None: - q = q.filter(ds.RecordInfo.name.in_(names)) - return q.all() - - def resolve_parameters( - root: ds.Ensemble, info: "ResolveInfo" - ) -> Iterable[ds.Record]: - session = info.context["session"] # type: ignore - return ( - session.query(ds.Record) - .join(ds.RecordInfo) - .filter_by(ensemble=root, record_class=ds.RecordClass.parameter) - .all() - ) - - -class CreateEnsemble(SQLAlchemyMutation): - class Meta: - model = ds.Ensemble - - class Arguments: - parameter_names = gr.List(gr.String) - size = gr.Int() - active_realizations = gr.List(gr.Int) - - @staticmethod - def mutate( - root: Optional["Experiment"], - info: "ResolveInfo", - parameter_names: List[str], - size: int, - active_realizations: Optional[List[int]] = None, - experiment_id: Optional[str] = None, - ) -> ds.Ensemble: - db = get_session(info.context) - - if experiment_id is not None: - experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() - elif hasattr(root, "id"): - experiment = root - else: - raise ValueError("ID is required") - - if active_realizations is None: - active_realizations = list(range(size)) - - ensemble = ds.Ensemble( - parameter_names=parameter_names, - response_names=[], - experiment=experiment, - size=size, - active_realizations=active_realizations, - ) - - db.add(ensemble) - db.commit() - return ensemble diff --git a/src/ert/dark_storage/graphql/experiments.py b/src/ert/dark_storage/graphql/experiments.py deleted file mode 100644 index e8654a777c1..00000000000 --- a/src/ert/dark_storage/graphql/experiments.py +++ /dev/null @@ -1,51 +0,0 @@ -from typing import List, Mapping, TYPE_CHECKING -import graphene as gr -from graphene_sqlalchemy.utils import get_session - -from ert_storage.ext.graphene_sqlalchemy import SQLAlchemyObjectType, SQLAlchemyMutation - -from ert_storage.graphql.ensembles import Ensemble, CreateEnsemble -from ert_storage.graphql.updates import Update -from ert_storage import database_schema as ds, json_schema as js -from ert_storage.endpoints.experiments import experiment_priors_to_dict - - -if TYPE_CHECKING: - from graphql.execution.base import ResolveInfo - - -class Experiment(SQLAlchemyObjectType): - class Meta: - model = ds.Experiment - - ensembles = gr.List(Ensemble) - priors = gr.JSONString() - - def resolve_ensembles( - root: ds.Experiment, info: "ResolveInfo" - ) -> List[ds.Ensemble]: - return root.ensembles - - def resolve_priors(root: ds.Experiment, info: "ResolveInfo") -> Mapping[str, dict]: - return experiment_priors_to_dict(root) - - -class CreateExperiment(SQLAlchemyMutation): - class Arguments: - name = gr.String() - - class Meta: - model = ds.Experiment - - create_ensemble = CreateEnsemble.Field() - - @staticmethod - def mutate(root: None, info: "ResolveInfo", name: str) -> ds.Experiment: - db = get_session(info.context) - - experiment = ds.Experiment(name=name) - - db.add(experiment) - db.commit() - - return experiment diff --git a/src/ert/dark_storage/graphql/parameters.py b/src/ert/dark_storage/graphql/parameters.py deleted file mode 100644 index 45e47fd0f8b..00000000000 --- a/src/ert/dark_storage/graphql/parameters.py +++ /dev/null @@ -1,26 +0,0 @@ -from typing import Any, List, Optional, TYPE_CHECKING -import graphene as gr -from sqlalchemy.orm import Session - -from ert_storage.ext.graphene_sqlalchemy import SQLAlchemyObjectType -from ert_storage import database_schema as ds -from ert_storage.endpoints.experiments import prior_to_dict - - -if TYPE_CHECKING: - from graphql.execution.base import ResolveInfo - - -class Parameter(SQLAlchemyObjectType): - class Meta: - model = ds.Record - - name = gr.String() - prior = gr.JSONString() - - def resolve_name(root: ds.Record, info: "ResolveInfo") -> str: - return root.name - - def resolve_prior(root: ds.Record, info: "ResolveInfo") -> Optional[dict]: - prior = root.record_info.prior - return prior_to_dict(prior) if prior is not None else None diff --git a/src/ert/dark_storage/graphql/responses.py b/src/ert/dark_storage/graphql/responses.py deleted file mode 100644 index 8154009ad51..00000000000 --- a/src/ert/dark_storage/graphql/responses.py +++ /dev/null @@ -1,19 +0,0 @@ -from typing import List, Optional, TYPE_CHECKING -import graphene as gr - -from ert_storage.ext.graphene_sqlalchemy import SQLAlchemyObjectType -from ert_storage import database_schema as ds - - -if TYPE_CHECKING: - from graphql.execution.base import ResolveInfo - - -class Response(SQLAlchemyObjectType): - class Meta: - model = ds.Record - - name = gr.String() - - def resolve_name(root: ds.Record, info: "ResolveInfo") -> str: - return root.name diff --git a/src/ert/dark_storage/graphql/unique_responses.py b/src/ert/dark_storage/graphql/unique_responses.py deleted file mode 100644 index 4d8acac50af..00000000000 --- a/src/ert/dark_storage/graphql/unique_responses.py +++ /dev/null @@ -1,13 +0,0 @@ -from typing import TYPE_CHECKING - -from ert_storage.ext.graphene_sqlalchemy import SQLAlchemyObjectType -from ert_storage import database_schema as ds - - -if TYPE_CHECKING: - from graphql.execution.base import ResolveInfo - - -class UniqueResponse(SQLAlchemyObjectType): - class Meta: - model = ds.RecordInfo diff --git a/src/ert/dark_storage/graphql/updates.py b/src/ert/dark_storage/graphql/updates.py deleted file mode 100644 index 84008735d6f..00000000000 --- a/src/ert/dark_storage/graphql/updates.py +++ /dev/null @@ -1,15 +0,0 @@ -from typing import List, TYPE_CHECKING -import graphene as gr -from graphene_sqlalchemy.utils import get_session - -from ert_storage.ext.graphene_sqlalchemy import SQLAlchemyObjectType, SQLAlchemyMutation -from ert_storage import database_schema as ds - - -if TYPE_CHECKING: - from graphql.execution.base import ResolveInfo - - -class Update(SQLAlchemyObjectType): - class Meta: - model = ds.Update diff --git a/src/ert/dark_storage/testing/testclient.py b/src/ert/dark_storage/testing/testclient.py index a8efcd9e58e..bee23d3caf5 100644 --- a/src/ert/dark_storage/testing/testclient.py +++ b/src/ert/dark_storage/testing/testclient.py @@ -4,14 +4,12 @@ Any, AsyncGenerator, Generator, - Mapping, - MutableMapping, Optional, TYPE_CHECKING, Tuple, Union, ) -from pprint import pformat + from contextlib import contextmanager from fastapi import Depends from sqlalchemy.orm.session import Session @@ -20,27 +18,12 @@ TestClient as StarletteTestClient, ASGI2App, ASGI3App, - Cookies, - Params, - DataType, - TimeOut, - FileType, ) from sqlalchemy.orm import sessionmaker -from graphene import Schema as GrapheneSchema -from graphene.test import Client as GrapheneClient from ert_storage.security import security -if TYPE_CHECKING: - from promise import Promise - from rx import Observable - from graphql.execution import ExecutionResult - - GraphQLResult = Union[ExecutionResult, Observable, Promise[ExecutionResult]] - - class ClientError(RuntimeError): pass @@ -52,7 +35,6 @@ def __init__( self, app: Union[ASGI2App, ASGI3App], session: sessionmaker, - gql_schema: GrapheneSchema, base_url: str = "http://testserver", raise_server_exceptions: bool = True, root_path: str = "", @@ -61,7 +43,6 @@ def __init__( self.http_client = StarletteTestClient( app, base_url, raise_server_exceptions, root_path ) - self.gql_client = GrapheneClient(gql_schema) self.session = session def get( @@ -126,18 +107,6 @@ def delete( self._check(check_status_code, resp) return resp - def gql_execute( - self, - request_string: str, - variable_values: Optional[Mapping[str, Any]] = None, - check: bool = True, - ) -> dict: - doc = self.gql_client.execute(request_string, variable_values=variable_values) - if self.raise_on_client_error and check and "errors" in doc: - raise ClientError(f"GraphQL query returned an error:\n{pformat(doc)}") - - return doc - def _check( self, check_status_code: Optional[int], response: requests.Response ) -> None: @@ -176,14 +145,11 @@ def testclient_factory() -> Generator[_TestClient, None, None]: rollback = True from ert_storage.app import app - from ert_storage.graphql import schema session, transaction, connection = _begin_transaction() - schema.override_session = session - yield _TestClient(app, session=session, gql_schema=schema) + yield _TestClient(app, session=session) - schema.override_session = None _end_transaction(transaction, connection, rollback) if env_unset: From 20bb1be7de9ee189c8b25bb863640d290f399f36 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=98ystein=20Olai=20Heggen?= Date: Tue, 16 Aug 2022 15:18:03 +0200 Subject: [PATCH 38/43] Handle specific dbapierror in get_db fastapi had breaking change in 0.74.0 where we now need to specifically handle only exceptions related to our database transaction. https://fastapi.tiangolo.com/release-notes/#breaking-changes --- src/ert/dark_storage/database.py | 3 ++- src/ert/dark_storage/testing/testclient.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/ert/dark_storage/database.py b/src/ert/dark_storage/database.py index ec2d412ff1f..f31352ffcb9 100644 --- a/src/ert/dark_storage/database.py +++ b/src/ert/dark_storage/database.py @@ -1,6 +1,7 @@ import os from typing import Any from fastapi import Depends +from sqlalchemy.exc import DBAPIError from sqlalchemy import create_engine from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker @@ -46,7 +47,7 @@ async def get_db(*, _: None = Depends(security)) -> Any: yield db db.commit() db.close() - except: + except DBAPIError: db.rollback() db.close() raise diff --git a/src/ert/dark_storage/testing/testclient.py b/src/ert/dark_storage/testing/testclient.py index bee23d3caf5..97bd4d2ca1c 100644 --- a/src/ert/dark_storage/testing/testclient.py +++ b/src/ert/dark_storage/testing/testclient.py @@ -5,13 +5,13 @@ AsyncGenerator, Generator, Optional, - TYPE_CHECKING, Tuple, Union, ) from contextlib import contextmanager from fastapi import Depends +from sqlalchemy.exc import DBAPIError from sqlalchemy.orm.session import Session from sqlalchemy.engine.base import Transaction from starlette.testclient import ( @@ -179,7 +179,7 @@ async def override_get_db( yield db db.commit() db.close() - except: + except DBAPIError: db.rollback() db.close() raise From 41bdca1a99b501aafb3a00debfe8dbc0c22628f5 Mon Sep 17 00:00:00 2001 From: Morten Bendiksen Date: Tue, 16 Aug 2022 11:47:38 +0200 Subject: [PATCH 39/43] Add endpoint /server/info Endpoint gives info about server, for now only a name. --- src/ert/dark_storage/endpoints/__init__.py | 2 ++ src/ert/dark_storage/endpoints/server.py | 13 +++++++++++++ 2 files changed, 15 insertions(+) create mode 100644 src/ert/dark_storage/endpoints/server.py diff --git a/src/ert/dark_storage/endpoints/__init__.py b/src/ert/dark_storage/endpoints/__init__.py index b63a5711241..da9e48b1659 100644 --- a/src/ert/dark_storage/endpoints/__init__.py +++ b/src/ert/dark_storage/endpoints/__init__.py @@ -6,6 +6,7 @@ from .updates import router as updates_router from .compute.misfits import router as misfits_router from .responses import router as response_router +from .server import router as server_router router = APIRouter() router.include_router(experiments_router) @@ -15,3 +16,4 @@ router.include_router(updates_router) router.include_router(misfits_router) router.include_router(response_router) +router.include_router(server_router) diff --git a/src/ert/dark_storage/endpoints/server.py b/src/ert/dark_storage/endpoints/server.py new file mode 100644 index 00000000000..c31a1ad42f3 --- /dev/null +++ b/src/ert/dark_storage/endpoints/server.py @@ -0,0 +1,13 @@ +from typing import Mapping, Any +from fastapi import APIRouter, Depends +from ert_storage.database import Session, get_db + +router = APIRouter(tags=["info"]) + + +@router.get("/server/info", response_model=Mapping[str, Any]) +def info( + *, + db: Session = Depends(get_db), +) -> Mapping[str, Any]: + return {"name": "Ert Storage Server"} From f1939970361e8b7cc93b1189ca915d5200ac5c29 Mon Sep 17 00:00:00 2001 From: xjules Date: Tue, 14 Feb 2023 13:46:49 +0100 Subject: [PATCH 40/43] Fix imports and text type in sql args --- src/ert/dark_storage/database.py | 9 +++-- src/ert/dark_storage/testing/testclient.py | 39 +++++++--------------- 2 files changed, 18 insertions(+), 30 deletions(-) diff --git a/src/ert/dark_storage/database.py b/src/ert/dark_storage/database.py index f31352ffcb9..17cbded4db4 100644 --- a/src/ert/dark_storage/database.py +++ b/src/ert/dark_storage/database.py @@ -1,12 +1,14 @@ import os from typing import Any + from fastapi import Depends -from sqlalchemy.exc import DBAPIError from sqlalchemy import create_engine +from sqlalchemy.exc import DBAPIError from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker -from ert_storage.security import security +from sqlalchemy.sql import text +from ert_storage.security import security ENV_RDBMS = "ERT_STORAGE_DATABASE_URL" ENV_BLOB = "ERT_STORAGE_AZURE_CONNECTION_STRING" @@ -42,7 +44,7 @@ async def get_db(*, _: None = Depends(security)) -> Any: # Make PostgreSQL return float8 columns with highest precision. If we don't # do this, we may lose up to 3 of the least significant digits. if IS_POSTGRES: - db.execute("SET extra_float_digits=3") + db.execute(text("SET extra_float_digits=3")) try: yield db db.commit() @@ -55,6 +57,7 @@ async def get_db(*, _: None = Depends(security)) -> Any: if HAS_AZURE_BLOB_STORAGE: import asyncio + from azure.core.exceptions import ResourceNotFoundError from azure.storage.blob.aio import ContainerClient diff --git a/src/ert/dark_storage/testing/testclient.py b/src/ert/dark_storage/testing/testclient.py index 97bd4d2ca1c..0cfe674a888 100644 --- a/src/ert/dark_storage/testing/testclient.py +++ b/src/ert/dark_storage/testing/testclient.py @@ -1,25 +1,16 @@ import os -import requests -from typing import ( - Any, - AsyncGenerator, - Generator, - Optional, - Tuple, - Union, -) - from contextlib import contextmanager +from typing import Any, AsyncGenerator, Generator, Optional, Tuple, Union + +import requests from fastapi import Depends -from sqlalchemy.exc import DBAPIError -from sqlalchemy.orm.session import Session from sqlalchemy.engine.base import Transaction -from starlette.testclient import ( - TestClient as StarletteTestClient, - ASGI2App, - ASGI3App, -) +from sqlalchemy.exc import DBAPIError from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm.session import Session +from sqlalchemy.sql import text +from starlette.testclient import ASGI2App, ASGI3App +from starlette.testclient import TestClient as StarletteTestClient from ert_storage.security import security @@ -161,10 +152,7 @@ def testclient_factory() -> Generator[_TestClient, None, None]: def _override_get_db(session: sessionmaker) -> None: from ert_storage.app import app - from ert_storage.database import ( - get_db, - IS_POSTGRES, - ) + from ert_storage.database import IS_POSTGRES, get_db async def override_get_db( *, _: None = Depends(security) @@ -174,7 +162,7 @@ async def override_get_db( # Make PostgreSQL return float8 columns with highest precision. If we don't # do this, we may lose up to 3 of the least significant digits. if IS_POSTGRES: - db.execute("SET extra_float_digits=3") + db.execute(text("SET extra_float_digits=3")) try: yield db db.commit() @@ -188,17 +176,14 @@ async def override_get_db( def _begin_transaction() -> _TransactionInfo: - from ert_storage.database import ( - engine, - IS_SQLITE, - HAS_AZURE_BLOB_STORAGE, - ) + from ert_storage.database import HAS_AZURE_BLOB_STORAGE, IS_SQLITE, engine from ert_storage.database_schema import Base if IS_SQLITE: Base.metadata.create_all(bind=engine) if HAS_AZURE_BLOB_STORAGE: import asyncio + from ert_storage.database import create_container_if_not_exist loop = asyncio.get_event_loop() From e934cfc8cbfbca0024623c0c9c4faf98b093258c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=98ystein=20Olai=20Heggen?= Date: Fri, 17 Feb 2023 15:01:15 +0100 Subject: [PATCH 41/43] Update formatting according to new version of black --- src/ert/dark_storage/endpoints/ensembles.py | 1 - src/ert/dark_storage/endpoints/updates.py | 1 - 2 files changed, 2 deletions(-) diff --git a/src/ert/dark_storage/endpoints/ensembles.py b/src/ert/dark_storage/endpoints/ensembles.py index d59eb483f70..f403cfe8f2c 100644 --- a/src/ert/dark_storage/endpoints/ensembles.py +++ b/src/ert/dark_storage/endpoints/ensembles.py @@ -14,7 +14,6 @@ def post_ensemble( *, db: Session = Depends(get_db), ens_in: js.EnsembleIn, experiment_id: UUID ) -> ds.Ensemble: - experiment = db.query(ds.Experiment).filter_by(id=experiment_id).one() active_reals = ( ens_in.active_realizations diff --git a/src/ert/dark_storage/endpoints/updates.py b/src/ert/dark_storage/endpoints/updates.py index 55f71637794..cdd4585a934 100644 --- a/src/ert/dark_storage/endpoints/updates.py +++ b/src/ert/dark_storage/endpoints/updates.py @@ -12,7 +12,6 @@ def create_update( db: Session = Depends(get_db), update: js.UpdateIn, ) -> js.UpdateOut: - ensemble = db.query(ds.Ensemble).filter_by(id=update.ensemble_reference_id).one() update_obj = ds.Update( algorithm=update.algorithm, From 029fa16c44d566b028f6f15f52266eb2a05aa9c5 Mon Sep 17 00:00:00 2001 From: Andreas Eknes Lie Date: Tue, 2 May 2023 11:53:42 +0200 Subject: [PATCH 42/43] Migrate declarative_base SQLalchemy --- src/ert/dark_storage/database.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/ert/dark_storage/database.py b/src/ert/dark_storage/database.py index 17cbded4db4..913391b7187 100644 --- a/src/ert/dark_storage/database.py +++ b/src/ert/dark_storage/database.py @@ -4,8 +4,7 @@ from fastapi import Depends from sqlalchemy import create_engine from sqlalchemy.exc import DBAPIError -from sqlalchemy.ext.declarative import declarative_base -from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm import declarative_base, sessionmaker from sqlalchemy.sql import text from ert_storage.security import security From fe5ec4081b0d0056c278a0bbf4e154ffa330a66f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A5vard=20Berland?= Date: Tue, 11 Jul 2023 07:29:38 +0200 Subject: [PATCH 43/43] Pin pydantic<2 and fastapi --- src/ert/dark_storage/database.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/ert/dark_storage/database.py b/src/ert/dark_storage/database.py index 913391b7187..7f80ebadb67 100644 --- a/src/ert/dark_storage/database.py +++ b/src/ert/dark_storage/database.py @@ -1,10 +1,11 @@ import os -from typing import Any +from typing import Any, Callable from fastapi import Depends from sqlalchemy import create_engine from sqlalchemy.exc import DBAPIError from sqlalchemy.orm import declarative_base, sessionmaker +from sqlalchemy.orm import Session as SessionType from sqlalchemy.sql import text from ert_storage.security import security @@ -31,7 +32,7 @@ def get_env_rdbms() -> str: engine = create_engine(URI_RDBMS, connect_args={"check_same_thread": False}) else: engine = create_engine(URI_RDBMS, pool_size=50, max_overflow=100) -Session = sessionmaker(autocommit=False, autoflush=False, bind=engine) +Session: Callable = sessionmaker(autocommit=False, autoflush=False, bind=engine) Base = declarative_base()