@@ -3080,3 +3080,161 @@ def test_partition_calling_stages_individually(
30803080 """ )
30813081 partitions = cur .fetchall ()
30823082 assert len (partitions ) > 0 , "Table should have partitions"
3083+
3084+
3085+ def test_pk_column_property_accessed_before_setup (
3086+ connection : _psycopg .Connection ,
3087+ ) -> None :
3088+ """
3089+ Test pk_column property when accessed before setup_repacking.
3090+ """
3091+ with _cur .get_cursor (connection , logged = True ) as cur :
3092+ # Create a simple table
3093+ cur .execute (
3094+ """
3095+ CREATE TABLE test_pk_early_access (
3096+ id bigint PRIMARY KEY,
3097+ data text
3098+ );
3099+ """
3100+ )
3101+ cur .execute ("INSERT INTO test_pk_early_access (id, data) VALUES (1, 'test');" )
3102+ connection .commit ()
3103+
3104+ try :
3105+ # Create Psycopack instance
3106+ repack = Psycopack (
3107+ table = "test_pk_early_access" ,
3108+ batch_size = 1 ,
3109+ conn = connection ,
3110+ cur = cur ,
3111+ )
3112+ # Access pk_column BEFORE calling setup_repacking
3113+ pk_col = repack .pk_column
3114+ assert pk_col == "id"
3115+ # Verify _pk_column was cached
3116+ assert repack .pk_column == "id"
3117+ finally :
3118+ connection .rollback ()
3119+ cur .execute ("DROP TABLE IF EXISTS test_pk_early_access CASCADE;" )
3120+ connection .commit ()
3121+
3122+
3123+ def test_pk_column_property_for_partitioned_table (
3124+ connection : _psycopg .Connection ,
3125+ ) -> None :
3126+ """
3127+ Test pk_column property when _pk_column is None for partitioned tables.
3128+
3129+ This tests the code path where pk_column is accessed before setup_repacking
3130+ and handles composite primary keys (original PK + partition column).
3131+ """
3132+ with _cur .get_cursor (connection , logged = True ) as cur :
3133+ # Create a simple table with date column for partitioning
3134+ cur .execute (
3135+ """
3136+ CREATE TABLE test_pk_column (
3137+ id bigint PRIMARY KEY,
3138+ datetime_field date NOT NULL,
3139+ data text
3140+ );
3141+ """
3142+ )
3143+ cur .execute (
3144+ "INSERT INTO test_pk_column (id, datetime_field, data) "
3145+ "VALUES (1, '2025-01-01', 'test');"
3146+ )
3147+ connection .commit ()
3148+ try :
3149+ # Create Psycopack instance with partitioning
3150+ repack = Psycopack (
3151+ table = "test_pk_column" ,
3152+ batch_size = 1 ,
3153+ conn = connection ,
3154+ cur = cur ,
3155+ partition_config = _partition .PartitionConfig (
3156+ column = "datetime_field" ,
3157+ num_of_extra_partitions_ahead = 1 ,
3158+ strategy = _partition .DateRangeStrategy (
3159+ partition_by = _partition .PartitionInterval .DAY
3160+ ),
3161+ ),
3162+ sync_strategy = SyncStrategy .CHANGE_LOG ,
3163+ )
3164+
3165+ # Validate and setup to create the partitioned table
3166+ repack .pre_validate ()
3167+ repack .setup_repacking ()
3168+
3169+ # Now test pk_column property - for partitioned tables,
3170+ # the PK becomes composite (original PK + partition column)
3171+ # but pk_column should return the first column (original PK)
3172+ pk_col = repack .pk_column
3173+ assert pk_col == "id"
3174+
3175+ # Verify the copy table was created as partitioned
3176+ cur .execute (
3177+ f"""
3178+ SELECT relkind FROM pg_class
3179+ WHERE relname = '{ repack .copy_table }
3180+ AND relnamespace = 'public'::regnamespace;
3181+ """
3182+ )
3183+ result = cur .fetchone ()
3184+ assert result is not None
3185+ assert result [0 ] == "p" # 'p' means partitioned table
3186+
3187+ finally :
3188+ connection .rollback ()
3189+ cur .execute ("DROP TABLE IF EXISTS test_pk_column CASCADE;" )
3190+ connection .commit ()
3191+
3192+
3193+ def test_full_method_resume_from_swap_stage (
3194+ connection : _psycopg .Connection ,
3195+ ) -> None :
3196+ """
3197+ Test the full() method when resuming from SWAP stage.
3198+ """
3199+ with _cur .get_cursor (connection , logged = True ) as cur :
3200+ # Create a test table
3201+ table_name = "test_resume_swap"
3202+ cur .execute (f"CREATE TABLE { table_name } )
3203+ cur .execute (f"INSERT INTO { table_name } )
3204+ connection .commit ()
3205+
3206+ try :
3207+ repack = Psycopack (
3208+ table = table_name ,
3209+ batch_size = 100 ,
3210+ conn = connection ,
3211+ cur = cur ,
3212+ )
3213+
3214+ # Run through to POST_SYNC_UPDATE stage (but not SWAP)
3215+ repack .pre_validate ()
3216+ repack .setup_repacking ()
3217+ repack .backfill ()
3218+ repack .sync_schemas ()
3219+ repack .post_sync_update ()
3220+
3221+ # Verify we're at SWAP stage (next to be completed)
3222+ current_stage = repack .tracker .get_current_stage ()
3223+ assert current_stage == _tracker .Stage .SWAP
3224+
3225+ # Now call full() - it should run swap and clean_up
3226+ repack .full ()
3227+
3228+ # Verify we've completed everything - after clean_up, the tracker row
3229+ # is deleted, so we're back at PRE_VALIDATION stage
3230+ current_stage = repack .tracker .get_current_stage ()
3231+ assert current_stage == _tracker .Stage .PRE_VALIDATION
3232+
3233+ finally :
3234+ connection .rollback ()
3235+ # Clean up any remaining tables
3236+ cur .execute (f"DROP TABLE IF EXISTS { table_name } )
3237+ cur .execute (
3238+ f"DROP TABLE IF EXISTS { table_name }
3239+ )
3240+ connection .commit ()
0 commit comments