Database utils updates

app-frame/macrostrat/app_frame/__init__.py

@@ -1,5 +1,6 @@
 from macrostrat.app_frame.compose.base import compose
+
+from .control_command import BackendType, CommandBase, ControlCommand
 from .core import Application
 from .exc import ApplicationError
 from .subsystems import Subsystem, SubsystemError, SubsystemManager
-from .control_command import ControlCommand, CommandBase, BackendType

app-frame/macrostrat/app_frame/compose/__init__.py

@@ -10,10 +10,11 @@
 from typer import Context, Typer
 
 from macrostrat.utils import get_logger
-from .base import check_status, compose
-from .follow_logs import Result, command_stream, follow_logs
+
 from ..core import Application
 from ..utils import add_click_command
+from .base import check_status, compose
+from .follow_logs import Result, command_stream, follow_logs
 
 # Typer command-line application
 

app-frame/macrostrat/app_frame/control_command.py

@@ -3,11 +3,11 @@
 from enum import Enum
 from os import environ
 
-from typer import Context, Option
-from typer import rich_utils
+from typer import Context, Option, rich_utils
 from typer.models import TyperInfo
 
 from macrostrat.utils import get_logger
+
 from .compose import add_docker_compose_commands
 from .core import Application
 from .utils import CommandBase, ControlCommandGroup  # noqa

app-frame/macrostrat/app_frame/core.py

@@ -83,7 +83,7 @@ def load_dotenv(self):
 
     def setup_environment(self, env: EnvironmentDependency):
         environ["DOCKER_SCAN_SUGGEST"] = "false"
-        #environ["DOCKER_BUILDKIT"] = "1"
+        # environ["DOCKER_BUILDKIT"] = "1"
 
         # Set up environment for docker-compose
         # We may need to move this to a context where it is only

app-frame/macrostrat/app_frame/subsystems/__init__.py

@@ -97,7 +97,9 @@ def order_plugins(self, store=None):
     def __load_plugin(self, plugin, app: ApplicationBase):
         if isinstance(plugin, Subsystem):
             if plugin.app is not app:
-                raise SubsystemError(f"Subsystem {plugin.name} was initialized against the wrong app")
+                raise SubsystemError(
+                    f"Subsystem {plugin.name} was initialized against the wrong app"
+                )
             return plugin
         if issubclass(plugin, Subsystem):
             return plugin(app)

app-frame/macrostrat/app_frame/utils/__init__.py

@@ -39,7 +39,6 @@ def command(self, name=None, **kwargs):
         _name = self._name_modifier(name, kwargs.pop("aliases", None))
         return super().command(_name, **kwargs)
 
