infinity-sdk 0.1.0__py3-none-any.whl

infinity/__init__.py

@@ -0,0 +1,29 @@
+# Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib.metadata
+
+__version__ = importlib.metadata.version("infinity_sdk")
+
+from infinity.common import URI, NetworkAddress, LOCAL_HOST
+from infinity.infinity import InfinityConnection
+from infinity.remote_thrift.infinity import RemoteThriftInfinityConnection
+
+def connect(
+        uri: URI = LOCAL_HOST
+) -> InfinityConnection:
+    if isinstance(uri, NetworkAddress) and (uri.port == 9090 or uri.port == 23817 or uri.port == 9070):
+        return RemoteThriftInfinityConnection(uri)
+    else:
+        raise Exception(f"unknown uri: {uri}")

infinity/common.py

@@ -0,0 +1,40 @@
+# Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+
+
+class NetworkAddress:
+    def __init__(self, ip, port):
+        self.ip = ip
+        self.port = port
+
+    def __str__(self):
+        return f'IP: {self.ip}, Port: {self.port}'
+
+
+URI = Union[NetworkAddress, Path]
+VEC = Union[list, np.ndarray]
+INSERT_DATA = dict[str, Union[str, int, float, list[Union[int, float]]]]
+
+REMOTE_HOST = NetworkAddress("127.0.0.1", 23817)
+LOCAL_HOST = NetworkAddress("0.0.0.0", 23817)
+
+
+class ConflictType(object):
+    Ignore = 0
+    Error = 1
+    Replace = 2

infinity/connection_pool.py

@@ -0,0 +1,57 @@
+from threading import Lock, Condition
+import infinity
+from infinity.common import NetworkAddress
+import logging
+from infinity.remote_thrift.infinity import RemoteThriftInfinityConnection
+
+
+class ConnectionPool(object):
+
+
+    def __init__(self, uri = NetworkAddress("127.0.0.1", 23817), min_size=4, max_size=16, timeout=10.0):
+        assert (min_size <= max_size)
+        self.uri_ = uri
+        self.min_size_ = min_size
+        self.max_size_ = max_size
+        self.curr_size_ = 0
+        self.timeout_ = timeout
+        self.free_pool_ = []
+        self.created_conns_ = []
+        self.mutex_ = Lock()
+        self.cond_ = Condition(self.mutex_)
+
+        while(self.curr_size_ < self.min_size_):
+            self._create_conn()
+
+
+    def _create_conn(self):
+        infinity_coon = infinity.connect(self.uri_)
+        self.curr_size_ += 1
+        self.free_pool_.append(infinity_coon)
+        self.created_conns_.append(infinity_coon)
+
+
+    def get_conn(self) -> RemoteThriftInfinityConnection:
+        with self.mutex_:
+            while(len(self.free_pool_) == 0 and self.curr_size_ == self.max_size_):
+                is_waken = self.cond_.wait(self.timeout_)
+                if (not is_waken):
+                    logging.warn("connection waiting...")
+            if (len(self.free_pool_) == 0):
+                self._create_conn()
+            conn = self.free_pool_.pop()
+            return conn
+
+
+    def release_conn(self, conn):
+        with self.mutex_:
+            if(self.free_pool_.count(conn)):
+                raise Exception("the connection has been released")
+            self.curr_size_ -= 1
+            self.free_pool_.append(conn)
+            self.cond_.notify(n=1)
+
+
+    def destroy(self):
+        for conn in iter(self.created_conns_):
+            conn.disconnect()

infinity/db.py

@@ -0,0 +1,45 @@
+# Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC, abstractmethod
+
+class Database(ABC):
+
+    @abstractmethod
+    def create_table(self, table_name, schema, options):
+        pass  # implement create table logic here
+
+    @abstractmethod
+    def drop_table(self, table_name):
+        pass  # implement drop table logic here
+
+    @abstractmethod
+    def list_tables(self):
+        pass  # implement list tables logic here
+
+    @abstractmethod
+    def show_table(self, table_name):
+        pass  # implement describe table logic here
+
+    @abstractmethod
+    def show_columns(self, table_name):
+        pass  # implement describe table logic here
+
+    @abstractmethod
+    def get_table(self, table_name):
+        pass  # implement get table logic here
+
+    @abstractmethod
+    def show_tables(self):
+        pass

infinity/errors.py

