black project

Author: Nospamas

pycds/alembic/versions/33179b5ae85a_add_network_key_column_to_meta_network.py

@@ -5,7 +5,7 @@
 Create Date: 2026-01-07 20:25:34.314026
 
 Notes: This process was made more complicated by some assumptions made by the history tracking code.
-In particular it assumes that the primary and history table have the same column order with the 
+In particular it assumes that the primary and history table have the same column order with the
 exception that history tables have additional columns at the end. When adding a new column it is added
 at the end and therefore breaks the assumption. To work around this, we have to recreate the history table
 with the correct column order. This involves renaming the existing history table, creating a new one with
@@ -55,19 +55,17 @@ def upgrade():
             """
         )
     )
-    
+
     # Drop existing triggers before modifying table structure so that we don't accidentally track
     # the intermediate states
     drop_history_triggers("meta_network")
-    
+
     # Rename the existing history table to preserve existing history data
     # We'll copy data from this into the new table with the correct column order
     op.execute(
-        text(
-            f"ALTER TABLE {schema_name}.meta_network_hx RENAME TO meta_network_hx_old"
-        )
+        text(f"ALTER TABLE {schema_name}.meta_network_hx RENAME TO meta_network_hx_old")
     )
-    
+
     op.add_column(
         "meta_network",
         sa.Column(
@@ -77,7 +75,7 @@ def upgrade():
         ),
         schema=schema_name,
     )
-    
+
     op.execute(
         text(
             f"""
@@ -86,14 +84,14 @@ def upgrade():
             """
         )
     )
-    
+
     op.create_unique_constraint(
         "uq_meta_network_network_key",
         "meta_network",
         ["network_key"],
         schema=schema_name,
     )
-    
+
     # Create a trigger function to auto-populate network_key on INSERT. Must be a trigger as
     # Deault values can't call functions that access other columns.
     op.execute(
@@ -113,7 +111,7 @@ def upgrade():
             """
         )
     )
-    
+
     # Create trigger to run before INSERT
     op.execute(
         text(
@@ -125,11 +123,11 @@ def upgrade():
             """
         )
     )
-    
+
     # Recreate the history table with the new column structure
     create_history_table("meta_network", foreign_tables=None)
     grant_standard_table_privileges(f"{schema_name}.meta_network_hx")
-    
+
     # Copy existing history data from the old table to the new one
     op.execute(
         text(
@@ -148,7 +146,7 @@ def upgrade():
             """
         )
     )
-    
+
     # Reset the sequence to continue from the last ID
     op.execute(
         text(
@@ -160,10 +158,10 @@ def upgrade():
             """
         )
     )
-    
+
     # Update foreign key references in dependent tables to point to the new history table
     # meta_station_hx and meta_vars_hx have foreign keys to meta_network_hx
-    
+
     # Drop the foreign key constraints from dependent tables
     op.execute(
         text(
@@ -175,10 +173,10 @@ def upgrade():
             f"ALTER TABLE {schema_name}.meta_vars_hx DROP CONSTRAINT meta_vars_hx_meta_network_hx_id_fkey"
         )
     )
-    
+
     # Drop the old history table now that data has been copied and FKs removed
     op.execute(text(f"DROP TABLE {schema_name}.meta_network_hx_old"))
-    
+
     # Recreate the foreign key constraints pointing to the new history table
     op.execute(
         text(
@@ -200,13 +198,15 @@ def upgrade():
             """
         )
     )
-    
+
     # Recreate the history tracking triggers
     create_primary_table_triggers("meta_network")
     create_history_table_triggers("meta_network", foreign_tables=None)
-    
+
     # Create indexes on the history table
-    create_history_table_indexes("meta_network", "network_id", foreign_tables=None, extras=None)
+    create_history_table_indexes(
+        "meta_network", "network_id", foreign_tables=None, extras=None
+    )
 
 
 def downgrade():
@@ -216,10 +216,8 @@ def downgrade():
             f"DROP TRIGGER IF EXISTS set_network_key_default_trigger ON {schema_name}.meta_network"
         )
     )
-    op.execute(
-        text(f"DROP FUNCTION IF EXISTS {schema_name}.set_network_key_default()")
-    )
-    
+    op.execute(text(f"DROP FUNCTION IF EXISTS {schema_name}.set_network_key_default()"))
+
     # Drop the constraint and column from primary table
     op.drop_constraint(
         "uq_meta_network_network_key",
@@ -228,13 +226,13 @@ def downgrade():
         schema=schema_name,
     )
 