-
     def _name_modifier(self, name, aliases: list[str] = None):
         """Return a function that modifies the name for a command to include aliases
         https://github.com/fastapi/typer/issues/132
@@ -69,7 +68,6 @@ def get_command(self, ctx, cmd_name):
         cmd_name = self._group_cmd_name(cmd_name)
         return super().get_command(ctx, cmd_name)
 
-
     def _group_cmd_name(self, default_name):
         for cmd in self.commands.values():
             name = cmd.name

auth-system/macrostrat/auth_system/core/__init__.py

@@ -1,7 +1,7 @@
 from .main import (
+    get_groups,
     get_groups_from_header_token,
+    get_user_id,
     get_user_token_from_cookie,
-    get_groups,
     has_access,
-    get_user_id
 )

auth-system/macrostrat/auth_system/core/database.py

@@ -2,11 +2,11 @@
 from contextvars import ContextVar
 from typing import Optional
 
-from sqlalchemy import Engine
-from sqlalchemy import select, update
-from sqlalchemy.orm import sessionmaker, declarative_base, Session
+from sqlalchemy import Engine, select, update
+from sqlalchemy.orm import Session, declarative_base, sessionmaker
 
 from macrostrat.database import Database
+
 from .schema import Token
 
 

auth-system/macrostrat/auth_system/core/schema.py

@@ -1,8 +1,24 @@
+import datetime
 import enum
 from typing import List
-import datetime
-from sqlalchemy import ForeignKey, func, DateTime, Enum, PrimaryKeyConstraint, UniqueConstraint
-from sqlalchemy.dialects.postgresql import VARCHAR, TEXT, INTEGER, ARRAY, BOOLEAN, JSON, JSONB
+
+from sqlalchemy import (
+    DateTime,
+    Enum,
+    ForeignKey,
+    PrimaryKeyConstraint,
+    UniqueConstraint,
+    func,
+)
+from sqlalchemy.dialects.postgresql import (
+    ARRAY,
+    BOOLEAN,
+    INTEGER,
+    JSON,
+    JSONB,
+    TEXT,
+    VARCHAR,
+)
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
 
 
@@ -12,36 +28,34 @@ class Base(DeclarativeBase):
 
 class GroupMembers(Base):
     __tablename__ = "group_members"
-    __table_args__ = {
-        'schema': 'macrostrat_auth'
-    }
+    __table_args__ = {"schema": "macrostrat_auth"}
     id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
     group_id: Mapped[int] = mapped_column(ForeignKey("macrostrat_auth.group.id"))
     user_id: Mapped[int] = mapped_column(ForeignKey("macrostrat_auth.user.id"))
 
 
 class Group(Base):
     __tablename__ = "group"
-    __table_args__ = {
-        'schema': 'macrostrat_auth'
-    }
+    __table_args__ = {"schema": "macrostrat_auth"}
     id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
     name: Mapped[str] = mapped_column(VARCHAR(255))
-    users: Mapped[List["User"]] = relationship(secondary="macrostrat_auth.group_members", lazy="joined",
-                                               back_populates="groups")
+    users: Mapped[List["User"]] = relationship(
+        secondary="macrostrat_auth.group_members",
+        lazy="joined",
+        back_populates="groups",
+    )
 
 
 class User(Base):
     __tablename__ = "user"
-    __table_args__ = {
-        'schema': 'macrostrat_auth'
-    }
+    __table_args__ = {"schema": "macrostrat_auth"}
     id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
     sub: Mapped[str] = mapped_column(VARCHAR(255))
     name: Mapped[str] = mapped_column(VARCHAR(255))
     email: Mapped[str] = mapped_column(VARCHAR(255))
-    groups: Mapped[List[Group]] = relationship(secondary="macrostrat_auth.group_members", lazy="joined",
-                                               back_populates="users")
+    groups: Mapped[List[Group]] = relationship(
+        secondary="macrostrat_auth.group_members", lazy="joined", back_populates="users"
+    )
     created_on: Mapped[datetime.datetime] = mapped_column(
         DateTime(timezone=True), server_default=func.now()
     )
@@ -52,20 +66,14 @@ class User(Base):
 
 class Token(Base):
     __tablename__ = "token"
-    __table_args__ = {
-        'schema': 'macrostrat_auth'
-    }
+    __table_args__ = {"schema": "macrostrat_auth"}
     id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
     token: Mapped[str] = mapped_column(VARCHAR(255), unique=True)
     group: Mapped[Group] = mapped_column(ForeignKey("macrostrat_auth.group.id"))
     used_on: Mapped[datetime.datetime] = mapped_column(
         DateTime(timezone=True), nullable=True
     )
-    expires_on: Mapped[datetime.datetime] = mapped_column(
-        DateTime(timezone=True)
-    )
+    expires_on: Mapped[datetime.datetime] = mapped_column(DateTime(timezone=True))
     created_on: Mapped[datetime.datetime] = mapped_column(
         DateTime(timezone=True), server_default=func.now()
     )
-
-

auth-system/macrostrat/auth_system/legacy/api.py

@@ -1,9 +1,10 @@
-from starlette.routing import Route, Router
+from starlette.authentication import AuthenticationError, requires
 from starlette.responses import JSONResponse
-from starlette.authentication import requires, AuthenticationError
+from starlette.routing import Route, Router
+from webargs_starlette import use_annotations
+
 from macrostrat.utils import get_logger
 
-from webargs_starlette import use_annotations
 from .context import get_backend, get_identity_provider
 
 log = get_logger(__name__)

auth-system/macrostrat/auth_system/legacy/backend.py

@@ -4,15 +4,15 @@
 """
 import time
 import warnings
