Skip to content

Commit 7107caf

Browse files
PijukatelPijukatel
committed
Draft proposal. TODO: Abstract and rpeeat in other storages as well.
1 parent 0e03e0e commit 7107caf

File tree

2 files changed

+64
-45
lines changed

2 files changed

+64
-45
lines changed

‎src/apify/storage_clients/_apify/_key_value_store_client.py‎

Lines changed: 64 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -106,54 +106,75 @@ async def open(
106106
async with AliasResolver(storage_type=KeyValueStore, alias=alias, configuration=configuration) as _alias:
107107
id = await _alias.resolve_id()
108108

109-
# There was no pre-existing alias in the mapping.
110-
# Create a new unnamed storage and store the mapping.
111-
if id is None:
112-
# Create a new storage and store the alias mapping
113-
new_storage_metadata = ApifyKeyValueStoreMetadata.model_validate(
114-
await apify_kvss_client.get_or_create(),
115-
)
116-
id = new_storage_metadata.id
117-
await _alias.store_mapping(storage_id=id)
109+
if id:
110+
# There was id, storage has to exist, fetch metadata to confirm it.
111+
apify_kvs_client = apify_client_async.key_value_store(key_value_store_id=id)
112+
raw_metadata = await apify_kvs_client.get()
113+
if raw_metadata:
114+
return cls(
115+
api_client=apify_kvs_client,
116+
api_public_base_url='',
117+
# Remove in version 4.0, https://github.com/apify/apify-sdk-python/issues/635
118+
lock=asyncio.Lock(),
119+
)
120+
121+
# There was no pre-existing alias in the mapping or the id did not point to existing storage.
122+
# Create a new unnamed storage and store the alias mapping.
123+
metadata = ApifyKeyValueStoreMetadata.model_validate(
124+
await apify_kvss_client.get_or_create(),
125+
)
126+
await _alias.store_mapping(storage_id=metadata.id)
127+
128+
# Return the client for the newly created storage directly.
129+
# It was just created, no need to refetch it.
130+
return cls(
131+
api_client=apify_client_async.key_value_store(key_value_store_id=metadata.id),
132+
api_public_base_url='', # Remove in version 4.0, https://github.com/apify/apify-sdk-python/issues/635
133+
lock=asyncio.Lock(),
134+
)
118135

119136
# If name is provided, get or create the storage by name.
120137
elif name:
121-
id = ApifyKeyValueStoreMetadata.model_validate(
122-
await apify_kvss_client.get_or_create(name=name),
123-
).id
124-
125-
# If none are provided, try to get the default storage ID from environment variables.
126-
elif id is None:
127-
id = configuration.default_key_value_store_id
128-
if not id:
129-
raise ValueError(
130-
'KeyValueStore "id", "name", or "alias" must be specified, '
131-
'or a default KeyValueStore ID must be set in the configuration.'
132-
)
138+
metadata = ApifyKeyValueStoreMetadata.model_validate(await apify_kvss_client.get_or_create(name=name))
133139

134-
# Now create the client for the determined ID
135-
apify_kvs_client = apify_client_async.key_value_store(key_value_store_id=id)
136-
137-
# Fetch its metadata.
138-
metadata = await apify_kvs_client.get()
139-
140-
# If metadata is None, it means the storage does not exist, so we create it.
141-
if metadata is None:
142-
id = ApifyKeyValueStoreMetadata.model_validate(
143-
await apify_kvss_client.get_or_create(),
144-
).id
140+
# Freshly fetched named storage. No need to fetch it again.
141+
return cls(
142+
api_client=apify_client_async.key_value_store(key_value_store_id=metadata.id),
143+
api_public_base_url='', # Remove in version 4.0, https://github.com/apify/apify-sdk-python/issues/635
144+
lock=asyncio.Lock(),
145+
)
146+
# If id is provided, then storage has to exists.
147+
elif id:
148+
# Now create the client for the determined ID
145149
apify_kvs_client = apify_client_async.key_value_store(key_value_store_id=id)
146-
147-
# Verify that the storage exists by fetching its metadata again.
148-
metadata = await apify_kvs_client.get()
149-
if metadata is None:
150-
raise ValueError(f'Opening key-value store with id={id}, name={name}, and alias={alias} failed.')
151-
152-
return cls(
153-
api_client=apify_kvs_client,
154-
api_public_base_url='', # Remove in version 4.0, https://github.com/apify/apify-sdk-python/issues/635
155-
lock=asyncio.Lock(),
156-
)
150+
# Fetch its metadata.
151+
raw_metadata = await apify_kvs_client.get()
152+
# If metadata is None, it means the storage does not exist.
153+
if raw_metadata is None:
154+
raise ValueError(f'Opening key-value store with id={id} failed.')
155+
return cls(
156+
api_client=apify_kvs_client,
157+
api_public_base_url='', # Remove in version 4.0, https://github.com/apify/apify-sdk-python/issues/635
158+
lock=asyncio.Lock(),
159+
)
160+
# Default key-value store ID from configuration
161+
elif configuration.default_key_value_store_id:
162+
# Now create the client for the determined ID
163+
apify_kvs_client = apify_client_async.key_value_store(
164+
key_value_store_id=configuration.default_key_value_store_id
165+
)
166+
# Fetch its metadata.
167+
raw_metadata = await apify_kvs_client.get()
168+
if not raw_metadata:
169+
metadata = ApifyKeyValueStoreMetadata.model_validate(await apify_kvss_client.get_or_create(name=name))
170+
apify_kvs_client = apify_client_async.key_value_store(key_value_store_id=metadata.id)
171+
172+
return cls(
173+
api_client=apify_kvs_client,
174+
api_public_base_url='', # Remove in version 4.0, https://github.com/apify/apify-sdk-python/issues/635
175+
lock=asyncio.Lock(),
176+
)
177+
raise RuntimeError('Will never happen')
157178

158179
@override
159180
async def purge(self) -> None:

‎src/apify/storage_clients/_apify/_utils.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -98,8 +98,6 @@ async def _get_alias_map(cls) -> dict[str, str]:
9898
async def resolve_id(self) -> str | None:
9999
"""Get id of the aliased storage.
100100
101-
Either locate the id in the in-memory mapping or create the new storage.
102-
103101
Returns:
104102
Storage id if it exists, None otherwise.
105103
"""

0 commit comments

Comments
 (0)