From 3ba7c979f68dcd9f4c5646522d5156b8340329a4 Mon Sep 17 00:00:00 2001 From: geetu040 Date: Mon, 26 Feb 2024 20:07:55 +0500 Subject: [PATCH] use grpc.Compression instead of creating new enum qdrant.grpc.Compression in qdrant/grpc/__init__.py --- qdrant_client/async_qdrant_remote.py | 7 +++++- qdrant_client/grpc/__init__.py | 25 +------------------ qdrant_client/qdrant_remote.py | 7 +++++- .../test_grpc_channel_compress.py | 15 +++-------- 4 files changed, 16 insertions(+), 38 deletions(-) diff --git a/qdrant_client/async_qdrant_remote.py b/qdrant_client/async_qdrant_remote.py index 0974a7b2f..cbc2bf644 100644 --- a/qdrant_client/async_qdrant_remote.py +++ b/qdrant_client/async_qdrant_remote.py @@ -29,6 +29,7 @@ import httpx import numpy as np +from grpc import Compression from urllib3.util import Url, parse_url from qdrant_client import grpc as grpc @@ -112,10 +113,14 @@ def __init__( self._rest_headers["api-key"] = api_key self._grpc_headers.append(("api-key", api_key)) grpc_compression = kwargs.pop("grpc_compression", None) - if grpc_compression is not None and (not isinstance(grpc_compression, grpc.Compression)): + if grpc_compression is not None and (not isinstance(grpc_compression, Compression)): raise TypeError( f"Expected 'grpc_compression' to be of type grpc.Compression or None, but got {type(grpc_compression).__name__}" ) + if grpc_compression == Compression.Deflate: + raise ValueError( + "grpc.Compression.Defalte is not supported. Try grpc.Compression.Gzip or grpc.Compression.NoCompression" + ) self._grpc_compression = grpc_compression address = f"{self._host}:{self._port}" if self._port is not None else self._host self.rest_uri = f"{self._scheme}://{address}{self._prefix}" diff --git a/qdrant_client/grpc/__init__.py b/qdrant_client/grpc/__init__.py index caf9ea18d..c730ba924 100644 --- a/qdrant_client/grpc/__init__.py +++ b/qdrant_client/grpc/__init__.py @@ -4,27 +4,4 @@ from .json_with_int_pb2 import * from .collections_service_pb2_grpc import * from .points_service_pb2_grpc import * -from .snapshots_service_pb2_grpc import * - - -# Compression Enum Implementation - -import enum -import grpc - -@enum.unique -class Compression(enum.IntEnum): - """Defines a custom compression enum tailored for Qdrant RUST Server RPCs. - - Qdrant RUST Server only supports specific compression algorithms for RPC communication. - This enum mirrors the supported compression options, excluding 'Deflate,' which is not - supported by the server. Use these compression options when specifying the desired - compression method for interactions with the Qdrant RUST Server. - - Attributes: - NoCompression: Do not use any compression algorithm. - Gzip: Use the "Gzip" compression algorithm. - """ - - NoCompression = grpc.Compression.NoCompression - Gzip = grpc.Compression.Gzip +from .snapshots_service_pb2_grpc import * \ No newline at end of file diff --git a/qdrant_client/qdrant_remote.py b/qdrant_client/qdrant_remote.py index 6bef89115..a6d5997d5 100644 --- a/qdrant_client/qdrant_remote.py +++ b/qdrant_client/qdrant_remote.py @@ -19,6 +19,7 @@ import httpx import numpy as np +from grpc import Compression from urllib3.util import Url, parse_url from qdrant_client import grpc as grpc @@ -128,10 +129,14 @@ def __init__( # GRPC Channel-Level Compression grpc_compression = kwargs.pop("grpc_compression", None) - if grpc_compression is not None and not isinstance(grpc_compression, grpc.Compression): + if grpc_compression is not None and not isinstance(grpc_compression, Compression): raise TypeError( f"Expected 'grpc_compression' to be of type grpc.Compression or None, but got {type(grpc_compression).__name__}" ) + if grpc_compression == Compression.Deflate: + raise ValueError( + "grpc.Compression.Defalte is not supported. Try grpc.Compression.Gzip or grpc.Compression.NoCompression" + ) self._grpc_compression = grpc_compression address = f"{self._host}:{self._port}" if self._port is not None else self._host diff --git a/tests/congruence_tests/test_grpc_channel_compress.py b/tests/congruence_tests/test_grpc_channel_compress.py index bee09f002..4c2d95f08 100644 --- a/tests/congruence_tests/test_grpc_channel_compress.py +++ b/tests/congruence_tests/test_grpc_channel_compress.py @@ -1,9 +1,8 @@ import pytest -from grpc import Compression as GRPCCompression +from grpc import Compression from qdrant_client import QdrantClient from qdrant_client.async_qdrant_client import AsyncQdrantClient -from qdrant_client.grpc import Compression from tests.congruence_tests.test_common import generate_fixtures, init_client @@ -49,21 +48,13 @@ def test_compression_none(): def test_compression_deflate(): """ Deflate algorithm is not supported by RUST GRPC Server - grpc.Compression provides Deflate option - qdrant.grpc.Compression is tailored to exclude Deflate option - Trying Deflate compression should raise particular errors + Trying Deflate compression should raise ValueError """ - with pytest.raises(AttributeError): + with pytest.raises(ValueError): # creates a grpc client with invalid enum attribute client = QdrantClient(prefer_grpc=True, grpc_compression=Compression.Deflate, timeout=30) - with pytest.raises(TypeError): - # creates a grpc client with grpc enum not that is excluded in qdrant.grpc enum - client = QdrantClient( - prefer_grpc=True, grpc_compression=GRPCCompression.Deflate, timeout=30 - ) - # Anticipated Errors were successfully raised