[Issue #1697] add more foreign table models and update setup_foreign_tables.py

api/src/constants/schema.py

@@ -3,3 +3,5 @@
 
 class Schemas(StrEnum):
     API = "api"
+    LEGACY = "legacy"
+    STAGING = "staging"

api/src/data_migration/copy_oracle_data.py

@@ -41,7 +41,7 @@ class SqlCommands:
                 last_upd_date,
                 creator_id,
                 created_date
-            from {}.foreign_topportunity
+            from {}.topportunity
             where is_draft = 'N'
     """
 
@@ -55,19 +55,21 @@ def copy_oracle_data(db_session: db.Session) -> None:
 
     try:
         with db_session.begin():
-            _run_copy_commands(db_session, Schemas.API)
+            _run_copy_commands(db_session, Schemas.API, Schemas.LEGACY)
     except Exception:
         logger.exception("Failed to run copy-oracle-data command")
         raise
 
     logger.info("Successfully ran copy-oracle-data")
 
 
-def _run_copy_commands(db_session: db.Session, api_schema: str) -> None:
+def _run_copy_commands(db_session: db.Session, api_schema: str, foreign_schema: str) -> None:
     logger.info("Running copy commands for TOPPORTUNITY")
 
     db_session.execute(text(SqlCommands.OPPORTUNITY_DELETE_QUERY.format(api_schema)))
-    db_session.execute(text(SqlCommands.OPPORTUNITY_INSERT_QUERY.format(api_schema, api_schema)))
+    db_session.execute(
+        text(SqlCommands.OPPORTUNITY_INSERT_QUERY.format(api_schema, foreign_schema))
+    )
     count = db_session.scalar(
         text(f"SELECT count(*) from {api_schema}.transfer_topportunity")  # nosec
     )

api/src/data_migration/setup_foreign_tables.py

@@ -1,11 +1,12 @@
 import logging
-from dataclasses import dataclass
 
+import sqlalchemy
 from pydantic import Field
-from sqlalchemy import text
 
 import src.adapters.db as db
 import src.adapters.db.flask_db as flask_db
+import src.db.foreign
+import src.db.foreign.dialect
 from src.constants.schema import Schemas
 from src.data_migration.data_migration_blueprint import data_migration_blueprint
 from src.util.env_config import PydanticBaseEnvConfig
@@ -15,38 +16,7 @@
 
 class ForeignTableConfig(PydanticBaseEnvConfig):
     is_local_foreign_table: bool = Field(False)
-    schema_name: str = Field(Schemas.API)
-
-
-@dataclass
-class Column:
-    column_name: str
-    postgres_type: str
-
-    is_nullable: bool = True
-    is_primary_key: bool = False
-
-
-OPPORTUNITY_COLUMNS: list[Column] = [
-    Column("OPPORTUNITY_ID", "numeric(20)", is_nullable=False, is_primary_key=True),
-    Column("OPPNUMBER", "character varying (40)"),
-    Column("REVISION_NUMBER", "numeric(20)"),
-    Column("OPPTITLE", "character varying (255)"),
-    Column("OWNINGAGENCY", "character varying (255)"),
-    Column("PUBLISHERUID", "character varying (255)"),
-    Column("LISTED", "CHAR(1)"),
-    Column("OPPCATEGORY", "CHAR(1)"),
-    Column("INITIAL_OPPORTUNITY_ID", "numeric(20)"),
-    Column("MODIFIED_COMMENTS", "character varying (2000)"),
-    Column("CREATED_DATE", "DATE"),
-    Column("LAST_UPD_DATE", "DATE"),
-    Column("CREATOR_ID", "character varying (50)"),
-    Column("LAST_UPD_ID", "character varying (50)"),
-    Column("FLAG_2006", "CHAR(1)"),
-    Column("CATEGORY_EXPLANATION", "character varying (255)"),
-    Column("PUBLISHER_PROFILE_ID", "numeric(20)"),
-    Column("IS_DRAFT", "character varying (1)"),
-]
+    schema_name: str = Field(Schemas.LEGACY)
 
 
 @data_migration_blueprint.cli.command(
@@ -64,66 +34,39 @@ def setup_foreign_tables(db_session: db.Session) -> None:
     logger.info("Successfully ran setup-foreign-tables")
 
 
-def build_sql(table_name: str, columns: list[Column], is_local: bool, schema_name: str) -> str:
+def build_sql(table: sqlalchemy.schema.Table, is_local: bool, schema_name: str) -> str:
     """
     Build the SQL for creating a possibly foreign data table. If running
     with is_local, it instead creates a regular table.
 
-    Assume you have a table with two columns, an "ID" primary key column, and a "description" text column,
-    you would call this as::
-
-        build_sql("EXAMPLE_TABLE", [Column("ID", "integer", is_nullable=False, is_primary_key=True), Column("DESCRIPTION", "text")], is_local)
-
-    Depending on whether the is_local bool is true or false would give two different outputs.
-
     is_local is True::
 
-        CREATE TABLE IF NOT EXISTS foreign_example_table (ID integer CONSTRAINT EXAMPLE_TABLE_pkey PRIMARY KEY NOT NULL,DESCRIPTION text)
+        CREATE TABLE IF NOT EXISTS foreign_example_table
+        (ID integer CONSTRAINT EXAMPLE_TABLE_pkey PRIMARY KEY NOT NULL,DESCRIPTION text)
 
     is_local is False::
 
-        CREATE FOREIGN TABLE IF NOT EXISTS foreign_example_table (ID integer OPTIONS (key 'true') NOT NULL,DESCRIPTION text) SERVER grants OPTIONS (schema 'EGRANTSADMIN', table 'EXAMPLE_TABLE')
+        CREATE FOREIGN TABLE IF NOT EXISTS foreign_example_table
+        (ID integer OPTIONS (key 'true') NOT NULL,DESCRIPTION text)
+        SERVER grants OPTIONS (schema 'EGRANTSADMIN', table 'EXAMPLE_TABLE')
     """
 
-    column_sql_parts = []
-    for column in columns:
-        column_sql = f"{column.column_name} {column.postgres_type}"
-
-        # Primary keys are defined as constraints in a regular table
-        # and as options in a foreign data table
-        if column.is_primary_key and is_local:
-            column_sql += f" CONSTRAINT {table_name}_pkey PRIMARY KEY"
-        elif column.is_primary_key and not is_local:
-            column_sql += " OPTIONS (key 'true')"
-
-        if not column.is_nullable:
-            column_sql += " NOT NULL"
-
-        column_sql_parts.append(column_sql)
-
-    create_table_command = "CREATE FOREIGN TABLE IF NOT EXISTS"
+    create_table = sqlalchemy.schema.CreateTable(table, if_not_exists=True)
     if is_local:
-        # Don't make a foreign table if running locally
-        create_table_command = "CREATE TABLE IF NOT EXISTS"
-
-    create_command_suffix = (
-        f" SERVER grants OPTIONS (schema 'EGRANTSADMIN', table '{table_name}')"  # noqa: B907
-    )
-    if is_local:
-        # We don't want the config at the end if we're running locally so unset it
-        create_command_suffix = ""
-
-    return f"{create_table_command} {schema_name}.foreign_{table_name.lower()} ({','.join(column_sql_parts)}){create_command_suffix}"
+        compiler = create_table.compile(
+            dialect=sqlalchemy.dialects.postgresql.dialect(),
+            schema_translate_map={Schemas.LEGACY.name: schema_name},
+        )
+    else:
+        compiler = create_table.compile(
+            dialect=src.db.foreign.dialect.ForeignTableDialect(),
+            schema_translate_map={Schemas.LEGACY.name: schema_name},
+        )
+    return str(compiler).strip()
 
 
 def _run_create_table_commands(db_session: db.Session, config: ForeignTableConfig) -> None:
-    db_session.execute(
-        text(
-            build_sql(
-                "TOPPORTUNITY",
-                OPPORTUNITY_COLUMNS,
-                config.is_local_foreign_table,
-                config.schema_name,
-            )
-        )
-    )
+    for table in src.db.foreign.metadata.tables.values():
+        sql = build_sql(table, config.is_local_foreign_table, config.schema_name)
+        logger.info("create table", extra={"table": table.name, "sql": sql})
+        db_session.execute(sqlalchemy.text(sql))

api/src/db/foreign/__init__.py

@@ -0,0 +1,9 @@
+#
+# SQLAlchemy models for foreign tables.
+#
+
+from . import forecast, foreignbase, opportunity, synopsis
+
+metadata = foreignbase.metadata
+
+__all__ = ["metadata", "forecast", "opportunity", "synopsis"]

api/src/db/foreign/dialect.py

@@ -0,0 +1,42 @@
+#
+# Support for generating SQL for "CREATE FOREIGN TABLE".
+#
+
+import re
+
+import sqlalchemy
+
+
+class ForeignTableDDLCompiler(sqlalchemy.sql.compiler.DDLCompiler):
+    """SQLAlchemy compiler for creating foreign tables."""
+
+    def create_table_constraints(self, _table, _include_foreign_key_constraints=None, **kw):
+        return ""  # Don't generate any constraints.
+
+    def visit_create_table(self, create, **kw):
+        table = create.element
+        table._prefixes = ("FOREIGN",)  # Add "FOREIGN" before "TABLE".
+        sql = super().visit_create_table(create, **kw)
+        table._prefixes = ()
+        return sql
+
+    def post_create_table(self, table):
+        # Add foreign options at the end.
+        return f" SERVER grants OPTIONS (schema 'EGRANTSADMIN', table '{table.name.upper()}')"
+
+    def visit_create_column(self, create, first_pk=False, **kw):
+        column = create.element
+        sql = super().visit_create_column(create, first_pk, **kw)
+        if sql and column.primary_key:
+            # Add "OPTIONS ..." to primary key column.
+            sql = re.sub(r"^(.*?)( NOT NULL)?$", r"\1 OPTIONS (key 'true')\2", sql)
+        return sql
+
+
+class ForeignTableDialect(sqlalchemy.dialects.postgresql.dialect):
+    """SQLAlchemy dialect for creating foreign tables.
+
+    See https://docs.sqlalchemy.org/en/20/dialects/
+    """
+
+    ddl_compiler = ForeignTableDDLCompiler