vastdb 0.0.5.3__py3-none-any.whl → 0.1.1__py3-none-any.whl

vast_flatbuf/tabular/GetTableStatsResponse.py

@@ -45,7 +45,39 @@ class GetTableStatsResponse(object):
             return bool(self._tab.Get(flatbuffers.number_types.BoolFlags, o + self._tab.Pos))
         return False
 
-def Start(builder): builder.StartObject(3)
+    # GetTableStatsResponse
+    def AddressType(self):
+        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(10))
+        if o != 0:
+            return self._tab.String(o + self._tab.Pos)
+        return None
+
+    # GetTableStatsResponse
+    def Vips(self, j):
+        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(12))
+        if o != 0:
+            x = self._tab.Vector(o)
+            x += flatbuffers.number_types.UOffsetTFlags.py_type(j) * 4
+            x = self._tab.Indirect(x)
+            from vast_flatbuf.tabular.VipRange import VipRange
+            obj = VipRange()
+            obj.Init(self._tab.Bytes, x)
+            return obj
+        return None
+
+    # GetTableStatsResponse
+    def VipsLength(self):
+        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(12))
+        if o != 0:
+            return self._tab.VectorLen(o)
+        return 0
+
+    # GetTableStatsResponse
+    def VipsIsNone(self):
+        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(12))
+        return o == 0
+
+def Start(builder): builder.StartObject(5)
 def GetTableStatsResponseStart(builder):
     """This method is deprecated. Please switch to Start."""
     return Start(builder)
@@ -61,6 +93,18 @@ def AddIsExternalRowidAlloc(builder, isExternalRowidAlloc): builder.PrependBoolS
 def GetTableStatsResponseAddIsExternalRowidAlloc(builder, isExternalRowidAlloc):
     """This method is deprecated. Please switch to AddIsExternalRowidAlloc."""
     return AddIsExternalRowidAlloc(builder, isExternalRowidAlloc)
+def AddAddressType(builder, addressType): builder.PrependUOffsetTRelativeSlot(3, flatbuffers.number_types.UOffsetTFlags.py_type(addressType), 0)
+def GetTableStatsResponseAddAddressType(builder, addressType):
+    """This method is deprecated. Please switch to AddAddressType."""
+    return AddAddressType(builder, addressType)
+def AddVips(builder, vips): builder.PrependUOffsetTRelativeSlot(4, flatbuffers.number_types.UOffsetTFlags.py_type(vips), 0)
+def GetTableStatsResponseAddVips(builder, vips):
+    """This method is deprecated. Please switch to AddVips."""
+    return AddVips(builder, vips)
+def StartVipsVector(builder, numElems): return builder.StartVector(4, numElems, 4)
+def GetTableStatsResponseStartVipsVector(builder, numElems):
+    """This method is deprecated. Please switch to Start."""
+    return StartVipsVector(builder, numElems)
 def End(builder): return builder.EndObject()
 def GetTableStatsResponseEnd(builder):
     """This method is deprecated. Please switch to End."""

vast_flatbuf/tabular/VipRange.py

@@ -0,0 +1,56 @@
+# automatically generated by the FlatBuffers compiler, do not modify
+
+# namespace: tabular
+
+import flatbuffers
+from flatbuffers.compat import import_numpy
+np = import_numpy()
+
+class VipRange(object):
+    __slots__ = ['_tab']
+
+    @classmethod
+    def GetRootAs(cls, buf, offset=0):
+        n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
+        x = VipRange()
+        x.Init(buf, n + offset)
+        return x
+
+    @classmethod
+    def GetRootAsVipRange(cls, buf, offset=0):
+        """This method is deprecated. Please switch to GetRootAs."""
+        return cls.GetRootAs(buf, offset)
+    # VipRange
+    def Init(self, buf, pos):
+        self._tab = flatbuffers.table.Table(buf, pos)
+
+    # VipRange
+    def StartAddress(self):
+        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
+        if o != 0:
+            return self._tab.String(o + self._tab.Pos)
+        return None
+
+    # VipRange
+    def AddressCount(self):
+        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(6))
+        if o != 0:
+            return self._tab.Get(flatbuffers.number_types.Uint16Flags, o + self._tab.Pos)
+        return 0
+
+def Start(builder): builder.StartObject(2)
+def VipRangeStart(builder):
+    """This method is deprecated. Please switch to Start."""
+    return Start(builder)
+def AddStartAddress(builder, startAddress): builder.PrependUOffsetTRelativeSlot(0, flatbuffers.number_types.UOffsetTFlags.py_type(startAddress), 0)
+def VipRangeAddStartAddress(builder, startAddress):
+    """This method is deprecated. Please switch to AddStartAddress."""
+    return AddStartAddress(builder, startAddress)
+def AddAddressCount(builder, addressCount): builder.PrependUint16Slot(1, addressCount, 0)
+def VipRangeAddAddressCount(builder, addressCount):
+    """This method is deprecated. Please switch to AddAddressCount."""
+    return AddAddressCount(builder, addressCount)
+def End(builder): return builder.EndObject()
+def VipRangeEnd(builder):
+    """This method is deprecated. Please switch to End."""
+    return End(builder)

