worqhat 3.5.0__py3-none-any.whl → 3.8.0__py3-none-any.whl

worqhat/resources/{db.py → db/db.py}

@@ -2,36 +2,53 @@
 
 from __future__ import annotations
 
+from typing import Dict, Iterable
+from typing_extensions import Literal
+
 import httpx
 
-from ..types import (
+from .tables import (
+    TablesResource,
+    AsyncTablesResource,
+    TablesResourceWithRawResponse,
+    AsyncTablesResourceWithRawResponse,
+    TablesResourceWithStreamingResponse,
+    AsyncTablesResourceWithStreamingResponse,
+)
+from ...types import (
+    db_execute_batch_params,
     db_execute_query_params,
     db_insert_record_params,
     db_delete_records_params,
     db_update_records_params,
     db_process_nl_query_params,
 )
-from .._types import NOT_GIVEN, Body, Query, Headers, NotGiven
-from .._utils import maybe_transform, async_maybe_transform
-from .._compat import cached_property
-from .._resource import SyncAPIResource, AsyncAPIResource
-from .._response import (
+from ..._types import Body, Omit, Query, Headers, NotGiven, omit, not_given
+from ..._utils import maybe_transform, async_maybe_transform
+from ..._compat import cached_property
+from ..._resource import SyncAPIResource, AsyncAPIResource
+from ..._response import (
     to_raw_response_wrapper,
     to_streamed_response_wrapper,
     async_to_raw_response_wrapper,
     async_to_streamed_response_wrapper,
 )
-from .._base_client import make_request_options
-from ..types.db_execute_query_response import DBExecuteQueryResponse
-from ..types.db_insert_record_response import DBInsertRecordResponse
-from ..types.db_delete_records_response import DBDeleteRecordsResponse
-from ..types.db_update_records_response import DBUpdateRecordsResponse
-from ..types.db_process_nl_query_response import DBProcessNlQueryResponse
+from ..._base_client import make_request_options
+from ...types.db_execute_batch_response import DBExecuteBatchResponse
+from ...types.db_execute_query_response import DBExecuteQueryResponse
+from ...types.db_insert_record_response import DBInsertRecordResponse
+from ...types.db_delete_records_response import DBDeleteRecordsResponse
+from ...types.db_update_records_response import DBUpdateRecordsResponse
+from ...types.db_process_nl_query_response import DBProcessNlQueryResponse
 
 __all__ = ["DBResource", "AsyncDBResource"]
 
 
 class DBResource(SyncAPIResource):
+    @cached_property
+    def tables(self) -> TablesResource:
+        return TablesResource(self._client)
+
     @cached_property
     def with_raw_response(self) -> DBResourceWithRawResponse:
         """
@@ -55,13 +72,14 @@ class DBResource(SyncAPIResource):
         self,
         *,
         table: str,
-        where: object,
+        where: Dict[str, object],
+        environment: Literal["development", "staging", "production"] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> DBDeleteRecordsResponse:
         """
         Deletes records from the specified table that match the where conditions.
@@ -72,6 +90,8 @@ class DBResource(SyncAPIResource):
 
           where: Where conditions
 
+          environment: Environment to delete from (development, staging, production)
+
           extra_headers: Send extra headers
 
           extra_query: Add additional query parameters to the request
@@ -86,6 +106,7 @@ class DBResource(SyncAPIResource):
                 {
                     "table": table,
                     "where": where,
+                    "environment": environment,
                 },
                 db_delete_records_params.DBDeleteRecordsParams,
             ),
@@ -95,16 +116,67 @@ class DBResource(SyncAPIResource):
             cast_to=DBDeleteRecordsResponse,
         )
 
