@@ -565,3 +565,132 @@ def get_rest_table_metadata_location(encoded_namespace, encoded_table_name, pg_c
565565 status , json_str , headers = res [0 ]
566566 metadata = json .loads (json_str )
567567 return metadata ["metadata" ]["location" ]
568+
569+
570+ def test_multi_table_different_rest_catalog_hosts_in_single_transaction (
571+ installcheck ,
572+ superuser_conn ,
573+ pg_conn ,
574+ s3 ,
575+ extension ,
576+ with_default_location ,
577+ polaris_session ,
578+ create_http_helper_functions ,
579+ ):
580+ """
581+ Tables from two REST catalog servers with different hosts are modified
582+ in the same transaction. PostAllRestCatalogRequests groups modifications
583+ by conn->host, so using 'localhost' vs '127.0.0.1' (same Polaris, different
584+ host strings) produces two separate batch commit requests.
585+ """
586+ if installcheck :
587+ return
588+
589+ server_a = "multi_host_catalog_a"
590+ server_b = "multi_host_catalog_b"
591+ table_a = "multi_host_tx_a"
592+ table_b = "multi_host_tx_b"
593+ ns = TABLE_NAMESPACE + "_multi_host"
594+
595+ _create_polaris_catalog_server (superuser_conn , server_a , "localhost" )
596+ _create_polaris_catalog_server (superuser_conn , server_b , "127.0.0.1" )
597+ superuser_conn .commit ()
598+
599+ run_command (f"CREATE SCHEMA IF NOT EXISTS { ns } , pg_conn )
600+ pg_conn .commit ()
601+
602+ run_command (
603+ f"CREATE TABLE { ns } { table_a } { server_a } ,
604+ pg_conn ,
605+ )
606+ pg_conn .commit ()
607+
608+ run_command (
609+ f"CREATE TABLE { ns } { table_b } { server_b } ,
610+ pg_conn ,
611+ )
612+ pg_conn .commit ()
613+
614+ # Insert into both tables (different hosts) within a single transaction
615+ run_command (
616+ f"INSERT INTO { ns } { table_a } ,
617+ pg_conn ,
618+ )
619+ run_command (
620+ f"INSERT INTO { ns } { table_b } ,
621+ pg_conn ,
622+ )
623+ pg_conn .commit ()
624+
625+ results_a = run_query (f"SELECT count(*) FROM { ns } { table_a } , pg_conn )
626+ assert results_a [0 ][0 ] == 50
627+
628+ results_b = run_query (f"SELECT count(*) FROM { ns } { table_b } , pg_conn )
629+ assert results_b [0 ][0 ] == 30
630+
631+ # Mixed DML across different hosts in a single transaction
632+ run_command (
633+ f"INSERT INTO { ns } { table_a } ,
634+ pg_conn ,
635+ )
636+ run_command (
637+ f"DELETE FROM { ns } { table_b } ,
638+ pg_conn ,
639+ )
640+ pg_conn .commit ()
641+
642+ results_a = run_query (f"SELECT count(*) FROM { ns } { table_a } , pg_conn )
643+ assert results_a [0 ][0 ] == 70
644+
645+ results_b = run_query (f"SELECT count(*) FROM { ns } { table_b } , pg_conn )
646+ assert results_b [0 ][0 ] == 20
647+
648+ # UPDATE on both hosts in a single transaction
649+ run_command (
650+ f"UPDATE { ns } { table_a } ,
651+ pg_conn ,
652+ )
653+ run_command (
654+ f"UPDATE { ns } { table_b } ,
655+ pg_conn ,
656+ )
657+ pg_conn .commit ()
658+
659+ results_a = run_query (
660+ f"SELECT count(*) FROM { ns } { table_a } , pg_conn
661+ )
662+ assert results_a [0 ][0 ] == 5
663+
664+ results_b = run_query (
665+ f"SELECT count(*) FROM { ns } { table_b } , pg_conn
666+ )
667+ assert results_b [0 ][0 ] == 10
668+
669+ # Cleanup
670+ pg_conn .rollback ()
671+ run_command (f"DROP SCHEMA { ns } , pg_conn )
672+ pg_conn .commit ()
673+ run_command (f"DROP SERVER { server_a } , superuser_conn )
674+ run_command (f"DROP SERVER { server_b } , superuser_conn )
675+ superuser_conn .commit ()
676+
677+
678+ def _create_polaris_catalog_server (conn , server_name , hostname ):
679+ """Create an iceberg_catalog server pointing to the Polaris instance via the given hostname."""
680+ creds = json .loads (Path (server_params .POLARIS_PRINCIPAL_CREDS_FILE ).read_text ())
681+ client_id = creds ["credentials" ]["clientId" ]
682+ client_secret = creds ["credentials" ]["clientSecret" ]
683+ endpoint = f"http://{ hostname } { server_params .POLARIS_PORT }
684+
685+ run_command (
686+ f"""
687+ CREATE SERVER { server_name }
688+ FOREIGN DATA WRAPPER iceberg_catalog
689+ OPTIONS (
690+ rest_endpoint '{ endpoint }
691+ client_id '{ client_id }
692+ client_secret '{ client_secret }
693+ )
694+ """ ,
695+ conn ,
696+ )
0 commit comments