feat: add publisher program to sync cli

ali-behjati · ali-behjati · commit eb2389714a16 · 2024-09-05T15:30:05.000+02:00
diff --git a/program_admin/__init__.py b/program_admin/__init__.py
@@ -16,6 +16,15 @@
 from program_admin import instructions as pyth_program
 from program_admin.keys import load_keypair
 from program_admin.parsing import parse_account
+from program_admin.publisher_program_instructions import (
+    config_account_pubkey as publisher_program_config_account_pubkey,
+)
+from program_admin.publisher_program_instructions import (
+    create_buffer_account,
+    initialize_publisher_config,
+    initialize_publisher_program,
+    publisher_config_account_pubkey,
+)
 from program_admin.types import (
     Network,
     PythAuthorityPermissionAccount,
@@ -56,6 +65,7 @@ class ProgramAdmin:
     rpc_endpoint: str
     key_dir: Path
     program_key: PublicKey
+    publisher_program_key: Optional[PublicKey]
     authority_permission_account: Optional[PythAuthorityPermissionAccount]
     _mapping_accounts: Dict[PublicKey, PythMappingAccount]
     _product_accounts: Dict[PublicKey, PythProductAccount]
@@ -66,13 +76,17 @@ def __init__(
         network: Network,
         key_dir: str,
         program_key: str,
+        publisher_program_key: Optional[str],
         commitment: Literal["confirmed", "finalized"],
         rpc_endpoint: str = "",
     ):
         self.network = network
         self.rpc_endpoint = rpc_endpoint or RPC_ENDPOINTS[network]
         self.key_dir = Path(key_dir)
         self.program_key = PublicKey(program_key)
+        self.publisher_program_key = (
+            PublicKey(publisher_program_key) if publisher_program_key else None
+        )
         self.commitment = Commitment(commitment)
         self.authority_permission_account = None
         self._mapping_accounts: Dict[PublicKey, PythMappingAccount] = {}
@@ -100,6 +114,12 @@ async def fetch_minimum_balance(self, size: int) -> int:
         async with AsyncClient(self.rpc_endpoint) as client:
             return (await client.get_minimum_balance_for_rent_exemption(size)).value
 
+    async def account_exists(self, key: PublicKey) -> bool:
+        async with AsyncClient(self.rpc_endpoint) as client:
+            response = await client.get_account_info(key)
+            # The RPC returns null if the account does not exist
+            return bool(response.value)
+
     async def refresh_program_accounts(self):
         async with AsyncClient(self.rpc_endpoint) as client:
             logger.info("Refreshing program accounts")
@@ -301,6 +321,19 @@ async def sync(
         if product_updates:
             await self.refresh_program_accounts()
 
+        # Sync publisher program
+        (
+            publisher_program_instructions,
+            publisher_program_signers,
+        ) = await self.sync_publisher_program(ref_publishers)
+
+        if publisher_program_instructions:
+            instructions.extend(publisher_program_instructions)
+            if send_transactions:
+                await self.send_transaction(
+                    publisher_program_instructions, publisher_program_signers
+                )
+
         # Sync publishers
 
         publisher_transactions = []
@@ -658,3 +691,54 @@ async def resize_price_accounts_v2(
 
         if send_transactions:
             await self.send_transaction(instructions, signers)
+
+    async def sync_publisher_program(
+        self, ref_publishers: ReferencePublishers
+    ) -> Tuple[List[TransactionInstruction], List[Keypair]]:
+        if self.publisher_program_key is None:
+            return [], []
+
+        instructions = []
+
+        authority = load_keypair("funding", key_dir=self.key_dir)
+
+        publisher_program_config = publisher_program_config_account_pubkey(
+            self.publisher_program_key
+        )
+
+        # Initialize the publisher program config if it does not exist
+        if not (await self.account_exists(publisher_program_config)):
+            initialize_publisher_program_instruction = initialize_publisher_program(
+                self.publisher_program_key, authority.public_key
+            )
+            instructions.append(initialize_publisher_program_instruction)
+
+        # Initialize publisher config accounts for new publishers
+        for publisher in ref_publishers["keys"].values():
+            publisher_config_account = publisher_config_account_pubkey(
+                publisher, self.publisher_program_key
+            )
+
+            if not (await self.account_exists(publisher_config_account)):
+                size = 100048  # This size is for a buffer supporting 5000 price updates
+                lamports = await self.fetch_minimum_balance(size)
+                buffer_account, create_buffer_instruction = create_buffer_account(
+                    self.publisher_program_key,
+                    authority.public_key,
+                    publisher,
+                    size,
+                    lamports,
+                )
+
+                initialize_publisher_config_instruction = initialize_publisher_config(
+                    self.publisher_program_key,
+                    publisher,
+                    authority.public_key,
+                    buffer_account,
+                )
+
+                instructions.extend(
+                    [create_buffer_instruction, initialize_publisher_config_instruction]
+                )
+
+        return (instructions, [authority])
diff --git a/program_admin/cli.py b/program_admin/cli.py
@@ -56,6 +56,7 @@ def delete_price(network, rpc_endpoint, program_key, keys, commitment, product,
         rpc_endpoint=rpc_endpoint,
         key_dir=keys,
         program_key=program_key,
+        publisher_program_key=None,
         commitment=commitment,
     )
     funding_keypair = load_keypair("funding", key_dir=keys)
@@ -236,6 +237,7 @@ def delete_product(
         rpc_endpoint=rpc_endpoint,
         key_dir=keys,
         program_key=program_key,
+        publisher_program_key=None,
         commitment=commitment,
     )
     funding_keypair = load_keypair("funding", key_dir=keys)
@@ -275,6 +277,7 @@ def list_accounts(network, rpc_endpoint, program_key, keys, publishers, commitme
         rpc_endpoint=rpc_endpoint,
         key_dir=keys,
         program_key=program_key,
+        publisher_program_key=None,
         commitment=commitment,
     )
 
@@ -333,6 +336,7 @@ def restore_links(network, rpc_endpoint, program_key, keys, products, commitment
         rpc_endpoint=rpc_endpoint,
         key_dir=keys,
         program_key=program_key,
+        publisher_program_key=None,
         commitment=commitment,
     )
     reference_products = parse_products_json(Path(products))
@@ -382,6 +386,12 @@ def restore_links(network, rpc_endpoint, program_key, keys, products, commitment
 @click.option("--network", help="Solana network", envvar="NETWORK")
 @click.option("--rpc-endpoint", help="Solana RPC endpoint", envvar="RPC_ENDPOINT")
 @click.option("--program-key", help="Pyth program key", envvar="PROGRAM_KEY")
+@click.option(
+    "--publisher-program-key",
+    help="Publisher program key",
+    envvar="PUBLISHER_PROGRAM_KEY",
+    default=None,
+)
 @click.option("--keys", help="Path to keys directory", envvar="KEYS")
 @click.option("--products", help="Path to reference products file", envvar="PRODUCTS")
 @click.option(
@@ -426,6 +436,7 @@ def sync(
     network,
     rpc_endpoint,
     program_key,
+    publisher_program_key,
     keys,
     products,
     publishers,
@@ -442,6 +453,7 @@ def sync(
         rpc_endpoint=rpc_endpoint,
         key_dir=keys,
         program_key=program_key,
+        publisher_program_key=publisher_program_key,
         commitment=commitment,
     )
 
@@ -495,6 +507,7 @@ def migrate_upgrade_authority(
         rpc_endpoint=rpc_endpoint,
         key_dir=keys,
         program_key=program_key,
+        publisher_program_key=None,
         commitment=commitment,
     )
     funding_keypair = load_keypair("funding", key_dir=keys)
@@ -544,6 +557,7 @@ def resize_price_accounts_v2(
         rpc_endpoint=rpc_endpoint,
         key_dir=keys,
         program_key=program_key,
+        publisher_program_key=None,
         commitment=commitment,
     )
 
diff --git a/program_admin/publisher_program_instructions.py b/program_admin/publisher_program_instructions.py
@@ -0,0 +1,156 @@
+from typing import Tuple
+
+from construct import Bytes, Int8ul, Struct
+from solana import system_program
+from solana.publickey import PublicKey
+from solana.system_program import SYS_PROGRAM_ID, CreateAccountWithSeedParams
+from solana.transaction import AccountMeta, TransactionInstruction
+
+
+def config_account_pubkey(program_key: PublicKey) -> PublicKey:
+    [config_account, _] = PublicKey.find_program_address(
+        [b"CONFIG"],
+        program_key,
+    )
+    return config_account
+
+
+def publisher_config_account_pubkey(
+    publisher_key: PublicKey, program_key: PublicKey
+) -> PublicKey:
+    [publisher_config_account, _] = PublicKey.find_program_address(
+        [b"PUBLISHER_CONFIG", bytes(publisher_key)],
+        program_key,
+    )
+    return publisher_config_account
+
+
+def initialize_publisher_program(
+    program_key: PublicKey,
+    authority: PublicKey,
+) -> TransactionInstruction:
+    """
+    Pyth publisher program initialize instruction with the given authority
+
+    accounts:
+    - payer account (signer, writable) - we pass the authority as the payer
+    - config account (writable)
+    - system program
+    """
+
+    [config_account, bump] = PublicKey.find_program_address(
+        [b"CONFIG"],
+        program_key,
+    )
+
+    ix_data_layout = Struct(
+        "bump" / Int8ul,
+        "authority" / Bytes(32),
+    )
+
+    ix_data = ix_data_layout.build(
+        dict(
+            bump=bump,
+            authority=bytes(authority),
+        )
+    )
+
+    return TransactionInstruction(
+        data=ix_data,
+        keys=[
+            AccountMeta(pubkey=authority, is_signer=True, is_writable=True),
+            AccountMeta(pubkey=config_account, is_signer=False, is_writable=True),
+            AccountMeta(pubkey=SYS_PROGRAM_ID, is_signer=False, is_writable=False),
+        ],
+        program_id=program_key,
+    )
+
+
+def create_buffer_account(
+    program_key: PublicKey,
+    base_pubkey: PublicKey,
+    publisher_pubkey: PublicKey,
+    space: int,
+    lamports: int,
+) -> Tuple[PublicKey, TransactionInstruction]:
+
+    seed = str(publisher_pubkey)
+    new_account_pubkey = PublicKey.create_with_seed(
+        base_pubkey,
+        seed,
+        program_key,
+    )
+
+    # space = 100048 # Required space to store 5000 price updates
+    # lamport =
+
+    return (
+        new_account_pubkey,
+        system_program.create_account_with_seed(
+            CreateAccountWithSeedParams(
+                from_pubkey=base_pubkey,
+                new_account_pubkey=new_account_pubkey,
+                base_pubkey=base_pubkey,
+                seed=seed,
+                program_id=program_key,
+                lamports=lamports,
+                space=space,
+            )
+        ),
+    )
+
+
+def initialize_publisher_config(
+    program_key: PublicKey,
+    publisher_key: PublicKey,
+    authority: PublicKey,
+    buffer_account: PublicKey,
+) -> TransactionInstruction:
+    """
+    Pyth publisher program initialize publisher config instruction with the given authority
+
+    accounts:
+    - authority account (signer, writable)
+    - config account
+    - publisher config account (writable)
+    - buffer account (writable)
+    - system program
+    """
+
+    [config_account, config_bump] = PublicKey.find_program_address(
+        [b"CONFIG"],
+        program_key,
+    )
+
+    [publisher_config_account, publisher_config_bump] = PublicKey.find_program_address(
+        [b"PUBLISHER_CONFIG", bytes(publisher_key)],
+        program_key,
+    )
+
+    ix_data_layout = Struct(
+        "config_bump" / Int8ul,
+        "publisher_config_bump" / Int8ul,
+        "publisher" / Bytes(32),
+    )
+
+    ix_data = ix_data_layout.build(
+        dict(
+            config_bump=config_bump,
+            publisher_config_bump=publisher_config_bump,
+            publisher=bytes(publisher_key),
+        )
+    )
+
+    return TransactionInstruction(
+        data=ix_data,
+        keys=[
+            AccountMeta(pubkey=authority, is_signer=True, is_writable=True),
+            AccountMeta(pubkey=config_account, is_signer=False, is_writable=True),
+            AccountMeta(
+                pubkey=publisher_config_account, is_signer=False, is_writable=True
+            ),
+            AccountMeta(pubkey=buffer_account, is_signer=False, is_writable=True),
+            AccountMeta(pubkey=SYS_PROGRAM_ID, is_signer=False, is_writable=False),
+        ],
+        program_id=program_key,
+    )
diff --git a/tests/test_sync.py b/tests/test_sync.py
@@ -421,6 +421,7 @@ async def test_sync(
         network=network,
         key_dir=key_dir,
         program_key=pyth_program,
+        publisher_program_key=None,
         commitment="confirmed",
     )
 

Original file line number	Diff line number	Diff line change
`@@ -421,6 +421,7 @@ async def test_sync(`
`421`	`421`	`network=network,`
`422`	`422`	`key_dir=key_dir,`
`423`	`423`	`program_key=pyth_program,`
	`424`	`+ publisher_program_key=None,`
`424`	`425`	`commitment="confirmed",`
`425`	`426`	`)`
`426`	`427`