+    def execute_batch(
+        self,
+        *,
+        operations: Iterable[db_execute_batch_params.Operation],
+        environment: Literal["development", "staging", "production"] | Omit = omit,
+        transactional: bool | Omit = omit,
+        # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
+        # The extra values given here take precedence over values defined on the client or passed to this method.
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
+    ) -> DBExecuteBatchResponse:
+        """
+        Executes multiple database operations (queries, inserts, updates, deletes) in a
+        single transaction. If transactional is true, all operations will be rolled back
+        if any operation fails.
+
+        Args:
+          operations: Array of database operations to execute
+
+          environment: Environment to execute operations in
+
+          transactional: Whether to execute all operations in a single transaction
+
+          extra_headers: Send extra headers
+
+          extra_query: Add additional query parameters to the request
+
+          extra_body: Add additional JSON properties to the request
+
+          timeout: Override the client-level default timeout for this request, in seconds
+        """
+        return self._post(
+            "/db/batch",
+            body=maybe_transform(
+                {
+                    "operations": operations,
+                    "environment": environment,
+                    "transactional": transactional,
+                },
+                db_execute_batch_params.DBExecuteBatchParams,
+            ),
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=DBExecuteBatchResponse,
+        )
+
     def execute_query(
         self,
         *,
         query: str,
+        environment: Literal["development", "staging", "production"] | Omit = omit,
+        params: Dict[str, object] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> DBExecuteQueryResponse:
         """Executes a raw SQL query directly against ClickHouse (WorqDB).
 
@@ -115,6 +187,10 @@ class DBResource(SyncAPIResource):
         Args:
           query: SQL query to execute
 
+          environment: Environment to query (development, staging, production)
+
+          params: Optional query parameters
+
           extra_headers: Send extra headers
 
           extra_query: Add additional query parameters to the request
@@ -125,7 +201,14 @@ class DBResource(SyncAPIResource):
         """
         return self._post(
             "/db/query",
-            body=maybe_transform({"query": query}, db_execute_query_params.DBExecuteQueryParams),
+            body=maybe_transform(
+                {
+                    "query": query,
+                    "environment": environment,
+                    "params": params,
+                },
+                db_execute_query_params.DBExecuteQueryParams,
+            ),
             options=make_request_options(
                 extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
             ),
@@ -135,14 +218,15 @@ class DBResource(SyncAPIResource):
     def insert_record(
         self,
         *,
-        data: object,
+        data: Dict[str, object],
         table: str,
+        environment: Literal["development", "staging", "production"] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> DBInsertRecordResponse:
         """Inserts a new record into the specified table.
 
@@ -154,6 +238,8 @@ class DBResource(SyncAPIResource):
 
           table: Table name to insert into
 
+          environment: Environment to insert into (development, staging, production)
+
           extra_headers: Send extra headers
 
           extra_query: Add additional query parameters to the request
@@ -168,6 +254,7 @@ class DBResource(SyncAPIResource):
                 {
                     "data": data,
                     "table": table,
+                    "environment": environment,
                 },
                 db_insert_record_params.DBInsertRecordParams,
             ),
@@ -182,12 +269,14 @@ class DBResource(SyncAPIResource):
         *,
         question: str,
         table: str,
+        context: Dict[str, object] | Omit = omit,
+        environment: Literal["development", "staging", "production"] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> DBProcessNlQueryResponse:
         """
         Converts a natural language question into a SQL query and executes it.
@@ -198,6 +287,10 @@ class DBResource(SyncAPIResource):
 
           table: Table name to query
 
+          context: Optional context for the query
+
+          environment: Environment to query (development, staging, production)
+
           extra_headers: Send extra headers
 
           extra_query: Add additional query parameters to the request
@@ -212,6 +305,8 @@ class DBResource(SyncAPIResource):
                 {
                     "question": question,
                     "table": table,
+                    "context": context,
+                    "environment": environment,
                 },
                 db_process_nl_query_params.DBProcessNlQueryParams,
             ),
@@ -224,15 +319,16 @@ class DBResource(SyncAPIResource):
     def update_records(
         self,
         *,
-        data: object,
+        data: Dict[str, object],
         table: str,
-        where: object,
+        where: Dict[str, object],
+        environment: Literal["development", "staging", "production"] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> DBUpdateRecordsResponse:
         """
         Updates records in the specified table that match the where conditions.
@@ -245,6 +341,8 @@ class DBResource(SyncAPIResource):
 
           where: Where conditions
 
+          environment: Environment to update in (development, staging, production)
+
           extra_headers: Send extra headers
 
           extra_query: Add additional query parameters to the request
@@ -260,6 +358,7 @@ class DBResource(SyncAPIResource):
                     "data": data,
                     "table": table,
                     "where": where,
+                    "environment": environment,
                 },
                 db_update_records_params.DBUpdateRecordsParams,
             ),
@@ -271,6 +370,10 @@ class DBResource(SyncAPIResource):
 
 
 class AsyncDBResource(AsyncAPIResource):
+    @cached_property
+    def tables(self) -> AsyncTablesResource:
+        return AsyncTablesResource(self._client)
+
     @cached_property
     def with_raw_response(self) -> AsyncDBResourceWithRawResponse:
         """
@@ -294,13 +397,14 @@ class AsyncDBResource(AsyncAPIResource):
         self,
         *,
         table: str,
-        where: object,
+        where: Dict[str, object],
+        environment: Literal["development", "staging", "production"] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> DBDeleteRecordsResponse:
         """
         Deletes records from the specified table that match the where conditions.
@@ -311,6 +415,8 @@ class AsyncDBResource(AsyncAPIResource):
 
           where: Where conditions
 
+          environment: Environment to delete from (development, staging, production)
+
           extra_headers: Send extra headers
 
           extra_query: Add additional query parameters to the request
@@ -325,6 +431,7 @@ class AsyncDBResource(AsyncAPIResource):
                 {
                     "table": table,
                     "where": where,
+                    "environment": environment,
                 },
                 db_delete_records_params.DBDeleteRecordsParams,
             ),
@@ -334,16 +441,67 @@ class AsyncDBResource(AsyncAPIResource):
             cast_to=DBDeleteRecordsResponse,
         )
 
+    async def execute_batch(
+        self,
+        *,
+        operations: Iterable[db_execute_batch_params.Operation],
+        environment: Literal["development", "staging", "production"] | Omit = omit,
+        transactional: bool | Omit = omit,
+        # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
+        # The extra values given here take precedence over values defined on the client or passed to this method.
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
+    ) -> DBExecuteBatchResponse:
+        """
+        Executes multiple database operations (queries, inserts, updates, deletes) in a
+        single transaction. If transactional is true, all operations will be rolled back
+        if any operation fails.
+
+        Args:
+          operations: Array of database operations to execute
+
+          environment: Environment to execute operations in
+
+          transactional: Whether to execute all operations in a single transaction
+
+          extra_headers: Send extra headers
+
+          extra_query: Add additional query parameters to the request
+
+          extra_body: Add additional JSON properties to the request
+
+          timeout: Override the client-level default timeout for this request, in seconds
+        """
+        return await self._post(
+            "/db/batch",
+            body=await async_maybe_transform(
+                {
+                    "operations": operations,
+                    "environment": environment,
+                    "transactional": transactional,
+                },
+                db_execute_batch_params.DBExecuteBatchParams,
+            ),
+            options=make_request_options(
+                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+            ),
+            cast_to=DBExecuteBatchResponse,
+        )
+
     async def execute_query(
         self,
         *,
         query: str,
+        environment: Literal["development", "staging", "production"] | Omit = omit,
+        params: Dict[str, object] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> DBExecuteQueryResponse:
         """Executes a raw SQL query directly against ClickHouse (WorqDB).
 
@@ -354,6 +512,10 @@ class AsyncDBResource(AsyncAPIResource):
         Args:
           query: SQL query to execute
 
+          environment: Environment to query (development, staging, production)
+
+          params: Optional query parameters
+
           extra_headers: Send extra headers
 
           extra_query: Add additional query parameters to the request
@@ -364,7 +526,14 @@ class AsyncDBResource(AsyncAPIResource):
         """
         return await self._post(
             "/db/query",
-            body=await async_maybe_transform({"query": query}, db_execute_query_params.DBExecuteQueryParams),
+            body=await async_maybe_transform(
+                {
+                    "query": query,
+                    "environment": environment,
+                    "params": params,
+                },
+                db_execute_query_params.DBExecuteQueryParams,
+            ),
             options=make_request_options(
                 extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
             ),
@@ -374,14 +543,15 @@ class AsyncDBResource(AsyncAPIResource):
     async def insert_record(
         self,
         *,
-        data: object,
+        data: Dict[str, object],
         table: str,
+        environment: Literal["development", "staging", "production"] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> DBInsertRecordResponse:
         """Inserts a new record into the specified table.
 
@@ -393,6 +563,8 @@ class AsyncDBResource(AsyncAPIResource):
 
           table: Table name to insert into
 
+          environment: Environment to insert into (development, staging, production)
+
           extra_headers: Send extra headers
 
           extra_query: Add additional query parameters to the request
@@ -407,6 +579,7 @@ class AsyncDBResource(AsyncAPIResource):
                 {
                     "data": data,
                     "table": table,
+                    "environment": environment,
                 },
                 db_insert_record_params.DBInsertRecordParams,
             ),
@@ -421,12 +594,14 @@ class AsyncDBResource(AsyncAPIResource):
         *,
         question: str,
         table: str,
+        context: Dict[str, object] | Omit = omit,
+        environment: Literal["development", "staging", "production"] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> DBProcessNlQueryResponse:
         """
         Converts a natural language question into a SQL query and executes it.
@@ -437,6 +612,10 @@ class AsyncDBResource(AsyncAPIResource):
 
           table: Table name to query
 
+          context: Optional context for the query
+
+          environment: Environment to query (development, staging, production)
+
           extra_headers: Send extra headers
 
           extra_query: Add additional query parameters to the request
@@ -451,6 +630,8 @@ class AsyncDBResource(AsyncAPIResource):
                 {
                     "question": question,
                     "table": table,
+                    "context": context,
+                    "environment": environment,
                 },
                 db_process_nl_query_params.DBProcessNlQueryParams,
             ),
@@ -463,15 +644,16 @@ class AsyncDBResource(AsyncAPIResource):
     async def update_records(
         self,
         *,
-        data: object,
+        data: Dict[str, object],
         table: str,
-        where: object,
+        where: Dict[str, object],
+        environment: Literal["development", "staging", "production"] | Omit = omit,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
-        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        timeout: float | httpx.Timeout | None | NotGiven = not_given,
     ) -> DBUpdateRecordsResponse:
         """
         Updates records in the specified table that match the where conditions.
@@ -484,6 +666,8 @@ class AsyncDBResource(AsyncAPIResource):
 
           where: Where conditions
 
+          environment: Environment to update in (development, staging, production)
+
           extra_headers: Send extra headers
 
           extra_query: Add additional query parameters to the request
@@ -499,6 +683,7 @@ class AsyncDBResource(AsyncAPIResource):
                     "data": data,
                     "table": table,
                     "where": where,
+                    "environment": environment,
                 },
                 db_update_records_params.DBUpdateRecordsParams,
             ),