vastdb/__init__.py

@@ -0,0 +1,7 @@
+"""VAST Database Python SDK."""
+
+from . import session
+
+# A helper function, useful as a short-hand for Session c-tor: `session = vastdb.connect(...)`
+connect = session.Session
+connect.__name__ = 'connect'

vastdb/bench/test_perf.py

@@ -0,0 +1,29 @@
+import logging
+import time
+
+import pyarrow as pa
+import pytest
+
+from vastdb import util
+from vastdb.table import ImportConfig, QueryConfig
+
+log = logging.getLogger(__name__)
+
+
+@pytest.mark.benchmark
+def test_bench(session, clean_bucket_name, parquets_path, crater_path):
+    files = [str(parquets_path/f) for f in (parquets_path.glob('**/*.pq'))]
+
+    with session.transaction() as tx:
+        b = tx.bucket(clean_bucket_name)
+        s = b.create_schema('s1')
+        t = util.create_table_from_files(s, 't1', files, config=ImportConfig(import_concurrency=8))
+        config = QueryConfig(num_splits=8, num_sub_splits=4)
+        s = time.time()
+        pa_table = pa.Table.from_batches(t.select(columns=['sid'], predicate=t['sid'] == 10033007, config=config))
+        e = time.time()
+        log.info("'SELECT sid from TABLE WHERE sid = 10033007' returned in %s seconds.", e-s)
+        if crater_path:
+            with open(f'{crater_path}/bench_results', 'a') as f:
+                f.write(f"'SELECT sid FROM TABLE WHERE sid = 10033007' returned in {e-s} seconds")
+        assert pa_table.num_rows == 255_075

vastdb/bucket.py

@@ -0,0 +1,85 @@
+"""VAST Database bucket.
+
+VAST S3 buckets can be used to create Database schemas and tables.
+It is possible to list and access VAST snapshots generated over a bucket.
+"""
+
+import logging
+from dataclasses import dataclass
+
+from . import errors, schema, transaction
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class Snapshot:
+    """VAST bucket-level snapshot."""
+
+    name: str
+    bucket: "Bucket"
+
+
+@dataclass
+class Bucket:
+    """VAST bucket."""
+
+    name: str
+    tx: "transaction.Transaction"
+
+    def create_schema(self, path: str, fail_if_exists=True) -> "schema.Schema":
+        """Create a new schema (a container of tables) under this bucket."""
+        if current := self.schema(path, fail_if_missing=False):
+            if fail_if_exists:
+                raise errors.SchemaExists(self.name, path)
+            else:
+                return current
+        self.tx._rpc.api.create_schema(self.name, path, txid=self.tx.txid)
+        log.info("Created schema: %s", path)
+        return self.schema(path)
+
+    def schema(self, path: str, fail_if_missing=True) -> "schema.Schema":
+        """Get a specific schema (a container of tables) under this bucket."""
+        s = self.schemas(path)
+        log.debug("schema: %s", s)
+        if not s:
+            if fail_if_missing:
+                raise errors.MissingSchema(self.name, path)
+            else:
+                return None
+        assert len(s) == 1, f"Expected to receive only a single schema, but got: {len(s)}. ({s})"
+        log.debug("Found schema: %s", s[0].name)
+        return s[0]
+
+    def schemas(self, name: str = None) -> ["schema.Schema"]:
+        """List bucket's schemas."""
+        schemas = []
+        next_key = 0
+        exact_match = bool(name)
+        log.debug("list schemas param: schema=%s, exact_match=%s", name, exact_match)
+        while True:
+            bucket_name, curr_schemas, next_key, is_truncated, _ = \
+                self.tx._rpc.api.list_schemas(bucket=self.name, next_key=next_key, txid=self.tx.txid,
+                                               name_prefix=name, exact_match=exact_match)
+            if not curr_schemas:
+                break
+            schemas.extend(curr_schemas)
+            if not is_truncated:
+                break
+
+        return [schema.Schema(name=name, bucket=self) for name, *_ in schemas]
+
+    def snapshots(self) -> [Snapshot]:
+        """List bucket's snapshots."""
+        snapshots = []
+        next_key = 0
+        while True:
+            curr_snapshots, is_truncated, next_key = \
+                self.tx._rpc.api.list_snapshots(bucket=self.name, next_token=next_key)
+            if not curr_snapshots:
+                break
+            snapshots.extend(curr_snapshots)
+            if not is_truncated:
+                break
+
+        return [Snapshot(name=snapshot, bucket=self) for snapshot in snapshots]