@@ -0,0 +1,137 @@
+# Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from enum import IntEnum
+
+
+class ErrorCode(IntEnum):
+    OK = 0,
+
+    INVALID_TIME_INFO = 1001,
+    EMPTY_CONFIG_PARAMETER = 1002,
+    MISMATCH_VERSION = 1003,
+    INVALID_TIMEZONE = 1004,
+    INVALID_BYTE_SIZE = 1005,
+    INVALID_IP_ADDR = 1006,
+    INVALID_LOG_LEVEL = 1007,
+
+    WRONG_PASSWD = 2001,
+    INSUFFICIENT_PRIVILEGE = 2002,
+
+    INVALID_USERNAME = 3001,
+    INVALID_PASSWD = 3002,
+    INVALID_DB_NAME = 3003,
+    INVALID_TABLE_NAME = 3004,
+    INVALID_COLUMN_NAME = 3005,
+    INVALID_INDEX_NAME = 3006,
+    INVALID_COLUMN_DEFINITION = 3007,
+    INVALID_TABLE_DEFINITION = 3008,
+    INVALID_INDEX_DEFINITION = 3009,
+    DATA_TYPE_MISMATCH = 3010,
+    NAME_TOO_LONG = 3011,
+    RESERVED_NAME = 3012,
+    SYNTAX_ERROR = 3013,
+    INVALID_PARAMETER_VALUE = 3014,
+    DUPLICATE_USER = 3015,
+    DUPLICATE_DATABASE_NAME = 3016,
+    DUPLICATE_TABLE_NAME = 3017,
+    DUPLICATE_INDEX_NAME = 3018,
+    DUPLICATE_INDEX_ON_COLUMN = 3019,
+    NO_SUCH_USER = 3020,
+    DB_NOT_EXIST = 3021,
+    TABLE_NOT_EXIST = 3022,
+    INDEX_NOT_EXIST = 3023,
+    COLUMN_NOT_EXIST = 3024,
+    AGG_NOT_ALLOW_IN_WHERE_CLAUSE = 3025,
+    COLUMN_IN_SELECT_NOT_IN_GROUP_BY = 3026,
+    NO_SUCH_SYSTEM_VAR = 3027,
+    INVALID_SYSTEM_VAR_VALUE = 3028,
+    SYSTEM_VAR_READ_ONLY = 3029,
+    FUNCTION_NOT_FOUND = 3030,
+    SPECIAL_FUNCTION_NOT_FOUND = 3031,
+    NOT_SUPPORTED = 3032,
+    DROPPING_USING_DB = 3033,
+    SESSION_NOT_FOUND = 3034,
+    RECURSIVE_AGG = 3035,
+    FUNCTION_ARGS_ERROR = 3036,
+    IMPORT_FILE_FORMAT_ERROR = 3037,
+    DATA_NOT_EXIST = 3038,
+    COLUMN_COUNT_MISMATCH = 3039,
+    EMPTY_DB_NAME = 3040,
+    EMPTY_TABLE_NAME = 3041,
+    EMPTY_COLUMN_NAME = 3042,
+    EMPTY_INDEX_NAME = 3043,
+    EXCEED_DB_NAME_LENGTH = 3044,
+    EXCEED_TABLE_NAME_LENGTH = 3045,
+    EXCEED_COLUMN_NAME_LENGTH = 3046,
+    EXCEED_INDEX_NAME_LENGTH = 3047,
+    NO_COLUMN_DEFINED = 3048,
+    NOT_SUPPORTED_TYPE_CONVERSION = 3049,
+    EMPTY_SELECT_FIELDS = 3050,
+    INVALID_DATA_TYPE = 3051,
+    PARSE_MATCH_EXPR_FAILED = 3052,
+    FTS_INDEX_NOT_EXIST = 3053,
+    UNKNOWN_FTS_FAULT = 3054,
+    INVALID_CONSTRAINT_TYPE = 3055,
+    INVALID_KNN_DISTANCE_TYPE = 3056,
+    INVALID_EMBEDDING_DATA_TYPE = 3057,
+    INVALID_CONSTANT_TYPE = 3058,
+    INVALID_PARSED_EXPR_TYPE = 3059,
+    INVALID_INDEX_TYPE = 3060,
+    INVALID_INDEX_PARAM = 3061,
+    LACK_INDEX_PARAM = 3062,
+    INVALID_FILTER_EXPRESSION = 3063,
+    MUTIPLE_FUNCTION_MATCHED = 3064,
+    INSERT_WITHOUT_VALUES = 3065,
+    INVALID_CONFLICT_TYPE = 3066,
+    INVALID_JSON_FORMAT = 3067,
+    DUPLICATE_COLUMN_NAME = 3068,
+    INVALID_EXPRESSION = 3069,
+    SEGMENT_NOT_EXIST = 3070,
+    AGGREGATE_FUNCTION_WITH_EMPTY_ARGS = 3071,
+    BLOCK_NOT_EXIST = 3072,
+    INVALID_TOPK_TYPE = 3073,
+    INVALID_CREATE_OPTION = 3074,
+    INVALID_DROP_OPTION = 3075,
+
+    TXN_ROLLBACK = 4001,
+    TXN_CONFLICT = 4002,
+
+    DISK_FULL = 5001,
+    OUT_OF_MEMORY = 5002,
+    TOO_MANY_CONNECTIONS = 5003,
+    CONFIGURATION_LIMIT_EXCEED = 5004,
+    QUERY_IS_TOO_COMPLEX = 5005,
+
+    QUERY_CANCELLED = 6001,
+    QUERY_NOT_SUPPORTED = 6002,
+    CLIENT_CLOSE = 6003,
+
+    DISK_IO_ERROR = 7001,
+    DUPLICATED_FILE = 7002,
+
+
+    CONFIG_FILE_ERROR = 7003,
+    LOCK_FILE_EXISTS = 7004,
+    CATALOG_CORRUPTED = 7005,
+    DATA_CORRUPTED = 7006,
+    INDEX_CORRUPTED = 7007,
+    FILE_NOT_FOUND = 7008,
+    DIR_NOT_FOUND = 7009,
+    DATA_IO_ERROR = 7010,
+    UNEXPECTED_ERROR = 7011,
+
+    INVALID_ENTRY = 8001,
+    NOT_FOUND_ENTRY = 8002,
+    EMPTY_ENTRY_LIST = 8003,

