Mirror of https://github.com/roostorg/osprey
github.com/roostorg/osprey
refactor: lazily initialize the bigtable client instance (#23)
+26
-23
+26
-23
osprey_worker/src/osprey/worker/lib/storage/bigtable.py
+26
-23
osprey_worker/src/osprey/worker/lib/storage/bigtable.py
···
7
7
from typing import Dict
8
8
9
9
from google.cloud.bigtable import Client
10
+
from google.cloud.bigtable.instance import Instance
10
11
from google.cloud.bigtable.table import Table
11
12
from osprey.worker.lib.config import Config
12
13
from osprey.worker.lib.ddtrace_utils import pin_override
···
24
25
25
26
26
27
class BigTableClient(ABC):
28
+
_gcp_project: str
29
+
_instance_id: str
30
+
_admin_enabled: bool
31
+
32
+
def _setup_client_instance(self) -> Instance:
33
+
client = Client(project=self._gcp_project, admin=self._admin_enabled)
34
+
fix_bigtable_client_if_using_emulator(client)
35
+
36
+
return client.instance(self._instance_id)
37
+
38
+
@property
39
+
def _instance(self) -> Instance:
40
+
if not getattr(self, '__instance', None):
41
+
self.__instance = self._setup_client_instance()
42
+
return self.__instance
43
+
27
44
@abstractmethod
28
45
def init_from_config(self, config: Config) -> None: ...
29
46
···
42
59
def init_from_config(self, config: Config) -> None:
43
60
"""Initialize this bigtable client once configuration is available."""
44
61
config = CONFIG.instance()
45
-
gcp_project = config.get_str('OSPREY_GCP_PROJECT_ID', 'osprey-dev')
46
-
bigtable_instance = config.get_str('OSPREY_BIGTABLE_INSTANCE_ID', 'osprey-bigtable')
47
-
bigtable_admin_enabled = config.get_bool('OSPREY_BIGTABLE_ADMIN_ENABLED', True)
48
-
49
-
client = Client(project=gcp_project, admin=bigtable_admin_enabled)
50
-
fix_bigtable_client_if_using_emulator(client)
51
-
52
-
instance = client.instance(bigtable_instance)
53
-
self._instance = instance
62
+
self._gcp_project = config.get_str('OSPREY_GCP_PROJECT_ID', 'osprey-dev')
63
+
self._instance_id = config.get_str('OSPREY_BIGTABLE_INSTANCE_ID', 'osprey-bigtable')
64
+
self._admin_enabled = config.get_bool('OSPREY_BIGTABLE_ADMIN_ENABLED', True)
54
65
55
66
# staging/production tables are managed through Terraform
56
67
def bootstrap(self) -> None:
···
89
100
def init_from_config(self, config: Config) -> None:
90
101
"""Initialize this bigtable client once configuration is available."""
91
102
config = CONFIG.instance()
92
-
gcp_project = config.get_str('DATA_SERVICES_GCP_PROJECT_ID', 'osprey-dev')
93
-
bigtable_instance = config.get_str('DATA_SERVICES_BIGTABLE_INSTANCE_ID', 'derived-sinks-ml-instance-dev')
94
-
bigtable_admin_enabled = config.get_bool('DATA_SERVICES_BIGTABLE_ADMIN_ENABLED', True)
95
-
96
-
client = Client(project=gcp_project, admin=bigtable_admin_enabled)
97
-
instance = client.instance(bigtable_instance)
98
-
self._instance = instance
103
+
self._gcp_project = config.get_str('DATA_SERVICES_GCP_PROJECT_ID', 'osprey-dev')
104
+
self._instance_id = config.get_str('DATA_SERVICES_BIGTABLE_INSTANCE_ID', 'derived-sinks-ml-instance-dev')
105
+
self._admin_enabled = config.get_bool('DATA_SERVICES_BIGTABLE_ADMIN_ENABLED', True)
99
106
100
107
def table(self, table_name: str) -> Table:
101
108
"""
···
122
129
def init_from_config(self, config: Config) -> None:
123
130
"""Initialize this bigtable client once configuration is available."""
124
131
config = CONFIG.instance()
125
-
gcp_project = config.get_str('DATA_GCP_PROJECT_ID', 'osprey-dev')
126
-
bigtable_instance = config.get_str('DATA_STREAM_BIGTABLE_INSTANCE_ID', 'stream')
127
-
bigtable_admin_enabled = config.get_bool('DATA_STREAM_BIGTABLE_ADMIN_ENABLED', True)
128
-
129
-
client = Client(project=gcp_project, admin=bigtable_admin_enabled)
130
-
instance = client.instance(bigtable_instance)
131
-
self._instance = instance
132
+
self._gcp_project = config.get_str('DATA_GCP_PROJECT_ID', 'osprey-dev')
133
+
self._instance_id = config.get_str('DATA_STREAM_BIGTABLE_INSTANCE_ID', 'stream')
134
+
self._admin_enabled = config.get_bool('DATA_STREAM_BIGTABLE_ADMIN_ENABLED', True)
132
135
133
136
def table(self, table_name: str) -> Table:
134
137
"""