vastdb/{tests/conftest.py → conftest.py}

@@ -1,33 +1,38 @@
-import pytest
+import os
+from pathlib import Path
+
 import boto3
+import pytest
 
-from vastdb import v2
+import vastdb
 
 
 def pytest_addoption(parser):
-    parser.addoption("--tabular-bucket-name", help="Name of the S3 bucket with Tabular enabled")
-    parser.addoption("--tabular-access-key", help="Access key with Tabular permissions")
-    parser.addoption("--tabular-secret-key", help="Secret key with Tabular permissions")
-    parser.addoption("--tabular-endpoint-url", help="Tabular server endpoint")
+    parser.addoption("--tabular-bucket-name", help="Name of the S3 bucket with Tabular enabled", default="vastdb")
+    parser.addoption("--tabular-access-key", help="Access key with Tabular permissions (AWS_ACCESS_KEY_ID)", default=os.environ.get("AWS_ACCESS_KEY_ID", None))
+    parser.addoption("--tabular-secret-key", help="Secret key with Tabular permissions (AWS_SECRET_ACCESS_KEY)", default=os.environ.get("AWS_SECRET_ACCESS_KEY", None))
+    parser.addoption("--tabular-endpoint-url", help="Tabular server endpoint", default="http://localhost:9090")
+    parser.addoption("--data-path", help="Data files location", default=None)
+    parser.addoption("--crater-path", help="Save benchmark results in a dedicated location", default=None)
 
 
-@pytest.fixture(scope="module")
-def rpc(request):
-    return v2.connect(
+@pytest.fixture(scope="session")
+def session(request):
+    return vastdb.connect(
         access=request.config.getoption("--tabular-access-key"),
         secret=request.config.getoption("--tabular-secret-key"),
         endpoint=request.config.getoption("--tabular-endpoint-url"),
     )
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="session")
 def test_bucket_name(request):
     return request.config.getoption("--tabular-bucket-name")
 
 
-@pytest.fixture(scope="module")
-def clean_bucket_name(request, test_bucket_name, rpc):
-    with rpc.transaction() as tx:
+@pytest.fixture(scope="function")
+def clean_bucket_name(request, test_bucket_name, session):
+    with session.transaction() as tx:
         b = tx.bucket(test_bucket_name)
         for s in b.schemas():
             for t in s.tables():
@@ -36,10 +41,20 @@ def clean_bucket_name(request, test_bucket_name, rpc):
     return test_bucket_name
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="session")
 def s3(request):
     return boto3.client(
         's3',
         aws_access_key_id=request.config.getoption("--tabular-access-key"),
         aws_secret_access_key=request.config.getoption("--tabular-secret-key"),
         endpoint_url=request.config.getoption("--tabular-endpoint-url"))
+
+
+@pytest.fixture(scope="function")
+def parquets_path(request):
+    return Path(request.config.getoption("--data-path"))
+
+
+@pytest.fixture(scope="function")
+def crater_path(request):
+    return request.config.getoption("--crater-path")

