code: #  Copyright 2023 LanceDB Developers
  #
  #  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
  #      http://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 functools
  import logging
  import os
  from typing import Any, Callable, Dict, List, Optional, Union
  from urllib.parse import urljoin
  
  import attrs
  import pyarrow as pa
  import requests
  from pydantic import BaseModel
  from requests.adapters import HTTPAdapter
  from urllib3 import Retry
  
  from lancedb.common import Credential
  from lancedb.remote import VectorQuery, VectorQueryResult
  from lancedb.remote.connection_timeout import LanceDBClientHTTPAdapterFactory
  from lancedb.remote.errors import LanceDBClientError
  
  ARROW_STREAM_CONTENT_TYPE = "application/vnd.apache.arrow.stream"
  
  
  def _check_not_closed(f):
      @functools.wraps(f)
      def wrapped(self, *args, **kwargs):
          if self.closed:
              raise ValueError("Connection is closed")
          return f(self, *args, **kwargs)
  
      return wrapped
  
  
  def _read_ipc(resp: requests.Response) -> pa.Table:
      resp_body = resp.content
      with pa.ipc.open_file(pa.BufferReader(resp_body)) as reader:
          return reader.read_all()
  
  
  @attrs.define(slots=False)
  class RestfulLanceDBClient:
      db_name: str
      region: str
      api_key: Credential
      host_override: Optional[str] = attrs.field(default=None)
  
      closed: bool = attrs.field(default=False, init=False)
  
      @functools.cached_property
      def session(self) -> requests.Session:
          sess = requests.Session()
  
          retry_adapter_instance = retry_adapter(retry_adapter_options())
          sess.mount(urljoin(self.url, "/v1/table/"), retry_adapter_instance)
  
          adapter_class = LanceDBClientHTTPAdapterFactory()
          sess.mount("https://", adapter_class())
          return sess
  
      @property
      def url(self) -> str:
          return (
              self.host_override
              or f"https://{self.db_name}.{self.region}.api.lancedb.com"
          )
  
      def close(self):
          self.session.close()
          self.closed = True
  
      @functools.cached_property
      def headers(self) -> Dict[str, str]:
          headers = {
              "x-api-key": self.api_key,
          }
          if self.region == "local":  # Local test mode
              headers["Host"] = f"{self.db_name}.{self.region}.api.lancedb.com"
          if self.host_override:
              headers["x-lancedb-database"] = self.db_name
          return headers
  
      @staticmethod
      def _check_status(resp: requests.Response):
          if resp.status_code == 404:
              raise LanceDBClientError(f"Not found: {resp.text}")
          elif 400 <= resp.status_code < 500:
              raise LanceDBClientError(
                  f"Bad Request: {resp.status_code}, error: {resp.text}"
              )
          elif 500 <= resp.status_code < 600:
              raise LanceDBClientError(
                  f"Internal Server Error: {resp.status_code}, error: {resp.text}"
              )
          elif resp.status_code != 200:
              raise LanceDBClientError(
                  f"Unknown Error: {resp.status_code}, error: {resp.text}"
              )
  
      @_check_not_closed
      def get(self, uri: str, params: Union[Dict[str, Any], BaseModel] = None):
          """Send a GET request and returns the deserialized response payload."""
          if isinstance(params, BaseModel):
              params: Dict[str, Any] = params.dict(exclude_none=True)
          with self.session.get(
              urljoin(self.url, uri),
              params=params,
              headers=self.headers,
              timeout=(120.0, 300.0),
          ) as resp:
              self._check_status(resp)
              return resp.json()
  
      @_check_not_closed
      def post(
          self,
          uri: str,
          data: Optional[Union[Dict[str, Any], BaseModel, bytes]] = None,
          params: Optional[Dict[str, Any]] = None,
          content_type: Optional[str] = None,
          deserialize: Callable = lambda resp: resp.json(),
          request_id: Optional[str] = None,
      ) -> Dict[str, Any]:
          """Send a POST request and returns the deserialized response payload.
  
          Parameters
          ----------
          uri : str
              The uri to send the POST request to.
          data: Union[Dict[str, Any], BaseModel]
          request_id: Optional[str]
              Optional client side request id to be sent in the request headers.
  
          """
          if isinstance(data, BaseModel):
              data: Dict[str, Any] = data.dict(exclude_none=True)
          if isinstance(data, bytes):
              req_kwargs = {"data": data}
          else:
              req_kwargs = {"json": data}
  
          headers = self.headers.copy()
          if content_type is not None:
              headers["content-type"] = content_type
          if request_id is not None:
              headers["x-request-id"] = request_id
          with self.session.post(
              urljoin(self.url, uri),
              headers=headers,
              params=params,
              timeout=(120.0, 300.0),
              **req_kwargs,
          ) as resp:
              self._check_status(resp)
              return deserialize(resp)
  
      @_check_not_closed
      def list_tables(self, limit: int, page_token: Optional[str] = None) -> List[str]:
          """List all tables in the database."""
          if page_token is None:
              page_token = ""
          json = self.get("/v1/table/", {"limit": limit, "page_token": page_token})
          return json["tables"]
  
      @_check_not_closed
      def query(self, table_name: str, query: VectorQuery) -> VectorQueryResult:
          """Query a table."""
          tbl = self.post(f"/v1/table/{table_name}/query/", query, deserialize=_read_ipc)
          return VectorQueryResult(tbl)
  
      def mount_retry_adapter_for_table(self, table_name: str) -> None:
          """
          Adds an http adapter to session that will retry retryable requests to the table.
          """
          retry_options = retry_adapter_options(methods=["GET", "POST"])
          retry_adapter_instance = retry_adapter(retry_options)
          session = self.session
  
          session.mount(
              urljoin(self.url, f"/v1/table/{table_name}/query/"), retry_adapter_instance
          )
          session.mount(
              urljoin(self.url, f"/v1/table/{table_name}/describe/"),
              retry_adapter_instance,
          )
          session.mount(
              urljoin(self.url, f"/v1/table/{table_name}/index/list/"),
              retry_adapter_instance,
          )
  
  
  def retry_adapter_options(methods=["GET"]) -> Dict[str, Any]:
      return {
          "retries": int(os.environ.get("LANCE_CLIENT_MAX_RETRIES", "3")),
          "connect_retries": int(os.environ.get("LANCE_CLIENT_CONNECT_RETRIES", "3")),
          "read_retries": int(os.environ.get("LANCE_CLIENT_READ_RETRIES", "3")),
          "backoff_factor": float(
              os.environ.get("LANCE_CLIENT_RETRY_BACKOFF_FACTOR", "0.25")
          ),
          "backoff_jitter": float(
              os.environ.get("LANCE_CLIENT_RETRY_BACKOFF_JITTER", "0.25")
          ),
          "statuses": [
              int(i.strip())
              for i in os.environ.get(
                  "LANCE_CLIENT_RETRY_STATUSES", "429, 500, 502, 503"
              ).split(",")
          ],
          "methods": methods,
      }
  
  
  def retry_adapter(options: Dict[str, Any]) -> HTTPAdapter:
      total_retries = options["retries"]
      connect_retries = options["connect_retries"]
      read_retries = options["read_retries"]
      backoff_factor = options["backoff_factor"]
      backoff_jitter = options["backoff_jitter"]
      statuses = options["statuses"]
      methods = frozenset(options["methods"])
      logging.debug(
          f"Setting up retry adapter with {total_retries} retries,"  # noqa G003
          + f"connect retries {connect_retries}, read retries {read_retries},"
          + f"backoff factor {backoff_factor}, statuses {statuses}, "
          + f"methods {methods}"
      )
  
      return HTTPAdapter(
          max_retries=Retry(
              total=total_retries,
              connect=connect_retries,
              read=read_retries,
              backoff_factor=backoff_factor,
              backoff_jitter=backoff_jitter,
              status_forcelist=statuses,
              allowed_methods=methods,
          )
      )
  
