Merge pull request #101 from UW-Macrostrat/strabospot-integration

Strabospot integration

cli/macrostrat/cli/cli.py

@@ -292,13 +292,22 @@ def update_tileserver(db):
             "Mapboard subsystem enabled, but no mapboard_database setting found"
         )
 
+from macrostrat.integrations import app as integrations_app
+
+main.add_typer(
+    integrations_app,
+    name="integrations",
+    short_help="Integrations with other data systems",
+    rich_help_panel="Subsystems",
+)
+
 
 app.subsystems.add(MacrostratAPISubsystem(app))
 
 if sgp_url := getattr(settings, "sgp_database", None):
     from .subsystems.sgp import sgp
 
-    main.add_typer(sgp, rich_help_panel="Subsystems")
+    main.add_typer(sgp, rich_help_panel="Integrations")
 
 # Mariadb CLI
 if mariadb_url := getattr(settings, "mysql_database", None):
@@ -317,7 +326,7 @@ def update_tileserver(db):
 main.add_typer(
     kg_cli,
     name="xdd",
-    rich_help_panel="Subsystems",
+    rich_help_panel="Integrations",
     short_help="Manage xDD integration",
 )
 

cli/macrostrat/cli/database/__init__.py

@@ -1,5 +1,4 @@
 from os import environ
-from pathlib import Path
 from sys import exit, stderr, stdin, stdout
 from typing import Any, Callable
 
@@ -10,13 +9,16 @@
 from typer import Argument, Option
 
 from macrostrat.core import MacrostratSubsystem, app
-from macrostrat.database import Database
+from macrostrat.core.migrations import run_migrations
 from macrostrat.utils import get_logger
 from macrostrat.utils.shell import run
 
 from .._dev.utils import raw_database_url
 from ._legacy import get_db
-from .migrations import run_migrations
+
+# First, register all migrations
+# NOTE: right now, this is quite implicit.
+from .migrations import *
 from .utils import engine_for_db_name
 
 log = get_logger(__name__)

cli/macrostrat/cli/database/_legacy.py

@@ -6,6 +6,7 @@
 from sqlalchemy import create_engine
 
 from macrostrat.core.config import MYSQL_DATABASE, PG_DATABASE
+from macrostrat.core.database import get_database, refresh_database
 
 
 # Connect to MySQL
@@ -34,29 +35,13 @@ def pgConnection():
 
 
 def get_pg_credentials():
-    engine = create_engine(PG_DATABASE)
-    return engine.url
-
-
-# Lazily initialize Database
-db = None
+    db = get_database()
+    return db.engine.url
 
 
 def get_db():
-    from macrostrat.database import Database
-
-    global db
-    if db is None:
-        db = Database(PG_DATABASE)
-    return db
+    return get_database()
 
 
 def refresh_db():
-    from macrostrat.database import Database, scoped_session
-
-    global db
-    if db is not None:
-        db.session.flush()
-        db.session.close()
-    db = Database(PG_DATABASE)
-    return db
+    return refresh_database()

cli/macrostrat/cli/database/mariadb/postgresql_migration/__init__.py

@@ -6,14 +6,14 @@
 from sqlalchemy.engine import Engine
 
 from macrostrat.core import app
+from macrostrat.core.migrations import migration_has_been_run
 from macrostrat.database import create_database, database_exists
 from macrostrat.database.utils import run_query, run_sql
 from macrostrat.utils import get_logger
 from macrostrat.utils.shell import run
 
 from ...._dev.utils import raw_database_url
 from ..._legacy import get_db
-from ...migrations import migration_has_been_run
 from ...utils import docker_internal_url, pg_temp_user
 from ..restore import copy_mariadb_database
 from ..utils import mariadb_engine

cli/macrostrat/cli/database/migrations/__init__.py

@@ -1,11 +1,8 @@
-from graphlib import TopologicalSorter
 from pathlib import Path
 
-from rich import print
-
+from macrostrat.core.migrations import ApplicationStatus, Migration
 from macrostrat.database import Database
 
-from .._legacy import get_db, refresh_db
 from . import (
     api_v3,
     baseline,
@@ -21,7 +18,6 @@
     tileserver,
     update_macrostrat,
 )
-from .base import ApplicationStatus, Migration
 
 __dir__ = Path(__file__).parent
 