vastdb/errors.py

@@ -0,0 +1,175 @@
+import logging
+import xml.etree.ElementTree
+from dataclasses import dataclass
+from enum import Enum
+
+import requests
+
+
+class HttpStatus(Enum):
+    SUCCESS = 200
+    BAD_REQUEST = 400
+    FOBIDDEN = 403
+    NOT_FOUND = 404
+    METHOD_NOT_ALLOWED = 405
+    REQUEST_TIMEOUT = 408
+    CONFLICT = 409
+    INTERNAL_SERVER_ERROR = 500
+    NOT_IMPLEMENTED = 501
+    SERVICE_UNAVAILABLE = 503
+
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class HttpError(Exception):
+    code: str
+    message: str
+    url: str
+    status: int  # HTTP status
+    headers: requests.structures.CaseInsensitiveDict  # HTTP response headers
+
+    def __post_init__(self):
+        self.args = [vars(self)]
+
+
+class NotFound(HttpError):
+    pass
+
+
+class Forbidden(HttpError):
+    pass
+
+
+class BadRequest(HttpError):
+    pass
+
+
+class MethodNotAllowed(HttpError):
+    pass
+
+
+class RequestTimeout(HttpError):
+    pass
+
+
+class Conflict(HttpError):
+    pass
+
+
+class InternalServerError(HttpError):
+    pass
+
+
+class NotImplemented(HttpError):
+    pass
+
+
+class ServiceUnavailable(HttpError):
+    pass
+
+
+class UnexpectedError(HttpError):
+    pass
+
+
+@dataclass
+class ImportFilesError(Exception):
+    message: str
+    error_dict: dict
+
+
+class InvalidArgument(Exception):
+    pass
+
+
+class Missing(Exception):
+    pass
+
+
+@dataclass
+class MissingBucket(Missing):
+    bucket: str
+
+
+@dataclass
+class MissingSchema(Missing):
+    bucket: str
+    schema: str
+
+
+@dataclass
+class MissingTable(Missing):
+    bucket: str
+    schema: str
+    table: str
+
+
+@dataclass
+class MissingProjection(Missing):
+    bucket: str
+    schema: str
+    table: str
+    projection: str
+
+
+class Exists(Exception):
+    pass
+
+
+@dataclass
+class SchemaExists(Exists):
+    bucket: str
+    schema: str
+
+
+@dataclass
+class TableExists(Exists):
+    bucket: str
+    schema: str
+    table: str
+
+
+ERROR_TYPES_MAP = {
+    HttpStatus.BAD_REQUEST: BadRequest,
+    HttpStatus.FOBIDDEN: Forbidden,
+    HttpStatus.NOT_FOUND: NotFound,
+    HttpStatus.METHOD_NOT_ALLOWED: MethodNotAllowed,
+    HttpStatus.REQUEST_TIMEOUT: RequestTimeout,
+    HttpStatus.CONFLICT: Conflict,
+    HttpStatus.INTERNAL_SERVER_ERROR: InternalServerError,
+    HttpStatus.NOT_IMPLEMENTED: NotImplemented,
+    HttpStatus.SERVICE_UNAVAILABLE: ServiceUnavailable,
+}
+
+
+def from_response(res: requests.Response):
+    if res.status_code == HttpStatus.SUCCESS.value:
+        return None
+
+    log.debug("response: url='%s', code=%s, headers=%s, body='%s'", res.request.url, res.status_code, res.headers, res.text)
+    # try to parse S3 XML response for the error details:
+    code = None
+    message = None
+    if res.text:
+        try:
+            root = xml.etree.ElementTree.fromstring(res.text)
+            code = root.find('Code')
+            code = code.text if code is not None else None
+            message = root.find('Message')
+            message = message.text if message is not None else None
+        except xml.etree.ElementTree.ParseError:
+            log.debug("invalid XML: %r", res.text)
+
+    kwargs = dict(
+        code=code,
+        message=message,
+        url=res.request.url,
+        status=res.status_code,
+        headers=res.headers,
+    )
+    log.warning("RPC failed: %s", kwargs)
+    status = HttpStatus(res.status_code)
+    error_type = ERROR_TYPES_MAP.get(status, UnexpectedError)
+    raise error_type(**kwargs)