-from typing import Tuple, Any
+from typing import Any, Tuple
 
 import jwt
 from starlette.authentication import (
-    BaseUser,
-    SimpleUser,
-    AuthenticationBackend,
     AuthCredentials,
+    AuthenticationBackend,
     AuthenticationError,
+    BaseUser,
+    SimpleUser,
     UnauthenticatedUser,
 )
 from starlette.requests import Request

auth-system/macrostrat/auth_system/legacy/identity.py

@@ -1,4 +1,4 @@
-from werkzeug.security import generate_password_hash, check_password_hash
+from werkzeug.security import check_password_hash, generate_password_hash
 
 
 class IdentityProvider:

auth-system/macrostrat/auth_system/legacy/test_auth.py

@@ -12,7 +12,7 @@
 
 from . import get_backend
 from .api import AuthAPI
-from .context import create_backend, set_identity_provider, get_identity_provider
+from .context import create_backend, get_identity_provider, set_identity_provider
 from .identity import BaseUser, IdentityProvider
 
 

database/macrostrat/database/__init__.py

@@ -5,13 +5,14 @@
 
 from psycopg2.errors import InvalidSavepointSpecification
 from psycopg2.sql import Identifier
-from sqlalchemy import URL, MetaData, create_engine, inspect, Engine
+from sqlalchemy import URL, Engine, MetaData, create_engine, inspect
 from sqlalchemy.exc import IntegrityError, InternalError
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.orm import Session, scoped_session, sessionmaker
 from sqlalchemy.sql.expression import Insert
 
 from macrostrat.utils import get_logger
+
 from .mapper import DatabaseMapper
 from .postgresql import on_conflict, prefix_inserts  # noqa
 from .utils import (  # noqa
@@ -184,6 +185,29 @@ def inspector(self):
             self.__inspector__ = inspect(self.engine)
         return self.__inspector__
 
+    def refresh_schema(self, *, automap=None):
+        """
+        Refresh the current database connection
+
+        - closes the session and flushes
+        - removes the inspector
+
+        If automap is True, will automap the database after refreshing.
+        If automap is False, will not automap the database after refreshing.
+        If automap is None, it will re-map the database if it was previously mapped.
+        """
+        # Close the session
+        self.session.flush()
+        self.session.close()
+        # Remove the inspector
+        self.__inspector__ = None
+
+        if automap is None:
+            automap = self.mapper is not None
+
+        if automap:
+            self.automap()
+
     def entity_names(self, **kwargs):
         """
         Returns an iterator of names of *schema objects*

database/macrostrat/database/transfer/dump_database.py

@@ -7,6 +7,7 @@
 from sqlalchemy.engine import Engine
 
 from macrostrat.utils import get_logger
+
 from .stream_utils import print_stdout, print_stream_progress
 from .utils import _create_command
 

database/macrostrat/database/transfer/move_tables.py

@@ -3,6 +3,7 @@
 from sqlalchemy.engine import Engine
 
 from macrostrat.utils import get_logger
+
 from .dump_database import pg_dump
 from .restore_database import pg_restore
 from .stream_utils import print_stdout, print_stream_progress

database/macrostrat/database/transfer/restore_database.py

@@ -7,6 +7,7 @@
 from sqlalchemy.engine import Engine
 
 from macrostrat.utils import get_logger
+
 from .stream_utils import print_stdout, print_stream_progress
 from .utils import _create_command, _create_database_if_not_exists
 

database/macrostrat/database/transfer/stream_utils.py

@@ -5,6 +5,7 @@
 from aiofiles.threadpool import AsyncBufferedIOBase
 
 from macrostrat.utils import get_logger
+
 from .utils import console
 
 log = get_logger(__name__)

database/macrostrat/database/transfer/utils.py

@@ -5,7 +5,7 @@
 from sqlalchemy.engine.url import URL
 from sqlalchemy_utils import create_database, database_exists, drop_database
 
-from macrostrat.utils import get_logger, ApplicationError
+from macrostrat.utils import ApplicationError, get_logger
 
 console = Console()