@@ -69,92 +65,3 @@ def has_enum(db: Database, name: str, schema: str = None):
     return db.run_query(
         f"select exists ({sql})", dict(name=name, schema=schema)
     ).scalar()
-
-
-def run_migrations(
-    apply: bool = False,
-    name: str = None,
-    force: bool = False,
-    data_changes: bool = False,
-):
-    """Apply database migrations"""
-    db = get_db()
-
-    # Check if migrations need to be run and if not, run them
-
-    if force and not name:
-        raise ValueError("--force can only be applied with --name")
-
-    # Find all subclasses of Migration among imported modules
-    migrations = Migration.__subclasses__()
-
-    # Instantiate each migration, then sort topologically according to dependency order
-    instances = [cls() for cls in migrations]
-    graph = {inst.name: inst.depends_on for inst in instances}
-    order = list(TopologicalSorter(graph).static_order())
-    instances.sort(key=lambda i: order.index(i.name))
-
-    # While iterating over migrations, keep track of which have already applied
-    completed_migrations = []
-
-    for _migration in instances:
-        _name = _migration.name
-
-        # Check whether the migration is capable of applying, or has already applied
-        apply_status = _migration.should_apply(db)
-        if apply_status == ApplicationStatus.APPLIED:
-            completed_migrations.append(_migration.name)
-
-        # If --name is specified, only run the migration with the matching name
-        if name is not None and name != _name:
-            continue
-
-        # By default, don't run migrations that depend on other non-applied migrations
-        dependencies_met = all(d in completed_migrations for d in _migration.depends_on)
-        if not dependencies_met and not force:
-            print(f"Dependencies not met for migration [cyan]{_name}[/cyan]")
-            continue
-
-        if force or apply_status == ApplicationStatus.CAN_APPLY:
-            if not apply:
-                print(f"Would apply migration [cyan]{_name}[/cyan]")
-            else:
-                if _migration.destructive and not data_changes and not force:
-                    print(
-                        f"Migration [cyan]{_name}[/cyan] would alter data in the database. Run with --force or --data-changes"
-                    )
-                    return
-
-                print(f"Applying migration [cyan]{_name}[/cyan]")
-                _migration.apply(db)
-                # After running migration, reload the database and confirm that application was sucessful
-                db = refresh_db()
-                if _migration.should_apply(db) == ApplicationStatus.APPLIED:
-                    completed_migrations.append(_migration.name)
-        elif apply_status == ApplicationStatus.APPLIED:
-            print(f"Migration [cyan]{_name}[/cyan] already applied")
-        else:
-            print(f"Migration [cyan]{_name}[/cyan] cannot apply")
-
-        # Short circuit after applying the migration specified by --name
-        if name is not None and name == _name:
-            break
-
-    # Notify PostgREST to reload the schema cache
-    db.run_sql("NOTIFY pgrst, 'reload schema';")
-
-
-def migration_has_been_run(*names: str):
-    db = get_db()
-    migrations = Migration.__subclasses__()
-
-    available_migrations = {m.name for m in migrations}
-    if not set(names).issubset(available_migrations):
-        raise ValueError(f"Unknown migrations: {set(names) - available_migrations}")
-
-    for _migration in migrations:
-        if _migration.name in names:
-            apply_status = _migration.should_apply(db)
-            if apply_status != ApplicationStatus.APPLIED:
-                return True
-    return False

cli/macrostrat/cli/database/migrations/api_v3/__init__.py

@@ -1,4 +1,4 @@
-from ..base import Migration, exists
+from macrostrat.core.migrations import Migration, exists
 
 
 class BaselineMigration(Migration):

cli/macrostrat/cli/database/migrations/baseline/__init__.py

@@ -1,4 +1,4 @@
-from ..base import Migration, exists, not_exists
+from macrostrat.core.migrations import Migration, exists
 
 
 class BaselineMigration(Migration):

cli/macrostrat/cli/database/migrations/column_builder/__init__.py

@@ -1,6 +1,4 @@
-from macrostrat.database import Database
-
-from ..base import Migration, schema_exists
+from macrostrat.core.migrations import Migration, schema_exists
 
 
 class ColumnBuilderMigration(Migration):

cli/macrostrat/cli/database/migrations/macrostrat_mariadb/__init__.py

@@ -1,8 +1,6 @@
 from pathlib import Path
 