@@ -516,6 +701,9 @@ class DBResourceWithRawResponse:
         self.delete_records = to_raw_response_wrapper(
             db.delete_records,
         )
+        self.execute_batch = to_raw_response_wrapper(
+            db.execute_batch,
+        )
         self.execute_query = to_raw_response_wrapper(
             db.execute_query,
         )
@@ -529,6 +717,10 @@ class DBResourceWithRawResponse:
             db.update_records,
         )
 
+    @cached_property
+    def tables(self) -> TablesResourceWithRawResponse:
+        return TablesResourceWithRawResponse(self._db.tables)
+
 
 class AsyncDBResourceWithRawResponse:
     def __init__(self, db: AsyncDBResource) -> None:
@@ -537,6 +729,9 @@ class AsyncDBResourceWithRawResponse:
         self.delete_records = async_to_raw_response_wrapper(
             db.delete_records,
         )
+        self.execute_batch = async_to_raw_response_wrapper(
+            db.execute_batch,
+        )
         self.execute_query = async_to_raw_response_wrapper(
             db.execute_query,
         )
@@ -550,6 +745,10 @@ class AsyncDBResourceWithRawResponse:
             db.update_records,
         )
 
+    @cached_property
+    def tables(self) -> AsyncTablesResourceWithRawResponse:
+        return AsyncTablesResourceWithRawResponse(self._db.tables)
+
 
 class DBResourceWithStreamingResponse:
     def __init__(self, db: DBResource) -> None:
