Skip to content

Commit e014bf8

Browse files
authored
Collection API completeness (#59)
* Adding /figures support * Adding support for /responsibleShard * Adding support for /shards * Adding support for /revision * Adding support for /checksum * Added support for /key-generators * Skipping part of test in 3.11 * Adding configure method * Adding renaming method * recalculate-the-document-count-of-a-collection * compact-a-collection
1 parent 52493cb commit e014bf8

File tree

7 files changed

+677
-17
lines changed

7 files changed

+677
-17
lines changed

arangoasync/collection.py

Lines changed: 315 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,16 @@
1616
HTTP_PRECONDITION_FAILED,
1717
)
1818
from arangoasync.exceptions import (
19+
CollectionChecksumError,
20+
CollectionCompactError,
21+
CollectionConfigureError,
1922
CollectionPropertiesError,
23+
CollectionRecalculateCountError,
24+
CollectionRenameError,
25+
CollectionResponsibleShardError,
26+
CollectionRevisionError,
27+
CollectionShardsError,
28+
CollectionStatisticsError,
2029
CollectionTruncateError,
2130
DocumentCountError,
2231
DocumentDeleteError,
@@ -40,7 +49,9 @@
4049
from arangoasync.result import Result
4150
from arangoasync.serialization import Deserializer, Serializer
4251
from arangoasync.typings import (
52+
CollectionInfo,
4353
CollectionProperties,
54+
CollectionStatistics,
4455
IndexProperties,
4556
Json,
4657
Jsons,
@@ -481,6 +492,26 @@ def response_handler(resp: Response) -> bool:
481492

482493
return await self._executor.execute(request, response_handler)
483494

495+
async def recalculate_count(self) -> None:
496+
"""Recalculate the document count.
497+
498+
Raises:
499+
CollectionRecalculateCountError: If re-calculation fails.
500+
501+
References:
502+
- `recalculate-the-document-count-of-a-collection <https://docs.arangodb.com/stable/develop/http-api/collections/#recalculate-the-document-count-of-a-collection>`__
503+
""" # noqa: E501
504+
request = Request(
505+
method=Method.PUT,
506+
endpoint=f"/_api/collection/{self.name}/recalculateCount",
507+
)
508+
509+
def response_handler(resp: Response) -> None:
510+
if not resp.is_success:
511+
raise CollectionRecalculateCountError(resp, request)
512+
513+
await self._executor.execute(request, response_handler)
514+
484515
async def properties(self) -> Result[CollectionProperties]:
485516
"""Return the full properties of the current collection.
486517
@@ -501,7 +532,129 @@ async def properties(self) -> Result[CollectionProperties]:
501532
def response_handler(resp: Response) -> CollectionProperties:
502533
if not resp.is_success:
503534
raise CollectionPropertiesError(resp, request)
504-
return CollectionProperties(self._executor.deserialize(resp.raw_body))
535+
return CollectionProperties(self.deserializer.loads(resp.raw_body))
536+
537+
return await self._executor.execute(request, response_handler)
538+
539+
async def configure(
540+
self,
541+
cache_enabled: Optional[bool] = None,
542+
computed_values: Optional[Jsons] = None,
543+
replication_factor: Optional[int | str] = None,
544+
schema: Optional[Json] = None,
545+
wait_for_sync: Optional[bool] = None,
546+
write_concern: Optional[int] = None,
547+
) -> Result[CollectionProperties]:
548+
"""Changes the properties of a collection.
549+
550+
Only the provided attributes are updated.
551+
552+
Args:
553+
cache_enabled (bool | None): Whether the in-memory hash cache
554+
for documents should be enabled for this collection.
555+
computed_values (list | None): An optional list of objects, each
556+
representing a computed value.
557+
replication_factor (int | None): In a cluster, this attribute determines
558+
how many copies of each shard are kept on different DB-Servers.
559+
For SatelliteCollections, it needs to be the string "satellite".
560+
schema (dict | None): The configuration of the collection-level schema
561+
validation for documents.
562+
wait_for_sync (bool | None): If set to `True`, the data is synchronized
563+
to disk before returning from a document create, update, replace or
564+
removal operation.
565+
write_concern (int | None): Determines how many copies of each shard are
566+
required to be in sync on the different DB-Servers.
567+
568+
Returns:
569+
CollectionProperties: Properties.
570+
571+
Raises:
572+
CollectionConfigureError: If configuration fails.
573+
574+
References:
575+
- `change-the-properties-of-a-collection <https://docs.arangodb.com/stable/develop/http-api/collections/#change-the-properties-of-a-collection>`__
576+
""" # noqa: E501
577+
data: Json = {}
578+
if cache_enabled is not None:
579+
data["cacheEnabled"] = cache_enabled
580+
if computed_values is not None:
581+
data["computedValues"] = computed_values
582+
if replication_factor is not None:
583+
data["replicationFactor"] = replication_factor
584+
if schema is not None:
585+
data["schema"] = schema
586+
if wait_for_sync is not None:
587+
data["waitForSync"] = wait_for_sync
588+
if write_concern is not None:
589+
data["writeConcern"] = write_concern
590+
request = Request(
591+
method=Method.PUT,
592+
endpoint=f"/_api/collection/{self.name}/properties",
593+
data=self.serializer.dumps(data),
594+
)
595+
596+
def response_handler(resp: Response) -> CollectionProperties:
597+
if not resp.is_success:
598+
raise CollectionConfigureError(resp, request)
599+
return CollectionProperties(self.deserializer.loads(resp.raw_body))
600+
601+
return await self._executor.execute(request, response_handler)
602+
603+
async def rename(self, new_name: str) -> None:
604+
"""Rename the collection.
605+
606+
Renames may not be reflected immediately in async execution, batch
607+
execution or transactions. It is recommended to initialize new API
608+
wrappers after a rename.
609+
610+
Note:
611+
Renaming collections is not supported in cluster deployments.
612+
613+
Args:
614+
new_name (str): New collection name.
615+
616+
Raises:
617+
CollectionRenameError: If rename fails.
618+
619+
References:
620+
- `rename-a-collection <https://docs.arangodb.com/stable/develop/http-api/collections/#rename-a-collection>`__
621+
""" # noqa: E501
622+
data: Json = {"name": new_name}
623+
request = Request(
624+
method=Method.PUT,
625+
endpoint=f"/_api/collection/{self.name}/rename",
626+
data=self.serializer.dumps(data),
627+
)
628+
629+
def response_handler(resp: Response) -> None:
630+
if not resp.is_success:
631+
raise CollectionRenameError(resp, request)
632+
self._name = new_name
633+
self._id_prefix = f"{new_name}/"
634+
635+
await self._executor.execute(request, response_handler)
636+
637+
async def compact(self) -> Result[CollectionInfo]:
638+
"""Compact a collection.
639+
640+
Returns:
641+
CollectionInfo: Collection information.
642+
643+
Raises:
644+
CollectionCompactError: If compaction fails.
645+
646+
References:
647+
- `compact-a-collection <https://docs.arangodb.com/stable/develop/http-api/collections/#compact-a-collection>`__
648+
""" # noqa: E501
649+
request = Request(
650+
method=Method.PUT,
651+
endpoint=f"/_api/collection/{self.name}/compact",
652+
)
653+
654+
def response_handler(resp: Response) -> CollectionInfo:
655+
if not resp.is_success:
656+
raise CollectionCompactError(resp, request)
657+
return CollectionInfo(self.deserializer.loads(resp.raw_body))
505658

506659
return await self._executor.execute(request, response_handler)
507660

@@ -552,7 +705,10 @@ async def count(self) -> Result[int]:
552705
553706
Raises:
554707
DocumentCountError: If retrieval fails.
555-
"""
708+
709+
References:
710+
- `get-the-document-count-of-a-collection <https://docs.arangodb.com/stable/develop/http-api/collections/#get-the-document-count-of-a-collection>`__
711+
""" # noqa: E501
556712
request = Request(
557713
method=Method.GET, endpoint=f"/_api/collection/{self.name}/count"
558714
)
@@ -565,6 +721,158 @@ def response_handler(resp: Response) -> int:
565721

566722
return await self._executor.execute(request, response_handler)
567723

724+
async def statistics(self) -> Result[CollectionStatistics]:
725+
"""Get additional statistical information about the collection.
726+
727+
Returns:
728+
CollectionStatistics: Collection statistics.
729+
730+
Raises:
731+
CollectionStatisticsError: If retrieval fails.
732+
733+
References:
734+
- `get-the-collection-statistics <https://docs.arangodb.com/stable/develop/http-api/collections/#get-the-collection-statistics>`__
735+
""" # noqa: E501
736+
request = Request(
737+
method=Method.GET,
738+
endpoint=f"/_api/collection/{self.name}/figures",
739+
)
740+
741+
def response_handler(resp: Response) -> CollectionStatistics:
742+
if not resp.is_success:
743+
raise CollectionStatisticsError(resp, request)
744+
return CollectionStatistics(self.deserializer.loads(resp.raw_body))
745+
746+
return await self._executor.execute(request, response_handler)
747+
748+
async def responsible_shard(self, document: Json) -> Result[str]:
749+
"""Return the ID of the shard responsible for given document.
750+
751+
If the document does not exist, return the shard that would be
752+
responsible.
753+
754+
Args:
755+
document (dict): Document body with "_key" field.
756+
757+
Returns:
758+
str: Shard ID.
759+
760+
Raises:
761+
CollectionResponsibleShardError: If retrieval fails.
762+
763+
References:
764+
- `get-the-responsible-shard-for-a-document <https://docs.arangodb.com/stable/develop/http-api/collections/#get-the-responsible-shard-for-a-document>`__
765+
""" # noqa: E501
766+
request = Request(
767+
method=Method.PUT,
768+
endpoint=f"/_api/collection/{self.name}/responsibleShard",
769+
data=self.serializer.dumps(document),
770+
)
771+
772+
def response_handler(resp: Response) -> str:
773+
if resp.is_success:
774+
body = self.deserializer.loads(resp.raw_body)
775+
return cast(str, body["shardId"])
776+
raise CollectionResponsibleShardError(resp, request)
777+
778+
return await self._executor.execute(request, response_handler)
779+
780+
async def shards(self, details: Optional[bool] = None) -> Result[Json]:
781+
"""Return collection shards and properties.
782+
783+
Available only in a cluster setup.
784+
785+
Args:
786+
details (bool | None): If set to `True`, include responsible
787+
servers for these shards.
788+
789+
Returns:
790+
dict: Collection shards.
791+
792+
Raises:
793+
CollectionShardsError: If retrieval fails.
794+
795+
References:
796+
- `get-the-shard-ids-of-a-collection <https://docs.arangodb.com/stable/develop/http-api/collections/#get-the-shard-ids-of-a-collection>`__
797+
""" # noqa: E501
798+
params: Params = {}
799+
if details is not None:
800+
params["details"] = details
801+
802+
request = Request(
803+
method=Method.GET,
804+
endpoint=f"/_api/collection/{self.name}/shards",
805+
params=params,
806+
)
807+
808+
def response_handler(resp: Response) -> Json:
809+
if not resp.is_success:
810+
raise CollectionShardsError(resp, request)
811+
return cast(Json, self.deserializer.loads(resp.raw_body)["shards"])
812+
813+
return await self._executor.execute(request, response_handler)
814+
815+
async def revision(self) -> Result[str]:
816+
"""Return collection revision.
817+
818+
Returns:
819+
str: Collection revision.
820+
821+
Raises:
822+
CollectionRevisionError: If retrieval fails.
823+
824+
References:
825+
- `get-the-collection-revision-id <https://docs.arangodb.com/stable/develop/http-api/collections/#get-the-collection-revision-id>`__
826+
""" # noqa: E501
827+
request = Request(
828+
method=Method.GET,
829+
endpoint=f"/_api/collection/{self.name}/revision",
830+
)
831+
832+
def response_handler(resp: Response) -> str:
833+
if not resp.is_success:
834+
raise CollectionRevisionError(resp, request)
835+
return cast(str, self.deserializer.loads(resp.raw_body)["revision"])
836+
837+
return await self._executor.execute(request, response_handler)
838+
839+
async def checksum(
840+
self, with_rev: Optional[bool] = None, with_data: Optional[bool] = None
841+
) -> Result[str]:
842+
"""Calculate collection checksum.
843+
844+
Args:
845+
with_rev (bool | None): Include document revisions in checksum calculation.
846+
with_data (bool | None): Include document data in checksum calculation.
847+
848+
Returns:
849+
str: Collection checksum.
850+
851+
Raises:
852+
CollectionChecksumError: If retrieval fails.
853+
854+
References:
855+
- `get-the-collection-checksum <https://docs.arangodb.com/stable/develop/http-api/collections/#get-the-collection-checksum>`__
856+
""" # noqa: E501
857+
params: Params = {}
858+
if with_rev is not None:
859+
params["withRevision"] = with_rev
860+
if with_data is not None:
861+
params["withData"] = with_data
862+
863+
request = Request(
864+
method=Method.GET,
865+
endpoint=f"/_api/collection/{self.name}/checksum",
866+
params=params,
867+
)
868+
869+
def response_handler(resp: Response) -> str:
870+
if not resp.is_success:
871+
raise CollectionChecksumError(resp, request)
872+
return cast(str, self.deserializer.loads(resp.raw_body)["checksum"])
873+
874+
return await self._executor.execute(request, response_handler)
875+
568876
async def has(
569877
self,
570878
document: str | Json,
@@ -1444,9 +1752,9 @@ async def insert(
14441752

14451753
def response_handler(resp: Response) -> bool | Json:
14461754
if resp.is_success:
1447-
if silent is True:
1755+
if silent:
14481756
return True
1449-
return self._executor.deserialize(resp.raw_body)
1757+
return self.deserializer.loads(resp.raw_body)
14501758
msg: Optional[str] = None
14511759
if resp.status_code == HTTP_BAD_PARAMETER:
14521760
msg = (
@@ -1551,7 +1859,7 @@ def response_handler(resp: Response) -> bool | Json:
15511859
if resp.is_success:
15521860
if silent is True:
15531861
return True
1554-
return self._executor.deserialize(resp.raw_body)
1862+
return self.deserializer.loads(resp.raw_body)
15551863
msg: Optional[str] = None
15561864
if resp.status_code == HTTP_PRECONDITION_FAILED:
15571865
raise DocumentRevisionError(resp, request)
@@ -1641,7 +1949,7 @@ def response_handler(resp: Response) -> bool | Json:
16411949
if resp.is_success:
16421950
if silent is True:
16431951
return True
1644-
return self._executor.deserialize(resp.raw_body)
1952+
return self.deserializer.loads(resp.raw_body)
16451953
msg: Optional[str] = None
16461954
if resp.status_code == HTTP_PRECONDITION_FAILED:
16471955
raise DocumentRevisionError(resp, request)
@@ -1726,7 +2034,7 @@ def response_handler(resp: Response) -> bool | Json:
17262034
if resp.is_success:
17272035
if silent is True:
17282036
return True
1729-
return self._executor.deserialize(resp.raw_body)
2037+
return self.deserializer.loads(resp.raw_body)
17302038
msg: Optional[str] = None
17312039
if resp.status_code == HTTP_PRECONDITION_FAILED:
17322040
raise DocumentRevisionError(resp, request)

0 commit comments

Comments
 (0)