-from macrostrat.database import Database
-
-from ..base import Migration, exists, has_fks
+from macrostrat.core.migrations import Migration, exists
 
 __dir__ = Path(__file__).parent
 

cli/macrostrat/cli/database/migrations/map_source_slugs/__init__.py

@@ -1,9 +1,8 @@
 from psycopg2.sql import Identifier
 
+from macrostrat.core.migrations import ApplicationStatus, Migration
 from macrostrat.database import Database
 
-from ..base import ApplicationStatus, Migration
-
 MATCHES_SLUG_SQL = """
 SELECT table_name
 FROM information_schema.tables

cli/macrostrat/cli/database/migrations/map_sources/__init__.py

@@ -1,6 +1,4 @@
-from macrostrat.database import Database
-
-from ..base import Migration, schema_exists, view_exists
+from macrostrat.core.migrations import Migration, view_exists
 
 
 class MapsSourcesMetadataMigration(Migration):

cli/macrostrat/cli/database/migrations/maps_scale_custom_type/__init__.py

@@ -1,7 +1,6 @@
+from macrostrat.core.migrations import Migration, custom_type_exists
 from macrostrat.database import Database
 
-from ..base import Migration, custom_type_exists
-
 
 class MapsScaleCustomTypeMigration(Migration):
     name = "maps-scale-type"

cli/macrostrat/cli/database/migrations/maps_source_operations/__init__.py

@@ -1,4 +1,4 @@
-from ..base import Migration, exists, view_exists
+from macrostrat.core.migrations import Migration, exists
 
 
 class MapsSoureOperationsMigration(Migration):

cli/macrostrat/cli/database/migrations/partition_carto/__init__.py

@@ -1,8 +1,6 @@
 from pathlib import Path
 
-from macrostrat.database import Database
-
-from ..base import Migration, not_exists, view_exists
+from macrostrat.core.migrations import Migration, view_exists
 
 __dir__ = Path(__file__).parent
 

cli/macrostrat/cli/database/migrations/partition_maps/__init__.py

@@ -1,8 +1,6 @@
 from pathlib import Path
 
-from macrostrat.database import Database
-
-from ..base import ApplicationStatus, Migration, exists, view_exists
+from macrostrat.core.migrations import Migration, exists, view_exists
 
 __dir__ = Path(__file__).parent
 

cli/macrostrat/cli/database/migrations/points/__init__.py

@@ -1,4 +1,4 @@
-from ..base import Migration, exists, view_exists
+from macrostrat.core.migrations import Migration, exists, view_exists
 
 
 class PointsMigration(Migration):

cli/macrostrat/cli/database/migrations/tileserver/__init__.py

@@ -4,9 +4,7 @@
 
 from pathlib import Path
 
-from macrostrat.database import Database
-
-from ..base import Migration, schema_exists
+from macrostrat.core.migrations import Migration, schema_exists
 
 __dir__ = Path(__file__).parent
 

cli/macrostrat/cli/database/migrations/update_macrostrat/__init__.py

@@ -1,8 +1,6 @@
 from pathlib import Path
 
-from macrostrat.database import Database
-
-from ..base import Migration, exists, has_fks
+from macrostrat.core.migrations import Migration, exists, has_fks
 
 __dir__ = Path(__file__).parent
 

cli/pyproject.toml

@@ -25,6 +25,7 @@ ipython = "^8.5.0"
 "macrostrat.utils" = "^1.2.2"
 "macrostrat.dinosaur" = "^3.0.1"
 "macrostrat.core" = { path = "../core", develop = true }
+"macrostrat.integrations" = { path = "../integrations", develop = true }
 "macrostrat.map-integration" = { path = "../map-integration", develop = true }
 "criticalmaas.ta1-geopackage" = "^0.2.0"
 numpy = "^1.23.4"

core/macrostrat/core/database.py

@@ -0,0 +1,25 @@
+from contextvars import ContextVar
+
+from macrostrat.database import Database
+
+from .config import PG_DATABASE
+
+db_ctx: ContextVar[Database | None] = ContextVar("db_ctx", default=None)
+
+
+def get_database():
+    from macrostrat.database import Database
+
+    db = db_ctx.get()
+    if db is None:
+        db = Database(PG_DATABASE)
+        db_ctx.set(db)
+    return db
+
+
+def refresh_database():
+    db = get_database()
+    db.session.flush()
+    db.session.close()
+    db_ctx.set(None)
+    return get_database()