-    # When dropping we don't have the same issues with column order so we can safely just drop the 
+    # When dropping we don't have the same issues with column order so we can safely just drop the
     # column to return to the pre-migration state
     op.drop_column("meta_network", "network_key", schema=schema_name)
-    
+
     # Drop the column from history table
     op.drop_column("meta_network_hx", "network_key", schema=schema_name)
-    
+
     # Drop the key generation function
     op.execute(
         text(f"DROP FUNCTION IF EXISTS {schema_name}.gen_network_key_from_name(text)")

pycds/orm/tables/__init__.py

@@ -21,11 +21,14 @@
 else:
     # Specific version requested - import from that version module
     import importlib
-    _version_module = importlib.import_module(f"pycds.orm.tables.version_{_requested_version}")
-    
+
+    _version_module = importlib.import_module(
+        f"pycds.orm.tables.version_{_requested_version}"
+    )
+
     # Import all public members from the version module
     for _name in dir(_version_module):
-        if not _name.startswith('_'):
+        if not _name.startswith("_"):
             globals()[_name] = getattr(_version_module, _name)
-    
+
     del importlib, _version_module, _name

pycds/orm/tables/base.py

@@ -46,7 +46,6 @@
 metadata = Base.metadata
 
 
-
 # string templating functions for check functions applied against multiple columns
 def no_newline_ck_name(column):
     return f"{column}_nolinebreak"

pycds/orm/tables/version_33179b5ae85a.py

@@ -120,5 +120,5 @@ class NetworkHistory(Base):
     def __str__(self):
         return f"<CRMP NetworkHistory {self.name}>"
 
-# import other tables from base
 
+# import other tables from base

pycds/orm/versioning.py

@@ -13,34 +13,36 @@
 class TableVersionManager:
     """
     Manages versioned ORM table definitions.
-    
+
     Usage:
         # Get tables at specific revision
         tables = get_tables_at_revision("a59d64cf16ca")
         Network = tables.Network
         NetworkHistory = tables.NetworkHistory
-        
+
         # Or use current version (default)
         tables = get_tables_at_revision()  # Returns current version
         Network = tables.Network
     """
-    
+
     def __init__(self, revision: Optional[str] = None):
         self.revision = revision
         self._module = None
-    
+
     def _load_module(self):
         """Lazily load the versioned module."""
         if self._module is None:
             if self.revision is None:
                 # Import current/head version from head version module
-                self._module = importlib.import_module("pycds.orm.tables.version_33179b5ae85a")
+                self._module = importlib.import_module(
+                    "pycds.orm.tables.version_33179b5ae85a"
+                )
             else:
                 # Import specific version
                 module_name = f"pycds.orm.tables.version_{self.revision}"
                 self._module = importlib.import_module(module_name)
         return self._module
-    
+
     def __getattr__(self, name):
         """Dynamically access classes from the versioned module."""
         module = self._load_module()
@@ -56,20 +58,20 @@ def __getattr__(self, name):
 def get_tables_at_revision(revision: Optional[str] = None) -> TableVersionManager:
     """
     Get ORM table classes at a specific migration revision.
-    
+
     Args:
         revision: Alembic revision ID (e.g., "a59d64cf16ca").
                  If None, returns current version.
-    
+
     Returns:
         TableVersionManager that provides access to versioned table classes.
-    
+
     Examples:
         # Get tables at old revision (before network_key)
         tables = get_tables_at_revision("a59d64cf16ca")
         Network = tables.Network
         assert not hasattr(Network, 'key')
-        
+
         # Get current tables
         tables = get_tables_at_revision()
         Network = tables.Network
@@ -85,30 +87,33 @@ def get_tables_at_revision(revision: Optional[str] = None) -> TableVersionManage
 def set_global_table_version(revision: Optional[str] = None):
     """
     Set the global table version that will be used by default.
-    
+
     This is useful for test fixtures that need to ensure all table
     references use a specific version.
-    
+
     IMPORTANT: This clears the module cache for pycds.orm.tables and related
     modules to ensure the new version is loaded on next import.
-    
+
     Args:
         revision: Alembic revision ID, or None for current version.
     """
     global _global_version
     _global_version = revision
-    
+
     # Clear module cache for tables/views/matviews modules to force reload with new version
     # We need to clear modules that import from pycds.orm.tables (like main pycds module)
     # but NOT modules that don't depend on tables (like pycds.database, pycds.context, etc.)
     # since they may have been imported by test files and clearing them breaks references.
     modules_to_clear = [
-        key for key in sys.modules.keys()
-        if (key.startswith('pycds.orm.tables') or 
-            key.startswith('pycds.orm.views') or 
-            key.startswith('pycds.orm.native_matviews') or
-            key.startswith('pycds.orm.manual_matviews') or
-            key == 'pycds')
+        key
+        for key in sys.modules.keys()
+        if (
+            key.startswith("pycds.orm.tables")
+            or key.startswith("pycds.orm.views")
+            or key.startswith("pycds.orm.native_matviews")
+            or key.startswith("pycds.orm.manual_matviews")
+            or key == "pycds"
+        )
     ]
     for module_name in modules_to_clear:
         del sys.modules[module_name]
@@ -122,7 +127,7 @@ def get_global_table_version() -> Optional[str]:
 def get_default_tables() -> TableVersionManager:
     """
     Get tables using the global version if set, otherwise current version.
-    
+
     This respects the global version set by set_global_table_version().
     """
     return TableVersionManager(_global_version)

tests/alembic_migrations/versions/v_33179b5ae85a_add_network_key_column_to_meta_network/test_smoke.py

@@ -38,7 +38,7 @@ def test_upgrade(
     )
     col = check_if_column_exists(column_name, meta_network_table)
     assert (col["nullable"] == True) and (col["name"] == column_name)
-    
+
     # Check that column has been added to meta_network_hx
     meta_network_hx_table = inspect(alembic_engine).get_columns(
         history_table_name, schema=schema_name
@@ -61,7 +61,7 @@ def test_downgrade(alembic_engine, alembic_runner, schema_name):
     )
     col = check_if_column_exists(column_name, meta_network_table)
     assert col == null
-    
+
     # Check that column has been removed from meta_network_hx
     meta_network_hx_table = inspect(alembic_engine).get_columns(
         history_table_name, schema=schema_name

tests/alembic_migrations/versions/v_7a3b247c577b_add_varsperhistory_native_matview/test_smoke.py

@@ -22,7 +22,9 @@
 @pytest.mark.update20
 @pytest.mark.parametrize("supports_matviews", [True, False])
 def test_mock(mocker, supports_matviews):
-    mock_func = mocker.patch("pycds.database.db_supports_matviews", return_value=supports_matviews)
+    mock_func = mocker.patch(
+        "pycds.database.db_supports_matviews", return_value=supports_matviews
+    )
     # Call the mock directly since pycds module reference may be stale
     assert mock_func() is supports_matviews
 

tests/alembic_migrations/versions/v_a59d64cf16ca_add_hx_tkg_to_main_metadata_tables/test_on_real_tables/test_on_real_tables.py

@@ -15,6 +15,7 @@
 # IMPORTANT: Set table version BEFORE any pycds imports
 # This test needs the table schema at revision a59d64cf16ca (before network_key was added)
 from pycds.orm.versioning import set_global_table_version
+
 set_global_table_version("a59d64cf16ca")
 
 import logging

tests/alembic_migrations/versions/v_bdc28573df56_add_obs_raw_indexes/test_smoke.py

@@ -23,7 +23,9 @@
 @pytest.mark.update20
 @pytest.mark.parametrize("item_names", [set(), {"alpha", "beta"}])
 def test_mock(mocker, item_names):
-    mock_func = mocker.patch("pycds.database.get_schema_item_names", return_value=item_names)
+    mock_func = mocker.patch(
+        "pycds.database.get_schema_item_names", return_value=item_names
+    )
     # Call the mock directly since pycds module reference may be stale
     assert mock_func() == item_names
 

tests/climate_baseline_helpers/test_climate_baseline_helpers.py

@@ -152,7 +152,9 @@ def it_creates_precip_variable(sesh_with_climate_baseline_variables):
     def it_creates_no_more_than_one_of_each(pycds_sesh):
         sesh = pycds_sesh
         get_or_create_pcic_climate_baseline_variables(sesh)
-        get_or_create_pcic_climate_baseline_variables(sesh) # TODO: Should this be run twice?
+        get_or_create_pcic_climate_baseline_variables(
+            sesh
+        )  # TODO: Should this be run twice?
         results = sesh.query(Variable).filter(Variable.name.like("%_Climatology"))
         assert results.count() == 3