Skip to content

Commit

Permalink
addressing #10
Browse files Browse the repository at this point in the history
  • Loading branch information
bh2smith committed Apr 22, 2024
1 parent fef6536 commit c73a4f8
Showing 1 changed file with 40 additions and 14 deletions.
54 changes: 40 additions & 14 deletions src/near_lake_framework/__init__.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
from __future__ import annotations
import asyncio
import itertools
from typing import Optional, Type
from enum import Enum
from typing import Optional

from aiobotocore.session import get_session # type: ignore

Expand All @@ -11,6 +13,27 @@
from dataclasses import dataclass


class Network(Enum):
MAINNET = "mainnet"
TESTNET = "testnet"

@staticmethod
def from_string(value: str):
try:
return Network(value.lower())
except ValueError as err:
valid_values = [v.value for v in Network]
raise ValueError(
f"Unknown network: {value}. Valid values are: {valid_values}"
) from err

def __str__(self) -> str:
return self.value

def __repr__(self) -> str:
return self.value


@dataclass
class LakeConfig:
s3_bucket_name: str
Expand All @@ -20,19 +43,22 @@ class LakeConfig:
start_block_height: near_primitives.BlockHeight
blocks_preload_pool_size: int = 200

@classmethod
def mainnet(cls) -> Type["LakeConfig"]:
cls.s3_bucket_name = "near-lake-data-mainnet"
cls.s3_region_name = "eu-central-1"

return cls

@classmethod
def testnet(cls) -> Type["LakeConfig"]:
cls.s3_bucket_name = "near-lake-data-testnet"
cls.s3_region_name = "eu-central-1"
def __init__(
self,
network: Network,
aws_access_key_id: str,
aws_secret_key: str,
start_block_height: near_primitives.BlockHeight,
preload_pool_size: int = 200,
):
# These are entirely determined by Network.
self.s3_bucket_name = f"near-lake-data-{network.value}"
self.s3_region_name = "eu-central-1"

return cls
self.aws_access_key_id = aws_access_key_id
self.aws_secret_key = aws_secret_key
self.start_block_height = start_block_height
self.preload_pool_size = preload_pool_size


async def start(config: LakeConfig, streamer_messages_queue: asyncio.Queue):
Expand Down Expand Up @@ -118,4 +144,4 @@ def streamer(config: LakeConfig):
maxsize=config.blocks_preload_pool_size
)
stream_handle = asyncio.create_task(start(config, streamer_messages_queue))
return (stream_handle, streamer_messages_queue)
return stream_handle, streamer_messages_queue

0 comments on commit c73a4f8

Please sign in to comment.