infinity/index.py

@@ -0,0 +1,77 @@
+# Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from enum import Enum
+
+import infinity.remote_thrift.infinity_thrift_rpc.ttypes as ttypes
+
+
+class IndexType(Enum):
+    IVFFlat = 1
+    HnswLVQ = 2
+    Hnsw = 3
+    FullText = 4
+
+    def to_ttype(self):
+        if self == IndexType.IVFFlat:
+            return ttypes.IndexType.IVFFlat
+        elif self == IndexType.HnswLVQ:
+            return ttypes.IndexType.HnswLVQ
+        elif self == IndexType.Hnsw:
+            return ttypes.IndexType.Hnsw
+        elif self == IndexType.FullText:
+            return ttypes.IndexType.FullText
+        else:
+            raise Exception("Unknown index type")
+
+
+class InitParameter:
+    def __init__(self, param_name: str, param_value: str):
+        self.param_name = param_name
+        self.param_value = param_value
+
+    def __str__(self):
+        return f"InitParameter({self.param_name}, {self.param_value})"
+
+    def __repr__(self):
+        return self.__str__()
+
+    def to_ttype(self):
+        return ttypes.InitParameter(self.param_name, self.param_value)
+
+
+class IndexInfo:
+    def __init__(self, column_name: str, index_type: IndexType, params: list[InitParameter]):
+        self.column_name = column_name
+        self.index_type = index_type
+        self.params = params
+
+    def __str__(self):
+        return f"IndexInfo({self.column_name}, {self.index_type}, {self.params})"
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __eq__(self, other):
+        return self.column_name == other.index_name and self.index_type == other.index_type and self.params == other.params
+
+    def __hash__(self):
+        return hash((self.column_name, self.index_type, self.params))
+
+    def to_ttype(self):
+        return ttypes.IndexInfo(
+            self.column_name,
+            self.index_type.to_ttype(),
+            [p.to_ttype() for p in self.params]
+        )

infinity/infinity.py

@@ -0,0 +1,68 @@
+# Copyright(C) 2023 InfiniFlow, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from abc import ABC, abstractmethod
+from enum import Enum
+
+from infinity import URI
+
+
+class ShowVariable(Enum):
+    QUERY_COUNT = "query_count"
+    SESSION_COUNT = "session_count"
+    BUFFER_POOL_USAGE = "buffer_pool_usage"
+    VERSION = "version"
+    QUERY_MEMORY_LIMIT = "query_memory_limit"
+    QUERY_CPU_LIMIT = "query_cpu_limit"
+    LOG_LEVEL = "log_level"
+    SCHEDULE_POLICY = "schedule_policy"
+    LISTEN_ADDRESS = "listen_address"
+    SQL_PORT = "sql_port"
+    SDK_PORT = "sdk_port"
+    HTTP_API_PORT = "http_api_port"
+    DATA_URL = "data_url"
+    TIME_ZONE = "time_zone"
+
+
+# abstract class
+class InfinityConnection(ABC):
+    def __init__(self, uri: URI):
+        self.uri = uri
+
+    @abstractmethod
+    def create_database(self, db_name, options=None):
+        pass
+
+    @abstractmethod
+    def list_databases(self):
+        pass
+
+    @abstractmethod
+    def show_database(self, db_name):
+        pass
+
+    @abstractmethod
+    def drop_database(self, db_name, options=None):
+        pass
+
+    @abstractmethod
+    def get_database(self, db_name):
+        pass
+
+    @abstractmethod
+    def disconnect(self):
+        pass
+
+    @abstractmethod
+    def show_variable(self, variable: ShowVariable):
+        pass