apis: lancedb.remote.connection_timeout.LanceDBClientHTTPAdapterFactory lancedb.remote.VectorQueryResult lancedb.remote.errors.LanceDBClientError
extract_api: [((1587, 1612), 'attrs.define', 'attrs.define', ([], {'slots': '(False)'}), '(slots=False)\n', (1599, 1612), False, 'import attrs\n'), ((1207, 1225), 'functools.wraps', 'functools.wraps', (['f'], {}), '(f)\n', (1222, 1225), False, 'import functools\n'), ((1733, 1758), 'attrs.field', 'attrs.field', ([], {'default': 'None'}), '(default=None)\n', (1744, 1758), False, 'import attrs\n'), ((1779, 1817), 'attrs.field', 'attrs.field', ([], {'default': '(False)', 'init': '(False)'}), '(default=False, init=False)\n', (1790, 1817), False, 'import attrs\n'), ((7965, 8201), 'logging.debug', 'logging.debug', (["(f'Setting up retry adapter with {total_retries} retries,' +\n    f'connect retries {connect_retries}, read retries {read_retries},' +\n    f'backoff factor {backoff_factor}, statuses {statuses}, ' +\n    f'methods {methods}')"], {}), "(f'Setting up retry adapter with {total_retries} retries,' +\n    f'connect retries {connect_retries}, read retries {read_retries},' +\n    f'backoff factor {backoff_factor}, statuses {statuses}, ' +\n    f'methods {methods}')\n", (7978, 8201), False, 'import logging\n'), ((1908, 1926), 'requests.Session', 'requests.Session', ([], {}), '()\n', (1924, 1926), False, 'import requests\n'), ((2101, 2134), 'lancedb.remote.connection_timeout.LanceDBClientHTTPAdapterFactory', 'LanceDBClientHTTPAdapterFactory', ([], {}), '()\n', (2132, 2134), False, 'from lancedb.remote.connection_timeout import LanceDBClientHTTPAdapterFactory\n'), ((6057, 6079), 'lancedb.remote.VectorQueryResult', 'VectorQueryResult', (['tbl'], {}), '(tbl)\n', (6074, 6079), False, 'from lancedb.remote import VectorQuery, VectorQueryResult\n'), ((1512, 1538), 'pyarrow.BufferReader', 'pa.BufferReader', (['resp_body'], {}), '(resp_body)\n', (1527, 1538), True, 'import pyarrow as pa\n'), ((2019, 2050), 'urllib.parse.urljoin', 'urljoin', (['self.url', '"""/v1/table/"""'], {}), "(self.url, '/v1/table/')\n", (2026, 2050), False, 'from urllib.parse import urljoin\n'), ((2957, 3002), 'lancedb.remote.errors.LanceDBClientError', 'LanceDBClientError', (['f"""Not found: {resp.text}"""'], {}), "(f'Not found: {resp.text}')\n", (2975, 3002), False, 'from lancedb.remote.errors import LanceDBClientError\n'), ((6464, 6515), 'urllib.parse.urljoin', 'urljoin', (['self.url', 'f"""/v1/table/{table_name}/query/"""'], {}), "(self.url, f'/v1/table/{table_name}/query/')\n", (6471, 6515), False, 'from urllib.parse import urljoin\n'), ((6585, 6639), 'urllib.parse.urljoin', 'urljoin', (['self.url', 'f"""/v1/table/{table_name}/describe/"""'], {}), "(self.url, f'/v1/table/{table_name}/describe/')\n", (6592, 6639), False, 'from urllib.parse import urljoin\n'), ((6722, 6778), 'urllib.parse.urljoin', 'urljoin', (['self.url', 'f"""/v1/table/{table_name}/index/list/"""'], {}), "(self.url, f'/v1/table/{table_name}/index/list/')\n", (6729, 6778), False, 'from urllib.parse import urljoin\n'), ((6926, 6973), 'os.environ.get', 'os.environ.get', (['"""LANCE_CLIENT_MAX_RETRIES"""', '"""3"""'], {}), "('LANCE_CLIENT_MAX_RETRIES', '3')\n", (6940, 6973), False, 'import os\n'), ((7007, 7058), 'os.environ.get', 'os.environ.get', (['"""LANCE_CLIENT_CONNECT_RETRIES"""', '"""3"""'], {}), "('LANCE_CLIENT_CONNECT_RETRIES', '3')\n", (7021, 7058), False, 'import os\n'), ((7089, 7137), 'os.environ.get', 'os.environ.get', (['"""LANCE_CLIENT_READ_RETRIES"""', '"""3"""'], {}), "('LANCE_CLIENT_READ_RETRIES', '3')\n", (7103, 7137), False, 'import os\n'), ((7185, 7244), 'os.environ.get', 'os.environ.get', (['"""LANCE_CLIENT_RETRY_BACKOFF_FACTOR"""', '"""0.25"""'], {}), "('LANCE_CLIENT_RETRY_BACKOFF_FACTOR', '0.25')\n", (7199, 7244), False, 'import os\n'), ((7301, 7360), 'os.environ.get', 'os.environ.get', (['"""LANCE_CLIENT_RETRY_BACKOFF_JITTER"""', '"""0.25"""'], {}), "('LANCE_CLIENT_RETRY_BACKOFF_JITTER', '0.25')\n", (7315, 7360), False, 'import os\n'), ((8286, 8478), 'urllib3.Retry', 'Retry', ([], {'total': 'total_retries', 'connect': 'connect_retries', 'read': 'read_retries', 'backoff_factor': 'backoff_factor', 'backoff_jitter': 'backoff_jitter', 'status_forcelist': 'statuses', 'allowed_methods': 'methods'}), '(total=total_retries, connect=connect_retries, read=read_retries,\n    backoff_factor=backoff_factor, backoff_jitter=backoff_jitter,\n    status_forcelist=statuses, allowed_methods=methods)\n', (8291, 8478), False, 'from urllib3 import Retry\n'), ((3065, 3139), 'lancedb.remote.errors.LanceDBClientError', 'LanceDBClientError', (['f"""Bad Request: {resp.status_code}, error: {resp.text}"""'], {}), "(f'Bad Request: {resp.status_code}, error: {resp.text}')\n", (3083, 3139), False, 'from lancedb.remote.errors import LanceDBClientError\n'), ((3845, 3867), 'urllib.parse.urljoin', 'urljoin', (['self.url', 'uri'], {}), '(self.url, uri)\n', (3852, 3867), False, 'from urllib.parse import urljoin\n'), ((5259, 5281), 'urllib.parse.urljoin', 'urljoin', (['self.url', 'uri'], {}), '(self.url, uri)\n', (5266, 5281), False, 'from urllib.parse import urljoin\n'), ((3232, 3321), 'lancedb.remote.errors.LanceDBClientError', 'LanceDBClientError', (['f"""Internal Server Error: {resp.status_code}, error: {resp.text}"""'], {}), "(\n    f'Internal Server Error: {resp.status_code}, error: {resp.text}')\n", (3250, 3321), False, 'from lancedb.remote.errors import LanceDBClientError\n'), ((3403, 3479), 'lancedb.remote.errors.LanceDBClientError', 'LanceDBClientError', (['f"""Unknown Error: {resp.status_code}, error: {resp.text}"""'], {}), "(f'Unknown Error: {resp.status_code}, error: {resp.text}')\n", (3421, 3479), False, 'from lancedb.remote.errors import LanceDBClientError\n'), ((7442, 7509), 'os.environ.get', 'os.environ.get', (['"""LANCE_CLIENT_RETRY_STATUSES"""', '"""429, 500, 502, 503"""'], {}), "('LANCE_CLIENT_RETRY_STATUSES', '429, 500, 502, 503')\n", (7456, 7509), False, 'import os\n')]