@@ -558,6 +757,9 @@ class DBResourceWithStreamingResponse:
         self.delete_records = to_streamed_response_wrapper(
             db.delete_records,
         )
+        self.execute_batch = to_streamed_response_wrapper(
+            db.execute_batch,
+        )
         self.execute_query = to_streamed_response_wrapper(
             db.execute_query,
         )
@@ -571,6 +773,10 @@ class DBResourceWithStreamingResponse:
             db.update_records,
         )
 
+    @cached_property
+    def tables(self) -> TablesResourceWithStreamingResponse:
+        return TablesResourceWithStreamingResponse(self._db.tables)
+
 
 class AsyncDBResourceWithStreamingResponse:
     def __init__(self, db: AsyncDBResource) -> None:
@@ -579,6 +785,9 @@ class AsyncDBResourceWithStreamingResponse:
         self.delete_records = async_to_streamed_response_wrapper(
             db.delete_records,
         )
+        self.execute_batch = async_to_streamed_response_wrapper(
+            db.execute_batch,
+        )
         self.execute_query = async_to_streamed_response_wrapper(
             db.execute_query,
         )
@@ -591,3 +800,7 @@ class AsyncDBResourceWithStreamingResponse:
         self.update_records = async_to_streamed_response_wrapper(
             db.update_records,
         )
+
+    @cached_property
+    def tables(self) -> AsyncTablesResourceWithStreamingResponse:
+        return AsyncTablesResourceWithStreamingResponse(self._db.tables)