def _get_database_name_for_state(state_code: StateCode,
instance: DirectIngestInstance) -> str:
"""Returns the database name for the given state and instance"""
return SQLAlchemyDatabaseKey.for_state_code(
state_code,
instance.database_version(SystemLevel.STATE,
state_code=state_code),
).db_name
def _main_database_key(cls) -> "SQLAlchemyDatabaseKey":
if cls.schema_type() == SchemaType.STATE:
state_code = StateCode(cls.region_code().upper())
return SQLAlchemyDatabaseKey.for_state_code(
state_code,
cls._main_ingest_instance().database_version(
SystemLevel.STATE, state_code=state_code),
)
return SQLAlchemyDatabaseKey.for_schema(cls.schema_type())
def test_for_state_code(self) -> None:
primary = SQLAlchemyDatabaseKey.for_state_code(
StateCode.US_MN, db_version=SQLAlchemyStateDatabaseVersion.PRIMARY)
secondary = SQLAlchemyDatabaseKey.for_state_code(
StateCode.US_MN,
db_version=SQLAlchemyStateDatabaseVersion.SECONDARY)
self.assertEqual(
SQLAlchemyDatabaseKey(schema_type=SchemaType.STATE,
db_name="us_mn_primary"),
primary,
)
self.assertEqual(
SQLAlchemyDatabaseKey(schema_type=SchemaType.STATE,
db_name="us_mn_secondary"),
secondary,
)
def ingest_database_key(self) -> SQLAlchemyDatabaseKey:
schema_type = self.system_level.schema_type()
if schema_type == SchemaType.STATE:
state_code = StateCode(self.region_code().upper())
return SQLAlchemyDatabaseKey.for_state_code(
state_code,
self.ingest_instance.database_version(self.system_level,
state_code=state_code),
)
return SQLAlchemyDatabaseKey.for_schema(schema_type)
def test_state_legacy_db(self) -> None:
db_key_1 = SQLAlchemyDatabaseKey(schema_type=SchemaType.STATE)
db_key_1_dup = SQLAlchemyDatabaseKey.canonical_for_schema(
schema_type=SchemaType.STATE)
self.assertEqual(db_key_1, db_key_1_dup)
# TODO(#7984): Once we have cut over all traffic to non-legacy state DBs and
# removed the LEGACY database version, remove this part of the test.
db_key_legacy = SQLAlchemyDatabaseKey.for_state_code(
StateCode.US_AK, SQLAlchemyStateDatabaseVersion.LEGACY)
self.assertEqual(db_key_1, db_key_legacy)
def test_key_attributes_state(self) -> None:
key = SQLAlchemyDatabaseKey.for_state_code(
StateCode.US_MN, db_version=SQLAlchemyStateDatabaseVersion.PRIMARY)
self.assertEqual(key.declarative_meta, StateBase)
self.assertTrue(os.path.exists(key.alembic_file))
self.assertTrue(
key.alembic_file.endswith("migrations/state_alembic.ini"))
self.assertTrue(os.path.exists(key.migrations_location))
self.assertTrue(key.migrations_location.endswith("/migrations/state"))
self.assertEqual(key.isolation_level, "SERIALIZABLE")
def _import_database_from_gcs() -> Tuple[str, HTTPStatus]:
try:
state_code = StateCode(request.json["stateCode"])
db_version = SQLAlchemyStateDatabaseVersion(
request.json["importToDatabaseVersion"].lower())
ingest_instance = DirectIngestInstance.for_state_database_version(
database_version=db_version, state_code=state_code)
exported_db_version = SQLAlchemyStateDatabaseVersion(
request.json["exportedDatabaseVersion"].lower())
except ValueError:
return "invalid parameters provided", HTTPStatus.BAD_REQUEST
if db_version == SQLAlchemyStateDatabaseVersion.LEGACY:
return "ingestInstance cannot be LEGACY", HTTPStatus.BAD_REQUEST
lock_manager = DirectIngestRegionLockManager.for_state_ingest(
state_code, ingest_instance=ingest_instance)
if not lock_manager.can_proceed():
return (
"other locks blocking ingest have been acquired; aborting operation",
HTTPStatus.CONFLICT,
)
db_key = SQLAlchemyDatabaseKey.for_state_code(state_code, db_version)
cloud_sql_client = CloudSQLClientImpl(project_id=project_id)
operation_id = cloud_sql_client.import_gcs_sql(
db_key,
GcsfsFilePath.from_absolute_path(
f"{STATE_INGEST_EXPORT_URI}/{exported_db_version.value}/{state_code.value}"
),
)
if operation_id is None:
return (
"Cloud SQL import operation was not started successfully.",
HTTPStatus.INTERNAL_SERVER_ERROR,
)
operation_succeeded = cloud_sql_client.wait_until_operation_completed(
operation_id, seconds_to_wait=GCS_IMPORT_EXPORT_TIMEOUT_SEC)
if not operation_succeeded:
return (
"Cloud SQL import did not complete within 60 seconds",
HTTPStatus.INTERNAL_SERVER_ERROR,
)
return operation_id, HTTPStatus.OK
def test_build_view_state_legacy(self) -> None:
table = one(t for t in StateBase.metadata.sorted_tables
if t.name == "state_person")
view_builder = FederatedCloudSQLTableBigQueryViewBuilder(
connection_region="us-east2",
table=table,
view_id=table.name,
cloud_sql_query="SELECT * FROM state_person;",
database_key=SQLAlchemyDatabaseKey.for_state_code(
StateCode.US_XX,
db_version=SQLAlchemyStateDatabaseVersion.LEGACY),
materialized_address_override=BigQueryAddress(
dataset_id="materialized_dataset",
table_id="materialized_table"),
)
expected_description = """View providing a connection to the [state_person]
table in the [postgres] database in the [STATE] schema. This view is
managed outside of regular view update operations and the results can be found in the
schema-specific datasets (`state`, `jails`, `justice_counts`, etc)."""
expected_view_query = f"""/*{expected_description}*/
SELECT
*
FROM EXTERNAL_QUERY(
"test-project.us-east2.state_cloudsql",
"SELECT * FROM state_person;"
)"""
# Build without dataset overrides
view = view_builder.build()
self.assertIsInstance(view, FederatedCloudSQLTableBigQueryView)
self.assertEqual(expected_view_query, view.view_query)
self.assertEqual(expected_description, view.description)
self.assertEqual(
BigQueryAddress(
dataset_id="state_cloudsql_connection",
table_id="state_person",
),
view.address,
)
self.assertEqual(
BigQueryAddress(dataset_id="materialized_dataset",
table_id="materialized_table"),
view.materialized_address,
)
def main(state_code: StateCode, ingest_instance: DirectIngestInstance) -> None:
"""Executes the main flow of the script."""
print(
f"RUN THE FOLLOWING COMMANDS IN ORDER TO DELETE ALL DATA FOR REGION [{state_code.value}]"
)
print(
"********************************************************************************"
)
db_version = ingest_instance.database_version(SystemLevel.STATE,
state_code)
db_key = SQLAlchemyDatabaseKey.for_state_code(state_code=state_code,
db_version=db_version)
# Connect to correct database for instance first
print(f"\\c {db_key.db_name}")
# Then run deletion commands
for cmd in generate_region_deletion_commands(state_code, db_version):
print(cmd)
print(
"********************************************************************************"
)
print("HOW TO PERFORM DELETION:")
print(
"1) Log into prod data client (`gcloud compute ssh prod-data-client --project=recidiviz-123`)"
)
print("\n> For production deletion:")
print(
"2) Go to secret manager to get login credentials stored in `state_db_user` and `state_db_password` secrets:"
"\n\thttps://console.cloud.google.com/security/secret-manager?organizationId=448885369991&"
"project=recidiviz-123")
print("3) Log into postgres database (`prod-state-psql`)")
print("\n> For staging deletion:")
print(
"2) Go to secret manager to get login credentials stored in `state_db_user` and `state_db_password` secrets:"
"\n\thttps://console.cloud.google.com/security/secret-manager?organizationId=448885369991&"
"project=recidiviz-staging")
print("3) Log into postgres database (`dev-state-psql`)")
print("\n> For all:")
print(
"4) Paste full list of commands listed above in postgres command line and run. Some commands may take a "
"while to run.")
def database_key_for_segment(
self, state_code: StateCode) -> SQLAlchemyDatabaseKey:
"""Returns a key for the database associated with a particular state segment.
Throws for unsegmented schemas.
"""
if not self.is_state_segmented_refresh_schema():
raise ValueError(
f"Only expect state-segmented schemas, found [{self.schema_type}]"
)
if self.schema_type == SchemaType.STATE:
if not self.direct_ingest_instance:
raise ValueError(
"Expected DirectIngestInstance to be non-None for STATE schema."
)
return SQLAlchemyDatabaseKey.for_state_code(
state_code=state_code,
db_version=self.direct_ingest_instance.database_version(
SystemLevel.STATE, state_code=state_code),
)
return SQLAlchemyDatabaseKey.for_schema(self.schema_type)
def main(
state_code: StateCode,
database_version: SQLAlchemyStateDatabaseVersion,
ssl_cert_path: str,
purge_schema: bool,
) -> None:
"""
Invokes the main code path for running a downgrade.
This checks for user validations that the database and branches are correct and then runs the downgrade
migration.
"""
# TODO(#7984): Once we have cut all traffic over to single-database traffic,
# delete this branch.
if database_version == SQLAlchemyStateDatabaseVersion.LEGACY:
logging.error(
"Should not invoke purge_state_db script with legacy database version."
)
sys.exit(1)
is_prod = metadata.project_id() == GCP_PROJECT_PRODUCTION
if is_prod:
logging.info("RUNNING AGAINST PRODUCTION\n")
purge_str = ("PURGE DATABASE STATE IN STAGING" if metadata.project_id()
== GCP_PROJECT_STAGING else "PURGE DATABASE STATE IN PROD")
prompt_for_confirmation(
f"This script will PURGE all data for for [{state_code.value}] in DB [{database_version.value}].",
purge_str,
)
if purge_schema:
purge_schema_str = ("RUN DOWNGRADE MIGRATIONS IN STAGING"
if metadata.project_id() == GCP_PROJECT_STAGING
else "RUN DOWNGRADE MIGRATIONS IN PROD")
prompt_for_confirmation(
f"This script will run all DOWNGRADE migrations for "
f"[{state_code.value}] in DB [{database_version.value}].",
purge_schema_str,
)
db_key = SQLAlchemyDatabaseKey.for_state_code(state_code, database_version)
with SessionFactory.for_prod_data_client(db_key, ssl_cert_path) as session:
truncate_commands = [
"TRUNCATE TABLE state_person CASCADE;",
"TRUNCATE TABLE state_agent CASCADE;",
]
for command in truncate_commands:
logging.info('Running query ["%s"]. This may take a while...',
command)
session.execute(command)
logging.info("Done running truncate commands.")
if purge_schema:
with SessionFactory.for_prod_data_client(
db_key, ssl_cert_path) as purge_session:
overriden_env_vars = None
try:
logging.info("Purging schema...")
overriden_env_vars = SQLAlchemyEngineManager.update_sqlalchemy_env_vars(
database_key=db_key,
ssl_cert_path=ssl_cert_path,
migration_user=True,
)
config = alembic.config.Config(db_key.alembic_file)
alembic.command.downgrade(config, "base")
# We need to manually delete alembic_version because it's leftover after
# the downgrade migrations
purge_session.execute("DROP TABLE alembic_version;")
finally:
if overriden_env_vars:
local_postgres_helpers.restore_local_env_vars(
overriden_env_vars)
logging.info("Purge complete.")